console: flow resolution command

This is our first built-in command, which will be used by very many other
commands.

Also add a --commands option to dump all commands, analogous to --options.

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)
 

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:])

mitmproxy/exceptions.py

@@ -93,6 +93,10 @@ class SetServerNotAllowedException(MitmproxyException):
     pass
 
 
+class CommandError(Exception):
+    pass
+
+
 class OptionsError(MitmproxyException):
     pass
 

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)

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,

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))

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
 
 

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:

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("")