add support for multiple scripts and script arguments. refs #76

5 files changed, 50 insertions, 44 deletions

diff --git a/libmproxy/cmdline.py b/libmproxy/cmdline.py
index b76792cf..75f28bca 100644
--- a/libmproxy/cmdline.py
+++ b/libmproxy/cmdline.py
@@ -16,6 +16,8 @@
 import proxy
 import re, filt
 import argparse
+import shlex
+import os
 
 class ParseException(Exception): pass
 class OptionException(Exception): pass
@@ -151,7 +153,7 @@ def get_common_options(options):
         replacements = reps,
         setheaders = setheaders,
         server_replay = options.server_replay,
-        script = options.script,
+        scripts = options.scripts,
         stickycookie = stickycookie,
         stickyauth = stickyauth,
         showhost = options.showhost,
@@ -209,8 +211,9 @@ def common_options(parser):
     )
     parser.add_argument(
         "-s",
-        action="store", dest="script", default=None,
-        help="Run a script."
+        action="append", type=lambda x: shlex.split(x,posix=(os.name != "nt")), dest="scripts", default=[],
+        metavar='"script.py --bar"',
+        help="Run a script. Surround with quotes to pass script arguments. Can be passed multiple times."
     )
     parser.add_argument(
         "-t",
diff --git a/libmproxy/console/__init__.py b/libmproxy/console/__init__.py
index ba699d33..bb57ba26 100644
--- a/libmproxy/console/__init__.py
+++ b/libmproxy/console/__init__.py
@@ -460,7 +460,7 @@ class ConsoleMaster(flow.FlowMaster):
             self._run_script_method("response", s, f)
         if f.error:
             self._run_script_method("error", s, f)
-        s.run("done")
+        s.unload()
         self.refresh_flow(f)
         self.state.last_script = path
 
diff --git a/libmproxy/dump.py b/libmproxy/dump.py
index b1022ef5..a4846da7 100644
--- a/libmproxy/dump.py
+++ b/libmproxy/dump.py
@@ -36,7 +36,7 @@ class Options(object):
         "rheaders",
         "setheaders",
         "server_replay",
-        "script",
+        "scripts",
         "showhost",
         "stickycookie",
         "stickyauth",
@@ -121,8 +121,8 @@ class DumpMaster(flow.FlowMaster):
                 not options.keepserving
             )
 
-        if options.script:
-            err = self.load_script(options.script)
+        for script_argv in options.scripts:
+            err = self.load_script(script_argv)
             if err:
                 raise DumpError(err)
 
@@ -230,8 +230,8 @@ class DumpMaster(flow.FlowMaster):
 
     def run(self):  # pragma: no cover
         if self.o.rfile and not self.o.keepserving:
-            if self.script:
-                self.load_script(None)
+            for script in self.scripts:
+                self.unload_script(script)
             return
         try:
             return flow.FlowMaster.run(self)
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 10c5da5d..61b2a241 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -1357,7 +1357,7 @@ class FlowMaster(controller.Master):
         self.server_playback = None
         self.client_playback = None
         self.kill_nonreplay = False
-        self.script = None
+        self.scripts = []
         self.pause_scripts = False
 
         self.stickycookie_state = False
@@ -1381,37 +1381,43 @@ class FlowMaster(controller.Master):
         """
         pass
 
-    def get_script(self, path):
+    def get_script(self, script_argv):
         """
             Returns an (error, script) tuple.
         """
-        s = script.Script(path, ScriptContext(self))
+        s = script.Script(script_argv, ScriptContext(self))
         try:
             s.load()
         except script.ScriptError, v:
             return (v.args[0], None)
-        ret = s.run("start")
-        if not ret[0] and ret[1]:
-            return ("Error in script start:\n\n" + ret[1][1], None)
         return (None, s)
 
-    def load_script(self, path):
+    def unload_script(self,script):
+        script.unload()
+        self.scripts.remove(script)
+
+    def load_script(self, script_argv):
         """
             Loads a script. Returns an error description if something went
-            wrong. If path is None, the current script is terminated.
+            wrong.
         """
-        if path is None:
-            self.run_script_hook("done")
-            self.script = None
+        r = self.get_script(script_argv)
+        if r[0]:
+            return r[0]
         else:
-            r = self.get_script(path)
-            if r[0]:
-                return r[0]
-            else:
-                if self.script:
-                    self.run_script_hook("done")
-                self.script = r[1]
+            self.scripts.append(r[1])
+
+    def run_single_script_hook(self, script, name, *args, **kwargs):
+        if script and not self.pause_scripts:
+            ret = script.run(name, *args, **kwargs)
+            if not ret[0] and ret[1]:
+                e = "Script error:\n" + ret[1][1]
+                self.add_event(e, "error")
 
+    def run_script_hook(self, name, *args, **kwargs):
+        for script in self.scripts:
+            self.run_single_script_hook(script, name, *args, **kwargs)
+      
     def set_stickycookie(self, txt):
         if txt:
             flt = filt.parse(txt)
@@ -1561,13 +1567,6 @@ class FlowMaster(controller.Master):
             if block:
                 rt.join()
 
-    def run_script_hook(self, name, *args, **kwargs):
-        if self.script and not self.pause_scripts:
-            ret = self.script.run(name, *args, **kwargs)
-            if not ret[0] and ret[1]:
-                e = "Script error:\n" + ret[1][1]
-                self.add_event(e, "error")
-
     def handle_clientconnect(self, cc):
         self.run_script_hook("clientconnect", cc)
         cc.reply()
@@ -1609,8 +1608,8 @@ class FlowMaster(controller.Master):
         return f
 
     def shutdown(self):
-        if self.script:
-            self.load_script(None)
+        for script in self.scripts:
+            self.unload_script(script)
         controller.Master.shutdown(self)
         if self.stream:
             for i in self.state._flow_list:
diff --git a/libmproxy/script.py b/libmproxy/script.py
index 4ffac71b..4676672b 100644
--- a/libmproxy/script.py
+++ b/libmproxy/script.py
@@ -23,12 +23,12 @@ class Script:
     """
         The instantiator should do something along this vein:
 
-            s = Script(path, master)
+            s = Script(argv, master)
             s.load()
-            s.run("start")
     """
-    def __init__(self, path, ctx):
-        self.path, self.ctx = path, ctx
+    def __init__(self, argv, ctx):
+        self.argv = argv
+        self.ctx = ctx
         self.ns = None
 
     def load(self):
@@ -38,17 +38,21 @@ class Script:
             Raises ScriptError on failure, with argument equal to an error
             message that may be a formatted traceback.
         """
-        path = os.path.expanduser(self.path)
+        path = os.path.expanduser(self.argv[0])
         if not os.path.exists(path):
-            raise ScriptError("No such file: %s"%self.path)
+            raise ScriptError("No such file: %s" % path)
         if not os.path.isfile(path):
-            raise ScriptError("Not a file: %s"%self.path)
+            raise ScriptError("Not a file: %s" % path)
         ns = {}
         try:
             execfile(path, ns, ns)
+            self.ns = ns
+            self.run("start", self.argv)
         except Exception, v:
             raise ScriptError(traceback.format_exc(v))
-        self.ns = ns
+
+    def unload(self):
+        return self.run("done")
 
     def run(self, name, *args, **kwargs):
         """