diff options
-rw-r--r-- | libmproxy/cmdline.py | 7 | ||||
-rw-r--r-- | libmproxy/console/__init__.py | 9 | ||||
-rw-r--r-- | libmproxy/flow.py | 15 | ||||
-rw-r--r-- | test/test_flow.py | 24 |
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) |