mitmproxy.websocket

  1"""
  2Mitmproxy used to have its own WebSocketFlow type until mitmproxy 6, but now WebSocket connections now are represented
  3as HTTP flows as well. They can be distinguished from regular HTTP requests by having the
  4`mitmproxy.http.HTTPFlow.websocket` attribute set.
  5
  6This module only defines the classes for individual `WebSocketMessage`s and the `WebSocketData` container.
  7"""
  8
  9import time
 10import warnings
 11from dataclasses import dataclass
 12from dataclasses import field
 13
 14from wsproto.frame_protocol import Opcode
 15
 16from mitmproxy.coretypes import serializable
 17
 18WebSocketMessageState = tuple[int, bool, bytes, float, bool, bool]
 19
 20
 21class WebSocketMessage(serializable.Serializable):
 22    """
 23    A single WebSocket message sent from one peer to the other.
 24
 25    Fragmented WebSocket messages are reassembled by mitmproxy and then
 26    represented as a single instance of this class.
 27
 28    The [WebSocket RFC](https://tools.ietf.org/html/rfc6455) specifies both
 29    text and binary messages. To avoid a whole class of nasty type confusion bugs,
 30    mitmproxy stores all message contents as `bytes`. If you need a `str`, you can access the `text` property
 31    on text messages:
 32
 33    >>> if message.is_text:
 34    >>>     text = message.text
 35    """
 36
 37    from_client: bool
 38    """True if this messages was sent by the client."""
 39    type: Opcode
 40    """
 41    The message type, as per RFC 6455's [opcode](https://tools.ietf.org/html/rfc6455#section-5.2).
 42
 43    Mitmproxy currently only exposes messages assembled from `TEXT` and `BINARY` frames.
 44    """
 45    content: bytes
 46    """A byte-string representing the content of this message."""
 47    timestamp: float
 48    """Timestamp of when this message was received or created."""
 49    dropped: bool
 50    """True if the message has not been forwarded by mitmproxy, False otherwise."""
 51    injected: bool
 52    """True if the message was injected and did not originate from a client/server, False otherwise"""
 53
 54    def __init__(
 55        self,
 56        type: int | Opcode,
 57        from_client: bool,
 58        content: bytes,
 59        timestamp: float | None = None,
 60        dropped: bool = False,
 61        injected: bool = False,
 62    ) -> None:
 63        self.from_client = from_client
 64        self.type = Opcode(type)
 65        self.content = content
 66        self.timestamp: float = timestamp or time.time()
 67        self.dropped = dropped
 68        self.injected = injected
 69
 70    @classmethod
 71    def from_state(cls, state: WebSocketMessageState):
 72        return cls(*state)
 73
 74    def get_state(self) -> WebSocketMessageState:
 75        return (
 76            int(self.type),
 77            self.from_client,
 78            self.content,
 79            self.timestamp,
 80            self.dropped,
 81            self.injected,
 82        )
 83
 84    def set_state(self, state: WebSocketMessageState) -> None:
 85        (
 86            typ,
 87            self.from_client,
 88            self.content,
 89            self.timestamp,
 90            self.dropped,
 91            self.injected,
 92        ) = state
 93        self.type = Opcode(typ)
 94
 95    def _format_ws_message(self) -> bytes:
 96        if self.from_client:
 97            return b"[OUTGOING] " + self.content
 98        else:
 99            return b"[INCOMING] " + self.content
100
101    def __repr__(self):
102        if self.type == Opcode.TEXT:
103            return repr(self.content.decode(errors="replace"))
104        else:
105            return repr(self.content)
106
107    @property
108    def is_text(self) -> bool:
109        """
110        `True` if this message is assembled from WebSocket `TEXT` frames,
111        `False` if it is assembled from `BINARY` frames.
112        """
113        return self.type == Opcode.TEXT
114
115    def drop(self):
116        """Drop this message, i.e. don't forward it to the other peer."""
117        self.dropped = True
118
119    def kill(self):  # pragma: no cover
120        """A deprecated alias for `.drop()`."""
121        warnings.warn(
122            "WebSocketMessage.kill() is deprecated, use .drop() instead.",
123            DeprecationWarning,
124            stacklevel=2,
125        )
126        self.drop()
127
128    @property
129    def text(self) -> str:
130        """
131        The message content as text.
132
133        This attribute is only available if `WebSocketMessage.is_text` is `True`.
134
135        *See also:* `WebSocketMessage.content`
136        """
137        if self.type != Opcode.TEXT:
138            raise AttributeError(
139                f"{self.type.name.title()} WebSocket frames do not have a 'text' attribute."
140            )
141
142        return self.content.decode()
143
144    @text.setter
145    def text(self, value: str) -> None:
146        if self.type != Opcode.TEXT:
147            raise AttributeError(
148                f"{self.type.name.title()} WebSocket frames do not have a 'text' attribute."
149            )
150
151        self.content = value.encode()
152
153
154@dataclass
155class WebSocketData(serializable.SerializableDataclass):
156    """
157    A data container for everything related to a single WebSocket connection.
158    This is typically accessed as `mitmproxy.http.HTTPFlow.websocket`.
159    """
160
161    messages: list[WebSocketMessage] = field(default_factory=list)
162    """All `WebSocketMessage`s transferred over this connection."""
163
164    closed_by_client: bool | None = None
165    """
166    `True` if the client closed the connection,
167    `False` if the server closed the connection,
168    `None` if the connection is active.
169    """
170    close_code: int | None = None
171    """[Close Code](https://tools.ietf.org/html/rfc6455#section-7.1.5)"""
172    close_reason: str | None = None
173    """[Close Reason](https://tools.ietf.org/html/rfc6455#section-7.1.6)"""
174
175    timestamp_end: float | None = None
176    """*Timestamp:* WebSocket connection closed."""
177
178    def __repr__(self):
179        return f"<WebSocketData ({len(self.messages)} messages)>"
180
181    def _get_formatted_messages(self) -> bytes:
182        return b"\n".join(m._format_ws_message() for m in self.messages)