diff --git a/mitmproxy/language/lexer.py b/mitmproxy/language/lexer.py index 44ccb0e29..67f77da45 100644 --- a/mitmproxy/language/lexer.py +++ b/mitmproxy/language/lexer.py @@ -1,3 +1,4 @@ +import re import typing import ply.lex as lex @@ -6,9 +7,10 @@ import ply.lex as lex class CommandLanguageLexer: tokens = ( "WHITESPACE", - "COMMAND", - "PLAIN_STR", - "QUOTED_STR" + "LPAREN", "RPAREN", + "LBRACE", "RBRACE", + "PLAIN_STR", "QUOTED_STR", + "COMMAND" ) states = ( ("interactive", "inclusive"), @@ -19,6 +21,13 @@ class CommandLanguageLexer: # Main(INITIAL) state t_ignore_WHITESPACE = r"\s+" + t_LPAREN = r"\(" + t_RPAREN = r"\)" + t_LBRACE = r"\[" + t_RBRACE = r"\]" + + special_symbols = re.escape("()[]") + plain_str = rf"[^{special_symbols}\s]+" def t_COMMAND(self, t): r"""\w+(\.\w+)+""" @@ -31,8 +40,8 @@ class CommandLanguageLexer: """ return t + @lex.TOKEN(plain_str) def t_PLAIN_STR(self, t): - r"""[^\s]+""" t.type = self.oneword_commands.get(t.value, "PLAIN_STR") return t diff --git a/mitmproxy/language/parser.py b/mitmproxy/language/parser.py index 34a81aeb2..14aa33c6f 100644 --- a/mitmproxy/language/parser.py +++ b/mitmproxy/language/parser.py @@ -16,16 +16,22 @@ class CommandLanguageParser: self.command_manager = command_manager def p_command_line(self, p): - """command_line : command_call""" + """command_line : command_call_no_parentheses + | command_call_with_parentheses""" p[0] = p[1] - - def p_command_call(self, p): - """command_call : COMMAND argument_list""" - p[0] = self.command_manager.call_strings(p[1], p[2]) self.return_value = p[0] + def p_command_call_no_parentheses(self, p): + """command_call_no_parentheses : COMMAND argument_list""" + p[0] = self.command_manager.call_strings(p[1], p[2]) + + def p_command_call_with_parentheses(self, p): + """command_call_with_parentheses : COMMAND LPAREN argument_list RPAREN""" + p[0] = self.command_manager.call_strings(p[1], p[3]) + def p_argument_list(self, p): """argument_list : empty + | argument | argument_list argument""" if len(p) == 2: p[0] = [] if p[1] is None else [p[1]] @@ -36,9 +42,15 @@ class CommandLanguageParser: def p_argument(self, p): """argument : PLAIN_STR | quoted_str - | COMMAND""" + | array + | COMMAND + | command_call_with_parentheses""" p[0] = p[1] + def p_array(self, p): + """array : LBRACE argument_list RBRACE""" + p[0] = ",".join(p[2]) if p[2] else "" + def p_quoted_str(self, p): """quoted_str : QUOTED_STR""" p[0] = p[1].strip("'\"") diff --git a/test/mitmproxy/test_command.py b/test/mitmproxy/test_command.py index b404d1196..6ec40fb9a 100644 --- a/test/mitmproxy/test_command.py +++ b/test/mitmproxy/test_command.py @@ -28,6 +28,10 @@ class TAddon: def cmd4(self, a: int, b: str, c: mitmproxy.types.Path) -> str: return "ok" + @command.command("cmd5") + def cmd5(self, choices: typing.Sequence[str]) -> typing.Sequence[str]: + return choices + @command.command("subcommand") def subcommand(self, cmd: mitmproxy.types.Cmd, *args: mitmproxy.types.Arg) -> str: return "ok" @@ -291,8 +295,11 @@ def test_simple(): c = command.CommandManager(tctx.master) a = TAddon() c.add("one.two", a.cmd1) + c.add("array.command", a.cmd5) assert c.commands["one.two"].help == "cmd1 help" assert(c.execute("one.two foo") == "ret foo") + assert (c.execute("one.two(foo)") == "ret foo") + assert (c.execute("array.command [1 2 3]") == ["1", "2", "3"]) assert(c.execute("one.two \"foo\"") == "ret foo") assert(c.execute("one.two 'foo'") == "ret foo") assert(c.call("one.two", "foo") == "ret foo")