ydb/contrib/python/websocket-client/py3/websocket/_abnf.py

import array
import os
import struct
import sys
from threading import Lock
from typing import Callable, Optional, Union

from ._exceptions import WebSocketPayloadException, WebSocketProtocolException
from ._utils import validate_utf8

"""
_abnf.py
websocket - WebSocket client library for Python

Copyright 2024 engn33r

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

try:
    # If wsaccel is available, use compiled routines to mask data.
    # wsaccel only provides around a 10% speed boost compared
    # to the websocket-client _mask() implementation.
    # Note that wsaccel is unmaintained.
    from wsaccel.xormask import XorMaskerSimple

    def _mask(mask_value: array.array, data_value: array.array) -> bytes:
        mask_result: bytes = XorMaskerSimple(mask_value).process(data_value)
        return mask_result

except ImportError:
    # wsaccel is not available, use websocket-client _mask()
    native_byteorder = sys.byteorder

    def _mask(mask_value: array.array, data_value: array.array) -> bytes:
        datalen = len(data_value)
        int_data_value = int.from_bytes(data_value, native_byteorder)
        int_mask_value = int.from_bytes(
            mask_value * (datalen // 4) + mask_value[: datalen % 4], native_byteorder
        )
        return (int_data_value ^ int_mask_value).to_bytes(datalen, native_byteorder)


__all__ = [
    "ABNF",
    "continuous_frame",
    "frame_buffer",
    "STATUS_NORMAL",
    "STATUS_GOING_AWAY",
    "STATUS_PROTOCOL_ERROR",
    "STATUS_UNSUPPORTED_DATA_TYPE",
    "STATUS_STATUS_NOT_AVAILABLE",
    "STATUS_ABNORMAL_CLOSED",
    "STATUS_INVALID_PAYLOAD",
    "STATUS_POLICY_VIOLATION",
    "STATUS_MESSAGE_TOO_BIG",
    "STATUS_INVALID_EXTENSION",
    "STATUS_UNEXPECTED_CONDITION",
    "STATUS_BAD_GATEWAY",
    "STATUS_TLS_HANDSHAKE_ERROR",
]

# closing frame status codes.
STATUS_NORMAL = 1000
STATUS_GOING_AWAY = 1001
STATUS_PROTOCOL_ERROR = 1002
STATUS_UNSUPPORTED_DATA_TYPE = 1003
STATUS_STATUS_NOT_AVAILABLE = 1005
STATUS_ABNORMAL_CLOSED = 1006
STATUS_INVALID_PAYLOAD = 1007
STATUS_POLICY_VIOLATION = 1008
STATUS_MESSAGE_TOO_BIG = 1009
STATUS_INVALID_EXTENSION = 1010
STATUS_UNEXPECTED_CONDITION = 1011
STATUS_SERVICE_RESTART = 1012
STATUS_TRY_AGAIN_LATER = 1013
STATUS_BAD_GATEWAY = 1014
STATUS_TLS_HANDSHAKE_ERROR = 1015

VALID_CLOSE_STATUS = (
    STATUS_NORMAL,
    STATUS_GOING_AWAY,
    STATUS_PROTOCOL_ERROR,
    STATUS_UNSUPPORTED_DATA_TYPE,
    STATUS_INVALID_PAYLOAD,
    STATUS_POLICY_VIOLATION,
    STATUS_MESSAGE_TOO_BIG,
    STATUS_INVALID_EXTENSION,
    STATUS_UNEXPECTED_CONDITION,
    STATUS_SERVICE_RESTART,
    STATUS_TRY_AGAIN_LATER,
    STATUS_BAD_GATEWAY,
)


class ABNF:
    """
    ABNF frame class.
    See http://tools.ietf.org/html/rfc5234
    and http://tools.ietf.org/html/rfc6455#section-5.2
    """

    # operation code values.
    OPCODE_CONT = 0x0
    OPCODE_TEXT = 0x1
    OPCODE_BINARY = 0x2
    OPCODE_CLOSE = 0x8
    OPCODE_PING = 0x9
    OPCODE_PONG = 0xA

    # available operation code value tuple
    OPCODES = (
        OPCODE_CONT,
        OPCODE_TEXT,
        OPCODE_BINARY,
        OPCODE_CLOSE,
        OPCODE_PING,
        OPCODE_PONG,
    )

    # opcode human readable string
    OPCODE_MAP = {
        OPCODE_CONT: "cont",
        OPCODE_TEXT: "text",
        OPCODE_BINARY: "binary",
        OPCODE_CLOSE: "close",
        OPCODE_PING: "ping",
        OPCODE_PONG: "pong",
    }

    # data length threshold.
    LENGTH_7 = 0x7E
    LENGTH_16 = 1 << 16
    LENGTH_63 = 1 << 63

    def __init__(
        self,
        fin: int = 0,
        rsv1: int = 0,
        rsv2: int = 0,
        rsv3: int = 0,
        opcode: int = OPCODE_TEXT,
        mask_value: int = 1,
        data: Union[str, bytes, None] = "",
    ) -> None:
        """
        Constructor for ABNF. Please check RFC for arguments.
        """
        self.fin = fin
        self.rsv1 = rsv1
        self.rsv2 = rsv2
        self.rsv3 = rsv3
        self.opcode = opcode
        self.mask_value = mask_value
        if data is None:
            data = ""
        self.data = data
        self.get_mask_key = os.urandom

    def validate(self, skip_utf8_validation: bool = False) -> None:
        """
        Validate the ABNF frame.

        Parameters
        ----------
        skip_utf8_validation: skip utf8 validation.
        """
        if self.rsv1 or self.rsv2 or self.rsv3:
            raise WebSocketProtocolException("rsv is not implemented, yet")

        if self.opcode not in ABNF.OPCODES:
            raise WebSocketProtocolException("Invalid opcode %r", self.opcode)

        if self.opcode == ABNF.OPCODE_PING and not self.fin:
            raise WebSocketProtocolException("Invalid ping frame.")

        if self.opcode == ABNF.OPCODE_CLOSE:
            l = len(self.data)
            if not l:
                return
            if l == 1 or l >= 126:
                raise WebSocketProtocolException("Invalid close frame.")
            if l > 2 and not skip_utf8_validation and not validate_utf8(self.data[2:]):
                raise WebSocketProtocolException("Invalid close frame.")

            code = 256 * int(self.data[0]) + int(self.data[1])
            if not self._is_valid_close_status(code):
                raise WebSocketProtocolException("Invalid close opcode %r", code)

    @staticmethod
    def _is_valid_close_status(code: int) -> bool:
        return code in VALID_CLOSE_STATUS or (3000 <= code < 5000)

    def __str__(self) -> str:
        return f"fin={self.fin} opcode={self.opcode} data={self.data}"

    @staticmethod
    def create_frame(data: Union[bytes, str], opcode: int, fin: int = 1) -> "ABNF":
        """
        Create frame to send text, binary and other data.

        Parameters
        ----------
        data: str
            data to send. This is string value(byte array).
            If opcode is OPCODE_TEXT and this value is unicode,
            data value is converted into unicode string, automatically.
        opcode: int
            operation code. please see OPCODE_MAP.
        fin: int
            fin flag. if set to 0, create continue fragmentation.
        """
        if opcode == ABNF.OPCODE_TEXT and isinstance(data, str):
            data = data.encode("utf-8")
        # mask must be set if send data from client
        return ABNF(fin, 0, 0, 0, opcode, 1, data)

    def format(self) -> bytes:
        """
        Format this object to string(byte array) to send data to server.
        """
        if any(x not in (0, 1) for x in [self.fin, self.rsv1, self.rsv2, self.rsv3]):
            raise ValueError("not 0 or 1")
        if self.opcode not in ABNF.OPCODES:
            raise ValueError("Invalid OPCODE")
        length = len(self.data)
        if length >= ABNF.LENGTH_63:
            raise ValueError("data is too long")

        frame_header = chr(
            self.fin << 7
            | self.rsv1 << 6
            | self.rsv2 << 5
            | self.rsv3 << 4
            | self.opcode
        ).encode("latin-1")
        if length < ABNF.LENGTH_7:
            frame_header += chr(self.mask_value << 7 | length).encode("latin-1")
        elif length < ABNF.LENGTH_16:
            frame_header += chr(self.mask_value << 7 | 0x7E).encode("latin-1")
            frame_header += struct.pack("!H", length)
        else:
            frame_header += chr(self.mask_value << 7 | 0x7F).encode("latin-1")
            frame_header += struct.pack("!Q", length)

        if not self.mask_value:
            if isinstance(self.data, str):
                self.data = self.data.encode("utf-8")
            return frame_header + self.data
        mask_key = self.get_mask_key(4)
        return frame_header + self._get_masked(mask_key)

    def _get_masked(self, mask_key: Union[str, bytes]) -> bytes:
        s = ABNF.mask(mask_key, self.data)

        if isinstance(mask_key, str):
            mask_key = mask_key.encode("utf-8")

        return mask_key + s

    @staticmethod
    def mask(mask_key: Union[str, bytes], data: Union[str, bytes]) -> bytes:
        """
        Mask or unmask data. Just do xor for each byte

        Parameters
        ----------
        mask_key: bytes or str
            4 byte mask.
        data: bytes or str
            data to mask/unmask.
        """
        if data is None:
            data = ""

        if isinstance(mask_key, str):
            mask_key = mask_key.encode("latin-1")

        if isinstance(data, str):
            data = data.encode("latin-1")

        return _mask(array.array("B", mask_key), array.array("B", data))


class frame_buffer:
    _HEADER_MASK_INDEX = 5
    _HEADER_LENGTH_INDEX = 6

    def __init__(
        self, recv_fn: Callable[[int], int], skip_utf8_validation: bool
    ) -> None:
        self.recv = recv_fn
        self.skip_utf8_validation = skip_utf8_validation
        # Buffers over the packets from the layer beneath until desired amount
        # bytes of bytes are received.
        self.recv_buffer: list = []
        self.clear()
        self.lock = Lock()

    def clear(self) -> None:
        self.header: Optional[tuple] = None
        self.length: Optional[int] = None
        self.mask_value: Union[bytes, str, None] = None

    def has_received_header(self) -> bool:
        return self.header is None

    def recv_header(self) -> None:
        header = self.recv_strict(2)
        b1 = header[0]
        fin = b1 >> 7 & 1
        rsv1 = b1 >> 6 & 1
        rsv2 = b1 >> 5 & 1
        rsv3 = b1 >> 4 & 1
        opcode = b1 & 0xF
        b2 = header[1]
        has_mask = b2 >> 7 & 1
        length_bits = b2 & 0x7F

        self.header = (fin, rsv1, rsv2, rsv3, opcode, has_mask, length_bits)

    def has_mask(self) -> Union[bool, int]:
        if not self.header:
            return False
        header_val: int = self.header[frame_buffer._HEADER_MASK_INDEX]
        return header_val

    def has_received_length(self) -> bool:
        return self.length is None

    def recv_length(self) -> None:
        bits = self.header[frame_buffer._HEADER_LENGTH_INDEX]
        length_bits = bits & 0x7F
        if length_bits == 0x7E:
            v = self.recv_strict(2)
            self.length = struct.unpack("!H", v)[0]
        elif length_bits == 0x7F:
            v = self.recv_strict(8)
            self.length = struct.unpack("!Q", v)[0]
        else:
            self.length = length_bits

    def has_received_mask(self) -> bool:
        return self.mask_value is None

    def recv_mask(self) -> None:
        self.mask_value = self.recv_strict(4) if self.has_mask() else ""

    def recv_frame(self) -> ABNF:
        with self.lock:
            # Header
            if self.has_received_header():
                self.recv_header()
            (fin, rsv1, rsv2, rsv3, opcode, has_mask, _) = self.header

            # Frame length
            if self.has_received_length():
                self.recv_length()
            length = self.length

            # Mask
            if self.has_received_mask():
                self.recv_mask()
            mask_value = self.mask_value

            # Payload
            payload = self.recv_strict(length)
            if has_mask:
                payload = ABNF.mask(mask_value, payload)

            # Reset for next frame
            self.clear()

            frame = ABNF(fin, rsv1, rsv2, rsv3, opcode, has_mask, payload)
            frame.validate(self.skip_utf8_validation)

        return frame

    def recv_strict(self, bufsize: int) -> bytes:
        shortage = bufsize - sum(map(len, self.recv_buffer))
        while shortage > 0:
            # Limit buffer size that we pass to socket.recv() to avoid
            # fragmenting the heap -- the number of bytes recv() actually
            # reads is limited by socket buffer and is relatively small,
            # yet passing large numbers repeatedly causes lots of large
            # buffers allocated and then shrunk, which results in
            # fragmentation.
            bytes_ = self.recv(min(16384, shortage))
            self.recv_buffer.append(bytes_)
            shortage -= len(bytes_)

        unified = b"".join(self.recv_buffer)

        if shortage == 0:
            self.recv_buffer = []
            return unified
        else:
            self.recv_buffer = [unified[bufsize:]]
            return unified[:bufsize]


class continuous_frame:
    def __init__(self, fire_cont_frame: bool, skip_utf8_validation: bool) -> None:
        self.fire_cont_frame = fire_cont_frame
        self.skip_utf8_validation = skip_utf8_validation
        self.cont_data: Optional[list] = None
        self.recving_frames: Optional[int] = None

    def validate(self, frame: ABNF) -> None:
        if not self.recving_frames and frame.opcode == ABNF.OPCODE_CONT:
            raise WebSocketProtocolException("Illegal frame")
        if self.recving_frames and frame.opcode in (
            ABNF.OPCODE_TEXT,
            ABNF.OPCODE_BINARY,
        ):
            raise WebSocketProtocolException("Illegal frame")

    def add(self, frame: ABNF) -> None:
        if self.cont_data:
            self.cont_data[1] += frame.data
        else:
            if frame.opcode in (ABNF.OPCODE_TEXT, ABNF.OPCODE_BINARY):
                self.recving_frames = frame.opcode
            self.cont_data = [frame.opcode, frame.data]

        if frame.fin:
            self.recving_frames = None

    def is_fire(self, frame: ABNF) -> Union[bool, int]:
        return frame.fin or self.fire_cont_frame

    def extract(self, frame: ABNF) -> tuple:
        data = self.cont_data
        self.cont_data = None
        frame.data = data[1]
        if (
            not self.fire_cont_frame
            and data[0] == ABNF.OPCODE_TEXT
            and not self.skip_utf8_validation
            and not validate_utf8(frame.data)
        ):
            raise WebSocketPayloadException(f"cannot decode: {repr(frame.data)}")
        return data[0], frame