From 179cf758624152a4988e04a767759cb15b9c358a Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Wed, 3 Aug 2011 16:36:20 +1200 Subject: [PATCH] Add script hooks, enable new engine for mitmdump. --- examples/simple_script | 6 +-- libmproxy/dump.py | 27 +++---------- libmproxy/flow.py | 84 ++++++++++++++++------------------------ libmproxy/script.py | 9 ++++- test/scripts/all.py | 20 ++++++++++ test/scripts/starterr.py | 3 ++ test/test_dump.py | 30 ++++++-------- test/test_flow.py | 36 ++++++++++++----- test/test_script.py | 12 ++++-- 9 files changed, 117 insertions(+), 110 deletions(-) create mode 100644 test/scripts/all.py create mode 100644 test/scripts/starterr.py diff --git a/examples/simple_script b/examples/simple_script index aed937ec7..de5621634 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 a47066d97..f0379b8bd 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 49cf0796a..89a52effc 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 da3131a85..03eff958f 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 000000000..e6da7e51f --- /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 000000000..456fecc0e --- /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 3abbad6d6..9a35772dc 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 + def test_script(self): + ret = self._dummy_cycle( + 1, None, "", + script="scripts/all.py", verbosity=0, eventlog=True + ) + 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, "", request_script="nonexistent" + self._dummy_cycle, 1, None, "", script="nonexistent" ) libpry.raises( dump.DumpError, - self._dummy_cycle, 1, None, "", request_script="scripts/err_return" - ) - - 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 - libpry.raises( - dump.DumpError, - self._dummy_cycle, 1, None, "", response_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 6e1de9afe..45aee311b 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 f6cfcced7..9a1c3062f 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(os.path.join("scripts", "syntaxerr.py"), fm) + s = script.Script("scripts", fm) + libpry.raises( + "not a file", + s.load + ) + + 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