From 635f7a971d4bb815c8963ac52187b0c0f4f143d7 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Wed, 26 Jan 2011 16:50:17 +1300 Subject: [PATCH] Add saving and loading of complete flows for later replay and analysis. --- libmproxy/console.py | 76 ++++++++++++++++++++++++-------------------- libmproxy/flow.py | 29 +++++++++++++---- test/test_console.py | 15 ++++++++- 3 files changed, 77 insertions(+), 43 deletions(-) diff --git a/libmproxy/console.py b/libmproxy/console.py index ff5383e04..ec17b2e90 100644 --- a/libmproxy/console.py +++ b/libmproxy/console.py @@ -73,8 +73,8 @@ class ConnectionItem(WWrap): self.intercepting = True self.w = self.get_text() - def get_text(self, nofocus=False): - return urwid.Text(self.flow.get_text(nofocus)) + def get_text(self): + return urwid.Text(self.flow.get_text()) def selectable(self): return True @@ -278,23 +278,6 @@ class ConnectionView(WWrap): self.flow.request.method = i[0].upper() self.master.refresh_connection(self.flow) - def save_connection(self, path): - if not path: - return - if self.viewing == self.REQ: - c = self.flow.request - else: - c = self.flow.response - path = os.path.expanduser(path) - try: - f = file(path, "w") - f.write(str(c.headers)) - f.write("\r\n") - f.write(str(c.content)) - f.close() - except IOError, v: - self.master.statusbar.message(str(v)) - def edit(self, part): if self.viewing == self.REQ: conn = self.flow.request @@ -368,10 +351,7 @@ class ConnectionView(WWrap): self.state.revert(self.flow) self.master.refresh_connection(self.flow) elif key == "S": - if self.viewing == self.REQ: - self.master.prompt("Save request: ", self.save_connection) - else: - self.master.prompt("Save response: ", self.save_connection) + self.master.prompt("Save all: ", self.save_flows) elif key == "v": if self.viewing == self.REQ: conn = self.flow.request @@ -519,7 +499,7 @@ class ConsoleFlow(flow.Flow): class ConsoleState(flow.State): def __init__(self): flow.State.__init__(self) - self.focus = False + self.focus = None self.beep = None def add_browserconnect(self, f): @@ -545,13 +525,6 @@ class ConsoleState(flow.State): self.set_focus(self.focus) return ret - @property - def view(self): - if self.limit: - return [i for i in self.flow_list if i.match(self.limit)] - else: - return self.flow_list[:] - def get_focus(self): if not self.view or self.focus is None: return None, None @@ -580,8 +553,6 @@ class ConsoleState(flow.State): return self.get_from_pos(pos-1) def delete_flow(self, f): - if not f.intercepting: - self.view[self.focus].focus = False ret = flow.State.delete_flow(self, f) self.set_focus(self.focus) return ret @@ -694,6 +665,33 @@ class ConsoleMaster(controller.Master): self.nested = True self.make_view() + def save_flows(self, path): + if not path: + return + data = self.state.dump_flows() + path = os.path.expanduser(path) + try: + f = file(path, "wb") + f.write(data) + f.close() + except IOError, v: + self.statusbar.message(str(v)) + + def load_flows(self, path): + if not path: + return + path = os.path.expanduser(path) + try: + f = file(path, "r") + data = f.read() + f.close() + except IOError, v: + self.statusbar.message(str(v)) + return + self.state.load_flows(data, ConsoleFlow) + self.conn_list_view.set_focus(0) + self.sync_list_view() + def helptext(self): text = [] text.extend([("head", "Global keys:\n")]) @@ -701,13 +699,15 @@ class ConsoleMaster(controller.Master): ("A", "accept all intercepted connections"), ("a", "accept this intercepted connection"), ("B", "set beep filter pattern"), + ("c", "set sticky cookie expression"), ("i", "set interception pattern"), ("j, k", "up, down"), ("l", "set limit filter pattern"), + ("L", "load saved flows"), ("q", "quit / return to connection list"), ("r", "replay request"), - ("s", "set sticky cookie expression"), ("R", "revert changes to request"), + ("S", "save flows matching current limit"), ("page up/down", "page up/down"), ("space", "page down"), ("enter", "view connection"), @@ -923,7 +923,13 @@ class ConsoleMaster(controller.Master): self.view_connlist() else: raise Stop - elif k == "s": + elif k == "S": + self.prompt("Save flows: ", self.save_flows) + k = None + elif k == "L": + self.prompt("Load flows: ", self.load_flows) + k = None + elif k == "c": self.prompt("Sticky cookie: ", self.set_stickycookie) k = None if k: diff --git a/libmproxy/flow.py b/libmproxy/flow.py index 410648296..31b7bc51a 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -2,6 +2,7 @@ This module provides more sophisticated flow tracking. These match requests with their responses, and provide filtering and interception facilities. """ +import json import proxy, threading class ReplayConnection: @@ -27,7 +28,6 @@ class Flow: def __init__(self, connection): self.connection = connection self.request, self.response, self.error = None, None, None - self.waiting = True self.intercepting = False self._backup = None @@ -40,7 +40,7 @@ class Flow: @classmethod def from_state(klass, state): - f = Flow(ReplayConnection) + f = klass(None) if state["request"]: f.request = proxy.Request.from_state(state["request"]) if state["response"]: @@ -63,7 +63,6 @@ class Flow: def revert(self): if self._backup: - self.waiting = False restore = [i.copy() if i else None for i in self._backup] self.connection, self.request, self.response, self.error = restore @@ -133,7 +132,6 @@ class State: if not f: return False f.response = resp - f.waiting = False f.backup() return f @@ -146,16 +144,31 @@ class State: if not f: return None f.error = err - f.waiting = False f.backup() return f + def dump_flows(self): + data = [i.get_state() for i in self.view] + return json.dumps(data) + + def load_flows(self, js, klass): + data = json.loads(js) + data = [klass.from_state(i) for i in data] + self.flow_list.extend(data) + def set_limit(self, limit): """ Limit is a compiled filter expression, or None. """ self.limit = limit + @property + def view(self): + if self.limit: + return tuple([i for i in self.flow_list if i.match(self.limit)]) + else: + return tuple(self.flow_list[:]) + def get_connection(self, itm): if isinstance(itm, (proxy.BrowserConnection, ReplayConnection)): return itm @@ -176,7 +189,8 @@ class State: def delete_flow(self, f): if not f.intercepting: c = self.get_connection(f) - del self.flow_map[c] + if c in self.flow_map: + del self.flow_map[c] self.flow_list.remove(f) return True return False @@ -214,7 +228,8 @@ class State: if f.request: f.backup() conn = self.get_connection(f) - del self.flow_map[conn] + if conn in self.flow_map: + del self.flow_map[conn] rp = ReplayConnection() f.connection = rp f.request.connection = rp diff --git a/test/test_console.py b/test/test_console.py index 399cc4856..c5c856f86 100644 --- a/test/test_console.py +++ b/test/test_console.py @@ -57,7 +57,6 @@ class uState(libpry.AutoTree): resp = tresp(req) assert c.add_response(resp) assert len(c.flow_list) == 1 - assert f.waiting == False assert c.lookup(resp) newresp = tresp() @@ -183,6 +182,20 @@ class uState(libpry.AutoTree): c.clear() assert len(c.flow_list) == 0 + def test_dump_flows(self): + c = console.ConsoleState() + self._add_request(c) + self._add_response(c) + self._add_request(c) + self._add_response(c) + self._add_request(c) + self._add_response(c) + + dump = c.dump_flows() + c.clear() + c.load_flows(dump, console.ConsoleFlow) + assert isinstance(c.flow_list[0], console.ConsoleFlow) + class uFlow(libpry.AutoTree): def test_match(self):