Merge pull request #2204 from cortesi/consoleopts

console options: load and save options to and from file

3 files changed, 64 insertions, 11 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 56d22715..64203f2b 100644
--- a/mitmproxy/tools/console/options.py
+++ b/mitmproxy/tools/console/options.py
@@ -5,6 +5,7 @@ import pprint
 from typing import Optional, Sequence
 
 from mitmproxy import exceptions
+from mitmproxy import optmanager
 from mitmproxy.tools.console import common
 from mitmproxy.tools.console import signals
 from mitmproxy.tools.console import overlay
@@ -31,7 +32,8 @@ def _mkhelp():
         ("enter", "edit option"),
         ("D", "reset all to defaults"),
         ("d", "reset this option to default"),
-        ("w", "save options"),
+        ("l", "load options from file"),
+        ("w", "save options to file"),
     ]
     text.extend(common.format_keyvals(keys, key="key", val="text", indent=4))
     return text
@@ -179,6 +181,18 @@ class OptionsList(urwid.ListBox):
         self.walker = OptionListWalker(master)
         super().__init__(self.walker)
 
+    def save_config(self, 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):
+        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:
             if key == "enter":
@@ -206,6 +220,16 @@ class OptionsList(urwid.ListBox):
             elif key == "G":
                 self.set_focus(len(self.walker.opts) - 1)
                 self.walker._modified()
+            elif key == "l":
+                signals.status_prompt_path.send(
+                    prompt = "Load config from",
+                    callback = self.load_config
+                )
+            elif key == "w":
+                signals.status_prompt_path.send(
+                    prompt = "Save config to",
+                    callback = self.save_config
+                )
             elif key == "enter":
                 foc, idx = self.get_focus()
                 if foc.opt.typespec == bool:
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()