diff options
Diffstat (limited to 'libmproxy/flow.py')
| -rw-r--r-- | libmproxy/flow.py | 33 |
1 files changed, 22 insertions, 11 deletions
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): |