|
8 | 8 | import sys |
9 | 9 | import zlib |
10 | 10 | from enum import IntEnum |
| 11 | +from functools import partial |
11 | 12 | from struct import Struct |
12 | 13 | from typing import ( |
13 | 14 | Any, |
@@ -103,6 +104,7 @@ class WSMsgType(IntEnum): |
103 | 104 | PACK_LEN2 = Struct("!BBH").pack |
104 | 105 | PACK_LEN3 = Struct("!BBQ").pack |
105 | 106 | PACK_CLOSE_CODE = Struct("!H").pack |
| 107 | +PACK_RANDBITS = Struct("!L").pack |
106 | 108 | MSG_SIZE: Final[int] = 2**14 |
107 | 109 | DEFAULT_LIMIT: Final[int] = 2**16 |
108 | 110 |
|
@@ -612,7 +614,7 @@ def __init__( |
612 | 614 | self.protocol = protocol |
613 | 615 | self.transport = transport |
614 | 616 | self.use_mask = use_mask |
615 | | - self.randrange = random.randrange |
| 617 | + self.get_random_bits = partial(random.getrandbits, 32) |
616 | 618 | self.compress = compress |
617 | 619 | self.notakeover = notakeover |
618 | 620 | self._closing = False |
@@ -668,8 +670,7 @@ async def _send_frame( |
668 | 670 | else: |
669 | 671 | header = PACK_LEN3(0x80 | rsv | opcode, 127 | mask_bit, msg_length) |
670 | 672 | if use_mask: |
671 | | - mask_int = self.randrange(0, 0xFFFFFFFF) |
672 | | - mask = mask_int.to_bytes(4, "big") |
| 673 | + mask = PACK_RANDBITS(self.get_random_bits()) |
673 | 674 | message = bytearray(message) |
674 | 675 | _websocket_mask(mask, message) |
675 | 676 | self._write(header + mask + message) |
|
0 commit comments