From f23818ceeaac88f266674ea518878a17a74a1d16 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Fri, 5 Aug 2011 14:03:10 +1200 Subject: [PATCH] Add a "done" event for scripts. Called exactly once after all other events. --- examples/stub.py | 15 ++++++++------- libmproxy/console.py | 16 +++++++++------- libmproxy/dump.py | 3 ++- libmproxy/flow.py | 33 ++++++++++++++++++++++----------- 4 files changed, 41 insertions(+), 26 deletions(-) diff --git a/examples/stub.py b/examples/stub.py index f235ea850..119298fc6 100644 --- a/examples/stub.py +++ b/examples/stub.py @@ -1,30 +1,31 @@ """ - This is a script stub, with empty definitions for all events. + This is a script stub, with definitions for all events. """ def start(ctx): """ Called once on script startup, before any other events. """ - pass + ctx.log("start") def clientconnect(ctx, client_connect): """ Called when a client initiates a connection to the proxy. Note that a connection can correspond to multiple HTTP requests """ - pass + ctx.log("clientconnect") def request(ctx, flow): """ Called when a client request has been received. """ + ctx.log("request") def response(ctx, flow): """ Called when a server response has been received. """ - pass + ctx.log("response") def error(ctx, flow): """ @@ -32,16 +33,16 @@ def error(ctx, flow): interrupted connections. This is distinct from a valid server HTTP error response, which is simply a response with an HTTP error code. """ - pass + ctx.log("error") def clientdisconnect(ctx, client_disconnect): """ Called when a client disconnects from the proxy. """ - pass + ctx.log("clientdisconnect") def done(ctx): """ Called once on script shutdown, after any other events. """ - pass + ctx.log("done") diff --git a/libmproxy/console.py b/libmproxy/console.py index a6501cc1e..4ee4018f6 100644 --- a/libmproxy/console.py +++ b/libmproxy/console.py @@ -964,12 +964,6 @@ class ConsoleMaster(flow.FlowMaster): print >> sys.stderr, "Sticky auth error:", r sys.exit(1) - if options.script: - err = self.load_script(options.script) - if err: - print >> sys.stderr, "Script load error:", r - sys.exit(1) - self.refresh_server_playback = options.refresh_server_playback self.anticache = options.anticache self.anticomp = options.anticomp @@ -987,6 +981,13 @@ class ConsoleMaster(flow.FlowMaster): self.debug = options.debug + if options.script: + err = self.load_script(options.script) + if err: + print >> sys.stderr, "Script load error:", err + sys.exit(1) + + def run_script_once(self, path, f): ret = self.get_script(path) if ret[0]: @@ -999,6 +1000,7 @@ class ConsoleMaster(flow.FlowMaster): s.run("response", f) if f.error: s.run("error", f) + s.run("done") self.refresh_connection(f) self.state.last_script = path @@ -1660,7 +1662,7 @@ class ConsoleMaster(flow.FlowMaster): k = None elif k == "s": if self.script: - self.script = None + self.load_script(None) else: self.path_prompt( "Set script: ", diff --git a/libmproxy/dump.py b/libmproxy/dump.py index 67f39a063..21ea038e5 100644 --- a/libmproxy/dump.py +++ b/libmproxy/dump.py @@ -191,10 +191,11 @@ class DumpMaster(flow.FlowMaster): self._process_flow(f) return f - # begin nocover def run(self): if self.o.rfile and not self.o.keepserving: + if self.script: + self.load_script(None) return try: return flow.FlowMaster.run(self) diff --git a/libmproxy/flow.py b/libmproxy/flow.py index c33a3f1ac..23a88cbad 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -1112,13 +1112,19 @@ class FlowMaster(controller.Master): def load_script(self, path): """ Loads a script. Returns an error description if something went - wrong. + wrong. If path is None, the current script is terminated. """ - r = self.get_script(path) - if r[0]: - return r[0] + if path is None: + self.run_script_hook("done") + self.script = None else: - self.script = r[1] + r = self.get_script(path) + if r[0]: + return r[0] + else: + if self.script: + self.run_script_hook("done") + self.script = r[1] def set_stickycookie(self, txt): if txt: @@ -1250,7 +1256,7 @@ class FlowMaster(controller.Master): rt.start() #end nocover - def run_script(self, name, *args, **kwargs): + def run_script_hook(self, name, *args, **kwargs): if self.script: ret = self.script.run(name, *args, **kwargs) if not ret[0] and ret[1]: @@ -1258,12 +1264,12 @@ class FlowMaster(controller.Master): self.add_event(e, "error") def handle_clientconnect(self, cc): - self.run_script("clientconnect", cc) + self.run_script_hook("clientconnect", cc) self.add_event("Connect from: %s:%s"%cc.address) cc._ack() def handle_clientdisconnect(self, r): - self.run_script("clientdisconnect", r) + self.run_script_hook("clientdisconnect", r) s = "Disconnect from: %s:%s"%r.client_conn.address self.add_event(s) if r.client_conn.requestcount: @@ -1278,7 +1284,7 @@ class FlowMaster(controller.Master): def handle_error(self, r): f = self.state.add_error(r) if f: - self.run_script("error", f) + self.run_script_hook("error", f) if self.client_playback: self.client_playback.clear(f) r._ack() @@ -1286,14 +1292,14 @@ class FlowMaster(controller.Master): def handle_request(self, r): f = self.state.add_request(r) - self.run_script("request", f) + self.run_script_hook("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) + self.run_script_hook("response", f) if self.client_playback: self.client_playback.clear(f) if not f: @@ -1301,6 +1307,11 @@ class FlowMaster(controller.Master): self.process_new_response(f) return f + def shutdown(self): + if self.script: + self.load_script(None) + controller.Master.shutdown(self) + class FlowWriter: def __init__(self, fo):