Merge pull request #1551 from cortesi/sreplay

console: fix/improve server replay

3 files changed, 24 insertions, 42 deletions

diff --git a/mitmproxy/console/flowlist.py b/mitmproxy/console/flowlist.py
index 8551f13c..4f70af07 100644
--- a/mitmproxy/console/flowlist.py
+++ b/mitmproxy/console/flowlist.py
@@ -18,14 +18,15 @@ def _mkhelp():
         ("d", "delete flow"),
         ("D", "duplicate flow"),
         ("e", "toggle eventlog"),
+        ("E", "export flow to file"),
         ("f", "filter view"),
         ("F", "toggle follow flow list"),
         ("L", "load saved flows"),
         ("m", "toggle flow mark"),
         ("M", "toggle marked flow view"),
         ("n", "create a new request"),
-        ("E", "export flow to file"),
         ("r", "replay request"),
+        ("S", "server replay request/s"),
         ("U", "unmark all marked flows"),
         ("V", "revert changes to request"),
         ("w", "save flows "),
@@ -140,36 +141,13 @@ class ConnectionItem(urwid.WidgetWrap):
                 args = (self.flow,)
             )
 
-    def stop_server_playback_prompt(self, a):
-        if a != "n":
-            self.master.stop_server_playback()
-
     def server_replay_prompt(self, k):
+        a = self.master.addons.get("serverplayback")
         if k == "a":
-            self.master.start_server_playback(
-                [i.copy() for i in self.master.state.view],
-                self.master.options.replay_kill_extra, self.master.options.server_replay_use_headers,
-                False, self.master.options.server_replay_nopop,
-                self.master.options.server_replay_ignore_params,
-                self.master.options.server_replay_ignore_content,
-                self.master.options.server_replay_ignore_payload_params,
-                self.master.options.server_replay_ignore_host
-            )
+            a.load([i.copy() for i in self.master.state.view])
         elif k == "t":
-            self.master.start_server_playback(
-                [self.flow.copy()],
-                self.master.options.replay_kill_extra, self.master.options.server_replay_use_headers,
-                False, self.master.options.server_replay_nopop,
-                self.master.options.server_replay_ignore_params,
-                self.master.options.server_replay_ignore_content,
-                self.master.options.server_replay_ignore_payload_params,
-                self.master.options.server_replay_ignore_host
-            )
-        else:
-            signals.status_prompt_path.send(
-                prompt = "Server replay path",
-                callback = self.master.server_playback_path
-            )
+            a.load([self.flow.copy()])
+        signals.update_settings.send(self)
 
     def mouse_event(self, size, event, button, col, row, focus):
         if event == "mouse press" and button == 1:
@@ -207,24 +185,27 @@ class ConnectionItem(urwid.WidgetWrap):
                 signals.status_message.send(message=r)
             signals.flowlist_change.send(self)
         elif key == "S":
-            if not self.master.server_playback:
+            def stop_server_playback(response):
+                if response == "y":
+                    self.master.options.server_replay = []
+            a = self.master.addons.get("serverplayback")
+            if a.count():
                 signals.status_prompt_onekey.send(
-                    prompt = "Server Replay",
+                    prompt = "Stop current server replay?",
                     keys = (
-                        ("all flows", "a"),
-                        ("this flow", "t"),
-                        ("file", "f"),
+                        ("yes", "y"),
+                        ("no", "n"),
                     ),
-                    callback = self.server_replay_prompt,
+                    callback = stop_server_playback,
                 )
             else:
                 signals.status_prompt_onekey.send(
-                    prompt = "Stop current server replay?",
+                    prompt = "Server Replay",
                     keys = (
-                        ("yes", "y"),
-                        ("no", "n"),
+                        ("all flows", "a"),
+                        ("this flow", "t"),
                     ),
-                    callback = self.stop_server_playback_prompt,
+                    callback = self.server_replay_prompt,
                 )
         elif key == "U":
             for f in self.state.flows:
diff --git a/mitmproxy/console/help.py b/mitmproxy/console/help.py
index 8024dc31..e3e2f54c 100644
--- a/mitmproxy/console/help.py
+++ b/mitmproxy/console/help.py
@@ -53,7 +53,7 @@ class HelpView(urwid.ListBox):
             ("o", "options"),
             ("q", "quit / return to previous page"),
             ("Q", "quit without confirm prompt"),
-            ("R", "replay of HTTP requests/responses"),
+            ("R", "replay of requests/responses from file"),
         ]
         text.extend(
             common.format_keyvals(keys, key="key", val="text", indent=4)
diff --git a/mitmproxy/console/statusbar.py b/mitmproxy/console/statusbar.py
index 28f29f03..8c667d7a 100644
--- a/mitmproxy/console/statusbar.py
+++ b/mitmproxy/console/statusbar.py
@@ -136,6 +136,8 @@ class StatusBar(urwid.WidgetWrap):
     def get_status(self):
         r = []
 
+        sreplay = self.master.addons.get("serverplayback")
+
         if len(self.master.options.setheaders):
             r.append("[")
             r.append(("heading_key", "H"))
@@ -148,11 +150,10 @@ class StatusBar(urwid.WidgetWrap):
             r.append("[")
             r.append(("heading_key", "cplayback"))
             r.append(":%s]" % self.master.client_playback.count())
-        if self.master.options.server_replay:
+        if sreplay.count():
             r.append("[")
             r.append(("heading_key", "splayback"))
-            a = self.master.addons.get("serverplayback")
-            r.append(":%s]" % a.count())
+            r.append(":%s]" % sreplay.count())
         if self.master.options.ignore_hosts:
             r.append("[")
             r.append(("heading_key", "I"))