From 1b44691d3345604f347ef2b0897914e932528a44 Mon Sep 17 00:00:00 2001 From: txrp0x9 <41499871+txrp0x9@users.noreply.github.com> Date: Fri, 5 Apr 2024 00:37:20 +0530 Subject: [PATCH] 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 * [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 --- CHANGELOG.md | 4 ++++ mitmproxy/addons/cut.py | 11 +++++++++++ mitmproxy/addons/export.py | 5 ++++- mitmproxy/websocket.py | 9 +++++++++ test/mitmproxy/addons/test_cut.py | 12 ++++++++++++ test/mitmproxy/addons/test_export.py | 14 ++++++++++++++ test/mitmproxy/test_websocket.py | 18 ++++++++++++++++++ 7 files changed, 72 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index aa70c89ea9..fa49236653 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/mitmproxy/addons/cut.py b/mitmproxy/addons/cut.py index c15ac3f539..c8be911272 100644 --- a/mitmproxy/addons/cut.py +++ b/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): diff --git a/mitmproxy/addons/export.py b/mitmproxy/addons/export.py index fb5ce2d5db..f1ac553f15 100644 --- a/mitmproxy/addons/export.py +++ b/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: diff --git a/mitmproxy/websocket.py b/mitmproxy/websocket.py index 6176cf0937..5558916dca 100644 --- a/mitmproxy/websocket.py +++ b/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"" + + def _get_formatted_messages(self) -> bytes: + return b"\n".join(m._format_ws_message() for m in self.messages) diff --git a/test/mitmproxy/addons/test_cut.py b/test/mitmproxy/addons/test_cut.py index 3f6dae3843..146bde170c 100644 --- a/test/mitmproxy/addons/test_cut.py +++ b/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" diff --git a/test/mitmproxy/addons/test_export.py b/test/mitmproxy/addons/test_export.py index f3dcb8d760..af1fcdf937 100644 --- a/test/mitmproxy/addons/test_export.py +++ b/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", diff --git a/test/mitmproxy/test_websocket.py b/test/mitmproxy/test_websocket.py index 08117b0fb0..80cb3d9b3e 100644 --- a/test/mitmproxy/test_websocket.py +++ b/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"