diff options
| -rw-r--r-- | mitmproxy/optmanager.py | 70 | ||||
| -rw-r--r-- | setup.py | 1 | ||||
| -rw-r--r-- | test/mitmproxy/test_optmanager.py | 55 |
3 files changed, 125 insertions, 1 deletions
diff --git a/mitmproxy/optmanager.py b/mitmproxy/optmanager.py index 304f5129..cbf656f5 100644 --- a/mitmproxy/optmanager.py +++ b/mitmproxy/optmanager.py @@ -5,6 +5,9 @@ import inspect import copy import functools import weakref +import os + +import ruamel.yaml from mitmproxy import exceptions from mitmproxy.utils import typecheck @@ -171,6 +174,73 @@ class OptManager(metaclass=_DefaultsMeta): if getattr(self, option) != self._defaults[option]: return True + 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): + data = open(path, "r").read() + else: + data = "" + data = self.serialize(data, defaults) + fp = open(path, "w") + fp.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._opts: + 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.Loader) + 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) + for k, v in data.items(): + setattr(self, k, 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: + if os.path.exists(p) and os.path.isfile(p): + txt = open(p, "r").read() + self.load(txt) + def __repr__(self): options = pprint.pformat(self._opts, indent=4).strip(" {}") if "\n" in options: @@ -79,6 +79,7 @@ setup( "pyparsing>=2.1.3, <2.2", "pyperclip>=1.5.22, <1.6", "requests>=2.9.1, <3", + "ruamel.yaml>=0.13.2, <0.14", "tornado>=4.3, <4.5", "urwid>=1.3.1, <1.4", "watchdog>=0.8.3, <0.9", diff --git a/test/mitmproxy/test_optmanager.py b/test/mitmproxy/test_optmanager.py index 385cf621..0c98daea 100644 --- a/test/mitmproxy/test_optmanager.py +++ b/test/mitmproxy/test_optmanager.py @@ -1,5 +1,7 @@ import copy +import os +from mitmproxy import options from mitmproxy import optmanager from mitmproxy import exceptions from mitmproxy.test import tutils @@ -24,7 +26,7 @@ class TD2(TD): def __init__(self, *, three="dthree", four="dfour", **kwargs): self.three = three self.four = four - super().__init__(**kwargs) + super().__init__(three=three, **kwargs) def test_defaults(): @@ -167,3 +169,54 @@ def test_repr(): 'one': 'xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx', 'two': None })""" + + +def test_serialize(): + o = TD2() + o.three = "set" + assert "dfour" in o.serialize(None, defaults=True) + + data = o.serialize(None) + assert "dfour" not in data + + o2 = TD2() + o2.load(data) + assert o2 == o + + t = """ + unknown: foo + """ + data = o.serialize(t) + o2 = TD2() + o2.load(data) + assert o2 == o + + t = "invalid: foo\ninvalid" + tutils.raises("config error", o2.load, t) + + t = "invalid" + tutils.raises("config error", o2.load, t) + + t = "" + o2.load(t) + + +def test_serialize_defaults(): + o = options.Options() + assert o.serialize(None, defaults=True) + + +def test_saving(): + o = TD2() + o.three = "set" + with tutils.tmpdir() as tdir: + dst = os.path.join(tdir, "conf") + o.save(dst, defaults=True) + + o2 = TD2() + o2.load_paths(dst) + o2.three = "foo" + o2.save(dst, defaults=True) + + o.load_paths(dst) + assert o.three == "foo" |