mirror of
https://github.com/vee1e/mitmproxy.git
synced 2026-09-04 19:57:16 +00:00
400 lines
13 KiB
Python
400 lines
13 KiB
Python
import typing
|
|
import inspect
|
|
from mitmproxy import command
|
|
from mitmproxy import flow
|
|
from mitmproxy import exceptions
|
|
from mitmproxy.test import tflow
|
|
from mitmproxy.test import taddons
|
|
import mitmproxy.types
|
|
import io
|
|
import pytest
|
|
|
|
|
|
class TAddon:
|
|
@command.command("cmd1")
|
|
def cmd1(self, foo: str) -> str:
|
|
"""cmd1 help"""
|
|
return "ret " + foo
|
|
|
|
@command.command("cmd2")
|
|
def cmd2(self, foo: str) -> str:
|
|
return 99
|
|
|
|
@command.command("cmd3")
|
|
def cmd3(self, foo: int) -> int:
|
|
return foo
|
|
|
|
@command.command("cmd4")
|
|
def cmd4(self, a: int, b: str, c: mitmproxy.types.Path) -> str:
|
|
return "ok"
|
|
|
|
@command.command("subcommand")
|
|
def subcommand(self, cmd: mitmproxy.types.Cmd, *args: mitmproxy.types.Arg) -> str:
|
|
return "ok"
|
|
|
|
@command.command("empty")
|
|
def empty(self) -> None:
|
|
pass
|
|
|
|
@command.command("varargs")
|
|
def varargs(self, one: str, *var: str) -> typing.Sequence[str]:
|
|
return list(var)
|
|
|
|
def choices(self) -> typing.Sequence[str]:
|
|
return ["one", "two", "three"]
|
|
|
|
@command.argument("arg", type=mitmproxy.types.Choice("choices"))
|
|
def choose(self, arg: str) -> typing.Sequence[str]:
|
|
return ["one", "two", "three"]
|
|
|
|
@command.command("path")
|
|
def path(self, arg: mitmproxy.types.Path) -> None:
|
|
pass
|
|
|
|
@command.command("flow")
|
|
def flow(self, f: flow.Flow, s: str) -> None:
|
|
pass
|
|
|
|
|
|
class Unsupported:
|
|
pass
|
|
|
|
|
|
class TypeErrAddon:
|
|
@command.command("noret")
|
|
def noret(self):
|
|
pass
|
|
|
|
@command.command("invalidret")
|
|
def invalidret(self) -> Unsupported:
|
|
pass
|
|
|
|
@command.command("invalidarg")
|
|
def invalidarg(self, u: Unsupported):
|
|
pass
|
|
|
|
|
|
class TestCommand:
|
|
def test_typecheck(self):
|
|
with taddons.context(loadcore=False) as tctx:
|
|
cm = command.CommandManager(tctx.master)
|
|
a = TypeErrAddon()
|
|
command.Command(cm, "noret", a.noret)
|
|
with pytest.raises(exceptions.CommandError):
|
|
command.Command(cm, "invalidret", a.invalidret)
|
|
with pytest.raises(exceptions.CommandError):
|
|
command.Command(cm, "invalidarg", a.invalidarg)
|
|
|
|
def test_varargs(self):
|
|
with taddons.context() as tctx:
|
|
cm = command.CommandManager(tctx.master)
|
|
a = TAddon()
|
|
c = command.Command(cm, "varargs", a.varargs)
|
|
assert c.signature_help() == "varargs str *str -> [str]"
|
|
assert c.call(["one", "two", "three"]) == ["two", "three"]
|
|
with pytest.raises(exceptions.CommandError):
|
|
c.call(["one", "two", 3])
|
|
|
|
def test_call(self):
|
|
with taddons.context() as tctx:
|
|
cm = command.CommandManager(tctx.master)
|
|
a = TAddon()
|
|
c = command.Command(cm, "cmd.path", a.cmd1)
|
|
assert c.call(["foo"]) == "ret foo"
|
|
assert c.signature_help() == "cmd.path str -> str"
|
|
|
|
c = command.Command(cm, "cmd.two", a.cmd2)
|
|
with pytest.raises(exceptions.CommandError):
|
|
c.call(["foo"])
|
|
|
|
c = command.Command(cm, "cmd.three", a.cmd3)
|
|
assert c.call(["1"]) == 1
|
|
|
|
def test_parse_partial(self):
|
|
tests = [
|
|
[
|
|
"foo bar",
|
|
[
|
|
command.ParseResult(
|
|
value = "foo", type = mitmproxy.types.Cmd, valid = False
|
|
),
|
|
command.ParseResult(
|
|
value = "bar", type = mitmproxy.types.Unknown, valid = False
|
|
)
|
|
],
|
|
[],
|
|
],
|
|
[
|
|
"cmd1 'bar",
|
|
[
|
|
command.ParseResult(value = "cmd1", type = mitmproxy.types.Cmd, valid = True),
|
|
command.ParseResult(value = "'bar", type = str, valid = True)
|
|
],
|
|
[],
|
|
],
|
|
[
|
|
"a",
|
|
[command.ParseResult(value = "a", type = mitmproxy.types.Cmd, valid = False)],
|
|
[],
|
|
],
|
|
[
|
|
"",
|
|
[command.ParseResult(value = "", type = mitmproxy.types.Cmd, valid = False)],
|
|
[]
|
|
],
|
|
[
|
|
"cmd3 1",
|
|
[
|
|
command.ParseResult(value = "cmd3", type = mitmproxy.types.Cmd, valid = True),
|
|
command.ParseResult(value = "1", type = int, valid = True),
|
|
],
|
|
[]
|
|
],
|
|
[
|
|
"cmd3 ",
|
|
[
|
|
command.ParseResult(value = "cmd3", type = mitmproxy.types.Cmd, valid = True),
|
|
command.ParseResult(value = "", type = int, valid = False),
|
|
],
|
|
[]
|
|
],
|
|
[
|
|
"subcommand ",
|
|
[
|
|
command.ParseResult(
|
|
value = "subcommand", type = mitmproxy.types.Cmd, valid = True,
|
|
),
|
|
command.ParseResult(value = "", type = mitmproxy.types.Cmd, valid = False),
|
|
],
|
|
["arg"],
|
|
],
|
|
[
|
|
"subcommand cmd3 ",
|
|
[
|
|
command.ParseResult(value = "subcommand", type = mitmproxy.types.Cmd, valid = True),
|
|
command.ParseResult(value = "cmd3", type = mitmproxy.types.Cmd, valid = True),
|
|
command.ParseResult(value = "", type = int, valid = False),
|
|
],
|
|
[]
|
|
],
|
|
[
|
|
"cmd4",
|
|
[
|
|
command.ParseResult(value = "cmd4", type = mitmproxy.types.Cmd, valid = True),
|
|
],
|
|
["int", "str", "path"]
|
|
],
|
|
[
|
|
"cmd4 ",
|
|
[
|
|
command.ParseResult(value = "cmd4", type = mitmproxy.types.Cmd, valid = True),
|
|
command.ParseResult(value = "", type = int, valid = False),
|
|
],
|
|
["str", "path"]
|
|
],
|
|
[
|
|
"cmd4 1",
|
|
[
|
|
command.ParseResult(value = "cmd4", type = mitmproxy.types.Cmd, valid = True),
|
|
command.ParseResult(value = "1", type = int, valid = True),
|
|
],
|
|
["str", "path"]
|
|
],
|
|
[
|
|
"cmd4 1",
|
|
[
|
|
command.ParseResult(value = "cmd4", type = mitmproxy.types.Cmd, valid = True),
|
|
command.ParseResult(value = "1", type = int, valid = True),
|
|
],
|
|
["str", "path"]
|
|
],
|
|
[
|
|
"flow",
|
|
[
|
|
command.ParseResult(value = "flow", type = mitmproxy.types.Cmd, valid = True),
|
|
],
|
|
["flow", "str"]
|
|
],
|
|
[
|
|
"flow ",
|
|
[
|
|
command.ParseResult(value = "flow", type = mitmproxy.types.Cmd, valid = True),
|
|
command.ParseResult(value = "", type = flow.Flow, valid = False),
|
|
],
|
|
["str"]
|
|
],
|
|
[
|
|
"flow x",
|
|
[
|
|
command.ParseResult(value = "flow", type = mitmproxy.types.Cmd, valid = True),
|
|
command.ParseResult(value = "x", type = flow.Flow, valid = False),
|
|
],
|
|
["str"]
|
|
],
|
|
[
|
|
"flow x ",
|
|
[
|
|
command.ParseResult(value = "flow", type = mitmproxy.types.Cmd, valid = True),
|
|
command.ParseResult(value = "x", type = flow.Flow, valid = False),
|
|
command.ParseResult(value = "", type = str, valid = True),
|
|
],
|
|
[]
|
|
],
|
|
[
|
|
"flow \"one two",
|
|
[
|
|
command.ParseResult(value = "flow", type = mitmproxy.types.Cmd, valid = True),
|
|
command.ParseResult(value = "\"one two", type = flow.Flow, valid = False),
|
|
],
|
|
["str"]
|
|
],
|
|
[
|
|
"flow \"one two\"",
|
|
[
|
|
command.ParseResult(value = "flow", type = mitmproxy.types.Cmd, valid = True),
|
|
command.ParseResult(value = "one two", type = flow.Flow, valid = False),
|
|
],
|
|
["str"]
|
|
],
|
|
]
|
|
with taddons.context() as tctx:
|
|
tctx.master.addons.add(TAddon())
|
|
for s, expected, expectedremain in tests:
|
|
current, remain = tctx.master.commands.parse_partial(s)
|
|
assert current == expected
|
|
assert expectedremain == remain
|
|
|
|
|
|
def test_simple():
|
|
with taddons.context() as tctx:
|
|
c = command.CommandManager(tctx.master)
|
|
a = TAddon()
|
|
c.add("one.two", a.cmd1)
|
|
assert c.commands["one.two"].help == "cmd1 help"
|
|
assert(c.execute("one.two foo") == "ret foo")
|
|
assert(c.call("one.two", "foo") == "ret foo")
|
|
with pytest.raises(exceptions.CommandError, match="Unknown"):
|
|
c.execute("nonexistent")
|
|
with pytest.raises(exceptions.CommandError, match="Invalid"):
|
|
c.execute("")
|
|
with pytest.raises(exceptions.CommandError, match="argument mismatch"):
|
|
c.execute("one.two too many args")
|
|
with pytest.raises(exceptions.CommandError, match="Unknown"):
|
|
c.call("nonexistent")
|
|
with pytest.raises(exceptions.CommandError, match="No escaped"):
|
|
c.execute("\\")
|
|
|
|
c.add("empty", a.empty)
|
|
c.execute("empty")
|
|
|
|
fp = io.StringIO()
|
|
c.dump(fp)
|
|
assert fp.getvalue()
|
|
|
|
|
|
def test_typename():
|
|
assert command.typename(str) == "str"
|
|
assert command.typename(typing.Sequence[flow.Flow]) == "[flow]"
|
|
|
|
assert command.typename(mitmproxy.types.Data) == "[data]"
|
|
assert command.typename(mitmproxy.types.CutSpec) == "[cut]"
|
|
|
|
assert command.typename(flow.Flow) == "flow"
|
|
assert command.typename(typing.Sequence[str]) == "[str]"
|
|
|
|
assert command.typename(mitmproxy.types.Choice("foo")) == "choice"
|
|
assert command.typename(mitmproxy.types.Path) == "path"
|
|
assert command.typename(mitmproxy.types.Cmd) == "cmd"
|
|
|
|
with pytest.raises(exceptions.CommandError, match="missing type annotation"):
|
|
command.typename(inspect._empty)
|
|
with pytest.raises(exceptions.CommandError, match="unsupported type"):
|
|
command.typename(None)
|
|
|
|
|
|
class DummyConsole:
|
|
@command.command("view.flows.resolve")
|
|
def resolve(self, spec: str) -> typing.Sequence[flow.Flow]:
|
|
n = int(spec)
|
|
return [tflow.tflow(resp=True)] * n
|
|
|
|
@command.command("cut")
|
|
def cut(self, spec: str) -> mitmproxy.types.Data:
|
|
return [["test"]]
|
|
|
|
|
|
def test_parsearg():
|
|
with taddons.context() as tctx:
|
|
tctx.master.addons.add(DummyConsole())
|
|
assert command.parsearg(tctx.master.commands, "foo", str) == "foo"
|
|
with pytest.raises(exceptions.CommandError, match="Unsupported"):
|
|
command.parsearg(tctx.master.commands, "foo", type)
|
|
with pytest.raises(exceptions.CommandError):
|
|
command.parsearg(tctx.master.commands, "foo", int)
|
|
|
|
|
|
class TDec:
|
|
@command.command("cmd1")
|
|
def cmd1(self, foo: str) -> str:
|
|
"""cmd1 help"""
|
|
return "ret " + foo
|
|
|
|
@command.command("cmd2")
|
|
def cmd2(self, foo: str) -> str:
|
|
return 99
|
|
|
|
@command.command("empty")
|
|
def empty(self) -> None:
|
|
pass
|
|
|
|
|
|
class TAttr:
|
|
def __getattr__(self, item):
|
|
raise IOError
|
|
|
|
|
|
class TCmds(TAttr):
|
|
def __init__(self):
|
|
self.TAttr = TAttr()
|
|
|
|
@command.command("empty")
|
|
def empty(self) -> None:
|
|
pass
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_collect_commands():
|
|
"""
|
|
This tests for the error thrown by hasattr()
|
|
"""
|
|
with taddons.context() as tctx:
|
|
c = command.CommandManager(tctx.master)
|
|
a = TCmds()
|
|
c.collect_commands(a)
|
|
assert "empty" in c.commands
|
|
|
|
a = TypeErrAddon()
|
|
c.collect_commands(a)
|
|
await tctx.master.await_log("Could not load")
|
|
|
|
|
|
def test_decorator():
|
|
with taddons.context() as tctx:
|
|
c = command.CommandManager(tctx.master)
|
|
a = TDec()
|
|
c.collect_commands(a)
|
|
assert "cmd1" in c.commands
|
|
assert c.execute("cmd1 bar") == "ret bar"
|
|
assert "empty" in c.commands
|
|
assert c.execute("empty") is None
|
|
|
|
with taddons.context() as tctx:
|
|
tctx.master.addons.add(a)
|
|
assert tctx.master.commands.execute("cmd1 bar") == "ret bar"
|
|
|
|
|
|
def test_verify_arg_signature():
|
|
with pytest.raises(exceptions.CommandError):
|
|
command.verify_arg_signature(lambda: None, [1, 2], {})
|
|
print('hello there')
|
|
command.verify_arg_signature(lambda a, b: None, [1, 2], {})
|