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.

1 files changed, 119 insertions, 41 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