bazarr/libs/websocket/_abnf.py

454 lines
14 KiB
Python
Raw Normal View History

2022-11-07 18:06:49 +00:00
import array
import os
import struct
import sys
from threading import Lock
from typing import Callable, Optional, Union
2022-11-07 18:06:49 +00:00
from ._exceptions import *
from ._utils import validate_utf8
2018-10-31 16:08:29 +00:00
"""
_abnf.py
2018-10-31 16:08:29 +00:00
websocket - WebSocket client library for Python
Copyright 2023 engn33r
2018-10-31 16:08:29 +00:00
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
2018-10-31 16:08:29 +00:00
http://www.apache.org/licenses/LICENSE-2.0
2018-10-31 16:08:29 +00:00
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.
2018-10-31 16:08:29 +00:00
"""
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
2018-10-31 16:08:29 +00:00
def _mask(mask_value: array.array, data_value: array.array) -> bytes:
mask_result: bytes = XorMaskerSimple(mask_value).process(data_value)
return mask_result
2018-10-31 16:08:29 +00:00
except ImportError:
# wsaccel is not available, use websocket-client _mask()
native_byteorder = sys.byteorder
2018-10-31 16:08:29 +00:00
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)
2018-10-31 16:08:29 +00:00
2018-10-31 16:08:29 +00:00
__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",
2018-10-31 16:08:29 +00:00
]
# 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
2022-11-07 18:06:49 +00:00
STATUS_SERVICE_RESTART = 1012
STATUS_TRY_AGAIN_LATER = 1013
2018-10-31 16:08:29 +00:00
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,
2022-11-07 18:06:49 +00:00
STATUS_SERVICE_RESTART,
STATUS_TRY_AGAIN_LATER,
2018-10-31 16:08:29 +00:00
STATUS_BAD_GATEWAY,
)
class ABNF:
2018-10-31 16:08:29 +00:00
"""
ABNF frame class.
See http://tools.ietf.org/html/rfc5234
2018-10-31 16:08:29 +00:00
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
2018-10-31 16:08:29 +00:00
# available operation code value tuple
OPCODES = (
OPCODE_CONT,
OPCODE_TEXT,
OPCODE_BINARY,
OPCODE_CLOSE,
OPCODE_PING,
OPCODE_PONG,
)
2018-10-31 16:08:29 +00:00
# opcode human readable string
OPCODE_MAP = {
OPCODE_CONT: "cont",
OPCODE_TEXT: "text",
OPCODE_BINARY: "binary",
OPCODE_CLOSE: "close",
OPCODE_PING: "ping",
OPCODE_PONG: "pong",
2018-10-31 16:08:29 +00:00
}
# data length threshold.
LENGTH_7 = 0x7E
2018-10-31 16:08:29 +00:00
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:
2018-10-31 16:08:29 +00:00
"""
Constructor for ABNF. Please check RFC for arguments.
2018-10-31 16:08:29 +00:00
"""
self.fin = fin
self.rsv1 = rsv1
self.rsv2 = rsv2
self.rsv3 = rsv3
self.opcode = opcode
self.mask_value = mask_value
2018-10-31 16:08:29 +00:00
if data is None:
data = ""
self.data = data
self.get_mask_key = os.urandom
def validate(self, skip_utf8_validation: bool = False) -> None:
2018-10-31 16:08:29 +00:00
"""
Validate the ABNF frame.
Parameters
----------
2018-10-31 16:08:29 +00:00
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])
2018-10-31 16:08:29 +00:00
if not self._is_valid_close_status(code):
2022-11-07 18:06:49 +00:00
raise WebSocketProtocolException("Invalid close opcode %r", code)
2018-10-31 16:08:29 +00:00
@staticmethod
2022-11-07 18:06:49 +00:00
def _is_valid_close_status(code: int) -> bool:
2018-10-31 16:08:29 +00:00
return code in VALID_CLOSE_STATUS or (3000 <= code < 5000)
2022-11-07 18:06:49 +00:00
def __str__(self) -> str:
return f"fin={self.fin} opcode={self.opcode} data={self.data}"
2018-10-31 16:08:29 +00:00
@staticmethod
def create_frame(data: Union[bytes, str], opcode: int, fin: int = 1) -> "ABNF":
2018-10-31 16:08:29 +00:00
"""
Create frame to send text, binary and other data.
2018-10-31 16:08:29 +00:00
Parameters
----------
data: str
data to send. This is string value(byte array).
If opcode is OPCODE_TEXT and this value is unicode,
2018-10-31 16:08:29 +00:00
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.
2018-10-31 16:08:29 +00:00
"""
if opcode == ABNF.OPCODE_TEXT and isinstance(data, str):
2018-10-31 16:08:29 +00:00
data = data.encode("utf-8")
# mask must be set if send data from client
return ABNF(fin, 0, 0, 0, opcode, 1, data)
2022-11-07 18:06:49 +00:00
def format(self) -> bytes:
2018-10-31 16:08:29 +00:00
"""
Format this object to string(byte array) to send data to server.
2018-10-31 16:08:29 +00:00
"""
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")
2018-10-31 16:08:29 +00:00
if length < ABNF.LENGTH_7:
frame_header += chr(self.mask_value << 7 | length).encode("latin-1")
2018-10-31 16:08:29 +00:00
elif length < ABNF.LENGTH_16:
frame_header += chr(self.mask_value << 7 | 0x7E).encode("latin-1")
2018-10-31 16:08:29 +00:00
frame_header += struct.pack("!H", length)
else:
frame_header += chr(self.mask_value << 7 | 0x7F).encode("latin-1")
2018-10-31 16:08:29 +00:00
frame_header += struct.pack("!Q", length)
if not self.mask_value:
if isinstance(self.data, str):
self.data = self.data.encode("utf-8")
2018-10-31 16:08:29 +00:00
return frame_header + self.data
mask_key = self.get_mask_key(4)
return frame_header + self._get_masked(mask_key)
2018-10-31 16:08:29 +00:00
def _get_masked(self, mask_key: Union[str, bytes]) -> bytes:
2018-10-31 16:08:29 +00:00
s = ABNF.mask(mask_key, self.data)
if isinstance(mask_key, str):
mask_key = mask_key.encode("utf-8")
2018-10-31 16:08:29 +00:00
return mask_key + s
@staticmethod
def mask(mask_key: Union[str, bytes], data: Union[str, bytes]) -> bytes:
2018-10-31 16:08:29 +00:00
"""
Mask or unmask data. Just do xor for each byte
Parameters
----------
2022-11-07 18:06:49 +00:00
mask_key: bytes or str
4 byte mask.
data: bytes or str
data to mask/unmask.
2018-10-31 16:08:29 +00:00
"""
if data is None:
data = ""
if isinstance(mask_key, str):
mask_key = mask_key.encode("latin-1")
2018-10-31 16:08:29 +00:00
if isinstance(data, str):
data = data.encode("latin-1")
2018-10-31 16:08:29 +00:00
return _mask(array.array("B", mask_key), array.array("B", data))
2018-10-31 16:08:29 +00:00
class frame_buffer:
2018-10-31 16:08:29 +00:00
_HEADER_MASK_INDEX = 5
_HEADER_LENGTH_INDEX = 6
def __init__(
self, recv_fn: Callable[[int], int], skip_utf8_validation: bool
) -> None:
2018-10-31 16:08:29 +00:00
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 = []
2018-10-31 16:08:29 +00:00
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
2018-10-31 16:08:29 +00:00
2022-11-07 18:06:49 +00:00
def has_received_header(self) -> bool:
2018-10-31 16:08:29 +00:00
return self.header is None
def recv_header(self) -> None:
2018-10-31 16:08:29 +00:00
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
2018-10-31 16:08:29 +00:00
b2 = header[1]
has_mask = b2 >> 7 & 1
length_bits = b2 & 0x7F
2018-10-31 16:08:29 +00:00
self.header = (fin, rsv1, rsv2, rsv3, opcode, has_mask, length_bits)
def has_mask(self) -> Union[bool, int]:
2018-10-31 16:08:29 +00:00
if not self.header:
return False
header_val: int = self.header[frame_buffer._HEADER_MASK_INDEX]
return header_val
2018-10-31 16:08:29 +00:00
2022-11-07 18:06:49 +00:00
def has_received_length(self) -> bool:
2018-10-31 16:08:29 +00:00
return self.length is None
def recv_length(self) -> None:
2018-10-31 16:08:29 +00:00
bits = self.header[frame_buffer._HEADER_LENGTH_INDEX]
length_bits = bits & 0x7F
if length_bits == 0x7E:
2018-10-31 16:08:29 +00:00
v = self.recv_strict(2)
self.length = struct.unpack("!H", v)[0]
elif length_bits == 0x7F:
2018-10-31 16:08:29 +00:00
v = self.recv_strict(8)
self.length = struct.unpack("!Q", v)[0]
else:
self.length = length_bits
2022-11-07 18:06:49 +00:00
def has_received_mask(self) -> bool:
return self.mask_value is None
2018-10-31 16:08:29 +00:00
def recv_mask(self) -> None:
self.mask_value = self.recv_strict(4) if self.has_mask() else ""
2018-10-31 16:08:29 +00:00
def recv_frame(self) -> ABNF:
2018-10-31 16:08:29 +00:00
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
2018-10-31 16:08:29 +00:00
# Payload
payload = self.recv_strict(length)
if has_mask:
payload = ABNF.mask(mask_value, payload)
2018-10-31 16:08:29 +00:00
# Reset for next frame
self.clear()
frame = ABNF(fin, rsv1, rsv2, rsv3, opcode, has_mask, payload)
frame.validate(self.skip_utf8_validation)
return frame
2022-11-07 18:06:49 +00:00
def recv_strict(self, bufsize: int) -> bytes:
shortage = bufsize - sum(map(len, self.recv_buffer))
2018-10-31 16:08:29 +00:00
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)
2018-10-31 16:08:29 +00:00
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:
2018-10-31 16:08:29 +00:00
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
2018-10-31 16:08:29 +00:00
def validate(self, frame: ABNF) -> None:
2018-10-31 16:08:29 +00:00
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,
):
2018-10-31 16:08:29 +00:00
raise WebSocketProtocolException("Illegal frame")
def add(self, frame: ABNF) -> None:
2018-10-31 16:08:29 +00:00
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]:
2018-10-31 16:08:29 +00:00
return frame.fin or self.fire_cont_frame
def extract(self, frame: ABNF) -> tuple:
2018-10-31 16:08:29 +00:00
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