Add primitive websocket interception and modification (#6766)

* Add primitive websocket interception and modification

* Update CHANGELOG.md

* Fix UI tab renaming on interception

* [autofix.ci] apply automated fixes

* Improve code readability

* [autofix.ci] apply automated fixes

* Improve code readability

* [autofix.ci] apply automated fixes

* Improve code readability

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

CHANGELOG.md

@@ -21,6 +21,8 @@
   ([#6749](https://github.com/mitmproxy/mitmproxy/pull/6749), @mhils)
   * Add button to close flow details panel
   ([#6734](https://github.com/mitmproxy/mitmproxy/pull/6734), @lups2000)
+* Add primitive websocket interception and modification
+  ([#6766](https://github.com/mitmproxy/mitmproxy/pull/6766), @errorxyz)
 
 ## 07 March 2024: mitmproxy 10.2.4
 

mitmproxy/addons/intercept.py

@@ -58,3 +58,6 @@ class Intercept:
 
     def dns_response(self, f):
         self.process_flow(f)
+
+    def websocket_message(self, f):
+        self.process_flow(f)

mitmproxy/tools/console/consoleaddons.py

@@ -398,6 +398,8 @@ class ConsoleAddon:
                 "set-cookies",
                 "url",
             ]
+            if flow.websocket:
+                focus_options.append("websocket-message")
         elif isinstance(flow, dns.DNSFlow):
             raise exceptions.CommandError(
                 "Cannot edit DNS flows yet, please submit a patch."
@@ -469,6 +471,10 @@ class ConsoleAddon:
             message = flow.messages[-1]
             c = self.master.spawn_editor(message.content or b"")
             message.content = c.rstrip(b"\n")
+        elif flow_part == "websocket-message":
+            message = flow.websocket.messages[-1]
+            c = self.master.spawn_editor(message.content or b"")
+            message.content = c.rstrip(b"\n")
 
     def _grideditor(self):
         gewidget = self.master.window.current("grideditor")

mitmproxy/tools/console/flowview.py

@@ -117,7 +117,14 @@ class FlowDetails(tabs.Tabs):
     def tab_http_response(self):
         flow = self.flow
         assert isinstance(flow, http.HTTPFlow)
-        if self.flow.intercepted and flow.response:
+
+        # there is no good way to detect what part of the flow is intercepted,
+        # so we apply some heuristics to see if it's the HTTP response.
+        websocket_started = flow.websocket and len(flow.websocket.messages) != 0
+        response_is_intercepted = (
+            self.flow.intercepted and flow.response and not websocket_started
+        )
+        if response_is_intercepted:
             return "Response intercepted"
         else:
             return "Response"
@@ -145,7 +152,14 @@ class FlowDetails(tabs.Tabs):
         return "UDP Stream"
 
     def tab_websocket_messages(self):
-        return "WebSocket Messages"
+        flow = self.flow
+        assert isinstance(flow, http.HTTPFlow)
+        assert flow.websocket
+
+        if self.flow.intercepted and len(flow.websocket.messages) != 0:
+            return "WebSocket Messages intercepted"
+        else:
+            return "WebSocket Messages"
 
     def tab_details(self):
         return "Detail"

test/mitmproxy/addons/test_intercept.py

@@ -89,3 +89,17 @@ async def test_udp():
         f = tflow.tudpflow()
         await tctx.cycle(r, f)
         assert not f.intercepted
+
+
+async def test_websocket_message():
+    r = intercept.Intercept()
+    with taddons.context(r) as tctx:
+        tctx.configure(r, intercept='~b "hello binary"')
+        f = tflow.twebsocketflow()
+        await tctx.cycle(r, f)
+        assert f.intercepted
+
+        tctx.configure(r, intercept_active=False)
+        f = tflow.twebsocketflow()
+        await tctx.cycle(r, f)
+        assert not f.intercepted