command: save.file flowspec path -> None

Our first user-facing command. The following commands do the obvious things:

save.file @marked /tmp/flows
save.file @focus /tmp/flows
save.file @hidden /tmp/flows
save.file "~m get" /tmp/flows

mitmproxy/addonmanager.py

@@ -145,7 +145,7 @@ class AddonManager:
         for a in traverse([addon]):
             name = _get_name(a)
             if name in self.lookup:
-                raise exceptions.AddonError(
+                raise exceptions.AddonManagerError(
                     "An addon called '%s' already exists." % name
                 )
         l = Loader(self.master)
@@ -175,7 +175,7 @@ class AddonManager:
         for a in traverse([addon]):
             n = _get_name(a)
             if n not in self.lookup:
-                raise exceptions.AddonError("No such addon: %s" % n)
+                raise exceptions.AddonManagerError("No such addon: %s" % n)
             self.chain = [i for i in self.chain if i is not a]
             del self.lookup[_get_name(a)]
         with self.master.handlecontext():
@@ -224,7 +224,7 @@ class AddonManager:
             func = getattr(a, name, None)
             if func:
                 if not callable(func):
-                    raise exceptions.AddonError(
+                    raise exceptions.AddonManagerError(
                         "Addon handler %s not callable" % name
                     )
                 func(*args, **kwargs)

mitmproxy/addons/save.py

@@ -1,9 +1,11 @@
 import os.path
+import typing
 
 from mitmproxy import exceptions
 from mitmproxy import flowfilter
 from mitmproxy import io
 from mitmproxy import ctx
+from mitmproxy import flow
 
 
 class Save:
@@ -12,10 +14,18 @@ class Save:
         self.filt = None
         self.active_flows = set()  # type: Set[flow.Flow]
 
-    def start_stream_to_path(self, path, mode, flt):
+    def open_file(self, path):
+        if path.startswith("+"):
+            path = path[1:]
+            mode = "ab"
+        else:
+            mode = "wb"
         path = os.path.expanduser(path)
+        return open(path, mode)
+
+    def start_stream_to_path(self, path, flt):
         try:
-            f = open(path, mode)
+            f = self.open_file(path)
         except IOError as v:
             raise exceptions.OptionsError(str(v))
         self.stream = io.FilteredFlowWriter(f, flt)
@@ -36,13 +46,19 @@ class Save:
             if self.stream:
                 self.done()
             if ctx.options.save_stream_file:
-                if ctx.options.save_stream_file.startswith("+"):
+                self.start_stream_to_path(ctx.options.save_stream_file, self.filt)
-                    path = ctx.options.save_stream_file[1:]
+
-                    mode = "ab"
+    def save(self, flows: typing.Sequence[flow.Flow], path: str) -> None:
-                else:
+        try:
-                    path = ctx.options.save_stream_file
+            f = self.open_file(path)
-                    mode = "wb"
+        except IOError as v:
-                self.start_stream_to_path(path, mode, self.filt)
+            raise exceptions.CommandError(v) from v
+        stream = io.FlowWriter(f)
+        for i in flows:
+            stream.add(i)
+
+    def load(self, l):
+        l.add_command("save.file", self.save)
 
     def tcp_start(self, flow):
         if self.stream:
@@ -64,8 +80,8 @@ class Save:
 
     def done(self):
         if self.stream:
-            for flow in self.active_flows:
+            for f in self.active_flows:
-                self.stream.add(flow)
+                self.stream.add(f)
             self.active_flows = set([])
             self.stream.fo.close()
             self.stream = None

mitmproxy/command.py

@@ -6,25 +6,19 @@ from mitmproxy import exceptions
 from mitmproxy import flow
 
 
-def typename(t: type) -> str:
+def typename(t: type, ret: bool) -> str:
+    """
+        Translates a type to an explanatory string. Ifl ret is True, we're
+        looking at a return type, else we're looking at a parameter type.
+    """
     if t in (str, int, bool):
         return t.__name__
     if t == typing.Sequence[flow.Flow]:
-        return "[flow]"
+        return "[flow]" if ret else "flowspec"
     else:  # pragma: no cover
         raise NotImplementedError(t)
 
 
-def parsearg(spec: str, argtype: type) -> typing.Any:
-    """
-        Convert a string to a argument to the appropriate type.
-    """
-    if argtype == str:
-        return spec
-    else:
-        raise exceptions.CommandError("Unsupported argument type: %s" % argtype)
-
-
 class Command:
     def __init__(self, manager, path, func) -> None:
         self.path = path
@@ -35,8 +29,8 @@ class Command:
         self.returntype = sig.return_annotation
 
     def signature_help(self) -> str:
-        params = " ".join([typename(i) for i in self.paramtypes])
+        params = " ".join([typename(i, False) for i in self.paramtypes])
-        ret = " -> " + typename(self.returntype) if self.returntype else ""
+        ret = " -> " + typename(self.returntype, True) if self.returntype else ""
         return "%s %s%s" % (self.path, params, ret)
 
     def call(self, args: typing.Sequence[str]):
@@ -46,10 +40,12 @@ class Command:
         if len(self.paramtypes) != len(args):
             raise exceptions.CommandError("Usage: %s" % self.signature_help())
 
-        args = [parsearg(args[i], self.paramtypes[i]) for i in range(len(args))]
+        pargs = []
+        for i in range(len(args)):
+            pargs.append(parsearg(self.manager, args[i], self.paramtypes[i]))
 
         with self.manager.master.handlecontext():
-            ret = self.func(*args)
+            ret = self.func(*pargs)
 
         if not typecheck.check_command_return_type(ret, self.returntype):
             raise exceptions.CommandError("Command returned unexpected data")
@@ -81,3 +77,15 @@ class CommandManager:
         if not len(parts) >= 1:
             raise exceptions.CommandError("Invalid command: %s" % cmdstr)
         return self.call_args(parts[0], parts[1:])
+
+
+def parsearg(manager: CommandManager, spec: str, argtype: type) -> typing.Any:
+    """
+        Convert a string to a argument to the appropriate type.
+    """
+    if argtype == str:
+        return spec
+    elif argtype == typing.Sequence[flow.Flow]:
+        return manager.call_args("console.resolve", [spec])
+    else:
+        raise exceptions.CommandError("Unsupported argument type: %s" % argtype)

mitmproxy/exceptions.py

@@ -101,7 +101,7 @@ class OptionsError(MitmproxyException):
     pass
 
 
-class AddonError(MitmproxyException):
+class AddonManagerError(MitmproxyException):
     pass
 
 

mitmproxy/tools/console/command.py

@@ -1,6 +1,6 @@
 import urwid
 
-from mitmproxy import command
+from mitmproxy import exceptions
 from mitmproxy.tools.console import signals
 
 
@@ -17,9 +17,10 @@ class CommandExecutor:
         self.master = master
 
     def __call__(self, cmd):
+        if cmd.strip():
             try:
                 ret = self.master.commands.call(cmd)
-        except command.CommandError as v:
+            except exceptions.CommandError as v:
                 signals.status_message.send(message=str(v))
             else:
                 if type(ret) == str:

mitmproxy/utils/typecheck.py

@@ -19,6 +19,8 @@ def check_command_return_type(value: typing.Any, typeinfo: typing.Any) -> bool:
         for v in value:
             if not check_command_return_type(v, T):
                 return False
+    elif value is None and typeinfo is None:
+        return True
     elif not isinstance(value, typeinfo):
         return False
     return True

test/mitmproxy/addons/test_save.py

@@ -7,6 +7,7 @@ from mitmproxy import io
 from mitmproxy import exceptions
 from mitmproxy import options
 from mitmproxy.addons import save
+from mitmproxy.addons import view
 
 
 def test_configure(tmpdir):
@@ -42,6 +43,26 @@ def test_tcp(tmpdir):
         assert rd(p)
 
 
+def test_save_command(tmpdir):
+    sa = save.Save()
+    with taddons.context() as tctx:
+        p = str(tmpdir.join("foo"))
+        sa.save([tflow.tflow(resp=True)], p)
+        assert len(rd(p)) == 1
+        sa.save([tflow.tflow(resp=True)], p)
+        assert len(rd(p)) == 1
+        sa.save([tflow.tflow(resp=True)], "+" + p)
+        assert len(rd(p)) == 2
+
+        with pytest.raises(exceptions.CommandError):
+            sa.save([tflow.tflow(resp=True)], str(tmpdir))
+
+        v = view.View()
+        tctx.master.addons.add(v)
+        tctx.master.addons.add(sa)
+        tctx.master.commands.call_args("save.file", ["@shown", p])
+
+
 def test_simple(tmpdir):
     sa = save.Save()
     with taddons.context() as tctx:

test/mitmproxy/test_addonmanager.py

@@ -61,9 +61,9 @@ def test_lifecycle():
     a = addonmanager.AddonManager(m)
     a.add(TAddon("one"))
 
-    with pytest.raises(exceptions.AddonError):
+    with pytest.raises(exceptions.AddonManagerError):
         a.add(TAddon("one"))
-    with pytest.raises(exceptions.AddonError):
+    with pytest.raises(exceptions.AddonManagerError):
         a.remove(TAddon("nonexistent"))
 
     f = tflow.tflow()