Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/websocket/_core.py: 20%
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1import socket
2import struct
3import threading
4import time
5from typing import Any, Callable, Optional, Type, Union
7# websocket modules
8from ._abnf import ABNF, STATUS_NORMAL, continuous_frame, frame_buffer
9from ._exceptions import (
10 WebSocketProtocolException,
11 WebSocketConnectionClosedException,
12 WebSocketTimeoutException,
13 WebSocketException,
14)
15from ._handshake import SUPPORTED_REDIRECT_STATUSES, handshake, handshake_response
16from ._http import connect, proxy_info
17from ._logging import debug, error, trace, isEnabledForError, isEnabledForTrace
18from ._socket import getdefaulttimeout, recv, send, sock_opt
19from ._ssl_compat import ssl
20from ._utils import NoLock
21from ._dispatcher import DispatcherBase, WrappedDispatcher
23"""
24_core.py
25websocket - WebSocket client library for Python
27Copyright 2026 engn33r
29Licensed under the Apache License, Version 2.0 (the "License");
30you may not use this file except in compliance with the License.
31You may obtain a copy of the License at
33 http://www.apache.org/licenses/LICENSE-2.0
35Unless required by applicable law or agreed to in writing, software
36distributed under the License is distributed on an "AS IS" BASIS,
37WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
38See the License for the specific language governing permissions and
39limitations under the License.
40"""
42__all__ = ["WebSocket", "create_connection"]
45def _normalize_close_reason(reason: Union[str, bytes, None]) -> bytes:
46 """Convert a close reason into the UTF-8 bytes for a close-frame payload."""
47 if reason is None:
48 return b""
49 if isinstance(reason, str):
50 return reason.encode("utf-8")
51 if isinstance(reason, bytes):
52 return reason
53 return bytes(reason)
56class WebSocket:
57 """
58 Low level WebSocket interface.
60 This class is based on the WebSocket protocol `draft-hixie-thewebsocketprotocol-76 <http://tools.ietf.org/html/draft-hixie-thewebsocketprotocol-76>`_
62 We can connect to the websocket server and send/receive data.
63 The following example is an echo client.
65 >>> import websocket
66 >>> ws = websocket.WebSocket()
67 >>> ws.connect("ws://websockets.chilkat.io/wsChilkatEcho.ashx")
68 >>> ws.send("Hello, Server")
69 19
70 >>> ws.recv()
71 'Hello, Server'
72 >>> ws.close()
74 Parameters
75 ----------
76 get_mask_key: func
77 A callable function to get new mask keys, see the
78 WebSocket.set_mask_key's docstring for more information.
79 sockopt: tuple
80 Values for socket.setsockopt.
81 sockopt must be tuple and each element is argument of sock.setsockopt.
82 sslopt: dict
83 Optional dict object for ssl socket options. See FAQ for details.
84 fire_cont_frame: bool
85 Fire recv event for each cont frame. Default is False.
86 enable_multithread: bool
87 If set to True, lock send method.
88 skip_utf8_validation: bool
89 Skip utf8 validation.
90 """
92 def __init__(
93 self,
94 get_mask_key: Optional[Callable] = None,
95 sockopt: Optional[list] = None,
96 sslopt: Optional[dict] = None,
97 fire_cont_frame: bool = False,
98 enable_multithread: bool = True,
99 skip_utf8_validation: bool = False,
100 dispatcher: Optional[Union[DispatcherBase, WrappedDispatcher]] = None,
101 **_: Any,
102 ) -> None:
103 """
104 Initialize WebSocket object.
106 Parameters
107 ----------
108 sslopt: dict
109 Optional dict object for ssl socket options. See FAQ for details.
110 """
111 self.sock_opt = sock_opt(sockopt, sslopt)
112 self.handshake_response: Optional[handshake_response] = None
113 self.sock: Optional[socket.socket] = None
115 self.connected = False
116 self.close_frame: Optional[ABNF] = None
117 self.get_mask_key = get_mask_key
118 # These buffer over the build-up of a single frame.
119 self.frame_buffer = frame_buffer(self._recv, skip_utf8_validation)
120 self.cont_frame = continuous_frame(fire_cont_frame, skip_utf8_validation)
121 self.dispatcher = dispatcher
123 if enable_multithread:
124 self.lock = threading.Lock()
125 self.readlock = threading.Lock()
126 else:
127 self.lock = NoLock() # type: ignore[assignment]
128 self.readlock = NoLock() # type: ignore[assignment]
130 def __iter__(self):
131 """
132 Allow iteration over websocket, implying sequential `recv` executions.
133 """
134 while True:
135 yield self.recv()
137 def __next__(self):
138 return self.recv()
140 def next(self):
141 return self.__next__()
143 def fileno(self):
144 if self.sock is None:
145 raise WebSocketException("Connection not established")
146 return self.sock.fileno()
148 def set_mask_key(self, func):
149 """
150 Set function to create mask key. You can customize mask key generator.
151 Mainly, this is for testing purpose.
153 Parameters
154 ----------
155 func: func
156 callable object. the func takes 1 argument as integer.
157 The argument means length of mask key.
158 This func must return string(byte array),
159 which length is argument specified.
160 """
161 self.get_mask_key = func
163 def gettimeout(self) -> Optional[Union[float, int]]:
164 """
165 Get the websocket timeout (in seconds) as an int or float
167 Returns
168 ----------
169 timeout: int or float
170 returns timeout value (in seconds). This value could be either float/integer.
171 """
172 return self.sock_opt.timeout
174 def settimeout(self, timeout: Optional[Union[float, int]]) -> None:
175 """
176 Set the timeout to the websocket.
178 Parameters
179 ----------
180 timeout: int or float
181 timeout time (in seconds). This value could be either float/integer.
182 """
183 self.sock_opt.timeout = timeout
184 if self.sock:
185 self.sock.settimeout(timeout)
187 timeout = property(gettimeout, settimeout)
189 def getsubprotocol(self) -> Optional[str]:
190 """
191 Get subprotocol
192 """
193 if self.handshake_response:
194 return self.handshake_response.subprotocol
195 else:
196 return None
198 subprotocol = property(getsubprotocol)
200 def getstatus(self) -> Optional[int]:
201 """
202 Get handshake status
203 """
204 if self.handshake_response:
205 return self.handshake_response.status
206 else:
207 return None
209 status = property(getstatus)
211 def getheaders(self) -> Optional[dict]:
212 """
213 Get handshake response header
214 """
215 if self.handshake_response:
216 return self.handshake_response.headers
217 else:
218 return None
220 def is_ssl(self):
221 try:
222 return isinstance(self.sock, ssl.SSLSocket)
223 except (AttributeError, NameError):
224 return False
226 headers = property(getheaders)
228 def connect(self, url, **options):
229 """
230 Connect to url. url is websocket url scheme.
231 ie. ws://host:port/resource
232 You can customize using 'options'.
233 If you set "header" list object, you can set your own custom header.
235 >>> ws = WebSocket()
236 >>> ws.connect("ws://websockets.chilkat.io/wsChilkatEcho.ashx",
237 ... header=["User-Agent: MyProgram",
238 ... "x-custom: header"])
240 Parameters
241 ----------
242 header: list or dict
243 Custom http header list or dict.
244 cookie: str
245 Cookie value.
246 origin: str
247 Custom origin url.
248 connection: str
249 Custom connection header value.
250 Default value "Upgrade" set in _handshake.py
251 suppress_origin: bool
252 Suppress outputting origin header.
253 suppress_host: bool
254 Suppress outputting host header.
255 host: str
256 Custom host header string.
257 timeout: int or float
258 Socket timeout time. This value is an integer or float.
259 If you set None for this value, it means "use default_timeout value"
260 http_proxy_host: str
261 HTTP proxy host name.
262 http_proxy_port: str or int
263 HTTP proxy port. Required when http_proxy_host is set. Proxies
264 from environment variables default to port 80.
265 http_no_proxy: list
266 Whitelisted host names that don't use the proxy.
267 http_proxy_auth: tuple
268 HTTP proxy auth information. Tuple of username and password. Default is None.
269 http_proxy_timeout: int or float
270 HTTP proxy timeout, default is 60 sec as per python-socks.
271 redirect_limit: int
272 Number of redirects to follow.
273 subprotocols: list
274 List of available subprotocols. Default is None.
275 socket: socket
276 Pre-initialized stream socket.
277 """
278 self.sock_opt.timeout = options.get("timeout", self.sock_opt.timeout)
279 self.sock, addrs = connect(
280 url, self.sock_opt, proxy_info(**options), options.pop("socket", None)
281 )
283 try:
284 self.handshake_response = handshake(self.sock, url, *addrs, **options)
285 for _ in range(options.pop("redirect_limit", 3)):
286 if (
287 self.handshake_response is not None
288 and self.handshake_response.status in SUPPORTED_REDIRECT_STATUSES
289 ):
290 url = self.handshake_response.headers.get("location")
291 if url is None:
292 raise WebSocketException(
293 "Redirect response without Location header, "
294 f"status {self.handshake_response.status}"
295 )
296 self.sock.close()
297 try:
298 self.sock, addrs = connect(
299 url,
300 self.sock_opt,
301 proxy_info(**options),
302 options.pop("socket", None),
303 )
304 except ValueError as e:
305 raise WebSocketException(
306 f"Invalid redirect target {url!r}: {e}"
307 ) from e
308 self.handshake_response = handshake(
309 self.sock, url, *addrs, **options
310 )
311 if (
312 self.handshake_response is not None
313 and self.handshake_response.status in SUPPORTED_REDIRECT_STATUSES
314 ):
315 raise WebSocketException("Redirect limit exhausted")
316 self.connected = True
317 except:
318 if self.sock:
319 self.sock.close()
320 self.sock = None
321 raise
323 def send(self, payload: Union[bytes, str], opcode: int = ABNF.OPCODE_TEXT) -> int:
324 """
325 Send the data as string.
327 Parameters
328 ----------
329 payload: str
330 Payload must be utf-8 string or unicode,
331 If the opcode is OPCODE_TEXT.
332 Otherwise, it must be string(byte array).
333 opcode: int
334 Operation code (opcode) to send.
335 """
337 frame = ABNF.create_frame(payload, opcode)
338 return self.send_frame(frame)
340 def send_text(self, text_data: str) -> int:
341 """
342 Sends UTF-8 encoded text.
343 """
344 return self.send(text_data, ABNF.OPCODE_TEXT)
346 def send_bytes(self, data: Union[bytes, bytearray]) -> int:
347 """
348 Sends a sequence of bytes.
349 """
350 return self.send(data, ABNF.OPCODE_BINARY)
352 def send_frame(self, frame: ABNF) -> int:
353 """
354 Send the data frame.
356 >>> ws = create_connection("ws://websockets.chilkat.io/wsChilkatEcho.ashx")
357 >>> frame = ABNF.create_frame("Hello", ABNF.OPCODE_TEXT)
358 >>> ws.send_frame(frame)
359 >>> cont_frame = ABNF.create_frame("My name is ", ABNF.OPCODE_CONT, 0)
360 >>> ws.send_frame(frame)
361 >>> cont_frame = ABNF.create_frame("Foo Bar", ABNF.OPCODE_CONT, 1)
362 >>> ws.send_frame(frame)
364 Parameters
365 ----------
366 frame: ABNF frame
367 frame data created by ABNF.create_frame
368 """
369 if self.get_mask_key:
370 frame.get_mask_key = self.get_mask_key
371 data = frame.format()
372 length = len(data)
373 if isEnabledForTrace():
374 trace(f"++Sent raw: {repr(data)}")
375 trace(f"++Sent decoded: {frame.__str__()}")
376 with self.lock:
377 while data:
378 bytes_sent = self._send(data)
379 data = data[bytes_sent:]
381 return length
383 def send_binary(self, payload: bytes) -> int:
384 """
385 Send a binary message (OPCODE_BINARY).
387 Parameters
388 ----------
389 payload: bytes
390 payload of message to send.
391 """
392 return self.send(payload, ABNF.OPCODE_BINARY)
394 def ping(self, payload: Union[str, bytes] = "") -> None:
395 """
396 Send ping data.
398 Parameters
399 ----------
400 payload: str
401 data payload to send server.
402 """
403 if isinstance(payload, str):
404 payload = payload.encode("utf-8")
405 self.send(payload, ABNF.OPCODE_PING)
407 def pong(self, payload: Union[str, bytes] = "") -> None:
408 """
409 Send pong data.
411 Parameters
412 ----------
413 payload: str
414 data payload to send server.
415 """
416 if isinstance(payload, str):
417 payload = payload.encode("utf-8")
418 self.send(payload, ABNF.OPCODE_PONG)
420 def recv(self) -> Union[str, bytes]:
421 """
422 Receive string data(byte array) from the server.
424 Returns
425 ----------
426 data: string (byte array) value.
427 """
428 with self.readlock:
429 opcode, data = self.recv_data()
430 if opcode == ABNF.OPCODE_TEXT:
431 data_received: Union[bytes, str] = data
432 if isinstance(data_received, bytes):
433 return data_received.decode("utf-8")
434 elif isinstance(data_received, str):
435 return data_received
436 elif opcode == ABNF.OPCODE_BINARY:
437 data_binary: bytes = data
438 return data_binary
439 else:
440 return ""
442 def recv_data(self, control_frame: bool = False) -> tuple:
443 """
444 Receive data with operation code.
446 Parameters
447 ----------
448 control_frame: bool
449 a boolean flag indicating whether to return control frame
450 data, defaults to False
452 Returns
453 -------
454 opcode, frame.data: tuple
455 tuple of operation code and string(byte array) value.
456 """
457 opcode, frame = self.recv_data_frame(control_frame)
458 return opcode, frame.data
460 def recv_data_frame(self, control_frame: bool = False) -> tuple:
461 """
462 Receive data with operation code.
464 If a valid ping message is received, a pong response is sent.
466 Parameters
467 ----------
468 control_frame: bool
469 a boolean flag indicating whether to return control frame
470 data, defaults to False
472 Returns
473 -------
474 frame.opcode, frame: tuple
475 tuple of operation code and string(byte array) value.
476 """
477 while True:
478 frame = self.recv_frame()
479 if isEnabledForTrace():
480 trace(f"++Rcv raw: {repr(frame.format())}")
481 trace(f"++Rcv decoded: {frame.__str__()}")
482 if frame.opcode in (
483 ABNF.OPCODE_TEXT,
484 ABNF.OPCODE_BINARY,
485 ABNF.OPCODE_CONT,
486 ):
487 self.cont_frame.validate(frame)
488 self.cont_frame.add(frame)
490 if self.cont_frame.is_fire(frame):
491 return self.cont_frame.extract(frame)
493 elif frame.opcode == ABNF.OPCODE_CLOSE:
494 self.send_close()
495 return frame.opcode, frame
496 elif frame.opcode == ABNF.OPCODE_PING:
497 if len(frame.data) < 126:
498 self.pong(frame.data)
499 else:
500 raise WebSocketProtocolException("Ping message is too long")
501 if control_frame:
502 return frame.opcode, frame
503 elif frame.opcode == ABNF.OPCODE_PONG:
504 if control_frame:
505 return frame.opcode, frame
507 def recv_frame(self):
508 """
509 Receive data as frame from server.
511 Returns
512 -------
513 self.frame_buffer.recv_frame(): ABNF frame object
514 """
515 return self.frame_buffer.recv_frame()
517 def send_close(
518 self, status: int = STATUS_NORMAL, reason: Union[str, bytes] = b""
519 ) -> None:
520 """
521 Send close data to the server.
523 Parameters
524 ----------
525 status: int
526 Status code to send. See STATUS_XXX.
527 reason: str or bytes
528 The reason to close. This must be string or UTF-8 bytes.
529 """
530 if status < 0 or status >= ABNF.LENGTH_16:
531 raise ValueError("code is invalid range")
533 reason_bytes = _normalize_close_reason(reason)
535 self.connected = False
536 self.send(struct.pack("!H", status) + reason_bytes, ABNF.OPCODE_CLOSE)
538 def close(
539 self,
540 status: int = STATUS_NORMAL,
541 reason: Union[str, bytes] = b"",
542 timeout: Optional[Union[int, float]] = 3,
543 ) -> None:
544 """
545 Close Websocket object
547 Parameters
548 ----------
549 status: int
550 Status code to send. See VALID_CLOSE_STATUS in ABNF.
551 reason: str or bytes
552 The reason to close in UTF-8.
553 timeout: int or float
554 Timeout until receive a close frame.
555 If None, it will wait forever until receive a close frame.
556 """
557 if not self.connected:
558 return
559 if status < 0 or status >= ABNF.LENGTH_16:
560 raise ValueError("code is invalid range")
562 # Reset close_frame to avoid stale data from previous connections
563 self.close_frame = None
565 try:
566 self.connected = False
567 self.send(
568 struct.pack("!H", status) + _normalize_close_reason(reason),
569 ABNF.OPCODE_CLOSE,
570 )
571 if self.sock is None:
572 return
573 sock_timeout = self.sock.gettimeout()
574 self.sock.settimeout(timeout)
575 start_time = time.time()
576 while timeout is None or time.time() - start_time < timeout:
577 try:
578 frame = self.recv_frame()
579 if frame.opcode != ABNF.OPCODE_CLOSE:
580 continue
581 # Store the peer's close frame for access by higher-level APIs
582 self.close_frame = frame
583 if isEnabledForError():
584 recv_status = struct.unpack("!H", frame.data[0:2])[0]
585 if recv_status >= 3000 and recv_status <= 4999:
586 debug(f"close status: {repr(recv_status)}")
587 elif recv_status != STATUS_NORMAL:
588 error(f"close status: {repr(recv_status)}")
589 break
590 except (
591 WebSocketConnectionClosedException,
592 WebSocketTimeoutException,
593 struct.error,
594 ):
595 break
596 if self.sock is not None:
597 self.sock.settimeout(sock_timeout)
598 self.sock.shutdown(socket.SHUT_RDWR)
599 except:
600 pass
602 self.shutdown()
604 def abort(self):
605 """
606 Low-level asynchronous abort, wakes up other threads that are waiting in recv_*
607 """
608 if self.connected and self.sock is not None:
609 try:
610 self.sock.shutdown(socket.SHUT_RDWR)
611 except (OSError, AttributeError):
612 # Socket already closed or never connected
613 # abort() is best-effort, like shutdown()
614 debug("Socket already closed during abort")
616 def shutdown(self):
617 """
618 close socket, immediately.
619 """
620 if self.sock:
621 try:
622 # Check if socket is still open before closing
623 if not self.sock._closed:
624 self.sock.close()
625 except (OSError, AttributeError):
626 # Socket already closed or invalid file descriptor - this can happen
627 # during reconnection scenarios when network failures occur
628 debug("Socket already closed during shutdown")
629 pass
630 finally:
631 self.sock = None
632 self.connected = False
634 def _send(self, data: Union[str, bytes]) -> int:
635 if self.sock is None:
636 raise WebSocketConnectionClosedException("socket is already closed.")
637 if self.dispatcher:
638 return self.dispatcher.send(self.sock, data)
639 return send(self.sock, data)
641 def _recv(self, bufsize):
642 if self.sock is None:
643 raise WebSocketConnectionClosedException("Connection is closed")
644 try:
645 return recv(self.sock, bufsize)
646 except WebSocketConnectionClosedException:
647 if self.sock:
648 self.sock.close()
649 self.sock = None
650 self.connected = False
651 raise
654def create_connection(
655 url: str,
656 timeout: Optional[Union[float, int]] = None,
657 class_: Type[WebSocket] = WebSocket,
658 **options: Any,
659) -> WebSocket:
660 """
661 Connect to url and return websocket object.
663 Connect to url and return the WebSocket object.
664 Passing optional timeout parameter will set the timeout on the socket.
665 If no timeout is supplied,
666 the global default timeout setting returned by getdefaulttimeout() is used.
667 You can customize using 'options'.
668 If you set "header" list object, you can set your own custom header.
670 >>> conn = create_connection("ws://websockets.chilkat.io/wsChilkatEcho.ashx",
671 ... header=["User-Agent: MyProgram",
672 ... "x-custom: header"])
674 Parameters
675 ----------
676 class_: class
677 class to instantiate when creating the connection. It has to implement
678 settimeout and connect. It's __init__ should be compatible with
679 WebSocket.__init__, i.e. accept all of it's kwargs.
680 header: list or dict
681 custom http header list or dict.
682 cookie: str
683 Cookie value.
684 origin: str
685 custom origin url.
686 suppress_origin: bool
687 suppress outputting origin header.
688 host: str
689 custom host header string.
690 timeout: int or float
691 socket timeout time. This value could be either float/integer.
692 If set to None, it uses the default_timeout value.
693 http_proxy_host: str
694 HTTP proxy host name.
695 http_proxy_port: str or int
696 HTTP proxy port. Required when http_proxy_host is set. Proxies
697 from environment variables default to port 80.
698 http_no_proxy: list
699 Whitelisted host names that don't use the proxy.
700 http_proxy_auth: tuple
701 HTTP proxy auth information. tuple of username and password. Default is None.
702 http_proxy_timeout: int or float
703 HTTP proxy timeout, default is 60 sec as per python-socks.
704 enable_multithread: bool
705 Enable lock for multithread.
706 redirect_limit: int
707 Number of redirects to follow.
708 sockopt: tuple
709 Values for socket.setsockopt.
710 sockopt must be a tuple and each element is an argument of sock.setsockopt.
711 sslopt: dict
712 Optional dict object for ssl socket options. See FAQ for details.
713 subprotocols: list
714 List of available subprotocols. Default is None.
715 skip_utf8_validation: bool
716 Skip utf8 validation.
717 socket: socket
718 Pre-initialized stream socket.
719 """
720 sockopt = options.pop("sockopt", [])
721 sslopt = options.pop("sslopt", {})
722 fire_cont_frame = options.pop("fire_cont_frame", False)
723 enable_multithread = options.pop("enable_multithread", True)
724 skip_utf8_validation = options.pop("skip_utf8_validation", False)
725 websock = class_(
726 sockopt=sockopt,
727 sslopt=sslopt,
728 fire_cont_frame=fire_cont_frame,
729 enable_multithread=enable_multithread,
730 skip_utf8_validation=skip_utf8_validation,
731 **options,
732 )
733 websock.settimeout(timeout if timeout is not None else getdefaulttimeout())
734 websock.connect(url, **options)
735 return websock