diff options
| -rw-r--r-- | mitmproxy/flow/master.py | 52 | ||||
| -rw-r--r-- | test/mitmproxy/mastertest.py | 2 | ||||
| -rw-r--r-- | test/mitmproxy/test_flow.py | 8 |
3 files changed, 40 insertions, 22 deletions
diff --git a/mitmproxy/flow/master.py b/mitmproxy/flow/master.py index 31475f5b..289102a1 100644 --- a/mitmproxy/flow/master.py +++ b/mitmproxy/flow/master.py @@ -103,9 +103,10 @@ class FlowMaster(controller.Master): except script.ScriptException as e: self.add_event("Script error:\n{}".format(e), "error") - def run_script_hook(self, name, *args, **kwargs): + def run_scripts(self, name, msg): for script_obj in self.scripts: - self._run_single_script_hook(script_obj, name, *args, **kwargs) + if not msg.reply.acked: + self._run_single_script_hook(script_obj, name, msg) def get_ignore_filter(self): return self.server.config.check_ignore.patterns @@ -373,28 +374,28 @@ class FlowMaster(controller.Master): @controller.handler def clientconnect(self, root_layer): - self.run_script_hook("clientconnect", root_layer) + self.run_scripts("clientconnect", root_layer) @controller.handler def clientdisconnect(self, root_layer): - self.run_script_hook("clientdisconnect", root_layer) + self.run_scripts("clientdisconnect", root_layer) @controller.handler def serverconnect(self, server_conn): - self.run_script_hook("serverconnect", server_conn) + self.run_scripts("serverconnect", server_conn) @controller.handler def serverdisconnect(self, server_conn): - self.run_script_hook("serverdisconnect", server_conn) + self.run_scripts("serverdisconnect", server_conn) @controller.handler def next_layer(self, top_layer): - self.run_script_hook("next_layer", top_layer) + self.run_scripts("next_layer", top_layer) @controller.handler def error(self, f): self.state.update_flow(f) - self.run_script_hook("error", f) + self.run_scripts("error", f) if self.client_playback: self.client_playback.clear(f) return f @@ -416,10 +417,14 @@ class FlowMaster(controller.Master): if f not in self.state.flows: # don't add again on replay self.state.add_flow(f) self.active_flows.add(f) - self.replacehooks.run(f) - self.setheaders.run(f) - self.process_new_request(f) - self.run_script_hook("request", f) + if not f.reply.acked: + self.replacehooks.run(f) + if not f.reply.acked: + self.setheaders.run(f) + if not f.reply.acked: + self.process_new_request(f) + if not f.reply.acked: + self.run_scripts("request", f) return f @controller.handler @@ -430,18 +435,21 @@ class FlowMaster(controller.Master): except netlib.exceptions.HttpException: f.reply.kill() return - self.run_script_hook("responseheaders", f) + self.run_scripts("responseheaders", f) return f @controller.handler def response(self, f): self.active_flows.discard(f) self.state.update_flow(f) - self.replacehooks.run(f) - self.setheaders.run(f) - self.run_script_hook("response", f) - if self.client_playback: - self.client_playback.clear(f) + if not f.reply.acked: + self.replacehooks.run(f) + if not f.reply.acked: + self.setheaders.run(f) + self.run_scripts("response", f) + if not f.reply.acked: + if self.client_playback: + self.client_playback.clear(f) self.process_new_response(f) if self.stream: self.stream.add(f) @@ -487,11 +495,11 @@ class FlowMaster(controller.Master): # TODO: This would break mitmproxy currently. # self.state.add_flow(flow) self.active_flows.add(flow) - self.run_script_hook("tcp_open", flow) + self.run_scripts("tcp_open", flow) @controller.handler def tcp_message(self, flow): - self.run_script_hook("tcp_message", flow) + self.run_scripts("tcp_message", flow) message = flow.messages[-1] direction = "->" if message.from_client else "<-" self.add_event("{client} {direction} tcp {direction} {server}".format( @@ -507,14 +515,14 @@ class FlowMaster(controller.Master): repr(flow.server_conn.address), flow.error ), "info") - self.run_script_hook("tcp_error", flow) + self.run_scripts("tcp_error", flow) @controller.handler def tcp_close(self, flow): self.active_flows.discard(flow) if self.stream: self.stream.add(flow) - self.run_script_hook("tcp_close", flow) + self.run_scripts("tcp_close", flow) def shutdown(self): super(FlowMaster, self).shutdown() diff --git a/test/mitmproxy/mastertest.py b/test/mitmproxy/mastertest.py index 9bb8826d..4d04f337 100644 --- a/test/mitmproxy/mastertest.py +++ b/test/mitmproxy/mastertest.py @@ -16,7 +16,9 @@ class MasterTest: master.request(f) if not f.error: f.response = models.HTTPResponse.wrap(netlib.tutils.tresp(content=content)) + f.reply.acked = False f = master.response(f) + f.client_conn.reply.acked = False master.clientdisconnect(f.client_conn) return f diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py index 1b1f03f9..af8256c4 100644 --- a/test/mitmproxy/test_flow.py +++ b/test/mitmproxy/test_flow.py @@ -807,17 +807,22 @@ class TestFlowMaster: fm.load_script(tutils.test_data.path("data/scripts/all.py")) f = tutils.tflow(resp=True) + f.client_conn.acked = False fm.clientconnect(f.client_conn) assert fm.scripts[0].ns["log"][-1] == "clientconnect" + f.server_conn.acked = False fm.serverconnect(f.server_conn) assert fm.scripts[0].ns["log"][-1] == "serverconnect" + f.reply.acked = False fm.request(f) assert fm.scripts[0].ns["log"][-1] == "request" + f.reply.acked = False fm.response(f) assert fm.scripts[0].ns["log"][-1] == "response" # load second script fm.load_script(tutils.test_data.path("data/scripts/all.py")) assert len(fm.scripts) == 2 + f.server_conn.reply.acked = False fm.clientdisconnect(f.server_conn) assert fm.scripts[0].ns["log"][-1] == "clientdisconnect" assert fm.scripts[1].ns["log"][-1] == "clientdisconnect" @@ -828,6 +833,7 @@ class TestFlowMaster: fm.load_script(tutils.test_data.path("data/scripts/all.py")) f.error = tutils.terr() + f.reply.acked = False fm.error(f) assert fm.scripts[0].ns["log"][-1] == "error" @@ -977,10 +983,12 @@ class TestFlowMaster: f = tutils.tflow(resp=True) f.response.headers["set-cookie"] = "foo=bar" fm.request(f) + f.reply.acked = False fm.response(f) assert fm.stickycookie_state.jar assert "cookie" not in f.request.headers f = f.copy() + f.reply.acked = False fm.request(f) assert f.request.headers["cookie"] == "foo=bar" |