Skip to content

Commit a69eea6

Browse files
kajojifykajojify
committed
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.
1 parent c78f10f commit a69eea6

File tree

7 files changed

+158
-32
lines changed

7 files changed

+158
-32
lines changed

‎mitmproxy/command.py‎

Lines changed: 9 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -38,12 +38,12 @@ def typename(t: type) -> str:
3838
"RunningCommand",
3939
[
4040
("cmdstr", str),
41-
("task", asyncio.Task)
41+
("task", asyncio.Future)
4242
],
4343
)
4444

4545

46-
class AsyncExectuionManager:
46+
class AsyncExecutionManager:
4747
def __init__(self) -> None:
4848
self.counter: int = 0
4949
self.running_cmds: typing.Dict[int, RunningCommand] = {}
@@ -114,7 +114,7 @@ def signature_help(self) -> str:
114114
ret = " -> " + ret
115115
return "%s %s%s" % (self.path, params, ret)
116116

117-
def prepare_args(self, args: typing.Sequence[str]) -> typing.List[typing.Any]:
117+
def prepare_args(self, args: typing.Sequence[typing.Any]) -> typing.List[typing.Any]:
118118
verify_arg_signature(self.func, list(args), {})
119119

120120
remainder: typing.Sequence[str] = []
@@ -131,7 +131,7 @@ def prepare_args(self, args: typing.Sequence[str]) -> typing.List[typing.Any]:
131131
pargs.append(arg)
132132
else:
133133
raise exceptions.CommandError(
134-
f"{arg} is unexpected data for {paramtype.display} type"
134+
f"{arg} is unexpected data for {t.display} type"
135135
)
136136
else:
137137
pargs.append(parsearg(self.manager, arg, paramtype))
@@ -149,9 +149,7 @@ def call(self, args: typing.Sequence[typing.Any]) -> typing.Any:
149149
typ = mitmproxy.types.CommandTypes.get(self.returntype)
150150
if not typ.is_valid(self.manager, typ, ret):
151151
raise exceptions.CommandError(
152-
"%s returned unexpected data - expected %s" % (
153-
self.path, typ.display
154-
)
152+
f"{self.path} returned unexpected data - expected {typ.display}"
155153
)
156154
return ret
157155

@@ -166,9 +164,7 @@ async def async_call(self, args: typing.Sequence[typing.Any]) -> typing.Any:
166164
typ = mitmproxy.types.CommandTypes.get(self.returntype)
167165
if not typ.is_valid(self.manager, typ, ret):
168166
raise exceptions.CommandError(
169-
"%s returned unexpected data - expected %s" % (
170-
self.path, typ.display
171-
)
167+
f"{self.path} returned unexpected data - expected {typ.display}"
172168
)
173169
return ret
174170

@@ -186,7 +182,7 @@ async def async_call(self, args: typing.Sequence[typing.Any]) -> typing.Any:
186182
class CommandManager(mitmproxy.types._CommandBase):
187183
def __init__(self, master):
188184
self.master = master
189-
self.async_manager = AsyncExectuionManager()
185+
self.async_manager = AsyncExecutionManager()
190186
self.command_parser = parser.create_parser(self)
191187
self.commands: typing.Dict[str, Command] = {}
192188
self.oneword_commands: typing.List[str] = []
@@ -288,7 +284,7 @@ def call_strings(self, path: str, args: typing.Sequence[str]) -> typing.Any:
288284
"""
289285
return self.get_command_by_path(path).call(args)
290286

291-
def async_execute(self, cmdstr: str) -> asyncio.Task:
287+
def async_execute(self, cmdstr: str) -> asyncio.Future:
292288
"""
293289
Schedule a command to be executed. May raise CommandError.
294290
"""
@@ -324,7 +320,7 @@ def parsearg(manager: CommandManager, spec: str, argtype: type) -> typing.Any:
324320
"""
325321
t = mitmproxy.types.CommandTypes.get(argtype, None)
326322
if not t:
327-
raise exceptions.CommandError("Unsupported argument type: %s" % argtype)
323+
raise exceptions.CommandError(f"Unsupported argument type: {argtype}")
328324
try:
329325
return t.parse(manager, argtype, spec) # type: ignore
330326
except exceptions.TypeError as e:

‎mitmproxy/language/lexer.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ class CommandLanguageLexer:
88
tokens = (
99
"WHITESPACE",
1010
"PIPE",
11+
"EQUAL_SIGN",
1112
"LPAREN", "RPAREN",
1213
"LBRACE", "RBRACE",
1314
"PLAIN_STR", "QUOTED_STR",
@@ -23,12 +24,13 @@ def __init__(self, oneword_commands: typing.Sequence[str]) -> None:
2324
# Main(INITIAL) state
2425
t_ignore_WHITESPACE = r"\s+"
2526
t_PIPE = r"\|"
27+
t_EQUAL_SIGN = r"\="
2628
t_LPAREN = r"\("
2729
t_RPAREN = r"\)"
2830
t_LBRACE = r"\["
2931
t_RBRACE = r"\]"
3032

31-
special_symbols = re.escape("()[]|")
33+
special_symbols = re.escape("()[]|=")
3234
plain_str = rf"[^{special_symbols}\s]+"
3335

3436
def t_COMMAND(self, t):

‎mitmproxy/language/parser.py‎

Lines changed: 15 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import typing
2+
import collections
23

34
import ply.lex as lex
45
import ply.yacc as yacc
@@ -11,12 +12,8 @@
1112
ParsedEntity = typing.Union[str, list, "ParsedCommand"]
1213

1314

14-
ParsedCommand = typing.NamedTuple(
15-
"ParsedCommand",
16-
[
17-
("command", "mitmproxy.command.Command"),
18-
("args", typing.List[ParsedEntity])
19-
]
15+
ParsedCommand = collections.namedtuple(
16+
"ParsedCommand", ["command", "args"]
2017
)
2118

2219

@@ -41,8 +38,7 @@ def p_starting_expression(self, p):
4138
"""starting_expression : PLAIN_STR
4239
| quoted_str
4340
| array
44-
| command_call_no_parentheses
45-
| command_call_with_parentheses"""
41+
| command_call"""
4642
p[0] = p[1]
4743
self._parsed_pipe_elem = p[0]
4844

@@ -62,6 +58,11 @@ def p_pipe_expression(self, p):
6258
p[0] = self._call_command(p[2], new_args)
6359
self._parsed_pipe_elem = p[0]
6460

61+
def p_command_call(self, p):
62+
"""command_call : command_call_no_parentheses
63+
| command_call_with_parentheses"""
64+
p[0] = p[1]
65+
6566
def p_command_call_no_parentheses(self, p):
6667
"""command_call_no_parentheses : COMMAND argument_list"""
6768
p[0] = self._call_command(p[1], p[2])
@@ -76,11 +77,17 @@ def p_argument_list(self, p):
7677
| argument_list argument"""
7778
p[0] = self._create_list(p)
7879

80+
def p_assignment(self, p):
81+
"""assignment : PLAIN_STR EQUAL_SIGN starting_expression
82+
| QUOTED_STR EQUAL_SIGN starting_expression"""
83+
p[0] = f"{p[1]}{p[2]}{p[3]}"
84+
7985
def p_argument(self, p):
8086
"""argument : PLAIN_STR
8187
| quoted_str
8288
| array
8389
| COMMAND
90+
| assignment
8491
| command_call_with_parentheses"""
8592
p[0] = p[1]
8693

‎mitmproxy/tools/console/keymap.py‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44
import ruamel.yaml
55

66
from mitmproxy import command
7-
from mitmproxy.language import lexer
87
from mitmproxy.tools.console import commandexecutor
98
from mitmproxy.tools.console import signals
109
from mitmproxy import ctx
@@ -56,7 +55,6 @@ def sortkey(self):
5655

5756
class Keymap:
5857
def __init__(self, master):
59-
self.oneword_commands = master.commands.oneword_commands
6058
self.executor = commandexecutor.CommandExecutor(master)
6159
self.keys = {}
6260
for c in Contexts:

‎test/mitmproxy/language/test_traversal.py‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
import typing
2+
import asyncio
3+
4+
from mitmproxy import command
5+
from mitmproxy.test import taddons
6+
from mitmproxy.language import lexer, parser, traversal
7+
8+
import pytest
9+
10+
11+
class TAddon:
12+
@command.command("cmd1")
13+
def cmd1(self, foo: typing.Sequence[str]) -> str:
14+
return " ".join(foo)
15+
16+
@command.command("cmd2")
17+
def cmd2(self, foo: str) -> str:
18+
return foo
19+
20+
@command.command("cmd3")
21+
async def cmd3(self, foo: str) -> str:
22+
await asyncio.sleep(0.01)
23+
return foo
24+
25+
26+
@pytest.mark.asyncio
27+
async def test_execute_parsed_line():
28+
test_commands = ["""join.cmd1 [str.cmd2(abc)
29+
str.cmd2(strasync.cmd3("def"))]""",
30+
"[1 2 3]", "str.cmd2 abc | strasync.cmd3()"]
31+
results = ["abc def", ['1', '2', '3'], "abc"]
32+
33+
with taddons.context() as tctx:
34+
cm = command.CommandManager(tctx.master)
35+
a = TAddon()
36+
cm.add("join.cmd1", a.cmd1)
37+
cm.add("str.cmd2", a.cmd2)
38+
cm.add("strasync.cmd3", a.cmd3)
39+
40+
command_parser = parser.create_parser(cm)
41+
for cmd, exp_res in zip(test_commands, results):
42+
lxr = lexer.create_lexer(cmd, cm.oneword_commands)
43+
parsed = command_parser.parse(lxr, async_exec=True)
44+
45+
result = await traversal.execute_parsed_line(parsed)
46+
assert result == exp_res

‎test/mitmproxy/test_command.py‎

Lines changed: 79 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,15 @@
11
import typing
22
import inspect
3+
import asyncio
4+
from unittest import mock
5+
6+
import mitmproxy.types
37
from mitmproxy import command
48
from mitmproxy import flow
59
from mitmproxy import exceptions
610
from mitmproxy.test import tflow
711
from mitmproxy.test import taddons
8-
import mitmproxy.types
12+
913
import io
1014
import pytest
1115

@@ -36,6 +40,15 @@ def cmd5(self, choices: typing.Sequence[str]) -> typing.Sequence[str]:
3640
def cmd6(self, pipe_value: str) -> str:
3741
return pipe_value
3842

43+
@command.command("cmd7")
44+
async def cmd7(self, foo: str) -> str:
45+
await asyncio.sleep(0.01)
46+
return foo
47+
48+
@command.command("cmd8")
49+
async def cmd8(self, foo: str) -> str:
50+
return 99
51+
3952
@command.command("subcommand")
4053
def subcommand(self, cmd: mitmproxy.types.Cmd, *args: mitmproxy.types.Arg) -> str:
4154
return "ok"
@@ -44,6 +57,10 @@ def subcommand(self, cmd: mitmproxy.types.Cmd, *args: mitmproxy.types.Arg) -> st
4457
def empty(self) -> None:
4558
pass
4659

60+
@command.command("empty")
61+
async def asyncempty(self) -> None:
62+
pass
63+
4764
@command.command("varargs")
4865
def varargs(self, one: str, *var: str) -> typing.Sequence[str]:
4966
return list(var)
@@ -82,6 +99,35 @@ def invalidarg(self, u: Unsupported):
8299
pass
83100

84101

102+
class TestAsyncExecutionManager:
103+
def test_add_command(self):
104+
aem = command.AsyncExecutionManager()
105+
dummy_command = command.RunningCommand("addon.command", mock.Mock())
106+
aem.add_command(dummy_command)
107+
assert aem.running_cmds == {1: dummy_command}
108+
109+
def test_stop_command(self):
110+
aem = command.AsyncExecutionManager()
111+
dummy_command = command.RunningCommand("addon.command", mock.Mock())
112+
aem.add_command(dummy_command)
113+
with pytest.raises(ValueError, match="There is not the command"):
114+
aem.stop_command(100)
115+
116+
aem.stop_command(1)
117+
assert aem.running_cmds == {}
118+
119+
def test_get_runnings(self):
120+
aem = command.AsyncExecutionManager()
121+
expected_res = []
122+
for i in range(3):
123+
cmd = f"addon.command{i}"
124+
dummy = command.RunningCommand(cmd, mock.Mock())
125+
aem.add_command(dummy)
126+
expected_res.append((i + 1, cmd))
127+
128+
assert aem.get_running() == expected_res
129+
130+
85131
class TestCommand:
86132
def test_typecheck(self):
87133
with taddons.context(loadcore=False) as tctx:
@@ -115,9 +161,25 @@ def test_call(self):
115161
with pytest.raises(exceptions.CommandError):
116162
c.call(["foo"])
117163

164+
with pytest.raises(exceptions.CommandError, match="unexpected data"):
165+
c.call([123])
166+
118167
c = command.Command(cm, "cmd.three", a.cmd3)
119168
assert c.call(["1"]) == 1
120169

170+
@pytest.mark.asyncio
171+
async def test_async_call(self):
172+
with taddons.context() as tctx:
173+
cm = command.CommandManager(tctx.master)
174+
a = TAddon()
175+
176+
c = command.Command(cm, "async.empty", a.asyncempty)
177+
await c.async_call([])
178+
179+
c = command.Command(cm, "asynccmd.two", a.cmd8)
180+
with pytest.raises(exceptions.CommandError, match="unexpected data"):
181+
await c.async_call(["foo"])
182+
121183
def test_parse_partial(self):
122184
tests = [
123185
[
@@ -301,6 +363,7 @@ def test_simple():
301363
c.add("one.two", a.cmd1)
302364
c.add("array.command", a.cmd5)
303365
c.add("pipe.command", a.cmd6)
366+
c.add("strasync.command", a.cmd7)
304367

305368
assert c.commands["one.two"].help == "cmd1 help"
306369
assert(c.execute("one.two foo") == "ret foo")
@@ -320,6 +383,8 @@ def test_simple():
320383
c.execute("")
321384
with pytest.raises(exceptions.CommandError, match="argument mismatch"):
322385
c.execute("one.two too many args")
386+
with pytest.raises(exceptions.ExecutionError, match="sync executor"):
387+
c.execute("strasync.command abc")
323388
with pytest.raises(exceptions.CommandError, match="Unknown"):
324389
c.call("nonexistent")
325390

@@ -331,6 +396,19 @@ def test_simple():
331396
assert fp.getvalue()
332397

333398

399+
@pytest.mark.asyncio
400+
async def test_async_execute():
401+
with taddons.context() as tctx:
402+
c = command.CommandManager(tctx.master)
403+
a = TAddon()
404+
c.add("strasync.command", a.cmd7)
405+
406+
c.async_execute("strasync.command abc")
407+
assert c.async_manager.get_running() == [(1, "strasync.command abc")]
408+
assert "abc" == await c.async_manager.running_cmds[1].task
409+
assert c.async_manager.get_running() == []
410+
411+
334412
def test_typename():
335413
assert command.typename(str) == "str"
336414
assert command.typename(typing.Sequence[flow.Flow]) == "[flow]"

‎test/mitmproxy/tools/console/test_defaultkeys.py‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
from mitmproxy.tools.console import defaultkeys
33
from mitmproxy.tools.console import keymap
44
from mitmproxy.tools.console import master
5-
from mitmproxy.language import lexer
5+
from mitmproxy.language import lexer, parser
66

77
import pytest
88

@@ -15,12 +15,11 @@ async def test_commands_exist():
1515
m = master.ConsoleMaster(None)
1616
await m.load_flow(tflow())
1717

18-
for binding in km.bindings:
19-
cmd, *args = lexer.get_tokens(binding.command, state="INITIAL")
20-
assert cmd in m.commands.commands
18+
command_parser = parser.create_parser(m.commands)
2119

22-
cmd_obj = m.commands.commands[cmd]
20+
for binding in km.bindings:
21+
lxr = lexer.create_lexer(binding.command, m.commands.oneword_commands)
2322
try:
24-
cmd_obj.prepare_args(args)
23+
command_parser.parse(lxr, async_exec=True)
2524
except Exception as e:
26-
raise ValueError("Invalid command: {}".format(binding.command)) from e
25+
raise ValueError(f"Invalid command: '{binding.command}'") from e

0 commit comments

Comments
 (0)