diff options
| -rw-r--r-- | libmproxy/dump.py | 13 | ||||
| -rw-r--r-- | libmproxy/flow.py | 21 | ||||
| -rw-r--r-- | libmproxy/proxy.py | 12 | ||||
| -rwxr-xr-x | mitmdump | 4 | ||||
| -rw-r--r-- | test/test_dump.py | 17 | ||||
| -rw-r--r-- | test/test_flow.py | 18 | ||||
| -rw-r--r-- | test/utils.py | 2 |
7 files changed, 85 insertions, 2 deletions
diff --git a/libmproxy/dump.py b/libmproxy/dump.py index f6a7ae7e..d43da44b 100644 --- a/libmproxy/dump.py +++ b/libmproxy/dump.py @@ -10,6 +10,7 @@ class Options(object): "wfile", "request_script", "response_script", + "replay", ] def __init__(self, **kwargs): for k, v in kwargs.items(): @@ -38,6 +39,15 @@ class DumpMaster(flow.FlowMaster): except IOError, v: raise DumpError(v.strerror) + if options.replay: + path = os.path.expanduser(options.replay) + try: + f = file(path, "r") + flows = list(flow.FlowReader(f).stream()) + except IOError, v: + raise DumpError(v.strerror) + self.start_playback(flows) + def _runscript(self, f, script): try: ret = f.run_script(script) @@ -56,7 +66,8 @@ class DumpMaster(flow.FlowMaster): f = flow.FlowMaster.handle_request(self, r) if self.o.request_script: self._runscript(f, self.o.request_script) - r.ack() + if not self.playback(f): + r.ack() def indent(self, n, t): l = str(t).strip().split("\n") diff --git a/libmproxy/flow.py b/libmproxy/flow.py index 16a2714c..4025a30d 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -328,6 +328,27 @@ class FlowMaster(controller.Master): def __init__(self, server, state): controller.Master.__init__(self, server) self.state = state + self._playback_state = None + + def start_playback(self, flows): + self._playback_state = ServerPlaybackState() + self._playback_state.load(flows) + + def 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_state: + rflow = self._playback_state.next_flow(flow) + if not rflow: + return None + response = proxy.Response.from_state(flow.request, rflow.response.get_state()) + response.set_replay() + flow.response = response + flow.request.ack(response) + return True + return None def handle_clientconnect(self, r): self.state.clientconnect(r) diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py index ed1c3d60..5e698c5f 100644 --- a/libmproxy/proxy.py +++ b/libmproxy/proxy.py @@ -265,6 +265,13 @@ class Response(controller.Msg): self.timestamp = timestamp or time.time() self.cached = False controller.Msg.__init__(self) + self.replay = False + + def set_replay(self): + self.replay = True + + def is_replay(self): + return self.replay def load_state(self, state): self.code = state["code"] @@ -308,7 +315,10 @@ class Response(controller.Msg): return self.cached def short(self): - return "%s %s"%(self.code, self.msg) + r = "%s %s"%(self.code, self.msg) + if self.is_replay(): + r = "[replay] " + r + return r def assemble(self): """ @@ -47,6 +47,9 @@ if __name__ == '__main__': parser.add_option("", "--respscript", action="store", dest="response_script", default=None, help="Script to run when a response is recieved.") + parser.add_option("-r", "--replay", + action="store", dest="replay", default=None, + help="Replay server responses from a saved file.") options, args = parser.parse_args() @@ -61,6 +64,7 @@ if __name__ == '__main__': wfile = options.wfile, request_script = options.request_script, response_script = options.response_script, + replay = options.replay, ) if args: filt = " ".join(args) diff --git a/test/test_dump.py b/test/test_dump.py index 46dd6dfd..90540408 100644 --- a/test/test_dump.py +++ b/test/test_dump.py @@ -23,6 +23,23 @@ class uDumpMaster(libpry.AutoTree): self._cycle(m, content) return cs.getvalue() + def test_replay(self): + cs = StringIO() + + o = dump.Options(replay="nonexistent") + libpry.raises(dump.DumpError, dump.DumpMaster, None, o, None, outfile=cs) + + t = self.tmpdir() + p = os.path.join(t, "rep") + f = open(p, "w") + fw = flow.FlowWriter(f) + t = utils.tflow() + fw.add(t) + f.close() + + o = dump.Options(replay=p) + m = dump.DumpMaster(None, o, None, outfile=cs) + def test_options(self): o = dump.Options(verbosity = 2) assert o.verbosity == 2 diff --git a/test/test_flow.py b/test/test_flow.py index adfeda6e..cd88464d 100644 --- a/test/test_flow.py +++ b/test/test_flow.py @@ -309,6 +309,24 @@ class uFlowMaster(libpry.AutoTree): err = proxy.Error(f.request, "msg") fm.handle_error(err) + def test_replay(self): + s = flow.State() + + f = utils.tflow() + f.response = utils.tresp(f.request) + pb = [f] + + fm = flow.FlowMaster(None, s) + assert not fm.playback(utils.tflow()) + + fm.start_playback(pb) + assert fm.playback(utils.tflow()) + + fm.start_playback(pb) + r = utils.tflow() + r.request.content = "gibble" + assert not fm.playback(r) + tests = [ diff --git a/test/utils.py b/test/utils.py index b1dc46d4..12646106 100644 --- a/test/utils.py +++ b/test/utils.py @@ -1,3 +1,4 @@ +import os.path from libmproxy import proxy, utils, filt, flow def treq(conn=None): @@ -20,3 +21,4 @@ def tflow(): r = treq() return flow.Flow(r) + |