diff --git a/mitmproxy/command.py b/mitmproxy/command.py index 2256e4ca5..c9776bc32 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 fb6d5cb96..dc6e1a35b 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 182980037..2d58218b1 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 6bd134295..c656c4371 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 06bcca9a7..5cb3c5793 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): - for c in binding.contexts: - d = self.keys.setdefault(c, {}) - d[binding.keyspec()] = binding.command + 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 get(self, context: str, key: str) -> typing.Optional[str]: + 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: + 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[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 49115a5d8..5d39d96a5 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 fdb2b0289..00e649917 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