Removed the custom lexer in favor of using pyparsing.

1 files changed, 51 insertions, 24 deletions

diff --git a/mitmproxy/command.py b/mitmproxy/command.py
index d9ba4055..4aa6fdb2 100644
--- a/mitmproxy/command.py
+++ b/mitmproxy/command.py
@@ -3,14 +3,13 @@
 """
 import inspect
 import types
-import io
 import typing
 import textwrap
 import functools
 import sys
+import pyparsing
 
 from mitmproxy import exceptions
-from mitmproxy import lexer
 import mitmproxy.types
 
 
@@ -22,10 +21,6 @@ def verify_arg_signature(f: typing.Callable, args: list, kwargs: dict) -> None:
         raise exceptions.CommandError("command argument mismatch: %s" % v.args[0])
 
 
-def get_lexer(s):
-    return lexer.Lexer(s)
-
-
 def typename(t: type) -> str:
     """
         Translates a type to an explanatory string.
@@ -79,6 +74,20 @@ class Command:
         return "%s %s%s" % (self.path, params, ret)
 
     def prepare_args(self, args: typing.Sequence[str]) -> typing.List[typing.Any]:
+
+        # Arguments that are just blank spaces aren't really arguments
+        # We need to get rid of those. If the user intended to pass a sequence
+        # of spaces, it would come between quotes
+        clean_args = []
+        for a in args:
+            if isinstance(a, str):
+                if a.strip() != '':
+                    clean_args.append(a)
+            else:
+                clean_args.append(a)
+
+        args = clean_args
+
         verify_arg_signature(self.func, list(args), {})
 
         remainder: typing.Sequence[str] = []
@@ -152,24 +161,36 @@ class CommandManager(mitmproxy.types._CommandBase):
         """
             Parse a possibly partial command. Return a sequence of ParseResults and a sequence of remainder type help items.
         """
-        buf = io.StringIO(cmdstr)
         parts: typing.List[str] = []
-        lex = get_lexer(buf)
-        while 1:
-            remainder = cmdstr[buf.tell():]
-            try:
-                t = lex.get_token()
-            except ValueError:
-                parts.append(remainder)
-                break
-            if not t:
-                break
-            parts.append(t)
+
+        rex = pyparsing.QuotedString("\"", escChar='\\', unquoteResults=False) |\
+            pyparsing.QuotedString("'", escChar='\\', unquoteResults=False) |\
+            pyparsing.Combine(pyparsing.Literal('"') + pyparsing.Word(pyparsing.printables + " ") + pyparsing.StringEnd()) |\
+            pyparsing.Word(pyparsing.printables) |\
+            pyparsing.Word(' ')
+
+        rex = rex.copy().leaveWhitespace()
+
+        remainder = cmdstr
+
+        for t, start, end in rex.scanString(cmdstr):
+
+            remainder = cmdstr[end:]
+            parts.append(t[0])
+
+        if remainder != '':
+            parts.append(remainder)
 
         if not parts:
             parts = []
-        elif cmdstr.endswith(" "):
-            parts.append("")
+
+        # First item in parts has always to be the command
+        # so we remove any blank tokens from the start of it
+        while True:
+            if parts and parts[0].strip() == '':
+                del parts[0]
+            else:
+                break
 
         parse: typing.List[ParseResult] = []
         params: typing.List[type] = []
@@ -180,10 +201,15 @@ class CommandManager(mitmproxy.types._CommandBase):
                 if parts[i] in self.commands:
                     params.extend(self.commands[parts[i]].paramtypes)
             elif params:
-                typ = params.pop(0)
-                if typ == mitmproxy.types.Cmd and params and params[0] == mitmproxy.types.Arg:
-                    if parts[i] in self.commands:
-                        params[:] = self.commands[parts[i]].paramtypes
+                if parts[i].strip() != '':
+                    typ = params.pop(0)
+                    if typ == mitmproxy.types.Cmd and params and params[0] == mitmproxy.types.Arg:
+                        if parts[i] in self.commands:
+                            params[:] = self.commands[parts[i]].paramtypes
+                else:
+                    # If the token is just a bunch of spaces, then we don't
+                    # want to count it against the arguments of the command
+                    typ = mitmproxy.types.Unknown
             else:
                 typ = mitmproxy.types.Unknown
 
@@ -228,6 +254,7 @@ class CommandManager(mitmproxy.types._CommandBase):
         """
         if path not in self.commands:
             raise exceptions.CommandError("Unknown command: %s" % path)
+
         return self.commands[path].call(args)
 
     def execute(self, cmdstr: str):