Basic outline of addons

Add addons.py, integrate with our event mechanism, and change the Master API so
options is the first init argument.

12 files changed, 138 insertions, 43 deletions

diff --git a/mitmproxy/addons.py b/mitmproxy/addons.py
new file mode 100644
index 00000000..7ac65a09
--- /dev/null
+++ b/mitmproxy/addons.py
@@ -0,0 +1,63 @@
+from __future__ import absolute_import, print_function, division
+from mitmproxy import exceptions
+import pprint
+
+
+def _get_name(itm):
+    return getattr(itm, "name", itm.__class__.__name__)
+
+
+class Addons(object):
+    def __init__(self, master):
+        self.chain = []
+        self.master = master
+        master.options.changed.connect(self.options_update)
+
+    def options_update(self, options):
+        for i in self.chain:
+            with self.master.handlecontext():
+                i.configure(options)
+
+    def add(self, *addons):
+        self.chain.extend(addons)
+        for i in addons:
+            self.invoke_with_context(i, "configure", self.master.options)
+
+    def remove(self, addon):
+        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 has_addon(self, name):
+        """
+            Is an addon with this name registered?
+        """
+        for i in self.chain:
+            if _get_name(i) == name:
+                return True
+
+    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:
+            self.invoke(i, name, *args, **kwargs)
diff --git a/mitmproxy/console/master.py b/mitmproxy/console/master.py
index 1daf1127..00905f36 100644
--- a/mitmproxy/console/master.py
+++ b/mitmproxy/console/master.py
@@ -216,10 +216,9 @@ class ConsoleMaster(flow.FlowMaster):
     palette = []
 
     def __init__(self, server, options):
-        flow.FlowMaster.__init__(self, server, ConsoleState())
+        flow.FlowMaster.__init__(self, options, server, ConsoleState())
 
         self.stream_path = None
-        self.options = options
         self.options.errored.connect(self.options_error)
 
         if options.replacements:
diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py
index e2be3a53..d09038f8 100644
--- a/mitmproxy/controller.py
+++ b/mitmproxy/controller.py
@@ -6,6 +6,8 @@ import contextlib
 
 from six.moves import queue
 
+from mitmproxy import addons
+from mitmproxy import options
 from . import ctx as mitmproxy_ctx
 from netlib import basethread
 from . import exceptions
@@ -49,7 +51,9 @@ class Master(object):
     """
         The master handles mitmproxy's main event loop.
     """
-    def __init__(self, *servers):
+    def __init__(self, opts, *servers):
+        self.options = opts or options.Options()
+        self.addons = addons.Addons(self)
         self.event_queue = queue.Queue()
         self.should_exit = threading.Event()
         self.servers = []
@@ -121,6 +125,7 @@ class Master(object):
         for server in self.servers:
             server.shutdown()
         self.should_exit.set()
+        self.addons.done()
 
 
 class ServerThread(basethread.BaseThread):
@@ -191,6 +196,10 @@ def handler(f):
 
         with master.handlecontext():
             ret = f(master, message)
+            if handling:
+                # Python2/3 compatibility hack
+                fn = getattr(f, "func_name", None) or getattr(f, "__name__")
+                master.addons(fn)
 
         if handling and not message.reply.acked and not message.reply.taken:
             message.reply.ack()
diff --git a/mitmproxy/dump.py b/mitmproxy/dump.py
index bfefb319..cd515945 100644
--- a/mitmproxy/dump.py
+++ b/mitmproxy/dump.py
@@ -58,7 +58,7 @@ class Options(options.Options):
 class DumpMaster(flow.FlowMaster):
 
     def __init__(self, server, options, outfile=None):
-        flow.FlowMaster.__init__(self, server, flow.State())
+        flow.FlowMaster.__init__(self, options, server, flow.State())
         self.outfile = outfile
         self.o = options
         self.anticache = options.anticache
@@ -137,8 +137,8 @@ class DumpMaster(flow.FlowMaster):
                 self.add_event("Flow file corrupted.", "error")
                 raise DumpError(v)
 
-        if self.o.app:
-            self.start_app(self.o.app_host, self.o.app_port)
+        if self.options.app:
+            self.start_app(self.options.app_host, self.options.app_port)
 
     def _readflow(self, paths):
         """
@@ -152,7 +152,7 @@ class DumpMaster(flow.FlowMaster):
 
     def add_event(self, e, level="info"):
         needed = dict(error=0, info=1, debug=2).get(level, 1)
-        if self.o.verbosity >= needed:
+        if self.options.verbosity >= needed:
             self.echo(
                 e,
                 fg="red" if level == "error" else None,
@@ -172,7 +172,7 @@ class DumpMaster(flow.FlowMaster):
         click.secho(text, file=self.outfile, **style)
 
     def _echo_message(self, message):
-        if self.o.flow_detail >= 2 and hasattr(message, "headers"):
+        if self.options.flow_detail >= 2 and hasattr(message, "headers"):
             headers = "\r\n".join(
                 "{}: {}".format(
                     click.style(strutils.bytes_to_escaped_str(k), fg="blue", bold=True),
@@ -180,7 +180,7 @@ class DumpMaster(flow.FlowMaster):
                 for k, v in message.headers.fields
             )
             self.echo(headers, indent=4)
-        if self.o.flow_detail >= 3:
+        if self.options.flow_detail >= 3:
             if message.content is None:
                 self.echo("(content missing)", indent=4)
             elif message.content:
@@ -213,7 +213,7 @@ class DumpMaster(flow.FlowMaster):
                     for (style, text) in line:
                         yield click.style(text, **styles.get(style, {}))
 
-                if self.o.flow_detail == 3:
+                if self.options.flow_detail == 3:
                     lines_to_echo = itertools.islice(lines, 70)
                 else:
                     lines_to_echo = lines
@@ -228,7 +228,7 @@ class DumpMaster(flow.FlowMaster):
                 if next(lines, None):
                     self.echo("(cut off)", indent=4, dim=True)
 
-        if self.o.flow_detail >= 2:
+        if self.options.flow_detail >= 2:
             self.echo("")
 
     def _echo_request_line(self, flow):
@@ -302,7 +302,7 @@ class DumpMaster(flow.FlowMaster):
         self.echo(line)
 
     def echo_flow(self, f):
-        if self.o.flow_detail == 0:
+        if self.options.flow_detail == 0:
             return
 
         if f.request:
@@ -350,7 +350,7 @@ class DumpMaster(flow.FlowMaster):
     def tcp_message(self, f):
         super(DumpMaster, self).tcp_message(f)
 
-        if self.o.flow_detail == 0:
+        if self.options.flow_detail == 0:
             return
         message = f.messages[-1]
         direction = "->" if message.from_client else "<-"
@@ -362,7 +362,7 @@ class DumpMaster(flow.FlowMaster):
         self._echo_message(message)
 
     def run(self):  # pragma: no cover
-        if self.o.rfile and not self.o.keepserving:
+        if self.options.rfile and not self.options.keepserving:
             self.unload_scripts()  # make sure to trigger script unload events.
             return
         super(DumpMaster, self).run()
diff --git a/mitmproxy/flow/master.py b/mitmproxy/flow/master.py
index 7590a3fa..b1951f94 100644
--- a/mitmproxy/flow/master.py
+++ b/mitmproxy/flow/master.py
@@ -27,8 +27,8 @@ class FlowMaster(controller.Master):
         if len(self.servers) > 0:
             return self.servers[0]
 
-    def __init__(self, server, state):
-        super(FlowMaster, self).__init__()
+    def __init__(self, options, server, state):
+        super(FlowMaster, self).__init__(options)
         if server:
             self.add_server(server)
         self.state = state
diff --git a/mitmproxy/options.py b/mitmproxy/options.py
index 7389df1f..0cc5fee1 100644
--- a/mitmproxy/options.py
+++ b/mitmproxy/options.py
@@ -42,7 +42,10 @@ class Options(object):
         return self.__class__(**self._opts)
 
     def __getattr__(self, attr):
-        return self._opts[attr]
+        if attr in self._opts:
+            return self._opts[attr]
+        else:
+            raise AttributeError()
 
     def __setattr__(self, attr, value):
         if attr not in self._opts:
diff --git a/mitmproxy/web/master.py b/mitmproxy/web/master.py
index 008b74f8..22972c14 100644
--- a/mitmproxy/web/master.py
+++ b/mitmproxy/web/master.py
@@ -147,9 +147,10 @@ class Options(options.Options):
 class WebMaster(flow.FlowMaster):
 
     def __init__(self, server, options):
-        self.options = options
-        super(WebMaster, self).__init__(server, WebState())
-        self.app = app.Application(self, self.options.wdebug, self.options.wauthenticator)
+        super(WebMaster, self).__init__(options, server, WebState())
+        self.app = app.Application(
+            self, self.options.wdebug, self.options.wauthenticator
+        )
         if options.rfile:
             try:
                 self.load_flows_file(options.rfile)
diff --git a/test/mitmproxy/test_addons.py b/test/mitmproxy/test_addons.py
new file mode 100644
index 00000000..1861d4ac
--- /dev/null
+++ b/test/mitmproxy/test_addons.py
@@ -0,0 +1,20 @@
+from __future__ import absolute_import, print_function, division
+from mitmproxy import addons
+from mitmproxy import controller
+from mitmproxy import options
+
+
+class TAddon:
+    def __init__(self, name):
+        self.name = name
+
+    def __repr__(self):
+        return "Addon(%s)" % self.name
+
+
+def test_simple():
+    m = controller.Master(options.Options())
+    a = addons.Addons(m)
+    a.add(TAddon("one"))
+    assert a.has_addon("one")
+    assert not a.has_addon("two")
diff --git a/test/mitmproxy/test_controller.py b/test/mitmproxy/test_controller.py
index 5a68e15b..6d4b8fe6 100644
--- a/test/mitmproxy/test_controller.py
+++ b/test/mitmproxy/test_controller.py
@@ -25,7 +25,7 @@ class TestMaster(object):
                 # Speed up test
                 super(DummyMaster, self).tick(0)
 
-        m = DummyMaster()
+        m = DummyMaster(None)
         assert not m.should_exit.is_set()
         msg = TMsg()
         msg.reply = controller.DummyReply()
@@ -34,7 +34,7 @@ class TestMaster(object):
         assert m.should_exit.is_set()
 
     def test_server_simple(self):
-        m = controller.Master()
+        m = controller.Master(None)
         s = DummyServer(None)
         m.add_server(s)
         m.start()
diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py
index 74b3f599..eda01ad9 100644
--- a/test/mitmproxy/test_flow.py
+++ b/test/mitmproxy/test_flow.py
@@ -139,7 +139,7 @@ class TestClientPlaybackState:
     def test_tick(self):
         first = tutils.tflow()
         s = flow.State()
-        fm = flow.FlowMaster(None, s)
+        fm = flow.FlowMaster(None, None, s)
         fm.start_client_playback([first, tutils.tflow()], True)
         c = fm.client_playback
         c.testing = True
@@ -470,7 +470,7 @@ class TestFlow(object):
 
     def test_kill(self):
         s = flow.State()
-        fm = flow.FlowMaster(None, s)
+        fm = flow.FlowMaster(None, None, s)
         f = tutils.tflow()
         f.intercept(mock.Mock())
         f.kill(fm)
@@ -479,7 +479,7 @@ class TestFlow(object):
 
     def test_killall(self):
         s = flow.State()
-        fm = flow.FlowMaster(None, s)
+        fm = flow.FlowMaster(None, None, s)
 
         f = tutils.tflow()
         f.intercept(fm)
@@ -714,7 +714,7 @@ class TestSerialize:
     def test_load_flows(self):
         r = self._treader()
         s = flow.State()
-        fm = flow.FlowMaster(None, s)
+        fm = flow.FlowMaster(None, None, s)
         fm.load_flows(r)
         assert len(s.flows) == 6
 
@@ -725,7 +725,7 @@ class TestSerialize:
             mode="reverse",
             upstream_server=("https", ("use-this-domain", 80))
         )
-        fm = flow.FlowMaster(DummyServer(conf), s)
+        fm = flow.FlowMaster(None, DummyServer(conf), s)
         fm.load_flows(r)
         assert s.flows[0].request.host == "use-this-domain"
 
@@ -772,7 +772,7 @@ class TestFlowMaster:
 
     def test_load_script(self):
         s = flow.State()
-        fm = flow.FlowMaster(None, s)
+        fm = flow.FlowMaster(None, None, s)
 
         fm.load_script(tutils.test_data.path("data/scripts/a.py"))
         fm.load_script(tutils.test_data.path("data/scripts/a.py"))
@@ -788,14 +788,14 @@ class TestFlowMaster:
     def test_getset_ignore(self):
         p = mock.Mock()
         p.config.check_ignore = HostMatcher()
-        fm = flow.FlowMaster(p, flow.State())
+        fm = flow.FlowMaster(None, p, flow.State())
         assert not fm.get_ignore_filter()
         fm.set_ignore_filter(["^apple\.com:", ":443$"])
         assert fm.get_ignore_filter()
 
     def test_replay(self):
         s = flow.State()
-        fm = flow.FlowMaster(None, s)
+        fm = flow.FlowMaster(None, None, s)
         f = tutils.tflow(resp=True)
         f.request.content = None
         assert "missing" in fm.replay_request(f)
@@ -808,7 +808,7 @@ class TestFlowMaster:
 
     def test_script_reqerr(self):
         s = flow.State()
-        fm = flow.FlowMaster(None, s)
+        fm = flow.FlowMaster(None, None, s)
         fm.load_script(tutils.test_data.path("data/scripts/reqerr.py"))
         f = tutils.tflow()
         fm.clientconnect(f.client_conn)
@@ -816,7 +816,7 @@ class TestFlowMaster:
 
     def test_script(self):
         s = flow.State()
-        fm = flow.FlowMaster(None, s)
+        fm = flow.FlowMaster(None, None, s)
         fm.load_script(tutils.test_data.path("data/scripts/all.py"))
         f = tutils.tflow(resp=True)
 
@@ -852,7 +852,7 @@ class TestFlowMaster:
 
     def test_duplicate_flow(self):
         s = flow.State()
-        fm = flow.FlowMaster(None, s)
+        fm = flow.FlowMaster(None, None, s)
         f = tutils.tflow(resp=True)
         fm.load_flow(f)
         assert s.flow_count() == 1
@@ -863,12 +863,12 @@ class TestFlowMaster:
 
     def test_create_flow(self):
         s = flow.State()
-        fm = flow.FlowMaster(None, s)
+        fm = flow.FlowMaster(None, None, s)
         assert fm.create_request("GET", "http", "example.com", 80, "/")
 
     def test_all(self):
         s = flow.State()
-        fm = flow.FlowMaster(None, s)
+        fm = flow.FlowMaster(None, None, s)
         fm.anticache = True
         fm.anticomp = True
         f = tutils.tflow(req=None)
@@ -895,7 +895,7 @@ class TestFlowMaster:
 
         f = tutils.tflow(resp=True)
         pb = [tutils.tflow(resp=True), f]
-        fm = flow.FlowMaster(DummyServer(ProxyConfig()), s)
+        fm = flow.FlowMaster(None, DummyServer(ProxyConfig()), s)
         assert not fm.start_server_playback(
             pb,
             False,
@@ -923,7 +923,7 @@ class TestFlowMaster:
         f.response = HTTPResponse.wrap(netlib.tutils.tresp(content=f.request))
         pb = [f]
 
-        fm = flow.FlowMaster(None, s)
+        fm = flow.FlowMaster(None, None, s)
         fm.refresh_server_playback = True
         assert not fm.do_server_playback(tutils.tflow())
 
@@ -965,7 +965,7 @@ class TestFlowMaster:
         f = tutils.tflow()
         f.response = HTTPResponse.wrap(netlib.tutils.tresp(content=f.request))
         pb = [f]
-        fm = flow.FlowMaster(None, s)
+        fm = flow.FlowMaster(None, None, s)
         fm.refresh_server_playback = True
         fm.start_server_playback(
             pb,
@@ -985,7 +985,7 @@ class TestFlowMaster:
 
     def test_stickycookie(self):
         s = flow.State()
-        fm = flow.FlowMaster(None, s)
+        fm = flow.FlowMaster(None, None, s)
         assert "Invalid" in fm.set_stickycookie("~h")
         fm.set_stickycookie(".*")
         assert fm.stickycookie_state
@@ -1007,7 +1007,7 @@ class TestFlowMaster:
 
     def test_stickyauth(self):
         s = flow.State()
-        fm = flow.FlowMaster(None, s)
+        fm = flow.FlowMaster(None, None, s)
         assert "Invalid" in fm.set_stickyauth("~h")
         fm.set_stickyauth(".*")
         assert fm.stickyauth_state
@@ -1035,7 +1035,7 @@ class TestFlowMaster:
                     return list(r.stream())
 
             s = flow.State()
-            fm = flow.FlowMaster(None, s)
+            fm = flow.FlowMaster(None, None, s)
             f = tutils.tflow(resp=True)
 
             with open(p, "ab") as tmpfile:
diff --git a/test/mitmproxy/test_script.py b/test/mitmproxy/test_script.py
index 81994780..1e8220f1 100644
--- a/test/mitmproxy/test_script.py
+++ b/test/mitmproxy/test_script.py
@@ -4,7 +4,7 @@ from . import tutils
 
 def test_duplicate_flow():
     s = flow.State()
-    fm = flow.FlowMaster(None, s)
+    fm = flow.FlowMaster(None, None, s)
     fm.load_script(tutils.test_data.path("data/scripts/duplicate_flow.py"))
     f = tutils.tflow()
     fm.request(f)
diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py
index 0760cb53..69a50b9d 100644
--- a/test/mitmproxy/tservers.py
+++ b/test/mitmproxy/tservers.py
@@ -34,7 +34,7 @@ class TestMaster(flow.FlowMaster):
         config.port = 0
         s = ProxyServer(config)
         state = flow.State()
-        flow.FlowMaster.__init__(self, s, state)
+        flow.FlowMaster.__init__(self, None, s, state)
         self.apps.add(testapp, "testapp", 80)
         self.apps.add(errapp, "errapp", 80)
         self.clear_log()