Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/websocket/_abnf.py: 29%

1import array 

2import os 

3import struct 

4import sys 

5from threading import Lock 

6from typing import Callable, NoReturn, Optional, Union, Any 

7 

8from ._exceptions import WebSocketPayloadException, WebSocketProtocolException 

9from ._utils import validate_utf8 

10 

11""" 

12_abnf.py 

13websocket - WebSocket client library for Python 

14 

15Copyright 2026 engn33r 

16 

17Licensed under the Apache License, Version 2.0 (the "License"); 

18you may not use this file except in compliance with the License. 

19You may obtain a copy of the License at 

20 

21 http://www.apache.org/licenses/LICENSE-2.0 

22 

23Unless required by applicable law or agreed to in writing, software 

24distributed under the License is distributed on an "AS IS" BASIS, 

25WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

26See the License for the specific language governing permissions and 

27limitations under the License. 

28""" 

29 

30try: 

31 # If wsaccel is available, use compiled routines to mask data. 

32 # wsaccel only provides around a 10% speed boost compared 

33 # to the websocket-client _mask() implementation. 

34 # Note that wsaccel is unmaintained. 

35 from wsaccel.xormask import XorMaskerSimple 

36 

37 def _mask(mask_value: array.array, data_value: array.array) -> bytes: 

38 mask_result: bytes = XorMaskerSimple(mask_value).process(data_value) 

39 return mask_result 

40 

41except ImportError: 

42 # wsaccel is not available, use websocket-client _mask() 

43 native_byteorder = sys.byteorder 

44 

45 def _mask(mask_value: array.array, data_value: array.array) -> bytes: 

46 datalen = len(data_value) 

47 int_data_value = int.from_bytes(data_value, native_byteorder) 

48 int_mask_value = int.from_bytes( 

49 mask_value * (datalen // 4) + mask_value[: datalen % 4], native_byteorder 

50 ) 

51 return (int_data_value ^ int_mask_value).to_bytes(datalen, native_byteorder) 

52 

53 

54__all__ = [ 

55 "ABNF", 

56 "continuous_frame", 

57 "frame_buffer", 

58 "STATUS_NORMAL", 

59 "STATUS_GOING_AWAY", 

60 "STATUS_PROTOCOL_ERROR", 

61 "STATUS_UNSUPPORTED_DATA_TYPE", 

62 "STATUS_STATUS_NOT_AVAILABLE", 

63 "STATUS_ABNORMAL_CLOSED", 

64 "STATUS_INVALID_PAYLOAD", 

65 "STATUS_POLICY_VIOLATION", 

66 "STATUS_MESSAGE_TOO_BIG", 

67 "STATUS_INVALID_EXTENSION", 

68 "STATUS_UNEXPECTED_CONDITION", 

69 "STATUS_SERVICE_RESTART", 

70 "STATUS_TRY_AGAIN_LATER", 

71 "STATUS_BAD_GATEWAY", 

72 "STATUS_TLS_HANDSHAKE_ERROR", 

73] 

74 

75# closing frame status codes. 

76STATUS_NORMAL = 1000 

77STATUS_GOING_AWAY = 1001 

78STATUS_PROTOCOL_ERROR = 1002 

79STATUS_UNSUPPORTED_DATA_TYPE = 1003 

80STATUS_STATUS_NOT_AVAILABLE = 1005 

81STATUS_ABNORMAL_CLOSED = 1006 

82STATUS_INVALID_PAYLOAD = 1007 

83STATUS_POLICY_VIOLATION = 1008 

84STATUS_MESSAGE_TOO_BIG = 1009 

85STATUS_INVALID_EXTENSION = 1010 

86STATUS_UNEXPECTED_CONDITION = 1011 

87STATUS_SERVICE_RESTART = 1012 

88STATUS_TRY_AGAIN_LATER = 1013 

89STATUS_BAD_GATEWAY = 1014 

90STATUS_TLS_HANDSHAKE_ERROR = 1015 

91 

92VALID_CLOSE_STATUS = ( 

93 STATUS_NORMAL, 

94 STATUS_GOING_AWAY, 

95 STATUS_PROTOCOL_ERROR, 

96 STATUS_UNSUPPORTED_DATA_TYPE, 

97 STATUS_INVALID_PAYLOAD, 

98 STATUS_POLICY_VIOLATION, 

99 STATUS_MESSAGE_TOO_BIG, 

100 STATUS_INVALID_EXTENSION, 

101 STATUS_UNEXPECTED_CONDITION, 

102 STATUS_SERVICE_RESTART, 

103 STATUS_TRY_AGAIN_LATER, 

104 STATUS_BAD_GATEWAY, 

105) 

106 

107 

108class ABNF: 

109 """ 

110 ABNF frame class. 

111 See http://tools.ietf.org/html/rfc5234 

112 and http://tools.ietf.org/html/rfc6455#section-5.2 

113 """ 

114 

115 # operation code values. 

116 OPCODE_CONT = 0x0 

117 OPCODE_TEXT = 0x1 

118 OPCODE_BINARY = 0x2 

119 OPCODE_CLOSE = 0x8 

120 OPCODE_PING = 0x9 

121 OPCODE_PONG = 0xA 

122 

123 # available operation code value tuple 

124 OPCODES = ( 

125 OPCODE_CONT, 

126 OPCODE_TEXT, 

127 OPCODE_BINARY, 

128 OPCODE_CLOSE, 

129 OPCODE_PING, 

130 OPCODE_PONG, 

131 ) 

132 

133 # opcode human readable string 

134 OPCODE_MAP = { 

135 OPCODE_CONT: "cont", 

136 OPCODE_TEXT: "text", 

137 OPCODE_BINARY: "binary", 

138 OPCODE_CLOSE: "close", 

139 OPCODE_PING: "ping", 

140 OPCODE_PONG: "pong", 

141 } 

142 

143 # data length threshold. 

144 LENGTH_7 = 0x7E 

145 LENGTH_16 = 1 << 16 

146 LENGTH_63 = 1 << 63 

147 

148 def __init__( 

149 self, 

150 fin: int = 0, 

151 rsv1: int = 0, 

152 rsv2: int = 0, 

153 rsv3: int = 0, 

154 opcode: int = OPCODE_TEXT, 

155 mask_value: int = 1, 

156 data: Optional[Union[str, bytes]] = "", 

157 ) -> None: 

158 """ 

159 Constructor for ABNF. Please check RFC for arguments. 

160 """ 

161 self.fin = fin 

162 self.rsv1 = rsv1 

163 self.rsv2 = rsv2 

164 self.rsv3 = rsv3 

165 self.opcode = opcode 

166 self.mask_value = mask_value 

167 if data is None: 

168 data = "" 

169 self.data = data 

170 self.get_mask_key = os.urandom 

171 

172 def validate(self, skip_utf8_validation: bool = False) -> None: 

173 """ 

174 Validate the ABNF frame. 

175 

176 Parameters 

177 ---------- 

178 skip_utf8_validation: skip utf8 validation. 

179 """ 

180 if self.rsv1 or self.rsv2 or self.rsv3: 

181 raise WebSocketProtocolException("rsv is not implemented, yet") 

182 

183 if self.opcode not in ABNF.OPCODES: 

184 raise WebSocketProtocolException(f"Invalid opcode {self.opcode!r}") 

185 

186 if self.opcode in (ABNF.OPCODE_CLOSE, ABNF.OPCODE_PING, ABNF.OPCODE_PONG): 

187 if not self.fin or len(self.data) >= ABNF.LENGTH_7: 

188 raise WebSocketProtocolException("Invalid control frame.") 

189 

190 if self.opcode == ABNF.OPCODE_CLOSE: 

191 data_length = len(self.data) 

192 if not data_length: 

193 return 

194 if data_length == 1: 

195 raise WebSocketProtocolException("Invalid close frame.") 

196 if ( 

197 data_length > 2 

198 and not skip_utf8_validation 

199 and not validate_utf8(self.data[2:]) 

200 ): 

201 raise WebSocketProtocolException("Invalid close frame.") 

202 

203 data_bytes = ( 

204 self.data[:2] 

205 if isinstance(self.data, bytes) 

206 else self.data[:2].encode("utf-8") 

207 ) 

208 code = struct.unpack("!H", data_bytes)[0] 

209 if not self._is_valid_close_status(code): 

210 raise WebSocketProtocolException(f"Invalid close opcode {code!r}") 

211 

212 @staticmethod 

213 def _is_valid_close_status(code: int) -> bool: 

214 return code in VALID_CLOSE_STATUS or (3000 <= code < 5000) 

215 

216 def __str__(self) -> str: 

217 data_repr = self.data if isinstance(self.data, str) else repr(self.data) 

218 return f"fin={self.fin} opcode={self.opcode} data={data_repr}" 

219 

220 @staticmethod 

221 def create_frame(data: Union[bytes, str], opcode: int, fin: int = 1) -> "ABNF": 

222 """ 

223 Create frame to send text, binary and other data. 

224 

225 Parameters 

226 ---------- 

227 data: str 

228 data to send. This is string value(byte array). 

229 If opcode is OPCODE_TEXT and this value is unicode, 

230 data value is converted into unicode string, automatically. 

231 opcode: int 

232 operation code. please see OPCODE_MAP. 

233 fin: int 

234 fin flag. if set to 0, create continue fragmentation. 

235 """ 

236 if opcode == ABNF.OPCODE_TEXT and isinstance(data, str): 

237 data = data.encode("utf-8") 

238 # mask must be set if send data from client 

239 return ABNF(fin, 0, 0, 0, opcode, 1, data) 

240 

241 def format(self) -> bytes: 

242 """ 

243 Format this object to string(byte array) to send data to server. 

244 """ 

245 if any(x not in (0, 1) for x in [self.fin, self.rsv1, self.rsv2, self.rsv3]): 

246 raise ValueError("not 0 or 1") 

247 if self.opcode not in ABNF.OPCODES: 

248 raise ValueError("Invalid OPCODE") 

249 length = len(self.data) 

250 if length >= ABNF.LENGTH_63: 

251 raise ValueError("data is too long") 

252 

253 frame_header = chr( 

254 self.fin << 7 

255 | self.rsv1 << 6 

256 | self.rsv2 << 5 

257 | self.rsv3 << 4 

258 | self.opcode 

259 ).encode("latin-1") 

260 if length < ABNF.LENGTH_7: 

261 frame_header += chr(self.mask_value << 7 | length).encode("latin-1") 

262 elif length < ABNF.LENGTH_16: 

263 frame_header += chr(self.mask_value << 7 | 0x7E).encode("latin-1") 

264 frame_header += struct.pack("!H", length) 

265 else: 

266 frame_header += chr(self.mask_value << 7 | 0x7F).encode("latin-1") 

267 frame_header += struct.pack("!Q", length) 

268 

269 if not self.mask_value: 

270 if isinstance(self.data, str): 

271 self.data = self.data.encode("utf-8") 

272 return frame_header + self.data 

273 mask_key = self.get_mask_key(4) 

274 return frame_header + self._get_masked(mask_key) 

275 

276 def _get_masked(self, mask_key: Union[str, bytes]) -> bytes: 

277 s = ABNF.mask(mask_key, self.data) 

278 

279 if isinstance(mask_key, str): 

280 mask_key = mask_key.encode("utf-8") 

281 

282 return mask_key + s 

283 

284 @staticmethod 

285 def mask(mask_key: Union[str, bytes], data: Union[str, bytes]) -> bytes: 

286 """ 

287 Mask or unmask data. Just do xor for each byte 

288 

289 Parameters 

290 ---------- 

291 mask_key: bytes or str 

292 4 byte mask. 

293 data: bytes or str 

294 data to mask/unmask. 

295 """ 

296 if data is None: 

297 data = "" 

298 

299 if isinstance(mask_key, str): 

300 mask_key = mask_key.encode("latin-1") 

301 

302 if isinstance(data, str): 

303 data = data.encode("latin-1") 

304 

305 return _mask(array.array("B", mask_key), array.array("B", data)) 

306 

307 

308class frame_buffer: 

309 _HEADER_MASK_INDEX = 5 

310 _HEADER_LENGTH_INDEX = 6 

311 

312 def __init__( 

313 self, recv_fn: Callable[[int], bytes], skip_utf8_validation: bool 

314 ) -> None: 

315 self.recv = recv_fn 

316 self.skip_utf8_validation = skip_utf8_validation 

317 # Buffers over the packets from the layer beneath until desired amount 

318 # bytes of bytes are received. 

319 self.recv_buffer: list = [] 

320 # Set once the stream has lost frame sync; see _lose_sync(). 

321 self.sync_error: Optional[str] = None 

322 self.clear() 

323 self.lock = Lock() 

324 

325 def clear(self) -> None: 

326 self.header: Optional[tuple] = None 

327 self.length: Optional[int] = None 

328 self.mask_value: Optional[Union[bytes, str]] = None 

329 

330 def needs_header(self) -> bool: 

331 return self.header is None 

332 

333 def recv_header(self) -> None: 

334 header = self.recv_strict(2) 

335 b1 = header[0] 

336 fin = b1 >> 7 & 1 

337 rsv1 = b1 >> 6 & 1 

338 rsv2 = b1 >> 5 & 1 

339 rsv3 = b1 >> 4 & 1 

340 opcode = b1 & 0xF 

341 b2 = header[1] 

342 has_mask = b2 >> 7 & 1 

343 length_bits = b2 & 0x7F 

344 

345 self.header = (fin, rsv1, rsv2, rsv3, opcode, has_mask, length_bits) 

346 

347 def has_mask(self) -> Union[bool, int]: 

348 if not self.header: 

349 return False 

350 header_val: int = self.header[frame_buffer._HEADER_MASK_INDEX] 

351 return header_val 

352 

353 def needs_length(self) -> bool: 

354 return self.length is None 

355 

356 def recv_length(self) -> None: 

357 if self.header is None: 

358 raise WebSocketProtocolException("Header not received") 

359 bits = self.header[frame_buffer._HEADER_LENGTH_INDEX] 

360 length_bits = bits & 0x7F 

361 if length_bits == 0x7E: 

362 v = self.recv_strict(2) 

363 length = struct.unpack("!H", v)[0] 

364 if length < ABNF.LENGTH_7: 

365 self._lose_sync("Invalid 16-bit frame length.") 

366 elif length_bits == 0x7F: 

367 v = self.recv_strict(8) 

368 length = struct.unpack("!Q", v)[0] 

369 if length < ABNF.LENGTH_16 or length >= ABNF.LENGTH_63: 

370 self._lose_sync("Invalid 64-bit frame length.") 

371 else: 

372 length = length_bits 

373 self.length = length 

374 

375 def _lose_sync(self, message: str) -> NoReturn: 

376 # The length field has been consumed but the payload has not, so the 

377 # stream is no longer at a frame boundary. Reading on would parse the 

378 # payload as frame headers, letting the peer hide frames inside what 

379 # intermediaries see as one frame's data. Fail the connection instead 

380 # (RFC 6455 7.1.7): this read and every later one raise. 

381 self.clear() 

382 self.recv_buffer = [] 

383 self.sync_error = message 

384 raise WebSocketProtocolException(message) 

385 

386 def needs_mask(self) -> bool: 

387 return self.mask_value is None 

388 

389 def recv_mask(self) -> None: 

390 self.mask_value = self.recv_strict(4) if self.has_mask() else "" 

391 

392 def recv_frame(self) -> ABNF: 

393 with self.lock: 

394 if self.sync_error is not None: 

395 raise WebSocketProtocolException( 

396 f"Frame stream out of sync after an earlier error: {self.sync_error}" 

397 ) 

398 # Header 

399 if self.needs_header(): 

400 self.recv_header() 

401 if self.header is None: 

402 raise WebSocketProtocolException("Header not received") 

403 fin, rsv1, rsv2, rsv3, opcode, has_mask, _ = self.header 

404 

405 # Frame length 

406 if self.needs_length(): 

407 self.recv_length() 

408 length = self.length 

409 

410 # Mask 

411 if self.needs_mask(): 

412 self.recv_mask() 

413 mask_value = self.mask_value 

414 

415 # Payload 

416 if length is None: 

417 raise WebSocketProtocolException("Length not received") 

418 payload = self.recv_strict(length) 

419 if has_mask: 

420 if mask_value is None: 

421 raise WebSocketProtocolException("Mask not received") 

422 payload = ABNF.mask(mask_value, payload) 

423 

424 # Reset for next frame 

425 self.clear() 

426 

427 frame = ABNF(fin, rsv1, rsv2, rsv3, opcode, has_mask, payload) 

428 frame.validate(self.skip_utf8_validation) 

429 

430 return frame 

431 

432 def recv_strict(self, bufsize: int) -> bytes: 

433 if not isinstance(bufsize, int): 

434 raise ValueError("bufsize must be an integer") 

435 shortage = bufsize - sum(len(buf) for buf in self.recv_buffer) 

436 while shortage > 0: 

437 # Limit buffer size that we pass to socket.recv() to avoid 

438 # fragmenting the heap -- the number of bytes recv() actually 

439 # reads is limited by socket buffer and is relatively small, 

440 # yet passing large numbers repeatedly causes lots of large 

441 # buffers allocated and then shrunk, which results in 

442 # fragmentation. 

443 bytes_ = self.recv(min(16384, shortage)) 

444 if isinstance(bytes_, bytes): 

445 self.recv_buffer.append(bytes_) 

446 shortage -= len(bytes_) 

447 else: 

448 # Handle case where recv returns int or other type 

449 break 

450 

451 unified = b"".join(self.recv_buffer) 

452 

453 if shortage == 0: 

454 self.recv_buffer = [] 

455 return unified 

456 else: 

457 self.recv_buffer = [unified[bufsize:]] 

458 return unified[:bufsize] 

459 

460 

461class continuous_frame: 

462 def __init__(self, fire_cont_frame: bool, skip_utf8_validation: bool) -> None: 

463 self.fire_cont_frame = fire_cont_frame 

464 self.skip_utf8_validation = skip_utf8_validation 

465 self.cont_data: Optional[list[Any]] = None 

466 self.recving_frames: Optional[int] = None 

467 

468 def validate(self, frame: ABNF) -> None: 

469 if not self.recving_frames and frame.opcode == ABNF.OPCODE_CONT: 

470 raise WebSocketProtocolException("Illegal frame") 

471 if self.recving_frames and frame.opcode in ( 

472 ABNF.OPCODE_TEXT, 

473 ABNF.OPCODE_BINARY, 

474 ): 

475 raise WebSocketProtocolException("Illegal frame") 

476 

477 def add(self, frame: ABNF) -> None: 

478 if self.cont_data: 

479 self.cont_data[1] += frame.data 

480 else: 

481 if frame.opcode in (ABNF.OPCODE_TEXT, ABNF.OPCODE_BINARY): 

482 self.recving_frames = frame.opcode 

483 self.cont_data = [frame.opcode, frame.data] 

484 

485 if frame.fin: 

486 self.recving_frames = None 

487 

488 def is_fire(self, frame: ABNF) -> Union[bool, int]: 

489 return frame.fin or self.fire_cont_frame 

490 

491 def extract(self, frame: ABNF) -> tuple: 

492 data = self.cont_data 

493 if data is None: 

494 raise WebSocketProtocolException("No continuation data available") 

495 self.cont_data = None 

496 frame.data = data[1] 

497 if ( 

498 not self.fire_cont_frame 

499 and data is not None 

500 and data[0] == ABNF.OPCODE_TEXT 

501 and not self.skip_utf8_validation 

502 and not validate_utf8(frame.data) 

503 ): 

504 raise WebSocketPayloadException(f"cannot decode: {repr(frame.data)}") 

505 if data is None: 

506 raise WebSocketProtocolException("No continuation data available") 

507 return data[0], frame