diff options
| -rw-r--r-- | mitmproxy/command.py | 2 | ||||
| -rw-r--r-- | mitmproxy/tools/console/flowlist.py | 4 | ||||
| -rw-r--r-- | mitmproxy/tools/console/keymap.py | 34 | ||||
| -rw-r--r-- | mitmproxy/tools/console/master.py | 61 | ||||
| -rw-r--r-- | mitmproxy/tools/console/window.py | 23 | ||||
| -rw-r--r-- | test/mitmproxy/test_command.py | 6 |
6 files changed, 104 insertions, 26 deletions
diff --git a/mitmproxy/command.py b/mitmproxy/command.py index 665e14cf..1c943cef 100644 --- a/mitmproxy/command.py +++ b/mitmproxy/command.py @@ -37,7 +37,7 @@ class Command: def paramnames(self) -> typing.Sequence[str]: return [typename(i, False) for i in self.paramtypes] - def retname(self) -> typing.Sequence[str]: + def retname(self) -> str: return typename(self.returntype, True) if self.returntype else "" def signature_help(self) -> str: diff --git a/mitmproxy/tools/console/flowlist.py b/mitmproxy/tools/console/flowlist.py index 00e5cf4e..e7570824 100644 --- a/mitmproxy/tools/console/flowlist.py +++ b/mitmproxy/tools/console/flowlist.py @@ -353,9 +353,7 @@ class FlowListBox(urwid.ListBox): def keypress(self, size, key): key = common.shortcuts(key) - if key == ":": - signals.status_prompt_command.send() - elif key == "A": + if key == "A": for f in self.master.view: if f.intercepted: f.resume() diff --git a/mitmproxy/tools/console/keymap.py b/mitmproxy/tools/console/keymap.py new file mode 100644 index 00000000..018d1bde --- /dev/null +++ b/mitmproxy/tools/console/keymap.py @@ -0,0 +1,34 @@ +import typing +from mitmproxy.tools.console import commandeditor + + +class Keymap: + def __init__(self, master): + self.executor = commandeditor.CommandExecutor(master) + self.keys = {} + + def add(self, key: str, command: str, context: str = "") -> None: + """ + Add a key to the key map. If context is empty, it's considered to be + a global binding. + """ + d = self.keys.setdefault(context, {}) + d[key] = command + + def get(self, context: str, key: str) -> typing.Optional[str]: + if context in self.keys: + return self.keys[context].get(key, None) + return None + + def handle(self, context: str, key: str) -> typing.Optional[str]: + """ + Returns the key if it has not been handled, or None. + """ + cmd = self.get(context, key) + if cmd: + return self.executor(cmd) + if cmd != "": + cmd = self.get("", key) + if cmd: + return self.executor(cmd) + return key diff --git a/mitmproxy/tools/console/master.py b/mitmproxy/tools/console/master.py index 47f021c2..d65562de 100644 --- a/mitmproxy/tools/console/master.py +++ b/mitmproxy/tools/console/master.py @@ -24,6 +24,7 @@ from mitmproxy.tools.console import flowlist from mitmproxy.tools.console import flowview from mitmproxy.tools.console import grideditor from mitmproxy.tools.console import help +from mitmproxy.tools.console import keymap from mitmproxy.tools.console import options from mitmproxy.tools.console import commands from mitmproxy.tools.console import overlay @@ -75,6 +76,58 @@ class UnsupportedLog: signals.add_log(strutils.bytes_to_escaped_str(message.content), "debug") +class ConsoleCommands: + """ + An addon that exposes console-specific commands. + """ + def __init__(self, master): + self.master = master + + def command(self) -> None: + """Prompt for a command.""" + signals.status_prompt_command.send() + + def view_commands(self) -> None: + """View the commands list.""" + self.master.view_commands() + + def view_options(self) -> None: + """View the options editor.""" + self.master.view_options() + + def view_help(self) -> None: + """View help.""" + self.master.view_help() + + def exit(self) -> None: + """Exit mitmproxy.""" + raise urwid.ExitMainLoop + + def view_pop(self) -> None: + """ + Pop a view off the console stack. At the top level, this prompts the + user to exit mitmproxy. + """ + signals.pop_view_state.send(self) + + def load(self, l): + l.add_command("console.command", self.command) + l.add_command("console.exit", self.exit) + l.add_command("console.view.commands", self.view_commands) + l.add_command("console.view.help", self.view_help) + l.add_command("console.view.options", self.view_options) + l.add_command("console.view.pop", self.view_pop) + + +def default_keymap(km): + km.add(":", "console.command") + km.add("?", "console.view.help") + km.add("C", "console.view.commands") + km.add("O", "console.view.options") + km.add("Q", "console.exit") + km.add("q", "console.view.pop") + + class ConsoleMaster(master.Master): def __init__(self, options, server): @@ -84,6 +137,8 @@ class ConsoleMaster(master.Master): self.stream_path = None # This line is just for type hinting self.options = self.options # type: Options + self.keymap = keymap.Keymap(self) + default_keymap(self.keymap) self.options.errored.connect(self.options_error) self.logbuffer = urwid.SimpleListWalker([]) @@ -102,6 +157,7 @@ class ConsoleMaster(master.Master): self.view, UnsupportedLog(), readfile.ReadFile(), + ConsoleCommands(self), ) def sigint_handler(*args, **kwargs): @@ -331,12 +387,13 @@ class ConsoleMaster(master.Master): ) ) - def view_help(self, helpctx): + def view_help(self): + hc = self.view_stack[0].helpctx signals.push_view_state.send( self, window = window.Window( self, - help.HelpView(helpctx), + help.HelpView(hc), None, statusbar.StatusBar(self, help.footer), None diff --git a/mitmproxy/tools/console/window.py b/mitmproxy/tools/console/window.py index 4555d564..0e41fb2a 100644 --- a/mitmproxy/tools/console/window.py +++ b/mitmproxy/tools/console/window.py @@ -82,8 +82,9 @@ class Window(urwid.Frame): def keypress(self, size, k): k = super().keypress(size, k) - if k == "?": - self.master.view_help(self.helpctx) + k = self.master.keymap.handle("", k) + if not k: + return elif k == "i": signals.status_prompt.send( self, @@ -91,23 +92,5 @@ class Window(urwid.Frame): text = self.master.options.intercept, callback = self.master.options.setter("intercept") ) - elif k == "C": - self.master.view_commands() - elif k == "O": - self.master.view_options() - elif k == "Q": - raise urwid.ExitMainLoop - elif k == "q": - signals.pop_view_state.send(self) - elif k == "R": - signals.status_prompt_onekey.send( - self, - prompt = "Replay", - keys = ( - ("client", "c"), - ("server", "s"), - ), - callback = self.handle_replay, - ) else: return k diff --git a/test/mitmproxy/test_command.py b/test/mitmproxy/test_command.py index b984bea6..0c272a2c 100644 --- a/test/mitmproxy/test_command.py +++ b/test/mitmproxy/test_command.py @@ -18,6 +18,9 @@ class TAddon: def cmd2(self, foo: str) -> str: return 99 + def empty(self) -> None: + pass + class TestCommand: def test_call(self): @@ -50,6 +53,9 @@ def test_simple(): with pytest.raises(exceptions.CommandError, match="Usage"): c.call("one.two too many args") + c.add("empty", a.empty) + c.call("empty") + def test_typename(): assert command.typename(str, True) == "str" |