Introduce cuts: a flow dissector

This PR introduces the cuts addon, a flow dissector that  allows you to select
and operate on specific components of flows. It also adds the first consumer
for cuts - the cuts.save command.

Save the content of the focus to /tmp/foo:

    cuts.save s.content|@focus /tmp/foo

Save the URL and response content-type headers for all flows currently shown to
file, comma-separated, one flow per line:

    cuts.save s.url,q.header[content-type]|@focus /tmp/foo

We also use this to replace the body save shortcut in the console flowlist.

10 files changed, 270 insertions, 5 deletions

diff --git a/mitmproxy/addons/__init__.py b/mitmproxy/addons/__init__.py
index fa194fc1..204d61e0 100644
--- a/mitmproxy/addons/__init__.py
+++ b/mitmproxy/addons/__init__.py
@@ -5,6 +5,7 @@ from mitmproxy.addons import check_ca
 from mitmproxy.addons import clientplayback
 from mitmproxy.addons import core_option_validation
 from mitmproxy.addons import core
+from mitmproxy.addons import cut
 from mitmproxy.addons import disable_h2c
 from mitmproxy.addons import onboarding
 from mitmproxy.addons import proxyauth
@@ -28,6 +29,7 @@ def default_addons():
         check_alpn.CheckALPN(),
         check_ca.CheckCA(),
         clientplayback.ClientPlayback(),
+        cut.Cut(),
         disable_h2c.DisableH2C(),
         onboarding.Onboarding(),
         proxyauth.ProxyAuth(),
diff --git a/mitmproxy/addons/cut.py b/mitmproxy/addons/cut.py
new file mode 100644
index 00000000..b4d9b522
--- /dev/null
+++ b/mitmproxy/addons/cut.py
@@ -0,0 +1,112 @@
+import os
+import typing
+from mitmproxy import command
+from mitmproxy import exceptions
+from mitmproxy import flow
+from mitmproxy import ctx
+from mitmproxy.utils import strutils
+
+
+def headername(spec: str):
+    if not (spec.startswith("header[") and spec.endswith("]")):
+        raise exceptions.CommandError("Invalid header spec: %s" % spec)
+    return spec[len("header["):-1].strip()
+
+
+def extract(cut: str, f: flow.Flow) -> typing.Union[str, bytes]:
+    if cut.startswith("q."):
+        req = getattr(f, "request", None)
+        if not req:
+            return ""
+        rem = cut[len("q."):]
+        if rem in ["method", "scheme", "host", "port", "path", "url"]:
+            return str(getattr(req, rem))
+        elif rem == "content":
+            return req.content
+        elif rem.startswith("header["):
+            return req.headers.get(headername(rem), "")
+    elif cut.startswith("s."):
+        resp = getattr(f, "response", None)
+        if not resp:
+            return ""
+        rem = cut[len("s."):]
+        if rem in ["status_code", "reason"]:
+            return str(getattr(resp, rem))
+        elif rem == "content":
+            return resp.content
+        elif rem.startswith("header["):
+            return resp.headers.get(headername(rem), "")
+    raise exceptions.CommandError("Invalid cut specification: %s" % cut)
+
+
+def parse_cutspec(s: str) -> typing.Tuple[str, typing.Sequence[str]]:
+    """
+        Returns (flowspec, [cuts]).
+
+        Raises exceptions.CommandError if input is invalid.
+    """
+    parts = s.split("|", maxsplit=1)
+    flowspec = "@all"
+    if len(parts) == 2:
+        flowspec = parts[1].strip()
+    cuts = parts[0]
+    cutparts = [i.strip() for i in cuts.split(",") if i.strip()]
+    if len(cutparts) == 0:
+        raise exceptions.CommandError("Invalid cut specification.")
+    return flowspec, cutparts
+
+
+class Cut:
+    @command.command("cut")
+    def cut(self, cutspec: str) -> command.Cuts:
+        """
+            Resolve a cut specification of the form "cuts|flowspec". The
+            flowspec is optional, and if it is not specified, it is assumed to
+            be @all. The cuts are a comma-separated list of cut snippets.
+
+            HTTP requests: q.method, q.scheme, q.host, q.port, q.path, q.url,
+            q.header[key], q.content
+
+            HTTP responses: s.status_code, s.reason, s.header[key], s.content
+
+            Client connections: cc.address, cc.sni, cc.cipher_name,
+            cc.alpn_proto, cc.tls_version
+
+            Server connections: sc.address, sc.ip, sc.cert, sc.sni,
+            sc.alpn_proto, sc.tls_version
+        """
+        flowspec, cuts = parse_cutspec(cutspec)
+        flows = ctx.master.commands.call_args("view.resolve", [flowspec])
+        ret = []
+        for f in flows:
+            ret.append([extract(c, f) for c in cuts])
+        return ret
+
+    @command.command("cut.save")
+    def save(self, cuts: command.Cuts, path: str) -> None:
+        """
+            Save cuts to file.
+
+                cut.save resp.content|@focus /tmp/foo
+
+                cut.save req.host,resp.header[content-type]|@focus /tmp/foo
+        """
+        mode = "wb"
+        if path.startswith("+"):
+            mode = "ab"
+            path = path[1:]
+        path = os.path.expanduser(path)
+        with open(path, mode) as fp:
+            if fp.tell() > 0:
+                # We're appending to a file that already exists and has content
+                fp.write(b"\n")
+            for ci, c in enumerate(cuts):
+                if ci > 0:
+                    fp.write(b"\n")
+                for vi, v in enumerate(c):
+                    if vi > 0:
+                        fp.write(b", ")
+                    if isinstance(v, str):
+                        v = strutils.always_bytes(v)
+                    fp.write(v)
+        ctx.log.alert("Saved %s cuts." % len(cuts))
diff --git a/mitmproxy/command.py b/mitmproxy/command.py
index 3b58cc25..8bf2794c 100644
--- a/mitmproxy/command.py
+++ b/mitmproxy/command.py
@@ -9,6 +9,11 @@ from mitmproxy import exceptions
 from mitmproxy import flow
 
 
+Cuts = typing.Sequence[
+    typing.Sequence[typing.Union[str, bytes]]
+]
+
+
 def typename(t: type, ret: bool) -> str:
     """
         Translates a type to an explanatory string. Ifl ret is True, we're
@@ -18,6 +23,8 @@ def typename(t: type, ret: bool) -> str:
         return t.__name__
     elif t == typing.Sequence[flow.Flow]:
         return "[flow]" if ret else "flowspec"
+    elif t == Cuts:
+        return "[cuts]" if ret else "cutspec"
     elif t == flow.Flow:
         return "flow"
     else:  # pragma: no cover
@@ -125,6 +132,8 @@ def parsearg(manager: CommandManager, spec: str, argtype: type) -> typing.Any:
             raise exceptions.CommandError("Expected an integer, got %s." % spec)
     elif argtype == typing.Sequence[flow.Flow]:
         return manager.call_args("view.resolve", [spec])
+    elif argtype == Cuts:
+        return manager.call_args("cut", [spec])
     elif argtype == flow.Flow:
         flows = manager.call_args("view.resolve", [spec])
         if len(flows) != 1:
diff --git a/mitmproxy/tools/console/flowlist.py b/mitmproxy/tools/console/flowlist.py
index 8f8f5493..898c1478 100644
--- a/mitmproxy/tools/console/flowlist.py
+++ b/mitmproxy/tools/console/flowlist.py
@@ -157,8 +157,6 @@ class FlowItem(urwid.WidgetWrap):
         #         callback = common.export_to_clip_or_file,
         #         args = (None, self.flow, common.copy_to_clipboard_or_prompt)
         #     )
-        elif key == "b":
-            common.ask_save_body(None, self.flow)
         else:
             return key
 
diff --git a/mitmproxy/tools/console/master.py b/mitmproxy/tools/console/master.py
index 29a61bb3..b86be7e1 100644
--- a/mitmproxy/tools/console/master.py
+++ b/mitmproxy/tools/console/master.py
@@ -151,6 +151,7 @@ def default_keymap(km):
 
     km.add("A", "flow.resume @all", context="flowlist")
     km.add("a", "flow.resume @focus", context="flowlist")
+    km.add("b", "console.command 'cut.save s.content|@focus '", context="flowlist")
     km.add("d", "view.remove @focus", context="flowlist")
     km.add("D", "view.duplicate @focus", context="flowlist")
     km.add("e", "set console_eventlog=toggle", context="flowlist")
diff --git a/mitmproxy/utils/strutils.py b/mitmproxy/utils/strutils.py
index 1b90c2e5..db0cfd2e 100644
--- a/mitmproxy/utils/strutils.py
+++ b/mitmproxy/utils/strutils.py
@@ -25,9 +25,10 @@ def always_str(str_or_bytes: Optional[AnyStr], *decode_args) -> Optional[str]:
         raise TypeError("Expected str or bytes, but got {}.".format(type(str_or_bytes).__name__))
 
 
-# Translate control characters to "safe" characters. This implementation initially
-# replaced them with the matching control pictures (http://unicode.org/charts/PDF/U2400.pdf),
-# but that turned out to render badly with monospace fonts. We are back to "." therefore.
+# Translate control characters to "safe" characters. This implementation
+# initially replaced them with the matching control pictures
+# (http://unicode.org/charts/PDF/U2400.pdf), but that turned out to render badly
+# with monospace fonts. We are back to "." therefore.
 _control_char_trans = {
     x: ord(".")  # x + 0x2400 for unicode control group pictures
     for x in range(32)
diff --git a/mitmproxy/utils/typecheck.py b/mitmproxy/utils/typecheck.py
index 20791e17..c97ff529 100644
--- a/mitmproxy/utils/typecheck.py
+++ b/mitmproxy/utils/typecheck.py
@@ -19,6 +19,16 @@ 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 typename.startswith("typing.Union"):
+        try:
+            types = typeinfo.__args__  # type: ignore
+        except AttributeError:
+            # Python 3.5.x
+            types = typeinfo.__union_params__  # type: ignore
+        for T in types:
+            checks = [check_command_return_type(value, T) for T in types]
+            if not any(checks):
+                return False
     elif value is None and typeinfo is None:
         return True
     elif not isinstance(value, typeinfo):
diff --git a/test/mitmproxy/addons/test_cut.py b/test/mitmproxy/addons/test_cut.py
new file mode 100644
index 00000000..5eb864c0
--- /dev/null
+++ b/test/mitmproxy/addons/test_cut.py
@@ -0,0 +1,109 @@
+
+from mitmproxy.addons import cut
+from mitmproxy.addons import view
+from mitmproxy import exceptions
+from mitmproxy.test import taddons
+from mitmproxy.test import tflow
+import pytest
+
+
+def test_parse_cutspec():
+    tests = [
+        ("", None, True),
+        ("req.method", ("@all", ["req.method"]), False),
+        (
+            "req.method,req.host",
+            ("@all", ["req.method", "req.host"]),
+            False
+        ),
+        (
+            "req.method,req.host|~b foo",
+            ("~b foo", ["req.method", "req.host"]),
+            False
+        ),
+        (
+            "req.method,req.host|~b foo | ~b bar",
+            ("~b foo | ~b bar", ["req.method", "req.host"]),
+            False
+        ),
+        (
+            "req.method, req.host | ~b foo | ~b bar",
+            ("~b foo | ~b bar", ["req.method", "req.host"]),
+            False
+        ),
+    ]
+    for cutspec, output, err in tests:
+        try:
+            assert cut.parse_cutspec(cutspec) == output
+        except exceptions.CommandError:
+            if not err:
+                raise
+        else:
+            if err:
+                raise AssertionError("Expected error.")
+
+
+def test_headername():
+    with pytest.raises(exceptions.CommandError):
+        cut.headername("header[foo.")
+
+
+def qr(f):
+    with open(f, "rb") as fp:
+        return fp.read()
+
+
+def test_cut_file(tmpdir):
+    f = str(tmpdir.join("path"))
+    v = view.View()
+    c = cut.Cut()
+    with taddons.context() as tctx:
+        tctx.master.addons.add(v, c)
+
+        v.add([tflow.tflow(resp=True)])
+
+        tctx.command(c.save, "q.method|@all", f)
+        assert qr(f) == b"GET"
+        tctx.command(c.save, "q.content|@all", f)
+        assert qr(f) == b"content"
+        tctx.command(c.save, "q.content|@all", "+" + f)
+        assert qr(f) == b"content\ncontent"
+
+        v.add([tflow.tflow(resp=True)])
+        tctx.command(c.save, "q.method|@all", f)
+        assert qr(f) == b"GET\nGET"
+        tctx.command(c.save, "q.method,q.path|@all", f)
+        assert qr(f) == b"GET, /path\nGET, /path"
+
+
+def test_cut():
+    v = view.View()
+    c = cut.Cut()
+    with taddons.context() as tctx:
+        v.add([tflow.tflow(resp=True)])
+        tctx.master.addons.add(v, c)
+        assert c.cut("q.method|@all") == [["GET"]]
+        assert c.cut("q.scheme|@all") == [["http"]]
+        assert c.cut("q.host|@all") == [["address"]]
+        assert c.cut("q.port|@all") == [["22"]]
+        assert c.cut("q.path|@all") == [["/path"]]
+        assert c.cut("q.url|@all") == [["http://address:22/path"]]
+        assert c.cut("q.content|@all") == [[b"content"]]
+        assert c.cut("q.header[header]|@all") == [["qvalue"]]
+        assert c.cut("q.header[unknown]|@all") == [[""]]
+
+        assert c.cut("s.status_code|@all") == [["200"]]
+        assert c.cut("s.reason|@all") == [["OK"]]
+        assert c.cut("s.content|@all") == [[b"message"]]
+        assert c.cut("s.header[header-response]|@all") == [["svalue"]]
+
+        with pytest.raises(exceptions.CommandError):
+            assert c.cut("moo") == [["svalue"]]
+
+    v = view.View()
+    c = cut.Cut()
+    with taddons.context() as tctx:
+        tctx.master.addons.add(v, c)
+        v.add([tflow.ttcpflow()])
+        assert c.cut("q.method|@all") == [[""]]
+        assert c.cut("s.status|@all") == [[""]]
diff --git a/test/mitmproxy/test_command.py b/test/mitmproxy/test_command.py
index 96d79dba..aef05adc 100644
--- a/test/mitmproxy/test_command.py
+++ b/test/mitmproxy/test_command.py
@@ -60,17 +60,25 @@ def test_typename():
     assert command.typename(str, True) == "str"
     assert command.typename(typing.Sequence[flow.Flow], True) == "[flow]"
     assert command.typename(typing.Sequence[flow.Flow], False) == "flowspec"
+
+    assert command.typename(command.Cuts, False) == "cutspec"
+    assert command.typename(command.Cuts, True) == "[cuts]"
+
     assert command.typename(flow.Flow, False) == "flow"
 
 
 class DummyConsole:
     def load(self, l):
         l.add_command("view.resolve", self.resolve)
+        l.add_command("cut", self.cut)
 
     def resolve(self, spec: str) -> typing.Sequence[flow.Flow]:
         n = int(spec)
         return [tflow.tflow(resp=True)] * n
 
+    def cut(self, spec: str) -> command.Cuts:
+        return [["test"]]
+
 
 def test_parsearg():
     with taddons.context() as tctx:
@@ -97,6 +105,10 @@ def test_parsearg():
         with pytest.raises(exceptions.CommandError):
             command.parsearg(tctx.master.commands, "foo", Exception)
 
+        assert command.parsearg(
+            tctx.master.commands, "foo", command.Cuts
+        ) == [["test"]]
+
 
 class TDec:
     @command.command("cmd1")
diff --git a/test/mitmproxy/utils/test_typecheck.py b/test/mitmproxy/utils/test_typecheck.py
index 22bd7c34..17f70d37 100644
--- a/test/mitmproxy/utils/test_typecheck.py
+++ b/test/mitmproxy/utils/test_typecheck.py
@@ -4,6 +4,7 @@ from unittest import mock
 import pytest
 
 from mitmproxy.utils import typecheck
+from mitmproxy import command
 
 
 class TBase:
@@ -93,9 +94,19 @@ def test_check_command_return_type():
     assert(typecheck.check_command_return_type(None, None))
     assert(not typecheck.check_command_return_type(["foo"], typing.Sequence[int]))
     assert(not typecheck.check_command_return_type("foo", typing.Sequence[int]))
+    assert(typecheck.check_command_return_type([["foo", b"bar"]], command.Cuts))
+    assert(not typecheck.check_command_return_type(["foo", b"bar"], command.Cuts))
+    assert(not typecheck.check_command_return_type([["foo", 22]], command.Cuts))
 
     # Python 3.5 only defines __parameters__
     m = mock.Mock()
     m.__str__ = lambda self: "typing.Sequence"
     m.__parameters__ = (int,)
+
     typecheck.check_command_return_type([10], m)
+
+    # Python 3.5 only defines __union_params__
+    m = mock.Mock()
+    m.__str__ = lambda self: "typing.Union"
+    m.__union_params__ = (int,)
+    assert not typecheck.check_command_return_type([22], m)