Revamp how addons work

- Addons now nest, which means that addons can manage addons. This has a number
of salutary effects - the scripts addon no longer has to poke into the global
addons list, we no longer have to replace/remove/boot-outof parent addons when
we load scripts, and this paves the way for making our top-level tools into
addons themselves.
- All addon calls are now wrapped in a safe execution environment where
exceptions are caught, and output to stdout/stderr are intercepted and turned
into logs.
- We no longer support script arguments in sys.argv - creating an option
properly is the only way to pass arguments. This means that all scripts are
always directly controllable from interctive tooling, and that arguments are
type-checked.

For now, I've disabled testing of the har dump example - it needs to be moved
to the new argument handling, and become a class addon. I'll address that in a
separate patch.

8 files changed, 216 insertions, 236 deletions

diff --git a/mitmproxy/addonmanager.py b/mitmproxy/addonmanager.py
index 241b3cde..b6d7adb6 100644
--- a/mitmproxy/addonmanager.py
+++ b/mitmproxy/addonmanager.py
@@ -1,4 +1,7 @@
 import typing
+import traceback
+import contextlib
+import sys
 
 from mitmproxy import exceptions
 from mitmproxy import eventsequence
@@ -11,13 +14,66 @@ def _get_name(itm):
     return getattr(itm, "name", itm.__class__.__name__.lower())
 
 
+def cut_traceback(tb, func_name):
+    """
+    Cut off a traceback at the function with the given name.
+    The func_name's frame is excluded.
+
+    Args:
+        tb: traceback object, as returned by sys.exc_info()[2]
+        func_name: function name
+
+    Returns:
+        Reduced traceback.
+    """
+    tb_orig = tb
+    for _, _, fname, _ in traceback.extract_tb(tb):
+        tb = tb.tb_next
+        if fname == func_name:
+            break
+    return tb or tb_orig
+
+
+class StreamLog:
+    """
+        A class for redirecting output using contextlib.
+    """
+    def __init__(self, log):
+        self.log = log
+
+    def write(self, buf):
+        if buf.strip():
+            self.log(buf)
+
+    def flush(self):  # pragma: no cover
+        # Click uses flush sometimes, so we dummy it up
+        pass
+
+
+@contextlib.contextmanager
+def safecall():
+    stdout_replacement = StreamLog(ctx.log.warn)
+    try:
+        with contextlib.redirect_stdout(stdout_replacement):
+            yield
+    except exceptions.AddonHalt:
+        raise
+    except Exception as e:
+        etype, value, tb = sys.exc_info()
+        tb = cut_traceback(tb, "invoke_addon").tb_next
+        ctx.log.error(
+            "Addon error: %s" % "".join(
+                traceback.format_exception(etype, value, tb)
+            )
+        )
+
+
 class Loader:
     """
         A loader object is passed to the load() event when addons start up.
     """
     def __init__(self, master):
         self.master = master
-        self.boot_into_addon = None
 
     def add_option(
         self,
@@ -35,25 +91,33 @@ class Loader:
             choices
         )
 
-    def boot_into(self, addon):
-        self.boot_into_addon = addon
-        func = getattr(addon, "load", None)
-        if func:
-            func(self)
+
+def traverse(chain):
+    """
+        Recursively traverse an addon chain.
+    """
+    for a in chain:
+        yield a
+        if hasattr(a, "addons"):
+            yield from traverse(a.addons)
 
 
 class AddonManager:
     def __init__(self, master):
+        self.lookup = {}
         self.chain = []
         self.master = master
-        master.options.changed.connect(self.configure_all)
+        master.options.changed.connect(self._configure_all)
+
+    def _configure_all(self, options, updated):
+        self.trigger("configure", options, updated)
 
     def clear(self):
         """
             Remove all addons.
         """
-        self.done()
-        self.chain = []
+        for i in self.chain:
+            self.remove(i)
 
     def get(self, name):
         """
@@ -61,36 +125,52 @@ class AddonManager:
             attribute on the instance, or the lower case class name if that
             does not exist.
         """
-        for i in self.chain:
-            if name == _get_name(i):
-                return i
+        return self.lookup.get(name, None)
 
-    def configure_all(self, options, updated):
-        self.trigger("configure", options, updated)
+    def register(self, addon):
+        """
+            Register an addon and all its sub-addons with the manager without
+            adding it to the chain. This should be used by addons that
+            dynamically manage addons. Must be called within a current context.
+        """
+        for a in traverse([addon]):
+            name = _get_name(a)
+            if name in self.lookup:
+                raise exceptions.AddonError(
+                    "An addon called '%s' already exists." % name
+                )
+        l = Loader(self.master)
+        self.invoke_addon(addon, "load", l)
+        for a in traverse([addon]):
+            name = _get_name(a)
+            self.lookup[name] = a
+        return addon
 
     def add(self, *addons):
         """
-            Add addons to the end of the chain, and run their startup events.
+            Add addons to the end of the chain, and run their load event.
+            If any addon has sub-addons, they are registered.
         """
         with self.master.handlecontext():
             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)
+                self.chain.append(self.register(i))
 
     def remove(self, addon):
         """
-            Remove an addon from the chain, and run its done events.
+            Remove an addon and all its sub-addons.
+
+            If the addon is not in the chain - that is, if it's managed by a
+            parent addon - it's the parent's responsibility to remove it from
+            its own addons attribute.
         """
-        self.chain = [i for i in self.chain if i is not addon]
+        for a in traverse([addon]):
+            n = _get_name(a)
+            if n not in self.lookup:
+                raise exceptions.AddonError("No such addon: %s" % n)
+            self.chain = [i for i in self.chain if i is not a]
+            del self.lookup[_get_name(a)]
         with self.master.handlecontext():
-            self.invoke_addon(addon, "done")
-
-    def done(self):
-        self.trigger("done")
+            self.invoke_addon(a, "done")
 
     def __len__(self):
         return len(self.chain)
@@ -126,22 +206,19 @@ class AddonManager:
 
     def invoke_addon(self, addon, name, *args, **kwargs):
         """
-            Invoke an event on an addon. This method must run within an
-            established handler context.
+            Invoke an event on an addon and all its children. This method must
+            run within an established handler context.
         """
-        if not ctx.master:
-            raise exceptions.AddonError(
-                "invoke_addon called without a handler context."
-            )
         if name not in eventsequence.Events:
             name = "event_" + name
-        func = getattr(addon, name, None)
-        if func:
-            if not callable(func):
-                raise exceptions.AddonError(
-                    "Addon handler %s not callable" % name
-                )
-            func(*args, **kwargs)
+        for a in traverse([addon]):
+            func = getattr(a, name, None)
+            if func:
+                if not callable(func):
+                    raise exceptions.AddonError(
+                        "Addon handler %s not callable" % name
+                    )
+                func(*args, **kwargs)
 
     def trigger(self, name, *args, **kwargs):
         """
@@ -150,6 +227,7 @@ class AddonManager:
         with self.master.handlecontext():
             for i in self.chain:
                 try:
-                    self.invoke_addon(i, name, *args, **kwargs)
+                    with safecall():
+                        self.invoke_addon(i, name, *args, **kwargs)
                 except exceptions.AddonHalt:
                     return
diff --git a/mitmproxy/addons/onboarding.py b/mitmproxy/addons/onboarding.py
index cb57990f..6552ec9e 100644
--- a/mitmproxy/addons/onboarding.py
+++ b/mitmproxy/addons/onboarding.py
@@ -3,6 +3,8 @@ from mitmproxy.addons.onboardingapp import app
 
 
 class Onboarding(wsgiapp.WSGIApp):
+    name = "onboarding"
+
     def __init__(self):
         super().__init__(app.Adapter(app.application), None, None)
         self.enabled = False
diff --git a/mitmproxy/addons/script.py b/mitmproxy/addons/script.py
index eef21293..bda823b4 100644
--- a/mitmproxy/addons/script.py
+++ b/mitmproxy/addons/script.py
@@ -1,123 +1,31 @@
-import contextlib
 import os
-import shlex
-import sys
+import importlib
 import threading
-import traceback
-import types
+import sys
 
 from mitmproxy import addonmanager
 from mitmproxy import exceptions
 from mitmproxy import ctx
-from mitmproxy import eventsequence
-
 
 import watchdog.events
 from watchdog.observers import polling
 
 
-def parse_command(command):
-    """
-        Returns a (path, args) tuple.
-    """
-    if not command or not command.strip():
-        raise ValueError("Empty script command.")
-    # Windows: escape all backslashes in the path.
-    if os.name == "nt":  # pragma: no cover
-        backslashes = shlex.split(command, posix=False)[0].count("\\")
-        command = command.replace("\\", "\\\\", backslashes)
-    args = shlex.split(command)  # pragma: no cover
-    args[0] = os.path.expanduser(args[0])
-    if not os.path.exists(args[0]):
-        raise ValueError(
-            ("Script file not found: %s.\r\n"
-             "If your script path contains spaces, "
-             "make sure to wrap it in additional quotes, e.g. -s \"'./foo bar/baz.py' --args\".") %
-            args[0])
-    elif os.path.isdir(args[0]):
-        raise ValueError("Not a file: %s" % args[0])
-    return args[0], args[1:]
-
-
-def cut_traceback(tb, func_name):
-    """
-    Cut off a traceback at the function with the given name.
-    The func_name's frame is excluded.
-
-    Args:
-        tb: traceback object, as returned by sys.exc_info()[2]
-        func_name: function name
-
-    Returns:
-        Reduced traceback.
-    """
-    tb_orig = tb
-
-    for _, _, fname, _ in traceback.extract_tb(tb):
-        tb = tb.tb_next
-        if fname == func_name:
-            break
-
-    if tb is None:
-        # We could not find the method, take the full stack trace.
-        # This may happen on some Python interpreters/flavors (e.g. PyInstaller).
-        return tb_orig
-    else:
-        return tb
-
-
-class StreamLog:
-    """
-        A class for redirecting output using contextlib.
-    """
-    def __init__(self, log):
-        self.log = log
-
-    def write(self, buf):
-        if buf.strip():
-            self.log(buf)
-
-
-@contextlib.contextmanager
-def scriptenv(path, args):
-    oldargs = sys.argv
-    sys.argv = [path] + args
-    script_dir = os.path.dirname(os.path.abspath(path))
-    sys.path.append(script_dir)
-    stdout_replacement = StreamLog(ctx.log.warn)
+def load_script(actx, path):
+    if not os.path.exists(path):
+        ctx.log.info("No such file: %s" % path)
+        return
+    loader = importlib.machinery.SourceFileLoader(os.path.basename(path), path)
     try:
-        with contextlib.redirect_stdout(stdout_replacement):
-            yield
-    except SystemExit as v:
-        ctx.log.error("Script exited with code %s" % v.code)
-    except Exception:
-        etype, value, tb = sys.exc_info()
-        tb = cut_traceback(tb, "scriptenv").tb_next
-        ctx.log.error(
-            "Script error: %s" % "".join(
-                traceback.format_exception(etype, value, tb)
-            )
-        )
+        oldpath = sys.path
+        sys.path.insert(0, os.path.dirname(path))
+        with addonmanager.safecall():
+            m = loader.load_module()
+            if not getattr(m, "name", None):
+                m.name = path
+            return m
     finally:
-        sys.argv = oldargs
-        sys.path.pop()
-
-
-def load_script(path, args):
-    with open(path, "rb") as f:
-        try:
-            code = compile(f.read(), path, 'exec')
-        except SyntaxError as e:
-            ctx.log.error(
-                "Script error: %s line %s: %s" % (
-                    e.filename, e.lineno, e.msg
-                )
-            )
-            return
-    ns = {'__file__': os.path.abspath(path)}
-    with scriptenv(path, args):
-        exec(code, ns)
-    return types.SimpleNamespace(**ns)
+        sys.path[:] = oldpath
 
 
 class ReloadHandler(watchdog.events.FileSystemEventHandler):
@@ -149,59 +57,39 @@ class Script:
     """
         An addon that manages a single script.
     """
-    def __init__(self, command):
-        self.name = command
-
-        self.command = command
-        self.path, self.args = parse_command(command)
+    def __init__(self, path):
+        self.name = "scriptmanager:" + path
+        self.path = path
         self.ns = None
         self.observer = None
-        self.dead = False
 
         self.last_options = None
         self.should_reload = threading.Event()
 
-        for i in eventsequence.Events:
-            if not hasattr(self, i):
-                def mkprox():
-                    evt = i
-
-                    def prox(*args, **kwargs):
-                        self.run(evt, *args, **kwargs)
-                    return prox
-                setattr(self, i, mkprox())
+    def load(self, l):
+        self.ns = load_script(ctx, self.path)
 
-    def run(self, name, *args, **kwargs):
-        # It's possible for ns to be un-initialised if we failed during
-        # configure
-        if self.ns is not None and not self.dead:
-            func = getattr(self.ns, name, None)
-            if func:
-                with scriptenv(self.path, self.args):
-                    return func(*args, **kwargs)
+    @property
+    def addons(self):
+        if self.ns is not None:
+            return [self.ns]
+        return []
 
     def reload(self):
         self.should_reload.set()
 
-    def load_script(self):
-        self.ns = load_script(self.path, self.args)
-        l = addonmanager.Loader(ctx.master)
-        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.configure(self.last_options, self.last_options.keys())
-        else:
-            self.run("tick")
-
-    def load(self, l):
-        self.last_options = ctx.master.options
-        self.load_script()
+            if self.ns:
+                ctx.master.addons.remove(self.ns)
+            self.ns = load_script(ctx, self.path)
+            if self.ns:
+                # We're already running, so we have to explicitly register and
+                # configure the addon
+                ctx.master.addons.register(self.ns)
+                self.configure(self.last_options, self.last_options.keys())
 
     def configure(self, options, updated):
         self.last_options = options
@@ -213,11 +101,6 @@ class Script:
                 os.path.dirname(self.path) or "."
             )
             self.observer.start()
-        self.run("configure", options, updated)
-
-    def done(self):
-        self.run("done")
-        self.dead = True
 
 
 class ScriptLoader:
@@ -226,21 +109,14 @@ class ScriptLoader:
     """
     def __init__(self):
         self.is_running = False
+        self.addons = []
 
     def running(self):
         self.is_running = True
 
     def run_once(self, command, flows):
-        try:
-            sc = Script(command)
-        except ValueError as e:
-            raise ValueError(str(e))
-        sc.load_script()
-        for f in flows:
-            for evt, o in eventsequence.iterate(f):
-                sc.run(evt, o)
-        sc.done()
-        return sc
+        # Returning once we have proper commands
+        raise NotImplementedError
 
     def configure(self, options, updated):
         if "scripts" in updated:
@@ -248,25 +124,21 @@ class ScriptLoader:
                 if options.scripts.count(s) > 1:
                     raise exceptions.OptionsError("Duplicate script: %s" % s)
 
-            for a in ctx.master.addons.chain[:]:
-                if isinstance(a, Script) and a.name not in options.scripts:
+            for a in self.addons[:]:
+                if a.path not in options.scripts:
                     ctx.log.info("Un-loading script: %s" % a.name)
                     ctx.master.addons.remove(a)
+                    self.addons.remove(a)
 
             # The machinations below are to ensure that:
             #   - Scripts remain in the same order
-            #   - Scripts are listed directly after the script addon. This is
-            #   needed to ensure that interactions with, for instance, flow
-            #   serialization remains correct.
             #   - Scripts are not initialized un-necessarily. If only a
-            #   script's order in the script list has changed, it should simply
-            #   be moved.
+            #   script's order in the script list has changed, it is just
+            #   moved.
 
             current = {}
-            for a in ctx.master.addons.chain[:]:
-                if isinstance(a, Script):
-                    current[a.name] = a
-                    ctx.master.addons.chain.remove(a)
+            for a in self.addons:
+                current[a.path] = a
 
             ordered = []
             newscripts = []
@@ -275,24 +147,15 @@ class ScriptLoader:
                     ordered.append(current[s])
                 else:
                     ctx.log.info("Loading script: %s" % s)
-                    try:
-                        sc = Script(s)
-                    except ValueError as e:
-                        raise exceptions.OptionsError(str(e))
+                    sc = Script(s)
                     ordered.append(sc)
                     newscripts.append(sc)
 
-            ochain = ctx.master.addons.chain
-            pos = ochain.index(self)
-            ctx.master.addons.chain = ochain[:pos + 1] + ordered + ochain[pos + 1:]
+            self.addons = ordered
 
             for s in newscripts:
-                l = addonmanager.Loader(ctx.master)
-                ctx.master.addons.invoke_addon(s, "load", l)
+                ctx.master.addons.register(s)
                 if self.is_running:
                     # If we're already running, we configure and tell the addon
                     # we're up and running.
-                    ctx.master.addons.invoke_addon(
-                        s, "configure", options, options.keys()
-                    )
                     ctx.master.addons.invoke_addon(s, "running")
diff --git a/mitmproxy/addons/termlog.py b/mitmproxy/addons/termlog.py
index 5fdb6245..0d42b18e 100644
--- a/mitmproxy/addons/termlog.py
+++ b/mitmproxy/addons/termlog.py
@@ -3,6 +3,11 @@ import click
 
 from mitmproxy import log
 
+# These get over-ridden by the save execution context. Keep them around so we
+# can log directly.
+realstdout = sys.stdout
+realstderr = sys.stderr
+
 
 class TermLog:
     def __init__(self, outfile=None):
@@ -14,9 +19,9 @@ class TermLog:
 
     def log(self, e):
         if log.log_tier(e.level) == log.log_tier("error"):
-            outfile = self.outfile or sys.stderr
+            outfile = self.outfile or realstderr
         else:
-            outfile = self.outfile or sys.stdout
+            outfile = self.outfile or realstdout
 
         if self.options.verbosity >= log.log_tier(e.level):
             click.secho(
diff --git a/mitmproxy/addons/wsgiapp.py b/mitmproxy/addons/wsgiapp.py
index c37fcb7b..155444fc 100644
--- a/mitmproxy/addons/wsgiapp.py
+++ b/mitmproxy/addons/wsgiapp.py
@@ -13,6 +13,10 @@ class WSGIApp:
     def __init__(self, app, host, port):
         self.app, self.host, self.port = app, host, port
 
+    @property
+    def name(self):
+        return "wsgiapp:%s:%s" % (self.host, self.port)
+
     def serve(self, app, flow):
         """
             Serves app on flow, and prevents further handling of the flow.
diff --git a/mitmproxy/master.py b/mitmproxy/master.py
index 946b25a4..46fdb585 100644
--- a/mitmproxy/master.py
+++ b/mitmproxy/master.py
@@ -103,7 +103,7 @@ class Master:
     def shutdown(self):
         self.server.shutdown()
         self.should_exit.set()
-        self.addons.done()
+        self.addons.trigger("done")
 
     def create_request(self, method, url):
         """
diff --git a/mitmproxy/test/taddons.py b/mitmproxy/test/taddons.py
index 8bc174c7..3dbccba2 100644
--- a/mitmproxy/test/taddons.py
+++ b/mitmproxy/test/taddons.py
@@ -1,3 +1,4 @@
+import sys
 import contextlib
 
 import mitmproxy.master
@@ -5,6 +6,7 @@ import mitmproxy.options
 from mitmproxy import proxy
 from mitmproxy import addonmanager
 from mitmproxy import eventsequence
+from mitmproxy.addons import script
 
 
 class TestAddons(addonmanager.AddonManager):
@@ -26,6 +28,10 @@ class RecordingMaster(mitmproxy.master.Master):
         self.events = []
         self.logs = []
 
+    def dump_log(self, outf=sys.stdout):
+        for i in self.logs:
+            print("%s: %s" % (i.level, i.msg), file=outf)
+
     def has_log(self, txt, level=None):
         for i in self.logs:
             if level and i.level != level:
@@ -51,14 +57,21 @@ class context:
         provides a number of helper methods for common testing scenarios.
     """
     def __init__(self, master = None, options = None):
-        self.options = options or mitmproxy.options.Options()
+        options = options or mitmproxy.options.Options()
         self.master = master or RecordingMaster(
             options, proxy.DummyServer(options)
         )
+        self.options = self.master.options
         self.wrapped = None
 
+    def ctx(self):
+        """
+            Returns a new handler context.
+        """
+        return self.master.handlecontext()
+
     def __enter__(self):
-        self.wrapped = self.master.handlecontext()
+        self.wrapped = self.ctx()
         self.wrapped.__enter__()
         return self
 
@@ -75,11 +88,13 @@ class context:
         """
         f.reply._state = "start"
         for evt, arg in eventsequence.iterate(f):
-            h = getattr(addon, evt, None)
-            if h:
-                h(arg)
-                if f.reply.state == "taken":
-                    return
+            self.master.addons.invoke_addon(
+                addon,
+                evt,
+                arg
+            )
+            if f.reply.state == "taken":
+                return
 
     def configure(self, addon, **kwargs):
         """
@@ -89,4 +104,17 @@ class context:
         """
         with self.options.rollback(kwargs.keys(), reraise=True):
             self.options.update(**kwargs)
-            addon.configure(self.options, kwargs.keys())
+            self.master.addons.invoke_addon(
+                addon,
+                "configure",
+                self.options,
+                kwargs.keys()
+            )
+
+    def script(self, path):
+        sc = script.Script(path)
+        loader = addonmanager.Loader(self.master)
+        sc.load(loader)
+        for a in addonmanager.traverse(sc.addons):
+            getattr(a, "load", lambda x: None)(loader)
+        return sc
diff --git a/mitmproxy/tools/main.py b/mitmproxy/tools/main.py
index 6db232fc..de9032c8 100644
--- a/mitmproxy/tools/main.py
+++ b/mitmproxy/tools/main.py
@@ -76,7 +76,7 @@ def run(MasterKlass, args, extra=None):  # pragma: no cover
         unknown = optmanager.load_paths(opts, args.conf)
         server = process_options(parser, opts, args)
         master = MasterKlass(opts, server)
-        master.addons.configure_all(opts, opts.keys())
+        master.addons.trigger("configure", opts, opts.keys())
         remaining = opts.update_known(**unknown)
         if remaining and opts.verbosity > 1:
             print("Ignored options: %s" % remaining)