1 files changed, 51 insertions, 12 deletions

diff --git a/libmproxy/script.py b/libmproxy/script.py
index f8a0d085..0912c9ae 100644
--- a/libmproxy/script.py
+++ b/libmproxy/script.py
@@ -1,10 +1,43 @@
-import os, traceback, threading
+import os, traceback, threading, shlex
 import controller
 
 class ScriptError(Exception):
     pass
 
 
+class ScriptContext:
+    def __init__(self, master):
+        self._master = master
+
+    def log(self, *args, **kwargs):
+        """
+            Logs an event.
+
+            How this is handled depends on the front-end. mitmdump will display
+            events if the eventlog flag ("-e") was passed. mitmproxy sends
+            output to the eventlog for display ("v" keyboard shortcut).
+        """
+        self._master.add_event(*args, **kwargs)
+
+    def duplicate_flow(self, f):
+        """
+            Returns a duplicate of the specified flow. The flow is also
+            injected into the current state, and is ready for editing, replay,
+            etc.
+        """
+        self._master.pause_scripts = True
+        f = self._master.duplicate_flow(f)
+        self._master.pause_scripts = False
+        return f
+
+    def replay_request(self, f):
+        """
+            Replay the request on the current flow. The response will be added
+            to the flow object.
+        """
+        self._master.replay_request(f)
+
+
 class Script:
     """
         The instantiator should do something along this vein:
@@ -12,10 +45,22 @@ class Script:
             s = Script(argv, master)
             s.load()
     """
-    def __init__(self, argv, ctx):
-        self.argv = argv
-        self.ctx = ctx
+    def __init__(self, command, master):
+        self.command = command
+        self.argv = self.parse_command(command)
+        self.ctx = ScriptContext(master)
         self.ns = None
+        self.load()
+
+    @classmethod
+    def parse_command(klass, command):
+        args = shlex.split(command, posix=(os.name != "nt"))
+        args[0] = os.path.expanduser(args[0])
+        if not os.path.exists(args[0]):
+            raise ScriptError("Command not found.")
+        elif not os.path.isfile(args[0]):
+            raise ScriptError("Not a file: %s" % args[0])
+        return args
 
     def load(self):
         """
@@ -24,14 +69,9 @@ class Script:
             Raises ScriptError on failure, with argument equal to an error
             message that may be a formatted traceback.
         """
-        path = os.path.expanduser(self.argv[0])
-        if not os.path.exists(path):
-            raise ScriptError("No such file: %s" % path)
-        if not os.path.isfile(path):
-            raise ScriptError("Not a file: %s" % path)
         ns = {}
         try:
-            execfile(path, ns, ns)
+            execfile(self.argv[0], ns, ns)
         except Exception, v:
             raise ScriptError(traceback.format_exc(v))
         self.ns = ns
@@ -65,7 +105,6 @@ class Script:
 def _handle_concurrent_reply(fn, o, args=[], kwargs={}):
     reply = o.reply
     o.reply = controller.DummyReply()
-
     def run():
         fn(*args, **kwargs)
         reply(o)
@@ -82,4 +121,4 @@ def concurrent(fn):
         def _concurrent(ctx, conn):
             _handle_concurrent_reply(fn, conn, [ctx, conn])
         return _concurrent
-    raise NotImplementedError("Concurrent decorator not supported for this method.")
\ No newline at end of file
+    raise NotImplementedError("Concurrent decorator not supported for this method.")