From a6e8b930c9aac350cd1701e5e7fe4e7ca7e1ba3c Mon Sep 17 00:00:00 2001
From: Michael McKeirnan <github@mckeimic.com>
Date: Sat, 16 Nov 2019 01:20:50 -0800
Subject: Adding raw_request and raw_response to export

This is a proposed change for
https://github.com/mitmproxy/mitmproxy/issues/3701 which alters the
behavior of a raw http export to include both the request and the
response. Additionally, this introduces two new export options
"raw_request" and "raw_response" which allow for exporting the raw HTTP
request or response individually.
---
 mitmproxy/addons/export.py           | 26 ++++++++++++++++++----
 test/mitmproxy/addons/test_export.py | 42 +++++++++++++++++++++++++++++++++++-
 2 files changed, 63 insertions(+), 5 deletions(-)

diff --git a/mitmproxy/addons/export.py b/mitmproxy/addons/export.py
index 704f68da..528ccbf6 100644
--- a/mitmproxy/addons/export.py
+++ b/mitmproxy/addons/export.py
@@ -12,7 +12,7 @@ import pyperclip
 
 
 def cleanup_request(f: flow.Flow):
-    if not hasattr(f, "request"):
+    if not hasattr(f, "request") or not f.request:  # type: ignore
         raise exceptions.CommandError("Can't export flow with no request.")
     request = f.request.copy()  # type: ignore
     request.decode(strict=False)
@@ -24,7 +24,7 @@ def cleanup_request(f: flow.Flow):
 
 
 def cleanup_response(f: flow.Flow):
-    if not hasattr(f, "response"):
+    if not hasattr(f, "response") or not f.response:  # type: ignore
         raise exceptions.CommandError("Can't export flow with no response.")
     response = f.response.copy()  # type: ignore
     response.decode(strict=False)
@@ -62,15 +62,33 @@ def httpie_command(f: flow.Flow) -> str:
 
 
 def raw_request(f: flow.Flow) -> bytes:
-    return assemble.assemble_request(cleanup_request(f))  # type: ignore
+    return assemble.assemble_request(cleanup_request(f))
+
 
 def raw_response(f: flow.Flow) -> bytes:
-    return assemble.assemble_response(cleanup_response(f))  # type: ignore
+    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,
 )
diff --git a/test/mitmproxy/addons/test_export.py b/test/mitmproxy/addons/test_export.py
index 1ea4f256..b0fbeb45 100644
--- a/test/mitmproxy/addons/test_export.py
+++ b/test/mitmproxy/addons/test_export.py
@@ -23,6 +23,13 @@ def get_response():
         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(
@@ -85,18 +92,51 @@ class TestExportHttpieCommand:
             export.httpie_command(tcp_flow)
 
 
+class TestRaw:
+    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_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)
@@ -111,7 +151,7 @@ def test_export(tmpdir):
     f = str(tmpdir.join("path"))
     e = export.Export()
     with taddons.context():
-        assert e.formats() == ["curl", "httpie", "raw_request", "raw_response"]
+        assert e.formats() == ["curl", "httpie", "raw", "raw_request", "raw_response"]
         with pytest.raises(exceptions.CommandError):
             e.file("nonexistent", tflow.tflow(resp=True), f)
 
-- 
cgit v1.2.3