codekingpro/portable-devtools

1"""Models for WebSocket protocol versions 13 and 8."""
2
3import json
4from enum import IntEnum
5from typing import Any, Callable, Final, NamedTuple, Optional, cast
6
7WS_DEFLATE_TRAILING: Final[bytes] = bytes([0x00, 0x00, 0xFF, 0xFF])
8
9
10class WSCloseCode(IntEnum):
11    OK = 1000
12    GOING_AWAY = 1001
13    PROTOCOL_ERROR = 1002
14    UNSUPPORTED_DATA = 1003
15    ABNORMAL_CLOSURE = 1006
16    INVALID_TEXT = 1007
17    POLICY_VIOLATION = 1008
18    MESSAGE_TOO_BIG = 1009
19    MANDATORY_EXTENSION = 1010
20    INTERNAL_ERROR = 1011
21    SERVICE_RESTART = 1012
22    TRY_AGAIN_LATER = 1013
23    BAD_GATEWAY = 1014
24
25
26class WSMsgType(IntEnum):
27    # websocket spec types
28    CONTINUATION = 0x0
29    TEXT = 0x1
30    BINARY = 0x2
31    PING = 0x9
32    PONG = 0xA
33    CLOSE = 0x8
34
35    # aiohttp specific types
36    CLOSING = 0x100
37    CLOSED = 0x101
38    ERROR = 0x102
39
40    text = TEXT
41    binary = BINARY
42    ping = PING
43    pong = PONG
44    close = CLOSE
45    closing = CLOSING
46    closed = CLOSED
47    error = ERROR
48
49
50class WSMessage(NamedTuple):
51    type: WSMsgType
52    # To type correctly, this would need some kind of tagged union for each type.
53    data: Any
54    extra: Optional[str]
55
56    def json(self, *, loads: Callable[[Any], Any] = json.loads) -> Any:
57        """Return parsed JSON data.
58
59        .. versionadded:: 0.22
60        """
61        return loads(self.data)
62
63
64# Constructing the tuple directly to avoid the overhead of
65# the lambda and arg processing since NamedTuples are constructed
66# with a run time built lambda
67# https://github.com/python/cpython/blob/d83fcf8371f2f33c7797bc8f5423a8bca8c46e5c/Lib/collections/__init__.py#L441
68WS_CLOSED_MESSAGE = tuple.__new__(WSMessage, (WSMsgType.CLOSED, None, None))
69WS_CLOSING_MESSAGE = tuple.__new__(WSMessage, (WSMsgType.CLOSING, None, None))
70
71
72class WebSocketError(Exception):
73    """WebSocket protocol parser error."""
74
75    def __init__(self, code: int, message: str) -> None:
76        self.code = code
77        super().__init__(code, message)
78
79    def __str__(self) -> str:
80        return cast(str, self.args[1])
81
82
83class WSHandshakeError(Exception):
84    """WebSocket protocol handshake error."""
85