diff options
author | Aldo Cortesi <aldo@nullcube.com> | 2016-05-26 12:31:29 +1200 |
---|---|---|
committer | Aldo Cortesi <aldo@nullcube.com> | 2016-05-28 11:00:00 +1200 |
commit | 7c4843857be0ad165fca46dcad24a213fc3bd284 (patch) | |
tree | 4a4ae7df627ca9700c2f1fab8b3a9f77c085a2ec /mitmproxy | |
parent | abf3cae54e715e53559b4c7103f8b2247ed791dc (diff) | |
download | mitmproxy-7c4843857be0ad165fca46dcad24a213fc3bd284.tar.gz mitmproxy-7c4843857be0ad165fca46dcad24a213fc3bd284.tar.bz2 mitmproxy-7c4843857be0ad165fca46dcad24a213fc3bd284.zip |
Sketch out a more solid core
- Decorator for handler methods
- Stricter checking for double-acks and non-acks
Diffstat (limited to 'mitmproxy')
-rw-r--r-- | mitmproxy/console/__init__.py | 2 | ||||
-rw-r--r-- | mitmproxy/controller.py | 60 | ||||
-rw-r--r-- | mitmproxy/flow.py | 14 |
3 files changed, 57 insertions, 19 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): |