From e943147fc392383ada7607124e84b1a7db4d4b28 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Thu, 8 Dec 2016 16:12:40 +1300 Subject: [PATCH] options: add a merge method to assimilate options Fixes #953 --- mitmproxy/addons/script.py | 6 ++-- mitmproxy/optmanager.py | 15 +++++++++ mitmproxy/tools/console/grideditor/editors.py | 2 +- mitmproxy/tools/main.py | 32 +++++++------------ test/mitmproxy/addons/test_script.py | 4 +-- test/mitmproxy/test_optmanager.py | 18 +++++++++++ test/mitmproxy/test_proxy.py | 6 ++-- 7 files changed, 54 insertions(+), 29 deletions(-) diff --git a/mitmproxy/addons/script.py b/mitmproxy/addons/script.py index 12544b27f..c89fa085a 100644 --- a/mitmproxy/addons/script.py +++ b/mitmproxy/addons/script.py @@ -32,7 +32,7 @@ def parse_command(command): Returns a (path, args) tuple. """ if not command or not command.strip(): - raise exceptions.AddonError("Empty script command.") + raise exceptions.OptionsError("Empty script command.") # Windows: escape all backslashes in the path. if os.name == "nt": # pragma: no cover backslashes = shlex.split(command, posix=False)[0].count("\\") @@ -40,13 +40,13 @@ def parse_command(command): args = shlex.split(command) # pragma: no cover args[0] = os.path.expanduser(args[0]) if not os.path.exists(args[0]): - raise exceptions.AddonError( + raise exceptions.OptionsError( ("Script file not found: %s.\r\n" "If your script path contains spaces, " "make sure to wrap it in additional quotes, e.g. -s \"'./foo bar/baz.py' --args\".") % args[0]) elif os.path.isdir(args[0]): - raise exceptions.AddonError("Not a file: %s" % args[0]) + raise exceptions.OptionsError("Not a file: %s" % args[0]) return args[0], args[1:] diff --git a/mitmproxy/optmanager.py b/mitmproxy/optmanager.py index e497c3b6a..78b358c9b 100644 --- a/mitmproxy/optmanager.py +++ b/mitmproxy/optmanager.py @@ -244,6 +244,21 @@ class OptManager(metaclass=_DefaultsMeta): txt = f.read() self.load(txt) + def merge(self, opts): + """ + Merge a dict of options into this object. Options that have None + value are ignored. Lists and tuples are appended to the current + option value. + """ + toset = {} + for k, v in opts.items(): + if v is not None: + if isinstance(v, (list, tuple)): + toset[k] = getattr(self, k) + v + else: + toset[k] = v + self.update(**toset) + def __repr__(self): options = pprint.pformat(self._opts, indent=4).strip(" {}") if "\n" in options: diff --git a/mitmproxy/tools/console/grideditor/editors.py b/mitmproxy/tools/console/grideditor/editors.py index 5e3f3d42d..c7c05c35d 100644 --- a/mitmproxy/tools/console/grideditor/editors.py +++ b/mitmproxy/tools/console/grideditor/editors.py @@ -164,7 +164,7 @@ class ScriptEditor(base.GridEditor): def is_error(self, col, val): try: script.parse_command(val) - except exceptions.AddonError as e: + except exceptions.OptionsError as e: return str(e) diff --git a/mitmproxy/tools/main.py b/mitmproxy/tools/main.py index 3c0e44bc9..d07ae6664 100644 --- a/mitmproxy/tools/main.py +++ b/mitmproxy/tools/main.py @@ -18,14 +18,6 @@ from mitmproxy.utils import version_check # noqa from mitmproxy.utils import debug # noqa -def notnone(d): - ret = {} - for k, v in d.items(): - if v is not None: - ret[k] = v - return ret - - def assert_utf8_env(): spec = "" for i in ["LANG", "LC_CTYPE", "LC_ALL"]: @@ -74,9 +66,9 @@ def mitmproxy(args=None): # pragma: no cover try: console_options = options.Options() console_options.load_paths(args.conf) - console_options.update(**notnone(cmdline.get_common_options(args))) - console_options.update( - **notnone(dict( + console_options.merge(cmdline.get_common_options(args)) + console_options.merge( + dict( palette = args.palette, palette_transparent = args.palette_transparent, eventlog = args.eventlog, @@ -85,7 +77,7 @@ def mitmproxy(args=None): # pragma: no cover filter = args.filter, no_mouse = args.no_mouse, order = args.order, - )) + ) ) server = process_options(parser, console_options, args) @@ -113,13 +105,13 @@ def mitmdump(args=None): # pragma: no cover try: dump_options = options.Options() dump_options.load_paths(args.conf) - dump_options.update(**notnone(cmdline.get_common_options(args))) - dump_options.update( - **notnone(dict( + dump_options.merge(cmdline.get_common_options(args)) + dump_options.merge( + dict( flow_detail = args.flow_detail, keepserving = args.keepserving, filtstr = " ".join(args.filter) if args.filter else None, - )) + ) ) server = process_options(parser, dump_options, args) @@ -152,15 +144,15 @@ def mitmweb(args=None): # pragma: no cover try: web_options = options.Options() web_options.load_paths(args.conf) - web_options.update(**notnone(cmdline.get_common_options(args))) - web_options.update( - **notnone(dict( + web_options.merge(cmdline.get_common_options(args)) + web_options.merge( + dict( intercept = args.intercept, open_browser = args.open_browser, wdebug = args.wdebug, wiface = args.wiface, wport = args.wport, - )) + ) ) server = process_options(parser, web_options, args) m = web.master.WebMaster(web_options, server) diff --git a/test/mitmproxy/addons/test_script.py b/test/mitmproxy/addons/test_script.py index a41f61035..06463fa36 100644 --- a/test/mitmproxy/addons/test_script.py +++ b/test/mitmproxy/addons/test_script.py @@ -64,10 +64,10 @@ def test_reloadhandler(): class TestParseCommand: def test_empty_command(self): - with tutils.raises(exceptions.AddonError): + with tutils.raises(exceptions.OptionsError): script.parse_command("") - with tutils.raises(exceptions.AddonError): + with tutils.raises(exceptions.OptionsError): script.parse_command(" ") def test_no_script_file(self): diff --git a/test/mitmproxy/test_optmanager.py b/test/mitmproxy/test_optmanager.py index c6eb25345..9e9387034 100644 --- a/test/mitmproxy/test_optmanager.py +++ b/test/mitmproxy/test_optmanager.py @@ -220,3 +220,21 @@ def test_saving(): o.load_paths(dst) assert o.three == "foo" + + +class TM(optmanager.OptManager): + def __init__(self, one="one", two=["foo"], three=None): + self.one = one + self.two = two + self.three = three + super().__init__() + + +def test_merge(): + m = TM() + m.merge(dict(one="two")) + assert m.one == "two" + m.merge(dict(one=None)) + assert m.one == "two" + m.merge(dict(two=["bar"])) + assert m.two == ["foo", "bar"] diff --git a/test/mitmproxy/test_proxy.py b/test/mitmproxy/test_proxy.py index 91da47a0b..0d63f1472 100644 --- a/test/mitmproxy/test_proxy.py +++ b/test/mitmproxy/test_proxy.py @@ -5,7 +5,6 @@ import argparse from OpenSSL import SSL from mitmproxy.tools import cmdline -from mitmproxy.tools import main from mitmproxy import options from mitmproxy.proxy import ProxyConfig from mitmproxy import connections @@ -76,8 +75,9 @@ class TestProcessProxyOptions: parser = MockParser() cmdline.common_options(parser) args = parser.parse_args(args=args) - opts = cmdline.get_common_options(args) - pconf = config.ProxyConfig(options.Options(**main.notnone(opts))) + opts = options.Options() + opts.merge(cmdline.get_common_options(args)) + pconf = config.ProxyConfig(opts) return parser, pconf def assert_err(self, err, *args):