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

1import array 

2import os 

3import struct 

4import sys 

5from threading import Lock 

6from typing import Callable, 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 == ABNF.OPCODE_PING and not self.fin: 

187 raise WebSocketProtocolException("Invalid ping frame.") 

188 

189 if self.opcode == ABNF.OPCODE_CLOSE: 

190 data_length = len(self.data) 

191 if not data_length: 

192 return 

193 if data_length == 1 or data_length >= 126: 

194 raise WebSocketProtocolException("Invalid close frame.") 

195 if ( 

196 data_length > 2 

197 and not skip_utf8_validation 

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

199 ): 

200 raise WebSocketProtocolException("Invalid close frame.") 

201 

202 data_bytes = ( 

203 self.data[:2] 

204 if isinstance(self.data, bytes) 

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

206 ) 

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

208 if not self._is_valid_close_status(code): 

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

210 

211 @staticmethod 

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

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

214 

215 def __str__(self) -> str: 

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

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

218 

219 @staticmethod 

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

221 """ 

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

223 

224 Parameters 

225 ---------- 

226 data: str 

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

228 If opcode is OPCODE_TEXT and this value is unicode, 

229 data value is converted into unicode string, automatically. 

230 opcode: int 

231 operation code. please see OPCODE_MAP. 

232 fin: int 

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

234 """ 

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

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

237 # mask must be set if send data from client 

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

239 

240 def format(self) -> bytes: 

241 """ 

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

243 """ 

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

245 raise ValueError("not 0 or 1") 

246 if self.opcode not in ABNF.OPCODES: 

247 raise ValueError("Invalid OPCODE") 

248 length = len(self.data) 

249 if length >= ABNF.LENGTH_63: 

250 raise ValueError("data is too long") 

251 

252 frame_header = chr( 

253 self.fin << 7 

254 | self.rsv1 << 6 

255 | self.rsv2 << 5 

256 | self.rsv3 << 4 

257 | self.opcode 

258 ).encode("latin-1") 

259 if length < ABNF.LENGTH_7: 

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

261 elif length < ABNF.LENGTH_16: 

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

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

264 else: 

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

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

267 

268 if not self.mask_value: 

269 if isinstance(self.data, str): 

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

271 return frame_header + self.data 

272 mask_key = self.get_mask_key(4) 

273 return frame_header + self._get_masked(mask_key) 

274 

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

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

277 

278 if isinstance(mask_key, str): 

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

280 

281 return mask_key + s 

282 

283 @staticmethod 

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

285 """ 

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

287 

288 Parameters 

289 ---------- 

290 mask_key: bytes or str 

291 4 byte mask. 

292 data: bytes or str 

293 data to mask/unmask. 

294 """ 

295 if data is None: 

296 data = "" 

297 

298 if isinstance(mask_key, str): 

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

300 

301 if isinstance(data, str): 

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

303 

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

305 

306 

307class frame_buffer: 

308 _HEADER_MASK_INDEX = 5 

309 _HEADER_LENGTH_INDEX = 6 

310 

311 def __init__( 

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

313 ) -> None: 

314 self.recv = recv_fn 

315 self.skip_utf8_validation = skip_utf8_validation 

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

317 # bytes of bytes are received. 

318 self.recv_buffer: list = [] 

319 self.clear() 

320 self.lock = Lock() 

321 

322 def clear(self) -> None: 

323 self.header: Optional[tuple] = None 

324 self.length: Optional[int] = None 

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

326 

327 def needs_header(self) -> bool: 

328 return self.header is None 

329 

330 def recv_header(self) -> None: 

331 header = self.recv_strict(2) 

332 b1 = header[0] 

333 fin = b1 >> 7 & 1 

334 rsv1 = b1 >> 6 & 1 

335 rsv2 = b1 >> 5 & 1 

336 rsv3 = b1 >> 4 & 1 

337 opcode = b1 & 0xF 

338 b2 = header[1] 

339 has_mask = b2 >> 7 & 1 

340 length_bits = b2 & 0x7F 

341 

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

343 

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

345 if not self.header: 

346 return False 

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

348 return header_val 

349 

350 def needs_length(self) -> bool: 

351 return self.length is None 

352 

353 def recv_length(self) -> None: 

354 if self.header is None: 

355 raise WebSocketProtocolException("Header not received") 

356 bits = self.header[frame_buffer._HEADER_LENGTH_INDEX] 

357 length_bits = bits & 0x7F 

358 if length_bits == 0x7E: 

359 v = self.recv_strict(2) 

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

361 elif length_bits == 0x7F: 

362 v = self.recv_strict(8) 

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

364 else: 

365 self.length = length_bits 

366 

367 def needs_mask(self) -> bool: 

368 return self.mask_value is None 

369 

370 def recv_mask(self) -> None: 

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

372 

373 def recv_frame(self) -> ABNF: 

374 with self.lock: 

375 # Header 

376 if self.needs_header(): 

377 self.recv_header() 

378 if self.header is None: 

379 raise WebSocketProtocolException("Header not received") 

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

381 

382 # Frame length 

383 if self.needs_length(): 

384 self.recv_length() 

385 length = self.length 

386 

387 # Mask 

388 if self.needs_mask(): 

389 self.recv_mask() 

390 mask_value = self.mask_value 

391 

392 # Payload 

393 if length is None: 

394 raise WebSocketProtocolException("Length not received") 

395 payload = self.recv_strict(length) 

396 if has_mask: 

397 if mask_value is None: 

398 raise WebSocketProtocolException("Mask not received") 

399 payload = ABNF.mask(mask_value, payload) 

400 

401 # Reset for next frame 

402 self.clear() 

403 

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

405 frame.validate(self.skip_utf8_validation) 

406 

407 return frame 

408 

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

410 if not isinstance(bufsize, int): 

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

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

413 while shortage > 0: 

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

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

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

417 # yet passing large numbers repeatedly causes lots of large 

418 # buffers allocated and then shrunk, which results in 

419 # fragmentation. 

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

421 if isinstance(bytes_, bytes): 

422 self.recv_buffer.append(bytes_) 

423 shortage -= len(bytes_) 

424 else: 

425 # Handle case where recv returns int or other type 

426 break 

427 

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

429 

430 if shortage == 0: 

431 self.recv_buffer = [] 

432 return unified 

433 else: 

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

435 return unified[:bufsize] 

436 

437 

438class continuous_frame: 

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

440 self.fire_cont_frame = fire_cont_frame 

441 self.skip_utf8_validation = skip_utf8_validation 

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

443 self.recving_frames: Optional[int] = None 

444 

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

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

447 raise WebSocketProtocolException("Illegal frame") 

448 if self.recving_frames and frame.opcode in ( 

449 ABNF.OPCODE_TEXT, 

450 ABNF.OPCODE_BINARY, 

451 ): 

452 raise WebSocketProtocolException("Illegal frame") 

453 

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

455 if self.cont_data: 

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

457 else: 

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

459 self.recving_frames = frame.opcode 

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

461 

462 if frame.fin: 

463 self.recving_frames = None 

464 

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

466 return frame.fin or self.fire_cont_frame 

467 

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

469 data = self.cont_data 

470 if data is None: 

471 raise WebSocketProtocolException("No continuation data available") 

472 self.cont_data = None 

473 frame.data = data[1] 

474 if ( 

475 not self.fire_cont_frame 

476 and data is not None 

477 and data[0] == ABNF.OPCODE_TEXT 

478 and not self.skip_utf8_validation 

479 and not validate_utf8(frame.data) 

480 ): 

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

482 if data is None: 

483 raise WebSocketProtocolException("No continuation data available") 

484 return data[0], frame