Apply masking to the data of a WebSocket message. Args: data: Data to mask. mask: 4-bytes mask.
(data: BytesLike, mask: bytes | bytearray)
| 36 | |
| 37 | |
| 38 | def 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) |