2025-12-25 upload

venv/Lib/site-packages/mitmproxy/websocket.py

@@ -0,0 +1,182 @@
+"""
+Mitmproxy used to have its own WebSocketFlow type until mitmproxy 6, but now WebSocket connections now are represented
+as HTTP flows as well. They can be distinguished from regular HTTP requests by having the
+`mitmproxy.http.HTTPFlow.websocket` attribute set.
+
+This module only defines the classes for individual `WebSocketMessage`s and the `WebSocketData` container.
+"""
+
+import time
+import warnings
+from dataclasses import dataclass
+from dataclasses import field
+
+from wsproto.frame_protocol import Opcode
+
+from mitmproxy.coretypes import serializable
+
+WebSocketMessageState = tuple[int, bool, bytes, float, bool, bool]
+
+
+class WebSocketMessage(serializable.Serializable):
+    """
+    A single WebSocket message sent from one peer to the other.
+
+    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 `bytes`. If you need a `str`, you can access the `text` property
+    on text messages:
+
+    >>> if message.is_text:
+    >>>     text = message.text
+    """
+
+    from_client: bool
+    """True if this messages was sent by the client."""
+    type: Opcode
+    """
+    The message type, as per RFC 6455's [opcode](https://tools.ietf.org/html/rfc6455#section-5.2).
+
+    Mitmproxy currently only exposes messages assembled from `TEXT` and `BINARY` frames.
+    """
+    content: bytes
+    """A byte-string representing the content of this message."""
+    timestamp: float
+    """Timestamp of when this message was received or created."""
+    dropped: bool
+    """True if the message has not been forwarded by mitmproxy, False otherwise."""
+    injected: bool
+    """True if the message was injected and did not originate from a client/server, False otherwise"""
+
+    def __init__(
+        self,
+        type: int | Opcode,
+        from_client: bool,
+        content: bytes,
+        timestamp: float | None = None,
+        dropped: bool = False,
+        injected: bool = False,
+    ) -> None:
+        self.from_client = from_client
+        self.type = Opcode(type)
+        self.content = content
+        self.timestamp: float = timestamp or time.time()
+        self.dropped = dropped
+        self.injected = injected
+
+    @classmethod
+    def from_state(cls, state: WebSocketMessageState):
+        return cls(*state)
+
+    def get_state(self) -> WebSocketMessageState:
+        return (
+            int(self.type),
+            self.from_client,
+            self.content,
+            self.timestamp,
+            self.dropped,
+            self.injected,
+        )
+
+    def set_state(self, state: WebSocketMessageState) -> None:
+        (
+            typ,
+            self.from_client,
+            self.content,
+            self.timestamp,
+            self.dropped,
+            self.injected,
+        ) = 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"))
+        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 drop(self):
+        """Drop this message, i.e. don't forward it to the other peer."""
+        self.dropped = True
+
+    def kill(self):  # pragma: no cover
+        """A deprecated alias for `.drop()`."""
+        warnings.warn(
+            "WebSocketMessage.kill() is deprecated, use .drop() instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        self.drop()
+
+    @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()
+
+
+@dataclass
+class WebSocketData(serializable.SerializableDataclass):
+    """
+    A data container for everything related to a single WebSocket connection.
+    This is typically accessed as `mitmproxy.http.HTTPFlow.websocket`.
+    """
+
+    messages: list[WebSocketMessage] = field(default_factory=list)
+    """All `WebSocketMessage`s transferred over this connection."""
+
+    closed_by_client: bool | None = None
+    """
+    `True` if the client closed the connection,
+    `False` if the server closed the connection,
+    `None` if the connection is active.
+    """
+    close_code: int | None = None
+    """[Close Code](https://tools.ietf.org/html/rfc6455#section-7.1.5)"""
+    close_reason: str | None = None
+    """[Close Reason](https://tools.ietf.org/html/rfc6455#section-7.1.6)"""
+
+    timestamp_end: float | None = None
+    """*Timestamp:* WebSocket connection closed."""
+
+    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)