diff options
| -rw-r--r-- | libmproxy/console.py | 36 | ||||
| -rw-r--r-- | libmproxy/flow.py | 6 | ||||
| -rw-r--r-- | test/test_flow.py | 6 |
3 files changed, 41 insertions, 7 deletions
diff --git a/libmproxy/console.py b/libmproxy/console.py index f13bc933..f068c163 100644 --- a/libmproxy/console.py +++ b/libmproxy/console.py @@ -675,6 +675,10 @@ class StatusBar(WWrap): def get_status(self): r = [] + if self.master.client_playback: + r.append("[") + r.append(("statusbar_highlight", "cplayback")) + r.append(":%s to go]"%self.master.client_playback.count()) if self.master.state.intercept_txt: r.append("[") r.append(("statusbar_highlight", "i")) @@ -872,6 +876,24 @@ class ConsoleMaster(flow.FlowMaster): self.stickyhosts = {} self.anticache = options.anticache + if options.client_replay: + self.client_playback_path(options.client_replay) + + def _readflow(self, path): + path = os.path.expanduser(path) + try: + f = file(path, "r") + flows = list(flow.FlowReader(f).stream()) + except (IOError, flow.FlowReadError), v: + return True, v.strerror + return False, flows + + def client_playback_path(self, path): + err, ret = self._readflow(path) + if err: + self.statusbar.message(ret) + else: + self.start_client_playback(ret, True) def spawn_external_viewer(self, data, contenttype): if contenttype: @@ -1246,6 +1268,13 @@ class ConsoleMaster(flow.FlowMaster): self.statusbar.message("") if k == "?": self.view_help() + elif k == "c": + self.path_prompt( + "Client replay: ", + self.state.last_saveload, + self.client_playback_path + ) + k = None elif k == "l": self.prompt("Limit: ", self.state.limit_txt, self.set_limit) self.sync_list_view() @@ -1329,7 +1358,7 @@ class ConsoleMaster(flow.FlowMaster): def shutdown(self): for i in self.state.flow_list: - i.kill() + i.kill(self) controller.Master.shutdown(self) def sync_list_view(self): @@ -1344,7 +1373,7 @@ class ConsoleMaster(flow.FlowMaster): self.sync_list_view() def kill_connection(self, f): - self.state.kill_flow(f) + f.kill(self) def refresh_connection(self, c): if hasattr(self.header, "refresh_connection"): @@ -1372,14 +1401,17 @@ class ConsoleMaster(flow.FlowMaster): f = flow.FlowMaster.handle_error(self, r) if f: self.process_flow(f, r) + return f def handle_request(self, r): f = flow.FlowMaster.handle_request(self, r) if f: self.process_flow(f, r) + return f def handle_response(self, r): f = flow.FlowMaster.handle_response(self, r) if f: self.process_flow(f, r) + return f diff --git a/libmproxy/flow.py b/libmproxy/flow.py index 3764fc95..fb5c82d7 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -287,13 +287,14 @@ class Flow: return pattern(self.request) return False - def kill(self): + def kill(self, master): self.error = proxy.Error(self.request, "Connection killed") if self.request and not self.request.acked: self.request.ack(None) elif self.response and not self.response.acked: self.response.ack(None) self.intercepting = False + master.handle_error(self.error) def intercept(self): self.intercepting = True @@ -552,8 +553,7 @@ class FlowMaster(controller.Master): pb = self.do_server_playback(f) if not pb: if self.kill_nonreplay: - f.kill() - self.handle_error(f.error) + f.kill(self) else: r.ack() return f diff --git a/test/test_flow.py b/test/test_flow.py index 947d0d59..91e3f5f4 100644 --- a/test/test_flow.py +++ b/test/test_flow.py @@ -168,18 +168,20 @@ class uFlow(libpry.AutoTree): assert f == f2 def test_kill(self): + s = flow.State() + fm = flow.FlowMaster(None, s) f = tutils.tflow() f.request = tutils.treq() f.intercept() assert not f.request.acked - f.kill() + f.kill(fm) assert f.request.acked f.intercept() f.response = tutils.tresp() f.request = f.response.request f.request.ack() assert not f.response.acked - f.kill() + f.kill(fm) assert f.response.acked def test_accept_intercept(self): |