diff options
| -rw-r--r-- | mitmproxy/addonmanager.py | 2 | ||||
| -rw-r--r-- | mitmproxy/addons/core.py | 5 | ||||
| -rw-r--r-- | mitmproxy/command.py | 18 | ||||
| -rw-r--r-- | mitmproxy/tools/console/master.py | 19 | ||||
| -rw-r--r-- | test/mitmproxy/test_addonmanager.py | 11 | ||||
| -rw-r--r-- | test/mitmproxy/test_command.py | 61 |
6 files changed, 87 insertions, 29 deletions
diff --git a/mitmproxy/addonmanager.py b/mitmproxy/addonmanager.py index 9779a31a..5d63b1b3 100644 --- a/mitmproxy/addonmanager.py +++ b/mitmproxy/addonmanager.py @@ -153,6 +153,8 @@ class AddonManager: for a in traverse([addon]): name = _get_name(a) self.lookup[name] = a + for a in traverse([addon]): + self.master.commands.collect_commands(a) return addon def add(self, *addons): diff --git a/mitmproxy/addons/core.py b/mitmproxy/addons/core.py index 1cbf4a01..3f9cb15e 100644 --- a/mitmproxy/addons/core.py +++ b/mitmproxy/addons/core.py @@ -1,8 +1,10 @@ from mitmproxy import ctx from mitmproxy import exceptions +from mitmproxy import command class Core: + @command.command("set") def set(self, spec: str) -> None: """ Set an option of the form "key[=value]". When the value is omitted, @@ -14,6 +16,3 @@ class Core: ctx.options.set(spec) except exceptions.OptionsError as e: raise exceptions.CommandError(e) from e - - def load(self, l): - l.add_command("set", self.set) diff --git a/mitmproxy/command.py b/mitmproxy/command.py index 5477cbdf..fa6e23ea 100644 --- a/mitmproxy/command.py +++ b/mitmproxy/command.py @@ -2,6 +2,7 @@ import inspect import typing import shlex import textwrap +import functools from mitmproxy.utils import typecheck from mitmproxy import exceptions @@ -74,6 +75,13 @@ class CommandManager: self.master = master self.commands = {} + def collect_commands(self, addon): + for i in dir(addon): + if not i.startswith("__"): + o = getattr(addon, i) + if hasattr(o, "command_path"): + self.add(o.command_path, o) + def add(self, path: str, func: typing.Callable): self.commands[path] = Command(self, path, func) @@ -112,3 +120,13 @@ def parsearg(manager: CommandManager, spec: str, argtype: type) -> typing.Any: return flows[0] else: raise exceptions.CommandError("Unsupported argument type: %s" % argtype) + + +def command(path): + def decorator(function): + @functools.wraps(function) + def wrapper(*args, **kwargs): + return function(*args, **kwargs) + wrapper.__dict__["command_path"] = path + return wrapper + return decorator diff --git a/mitmproxy/tools/console/master.py b/mitmproxy/tools/console/master.py index 4b34b75b..7787ba11 100644 --- a/mitmproxy/tools/console/master.py +++ b/mitmproxy/tools/console/master.py @@ -14,6 +14,7 @@ import urwid from mitmproxy import addons from mitmproxy import exceptions +from mitmproxy import command from mitmproxy import master from mitmproxy import io from mitmproxy import log @@ -85,34 +86,41 @@ class ConsoleCommands: self.master = master self.started = False - def command(self, partial: str) -> None: + @command.command("console.command") + def console_command(self, partial: str) -> None: """ Prompt the user to edit a command with a (possilby empty) starting value. """ signals.status_prompt_command.send(partial=partial) + @command.command("console.view.commands") def view_commands(self) -> None: """View the commands list.""" self.master.view_commands() + @command.command("console.view.options") def view_options(self) -> None: """View the options editor.""" self.master.view_options() + @command.command("console.view.help") def view_help(self) -> None: """View help.""" self.master.view_help() + @command.command("console.view.flow") def view_flow(self, flow: flow.Flow) -> None: """View a flow.""" if hasattr(flow, "request"): # FIME: Also set focus? self.master.view_flow(flow) + @command.command("console.exit") def exit(self) -> None: """Exit mitmproxy.""" raise urwid.ExitMainLoop + @command.command("console.view.pop") def view_pop(self) -> None: """ Pop a view off the console stack. At the top level, this prompts the @@ -120,15 +128,6 @@ class ConsoleCommands: """ signals.pop_view_state.send(self) - def load(self, l): - l.add_command("console.command", self.command) - l.add_command("console.exit", self.exit) - l.add_command("console.view.commands", self.view_commands) - l.add_command("console.view.help", self.view_help) - l.add_command("console.view.options", self.view_options) - l.add_command("console.view.pop", self.view_pop) - l.add_command("console.view.flow", self.view_flow) - def running(self): self.started = True diff --git a/test/mitmproxy/test_addonmanager.py b/test/mitmproxy/test_addonmanager.py index 034182a6..678bc1b7 100644 --- a/test/mitmproxy/test_addonmanager.py +++ b/test/mitmproxy/test_addonmanager.py @@ -4,6 +4,7 @@ from mitmproxy import addons from mitmproxy import addonmanager from mitmproxy import exceptions from mitmproxy import options +from mitmproxy import command from mitmproxy import master from mitmproxy import proxy from mitmproxy.test import taddons @@ -18,6 +19,10 @@ class TAddon: if addons: self.addons = addons + @command.command("test.command") + def testcommand(self) -> str: + return "here" + def __repr__(self): return "Addon(%s)" % self.name @@ -38,6 +43,12 @@ class AOption: l.add_option("custom_option", bool, False, "help") +def test_command(): + with taddons.context() as tctx: + tctx.master.addons.add(TAddon("test")) + assert tctx.master.commands.call("test.command") == "here" + + def test_halt(): o = options.Options() m = master.Master(o, proxy.DummyServer(o)) diff --git a/test/mitmproxy/test_command.py b/test/mitmproxy/test_command.py index 16bc8a90..64928dbf 100644 --- a/test/mitmproxy/test_command.py +++ b/test/mitmproxy/test_command.py @@ -39,22 +39,21 @@ class TestCommand: def test_simple(): - o = options.Options() - m = master.Master(o, proxy.DummyServer(o)) - c = command.CommandManager(m) - a = TAddon() - c.add("one.two", a.cmd1) - assert c.commands["one.two"].help == "cmd1 help" - assert(c.call("one.two foo") == "ret foo") - with pytest.raises(exceptions.CommandError, match="Unknown"): - c.call("nonexistent") - with pytest.raises(exceptions.CommandError, match="Invalid"): - c.call("") - with pytest.raises(exceptions.CommandError, match="Usage"): - c.call("one.two too many args") - - c.add("empty", a.empty) - c.call("empty") + 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.call("one.two foo") == "ret foo") + with pytest.raises(exceptions.CommandError, match="Unknown"): + c.call("nonexistent") + with pytest.raises(exceptions.CommandError, match="Invalid"): + c.call("") + with pytest.raises(exceptions.CommandError, match="Usage"): + c.call("one.two too many args") + + c.add("empty", a.empty) + c.call("empty") def test_typename(): @@ -87,3 +86,33 @@ def test_parsearg(): command.parsearg(tctx.master.commands, "0", flow.Flow) with pytest.raises(exceptions.CommandError): command.parsearg(tctx.master.commands, "foo", Exception) + + +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 + + +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.call("cmd1 bar") == "ret bar" + assert "empty" in c.commands + assert c.call("empty") is None + + with taddons.context() as tctx: + tctx.master.addons.add(a) + assert tctx.master.commands.call("cmd1 bar") == "ret bar" |