Merge remote-tracking branch 'taiste/server-replay-pop'

4 files changed, 42 insertions, 13 deletions

diff --git a/libmproxy/cmdline.py b/libmproxy/cmdline.py
index ffd1826a..27819294 100644
--- a/libmproxy/cmdline.py
+++ b/libmproxy/cmdline.py
@@ -45,6 +45,7 @@ def get_common_options(options):
         stickyauth = stickyauth,
         wfile = options.wfile,
         verbosity = options.verbose,
+        nopop = options.nopop,
     )
 
 
@@ -183,6 +184,12 @@ def common_options(parser):
         help= "Disable response refresh, "
         "which updates times in cookies and headers for replayed responses."
     )
+    group.add_option(
+        "--no-pop",
+        action="store_true", dest="nopop", default=False,
+        help="Disable response pop from response flow. "
+        "This makes it possible to replay same response multiple times."
+    )
     parser.add_option_group(group)
 
     proxy.certificate_option_group(parser)
diff --git a/libmproxy/console/__init__.py b/libmproxy/console/__init__.py
index 7d936892..5d9c5da2 100644
--- a/libmproxy/console/__init__.py
+++ b/libmproxy/console/__init__.py
@@ -126,7 +126,10 @@ class StatusBar(common.WWrap):
         if self.master.server_playback:
             r.append("[")
             r.append(("heading_key", "splayback"))
-            r.append(":%s to go]"%self.master.server_playback.count())
+            if self.master.nopop:
+                r.append(":%s in file]"%self.master.server_playback.count())
+            else:
+                r.append(":%s to go]"%self.master.server_playback.count())
         if self.master.state.intercept_txt:
             r.append("[")
             r.append(("heading_key", "i"))
@@ -297,6 +300,7 @@ class Options(object):
         "stickyauth",
         "verbosity",
         "wfile",
+        "nopop",
     ]
     def __init__(self, **kwargs):
         for k, v in kwargs.items():
@@ -350,6 +354,7 @@ class ConsoleMaster(flow.FlowMaster):
         self.anticomp = options.anticomp
         self.killextra = options.kill
         self.rheaders = options.rheaders
+        self.nopop = options.nopop
 
         self.eventlog = options.eventlog
         self.eventlist = urwid.SimpleListWalker([])
@@ -422,7 +427,7 @@ class ConsoleMaster(flow.FlowMaster):
             self.start_server_playback(
                 ret,
                 self.killextra, self.rheaders,
-                False
+                False, self.nopop
             )
 
     def spawn_editor(self, data):
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index c4bf35a5..e7af996c 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -768,12 +768,12 @@ class ClientPlaybackState:
 
 
 class ServerPlaybackState:
-    def __init__(self, headers, flows, exit):
+    def __init__(self, headers, flows, exit, nopop):
         """
             headers: Case-insensitive list of request headers that should be
             included in request-response matching.
         """
-        self.headers, self.exit = headers, exit
+        self.headers, self.exit, self.nopop = headers, exit, nopop
         self.fmap = {}
         for i in flows:
             if i.response:
@@ -815,7 +815,12 @@ class ServerPlaybackState:
         l = self.fmap.get(self._hash(request))
         if not l:
             return None
-        return l.pop(0)
+
+        if self.nopop: 
+            return l[0]
+        else: 
+            return l.pop(0)
+
 
 
 class StickyCookieState:
@@ -1251,12 +1256,12 @@ class FlowMaster(controller.Master):
     def stop_client_playback(self):
         self.client_playback = None
 
-    def start_server_playback(self, flows, kill, headers, exit):
+    def start_server_playback(self, flows, kill, headers, exit, nopop):
         """
             flows: List of flows.
             kill: Boolean, should we kill requests not part of the replay?
         """
-        self.server_playback = ServerPlaybackState(headers, flows, exit)
+        self.server_playback = ServerPlaybackState(headers, flows, exit, nopop)
         self.kill_nonreplay = kill
 
     def stop_server_playback(self):
diff --git a/test/test_flow.py b/test/test_flow.py
index 67dfe3c2..b6818960 100644
--- a/test/test_flow.py
+++ b/test/test_flow.py
@@ -87,7 +87,7 @@ class uClientPlaybackState(libpry.AutoTree):
 
 class uServerPlaybackState(libpry.AutoTree):
     def test_hash(self):
-        s = flow.ServerPlaybackState(None, [], False)
+        s = flow.ServerPlaybackState(None, [], False, False)
         r = tutils.tflow()
         r2 = tutils.tflow()
 
@@ -99,7 +99,7 @@ class uServerPlaybackState(libpry.AutoTree):
         assert s._hash(r) != s._hash(r2)
 
     def test_headers(self):
-        s = flow.ServerPlaybackState(["foo"], [], False)
+        s = flow.ServerPlaybackState(["foo"], [], False, False)
         r = tutils.tflow_full()
         r.request.headers["foo"] = ["bar"]
         r2 = tutils.tflow_full()
@@ -120,7 +120,7 @@ class uServerPlaybackState(libpry.AutoTree):
         r2 = tutils.tflow_full()
         r2.request.headers["key"] = ["two"]
 
-        s = flow.ServerPlaybackState(None, [r, r2], False)
+        s = flow.ServerPlaybackState(None, [r, r2], False, False)
         assert s.count() == 2
         assert len(s.fmap.keys()) == 1
 
@@ -134,6 +134,18 @@ class uServerPlaybackState(libpry.AutoTree):
 
         assert not s.next_flow(r)
 
+    def test_load_with_nopop(self):
+        r = tutils.tflow_full()
+        r.request.headers["key"] = ["one"]
+
+        r2 = tutils.tflow_full()
+        r2.request.headers["key"] = ["two"]
+
+        s = flow.ServerPlaybackState(None, [r, r2], False, True)
+
+        assert s.count() == 2
+        n = s.next_flow(r)
+        assert s.count() == 2
 
 class uFlow(libpry.AutoTree):
     def test_copy(self):
@@ -547,7 +559,7 @@ class uFlowMaster(libpry.AutoTree):
         f = tutils.tflow_full()
         pb = [tutils.tflow_full(), f]
         fm = flow.FlowMaster(None, s)
-        assert not fm.start_server_playback(pb, False, [], False)
+        assert not fm.start_server_playback(pb, False, [], False, False)
         assert not fm.start_client_playback(pb, False)
 
         q = Queue.Queue()
@@ -568,10 +580,10 @@ class uFlowMaster(libpry.AutoTree):
         fm.refresh_server_playback = True
         assert not fm.do_server_playback(tutils.tflow())
 
-        fm.start_server_playback(pb, False, [], False)
+        fm.start_server_playback(pb, False, [], False, False)
         assert fm.do_server_playback(tutils.tflow())
 
-        fm.start_server_playback(pb, False, [], True)
+        fm.start_server_playback(pb, False, [], True, False)
         r = tutils.tflow()
         r.request.content = "gibble"
         assert not fm.do_server_playback(r)