codekingpro/portable-devtools

1"""2Mitmproxy used to have its own WebSocketFlow type until mitmproxy 6, but now WebSocket connections now are represented3as HTTP flows as well. They can be distinguished from regular HTTP requests by having the4`mitmproxy.http.HTTPFlow.websocket` attribute set.5 6This module only defines the classes for individual `WebSocketMessage`s and the `WebSocketData` container.7"""8 9import time10import warnings11from dataclasses import dataclass12from dataclasses import field13 14from wsproto.frame_protocol import Opcode15 16from mitmproxy.coretypes import serializable17 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 then26    represented as a single instance of this class.27 28    The [WebSocket RFC](https://tools.ietf.org/html/rfc6455) specifies both29    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` property31    on text messages:32 33    >>> if message.is_text:34    >>>     text = message.text35    """36 37    from_client: bool38    """True if this messages was sent by the client."""39    type: Opcode40    """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: bytes46    """A byte-string representing the content of this message."""47    timestamp: float48    """Timestamp of when this message was received or created."""49    dropped: bool50    """True if the message has not been forwarded by mitmproxy, False otherwise."""51    injected: bool52    """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_client64        self.type = Opcode(type)65        self.content = content66        self.timestamp: float = timestamp or time.time()67        self.dropped = dropped68        self.injected = injected69 70    @classmethod71    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        ) = state93        self.type = Opcode(typ)94 95    def _format_ws_message(self) -> bytes:96        if self.from_client:97            return b"[OUTGOING] " + self.content98        else:99            return b"[INCOMING] " + self.content100 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    @property108    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.TEXT114 115    def drop(self):116        """Drop this message, i.e. don't forward it to the other peer."""117        self.dropped = True118 119    def kill(self):  # pragma: no cover120        """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    @property129    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.setter145    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@dataclass155class 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 = None165    """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 = None171    """[Close Code](https://tools.ietf.org/html/rfc6455#section-7.1.5)"""172    close_reason: str | None = None173    """[Close Reason](https://tools.ietf.org/html/rfc6455#section-7.1.6)"""174 175    timestamp_end: float | None = None176    """*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)183