Merge pull request #2287 from cortesi/cmdmark

commands: marking

10 files changed, 152 insertions, 82 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 5194c4ee..3530b108 100644
--- a/mitmproxy/addons/core.py
+++ b/mitmproxy/addons/core.py
@@ -29,3 +29,38 @@ class Core:
         for f in intercepted:
             f.resume()
         ctx.master.addons.trigger("update", intercepted)
+
+    # FIXME: this will become view.mark later
+    @command.command("flow.mark")
+    def mark(self, flows: typing.Sequence[flow.Flow], val: bool) -> None:
+        """
+            Mark flows.
+        """
+        updated = []
+        for i in flows:
+            if i.marked != val:
+                i.marked = val
+                updated.append(i)
+        ctx.master.addons.trigger("update", updated)
+
+    # FIXME: this will become view.mark.toggle later
+    @command.command("flow.mark.toggle")
+    def mark_toggle(self, flows: typing.Sequence[flow.Flow]) -> None:
+        """
+            Toggle mark for flows.
+        """
+        for i in flows:
+            i.marked = not i.marked
+        ctx.master.addons.trigger("update", flows)
+
+    @command.command("flow.kill")
+    def kill(self, flows: typing.Sequence[flow.Flow]) -> None:
+        """
+            Kill running flows.
+        """
+        updated = []
+        for f in flows:
+            if f.killable:
+                f.kill()
+                updated.append(f)
+        ctx.master.addons.trigger("update", updated)
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/command.py b/mitmproxy/command.py
index 8a0a0bc8..3b58cc25 100644
--- a/mitmproxy/command.py
+++ b/mitmproxy/command.py
@@ -109,7 +109,16 @@ def parsearg(manager: CommandManager, spec: str, argtype: type) -> typing.Any:
     """
     if argtype == str:
         return spec
-    if argtype == int:
+    elif argtype == bool:
+        if spec == "true":
+            return True
+        elif spec == "false":
+            return False
+        else:
+            raise exceptions.CommandError(
+                "Booleans are 'true' or 'false', got %s" % spec
+            )
+    elif argtype == int:
         try:
             return int(spec)
         except ValueError as e:
diff --git a/mitmproxy/tools/console/flowlist.py b/mitmproxy/tools/console/flowlist.py
index 0257bab6..bf8e2eee 100644
--- a/mitmproxy/tools/console/flowlist.py
+++ b/mitmproxy/tools/console/flowlist.py
@@ -1,6 +1,5 @@
 import urwid
 
-from mitmproxy import exceptions
 from mitmproxy.tools.console import common
 from mitmproxy.tools.console import signals
 from mitmproxy.addons import view
@@ -133,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:
@@ -150,53 +141,13 @@ class FlowItem(urwid.WidgetWrap):
     def keypress(self, xxx_todo_changeme, key):
         (maxcol,) = xxx_todo_changeme
         key = common.shortcuts(key)
-        if key == "m":
-            self.flow.marked = not self.flow.marked
-            signals.flowlist_change.send(self)
-        elif key == "r":
-            try:
-                self.master.replay_request(self.flow)
-            except exceptions.ReplayException as e:
-                signals.add_log("Replay error: %s" % e, "warn")
-            signals.flowlist_change.send(self)
-        elif 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 == "U":
-            for f in self.master.view:
-                f.marked = False
-            signals.flowlist_change.send(self)
-        elif key == "V":
+        if key == "V":
             if not self.flow.modified():
                 signals.status_message.send(message="Flow not modified.")
                 return
             self.flow.revert()
             signals.flowlist_change.send(self)
             signals.status_message.send(message="Reverted.")
-        elif key == "X":
-            if self.flow.killable:
-                self.flow.kill()
-                self.master.view.update(self.flow)
         elif key == "|":
             signals.status_prompt_path.send(
                 prompt = "Send flow to script",
@@ -310,9 +261,7 @@ class FlowListBox(urwid.ListBox):
 
     def keypress(self, size, key):
         key = common.shortcuts(key)
-        if key == "Z":
-            self.master.view.clear_not_marked()
-        elif key == "L":
+        if key == "L":
             signals.status_prompt_path.send(
                 self,
                 prompt = "Load flows",
diff --git a/mitmproxy/tools/console/master.py b/mitmproxy/tools/console/master.py
index c0242a57..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:
@@ -151,14 +155,20 @@ def default_keymap(km):
     km.add("a", "flow.resume @focus", context="flowlist")
     km.add("d", "view.remove @focus", context="flowlist")
     km.add("D", "view.duplicate @focus", context="flowlist")
+    km.add("e", "set console_eventlog=toggle", context="flowlist")
+    km.add("f", "console.command 'set view_filter='", context="flowlist")
     km.add("F", "set console_focus_follow=toggle", context="flowlist")
     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", "replay.client @focus", context="flowlist")
+    km.add("S", "console.command 'replay.server '")
     km.add("v", "set console_order_reversed=toggle", context="flowlist")
-    km.add("f", "console.command 'set view_filter='", context="flowlist")
-    km.add("e", "set console_eventlog=toggle", context="flowlist")
+    km.add("U", "flow.mark @all false", context="flowlist")
     km.add("w", "console.command 'save.file @shown '", context="flowlist")
+    km.add("X", "flow.kill @focus", context="flowlist")
     km.add("z", "view.remove @all", context="flowlist")
+    km.add("Z", "view.remove @hidden", context="flowlist")
     km.add("enter", "console.view.flow @focus", context="flowlist")
 
 
@@ -191,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 39aac935..25fefb5d 100644
--- a/test/mitmproxy/addons/test_core.py
+++ b/test/mitmproxy/addons/test_core.py
@@ -26,3 +26,27 @@ def test_resume():
         f.intercept()
         sa.resume([f])
         assert not f.reply.state == "taken"
+
+
+def test_mark():
+    sa = core.Core()
+    with taddons.context():
+        f = tflow.tflow()
+        assert not f.marked
+        sa.mark([f], True)
+        assert f.marked
+
+        sa.mark_toggle([f])
+        assert not f.marked
+        sa.mark_toggle([f])
+        assert f.marked
+
+
+def test_kill():
+    sa = core.Core()
+    with taddons.context():
+        f = tflow.tflow()
+        f.intercept()
+        assert f.killable
+        sa.kill([f])
+        assert not f.killable
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)])
 
diff --git a/test/mitmproxy/test_command.py b/test/mitmproxy/test_command.py
index 306650c7..96d79dba 100644
--- a/test/mitmproxy/test_command.py
+++ b/test/mitmproxy/test_command.py
@@ -76,10 +76,16 @@ def test_parsearg():
     with taddons.context() as tctx:
         tctx.master.addons.add(DummyConsole())
         assert command.parsearg(tctx.master.commands, "foo", str) == "foo"
+
         assert command.parsearg(tctx.master.commands, "1", int) == 1
         with pytest.raises(exceptions.CommandError):
             command.parsearg(tctx.master.commands, "foo", int)
 
+        assert command.parsearg(tctx.master.commands, "true", bool) is True
+        assert command.parsearg(tctx.master.commands, "false", bool) is False
+        with pytest.raises(exceptions.CommandError):
+            command.parsearg(tctx.master.commands, "flobble", bool)
+
         assert len(command.parsearg(
             tctx.master.commands, "2", typing.Sequence[flow.Flow]
         )) == 2