command cuts: add completion

- Remove shortcuts for request, response, etc. - we don't need them if we have completion
- Restrict cuts specification to a set of prefixes
- Extend cuts to add a few more items

6 files changed, 79 insertions, 31 deletions

diff --git a/mitmproxy/addons/cut.py b/mitmproxy/addons/cut.py
index 6b9dc723..efc9e5df 100644
--- a/mitmproxy/addons/cut.py
+++ b/mitmproxy/addons/cut.py
@@ -17,14 +17,6 @@ def headername(spec: str):
     return spec[len("header["):-1].strip()
 
 
-flow_shortcuts = {
-    "q": "request",
-    "s": "response",
-    "cc": "client_conn",
-    "sc": "server_conn",
-}
-
-
 def is_addr(v):
     return isinstance(v, tuple) and len(v) > 1
 
@@ -35,8 +27,6 @@ def extract(cut: str, f: flow.Flow) -> typing.Union[str, bytes]:
     for i, spec in enumerate(path):
         if spec.startswith("_"):
             raise exceptions.CommandError("Can't access internal attribute %s" % spec)
-        if isinstance(current, flow.Flow):
-            spec = flow_shortcuts.get(spec, spec)
 
         part = getattr(current, spec, None)
         if i == len(path) - 1:
@@ -65,13 +55,12 @@ class Cut:
     ) -> command.Cuts:
         """
             Cut data from a set of flows. Cut specifications are attribute paths
-            from the base of the flow object, with a few conveniences - "q",
-            "s", "cc" and "sc" are shortcuts for request, response, client_conn
-            and server_conn, "port" and "host" retrieve parts of an address
-            tuple, ".header[key]" retrieves a header value. Return values
-            converted to strings or bytes: SSL certicates are converted to PEM
-            format, bools are "true" or "false", "bytes" are preserved, and all
-            other values are converted to strings.
+            from the base of the flow object, with a few conveniences - "port"
+            and "host" retrieve parts of an address tuple, ".header[key]"
+            retrieves a header value. Return values converted to strings or
+            bytes: SSL certicates are converted to PEM format, bools are "true"
+            or "false", "bytes" are preserved, and all other values are
+            converted to strings.
         """
         ret = []
         for f in flows:
@@ -101,11 +90,11 @@ class Cut:
                 if fp.tell() > 0:
                     # We're appending to a file that already exists and has content
                     fp.write(b"\n")
-                for v in [extract(cuts[0], f) for f in flows]:
-                    if isinstance(v, bytes):
-                        fp.write(v)
-                    else:
-                        fp.write(v.encode("utf8"))
+                v = extract(cuts[0], flows[0])
+                if isinstance(v, bytes):
+                    fp.write(v)
+                else:
+                    fp.write(v.encode("utf8"))
             ctx.log.alert("Saved single cut.")
         else:
             with open(path, "a" if append else "w", newline='', encoding="utf8") as fp:
@@ -129,8 +118,8 @@ class Cut:
             column, the data is written to file as-is, with raw bytes preserved.
         """
         fp = io.StringIO(newline="")
-        if len(cuts) == 1 and len(cuts[0]) == 1:
-            v = cuts[0][0]
+        if len(cuts) == 1 and len(flows) == 1:
+            v = extract(cuts[0], flows[0])
             if isinstance(v, bytes):
                 fp.write(strutils.always_str(v))
             else:
@@ -138,9 +127,10 @@ class Cut:
             ctx.log.alert("Clipped single cut.")
         else:
             writer = csv.writer(fp)
-            for r in cuts:
+            for f in flows:
+                vals = [extract(c, f) for c in cuts]
                 writer.writerow(
-                    [strutils.always_str(c) or "" for c in r]  # type: ignore
+                    [strutils.always_str(v) or "" for v in vals]  # type: ignore
                 )
             ctx.log.alert("Clipped %s cuts as CSV." % len(cuts))
         pyperclip.copy(fp.getvalue())
diff --git a/mitmproxy/command.py b/mitmproxy/command.py
index 82bad4fa..7d7fa735 100644
--- a/mitmproxy/command.py
+++ b/mitmproxy/command.py
@@ -30,7 +30,46 @@ Cuts = typing.Sequence[
 
 
 class Cut(str):
-    pass
+    # This is an awkward location for these values, but it's better than having
+    # the console core import and depend on an addon. FIXME: Add a way for
+    # addons to add custom types and manage their completion and validation.
+    valid_prefixes = [
+        "request.method",
+        "request.scheme",
+        "request.host",
+        "request.http_version",
+        "request.port",
+        "request.path",
+        "request.url",
+        "request.text",
+        "request.content",
+        "request.raw_content",
+        "request.timestamp_start",
+        "request.timestamp_end",
+        "request.header[",
+
+        "response.status_code",
+        "response.reason",
+        "response.text",
+        "response.content",
+        "response.timestamp_start",
+        "response.timestamp_end",
+        "response.raw_content",
+        "response.header[",
+
+        "client_conn.address.port",
+        "client_conn.address.host",
+        "client_conn.tls_version",
+        "client_conn.sni",
+        "client_conn.ssl_established",
+
+        "server_conn.address.port",
+        "server_conn.address.host",
+        "server_conn.ip_address.host",
+        "server_conn.tls_version",
+        "server_conn.sni",
+        "server_conn.ssl_established",
+    ]
 
 
 class Path(str):
diff --git a/mitmproxy/test/tflow.py b/mitmproxy/test/tflow.py
index e754cb54..c3dab30c 100644
--- a/mitmproxy/test/tflow.py
+++ b/mitmproxy/test/tflow.py
@@ -53,6 +53,8 @@ def twebsocketflow(client_conn=True, server_conn=True, messages=True, err=None,
                 sec_websocket_version="13",
                 sec_websocket_key="1234",
             ),
+            timestamp_start=1,
+            timestamp_end=2,
             content=b''
         )
         resp = http.HTTPResponse(
@@ -64,6 +66,8 @@ def twebsocketflow(client_conn=True, server_conn=True, messages=True, err=None,
                 upgrade='websocket',
                 sec_websocket_accept=b'',
             ),
+            timestamp_start=1,
+            timestamp_end=2,
             content=b'',
         )
         handshake_flow = http.HTTPFlow(client_conn, server_conn)
diff --git a/mitmproxy/test/tutils.py b/mitmproxy/test/tutils.py
index 80e5b6fd..bcce547a 100644
--- a/mitmproxy/test/tutils.py
+++ b/mitmproxy/test/tutils.py
@@ -31,7 +31,9 @@ def treq(**kwargs):
         path=b"/path",
         http_version=b"HTTP/1.1",
         headers=http.Headers(((b"header", b"qvalue"), (b"content-length", b"7"))),
-        content=b"content"
+        content=b"content",
+        timestamp_start=1,
+        timestamp_end=2,
     )
     default.update(kwargs)
     return http.Request(**default)
@@ -48,8 +50,8 @@ def tresp(**kwargs):
         reason=b"OK",
         headers=http.Headers(((b"header-response", b"svalue"), (b"content-length", b"7"))),
         content=b"message",
-        timestamp_start=time.time(),
-        timestamp_end=time.time(),
+        timestamp_start=1,
+        timestamp_end=2,
     )
     default.update(kwargs)
     return http.Response(**default)
diff --git a/mitmproxy/tools/console/commander/commander.py b/mitmproxy/tools/console/commander/commander.py
index 5fc7dd12..b94d6f69 100644
--- a/mitmproxy/tools/console/commander/commander.py
+++ b/mitmproxy/tools/console/commander/commander.py
@@ -113,6 +113,19 @@ class CommandBuffer():
                     ),
                     parse = parts,
                 )
+            if last.type == typing.Sequence[mitmproxy.command.Cut]:
+                spec = parts[-1].value.split(",")
+                opts = []
+                for pref in mitmproxy.command.Cut.valid_prefixes:
+                    spec[-1] = pref
+                    opts.append(",".join(spec))
+                self.completion = CompletionState(
+                    completer = ListCompleter(
+                        parts[-1].value,
+                        opts,
+                    ),
+                    parse = parts,
+                )
             elif isinstance(last.type, mitmproxy.command.Choice):
                 self.completion = CompletionState(
                     completer = ListCompleter(
diff --git a/mitmproxy/tools/console/defaultkeys.py b/mitmproxy/tools/console/defaultkeys.py
index b845a3ae..7e078bbf 100644
--- a/mitmproxy/tools/console/defaultkeys.py
+++ b/mitmproxy/tools/console/defaultkeys.py
@@ -31,7 +31,7 @@ def map(km):
     km.add("A", "flow.resume @all", ["flowlist", "flowview"], "Resume all intercepted flows")
     km.add("a", "flow.resume @focus", ["flowlist", "flowview"], "Resume this intercepted flow")
     km.add(
-        "b", "console.command cut.save @focus s.content ",
+        "b", "console.command cut.save @focus response.content ",
         ["flowlist", "flowview"],
         "Save response body to file"
     )