scripts: keep scripts just after the ScriptLoader in addon chain

We need scripts to run _before_ filestreamer, so we can't just add them to the
end of the chain. This patch also fixes an issue that could cause scripts to be
initialised un-necessarily if only the order of scripts in options changed.

4 files changed, 43 insertions, 16 deletions

diff --git a/mitmproxy/addons.py b/mitmproxy/addons.py
index db126dd9..b575b607 100644
--- a/mitmproxy/addons.py
+++ b/mitmproxy/addons.py
@@ -28,20 +28,32 @@ class Addons(object):
             with self.master.handlecontext():
                 i.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()
+        )
+
     def add(self, *addons):
+        """
+            Add addons to the end of the chain, and run their startup events.
+        """
         if not addons:
             raise ValueError("No addons specified.")
         self.chain.extend(addons)
         for i in addons:
-            self.invoke_with_context(i, "start")
-            self.invoke_with_context(
-                i,
-                "configure",
-                self.master.options,
-                self.master.options.keys()
-            )
+            self.startup(i)
 
     def remove(self, addon):
+        """
+            Remove an addon from the chain, and run its done events.
+        """
         self.chain = [i for i in self.chain if i is not addon]
         self.invoke_with_context(addon, "done")
 
diff --git a/mitmproxy/builtins/script.py b/mitmproxy/builtins/script.py
index 15ee4936..5c375d27 100644
--- a/mitmproxy/builtins/script.py
+++ b/mitmproxy/builtins/script.py
@@ -239,16 +239,35 @@ class ScriptLoader():
                         ctx.log.info("Un-loading script: %s" % a.name)
                         ctx.master.addons.remove(a)
 
+            # The machinations below are to ensure that:
+            #   - Scripts remain in the same order
+            #   - Scripts are listed directly after the script addon. This is
+            #   needed to ensure that interactions with, for instance, flow
+            #   serialization remains correct.
+            #   - Scripts are not initialized un-necessarily. If only a
+            #   script's order in the script list has changed, it should simply
+            #   be moved.
+
             current = {}
             for a in ctx.master.addons.chain[:]:
                 if isinstance(a, Script):
                     current[a.name] = a
                     ctx.master.addons.chain.remove(a)
 
+            ordered = []
+            newscripts = []
             for s in options.scripts:
                 if s in current:
-                    ctx.master.addons.chain.append(current[s])
+                    ordered.append(current[s])
                 else:
                     ctx.log.info("Loading script: %s" % s)
                     sc = Script(s)
-                    ctx.master.addons.add(sc)
+                    ordered.append(sc)
+                    newscripts.append(sc)
+
+            ochain = ctx.master.addons.chain
+            pos = ochain.index(self)
+            ctx.master.addons.chain = ochain[:pos+1] + ordered + ochain[pos+1:]
+
+            for s in newscripts:
+                ctx.master.addons.startup(s)
diff --git a/test/mitmproxy/builtins/test_script.py b/test/mitmproxy/builtins/test_script.py
index 261adb65..54451313 100644
--- a/test/mitmproxy/builtins/test_script.py
+++ b/test/mitmproxy/builtins/test_script.py
@@ -237,12 +237,8 @@ class TestScriptLoader(mastertest.MasterTest):
             "%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[:] = []
+        # No events, only order has changed
+        assert debug == []
 
         o.scripts = [
             "%s %s" % (rec, "x"),
diff --git a/test/mitmproxy/test_addons.py b/test/mitmproxy/test_addons.py
index 52d7f07f..c5d54e8c 100644
--- a/test/mitmproxy/test_addons.py
+++ b/test/mitmproxy/test_addons.py
@@ -16,6 +16,6 @@ def test_simple():
     o = options.Options()
     m = controller.Master(o)
     a = addons.Addons(m)
-    a.add(o, TAddon("one"))
+    a.add(TAddon("one"))
     assert a.get("one")
     assert not a.get("two")