improve curl/httpie export

1 files changed, 40 insertions, 33 deletions

diff --git a/mitmproxy/addons/export.py b/mitmproxy/addons/export.py
index 20dac655..638a1e2e 100644
--- a/mitmproxy/addons/export.py
+++ b/mitmproxy/addons/export.py
@@ -1,47 +1,60 @@
-import typing
 import shlex
+import typing
+
+import pyperclip
 
-from mitmproxy import ctx
+import mitmproxy.types
 from mitmproxy import command
-from mitmproxy import flow
+from mitmproxy import ctx, http
 from mitmproxy import exceptions
-from mitmproxy.utils import strutils
+from mitmproxy import flow
 from mitmproxy.net.http.http1 import assemble
-import mitmproxy.types
-
-import pyperclip
+from mitmproxy.utils import strutils
 
 
-def cleanup_request(f: flow.Flow):
+def cleanup_request(f: flow.Flow) -> http.HTTPRequest:
     if not hasattr(f, "request"):
         raise exceptions.CommandError("Can't export flow with no request.")
-    request = f.request.copy()  # type: ignore
+    assert isinstance(f, http.HTTPFlow)
+    request = f.request.copy()
     request.decode(strict=False)
-    # a bit of clean-up
-    if request.method == 'GET' and request.headers.get("content-length", None) == "0":
-        request.headers.pop('content-length')
-    request.headers.pop(':authority', None)
+    # a bit of clean-up - these headers should be automatically set by curl/httpie
+    request.headers.pop('content-length')
+    if request.headers.get("host", "") == request.host:
+        request.headers.pop("host")
+    if request.headers.get(":authority", "") == request.host:
+        request.headers.pop(":authority")
     return request
 
 
+def request_content_for_console(request: http.HTTPRequest) -> str:
+    try:
+        text = request.get_text(strict=True)
+    except ValueError:
+        # shlex.quote doesn't support a bytes object
+        # see https://github.com/python/cpython/pull/10871
+        raise exceptions.CommandError("Request content must be valid unicode")
+    escape_control_chars = {chr(i): f"\\x{i:02x}" for i in range(32)}
+    return "".join(
+        escape_control_chars.get(x, x)
+        for x in text
+    )
+
+
 def curl_command(f: flow.Flow) -> str:
     request = cleanup_request(f)
     args = ["curl"]
     for k, v in request.headers.items(multi=True):
-        args += ["--compressed "] if k == 'accept-encoding' else []
-        args += ["-H", "%s:%s" % (k, v)]
+        if k.lower() == "accept-encoding":
+            args.append("--compressed")
+        else:
+            args += ["-H", f"{k}: {v}"]
 
     if request.method != "GET":
         args += ["-X", request.method]
     args.append(request.url)
     if request.content:
-        try:
-            content = strutils.always_str(request.content)
-        except UnicodeDecodeError:
-            # shlex.quote doesn't support a bytes object
-            # see https://github.com/python/cpython/pull/10871
-            raise exceptions.CommandError("Request content must be valid unicode")
-        args += ["--data-binary", content]
+        args += ["-d", request_content_for_console(request)]
     return ' '.join(shlex.quote(arg) for arg in args)
 
 
@@ -49,16 +62,10 @@ def httpie_command(f: flow.Flow) -> str:
     request = cleanup_request(f)
     args = ["http", request.method, request.url]
     for k, v in request.headers.items(multi=True):
-        args.append("%s:%s" % (k, v))
+        args.append(f"{k}: {v}")
     cmd = ' '.join(shlex.quote(arg) for arg in args)
     if request.content:
-        try:
-            content = strutils.always_str(request.content)
-        except UnicodeDecodeError:
-            # shlex.quote doesn't support a bytes object
-            # see https://github.com/python/cpython/pull/10871
-            raise exceptions.CommandError("Request content must be valid unicode")
-        cmd += " <<< " + shlex.quote(content)
+        cmd += " <<< " + shlex.quote(request_content_for_console(request))
     return cmd
 
 
@@ -67,9 +74,9 @@ def raw(f: flow.Flow) -> bytes:
 
 
 formats = dict(
-    curl = curl_command,
-    httpie = httpie_command,
-    raw = raw,
+    curl=curl_command,
+    httpie=httpie_command,
+    raw=raw,
 )