optmanager: make serialization and loading more robust

Also handle errors in console options manager.

3 files changed, 47 insertions, 13 deletions

diff --git a/mitmproxy/optmanager.py b/mitmproxy/optmanager.py
index c58bc7d0..5ab3496c 100644
--- a/mitmproxy/optmanager.py
+++ b/mitmproxy/optmanager.py
@@ -422,11 +422,14 @@ def parse(text):
     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 hasattr(v, "problem_mark"):
+            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)
+            )
+        else:
+            raise exceptions.OptionsError("Could not parse options.")
     if isinstance(data, str):
         raise exceptions.OptionsError("Config error - no keys found.")
     return data
@@ -455,8 +458,13 @@ def load_paths(opts, *paths):
     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()
+            with open(p, "rt", encoding="utf8") as f:
+                try:
+                    txt = f.read()
+                except UnicodeDecodeError as e:
+                    raise exceptions.OptionsError(
+                        "Error reading %s: %s" % (p, e)
+                    )
             try:
                 ret.update(load(opts, txt))
             except exceptions.OptionsError as e:
@@ -490,12 +498,19 @@ def serialize(opts, text, defaults=False):
 def save(opts, path, defaults=False):
     """
         Save to path. If the destination file exists, modify it in-place.
+
+        Raises OptionsError if the existing data is corrupt.
     """
     if os.path.exists(path) and os.path.isfile(path):
-        with open(path, "r") as f:
-            data = f.read()
+        with open(path, "rt", encoding="utf8") as f:
+            try:
+                data = f.read()
+            except UnicodeDecodeError as e:
+                raise exceptions.OptionsError(
+                    "Error trying to modify %s: %s" % (path, e)
+                )
     else:
         data = ""
     data = serialize(opts, data, defaults)
-    with open(path, "w") as f:
+    with open(path, "wt", encoding="utf8") as f:
         f.write(data)
diff --git a/mitmproxy/tools/console/options.py b/mitmproxy/tools/console/options.py
index 98606b9c..64203f2b 100644
--- a/mitmproxy/tools/console/options.py
+++ b/mitmproxy/tools/console/options.py
@@ -182,11 +182,16 @@ class OptionsList(urwid.ListBox):
         super().__init__(self.walker)
 
     def save_config(self, path):
-        optmanager.save(self.master.options, path)
+        try:
+            optmanager.save(self.master.options, path)
+        except exceptions.OptionsError as e:
+            signals.status_message.send(message=str(e))
 
     def load_config(self, path):
-        txt = open(path, "r").read()
-        optmanager.load(self.master.options, txt)
+        try:
+            optmanager.load_paths(self.master.options, path)
+        except exceptions.OptionsError as e:
+            signals.status_message.send(message=str(e))
 
     def keypress(self, size, key):
         if self.walker.editing:
diff --git a/test/mitmproxy/test_optmanager.py b/test/mitmproxy/test_optmanager.py
index 01636640..31b6e52b 100644
--- a/test/mitmproxy/test_optmanager.py
+++ b/test/mitmproxy/test_optmanager.py
@@ -296,6 +296,20 @@ def test_saving(tmpdir):
     with pytest.raises(exceptions.OptionsError):
         optmanager.load_paths(o, dst)
 
+    with open(dst, 'wb') as f:
+        f.write(b"\x01\x02\x03")
+    with pytest.raises(exceptions.OptionsError):
+        optmanager.load_paths(o, dst)
+    with pytest.raises(exceptions.OptionsError):
+        optmanager.save(o, dst)
+
+    with open(dst, 'wb') as f:
+        f.write(b"\xff\xff\xff")
+    with pytest.raises(exceptions.OptionsError):
+        optmanager.load_paths(o, dst)
+    with pytest.raises(exceptions.OptionsError):
+        optmanager.save(o, dst)
+
 
 def test_merge():
     m = TM()