Add saving and loading of complete flows for later replay and analysis.

3 files changed, 77 insertions, 43 deletions

diff --git a/libmproxy/console.py b/libmproxy/console.py
index ff5383e0..ec17b2e9 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 41064829..31b7bc51 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 399cc485..c5c856f8 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):