From f21a970f294486e47b183472ef9f535e2a661604 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Sat, 29 Apr 2017 09:14:44 +1200 Subject: [PATCH] commands: marking Add "view.mark [flows] bool" and "view.mark.toggle [flows]". Use this to rebind marking keys in flowlist. --- mitmproxy/addons/core.py | 23 +++++++++++++++++++++++ mitmproxy/command.py | 11 ++++++++++- mitmproxy/tools/console/flowlist.py | 9 +-------- mitmproxy/tools/console/master.py | 6 ++++-- test/mitmproxy/addons/test_core.py | 14 ++++++++++++++ test/mitmproxy/test_command.py | 6 ++++++ 6 files changed, 58 insertions(+), 11 deletions(-) diff --git a/mitmproxy/addons/core.py b/mitmproxy/addons/core.py index 5194c4ee1..5728b9cc1 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 8a0a0bc86..3b58cc257 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 0257bab60..5ca93a280 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 c0242a576..50a64006d 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 39aac9359..2d9ee751c 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 306650c79..96d79dbae 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