Enable custom options for addons

- Add an options parameter to the start() event. This is to be used by addons
on startup to add custom options.
- Add a running() event that is called once the proxy is up and running.
- With the new paradigm we can't log during master __init__, so add a tiny
termstatus addon to print proxy status to terminal once we're running.

examples/complex/dns_spoofing.py

@@ -1,11 +1,12 @@
 """
-This script makes it possible to use mitmproxy in scenarios where IP spoofing has been used to redirect
-connections to mitmproxy. The way this works is that we rely on either the TLS Server Name Indication (SNI) or the
-Host header of the HTTP request.
-Of course, this is not foolproof - if an HTTPS connection comes without SNI, we don't
-know the actual target and cannot construct a certificate that looks valid.
-Similarly, if there's no Host header or a spoofed Host header, we're out of luck as well.
-Using transparent mode is the better option most of the time.
+This script makes it possible to use mitmproxy in scenarios where IP spoofing
+has been used to redirect connections to mitmproxy. The way this works is that
+we rely on either the TLS Server Name Indication (SNI) or the Host header of the
+HTTP request. Of course, this is not foolproof - if an HTTPS connection comes
+without SNI, we don't know the actual target and cannot construct a certificate
+that looks valid. Similarly, if there's no Host header or a spoofed Host header,
+we're out of luck as well. Using transparent mode is the better option most of
+the time.
 
 Usage:
     mitmproxy
@@ -53,5 +54,5 @@ class Rerouter:
         flow.request.port = port
 
 
-def start():
+def start(opts):
     return Rerouter()

examples/complex/har_dump.py

@@ -25,7 +25,7 @@ HAR = {}
 SERVERS_SEEN = set()
 
 
-def start():
+def start(opts):
     """
         Called once on script startup before any other events.
     """

examples/complex/remote_debug.py

@@ -14,6 +14,6 @@ Usage:
 """
 
 
-def start():
+def start(opts):
     import pydevd
     pydevd.settrace("localhost", port=5678, stdoutToServer=True, stderrToServer=True)

examples/complex/tls_passthrough.py

@@ -112,7 +112,7 @@ class TlsFeedback(TlsLayer):
 tls_strategy = None
 
 
-def start():
+def start(opts):
     global tls_strategy
     if len(sys.argv) == 2:
         tls_strategy = ProbabilisticStrategy(float(sys.argv[1]))

examples/simple/add_header_class.py

@@ -3,5 +3,5 @@ class AddHeader:
         flow.response.headers["newheader"] = "foo"
 
 
-def start():
+def start(opts):
     return AddHeader()

examples/simple/custom_contentview.py

@@ -20,7 +20,7 @@ class ViewSwapCase(contentviews.View):
 view = ViewSwapCase()
 
 
-def start():
+def start(opts):
     contentviews.add(view)
 
 

examples/simple/custom_option.py

@@ -0,0 +1,10 @@
+from mitmproxy import ctx
+
+
+def start(options):
+    ctx.log.info("Registering option 'custom'")
+    options.add_option("custom", str, "default", "A custom option")
+
+
+def configure(options, updated):
+    ctx.log.info("custom option value: %s" % options.custom)

examples/simple/filter_flows.py

@@ -17,7 +17,7 @@ class Filter:
             print(flow)
 
 
-def start():
+def start(opts):
     if len(sys.argv) != 2:
         raise ValueError("Usage: -s 'filt.py FILTER'")
     return Filter(sys.argv[1])

examples/simple/io_write_dumpfile.py

@@ -23,7 +23,7 @@ class Writer:
             self.w.add(flow)
 
 
-def start():
+def start(opts):
     if len(sys.argv) != 2:
         raise ValueError('Usage: -s "flowriter.py filename"')
     return Writer(sys.argv[1])

examples/simple/log_events.py

@@ -7,6 +7,6 @@ If you want to help us out: https://github.com/mitmproxy/mitmproxy/issues/1530 :
 from mitmproxy import ctx
 
 
-def start():
+def start(opts):
     ctx.log.info("This is some informative text.")
     ctx.log.error("This is an error.")

examples/simple/modify_body_inject_iframe.py

@@ -23,7 +23,7 @@ class Injector:
             flow.response.content = str(html).encode("utf8")
 
 
-def start():
+def start(opts):
     if len(sys.argv) != 2:
         raise ValueError('Usage: -s "iframe_injector.py url"')
     return Injector(sys.argv[1])

examples/simple/script_arguments.py

@@ -9,7 +9,7 @@ class Replacer:
         flow.response.replace(self.src, self.dst)
 
 
-def start():
+def start(opts):
     parser = argparse.ArgumentParser()
     parser.add_argument("src", type=str)
     parser.add_argument("dst", type=str)

examples/simple/wsgi_flask_app.py

@@ -14,7 +14,7 @@ def hello_world():
     return 'Hello World!'
 
 
-def start():
+def start(opts):
     # Host app at the magic domain "proxapp" on port 80. Requests to this
     # domain and port combination will now be routed to the WSGI app instance.
     return wsgiapp.WSGIApp(app, "proxapp", 80)

mitmproxy/addonmanager.py

@@ -1,4 +1,5 @@
 from mitmproxy import exceptions
+from mitmproxy import eventsequence
 import pprint
 
 
@@ -10,7 +11,7 @@ class AddonManager:
     def __init__(self, master):
         self.chain = []
         self.master = master
-        master.options.changed.connect(self._options_update)
+        master.options.changed.connect(self.configure_all)
 
     def clear(self):
         """
@@ -29,22 +30,14 @@ class AddonManager:
             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 configure_all(self, options, updated):
+        self.invoke_all_with_context("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()
-        )
+        self.invoke_with_context(s, "start", self.master.options)
 
     def add(self, *addons):
         """
@@ -62,8 +55,7 @@ class AddonManager:
         self.invoke_with_context(addon, "done")
 
     def done(self):
-        for i in self.chain:
-            self.invoke_with_context(i, "done")
+        self.invoke_all_with_context("done")
 
     def __len__(self):
         return len(self.chain)
@@ -75,7 +67,14 @@ class AddonManager:
         with self.master.handlecontext():
             self.invoke(addon, name, *args, **kwargs)
 
+    def invoke_all_with_context(self, name, *args, **kwargs):
+        with self.master.handlecontext():
+            for i in self.chain:
+                self.invoke(i, name, *args, **kwargs)
+
     def invoke(self, addon, name, *args, **kwargs):
+        if name not in eventsequence.Events:  # prama: no cover
+            raise NotImplementedError("Unknown event")
         func = getattr(addon, name, None)
         if func:
             if not callable(func):

mitmproxy/addons/script.py

@@ -170,22 +170,23 @@ class Script:
 
     def load_script(self):
         self.ns = load_script(self.path, self.args)
-        ret = self.run("start")
+        ret = self.run("start", self.last_options)
         if ret:
             self.ns = ret
-            self.run("start")
+            self.run("start", self.last_options)
 
     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.start()
+            self.start(self.last_options)
             self.configure(self.last_options, self.last_options.keys())
         else:
             self.run("tick")
 
-    def start(self):
+    def start(self, opts):
+        self.last_options = opts
         self.load_script()
 
     def configure(self, options, updated):
@@ -209,6 +210,12 @@ class ScriptLoader:
     """
         An addon that manages loading scripts from options.
     """
+    def __init__(self):
+        self.is_running = False
+
+    def running(self):
+        self.is_running = True
+
     def run_once(self, command, flows):
         try:
             sc = Script(command)
@@ -267,3 +274,10 @@ class ScriptLoader:
 
             for s in newscripts:
                 ctx.master.addons.startup(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_with_context(
+                        s, "configure", options, options.keys()
+                    )
+                    ctx.master.addons.invoke_with_context(s, "running")

mitmproxy/addons/termstatus.py

@@ -0,0 +1,23 @@
+from mitmproxy import ctx
+
+"""
+    A tiny addon to print the proxy status to terminal. Eventually this could
+    also print some stats on exit.
+"""
+
+
+class TermStatus:
+    def __init__(self):
+        self.server = False
+
+    def configure(self, options, updated):
+        if "server" in updated:
+            self.server = options.server
+
+    def running(self):
+        if self.server:
+            ctx.log.info(
+                "Proxy server listening at http://{}:{}".format(
+                    *ctx.master.server.address,
+                )
+            )

mitmproxy/eventsequence.py

@@ -33,6 +33,7 @@ Events = frozenset([
     "done",
     "log",
     "start",
+    "running",
     "tick",
 ])
 

mitmproxy/master.py

@@ -76,12 +76,16 @@ class Master:
 
     def run(self):
         self.start()
+        running = False
         try:
             while not self.should_exit.is_set():
                 # Don't choose a very small timeout in Python 2:
                 # https://github.com/mitmproxy/mitmproxy/issues/443
                 # TODO: Lower the timeout value if we move to Python 3.
                 self.tick(0.1)
+                if not running:
+                    running = True
+                    self.addons.invoke_all_with_context("running")
         finally:
             self.shutdown()
 

mitmproxy/optmanager.py

@@ -324,7 +324,15 @@ class OptManager:
             options=options
         )
 
-    def set(self, spec):
+    def set(self, *spec):
+        vals = {}
+        for i in spec:
+            vals.update(self._setspec(i))
+        self.update(**vals)
+
+    def _setspec(self, spec):
+        d = {}
+
         parts = spec.split("=", maxsplit=1)
         if len(parts) == 1:
             optname, optval = parts[0], None
@@ -333,14 +341,14 @@ class OptManager:
         o = self._options[optname]
 
         if o.typespec in (str, typing.Optional[str]):
-            setattr(self, optname, optval)
+            d[optname] = optval
         elif o.typespec in (int, typing.Optional[int]):
             if optval:
                 try:
                     optval = int(optval)
                 except ValueError:
                     raise exceptions.OptionsError("Not an integer: %s" % optval)
-            setattr(self, optname, optval)
+            d[optname] = optval
         elif o.typespec == bool:
             if not optval or optval == "true":
                 v = True
@@ -350,18 +358,15 @@ class OptManager:
                 raise exceptions.OptionsError(
                     "Boolean must be \"true\", \"false\", or have the value " "omitted (a synonym for \"true\")."
                 )
-            setattr(self, optname, v)
+            d[optname] = v
         elif o.typespec == typing.Sequence[str]:
             if not optval:
-                setattr(self, optname, [])
+                d[optname] = []
             else:
-                setattr(
-                    self,
-                    optname,
-                    getattr(self, optname) + [optval]
-                )
+                d[optname] = getattr(self, optname) + [optval]
         else:  # pragma: no cover
             raise NotImplementedError("Unsupported option type: %s", o.typespec)
+        return d
 
     def make_parser(self, parser, optname, metavar=None, short=None):
         o = self._options[optname]

mitmproxy/tools/dump.py

@@ -3,7 +3,7 @@ from mitmproxy import exceptions
 from mitmproxy import addons
 from mitmproxy import options
 from mitmproxy import master
-from mitmproxy.addons import dumper, termlog
+from mitmproxy.addons import dumper, termlog, termstatus
 
 
 class DumpMaster(master.Master):
@@ -18,17 +18,11 @@ class DumpMaster(master.Master):
         master.Master.__init__(self, options, server)
         self.has_errored = False
         if with_termlog:
-            self.addons.add(termlog.TermLog())
+            self.addons.add(termlog.TermLog(), termstatus.TermStatus())
         self.addons.add(*addons.default_addons())
         if with_dumper:
             self.addons.add(dumper.Dumper())
 
-        if self.options.server:
-            self.add_log(
-                "Proxy server listening at http://{}:{}".format(server.address[0], server.address[1]),
-                "info"
-            )
-
         if options.rfile:
             try:
                 self.load_flows_file(options.rfile)

mitmproxy/tools/main.py

@@ -45,9 +45,6 @@ def process_options(parser, opts, args):
     if args.quiet:
         args.flow_detail = 0
 
-    for i in args.setoptions:
-        opts.set(i)
-
     adict = {}
     for n in dir(args):
         if n in opts:
@@ -77,6 +74,8 @@ def run(MasterKlass, args):  # pragma: no cover
         opts.load_paths(args.conf)
         server = process_options(parser, opts, args)
         master = MasterKlass(opts, server)
+        master.addons.configure_all(opts, opts.keys())
+        opts.set(*args.setoptions)
 
         def cleankill(*args, **kwargs):
             master.shutdown()

test/mitmproxy/addons/test_script.py

@@ -116,9 +116,7 @@ class TestScript:
                 )
             )
             sc.load_script()
-            assert sc.ns.call_log == [
-                ("solo", "start", (), {}),
-            ]
+            assert sc.ns.call_log[0][0:2] == ("solo", "start")
 
             sc.ns.call_log = []
             f = tflow.tflow(resp=True)
@@ -146,7 +144,7 @@ class TestScript:
             sc = script.Script(
                 tutils.test_data.path("mitmproxy/data/addonscripts/error.py")
             )
-            sc.start()
+            sc.start(tctx.options)
             f = tflow.tflow(resp=True)
             sc.request(f)
             assert tctx.master.event_log[0][0] == "error"
@@ -162,7 +160,7 @@ class TestScript:
                     "mitmproxy/data/addonscripts/addon.py"
                 )
             )
-            sc.start()
+            sc.start(tctx.options)
             tctx.configure(sc)
             assert sc.ns.event_log == [
                 'scriptstart', 'addonstart', 'addonconfigure'
@@ -225,24 +223,31 @@ class TestScriptLoader:
         assert len(m.addons) == 1
 
     def test_dupes(self):
-        o = options.Options(scripts=["one", "one"])
-        m = master.Master(o, proxy.DummyServer())
         sc = script.ScriptLoader()
+        with taddons.context() as tctx:
+            tctx.master.addons.add(sc)
             with pytest.raises(exceptions.OptionsError):
-            m.addons.add(o, sc)
+                tctx.configure(
+                    sc,
+                    scripts = ["one", "one"]
+                )
 
     def test_nonexistent(self):
-        o = options.Options(scripts=["nonexistent"])
-        m = master.Master(o, proxy.DummyServer())
         sc = script.ScriptLoader()
+        with taddons.context() as tctx:
+            tctx.master.addons.add(sc)
             with pytest.raises(exceptions.OptionsError):
-            m.addons.add(o, sc)
+                tctx.configure(
+                    sc,
+                    scripts = ["nonexistent"]
+                )
 
     def test_order(self):
         rec = tutils.test_data.path("mitmproxy/data/addonscripts/recorder.py")
         sc = script.ScriptLoader()
         with taddons.context() as tctx:
             tctx.master.addons.add(sc)
+            sc.running()
             tctx.configure(
                 sc,
                 scripts = [
@@ -253,9 +258,17 @@ class TestScriptLoader:
             )
             debug = [(i[0], i[1]) for i in tctx.master.event_log if i[0] == "debug"]
             assert debug == [
-                ('debug', 'a start'), ('debug', 'a configure'),
-                ('debug', 'b start'), ('debug', 'b configure'),
-                ('debug', 'c start'), ('debug', 'c configure')
+                ('debug', 'a start'),
+                ('debug', 'a configure'),
+                ('debug', 'a running'),
+
+                ('debug', 'b start'),
+                ('debug', 'b configure'),
+                ('debug', 'b running'),
+
+                ('debug', 'c start'),
+                ('debug', 'c configure'),
+                ('debug', 'c running'),
             ]
             tctx.master.event_log = []
             tctx.configure(
@@ -284,4 +297,5 @@ class TestScriptLoader:
                 ('debug', 'b done'),
                 ('debug', 'x start'),
                 ('debug', 'x configure'),
+                ('debug', 'x running'),
             ]

test/mitmproxy/addons/test_termstatus.py

@@ -0,0 +1,12 @@
+from mitmproxy.addons import termstatus
+from mitmproxy.test import taddons
+
+
+def test_configure():
+    ts = termstatus.TermStatus()
+    with taddons.context() as ctx:
+        ts.running()
+        assert not ctx.master.event_log
+        ctx.configure(ts, server=True)
+        ts.running()
+        assert ctx.master.event_log

test/mitmproxy/data/addonscripts/addon.py

@@ -6,7 +6,7 @@ class Addon:
     def event_log(self):
         return event_log
 
-    def start(self):
+    def start(self, opts):
         event_log.append("addonstart")
 
     def configure(self, options, updated):
@@ -17,6 +17,6 @@ def configure(options, updated):
     event_log.append("addonconfigure")
 
 
-def start():
+def start(opts):
     event_log.append("scriptstart")
     return Addon()

test/mitmproxy/data/addonscripts/concurrent_decorator_class.py

@@ -9,5 +9,5 @@ class ConcurrentClass:
         time.sleep(0.1)
 
 
-def start():
+def start(opts):
     return ConcurrentClass()

test/mitmproxy/data/addonscripts/concurrent_decorator_err.py

@@ -2,5 +2,5 @@ from mitmproxy.script import concurrent
 
 
 @concurrent
-def start():
+def start(opts):
     pass

test/mitmproxy/data/addonscripts/recorder.py

@@ -22,5 +22,5 @@ class CallLogger:
         raise AttributeError
 
 
-def start():
+def start(opts):
     return CallLogger(*sys.argv[1:])

test/mitmproxy/proxy/test_server.py

@@ -302,6 +302,9 @@ class TestHTTP(tservers.HTTPProxyTest, CommonMixin):
 class TestHTTPAuth(tservers.HTTPProxyTest):
     def test_auth(self):
         self.master.addons.add(proxyauth.ProxyAuth())
+        self.master.addons.configure_all(
+            self.master.options, self.master.options.keys()
+        )
         self.master.options.proxyauth = "test:test"
         assert self.pathod("202").status_code == 407
         p = self.pathoc()

test/mitmproxy/script/test_concurrent.py

@@ -24,7 +24,7 @@ class TestConcurrent(tservers.MasterTest):
                     "mitmproxy/data/addonscripts/concurrent_decorator.py"
                 )
             )
-            sc.start()
+            sc.start(tctx.options)
 
             f1, f2 = tflow.tflow(), tflow.tflow()
             tctx.cycle(sc, f1)
@@ -42,7 +42,7 @@ class TestConcurrent(tservers.MasterTest):
                     "mitmproxy/data/addonscripts/concurrent_decorator_err.py"
                 )
             )
-            sc.start()
+            sc.start(tctx.options)
             assert "decorator not supported" in tctx.master.event_log[0][1]
 
     def test_concurrent_class(self):
@@ -52,7 +52,7 @@ class TestConcurrent(tservers.MasterTest):
                         "mitmproxy/data/addonscripts/concurrent_decorator_class.py"
                     )
                 )
-                sc.start()
+                sc.start(tctx.options)
 
                 f1, f2 = tflow.tflow(), tflow.tflow()
                 tctx.cycle(sc, f1)

test/mitmproxy/test_addonmanager.py

@@ -10,12 +10,12 @@ from mitmproxy import proxy
 class TAddon:
     def __init__(self, name):
         self.name = name
-        self.noop_member = True
+        self.tick = True
 
     def __repr__(self):
         return "Addon(%s)" % self.name
 
-    def noop(self):
+    def done(self):
         pass
 
 
@@ -30,6 +30,6 @@ def test_simple():
     assert not a.chain
 
     a.add(TAddon("one"))
-    a("noop")
+    a("done")
     with pytest.raises(exceptions.AddonError):
-        a("noop_member")
+        a("tick")

test/mitmproxy/tools/console/test_master.py

@@ -28,7 +28,9 @@ class TestMaster(tservers.MasterTest):
         if "verbosity" not in opts:
             opts["verbosity"] = 1
         o = options.Options(**opts)
-        return console.master.ConsoleMaster(o, proxy.DummyServer())
+        m = console.master.ConsoleMaster(o, proxy.DummyServer())
+        m.addons.configure_all(o, o.keys())
+        return m
 
     def test_basic(self):
         m = self.mkmaster()

test/mitmproxy/tservers.py

@@ -79,6 +79,8 @@ class TestMaster(master.Master):
         self.state = TestState()
         self.addons.add(self.state)
         self.addons.add(*addons)
+        self.addons.configure_all(self.options, self.options.keys())
+        self.addons.invoke_all_with_context("running")
 
     def clear_log(self):
         self.tlog = []