From 613e9a298ec2fc82a4811c35f806d129ff883f2a Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Sun, 15 May 2011 11:22:35 +1200
Subject: [PATCH] Add a new flow loading mechanism.

We now simulate the normal connection flow when we load flows. That means
that we can run scripts, hooks, sticky cookies, etc.
---
 libmproxy/console.py |  5 ++---
 libmproxy/flow.py    | 25 +++++++++++++++++++++----
 test/test_flow.py    | 21 +++++++++++++++++++++
 test/tutils.py       |  7 +++++++
 4 files changed, 51 insertions(+), 7 deletions(-)

diff --git a/libmproxy/console.py b/libmproxy/console.py
index 63dbf4f03..9e065e4d8 100644
--- a/libmproxy/console.py
+++ b/libmproxy/console.py
@@ -1127,11 +1127,10 @@ class ConsoleMaster(flow.FlowMaster):
         try:
             f = file(path, "r")
             fr = flow.FlowReader(f)
-            data = list(fr.stream())
-            f.close()
         except IOError, v:
             return v.strerror
-        self.state.load_flows(data)
+        flow.FlowMaster.load_flows(self, fr)
+        f.close()
         if self.conn_list_view:
             self.sync_list_view()
             self.focus_current()
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 387c49f0b..b53794c9b 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -542,6 +542,20 @@ class FlowMaster(controller.Master):
 
         controller.Master.tick(self, q)
 
+    def load_flows(self, fr):
+        """
+            Load flows from a FlowReader object.
+        """
+        for i in fr.stream():
+            if i.request:
+                f = self.state.add_request(i.request)
+                self.process_new_request(f)
+            if i.response:
+                f = self.state.add_response(i.response)
+                self.process_new_response(f)
+            if i.error:
+                f = self.state.add_error(i.error)
+
     def process_new_request(self, f):
         if self.stickycookie_state:
             self.stickycookie_state.handle_request(f)
@@ -560,6 +574,12 @@ class FlowMaster(controller.Master):
                 else:
                     f.request.ack()
 
+    def process_new_response(self, f):
+        if self.stickycookie_state:
+            self.stickycookie_state.handle_response(f)
+        if "response" in self.scripts:
+            self._runscript(f, self.scripts["response"])
+
     def replay_request(self, f):
         """
             Returns None if successful, or error message if not.
@@ -605,10 +625,7 @@ class FlowMaster(controller.Master):
             self.client_playback.clear(f)
         if not f:
             r.ack()
-        if self.stickycookie_state:
-            self.stickycookie_state.handle_response(f)
-        if "response" in self.scripts:
-            self._runscript(f, self.scripts["response"])
+        self.process_new_response(f)
         return f
 
 
diff --git a/test/test_flow.py b/test/test_flow.py
index a64894622..b7a368d27 100644
--- a/test/test_flow.py
+++ b/test/test_flow.py
@@ -369,6 +369,19 @@ class uState(libpry.AutoTree):
 
 
 class uSerialize(libpry.AutoTree):
+    def _treader(self):
+        sio = StringIO()
+        w = flow.FlowWriter(sio)
+        for i in range(3):
+            f = tutils.tflow_full()
+            w.add(f)
+        for i in range(3):
+            f = tutils.tflow_err()
+            w.add(f)
+
+        sio.seek(0)
+        return flow.FlowReader(sio)
+
     def test_roundtrip(self):
         sio = StringIO()
         f = tutils.tflow()
@@ -381,6 +394,14 @@ class uSerialize(libpry.AutoTree):
         assert len(l) == 1
         assert l[0] == f
 
+    def test_load_flows(self):
+        r = self._treader()
+        s = flow.State()
+        fm = flow.FlowMaster(None, s)
+        fm.load_flows(r)
+        assert len(s.flow_list) == 6
+
+
     def test_error(self):
         sio = StringIO()
         sio.write("bogus")
diff --git a/test/tutils.py b/test/tutils.py
index fea7f224e..2327bdb98 100644
--- a/test/tutils.py
+++ b/test/tutils.py
@@ -32,6 +32,13 @@ def tflow_full():
     return f
 
 
+def tflow_err():
+    r = treq()
+    f = flow.Flow(r)
+    f.error = proxy.Error(r, "error")
+    return f
+
+
 # Yes, the random ports are horrible. During development, sockets are often not
 # properly closed during error conditions, which means you have to wait until
 # you can re-bind to the same port. This is a pain in the ass, so we just pick