diff options
| -rw-r--r-- | libmproxy/console/__init__.py | 4 | ||||
| -rw-r--r-- | libmproxy/console/connlist.py | 5 | ||||
| -rw-r--r-- | libmproxy/console/connview.py | 6 | ||||
| -rw-r--r-- | libmproxy/flow.py | 37 | ||||
| -rw-r--r-- | test/test_flow.py | 20 |
5 files changed, 63 insertions, 9 deletions
diff --git a/libmproxy/console/__init__.py b/libmproxy/console/__init__.py index 6fc37a47..ca80d013 100644 --- a/libmproxy/console/__init__.py +++ b/libmproxy/console/__init__.py @@ -516,9 +516,7 @@ class ConsoleMaster(flow.FlowMaster): def focus_current(self): if self.currentflow: try: - ids = [id(i) for i in self.state.view] - idx = ids.index(id(self.currentflow)) - self.conn_list_view.set_focus(idx) + self.conn_list_view.set_focus(self.state.index(self.currentflow)) except (IndexError, ValueError): pass diff --git a/libmproxy/console/connlist.py b/libmproxy/console/connlist.py index 957e797b..f54b8061 100644 --- a/libmproxy/console/connlist.py +++ b/libmproxy/console/connlist.py @@ -8,6 +8,7 @@ def _mkhelp(): ("a", "accept this intercepted connection"), ("C", "clear connection list or eventlog"), ("d", "delete flow"), + ("D", "duplicate flow"), ("e", "toggle eventlog"), ("l", "set limit filter pattern"), ("L", "load saved flows"), @@ -101,6 +102,10 @@ class ConnectionItem(common.WWrap): self.flow.kill(self.master) self.state.delete_flow(self.flow) self.master.sync_list_view() + elif key == "D": + f = self.master.duplicate_flow(self.flow) + self.master.currentflow = f + self.master.focus_current() elif key == "l": self.master.prompt("Limit: ", self.state.limit_txt, self.master.set_limit) self.master.sync_list_view() diff --git a/libmproxy/console/connview.py b/libmproxy/console/connview.py index d668bfe2..ed038a11 100644 --- a/libmproxy/console/connview.py +++ b/libmproxy/console/connview.py @@ -10,6 +10,7 @@ def _mkhelp(): ("a", "accept this intercepted connection"), ("b", "save request/response body"), ("d", "delete flow"), + ("D", "duplicate flow"), ("e", "edit request/response"), ("m", "change body display mode"), (None, @@ -463,6 +464,11 @@ class ConnectionView(common.WWrap): f = self.flow f.kill(self.master) self.state.delete_flow(f) + elif key == "D": + f = self.master.duplicate_flow(self.flow) + self.master.view_flow(f) + self.master.currentflow = f + self.master.statusbar.message("Duplicated.") elif key == "e": if self.state.view_flow_mode == common.VIEW_FLOW_REQUEST: self.master.prompt_onekey( diff --git a/libmproxy/flow.py b/libmproxy/flow.py index 81199035..db4cde5e 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -851,6 +851,17 @@ class Flow: self.intercepting = False self._backup = None + def copy(self): + rc = self.request.copy() + f = Flow(rc) + if self.response: + f.response = self.response.copy() + f.response.request = rc + if self.error: + f.error = self.error.copy() + f.error.request = rc + return f + @classmethod def _from_state(klass, state): f = klass(None) @@ -997,6 +1008,9 @@ class State(object): def flow_count(self): return len(self._flow_map) + def index(self, f): + return self._flow_list.index(f) + def active_flow_count(self): c = 0 for i in self._flow_list: @@ -1011,6 +1025,7 @@ class State(object): f = Flow(req) self._flow_list.append(f) self._flow_map[req] = f + assert len(self._flow_list) == len(self._flow_map) if f.match(self._limit): self.view.append(f) return f @@ -1235,17 +1250,27 @@ class FlowMaster(controller.Master): return controller.Master.tick(self, q) + def duplicate_flow(self, f): + return self.load_flow(f.copy()) + + def load_flow(self, f): + """ + Loads a flow, and returns a new flow object. + """ + if f.request: + fr = self.handle_request(f.request) + if f.response: + self.handle_response(f.response) + if f.error: + self.handle_error(f.error) + return fr + def load_flows(self, fr): """ Load flows from a FlowReader object. """ for i in fr.stream(): - if i.request: - self.handle_request(i.request) - if i.response: - self.handle_response(i.response) - if i.error: - self.handle_error(i.error) + self.load_flow(i) def process_new_request(self, f): if self.stickycookie_state: diff --git a/test/test_flow.py b/test/test_flow.py index 94c4614d..3e126263 100644 --- a/test/test_flow.py +++ b/test/test_flow.py @@ -136,6 +136,14 @@ class uServerPlaybackState(libpry.AutoTree): class uFlow(libpry.AutoTree): + def test_copy(self): + f = tutils.tflow_full() + f2 = f.copy() + assert not f is f2 + assert not f.request is f2.request + assert f.request.headers == f2.request.headers + assert not f.request.headers is f2.request.headers + def test_match(self): f = tutils.tflow() f.response = tutils.tresp() @@ -485,6 +493,17 @@ class uFlowMaster(libpry.AutoTree): fm.handle_error(err) assert fm.script.ns["log"][-1] == "error" + def test_duplicate_flow(self): + s = flow.State() + fm = flow.FlowMaster(None, s) + f = tutils.tflow_full() + fm.load_flow(f) + assert s.flow_count() == 1 + f2 = fm.duplicate_flow(f) + assert f2.response + assert s.flow_count() == 2 + assert s.index(f2) + def test_all(self): s = flow.State() fm = flow.FlowMaster(None, s) @@ -572,6 +591,7 @@ class uFlowMaster(libpry.AutoTree): fm.handle_response(tf.response) assert fm.stickycookie_state.jar assert not "cookie" in tf.request.headers + tf = tf.copy() fm.handle_request(tf.request) assert tf.request.headers["cookie"] == ["foo=bar"] |