Add script hooks, enable new engine for mitmdump.

9 files changed, 116 insertions, 109 deletions

diff --git a/examples/simple_script b/examples/simple_script
index aed937ec..de562163 100755
--- a/examples/simple_script
+++ b/examples/simple_script
@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 from libmproxy import script
 
-f = script.load_flow()
-f.request.headers["newheader"] = ["foo"]
-script.return_flow(f)
+def response(ctx, f):
+    ctx.log("processing a response")
+    f.response.headers["newheader"] = ["foo"]
diff --git a/libmproxy/dump.py b/libmproxy/dump.py
index a47066d9..f0379b8b 100644
--- a/libmproxy/dump.py
+++ b/libmproxy/dump.py
@@ -1,5 +1,5 @@
 import sys, os
-import flow, filt, utils
+import flow, filt, utils, script
 
 class DumpError(Exception): pass
 
@@ -15,8 +15,6 @@ class Options(object):
         "kill",
         "no_server",
         "refresh_server_playback",
-        "request_script",
-        "response_script",
         "rfile",
         "rheaders",
         "server_replay",
@@ -68,11 +66,6 @@ class DumpMaster(flow.FlowMaster):
         else:
             self.filt = None
 
-        if self.o.response_script:
-            self.set_response_script(self.o.response_script)
-        if self.o.request_script:
-            self.set_request_script(self.o.request_script)
-
         if options.stickycookie:
             self.set_stickycookie(options.stickycookie)
 
@@ -109,6 +102,10 @@ class DumpMaster(flow.FlowMaster):
                 not options.keepserving
             )
 
+        if options.script:
+            err = self.load_script(options.script)
+            if err:
+                raise DumpError(err)
 
     def _readflow(self, path):
         path = os.path.expanduser(path)
@@ -119,20 +116,6 @@ class DumpMaster(flow.FlowMaster):
             raise DumpError(v.strerror)
         return flows
 
-    def _runscript(self, f, script):
-        try:
-            ret = f.run_script(script)
-            if self.o.verbosity > 0:
-                print >> self.outfile, ret
-        except flow.RunException, e:
-            if e.errout:
-                eout = "Script output:\n" + self.indent(4, e.errout) + "\n"
-            else:
-                eout = ""
-            raise DumpError(
-                    "%s: %s\n%s"%(script, e.args[0], eout)
-                )
-
     def add_event(self, e, level="info"):
         if self.eventlog:
             print >> self.outfile, e
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 49cf0796..89a52eff 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -3,7 +3,7 @@
     with their responses, and provide filtering and interception facilities.
 """
 import subprocess, sys, json, hashlib, Cookie, cookielib
-import proxy, threading, netstring, filt
+import proxy, threading, netstring, filt, script
 import controller, version
 
 class RunException(Exception):
@@ -191,10 +191,6 @@ class Flow:
         f.load_state(state)
         return f
 
-    def script_serialize(self):
-        data = self.get_state()
-        return json.dumps(data)
-
     @classmethod
     def script_deserialize(klass, data):
         try:
@@ -203,41 +199,6 @@ class Flow:
             return None
         return klass.from_state(data)
 
-    def run_script(self, path):
-        """
-            Run a script on a flow.
-
-            Returns a (flow, stderr output) tuple, or raises RunException if
-            there's an error.
-        """
-        self.backup()
-        data = self.script_serialize()
-        try:
-            p = subprocess.Popen(
-                    [path],
-                    stdout=subprocess.PIPE,
-                    stdin=subprocess.PIPE,
-                    stderr=subprocess.PIPE,
-                )
-        except OSError, e:
-            raise RunException(e.args[1], None, None)
-        so, se = p.communicate(data)
-        if p.returncode:
-            raise RunException(
-                "Script returned error code %s"%p.returncode,
-                p.returncode,
-                se
-            )
-        f = Flow.script_deserialize(so)
-        if not f:
-            raise RunException(
-                    "Invalid response from script.",
-                    p.returncode,
-                    se
-                )
-        self.load_state(f.get_state())
-        return se
-
     def get_state(self, nobackup=False):
         d = dict(
             request = self.request.get_state() if self.request else None,
@@ -463,7 +424,7 @@ class FlowMaster(controller.Master):
         self.server_playback = None
         self.client_playback = None
         self.kill_nonreplay = False
-        self.plugin = None
+        self.script = None
 
         self.stickycookie_state = False
         self.stickycookie_txt = None
@@ -476,19 +437,26 @@ class FlowMaster(controller.Master):
         self.autodecode = False
         self.refresh_server_playback = False
 
-    def _runscript(self, f, script):
-        #begin nocover
-        raise NotImplementedError
-        #end nocover
-
     def add_event(self, e, level="info"):
         """
             level: info, error
         """
         pass
 
-    def set_plugin(self, p):
-        self.plugin = p
+    def load_script(self, path):
+        """
+            Loads a script. Returns an error description if something went
+            wrong.
+        """
+        s = script.Script(path, self)
+        try:
+            s.load()
+        except script.ScriptError, v:
+            return v.args[0]
+        ret = s.run("start")
+        if not ret[0] and ret[1]:
+            return "Error in script start:\n\n" + ret[1][1]
+        self.script = s
 
     def set_stickycookie(self, txt):
         if txt:
@@ -620,11 +588,20 @@ class FlowMaster(controller.Master):
             rt.start()
         #end nocover
 
-    def handle_clientconnect(self, r):
-        self.add_event("Connect from: %s:%s"%r.address)
-        r.ack()
+    def run_script(self, name, *args, **kwargs):
+        if self.script:
+            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("clientconnect", cc)
+        self.add_event("Connect from: %s:%s"%cc.address)
+        cc.ack()
 
     def handle_clientdisconnect(self, r):
+        self.run_script("clientdisconnect", r)
         s = "Disconnect from: %s:%s"%r.client_conn.address
         self.add_event(s)
         if r.client_conn.requestcount:
@@ -638,6 +615,8 @@ class FlowMaster(controller.Master):
 
     def handle_error(self, r):
         f = self.state.add_error(r)
+        if f:
+            self.run_script("error", f)
         if self.client_playback:
             self.client_playback.clear(f)
         r.ack()
@@ -645,11 +624,14 @@ class FlowMaster(controller.Master):
 
     def handle_request(self, r):
         f = self.state.add_request(r)
+        self.run_script("request", f)
         self.process_new_request(f)
         return f
 
     def handle_response(self, r):
         f = self.state.add_response(r)
+        if f:
+            self.run_script("response", f)
         if self.client_playback:
             self.client_playback.clear(f)
         if not f:
diff --git a/libmproxy/script.py b/libmproxy/script.py
index da3131a8..03eff958 100644
--- a/libmproxy/script.py
+++ b/libmproxy/script.py
@@ -8,7 +8,7 @@ class Context:
         self.master, self.state = master, state
 
     def log(self, *args, **kwargs):
-        self.master.log(*args, **kwargs)
+        self.master.add_event(*args, **kwargs)
 
 
 class Script:
@@ -32,9 +32,14 @@ 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)
+        if not os.path.exists(path):
+            raise ScriptError("No such file: %s"%self.path)
+        if not os.path.isfile(path):
+            raise ScriptError("Not a file: %s"%self.path)
         ns = {}
         try:
-            self.mod = execfile(os.path.expanduser(self.path), {}, ns)
+            self.mod = execfile(path, ns, ns)
         except Exception, v:
             raise ScriptError(traceback.format_exc(v))
         self.ns = ns
diff --git a/test/scripts/all.py b/test/scripts/all.py
new file mode 100644
index 00000000..e6da7e51
--- /dev/null
+++ b/test/scripts/all.py
@@ -0,0 +1,20 @@
+log = []
+def clientconnect(ctx, cc):
+    ctx.log("XCLIENTCONNECT")
+    log.append("clientconnect")
+
+def request(ctx, r):
+    ctx.log("XREQUEST")
+    log.append("request")
+
+def response(ctx, r):
+    ctx.log("XRESPONSE")
+    log.append("response")
+
+def clientdisconnect(ctx, cc):
+    ctx.log("XCLIENTDISCONNECT")
+    log.append("clientdisconnect")
+
+def error(ctx, cc):
+    ctx.log("XERROR")
+    log.append("error")
diff --git a/test/scripts/starterr.py b/test/scripts/starterr.py
new file mode 100644
index 00000000..456fecc0
--- /dev/null
+++ b/test/scripts/starterr.py
@@ -0,0 +1,3 @@
+
+def start(ctx):
+    raise ValueError
diff --git a/test/test_dump.py b/test/test_dump.py
index 3abbad6d..9a35772d 100644
--- a/test/test_dump.py
+++ b/test/test_dump.py
@@ -111,30 +111,22 @@ class uDumpMaster(libpry.AutoTree):
             wfile = "nonexistentdir/foo"
         )
 
-    def test_request_script(self):
-        ret = self._dummy_cycle(1, None, "", request_script="scripts/a", verbosity=1)
-        assert "TESTOK" in ret
-        assert "DEBUG" in ret
-        libpry.raises(
-            dump.DumpError,
-            self._dummy_cycle, 1, None, "", request_script="nonexistent"
-        )
-        libpry.raises(
-            dump.DumpError,
-            self._dummy_cycle, 1, None, "", request_script="scripts/err_return"
+    def test_script(self):
+        ret = self._dummy_cycle(
+            1, None, "",
+            script="scripts/all.py", verbosity=0, eventlog=True
         )
-
-    def test_response_script(self):
-        ret = self._dummy_cycle(1, None, "", response_script="scripts/a", verbosity=1)
-        assert "TESTOK" in ret
-        assert "DEBUG" in ret
+        assert "XCLIENTCONNECT" in ret
+        assert "XREQUEST" in ret
+        assert "XRESPONSE" in ret
+        assert "XCLIENTDISCONNECT" in ret
         libpry.raises(
             dump.DumpError,
-            self._dummy_cycle, 1, None, "", response_script="nonexistent"
+            self._dummy_cycle, 1, None, "", script="nonexistent"
         )
         libpry.raises(
             dump.DumpError,
-            self._dummy_cycle, 1, None, "", response_script="scripts/err_return"
+            self._dummy_cycle, 1, None, "", script="starterr.py"
         )
 
     def test_stickycookie(self):
diff --git a/test/test_flow.py b/test/test_flow.py
index 6e1de9af..45aee311 100644
--- a/test/test_flow.py
+++ b/test/test_flow.py
@@ -138,15 +138,6 @@ class uFlow(libpry.AutoTree):
         assert "DEBUG" == se.strip()
         assert f.request.host == "TESTOK"
 
-    def test_run_script_err(self):
-        f = tutils.tflow()
-        f.response = tutils.tresp()
-        f.request = f.response.request
-        libpry.raises("returned error", f.run_script,"scripts/err_return")
-        libpry.raises("invalid response", f.run_script,"scripts/err_data")
-        libpry.raises("no such file", f.run_script,"nonexistent")
-        libpry.raises("permission denied", f.run_script,"scripts/nonexecutable")
-
     def test_match(self):
         f = tutils.tflow()
         f.response = tutils.tresp()
@@ -449,13 +440,38 @@ class uSerialize(libpry.AutoTree):
 
 
 class uFlowMaster(libpry.AutoTree):
+    def test_load_script(self):
+        s = flow.State()
+        fm = flow.FlowMaster(None, s)
+        assert not fm.load_script("scripts/a.py")
+        assert fm.load_script("nonexistent")
+        assert "ValueError" in fm.load_script("scripts/starterr.py")
+
+    def test_script(self):
+        s = flow.State()
+        fm = flow.FlowMaster(None, s)
+        assert not fm.load_script("scripts/all.py")
+        req = tutils.treq()
+        fm.handle_clientconnect(req.client_conn)
+        assert fm.script.ns["log"][-1] == "clientconnect"
+        f = fm.handle_request(req)
+        assert fm.script.ns["log"][-1] == "request"
+        resp = tutils.tresp(req)
+        fm.handle_response(resp)
+        assert fm.script.ns["log"][-1] == "response"
+        dc = proxy.ClientDisconnect(req.client_conn)
+        fm.handle_clientdisconnect(dc)
+        assert fm.script.ns["log"][-1] == "clientdisconnect"
+        err = proxy.Error(f.request, "msg")
+        fm.handle_error(err)
+        assert fm.script.ns["log"][-1] == "error"
+
     def test_all(self):
         s = flow.State()
         fm = flow.FlowMaster(None, s)
         fm.anticache = True
         fm.anticomp = True
         req = tutils.treq()
-
         fm.handle_clientconnect(req.client_conn)
 
         f = fm.handle_request(req)
diff --git a/test/test_script.py b/test/test_script.py
index f6cfcced..9a1c3062 100644
--- a/test/test_script.py
+++ b/test/test_script.py
@@ -27,17 +27,23 @@ class uScript(libpry.AutoTree):
 
         s = script.Script("nonexistent", fm)
         libpry.raises(
-            script.ScriptError,
+            "no such file",
+            s.load
+        )
+
+        s = script.Script("scripts", fm)
+        libpry.raises(
+            "not a file",
             s.load
         )
 
-        s = script.Script(os.path.join("scripts", "syntaxerr.py"), fm)
+        s = script.Script("scripts/syntaxerr.py", fm)
         libpry.raises(
             script.ScriptError,
             s.load
         )
 
-        s = script.Script(os.path.join("scripts", "loaderr.py"), fm)
+        s = script.Script("scripts/loaderr.py", fm)
         libpry.raises(
             script.ScriptError,
             s.load