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

3 files changed, 24 insertions, 7 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):