diff options
| -rw-r--r-- | mitmproxy/command.py | 51 | ||||
| -rw-r--r-- | mitmproxy/tools/console/commander/commander.py | 11 | ||||
| -rw-r--r-- | test/mitmproxy/test_command.py | 24 |
3 files changed, 66 insertions, 20 deletions
diff --git a/mitmproxy/command.py b/mitmproxy/command.py index b909dfd5..7374a19a 100644 --- a/mitmproxy/command.py +++ b/mitmproxy/command.py @@ -15,6 +15,14 @@ from mitmproxy import exceptions from mitmproxy import flow +def lexer(s): + # mypy mis-identifies shlex.shlex as abstract + lex = shlex.shlex(s, punctuation_chars=True) # type: ignore + lex.whitespace_split = True + lex.commenters = '' + return lex + + Cuts = typing.Sequence[ typing.Sequence[typing.Union[str, bytes]] ] @@ -123,9 +131,10 @@ class Command: return ret -class ParseResult(typing.NamedTuple): - value: str - type: typing.Type +ParseResult = typing.NamedTuple( + "ParseResult", + [("value", str), ("type", typing.Type)], +) class CommandManager: @@ -147,27 +156,37 @@ class CommandManager: """ Parse a possibly partial command. Return a sequence of (part, type) tuples. """ - parts: typing.List[ParseResult] = [] buf = io.StringIO(cmdstr) - # mypy mis-identifies shlex.shlex as abstract - lex = shlex.shlex(buf) # type: ignore + parts: typing.List[str] = [] + lex = lexer(buf) while 1: remainder = cmdstr[buf.tell():] try: t = lex.get_token() except ValueError: - parts.append(ParseResult(value = remainder, type = str)) + parts.append(remainder) break if not t: break - typ: type = str - # First value is a special case: it has to be a command - if not parts: - typ = Cmd - parts.append(ParseResult(value = t, type = typ)) + parts.append(t) if not parts: - return [ParseResult(value = "", type = Cmd)] - return parts + parts = [""] + elif cmdstr.endswith(" "): + parts.append("") + + parse: typing.List[ParseResult] = [] + params: typing.List[type] = [] + for i in range(len(parts)): + if i == 0: + params[:] = [Cmd] + if parts[i] in self.commands: + params.extend(self.commands[parts[i]].paramtypes) + if params: + typ = params.pop(0) + else: + typ = str + parse.append(ParseResult(value=parts[i], type=typ)) + return parse def call_args(self, path, args): """ @@ -181,7 +200,7 @@ class CommandManager: """ Call a command using a string. May raise CommandError. """ - parts = shlex.split(cmdstr) + parts = list(lexer(cmdstr)) if not len(parts) >= 1: raise exceptions.CommandError("Invalid command: %s" % cmdstr) return self.call_args(parts[0], parts[1:]) @@ -208,8 +227,6 @@ 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/commander.py b/mitmproxy/tools/console/commander/commander.py index a0c3a3b2..f82ce9ce 100644 --- a/mitmproxy/tools/console/commander/commander.py +++ b/mitmproxy/tools/console/commander/commander.py @@ -72,7 +72,8 @@ class CommandBuffer(): def cycle_completion(self) -> None: if not self.completion: parts = self.master.commands.parse_partial(self.buf[:self.cursor]) - if parts[-1].type == mitmproxy.command.Cmd: + last = parts[-1] + if last.type == mitmproxy.command.Cmd: self.completion = CompletionState( completer = ListCompleter( parts[-1].value, @@ -80,6 +81,14 @@ class CommandBuffer(): ), parse = parts, ) + elif isinstance(last.type, mitmproxy.command.Choice): + self.completion = CompletionState( + completer = ListCompleter( + parts[-1].value, + self.master.commands.call(last.type.options_command), + ), + parse = parts, + ) if self.completion: nxt = self.completion.completer.cycle() buf = " ".join([i.value for i in self.completion.parse[:-1]]) + " " + nxt diff --git a/test/mitmproxy/test_command.py b/test/mitmproxy/test_command.py index b4711236..76ce2245 100644 --- a/test/mitmproxy/test_command.py +++ b/test/mitmproxy/test_command.py @@ -11,19 +11,24 @@ from mitmproxy.utils import typecheck class TAddon: + @command.command("cmd1") def cmd1(self, foo: str) -> str: """cmd1 help""" return "ret " + foo + @command.command("cmd2") def cmd2(self, foo: str) -> str: return 99 + @command.command("cmd3") def cmd3(self, foo: int) -> int: return foo + @command.command("empty") def empty(self) -> None: pass + @command.command("varargs") def varargs(self, one: str, *var: str) -> typing.Sequence[str]: return list(var) @@ -34,6 +39,7 @@ class TAddon: def choose(self, arg: str) -> typing.Sequence[str]: return ["one", "two", "three"] + @command.command("path") def path(self, arg: command.Path) -> None: pass @@ -82,11 +88,25 @@ class TestCommand: ], ["a", [command.ParseResult(value = "a", type = command.Cmd)]], ["", [command.ParseResult(value = "", type = command.Cmd)]], + [ + "cmd3 1", + [ + command.ParseResult(value = "cmd3", type = command.Cmd), + command.ParseResult(value = "1", type = int), + ] + ], + [ + "cmd3 ", + [ + command.ParseResult(value = "cmd3", type = command.Cmd), + command.ParseResult(value = "", type = int), + ] + ], ] with taddons.context() as tctx: - cm = command.CommandManager(tctx.master) + tctx.master.addons.add(TAddon()) for s, expected in tests: - assert cm.parse_partial(s) == expected + assert tctx.master.commands.parse_partial(s) == expected def test_simple(): |