From 7779eef572e8deeae895ea6d700265e6f9b432c8 Mon Sep 17 00:00:00 2001 From: Henrique Date: Sat, 16 Nov 2019 17:01:47 -0500 Subject: [PATCH] Various changes to address PR comments Made a change to make `CommandManager.execute` the main entry point for executing commands and made `call_strings` into a private method. --- mitmproxy/command.py | 15 ++++++++------- mitmproxy/tools/console/consoleaddons.py | 19 ++++++------------- mitmproxy/types.py | 6 +++--- test/mitmproxy/addons/test_save.py | 2 +- .../tools/console/test_defaultkeys.py | 2 +- 5 files changed, 19 insertions(+), 25 deletions(-) diff --git a/mitmproxy/command.py b/mitmproxy/command.py index 32a9c9c1a..609e288c4 100644 --- a/mitmproxy/command.py +++ b/mitmproxy/command.py @@ -128,10 +128,10 @@ class CommandManager(mitmproxy.types._CommandBase): self.commands: typing.Dict[str, Command] = {} self.regex = 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(" \r\n\t") + pyparsing.QuotedString("'", escChar='\\', unquoteResults=False) |\ + pyparsing.Combine(pyparsing.Literal('"') + pyparsing.Word(pyparsing.printables + " ") + pyparsing.StringEnd()) |\ + pyparsing.Word(pyparsing.printables) |\ + pyparsing.Word(" \r\n\t") self.regex = self.regex.leaveWhitespace() def collect_commands(self, addon): @@ -220,7 +220,7 @@ class CommandManager(mitmproxy.types._CommandBase): raise exceptions.CommandError("Unknown command: %s" % path) return self.commands[path].func(*args) - def call_strings(self, path: str, args: typing.Sequence[str]) -> typing.Any: + def _call_strings(self, path: str, args: typing.Sequence[str]) -> typing.Any: """ Call a command using a list of string arguments. May raise CommandError. """ @@ -236,12 +236,13 @@ class CommandManager(mitmproxy.types._CommandBase): parts, _ = self.parse_partial(cmdstr) params = [] for p in parts: - params.append(p.value) + if p.value.strip() != '': + params.append(p.value) if len(parts) == 0: raise exceptions.CommandError("Invalid command: %s" % cmdstr) - return self.call_strings(params[0], params[1:]) + return self._call_strings(params[0], params[1:]) def dump(self, out=sys.stdout) -> None: cmds = list(self.commands.values()) diff --git a/mitmproxy/tools/console/consoleaddons.py b/mitmproxy/tools/console/consoleaddons.py index 9f595b422..967c2a352 100644 --- a/mitmproxy/tools/console/consoleaddons.py +++ b/mitmproxy/tools/console/consoleaddons.py @@ -271,7 +271,7 @@ class ConsoleAddon: command, then invoke another command with all occurrences of {choice} replaced by the choice the user made. """ - choices = ctx.master.commands.call_strings(choicecmd, []) + choices = ctx.master.commands.execute(choicecmd) def callback(opt): # We're now outside of the call context... @@ -535,10 +535,8 @@ class ConsoleAddon: raise exceptions.CommandError("Invalid flowview mode.") try: - self.master.commands.call_strings( - "view.settings.setval", - ["@focus", "flowview_mode_%s" % idx, mode] - ) + cmd = 'view.settings.setval @focus flowview_mode_%s %s' % (idx, mode) + self.master.commands.execute(cmd) except exceptions.CommandError as e: signals.status_message.send(message=str(e)) @@ -558,14 +556,9 @@ class ConsoleAddon: if not fv: raise exceptions.CommandError("Not viewing a flow.") idx = fv.body.tab_offset - return self.master.commands.call_strings( - "view.settings.getval", - [ - "@focus", - "flowview_mode_%s" % idx, - self.master.options.console_default_contentview, - ] - ) + + cmd = 'view.settings.getval @focus flowview_mode_%s %s' % (idx, self.master.options.console_default_contentview) + return self.master.commands.execute(cmd) @command.command("console.key.contexts") def key_contexts(self) -> typing.Sequence[str]: diff --git a/mitmproxy/types.py b/mitmproxy/types.py index 0634e4d77..b48aef846 100644 --- a/mitmproxy/types.py +++ b/mitmproxy/types.py @@ -47,7 +47,7 @@ class Choice: class _CommandBase: commands: typing.MutableMapping[str, typing.Any] = {} - def call_strings(self, path: str, args: typing.Sequence[str]) -> typing.Any: + def _call_strings(self, path: str, args: typing.Sequence[str]) -> typing.Any: raise NotImplementedError def execute(self, cmd: str) -> typing.Any: @@ -337,7 +337,7 @@ class _FlowType(_BaseFlowType): def parse(self, manager: _CommandBase, t: type, s: str) -> flow.Flow: try: - flows = manager.call_strings("view.flows.resolve", [s]) + flows = manager.execute("view.flows.resolve %s" % (s)) except exceptions.CommandError as e: raise exceptions.TypeError from e if len(flows) != 1: @@ -356,7 +356,7 @@ class _FlowsType(_BaseFlowType): def parse(self, manager: _CommandBase, t: type, s: str) -> typing.Sequence[flow.Flow]: try: - return manager.call_strings("view.flows.resolve", [s]) + return manager.execute("view.flows.resolve %s" % (s)) except exceptions.CommandError as e: raise exceptions.TypeError from e diff --git a/test/mitmproxy/addons/test_save.py b/test/mitmproxy/addons/test_save.py index 4aa1f6488..6727a96f7 100644 --- a/test/mitmproxy/addons/test_save.py +++ b/test/mitmproxy/addons/test_save.py @@ -73,7 +73,7 @@ def test_save_command(tmpdir): v = view.View() tctx.master.addons.add(v) tctx.master.addons.add(sa) - tctx.master.commands.call_strings("save.file", ["@shown", p]) + tctx.master.commands.execute("save.file @shown %s" % p) def test_simple(tmpdir): diff --git a/test/mitmproxy/tools/console/test_defaultkeys.py b/test/mitmproxy/tools/console/test_defaultkeys.py index 40e536b03..9c79525ba 100644 --- a/test/mitmproxy/tools/console/test_defaultkeys.py +++ b/test/mitmproxy/tools/console/test_defaultkeys.py @@ -18,7 +18,7 @@ async def test_commands_exist(): await m.load_flow(tflow()) for binding in km.bindings: - results = command_manager.parse_partial(binding.command) + results = command_manager.parse_partial(binding.command.strip()) cmd = results[0][0].value args = [a.value for a in results[0][1:]]