| 95 | |
| 96 | |
| 97 | class WsFrameReader: |
| 98 | |
| 99 | def __init__(self, data: bytes): |
| 100 | self.data = data |
| 101 | |
| 102 | def _read(self, n: int): |
| 103 | if len(self.data) < n: |
| 104 | raise EOFError(f'have {len(self.data)} bytes left, but {n} requested') |
| 105 | elif n == 0: |
| 106 | return b'' |
| 107 | chunk = self.data[:n] |
| 108 | del self.data[:n] |
| 109 | return chunk |
| 110 | |
| 111 | def next_frame(self): |
| 112 | data = self._read(2) |
| 113 | h1, h2 = struct.unpack("!BB", data) |
| 114 | log.debug(f'parsed h1={h1} h2={h2} from {data}') |
| 115 | fin = True if h1 & 0x80 else False |
| 116 | opcode = h1 & 0xf |
| 117 | has_mask = True if h2 & 0x80 else False |
| 118 | mask = None |
| 119 | dlen = h2 & 0x7f |
| 120 | if dlen == 126: |
| 121 | (dlen,) = struct.unpack("!H", self._read(2)) |
| 122 | elif dlen == 127: |
| 123 | (dlen,) = struct.unpack("!Q", self._read(8)) |
| 124 | if has_mask: |
| 125 | mask = self._read(4) |
| 126 | return WsFrame(opcode=opcode, fin=fin, mask=mask, data=self._read(dlen)) |
| 127 | |
| 128 | def eof(self): |
| 129 | return len(self.data) == 0 |
| 130 | |
| 131 | @classmethod |
| 132 | def parse(cls, data: bytes): |
| 133 | frames = [] |
| 134 | reader = WsFrameReader(data=data) |
| 135 | while not reader.eof(): |
| 136 | frames.append(reader.next_frame()) |
| 137 | return frames |