Removed the custom lexer in favor of using pyparsing.

3 files changed, 59 insertions, 131 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):
diff --git a/mitmproxy/lexer.py b/mitmproxy/lexer.py
deleted file mode 100644
index f123a838..00000000
--- a/mitmproxy/lexer.py
+++ /dev/null
@@ -1,99 +0,0 @@
-from enum import Enum
-import io
-from typing import Union, List
-
-
-class State(Enum):
-    QUOTE = 1
-    ESCAPE = 2
-    TEXT = 3
-
-
-class Lexer:
-
-    def __init__(self, text: Union[str, io.StringIO]) -> None:
-        self._count = 0
-        self._parsed = False
-
-        self._state = State.TEXT
-        self._states: List[State] = []
-        self._text_pos = 0
-        self._quote_start_pos = 0
-
-        if isinstance(text, str):
-            self.text = io.StringIO(text)
-        else:
-            self.text = text
-
-    def __iter__(self):
-        return self
-
-    def __next__(self):
-        t = self.get_token()
-
-        if t == '':
-            raise StopIteration
-
-        return t
-
-    def get_token(self):
-        try:
-            return self.parse()
-        except ValueError:
-            raise
-
-    def parse(self):
-        acc = ''
-        quote = ''
-        self._state = State.TEXT
-
-        whitespace = "\r\n\t "
-
-        self.text.seek(self._text_pos)
-
-        while True:
-            ch = self.text.read(1)
-            self._text_pos += 1
-
-            # If this is the last char of the string, let's save the token
-            if ch == '' or ch is None:
-                break
-
-            if self._state == State.QUOTE:
-                if ch == '\\':
-                    self._states.append(self._state)
-                    self._state = State.ESCAPE
-                    acc += ch
-                elif ch == quote:
-                    self._state = self._states.pop()
-                    acc += ch
-                else:
-                    acc += ch
-
-            elif self._state == State.ESCAPE:
-                acc += ch
-                self._state = self._states.pop()
-
-            elif self._state == State.TEXT:
-                if ch in whitespace:
-                    if acc != '':
-                        break
-                elif ch == '"' or ch == "'":
-                    quote = ch
-                    self._quote_start_pos = self._text_pos
-                    self._states.append(self._state)
-                    self._state = State.QUOTE
-                    acc += ch
-                elif ch == '\\':
-                    self._states.append(self._state)
-                    self._state = State.ESCAPE
-                    acc += ch
-                else:
-                    acc += ch
-
-        self._token = acc
-
-        if self._state == State.QUOTE:
-            raise ValueError("No closing quotation for quote in position %d" % self._quote_start_pos)
-
-        return self._token
diff --git a/mitmproxy/tools/console/commander/commander.py b/mitmproxy/tools/console/commander/commander.py
index ee31e1e9..fa67407e 100644
--- a/mitmproxy/tools/console/commander/commander.py
+++ b/mitmproxy/tools/console/commander/commander.py
@@ -52,7 +52,7 @@ CompletionState = typing.NamedTuple(
 class CommandBuffer:
     def __init__(self, master: mitmproxy.master.Master, start: str = "") -> None:
         self.master = master
-        self.text = self.flatten(start)
+        self.text = start
         # Cursor is always within the range [0:len(buffer)].
         self._cursor = len(self.text)
         self.completion: CompletionState = None
@@ -105,7 +105,7 @@ class CommandBuffer:
                     ret.append(("commander_invalid", p.value))
                 else:
                     ret.append(("text", ""))
-                ret.append(("text", " "))
+
             if remhelp:
                 ret.append(("text", " "))
                 for v in remhelp:
@@ -113,11 +113,6 @@ class CommandBuffer:
 
         return ret
 
-    def flatten(self, txt):
-        parts, _ = self.parse_quoted(txt)
-        ret = [x.value for x in parts]
-        return " ".join(ret)
-
     def left(self) -> None:
         self.cursor = self.cursor - 1
 
@@ -141,7 +136,7 @@ class CommandBuffer:
             nxt = self.completion.completer.cycle()
             buf = " ".join([i.value for i in self.completion.parse[:-1]]) + " " + nxt
             buf = buf.strip()
-            self.text = self.flatten(buf)
+            self.text = buf
             self.cursor = len(self.text)
 
     def backspace(self) -> None:
@@ -155,6 +150,11 @@ class CommandBuffer:
         """
             Inserts text at the cursor.
         """
+
+        # We don't want to insert a space before the command
+        if k == ' ' and self.text[0:self.cursor].strip() == '':
+            return
+
         self.text = self.text[:self.cursor] + k + self.text[self.cursor:]
         self.cursor += len(k)
         self.completion = None