Grasscutters/mitmproxy
3
0
Fork 0
mirror of https://github.com/Grasscutters/mitmproxy.git synced 2025-02-02 00:05:27 +00:00
Code Issues Projects Releases Wiki Activity

types: cleanups and minor refactorings

Browse Source
This commit is contained in:
Aldo Cortesi 2017-12-19 07:38:21 +13:00 committed by Aldo Cortesi
parent 9faad6bc9c
commit 38b37ba7f5
5 changed files with 80 additions and 76 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mitmproxy/command.py
View File

@ -114,7 +114,7 @@ ParseResult = typing.NamedTuple(
)
class CommandManager:
class CommandManager(mitmproxy.types._CommandBase):
def __init__(self, master):
self.master = master
self.commands = {}

12
mitmproxy/tools/console/consoleaddons.py
View File

@ -104,7 +104,7 @@ class ConsoleAddon:
@command.command("console.layout.options")
def layout_options(self) -> typing.Sequence[str]:
"""
Returns the available options for the consoler_layout option.
Returns the available options for the console_layout option.
"""
return ["single", "vertical", "horizontal"]
@ -243,7 +243,11 @@ class ConsoleAddon:
@command.command("console.choose.cmd")
def console_choose_cmd(
self, prompt: str, choicecmd: mitmproxy.types.Cmd, *cmd: mitmproxy.types.Arg
self,
prompt: str,
choicecmd: mitmproxy.types.Cmd,
subcmd: mitmproxy.types.Cmd,
*args: mitmproxy.types.Arg
) -> None:
"""
Prompt the user to choose from a list of strings returned by a
@ -254,10 +258,10 @@ class ConsoleAddon:
def callback(opt):
# We're now outside of the call context...
repl = " ".join(cmd)
repl = " ".join(args)
repl = repl.replace("{choice}", opt)
try:
self.master.commands.call(repl)
self.master.commands.call(subcmd + " " + repl)
except exceptions.CommandError as e:
signals.status_message.send(message=str(e))

2
mitmproxy/tools/console/defaultkeys.py
View File

@ -41,7 +41,7 @@ def map(km):
"e",
"""
console.choose.cmd Format export.formats
console.command export.file {choice} @focus ''
console.command export.file {choice} @focus
""",
["flowlist", "flowview"],
"Export this flow to file"

116
mitmproxy/types.py
View File

@ -40,39 +40,39 @@ class Choice:
# annotations can cause circular dependencies where there were none before.
# Rather than putting types and the CommandManger in the same file, we introduce
# a stub type with the signature we use.
class _CommandStub:
commands = {} # type: typing.Mapping[str, typing.Any]
class _CommandBase:
commands = {} # type: typing.MutableMapping[str, typing.Any]
def call_args(self, path: str, args: typing.Sequence[str]) -> typing.Any: # pragma: no cover
pass
def call_args(self, path: str, args: typing.Sequence[str]) -> typing.Any:
raise NotImplementedError
def call(self, args: typing.Sequence[str]) -> typing.Any: # pragma: no cover
pass
def call(self, cmd: str) -> typing.Any:
raise NotImplementedError
class BaseType:
class _BaseType:
typ = object # type: typing.Type
display = "" # type: str
def completion(
self, manager: _CommandStub, t: type, s: str
self, manager: _CommandBase, t: typing.Any, s: str
) -> typing.Sequence[str]: # pragma: no cover
pass
def parse(
self, manager: _CommandStub, t: type, s: str
self, manager: _CommandBase, t: typing.Any, s: str
) -> typing.Any: # pragma: no cover
pass
class Bool(BaseType):
class _BoolType(_BaseType):
typ = bool
display = "bool"
def completion(self, manager: _CommandStub, t: type, s: str) -> typing.Sequence[str]:
def completion(self, manager: _CommandBase, t: type, s: str) -> typing.Sequence[str]:
return ["false", "true"]
def parse(self, manager: _CommandStub, t: type, s: str) -> bool:
def parse(self, manager: _CommandBase, t: type, s: str) -> bool:
if s == "true":
return True
elif s == "false":
@ -83,36 +83,36 @@ class Bool(BaseType):
)
class Str(BaseType):
class _StrType(_BaseType):
typ = str
display = "str"
def completion(self, manager: _CommandStub, t: type, s: str) -> typing.Sequence[str]:
def completion(self, manager: _CommandBase, t: type, s: str) -> typing.Sequence[str]:
return []
def parse(self, manager: _CommandStub, t: type, s: str) -> str:
def parse(self, manager: _CommandBase, t: type, s: str) -> str:
return s
class Int(BaseType):
class _IntType(_BaseType):
typ = int
display = "int"
def completion(self, manager: _CommandStub, t: type, s: str) -> typing.Sequence[str]:
def completion(self, manager: _CommandBase, t: type, s: str) -> typing.Sequence[str]:
return []
def parse(self, manager: _CommandStub, t: type, s: str) -> int:
def parse(self, manager: _CommandBase, t: type, s: str) -> int:
try:
return int(s)
except ValueError as e:
raise exceptions.TypeError from e
class PathType(BaseType):
class _PathType(_BaseType):
typ = Path
display = "path"
def completion(self, manager: _CommandStub, t: type, start: str) -> typing.Sequence[str]:
def completion(self, manager: _CommandBase, t: type, start: str) -> typing.Sequence[str]:
if not start:
start = "./"
path = os.path.expanduser(start)
@ -134,44 +134,44 @@ class PathType(BaseType):
ret.sort()
return ret
def parse(self, manager: _CommandStub, t: type, s: str) -> str:
def parse(self, manager: _CommandBase, t: type, s: str) -> str:
return s
class CmdType(BaseType):
class _CmdType(_BaseType):
typ = Cmd
display = "cmd"
def completion(self, manager: _CommandStub, t: type, s: str) -> typing.Sequence[str]:
def completion(self, manager: _CommandBase, t: type, s: str) -> typing.Sequence[str]:
return list(manager.commands.keys())
def parse(self, manager: _CommandStub, t: type, s: str) -> str:
def parse(self, manager: _CommandBase, t: type, s: str) -> str:
return s
class ArgType(BaseType):
class _ArgType(_BaseType):
typ = Arg
display = "arg"
def completion(self, manager: _CommandStub, t: type, s: str) -> typing.Sequence[str]:
def completion(self, manager: _CommandBase, t: type, s: str) -> typing.Sequence[str]:
return []
def parse(self, manager: _CommandStub, t: type, s: str) -> str:
def parse(self, manager: _CommandBase, t: type, s: str) -> str:
return s
class StrSeq(BaseType):
class _StrSeqType(_BaseType):
typ = typing.Sequence[str]
display = "[str]"
def completion(self, manager: _CommandStub, t: type, s: str) -> typing.Sequence[str]:
def completion(self, manager: _CommandBase, t: type, s: str) -> typing.Sequence[str]:
return []
def parse(self, manager: _CommandStub, t: type, s: str) -> typing.Sequence[str]:
def parse(self, manager: _CommandBase, t: type, s: str) -> typing.Sequence[str]:
return [x.strip() for x in s.split(",")]
class CutSpecType(BaseType):
class _CutSpecType(_BaseType):
typ = CutSpec
display = "[cut]"
valid_prefixes = [
@ -212,7 +212,7 @@ class CutSpecType(BaseType):
"server_conn.ssl_established",
]
def completion(self, manager: _CommandStub, t: type, s: str) -> typing.Sequence[str]:
def completion(self, manager: _CommandBase, t: type, s: str) -> typing.Sequence[str]:
spec = s.split(",")
opts = []
for pref in self.valid_prefixes:
@ -220,12 +220,12 @@ class CutSpecType(BaseType):
opts.append(",".join(spec))
return opts
def parse(self, manager: _CommandStub, t: type, s: str) -> CutSpec:
def parse(self, manager: _CommandBase, t: type, s: str) -> CutSpec:
parts = s.split(",") # type: typing.Any
return parts
class BaseFlowType(BaseType):
class _BaseFlowType(_BaseType):
valid_prefixes = [
"@all",
"@focus",
@ -248,15 +248,15 @@ class BaseFlowType(BaseType):
"~c",
]
def completion(self, manager: _CommandStub, t: type, s: str) -> typing.Sequence[str]:
def completion(self, manager: _CommandBase, t: type, s: str) -> typing.Sequence[str]:
return self.valid_prefixes
class FlowType(BaseFlowType):
class _FlowType(_BaseFlowType):
typ = flow.Flow
display = "flow"
def parse(self, manager: _CommandStub, t: type, s: str) -> flow.Flow:
def parse(self, manager: _CommandBase, t: type, s: str) -> flow.Flow:
flows = manager.call_args("view.resolve", [s])
if len(flows) != 1:
raise exceptions.TypeError(
@ -265,37 +265,37 @@ class FlowType(BaseFlowType):
return flows[0]
class FlowsType(BaseFlowType):
class _FlowsType(_BaseFlowType):
typ = typing.Sequence[flow.Flow]
display = "[flow]"
def parse(self, manager: _CommandStub, t: type, s: str) -> typing.Sequence[flow.Flow]:
def parse(self, manager: _CommandBase, t: type, s: str) -> typing.Sequence[flow.Flow]:
return manager.call_args("view.resolve", [s])
class DataType:
class _DataType(_BaseType):
typ = Data
display = "[data]"
def completion(
self, manager: _CommandStub, t: type, s: str
self, manager: _CommandBase, t: type, s: str
) -> typing.Sequence[str]: # pragma: no cover
raise exceptions.TypeError("data cannot be passed as argument")
def parse(
self, manager: _CommandStub, t: type, s: str
self, manager: _CommandBase, t: type, s: str
) -> typing.Any: # pragma: no cover
raise exceptions.TypeError("data cannot be passed as argument")
class ChoiceType:
class _ChoiceType(_BaseType):
typ = Choice
display = "choice"
def completion(self, manager: _CommandStub, t: Choice, s: str) -> typing.Sequence[str]:
def completion(self, manager: _CommandBase, t: Choice, s: str) -> typing.Sequence[str]:
return manager.call(t.options_command)
def parse(self, manager: _CommandStub, t: Choice, s: str) -> str:
def parse(self, manager: _CommandBase, t: Choice, s: str) -> str:
opts = manager.call(t.options_command)
if s not in opts:
raise exceptions.TypeError("Invalid choice.")
@ -308,23 +308,23 @@ class TypeManager:
for t in types:
self.typemap[t.typ] = t()
def get(self, t: type, default=None) -> BaseType:
def get(self, t: type, default=None) -> _BaseType:
if type(t) in self.typemap:
return self.typemap[type(t)]
return self.typemap.get(t, default)
CommandTypes = TypeManager(
ArgType,
Bool,
ChoiceType,
CmdType,
CutSpecType,
DataType,
FlowType,
FlowsType,
Int,
PathType,
Str,
StrSeq,
_ArgType,
_BoolType,
_ChoiceType,
_CmdType,
_CutSpecType,
_DataType,
_FlowType,
_FlowsType,
_IntType,
_PathType,
_StrType,
_StrSeqType,
)

24
test/mitmproxy/test_types.py
View File

@ -24,7 +24,7 @@ def chdir(path: str):
def test_bool():
with taddons.context() as tctx:
b = mitmproxy.types.Bool()
b = mitmproxy.types._BoolType()
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
@ -34,14 +34,14 @@ def test_bool():
def test_str():
with taddons.context() as tctx:
b = mitmproxy.types.Str()
b = mitmproxy.types._StrType()
assert b.completion(tctx.master.commands, str, "") == []
assert b.parse(tctx.master.commands, str, "foo") == "foo"
def test_int():
with taddons.context() as tctx:
b = mitmproxy.types.Int()
b = mitmproxy.types._IntType()
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
@ -51,7 +51,7 @@ def test_int():
def test_path():
with taddons.context() as tctx:
b = mitmproxy.types.PathType()
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"
@ -77,7 +77,7 @@ def test_path():
def test_cmd():
with taddons.context() as tctx:
tctx.master.addons.add(test_command.TAddon())
b = mitmproxy.types.CmdType()
b = mitmproxy.types._CmdType()
assert b.parse(tctx.master.commands, mitmproxy.types.Cmd, "foo") == "foo"
assert len(
b.completion(tctx.master.commands, mitmproxy.types.Cmd, "")
@ -86,7 +86,7 @@ def test_cmd():
def test_cutspec():
with taddons.context() as tctx:
b = mitmproxy.types.CutSpecType()
b = mitmproxy.types._CutSpecType()
b.parse(tctx.master.commands, mitmproxy.types.CutSpec, "foo,bar") == ["foo", "bar"]
assert b.completion(
tctx.master.commands, mitmproxy.types.CutSpec, "request.p"
@ -98,14 +98,14 @@ def test_cutspec():
def test_arg():
with taddons.context() as tctx:
b = mitmproxy.types.ArgType()
b = mitmproxy.types._ArgType()
assert b.completion(tctx.master.commands, mitmproxy.types.Arg, "") == []
assert b.parse(tctx.master.commands, mitmproxy.types.Arg, "foo") == "foo"
def test_strseq():
with taddons.context() as tctx:
b = mitmproxy.types.StrSeq()
b = mitmproxy.types._StrSeqType()
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"]
@ -129,7 +129,7 @@ class DummyConsole:
def test_flow():
with taddons.context() as tctx:
tctx.master.addons.add(DummyConsole())
b = mitmproxy.types.FlowType()
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")
with pytest.raises(mitmproxy.exceptions.TypeError):
@ -141,7 +141,7 @@ def test_flow():
def test_flows():
with taddons.context() as tctx:
tctx.master.addons.add(DummyConsole())
b = mitmproxy.types.FlowsType()
b = mitmproxy.types._FlowsType()
assert len(
b.completion(tctx.master.commands, typing.Sequence[flow.Flow], "")
) == len(b.valid_prefixes)
@ -152,7 +152,7 @@ def test_flows():
def test_data():
with taddons.context() as tctx:
b = mitmproxy.types.DataType()
b = mitmproxy.types._DataType()
with pytest.raises(mitmproxy.exceptions.TypeError):
b.parse(tctx.master.commands, mitmproxy.types.Data, "foo")
with pytest.raises(mitmproxy.exceptions.TypeError):
@ -162,7 +162,7 @@ def test_data():
def test_choice():
with taddons.context() as tctx:
tctx.master.addons.add(DummyConsole())
b = mitmproxy.types.ChoiceType()
b = mitmproxy.types._ChoiceType()
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"