diff options
Diffstat (limited to 'libmproxy')
| -rw-r--r-- | libmproxy/console.py | 16 | ||||
| -rw-r--r-- | libmproxy/dump.py | 3 | ||||
| -rw-r--r-- | libmproxy/flow.py | 33 |
3 files changed, 33 insertions, 19 deletions
diff --git a/libmproxy/console.py b/libmproxy/console.py index a6501cc1..4ee4018f 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 67f39a06..21ea038e 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 c33a3f1a..23a88cba 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): |