From e337682d8e4b36c701ed2c6afede78907a0fe8f2 Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Wed, 3 Aug 2011 17:35:18 +1200
Subject: Enable "|" command to run a oneshot script on a single flow.

---
 libmproxy/console.py | 19 +++++++++++++++++--
 libmproxy/flow.py    | 22 ++++++++++++++++------
 test/test_flow.py    |  8 --------
 3 files changed, 33 insertions(+), 16 deletions(-)

diff --git a/libmproxy/console.py b/libmproxy/console.py
index e973ff7f..90075ee1 100644
--- a/libmproxy/console.py
+++ b/libmproxy/console.py
@@ -194,7 +194,7 @@ class ConnectionItem(WWrap):
         elif key == "|":
             self.master.path_prompt(
                 "Send flow to script: ", self.state.last_script,
-                self.master.run_script, self.flow
+                self.master.run_script_once, self.flow
             )
         return key
 
@@ -551,7 +551,7 @@ class ConnectionView(WWrap):
         elif key == "|":
             self.master.path_prompt(
                 "Send flow to script: ", self.state.last_script,
-                self.master.run_script, self.flow
+                self.master.run_script_once, self.flow
             )
         elif key == "z":
             if self.state.view_flow_mode == VIEW_FLOW_RESPONSE:
@@ -989,6 +989,21 @@ class ConsoleMaster(flow.FlowMaster):
 
         self.debug = options.debug
 
+    def run_script_once(self, path, f):
+        ret = self.get_script(path)
+        if ret[0]:
+            self.statusbar.message(ret[0])
+        s = ret[1]
+
+        if f.request:
+            s.run("request", f)
+        if f.response:
+            s.run("response", f)
+        if f.error:
+            s.run("error", f)
+        self.refresh_connection(f)
+        self.state.last_script = path
+
     def set_script(self, path):
         if not path:
             return
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 89a52eff..1043cb21 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -443,20 +443,30 @@ class FlowMaster(controller.Master):
         """
         pass
 
-    def load_script(self, path):
+    def get_script(self, path):
         """
-            Loads a script. Returns an error description if something went
-            wrong.
+            Returns an (error, script) tuple.
         """
         s = script.Script(path, self)
         try:
             s.load()
         except script.ScriptError, v:
-            return v.args[0]
+            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]
-        self.script = s
+            return ("Error in script start:\n\n" + ret[1][1], None)
+        return (None, s)
+
+    def load_script(self, path):
+        """
+            Loads a script. Returns an error description if something went
+            wrong.
+        """
+        r = self.get_script(path)
+        if r[0]:
+            return r[0]
+        else:
+            self.script = r[1]
 
     def set_stickycookie(self, txt):
         if txt:
diff --git a/test/test_flow.py b/test/test_flow.py
index 45aee311..b21f84cd 100644
--- a/test/test_flow.py
+++ b/test/test_flow.py
@@ -130,14 +130,6 @@ class uServerPlaybackState(libpry.AutoTree):
 
 
 class uFlow(libpry.AutoTree):
-    def test_run_script(self):
-        f = tutils.tflow()
-        f.response = tutils.tresp()
-        f.request = f.response.request
-        se = f.run_script("scripts/a")
-        assert "DEBUG" == se.strip()
-        assert f.request.host == "TESTOK"
-
     def test_match(self):
         f = tutils.tflow()
         f.response = tutils.tresp()
-- 
cgit v1.2.3