diff options
| -rw-r--r-- | mitmproxy/addons/export.py | 16 | ||||
| -rw-r--r-- | 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 2776118a..704f68da 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 c86e0c7d..1ea4f256 100644 --- a/test/mitmproxy/addons/test_export.py +++ b/test/mitmproxy/addons/test_export.py @@ -18,6 +18,12 @@ def get_request(): @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( req=tutils.treq(method=b'POST', headers=(), content=bytes(range(256)))) @@ -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_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(tcp_flow) + 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") |