diff options
Diffstat (limited to 'libmproxy')
| -rw-r--r-- | libmproxy/flow.py | 16 | ||||
| -rw-r--r-- | libmproxy/protocol/http.py | 38 | ||||
| -rw-r--r-- | libmproxy/script.py | 16 |
3 files changed, 47 insertions, 23 deletions
diff --git a/libmproxy/flow.py b/libmproxy/flow.py index f9e2b94d..49ec5a0b 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -309,10 +309,10 @@ class StickyCookieState: # FIXME: We now know that Cookie.py screws up some cookies with # valid RFC 822/1123 datetime specifications for expiry. Sigh. c = Cookie.SimpleCookie(str(i)) - m = c.values()[0] - k = self.ckey(m, f) - if self.domain_match(f.request.host, k[0]): - self.jar[self.ckey(m, f)] = m + for m in c.values(): + k = self.ckey(m, f) + if self.domain_match(f.request.host, k[0]): + self.jar[k] = m def handle_request(self, f): l = [] @@ -824,12 +824,12 @@ class FlowMaster(controller.Master): if self.stickycookie_state: self.stickycookie_state.handle_response(f) - def replay_request(self, f, block=False): + def replay_request(self, f, block=False, run_scripthooks=True): """ Returns None if successful, or error message if not. """ - if f.live: - return "Can't replay request which is still live..." + if f.live and run_scripthooks: + return "Can't replay live request." if f.intercepted: return "Can't replay while intercepting..." if f.request.content == http.CONTENT_MISSING: @@ -845,7 +845,7 @@ class FlowMaster(controller.Master): rt = http.RequestReplayThread( self.server.config, f, - self.masterq, + self.masterq if run_scripthooks else False, self.should_exit ) rt.start() # pragma: no cover diff --git a/libmproxy/protocol/http.py b/libmproxy/protocol/http.py index bebb4f7b..2f858a7c 100644 --- a/libmproxy/protocol/http.py +++ b/libmproxy/protocol/http.py @@ -1416,20 +1416,31 @@ class RequestReplayThread(threading.Thread): name = "RequestReplayThread" def __init__(self, config, flow, masterq, should_exit): - self.config, self.flow, self.channel = config, flow, controller.Channel(masterq, should_exit) - threading.Thread.__init__(self) + """ + masterqueue can be a queue or None, if no scripthooks should be processed. + """ + self.config, self.flow = config, flow + if masterq: + self.channel = controller.Channel(masterq, should_exit) + else: + self.channel = None + super(RequestReplayThread, self).__init__() def run(self): r = self.flow.request form_out_backup = r.form_out try: self.flow.response = None - request_reply = self.channel.ask("request", self.flow) - if request_reply is None or request_reply == KILL: - raise KillSignal() - elif isinstance(request_reply, HTTPResponse): - self.flow.response = request_reply - else: + + # If we have a channel, run script hooks. + if self.channel: + request_reply = self.channel.ask("request", self.flow) + if request_reply is None or request_reply == KILL: + raise KillSignal() + elif isinstance(request_reply, HTTPResponse): + self.flow.response = request_reply + + if not self.flow.response: # In all modes, we directly connect to the server displayed if self.config.mode == "upstream": server_address = self.config.mode.get_upstream_server(self.flow.client_conn)[2:] @@ -1453,13 +1464,16 @@ class RequestReplayThread(threading.Thread): self.flow.server_conn = server self.flow.response = HTTPResponse.from_stream(server.rfile, r.method, body_size_limit=self.config.body_size_limit) - response_reply = self.channel.ask("response", self.flow) - if response_reply is None or response_reply == KILL: - raise KillSignal() + if self.channel: + response_reply = self.channel.ask("response", self.flow) + if response_reply is None or response_reply == KILL: + raise KillSignal() except (proxy.ProxyError, http.HttpError, tcp.NetLibError) as v: self.flow.error = Error(repr(v)) - self.channel.ask("error", self.flow) + if self.channel: + self.channel.ask("error", self.flow) except KillSignal: + # KillSignal should only be raised if there's a channel in the first place. self.channel.tell("log", proxy.Log("Connection killed", "info")) finally: r.form_out = form_out_backup diff --git a/libmproxy/script.py b/libmproxy/script.py index b559615b..be226004 100644 --- a/libmproxy/script.py +++ b/libmproxy/script.py @@ -20,6 +20,12 @@ class ScriptContext: """ self._master.add_event(message, level) + def kill_flow(self, f): + """ + Kills a flow immediately. No further data will be sent to the client or the server. + """ + f.kill(self._master) + def duplicate_flow(self, f): """ Returns a duplicate of the specified flow. The flow is also @@ -36,7 +42,7 @@ class ScriptContext: Replay the request on the current flow. The response will be added to the flow object. """ - self._master.replay_request(f) + return self._master.replay_request(f, block=True, run_scripthooks=False) @property def app_registry(self): @@ -139,8 +145,12 @@ def _handle_concurrent_reply(fn, o, *args, **kwargs): def run(): fn(*args, **kwargs) - o.reply() # If the script did not call .reply(), we have to do it now. - threading.Thread(target=run, name="ScriptThread").start() + reply_proxy() # If the script did not call .reply(), we have to do it now. + ScriptThread(target=run).start() + + +class ScriptThread(threading.Thread): + name = "ScriptThread" def concurrent(fn): |