From 7f8fd3cdffedb537f95773110d8ef2be60666133 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Wed, 13 Jul 2016 23:26:04 +1200 Subject: [PATCH] Basic outline of addons Add addons.py, integrate with our event mechanism, and change the Master API so options is the first init argument. --- mitmproxy/addons.py | 63 +++++++++++++++++++++++++++++++ mitmproxy/console/master.py | 3 +- mitmproxy/controller.py | 11 +++++- mitmproxy/dump.py | 22 +++++------ mitmproxy/flow/master.py | 4 +- mitmproxy/options.py | 5 ++- mitmproxy/web/master.py | 7 ++-- test/mitmproxy/test_addons.py | 20 ++++++++++ test/mitmproxy/test_controller.py | 4 +- test/mitmproxy/test_flow.py | 38 +++++++++---------- test/mitmproxy/test_script.py | 2 +- test/mitmproxy/tservers.py | 2 +- 12 files changed, 138 insertions(+), 43 deletions(-) create mode 100644 mitmproxy/addons.py create mode 100644 test/mitmproxy/test_addons.py diff --git a/mitmproxy/addons.py b/mitmproxy/addons.py new file mode 100644 index 000000000..7ac65a09d --- /dev/null +++ b/mitmproxy/addons.py @@ -0,0 +1,63 @@ +from __future__ import absolute_import, print_function, division +from mitmproxy import exceptions +import pprint + + +def _get_name(itm): + return getattr(itm, "name", itm.__class__.__name__) + + +class Addons(object): + def __init__(self, master): + self.chain = [] + self.master = master + master.options.changed.connect(self.options_update) + + def options_update(self, options): + for i in self.chain: + with self.master.handlecontext(): + i.configure(options) + + def add(self, *addons): + self.chain.extend(addons) + for i in addons: + self.invoke_with_context(i, "configure", self.master.options) + + def remove(self, addon): + self.chain = [i for i in self.chain if i is not addon] + self.invoke_with_context(addon, "done") + + def done(self): + for i in self.chain: + self.invoke_with_context(i, "done") + + def has_addon(self, name): + """ + Is an addon with this name registered? + """ + for i in self.chain: + if _get_name(i) == name: + return True + + def __len__(self): + return len(self.chain) + + def __str__(self): + return pprint.pformat([str(i) for i in self.chain]) + + def invoke_with_context(self, addon, name, *args, **kwargs): + with self.master.handlecontext(): + self.invoke(addon, name, *args, **kwargs) + + def invoke(self, addon, name, *args, **kwargs): + func = getattr(addon, name, None) + if func: + if not callable(func): + raise exceptions.AddonError( + "Addon handler %s not callable" % name + ) + func(*args, **kwargs) + + def __call__(self, name, *args, **kwargs): + for i in self.chain: + self.invoke(i, name, *args, **kwargs) diff --git a/mitmproxy/console/master.py b/mitmproxy/console/master.py index 1daf11270..00905f366 100644 --- a/mitmproxy/console/master.py +++ b/mitmproxy/console/master.py @@ -216,10 +216,9 @@ class ConsoleMaster(flow.FlowMaster): palette = [] def __init__(self, server, options): - flow.FlowMaster.__init__(self, server, ConsoleState()) + flow.FlowMaster.__init__(self, options, server, ConsoleState()) self.stream_path = None - self.options = options self.options.errored.connect(self.options_error) if options.replacements: diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py index e2be3a532..d09038f8f 100644 --- a/mitmproxy/controller.py +++ b/mitmproxy/controller.py @@ -6,6 +6,8 @@ import contextlib from six.moves import queue +from mitmproxy import addons +from mitmproxy import options from . import ctx as mitmproxy_ctx from netlib import basethread from . import exceptions @@ -49,7 +51,9 @@ class Master(object): """ The master handles mitmproxy's main event loop. """ - def __init__(self, *servers): + def __init__(self, opts, *servers): + self.options = opts or options.Options() + self.addons = addons.Addons(self) self.event_queue = queue.Queue() self.should_exit = threading.Event() self.servers = [] @@ -121,6 +125,7 @@ class Master(object): for server in self.servers: server.shutdown() self.should_exit.set() + self.addons.done() class ServerThread(basethread.BaseThread): @@ -191,6 +196,10 @@ def handler(f): with master.handlecontext(): ret = f(master, message) + if handling: + # Python2/3 compatibility hack + fn = getattr(f, "func_name", None) or getattr(f, "__name__") + master.addons(fn) if handling and not message.reply.acked and not message.reply.taken: message.reply.ack() diff --git a/mitmproxy/dump.py b/mitmproxy/dump.py index bfefb319d..cd5159450 100644 --- a/mitmproxy/dump.py +++ b/mitmproxy/dump.py @@ -58,7 +58,7 @@ class Options(options.Options): class DumpMaster(flow.FlowMaster): def __init__(self, server, options, outfile=None): - flow.FlowMaster.__init__(self, server, flow.State()) + flow.FlowMaster.__init__(self, options, server, flow.State()) self.outfile = outfile self.o = options self.anticache = options.anticache @@ -137,8 +137,8 @@ class DumpMaster(flow.FlowMaster): self.add_event("Flow file corrupted.", "error") raise DumpError(v) - if self.o.app: - self.start_app(self.o.app_host, self.o.app_port) + if self.options.app: + self.start_app(self.options.app_host, self.options.app_port) def _readflow(self, paths): """ @@ -152,7 +152,7 @@ class DumpMaster(flow.FlowMaster): def add_event(self, e, level="info"): needed = dict(error=0, info=1, debug=2).get(level, 1) - if self.o.verbosity >= needed: + if self.options.verbosity >= needed: self.echo( e, fg="red" if level == "error" else None, @@ -172,7 +172,7 @@ class DumpMaster(flow.FlowMaster): click.secho(text, file=self.outfile, **style) def _echo_message(self, message): - if self.o.flow_detail >= 2 and hasattr(message, "headers"): + if self.options.flow_detail >= 2 and hasattr(message, "headers"): headers = "\r\n".join( "{}: {}".format( click.style(strutils.bytes_to_escaped_str(k), fg="blue", bold=True), @@ -180,7 +180,7 @@ class DumpMaster(flow.FlowMaster): for k, v in message.headers.fields ) self.echo(headers, indent=4) - if self.o.flow_detail >= 3: + if self.options.flow_detail >= 3: if message.content is None: self.echo("(content missing)", indent=4) elif message.content: @@ -213,7 +213,7 @@ class DumpMaster(flow.FlowMaster): for (style, text) in line: yield click.style(text, **styles.get(style, {})) - if self.o.flow_detail == 3: + if self.options.flow_detail == 3: lines_to_echo = itertools.islice(lines, 70) else: lines_to_echo = lines @@ -228,7 +228,7 @@ class DumpMaster(flow.FlowMaster): if next(lines, None): self.echo("(cut off)", indent=4, dim=True) - if self.o.flow_detail >= 2: + if self.options.flow_detail >= 2: self.echo("") def _echo_request_line(self, flow): @@ -302,7 +302,7 @@ class DumpMaster(flow.FlowMaster): self.echo(line) def echo_flow(self, f): - if self.o.flow_detail == 0: + if self.options.flow_detail == 0: return if f.request: @@ -350,7 +350,7 @@ class DumpMaster(flow.FlowMaster): def tcp_message(self, f): super(DumpMaster, self).tcp_message(f) - if self.o.flow_detail == 0: + if self.options.flow_detail == 0: return message = f.messages[-1] direction = "->" if message.from_client else "<-" @@ -362,7 +362,7 @@ class DumpMaster(flow.FlowMaster): self._echo_message(message) def run(self): # pragma: no cover - if self.o.rfile and not self.o.keepserving: + if self.options.rfile and not self.options.keepserving: self.unload_scripts() # make sure to trigger script unload events. return super(DumpMaster, self).run() diff --git a/mitmproxy/flow/master.py b/mitmproxy/flow/master.py index 7590a3fa5..b1951f944 100644 --- a/mitmproxy/flow/master.py +++ b/mitmproxy/flow/master.py @@ -27,8 +27,8 @@ class FlowMaster(controller.Master): if len(self.servers) > 0: return self.servers[0] - def __init__(self, server, state): - super(FlowMaster, self).__init__() + def __init__(self, options, server, state): + super(FlowMaster, self).__init__(options) if server: self.add_server(server) self.state = state diff --git a/mitmproxy/options.py b/mitmproxy/options.py index 7389df1f5..0cc5fee1e 100644 --- a/mitmproxy/options.py +++ b/mitmproxy/options.py @@ -42,7 +42,10 @@ class Options(object): return self.__class__(**self._opts) def __getattr__(self, attr): - return self._opts[attr] + if attr in self._opts: + return self._opts[attr] + else: + raise AttributeError() def __setattr__(self, attr, value): if attr not in self._opts: diff --git a/mitmproxy/web/master.py b/mitmproxy/web/master.py index 008b74f83..22972c142 100644 --- a/mitmproxy/web/master.py +++ b/mitmproxy/web/master.py @@ -147,9 +147,10 @@ class Options(options.Options): class WebMaster(flow.FlowMaster): def __init__(self, server, options): - self.options = options - super(WebMaster, self).__init__(server, WebState()) - self.app = app.Application(self, self.options.wdebug, self.options.wauthenticator) + super(WebMaster, self).__init__(options, server, WebState()) + self.app = app.Application( + self, self.options.wdebug, self.options.wauthenticator + ) if options.rfile: try: self.load_flows_file(options.rfile) diff --git a/test/mitmproxy/test_addons.py b/test/mitmproxy/test_addons.py new file mode 100644 index 000000000..1861d4acd --- /dev/null +++ b/test/mitmproxy/test_addons.py @@ -0,0 +1,20 @@ +from __future__ import absolute_import, print_function, division +from mitmproxy import addons +from mitmproxy import controller +from mitmproxy import options + + +class TAddon: + def __init__(self, name): + self.name = name + + def __repr__(self): + return "Addon(%s)" % self.name + + +def test_simple(): + m = controller.Master(options.Options()) + a = addons.Addons(m) + a.add(TAddon("one")) + assert a.has_addon("one") + assert not a.has_addon("two") diff --git a/test/mitmproxy/test_controller.py b/test/mitmproxy/test_controller.py index 5a68e15b4..6d4b8fe63 100644 --- a/test/mitmproxy/test_controller.py +++ b/test/mitmproxy/test_controller.py @@ -25,7 +25,7 @@ class TestMaster(object): # Speed up test super(DummyMaster, self).tick(0) - m = DummyMaster() + m = DummyMaster(None) assert not m.should_exit.is_set() msg = TMsg() msg.reply = controller.DummyReply() @@ -34,7 +34,7 @@ class TestMaster(object): assert m.should_exit.is_set() def test_server_simple(self): - m = controller.Master() + m = controller.Master(None) s = DummyServer(None) m.add_server(s) m.start() diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py index 74b3f5992..eda01ad9d 100644 --- a/test/mitmproxy/test_flow.py +++ b/test/mitmproxy/test_flow.py @@ -139,7 +139,7 @@ class TestClientPlaybackState: def test_tick(self): first = tutils.tflow() s = flow.State() - fm = flow.FlowMaster(None, s) + fm = flow.FlowMaster(None, None, s) fm.start_client_playback([first, tutils.tflow()], True) c = fm.client_playback c.testing = True @@ -470,7 +470,7 @@ class TestFlow(object): def test_kill(self): s = flow.State() - fm = flow.FlowMaster(None, s) + fm = flow.FlowMaster(None, None, s) f = tutils.tflow() f.intercept(mock.Mock()) f.kill(fm) @@ -479,7 +479,7 @@ class TestFlow(object): def test_killall(self): s = flow.State() - fm = flow.FlowMaster(None, s) + fm = flow.FlowMaster(None, None, s) f = tutils.tflow() f.intercept(fm) @@ -714,7 +714,7 @@ class TestSerialize: def test_load_flows(self): r = self._treader() s = flow.State() - fm = flow.FlowMaster(None, s) + fm = flow.FlowMaster(None, None, s) fm.load_flows(r) assert len(s.flows) == 6 @@ -725,7 +725,7 @@ class TestSerialize: mode="reverse", upstream_server=("https", ("use-this-domain", 80)) ) - fm = flow.FlowMaster(DummyServer(conf), s) + fm = flow.FlowMaster(None, DummyServer(conf), s) fm.load_flows(r) assert s.flows[0].request.host == "use-this-domain" @@ -772,7 +772,7 @@ class TestFlowMaster: def test_load_script(self): s = flow.State() - fm = flow.FlowMaster(None, s) + fm = flow.FlowMaster(None, None, s) fm.load_script(tutils.test_data.path("data/scripts/a.py")) fm.load_script(tutils.test_data.path("data/scripts/a.py")) @@ -788,14 +788,14 @@ class TestFlowMaster: def test_getset_ignore(self): p = mock.Mock() p.config.check_ignore = HostMatcher() - fm = flow.FlowMaster(p, flow.State()) + fm = flow.FlowMaster(None, p, flow.State()) assert not fm.get_ignore_filter() fm.set_ignore_filter(["^apple\.com:", ":443$"]) assert fm.get_ignore_filter() def test_replay(self): s = flow.State() - fm = flow.FlowMaster(None, s) + fm = flow.FlowMaster(None, None, s) f = tutils.tflow(resp=True) f.request.content = None assert "missing" in fm.replay_request(f) @@ -808,7 +808,7 @@ class TestFlowMaster: def test_script_reqerr(self): s = flow.State() - fm = flow.FlowMaster(None, s) + fm = flow.FlowMaster(None, None, s) fm.load_script(tutils.test_data.path("data/scripts/reqerr.py")) f = tutils.tflow() fm.clientconnect(f.client_conn) @@ -816,7 +816,7 @@ class TestFlowMaster: def test_script(self): s = flow.State() - fm = flow.FlowMaster(None, s) + fm = flow.FlowMaster(None, None, s) fm.load_script(tutils.test_data.path("data/scripts/all.py")) f = tutils.tflow(resp=True) @@ -852,7 +852,7 @@ class TestFlowMaster: def test_duplicate_flow(self): s = flow.State() - fm = flow.FlowMaster(None, s) + fm = flow.FlowMaster(None, None, s) f = tutils.tflow(resp=True) fm.load_flow(f) assert s.flow_count() == 1 @@ -863,12 +863,12 @@ class TestFlowMaster: def test_create_flow(self): s = flow.State() - fm = flow.FlowMaster(None, s) + fm = flow.FlowMaster(None, None, s) assert fm.create_request("GET", "http", "example.com", 80, "/") def test_all(self): s = flow.State() - fm = flow.FlowMaster(None, s) + fm = flow.FlowMaster(None, None, s) fm.anticache = True fm.anticomp = True f = tutils.tflow(req=None) @@ -895,7 +895,7 @@ class TestFlowMaster: f = tutils.tflow(resp=True) pb = [tutils.tflow(resp=True), f] - fm = flow.FlowMaster(DummyServer(ProxyConfig()), s) + fm = flow.FlowMaster(None, DummyServer(ProxyConfig()), s) assert not fm.start_server_playback( pb, False, @@ -923,7 +923,7 @@ class TestFlowMaster: f.response = HTTPResponse.wrap(netlib.tutils.tresp(content=f.request)) pb = [f] - fm = flow.FlowMaster(None, s) + fm = flow.FlowMaster(None, None, s) fm.refresh_server_playback = True assert not fm.do_server_playback(tutils.tflow()) @@ -965,7 +965,7 @@ class TestFlowMaster: f = tutils.tflow() f.response = HTTPResponse.wrap(netlib.tutils.tresp(content=f.request)) pb = [f] - fm = flow.FlowMaster(None, s) + fm = flow.FlowMaster(None, None, s) fm.refresh_server_playback = True fm.start_server_playback( pb, @@ -985,7 +985,7 @@ class TestFlowMaster: def test_stickycookie(self): s = flow.State() - fm = flow.FlowMaster(None, s) + fm = flow.FlowMaster(None, None, s) assert "Invalid" in fm.set_stickycookie("~h") fm.set_stickycookie(".*") assert fm.stickycookie_state @@ -1007,7 +1007,7 @@ class TestFlowMaster: def test_stickyauth(self): s = flow.State() - fm = flow.FlowMaster(None, s) + fm = flow.FlowMaster(None, None, s) assert "Invalid" in fm.set_stickyauth("~h") fm.set_stickyauth(".*") assert fm.stickyauth_state @@ -1035,7 +1035,7 @@ class TestFlowMaster: return list(r.stream()) s = flow.State() - fm = flow.FlowMaster(None, s) + fm = flow.FlowMaster(None, None, s) f = tutils.tflow(resp=True) with open(p, "ab") as tmpfile: diff --git a/test/mitmproxy/test_script.py b/test/mitmproxy/test_script.py index 819947808..1e8220f10 100644 --- a/test/mitmproxy/test_script.py +++ b/test/mitmproxy/test_script.py @@ -4,7 +4,7 @@ from . import tutils def test_duplicate_flow(): s = flow.State() - fm = flow.FlowMaster(None, s) + fm = flow.FlowMaster(None, None, s) fm.load_script(tutils.test_data.path("data/scripts/duplicate_flow.py")) f = tutils.tflow() fm.request(f) diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py index 0760cb53f..69a50b9d5 100644 --- a/test/mitmproxy/tservers.py +++ b/test/mitmproxy/tservers.py @@ -34,7 +34,7 @@ class TestMaster(flow.FlowMaster): config.port = 0 s = ProxyServer(config) state = flow.State() - flow.FlowMaster.__init__(self, s, state) + flow.FlowMaster.__init__(self, None, s, state) self.apps.add(testapp, "testapp", 80) self.apps.add(errapp, "errapp", 80) self.clear_log()