diff options
| -rw-r--r-- | mitmproxy/tools/console/keymap.py | 91 | ||||
| -rw-r--r-- | test/mitmproxy/tools/console/test_keymap.py | 94 |
2 files changed, 185 insertions, 0 deletions
diff --git a/mitmproxy/tools/console/keymap.py b/mitmproxy/tools/console/keymap.py index fbb569a4..5f7cdb11 100644 --- a/mitmproxy/tools/console/keymap.py +++ b/mitmproxy/tools/console/keymap.py @@ -1,6 +1,18 @@ import typing +import os + +import ruamel.yaml + +from mitmproxy import command from mitmproxy.tools.console import commandexecutor from mitmproxy.tools.console import signals +from mitmproxy import ctx +from mitmproxy import exceptions +import mitmproxy.types + + +class KeyBindingError(Exception): + pass Contexts = { @@ -139,3 +151,82 @@ class Keymap: if b: return self.executor(b.command) return key + + +keyAttrs = { + "key": lambda x: isinstance(x, str), + "cmd": lambda x: isinstance(x, str), + "ctx": lambda x: isinstance(x, list) and [isinstance(v, str) for v in x], + "help": lambda x: isinstance(x, str), +} +requiredKeyAttrs = set(["key", "cmd"]) + + +class KeymapConfig: + @command.command("console.keymap.load") + def keymap_load_path(self, path: mitmproxy.types.Path) -> None: + try: + master = ctx.master # type: mitmproxy.tools.console.master.ConsoleMaster + self.load_path(master.keymap, path) + except (OSError, KeyBindingError) as e: + raise exceptions.CommandError( + "Could not load key bindings - %s" % e + ) from e + + def load_path(self, km, p): + if os.path.exists(p) and os.path.isfile(p): + with open(p, "rt", encoding="utf8") as f: + try: + txt = f.read() + except UnicodeDecodeError as e: + raise KeyBindingError( + "Encoding error - expected UTF8: %s: %s" % (p, e) + ) + try: + vals = self.parse(txt) + except KeyBindingError as e: + raise KeyBindingError( + "Error reading %s: %s" % (p, e) + ) from e + for v in vals: + try: + km.add( + key = v["key"], + command = v["cmd"], + contexts = v.get("ctx", None), + help = v.get("help", None), + ) + except ValueError as e: + raise KeyBindingError( + "Error reading %s: %s" % (p, e) + ) from e + + def parse(self, text): + try: + data = ruamel.yaml.safe_load(text) + except ruamel.yaml.error.YAMLError as v: + if hasattr(v, "problem_mark"): + snip = v.problem_mark.get_snippet() + raise KeyBindingError( + "Key binding config error at line %s:\n%s\n%s" % + (v.problem_mark.line + 1, snip, v.problem) + ) + else: + raise KeyBindingError("Could not parse key bindings.") + if not data: + return [] + if not isinstance(data, list): + raise KeyBindingError("Inalid keybinding config - expected a list of keys") + + for k in data: + unknown = k.keys() - keyAttrs.keys() + if unknown: + raise KeyBindingError("Unknown key attributes: %s" % unknown) + missing = requiredKeyAttrs - k.keys() + if missing: + raise KeyBindingError("Missing required key attributes: %s" % unknown) + for attr in k.keys(): + if not keyAttrs[attr](k[attr]): + raise KeyBindingError("Invalid type for %s" % attr) + + return data diff --git a/test/mitmproxy/tools/console/test_keymap.py b/test/mitmproxy/tools/console/test_keymap.py index 7b475ff8..901cd795 100644 --- a/test/mitmproxy/tools/console/test_keymap.py +++ b/test/mitmproxy/tools/console/test_keymap.py @@ -70,3 +70,97 @@ def test_remove(): km.remove("key", ["commands"]) assert len(km.bindings) == 0 + + +def test_load_path(tmpdir): + dst = str(tmpdir.join("conf")) + + kmc = keymap.KeymapConfig() + with taddons.context(kmc) as tctx: + km = keymap.Keymap(tctx.master) + + with open(dst, 'wb') as f: + f.write(b"\xff\xff\xff") + with pytest.raises(keymap.KeyBindingError, match="expected UTF8"): + kmc.load_path(km, dst) + + with open(dst, 'w') as f: + f.write("'''") + with pytest.raises(keymap.KeyBindingError): + kmc.load_path(km, dst) + + with open(dst, 'w') as f: + f.write( + """ + - key: key1 + ctx: [unknown] + cmd: > + foo bar + foo bar + """ + ) + with pytest.raises(keymap.KeyBindingError): + kmc.load_path(km, dst) + + with open(dst, 'w') as f: + f.write( + """ + - key: key1 + ctx: [chooser] + help: one + cmd: > + foo bar + foo bar + """ + ) + kmc.load_path(km, dst) + assert(km.get("chooser", "key1")) + + +def test_parse(): + kmc = keymap.KeymapConfig() + with taddons.context(kmc): + assert kmc.parse("") == [] + assert kmc.parse("\n\n\n \n") == [] + with pytest.raises(keymap.KeyBindingError, match="expected a list of keys"): + kmc.parse("key: val") + with pytest.raises(keymap.KeyBindingError, match="expected a list of keys"): + kmc.parse("val") + with pytest.raises(keymap.KeyBindingError, match="Unknown key attributes"): + kmc.parse( + """ + - key: key1 + nonexistent: bar + """ + ) + with pytest.raises(keymap.KeyBindingError, match="Missing required key attributes"): + kmc.parse( + """ + - help: key1 + """ + ) + with pytest.raises(keymap.KeyBindingError, match="Invalid type for cmd"): + kmc.parse( + """ + - key: key1 + cmd: [ cmd ] + """ + ) + with pytest.raises(keymap.KeyBindingError, match="Invalid type for ctx"): + kmc.parse( + """ + - key: key1 + ctx: foo + cmd: cmd + """ + ) + assert kmc.parse( + """ + - key: key1 + ctx: [one, two] + help: one + cmd: > + foo bar + foo bar + """ + ) == [{"key": "key1", "ctx": ["one", "two"], "help": "one", "cmd": "foo bar foo bar\n"}]
\ No newline at end of file |