Add an external script API.

External scripts can read a flow, modify it, and then return it to mitmproxy
using a simple API.

The "|" keyboard shortcut within mitmproxy prompts the user for a script.

3 files changed, 107 insertions, 33 deletions

diff --git a/libmproxy/console.py b/libmproxy/console.py
index 12f6dec9..c48a41bc 100644
--- a/libmproxy/console.py
+++ b/libmproxy/console.py
@@ -44,46 +44,46 @@ def format_keyvals(lst, key="key", val="text", space=5, indent=0):
     return ret
 
 
-def format_flow(flow, focus, padding=3):
-    if not flow.request and not flow.response:
+def format_flow(f, focus, padding=3):
+    if not f.request and not f.response:
         txt = [
-            ("title", " Connection from %s..."%(flow.connection.address)),
+            ("title", " Connection from %s..."%(f.connection.address)),
         ]
     else:
         txt = [
-            ("ack", "!") if flow.intercepting and not flow.request.acked else " ",
-            ("method", flow.request.method),
+            ("ack", "!") if f.intercepting and not f.request.acked else " ",
+            ("method", f.request.method),
             " ",
             (
-                "text" if (flow.response or flow.error) else "title",
-                flow.request.url(),
+                "text" if (f.response or f.error) else "title",
+                f.request.url(),
             ),
         ]
-        if flow.response or flow.error or flow.is_replay():
+        if f.response or f.error or f.is_replay():
             txt.append("\n" + " "*(padding+2))
-            if flow.is_replay():
+            if f.is_replay():
                 txt.append(("method", "[replay] "))
-            if not (flow.response or flow.error):
+            if not (f.response or f.error):
                 txt.append(("text", "waiting for response..."))
 
-        if flow.response:
+        if f.response:
             txt.append(
-               ("ack", "!") if flow.intercepting and not flow.response.acked else " "
+               ("ack", "!") if f.intercepting and not f.response.acked else " "
             )
             txt.append("<- ")
-            if flow.response.code in [200, 304]:
-                txt.append(("goodcode", str(flow.response.code)))
+            if f.response.code in [200, 304]:
+                txt.append(("goodcode", str(f.response.code)))
             else:
-                txt.append(("error", str(flow.response.code)))
-            t = flow.response.headers.get("content-type")
+                txt.append(("error", str(f.response.code)))
+            t = f.response.headers.get("content-type")
             if t:
                 t = t[0].split(";")[0]
                 txt.append(("text", " %s"%t))
-            if flow.response.content:
-                txt.append(", %s"%utils.pretty_size(len(flow.response.content)))
-        elif flow.error:
+            if f.response.content:
+                txt.append(", %s"%utils.pretty_size(len(f.response.content)))
+        elif f.error:
             txt.append(
-               ("error", flow.error.msg)
+               ("error", f.error.msg)
             )
     if focus:
         txt.insert(0, ("focus", ">>" + " "*(padding-2)))
@@ -193,13 +193,13 @@ class ConnectionListView(urwid.ListWalker):
 
 
 class ConnectionViewHeader(WWrap):
-    def __init__(self, master, flow):
-        self.master, self.flow = master, flow
-        self.w = urwid.Text(format_flow(flow, False, padding=0))
+    def __init__(self, master, f):
+        self.master, self.flow = master, f
+        self.w = urwid.Text(format_flow(f, False, padding=0))
 
-    def refresh_connection(self, flow):
+    def refresh_connection(self, f):
         if f == self.flow:
-            self.w = urwid.Text(format_flow(flow, False, padding=0))
+            self.w = urwid.Text(format_flow(f, False, padding=0))
 
 
 VIEW_BODY_RAW = 0
@@ -520,8 +520,20 @@ class ConnectionView(WWrap):
                 self.master.prompt("Save response body: ", self.save_body)
         elif key == " ":
             self.master.view_next_flow(self.flow)
+        elif key == "|":
+            self.master.path_prompt("Script:", self.run_script)
         return key
 
+    def run_script(self, path):
+        path = os.path.expanduser(path)
+        try:
+            newflow = self.flow.run_script(path)
+        except flow.RunException, e:
+            self.master.statusbar.message("Script error: %s"%e)
+            return
+        self.flow.load_state(newflow.get_state())
+        self.master.refresh_connection(self.flow)
+
 
 class _PathCompleter:
     DEFAULTPATH = "/bin:/usr/bin:/usr/local/bin"
@@ -942,6 +954,7 @@ class ConsoleMaster(controller.Master):
             ("s", "save this flow"),
             ("v", "view contents in external viewer"),
             ("w", "save request or response body"),
+            ("|", "run script"),
             ("tab", "toggle response/request view"),
             ("space", "next flow"),
         ]
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index a014f8cb..c0b96c90 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -2,9 +2,12 @@
     This module provides more sophisticated flow tracking. These match requests
     with their responses, and provide filtering and interception facilities.
 """
+import subprocess, base64, sys
 from contrib import bson
 import proxy, threading
 
+class RunException(Exception): pass
+
 class ReplayConnection:
     pass
 
@@ -33,11 +36,39 @@ class Flow:
         self.intercepting = False
         self._backup = None
 
-    def run_script(self):
+    def script_serialize(self):
+        data = self.get_state()
+        data = bson.dumps(data)
+        return base64.encodestring(data)
+
+    @classmethod
+    def script_deserialize(klass, data):
+        data = base64.decodestring(data)
+        try:
+            data = bson.loads(data)
+        # bson.loads doesn't define a particular exception on error...
+        except Exception:
+            return None
+        return klass.from_state(data)
+
+    def run_script(self, path):
         """
             Run a script on a flow, returning the modified flow.
+
+            Raises RunException if there's an error.
         """
-        pass
+        data = self.script_serialize()
+        try:
+            p = subprocess.Popen([path], stdout=subprocess.PIPE, stdin=subprocess.PIPE)
+        except OSError, e:
+            raise RunException(e.args[1])
+        so, se = p.communicate(data)
+        if p.returncode:
+            raise RunException("Script returned error code %s"%p.returncode)
+        f = Flow.script_deserialize(so)
+        if not f:
+            raise RunException("Invalid response from script.")
+        return f
 
     def dump(self):
         data = dict(
@@ -52,15 +83,18 @@ class Flow:
             error = self.error.get_state() if self.error else None,
         )
 
-    @classmethod
-    def from_state(klass, state):
-        f = klass(None)
+    def load_state(self, state):
         if state["request"]:
-            f.request = proxy.Request.from_state(state["request"])
+            self.request = proxy.Request.from_state(state["request"])
         if state["response"]:
-            f.response = proxy.Response.from_state(f.request, state["response"])
+            self.response = proxy.Response.from_state(self.request, state["response"])
         if state["error"]:
-            f.error = proxy.Error.from_state(state["error"])
+            self.error = proxy.Error.from_state(state["error"])
+
+    @classmethod
+    def from_state(klass, state):
+        f = klass(None)
+        f.load_state(state)
         return f
 
     def __eq__(self, other):
diff --git a/libmproxy/script.py b/libmproxy/script.py
new file mode 100644
index 00000000..9ff861e9
--- /dev/null
+++ b/libmproxy/script.py
@@ -0,0 +1,27 @@
+"""
+    The mitmproxy scripting interface is simple - a serialized representation
+    of a flow is passed to the script on stdin, and a possibly modified flow is
+    then read by mitmproxy from the scripts stdout. This module provides two
+    convenience functions to make loading and returning data from scripts
+    simple.
+"""
+import sys, base64
+from contrib import bson
+import flow
+
+
+def load_flow():
+    """
+        Load a flow from the stdin. Returns a Flow object.
+    """
+    data = sys.stdin.read()
+    return flow.Flow.script_deserialize(data)
+
+
+def return_flow(f):
+    """
+        Print a flow to stdout. 
+    """
+    print >> sys.stdout, f.script_serialize()
+    
+