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