Merge pull request #1542 from cortesi/replayconsole

Fix server replay in mitmproxy console

6 files changed, 36 insertions, 49 deletions

diff --git a/mitmproxy/addons.py b/mitmproxy/addons.py
index 329d1215..2658c0af 100644
--- a/mitmproxy/addons.py
+++ b/mitmproxy/addons.py
@@ -4,7 +4,7 @@ import pprint
 
 
 def _get_name(itm):
-    return getattr(itm, "name", itm.__class__.__name__)
+    return getattr(itm, "name", itm.__class__.__name__.lower())
 
 
 class Addons(object):
@@ -13,6 +13,16 @@ class Addons(object):
         self.master = master
         master.options.changed.connect(self.options_update)
 
+    def get(self, name):
+        """
+            Retrieve an addon by name. Addon names are equal to the .name
+            attribute on the instance, or the lower case class name if that
+            does not exist.
+        """
+        for i in self.chain:
+            if name == _get_name(i):
+                return i
+
     def options_update(self, options, updated):
         for i in self.chain:
             with self.master.handlecontext():
@@ -39,14 +49,6 @@ class Addons(object):
         for i in self.chain:
             self.invoke_with_context(i, "done")
 
-    def has_addon(self, name):
-        """
-            Is an addon with this name registered?
-        """
-        for i in self.chain:
-            if _get_name(i) == name:
-                return True
-
     def __len__(self):
         return len(self.chain)
 
diff --git a/mitmproxy/builtins/serverplayback.py b/mitmproxy/builtins/serverplayback.py
index fe56d68b..be82cad9 100644
--- a/mitmproxy/builtins/serverplayback.py
+++ b/mitmproxy/builtins/serverplayback.py
@@ -88,13 +88,14 @@ class ServerPlayback(object):
 
     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))
+        if "server_replay" in updated:
             self.clear()
-            self.load(flows)
+            if options.server_replay:
+                try:
+                    flows = flow.read_flows_from_paths(options.server_replay)
+                except exceptions.FlowReadException as e:
+                    raise exceptions.OptionsError(str(e))
+                self.load(flows)
 
         # FIXME: These options have to be renamed to something more sensible -
         # prefixed with serverplayback_ where appropriate, and playback_ where
diff --git a/mitmproxy/console/master.py b/mitmproxy/console/master.py
index a6942ca4..1cb3a32b 100644
--- a/mitmproxy/console/master.py
+++ b/mitmproxy/console/master.py
@@ -248,9 +248,6 @@ class ConsoleMaster(flow.FlowMaster):
         if options.client_replay:
             self.client_playback_path(options.client_replay)
 
-        if options.server_replay:
-            self.server_playback_path(options.server_replay)
-
         self.view_stack = []
 
         if options.app:
@@ -391,21 +388,6 @@ class ConsoleMaster(flow.FlowMaster):
         if flows:
             self.start_client_playback(flows, False)
 
-    def server_playback_path(self, path):
-        if not isinstance(path, list):
-            path = [path]
-        flows = self._readflows(path)
-        if flows:
-            self.start_server_playback(
-                flows,
-                self.options.kill, self.options.rheaders,
-                False, self.options.nopop,
-                self.options.replay_ignore_params,
-                self.options.replay_ignore_content,
-                self.options.replay_ignore_payload_params,
-                self.options.replay_ignore_host
-            )
-
     def spawn_editor(self, data):
         text = not isinstance(data, bytes)
         fd, name = tempfile.mkstemp('', "mproxy", text=text)
diff --git a/mitmproxy/console/statusbar.py b/mitmproxy/console/statusbar.py
index 43d68d51..6c4cc8b5 100644
--- a/mitmproxy/console/statusbar.py
+++ b/mitmproxy/console/statusbar.py
@@ -147,14 +147,12 @@ class StatusBar(urwid.WidgetWrap):
         if self.master.client_playback:
             r.append("[")
             r.append(("heading_key", "cplayback"))
-            r.append(":%s to go]" % self.master.client_playback.count())
-        if self.master.server_playback:
+            r.append(":%s]" % self.master.client_playback.count())
+        if self.master.options.server_replay:
             r.append("[")
             r.append(("heading_key", "splayback"))
-            if self.master.options.nopop:
-                r.append(":%s in file]" % self.master.server_playback.count())
-            else:
-                r.append(":%s to go]" % self.master.server_playback.count())
+            a = self.master.addons.get("serverplayback")
+            r.append(":%s]" % a.count())
         if self.master.options.ignore_hosts:
             r.append("[")
             r.append(("heading_key", "I"))
diff --git a/mitmproxy/console/window.py b/mitmproxy/console/window.py
index 35593643..159f68ed 100644
--- a/mitmproxy/console/window.py
+++ b/mitmproxy/console/window.py
@@ -57,13 +57,11 @@ class Window(urwid.Frame):
                     callback = self.master.stop_client_playback_prompt,
                 )
         elif k == "s":
-            if not self.master.server_playback:
-                signals.status_prompt_path.send(
-                    self,
-                    prompt = "Server replay path",
-                    callback = self.master.server_playback_path
-                )
-            else:
+            a = self.master.addons.get("serverplayback")
+            if a.count():
+                def stop_server_playback(response):
+                    if response == "y":
+                        self.master.options.server_replay = []
                 signals.status_prompt_onekey.send(
                     self,
                     prompt = "Stop current server replay?",
@@ -71,7 +69,13 @@ class Window(urwid.Frame):
                         ("yes", "y"),
                         ("no", "n"),
                     ),
-                    callback = self.master.stop_server_playback_prompt,
+                    callback = stop_server_playback
+                )
+            else:
+                signals.status_prompt_path.send(
+                    self,
+                    prompt = "Server playback path",
+                    callback = lambda x: self.master.options.setter("server_replay")([x])
                 )
 
     def keypress(self, size, k):
diff --git a/test/mitmproxy/test_addons.py b/test/mitmproxy/test_addons.py
index a5085ea0..52d7f07f 100644
--- a/test/mitmproxy/test_addons.py
+++ b/test/mitmproxy/test_addons.py
@@ -17,5 +17,5 @@ def test_simple():
     m = controller.Master(o)
     a = addons.Addons(m)
     a.add(o, TAddon("one"))
-    assert a.has_addon("one")
-    assert not a.has_addon("two")
+    assert a.get("one")
+    assert not a.get("two")