From f4813d56492a32ef7856002c8be25cf7012bab4e Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Sat, 28 May 2016 12:38:07 +1200 Subject: Make @controller.handler inheritance-friendly Use this to adapt mitmweb and mitproxy console --- mitmproxy/console/__init__.py | 20 +++++++++++++------- mitmproxy/controller.py | 16 ++++++++++++++-- mitmproxy/flow.py | 3 ++- mitmproxy/proxy/root_context.py | 1 - mitmproxy/web/__init__.py | 9 ++++++--- test/mitmproxy/test_flow.py | 5 +---- 6 files changed, 36 insertions(+), 18 deletions(-) diff --git a/mitmproxy/console/__init__.py b/mitmproxy/console/__init__.py index ce202b39..a87e691e 100644 --- a/mitmproxy/console/__init__.py +++ b/mitmproxy/console/__init__.py @@ -713,14 +713,15 @@ class ConsoleMaster(flow.FlowMaster): ) def process_flow(self, f): - if self.state.intercept and f.match(self.state.intercept) and not f.request.is_replay: + should_intercept = any( + [ + self.state.intercept and f.match(self.state.intercept) and not f.request.is_replay, + f.intercepted, + ] + ) + if should_intercept: f.intercept(self) - else: - # check if flow was intercepted within an inline script by flow.intercept() - if f.intercepted: - f.intercept(self) - else: - f.reply() + f.reply.take() signals.flowlist_change.send(self) signals.flow_change.send(self, flow = f) @@ -728,24 +729,29 @@ class ConsoleMaster(flow.FlowMaster): self.eventlist[:] = [] # Handlers + @controller.handler def handle_error(self, f): f = flow.FlowMaster.handle_error(self, f) if f: self.process_flow(f) return f + @controller.handler def handle_request(self, f): f = flow.FlowMaster.handle_request(self, f) + self.add_event(f.reply.acked, "info") if f: self.process_flow(f) return f + @controller.handler def handle_response(self, f): f = flow.FlowMaster.handle_response(self, f) if f: self.process_flow(f) return f + @controller.handler def handle_script_change(self, script): if super(ConsoleMaster, self).handle_script_change(script): signals.status_message.send(message='"{}" reloaded.'.format(script.filename)) diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py index c125ee4f..57c01f59 100644 --- a/mitmproxy/controller.py +++ b/mitmproxy/controller.py @@ -125,7 +125,6 @@ class Channel(object): if g == exceptions.Kill: raise exceptions.Kill() return g - raise exceptions.Kill() def tell(self, mtype, m): @@ -146,6 +145,7 @@ class DummyReply(object): def __init__(self): self.acked = False self.taken = False + self.handled = False def take(self): self.taken = True @@ -164,8 +164,16 @@ def handler(f): message = args[-1] if not hasattr(message, "reply"): raise ControlError("Message %s has no reply attribute"%message) + + handling = False + # We're the first handler - ack responsibility is ours + if not message.reply.handled: + handling = True + message.reply.handled = True + ret = f(*args, **kwargs) - if not message.reply.acked and not message.reply.taken: + + if handling and not message.reply.acked and not message.reply.taken: message.reply() return ret wrapper.func_dict["handler"] = True @@ -181,8 +189,12 @@ class Reply(object): def __init__(self, obj): self.obj = obj self.q = queue.Queue() + # Has this message been acked? self.acked = False + # Has the user taken responsibility for ack-ing? self.taken = False + # Has a handler taken responsibility for ack-ing? + self.handled = False def take(self): self.taken = True diff --git a/mitmproxy/flow.py b/mitmproxy/flow.py index 2a324993..ba8dec5d 100644 --- a/mitmproxy/flow.py +++ b/mitmproxy/flow.py @@ -546,7 +546,8 @@ class FlowStore(FlowList): def kill_all(self, master): for f in self._list: - f.kill(master) + if not f.reply.acked: + f.kill(master) class State(object): diff --git a/mitmproxy/proxy/root_context.py b/mitmproxy/proxy/root_context.py index 96e7aab6..9b4e2963 100644 --- a/mitmproxy/proxy/root_context.py +++ b/mitmproxy/proxy/root_context.py @@ -132,7 +132,6 @@ class RootContext(object): class Log(object): - def __init__(self, msg, level="info"): self.msg = msg self.level = level diff --git a/mitmproxy/web/__init__.py b/mitmproxy/web/__init__.py index 956d221d..d90830d6 100644 --- a/mitmproxy/web/__init__.py +++ b/mitmproxy/web/__init__.py @@ -6,7 +6,7 @@ import sys from netlib.http import authentication -from .. import flow +from .. import flow, controller from ..exceptions import FlowReadException from . import app @@ -194,21 +194,24 @@ class WebMaster(flow.FlowMaster): if self.state.intercept and self.state.intercept( f) and not f.request.is_replay: f.intercept(self) - else: - f.reply() + f.reply.take() + @controller.handler def handle_request(self, f): super(WebMaster, self).handle_request(f) self._process_flow(f) + @controller.handler def handle_response(self, f): super(WebMaster, self).handle_response(f) self._process_flow(f) + @controller.handler def handle_error(self, f): super(WebMaster, self).handle_error(f) self._process_flow(f) + @controller.handler def add_event(self, e, level="info"): super(WebMaster, self).add_event(e, level) self.state.add_event(e, level) diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py index 62f23ac8..da8b8ddd 100644 --- a/test/mitmproxy/test_flow.py +++ b/test/mitmproxy/test_flow.py @@ -470,10 +470,7 @@ class TestFlow(object): fm = flow.FlowMaster(None, s) f = tutils.tflow() - fm.handle_request(f) - - f = tutils.tflow() - fm.handle_request(f) + f.intercept(fm) s.killall(fm) for i in s.view: -- cgit v1.2.3