From 8ddc3b4ef275d70cdbc9bc49add02ea21cca3583 Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Sat, 18 Feb 2012 23:56:40 +1300
Subject: Add API for duplicating flows.

---
 libmproxy/console/common.py | 28 ++++++++++++++--------------
 libmproxy/flow.py           | 29 +++++++++++++++++++++++------
 test/test_flow.py           | 18 ++++++++++++++++++
 3 files changed, 55 insertions(+), 20 deletions(-)

diff --git a/libmproxy/console/common.py b/libmproxy/console/common.py
index 5f80c250..fbeb83d7 100644
--- a/libmproxy/console/common.py
+++ b/libmproxy/console/common.py
@@ -43,20 +43,20 @@ def format_keyvals(lst, key="key", val="text", indent=0):
             if kv is None:
                 ret.append(urwid.Text(""))
             else:
-                ret.append(
-                    urwid.Columns(
-                        [
-                            ("fixed", indent, urwid.Text("")),
-                            (
-                                "fixed",
-                                maxk,
-                                urwid.Text([(key, kv[0] or "")])
-                            ),
-                            urwid.Text([(val, kv[1])])
-                        ],
-                        dividechars = 2
-                    )
-                )
+                cols = []
+                # This cumbersome construction process is here for a reason:
+                # Urwid < 1.0 barfs if given a fixed size column of size zero.
+                if indent:
+                    cols.append(("fixed", indent, urwid.Text("")))
+                cols.extend([
+                    (
+                        "fixed",
+                        maxk,
+                        urwid.Text([(key, kv[0] or "")])
+                    ),
+                    urwid.Text([(val, kv[1])])
+                ])
+                ret.append(urwid.Columns(cols, dividechars = 2))
     return ret
 
 
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 81199035..8ce25c5c 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -851,6 +851,15 @@ 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()
+        if self.error:
+            f.error = self.error.copy()
+        return f
+
     @classmethod
     def _from_state(klass, state):
         f = klass(None)
@@ -1011,6 +1020,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 +1245,24 @@ 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):
+        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..fe84be53 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,15 @@ 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 s.flow_count() == 2
+
     def test_all(self):
         s = flow.State()
         fm = flow.FlowMaster(None, s)
@@ -572,6 +589,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"]
 
-- 
cgit v1.2.3