diff options
| -rw-r--r-- | libmproxy/dump.py | 21 | ||||
| -rw-r--r-- | libmproxy/flow.py | 46 | ||||
| -rw-r--r-- | test/test_flow.py | 16 |
3 files changed, 52 insertions, 31 deletions
diff --git a/libmproxy/dump.py b/libmproxy/dump.py index 7f6f1e7c..66bb5206 100644 --- a/libmproxy/dump.py +++ b/libmproxy/dump.py @@ -32,6 +32,11 @@ class DumpMaster(flow.FlowMaster): else: self.filt = None + if self.o.response_script: + self.set_response_script(self.o.response_script) + if self.o.request_script: + self.set_request_script(self.o.request_script) + if options.wfile: path = os.path.expanduser(options.wfile) try: @@ -47,7 +52,7 @@ class DumpMaster(flow.FlowMaster): flows = list(flow.FlowReader(f).stream()) except IOError, v: raise DumpError(v.strerror) - self.start_playback(flows) + self.start_playback(flows, options.kill) def _runscript(self, f, script): try: @@ -65,17 +70,7 @@ class DumpMaster(flow.FlowMaster): def handle_request(self, r): f = flow.FlowMaster.handle_request(self, r) - if self.o.request_script: - self._runscript(f, self.o.request_script) - - if self.o.replay: - pb = self.playback(f) - if not pb: - if self.o.kill: - self.state.kill_flow(f) - else: - r.ack() - else: + if f: r.ack() def indent(self, n, t): @@ -85,8 +80,6 @@ class DumpMaster(flow.FlowMaster): def handle_response(self, msg): f = flow.FlowMaster.handle_response(self, msg) if f: - if self.o.response_script: - self._runscript(f, self.o.response_script) msg.ack() if self.filt and not f.match(self.filt): return diff --git a/libmproxy/flow.py b/libmproxy/flow.py index 996b45cd..80a88708 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -35,7 +35,7 @@ class ServerPlaybackState: def __init__(self): self.fmap = {} - def __len__(self): + def count(self): return sum([len(i) for i in self.fmap.values()]) def load(self, flows): @@ -329,19 +329,35 @@ class FlowMaster(controller.Master): def __init__(self, server, state): controller.Master.__init__(self, server) self.state = state - self._playback_state = None + self.playback = None + self.scripts = {} + self.kill_nonreplay = False - def start_playback(self, flows): - self._playback_state = ServerPlaybackState() - self._playback_state.load(flows) + def _runscript(self, f, script): + return f.run_script(script) - def playback(self, flow): + def set_response_script(self, s): + self.scripts["response"] = s + + def set_request_script(self, s): + self.scripts["request"] = s + + def start_playback(self, flows, kill): + """ + flows: A list of flows. + kill: Boolean, should we kill requests not part of the replay? + """ + self.playback = ServerPlaybackState() + self.playback.load(flows) + self.kill_nonreplay = kill + + def do_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 self.playback: + rflow = self.playback.next_flow(flow) if not rflow: return None response = proxy.Response.from_state(flow.request, rflow.response.get_state()) @@ -365,12 +381,24 @@ class FlowMaster(controller.Master): return f def handle_request(self, r): - return self.state.add_request(r) + f = self.state.add_request(r) + if "request" in self.scripts: + self._runscript(f, self.scripts["request"]) + if self.playback: + pb = self.do_playback(f) + if not pb: + if self.kill_nonreplay: + self.state.kill_flow(f) + else: + r.ack() + return f def handle_response(self, r): f = self.state.add_response(r) if not f: r.ack() + if "response" in self.scripts: + self._runscript(f, self.scripts["response"]) return f diff --git a/test/test_flow.py b/test/test_flow.py index 4bfd5ad9..9963ad18 100644 --- a/test/test_flow.py +++ b/test/test_flow.py @@ -26,16 +26,16 @@ class uServerPlaybackState(libpry.AutoTree): r2.request.headers["key"] = ["two"] s.load([r, r2]) - assert len(s) == 2 + assert s.count() == 2 assert len(s.fmap.keys()) == 1 n = s.next_flow(r) assert n.request.headers["key"] == ["one"] - assert len(s) == 1 + assert s.count() == 1 n = s.next_flow(r) assert n.request.headers["key"] == ["two"] - assert len(s) == 0 + assert s.count() == 0 assert not s.next_flow(r) @@ -317,15 +317,15 @@ class uFlowMaster(libpry.AutoTree): pb = [f] fm = flow.FlowMaster(None, s) - assert not fm.playback(utils.tflow()) + assert not fm.do_playback(utils.tflow()) - fm.start_playback(pb) - assert fm.playback(utils.tflow()) + fm.start_playback(pb, False) + assert fm.do_playback(utils.tflow()) - fm.start_playback(pb) + fm.start_playback(pb, False) r = utils.tflow() r.request.content = "gibble" - assert not fm.playback(r) + assert not fm.do_playback(r) |