addons.Addons -> addonmanager, builtins -> addons

1 files changed, 92 insertions, 0 deletions

diff --git a/mitmproxy/addonmanager.py b/mitmproxy/addonmanager.py
new file mode 100644
index 00000000..db8e0cd7
--- /dev/null
+++ b/mitmproxy/addonmanager.py
@@ -0,0 +1,92 @@
+from mitmproxy import exceptions
+import pprint
+
+
+def _get_name(itm):
+    return getattr(itm, "name", itm.__class__.__name__.lower())
+
+
+class AddonManager:
+    def __init__(self, master):
+        self.chain = []
+        self.master = master
+        master.options.changed.connect(self._options_update)
+
+    def clear(self):
+        """
+            Remove all addons.
+        """
+        self.done()
+        self.chain = []
+
+    def get(self, name):
+        """
+            Retrieve an addon by name. Addon names are equal to the .name
+            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
+
+    def _options_update(self, options, updated):
+        for i in self.chain:
+            with self.master.handlecontext():
+                self.invoke_with_context(i, "configure", options, updated)
+
+    def startup(self, s):
+        """
+            Run startup events on addon.
+        """
+        self.invoke_with_context(s, "start")
+        self.invoke_with_context(
+            s,
+            "configure",
+            self.master.options,
+            self.master.options.keys()
+        )
+
+    def add(self, *addons):
+        """
+            Add addons to the end of the chain, and run their startup events.
+        """
+        self.chain.extend(addons)
+        for i in addons:
+            self.startup(i)
+
+    def remove(self, addon):
+        """
+            Remove an addon from the chain, and run its done events.
+        """
+        self.chain = [i for i in self.chain if i is not addon]
+        self.invoke_with_context(addon, "done")
+
+    def done(self):
+        for i in self.chain:
+            self.invoke_with_context(i, "done")
+
+    def __len__(self):
+        return len(self.chain)
+
+    def __str__(self):
+        return pprint.pformat([str(i) for i in self.chain])
+
+    def invoke_with_context(self, addon, name, *args, **kwargs):
+        with self.master.handlecontext():
+            self.invoke(addon, name, *args, **kwargs)
+
+    def invoke(self, addon, name, *args, **kwargs):
+        func = getattr(addon, name, None)
+        if func:
+            if not callable(func):
+                raise exceptions.AddonError(
+                    "Addon handler %s not callable" % name
+                )
+            func(*args, **kwargs)
+
+    def __call__(self, name, *args, **kwargs):
+        for i in self.chain:
+            try:
+                self.invoke(i, name, *args, **kwargs)
+            except exceptions.AddonHalt:
+                return