diff options
author | Aldo Cortesi <aldo@corte.si> | 2017-03-14 13:02:58 +1300 |
---|---|---|
committer | GitHub <noreply@github.com> | 2017-03-14 13:02:58 +1300 |
commit | 124a6c9e5af44121208c3362215cc3ea895ffad7 (patch) | |
tree | 8c8f24f119aae9b22832bfa0e6b5c0a397133d17 | |
parent | ee65894d40f5a9f73125a8d3e73ba50540939e5b (diff) | |
parent | 1b301ad5bbe5765b608bf4f8480720065d3343c8 (diff) | |
download | mitmproxy-124a6c9e5af44121208c3362215cc3ea895ffad7.tar.gz mitmproxy-124a6c9e5af44121208c3362215cc3ea895ffad7.tar.bz2 mitmproxy-124a6c9e5af44121208c3362215cc3ea895ffad7.zip |
Merge pull request #2129 from cortesi/addonopts
Custom options for addons
35 files changed, 321 insertions, 204 deletions
diff --git a/examples/complex/dns_spoofing.py b/examples/complex/dns_spoofing.py index 2fd6b699..ca2bcd35 100644 --- a/examples/complex/dns_spoofing.py +++ b/examples/complex/dns_spoofing.py @@ -1,11 +1,12 @@ """ -This script makes it possible to use mitmproxy in scenarios where IP spoofing has been used to redirect -connections to mitmproxy. The way this works is that we rely on either the TLS Server Name Indication (SNI) or the -Host header of the HTTP request. -Of course, this is not foolproof - if an HTTPS connection comes without SNI, we don't -know the actual target and cannot construct a certificate that looks valid. -Similarly, if there's no Host header or a spoofed Host header, we're out of luck as well. -Using transparent mode is the better option most of the time. +This script makes it possible to use mitmproxy in scenarios where IP spoofing +has been used to redirect connections to mitmproxy. The way this works is that +we rely on either the TLS Server Name Indication (SNI) or the Host header of the +HTTP request. Of course, this is not foolproof - if an HTTPS connection comes +without SNI, we don't know the actual target and cannot construct a certificate +that looks valid. Similarly, if there's no Host header or a spoofed Host header, +we're out of luck as well. Using transparent mode is the better option most of +the time. Usage: mitmproxy @@ -53,5 +54,5 @@ class Rerouter: flow.request.port = port -def start(): +def start(opts): return Rerouter() diff --git a/examples/complex/har_dump.py b/examples/complex/har_dump.py index 86a33684..9a86e45e 100644 --- a/examples/complex/har_dump.py +++ b/examples/complex/har_dump.py @@ -25,7 +25,7 @@ HAR = {} SERVERS_SEEN = set() -def start(): +def start(opts): """ Called once on script startup before any other events. """ diff --git a/examples/complex/remote_debug.py b/examples/complex/remote_debug.py index fb864f78..ae0dffc1 100644 --- a/examples/complex/remote_debug.py +++ b/examples/complex/remote_debug.py @@ -14,6 +14,6 @@ Usage: """ -def start(): +def start(opts): 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 40c1051d..6dba7ca1 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(): +def start(opts): 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 6443798a..9270be09 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(): +def start(opts): return AddHeader() diff --git a/examples/simple/custom_contentview.py b/examples/simple/custom_contentview.py index 1f3a38ec..4bc17af0 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(): +def start(opts): contentviews.add(view) diff --git a/examples/simple/custom_option.py b/examples/simple/custom_option.py new file mode 100644 index 00000000..324d27e7 --- /dev/null +++ b/examples/simple/custom_option.py @@ -0,0 +1,11 @@ +from mitmproxy import ctx + + +def start(options): + ctx.log.info("Registering option 'custom'") + options.add_option("custom", bool, False, "A custom option") + + +def configure(options, updated): + if "custom" in updated: + ctx.log.info("custom option value: %s" % options.custom) diff --git a/examples/simple/filter_flows.py b/examples/simple/filter_flows.py index 29d0a9b8..24e8b6c1 100644 --- a/examples/simple/filter_flows.py +++ b/examples/simple/filter_flows.py @@ -17,7 +17,7 @@ class Filter: print(flow) -def start(): +def start(opts): if len(sys.argv) != 2: raise ValueError("Usage: -s 'filt.py FILTER'") return Filter(sys.argv[1]) diff --git a/examples/simple/io_write_dumpfile.py b/examples/simple/io_write_dumpfile.py index ff1fd0f4..311950af 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(): +def start(opts): if len(sys.argv) != 2: raise ValueError('Usage: -s "flowriter.py filename"') return Writer(sys.argv[1]) diff --git a/examples/simple/log_events.py b/examples/simple/log_events.py index ab1baf75..a81892aa 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(): +def start(opts): 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 e3d5fee9..ab5abf27 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(): +def start(opts): if len(sys.argv) != 2: raise ValueError('Usage: -s "iframe_injector.py url"') return Injector(sys.argv[1]) diff --git a/examples/simple/script_arguments.py b/examples/simple/script_arguments.py index 70851192..b46a1960 100644 --- a/examples/simple/script_arguments.py +++ b/examples/simple/script_arguments.py @@ -9,7 +9,7 @@ class Replacer: flow.response.replace(self.src, self.dst) -def start(): +def start(opts): parser = argparse.ArgumentParser() parser.add_argument("src", type=str) parser.add_argument("dst", type=str) diff --git a/examples/simple/wsgi_flask_app.py b/examples/simple/wsgi_flask_app.py index f95c41e5..db3b1adf 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(): +def start(opts): # 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 db8e0cd7..43e76510 100644 --- a/mitmproxy/addonmanager.py +++ b/mitmproxy/addonmanager.py @@ -1,4 +1,5 @@ from mitmproxy import exceptions +from mitmproxy import eventsequence import pprint @@ -10,7 +11,7 @@ class AddonManager: def __init__(self, master): self.chain = [] self.master = master - master.options.changed.connect(self._options_update) + master.options.changed.connect(self.configure_all) def clear(self): """ @@ -29,22 +30,14 @@ class AddonManager: if name == _get_name(i): return i - def _options_update(self, options, updated): - for i in self.chain: - with self.master.handlecontext(): - self.invoke_with_context(i, "configure", options, updated) + def configure_all(self, options, updated): + self.invoke_all_with_context("configure", options, updated) def startup(self, s): """ Run startup events on addon. """ - self.invoke_with_context(s, "start") - self.invoke_with_context( - s, - "configure", - self.master.options, - self.master.options.keys() - ) + self.invoke_with_context(s, "start", self.master.options) def add(self, *addons): """ @@ -62,8 +55,7 @@ class AddonManager: self.invoke_with_context(addon, "done") def done(self): - for i in self.chain: - self.invoke_with_context(i, "done") + self.invoke_all_with_context("done") def __len__(self): return len(self.chain) @@ -75,7 +67,14 @@ class AddonManager: with self.master.handlecontext(): self.invoke(addon, name, *args, **kwargs) + def invoke_all_with_context(self, name, *args, **kwargs): + with self.master.handlecontext(): + for i in self.chain: + self.invoke(i, name, *args, **kwargs) + def invoke(self, addon, name, *args, **kwargs): + if name not in eventsequence.Events: # prama: no cover + raise NotImplementedError("Unknown event") func = getattr(addon, name, None) if func: if not callable(func): diff --git a/mitmproxy/addons/script.py b/mitmproxy/addons/script.py index a7d3a312..cfbe5284 100644 --- a/mitmproxy/addons/script.py +++ b/mitmproxy/addons/script.py @@ -170,22 +170,23 @@ class Script: def load_script(self): self.ns = load_script(self.path, self.args) - ret = self.run("start") + ret = self.run("start", self.last_options) if ret: self.ns = ret - self.run("start") + self.run("start", self.last_options) 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.start(self.last_options) self.configure(self.last_options, self.last_options.keys()) else: self.run("tick") - def start(self): + def start(self, opts): + self.last_options = opts self.load_script() def configure(self, options, updated): @@ -209,6 +210,12 @@ class ScriptLoader: """ An addon that manages loading scripts from options. """ + def __init__(self): + self.is_running = False + + def running(self): + self.is_running = True + def run_once(self, command, flows): try: sc = Script(command) @@ -267,3 +274,10 @@ class ScriptLoader: for s in newscripts: ctx.master.addons.startup(s) + if self.is_running: + # If we're already running, we configure and tell the addon + # we're up and running. + ctx.master.addons.invoke_with_context( + s, "configure", options, options.keys() + ) + ctx.master.addons.invoke_with_context(s, "running") diff --git a/mitmproxy/addons/termstatus.py b/mitmproxy/addons/termstatus.py new file mode 100644 index 00000000..7b05f409 --- /dev/null +++ b/mitmproxy/addons/termstatus.py @@ -0,0 +1,23 @@ +from mitmproxy import ctx + +""" + A tiny addon to print the proxy status to terminal. Eventually this could + also print some stats on exit. +""" + + +class TermStatus: + def __init__(self): + self.server = False + + def configure(self, options, updated): + if "server" in updated: + self.server = options.server + + def running(self): + if self.server: + ctx.log.info( + "Proxy server listening at http://{}:{}".format( + *ctx.master.server.address, + ) + ) diff --git a/mitmproxy/eventsequence.py b/mitmproxy/eventsequence.py index 5872f607..bc6660e0 100644 --- a/mitmproxy/eventsequence.py +++ b/mitmproxy/eventsequence.py @@ -33,6 +33,7 @@ Events = frozenset([ "done", "log", "start", + "running", "tick", ]) diff --git a/mitmproxy/master.py b/mitmproxy/master.py index 8855452c..79747a97 100644 --- a/mitmproxy/master.py +++ b/mitmproxy/master.py @@ -42,6 +42,7 @@ class Master: self.event_queue = queue.Queue() self.should_exit = threading.Event() self.server = server + self.first_tick = True channel = controller.Channel(self.event_queue, self.should_exit) server.set_channel(channel) @@ -86,6 +87,9 @@ class Master: self.shutdown() def tick(self, timeout): + if self.first_tick: + self.first_tick = False + self.addons.invoke_all_with_context("running") with self.handlecontext(): self.addons("tick") changed = False diff --git a/mitmproxy/optmanager.py b/mitmproxy/optmanager.py index 9553bd32..495354f4 100644 --- a/mitmproxy/optmanager.py +++ b/mitmproxy/optmanager.py @@ -81,8 +81,6 @@ class _Option: class OptManager: """ OptManager is the base class from which Options objects are derived. - Note that the __init__ method of all child classes must force all - arguments to be positional only, by including a "*" argument. .changed is a blinker Signal that triggers whenever options are updated. If any handler in the chain raises an exceptions.OptionsError @@ -176,15 +174,29 @@ class OptManager: o.reset() self.changed.send(self._options.keys()) + def update_known(self, **kwargs): + """ + Update and set all known options from kwargs. Returns a dictionary + of unknown options. + """ + known, unknown = {}, {} + for k, v in kwargs.items(): + if k in self._options: + known[k] = v + else: + unknown[k] = v + updated = set(known.keys()) + if updated: + with self.rollback(updated): + for k, v in known.items(): + self._options[k].set(v) + self.changed.send(self, updated=updated) + return unknown + def update(self, **kwargs): - updated = set(kwargs.keys()) - with self.rollback(updated): - for k, v in kwargs.items(): - if k not in self._options: - raise KeyError("No such option: %s" % k) - self._options[k].set(v) - self.changed.send(self, updated=updated) - return self + u = self.update_known(**kwargs) + if u: + raise KeyError("Unknown options: %s" % ", ".join(u.keys())) def setter(self, attr): """ @@ -222,83 +234,6 @@ class OptManager: """ return self._options[option].has_changed() - def save(self, path, defaults=False): - """ - Save to path. If the destination file exists, modify it in-place. - """ - if os.path.exists(path) and os.path.isfile(path): - with open(path, "r") as f: - data = f.read() - else: - data = "" - data = self.serialize(data, defaults) - with open(path, "w") as f: - f.write(data) - - def serialize(self, text, defaults=False): - """ - Performs a round-trip serialization. If text is not None, it is - treated as a previous serialization that should be modified - in-place. - - - If "defaults" is False, only options with non-default values are - serialized. Default values in text are preserved. - - Unknown options in text are removed. - - Raises OptionsError if text is invalid. - """ - data = self._load(text) - for k in self.keys(): - if defaults or self.has_changed(k): - data[k] = getattr(self, k) - for k in list(data.keys()): - if k not in self._options: - del data[k] - return ruamel.yaml.round_trip_dump(data) - - def _load(self, text): - if not text: - return {} - try: - data = ruamel.yaml.load(text, ruamel.yaml.RoundTripLoader) - except ruamel.yaml.error.YAMLError as v: - snip = v.problem_mark.get_snippet() - raise exceptions.OptionsError( - "Config error at line %s:\n%s\n%s" % - (v.problem_mark.line + 1, snip, v.problem) - ) - if isinstance(data, str): - raise exceptions.OptionsError("Config error - no keys found.") - return data - - def load(self, text): - """ - Load configuration from text, over-writing options already set in - this object. May raise OptionsError if the config file is invalid. - """ - data = self._load(text) - try: - self.update(**data) - except KeyError as v: - raise exceptions.OptionsError(v) - - def load_paths(self, *paths): - """ - Load paths in order. Each path takes precedence over the previous - path. Paths that don't exist are ignored, errors raise an - OptionsError. - """ - for p in paths: - p = os.path.expanduser(p) - if os.path.exists(p) and os.path.isfile(p): - with open(p, "r") as f: - txt = f.read() - try: - self.load(txt) - except exceptions.OptionsError as e: - raise exceptions.OptionsError( - "Error reading %s: %s" % (p, e) - ) - def merge(self, opts): """ Merge a dict of options into this object. Options that have None @@ -324,23 +259,33 @@ class OptManager: options=options ) - def set(self, spec): + def set(self, *spec): + vals = {} + for i in spec: + vals.update(self._setspec(i)) + self.update(**vals) + + def _setspec(self, spec): + d = {} + parts = spec.split("=", maxsplit=1) if len(parts) == 1: optname, optval = parts[0], None else: optname, optval = parts[0], parts[1] + if optname not in self._options: + raise exceptions.OptionsError("No such option %s" % optname) o = self._options[optname] if o.typespec in (str, typing.Optional[str]): - setattr(self, optname, optval) + d[optname] = optval elif o.typespec in (int, typing.Optional[int]): if optval: try: optval = int(optval) except ValueError: raise exceptions.OptionsError("Not an integer: %s" % optval) - setattr(self, optname, optval) + d[optname] = optval elif o.typespec == bool: if not optval or optval == "true": v = True @@ -350,18 +295,15 @@ class OptManager: raise exceptions.OptionsError( "Boolean must be \"true\", \"false\", or have the value " "omitted (a synonym for \"true\")." ) - setattr(self, optname, v) + d[optname] = v elif o.typespec == typing.Sequence[str]: if not optval: - setattr(self, optname, []) + d[optname] = [] else: - setattr( - self, - optname, - getattr(self, optname) + [optval] - ) + d[optname] = getattr(self, optname) + [optval] else: # pragma: no cover raise NotImplementedError("Unsupported option type: %s", o.typespec) + return d def make_parser(self, parser, optname, metavar=None, short=None): o = self._options[optname] @@ -430,7 +372,7 @@ class OptManager: raise ValueError("Unsupported option type: %s", o.typespec) -def dump(opts): +def dump_defaults(opts): """ Dumps an annotated file with all options. """ @@ -461,3 +403,88 @@ def dump(opts): ) s.yaml_set_comment_before_after_key(k, before = "\n" + txt) return ruamel.yaml.round_trip_dump(s) + + +def parse(text): + if not text: + return {} + try: + data = ruamel.yaml.load(text, ruamel.yaml.RoundTripLoader) + except ruamel.yaml.error.YAMLError as v: + snip = v.problem_mark.get_snippet() + raise exceptions.OptionsError( + "Config error at line %s:\n%s\n%s" % + (v.problem_mark.line + 1, snip, v.problem) + ) + if isinstance(data, str): + raise exceptions.OptionsError("Config error - no keys found.") + return data + + +def load(opts, text): + """ + Load configuration from text, over-writing options already set in + this object. May raise OptionsError if the config file is invalid. + + Returns a dictionary of all unknown options. + """ + data = parse(text) + return opts.update_known(**data) + + +def load_paths(opts, *paths): + """ + Load paths in order. Each path takes precedence over the previous + path. Paths that don't exist are ignored, errors raise an + OptionsError. + + Returns a dictionary of unknown options. + """ + ret = {} + for p in paths: + p = os.path.expanduser(p) + if os.path.exists(p) and os.path.isfile(p): + with open(p, "r") as f: + txt = f.read() + try: + ret.update(load(opts, txt)) + except exceptions.OptionsError as e: + raise exceptions.OptionsError( + "Error reading %s: %s" % (p, e) + ) + return ret + + +def serialize(opts, text, defaults=False): + """ + Performs a round-trip serialization. If text is not None, it is + treated as a previous serialization that should be modified + in-place. + + - If "defaults" is False, only options with non-default values are + serialized. Default values in text are preserved. + - Unknown options in text are removed. + - Raises OptionsError if text is invalid. + """ + data = parse(text) + for k in opts.keys(): + if defaults or opts.has_changed(k): + data[k] = getattr(opts, k) + for k in list(data.keys()): + if k not in opts._options: + del data[k] + return ruamel.yaml.round_trip_dump(data) + + +def save(opts, path, defaults=False): + """ + Save to path. If the destination file exists, modify it in-place. + """ + if os.path.exists(path) and os.path.isfile(path): + with open(path, "r") as f: + data = f.read() + else: + data = "" + data = serialize(opts, data, defaults) + with open(path, "w") as f: + f.write(data) diff --git a/mitmproxy/tools/cmdline.py b/mitmproxy/tools/cmdline.py index aaefd10a..c5d2bbd7 100644 --- a/mitmproxy/tools/cmdline.py +++ b/mitmproxy/tools/cmdline.py @@ -23,7 +23,7 @@ def common_options(parser, opts): parser.add_argument( '--options', action='store_true', - help="Dump all options", + help="Show all options and their default values", ) parser.add_argument( "--conf", diff --git a/mitmproxy/tools/console/options.py b/mitmproxy/tools/console/options.py index 33e3ec38..79bb53c2 100644 --- a/mitmproxy/tools/console/options.py +++ b/mitmproxy/tools/console/options.py @@ -1,6 +1,7 @@ import urwid from mitmproxy import contentviews +from mitmproxy import optmanager from mitmproxy.tools.console import common from mitmproxy.tools.console import grideditor from mitmproxy.tools.console import select @@ -173,7 +174,7 @@ class Options(urwid.WidgetWrap): return super().keypress(size, key) def do_save(self, path): - self.master.options.save(path) + optmanager.save(self.master.options, path) return "Saved" def save(self): diff --git a/mitmproxy/tools/dump.py b/mitmproxy/tools/dump.py index e70ce2f9..4bfe2dc4 100644 --- a/mitmproxy/tools/dump.py +++ b/mitmproxy/tools/dump.py @@ -3,7 +3,7 @@ from mitmproxy import exceptions from mitmproxy import addons from mitmproxy import options from mitmproxy import master -from mitmproxy.addons import dumper, termlog +from mitmproxy.addons import dumper, termlog, termstatus class DumpMaster(master.Master): @@ -18,17 +18,11 @@ class DumpMaster(master.Master): master.Master.__init__(self, options, server) self.has_errored = False if with_termlog: - self.addons.add(termlog.TermLog()) + self.addons.add(termlog.TermLog(), termstatus.TermStatus()) self.addons.add(*addons.default_addons()) if with_dumper: self.addons.add(dumper.Dumper()) - if self.options.server: - self.add_log( - "Proxy server listening at http://{}:{}".format(server.address[0], server.address[1]), - "info" - ) - if options.rfile: try: self.load_flows_file(options.rfile) diff --git a/mitmproxy/tools/main.py b/mitmproxy/tools/main.py index 17c1abbb..35567b62 100644 --- a/mitmproxy/tools/main.py +++ b/mitmproxy/tools/main.py @@ -39,15 +39,10 @@ def process_options(parser, opts, args): if args.version: print(debug.dump_system_info()) sys.exit(0) - if args.options: - print(optmanager.dump(opts)) - sys.exit(0) - if args.quiet: + if args.quiet or args.options: + args.verbosity = 0 args.flow_detail = 0 - for i in args.setoptions: - opts.set(i) - adict = {} for n in dir(args): if n in opts: @@ -74,9 +69,17 @@ def run(MasterKlass, args): # pragma: no cover args = parser.parse_args(args) master = None try: - opts.load_paths(args.conf) + unknown = optmanager.load_paths(opts, args.conf) server = process_options(parser, opts, args) master = MasterKlass(opts, server) + master.addons.configure_all(opts, opts.keys()) + remaining = opts.update_known(**unknown) + if remaining and opts.verbosity > 1: + print("Ignored options: %s" % remaining) + if args.options: + print(optmanager.dump_defaults(opts)) + sys.exit(0) + opts.set(*args.setoptions) def cleankill(*args, **kwargs): master.shutdown() diff --git a/test/mitmproxy/addons/test_script.py b/test/mitmproxy/addons/test_script.py index d79ed4ef..4c1b2e43 100644 --- a/test/mitmproxy/addons/test_script.py +++ b/test/mitmproxy/addons/test_script.py @@ -116,9 +116,7 @@ class TestScript: ) ) sc.load_script() - assert sc.ns.call_log == [ - ("solo", "start", (), {}), - ] + assert sc.ns.call_log[0][0:2] == ("solo", "start") sc.ns.call_log = [] f = tflow.tflow(resp=True) @@ -146,7 +144,7 @@ class TestScript: sc = script.Script( tutils.test_data.path("mitmproxy/data/addonscripts/error.py") ) - sc.start() + sc.start(tctx.options) f = tflow.tflow(resp=True) sc.request(f) assert tctx.master.event_log[0][0] == "error" @@ -162,7 +160,7 @@ class TestScript: "mitmproxy/data/addonscripts/addon.py" ) ) - sc.start() + sc.start(tctx.options) tctx.configure(sc) assert sc.ns.event_log == [ 'scriptstart', 'addonstart', 'addonconfigure' @@ -225,24 +223,31 @@ class TestScriptLoader: assert len(m.addons) == 1 def test_dupes(self): - o = options.Options(scripts=["one", "one"]) - m = master.Master(o, proxy.DummyServer()) sc = script.ScriptLoader() - with pytest.raises(exceptions.OptionsError): - m.addons.add(o, sc) + with taddons.context() as tctx: + tctx.master.addons.add(sc) + with pytest.raises(exceptions.OptionsError): + tctx.configure( + sc, + scripts = ["one", "one"] + ) def test_nonexistent(self): - o = options.Options(scripts=["nonexistent"]) - m = master.Master(o, proxy.DummyServer()) sc = script.ScriptLoader() - with pytest.raises(exceptions.OptionsError): - m.addons.add(o, sc) + with taddons.context() as tctx: + tctx.master.addons.add(sc) + with pytest.raises(exceptions.OptionsError): + tctx.configure( + sc, + scripts = ["nonexistent"] + ) def test_order(self): rec = tutils.test_data.path("mitmproxy/data/addonscripts/recorder.py") sc = script.ScriptLoader() with taddons.context() as tctx: tctx.master.addons.add(sc) + sc.running() tctx.configure( sc, scripts = [ @@ -253,9 +258,17 @@ class TestScriptLoader: ) debug = [(i[0], i[1]) for i in tctx.master.event_log if i[0] == "debug"] assert debug == [ - ('debug', 'a start'), ('debug', 'a configure'), - ('debug', 'b start'), ('debug', 'b configure'), - ('debug', 'c start'), ('debug', 'c configure') + ('debug', 'a start'), + ('debug', 'a configure'), + ('debug', 'a running'), + + ('debug', 'b start'), + ('debug', 'b configure'), + ('debug', 'b running'), + + ('debug', 'c start'), + ('debug', 'c configure'), + ('debug', 'c running'), ] tctx.master.event_log = [] tctx.configure( @@ -284,4 +297,5 @@ class TestScriptLoader: ('debug', 'b done'), ('debug', 'x start'), ('debug', 'x configure'), + ('debug', 'x running'), ] diff --git a/test/mitmproxy/addons/test_termstatus.py b/test/mitmproxy/addons/test_termstatus.py new file mode 100644 index 00000000..01c14814 --- /dev/null +++ b/test/mitmproxy/addons/test_termstatus.py @@ -0,0 +1,12 @@ +from mitmproxy.addons import termstatus +from mitmproxy.test import taddons + + +def test_configure(): + ts = termstatus.TermStatus() + with taddons.context() as ctx: + ts.running() + assert not ctx.master.event_log + ctx.configure(ts, server=True) + ts.running() + assert ctx.master.event_log diff --git a/test/mitmproxy/data/addonscripts/addon.py b/test/mitmproxy/data/addonscripts/addon.py index 84173cb6..f34f41cb 100644 --- a/test/mitmproxy/data/addonscripts/addon.py +++ b/test/mitmproxy/data/addonscripts/addon.py @@ -6,7 +6,7 @@ class Addon: def event_log(self): return event_log - def start(self): + def start(self, opts): event_log.append("addonstart") def configure(self, options, updated): @@ -17,6 +17,6 @@ def configure(options, updated): event_log.append("addonconfigure") -def start(): +def start(opts): event_log.append("scriptstart") return Addon() diff --git a/test/mitmproxy/data/addonscripts/concurrent_decorator_class.py b/test/mitmproxy/data/addonscripts/concurrent_decorator_class.py index bd047c99..10ba24cd 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(): +def start(opts): return ConcurrentClass() diff --git a/test/mitmproxy/data/addonscripts/concurrent_decorator_err.py b/test/mitmproxy/data/addonscripts/concurrent_decorator_err.py index 756869c8..7bc28182 100644 --- a/test/mitmproxy/data/addonscripts/concurrent_decorator_err.py +++ b/test/mitmproxy/data/addonscripts/concurrent_decorator_err.py @@ -2,5 +2,5 @@ from mitmproxy.script import concurrent @concurrent -def start(): +def start(opts): pass diff --git a/test/mitmproxy/data/addonscripts/recorder.py b/test/mitmproxy/data/addonscripts/recorder.py index 6b9b6ea8..aff524a8 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(): +def start(opts): return CallLogger(*sys.argv[1:]) diff --git a/test/mitmproxy/proxy/test_server.py b/test/mitmproxy/proxy/test_server.py index aa45761a..16efe415 100644 --- a/test/mitmproxy/proxy/test_server.py +++ b/test/mitmproxy/proxy/test_server.py @@ -302,6 +302,9 @@ class TestHTTP(tservers.HTTPProxyTest, CommonMixin): class TestHTTPAuth(tservers.HTTPProxyTest): def test_auth(self): self.master.addons.add(proxyauth.ProxyAuth()) + self.master.addons.configure_all( + self.master.options, self.master.options.keys() + ) self.master.options.proxyauth = "test:test" assert self.pathod("202").status_code == 407 p = self.pathoc() diff --git a/test/mitmproxy/script/test_concurrent.py b/test/mitmproxy/script/test_concurrent.py index e81c023d..a9b6f0c4 100644 --- a/test/mitmproxy/script/test_concurrent.py +++ b/test/mitmproxy/script/test_concurrent.py @@ -24,7 +24,7 @@ class TestConcurrent(tservers.MasterTest): "mitmproxy/data/addonscripts/concurrent_decorator.py" ) ) - sc.start() + sc.start(tctx.options) f1, f2 = tflow.tflow(), tflow.tflow() tctx.cycle(sc, f1) @@ -42,7 +42,7 @@ class TestConcurrent(tservers.MasterTest): "mitmproxy/data/addonscripts/concurrent_decorator_err.py" ) ) - sc.start() + sc.start(tctx.options) assert "decorator not supported" in tctx.master.event_log[0][1] def test_concurrent_class(self): @@ -52,7 +52,7 @@ class TestConcurrent(tservers.MasterTest): "mitmproxy/data/addonscripts/concurrent_decorator_class.py" ) ) - sc.start() + sc.start(tctx.options) 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 17402e26..3e5f71c6 100644 --- a/test/mitmproxy/test_addonmanager.py +++ b/test/mitmproxy/test_addonmanager.py @@ -10,12 +10,12 @@ from mitmproxy import proxy class TAddon: def __init__(self, name): self.name = name - self.noop_member = True + self.tick = True def __repr__(self): return "Addon(%s)" % self.name - def noop(self): + def done(self): pass @@ -30,6 +30,6 @@ def test_simple(): assert not a.chain a.add(TAddon("one")) - a("noop") + a("done") with pytest.raises(exceptions.AddonError): - a("noop_member") + a("tick") diff --git a/test/mitmproxy/test_optmanager.py b/test/mitmproxy/test_optmanager.py index 0ecfc4e5..df392829 100644 --- a/test/mitmproxy/test_optmanager.py +++ b/test/mitmproxy/test_optmanager.py @@ -83,10 +83,11 @@ def test_options(): with pytest.raises(TypeError): TO(nonexistent = "value") - with pytest.raises(Exception, match="No such option"): + with pytest.raises(Exception, match="Unknown options"): o.nonexistent = "value" - with pytest.raises(Exception, match="No such option"): + with pytest.raises(Exception, match="Unknown options"): o.update(nonexistent = "value") + assert o.update_known(nonexistent = "value") == {"nonexistent": "value"} rec = [] @@ -199,61 +200,63 @@ def test_simple(): def test_serialize(): o = TD2() o.three = "set" - assert "dfour" in o.serialize(None, defaults=True) + assert "dfour" in optmanager.serialize(o, None, defaults=True) - data = o.serialize(None) + data = optmanager.serialize(o, None) assert "dfour" not in data o2 = TD2() - o2.load(data) + optmanager.load(o2, data) assert o2 == o t = """ unknown: foo """ - data = o.serialize(t) + data = optmanager.serialize(o, t) o2 = TD2() - o2.load(data) + optmanager.load(o2, data) assert o2 == o t = "invalid: foo\ninvalid" with pytest.raises(Exception, match="Config error"): - o2.load(t) + optmanager.load(o2, t) t = "invalid" with pytest.raises(Exception, match="Config error"): - o2.load(t) + optmanager.load(o2, t) t = "" - o2.load(t) - - with pytest.raises(exceptions.OptionsError, matches='No such option: foobar'): - o2.load("foobar: '123'") + optmanager.load(o2, t) + assert optmanager.load(o2, "foobar: '123'") == {"foobar": "123"} def test_serialize_defaults(): o = options.Options() - assert o.serialize(None, defaults=True) + assert optmanager.serialize(o, None, defaults=True) def test_saving(tmpdir): o = TD2() o.three = "set" dst = str(tmpdir.join("conf")) - o.save(dst, defaults=True) + optmanager.save(o, dst, defaults=True) o2 = TD2() - o2.load_paths(dst) + optmanager.load_paths(o2, dst) o2.three = "foo" - o2.save(dst, defaults=True) + optmanager.save(o2, dst, defaults=True) - o.load_paths(dst) + optmanager.load_paths(o, dst) assert o.three == "foo" with open(dst, 'a') as f: f.write("foobar: '123'") - with pytest.raises(exceptions.OptionsError, matches=''): - o.load_paths(dst) + assert optmanager.load_paths(o, dst) == {"foobar": "123"} + + with open(dst, 'a') as f: + f.write("'''") + with pytest.raises(exceptions.OptionsError): + optmanager.load_paths(o, dst) def test_merge(): @@ -280,9 +283,9 @@ def test_option(): assert o2 != o -def test_dump(): +def test_dump_defaults(): o = options.Options() - assert optmanager.dump(o) + assert optmanager.dump_defaults(o) class TTypes(optmanager.OptManager): @@ -346,3 +349,6 @@ def test_set(): assert opts.seqstr == ["foo", "bar"] opts.set("seqstr") assert opts.seqstr == [] + + with pytest.raises(exceptions.OptionsError): + opts.set("nonexistent=wobble") diff --git a/test/mitmproxy/tools/console/test_master.py b/test/mitmproxy/tools/console/test_master.py index 0bf3734b..6c716ad1 100644 --- a/test/mitmproxy/tools/console/test_master.py +++ b/test/mitmproxy/tools/console/test_master.py @@ -28,7 +28,9 @@ class TestMaster(tservers.MasterTest): if "verbosity" not in opts: opts["verbosity"] = 1 o = options.Options(**opts) - return console.master.ConsoleMaster(o, proxy.DummyServer()) + m = console.master.ConsoleMaster(o, proxy.DummyServer()) + m.addons.configure_all(o, o.keys()) + return m def test_basic(self): m = self.mkmaster() diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py index a8aaa358..c47411ee 100644 --- a/test/mitmproxy/tservers.py +++ b/test/mitmproxy/tservers.py @@ -79,6 +79,8 @@ class TestMaster(master.Master): self.state = TestState() self.addons.add(self.state) self.addons.add(*addons) + self.addons.configure_all(self.options, self.options.keys()) + self.addons.invoke_all_with_context("running") def clear_log(self): self.tlog = [] |