utils.py 1.2 KB

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