Merge pull request #1532 from cortesi/playback

Playback and fix construct breakage

7 files changed, 141 insertions, 190 deletions

diff --git a/mitmproxy/builtins/__init__.py b/mitmproxy/builtins/__init__.py
index 3974d736..5f668570 100644
--- a/mitmproxy/builtins/__init__.py
+++ b/mitmproxy/builtins/__init__.py
@@ -8,6 +8,7 @@ from mitmproxy.builtins import stickycookie
 from mitmproxy.builtins import script
 from mitmproxy.builtins import replace
 from mitmproxy.builtins import setheaders
+from mitmproxy.builtins import serverplayback
 
 
 def default_addons():
@@ -20,4 +21,5 @@ def default_addons():
         filestreamer.FileStreamer(),
         replace.Replace(),
         setheaders.SetHeaders(),
+        serverplayback.ServerPlayback()
     ]
diff --git a/mitmproxy/builtins/serverplayback.py b/mitmproxy/builtins/serverplayback.py
new file mode 100644
index 00000000..fe56d68b
--- /dev/null
+++ b/mitmproxy/builtins/serverplayback.py
@@ -0,0 +1,133 @@
+from __future__ import absolute_import, print_function, division
+from six.moves import urllib
+import hashlib
+
+from netlib import strutils
+from mitmproxy import exceptions, flow, ctx
+
+
+class ServerPlayback(object):
+    def __init__(self):
+        self.options = None
+
+        self.flowmap = {}
+        self.stop = False
+        self.final_flow = None
+
+    def load(self, flows):
+        for i in flows:
+            if i.response:
+                l = self.flowmap.setdefault(self._hash(i), [])
+                l.append(i)
+
+    def clear(self):
+        self.flowmap = {}
+
+    def count(self):
+        return sum([len(i) for i in self.flowmap.values()])
+
+    def _hash(self, flow):
+        """
+            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 = [str(r.port), str(r.scheme), str(r.method), str(path)]
+        if not self.options.replay_ignore_content:
+            form_contents = r.urlencoded_form or r.multipart_form
+            if self.options.replay_ignore_payload_params and form_contents:
+                params = [
+                    strutils.always_bytes(i)
+                    for i in self.options.replay_ignore_payload_params
+                ]
+                for p in form_contents.items(multi=True):
+                    if p[0] not in params:
+                        key.append(p)
+            else:
+                key.append(str(r.raw_content))
+
+        if not self.options.replay_ignore_host:
+            key.append(r.host)
+
+        filtered = []
+        ignore_params = self.options.replay_ignore_params or []
+        for p in queriesArray:
+            if p[0] not in ignore_params:
+                filtered.append(p)
+        for p in filtered:
+            key.append(p[0])
+            key.append(p[1])
+
+        if self.options.rheaders:
+            headers = []
+            for i in self.options.rheaders:
+                v = r.headers.get(i)
+                headers.append((i, v))
+            key.append(headers)
+        return hashlib.sha256(
+            repr(key).encode("utf8", "surrogateescape")
+        ).digest()
+
+    def next_flow(self, request):
+        """
+            Returns the next flow object, or None if no matching flow was
+            found.
+        """
+        hsh = self._hash(request)
+        if hsh in self.flowmap:
+            if self.options.nopop:
+                return self.flowmap[hsh][0]
+            else:
+                ret = self.flowmap[hsh].pop(0)
+                if not self.flowmap[hsh]:
+                    del self.flowmap[hsh]
+                return ret
+
+    def configure(self, options, updated):
+        self.options = options
+        if options.server_replay and "server_replay" in updated:
+            try:
+                flows = flow.read_flows_from_paths(options.server_replay)
+            except exceptions.FlowReadException as e:
+                raise exceptions.OptionsError(str(e))
+            self.clear()
+            self.load(flows)
+
+        # FIXME: These options have to be renamed to something more sensible -
+        # prefixed with serverplayback_ where appropriate, and playback_ where
+        # they're shared with client playback.
+        #
+        # options.kill
+        # options.rheaders,
+        # options.nopop,
+        # options.replay_ignore_params,
+        # options.replay_ignore_content,
+        # options.replay_ignore_payload_params,
+        # options.replay_ignore_host
+
+    def tick(self):
+        if self.stop and not self.final_flow.live:
+            ctx.master.shutdown()
+
+    def request(self, f):
+        if self.flowmap:
+            rflow = self.next_flow(f)
+            if rflow:
+                response = rflow.response.copy()
+                response.is_replay = True
+                if self.options.refresh_server_playback:
+                    response.refresh()
+                f.response = response
+                if not self.flowmap and not self.options.keepserving:
+                    self.final_flow = f
+                    self.stop = True
+            elif self.options.kill:
+                ctx.log.warn(
+                    "server_playback: killed non-replay request {}".format(
+                        f.request.url
+                    )
+                )
+                f.reply.kill()
diff --git a/mitmproxy/dump.py b/mitmproxy/dump.py
index 51124224..49215b3a 100644
--- a/mitmproxy/dump.py
+++ b/mitmproxy/dump.py
@@ -59,18 +59,6 @@ class DumpMaster(flow.FlowMaster):
                   "HTTP/2 is disabled. Use --no-http2 to silence this warning.",
                   file=sys.stderr)
 
-        if options.server_replay:
-            self.start_server_playback(
-                self._readflow(options.server_replay),
-                options.kill, options.rheaders,
-                not options.keepserving,
-                options.nopop,
-                options.replay_ignore_params,
-                options.replay_ignore_content,
-                options.replay_ignore_payload_params,
-                options.replay_ignore_host
-            )
-
         if options.client_replay:
             self.start_client_playback(
                 self._readflow(options.client_replay),
diff --git a/mitmproxy/flow/__init__.py b/mitmproxy/flow/__init__.py
index 8a64180e..10e66f08 100644
--- a/mitmproxy/flow/__init__.py
+++ b/mitmproxy/flow/__init__.py
@@ -4,16 +4,14 @@ from mitmproxy.flow import export, modules
 from mitmproxy.flow.io import FlowWriter, FilteredFlowWriter, FlowReader, read_flows_from_paths
 from mitmproxy.flow.master import FlowMaster
 from mitmproxy.flow.modules import (
-    AppRegistry, StreamLargeBodies, ClientPlaybackState, ServerPlaybackState
+    AppRegistry, StreamLargeBodies, ClientPlaybackState
 )
 from mitmproxy.flow.state import State, FlowView
 
-# TODO: We may want to remove the imports from .modules and just expose "modules"
-
 __all__ = [
     "export", "modules",
     "FlowWriter", "FilteredFlowWriter", "FlowReader", "read_flows_from_paths",
     "FlowMaster",
     "AppRegistry", "StreamLargeBodies", "ClientPlaybackState",
-    "ServerPlaybackState", "State", "FlowView",
+    "State", "FlowView",
 ]
diff --git a/mitmproxy/flow/master.py b/mitmproxy/flow/master.py
index 9cdcc8dd..b71c2c8d 100644
--- a/mitmproxy/flow/master.py
+++ b/mitmproxy/flow/master.py
@@ -29,15 +29,8 @@ class FlowMaster(controller.Master):
         if server:
             self.add_server(server)
         self.state = state
-        self.server_playback = None  # type: Optional[modules.ServerPlaybackState]
         self.client_playback = None  # type: Optional[modules.ClientPlaybackState]
-        self.kill_nonreplay = False
-
         self.stream_large_bodies = None  # type: Optional[modules.StreamLargeBodies]
-        self.replay_ignore_params = False
-        self.replay_ignore_content = None
-        self.replay_ignore_host = False
-
         self.apps = modules.AppRegistry()
 
     def start_app(self, host, port):
@@ -62,56 +55,6 @@ class FlowMaster(controller.Master):
     def stop_client_playback(self):
         self.client_playback = None
 
-    def start_server_playback(
-            self,
-            flows,
-            kill,
-            headers,
-            exit,
-            nopop,
-            ignore_params,
-            ignore_content,
-            ignore_payload_params,
-            ignore_host):
-        """
-            flows: List of flows.
-            kill: Boolean, should we kill requests not part of the replay?
-            ignore_params: list of parameters to ignore in server replay
-            ignore_content: true if request content should be ignored in server replay
-            ignore_payload_params: list of content params to ignore in server replay
-            ignore_host: true if request host should be ignored in server replay
-        """
-        self.server_playback = modules.ServerPlaybackState(
-            headers,
-            flows,
-            exit,
-            nopop,
-            ignore_params,
-            ignore_content,
-            ignore_payload_params,
-            ignore_host)
-        self.kill_nonreplay = kill
-
-    def stop_server_playback(self):
-        self.server_playback = None
-
-    def do_server_playback(self, flow):
-        """
-            This method should be called by child classes in the request
-            handler. Returns True if playback has taken place, None if not.
-        """
-        if self.server_playback:
-            rflow = self.server_playback.next_flow(flow)
-            if not rflow:
-                return None
-            response = rflow.response.copy()
-            response.is_replay = True
-            if self.options.refresh_server_playback:
-                response.refresh()
-            flow.response = response
-            return True
-        return None
-
     def tick(self, timeout):
         if self.client_playback:
             stop = (
@@ -126,17 +69,6 @@ class FlowMaster(controller.Master):
             else:
                 self.client_playback.tick(self)
 
-        if self.server_playback:
-            stop = (
-                self.server_playback.count() == 0 and
-                self.state.active_flow_count() == 0 and
-                not self.kill_nonreplay
-            )
-            exit = self.server_playback.exit
-            if stop:
-                self.stop_server_playback()
-                if exit:
-                    self.shutdown()
         return super(FlowMaster, self).tick(timeout)
 
     def duplicate_flow(self, f):
@@ -229,13 +161,6 @@ class FlowMaster(controller.Master):
         except IOError as v:
             raise exceptions.FlowReadException(v.strerror)
 
-    def process_new_request(self, f):
-        if self.server_playback:
-            pb = self.do_server_playback(f)
-            if not pb and self.kill_nonreplay:
-                self.add_log("Killed {}".format(f.request.url), "info")
-                f.reply.kill()
-
     def replay_request(self, f, block=False):
         """
             Returns None if successful, or error message if not.
@@ -256,7 +181,8 @@ class FlowMaster(controller.Master):
 
             f.response = None
             f.error = None
-            self.process_new_request(f)
+            # FIXME: process through all addons?
+            # self.process_new_request(f)
             rt = http_replay.RequestReplayThread(
                 self.server.config,
                 f,
@@ -314,7 +240,6 @@ class FlowMaster(controller.Master):
                 return
         if f not in self.state.flows:  # don't add again on replay
             self.state.add_flow(f)
-        self.process_new_request(f)
         return f
 
     @controller.handler
diff --git a/mitmproxy/flow/modules.py b/mitmproxy/flow/modules.py
index fb3c52da..e44416c3 100644
--- a/mitmproxy/flow/modules.py
+++ b/mitmproxy/flow/modules.py
@@ -1,13 +1,8 @@
 from __future__ import absolute_import, print_function, division
 
-import hashlib
-
-from six.moves import urllib
-
 from mitmproxy import controller
 from netlib import wsgi
 from netlib import version
-from netlib import strutils
 from netlib.http import http1
 
 
@@ -84,95 +79,3 @@ class ClientPlaybackState:
                 master.request(self.current)
                 if self.current.response:
                     master.response(self.current)
-
-
-class ServerPlaybackState:
-    def __init__(
-            self,
-            headers,
-            flows,
-            exit,
-            nopop,
-            ignore_params,
-            ignore_content,
-            ignore_payload_params,
-            ignore_host):
-        """
-            headers: Case-insensitive list of request headers that should be
-            included in request-response matching.
-        """
-        self.headers = headers
-        self.exit = exit
-        self.nopop = nopop
-        self.ignore_params = ignore_params
-        self.ignore_content = ignore_content
-        self.ignore_payload_params = [strutils.always_bytes(x) for x in (ignore_payload_params or ())]
-        self.ignore_host = ignore_host
-        self.fmap = {}
-        for i in flows:
-            if i.response:
-                l = self.fmap.setdefault(self._hash(i), [])
-                l.append(i)
-
-    def count(self):
-        return sum(len(i) for i in self.fmap.values())
-
-    def _hash(self, flow):
-        """
-            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 = [
-            str(r.port),
-            str(r.scheme),
-            str(r.method),
-            str(path),
-        ]
-
-        if not self.ignore_content:
-            form_contents = r.urlencoded_form or r.multipart_form
-            if self.ignore_payload_params and form_contents:
-                key.extend(
-                    p for p in form_contents.items(multi=True)
-                    if p[0] not in self.ignore_payload_params
-                )
-            else:
-                key.append(str(r.raw_content))
-
-        if not self.ignore_host:
-            key.append(r.host)
-
-        filtered = []
-        ignore_params = self.ignore_params or []
-        for p in queriesArray:
-            if p[0] not in ignore_params:
-                filtered.append(p)
-        for p in filtered:
-            key.append(p[0])
-            key.append(p[1])
-
-        if self.headers:
-            headers = []
-            for i in self.headers:
-                v = r.headers.get(i)
-                headers.append((i, v))
-            key.append(headers)
-        return hashlib.sha256(repr(key).encode("utf8", "surrogateescape")).digest()
-
-    def next_flow(self, request):
-        """
-            Returns the next flow object, or None if no matching flow was
-            found.
-        """
-        l = self.fmap.get(self._hash(request))
-        if not l:
-            return None
-
-        if self.nopop:
-            return l[0]
-        else:
-            return l.pop(0)
diff --git a/mitmproxy/options.py b/mitmproxy/options.py
index 75798381..c4974839 100644
--- a/mitmproxy/options.py
+++ b/mitmproxy/options.py
@@ -31,6 +31,7 @@ class Options(optmanager.OptManager):
             anticomp=False,  # type: bool
             client_replay=None,  # type: Optional[str]
             kill=False,  # type: bool
+            keepserving=True,  # type: bool
             no_server=False,  # type: bool
             nopop=False,  # type: bool
             refresh_server_playback=False,  # type: bool
@@ -87,6 +88,7 @@ class Options(optmanager.OptManager):
         self.anticache = anticache
         self.anticomp = anticomp
         self.client_replay = client_replay
+        self.keepserving = keepserving
         self.kill = kill
         self.no_server = no_server
         self.nopop = nopop