Script cleanups

- Preserve script order on config change
- Prohibit script duplicates (i.e. identical script + args)
- Various cleanups and tweaks

3 files changed, 100 insertions, 18 deletions

diff --git a/mitmproxy/builtins/script.py b/mitmproxy/builtins/script.py
index bcb756f6..14f7dd5f 100644
--- a/mitmproxy/builtins/script.py
+++ b/mitmproxy/builtins/script.py
@@ -99,23 +99,25 @@ class Script:
         self.path, self.args = parse_command(command)
         self.ns = None
         self.observer = None
+        self.dead = False
 
         self.last_options = None
         self.should_reload = threading.Event()
 
-        for i in controller.Events - set(["start", "configure", "tick"]):
-            def mkprox():
-                evt = i
+        for i in controller.Events:
+            if not hasattr(self, i):
+                def mkprox():
+                    evt = i
 
-                def prox(*args, **kwargs):
-                    self.run(evt, *args, **kwargs)
-                return prox
-            setattr(self, i, mkprox())
+                    def prox(*args, **kwargs):
+                        self.run(evt, *args, **kwargs)
+                    return prox
+                setattr(self, i, mkprox())
 
     def run(self, name, *args, **kwargs):
         # It's possible for ns to be un-initialised if we failed during
         # configure
-        if self.ns is not None:
+        if self.ns is not None and not self.dead:
             func = self.ns.get(name)
             if func:
                 with scriptenv(self.path, self.args):
@@ -149,18 +151,35 @@ class Script:
             self.observer.start()
         self.run("configure", options)
 
+    def done(self):
+        self.run("done")
+        self.dead = True
+
 
 class ScriptLoader():
     """
         An addon that manages loading scripts from options.
     """
     def configure(self, options):
-        for s in options.scripts or []:
-            if not ctx.master.addons.has_addon(s):
+        for s in options.scripts:
+            if options.scripts.count(s) > 1:
+                raise exceptions.OptionsError("Duplicate script: %s" % s)
+
+        for a in ctx.master.addons.chain[:]:
+            if isinstance(a, Script) and a.name not in options.scripts:
+                    ctx.log.info("Un-loading script: %s" % a.name)
+                    ctx.master.addons.remove(a)
+
+        current = {}
+        for a in ctx.master.addons.chain[:]:
+            if isinstance(a, Script):
+                current[a.name] = a
+                ctx.master.addons.chain.remove(a)
+
+        for s in options.scripts:
+            if s in current:
+                ctx.master.addons.chain.append(current[s])
+            else:
                 ctx.log.info("Loading script: %s" % s)
                 sc = Script(s)
                 ctx.master.addons.add(sc)
-        for a in ctx.master.addons.chain:
-            if isinstance(a, Script):
-                if a.name not in options.scripts or []:
-                    ctx.master.addons.remove(a)
diff --git a/test/mitmproxy/builtins/test_script.py b/test/mitmproxy/builtins/test_script.py
index 2c2568ed..da60b64c 100644
--- a/test/mitmproxy/builtins/test_script.py
+++ b/test/mitmproxy/builtins/test_script.py
@@ -58,8 +58,8 @@ class TestScript(mastertest.MasterTest):
         )
         m.addons.add(sc)
         assert sc.ns["call_log"] == [
-            ("start", (), {}),
-            ("configure", (options.Options(),), {})
+            ("solo", "start", (), {}),
+            ("solo", "configure", (options.Options(),), {})
         ]
 
         sc.ns["call_log"] = []
@@ -67,7 +67,7 @@ class TestScript(mastertest.MasterTest):
         self.invoke(m, "request", f)
 
         recf = sc.ns["call_log"][0]
-        assert recf[0] == "request"
+        assert recf[1] == "request"
 
     def test_reload(self):
         s = state.State()
@@ -129,3 +129,59 @@ class TestScriptLoader(mastertest.MasterTest):
         assert len(m.addons) == 2
         o.update(scripts = [])
         assert len(m.addons) == 1
+
+    def test_dupes(self):
+        s = state.State()
+        o = options.Options(scripts=["one", "one"])
+        m = master.FlowMaster(o, None, s)
+        sc = script.ScriptLoader()
+        tutils.raises(exceptions.OptionsError, m.addons.add, sc)
+
+    def test_order(self):
+        rec = tutils.test_data.path("data/addonscripts/recorder.py")
+
+        s = state.State()
+        o = options.Options(
+            scripts = [
+                "%s %s" % (rec, "a"),
+                "%s %s" % (rec, "b"),
+                "%s %s" % (rec, "c"),
+            ]
+        )
+        m = mastertest.RecordingMaster(o, None, s)
+        sc = script.ScriptLoader()
+        m.addons.add(sc)
+
+        debug = [(i[0], i[1]) for i in m.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')
+        ]
+        m.event_log[:] = []
+
+        o.scripts = [
+            "%s %s" % (rec, "c"),
+            "%s %s" % (rec, "a"),
+            "%s %s" % (rec, "b"),
+        ]
+        debug = [(i[0], i[1]) for i in m.event_log if i[0] == "debug"]
+        assert debug == [
+            ('debug', 'c configure'),
+            ('debug', 'a configure'),
+            ('debug', 'b configure'),
+        ]
+        m.event_log[:] = []
+
+        o.scripts = [
+            "%s %s" % (rec, "x"),
+            "%s %s" % (rec, "a"),
+        ]
+        debug = [(i[0], i[1]) for i in m.event_log if i[0] == "debug"]
+        assert debug == [
+            ('debug', 'c done'),
+            ('debug', 'b done'),
+            ('debug', 'x start'),
+            ('debug', 'x configure'),
+            ('debug', 'a configure'),
+        ]
diff --git a/test/mitmproxy/data/addonscripts/recorder.py b/test/mitmproxy/data/addonscripts/recorder.py
index 728203e3..b6ac8d89 100644
--- a/test/mitmproxy/data/addonscripts/recorder.py
+++ b/test/mitmproxy/data/addonscripts/recorder.py
@@ -1,8 +1,14 @@
 from mitmproxy import controller
 from mitmproxy import ctx
+import sys
 
 call_log = []
 
+if len(sys.argv) > 1:
+    name = sys.argv[1]
+else:
+    name = "solo"
+
 # Keep a log of all possible event calls
 evts = list(controller.Events) + ["configure"]
 for i in evts:
@@ -10,9 +16,10 @@ for i in evts:
         evt = i
 
         def prox(*args, **kwargs):
-            lg = (evt, args, kwargs)
+            lg = (name, evt, args, kwargs)
             if evt != "log":
                 ctx.log.info(str(lg))
             call_log.append(lg)
+            ctx.log.debug("%s %s" % (name, evt))
         return prox
     globals()[i] = mkprox()