From dae01ad62329c74fa9373d2052c6bf7f879fa523 Mon Sep 17 00:00:00 2001
From: Michael McKeirnan <github@mckeimic.com>
Date: Thu, 14 Nov 2019 23:52:45 -0800
Subject: [PATCH] Adding export raw http response

Adding a new export type for raw http response, and changing export raw to export
raw_request to distinguish between the two. This is a proposed change for https://github.com/mitmproxy/mitmproxy/issues/3701
---
 mitmproxy/addons/export.py           | 16 ++++++++++--
 test/mitmproxy/addons/test_export.py | 38 ++++++++++++++++++++++------
 2 files changed, 44 insertions(+), 10 deletions(-)

diff --git a/mitmproxy/addons/export.py b/mitmproxy/addons/export.py
index 2776118ab..704f68da0 100644
--- a/mitmproxy/addons/export.py
+++ b/mitmproxy/addons/export.py
@@ -23,6 +23,14 @@ def cleanup_request(f: flow.Flow):
     return request
 
 
+def cleanup_response(f: flow.Flow):
+    if not hasattr(f, "response"):
+        raise exceptions.CommandError("Can't export flow with no response.")
+    response = f.response.copy()  # type: ignore
+    response.decode(strict=False)
+    return response
+
+
 def curl_command(f: flow.Flow) -> str:
     data = "curl "
     request = cleanup_request(f)
@@ -53,14 +61,18 @@ def httpie_command(f: flow.Flow) -> str:
     return data
 
 
-def raw(f: flow.Flow) -> bytes:
+def raw_request(f: flow.Flow) -> bytes:
     return assemble.assemble_request(cleanup_request(f))  # type: ignore
 
+def raw_response(f: flow.Flow) -> bytes:
+    return assemble.assemble_response(cleanup_response(f))  # type: ignore
+
 
 formats = dict(
     curl = curl_command,
     httpie = httpie_command,
-    raw = raw,
+    raw_request = raw_request,
+    raw_response = raw_response,
 )
 
 
diff --git a/test/mitmproxy/addons/test_export.py b/test/mitmproxy/addons/test_export.py
index c86e0c7da..1ea4f2563 100644
--- a/test/mitmproxy/addons/test_export.py
+++ b/test/mitmproxy/addons/test_export.py
@@ -17,6 +17,12 @@ def get_request():
         req=tutils.treq(method=b'GET', content=b'', path=b"/path?a=foo&a=bar&b=baz"))
 
 
+@pytest.fixture
+def get_response():
+    return tflow.tflow(
+        resp=tutils.tresp(status_code=404, content=b"Test Response Body"))
+
+
 @pytest.fixture
 def post_request():
     return tflow.tflow(
@@ -79,13 +85,21 @@ class TestExportHttpieCommand:
             export.httpie_command(tcp_flow)
 
 
-class TestRaw:
+class TestRawRequest:
     def test_get(self, get_request):
-        assert b"header: qvalue" in export.raw(get_request)
+        assert b"header: qvalue" in export.raw_request(get_request)
 
     def test_tcp(self, tcp_flow):
         with pytest.raises(exceptions.CommandError):
-            export.raw(tcp_flow)
+            export.raw_request(tcp_flow)
+
+class TestRawResponse:
+    def test_get(self, get_response):
+        assert b"header-response: svalue" in export.raw_response(get_response)
+
+    def test_tcp(self, tcp_flow):
+        with pytest.raises(exceptions.CommandError):
+            export.raw_response(tcp_flow)
 
 
 def qr(f):
@@ -97,11 +111,15 @@ def test_export(tmpdir):
     f = str(tmpdir.join("path"))
     e = export.Export()
     with taddons.context():
-        assert e.formats() == ["curl", "httpie", "raw"]
+        assert e.formats() == ["curl", "httpie", "raw_request", "raw_response"]
         with pytest.raises(exceptions.CommandError):
             e.file("nonexistent", tflow.tflow(resp=True), f)
 
-        e.file("raw", tflow.tflow(resp=True), f)
+        e.file("raw_request", tflow.tflow(resp=True), f)
+        assert qr(f)
+        os.unlink(f)
+
+        e.file("raw_response", tflow.tflow(resp=True), f)
         assert qr(f)
         os.unlink(f)
 
@@ -126,7 +144,7 @@ async def test_export_open(exception, log_message, tmpdir):
     with taddons.context() as tctx:
         with mock.patch("mitmproxy.addons.export.open") as m:
             m.side_effect = exception(log_message)
-            e.file("raw", tflow.tflow(resp=True), f)
+            e.file("raw_request", tflow.tflow(resp=True), f)
             assert await tctx.master.await_log(log_message, level="error")
 
 
@@ -138,7 +156,11 @@ async def test_clip(tmpdir):
             e.clip("nonexistent", tflow.tflow(resp=True))
 
         with mock.patch('pyperclip.copy') as pc:
-            e.clip("raw", tflow.tflow(resp=True))
+            e.clip("raw_request", tflow.tflow(resp=True))
+            assert pc.called
+
+        with mock.patch('pyperclip.copy') as pc:
+            e.clip("raw_response", tflow.tflow(resp=True))
             assert pc.called
 
         with mock.patch('pyperclip.copy') as pc:
@@ -153,5 +175,5 @@ async def test_clip(tmpdir):
             log_message = "Pyperclip could not find a " \
                           "copy/paste mechanism for your system."
             pc.side_effect = pyperclip.PyperclipException(log_message)
-            e.clip("raw", tflow.tflow(resp=True))
+            e.clip("raw_request", tflow.tflow(resp=True))
             assert await tctx.master.await_log(log_message, level="error")