From a69eea67fbb7a58aad0a3fd513bd2a70a5d581ce Mon Sep 17 00:00:00 2001 From: Miroslav Date: Sun, 12 Aug 2018 16:00:39 +0300 Subject: [PATCH] Going into mergeable state. Some typing fixes since typing module doesn't fully support recursive types. a few typos were fixed test coverage, new tests for the new code new syntax feature - assignment for 'set' command. --- mitmproxy/command.py | 22 +++-- mitmproxy/language/lexer.py | 4 +- mitmproxy/language/parser.py | 23 ++++-- mitmproxy/tools/console/keymap.py | 2 - test/mitmproxy/language/test_traversal.py | 46 +++++++++++ test/mitmproxy/test_command.py | 80 ++++++++++++++++++- .../tools/console/test_defaultkeys.py | 13 ++- 7 files changed, 158 insertions(+), 32 deletions(-) create mode 100644 test/mitmproxy/language/test_traversal.py diff --git a/mitmproxy/command.py b/mitmproxy/command.py index df32f9fea..5c07f0ad5 100644 --- a/mitmproxy/command.py +++ b/mitmproxy/command.py @@ -38,12 +38,12 @@ RunningCommand = typing.NamedTuple( "RunningCommand", [ ("cmdstr", str), - ("task", asyncio.Task) + ("task", asyncio.Future) ], ) -class AsyncExectuionManager: +class AsyncExecutionManager: def __init__(self) -> None: self.counter: int = 0 self.running_cmds: typing.Dict[int, RunningCommand] = {} @@ -114,7 +114,7 @@ class Command: ret = " -> " + ret return "%s %s%s" % (self.path, params, ret) - def prepare_args(self, args: typing.Sequence[str]) -> typing.List[typing.Any]: + def prepare_args(self, args: typing.Sequence[typing.Any]) -> typing.List[typing.Any]: verify_arg_signature(self.func, list(args), {}) remainder: typing.Sequence[str] = [] @@ -131,7 +131,7 @@ class Command: pargs.append(arg) else: raise exceptions.CommandError( - f"{arg} is unexpected data for {paramtype.display} type" + f"{arg} is unexpected data for {t.display} type" ) else: pargs.append(parsearg(self.manager, arg, paramtype)) @@ -149,9 +149,7 @@ class Command: typ = mitmproxy.types.CommandTypes.get(self.returntype) if not typ.is_valid(self.manager, typ, ret): raise exceptions.CommandError( - "%s returned unexpected data - expected %s" % ( - self.path, typ.display - ) + f"{self.path} returned unexpected data - expected {typ.display}" ) return ret @@ -166,9 +164,7 @@ class Command: typ = mitmproxy.types.CommandTypes.get(self.returntype) if not typ.is_valid(self.manager, typ, ret): raise exceptions.CommandError( - "%s returned unexpected data - expected %s" % ( - self.path, typ.display - ) + f"{self.path} returned unexpected data - expected {typ.display}" ) return ret @@ -186,7 +182,7 @@ ParseResult = typing.NamedTuple( class CommandManager(mitmproxy.types._CommandBase): def __init__(self, master): self.master = master - self.async_manager = AsyncExectuionManager() + self.async_manager = AsyncExecutionManager() self.command_parser = parser.create_parser(self) self.commands: typing.Dict[str, Command] = {} self.oneword_commands: typing.List[str] = [] @@ -288,7 +284,7 @@ class CommandManager(mitmproxy.types._CommandBase): """ return self.get_command_by_path(path).call(args) - def async_execute(self, cmdstr: str) -> asyncio.Task: + def async_execute(self, cmdstr: str) -> asyncio.Future: """ Schedule a command to be executed. May raise CommandError. """ @@ -324,7 +320,7 @@ def parsearg(manager: CommandManager, spec: str, argtype: type) -> typing.Any: """ t = mitmproxy.types.CommandTypes.get(argtype, None) if not t: - raise exceptions.CommandError("Unsupported argument type: %s" % argtype) + raise exceptions.CommandError(f"Unsupported argument type: {argtype}") try: return t.parse(manager, argtype, spec) # type: ignore except exceptions.TypeError as e: diff --git a/mitmproxy/language/lexer.py b/mitmproxy/language/lexer.py index 2e44d7df2..bc76e6f05 100644 --- a/mitmproxy/language/lexer.py +++ b/mitmproxy/language/lexer.py @@ -8,6 +8,7 @@ class CommandLanguageLexer: tokens = ( "WHITESPACE", "PIPE", + "EQUAL_SIGN", "LPAREN", "RPAREN", "LBRACE", "RBRACE", "PLAIN_STR", "QUOTED_STR", @@ -23,12 +24,13 @@ class CommandLanguageLexer: # Main(INITIAL) state t_ignore_WHITESPACE = r"\s+" t_PIPE = r"\|" + t_EQUAL_SIGN = r"\=" t_LPAREN = r"\(" t_RPAREN = r"\)" t_LBRACE = r"\[" t_RBRACE = r"\]" - special_symbols = re.escape("()[]|") + special_symbols = re.escape("()[]|=") plain_str = rf"[^{special_symbols}\s]+" def t_COMMAND(self, t): diff --git a/mitmproxy/language/parser.py b/mitmproxy/language/parser.py index c3923ccfa..cc7824095 100644 --- a/mitmproxy/language/parser.py +++ b/mitmproxy/language/parser.py @@ -1,4 +1,5 @@ import typing +import collections import ply.lex as lex import ply.yacc as yacc @@ -11,12 +12,8 @@ from mitmproxy.language.lexer import CommandLanguageLexer ParsedEntity = typing.Union[str, list, "ParsedCommand"] -ParsedCommand = typing.NamedTuple( - "ParsedCommand", - [ - ("command", "mitmproxy.command.Command"), - ("args", typing.List[ParsedEntity]) - ] +ParsedCommand = collections.namedtuple( + "ParsedCommand", ["command", "args"] ) @@ -41,8 +38,7 @@ class CommandLanguageParser: """starting_expression : PLAIN_STR | quoted_str | array - | command_call_no_parentheses - | command_call_with_parentheses""" + | command_call""" p[0] = p[1] self._parsed_pipe_elem = p[0] @@ -62,6 +58,11 @@ class CommandLanguageParser: p[0] = self._call_command(p[2], new_args) self._parsed_pipe_elem = p[0] + def p_command_call(self, p): + """command_call : command_call_no_parentheses + | command_call_with_parentheses""" + p[0] = p[1] + def p_command_call_no_parentheses(self, p): """command_call_no_parentheses : COMMAND argument_list""" p[0] = self._call_command(p[1], p[2]) @@ -76,11 +77,17 @@ class CommandLanguageParser: | argument_list argument""" p[0] = self._create_list(p) + def p_assignment(self, p): + """assignment : PLAIN_STR EQUAL_SIGN starting_expression + | QUOTED_STR EQUAL_SIGN starting_expression""" + p[0] = f"{p[1]}{p[2]}{p[3]}" + def p_argument(self, p): """argument : PLAIN_STR | quoted_str | array | COMMAND + | assignment | command_call_with_parentheses""" p[0] = p[1] diff --git a/mitmproxy/tools/console/keymap.py b/mitmproxy/tools/console/keymap.py index d85d7d68b..d22420bfd 100644 --- a/mitmproxy/tools/console/keymap.py +++ b/mitmproxy/tools/console/keymap.py @@ -4,7 +4,6 @@ import os import ruamel.yaml from mitmproxy import command -from mitmproxy.language import lexer from mitmproxy.tools.console import commandexecutor from mitmproxy.tools.console import signals from mitmproxy import ctx @@ -56,7 +55,6 @@ class Binding: class Keymap: def __init__(self, master): - self.oneword_commands = master.commands.oneword_commands self.executor = commandexecutor.CommandExecutor(master) self.keys = {} for c in Contexts: diff --git a/test/mitmproxy/language/test_traversal.py b/test/mitmproxy/language/test_traversal.py new file mode 100644 index 000000000..5f9e1c9b3 --- /dev/null +++ b/test/mitmproxy/language/test_traversal.py @@ -0,0 +1,46 @@ +import typing +import asyncio + +from mitmproxy import command +from mitmproxy.test import taddons +from mitmproxy.language import lexer, parser, traversal + +import pytest + + +class TAddon: + @command.command("cmd1") + def cmd1(self, foo: typing.Sequence[str]) -> str: + return " ".join(foo) + + @command.command("cmd2") + def cmd2(self, foo: str) -> str: + return foo + + @command.command("cmd3") + async def cmd3(self, foo: str) -> str: + await asyncio.sleep(0.01) + return foo + + +@pytest.mark.asyncio +async def test_execute_parsed_line(): + test_commands = ["""join.cmd1 [str.cmd2(abc) + str.cmd2(strasync.cmd3("def"))]""", + "[1 2 3]", "str.cmd2 abc | strasync.cmd3()"] + results = ["abc def", ['1', '2', '3'], "abc"] + + with taddons.context() as tctx: + cm = command.CommandManager(tctx.master) + a = TAddon() + cm.add("join.cmd1", a.cmd1) + cm.add("str.cmd2", a.cmd2) + cm.add("strasync.cmd3", a.cmd3) + + command_parser = parser.create_parser(cm) + for cmd, exp_res in zip(test_commands, results): + lxr = lexer.create_lexer(cmd, cm.oneword_commands) + parsed = command_parser.parse(lxr, async_exec=True) + + result = await traversal.execute_parsed_line(parsed) + assert result == exp_res diff --git a/test/mitmproxy/test_command.py b/test/mitmproxy/test_command.py index b26ddac96..44fb1a06e 100644 --- a/test/mitmproxy/test_command.py +++ b/test/mitmproxy/test_command.py @@ -1,11 +1,15 @@ import typing import inspect +import asyncio +from unittest import mock + +import mitmproxy.types from mitmproxy import command from mitmproxy import flow from mitmproxy import exceptions from mitmproxy.test import tflow from mitmproxy.test import taddons -import mitmproxy.types + import io import pytest @@ -36,6 +40,15 @@ class TAddon: def cmd6(self, pipe_value: str) -> str: return pipe_value + @command.command("cmd7") + async def cmd7(self, foo: str) -> str: + await asyncio.sleep(0.01) + return foo + + @command.command("cmd8") + async def cmd8(self, foo: str) -> str: + return 99 + @command.command("subcommand") def subcommand(self, cmd: mitmproxy.types.Cmd, *args: mitmproxy.types.Arg) -> str: return "ok" @@ -44,6 +57,10 @@ class TAddon: def empty(self) -> None: pass + @command.command("empty") + async def asyncempty(self) -> None: + pass + @command.command("varargs") def varargs(self, one: str, *var: str) -> typing.Sequence[str]: return list(var) @@ -82,6 +99,35 @@ class TypeErrAddon: pass +class TestAsyncExecutionManager: + def test_add_command(self): + aem = command.AsyncExecutionManager() + dummy_command = command.RunningCommand("addon.command", mock.Mock()) + aem.add_command(dummy_command) + assert aem.running_cmds == {1: dummy_command} + + def test_stop_command(self): + aem = command.AsyncExecutionManager() + dummy_command = command.RunningCommand("addon.command", mock.Mock()) + aem.add_command(dummy_command) + with pytest.raises(ValueError, match="There is not the command"): + aem.stop_command(100) + + aem.stop_command(1) + assert aem.running_cmds == {} + + def test_get_runnings(self): + aem = command.AsyncExecutionManager() + expected_res = [] + for i in range(3): + cmd = f"addon.command{i}" + dummy = command.RunningCommand(cmd, mock.Mock()) + aem.add_command(dummy) + expected_res.append((i + 1, cmd)) + + assert aem.get_running() == expected_res + + class TestCommand: def test_typecheck(self): with taddons.context(loadcore=False) as tctx: @@ -115,9 +161,25 @@ class TestCommand: with pytest.raises(exceptions.CommandError): c.call(["foo"]) + with pytest.raises(exceptions.CommandError, match="unexpected data"): + c.call([123]) + c = command.Command(cm, "cmd.three", a.cmd3) assert c.call(["1"]) == 1 + @pytest.mark.asyncio + async def test_async_call(self): + with taddons.context() as tctx: + cm = command.CommandManager(tctx.master) + a = TAddon() + + c = command.Command(cm, "async.empty", a.asyncempty) + await c.async_call([]) + + c = command.Command(cm, "asynccmd.two", a.cmd8) + with pytest.raises(exceptions.CommandError, match="unexpected data"): + await c.async_call(["foo"]) + def test_parse_partial(self): tests = [ [ @@ -301,6 +363,7 @@ def test_simple(): c.add("one.two", a.cmd1) c.add("array.command", a.cmd5) c.add("pipe.command", a.cmd6) + c.add("strasync.command", a.cmd7) assert c.commands["one.two"].help == "cmd1 help" assert(c.execute("one.two foo") == "ret foo") @@ -320,6 +383,8 @@ def test_simple(): c.execute("") with pytest.raises(exceptions.CommandError, match="argument mismatch"): c.execute("one.two too many args") + with pytest.raises(exceptions.ExecutionError, match="sync executor"): + c.execute("strasync.command abc") with pytest.raises(exceptions.CommandError, match="Unknown"): c.call("nonexistent") @@ -331,6 +396,19 @@ def test_simple(): assert fp.getvalue() +@pytest.mark.asyncio +async def test_async_execute(): + with taddons.context() as tctx: + c = command.CommandManager(tctx.master) + a = TAddon() + c.add("strasync.command", a.cmd7) + + c.async_execute("strasync.command abc") + assert c.async_manager.get_running() == [(1, "strasync.command abc")] + assert "abc" == await c.async_manager.running_cmds[1].task + assert c.async_manager.get_running() == [] + + def test_typename(): assert command.typename(str) == "str" assert command.typename(typing.Sequence[flow.Flow]) == "[flow]" diff --git a/test/mitmproxy/tools/console/test_defaultkeys.py b/test/mitmproxy/tools/console/test_defaultkeys.py index f60174aa7..fe32a1646 100644 --- a/test/mitmproxy/tools/console/test_defaultkeys.py +++ b/test/mitmproxy/tools/console/test_defaultkeys.py @@ -2,7 +2,7 @@ from mitmproxy.test.tflow import tflow from mitmproxy.tools.console import defaultkeys from mitmproxy.tools.console import keymap from mitmproxy.tools.console import master -from mitmproxy.language import lexer +from mitmproxy.language import lexer, parser import pytest @@ -15,12 +15,11 @@ async def test_commands_exist(): m = master.ConsoleMaster(None) await m.load_flow(tflow()) - for binding in km.bindings: - cmd, *args = lexer.get_tokens(binding.command, state="INITIAL") - assert cmd in m.commands.commands + command_parser = parser.create_parser(m.commands) - cmd_obj = m.commands.commands[cmd] + for binding in km.bindings: + lxr = lexer.create_lexer(binding.command, m.commands.oneword_commands) try: - cmd_obj.prepare_args(args) + command_parser.parse(lxr, async_exec=True) except Exception as e: - raise ValueError("Invalid command: {}".format(binding.command)) from e + raise ValueError(f"Invalid command: '{binding.command}'") from e