Rework client and server replay

- Add client.replay [flows], client.replay.stop
- Add server.replay [flows], server.replay.stop
- The corresponding options for file loading are only read on startup, further
changes are ignored. In interactive contexts, replay is started with the
commands, not through option changes.
- Deprecate flow.replay, use replay.client instead

8 files changed, 71 insertions, 81 deletions

diff --git a/mitmproxy/addons/clientplayback.py b/mitmproxy/addons/clientplayback.py
index acb77bb2..322933f9 100644
--- a/mitmproxy/addons/clientplayback.py
+++ b/mitmproxy/addons/clientplayback.py
@@ -2,6 +2,7 @@ from mitmproxy import exceptions
 from mitmproxy import ctx
 from mitmproxy import io
 from mitmproxy import flow
+from mitmproxy import command
 
 import typing
 
@@ -11,32 +12,46 @@ class ClientPlayback:
         self.flows = None
         self.current_thread = None
         self.has_replayed = False
+        self.configured = False
 
     def count(self) -> int:
         if self.flows:
             return len(self.flows)
         return 0
 
-    def load(self, flows: typing.Sequence[flow.Flow]):
+    @command.command("replay.client.stop")
+    def stop_replay(self) -> None:
+        """
+            Stop client replay.
+        """
+        self.flows = []
+        ctx.master.addons.trigger("update", [])
+
+    @command.command("replay.client")
+    def start_replay(self, flows: typing.Sequence[flow.Flow]) -> None:
+        """
+            Replay requests from flows.
+        """
         self.flows = flows
+        ctx.master.addons.trigger("update", [])
 
     def configure(self, updated):
-        if "client_replay" in updated:
-            if ctx.options.client_replay:
-                ctx.log.info("Client Replay: {}".format(ctx.options.client_replay))
-                try:
-                    flows = io.read_flows_from_paths(ctx.options.client_replay)
-                except exceptions.FlowReadException as e:
-                    raise exceptions.OptionsError(str(e))
-                self.load(flows)
-            else:
-                self.flows = None
+        if not self.configured and ctx.options.client_replay:
+            self.configured = True
+            ctx.log.info("Client Replay: {}".format(ctx.options.client_replay))
+            try:
+                flows = io.read_flows_from_paths(ctx.options.client_replay)
+            except exceptions.FlowReadException as e:
+                raise exceptions.OptionsError(str(e))
+            self.start_replay(flows)
 
     def tick(self):
         if self.current_thread and not self.current_thread.is_alive():
             self.current_thread = None
         if self.flows and not self.current_thread:
-            self.current_thread = ctx.master.replay_request(self.flows.pop(0))
+            f = self.flows.pop(0)
+            self.current_thread = ctx.master.replay_request(f)
+            ctx.master.addons.trigger("update", [f])
             self.has_replayed = True
         if self.has_replayed:
             if not self.flows and not self.current_thread:
diff --git a/mitmproxy/addons/core.py b/mitmproxy/addons/core.py
index 215cbb4c..3530b108 100644
--- a/mitmproxy/addons/core.py
+++ b/mitmproxy/addons/core.py
@@ -47,23 +47,12 @@ class Core:
     @command.command("flow.mark.toggle")
     def mark_toggle(self, flows: typing.Sequence[flow.Flow]) -> None:
         """
-            Mark flows.
+            Toggle mark for flows.
         """
         for i in flows:
             i.marked = not i.marked
         ctx.master.addons.trigger("update", flows)
 
-    @command.command("flow.replay")
-    def replay(self, f: flow.Flow) -> None:
-        """
-            Replay an HTTP flow request.
-        """
-        try:
-            ctx.master.replay_request(f)  # type: ignore
-        except exceptions.ReplayException as e:
-            raise exceptions.CommandError("Replay error: %s" % e) from e
-        ctx.master.addons.trigger("update", [f])
-
     @command.command("flow.kill")
     def kill(self, flows: typing.Sequence[flow.Flow]) -> None:
         """
diff --git a/mitmproxy/addons/serverplayback.py b/mitmproxy/addons/serverplayback.py
index 2255aaf2..ab318d4e 100644
--- a/mitmproxy/addons/serverplayback.py
+++ b/mitmproxy/addons/serverplayback.py
@@ -1,11 +1,14 @@
 import hashlib
 import urllib
+import typing
 from typing import Any  # noqa
 from typing import List  # noqa
 
 from mitmproxy import ctx
+from mitmproxy import flow
 from mitmproxy import exceptions
 from mitmproxy import io
+from mitmproxy import command
 
 
 class ServerPlayback:
@@ -13,15 +16,27 @@ class ServerPlayback:
         self.flowmap = {}
         self.stop = False
         self.final_flow = None
+        self.configured = False
 
-    def load_flows(self, flows):
+    @command.command("replay.server")
+    def load_flows(self, flows: typing.Sequence[flow.Flow]) -> None:
+        """
+            Replay server responses from flows.
+        """
+        self.flowmap = {}
         for i in flows:
-            if i.response:
+            if i.response:  # type: ignore
                 l = self.flowmap.setdefault(self._hash(i), [])
                 l.append(i)
+        ctx.master.addons.trigger("update", [])
 
-    def clear(self):
+    @command.command("replay.server.stop")
+    def clear(self) -> None:
+        """
+            Stop server replay.
+        """
         self.flowmap = {}
+        ctx.master.addons.trigger("update", [])
 
     def count(self):
         return sum([len(i) for i in self.flowmap.values()])
@@ -90,14 +105,13 @@ class ServerPlayback:
                 return ret
 
     def configure(self, updated):
-        if "server_replay" in updated:
-            self.clear()
-            if ctx.options.server_replay:
-                try:
-                    flows = io.read_flows_from_paths(ctx.options.server_replay)
-                except exceptions.FlowReadException as e:
-                    raise exceptions.OptionsError(str(e))
-                self.load_flows(flows)
+        if not self.configured and ctx.options.server_replay:
+            self.configured = True
+            try:
+                flows = io.read_flows_from_paths(ctx.options.server_replay)
+            except exceptions.FlowReadException as e:
+                raise exceptions.OptionsError(str(e))
+            self.load_flows(flows)
 
     def tick(self):
         if self.stop and not self.final_flow.live:
diff --git a/mitmproxy/tools/console/flowlist.py b/mitmproxy/tools/console/flowlist.py
index 9f652bd9..bf8e2eee 100644
--- a/mitmproxy/tools/console/flowlist.py
+++ b/mitmproxy/tools/console/flowlist.py
@@ -132,14 +132,6 @@ class FlowItem(urwid.WidgetWrap):
     def selectable(self):
         return True
 
-    def server_replay_prompt(self, k):
-        a = self.master.addons.get("serverplayback")
-        if k == "a":
-            a.load([i.copy() for i in self.master.view])
-        elif k == "t":
-            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:
             if self.flow.request:
@@ -149,30 +141,7 @@ class FlowItem(urwid.WidgetWrap):
     def keypress(self, xxx_todo_changeme, key):
         (maxcol,) = xxx_todo_changeme
         key = common.shortcuts(key)
-        if key == "S":
-            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 = "Stop current server replay?",
-                    keys = (
-                        ("yes", "y"),
-                        ("no", "n"),
-                    ),
-                    callback = stop_server_playback,
-                )
-            else:
-                signals.status_prompt_onekey.send(
-                    prompt = "Server Replay",
-                    keys = (
-                        ("all flows", "a"),
-                        ("this flow", "t"),
-                    ),
-                    callback = self.server_replay_prompt,
-                )
-        elif key == "V":
+        if key == "V":
             if not self.flow.modified():
                 signals.status_message.send(message="Flow not modified.")
                 return
diff --git a/mitmproxy/tools/console/master.py b/mitmproxy/tools/console/master.py
index 6ddfcb81..de58f43f 100644
--- a/mitmproxy/tools/console/master.py
+++ b/mitmproxy/tools/console/master.py
@@ -78,7 +78,7 @@ class UnsupportedLog:
         signals.add_log(strutils.bytes_to_escaped_str(message.content), "debug")
 
 
-class ConsoleCommands:
+class ConsoleAddon:
     """
         An addon that exposes console-specific commands.
     """
@@ -131,6 +131,10 @@ class ConsoleCommands:
     def running(self):
         self.started = True
 
+    def update(self, flows):
+        if not flows:
+            signals.update_settings.send(self)
+
     def configure(self, updated):
         if self.started:
             if "console_eventlog" in updated:
@@ -157,7 +161,8 @@ def default_keymap(km):
     km.add("g", "view.go 0", context="flowlist")
     km.add("G", "view.go -1", context="flowlist")
     km.add("m", "flow.mark.toggle @focus", context="flowlist")
-    km.add("r", "flow.replay @focus", context="flowlist")
+    km.add("r", "replay.client @focus", context="flowlist")
+    km.add("S", "console.command 'replay.server '")
     km.add("v", "set console_order_reversed=toggle", context="flowlist")
     km.add("U", "flow.mark @all false", context="flowlist")
     km.add("w", "console.command 'save.file @shown '", context="flowlist")
@@ -196,7 +201,7 @@ class ConsoleMaster(master.Master):
             self.view,
             UnsupportedLog(),
             readfile.ReadFile(),
-            ConsoleCommands(self),
+            ConsoleAddon(self),
         )
 
         def sigint_handler(*args, **kwargs):
diff --git a/test/mitmproxy/addons/test_clientplayback.py b/test/mitmproxy/addons/test_clientplayback.py
index f71662f0..843d7409 100644
--- a/test/mitmproxy/addons/test_clientplayback.py
+++ b/test/mitmproxy/addons/test_clientplayback.py
@@ -26,7 +26,7 @@ class TestClientPlayback:
         with taddons.context() as tctx:
             assert cp.count() == 0
             f = tflow.tflow(resp=True)
-            cp.load([f])
+            cp.start_replay([f])
             assert cp.count() == 1
             RP = "mitmproxy.proxy.protocol.http_replay.RequestReplayThread"
             with mock.patch(RP) as rp:
@@ -44,13 +44,20 @@ class TestClientPlayback:
             cp.tick()
             assert cp.current_thread is None
 
+            cp.start_replay([f])
+            cp.stop_replay()
+            assert not cp.flows
+
     def test_configure(self, tmpdir):
         cp = clientplayback.ClientPlayback()
         with taddons.context() as tctx:
             path = str(tmpdir.join("flows"))
             tdump(path, [tflow.tflow()])
             tctx.configure(cp, client_replay=[path])
+            cp.configured = False
             tctx.configure(cp, client_replay=[])
+            cp.configured = False
             tctx.configure(cp)
+            cp.configured = False
             with pytest.raises(exceptions.OptionsError):
                 tctx.configure(cp, client_replay=["nonexistent"])
diff --git a/test/mitmproxy/addons/test_core.py b/test/mitmproxy/addons/test_core.py
index 9501fb22..25fefb5d 100644
--- a/test/mitmproxy/addons/test_core.py
+++ b/test/mitmproxy/addons/test_core.py
@@ -3,7 +3,6 @@ from mitmproxy.test import taddons
 from mitmproxy.test import tflow
 from mitmproxy import exceptions
 import pytest
-from unittest import mock
 
 
 def test_set():
@@ -43,15 +42,6 @@ def test_mark():
         assert f.marked
 
 
-def test_replay():
-    sa = core.Core()
-    with taddons.context():
-        f = tflow.tflow()
-        with mock.patch("mitmproxy.master.Master.replay_request") as rp:
-            sa.replay(f)
-            assert rp.called
-
-
 def test_kill():
     sa = core.Core()
     with taddons.context():
diff --git a/test/mitmproxy/addons/test_serverplayback.py b/test/mitmproxy/addons/test_serverplayback.py
index 29de48a0..e0c025fe 100644
--- a/test/mitmproxy/addons/test_serverplayback.py
+++ b/test/mitmproxy/addons/test_serverplayback.py
@@ -22,6 +22,7 @@ def test_config(tmpdir):
         fpath = str(tmpdir.join("flows"))
         tdump(fpath, [tflow.tflow(resp=True)])
         tctx.configure(s, server_replay=[fpath])
+        s.configured = False
         with pytest.raises(exceptions.OptionsError):
             tctx.configure(s, server_replay=[str(tmpdir)])