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.

8 files changed, 78 insertions, 39 deletions

diff --git a/mitmproxy/addonmanager.py b/mitmproxy/addonmanager.py
index db8e0cd7..43e76510 100644
--- a/mitmproxy/addonmanager.py
+++ b/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):
diff --git a/mitmproxy/addons/script.py b/mitmproxy/addons/script.py
index a7d3a312..cfbe5284 100644
--- a/mitmproxy/addons/script.py
+++ b/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")
diff --git a/mitmproxy/addons/termstatus.py b/mitmproxy/addons/termstatus.py
new file mode 100644
index 00000000..7b05f409
--- /dev/null
+++ b/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,
+                )
+            )
diff --git a/mitmproxy/eventsequence.py b/mitmproxy/eventsequence.py
index 5872f607..bc6660e0 100644
--- a/mitmproxy/eventsequence.py
+++ b/mitmproxy/eventsequence.py
@@ -33,6 +33,7 @@ Events = frozenset([
     "done",
     "log",
     "start",
+    "running",
     "tick",
 ])
 
diff --git a/mitmproxy/master.py b/mitmproxy/master.py
index 8855452c..88c64a9f 100644
--- a/mitmproxy/master.py
+++ b/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()
 
diff --git a/mitmproxy/optmanager.py b/mitmproxy/optmanager.py
index 9553bd32..698c9bb8 100644
--- a/mitmproxy/optmanager.py
+++ b/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]
diff --git a/mitmproxy/tools/dump.py b/mitmproxy/tools/dump.py
index e70ce2f9..4bfe2dc4 100644
--- a/mitmproxy/tools/dump.py
+++ b/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)
diff --git a/mitmproxy/tools/main.py b/mitmproxy/tools/main.py
index 17c1abbb..7b2d2a7b 100644
--- a/mitmproxy/tools/main.py
+++ b/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()