Merge branch 'master' into master

3 files changed, 105 insertions, 64 deletions

diff --git a/mitmproxy/addons/clientplayback.py b/mitmproxy/addons/clientplayback.py
index 7bdaeb33..7adefd7a 100644
--- a/mitmproxy/addons/clientplayback.py
+++ b/mitmproxy/addons/clientplayback.py
@@ -1,23 +1,23 @@
 import queue
 import threading
-import typing
 import time
+import typing
 
-from mitmproxy import log
+import mitmproxy.types
+from mitmproxy import command
+from mitmproxy import connections
 from mitmproxy import controller
+from mitmproxy import ctx
 from mitmproxy import exceptions
-from mitmproxy import http
 from mitmproxy import flow
+from mitmproxy import http
+from mitmproxy import io
+from mitmproxy import log
 from mitmproxy import options
-from mitmproxy import connections
+from mitmproxy.coretypes import basethread
 from mitmproxy.net import server_spec, tls
 from mitmproxy.net.http import http1
-from mitmproxy.coretypes import basethread
 from mitmproxy.utils import human
-from mitmproxy import ctx
-from mitmproxy import io
-from mitmproxy import command
-import mitmproxy.types
 
 
 class RequestReplayThread(basethread.BaseThread):
@@ -117,7 +117,7 @@ class RequestReplayThread(basethread.BaseThread):
         finally:
             r.first_line_format = first_line_format_backup
             f.live = False
-            if server.connected():
+            if server and server.connected():
                 server.finish()
                 server.close()
 
diff --git a/mitmproxy/addons/export.py b/mitmproxy/addons/export.py
index 528ccbf6..d87cd787 100644
--- a/mitmproxy/addons/export.py
+++ b/mitmproxy/addons/export.py
@@ -1,25 +1,29 @@
+import shlex
 import typing
 
-from mitmproxy import ctx
+import pyperclip
+
+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):
     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
+    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
 
 
@@ -30,35 +34,47 @@ def cleanup_response(f: flow.Flow):
     response.decode(strict=False)
     return response
 
+def request_content_for_console(request: http.HTTPRequest) -> str:
+    try:
+        text = request.get_text(strict=True)
+        assert text
+    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:
-    data = "curl "
     request = cleanup_request(f)
+    args = ["curl"]
     for k, v in request.headers.items(multi=True):
-        data += "--compressed " if k == 'accept-encoding' else ""
-        data += "-H '%s:%s' " % (k, v)
+        if k.lower() == "accept-encoding":
+            args.append("--compressed")
+        else:
+            args += ["-H", f"{k}: {v}"]
+
     if request.method != "GET":
-        data += "-X %s " % request.method
-    data += "'%s'" % request.url
+        args += ["-X", request.method]
+    args.append(request.url)
     if request.content:
-        data += " --data-binary '%s'" % strutils.bytes_to_escaped_str(
-            request.content,
-            escape_single_quotes=True
-        )
-    return data
+        args += ["-d", request_content_for_console(request)]
+    return ' '.join(shlex.quote(arg) for arg in args)
 
 
 def httpie_command(f: flow.Flow) -> str:
     request = cleanup_request(f)
-    data = "http %s %s" % (request.method, request.url)
+    args = ["http", request.method, request.url]
     for k, v in request.headers.items(multi=True):
-        data += " '%s:%s'" % (k, v)
+        args.append(f"{k}: {v}")
+    cmd = ' '.join(shlex.quote(arg) for arg in args)
     if request.content:
-        data += " <<< '%s'" % strutils.bytes_to_escaped_str(
-            request.content,
-            escape_single_quotes=True
-        )
-    return data
+        cmd += " <<< " + shlex.quote(request_content_for_console(request))
+    return cmd
 
 
 def raw_request(f: flow.Flow) -> bytes:
@@ -86,11 +102,11 @@ def raw(f: flow.Flow, separator=b"\r\n\r\n") -> bytes:
 
 
 formats = dict(
-    curl = curl_command,
-    httpie = httpie_command,
-    raw = raw,
-    raw_request = raw_request,
-    raw_response = raw_response,
+    curl=curl_command,
+    httpie=httpie_command,
+    raw=raw,
+    raw_request=raw_request,
+    raw_response=raw_response,
 )
 
 
diff --git a/mitmproxy/addons/serverplayback.py b/mitmproxy/addons/serverplayback.py
index 51ba60b4..7f642585 100644
--- a/mitmproxy/addons/serverplayback.py
+++ b/mitmproxy/addons/serverplayback.py
@@ -1,16 +1,19 @@
 import hashlib
-import urllib
 import typing
+import urllib
 
-from mitmproxy import ctx
-from mitmproxy import flow
+import mitmproxy.types
+from mitmproxy import command
+from mitmproxy import ctx, http
 from mitmproxy import exceptions
+from mitmproxy import flow
 from mitmproxy import io
-from mitmproxy import command
-import mitmproxy.types
 
 
 class ServerPlayback:
+    flowmap: typing.Dict[typing.Hashable, typing.List[http.HTTPFlow]]
+    configured: bool
+
     def __init__(self):
         self.flowmap = {}
         self.configured = False
@@ -68,6 +71,13 @@ class ServerPlayback:
             to replay.
             """
         )
+        loader.add_option(
+            "server_replay_ignore_port", bool, False,
+            """
+            Ignore request's destination port while searching for a saved flow
+            to replay.
+            """
+        )
 
     @command.command("replay.server")
     def load_flows(self, flows: typing.Sequence[flow.Flow]) -> None:
@@ -75,10 +85,10 @@ class ServerPlayback:
             Replay server responses from flows.
         """
         self.flowmap = {}
-        for i in flows:
-            if i.response:  # type: ignore
-                l = self.flowmap.setdefault(self._hash(i), [])
-                l.append(i)
+        for f in flows:
+            if isinstance(f, http.HTTPFlow):
+                lst = self.flowmap.setdefault(self._hash(f), [])
+                lst.append(f)
         ctx.master.addons.trigger("update", [])
 
     @command.command("replay.server.file")
@@ -101,16 +111,15 @@ class ServerPlayback:
     def count(self) -> int:
         return sum([len(i) for i in self.flowmap.values()])
 
-    def _hash(self, flow):
+    def _hash(self, flow: http.HTTPFlow) -> typing.Hashable:
         """
             Calculates a loose hash of the flow request.
         """
         r = flow.request
-
         _, _, path, _, query, _ = urllib.parse.urlparse(r.url)
         queriesArray = urllib.parse.parse_qsl(query, keep_blank_values=True)
 
-        key: typing.List[typing.Any] = [str(r.port), str(r.scheme), str(r.method), str(path)]
+        key: typing.List[typing.Any] = [str(r.scheme), str(r.method), str(path)]
         if not ctx.options.server_replay_ignore_content:
             if ctx.options.server_replay_ignore_payload_params and r.multipart_form:
                 key.extend(
@@ -128,7 +137,9 @@ class ServerPlayback:
                 key.append(str(r.raw_content))
 
         if not ctx.options.server_replay_ignore_host:
-            key.append(r.host)
+            key.append(r.pretty_host)
+        if not ctx.options.server_replay_ignore_port:
+            key.append(r.port)
 
         filtered = []
         ignore_params = ctx.options.server_replay_ignore_params or []
@@ -149,20 +160,32 @@ class ServerPlayback:
             repr(key).encode("utf8", "surrogateescape")
         ).digest()
 
-    def next_flow(self, request):
+    def next_flow(self, flow: http.HTTPFlow) -> typing.Optional[http.HTTPFlow]:
         """
             Returns the next flow object, or None if no matching flow was
             found.
         """
-        hsh = self._hash(request)
-        if hsh in self.flowmap:
+        hash = self._hash(flow)
+        if hash in self.flowmap:
             if ctx.options.server_replay_nopop:
-                return self.flowmap[hsh][0]
+                return next((
+                    flow
+                    for flow in self.flowmap[hash]
+                    if flow.response
+                ), None)
             else:
-                ret = self.flowmap[hsh].pop(0)
-                if not self.flowmap[hsh]:
-                    del self.flowmap[hsh]
+                ret = self.flowmap[hash].pop(0)
+                while not ret.response:
+                    if self.flowmap[hash]:
+                        ret = self.flowmap[hash].pop(0)
+                    else:
+                        del self.flowmap[hash]
+                        return None
+                if not self.flowmap[hash]:
+                    del self.flowmap[hash]
                 return ret
+        else:
+            return None
 
     def configure(self, updated):
         if not self.configured and ctx.options.server_replay:
@@ -173,10 +196,11 @@ class ServerPlayback:
                 raise exceptions.OptionsError(str(e))
             self.load_flows(flows)
 
-    def request(self, f):
+    def request(self, f: http.HTTPFlow) -> None:
         if self.flowmap:
             rflow = self.next_flow(f)
             if rflow:
+                assert rflow.response
                 response = rflow.response.copy()
                 response.is_replay = True
                 if ctx.options.server_replay_refresh:
@@ -188,4 +212,5 @@ class ServerPlayback:
                         f.request.url
                     )
                 )
+                assert f.reply
                 f.reply.kill()