Add script hooks, enable new engine for mitmdump.

5 files changed, 68 insertions, 31 deletions

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