From f96b41b6e624852613ae0daced8bc8076246d243 Mon Sep 17 00:00:00 2001
From: Maximilian Hils <git@maximilianhils.com>
Date: Sat, 11 Apr 2020 23:58:33 +0200
Subject: view addon: systematize different flow types

---
 mitmproxy/addons/view.py | 47 +++++++++++++++++++++++++++--------------------
 mitmproxy/flow.py        |  5 +++++
 mitmproxy/http.py        |  4 ++++
 mitmproxy/tcp.py         |  1 -
 4 files changed, 36 insertions(+), 21 deletions(-)

diff --git a/mitmproxy/addons/view.py b/mitmproxy/addons/view.py
index 9266fb9f..4b7e076b 100644
--- a/mitmproxy/addons/view.py
+++ b/mitmproxy/addons/view.py
@@ -21,8 +21,10 @@ from mitmproxy import command
 from mitmproxy import connections
 from mitmproxy import ctx
 from mitmproxy import io
-from mitmproxy import http  # noqa
-from mitmproxy import tcp  # noqa
+from mitmproxy import http
+from mitmproxy import tcp
+from mitmproxy.utils import human
+
 
 # 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,
@@ -39,7 +41,7 @@ class _OrderKey:
     def __init__(self, view):
         self.view = view
 
-    def generate(self, f: http.HTTPFlow) -> typing.Any:  # pragma: no cover
+    def generate(self, f: mitmproxy.flow.Flow) -> typing.Any:  # pragma: no cover
         pass
 
     def refresh(self, f):
@@ -69,44 +71,49 @@ class _OrderKey:
 
 
 class OrderRequestStart(_OrderKey):
-    def generate(self, f: http.HTTPFlow) -> int:
-        if isinstance(f, http.HTTPFlow):
-            return f.request.timestamp_start or 0
-        else:
-            return f.timestamp_start
+    def generate(self, f: mitmproxy.flow.Flow) -> float:
+        return f.timestamp_start
 
 
 class OrderRequestMethod(_OrderKey):
-    def generate(self, f: http.HTTPFlow) -> str:
+    def generate(self, f: mitmproxy.flow.Flow) -> str:
         if isinstance(f, http.HTTPFlow):
             return f.request.method
+        elif isinstance(f, tcp.TCPFlow):
+            return "TCP"
         else:
-            return "TCP" # Stub
+            raise NotImplementedError()
 
 
 class OrderRequestURL(_OrderKey):
-    def generate(self, f: http.HTTPFlow) -> str:
+    def generate(self, f: mitmproxy.flow.Flow) -> str:
         if isinstance(f, http.HTTPFlow):
             return f.request.url
+        elif isinstance(f, tcp.TCPFlow):
+            return human.format_address(f.server_conn.address)
         else:
-            return "f.server" # Stub
+            raise NotImplementedError()
 
 
 class OrderKeySize(_OrderKey):
-    def generate(self, f: http.HTTPFlow) -> int:
-        s = 0
+    def generate(self, f: mitmproxy.flow.Flow) -> int:
         if isinstance(f, http.HTTPFlow):
+            size = 0
             if f.request.raw_content:
-                s += len(f.request.raw_content)
+                size += len(f.request.raw_content)
             if f.response and f.response.raw_content:
-                s += len(f.response.raw_content)
+                size += len(f.response.raw_content)
+            return size
+        elif isinstance(f, tcp.TCPFlow):
+            size = 0
+            for message in f.messages:
+                size += len(message.content)
+            return size
         else:
-            s = 1337 # Stub
-        return s
-
+            raise NotImplementedError()
 
-matchall = flowfilter.parse(". | ~tcp")
 
+matchall = flowfilter.parse("~http | ~tcp")
 
 orders = [
     ("t", "time"),
diff --git a/mitmproxy/flow.py b/mitmproxy/flow.py
index 35d1a688..450667a6 100644
--- a/mitmproxy/flow.py
+++ b/mitmproxy/flow.py
@@ -180,3 +180,8 @@ class Flow(stateobject.StateObject):
         if self.reply.state == "taken":
             self.reply.ack()
             self.reply.commit()
+
+    @property
+    def timestamp_start(self) -> float:
+        """Start time of the flow."""
+        return self.client_conn.timestamp_start
diff --git a/mitmproxy/http.py b/mitmproxy/http.py
index 6b527e75..e9902224 100644
--- a/mitmproxy/http.py
+++ b/mitmproxy/http.py
@@ -173,6 +173,10 @@ class HTTPFlow(flow.Flow):
         s += ">"
         return s.format(flow=self)
 
+    @property
+    def timestamp_start(self) -> float:
+        return self.request.timestamp_start
+
     def copy(self):
         f = super().copy()
         if self.request:
diff --git a/mitmproxy/tcp.py b/mitmproxy/tcp.py
index d32aa3ef..264e46b3 100644
--- a/mitmproxy/tcp.py
+++ b/mitmproxy/tcp.py
@@ -39,7 +39,6 @@ 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]
-- 
cgit v1.2.3