diff --git a/mitmproxy/addons/view.py b/mitmproxy/addons/view.py index 341958c27..63416b9fd 100644 --- a/mitmproxy/addons/view.py +++ b/mitmproxy/addons/view.py @@ -111,10 +111,8 @@ class View(collections.Sequence): self.default_order = OrderRequestStart(self) self.orders = dict( - time = OrderRequestStart(self), - method = OrderRequestMethod(self), - url = OrderRequestURL(self), - size = OrderKeySize(self), + time = OrderRequestStart(self), method = OrderRequestMethod(self), + url = OrderRequestURL(self), size = OrderKeySize(self), ) self.order_key = self.default_order self.order_reversed = False @@ -324,6 +322,26 @@ class View(collections.Sequence): if "console_focus_follow" in updated: self.focus_follow = ctx.options.console_focus_follow + def resolve(self, spec: str) -> typing.Sequence[mitmproxy.flow.Flow]: + if spec == "@focus": + return [self.focus.flow] if self.focus.flow else [] + elif spec == "@shown": + return [i for i in self] + elif spec == "@hidden": + return [i for i in self._store.values() if i not in self._view] + elif spec == "@marked": + return [i for i in self._store.values() if i.marked] + elif spec == "@unmarked": + return [i for i in self._store.values() if not i.marked] + else: + filt = flowfilter.parse(spec) + if not filt: + raise exceptions.CommandError("Invalid flow filter: %s" % spec) + return [i for i in self._store.values() if filt(i)] + + def load(self, l): + l.add_command("console.resolve", self.resolve) + def request(self, f): self.add(f) diff --git a/mitmproxy/command.py b/mitmproxy/command.py index e964353b7..3d24675b2 100644 --- a/mitmproxy/command.py +++ b/mitmproxy/command.py @@ -2,17 +2,17 @@ import inspect import typing import shlex from mitmproxy.utils import typecheck - - -class CommandError(Exception): - pass +from mitmproxy import exceptions +from mitmproxy import flow def typename(t: type) -> str: if t in (str, int, bool): return t.__name__ + if t == typing.Sequence[flow.Flow]: + return "[flow]" else: # pragma: no cover - raise NotImplementedError + raise NotImplementedError(t) def parsearg(spec: str, argtype: type) -> typing.Any: @@ -22,7 +22,7 @@ def parsearg(spec: str, argtype: type) -> typing.Any: if argtype == str: return spec else: - raise CommandError("Unsupported argument type: %s" % argtype) + raise exceptions.CommandError("Unsupported argument type: %s" % argtype) class Command: @@ -44,7 +44,7 @@ class Command: Call the command with a set of arguments. At this point, all argumets are strings. """ if len(self.paramtypes) != len(args): - raise CommandError("Usage: %s" % self.signature_help()) + raise exceptions.CommandError("Usage: %s" % self.signature_help()) args = [parsearg(args[i], self.paramtypes[i]) for i in range(len(args))] @@ -52,7 +52,7 @@ class Command: ret = self.func(*args) if not typecheck.check_command_return_type(ret, self.returntype): - raise CommandError("Command returned unexpected data") + raise exceptions.CommandError("Command returned unexpected data") return ret @@ -65,14 +65,19 @@ class CommandManager: def add(self, path: str, func: typing.Callable): self.commands[path] = Command(self, path, func) + def call_args(self, path, args): + """ + Call a command using a list of string arguments. May raise CommandError. + """ + if path not in self.commands: + raise exceptions.CommandError("Unknown command: %s" % path) + return self.commands[path].call(args) + def call(self, cmdstr: str): """ Call a command using a string. May raise CommandError. """ parts = shlex.split(cmdstr) if not len(parts) >= 1: - raise CommandError("Invalid command: %s" % cmdstr) - path = parts[0] - if path not in self.commands: - raise CommandError("Unknown command: %s" % path) - return self.commands[path].call(parts[1:]) + raise exceptions.CommandError("Invalid command: %s" % cmdstr) + return self.call_args(parts[0], parts[1:]) diff --git a/mitmproxy/exceptions.py b/mitmproxy/exceptions.py index 9b6328aca..04525d1fe 100644 --- a/mitmproxy/exceptions.py +++ b/mitmproxy/exceptions.py @@ -93,6 +93,10 @@ class SetServerNotAllowedException(MitmproxyException): pass +class CommandError(Exception): + pass + + class OptionsError(MitmproxyException): pass diff --git a/mitmproxy/test/taddons.py b/mitmproxy/test/taddons.py index ea9534af0..9e2c9838b 100644 --- a/mitmproxy/test/taddons.py +++ b/mitmproxy/test/taddons.py @@ -6,6 +6,7 @@ import mitmproxy.options from mitmproxy import proxy from mitmproxy import addonmanager from mitmproxy import eventsequence +from mitmproxy import command from mitmproxy.addons import script @@ -126,3 +127,10 @@ class context: Recursively invoke an event on an addon and all its children. """ return self.master.addons.invoke_addon(addon, event, *args, **kwargs) + + def command(self, func, *args): + """ + Invoke a command function within a command context, mimicing the actual command environment. + """ + cmd = command.Command(self.master.commands, "test.command", func) + return cmd.call(args) diff --git a/mitmproxy/tools/cmdline.py b/mitmproxy/tools/cmdline.py index fbdbce52a..2714fed63 100644 --- a/mitmproxy/tools/cmdline.py +++ b/mitmproxy/tools/cmdline.py @@ -25,6 +25,11 @@ def common_options(parser, opts): action='store_true', help="Show all options and their default values", ) + parser.add_argument( + '--commands', + action='store_true', + help="Show all commands and their signatures", + ) parser.add_argument( "--conf", type=str, dest="conf", default=CONFIG_PATH, diff --git a/mitmproxy/tools/main.py b/mitmproxy/tools/main.py index b83a35d11..9621dbbc6 100644 --- a/mitmproxy/tools/main.py +++ b/mitmproxy/tools/main.py @@ -39,7 +39,7 @@ def process_options(parser, opts, args): if args.version: print(debug.dump_system_info()) sys.exit(0) - if args.quiet or args.options: + if args.quiet or args.options or args.commands: args.verbosity = 0 args.flow_detail = 0 @@ -84,6 +84,13 @@ def run(MasterKlass, args, extra=None): # pragma: no cover if args.options: print(optmanager.dump_defaults(opts)) sys.exit(0) + if args.commands: + cmds = [] + for c in master.commands.commands.values(): + cmds.append(c.signature_help()) + for i in sorted(cmds): + print(i) + sys.exit(0) opts.set(*args.setoptions) if extra: opts.update(**extra(args)) diff --git a/mitmproxy/utils/typecheck.py b/mitmproxy/utils/typecheck.py index 7199f2fb3..33dd70b08 100644 --- a/mitmproxy/utils/typecheck.py +++ b/mitmproxy/utils/typecheck.py @@ -7,6 +7,20 @@ def check_command_return_type(value: typing.Any, typeinfo: typing.Any) -> bool: types match, False otherwise. This function supports only those types required for command return values. """ + typename = str(typeinfo) + if typename.startswith("typing.Sequence"): + try: + T = typeinfo.__args__[0] # type: ignore + except AttributeError: + # Python 3.5.0 + T = typeinfo.__parameters__[0] # type: ignore + if not isinstance(value, (tuple, list)): + return False + for v in value: + if not check_command_return_type(v, T): + return False + elif not isinstance(value, typeinfo): + return False return True diff --git a/test/mitmproxy/addons/test_view.py b/test/mitmproxy/addons/test_view.py index 7fa3819ef..aca357a2e 100644 --- a/test/mitmproxy/addons/test_view.py +++ b/test/mitmproxy/addons/test_view.py @@ -130,6 +130,46 @@ def test_filter(): assert len(v) == 4 +def test_resolve(): + v = view.View() + with taddons.context(options=options.Options()) as tctx: + assert tctx.command(v.resolve, "@focus") == [] + assert tctx.command(v.resolve, "@shown") == [] + assert tctx.command(v.resolve, "@hidden") == [] + assert tctx.command(v.resolve, "@marked") == [] + assert tctx.command(v.resolve, "@unmarked") == [] + assert tctx.command(v.resolve, "~m get") == [] + v.request(tft(method="get")) + assert len(tctx.command(v.resolve, "~m get")) == 1 + assert len(tctx.command(v.resolve, "@focus")) == 1 + assert len(tctx.command(v.resolve, "@shown")) == 1 + assert len(tctx.command(v.resolve, "@unmarked")) == 1 + assert tctx.command(v.resolve, "@hidden") == [] + assert tctx.command(v.resolve, "@marked") == [] + v.request(tft(method="put")) + assert len(tctx.command(v.resolve, "@focus")) == 1 + assert len(tctx.command(v.resolve, "@shown")) == 2 + assert tctx.command(v.resolve, "@hidden") == [] + assert tctx.command(v.resolve, "@marked") == [] + + v.request(tft(method="get")) + v.request(tft(method="put")) + + f = flowfilter.parse("~m get") + v.set_filter(f) + v[0].marked = True + + def m(l): + return [i.request.method for i in l] + + assert m(tctx.command(v.resolve, "~m get")) == ["GET", "GET"] + assert m(tctx.command(v.resolve, "~m put")) == ["PUT", "PUT"] + assert m(tctx.command(v.resolve, "@shown")) == ["GET", "GET"] + assert m(tctx.command(v.resolve, "@hidden")) == ["PUT", "PUT"] + assert m(tctx.command(v.resolve, "@marked")) == ["GET"] + assert m(tctx.command(v.resolve, "@unmarked")) == ["PUT", "GET", "PUT"] + + def test_order(): v = view.View() with taddons.context(options=options.Options()) as tctx: diff --git a/test/mitmproxy/test_command.py b/test/mitmproxy/test_command.py index d4da7c326..7d0359fa6 100644 --- a/test/mitmproxy/test_command.py +++ b/test/mitmproxy/test_command.py @@ -2,6 +2,7 @@ from mitmproxy import command from mitmproxy import master from mitmproxy import options from mitmproxy import proxy +from mitmproxy import exceptions import pytest @@ -29,7 +30,7 @@ def test_simple(): a = TAddon() c.add("one.two", a.cmd1) assert(c.call("one.two foo") == "ret foo") - with pytest.raises(command.CommandError, match="Unknown"): + with pytest.raises(exceptions.CommandError, match="Unknown"): c.call("nonexistent") - with pytest.raises(command.CommandError, match="Invalid"): + with pytest.raises(exceptions.CommandError, match="Invalid"): c.call("")