diff --git a/docs/scripts/pdoc-template/module.html.jinja2 b/docs/scripts/pdoc-template/module.html.jinja2 index dbe3baf15..268c393fb 100644 --- a/docs/scripts/pdoc-template/module.html.jinja2 +++ b/docs/scripts/pdoc-template/module.html.jinja2 @@ -55,6 +55,10 @@ To document all event hooks, we do a bit of hackery: {% if doc.qualname.startswith("ServerConnectionHookData") and doc.name != "__init__" %} {{ default_is_public(doc) }} {% endif %} + {% elif doc.modulename == "mitmproxy.websocket" %} + {% if doc.qualname != "WebSocketMessage.type" %} + {{ default_is_public(doc) }} + {% endif %} {% else %} {{ default_is_public(doc) }} {% endif %} diff --git a/mitmproxy/websocket.py b/mitmproxy/websocket.py index 7fd0fcb02..00accbf55 100644 --- a/mitmproxy/websocket.py +++ b/mitmproxy/websocket.py @@ -20,18 +20,16 @@ class WebSocketMessage(serializable.Serializable): """ A single WebSocket message sent from one peer to the other. - Fragmented WebSocket messages are reassembled by mitmproxy and the + Fragmented WebSocket messages are reassembled by mitmproxy and then represented as a single instance of this class. The [WebSocket RFC](https://tools.ietf.org/html/rfc6455) specifies both text and binary messages. To avoid a whole class of nasty type confusion bugs, - mitmproxy stores all message contents as binary. If you need text, you can decode the `content` property: + mitmproxy stores all message contents as `bytes`. If you need a `str`, you can access the `text` property + on text messages: - >>> from wsproto.frame_protocol import Opcode - >>> if message.type == Opcode.TEXT: - >>> text = message.content.decode() - - Per the WebSocket spec, text messages always use UTF-8 encoding. + >>> if message.is_text: + >>> text = message.text """ from_client: bool @@ -40,8 +38,7 @@ class WebSocketMessage(serializable.Serializable): """ The message type, as per RFC 6455's [opcode](https://tools.ietf.org/html/rfc6455#section-5.2). - Note that mitmproxy will always store the message contents as *bytes*. - A dedicated `.text` property for text messages is planned, see https://github.com/mitmproxy/mitmproxy/pull/4486. + Mitmproxy currently only exposes messages assembled from `TEXT` and `BINARY` frames. """ content: bytes """A byte-string representing the content of this message.""" @@ -81,10 +78,39 @@ class WebSocketMessage(serializable.Serializable): else: return repr(self.content) + @property + def is_text(self) -> bool: + """ + `True` if this message is assembled from WebSocket `TEXT` frames, + `False` if it is assembled from `BINARY` frames. + """ + return self.type == Opcode.TEXT + def kill(self): # Likely to be replaced with .drop() in the future, see https://github.com/mitmproxy/mitmproxy/pull/4486 self.killed = True + @property + def text(self) -> str: + """ + The message content as text. + + This attribute is only available if `WebSocketMessage.is_text` is `True`. + + *See also:* `WebSocketMessage.content` + """ + if self.type != Opcode.TEXT: + raise AttributeError(f"{self.type.name.title()} WebSocket frames do not have a 'text' attribute.") + + return self.content.decode() + + @text.setter + def text(self, value: str) -> None: + if self.type != Opcode.TEXT: + raise AttributeError(f"{self.type.name.title()} WebSocket frames do not have a 'text' attribute.") + + self.content = value.encode() + class WebSocketData(stateobject.StateObject): """ @@ -97,9 +123,9 @@ class WebSocketData(stateobject.StateObject): closed_by_client: Optional[bool] = None """ - True if the client closed the connection, - False if the server closed the connection, - None if the connection is active. + `True` if the client closed the connection, + `False` if the server closed the connection, + `None` if the connection is active. """ close_code: Optional[int] = None """[Close Code](https://tools.ietf.org/html/rfc6455#section-7.1.5)""" diff --git a/test/mitmproxy/test_websocket.py b/test/mitmproxy/test_websocket.py index 26d7b0f6a..a47f6e4c3 100644 --- a/test/mitmproxy/test_websocket.py +++ b/test/mitmproxy/test_websocket.py @@ -1,3 +1,5 @@ +import pytest + from mitmproxy import http from mitmproxy import websocket from mitmproxy.test import tflow @@ -26,3 +28,18 @@ class TestWebSocketMessage: assert not m.killed m.kill() assert m.killed + + def test_text(self): + txt = websocket.WebSocketMessage(Opcode.TEXT, True, b"foo") + bin = websocket.WebSocketMessage(Opcode.BINARY, True, b"foo") + + assert txt.is_text + assert txt.text + txt.text = "bar" + assert txt.content == b"bar" + + assert not bin.is_text + with pytest.raises(AttributeError, match="do not have a 'text' attribute."): + _ = bin.text + with pytest.raises(AttributeError, match="do not have a 'text' attribute."): + bin.text = "bar"