From 204726349df6f24c6a542902974e1cca7e20ce3a Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Tue, 1 May 2018 11:31:42 +1200 Subject: [PATCH 1/6] client replay: replaying flows in-flight should be added to count() --- mitmproxy/addons/clientplayback.py | 7 ++++++- mitmproxy/tools/console/commandexecutor.py | 2 +- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/mitmproxy/addons/clientplayback.py b/mitmproxy/addons/clientplayback.py index 11d2453ba..e9fbadcef 100644 --- a/mitmproxy/addons/clientplayback.py +++ b/mitmproxy/addons/clientplayback.py @@ -1,4 +1,5 @@ import queue +import threading import typing from mitmproxy import log @@ -30,12 +31,15 @@ class RequestReplayThread(basethread.BaseThread): self.options = opts self.channel = channel self.queue = queue + self.inflight = threading.Event() super().__init__("RequestReplayThread") def run(self): while True: f = self.queue.get() + self.inflight.set() self.replay(f) + self.inflight.clear() def replay(self, f): # pragma: no cover f.live = True @@ -163,7 +167,8 @@ class ClientPlayback: """ Approximate number of flows queued for replay. """ - return self.q.qsize() + inflight = 1 if self.thread and self.thread.inflight.is_set() else 0 + return self.q.qsize() + inflight @command.command("replay.client.stop") def stop_replay(self) -> None: diff --git a/mitmproxy/tools/console/commandexecutor.py b/mitmproxy/tools/console/commandexecutor.py index 26f92238d..3db03d3e6 100644 --- a/mitmproxy/tools/console/commandexecutor.py +++ b/mitmproxy/tools/console/commandexecutor.py @@ -14,7 +14,7 @@ class CommandExecutor: def __call__(self, cmd): if cmd.strip(): try: - ret = self.master.commands.call(cmd) + ret = self.master.commands.execute(cmd) except exceptions.CommandError as v: signals.status_message.send(message=str(v)) else: From 42b04a54131c9abeecf3c8a3a7207645e01a9fc1 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Wed, 2 May 2018 09:02:27 +1200 Subject: [PATCH 2/6] server replay: expose the replay.server.count command --- mitmproxy/addons/serverplayback.py | 1 + 1 file changed, 1 insertion(+) diff --git a/mitmproxy/addons/serverplayback.py b/mitmproxy/addons/serverplayback.py index 35f41c650..ce0705ef6 100644 --- a/mitmproxy/addons/serverplayback.py +++ b/mitmproxy/addons/serverplayback.py @@ -99,6 +99,7 @@ class ServerPlayback: self.flowmap = {} ctx.master.addons.trigger("update", []) + @command.command("replay.server.count") def count(self): return sum([len(i) for i in self.flowmap.values()]) From e96340843401d677d5e8272d562b48fe4358362a Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Wed, 2 May 2018 09:32:27 +1200 Subject: [PATCH 3/6] console: use replay count commands in statusbar Also add a periodic refresh every 0.5 seconds for the statusbar. This is in addition to refreshes upon event update notifications, and picks up replay status changes not linked to flow events. --- mitmproxy/addons/serverplayback.py | 2 +- mitmproxy/tools/console/statusbar.py | 11 ++++++++--- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/mitmproxy/addons/serverplayback.py b/mitmproxy/addons/serverplayback.py index ce0705ef6..efc7d3598 100644 --- a/mitmproxy/addons/serverplayback.py +++ b/mitmproxy/addons/serverplayback.py @@ -100,7 +100,7 @@ class ServerPlayback: ctx.master.addons.trigger("update", []) @command.command("replay.server.count") - def count(self): + def count(self) -> int: return sum([len(i) for i in self.flowmap.values()]) def _hash(self, flow): diff --git a/mitmproxy/tools/console/statusbar.py b/mitmproxy/tools/console/statusbar.py index 1e1c0b927..ccf5e2e00 100644 --- a/mitmproxy/tools/console/statusbar.py +++ b/mitmproxy/tools/console/statusbar.py @@ -158,6 +158,7 @@ class ActionBar(urwid.WidgetWrap): class StatusBar(urwid.WidgetWrap): + REFRESHTIME = 0.5 # Timed refresh time in seconds keyctx = "" def __init__( @@ -173,7 +174,11 @@ class StatusBar(urwid.WidgetWrap): master.options.changed.connect(self.sig_update) master.view.focus.sig_change.connect(self.sig_update) master.view.sig_view_add.connect(self.sig_update) + self.refresh() + + def refresh(self): self.redraw() + signals.call_in.send(seconds=self.REFRESHTIME, callback=self.refresh) def sig_update(self, sender, flow=None, updated=None): self.redraw() @@ -184,7 +189,7 @@ class StatusBar(urwid.WidgetWrap): def get_status(self): r = [] - sreplay = self.master.addons.get("serverplayback") + sreplay = self.master.commands.call("replay.server.count") creplay = self.master.commands.call("replay.client.count") if len(self.master.options.setheaders): @@ -197,10 +202,10 @@ class StatusBar(urwid.WidgetWrap): r.append("[") r.append(("heading_key", "cplayback")) r.append(":%s]" % creplay) - if sreplay.count(): + if sreplay: r.append("[") r.append(("heading_key", "splayback")) - r.append(":%s]" % sreplay.count()) + r.append(":%s]" % sreplay) if self.master.options.ignore_hosts: r.append("[") r.append(("heading_key", "I")) From 22a4b1d5d4f315ed013332e4219f105e6d928615 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Wed, 2 May 2018 11:20:20 +1200 Subject: [PATCH 4/6] Redesign keepserving - Instead of listening for a pseudo-event, we periodically check whether client replay, server replay or file reading is active. - Adjust server replay not to use tick. - Adjust readfile to expose a command to check whether reading is in progress. --- mitmproxy/addons/keepserving.py | 29 ++++++++++-- mitmproxy/addons/readfile.py | 8 ++++ mitmproxy/addons/serverplayback.py | 9 ---- mitmproxy/master.py | 5 +-- mitmproxy/tools/main.py | 2 - test/mitmproxy/addons/test_keepserving.py | 44 +++++++++++++++++-- test/mitmproxy/addons/test_readfile.py | 2 + test/mitmproxy/addons/test_serverplayback.py | 18 -------- .../mitmproxy/tools/console/test_statusbar.py | 1 + 9 files changed, 78 insertions(+), 40 deletions(-) diff --git a/mitmproxy/addons/keepserving.py b/mitmproxy/addons/keepserving.py index 6413299d8..161f33ff0 100644 --- a/mitmproxy/addons/keepserving.py +++ b/mitmproxy/addons/keepserving.py @@ -1,3 +1,4 @@ +import asyncio from mitmproxy import ctx @@ -12,6 +13,28 @@ class KeepServing: """ ) - def event_processing_complete(self): - if not ctx.master.options.keepserving: - ctx.master.shutdown() + def keepgoing(self) -> bool: + checks = [ + "readfile.reading", + "replay.client.count", + "replay.server.count", + ] + return any([ctx.master.commands.call(c) for c in checks]) + + def shutdown(self): # pragma: no cover + ctx.master.shutdown() + + async def watch(self): + while True: + await asyncio.sleep(0.1) + if not self.keepgoing(): + self.shutdown() + + def running(self): + opts = [ + ctx.options.client_replay, + ctx.options.server_replay, + ctx.options.rfile, + ] + if any(opts) and not ctx.options.keepserving: + asyncio.get_event_loop().create_task(self.watch()) diff --git a/mitmproxy/addons/readfile.py b/mitmproxy/addons/readfile.py index 2ae81fae9..b0a279ba9 100644 --- a/mitmproxy/addons/readfile.py +++ b/mitmproxy/addons/readfile.py @@ -7,6 +7,7 @@ from mitmproxy import ctx from mitmproxy import exceptions from mitmproxy import flowfilter from mitmproxy import io +from mitmproxy import command class ReadFile: @@ -15,6 +16,7 @@ class ReadFile: """ def __init__(self): self.filter = None + self.is_reading = False def load(self, loader): loader.add_option( @@ -65,17 +67,23 @@ class ReadFile: raise exceptions.FlowReadException(str(e)) from e async def doread(self, rfile): + self.is_reading = True try: await self.load_flows_from_path(ctx.options.rfile) except exceptions.FlowReadException as e: raise exceptions.OptionsError(e) from e finally: + self.is_reading = False ctx.master.addons.trigger("processing_complete") def running(self): if ctx.options.rfile: asyncio.get_event_loop().create_task(self.doread(ctx.options.rfile)) + @command.command("readfile.reading") + def reading(self) -> bool: + return self.is_reading + class ReadFileStdin(ReadFile): """Support the special case of "-" for reading from stdin""" diff --git a/mitmproxy/addons/serverplayback.py b/mitmproxy/addons/serverplayback.py index efc7d3598..51ba60b4a 100644 --- a/mitmproxy/addons/serverplayback.py +++ b/mitmproxy/addons/serverplayback.py @@ -13,8 +13,6 @@ import mitmproxy.types class ServerPlayback: def __init__(self): self.flowmap = {} - self.stop = False - self.final_flow = None self.configured = False def load(self, loader): @@ -175,10 +173,6 @@ class ServerPlayback: raise exceptions.OptionsError(str(e)) self.load_flows(flows) - def tick(self): - if self.stop and not self.final_flow.live: - ctx.master.addons.trigger("processing_complete") - def request(self, f): if self.flowmap: rflow = self.next_flow(f) @@ -188,9 +182,6 @@ class ServerPlayback: if ctx.options.server_replay_refresh: response.refresh() f.response = response - if not self.flowmap: - self.final_flow = f - self.stop = True elif ctx.options.server_replay_kill_extra: ctx.log.warn( "server_playback: killed non-replay request {}".format( diff --git a/mitmproxy/master.py b/mitmproxy/master.py index 7f81d1857..c0f6e86ff 100644 --- a/mitmproxy/master.py +++ b/mitmproxy/master.py @@ -127,10 +127,7 @@ class Master: """ if not self.should_exit.is_set(): self.should_exit.set() - asyncio.run_coroutine_threadsafe( - self._shutdown(), - loop = self.channel.loop, - ) + asyncio.run_coroutine_threadsafe(self._shutdown(), loop = self.channel.loop) def _change_reverse_host(self, f): """ diff --git a/mitmproxy/tools/main.py b/mitmproxy/tools/main.py index 5548509e4..cf370f29c 100644 --- a/mitmproxy/tools/main.py +++ b/mitmproxy/tools/main.py @@ -103,8 +103,6 @@ def run( server = proxy.server.DummyServer(pconf) master.server = server - master.addons.trigger("configure", opts.keys()) - master.addons.trigger("tick") opts.update_known(**unknown) if args.options: print(optmanager.dump_defaults(opts)) diff --git a/test/mitmproxy/addons/test_keepserving.py b/test/mitmproxy/addons/test_keepserving.py index 5eafa7929..01b0d09c2 100644 --- a/test/mitmproxy/addons/test_keepserving.py +++ b/test/mitmproxy/addons/test_keepserving.py @@ -3,12 +3,48 @@ import pytest from mitmproxy.addons import keepserving from mitmproxy.test import taddons +from mitmproxy import command + + +class Dummy: + def __init__(self, val: bool): + self.val = val + + def load(self, loader): + loader.add_option("client_replay", bool, self.val, "test") + loader.add_option("server_replay", bool, self.val, "test") + loader.add_option("rfile", bool, self.val, "test") + + @command.command("readfile.reading") + def readfile(self) -> bool: + return self.val + + @command.command("replay.client.count") + def creplay(self) -> int: + return 1 if self.val else 0 + + @command.command("replay.server.count") + def sreplay(self) -> int: + return 1 if self.val else 0 + + +class TKS(keepserving.KeepServing): + _is_shutdown = False + + def shutdown(self): + self.is_shutdown = True @pytest.mark.asyncio async def test_keepserving(): - ks = keepserving.KeepServing() + ks = TKS() + d = Dummy(True) with taddons.context(ks) as tctx: - ks.event_processing_complete() - asyncio.sleep(0.1) - assert tctx.master.should_exit.is_set() + tctx.master.addons.add(d) + ks.running() + assert ks.keepgoing() + + d.val = False + assert not ks.keepgoing() + await asyncio.sleep(0.3) + assert ks.is_shutdown diff --git a/test/mitmproxy/addons/test_readfile.py b/test/mitmproxy/addons/test_readfile.py index d22382a84..62f4d917f 100644 --- a/test/mitmproxy/addons/test_readfile.py +++ b/test/mitmproxy/addons/test_readfile.py @@ -51,6 +51,8 @@ class TestReadFile: async def test_read(self, tmpdir, data, corrupt_data): rf = readfile.ReadFile() with taddons.context(rf) as tctx: + assert not rf.reading() + tf = tmpdir.join("tfile") with asynctest.patch('mitmproxy.master.Master.load_flow') as mck: diff --git a/test/mitmproxy/addons/test_serverplayback.py b/test/mitmproxy/addons/test_serverplayback.py index 0bc28ac89..c6a0c1f48 100644 --- a/test/mitmproxy/addons/test_serverplayback.py +++ b/test/mitmproxy/addons/test_serverplayback.py @@ -39,16 +39,6 @@ def test_config(tmpdir): tctx.configure(s, server_replay=[str(tmpdir)]) -def test_tick(): - s = serverplayback.ServerPlayback() - with taddons.context(s) as tctx: - s.stop = True - s.final_flow = tflow.tflow() - s.final_flow.live = False - s.tick() - assert tctx.master.has_event("processing_complete") - - def test_server_playback(): sp = serverplayback.ServerPlayback() with taddons.context(sp) as tctx: @@ -349,14 +339,6 @@ def test_server_playback_full(): s.request(tf) assert not tf.response - assert not s.stop - s.tick() - assert not s.stop - - tf = tflow.tflow() - s.request(tflow.tflow()) - assert s.stop - def test_server_playback_kill(): s = serverplayback.ServerPlayback() diff --git a/test/mitmproxy/tools/console/test_statusbar.py b/test/mitmproxy/tools/console/test_statusbar.py index 108f238e8..f1cc67d45 100644 --- a/test/mitmproxy/tools/console/test_statusbar.py +++ b/test/mitmproxy/tools/console/test_statusbar.py @@ -30,6 +30,7 @@ def test_statusbar(monkeypatch): m.options.update(view_order='url', console_focus_follow=True) monkeypatch.setattr(m.addons.get("clientplayback"), "count", lambda: 42) monkeypatch.setattr(m.addons.get("serverplayback"), "count", lambda: 42) + monkeypatch.setattr(statusbar.StatusBar, "refresh", lambda x: None) bar = statusbar.StatusBar(m) # this already causes a redraw assert bar.ib._w From 38ff8109fb814d96c271e27caa65b0c98a169ee2 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Wed, 2 May 2018 11:31:28 +1200 Subject: [PATCH 5/6] taddons: remove has_event We no longer use this anywhere, so ditch it. --- mitmproxy/test/taddons.py | 11 ----------- test/mitmproxy/test_taddons.py | 1 - 2 files changed, 12 deletions(-) diff --git a/mitmproxy/test/taddons.py b/mitmproxy/test/taddons.py index 67c15f753..509f8d530 100644 --- a/mitmproxy/test/taddons.py +++ b/mitmproxy/test/taddons.py @@ -17,10 +17,6 @@ class TestAddons(addonmanager.AddonManager): def trigger(self, event, *args, **kwargs): if event == "log": self.master.logs.append(args[0]) - elif event == "tick" and not args and not kwargs: - pass - else: - self.master.events.append((event, args, kwargs)) super().trigger(event, *args, **kwargs) @@ -28,7 +24,6 @@ class RecordingMaster(mitmproxy.master.Master): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.addons = TestAddons(self) - self.events = [] self.logs = [] def dump_log(self, outf=sys.stdout): @@ -51,12 +46,6 @@ class RecordingMaster(mitmproxy.master.Master): await asyncio.sleep(0.1) return False - def has_event(self, name): - for i in self.events: - if i[0] == name: - return True - return False - def clear(self): self.logs = [] diff --git a/test/mitmproxy/test_taddons.py b/test/mitmproxy/test_taddons.py index 5266e0383..53091bc14 100644 --- a/test/mitmproxy/test_taddons.py +++ b/test/mitmproxy/test_taddons.py @@ -10,7 +10,6 @@ from mitmproxy import ctx async def test_recordingmaster(): with taddons.context() as tctx: assert not tctx.master.has_log("nonexistent") - assert not tctx.master.has_event("nonexistent") ctx.log.error("foo") assert not tctx.master.has_log("foo", level="debug") assert await tctx.master.await_log("foo", level="error") From 2f3ba1f66dca70764f7d081aa6836d34bbe963c8 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Wed, 2 May 2018 12:08:15 +1200 Subject: [PATCH 6/6] Catch some stray command.call invocations --- mitmproxy/tools/console/flowlist.py | 6 +++--- mitmproxy/tools/console/flowview.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/mitmproxy/tools/console/flowlist.py b/mitmproxy/tools/console/flowlist.py index 78892fe18..edb7004fc 100644 --- a/mitmproxy/tools/console/flowlist.py +++ b/mitmproxy/tools/console/flowlist.py @@ -79,11 +79,11 @@ class FlowListBox(urwid.ListBox, layoutwidget.LayoutWidget): def keypress(self, size, key): if key == "m_start": - self.master.commands.call("view.go 0") + self.master.commands.execute("view.go 0") elif key == "m_end": - self.master.commands.call("view.go -1") + self.master.commands.execute("view.go -1") elif key == "m_select": - self.master.commands.call("console.view.flow @focus") + self.master.commands.execute("console.view.flow @focus") return urwid.ListBox.keypress(self, size, key) def view_changed(self): diff --git a/mitmproxy/tools/console/flowview.py b/mitmproxy/tools/console/flowview.py index 837d71a84..b8ba7f12f 100644 --- a/mitmproxy/tools/console/flowview.py +++ b/mitmproxy/tools/console/flowview.py @@ -98,7 +98,7 @@ class FlowDetails(tabs.Tabs): msg, body = "", [urwid.Text([("error", "[content missing]")])] return msg, body else: - full = self.master.commands.call("view.getval @focus fullcontents false") + full = self.master.commands.execute("view.getval @focus fullcontents false") if full == "true": limit = sys.maxsize else: