From f75a95acea6772457b25395b0fdd2c97bfebb936 Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Mon, 18 Nov 2019 03:45:16 +0100 Subject: [PATCH] fix vararg handling --- mitmproxy/command.py | 36 +++++++++++-------- .../tools/console/commander/commander.py | 2 +- mitmproxy/types.py | 2 +- test/mitmproxy/test_types.py | 3 +- .../mitmproxy/tools/console/test_commander.py | 14 ++++---- .../tools/console/test_defaultkeys.py | 17 +++++---- 6 files changed, 43 insertions(+), 31 deletions(-) diff --git a/mitmproxy/command.py b/mitmproxy/command.py index a64c74045..00238f46b 100644 --- a/mitmproxy/command.py +++ b/mitmproxy/command.py @@ -41,8 +41,15 @@ def _empty_as_none(x: typing.Any) -> typing.Any: class CommandParameter(typing.NamedTuple): - display_name: str + name: str type: typing.Type + kind: inspect._ParameterKind = inspect.Parameter.POSITIONAL_OR_KEYWORD + + def __str__(self): + if self.kind is inspect.Parameter.VAR_POSITIONAL: + return f"*{self.name}" + else: + return self.name class Command: @@ -77,16 +84,14 @@ class Command: @property def parameters(self) -> typing.List[CommandParameter]: - """Returns a list of (display name, type) tuples.""" + """Returns a list of CommandParameters.""" ret = [] for name, param in self.signature.parameters.items(): - if param.kind is param.VAR_POSITIONAL: - name = f"*{name}" - ret.append(CommandParameter(name, param.annotation)) + ret.append(CommandParameter(name, param.annotation, param.kind)) return ret def signature_help(self) -> str: - params = " ".join(name for name, t in self.parameters) + params = " ".join(str(param) for param in self.parameters) if self.return_type: ret = f" -> {typename(self.return_type)}" else: @@ -184,6 +189,7 @@ class CommandManager: CommandParameter("", mitmproxy.types.Cmd), CommandParameter("", mitmproxy.types.CmdArgs), ] + expected: typing.Optional[CommandParameter] = None for part in parts: if part.isspace(): parsed.append( @@ -195,16 +201,18 @@ class CommandManager: ) continue - if next_params: - expected_type: typing.Type = next_params.pop(0).type + if expected and expected.kind is inspect.Parameter.VAR_POSITIONAL: + assert not next_params + elif next_params: + expected = next_params.pop(0) else: - expected_type = mitmproxy.types.Unknown + expected = CommandParameter("", mitmproxy.types.Unknown) arg_is_known_command = ( - expected_type == mitmproxy.types.Cmd and part in self.commands + expected.type == mitmproxy.types.Cmd and part in self.commands ) arg_is_unknown_command = ( - expected_type == mitmproxy.types.Cmd and part not in self.commands + expected.type == mitmproxy.types.Cmd and part not in self.commands ) command_args_following = ( next_params and next_params[0].type == mitmproxy.types.CmdArgs @@ -214,11 +222,11 @@ class CommandManager: if arg_is_unknown_command and command_args_following: next_params.pop(0) - to = mitmproxy.types.CommandTypes.get(expected_type, None) + to = mitmproxy.types.CommandTypes.get(expected.type, None) valid = False if to: try: - to.parse(self, expected_type, part) + to.parse(self, expected.type, part) except exceptions.TypeError: valid = False else: @@ -227,7 +235,7 @@ class CommandManager: parsed.append( ParseResult( value=part, - type=expected_type, + type=expected.type, valid=valid, ) ) diff --git a/mitmproxy/tools/console/commander/commander.py b/mitmproxy/tools/console/commander/commander.py index ba3e601e3..f826b9842 100644 --- a/mitmproxy/tools/console/commander/commander.py +++ b/mitmproxy/tools/console/commander/commander.py @@ -88,7 +88,7 @@ class CommandBuffer: if parts[-1].type != mitmproxy.types.Space: ret.append(("text", " ")) for param in remaining: - ret.append(("commander_hint", f"{param.display_name} ")) + ret.append(("commander_hint", f"{param} ")) return ret diff --git a/mitmproxy/types.py b/mitmproxy/types.py index 1f1c503b8..24a1172b8 100644 --- a/mitmproxy/types.py +++ b/mitmproxy/types.py @@ -197,7 +197,7 @@ class _ArgType(_BaseType): return [] def parse(self, manager: "CommandManager", t: type, s: str) -> str: - raise exceptions.TypeError("Arguments for unknown command.") + return s def is_valid(self, manager: "CommandManager", typ: typing.Any, val: typing.Any) -> bool: return isinstance(val, str) diff --git a/test/mitmproxy/test_types.py b/test/mitmproxy/test_types.py index c8f7afde5..2cd17d87f 100644 --- a/test/mitmproxy/test_types.py +++ b/test/mitmproxy/test_types.py @@ -128,8 +128,7 @@ def test_arg(): with taddons.context() as tctx: b = mitmproxy.types._ArgType() assert b.completion(tctx.master.commands, mitmproxy.types.CmdArgs, "") == [] - with pytest.raises(mitmproxy.exceptions.TypeError): - b.parse(tctx.master.commands, mitmproxy.types.CmdArgs, "foo") + assert b.parse(tctx.master.commands, mitmproxy.types.CmdArgs, "foo") == "foo" assert b.is_valid(tctx.master.commands, mitmproxy.types.CmdArgs, 1) is False diff --git a/test/mitmproxy/tools/console/test_commander.py b/test/mitmproxy/tools/console/test_commander.py index 9a2ec1024..ce789d30c 100644 --- a/test/mitmproxy/tools/console/test_commander.py +++ b/test/mitmproxy/tools/console/test_commander.py @@ -25,7 +25,7 @@ class TestListCompleter: for start, options, cycle in tests: c = commander.ListCompleter(start, options) for expected in cycle: - assert c.cycle() == expected + assert c.cycle(True) == expected class TestCommandEdit: @@ -252,7 +252,7 @@ class TestCommandBuffer: cb = commander.CommandBuffer(tctx.master) cb.text = "foo bar" cb.cursor = len(cb.text) - cb.cycle_completion() + cb.cycle_completion(True) ch = commander.CommandHistory(tctx.master, 30) ce = commander.CommandEdit(tctx.master, "se", ch) @@ -261,7 +261,7 @@ class TestCommandBuffer: ret = ce.cbuf.render() assert ret[0] == ('commander_command', 'set') assert ret[1] == ('text', ' ') - assert ret[2] == ('commander_hint', 'str ') + assert ret[2] == ('commander_hint', '*options ') def test_render(self): with taddons.context() as tctx: @@ -272,13 +272,13 @@ class TestCommandBuffer: cb.text = 'set view_filter=~bq test' ret = cb.render() assert ret[0] == ('commander_command', 'set') - assert ret[1] == ('commander_invalid', ' ') + assert ret[1] == ('text', ' ') assert ret[2] == ('text', 'view_filter=~bq') - assert ret[3] == ('commander_invalid', ' ') - assert ret[4] == ('commander_invalid', 'test') + assert ret[3] == ('text', ' ') + assert ret[4] == ('text', 'test') cb.text = "set" ret = cb.render() assert ret[0] == ('commander_command', 'set') assert ret[1] == ('text', ' ') - assert ret[2] == ('commander_hint', 'str ') + assert ret[2] == ('commander_hint', '*options ') diff --git a/test/mitmproxy/tools/console/test_defaultkeys.py b/test/mitmproxy/tools/console/test_defaultkeys.py index 9c79525ba..58a0a585d 100644 --- a/test/mitmproxy/tools/console/test_defaultkeys.py +++ b/test/mitmproxy/tools/console/test_defaultkeys.py @@ -1,10 +1,12 @@ +import pytest + +import mitmproxy.types +from mitmproxy import command +from mitmproxy import ctx 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 import command -from mitmproxy import ctx -import pytest @pytest.mark.asyncio @@ -18,10 +20,13 @@ async def test_commands_exist(): await m.load_flow(tflow()) for binding in km.bindings: - results = command_manager.parse_partial(binding.command.strip()) + parsed, _ = command_manager.parse_partial(binding.command.strip()) - cmd = results[0][0].value - args = [a.value for a in results[0][1:]] + cmd = parsed[0].value + args = [ + a.value for a in parsed[1:] + if a.type != mitmproxy.types.Space + ] assert cmd in m.commands.commands