diff options
| -rw-r--r-- | mitmproxy/console/__init__.py | 2 | ||||
| -rw-r--r-- | mitmproxy/controller.py | 60 | ||||
| -rw-r--r-- | mitmproxy/flow.py | 14 | ||||
| -rw-r--r-- | netlib/tcp.py | 2 | ||||
| -rw-r--r-- | test/mitmproxy/test_controller.py | 31 | ||||
| -rw-r--r-- | test/mitmproxy/test_flow.py | 1 | ||||
| -rw-r--r-- | test/mitmproxy/tservers.py | 4 |
7 files changed, 79 insertions, 35 deletions
diff --git a/mitmproxy/console/__init__.py b/mitmproxy/console/__init__.py index 1dd032be..ce202b39 100644 --- a/mitmproxy/console/__init__.py +++ b/mitmproxy/console/__init__.py @@ -16,7 +16,7 @@ import weakref from netlib import tcp -from .. import flow, script, contentviews +from .. import flow, script, contentviews, controller from . import flowlist, flowview, help, window, signals, options from . import grideditor, palettes, statusbar, palettepicker from ..exceptions import FlowReadException, ScriptException diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py index 81978a09..60d4e71e 100644 --- a/mitmproxy/controller.py +++ b/mitmproxy/controller.py @@ -1,8 +1,14 @@ from __future__ import absolute_import from six.moves import queue import threading +import functools +import sys -from .exceptions import Kill +from . import exceptions + + +class ControlError(Exception): + pass class Master(object): @@ -36,6 +42,12 @@ class Master(object): while True: mtype, obj = self.event_queue.get(timeout=timeout) handle_func = getattr(self, "handle_" + mtype) + # if not handle_func.func_dict.get("handler"): + # raise ControlError( + # "Handler function %s is not decorated with controller.handler"%( + # handle_func + # ) + # ) handle_func(obj) self.event_queue.task_done() changed = True @@ -100,7 +112,7 @@ class Channel(object): master. Then wait for a response. Raises: - Kill: All connections should be closed immediately. + exceptions.Kill: All connections should be closed immediately. """ m.reply = Reply(m) self.q.put((mtype, m)) @@ -110,11 +122,11 @@ class Channel(object): g = m.reply.q.get(timeout=0.5) except queue.Empty: # pragma: no cover continue - if g == Kill: - raise Kill() + if g == exceptions.Kill: + raise exceptions.Kill() return g - raise Kill() + raise exceptions.Kill() def tell(self, mtype, m): """ @@ -133,6 +145,10 @@ class DummyReply(object): def __init__(self): self.acked = False + self.taken = False + + def take(self): + self.taken = True def __call__(self, msg=False): self.acked = True @@ -142,22 +158,44 @@ class DummyReply(object): NO_REPLY = object() +def handler(f): + @functools.wraps(f) + def wrapper(obj, message, *args, **kwargs): + if not hasattr(message, "reply"): + raise ControlError("Message %s has no reply attribute"%message) + ret = f(obj, message, *args, **kwargs) + if not message.reply.acked and not message.reply.taken: + message.reply() + return ret + wrapper.func_dict["handler"] = True + return wrapper + + class Reply(object): """ Messages sent through a channel are decorated with a "reply" attribute. This object is used to respond to the message through the return channel. """ - def __init__(self, obj): self.obj = obj self.q = queue.Queue() self.acked = False + self.taken = False + + def take(self): + self.taken = True def __call__(self, msg=NO_REPLY): + if self.acked: + raise ControlError("Message already acked.") + self.acked = True + if msg is NO_REPLY: + self.q.put(self.obj) + else: + self.q.put(msg) + + def __del__(self): if not self.acked: - self.acked = True - if msg is NO_REPLY: - self.q.put(self.obj) - else: - self.q.put(msg) + # This will be ignored by the interpreter, but emit a warning + raise ControlError("Un-acked message") diff --git a/mitmproxy/flow.py b/mitmproxy/flow.py index a09a81a7..21aa3051 100644 --- a/mitmproxy/flow.py +++ b/mitmproxy/flow.py @@ -985,36 +985,36 @@ class FlowMaster(controller.ServerMaster): if block: rt.join() + @controller.handler def handle_log(self, l): self.add_event(l.msg, l.level) - l.reply() + @controller.handler def handle_clientconnect(self, root_layer): self.run_script_hook("clientconnect", root_layer) - root_layer.reply() + @controller.handler def handle_clientdisconnect(self, root_layer): self.run_script_hook("clientdisconnect", root_layer) - root_layer.reply() + @controller.handler def handle_serverconnect(self, server_conn): self.run_script_hook("serverconnect", server_conn) - server_conn.reply() + @controller.handler def handle_serverdisconnect(self, server_conn): self.run_script_hook("serverdisconnect", server_conn) - server_conn.reply() + @controller.handler def handle_next_layer(self, top_layer): self.run_script_hook("next_layer", top_layer) - top_layer.reply() + @controller.handler def handle_error(self, f): self.state.update_flow(f) self.run_script_hook("error", f) if self.client_playback: self.client_playback.clear(f) - f.reply() return f def handle_request(self, f): diff --git a/netlib/tcp.py b/netlib/tcp.py index d26bb5f7..8e3ae65e 100644 --- a/netlib/tcp.py +++ b/netlib/tcp.py @@ -900,7 +900,7 @@ class TCPServer(object): """ # If a thread has persisted after interpreter exit, the module might be # none. - if traceback: + if traceback and six: exc = six.text_type(traceback.format_exc()) print(u'-' * 40, file=fp) print( diff --git a/test/mitmproxy/test_controller.py b/test/mitmproxy/test_controller.py index f7bf615a..c9a8e2f4 100644 --- a/test/mitmproxy/test_controller.py +++ b/test/mitmproxy/test_controller.py @@ -2,7 +2,7 @@ from threading import Thread, Event from mock import Mock -from mitmproxy.controller import Reply, DummyReply, Channel, ServerThread, ServerMaster, Master +from mitmproxy import controller from six.moves import queue from mitmproxy.exceptions import Kill @@ -10,10 +10,15 @@ from mitmproxy.proxy import DummyServer from netlib.tutils import raises +class TMsg: + pass + + class TestMaster(object): def test_simple(self): - class DummyMaster(Master): + class DummyMaster(controller.Master): + @controller.handler def handle_panic(self, _): m.should_exit.set() @@ -23,14 +28,16 @@ class TestMaster(object): m = DummyMaster() assert not m.should_exit.is_set() - m.event_queue.put(("panic", 42)) + msg = TMsg() + msg.reply = controller.DummyReply() + m.event_queue.put(("panic", msg)) m.run() assert m.should_exit.is_set() class TestServerMaster(object): def test_simple(self): - m = ServerMaster() + m = controller.ServerMaster() s = DummyServer(None) m.add_server(s) m.start() @@ -42,7 +49,7 @@ class TestServerMaster(object): class TestServerThread(object): def test_simple(self): m = Mock() - t = ServerThread(m) + t = controller.ServerThread(m) t.run() assert m.serve_forever.called @@ -50,7 +57,7 @@ class TestServerThread(object): class TestChannel(object): def test_tell(self): q = queue.Queue() - channel = Channel(q, Event()) + channel = controller.Channel(q, Event()) m = Mock() channel.tell("test", m) assert q.get() == ("test", m) @@ -66,21 +73,21 @@ class TestChannel(object): Thread(target=reply).start() - channel = Channel(q, Event()) + channel = controller.Channel(q, Event()) assert channel.ask("test", Mock()) == 42 def test_ask_shutdown(self): q = queue.Queue() done = Event() done.set() - channel = Channel(q, done) + channel = controller.Channel(q, done) with raises(Kill): channel.ask("test", Mock()) class TestDummyReply(object): def test_simple(self): - reply = DummyReply() + reply = controller.DummyReply() assert not reply.acked reply() assert reply.acked @@ -88,18 +95,18 @@ class TestDummyReply(object): class TestReply(object): def test_simple(self): - reply = Reply(42) + reply = controller.Reply(42) assert not reply.acked reply("foo") assert reply.acked assert reply.q.get() == "foo" def test_default(self): - reply = Reply(42) + reply = controller.Reply(42) reply() assert reply.q.get() == 42 def test_reply_none(self): - reply = Reply(42) + reply = controller.Reply(42) reply(None) assert reply.q.get() is None diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py index 3e78a5c4..62ea7d2d 100644 --- a/test/mitmproxy/test_flow.py +++ b/test/mitmproxy/test_flow.py @@ -866,7 +866,6 @@ class TestFlowMaster: f.response = HTTPResponse.wrap(netlib.tutils.tresp()) fm.handle_response(f) - assert not fm.handle_response(None) assert s.flow_count() == 1 fm.handle_clientdisconnect(f.client_conn) diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py index c9d68cfd..51f4109d 100644 --- a/test/mitmproxy/tservers.py +++ b/test/mitmproxy/tservers.py @@ -39,13 +39,13 @@ class TestMaster(flow.FlowMaster): self.apps.add(errapp, "errapp", 80) self.clear_log() + @controller.handler def handle_request(self, f): flow.FlowMaster.handle_request(self, f) - f.reply() + @controller.handler def handle_response(self, f): flow.FlowMaster.handle_response(self, f) - f.reply() def clear_log(self): self.log = [] |