From 65f0885bd6809966f694d1ffb965401b8ab2cffc Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@corte.si>
Date: Fri, 24 Mar 2017 11:29:36 +1300
Subject: addon loader: add boot_into, which replaces returning from start()

While we're here, expand test coverage for addonmanager to 100%, and promote to
individual coverage.
---
 examples/complex/dns_spoofing.py                   |  4 +-
 examples/complex/har_dump.py                       |  2 +-
 examples/complex/remote_debug.py                   |  2 +-
 examples/complex/tls_passthrough.py                |  2 +-
 examples/simple/add_header_class.py                |  4 +-
 examples/simple/custom_contentview.py              |  2 +-
 examples/simple/custom_option.py                   |  4 +-
 examples/simple/filter_flows.py                    |  4 +-
 examples/simple/io_write_dumpfile.py               |  4 +-
 examples/simple/log_events.py                      |  2 +-
 examples/simple/modify_body_inject_iframe.py       |  4 +-
 examples/simple/script_arguments.py                |  4 +-
 examples/simple/wsgi_flask_app.py                  |  2 +-
 mitmproxy/addonmanager.py                          | 14 +++-
 mitmproxy/addons/script.py                         | 15 ++--
 mitmproxy/tools/dump.py                            | 10 +--
 setup.cfg                                          |  1 -
 test/mitmproxy/addons/test_script.py               |  9 +-
 test/mitmproxy/data/addonscripts/addon.py          |  4 +-
 .../addonscripts/concurrent_decorator_class.py     |  4 +-
 test/mitmproxy/data/addonscripts/recorder.py       |  4 +-
 test/mitmproxy/script/test_concurrent.py           | 10 ++-
 test/mitmproxy/test_addonmanager.py                | 95 ++++++++++++++++++++++
 23 files changed, 158 insertions(+), 48 deletions(-)

diff --git a/examples/complex/dns_spoofing.py b/examples/complex/dns_spoofing.py
index 9f6af522..01e036b2 100644
--- a/examples/complex/dns_spoofing.py
+++ b/examples/complex/dns_spoofing.py
@@ -54,5 +54,5 @@ class Rerouter:
         flow.request.port = port
 
 
-def load(opts):
-    return Rerouter()
+def load(l):
+    l.boot_into(Rerouter())
diff --git a/examples/complex/har_dump.py b/examples/complex/har_dump.py
index f9f8cd95..0515d0f5 100644
--- a/examples/complex/har_dump.py
+++ b/examples/complex/har_dump.py
@@ -25,7 +25,7 @@ HAR = {}
 SERVERS_SEEN = set()
 
 
-def load(opts):
+def load(l):
     """
         Called once on script startup before any other events.
     """
diff --git a/examples/complex/remote_debug.py b/examples/complex/remote_debug.py
index 0b7049d2..fa6f3d33 100644
--- a/examples/complex/remote_debug.py
+++ b/examples/complex/remote_debug.py
@@ -14,6 +14,6 @@ Usage:
 """
 
 
-def load(opts):
+def load(l):
     import pydevd
     pydevd.settrace("localhost", port=5678, stdoutToServer=True, stderrToServer=True)
diff --git a/examples/complex/tls_passthrough.py b/examples/complex/tls_passthrough.py
index 3b3c8936..72c0244b 100644
--- a/examples/complex/tls_passthrough.py
+++ b/examples/complex/tls_passthrough.py
@@ -112,7 +112,7 @@ class TlsFeedback(TlsLayer):
 tls_strategy = None
 
 
-def load(opts):
+def load(l):
     global tls_strategy
     if len(sys.argv) == 2:
         tls_strategy = ProbabilisticStrategy(float(sys.argv[1]))
diff --git a/examples/simple/add_header_class.py b/examples/simple/add_header_class.py
index 7a0dbafa..69b64163 100644
--- a/examples/simple/add_header_class.py
+++ b/examples/simple/add_header_class.py
@@ -3,5 +3,5 @@ class AddHeader:
         flow.response.headers["newheader"] = "foo"
 
 
-def load(opts):
-    return AddHeader()
+def load(l):
+    return l.boot_into(AddHeader())
diff --git a/examples/simple/custom_contentview.py b/examples/simple/custom_contentview.py
index b159154f..34fa5541 100644
--- a/examples/simple/custom_contentview.py
+++ b/examples/simple/custom_contentview.py
@@ -20,7 +20,7 @@ class ViewSwapCase(contentviews.View):
 view = ViewSwapCase()
 
 
-def load(opts):
+def load(l):
     contentviews.add(view)
 
 
diff --git a/examples/simple/custom_option.py b/examples/simple/custom_option.py
index 3190b840..c8bc98d4 100644
--- a/examples/simple/custom_option.py
+++ b/examples/simple/custom_option.py
@@ -1,9 +1,9 @@
 from mitmproxy import ctx
 
 
-def load(options):
+def load(l):
     ctx.log.info("Registering option 'custom'")
-    options.add_option("custom", bool, False, "A custom option")
+    l.add_option("custom", bool, False, "A custom option")
 
 
 def configure(options, updated):
diff --git a/examples/simple/filter_flows.py b/examples/simple/filter_flows.py
index 4d3e264e..d2b735be 100644
--- a/examples/simple/filter_flows.py
+++ b/examples/simple/filter_flows.py
@@ -17,7 +17,7 @@ class Filter:
             print(flow)
 
 
-def load(opts):
+def load(l):
     if len(sys.argv) != 2:
         raise ValueError("Usage: -s 'filt.py FILTER'")
-    return Filter(sys.argv[1])
+    l.boot_into(Filter(sys.argv[1]))
diff --git a/examples/simple/io_write_dumpfile.py b/examples/simple/io_write_dumpfile.py
index 3204e566..15e7693c 100644
--- a/examples/simple/io_write_dumpfile.py
+++ b/examples/simple/io_write_dumpfile.py
@@ -23,7 +23,7 @@ class Writer:
             self.w.add(flow)
 
 
-def load(opts):
+def load(l):
     if len(sys.argv) != 2:
         raise ValueError('Usage: -s "flowriter.py filename"')
-    return Writer(sys.argv[1])
+    l.boot_into(Writer(sys.argv[1]))
diff --git a/examples/simple/log_events.py b/examples/simple/log_events.py
index 518a202f..581b99f3 100644
--- a/examples/simple/log_events.py
+++ b/examples/simple/log_events.py
@@ -7,6 +7,6 @@ If you want to help us out: https://github.com/mitmproxy/mitmproxy/issues/1530 :
 from mitmproxy import ctx
 
 
-def load(opts):
+def load(l):
     ctx.log.info("This is some informative text.")
     ctx.log.error("This is an error.")
diff --git a/examples/simple/modify_body_inject_iframe.py b/examples/simple/modify_body_inject_iframe.py
index f92aac04..442a5118 100644
--- a/examples/simple/modify_body_inject_iframe.py
+++ b/examples/simple/modify_body_inject_iframe.py
@@ -23,7 +23,7 @@ class Injector:
             flow.response.content = str(html).encode("utf8")
 
 
-def load(opts):
+def load(l):
     if len(sys.argv) != 2:
         raise ValueError('Usage: -s "iframe_injector.py url"')
-    return Injector(sys.argv[1])
+    return l.boot_into(Injector(sys.argv[1]))
diff --git a/examples/simple/script_arguments.py b/examples/simple/script_arguments.py
index e16f6a22..84292eb9 100644
--- a/examples/simple/script_arguments.py
+++ b/examples/simple/script_arguments.py
@@ -9,9 +9,9 @@ class Replacer:
         flow.response.replace(self.src, self.dst)
 
 
-def load(opts):
+def load(l):
     parser = argparse.ArgumentParser()
     parser.add_argument("src", type=str)
     parser.add_argument("dst", type=str)
     args = parser.parse_args()
-    return Replacer(args.src, args.dst)
+    l.boot_into(Replacer(args.src, args.dst))
diff --git a/examples/simple/wsgi_flask_app.py b/examples/simple/wsgi_flask_app.py
index 5def17b4..30008b16 100644
--- a/examples/simple/wsgi_flask_app.py
+++ b/examples/simple/wsgi_flask_app.py
@@ -14,7 +14,7 @@ def hello_world():
     return 'Hello World!'
 
 
-def load(opts):
+def load(l):
     # Host app at the magic domain "proxapp" on port 80. Requests to this
     # domain and port combination will now be routed to the WSGI app instance.
     return wsgiapp.WSGIApp(app, "proxapp", 80)
diff --git a/mitmproxy/addonmanager.py b/mitmproxy/addonmanager.py
index 1fc8e4dc..241b3cde 100644
--- a/mitmproxy/addonmanager.py
+++ b/mitmproxy/addonmanager.py
@@ -17,6 +17,7 @@ class Loader:
     """
     def __init__(self, master):
         self.master = master
+        self.boot_into_addon = None
 
     def add_option(
         self,
@@ -34,6 +35,12 @@ class Loader:
             choices
         )
 
+    def boot_into(self, addon):
+        self.boot_into_addon = addon
+        func = getattr(addon, "load", None)
+        if func:
+            func(self)
+
 
 class AddonManager:
     def __init__(self, master):
@@ -65,11 +72,14 @@ class AddonManager:
         """
             Add addons to the end of the chain, and run their startup events.
         """
-        self.chain.extend(addons)
         with self.master.handlecontext():
-            l = Loader(self.master)
             for i in addons:
+                l = Loader(self.master)
                 self.invoke_addon(i, "load", l)
+                if l.boot_into_addon:
+                    self.chain.append(l.boot_into_addon)
+                else:
+                    self.chain.append(i)
 
     def remove(self, addon):
         """
diff --git a/mitmproxy/addons/script.py b/mitmproxy/addons/script.py
index ca1085d0..eef21293 100644
--- a/mitmproxy/addons/script.py
+++ b/mitmproxy/addons/script.py
@@ -186,23 +186,21 @@ class Script:
     def load_script(self):
         self.ns = load_script(self.path, self.args)
         l = addonmanager.Loader(ctx.master)
-        ret = self.run("load", l)
-        if ret:
-            self.ns = ret
-            self.run("load", l)
+        self.run("load", l)
+        if l.boot_into_addon:
+            self.ns = l.boot_into_addon
 
     def tick(self):
         if self.should_reload.is_set():
             self.should_reload.clear()
             ctx.log.info("Reloading script: %s" % self.name)
             self.ns = load_script(self.path, self.args)
-            self.load(self.last_options)
             self.configure(self.last_options, self.last_options.keys())
         else:
             self.run("tick")
 
-    def load(self, opts):
-        self.last_options = opts
+    def load(self, l):
+        self.last_options = ctx.master.options
         self.load_script()
 
     def configure(self, options, updated):
@@ -289,7 +287,8 @@ class ScriptLoader:
             ctx.master.addons.chain = ochain[:pos + 1] + ordered + ochain[pos + 1:]
 
             for s in newscripts:
-                ctx.master.addons.invoke_addon(s, "load", options)
+                l = addonmanager.Loader(ctx.master)
+                ctx.master.addons.invoke_addon(s, "load", l)
                 if self.is_running:
                     # If we're already running, we configure and tell the addon
                     # we're up and running.
diff --git a/mitmproxy/tools/dump.py b/mitmproxy/tools/dump.py
index e6f0c3df..6329f6b7 100644
--- a/mitmproxy/tools/dump.py
+++ b/mitmproxy/tools/dump.py
@@ -16,11 +16,11 @@ class ErrorCheck:
 class DumpMaster(master.Master):
 
     def __init__(
-            self,
-            options: options.Options,
-            server,
-            with_termlog=True,
-            with_dumper=True,
+        self,
+        options: options.Options,
+        server,
+        with_termlog=True,
+        with_dumper=True,
     ) -> None:
         master.Master.__init__(self, options, server)
         self.errorcheck = ErrorCheck()
diff --git a/setup.cfg b/setup.cfg
index 8e231f28..d0307bc8 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -46,7 +46,6 @@ exclude =
 
 [tool:individual_coverage]
 exclude =
-    mitmproxy/addonmanager.py
     mitmproxy/addons/onboardingapp/app.py
     mitmproxy/addons/termlog.py
     mitmproxy/contentviews/base.py
diff --git a/test/mitmproxy/addons/test_script.py b/test/mitmproxy/addons/test_script.py
index 79416161..c68981de 100644
--- a/test/mitmproxy/addons/test_script.py
+++ b/test/mitmproxy/addons/test_script.py
@@ -9,6 +9,7 @@ from unittest import mock
 from mitmproxy.test import tflow
 from mitmproxy.test import tutils
 from mitmproxy.test import taddons
+from mitmproxy import addonmanager
 from mitmproxy import exceptions
 from mitmproxy import options
 from mitmproxy import proxy
@@ -116,7 +117,7 @@ def test_script_print_stdout():
                         "mitmproxy/data/addonscripts/print.py"
                     ), []
                 )
-                ns.start(tctx.options)
+                ns.load(addonmanager.Loader(tctx.master))
         mock_warn.assert_called_once_with("stdoutprint")
 
 
@@ -157,7 +158,8 @@ class TestScript:
             sc = script.Script(
                 tutils.test_data.path("mitmproxy/data/addonscripts/error.py")
             )
-            sc.load(tctx.options)
+            l = addonmanager.Loader(tctx.master)
+            sc.load(l)
             f = tflow.tflow(resp=True)
             sc.request(f)
             assert tctx.master.logs[0].level == "error"
@@ -173,7 +175,8 @@ class TestScript:
                     "mitmproxy/data/addonscripts/addon.py"
                 )
             )
-            sc.load(tctx.options)
+            l = addonmanager.Loader(tctx.master)
+            sc.load(l)
             tctx.configure(sc)
             assert sc.ns.event_log == [
                 'scriptload', 'addonload', 'addonconfigure'
diff --git a/test/mitmproxy/data/addonscripts/addon.py b/test/mitmproxy/data/addonscripts/addon.py
index 3186400a..beef2ce7 100644
--- a/test/mitmproxy/data/addonscripts/addon.py
+++ b/test/mitmproxy/data/addonscripts/addon.py
@@ -17,6 +17,6 @@ def configure(options, updated):
     event_log.append("addonconfigure")
 
 
-def load(opts):
+def load(l):
     event_log.append("scriptload")
-    return Addon()
+    l.boot_into(Addon())
diff --git a/test/mitmproxy/data/addonscripts/concurrent_decorator_class.py b/test/mitmproxy/data/addonscripts/concurrent_decorator_class.py
index d486923a..8e6988d4 100644
--- a/test/mitmproxy/data/addonscripts/concurrent_decorator_class.py
+++ b/test/mitmproxy/data/addonscripts/concurrent_decorator_class.py
@@ -9,5 +9,5 @@ class ConcurrentClass:
         time.sleep(0.1)
 
 
-def load(opts):
-    return ConcurrentClass()
+def load(l):
+    l.boot_into(ConcurrentClass())
diff --git a/test/mitmproxy/data/addonscripts/recorder.py b/test/mitmproxy/data/addonscripts/recorder.py
index 5e0cdfc2..fe497b05 100644
--- a/test/mitmproxy/data/addonscripts/recorder.py
+++ b/test/mitmproxy/data/addonscripts/recorder.py
@@ -22,5 +22,5 @@ class CallLogger:
         raise AttributeError
 
 
-def load(opts):
-    return CallLogger(*sys.argv[1:])
+def load(l):
+    l.boot_into(CallLogger(*sys.argv[1:]))
diff --git a/test/mitmproxy/script/test_concurrent.py b/test/mitmproxy/script/test_concurrent.py
index ccae0d31..86efdfc2 100644
--- a/test/mitmproxy/script/test_concurrent.py
+++ b/test/mitmproxy/script/test_concurrent.py
@@ -2,6 +2,7 @@ from mitmproxy.test import tflow
 from mitmproxy.test import tutils
 from mitmproxy.test import taddons
 
+from mitmproxy import addonmanager
 from mitmproxy import controller
 from mitmproxy.addons import script
 
@@ -24,7 +25,8 @@ class TestConcurrent(tservers.MasterTest):
                     "mitmproxy/data/addonscripts/concurrent_decorator.py"
                 )
             )
-            sc.load(tctx.options)
+            l = addonmanager.Loader(tctx.master)
+            sc.load(l)
 
             f1, f2 = tflow.tflow(), tflow.tflow()
             tctx.cycle(sc, f1)
@@ -42,7 +44,8 @@ class TestConcurrent(tservers.MasterTest):
                     "mitmproxy/data/addonscripts/concurrent_decorator_err.py"
                 )
             )
-            sc.load(tctx.options)
+            l = addonmanager.Loader(tctx.master)
+            sc.load(l)
             assert tctx.master.has_log("decorator not supported")
 
     def test_concurrent_class(self):
@@ -52,7 +55,8 @@ class TestConcurrent(tservers.MasterTest):
                         "mitmproxy/data/addonscripts/concurrent_decorator_class.py"
                     )
                 )
-                sc.load(tctx.options)
+                l = addonmanager.Loader(tctx.master)
+                sc.load(l)
 
                 f1, f2 = tflow.tflow(), tflow.tflow()
                 tctx.cycle(sc, f1)
diff --git a/test/mitmproxy/test_addonmanager.py b/test/mitmproxy/test_addonmanager.py
index e7be25b8..5bb88eb6 100644
--- a/test/mitmproxy/test_addonmanager.py
+++ b/test/mitmproxy/test_addonmanager.py
@@ -5,6 +5,7 @@ from mitmproxy import exceptions
 from mitmproxy import options
 from mitmproxy import master
 from mitmproxy import proxy
+from mitmproxy.test import tflow
 
 
 class TAddon:
@@ -23,6 +24,58 @@ class TAddon:
         self.custom_called = True
 
 
+class THalt:
+    def event_custom(self):
+        raise exceptions.AddonHalt
+
+
+class AOption:
+    def load(self, l):
+        l.add_option("custom_option", bool, False, "help")
+
+
+class AChain:
+    def __init__(self, name, next):
+        self.name = name
+        self.next = next
+
+    def load(self, l):
+        if self.next:
+            l.boot_into(self.next)
+
+    def __repr__(self):
+        return "<%s>" % self.name
+
+
+def test_halt():
+    o = options.Options()
+    m = master.Master(o, proxy.DummyServer(o))
+    a = addonmanager.AddonManager(m)
+    halt = THalt()
+    end = TAddon("end")
+    a.add(halt)
+    a.add(end)
+
+    a.trigger("custom")
+    assert not end.custom_called
+
+    a.remove(halt)
+    a.trigger("custom")
+    assert end.custom_called
+
+
+def test_lifecycle():
+    o = options.Options()
+    m = master.Master(o, proxy.DummyServer(o))
+    a = addonmanager.AddonManager(m)
+    a.add(TAddon("one"))
+
+    f = tflow.tflow()
+    a.handle_lifecycle("request", f)
+
+    a.configure_all(o, o.keys())
+
+
 def test_simple():
     o = options.Options()
     m = master.Master(o, proxy.DummyServer(o))
@@ -30,10 +83,13 @@ def test_simple():
     with pytest.raises(exceptions.AddonError):
         a.invoke_addon(TAddon("one"), "done")
 
+    assert len(a) == 0
     a.add(TAddon("one"))
     assert a.get("one")
     assert not a.get("two")
+    assert len(a) == 1
     a.clear()
+    assert len(a) == 0
     assert not a.chain
 
     a.add(TAddon("one"))
@@ -41,7 +97,46 @@ def test_simple():
     with pytest.raises(exceptions.AddonError):
         a.trigger("tick")
 
+    a.remove(a.get("one"))
+    assert not a.get("one")
+
     ta = TAddon("one")
     a.add(ta)
     a.trigger("custom")
     assert ta.custom_called
+
+
+def test_load_option():
+    o = options.Options()
+    m = master.Master(o, proxy.DummyServer(o))
+    a = addonmanager.AddonManager(m)
+    a.add(AOption())
+    assert "custom_option" in m.options._options
+
+
+def test_loadchain():
+    o = options.Options()
+    m = master.Master(o, proxy.DummyServer(o))
+    a = addonmanager.AddonManager(m)
+
+    a.add(AChain("one", None))
+    assert a.get("one")
+    a.clear()
+
+    a.add(AChain("one", AChain("two", None)))
+    assert not a.get("one")
+    assert a.get("two")
+    a.clear()
+
+    a.add(AChain("one", AChain("two", AChain("three", None))))
+    assert not a.get("one")
+    assert not a.get("two")
+    assert a.get("three")
+    a.clear()
+
+    a.add(AChain("one", AChain("two", AChain("three", AChain("four", None)))))
+    assert not a.get("one")
+    assert not a.get("two")
+    assert not a.get("three")
+    assert a.get("four")
+    a.clear()
-- 
cgit v1.2.3