From 9a469806eb6a202f5b35c011368fd240a16f5397 Mon Sep 17 00:00:00 2001 From: shindexro Date: Fri, 19 Nov 2021 12:04:20 +0000 Subject: [PATCH] quote argument of view.flows.resolve (#4910) * Fix #4902 * Update type signature * Switch to None check * Fix spacing * Quote view.flows.resolve argument * Switch to call_strings --- mitmproxy/types.py | 4 ++-- test/mitmproxy/test_types.py | 7 ++++++- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/mitmproxy/types.py b/mitmproxy/types.py index 01d7ac47f..b4b17d80e 100644 --- a/mitmproxy/types.py +++ b/mitmproxy/types.py @@ -369,7 +369,7 @@ class _FlowType(_BaseFlowType): def parse(self, manager: "CommandManager", t: type, s: str) -> flow.Flow: try: - flows = manager.execute("view.flows.resolve %s" % (s)) + flows = manager.call_strings("view.flows.resolve", [s]) except exceptions.CommandError as e: raise exceptions.TypeError(str(e)) from e if len(flows) != 1: @@ -388,7 +388,7 @@ class _FlowsType(_BaseFlowType): def parse(self, manager: "CommandManager", t: type, s: str) -> typing.Sequence[flow.Flow]: try: - return manager.execute("view.flows.resolve %s" % (s)) + return manager.call_strings("view.flows.resolve", [s]) except exceptions.CommandError as e: raise exceptions.TypeError(str(e)) from e diff --git a/test/mitmproxy/test_types.py b/test/mitmproxy/test_types.py index e0f56995f..7b7d203cd 100644 --- a/test/mitmproxy/test_types.py +++ b/test/mitmproxy/test_types.py @@ -183,7 +183,10 @@ class DummyConsole: def resolve(self, spec: str) -> typing.Sequence[flow.Flow]: if spec == "err": raise mitmproxy.exceptions.CommandError() - n = int(spec) + try: + n = int(spec) + except ValueError: + n = 1 return [tflow.tflow(resp=True)] * n @command.command("cut") @@ -201,6 +204,7 @@ def test_flow(): b = mitmproxy.types._FlowType() assert len(b.completion(tctx.master.commands, flow.Flow, "")) == len(b.valid_prefixes) assert b.parse(tctx.master.commands, flow.Flow, "1") + assert b.parse(tctx.master.commands, flow.Flow, "has space") assert b.is_valid(tctx.master.commands, flow.Flow, tflow.tflow()) is True assert b.is_valid(tctx.master.commands, flow.Flow, "xx") is False with pytest.raises(mitmproxy.exceptions.TypeError): @@ -224,6 +228,7 @@ def test_flows(): assert len(b.parse(tctx.master.commands, typing.Sequence[flow.Flow], "0")) == 0 assert len(b.parse(tctx.master.commands, typing.Sequence[flow.Flow], "1")) == 1 assert len(b.parse(tctx.master.commands, typing.Sequence[flow.Flow], "2")) == 2 + assert len(b.parse(tctx.master.commands, typing.Sequence[flow.Flow], "has space")) == 1 with pytest.raises(mitmproxy.exceptions.TypeError): b.parse(tctx.master.commands, typing.Sequence[flow.Flow], "err")