console: console.key.bind console.key.unbind commands

7 files changed, 147 insertions, 19 deletions

diff --git a/mitmproxy/command.py b/mitmproxy/command.py
index 2256e4ca..c9776bc3 100644
--- a/mitmproxy/command.py
+++ b/mitmproxy/command.py
@@ -74,7 +74,8 @@ class Command:
 
     def call(self, args: typing.Sequence[str]):
         """
-            Call the command with a set of arguments. At this point, all argumets are strings.
+            Call the command with a list of arguments. At this point, all
+            arguments are strings.
         """
         if not self.has_positional and (len(self.paramtypes) != len(args)):
             raise exceptions.CommandError("Usage: %s" % self.signature_help())
diff --git a/mitmproxy/tools/console/consoleaddons.py b/mitmproxy/tools/console/consoleaddons.py
index fb6d5cb9..dc6e1a35 100644
--- a/mitmproxy/tools/console/consoleaddons.py
+++ b/mitmproxy/tools/console/consoleaddons.py
@@ -4,10 +4,12 @@ from mitmproxy import ctx
 from mitmproxy import command
 from mitmproxy import exceptions
 from mitmproxy import flow
-from mitmproxy.tools.console import overlay
 from mitmproxy import contentviews
 from mitmproxy.utils import strutils
+
+from mitmproxy.tools.console import overlay
 from mitmproxy.tools.console import signals
+from mitmproxy.tools.console import keymap
 
 
 class Logger:
@@ -415,6 +417,39 @@ class ConsoleAddon:
         """
         signals.sig_clear_log.send(self)
 
+    @command.command("console.key.contexts")
+    def key_contexts(self) -> typing.Sequence[str]:
+        """
+            The available contexts for key binding.
+        """
+        return list(sorted(keymap.Contexts))
+
+    @command.command("console.key.bind")
+    def key_bind(self, context: str, key: str, command: str) -> None:
+        """
+            Bind a shortcut key.
+        """
+        try:
+            self.master.keymap.add(
+                key,
+                command,
+                [context],
+                command
+            )
+        except ValueError as v:
+            raise exceptions.CommandError(v)
+        signals.keybindings_change.send(self)
+
+    @command.command("console.key.unbind")
+    def key_unbind(self, contexts: typing.Sequence[str], key: str) -> None:
+        """
+            Un-bind a shortcut key.
+        """
+        try:
+            self.master.keymap.remove(key, contexts)
+        except ValueError as v:
+            raise exceptions.CommandError(v)
+
     def running(self):
         self.started = True
 
@@ -422,4 +457,4 @@ class ConsoleAddon:
         if not flows:
             signals.update_settings.send(self)
         for f in flows:
-            signals.flow_change.send(self, flow=f)
\ No newline at end of file
+            signals.flow_change.send(self, flow=f)
diff --git a/mitmproxy/tools/console/defaultkeys.py b/mitmproxy/tools/console/defaultkeys.py
index 18298003..2d58218b 100644
--- a/mitmproxy/tools/console/defaultkeys.py
+++ b/mitmproxy/tools/console/defaultkeys.py
@@ -148,3 +148,13 @@ def map(km):
     km.add("e", "console.grideditor.editor", ["grideditor"], "Edit in external editor")
 
     km.add("z", "console.eventlog.clear", ["eventlog"], "Clear")
+
+    km.add(
+        "a",
+        """
+        console.choose.cmd "Context" console.key.contexts
+        console.command console.key.bind {choice}
+        """,
+        ["keybindings"],
+        "Add a key binding"
+    )
diff --git a/mitmproxy/tools/console/keybindings.py b/mitmproxy/tools/console/keybindings.py
index 6bd13429..c656c437 100644
--- a/mitmproxy/tools/console/keybindings.py
+++ b/mitmproxy/tools/console/keybindings.py
@@ -2,6 +2,7 @@ import urwid
 import blinker
 import textwrap
 from mitmproxy.tools.console import layoutwidget
+from mitmproxy.tools.console import signals
 
 HELP_HEIGHT = 5
 
@@ -43,6 +44,11 @@ class KeyListWalker(urwid.ListWalker):
         self.focusobj = None
         self.bindings = list(master.keymap.list("all"))
         self.set_focus(0)
+        signals.keybindings_change.connect(self.sig_modified)
+
+    def sig_modified(self, sender):
+        self.bindings = list(self.master.keymap.list("all"))
+        self._modified()
 
     def get_edit_text(self):
         return self.focus_obj.get_edit_text()
diff --git a/mitmproxy/tools/console/keymap.py b/mitmproxy/tools/console/keymap.py
index 06bcca9a..5cb3c579 100644
--- a/mitmproxy/tools/console/keymap.py
+++ b/mitmproxy/tools/console/keymap.py
@@ -2,7 +2,7 @@ import typing
 from mitmproxy.tools.console import commandeditor
 
 
-SupportedContexts = {
+Contexts = {
     "chooser",
     "commands",
     "eventlog",
@@ -11,6 +11,7 @@ SupportedContexts = {
     "global",
     "grideditor",
     "help",
+    "keybindings",
     "options",
 }
 
@@ -32,29 +33,68 @@ class Keymap:
     def __init__(self, master):
         self.executor = commandeditor.CommandExecutor(master)
         self.keys = {}
+        for c in Contexts:
+            self.keys[c] = {}
         self.bindings = []
 
-    def add(self, key: str, command: str, contexts: typing.Sequence[str], help="") -> None:
-        """
-            Add a key to the key map. If context is empty, it's considered to be
-            a global binding.
-        """
+    def _check_contexts(self, contexts):
         if not contexts:
             raise ValueError("Must specify at least one context.")
         for c in contexts:
-            if c not in SupportedContexts:
+            if c not in Contexts:
                 raise ValueError("Unsupported context: %s" % c)
 
+    def add(
+        self,
+        key: str,
+        command: str,
+        contexts: typing.Sequence[str],
+        help=""
+    ) -> None:
+        """
+            Add a key to the key map.
+        """
+        self._check_contexts(contexts)
+        self.remove(key, contexts)
+
+        for b in self.bindings:
+            if b.key == key and b.command == command:
+                b.contexts = list(set(b.contexts + contexts))
+                if help:
+                    b.help = help
+                self.bind(b)
+                return
+
         b = Binding(key=key, command=command, contexts=contexts, help=help)
         self.bindings.append(b)
         self.bind(b)
 
-    def bind(self, binding):
+    def remove(self, key: str, contexts: typing.Sequence[str]) -> None:
+        """
+            Remove a key from the key map.
+        """
+        self._check_contexts(contexts)
+        for c in contexts:
+            b = self.get(c, key)
+            if b:
+                self.unbind(b)
+                b.contexts = [x for x in b.contexts if x != c]
+                if b.contexts:
+                    self.bind(b)
+
+    def bind(self, binding: Binding) -> None:
+        for c in binding.contexts:
+            self.keys[c][binding.keyspec()] = binding
+
+    def unbind(self, binding: Binding) -> None:
+        """
+            Unbind also removes the binding from the list.
+        """
         for c in binding.contexts:
-            d = self.keys.setdefault(c, {})
-            d[binding.keyspec()] = binding.command
+            del self.keys[c][binding.keyspec()]
+            self.bindings = [b for b in self.bindings if b != binding]
 
-    def get(self, context: str, key: str) -> typing.Optional[str]:
+    def get(self, context: str, key: str) -> typing.Optional[Binding]:
         if context in self.keys:
             return self.keys[context].get(key, None)
         return None
@@ -71,9 +111,7 @@ class Keymap:
         """
             Returns the key if it has not been handled, or None.
         """
-        cmd = self.get(context, key)
-        if not cmd:
-            cmd = self.get("global", key)
-        if cmd:
-            return self.executor(cmd)
+        b = self.get(context, key) or self.get("global", key)
+        if b:
+            return self.executor(b.command)
         return key
diff --git a/mitmproxy/tools/console/signals.py b/mitmproxy/tools/console/signals.py
index 49115a5d..5d39d96a 100644
--- a/mitmproxy/tools/console/signals.py
+++ b/mitmproxy/tools/console/signals.py
@@ -48,3 +48,6 @@ flowlist_change = blinker.Signal()
 # Pop and push view state onto a stack
 pop_view_state = blinker.Signal()
 push_view_state = blinker.Signal()
+
+# Fired when the key bindings change
+keybindings_change = blinker.Signal()
diff --git a/test/mitmproxy/tools/console/test_keymap.py b/test/mitmproxy/tools/console/test_keymap.py
index fdb2b028..00e64991 100644
--- a/test/mitmproxy/tools/console/test_keymap.py
+++ b/test/mitmproxy/tools/console/test_keymap.py
@@ -35,3 +35,38 @@ def test_bind():
         assert km.executor.called
 
         assert len((km.list("global"))) == 1
+
+
+def test_join():
+    with taddons.context() as tctx:
+        km = keymap.Keymap(tctx.master)
+        km.add("key", "str", ["options"], "help1")
+        km.add("key", "str", ["commands"])
+        return
+        assert len(km.bindings) == 1
+        assert len(km.bindings[0].contexts) == 2
+        assert km.bindings[0].help == "help1"
+        km.add("key", "str", ["commands"], "help2")
+        assert len(km.bindings) == 1
+        assert len(km.bindings[0].contexts) == 2
+        assert km.bindings[0].help == "help2"
+
+        assert km.get("commands", "key")
+        km.unbind(km.bindings[0])
+        assert len(km.bindings) == 0
+        assert not km.get("commands", "key")
+
+
+def test_remove():
+    with taddons.context() as tctx:
+        km = keymap.Keymap(tctx.master)
+        km.add("key", "str", ["options", "commands"], "help1")
+        assert len(km.bindings) == 1
+        assert "options" in km.bindings[0].contexts
+
+        km.remove("key", ["options"])
+        assert len(km.bindings) == 1
+        assert "options" not in km.bindings[0].contexts
+
+        km.remove("key", ["commands"])
+        assert len(km.bindings) == 0