diff --git a/mitmproxy/addons/core.py b/mitmproxy/addons/core.py index 5728b9cc1..215cbb4c5 100644 --- a/mitmproxy/addons/core.py +++ b/mitmproxy/addons/core.py @@ -52,3 +52,26 @@ class Core: 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: + """ + 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/tools/console/flowlist.py b/mitmproxy/tools/console/flowlist.py index 5ca93a280..9f652bd9e 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 @@ -150,13 +149,7 @@ class FlowItem(urwid.WidgetWrap): def keypress(self, xxx_todo_changeme, key): (maxcol,) = xxx_todo_changeme key = common.shortcuts(key) - if 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": + if key == "S": def stop_server_playback(response): if response == "y": self.master.options.server_replay = [] @@ -186,10 +179,6 @@ class FlowItem(urwid.WidgetWrap): 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", @@ -303,9 +292,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 50a64006d..6ddfcb812 100644 --- a/mitmproxy/tools/console/master.py +++ b/mitmproxy/tools/console/master.py @@ -157,10 +157,13 @@ 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("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") + 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") diff --git a/test/mitmproxy/addons/test_core.py b/test/mitmproxy/addons/test_core.py index 2d9ee751c..9501fb22e 100644 --- a/test/mitmproxy/addons/test_core.py +++ b/test/mitmproxy/addons/test_core.py @@ -3,6 +3,7 @@ from mitmproxy.test import taddons from mitmproxy.test import tflow from mitmproxy import exceptions import pytest +from unittest import mock def test_set(): @@ -40,3 +41,22 @@ def test_mark(): assert not f.marked sa.mark_toggle([f]) 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(): + f = tflow.tflow() + f.intercept() + assert f.killable + sa.kill([f]) + assert not f.killable