diff options
| -rw-r--r-- | mitmproxy/types.py | 93 | ||||
| -rw-r--r-- | test/mitmproxy/test_types.py | 42 |
2 files changed, 129 insertions, 6 deletions
diff --git a/mitmproxy/types.py b/mitmproxy/types.py index b6b414ba..713a0ae5 100644 --- a/mitmproxy/types.py +++ b/mitmproxy/types.py @@ -56,13 +56,29 @@ class _BaseType: def completion( self, manager: _CommandBase, t: typing.Any, s: str - ) -> typing.Sequence[str]: # pragma: no cover - pass + ) -> typing.Sequence[str]: + """ + Returns a list of completion strings for a given prefix. The strings + returned don't necessarily need to be suffixes of the prefix, since + completers will do prefix filtering themselves.. + """ + raise NotImplementedError def parse( - self, manager: _CommandBase, t: typing.Any, s: str - ) -> typing.Any: # pragma: no cover - pass + self, manager: _CommandBase, typ: typing.Any, s: str + ) -> typing.Any: + """ + Parse a string, given the specific type instance (to allow rich type annotations like Choice) and a string. + + Raises exceptions.TypeError if the value is invalid. + """ + raise NotImplementedError + + def is_valid(self, manager: _CommandBase, typ: typing.Any, val: typing.Any) -> bool: + """ + Check if data is valid for this type. + """ + raise NotImplementedError class _BoolType(_BaseType): @@ -82,6 +98,9 @@ class _BoolType(_BaseType): "Booleans are 'true' or 'false', got %s" % s ) + def is_valid(self, manager: _CommandBase, typ: typing.Any, val: typing.Any) -> bool: + return val in [True, False] + class _StrType(_BaseType): typ = str @@ -93,6 +112,9 @@ class _StrType(_BaseType): def parse(self, manager: _CommandBase, t: type, s: str) -> str: return s + def is_valid(self, manager: _CommandBase, typ: typing.Any, val: typing.Any) -> bool: + return isinstance(val, str) + class _IntType(_BaseType): typ = int @@ -107,6 +129,9 @@ class _IntType(_BaseType): except ValueError as e: raise exceptions.TypeError from e + def is_valid(self, manager: _CommandBase, typ: typing.Any, val: typing.Any) -> bool: + return isinstance(val, int) + class _PathType(_BaseType): typ = Path @@ -137,6 +162,9 @@ class _PathType(_BaseType): def parse(self, manager: _CommandBase, t: type, s: str) -> str: return s + def is_valid(self, manager: _CommandBase, typ: typing.Any, val: typing.Any) -> bool: + return isinstance(val, str) + class _CmdType(_BaseType): typ = Cmd @@ -148,6 +176,9 @@ class _CmdType(_BaseType): def parse(self, manager: _CommandBase, t: type, s: str) -> str: return s + def is_valid(self, manager: _CommandBase, typ: typing.Any, val: typing.Any) -> bool: + return val in manager.commands + class _ArgType(_BaseType): typ = Arg @@ -159,6 +190,9 @@ class _ArgType(_BaseType): def parse(self, manager: _CommandBase, t: type, s: str) -> str: return s + def is_valid(self, manager: _CommandBase, typ: typing.Any, val: typing.Any) -> bool: + return isinstance(val, str) + class _StrSeqType(_BaseType): typ = typing.Sequence[str] @@ -170,6 +204,9 @@ class _StrSeqType(_BaseType): def parse(self, manager: _CommandBase, t: type, s: str) -> typing.Sequence[str]: return [x.strip() for x in s.split(",")] + def is_valid(self, manager: _CommandBase, typ: typing.Any, val: typing.Any) -> bool: + return isinstance(val, str) + class _CutSpecType(_BaseType): typ = CutSpec @@ -224,15 +261,29 @@ class _CutSpecType(_BaseType): parts = s.split(",") # type: typing.Any return parts + def is_valid(self, manager: _CommandBase, typ: typing.Any, val: typing.Any) -> bool: + if not isinstance(val, str): + return False + parts = [x.strip() for x in val.split(",")] + for p in parts: + for pref in self.valid_prefixes: + if p.startswith(pref): + break + else: + return False + return True + class _BaseFlowType(_BaseType): - valid_prefixes = [ + viewmarkers = [ "@all", "@focus", "@shown", "@hidden", "@marked", "@unmarked", + ] + valid_prefixes = viewmarkers + [ "~q", "~s", "~a", @@ -264,6 +315,9 @@ class _FlowType(_BaseFlowType): ) return flows[0] + def is_valid(self, manager: _CommandBase, typ: typing.Any, val: typing.Any) -> bool: + return isinstance(val, flow.Flow) + class _FlowsType(_BaseFlowType): typ = typing.Sequence[flow.Flow] @@ -272,6 +326,15 @@ class _FlowsType(_BaseFlowType): def parse(self, manager: _CommandBase, t: type, s: str) -> typing.Sequence[flow.Flow]: return manager.call_args("view.resolve", [s]) + def is_valid(self, manager: _CommandBase, typ: typing.Any, val: typing.Any) -> bool: + try: + for v in val: + if not isinstance(v, flow.Flow): + return False + except TypeError: + return False + return True + class _DataType(_BaseType): typ = Data @@ -287,6 +350,17 @@ class _DataType(_BaseType): ) -> typing.Any: # pragma: no cover raise exceptions.TypeError("data cannot be passed as argument") + def is_valid(self, manager: _CommandBase, typ: typing.Any, val: typing.Any) -> bool: + # FIXME: validate that all rows have equal length, and all columns have equal types + try: + for row in val: + for cell in row: + if not (isinstance(cell, str) or isinstance(cell, bytes)): + return False + except TypeError: + return False + return True + class _ChoiceType(_BaseType): typ = Choice @@ -301,6 +375,13 @@ class _ChoiceType(_BaseType): raise exceptions.TypeError("Invalid choice.") return s + def is_valid(self, manager: _CommandBase, typ: typing.Any, val: typing.Any) -> bool: + try: + opts = manager.call(typ.options_command) + except exceptions.CommandError: + return False + return val in opts + class TypeManager: def __init__(self, *types): diff --git a/test/mitmproxy/test_types.py b/test/mitmproxy/test_types.py index 90032204..dcf7a1d2 100644 --- a/test/mitmproxy/test_types.py +++ b/test/mitmproxy/test_types.py @@ -28,6 +28,8 @@ def test_bool(): assert b.completion(tctx.master.commands, bool, "b") == ["false", "true"] assert b.parse(tctx.master.commands, bool, "true") is True assert b.parse(tctx.master.commands, bool, "false") is False + assert b.is_valid(tctx.master.commands, bool, True) is True + assert b.is_valid(tctx.master.commands, bool, "foo") is False with pytest.raises(mitmproxy.exceptions.TypeError): b.parse(tctx.master.commands, bool, "foo") @@ -35,6 +37,8 @@ def test_bool(): def test_str(): with taddons.context() as tctx: b = mitmproxy.types._StrType() + assert b.is_valid(tctx.master.commands, str, "foo") is True + assert b.is_valid(tctx.master.commands, str, 1) is False assert b.completion(tctx.master.commands, str, "") == [] assert b.parse(tctx.master.commands, str, "foo") == "foo" @@ -42,6 +46,8 @@ def test_str(): def test_int(): with taddons.context() as tctx: b = mitmproxy.types._IntType() + assert b.is_valid(tctx.master.commands, int, "foo") is False + assert b.is_valid(tctx.master.commands, int, 1) is True assert b.completion(tctx.master.commands, int, "b") == [] assert b.parse(tctx.master.commands, int, "1") == 1 assert b.parse(tctx.master.commands, int, "999") == 999 @@ -54,6 +60,8 @@ def test_path(): b = mitmproxy.types._PathType() assert b.parse(tctx.master.commands, mitmproxy.types.Path, "/foo") == "/foo" assert b.parse(tctx.master.commands, mitmproxy.types.Path, "/bar") == "/bar" + assert b.is_valid(tctx.master.commands, mitmproxy.types.Path, "foo") is True + assert b.is_valid(tctx.master.commands, mitmproxy.types.Path, 3) is False def normPathOpts(prefix, match): ret = [] @@ -78,6 +86,8 @@ def test_cmd(): with taddons.context() as tctx: tctx.master.addons.add(test_command.TAddon()) b = mitmproxy.types._CmdType() + assert b.is_valid(tctx.master.commands, mitmproxy.types.Cmd, "foo") is False + assert b.is_valid(tctx.master.commands, mitmproxy.types.Cmd, "cmd1") is True assert b.parse(tctx.master.commands, mitmproxy.types.Cmd, "foo") == "foo" assert len( b.completion(tctx.master.commands, mitmproxy.types.Cmd, "") @@ -88,6 +98,10 @@ def test_cutspec(): with taddons.context() as tctx: b = mitmproxy.types._CutSpecType() b.parse(tctx.master.commands, mitmproxy.types.CutSpec, "foo,bar") == ["foo", "bar"] + assert b.is_valid(tctx.master.commands, mitmproxy.types.CutSpec, 1) is False + assert b.is_valid(tctx.master.commands, mitmproxy.types.CutSpec, "foo") is False + assert b.is_valid(tctx.master.commands, mitmproxy.types.CutSpec, "request.path") is True + assert b.completion( tctx.master.commands, mitmproxy.types.CutSpec, "request.p" ) == b.valid_prefixes @@ -101,6 +115,8 @@ def test_arg(): b = mitmproxy.types._ArgType() assert b.completion(tctx.master.commands, mitmproxy.types.Arg, "") == [] assert b.parse(tctx.master.commands, mitmproxy.types.Arg, "foo") == "foo" + assert b.is_valid(tctx.master.commands, mitmproxy.types.Arg, "foo") is True + assert b.is_valid(tctx.master.commands, mitmproxy.types.Arg, 1) is False def test_strseq(): @@ -109,6 +125,8 @@ def test_strseq(): assert b.completion(tctx.master.commands, typing.Sequence[str], "") == [] assert b.parse(tctx.master.commands, typing.Sequence[str], "foo") == ["foo"] assert b.parse(tctx.master.commands, typing.Sequence[str], "foo,bar") == ["foo", "bar"] + assert b.is_valid(tctx.master.commands, typing.Sequence[str], "foo") is True + assert b.is_valid(tctx.master.commands, typing.Sequence[str], 1) is False class DummyConsole: @@ -132,6 +150,8 @@ 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.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): assert b.parse(tctx.master.commands, flow.Flow, "0") with pytest.raises(mitmproxy.exceptions.TypeError): @@ -145,6 +165,9 @@ def test_flows(): assert len( b.completion(tctx.master.commands, typing.Sequence[flow.Flow], "") ) == len(b.valid_prefixes) + assert b.is_valid(tctx.master.commands, typing.Sequence[flow.Flow], [tflow.tflow()]) is True + assert b.is_valid(tctx.master.commands, typing.Sequence[flow.Flow], "xx") is False + assert b.is_valid(tctx.master.commands, typing.Sequence[flow.Flow], 0) is False 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 @@ -153,6 +176,10 @@ def test_flows(): def test_data(): with taddons.context() as tctx: b = mitmproxy.types._DataType() + assert b.is_valid(tctx.master.commands, mitmproxy.types.Data, 0) is False + assert b.is_valid(tctx.master.commands, mitmproxy.types.Data, []) is True + assert b.is_valid(tctx.master.commands, mitmproxy.types.Data, [["x"]]) is True + assert b.is_valid(tctx.master.commands, mitmproxy.types.Data, [[b"x"]]) is True with pytest.raises(mitmproxy.exceptions.TypeError): b.parse(tctx.master.commands, mitmproxy.types.Data, "foo") with pytest.raises(mitmproxy.exceptions.TypeError): @@ -163,6 +190,21 @@ def test_choice(): with taddons.context() as tctx: tctx.master.addons.add(DummyConsole()) b = mitmproxy.types._ChoiceType() + assert b.is_valid( + tctx.master.commands, + mitmproxy.types.Choice("options"), + "one", + ) is True + assert b.is_valid( + tctx.master.commands, + mitmproxy.types.Choice("options"), + "invalid", + ) is False + assert b.is_valid( + tctx.master.commands, + mitmproxy.types.Choice("nonexistent"), + "invalid", + ) is False comp = b.completion(tctx.master.commands, mitmproxy.types.Choice("options"), "") assert comp == ["one", "two", "three"] assert b.parse(tctx.master.commands, mitmproxy.types.Choice("options"), "one") == "one" |