From b886f808beaba097066a1b82fe560b1e70099df0 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Mon, 31 Jan 2011 13:26:56 +1300 Subject: [PATCH] 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. --- libmproxy/console.py | 63 +++++++++++++++++++++++--------------- libmproxy/flow.py | 50 +++++++++++++++++++++++++----- libmproxy/script.py | 27 ++++++++++++++++ test/scripts/a | 8 +++++ test/scripts/err_data | 2 ++ test/scripts/err_return | 5 +++ test/scripts/nonexecutable | 0 test/test_flow.py | 16 ++++++++++ 8 files changed, 138 insertions(+), 33 deletions(-) create mode 100644 libmproxy/script.py create mode 100755 test/scripts/a create mode 100755 test/scripts/err_data create mode 100755 test/scripts/err_return create mode 100644 test/scripts/nonexecutable diff --git a/libmproxy/console.py b/libmproxy/console.py index 12f6dec93..c48a41bc9 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 a014f8cb5..c0b96c90b 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, ) + def load_state(self, state): + if state["request"]: + self.request = proxy.Request.from_state(state["request"]) + if state["response"]: + self.response = proxy.Response.from_state(self.request, state["response"]) + if state["error"]: + self.error = proxy.Error.from_state(state["error"]) + @classmethod def from_state(klass, state): f = klass(None) - if state["request"]: - f.request = proxy.Request.from_state(state["request"]) - if state["response"]: - f.response = proxy.Response.from_state(f.request, state["response"]) - if state["error"]: - f.error = proxy.Error.from_state(state["error"]) + 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 000000000..9ff861e93 --- /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() + + diff --git a/test/scripts/a b/test/scripts/a new file mode 100755 index 000000000..6973a44f7 --- /dev/null +++ b/test/scripts/a @@ -0,0 +1,8 @@ +#!/usr/bin/env python +import sys +sys.path.insert(0, "..") +from libmproxy import script + +f = script.load_flow() +f.request.host = "TESTOK" +script.return_flow(f) diff --git a/test/scripts/err_data b/test/scripts/err_data new file mode 100755 index 000000000..4bb1b1ba6 --- /dev/null +++ b/test/scripts/err_data @@ -0,0 +1,2 @@ +#!/usr/bin/env python +print "NONSENSE" diff --git a/test/scripts/err_return b/test/scripts/err_return new file mode 100755 index 000000000..a45926b5d --- /dev/null +++ b/test/scripts/err_return @@ -0,0 +1,5 @@ +#!/usr/bin/env python +import sys +sys.path.insert(0, "..") +sys.exit(1) + diff --git a/test/scripts/nonexecutable b/test/scripts/nonexecutable new file mode 100644 index 000000000..e69de29bb diff --git a/test/test_flow.py b/test/test_flow.py index 1867616e8..c97cc030f 100644 --- a/test/test_flow.py +++ b/test/test_flow.py @@ -3,6 +3,22 @@ import utils import libpry class uFlow(libpry.AutoTree): + def test_run_script(self): + f = utils.tflow() + f.response = utils.tresp() + f.request = f.response.request + f = f.run_script("scripts/a") + assert f.request.host == "TESTOK" + + def test_run_script_err(self): + f = utils.tflow() + f.response = utils.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 = utils.tflow() f.response = utils.tresp()