codekingpro/portable-devtools

1from __future__ import annotations
2
3import base64
4import hashlib
5import secrets
6import sys
7
8from .typing import BytesLike
9
10
11__all__ = ["accept_key", "apply_mask"]
12
13
14GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
15
16
17def generate_key() -> str:
18    """
19    Generate a random key for the Sec-WebSocket-Key header.
20
21    """
22    key = secrets.token_bytes(16)
23    return base64.b64encode(key).decode()
24
25
26def accept_key(key: str) -> str:
27    """
28    Compute the value of the Sec-WebSocket-Accept header.
29
30    Args:
31        key: Value of the Sec-WebSocket-Key header.
32
33    """
34    sha1 = hashlib.sha1((key + GUID).encode()).digest()
35    return base64.b64encode(sha1).decode()
36
37
38def apply_mask(data: BytesLike, mask: bytes | bytearray) -> bytes:
39    """
40    Apply masking to the data of a WebSocket message.
41
42    Args:
43        data: Data to mask.
44        mask: 4-bytes mask.
45
46    """
47    if len(mask) != 4:
48        raise ValueError("mask must contain 4 bytes")
49
50    data_int = int.from_bytes(data, sys.byteorder)
51    mask_repeated = mask * (len(data) // 4) + mask[: len(data) % 4]
52    mask_int = int.from_bytes(mask_repeated, sys.byteorder)
53    return (data_int ^ mask_int).to_bytes(len(data), sys.byteorder)
54