Add a light-weight custom event system, use it for keepserving

This patch implements the lightweight event system I propose in #2144, adds a
custom event "processing_complete" that is triggered after file read, client
replay and server replay, and introduces a KeepServing addon to handle this for
mitmdump.

22 files changed, 102 insertions, 75 deletions

diff --git a/mitmproxy/addonmanager.py b/mitmproxy/addonmanager.py
index 097f87b7..123f64b2 100644
--- a/mitmproxy/addonmanager.py
+++ b/mitmproxy/addonmanager.py
@@ -69,8 +69,8 @@ class AddonManager:
             raise exceptions.AddonError(
                 "invoke_addon called without a handler context."
             )
-        if name not in eventsequence.Events:  # prama: no cover
-            raise NotImplementedError("Unknown event")
+        if name not in eventsequence.Events:
+            name = "event_" + name
         func = getattr(addon, name, None)
         if func:
             if not callable(func):
@@ -89,4 +89,3 @@ class AddonManager:
                     self.invoke_addon(i, name, *args, **kwargs)
                 except exceptions.AddonHalt:
                     return
-
diff --git a/mitmproxy/addons/clientplayback.py b/mitmproxy/addons/clientplayback.py
index 34c6c9c9..3345e65a 100644
--- a/mitmproxy/addons/clientplayback.py
+++ b/mitmproxy/addons/clientplayback.py
@@ -10,7 +10,6 @@ class ClientPlayback:
     def __init__(self):
         self.flows = None
         self.current_thread = None
-        self.keepserving = False
         self.has_replayed = False
 
     def count(self) -> int:
@@ -32,7 +31,6 @@ class ClientPlayback:
                 self.load(flows)
             else:
                 self.flows = None
-        self.keepserving = options.keepserving
 
     def tick(self):
         if self.current_thread and not self.current_thread.is_alive():
@@ -41,5 +39,5 @@ class ClientPlayback:
             self.current_thread = ctx.master.replay_request(self.flows.pop(0))
             self.has_replayed = True
         if self.has_replayed:
-            if not self.flows and not self.current_thread and not self.keepserving:
-                ctx.master.shutdown()
+            if not self.flows and not self.current_thread:
+                ctx.master.addons.trigger("processing_complete")
diff --git a/mitmproxy/addons/keepserving.py b/mitmproxy/addons/keepserving.py
new file mode 100644
index 00000000..9c975a7b
--- /dev/null
+++ b/mitmproxy/addons/keepserving.py
@@ -0,0 +1,7 @@
+from mitmproxy import ctx
+
+
+class KeepServing:
+    def event_processing_complete(self):
+        if not ctx.master.options.keepserving:
+            ctx.master.shutdown()
diff --git a/mitmproxy/addons/readfile.py b/mitmproxy/addons/readfile.py
index a4b92444..03dcd084 100644
--- a/mitmproxy/addons/readfile.py
+++ b/mitmproxy/addons/readfile.py
@@ -11,7 +11,6 @@ class ReadFile:
     """
     def __init__(self):
         self.path = None
-        self.keepserving = False
 
     def load_flows_file(self, path: str) -> int:
         path = os.path.expanduser(path)
@@ -33,8 +32,6 @@ class ReadFile:
             raise exceptions.FlowReadException(v)
 
     def configure(self, options, updated):
-        if "keepserving" in updated:
-            self.keepserving = options.keepserving
         if "rfile" in updated and options.rfile:
             self.path = options.rfile
 
@@ -46,5 +43,4 @@ class ReadFile:
                 raise exceptions.OptionsError(v)
             finally:
                 self.path = None
-                if not self.keepserving:
-                    ctx.master.shutdown()
+                ctx.master.addons.trigger("processing_complete")
diff --git a/mitmproxy/addons/readstdin.py b/mitmproxy/addons/readstdin.py
index e45d25b8..93a99f01 100644
--- a/mitmproxy/addons/readstdin.py
+++ b/mitmproxy/addons/readstdin.py
@@ -9,13 +9,6 @@ class ReadStdin:
         An addon that reads from stdin if we're not attached to (someting like)
         a tty.
     """
-    def __init__(self):
-        self.keepserving = False
-
-    def configure(self, options, updated):
-        if "keepserving" in updated:
-            self.keepserving = options.keepserving
-
     def running(self, stdin = sys.stdin):
         if not stdin.isatty():
             ctx.log.info("Reading from stdin")
@@ -30,5 +23,4 @@ class ReadStdin:
                     ctx.master.load_flow(i)
             except exceptions.FlowReadException as e:
                 ctx.log.error("Error reading from stdin: %s" % e)
-            if not self.keepserving:
-                ctx.master.shutdown()
+            ctx.master.addons.trigger("processing_complete")
diff --git a/mitmproxy/addons/serverplayback.py b/mitmproxy/addons/serverplayback.py
index f2b5f206..be2d6f2b 100644
--- a/mitmproxy/addons/serverplayback.py
+++ b/mitmproxy/addons/serverplayback.py
@@ -104,7 +104,7 @@ class ServerPlayback:
 
     def tick(self):
         if self.stop and not self.final_flow.live:
-            ctx.master.shutdown()
+            ctx.master.addons.trigger("processing_complete")
 
     def request(self, f):
         if self.flowmap:
@@ -115,7 +115,7 @@ class ServerPlayback:
                 if self.options.refresh_server_playback:
                     response.refresh()
                 f.response = response
-                if not self.flowmap and not self.options.keepserving:
+                if not self.flowmap:
                     self.final_flow = f
                     self.stop = True
             elif self.options.replay_kill_extra:
diff --git a/mitmproxy/options.py b/mitmproxy/options.py
index 036b3d29..5b84ac93 100644
--- a/mitmproxy/options.py
+++ b/mitmproxy/options.py
@@ -79,7 +79,10 @@ class Options(optmanager.OptManager):
         )
         self.add_option(
             "keepserving", bool, False,
-            "Continue serving after client playback or file read."
+            """
+                Instructs mitmdump to continue serving after client playback,
+                server playback or file read. This option is ignored by interactive tools, which always keep serving.
+            """
         )
         self.add_option(
             "server", bool, True,
diff --git a/mitmproxy/test/taddons.py b/mitmproxy/test/taddons.py
index 8d6baa12..c3b65e92 100644
--- a/mitmproxy/test/taddons.py
+++ b/mitmproxy/test/taddons.py
@@ -6,16 +6,36 @@ from mitmproxy import proxy
 from mitmproxy import eventsequence
 
 
+class _AddonWrapper:
+    def __init__(self, master, addons):
+        self.master = master
+        self.addons = addons
+
+    def trigger(self, event, *args, **kwargs):
+        self.master.events.append((event, args, kwargs))
+        return self.addons.trigger(event, *args, **kwargs)
+
+    def __getattr__(self, attr):
+        return getattr(self.addons, attr)
+
+
 class RecordingMaster(mitmproxy.master.Master):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.event_log = []
+        self.addons = _AddonWrapper(self, self.addons)
+        self.events = []
+        self.logs = []
+
+    def has_event(self, name):
+        for i in self.events:
+            if i[0] == name:
+                return True
 
     def add_log(self, e, level):
-        self.event_log.append((level, e))
+        self.logs.append((level, e))
 
     def clear(self):
-        self.event_log = []
+        self.logs = []
 
 
 class context:
diff --git a/mitmproxy/tools/dump.py b/mitmproxy/tools/dump.py
index be83fb1d..42930a7e 100644
--- a/mitmproxy/tools/dump.py
+++ b/mitmproxy/tools/dump.py
@@ -2,7 +2,7 @@ from mitmproxy import controller
 from mitmproxy import addons
 from mitmproxy import options
 from mitmproxy import master
-from mitmproxy.addons import dumper, termlog, termstatus, readstdin
+from mitmproxy.addons import dumper, termlog, termstatus, readstdin, keepserving
 
 
 class DumpMaster(master.Master):
@@ -21,7 +21,7 @@ class DumpMaster(master.Master):
         self.addons.add(*addons.default_addons())
         if with_dumper:
             self.addons.add(dumper.Dumper())
-        self.addons.add(readstdin.ReadStdin())
+        self.addons.add(readstdin.ReadStdin(), keepserving.KeepServing())
 
     @controller.handler
     def log(self, e):
diff --git a/test/mitmproxy/addons/test_check_alpn.py b/test/mitmproxy/addons/test_check_alpn.py
index 2dc0c835..ffaf6cff 100644
--- a/test/mitmproxy/addons/test_check_alpn.py
+++ b/test/mitmproxy/addons/test_check_alpn.py
@@ -12,7 +12,7 @@ class TestCheckALPN:
         with taddons.context() as tctx:
             a = check_alpn.CheckALPN()
             tctx.configure(a)
-            assert not any(msg in m for l, m in tctx.master.event_log)
+            assert not any(msg in m for l, m in tctx.master.logs)
 
     def test_check_no_alpn(self, disable_alpn):
         msg = 'ALPN support missing'
@@ -20,4 +20,4 @@ class TestCheckALPN:
         with taddons.context() as tctx:
             a = check_alpn.CheckALPN()
             tctx.configure(a)
-            assert any(msg in m for l, m in tctx.master.event_log)
+            assert any(msg in m for l, m in tctx.master.logs)
diff --git a/test/mitmproxy/addons/test_check_ca.py b/test/mitmproxy/addons/test_check_ca.py
index fc64621c..859b6d8d 100644
--- a/test/mitmproxy/addons/test_check_ca.py
+++ b/test/mitmproxy/addons/test_check_ca.py
@@ -16,4 +16,4 @@ class TestCheckCA:
             tctx.master.server.config.certstore.default_ca.has_expired = mock.MagicMock(return_value=expired)
             a = check_ca.CheckCA()
             tctx.configure(a)
-            assert any(msg in m for l, m in tctx.master.event_log) is expired
+            assert any(msg in m for l, m in tctx.master.logs) is expired
diff --git a/test/mitmproxy/addons/test_clientplayback.py b/test/mitmproxy/addons/test_clientplayback.py
index c22b3589..f71662f0 100644
--- a/test/mitmproxy/addons/test_clientplayback.py
+++ b/test/mitmproxy/addons/test_clientplayback.py
@@ -23,7 +23,7 @@ class MockThread():
 class TestClientPlayback:
     def test_playback(self):
         cp = clientplayback.ClientPlayback()
-        with taddons.context():
+        with taddons.context() as tctx:
             assert cp.count() == 0
             f = tflow.tflow(resp=True)
             cp.load([f])
@@ -35,17 +35,14 @@ class TestClientPlayback:
                 assert rp.called
                 assert cp.current_thread
 
-            cp.keepserving = False
             cp.flows = None
             cp.current_thread = None
-            with mock.patch("mitmproxy.master.Master.shutdown") as sd:
-                cp.tick()
-                assert sd.called
+            cp.tick()
+            assert tctx.master.has_event("processing_complete")
 
             cp.current_thread = MockThread()
-            with mock.patch("mitmproxy.master.Master.shutdown") as sd:
-                cp.tick()
-                assert cp.current_thread is None
+            cp.tick()
+            assert cp.current_thread is None
 
     def test_configure(self, tmpdir):
         cp = clientplayback.ClientPlayback()
diff --git a/test/mitmproxy/addons/test_dumper.py b/test/mitmproxy/addons/test_dumper.py
index 47374617..23299431 100644
--- a/test/mitmproxy/addons/test_dumper.py
+++ b/test/mitmproxy/addons/test_dumper.py
@@ -151,7 +151,7 @@ class TestContentView:
         with taddons.context(options=options.Options()) as ctx:
             ctx.configure(d, flow_detail=4, verbosity=3)
             d.response(tflow.tflow())
-            assert "Content viewer failed" in ctx.master.event_log[0][1]
+            assert "Content viewer failed" in ctx.master.logs[0][1]
 
 
 def test_tcp():
diff --git a/test/mitmproxy/addons/test_keepserving.py b/test/mitmproxy/addons/test_keepserving.py
new file mode 100644
index 00000000..70f7e1e6
--- /dev/null
+++ b/test/mitmproxy/addons/test_keepserving.py
@@ -0,0 +1,10 @@
+from mitmproxy.addons import keepserving
+from mitmproxy.test import taddons
+
+
+def test_keepserving():
+    ks = keepserving.KeepServing()
+
+    with taddons.context() as tctx:
+        ks.event_processing_complete()
+        assert tctx.master.should_exit.is_set()
diff --git a/test/mitmproxy/addons/test_readfile.py b/test/mitmproxy/addons/test_readfile.py
index c0cf97ae..b30c147b 100644
--- a/test/mitmproxy/addons/test_readfile.py
+++ b/test/mitmproxy/addons/test_readfile.py
@@ -32,13 +32,13 @@ def test_configure(mck, tmpdir):
     with taddons.context() as tctx:
         tf = str(tmpdir.join("tfile"))
         write_data(tf)
-        tctx.configure(rf, rfile=str(tf), keepserving=False)
+        tctx.configure(rf, rfile=str(tf))
         assert not mck.called
         rf.running()
         assert mck.called
 
         write_data(tf, corrupt=True)
-        tctx.configure(rf, rfile=str(tf), keepserving=False)
+        tctx.configure(rf, rfile=str(tf))
         with pytest.raises(exceptions.OptionsError):
             rf.running()
 
@@ -51,7 +51,7 @@ def test_corruption(mck, tmpdir):
         with pytest.raises(exceptions.FlowReadException):
             rf.load_flows_file("nonexistent")
         assert not mck.called
-        assert len(tctx.master.event_log) == 1
+        assert len(tctx.master.logs) == 1
 
         tfc = str(tmpdir.join("tfile"))
         write_data(tfc, corrupt=True)
@@ -59,4 +59,4 @@ def test_corruption(mck, tmpdir):
         with pytest.raises(exceptions.FlowReadException):
             rf.load_flows_file(tfc)
         assert mck.called
-        assert len(tctx.master.event_log) == 2
+        assert len(tctx.master.logs) == 2
diff --git a/test/mitmproxy/addons/test_readstdin.py b/test/mitmproxy/addons/test_readstdin.py
index bbef81fc..76b01f4f 100644
--- a/test/mitmproxy/addons/test_readstdin.py
+++ b/test/mitmproxy/addons/test_readstdin.py
@@ -26,12 +26,6 @@ def gen_data(corrupt=False):
     return tf
 
 
-def test_configure(tmpdir):
-    rf = readstdin.ReadStdin()
-    with taddons.context() as tctx:
-        tctx.configure(rf, keepserving=False)
-
-
 class mStdin:
     def __init__(self, d):
         self.buffer = d
@@ -49,11 +43,11 @@ def test_read(m, tmpdir):
         assert m.called
 
         rf.running(stdin=mStdin(None))
-        assert tctx.master.event_log
+        assert tctx.master.logs
         tctx.master.clear()
 
         m.reset_mock()
         assert not m.called
         rf.running(stdin=mStdin(gen_data(corrupt=True)))
         assert m.called
-        assert tctx.master.event_log
+        assert tctx.master.logs
diff --git a/test/mitmproxy/addons/test_replace.py b/test/mitmproxy/addons/test_replace.py
index 7d590b35..9002afb5 100644
--- a/test/mitmproxy/addons/test_replace.py
+++ b/test/mitmproxy/addons/test_replace.py
@@ -97,6 +97,6 @@ class TestReplaceFile:
             tmpfile.remove()
             f = tflow.tflow()
             f.request.content = b"foo"
-            assert not tctx.master.event_log
+            assert not tctx.master.logs
             r.request(f)
-            assert tctx.master.event_log
+            assert tctx.master.logs
diff --git a/test/mitmproxy/addons/test_script.py b/test/mitmproxy/addons/test_script.py
index 4c1b2e43..ad3c9a1a 100644
--- a/test/mitmproxy/addons/test_script.py
+++ b/test/mitmproxy/addons/test_script.py
@@ -22,14 +22,14 @@ def test_scriptenv():
     with taddons.context() as tctx:
         with script.scriptenv("path", []):
             raise SystemExit
-        assert tctx.master.event_log[0][0] == "error"
-        assert "exited" in tctx.master.event_log[0][1]
+        assert tctx.master.logs[0][0] == "error"
+        assert "exited" in tctx.master.logs[0][1]
 
         tctx.master.clear()
         with script.scriptenv("path", []):
             raise ValueError("fooo")
-        assert tctx.master.event_log[0][0] == "error"
-        assert "foo" in tctx.master.event_log[0][1]
+        assert tctx.master.logs[0][0] == "error"
+        assert "foo" in tctx.master.logs[0][1]
 
 
 class Called:
@@ -135,7 +135,7 @@ class TestScript:
                 f.write(".")
                 sc.tick()
                 time.sleep(0.1)
-                if tctx.master.event_log:
+                if tctx.master.logs:
                     return
             raise AssertionError("Change event not detected.")
 
@@ -147,11 +147,11 @@ class TestScript:
             sc.start(tctx.options)
             f = tflow.tflow(resp=True)
             sc.request(f)
-            assert tctx.master.event_log[0][0] == "error"
-            assert len(tctx.master.event_log[0][1].splitlines()) == 6
-            assert re.search(r'addonscripts[\\/]error.py", line \d+, in request', tctx.master.event_log[0][1])
-            assert re.search(r'addonscripts[\\/]error.py", line \d+, in mkerr', tctx.master.event_log[0][1])
-            assert tctx.master.event_log[0][1].endswith("ValueError: Error!\n")
+            assert tctx.master.logs[0][0] == "error"
+            assert len(tctx.master.logs[0][1].splitlines()) == 6
+            assert re.search(r'addonscripts[\\/]error.py", line \d+, in request', tctx.master.logs[0][1])
+            assert re.search(r'addonscripts[\\/]error.py", line \d+, in mkerr', tctx.master.logs[0][1])
+            assert tctx.master.logs[0][1].endswith("ValueError: Error!\n")
 
     def test_addon(self):
         with taddons.context() as tctx:
@@ -256,7 +256,7 @@ class TestScriptLoader:
                     "%s %s" % (rec, "c"),
                 ]
             )
-            debug = [(i[0], i[1]) for i in tctx.master.event_log if i[0] == "debug"]
+            debug = [(i[0], i[1]) for i in tctx.master.logs if i[0] == "debug"]
             assert debug == [
                 ('debug', 'a start'),
                 ('debug', 'a configure'),
@@ -270,7 +270,7 @@ class TestScriptLoader:
                 ('debug', 'c configure'),
                 ('debug', 'c running'),
             ]
-            tctx.master.event_log = []
+            tctx.master.logs = []
             tctx.configure(
                 sc,
                 scripts = [
@@ -279,11 +279,11 @@ class TestScriptLoader:
                     "%s %s" % (rec, "b"),
                 ]
             )
-            debug = [(i[0], i[1]) for i in tctx.master.event_log if i[0] == "debug"]
+            debug = [(i[0], i[1]) for i in tctx.master.logs if i[0] == "debug"]
             # No events, only order has changed
             assert debug == []
 
-            tctx.master.event_log = []
+            tctx.master.logs = []
             tctx.configure(
                 sc,
                 scripts = [
@@ -291,7 +291,7 @@ class TestScriptLoader:
                     "%s %s" % (rec, "a"),
                 ]
             )
-            debug = [(i[0], i[1]) for i in tctx.master.event_log if i[0] == "debug"]
+            debug = [(i[0], i[1]) for i in tctx.master.logs if i[0] == "debug"]
             assert debug == [
                 ('debug', 'c done'),
                 ('debug', 'b done'),
diff --git a/test/mitmproxy/addons/test_serverplayback.py b/test/mitmproxy/addons/test_serverplayback.py
index 54e4d281..02642c35 100644
--- a/test/mitmproxy/addons/test_serverplayback.py
+++ b/test/mitmproxy/addons/test_serverplayback.py
@@ -34,7 +34,7 @@ def test_tick():
         s.final_flow = tflow.tflow()
         s.final_flow.live = False
         s.tick()
-        assert tctx.master.should_exit.is_set()
+        assert tctx.master.has_event("processing_complete")
 
 
 def test_server_playback():
@@ -315,7 +315,6 @@ def test_server_playback_full():
         tctx.configure(
             s,
             refresh_server_playback = True,
-            keepserving=False
         )
 
         f = tflow.tflow()
diff --git a/test/mitmproxy/addons/test_termstatus.py b/test/mitmproxy/addons/test_termstatus.py
index 01c14814..7becc857 100644
--- a/test/mitmproxy/addons/test_termstatus.py
+++ b/test/mitmproxy/addons/test_termstatus.py
@@ -6,7 +6,7 @@ def test_configure():
     ts = termstatus.TermStatus()
     with taddons.context() as ctx:
         ts.running()
-        assert not ctx.master.event_log
+        assert not ctx.master.logs
         ctx.configure(ts, server=True)
         ts.running()
-        assert ctx.master.event_log
+        assert ctx.master.logs
diff --git a/test/mitmproxy/script/test_concurrent.py b/test/mitmproxy/script/test_concurrent.py
index a9b6f0c4..206482e2 100644
--- a/test/mitmproxy/script/test_concurrent.py
+++ b/test/mitmproxy/script/test_concurrent.py
@@ -43,7 +43,7 @@ class TestConcurrent(tservers.MasterTest):
                 )
             )
             sc.start(tctx.options)
-            assert "decorator not supported" in tctx.master.event_log[0][1]
+            assert "decorator not supported" in tctx.master.logs[0][1]
 
     def test_concurrent_class(self):
             with taddons.context() as tctx:
diff --git a/test/mitmproxy/test_addonmanager.py b/test/mitmproxy/test_addonmanager.py
index ef34371f..e7be25b8 100644
--- a/test/mitmproxy/test_addonmanager.py
+++ b/test/mitmproxy/test_addonmanager.py
@@ -11,6 +11,7 @@ class TAddon:
     def __init__(self, name):
         self.name = name
         self.tick = True
+        self.custom_called = False
 
     def __repr__(self):
         return "Addon(%s)" % self.name
@@ -18,11 +19,17 @@ class TAddon:
     def done(self):
         pass
 
+    def event_custom(self):
+        self.custom_called = True
+
 
 def test_simple():
     o = options.Options()
     m = master.Master(o, proxy.DummyServer(o))
     a = addonmanager.AddonManager(m)
+    with pytest.raises(exceptions.AddonError):
+        a.invoke_addon(TAddon("one"), "done")
+
     a.add(TAddon("one"))
     assert a.get("one")
     assert not a.get("two")
@@ -33,3 +40,8 @@ def test_simple():
     a.trigger("done")
     with pytest.raises(exceptions.AddonError):
         a.trigger("tick")
+
+    ta = TAddon("one")
+    a.add(ta)
+    a.trigger("custom")
+    assert ta.custom_called