commands: marking

Add "view.mark [flows] bool" and "view.mark.toggle [flows]". Use this to rebind
marking keys in flowlist.

6 files changed, 58 insertions, 11 deletions

diff --git a/mitmproxy/addons/core.py b/mitmproxy/addons/core.py
index 5194c4ee..5728b9cc 100644
--- a/mitmproxy/addons/core.py
+++ b/mitmproxy/addons/core.py
@@ -29,3 +29,26 @@ 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:
+        """
+            Mark flows.
+        """
+        for i in flows:
+            i.marked = not i.marked
+        ctx.master.addons.trigger("update", flows)
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..5ca93a28 100644
--- a/mitmproxy/tools/console/flowlist.py
+++ b/mitmproxy/tools/console/flowlist.py
@@ -150,10 +150,7 @@ 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":
+        if key == "r":
             try:
                 self.master.replay_request(self.flow)
             except exceptions.ReplayException as e:
@@ -182,10 +179,6 @@ class FlowItem(urwid.WidgetWrap):
                     ),
                     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 not self.flow.modified():
                 signals.status_message.send(message="Flow not modified.")
diff --git a/mitmproxy/tools/console/master.py b/mitmproxy/tools/console/master.py
index c0242a57..50a64006 100644
--- a/mitmproxy/tools/console/master.py
+++ b/mitmproxy/tools/console/master.py
@@ -151,12 +151,14 @@ 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("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("z", "view.remove @all", context="flowlist")
     km.add("enter", "console.view.flow @focus", context="flowlist")
diff --git a/test/mitmproxy/addons/test_core.py b/test/mitmproxy/addons/test_core.py
index 39aac935..2d9ee751 100644
--- a/test/mitmproxy/addons/test_core.py
+++ b/test/mitmproxy/addons/test_core.py
@@ -26,3 +26,17 @@ 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
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