From 53fadd05f4e015657cf8815d9dedc6486a4b3f02 Mon Sep 17 00:00:00 2001
From: Maximilian Hils <git@maximilianhils.com>
Date: Sat, 7 Feb 2015 16:26:19 +0100
Subject: fix #451

---
 libmproxy/flow.py          |  6 ++----
 libmproxy/protocol/http.py | 38 ++++++++++++++++++++++++++------------
 libmproxy/script.py        | 10 +++++++---
 3 files changed, 35 insertions(+), 19 deletions(-)

(limited to 'libmproxy')

diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index d42bbb12..dd9900e9 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -824,12 +824,10 @@ 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.intercepted:
             return "Can't replay while intercepting..."
         if f.request.content == http.CONTENT_MISSING:
@@ -845,7 +843,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..9cf402b7 100644
--- a/libmproxy/script.py
+++ b/libmproxy/script.py
@@ -36,7 +36,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 +139,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):
-- 
cgit v1.2.3