diff options
28 files changed, 209 insertions, 70 deletions
diff --git a/examples/complex/dns_spoofing.py b/examples/complex/dns_spoofing.py index ca2bcd35..01e036b2 100644 --- a/examples/complex/dns_spoofing.py +++ b/examples/complex/dns_spoofing.py @@ -54,5 +54,5 @@ class Rerouter: flow.request.port = port -def start(opts): - return Rerouter() +def load(l): + l.boot_into(Rerouter()) diff --git a/examples/complex/har_dump.py b/examples/complex/har_dump.py index 9a86e45e..0515d0f5 100644 --- a/examples/complex/har_dump.py +++ b/examples/complex/har_dump.py @@ -25,7 +25,7 @@ HAR = {} SERVERS_SEEN = set() -def start(opts): +def load(l): """ Called once on script startup before any other events. """ diff --git a/examples/complex/remote_debug.py b/examples/complex/remote_debug.py index ae0dffc1..fa6f3d33 100644 --- a/examples/complex/remote_debug.py +++ b/examples/complex/remote_debug.py @@ -14,6 +14,6 @@ Usage: """ -def start(opts): +def load(l): import pydevd pydevd.settrace("localhost", port=5678, stdoutToServer=True, stderrToServer=True) diff --git a/examples/complex/tls_passthrough.py b/examples/complex/tls_passthrough.py index 6dba7ca1..72c0244b 100644 --- a/examples/complex/tls_passthrough.py +++ b/examples/complex/tls_passthrough.py @@ -112,7 +112,7 @@ class TlsFeedback(TlsLayer): tls_strategy = None -def start(opts): +def load(l): global tls_strategy if len(sys.argv) == 2: tls_strategy = ProbabilisticStrategy(float(sys.argv[1])) diff --git a/examples/simple/add_header_class.py b/examples/simple/add_header_class.py index 9270be09..69b64163 100644 --- a/examples/simple/add_header_class.py +++ b/examples/simple/add_header_class.py @@ -3,5 +3,5 @@ class AddHeader: flow.response.headers["newheader"] = "foo" -def start(opts): - return AddHeader() +def load(l): + return l.boot_into(AddHeader()) diff --git a/examples/simple/custom_contentview.py b/examples/simple/custom_contentview.py index 4bc17af0..34fa5541 100644 --- a/examples/simple/custom_contentview.py +++ b/examples/simple/custom_contentview.py @@ -20,7 +20,7 @@ class ViewSwapCase(contentviews.View): view = ViewSwapCase() -def start(opts): +def load(l): contentviews.add(view) diff --git a/examples/simple/custom_option.py b/examples/simple/custom_option.py index 324d27e7..c8bc98d4 100644 --- a/examples/simple/custom_option.py +++ b/examples/simple/custom_option.py @@ -1,9 +1,9 @@ from mitmproxy import ctx -def start(options): +def load(l): ctx.log.info("Registering option 'custom'") - options.add_option("custom", bool, False, "A custom option") + l.add_option("custom", bool, False, "A custom option") def configure(options, updated): diff --git a/examples/simple/filter_flows.py b/examples/simple/filter_flows.py index 24e8b6c1..d2b735be 100644 --- a/examples/simple/filter_flows.py +++ b/examples/simple/filter_flows.py @@ -17,7 +17,7 @@ class Filter: print(flow) -def start(opts): +def load(l): if len(sys.argv) != 2: raise ValueError("Usage: -s 'filt.py FILTER'") - return Filter(sys.argv[1]) + l.boot_into(Filter(sys.argv[1])) diff --git a/examples/simple/io_write_dumpfile.py b/examples/simple/io_write_dumpfile.py index 311950af..15e7693c 100644 --- a/examples/simple/io_write_dumpfile.py +++ b/examples/simple/io_write_dumpfile.py @@ -23,7 +23,7 @@ class Writer: self.w.add(flow) -def start(opts): +def load(l): if len(sys.argv) != 2: raise ValueError('Usage: -s "flowriter.py filename"') - return Writer(sys.argv[1]) + l.boot_into(Writer(sys.argv[1])) diff --git a/examples/simple/log_events.py b/examples/simple/log_events.py index a81892aa..581b99f3 100644 --- a/examples/simple/log_events.py +++ b/examples/simple/log_events.py @@ -7,6 +7,6 @@ If you want to help us out: https://github.com/mitmproxy/mitmproxy/issues/1530 : from mitmproxy import ctx -def start(opts): +def load(l): ctx.log.info("This is some informative text.") ctx.log.error("This is an error.") diff --git a/examples/simple/modify_body_inject_iframe.py b/examples/simple/modify_body_inject_iframe.py index ab5abf27..442a5118 100644 --- a/examples/simple/modify_body_inject_iframe.py +++ b/examples/simple/modify_body_inject_iframe.py @@ -23,7 +23,7 @@ class Injector: flow.response.content = str(html).encode("utf8") -def start(opts): +def load(l): if len(sys.argv) != 2: raise ValueError('Usage: -s "iframe_injector.py url"') - return Injector(sys.argv[1]) + return l.boot_into(Injector(sys.argv[1])) diff --git a/examples/simple/script_arguments.py b/examples/simple/script_arguments.py index b46a1960..84292eb9 100644 --- a/examples/simple/script_arguments.py +++ b/examples/simple/script_arguments.py @@ -9,9 +9,9 @@ class Replacer: flow.response.replace(self.src, self.dst) -def start(opts): +def load(l): parser = argparse.ArgumentParser() parser.add_argument("src", type=str) parser.add_argument("dst", type=str) args = parser.parse_args() - return Replacer(args.src, args.dst) + l.boot_into(Replacer(args.src, args.dst)) diff --git a/examples/simple/wsgi_flask_app.py b/examples/simple/wsgi_flask_app.py index db3b1adf..30008b16 100644 --- a/examples/simple/wsgi_flask_app.py +++ b/examples/simple/wsgi_flask_app.py @@ -14,7 +14,7 @@ def hello_world(): return 'Hello World!' -def start(opts): +def load(l): # Host app at the magic domain "proxapp" on port 80. Requests to this # domain and port combination will now be routed to the WSGI app instance. return wsgiapp.WSGIApp(app, "proxapp", 80) diff --git a/mitmproxy/addonmanager.py b/mitmproxy/addonmanager.py index 670c4f24..241b3cde 100644 --- a/mitmproxy/addonmanager.py +++ b/mitmproxy/addonmanager.py @@ -1,3 +1,5 @@ +import typing + from mitmproxy import exceptions from mitmproxy import eventsequence from mitmproxy import controller @@ -9,6 +11,37 @@ def _get_name(itm): return getattr(itm, "name", itm.__class__.__name__.lower()) +class Loader: + """ + A loader object is passed to the load() event when addons start up. + """ + def __init__(self, master): + self.master = master + self.boot_into_addon = None + + def add_option( + self, + name: str, + typespec: type, + default: typing.Any, + help: str, + choices: typing.Optional[typing.Sequence[str]] = None + ) -> None: + self.master.options.add_option( + name, + typespec, + default, + help, + choices + ) + + def boot_into(self, addon): + self.boot_into_addon = addon + func = getattr(addon, "load", None) + if func: + func(self) + + class AddonManager: def __init__(self, master): self.chain = [] @@ -39,10 +72,14 @@ class AddonManager: """ Add addons to the end of the chain, and run their startup events. """ - self.chain.extend(addons) with self.master.handlecontext(): for i in addons: - self.invoke_addon(i, "start", self.master.options) + l = Loader(self.master) + self.invoke_addon(i, "load", l) + if l.boot_into_addon: + self.chain.append(l.boot_into_addon) + else: + self.chain.append(i) def remove(self, addon): """ diff --git a/mitmproxy/addons/script.py b/mitmproxy/addons/script.py index 3b7a4f79..eef21293 100644 --- a/mitmproxy/addons/script.py +++ b/mitmproxy/addons/script.py @@ -6,6 +6,7 @@ import threading import traceback import types +from mitmproxy import addonmanager from mitmproxy import exceptions from mitmproxy import ctx from mitmproxy import eventsequence @@ -184,23 +185,22 @@ class Script: def load_script(self): self.ns = load_script(self.path, self.args) - ret = self.run("start", self.last_options) - if ret: - self.ns = ret - self.run("start", self.last_options) + l = addonmanager.Loader(ctx.master) + self.run("load", l) + if l.boot_into_addon: + self.ns = l.boot_into_addon def tick(self): if self.should_reload.is_set(): self.should_reload.clear() ctx.log.info("Reloading script: %s" % self.name) self.ns = load_script(self.path, self.args) - self.start(self.last_options) self.configure(self.last_options, self.last_options.keys()) else: self.run("tick") - def start(self, opts): - self.last_options = opts + def load(self, l): + self.last_options = ctx.master.options self.load_script() def configure(self, options, updated): @@ -287,7 +287,8 @@ class ScriptLoader: ctx.master.addons.chain = ochain[:pos + 1] + ordered + ochain[pos + 1:] for s in newscripts: - ctx.master.addons.invoke_addon(s, "start", options) + l = addonmanager.Loader(ctx.master) + ctx.master.addons.invoke_addon(s, "load", l) if self.is_running: # If we're already running, we configure and tell the addon # we're up and running. diff --git a/mitmproxy/addons/serverplayback.py b/mitmproxy/addons/serverplayback.py index be2d6f2b..7bb0c716 100644 --- a/mitmproxy/addons/serverplayback.py +++ b/mitmproxy/addons/serverplayback.py @@ -16,7 +16,7 @@ class ServerPlayback: self.stop = False self.final_flow = None - def load(self, flows): + def load_flows(self, flows): for i in flows: if i.response: l = self.flowmap.setdefault(self._hash(i), []) @@ -100,7 +100,7 @@ class ServerPlayback: flows = io.read_flows_from_paths(options.server_replay) except exceptions.FlowReadException as e: raise exceptions.OptionsError(str(e)) - self.load(flows) + self.load_flows(flows) def tick(self): if self.stop and not self.final_flow.live: diff --git a/mitmproxy/eventsequence.py b/mitmproxy/eventsequence.py index bc6660e0..30c037f1 100644 --- a/mitmproxy/eventsequence.py +++ b/mitmproxy/eventsequence.py @@ -32,7 +32,7 @@ Events = frozenset([ "configure", "done", "log", - "start", + "load", "running", "tick", ]) diff --git a/mitmproxy/script/concurrent.py b/mitmproxy/script/concurrent.py index 7573f2a5..cbb3beb0 100644 --- a/mitmproxy/script/concurrent.py +++ b/mitmproxy/script/concurrent.py @@ -12,7 +12,7 @@ class ScriptThread(basethread.BaseThread): def concurrent(fn): - if fn.__name__ not in eventsequence.Events - {"start", "configure", "tick"}: + if fn.__name__ not in eventsequence.Events - {"load", "configure", "tick"}: raise NotImplementedError( "Concurrent decorator not supported for '%s' method." % fn.__name__ ) diff --git a/mitmproxy/tools/dump.py b/mitmproxy/tools/dump.py index e6f0c3df..6329f6b7 100644 --- a/mitmproxy/tools/dump.py +++ b/mitmproxy/tools/dump.py @@ -16,11 +16,11 @@ class ErrorCheck: class DumpMaster(master.Master): def __init__( - self, - options: options.Options, - server, - with_termlog=True, - with_dumper=True, + self, + options: options.Options, + server, + with_termlog=True, + with_dumper=True, ) -> None: master.Master.__init__(self, options, server) self.errorcheck = ErrorCheck() @@ -46,7 +46,6 @@ exclude = [tool:individual_coverage] exclude = - mitmproxy/addonmanager.py mitmproxy/addons/onboardingapp/app.py mitmproxy/addons/termlog.py mitmproxy/contentviews/base.py diff --git a/test/mitmproxy/addons/test_script.py b/test/mitmproxy/addons/test_script.py index 84d36e2a..c68981de 100644 --- a/test/mitmproxy/addons/test_script.py +++ b/test/mitmproxy/addons/test_script.py @@ -9,6 +9,7 @@ from unittest import mock from mitmproxy.test import tflow from mitmproxy.test import tutils from mitmproxy.test import taddons +from mitmproxy import addonmanager from mitmproxy import exceptions from mitmproxy import options from mitmproxy import proxy @@ -104,7 +105,7 @@ def test_load_script(): "mitmproxy/data/addonscripts/recorder.py" ), [] ) - assert ns.start + assert ns.load def test_script_print_stdout(): @@ -116,7 +117,7 @@ def test_script_print_stdout(): "mitmproxy/data/addonscripts/print.py" ), [] ) - ns.start(tctx.options) + ns.load(addonmanager.Loader(tctx.master)) mock_warn.assert_called_once_with("stdoutprint") @@ -129,7 +130,7 @@ class TestScript: ) ) sc.load_script() - assert sc.ns.call_log[0][0:2] == ("solo", "start") + assert sc.ns.call_log[0][0:2] == ("solo", "load") sc.ns.call_log = [] f = tflow.tflow(resp=True) @@ -157,7 +158,8 @@ class TestScript: sc = script.Script( tutils.test_data.path("mitmproxy/data/addonscripts/error.py") ) - sc.start(tctx.options) + l = addonmanager.Loader(tctx.master) + sc.load(l) f = tflow.tflow(resp=True) sc.request(f) assert tctx.master.logs[0].level == "error" @@ -173,10 +175,11 @@ class TestScript: "mitmproxy/data/addonscripts/addon.py" ) ) - sc.start(tctx.options) + l = addonmanager.Loader(tctx.master) + sc.load(l) tctx.configure(sc) assert sc.ns.event_log == [ - 'scriptstart', 'addonstart', 'addonconfigure' + 'scriptload', 'addonload', 'addonconfigure' ] @@ -213,7 +216,7 @@ class TestScriptLoader: ), [f] ) evts = [i[1] for i in sc.ns.call_log] - assert evts == ['start', 'requestheaders', 'request', 'responseheaders', 'response', 'done'] + assert evts == ['load', 'requestheaders', 'request', 'responseheaders', 'response', 'done'] f = tflow.tflow(resp=True) with m.handlecontext(): @@ -271,15 +274,15 @@ class TestScriptLoader: ) debug = [i.msg for i in tctx.master.logs if i.level == "debug"] assert debug == [ - 'a start', + 'a load', 'a configure', 'a running', - 'b start', + 'b load', 'b configure', 'b running', - 'c start', + 'c load', 'c configure', 'c running', ] @@ -307,7 +310,7 @@ class TestScriptLoader: assert debug == [ 'c done', 'b done', - 'x start', + 'x load', 'x configure', 'x running', ] diff --git a/test/mitmproxy/addons/test_serverplayback.py b/test/mitmproxy/addons/test_serverplayback.py index 02642c35..7078b66e 100644 --- a/test/mitmproxy/addons/test_serverplayback.py +++ b/test/mitmproxy/addons/test_serverplayback.py @@ -44,12 +44,12 @@ def test_server_playback(): assert not sp.flowmap - sp.load([f]) + sp.load_flows([f]) assert sp.flowmap assert sp.next_flow(f) assert not sp.flowmap - sp.load([f]) + sp.load_flows([f]) assert sp.flowmap sp.clear() assert not sp.flowmap @@ -192,7 +192,7 @@ def test_load(): r2 = tflow.tflow(resp=True) r2.request.headers["key"] = "two" - s.load([r, r2]) + s.load_flows([r, r2]) assert s.count() == 2 @@ -218,7 +218,7 @@ def test_load_with_server_replay_nopop(): r2 = tflow.tflow(resp=True) r2.request.headers["key"] = "two" - s.load([r, r2]) + s.load_flows([r, r2]) assert s.count() == 2 s.next_flow(r) @@ -319,7 +319,7 @@ def test_server_playback_full(): f = tflow.tflow() f.response = mitmproxy.test.tutils.tresp(content=f.request.content) - s.load([f, f]) + s.load_flows([f, f]) tf = tflow.tflow() assert not tf.response @@ -352,7 +352,7 @@ def test_server_playback_kill(): f = tflow.tflow() f.response = mitmproxy.test.tutils.tresp(content=f.request.content) - s.load([f]) + s.load_flows([f]) f = tflow.tflow() f.request.host = "nonexistent" diff --git a/test/mitmproxy/data/addonscripts/addon.py b/test/mitmproxy/data/addonscripts/addon.py index f34f41cb..beef2ce7 100644 --- a/test/mitmproxy/data/addonscripts/addon.py +++ b/test/mitmproxy/data/addonscripts/addon.py @@ -6,8 +6,8 @@ class Addon: def event_log(self): return event_log - def start(self, opts): - event_log.append("addonstart") + def load(self, opts): + event_log.append("addonload") def configure(self, options, updated): event_log.append("addonconfigure") @@ -17,6 +17,6 @@ def configure(options, updated): event_log.append("addonconfigure") -def start(opts): - event_log.append("scriptstart") - return Addon() +def load(l): + event_log.append("scriptload") + l.boot_into(Addon()) diff --git a/test/mitmproxy/data/addonscripts/concurrent_decorator_class.py b/test/mitmproxy/data/addonscripts/concurrent_decorator_class.py index 10ba24cd..8e6988d4 100644 --- a/test/mitmproxy/data/addonscripts/concurrent_decorator_class.py +++ b/test/mitmproxy/data/addonscripts/concurrent_decorator_class.py @@ -9,5 +9,5 @@ class ConcurrentClass: time.sleep(0.1) -def start(opts): - return ConcurrentClass() +def load(l): + l.boot_into(ConcurrentClass()) diff --git a/test/mitmproxy/data/addonscripts/print.py b/test/mitmproxy/data/addonscripts/print.py index fdfcc3d9..93b65a64 100644 --- a/test/mitmproxy/data/addonscripts/print.py +++ b/test/mitmproxy/data/addonscripts/print.py @@ -1,2 +1,2 @@ -def start(opts): +def load(l): print("stdoutprint") diff --git a/test/mitmproxy/data/addonscripts/recorder.py b/test/mitmproxy/data/addonscripts/recorder.py index aff524a8..fe497b05 100644 --- a/test/mitmproxy/data/addonscripts/recorder.py +++ b/test/mitmproxy/data/addonscripts/recorder.py @@ -22,5 +22,5 @@ class CallLogger: raise AttributeError -def start(opts): - return CallLogger(*sys.argv[1:]) +def load(l): + l.boot_into(CallLogger(*sys.argv[1:])) diff --git a/test/mitmproxy/script/test_concurrent.py b/test/mitmproxy/script/test_concurrent.py index 0e397b8f..86efdfc2 100644 --- a/test/mitmproxy/script/test_concurrent.py +++ b/test/mitmproxy/script/test_concurrent.py @@ -2,6 +2,7 @@ from mitmproxy.test import tflow from mitmproxy.test import tutils from mitmproxy.test import taddons +from mitmproxy import addonmanager from mitmproxy import controller from mitmproxy.addons import script @@ -24,7 +25,8 @@ class TestConcurrent(tservers.MasterTest): "mitmproxy/data/addonscripts/concurrent_decorator.py" ) ) - sc.start(tctx.options) + l = addonmanager.Loader(tctx.master) + sc.load(l) f1, f2 = tflow.tflow(), tflow.tflow() tctx.cycle(sc, f1) @@ -42,7 +44,8 @@ class TestConcurrent(tservers.MasterTest): "mitmproxy/data/addonscripts/concurrent_decorator_err.py" ) ) - sc.start(tctx.options) + l = addonmanager.Loader(tctx.master) + sc.load(l) assert tctx.master.has_log("decorator not supported") def test_concurrent_class(self): @@ -52,7 +55,8 @@ class TestConcurrent(tservers.MasterTest): "mitmproxy/data/addonscripts/concurrent_decorator_class.py" ) ) - sc.start(tctx.options) + l = addonmanager.Loader(tctx.master) + sc.load(l) f1, f2 = tflow.tflow(), tflow.tflow() tctx.cycle(sc, f1) diff --git a/test/mitmproxy/test_addonmanager.py b/test/mitmproxy/test_addonmanager.py index e7be25b8..5bb88eb6 100644 --- a/test/mitmproxy/test_addonmanager.py +++ b/test/mitmproxy/test_addonmanager.py @@ -5,6 +5,7 @@ from mitmproxy import exceptions from mitmproxy import options from mitmproxy import master from mitmproxy import proxy +from mitmproxy.test import tflow class TAddon: @@ -23,6 +24,58 @@ class TAddon: self.custom_called = True +class THalt: + def event_custom(self): + raise exceptions.AddonHalt + + +class AOption: + def load(self, l): + l.add_option("custom_option", bool, False, "help") + + +class AChain: + def __init__(self, name, next): + self.name = name + self.next = next + + def load(self, l): + if self.next: + l.boot_into(self.next) + + def __repr__(self): + return "<%s>" % self.name + + +def test_halt(): + o = options.Options() + m = master.Master(o, proxy.DummyServer(o)) + a = addonmanager.AddonManager(m) + halt = THalt() + end = TAddon("end") + a.add(halt) + a.add(end) + + a.trigger("custom") + assert not end.custom_called + + a.remove(halt) + a.trigger("custom") + assert end.custom_called + + +def test_lifecycle(): + o = options.Options() + m = master.Master(o, proxy.DummyServer(o)) + a = addonmanager.AddonManager(m) + a.add(TAddon("one")) + + f = tflow.tflow() + a.handle_lifecycle("request", f) + + a.configure_all(o, o.keys()) + + def test_simple(): o = options.Options() m = master.Master(o, proxy.DummyServer(o)) @@ -30,10 +83,13 @@ def test_simple(): with pytest.raises(exceptions.AddonError): a.invoke_addon(TAddon("one"), "done") + assert len(a) == 0 a.add(TAddon("one")) assert a.get("one") assert not a.get("two") + assert len(a) == 1 a.clear() + assert len(a) == 0 assert not a.chain a.add(TAddon("one")) @@ -41,7 +97,46 @@ def test_simple(): with pytest.raises(exceptions.AddonError): a.trigger("tick") + a.remove(a.get("one")) + assert not a.get("one") + ta = TAddon("one") a.add(ta) a.trigger("custom") assert ta.custom_called + + +def test_load_option(): + o = options.Options() + m = master.Master(o, proxy.DummyServer(o)) + a = addonmanager.AddonManager(m) + a.add(AOption()) + assert "custom_option" in m.options._options + + +def test_loadchain(): + o = options.Options() + m = master.Master(o, proxy.DummyServer(o)) + a = addonmanager.AddonManager(m) + + a.add(AChain("one", None)) + assert a.get("one") + a.clear() + + a.add(AChain("one", AChain("two", None))) + assert not a.get("one") + assert a.get("two") + a.clear() + + a.add(AChain("one", AChain("two", AChain("three", None)))) + assert not a.get("one") + assert not a.get("two") + assert a.get("three") + a.clear() + + a.add(AChain("one", AChain("two", AChain("three", AChain("four", None))))) + assert not a.get("one") + assert not a.get("two") + assert not a.get("three") + assert a.get("four") + a.clear() |