diff --git a/mitmproxy/command.py b/mitmproxy/command.py index f73eeb68e..aab721b53 100644 --- a/mitmproxy/command.py +++ b/mitmproxy/command.py @@ -123,6 +123,11 @@ class Command: return ret +class ParseResult(typing.NamedTuple): + value: str + type: type + + class CommandManager: def __init__(self, master): self.master = master @@ -138,11 +143,11 @@ class CommandManager: def add(self, path: str, func: typing.Callable): self.commands[path] = Command(self, path, func) - def parse_partial(self, cmdstr: str) -> typing.Sequence[typing.Tuple[str, type]]: + def parse_partial(self, cmdstr: str) -> typing.Sequence[ParseResult]: """ Parse a possibly partial command. Return a sequence of (part, type) tuples. """ - parts: typing.List[typing.Tuple[str, type]] = [] + parts: typing.List[ParseResult] = [] buf = io.StringIO(cmdstr) # mypy mis-identifies shlex.shlex as abstract lex = shlex.shlex(buf) # type: ignore @@ -151,7 +156,7 @@ class CommandManager: try: t = lex.get_token() except ValueError: - parts.append((remainder, str)) + parts.append(ParseResult(value = remainder, type = str)) break if not t: break @@ -159,7 +164,9 @@ class CommandManager: # First value is a special case: it has to be a command if not parts: typ = Cmd - parts.append((t, typ)) + parts.append(ParseResult(value = t, type = typ)) + if not parts: + return [ParseResult(value = "", type = Cmd)] return parts def call_args(self, path, args): diff --git a/mitmproxy/tools/console/commander/commander.py b/mitmproxy/tools/console/commander/commander.py index d07910a6e..d961d421c 100644 --- a/mitmproxy/tools/console/commander/commander.py +++ b/mitmproxy/tools/console/commander/commander.py @@ -1,9 +1,19 @@ import urwid from urwid.text_layout import calc_coords +import typing + +import mitmproxy.master +import mitmproxy.command + + +class CompletionState: + def __init__(self, parts: typing.Sequence[mitmproxy.command.ParseResult]) -> None: + self.parts = parts class CommandBuffer(): - def __init__(self, start: str = "") -> None: + def __init__(self, master: mitmproxy.master.Master, start: str = "") -> None: + self.master = master 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)]. @@ -31,6 +41,12 @@ class CommandBuffer(): def right(self) -> None: self.cursor = self.cursor + 1 + def cycle_completion(self) -> None: + parts = self.master.commands.parse_partial(self.buf[:self.cursor]) + if parts[-1][1] == str: + return + raise ValueError + def backspace(self) -> None: if self.cursor == 0: return @@ -48,8 +64,9 @@ class CommandBuffer(): class CommandEdit(urwid.WidgetWrap): leader = ": " - def __init__(self, text) -> None: - self.cbuf = CommandBuffer(text) + def __init__(self, master: mitmproxy.master.Master, text: str) -> None: + self.master = master + self.cbuf = CommandBuffer(master, text) self._w = urwid.Text(self.leader) self.update() @@ -60,6 +77,8 @@ class CommandEdit(urwid.WidgetWrap): self.cbuf.left() elif key == "right": self.cbuf.right() + elif key == "tab": + self.cbuf.cycle_completion() elif len(key) == 1: self.cbuf.insert(key) self.update() diff --git a/mitmproxy/tools/console/statusbar.py b/mitmproxy/tools/console/statusbar.py index a59fc92e1..6a1f07a93 100644 --- a/mitmproxy/tools/console/statusbar.py +++ b/mitmproxy/tools/console/statusbar.py @@ -67,7 +67,7 @@ class ActionBar(urwid.WidgetWrap): def sig_prompt_command(self, sender, partial=""): signals.focus.send(self, section="footer") - self._w = commander.CommandEdit(partial) + self._w = commander.CommandEdit(self.master, partial) self.prompting = commandeditor.CommandExecutor(self.master) def sig_prompt_onekey(self, sender, prompt, keys, callback, args=()): diff --git a/test/mitmproxy/test_command.py b/test/mitmproxy/test_command.py index 5218042c9..b47112363 100644 --- a/test/mitmproxy/test_command.py +++ b/test/mitmproxy/test_command.py @@ -66,8 +66,22 @@ class TestCommand: def test_parse_partial(self): tests = [ - ["foo bar", [("foo", command.Cmd), ("bar", str)]], - ["foo 'bar", [("foo", command.Cmd), ("'bar", str)]], + [ + "foo bar", + [ + command.ParseResult(value = "foo", type = command.Cmd), + command.ParseResult(value = "bar", type = str) + ], + ], + [ + "foo 'bar", + [ + command.ParseResult(value = "foo", type = command.Cmd), + command.ParseResult(value = "'bar", type = str) + ] + ], + ["a", [command.ParseResult(value = "a", type = command.Cmd)]], + ["", [command.ParseResult(value = "", type = command.Cmd)]], ] with taddons.context() as tctx: cm = command.CommandManager(tctx.master) diff --git a/test/mitmproxy/tools/console/test_commander.py b/test/mitmproxy/tools/console/test_commander.py index fdf54897e..9ef4a318b 100644 --- a/test/mitmproxy/tools/console/test_commander.py +++ b/test/mitmproxy/tools/console/test_commander.py @@ -1,5 +1,5 @@ - from mitmproxy.tools.console.commander import commander +from mitmproxy.test import taddons class TestCommandBuffer: @@ -13,12 +13,13 @@ class TestCommandBuffer: [("123", 2), ("13", 1)], [("123", 0), ("123", 0)], ] - for start, output in tests: - cb = commander.CommandBuffer() - cb.buf, cb.cursor = start[0], start[1] - cb.backspace() - assert cb.buf == output[0] - assert cb.cursor == output[1] + with taddons.context() as tctx: + for start, output in tests: + cb = commander.CommandBuffer(tctx.master) + cb.buf, cb.cursor = start[0], start[1] + cb.backspace() + assert cb.buf == output[0] + assert cb.cursor == output[1] def test_insert(self): tests = [ @@ -26,9 +27,17 @@ class TestCommandBuffer: [("a", 0), ("xa", 1)], [("xa", 2), ("xax", 3)], ] - for start, output in tests: - cb = commander.CommandBuffer() - cb.buf, cb.cursor = start[0], start[1] - cb.insert("x") - assert cb.buf == output[0] - assert cb.cursor == output[1] + with taddons.context() as tctx: + for start, output in tests: + cb = commander.CommandBuffer(tctx.master) + cb.buf, cb.cursor = start[0], start[1] + cb.insert("x") + assert cb.buf == output[0] + assert cb.cursor == output[1] + + def test_cycle_completion(self): + with taddons.context() as tctx: + cb = commander.CommandBuffer(tctx.master) + cb.buf = "foo bar" + cb.cursor = len(cb.buf) + cb.cycle_completion()