utils.py 928 B

1234567891011121314151617181920212223242526272829303132333435363738394041424344
  1. import base64
  2. import hashlib
  3. import itertools
  4. import secrets
  5. __all__ = ["accept_key", "apply_mask"]
  6. GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
  7. def generate_key() -> str:
  8. """
  9. Generate a random key for the Sec-WebSocket-Key header.
  10. """
  11. key = secrets.token_bytes(16)
  12. return base64.b64encode(key).decode()
  13. def accept_key(key: str) -> str:
  14. """
  15. Compute the value of the Sec-WebSocket-Accept header.
  16. :param key: value of the Sec-WebSocket-Key header
  17. """
  18. sha1 = hashlib.sha1((key + GUID).encode()).digest()
  19. return base64.b64encode(sha1).decode()
  20. def apply_mask(data: bytes, mask: bytes) -> bytes:
  21. """
  22. Apply masking to the data of a WebSocket message.
  23. :param data: Data to mask
  24. :param mask: 4-bytes mask
  25. """
  26. if len(mask) != 4:
  27. raise ValueError("mask must contain 4 bytes")
  28. return bytes(b ^ m for b, m in zip(data, itertools.cycle(mask)))