Repository: mitmproxy/mitmproxy

Issue:
"Cannot assemble flow with missing content" when exporting streamed flow in raw format
#### Problem Description
When a response (and possibly request, haven't checked that) is streamed and user tries to command `:export.clip raw @focus`, there is an exception:

```
Traceback (most recent call last):
  File "/Users/korran/Projects/mitmproxy/mitmproxy/master.py", line 54, in run_loop
    loop()
  File "/Users/korran/Projects/mitmproxy/venv/lib/python3.9/site-packages/urwid/main_loop.py", line 287, in run
    self._run()
  File "/Users/korran/Projects/mitmproxy/venv/lib/python3.9/site-packages/urwid/main_loop.py", line 385, in _run
    self.event_loop.run()
  File "/Users/korran/Projects/mitmproxy/venv/lib/python3.9/site-packages/urwid/main_loop.py", line 1494, in run
    reraise(*exc_info)
  File "/Users/korran/Projects/mitmproxy/venv/lib/python3.9/site-packages/urwid/compat.py", line 58, in reraise
    raise value
  File "/usr/local/Cellar/python@3.9/3.9.6/Frameworks/Python.framework/Versions/3.9/lib/python3.9/asyncio/events.py", line 80, in _run
    self._context.run(self._callback, *self._args)
  File "/Users/korran/Projects/mitmproxy/venv/lib/python3.9/site-packages/urwid/raw_display.py", line 416, in <lambda>
    wrapper = lambda: self.parse_input(
  File "/Users/korran/Projects/mitmproxy/venv/lib/python3.9/site-packages/urwid/raw_display.py", line 515, in parse_input
    callback(processed, processed_codes)
  File "/Users/korran/Projects/mitmproxy/venv/lib/python3.9/site-packages/urwid/main_loop.py", line 412, in _update
    self.process_input(keys)
  File "/Users/korran/Projects/mitmproxy/venv/lib/python3.9/site-packages/urwid/main_loop.py", line 513, in process_input
    k = self._topmost_widget.keypress(self.screen_size, k)
  File "/Users/korran/Projects/mitmproxy/mitmproxy/tools/console/window.py", line 316, in keypress
    k = super().keypress(size, k)
  File "/Users/korran/Projects/mitmproxy/venv/lib/python3.9/site-packages/urwid/container.py", line 1123, in keypress
    return self.footer.keypress((maxcol,),key)
  File "/Users/korran/Projects/mitmproxy/mitmproxy/tools/console/statusbar.py", line 201, in keypress
    return self.ab.keypress(*args, **kwargs)
  File "/Users/korran/Projects/mitmproxy/mitmproxy/tools/console/statusbar.py", line 149, in keypress
    self.prompt_execute(text)
  File "/Users/korran/Projects/mitmproxy/mitmproxy/tools/console/statusbar.py", line 169, in prompt_execute
    msg = p(txt)
  File "/Users/korran/Projects/mitmproxy/mitmproxy/tools/console/statusbar.py", line 115, in execute_command
    execute(txt)
  File "/Users/korran/Projects/mitmproxy/mitmproxy/tools/console/commandexecutor.py", line 18, in __call__
    ret = self.master.commands.execute(cmd)
  File "/Users/korran/Projects/mitmproxy/mitmproxy/command.py", line 273, in execute
    return self.call_strings(command_name, args)
  File "/Users/korran/Projects/mitmproxy/mitmproxy/command.py", line 259, in call_strings
    return self.commands[command_name].call(args)
  File "/Users/korran/Projects/mitmproxy/mitmproxy/command.py", line 129, in call
    ret = self.func(*bound_args.args, **bound_args.kwargs)
  File "/Users/korran/Projects/mitmproxy/mitmproxy/command.py", line 303, in wrapper
    return function(*args, **kwargs)
  File "/Users/korran/Projects/mitmproxy/mitmproxy/addons/export.py", line 189, in clip
    v = strutils.always_str(func(flow))
  File "/Users/korran/Projects/mitmproxy/mitmproxy/addons/export.py", line 121, in raw
    return b"".join([raw_request(f), separator, raw_response(f)])
  File "/Users/korran/Projects/mitmproxy/mitmproxy/addons/export.py", line 109, in raw_response
    return assemble.assemble_response(cleanup_response(f))
  File "/Users/korran/Projects/mitmproxy/mitmproxy/net/http/http1/assemble.py", line 17, in assemble_response
    raise ValueError("Cannot assemble flow with missing content")
ValueError: Cannot assemble flow with missing content
```
#### Steps to reproduce the behavior:
1. Run mitmproxy with `--set stream_large_bodies=10k`
2. `curl -x 127.0.0.1 www.google.com 1>/dev/null`
3. `:export.clip raw @focus`

#### System Information
```
Mitmproxy: 8.0.0.dev (+17, commit 13131e2)
Python:    3.9.6
OpenSSL:   OpenSSL 1.1.1i  8 Dec 2020
Platform:  macOS-11.5.1-x86_64-i386-64bit
```
This exception is explicitly raised when `flow.response.content` is `None`, but I think it's still valuable to export whatever is available in a flow. I was thinking about setting some artificial content (or `b''`) for streamed request/response, but it doesn't seem like a good idea.


Discussion:
Thanks for the clear report! 😃 

Resolving patch:
diff --git a/mitmproxy/addons/export.py b/mitmproxy/addons/export.py
--- a/mitmproxy/addons/export.py
+++ b/mitmproxy/addons/export.py
@@ -55,14 +55,14 @@ def request_content_for_console(request: http.Request) -> str:
     )
 
 
-def curl_command(f: flow.Flow, preserve_ip: bool = False) -> str:
+def curl_command(f: flow.Flow) -> str:
     request = cleanup_request(f)
     request = pop_headers(request)
     args = ["curl"]
 
     server_addr = f.server_conn.peername[0] if f.server_conn.peername else None
 
-    if preserve_ip and server_addr and request.pretty_host != server_addr:
+    if ctx.options.export_preserve_original_ip and server_addr and request.pretty_host != server_addr:
         resolve = "{}:{}:[{}]".format(request.pretty_host, request.port, server_addr)
         args.append("--resolve")
         args.append(resolve)
@@ -102,30 +102,35 @@ def httpie_command(f: flow.Flow) -> str:
 
 
 def raw_request(f: flow.Flow) -> bytes:
-    return assemble.assemble_request(cleanup_request(f))
+    request = cleanup_request(f)
+    if request.raw_content is None:
+        raise exceptions.CommandError("Request content missing.")
+    return assemble.assemble_request(request)
 
 
 def raw_response(f: flow.Flow) -> bytes:
-    return assemble.assemble_response(cleanup_response(f))
+    response = cleanup_response(f)
+    if response.raw_content is None:
+        raise exceptions.CommandError("Response content missing.")
+    return assemble.assemble_response(response)
 
 
 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.")
+    request_present = isinstance(f, http.HTTPFlow) and f.request and f.request.raw_content is not None
+    response_present = isinstance(f, http.HTTPFlow) and f.response and f.response.raw_content is not None
 
     if request_present and response_present:
         return b"".join([raw_request(f), separator, raw_response(f)])
-    elif not request_present:
+    elif request_present:
+        return raw_request(f)
+    elif response_present:
         return raw_response(f)
     else:
-        return raw_request(f)
+        raise exceptions.CommandError("Can't export flow with no request or response.")
 
 
-formats = dict(
+formats: typing.Dict[str, typing.Callable[[flow.Flow], typing.Union[str, bytes]]] = dict(
     curl=curl_command,
     httpie=httpie_command,
     raw=raw,
@@ -134,7 +139,7 @@ def raw(f: flow.Flow, separator=b"\r\n\r\n") -> bytes:
 )
 
 
-class Export():
+class Export:
     def load(self, loader):
         loader.add_option(
             "export_preserve_original_ip", bool, False,
@@ -162,10 +167,7 @@ def file(self, format: str, flow: flow.Flow, path: mitmproxy.types.Path) -> None
         if format not in formats:
             raise exceptions.CommandError("No such export format: %s" % format)
         func: typing.Any = formats[format]
-        if format == "curl":
-            v = func(flow, preserve_ip=ctx.options.export_preserve_original_ip)
-        else:
-            v = func(flow)
+        v = func(flow)
         try:
             with open(path, "wb") as fp:
                 if isinstance(v, bytes):
@@ -176,18 +178,16 @@ def file(self, format: str, flow: flow.Flow, path: mitmproxy.types.Path) -> None
             ctx.log.error(str(e))
 
     @command.command("export.clip")
-    def clip(self, format: str, flow: flow.Flow) -> None:
+    def clip(self, format: str, f: flow.Flow) -> None:
         """
             Export a flow to the system clipboard.
         """
         if format not in formats:
             raise exceptions.CommandError("No such export format: %s" % format)
-        func: typing.Any = formats[format]
-        if format == "curl":
-            v = strutils.always_str(func(flow, preserve_ip=ctx.options.export_preserve_original_ip))
-        else:
-            v = strutils.always_str(func(flow))
+        func = formats[format]
+
+        val = strutils.always_str(func(f), "utf8", "backslashreplace")
         try:
-            pyperclip.copy(v)
+            pyperclip.copy(val)
         except pyperclip.PyperclipException as e:
             ctx.log.error(str(e))


Test patch:
diff --git a/test/mitmproxy/addons/test_export.py b/test/mitmproxy/addons/test_export.py
--- a/test/mitmproxy/addons/test_export.py
+++ b/test/mitmproxy/addons/test_export.py
@@ -53,32 +53,40 @@ def tcp_flow():
     return tflow.ttcpflow()
 
 
+@pytest.fixture(scope="module")
+def export_curl():
+    e = export.Export()
+    with taddons.context() as tctx:
+        tctx.configure(e)
+        yield export.curl_command
+
+
 class TestExportCurlCommand:
-    def test_get(self, get_request):
+    def test_get(self, export_curl, get_request):
         result = """curl -H 'header: qvalue' 'http://address:22/path?a=foo&a=bar&b=baz'"""
-        assert export.curl_command(get_request) == result
+        assert export_curl(get_request) == result
 
-    def test_post(self, post_request):
+    def test_post(self, export_curl, post_request):
         post_request.request.content = b'nobinarysupport'
         result = "curl -X POST http://address:22/path -d nobinarysupport"
-        assert export.curl_command(post_request) == result
+        assert export_curl(post_request) == result
 
-    def test_fails_with_binary_data(self, post_request):
+    def test_fails_with_binary_data(self, export_curl, post_request):
         # shlex.quote doesn't support a bytes object
         # see https://github.com/python/cpython/pull/10871
         post_request.request.headers["Content-Type"] = "application/json; charset=utf-8"
         with pytest.raises(exceptions.CommandError):
-            export.curl_command(post_request)
+            export_curl(post_request)
 
-    def test_patch(self, patch_request):
+    def test_patch(self, export_curl, patch_request):
         result = """curl -H 'header: qvalue' -X PATCH 'http://address:22/path?query=param' -d content"""
-        assert export.curl_command(patch_request) == result
+        assert export_curl(patch_request) == result
 
-    def test_tcp(self, tcp_flow):
+    def test_tcp(self, export_curl, tcp_flow):
         with pytest.raises(exceptions.CommandError):
-            export.curl_command(tcp_flow)
+            export_curl(tcp_flow)
 
-    def test_escape_single_quotes_in_body(self):
+    def test_escape_single_quotes_in_body(self, export_curl):
         request = tflow.tflow(
             req=tutils.treq(
                 method=b'POST',
@@ -86,31 +94,36 @@ def test_escape_single_quotes_in_body(self):
                 content=b"'&#"
             )
         )
-        command = export.curl_command(request)
+        command = export_curl(request)
         assert shlex.split(command)[-2] == '-d'
         assert shlex.split(command)[-1] == "'&#"
 
-    def test_strip_unnecessary(self, get_request):
+    def test_strip_unnecessary(self, export_curl, get_request):
         get_request.request.headers.clear()
         get_request.request.headers["host"] = "address"
         get_request.request.headers[":authority"] = "address"
         get_request.request.headers["accept-encoding"] = "br"
         result = """curl --compressed 'http://address:22/path?a=foo&a=bar&b=baz'"""
-        assert export.curl_command(get_request) == result
+        assert export_curl(get_request) == result
 
     # This tests that we always specify the original host in the URL, which is
     # important for SNI. If option `export_preserve_original_ip` is true, we
     # ensure that we still connect to the same IP by using curl's `--resolve`
     # option.
     def test_correct_host_used(self, get_request):
-        get_request.request.headers["host"] = "domain:22"
+        e = export.Export()
+        with taddons.context() as tctx:
+            tctx.configure(e)
 
-        result = """curl -H 'header: qvalue' -H 'host: domain:22' 'http://domain:22/path?a=foo&a=bar&b=baz'"""
-        assert export.curl_command(get_request) == result
+            get_request.request.headers["host"] = "domain:22"
 
-        result = """curl --resolve 'domain:22:[192.168.0.1]' -H 'header: qvalue' -H 'host: domain:22' """ \
-                 """'http://domain:22/path?a=foo&a=bar&b=baz'"""
-        assert export.curl_command(get_request, preserve_ip=True) == result
+            result = """curl -H 'header: qvalue' -H 'host: domain:22' 'http://domain:22/path?a=foo&a=bar&b=baz'"""
+            assert export.curl_command(get_request) == result
+
+            tctx.options.export_preserve_original_ip = True
+            result = """curl --resolve 'domain:22:[192.168.0.1]' -H 'header: qvalue' -H 'host: domain:22' """ \
+                     """'http://domain:22/path?a=foo&a=bar&b=baz'"""
+            assert export.curl_command(get_request) == result
 
 
 class TestExportHttpieCommand:
@@ -174,18 +187,12 @@ def test_get_request_present(self, get_request):
         assert b"content-length: 0" in export.raw_request(get_request)
 
     def test_get_response_present(self, get_response):
-        delattr(get_response, 'request')
+        get_response.request.content = None
         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)
+        with pytest.raises(exceptions.CommandError, match="Can't export flow with no request or response"):
+            export.raw(tcp_flow)
 
 
 class TestRawRequest:
@@ -193,10 +200,10 @@ def test_get(self, get_request):
         assert b"header: qvalue" in export.raw_request(get_request)
         assert b"content-length: 0" in export.raw_request(get_request)
 
-    def test_no_request(self, get_response):
-        delattr(get_response, 'request')
+    def test_no_content(self, get_request):
+        get_request.request.content = None
         with pytest.raises(exceptions.CommandError):
-            export.raw_request(get_response)
+            export.raw_request(get_request)
 
     def test_tcp(self, tcp_flow):
         with pytest.raises(exceptions.CommandError):
@@ -207,9 +214,10 @@ 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):
+    def test_no_content(self, get_response):
+        get_response.response.content = None
         with pytest.raises(exceptions.CommandError):
-            export.raw_response(get_request)
+            export.raw_response(get_response)
 
     def test_tcp(self, tcp_flow):
         with pytest.raises(exceptions.CommandError):
@@ -221,8 +229,8 @@ def qr(f):
         return fp.read()
 
 
-def test_export(tmpdir):
-    f = str(tmpdir.join("path"))
+def test_export(tmp_path) -> None:
+    f = tmp_path / "outfile"
     e = export.Export()
     with taddons.context() as tctx:
         tctx.configure(e)
diff --git a/test/mitmproxy/test_optmanager.py b/test/mitmproxy/test_optmanager.py
--- a/test/mitmproxy/test_optmanager.py
+++ b/test/mitmproxy/test_optmanager.py
@@ -136,7 +136,7 @@ def test_toggler():
         o.toggler("one")
 
 
-class Rec():
+class Rec:
     def __init__(self):
         self.called = None