From 5da4099ddfee61b717c23a8c4227947ecc2f6c53 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Sat, 5 Mar 2011 13:03:26 +1300 Subject: [PATCH] Massage namespace to make room for client replay. Mostly replay -> server_replay --- libmproxy/dump.py | 8 ++++---- libmproxy/flow.py | 32 ++++++++++++++++++++++++-------- mitmdump | 4 ++-- test/test_dump.py | 6 +++--- test/test_flow.py | 19 ++++++++++++------- 5 files changed, 45 insertions(+), 24 deletions(-) diff --git a/libmproxy/dump.py b/libmproxy/dump.py index 50d4dc4b2..54be4933f 100644 --- a/libmproxy/dump.py +++ b/libmproxy/dump.py @@ -9,7 +9,7 @@ class Options(object): "kill", "request_script", "response_script", - "replay", + "server_replay", "verbosity", "wfile", "rheaders", @@ -64,14 +64,14 @@ class DumpMaster(flow.FlowMaster): except IOError, v: raise DumpError(v.strerror) - if options.replay: - path = os.path.expanduser(options.replay) + if options.server_replay: + path = os.path.expanduser(options.server_replay) try: f = file(path, "r") flows = list(flow.FlowReader(f).stream()) except IOError, v: raise DumpError(v.strerror) - self.start_playback(flows, options.kill, options.rheaders) + self.start_server_playback(flows, options.kill, options.rheaders) def _runscript(self, f, script): diff --git a/libmproxy/flow.py b/libmproxy/flow.py index a4e0a79c8..2e734500f 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -409,7 +409,8 @@ class FlowMaster(controller.Master): def __init__(self, server, state): controller.Master.__init__(self, server) self.state = state - self.playback = None + self.server_playback = None + self.client_playback = None self.scripts = {} self.kill_nonreplay = False self.stickycookie_state = False @@ -434,21 +435,27 @@ class FlowMaster(controller.Master): else: self.stickycookie_state = None - def start_playback(self, flows, kill, headers): + def start_client_playback(self, flows): + """ + flows: A list of flows. + """ + self.client_playback = ClientPlaybackState(flows) + + def start_server_playback(self, flows, kill, headers): """ flows: A list of flows. kill: Boolean, should we kill requests not part of the replay? """ - self.playback = ServerPlaybackState(headers, flows) + self.server_playback = ServerPlaybackState(headers, flows) self.kill_nonreplay = kill - def do_playback(self, flow): + def do_server_playback(self, flow): """ This method should be called by child classes in the handle_request handler. Returns True if playback has taken place, None if not. """ - if self.playback: - rflow = self.playback.next_flow(flow) + if self.server_playback: + rflow = self.server_playback.next_flow(flow) if not rflow: return None response = proxy.Response.from_state(flow.request, rflow.response.get_state()) @@ -458,6 +465,11 @@ class FlowMaster(controller.Master): return True return None + def tick(self, q): + if self.client_playback: + self.client_playback.tick() + controller.Master.tick(self, q) + def handle_clientconnect(self, r): self.state.clientconnect(r) r.ack() @@ -468,6 +480,8 @@ class FlowMaster(controller.Master): def handle_error(self, r): f = self.state.add_error(r) + if self.client_playback: + self.client_playback.clear(f) r.ack() return f @@ -477,8 +491,8 @@ class FlowMaster(controller.Master): self.stickycookie_state.handle_request(f) if "request" in self.scripts: self._runscript(f, self.scripts["request"]) - if self.playback: - pb = self.do_playback(f) + if self.server_playback: + pb = self.do_server_playback(f) if not pb: if self.kill_nonreplay: self.state.kill_flow(f) @@ -488,6 +502,8 @@ class FlowMaster(controller.Master): def handle_response(self, r): f = self.state.add_response(r) + if self.client_playback: + self.client_playback.clear(f) if not f: r.ack() if self.stickycookie_state: diff --git a/mitmdump b/mitmdump index f92444b8e..63d942af8 100755 --- a/mitmdump +++ b/mitmdump @@ -58,7 +58,7 @@ if __name__ == '__main__': group = OptionGroup(parser, "Server Replay") - group.add_option("-r", action="store", dest="replay", default=None, metavar="PATH", + group.add_option("-r", action="store", dest="server_replay", default=None, metavar="PATH", help="Replay server responses from a saved file.") group.add_option("-k", "--kill", action="store_true", dest="kill", default=False, @@ -89,7 +89,7 @@ if __name__ == '__main__': wfile = options.wfile, request_script = options.request_script, response_script = options.response_script, - replay = options.replay, + server_replay = options.server_replay, kill = options.kill, rheaders = options.rheaders, stickycookie = stickycookie diff --git a/test/test_dump.py b/test/test_dump.py index cfa690b17..0a42c8c3c 100644 --- a/test/test_dump.py +++ b/test/test_dump.py @@ -36,7 +36,7 @@ class uDumpMaster(libpry.AutoTree): def test_replay(self): cs = StringIO() - o = dump.Options(replay="nonexistent", kill=True) + o = dump.Options(server_replay="nonexistent", kill=True) libpry.raises(dump.DumpError, dump.DumpMaster, None, o, None, outfile=cs) t = self.tmpdir() @@ -48,13 +48,13 @@ class uDumpMaster(libpry.AutoTree): fw.add(t) f.close() - o = dump.Options(replay=p, kill=True) + o = dump.Options(server_replay=p, kill=True) m = dump.DumpMaster(None, o, None, outfile=cs) self._cycle(m, "content") self._cycle(m, "content") - o = dump.Options(replay=p, kill=False) + o = dump.Options(server_replay=p, kill=False) m = dump.DumpMaster(None, o, None, outfile=cs) self._cycle(m, "nonexistent") diff --git a/test/test_flow.py b/test/test_flow.py index f9fb58e83..bd8c01a32 100644 --- a/test/test_flow.py +++ b/test/test_flow.py @@ -375,7 +375,7 @@ class uFlowMaster(libpry.AutoTree): err = proxy.Error(f.request, "msg") fm.handle_error(err) - def test_replay(self): + def test_server_playback(self): s = flow.State() f = utils.tflow() @@ -383,15 +383,21 @@ class uFlowMaster(libpry.AutoTree): pb = [f] fm = flow.FlowMaster(None, s) - assert not fm.do_playback(utils.tflow()) + assert not fm.do_server_playback(utils.tflow()) - fm.start_playback(pb, False, []) - assert fm.do_playback(utils.tflow()) + fm.start_server_playback(pb, False, []) + assert fm.do_server_playback(utils.tflow()) - fm.start_playback(pb, False, []) + fm.start_server_playback(pb, False, []) r = utils.tflow() r.request.content = "gibble" - assert not fm.do_playback(r) + assert not fm.do_server_playback(r) + + def test_client_playback(self): + s = flow.State() + fm = flow.FlowMaster(None, s) + pb = [utils.tflow_full()] + fm.start_client_playback(pb) def test_stickycookie(self): s = flow.State() @@ -424,5 +430,4 @@ tests = [ uState(), uSerialize(), uFlowMaster() - ]