From 743cb949e8f5df5029b4b364d3d9692d1b6eb5d4 Mon Sep 17 00:00:00 2001
From: Nikita Stupin <18281368+nikitastupin@users.noreply.github.com>
Date: Fri, 10 Apr 2020 21:07:21 +0300
Subject: Display TCP flows at flow list!

---
 mitmproxy/addons/view.py          | 41 +++++++++++++++++++++------
 mitmproxy/tcp.py                  |  1 +
 mitmproxy/tools/console/common.py | 58 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 92 insertions(+), 8 deletions(-)

diff --git a/mitmproxy/addons/view.py b/mitmproxy/addons/view.py
index 1d57d781..9266fb9f 100644
--- a/mitmproxy/addons/view.py
+++ b/mitmproxy/addons/view.py
@@ -22,6 +22,7 @@ from mitmproxy import connections
 from mitmproxy import ctx
 from mitmproxy import io
 from mitmproxy import http  # noqa
+from mitmproxy import tcp  # noqa
 
 # The underlying sorted list implementation expects the sort key to be stable
 # for the lifetime of the object. However, if we sort by size, for instance,
@@ -69,30 +70,42 @@ class _OrderKey:
 
 class OrderRequestStart(_OrderKey):
     def generate(self, f: http.HTTPFlow) -> int:
-        return f.request.timestamp_start or 0
+        if isinstance(f, http.HTTPFlow):
+            return f.request.timestamp_start or 0
+        else:
+            return f.timestamp_start
 
 
 class OrderRequestMethod(_OrderKey):
     def generate(self, f: http.HTTPFlow) -> str:
-        return f.request.method
+        if isinstance(f, http.HTTPFlow):
+            return f.request.method
+        else:
+            return "TCP" # Stub
 
 
 class OrderRequestURL(_OrderKey):
     def generate(self, f: http.HTTPFlow) -> str:
-        return f.request.url
+        if isinstance(f, http.HTTPFlow):
+            return f.request.url
+        else:
+            return "f.server" # Stub
 
 
 class OrderKeySize(_OrderKey):
     def generate(self, f: http.HTTPFlow) -> int:
         s = 0
-        if f.request.raw_content:
-            s += len(f.request.raw_content)
-        if f.response and f.response.raw_content:
-            s += len(f.response.raw_content)
+        if isinstance(f, http.HTTPFlow):
+            if f.request.raw_content:
+                s += len(f.request.raw_content)
+            if f.response and f.response.raw_content:
+                s += len(f.response.raw_content)
+        else:
+            s = 1337 # Stub
         return s
 
 
-matchall = flowfilter.parse(".")
+matchall = flowfilter.parse(". | ~tcp")
 
 
 orders = [
@@ -555,6 +568,18 @@ class View(collections.abc.Sequence):
     def kill(self, f):
         self.update([f])
 
+    def tcp_start(self,f):
+        self.add([f])
+
+    def tcp_message(self, f):
+        self.update([f])
+
+    def tcp_error(self, f):
+        self.update([f])
+
+    def tcp_end(self, f):
+        self.update([f])
+
     def update(self, flows: typing.Sequence[mitmproxy.flow.Flow]) -> None:
         """
             Updates a list of flows. If flow is not in the state, it's ignored.
diff --git a/mitmproxy/tcp.py b/mitmproxy/tcp.py
index 264e46b3..d32aa3ef 100644
--- a/mitmproxy/tcp.py
+++ b/mitmproxy/tcp.py
@@ -39,6 +39,7 @@ class TCPFlow(flow.Flow):
     def __init__(self, client_conn, server_conn, live=None):
         super().__init__("tcp", client_conn, server_conn, live)
         self.messages: List[TCPMessage] = []
+        self.timestamp_start: float = time.time()
 
     _stateobject_attributes = flow.Flow._stateobject_attributes.copy()
     _stateobject_attributes["messages"] = List[TCPMessage]
diff --git a/mitmproxy/tools/console/common.py b/mitmproxy/tools/console/common.py
index 3dce8363..9ed5795f 100644
--- a/mitmproxy/tools/console/common.py
+++ b/mitmproxy/tools/console/common.py
@@ -10,6 +10,7 @@ import urwid
 import urwid.util
 
 from mitmproxy.utils import human
+from mitmproxy.tcp import TCPFlow
 
 # Detect Windows Subsystem for Linux
 IS_WSL = "Microsoft" in platform.platform()
@@ -507,7 +508,64 @@ def raw_format_table(f):
     return urwid.Pile(pile)
 
 
+# TODO: this function can replace repeated code in raw_format_table() in the future
+def raw_format_cursor(f):
+    cursor = [" ", "focus"]
+    if f["focus"]:
+        cursor[0] = ">"
+    return fcol(*cursor)
+
+
+# TODO: this function can replace repeated code in raw_format_table() in the future
+def raw_format_timestamp(timestamp, extended):
+    if extended:
+        s = human.format_timestamp(timestamp)
+    else:
+        s = datetime.datetime.fromtimestamp(time.mktime(time.localtime(timestamp))).strftime("%H:%M:%S")
+    return fcol(s, "title")
+
+
+@lru_cache(maxsize=800)
+def raw_format_tcp_table(f):
+    # If you'll remove this line TypeError: unhashable type: 'dict' will occur
+    # because of @lru_cache
+    f = dict(f)
+
+    pile = []
+
+    columns = [
+        raw_format_cursor(f),
+        raw_format_timestamp(f["timestamp_start"], f["extended"]),
+        fcol("TCP", "tcp"),
+        fcol(f["client"], "client"),
+        fcol("---", "direction"),
+        fcol(f["server"], "server"),
+    ]
+
+    m = [c for c in columns]
+
+    pile.append(urwid.Columns(m, dividechars=1))
+
+    return urwid.Pile(pile)
+
+
 def format_flow(f, focus, extended=False, hostheader=False, cols=False, layout='default'):
+
+    if isinstance(f, TCPFlow):
+        d = dict(
+            focus=focus,
+            extended=extended,
+            timestamp_start=f.timestamp_start,
+            client=human.format_address(f.client_conn.address),
+            server=human.format_address(f.server_conn.address),
+        )
+
+        # If you'll remove this line TypeError: unhashable type: 'dict' will occur
+        # because of @lru_cache.
+        t = tuple(sorted(d.items()))
+
+        return raw_format_tcp_table(t)
+
     acked = False
     if f.reply and f.reply.state == "committed":
         acked = True
-- 
cgit v1.2.3