diff options
| -rw-r--r-- | examples/redirect_requests.py | 2 | ||||
| -rw-r--r-- | examples/tls_passthrough.py | 2 | ||||
| -rw-r--r-- | mitmproxy/controller.py | 21 | ||||
| -rw-r--r-- | mitmproxy/script/concurrent.py | 31 | ||||
| -rw-r--r-- | test/mitmproxy/test_controller.py | 10 | ||||
| -rw-r--r-- | test/mitmproxy/test_server.py | 4 |
6 files changed, 29 insertions, 41 deletions
diff --git a/examples/redirect_requests.py b/examples/redirect_requests.py index 3ff8f9e4..d7db3f1c 100644 --- a/examples/redirect_requests.py +++ b/examples/redirect_requests.py @@ -16,7 +16,7 @@ def request(context, flow): "HTTP/1.1", 200, "OK", Headers(Content_Type="text/html"), "helloworld") - flow.reply(resp) + flow.reply.send(resp) # Method 2: Redirect the request to a different server if flow.request.pretty_host.endswith("example.org"): diff --git a/examples/tls_passthrough.py b/examples/tls_passthrough.py index 23afe3ff..0c6d450d 100644 --- a/examples/tls_passthrough.py +++ b/examples/tls_passthrough.py @@ -134,5 +134,5 @@ def next_layer(context, next_layer): # We don't intercept - reply with a pass-through layer and add a "skipped" entry. context.log("TLS passthrough for %s" % repr(next_layer.server_conn.address), "info") next_layer_replacement = RawTCPLayer(next_layer.ctx, logging=False) - next_layer.reply(next_layer_replacement) + next_layer.reply.send(next_layer_replacement) context.tls_strategy.record_skipped(server_address) diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py index 1aac82db..084702a6 100644 --- a/mitmproxy/controller.py +++ b/mitmproxy/controller.py @@ -145,10 +145,6 @@ class Channel(object): self.q.put((mtype, m)) -# Special value to distinguish the case where no reply was sent -NO_REPLY = object() - - def handler(f): @functools.wraps(f) def wrapper(*args, **kwargs): @@ -199,22 +195,19 @@ class Reply(object): self.handled = False def ack(self): - self(NO_REPLY) + self.send(self.obj) def kill(self): - self(exceptions.Kill) + self.send(exceptions.Kill) def take(self): self.taken = True - def __call__(self, msg=NO_REPLY): + def send(self, msg): if self.acked: raise exceptions.ControlException("Message already acked.") self.acked = True - if msg is NO_REPLY: - self.q.put(self.obj) - else: - self.q.put(msg) + self.q.put(msg) def __del__(self): if not self.acked: @@ -233,13 +226,13 @@ class DummyReply(object): self.handled = False def kill(self): - self() + self.send(None) def ack(self): - self() + self.send(None) def take(self): self.taken = True - def __call__(self, msg=False): + def send(self, msg): self.acked = True diff --git a/mitmproxy/script/concurrent.py b/mitmproxy/script/concurrent.py index 43d0d328..b81f2ab1 100644 --- a/mitmproxy/script/concurrent.py +++ b/mitmproxy/script/concurrent.py @@ -4,6 +4,7 @@ offload computations from mitmproxy's main master thread. """ from __future__ import absolute_import, print_function, division +from mitmproxy import controller import threading @@ -14,15 +15,15 @@ class ReplyProxy(object): self.script_thread = script_thread self.master_reply = None - def __call__(self, *args): + def send(self, message): if self.master_reply is None: - self.master_reply = args + self.master_reply = message self.script_thread.start() return - self.reply_func(*args) + self.reply_func(message) def done(self): - self.reply_func(*self.master_reply) + self.reply_func.send(self.master_reply) def __getattr__(self, k): return getattr(self.reply_func, k) @@ -49,17 +50,11 @@ class ScriptThread(threading.Thread): def concurrent(fn): - if fn.__name__ in ( - "request", - "response", - "error", - "clientconnect", - "serverconnect", - "clientdisconnect", - "next_layer"): - def _concurrent(ctx, obj): - _handle_concurrent_reply(fn, obj, ctx, obj) - - return _concurrent - raise NotImplementedError( - "Concurrent decorator not supported for '%s' method." % fn.__name__) + if fn.__name__ not in controller.Events: + raise NotImplementedError( + "Concurrent decorator not supported for '%s' method." % fn.__name__ + ) + + def _concurrent(ctx, obj): + _handle_concurrent_reply(fn, obj, ctx, obj) + return _concurrent diff --git a/test/mitmproxy/test_controller.py b/test/mitmproxy/test_controller.py index 83ad428e..5a68e15b 100644 --- a/test/mitmproxy/test_controller.py +++ b/test/mitmproxy/test_controller.py @@ -66,7 +66,7 @@ class TestChannel(object): def reply(): m, obj = q.get() assert m == "test" - obj.reply(42) + obj.reply.send(42) Thread(target=reply).start() @@ -86,7 +86,7 @@ class TestDummyReply(object): def test_simple(self): reply = controller.DummyReply() assert not reply.acked - reply() + reply.ack() assert reply.acked @@ -94,16 +94,16 @@ class TestReply(object): def test_simple(self): reply = controller.Reply(42) assert not reply.acked - reply("foo") + reply.send("foo") assert reply.acked assert reply.q.get() == "foo" def test_default(self): reply = controller.Reply(42) - reply() + reply.ack() assert reply.q.get() == 42 def test_reply_none(self): reply = controller.Reply(42) - reply(None) + reply.send(None) assert reply.q.get() is None diff --git a/test/mitmproxy/test_server.py b/test/mitmproxy/test_server.py index 1cd6cb0c..432340c0 100644 --- a/test/mitmproxy/test_server.py +++ b/test/mitmproxy/test_server.py @@ -743,7 +743,7 @@ class MasterFakeResponse(tservers.TestMaster): @controller.handler def request(self, f): resp = HTTPResponse.wrap(netlib.tutils.tresp()) - f.reply(resp) + f.reply.send(resp) class TestFakeResponse(tservers.HTTPProxyTest): @@ -819,7 +819,7 @@ class MasterIncomplete(tservers.TestMaster): def request(self, f): resp = HTTPResponse.wrap(netlib.tutils.tresp()) resp.content = None - f.reply(resp) + f.reply.send(resp) class TestIncompleteResponse(tservers.HTTPProxyTest): |