diff options
| -rw-r--r-- | mitmproxy/optmanager.py | 159 | ||||
| -rw-r--r-- | mitmproxy/tools/console/options.py | 3 | ||||
| -rw-r--r-- | mitmproxy/tools/main.py | 2 | ||||
| -rw-r--r-- | test/mitmproxy/test_optmanager.py | 30 |
4 files changed, 100 insertions, 94 deletions
diff --git a/mitmproxy/optmanager.py b/mitmproxy/optmanager.py index fc540b74..225a2f73 100644 --- a/mitmproxy/optmanager.py +++ b/mitmproxy/optmanager.py @@ -222,83 +222,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 @@ -468,3 +391,85 @@ def dump_defaults(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. + """ + data = parse(text) + try: + opts.update(**data) + except KeyError as v: + raise exceptions.OptionsError(v) + + +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. + """ + 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: + load(opts, txt) + except exceptions.OptionsError as e: + raise exceptions.OptionsError( + "Error reading %s: %s" % (p, e) + ) + + +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/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/main.py b/mitmproxy/tools/main.py index cc2310c2..8210580c 100644 --- a/mitmproxy/tools/main.py +++ b/mitmproxy/tools/main.py @@ -69,7 +69,7 @@ def run(MasterKlass, args): # pragma: no cover args = parser.parse_args(args) master = None try: - opts.load_paths(args.conf) + optmanager.load_paths(opts, args.conf) server = process_options(parser, opts, args) master = MasterKlass(opts, server) master.addons.configure_all(opts, opts.keys()) diff --git a/test/mitmproxy/test_optmanager.py b/test/mitmproxy/test_optmanager.py index 8ca35984..012c463c 100644 --- a/test/mitmproxy/test_optmanager.py +++ b/test/mitmproxy/test_optmanager.py @@ -199,61 +199,61 @@ 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) + optmanager.load(o2, t) with pytest.raises(exceptions.OptionsError, matches='No such option: foobar'): - o2.load("foobar: '123'") + optmanager.load(o2, "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) + optmanager.load_paths(o, dst) def test_merge(): |