From a92017a6c1a0d7cf20f53e1cb4eb24716aaf55e6 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Sat, 29 Apr 2017 11:02:36 +1200 Subject: 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 --- mitmproxy/addons/clientplayback.py | 39 +++++++++++++++++++--------- mitmproxy/addons/core.py | 13 +--------- mitmproxy/addons/serverplayback.py | 36 +++++++++++++++++-------- mitmproxy/tools/console/flowlist.py | 33 +---------------------- mitmproxy/tools/console/master.py | 11 +++++--- test/mitmproxy/addons/test_clientplayback.py | 9 ++++++- test/mitmproxy/addons/test_core.py | 10 ------- test/mitmproxy/addons/test_serverplayback.py | 1 + 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)]) -- cgit v1.2.3