From d74cac265a9d1d8ce176a7ef96be2d91c4f40819 Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Mon, 5 Dec 2016 07:18:53 +1300
Subject: [PATCH] Add YAML serialization of options

This uses ruamel.yaml. The library seems well-supported, and can do in-place
modification of config files that retains user comments and file structure.
---
 mitmproxy/optmanager.py           | 70 +++++++++++++++++++++++++++++++
 setup.py                          |  1 +
 test/mitmproxy/test_optmanager.py | 55 +++++++++++++++++++++++-
 3 files changed, 125 insertions(+), 1 deletion(-)

diff --git a/mitmproxy/optmanager.py b/mitmproxy/optmanager.py
index 304f51291..cbf656f57 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:
diff --git a/setup.py b/setup.py
index 56ba46fc8..35f7edb36 100644
--- a/setup.py
+++ b/setup.py
@@ -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 385cf621c..0c98daea0 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"