Fix saving of websocket flows (#6767)

* fix websocket message saving

* [autofix.ci] apply automated fixes

* added websocket as a supported format to export command, with tests for the same

* [autofix.ci] apply automated fixes

* added websocket message serialization in raw export, with test coverage

* [autofix.ci] apply automated fixes

* code suggestion fixes

Co-authored-by: Maximilian Hils <github@maximilianhils.com>

* [autofix.ci] apply automated fixes

* suggestion fixes

* fix merged code

* added tests for websocket export and cut.save

* [autofix.ci] apply automated fixes

* fix tests and add changes to changelog

* fix tests and add changes to changelog

* fix changelog

* fix changelog

* changelog addition

* changelog revert

* test fix

* [autofix.ci] apply automated fixes

* more test coverage

* [autofix.ci] apply automated fixes

* add changes to changelog

* add more test coverage

* [autofix.ci] apply automated fixes

* simplify

* [autofix.ci] apply automated fixes

---------

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

CHANGELOG.md

@@ -29,6 +29,10 @@
   ([#6764](https://github.com/mitmproxy/mitmproxy/pull/6764), @changsin)
 * Add primitive websocket interception and modification
   ([#6766](https://github.com/mitmproxy/mitmproxy/pull/6766), @errorxyz)
+* Add support for exporting websocket messages when using "raw" export.
+  ([#6767](https://github.com/mitmproxy/mitmproxy/pull/6767), @txrp0x9)
+* The "save body" feature now also includes WebSocket messages.
+  ([#6767](https://github.com/mitmproxy/mitmproxy/pull/6767), @txrp0x9)
 
 ## 07 March 2024: mitmproxy 10.2.4
 

mitmproxy/addons/cut.py

@@ -12,6 +12,7 @@
 from mitmproxy import command
 from mitmproxy import exceptions
 from mitmproxy import flow
+from mitmproxy import http
 from mitmproxy.log import ALERT
 
 logger = logging.getLogger(__name__)
@@ -28,6 +29,16 @@ def is_addr(v):
 
 
 def extract(cut: str, f: flow.Flow) -> str | bytes:
+    # Hack for https://github.com/mitmproxy/mitmproxy/issues/6721:
+    # Make "save body" keybind work for WebSocket flows.
+    # Ideally the keybind would be smarter and this here can get removed.
+    if (
+        isinstance(f, http.HTTPFlow)
+        and f.websocket
+        and cut in ("request.content", "response.content")
+    ):
+        return f.websocket._get_formatted_messages()
+
     path = cut.split(".")
     current: Any = f
     for i, spec in enumerate(path):

mitmproxy/addons/export.py

@@ -132,7 +132,10 @@ def raw(f: flow.Flow, separator=b"\r\n\r\n") -> bytes:
     )
 
     if request_present and response_present:
-        return b"".join([raw_request(f), separator, raw_response(f)])
+        parts = [raw_request(f), raw_response(f)]
+        if isinstance(f, http.HTTPFlow) and f.websocket:
+            parts.append(f.websocket._get_formatted_messages())
+        return separator.join(parts)
     elif request_present:
         return raw_request(f)
     elif response_present:

mitmproxy/websocket.py

@@ -92,6 +92,12 @@ def set_state(self, state: WebSocketMessageState) -> None:
         ) = state
         self.type = Opcode(typ)
 
+    def _format_ws_message(self) -> bytes:
+        if self.from_client:
+            return b"[OUTGOING] " + self.content
+        else:
+            return b"[INCOMING] " + self.content
+
     def __repr__(self):
         if self.type == Opcode.TEXT:
             return repr(self.content.decode(errors="replace"))
@@ -171,3 +177,6 @@ class WebSocketData(serializable.SerializableDataclass):
 
     def __repr__(self):
         return f"<WebSocketData ({len(self.messages)} messages)>"
+
+    def _get_formatted_messages(self) -> bytes:
+        return b"\n".join(m._format_ws_message() for m in self.messages)

test/mitmproxy/addons/test_cut.py

@@ -58,6 +58,18 @@ def test_extract(tdata):
     assert "CERTIFICATE" in cut.extract("server_conn.certificate_list", tf)
 
 
+def test_extract_websocket():
+    tf = tflow.twebsocketflow(messages=True)
+    extracted_request_content = cut.extract("request.content", tf)
+    extracted_response_content = cut.extract("response.content", tf)
+    assert b"hello binary" in extracted_request_content
+    assert b"hello text" in extracted_request_content
+    assert b"it's me" in extracted_request_content
+    assert b"hello binary" in extracted_response_content
+    assert b"hello text" in extracted_response_content
+    assert b"it's me" in extracted_response_content
+
+
 def test_extract_str():
     tf = tflow.tflow()
     tf.request.raw_content = b"\xff"

test/mitmproxy/addons/test_export.py

@@ -58,6 +58,11 @@ def udp_flow():
     return tflow.tudpflow()
 
 
+@pytest.fixture
+def websocket_flow():
+    return tflow.twebsocketflow()
+
+
 @pytest.fixture(scope="module")
 def export_curl():
     e = export.Export()
@@ -217,6 +222,11 @@ def test_udp(self, udp_flow):
         ):
             export.raw(udp_flow)
 
+    def test_websocket(self, websocket_flow):
+        assert b"hello binary" in export.raw(websocket_flow)
+        assert b"hello text" in export.raw(websocket_flow)
+        assert b"it's me" in export.raw(websocket_flow)
+
 
 class TestRawRequest:
     def test_get(self, get_request):
@@ -286,6 +296,10 @@ def test_export(tmp_path) -> None:
         assert qr(f)
         os.unlink(f)
 
+        e.file("raw", tflow.twebsocketflow(), f)
+        assert qr(f)
+        os.unlink(f)
+
 
 @pytest.mark.parametrize(
     "exception, log_message",

test/mitmproxy/test_websocket.py

@@ -15,6 +15,13 @@ def test_state(self):
         f2 = http.HTTPFlow.from_state(f.get_state())
         f2.set_state(f.get_state())
 
+    def test_formatting(self):
+        tf = tflow.twebsocketflow().websocket
+        formatted_messages = tf._get_formatted_messages()
+        assert b"[OUTGOING] hello binary" in formatted_messages
+        assert b"[OUTGOING] hello text" in formatted_messages
+        assert b"[INCOMING] it's me" in formatted_messages
+
 
 class TestWebSocketMessage:
     def test_basic(self):
@@ -43,3 +50,14 @@ def test_text(self):
             _ = bin.text
         with pytest.raises(AttributeError, match="do not have a 'text' attribute."):
             bin.text = "bar"
+
+    def test_message_formatting(self):
+        incoming_message = websocket.WebSocketMessage(
+            Opcode.BINARY, False, b"Test Incoming"
+        )
+        outgoing_message = websocket.WebSocketMessage(
+            Opcode.BINARY, True, b"Test OutGoing"
+        )
+
+        assert incoming_message._format_ws_message() == b"[INCOMING] Test Incoming"
+        assert outgoing_message._format_ws_message() == b"[OUTGOING] Test OutGoing"