diff options
| -rw-r--r-- | mitmproxy/flow.py | 13 | ||||
| -rw-r--r-- | mitmproxy/script/concurrent.py | 27 |
2 files changed, 19 insertions, 21 deletions
diff --git a/mitmproxy/flow.py b/mitmproxy/flow.py index 8ba9b7f9..eefe1b5e 100644 --- a/mitmproxy/flow.py +++ b/mitmproxy/flow.py @@ -806,7 +806,7 @@ class FlowMaster(controller.Master): response.is_replay = True if self.refresh_server_playback: response.refresh() - flow.reply(response) + flow.response = response if self.server_playback.count() == 0: self.stop_server_playback() return True @@ -906,11 +906,11 @@ class FlowMaster(controller.Master): def load_flows_file(self, path): path = os.path.expanduser(path) try: - f = file(path, "rb") - freader = FlowReader(f) + with open(path, "rb") as f: + freader = FlowReader(f) + return self.load_flows(freader) except IOError as v: raise FlowReadError(v.strerror) - return self.load_flows(freader) def process_new_request(self, f): if self.stickycookie_state: @@ -925,11 +925,8 @@ class FlowMaster(controller.Master): if self.server_playback: pb = self.do_server_playback(f) - if not pb: - if self.kill_nonreplay: + if not pb and self.kill_nonreplay: f.kill(self) - else: - f.reply() def process_new_response(self, f): if self.stickycookie_state: diff --git a/mitmproxy/script/concurrent.py b/mitmproxy/script/concurrent.py index 57ee37de..2f25e78c 100644 --- a/mitmproxy/script/concurrent.py +++ b/mitmproxy/script/concurrent.py @@ -8,22 +8,23 @@ import threading class ReplyProxy(object): - def __init__(self, original_reply, script_thread): - self.original_reply = original_reply + def __init__(self, reply_func, script_thread): + self.reply_func = reply_func self.script_thread = script_thread - self._ignore_call = True - self.lock = threading.Lock() + self.master_reply = None - def __call__(self, *args, **kwargs): - with self.lock: - if self._ignore_call: - self.script_thread.start() - self._ignore_call = False - return - self.original_reply(*args, **kwargs) + def __call__(self, *args): + if self.master_reply is None: + self.master_reply = args + self.script_thread.start() + return + self.reply_func(*args) + + def done(self): + self.reply_func(*self.master_reply) def __getattr__(self, k): - return getattr(self.original_reply, k) + return getattr(self.reply_func, k) def _handle_concurrent_reply(fn, o, *args, **kwargs): @@ -34,7 +35,7 @@ def _handle_concurrent_reply(fn, o, *args, **kwargs): def run(): fn(*args, **kwargs) # If the script did not call .reply(), we have to do it now. - reply_proxy() + reply_proxy.done() script_thread = ScriptThread(target=run) |