diff --git a/mitmproxy/addons/clientplayback.py b/mitmproxy/addons/clientplayback.py index 9e012b674..fcc3209b9 100644 --- a/mitmproxy/addons/clientplayback.py +++ b/mitmproxy/addons/clientplayback.py @@ -37,7 +37,7 @@ class ClientPlayback: ctx.master.addons.trigger("update", []) @command.command("replay.client.file") - def load_file(self, path: str) -> None: + def load_file(self, path: command.Path) -> None: try: flows = io.read_flows_from_paths([path]) except exceptions.FlowReadException as e: diff --git a/mitmproxy/addons/serverplayback.py b/mitmproxy/addons/serverplayback.py index 927f6e15c..46968a8d6 100644 --- a/mitmproxy/addons/serverplayback.py +++ b/mitmproxy/addons/serverplayback.py @@ -31,7 +31,7 @@ class ServerPlayback: ctx.master.addons.trigger("update", []) @command.command("replay.server.file") - def load_file(self, path: str) -> None: + def load_file(self, path: command.Path) -> None: try: flows = io.read_flows_from_paths([path]) except exceptions.FlowReadException as e: diff --git a/mitmproxy/command.py b/mitmproxy/command.py index c48219730..a638b7ee1 100644 --- a/mitmproxy/command.py +++ b/mitmproxy/command.py @@ -23,6 +23,10 @@ class Path(str): pass +class Cmd(str): + pass + + def typename(t: type, ret: bool) -> str: """ Translates a type to an explanatory string. If ret is True, we're @@ -172,6 +176,8 @@ def parsearg(manager: CommandManager, spec: str, argtype: type) -> typing.Any: "Invalid choice: see %s for options" % cmd ) return spec + if argtype in (Path, Cmd): + return spec elif issubclass(argtype, str): return spec elif argtype == bool: diff --git a/mitmproxy/tools/console/commander/__init__.py b/mitmproxy/tools/console/commander/__init__.py index 8b1378917..e69de29bb 100644 --- a/mitmproxy/tools/console/commander/__init__.py +++ b/mitmproxy/tools/console/commander/__init__.py @@ -1 +0,0 @@ - diff --git a/mitmproxy/tools/console/commander/commander.py b/mitmproxy/tools/console/commander/commander.py index 74855e4de..d07910a6e 100644 --- a/mitmproxy/tools/console/commander/commander.py +++ b/mitmproxy/tools/console/commander/commander.py @@ -3,18 +3,18 @@ from urwid.text_layout import calc_coords class CommandBuffer(): - def __init__(self, start: str = ""): + def __init__(self, start: str = "") -> None: self.buf = start # This is the logical cursor position - the display cursor is one # character further on. Cursor is always within the range [0:len(buffer)]. self._cursor = len(self.buf) @property - def cursor(self): + def cursor(self) -> int: return self._cursor @cursor.setter - def cursor(self, x): + def cursor(self, x) -> None: if x < 0: self._cursor = 0 elif x > len(self.buf): @@ -25,19 +25,19 @@ class CommandBuffer(): def render(self): return self.buf - def left(self): + def left(self) -> None: self.cursor = self.cursor - 1 - def right(self): + def right(self) -> None: self.cursor = self.cursor + 1 - def backspace(self): + def backspace(self) -> None: if self.cursor == 0: return self.buf = self.buf[:self.cursor - 1] + self.buf[self.cursor:] self.cursor = self.cursor - 1 - def insert(self, k: str): + def insert(self, k: str) -> None: """ Inserts text at the cursor. """ @@ -48,7 +48,7 @@ class CommandBuffer(): class CommandEdit(urwid.WidgetWrap): leader = ": " - def __init__(self, text): + def __init__(self, text) -> None: self.cbuf = CommandBuffer(text) self._w = urwid.Text(self.leader) self.update() @@ -82,4 +82,3 @@ class CommandEdit(urwid.WidgetWrap): def get_value(self): return self.cbuf.buf - diff --git a/mitmproxy/tools/console/consoleaddons.py b/mitmproxy/tools/console/consoleaddons.py index 06ee33418..d79e39472 100644 --- a/mitmproxy/tools/console/consoleaddons.py +++ b/mitmproxy/tools/console/consoleaddons.py @@ -224,7 +224,11 @@ class ConsoleAddon: @command.command("console.choose") def console_choose( - self, prompt: str, choices: typing.Sequence[str], *cmd: str + self, + prompt: str, + choices: typing.Sequence[str], + cmd: command.Cmd, + *args: str, ) -> None: """ Prompt the user to choose from a specified list of strings, then @@ -233,7 +237,7 @@ class ConsoleAddon: """ def callback(opt): # We're now outside of the call context... - repl = " ".join(cmd) + repl = cmd + " " + " ".join(args) repl = repl.replace("{choice}", opt) try: self.master.commands.call(repl) @@ -246,7 +250,7 @@ class ConsoleAddon: @command.command("console.choose.cmd") def console_choose_cmd( - self, prompt: str, choicecmd: str, *cmd: str + self, prompt: str, choicecmd: command.Cmd, *cmd: str ) -> None: """ Prompt the user to choose from a list of strings returned by a @@ -492,14 +496,20 @@ class ConsoleAddon: return list(sorted(keymap.Contexts)) @command.command("console.key.bind") - def key_bind(self, contexts: typing.Sequence[str], key: str, *command: str) -> None: + def key_bind( + self, + contexts: typing.Sequence[str], + key: str, + cmd: command.Cmd, + *args: str, + ) -> None: """ Bind a shortcut key. """ try: self.master.keymap.add( key, - " ".join(command), + cmd + " " + " ".join(args), contexts, "" ) diff --git a/test/mitmproxy/test_command.py b/test/mitmproxy/test_command.py index e1879ba28..c8007463a 100644 --- a/test/mitmproxy/test_command.py +++ b/test/mitmproxy/test_command.py @@ -100,6 +100,7 @@ def test_typename(): assert command.typename(command.Choice("foo"), False) == "choice" assert command.typename(command.Path, False) == "path" + assert command.typename(command.Cmd, False) == "cmd" class DummyConsole: @@ -162,6 +163,9 @@ def test_parsearg(): assert command.parsearg( tctx.master.commands, "foo", command.Path ) == "foo" + assert command.parsearg( + tctx.master.commands, "foo", command.Cmd + ) == "foo" class TDec: diff --git a/test/mitmproxy/tools/console/test_commander.py b/test/mitmproxy/tools/console/test_commander.py index b1f23df42..fdf54897e 100644 --- a/test/mitmproxy/tools/console/test_commander.py +++ b/test/mitmproxy/tools/console/test_commander.py @@ -32,6 +32,3 @@ class TestCommandBuffer: cb.insert("x") assert cb.buf == output[0] assert cb.cursor == output[1] - - -