Merge pull request #1348 from cortesi/addons

Addons

18 files changed, 194 insertions, 51 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/builtins/__init__.py b/mitmproxy/builtins/__init__.py
new file mode 100644
index 00000000..59173145
--- /dev/null
+++ b/mitmproxy/builtins/__init__.py
@@ -0,0 +1,9 @@
+from __future__ import absolute_import, print_function, division
+
+from mitmproxy.builtins import anticomp
+
+
+def default_addons():
+    return [
+        anticomp.AntiComp(),
+    ]
diff --git a/mitmproxy/builtins/anticomp.py b/mitmproxy/builtins/anticomp.py
new file mode 100644
index 00000000..50bd1b73
--- /dev/null
+++ b/mitmproxy/builtins/anticomp.py
@@ -0,0 +1,13 @@
+from __future__ import absolute_import, print_function, division
+
+
+class AntiComp:
+    def __init__(self):
+        self.enabled = False
+
+    def configure(self, options):
+        self.enabled = options.anticomp
+
+    def request(self, flow):
+        if self.enabled:
+            flow.request.anticomp()
diff --git a/mitmproxy/console/master.py b/mitmproxy/console/master.py
index 1daf1127..06a05464 100644
--- a/mitmproxy/console/master.py
+++ b/mitmproxy/console/master.py
@@ -15,6 +15,7 @@ import weakref
 
 import urwid
 
+from mitmproxy import builtins
 from mitmproxy import contentviews
 from mitmproxy import controller
 from mitmproxy import exceptions
@@ -216,10 +217,10 @@ class ConsoleMaster(flow.FlowMaster):
     palette = []
 
     def __init__(self, server, options):
-        flow.FlowMaster.__init__(self, server, ConsoleState())
+        flow.FlowMaster.__init__(self, options, server, ConsoleState())
+        self.addons.add(*builtins.default_addons())
 
         self.stream_path = None
-        self.options = options
         self.options.errored.connect(self.options_error)
 
         if options.replacements:
@@ -252,7 +253,6 @@ class ConsoleMaster(flow.FlowMaster):
 
         self.refresh_server_playback = options.refresh_server_playback
         self.anticache = options.anticache
-        self.anticomp = options.anticomp
         self.killextra = options.kill
         self.rheaders = options.rheaders
         self.nopop = options.nopop
diff --git a/mitmproxy/console/options.py b/mitmproxy/console/options.py
index 6a4b8dd6..db6d405a 100644
--- a/mitmproxy/console/options.py
+++ b/mitmproxy/console/options.py
@@ -102,7 +102,7 @@ class Options(urwid.WidgetWrap):
                 select.Option(
                     "Anti-Compression",
                     "o",
-                    lambda: master.anticomp,
+                    lambda: master.options.anticomp,
                     self.toggle_anticomp
                 ),
                 select.Option(
@@ -177,7 +177,7 @@ class Options(urwid.WidgetWrap):
         self.master.anticache = not self.master.anticache
 
     def toggle_anticomp(self):
-        self.master.anticomp = not self.master.anticomp
+        self.master.options.anticomp = not self.master.options.anticomp
 
     def toggle_killextra(self):
         self.master.killextra = not self.master.killextra
diff --git a/mitmproxy/console/statusbar.py b/mitmproxy/console/statusbar.py
index d1ab5906..543c6771 100644
--- a/mitmproxy/console/statusbar.py
+++ b/mitmproxy/console/statusbar.py
@@ -189,7 +189,7 @@ class StatusBar(urwid.WidgetWrap):
         opts = []
         if self.master.anticache:
             opts.append("anticache")
-        if self.master.anticomp:
+        if self.master.options.anticomp:
             opts.append("anticomp")
         if self.master.showhost:
             opts.append("showhost")
diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py
index e2be3a53..8b968eb5 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, message)
 
         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..cf7aa5c9 100644
--- a/mitmproxy/dump.py
+++ b/mitmproxy/dump.py
@@ -12,6 +12,7 @@ from mitmproxy import exceptions
 from mitmproxy import filt
 from mitmproxy import flow
 from mitmproxy import options
+from mitmproxy import builtins
 from netlib import human
 from netlib import tcp
 from netlib import strutils
@@ -58,7 +59,8 @@ 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.addons.add(*builtins.default_addons())
         self.outfile = outfile
         self.o = options
         self.anticache = options.anticache
@@ -137,8 +139,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 +154,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 +174,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 +182,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 +215,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 +230,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 +304,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 +352,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 +364,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..ed2ee138 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
@@ -46,7 +46,6 @@ class FlowMaster(controller.Master):
         self.stickyauth_txt = None
 
         self.anticache = False
-        self.anticomp = False
         self.stream_large_bodies = None  # type: Optional[modules.StreamLargeBodies]
         self.refresh_server_playback = False
         self.replacehooks = modules.ReplaceHooks()
@@ -332,8 +331,6 @@ class FlowMaster(controller.Master):
 
         if self.anticache:
             f.request.anticache()
-        if self.anticomp:
-            f.request.anticomp()
 
         if self.server_playback:
             pb = self.do_server_playback(f)
diff --git a/mitmproxy/web/master.py b/mitmproxy/web/master.py
index 008b74f8..7c775c32 100644
--- a/mitmproxy/web/master.py
+++ b/mitmproxy/web/master.py
@@ -6,6 +6,7 @@ import collections
 import tornado.httpserver
 import tornado.ioloop
 
+from mitmproxy import builtins
 from mitmproxy import controller
 from mitmproxy import exceptions
 from mitmproxy import flow
@@ -147,9 +148,11 @@ 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.addons.add(*builtins.default_addons())
+        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/builtins/__init__.py b/test/mitmproxy/builtins/__init__.py
new file mode 100644
index 00000000..e69de29b
--- /dev/null
+++ b/test/mitmproxy/builtins/__init__.py
diff --git a/test/mitmproxy/builtins/test_anticomp.py b/test/mitmproxy/builtins/test_anticomp.py
new file mode 100644
index 00000000..6bfd54bb
--- /dev/null
+++ b/test/mitmproxy/builtins/test_anticomp.py
@@ -0,0 +1,22 @@
+from .. import tutils, mastertest
+from mitmproxy.builtins import anticomp
+from mitmproxy.flow import master
+from mitmproxy.flow import state
+from mitmproxy import options
+
+
+class TestAntiComp(mastertest.MasterTest):
+    def test_simple(self):
+        s = state.State()
+        m = master.FlowMaster(options.Options(anticomp = True), None, s)
+        sa = anticomp.AntiComp()
+        m.addons.add(sa)
+
+        f = tutils.tflow(resp=True)
+        self.invoke(m, "request", f)
+
+        f = tutils.tflow(resp=True)
+
+        f.request.headers["Accept-Encoding"] = "foobar"
+        self.invoke(m, "request", f)
+        assert f.request.headers["Accept-Encoding"] == "identity"
diff --git a/test/mitmproxy/mastertest.py b/test/mitmproxy/mastertest.py
index 9e726a32..06854e25 100644
--- a/test/mitmproxy/mastertest.py
+++ b/test/mitmproxy/mastertest.py
@@ -3,10 +3,16 @@ import mock
 from . import tutils
 import netlib.tutils
 
-from mitmproxy import flow, proxy, models
+from mitmproxy import flow, proxy, models, controller
 
 
 class MasterTest:
+    def invoke(self, master, handler, message):
+        with master.handlecontext():
+            func = getattr(master, handler)
+            func(message)
+        message.reply = controller.DummyReply()
+
     def cycle(self, master, content):
         f = tutils.tflow(req=netlib.tutils.treq(content=content))
         l = proxy.Log("connect")
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..6b270872 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,14 +863,13 @@ 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)
         fm.clientconnect(f.client_conn)
         f.request = HTTPRequest.wrap(netlib.tutils.treq())
@@ -895,7 +894,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 +922,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 +964,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 +984,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 +1006,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 +1034,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()