Merge pull request #3260 from kajojify/language

Brackets support

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
 

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("'\"")

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")