Merge pull request #3702 from mckeimic/master

Adding export raw http response

mitmproxy/addons/export.py

@@ -13,7 +13,7 @@ from mitmproxy.utils import strutils
 
 
 def cleanup_request(f: flow.Flow) -> http.HTTPRequest:
-    if not hasattr(f, "request"):
+    if not hasattr(f, "request") or not f.request:
         raise exceptions.CommandError("Can't export flow with no request.")
     assert isinstance(f, http.HTTPFlow)
     request = f.request.copy()
@@ -27,6 +27,15 @@ def cleanup_request(f: flow.Flow) -> http.HTTPRequest:
     return request
 
 
+def cleanup_response(f: flow.Flow) -> http.HTTPResponse:
+    if not hasattr(f, "response") or not f.response:
+        raise exceptions.CommandError("Can't export flow with no response.")
+    assert isinstance(f, http.HTTPFlow)
+    response = f.response.copy()  # type: ignore
+    response.decode(strict=False)
+    return response
+
+
 def request_content_for_console(request: http.HTTPRequest) -> str:
     try:
         text = request.get_text(strict=True)
@@ -70,14 +79,36 @@ def httpie_command(f: flow.Flow) -> str:
     return cmd
 
 
-def raw(f: flow.Flow) -> bytes:
-    return assemble.assemble_request(cleanup_request(f))  # type: ignore
+def raw_request(f: flow.Flow) -> bytes:
+    return assemble.assemble_request(cleanup_request(f))
+
+
+def raw_response(f: flow.Flow) -> bytes:
+    return assemble.assemble_response(cleanup_response(f))
+
+
+def raw(f: flow.Flow, separator=b"\r\n\r\n") -> bytes:
+    """Return either the request or response if only one exists, otherwise return both"""
+    request_present = hasattr(f, "request") and f.request  # type: ignore
+    response_present = hasattr(f, "response") and f.response  # type: ignore
+
+    if not (request_present or response_present):
+        raise exceptions.CommandError("Can't export flow with no request or response.")
+
+    if request_present and response_present:
+        return b"".join([raw_request(f), separator, raw_response(f)])
+    elif not request_present:
+        return raw_response(f)
+    else:
+        return raw_request(f)
 
 
 formats = dict(
     curl=curl_command,
     httpie=httpie_command,
     raw=raw,
+    raw_request=raw_request,
+    raw_response=raw_response,
 )
 
 

test/mitmproxy/addons/test_export.py

@@ -18,6 +18,19 @@ 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 get_flow():
+    return tflow.tflow(
+        req=tutils.treq(method=b'GET', content=b'', path=b"/path?a=foo&a=bar&b=baz"),
+        resp=tutils.tresp(status_code=404, content=b"Test Response Body"))
+
+
 @pytest.fixture
 def post_request():
     return tflow.tflow(
@@ -125,12 +138,53 @@ class TestExportHttpieCommand:
 
 
 class TestRaw:
-    def test_get(self, get_request):
+    def test_req_and_resp_present(self, get_flow):
+        assert b"header: qvalue" in export.raw(get_flow)
+        assert b"header-response: svalue" in export.raw(get_flow)
+
+    def test_get_request_present(self, get_request):
         assert b"header: qvalue" in export.raw(get_request)
 
+    def test_get_response_present(self, get_response):
+        delattr(get_response, 'request')
+        assert b"header-response: svalue" in export.raw(get_response)
+
+    def test_missing_both(self, get_request):
+        delattr(get_request, 'request')
+        delattr(get_request, 'response')
+        with pytest.raises(exceptions.CommandError):
+            export.raw(get_request)
+
     def test_tcp(self, tcp_flow):
         with pytest.raises(exceptions.CommandError):
-            export.raw(tcp_flow)
+            export.raw_request(tcp_flow)
+
+
+class TestRawRequest:
+    def test_get(self, get_request):
+        assert b"header: qvalue" in export.raw_request(get_request)
+
+    def test_no_request(self, get_response):
+        delattr(get_response, 'request')
+        with pytest.raises(exceptions.CommandError):
+            export.raw_request(get_response)
+
+    def test_tcp(self, tcp_flow):
+        with pytest.raises(exceptions.CommandError):
+            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_no_response(self, get_request):
+        with pytest.raises(exceptions.CommandError):
+            export.raw_response(get_request)
+
+    def test_tcp(self, tcp_flow):
+        with pytest.raises(exceptions.CommandError):
+            export.raw_response(tcp_flow)
 
 
 def qr(f):
@@ -142,11 +196,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", "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)
 
@@ -171,7 +229,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")
 
 
@@ -183,7 +241,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:
@@ -198,5 +260,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")