Merge pull request #3970 from nikitastupin/master

Integrate contentviews to TCP flow details

4 files changed, 77 insertions, 9 deletions

diff --git a/mitmproxy/contentviews/__init__.py b/mitmproxy/contentviews/__init__.py
index fc38543d..f5c31eef 100644
--- a/mitmproxy/contentviews/__init__.py
+++ b/mitmproxy/contentviews/__init__.py
@@ -116,6 +116,19 @@ def get_message_content_view(viewname, message, flow):
     return description, lines, error
 
 
+def get_tcp_content_view(viewname: str, data: bytes):
+    viewmode = get(viewname)
+    if not viewmode:
+        viewmode = get("auto")
+
+    # https://github.com/mitmproxy/mitmproxy/pull/3970#issuecomment-623024447
+    assert viewmode
+
+    description, lines, error = get_content_view(viewmode, data)
+
+    return description, lines, error
+
+
 def get_content_view(viewmode: View, data: bytes, **metadata):
     """
         Args:
diff --git a/mitmproxy/tools/console/common.py b/mitmproxy/tools/console/common.py
index cba3a355..46800fdd 100644
--- a/mitmproxy/tools/console/common.py
+++ b/mitmproxy/tools/console/common.py
@@ -104,6 +104,8 @@ if urwid.util.detected_encoding:
     SYMBOL_UP = u"\u21E7"
     SYMBOL_DOWN = u"\u21E9"
     SYMBOL_ELLIPSIS = u"\u2026"
+    SYMBOL_FROM_CLIENT = u"\u21d2"
+    SYMBOL_TO_CLIENT = u"\u21d0"
 else:
     SYMBOL_REPLAY = u"[r]"
     SYMBOL_RETURN = u"<-"
@@ -111,6 +113,8 @@ else:
     SYMBOL_UP = "^"
     SYMBOL_DOWN = " "
     SYMBOL_ELLIPSIS = "~"
+    SYMBOL_FROM_CLIENT = u"->"
+    SYMBOL_TO_CLIENT = u"<-"
 
 SCHEME_STYLES = {
     'http': 'scheme_http',
diff --git a/mitmproxy/tools/console/flowview.py b/mitmproxy/tools/console/flowview.py
index 3fef70ce..466ff0c2 100644
--- a/mitmproxy/tools/console/flowview.py
+++ b/mitmproxy/tools/console/flowview.py
@@ -108,6 +108,26 @@ class FlowDetails(tabs.Tabs):
         assert isinstance(flow, http.HTTPFlow)
         return self.conn_text(flow.response)
 
+    def _contentview_status_bar(self, description: str, viewmode: str):
+        cols = [
+            urwid.Text(
+                [
+                    ("heading", description),
+                ]
+            ),
+            urwid.Text(
+                [
+                    " ",
+                    ('heading', "["),
+                    ('heading_key', "m"),
+                    ('heading', (":%s]" % viewmode)),
+                ],
+                align="right"
+            )
+        ]
+        contentview_status_bar = urwid.AttrWrap(urwid.Columns(cols), "heading")
+        return contentview_status_bar
+
     def view_tcp_stream(self) -> urwid.Widget:
         flow = self.flow
         assert isinstance(flow, tcp.TCPFlow)
@@ -115,6 +135,10 @@ class FlowDetails(tabs.Tabs):
         if not flow.messages:
             return searchable.Searchable([urwid.Text(("highlight", "No messages."))])
 
+        viewmode = self.master.commands.call("console.flowview.mode")
+
+        # Merge adjacent TCP "messages". For detailed explanation of this code block see:
+        # https://github.com/mitmproxy/mitmproxy/pull/3970/files/469bd32582f764f9a29607efa4f5b04bd87961fb#r418670880
         from_client = None
         messages = []
         for message in flow.messages:
@@ -124,17 +148,25 @@ class FlowDetails(tabs.Tabs):
             else:
                 messages[-1] += message.content
 
+        widget_lines = []
+
         from_client = flow.messages[0].from_client
-        parts = []
-        for message in messages:
-            parts.append(
-                (
-                    "head" if from_client else "key",
-                    message
-                )
-            )
+        for m in messages:
+            _, lines, _ = contentviews.get_tcp_content_view(viewmode, m)
+
+            for line in lines:
+                if from_client:
+                    line.insert(0, ("from_client", f"{common.SYMBOL_FROM_CLIENT} "))
+                else:
+                    line.insert(0, ("to_client", f"{common.SYMBOL_TO_CLIENT} "))
+
+                widget_lines.append(urwid.Text(line))
+
             from_client = not from_client
-        return searchable.Searchable([urwid.Text(parts)])
+
+        widget_lines.insert(0, self._contentview_status_bar(viewmode.capitalize(), viewmode))
+
+        return searchable.Searchable(widget_lines)
 
     def view_details(self):
         return flowdetailview.flowdetails(self.view, self.flow)
diff --git a/mitmproxy/tools/console/palettes.py b/mitmproxy/tools/console/palettes.py
index a680a5a7..15745df5 100644
--- a/mitmproxy/tools/console/palettes.py
+++ b/mitmproxy/tools/console/palettes.py
@@ -37,6 +37,9 @@ class Palette:
         # JSON view
         'json_string', 'json_number', 'json_boolean',
 
+        # TCP flow details
+        'from_client', 'to_client',
+
         # Grid Editor
         'focusfield', 'focusfield_error', 'field_error', 'editfield',
 
@@ -177,6 +180,10 @@ class LowDark(Palette):
         json_number = ('light magenta', 'default'),
         json_boolean = ('dark magenta', 'default'),
 
+        # TCP flow details
+        from_client = ('light blue', 'default'),
+        to_client = ('light red', 'default'),
+
         # Grid Editor
         focusfield = ('black', 'light gray'),
         focusfield_error = ('dark red', 'light gray'),
@@ -282,6 +289,10 @@ class LowLight(Palette):
         json_number = ('light magenta', 'default'),
         json_boolean = ('dark magenta', 'default'),
 
+        # TCP flow details
+        from_client = ('dark blue', 'default'),
+        to_client = ('dark red', 'default'),
+
         # Grid Editor
         focusfield = ('black', 'light gray'),
         focusfield_error = ('dark red', 'light gray'),
@@ -397,6 +408,10 @@ class SolarizedLight(LowLight):
         json_number = (sol_blue, 'default'),
         json_boolean = (sol_magenta, 'default'),
 
+        # TCP flow details
+        from_client = (sol_blue, 'default'),
+        to_client = (sol_red, 'default'),
+
         # Grid Editor
         focusfield = (sol_base00, sol_base2),
         focusfield_error = (sol_red, sol_base2),
@@ -475,6 +490,10 @@ class SolarizedDark(LowDark):
         json_number = (sol_blue, 'default'),
         json_boolean = (sol_magenta, 'default'),
 
+        # TCP flow details
+        from_client = (sol_blue, 'default'),
+        to_client = (sol_red, 'default'),
+
         # Grid Editor
         focusfield = (sol_base0, sol_base02),
         focusfield_error = (sol_red, sol_base02),