diff options
Diffstat (limited to 'libmproxy')
| -rw-r--r-- | libmproxy/dump.py | 21 | ||||
| -rw-r--r-- | libmproxy/flow.py | 46 |
2 files changed, 44 insertions, 23 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 |