chore: 添加虚拟环境到仓库
- 添加 backend_service/venv 虚拟环境 - 包含所有Python依赖包 - 注意:虚拟环境约393MB,包含12655个文件
This commit is contained in:
@@ -0,0 +1,30 @@
|
||||
"""
|
||||
__init__.py
|
||||
websocket - WebSocket client library for Python
|
||||
|
||||
Copyright 2025 engn33r
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
from ._abnf import * # noqa: F401,F403
|
||||
from ._app import ( # noqa: F401
|
||||
WebSocketApp as WebSocketApp,
|
||||
set_reconnect as set_reconnect,
|
||||
)
|
||||
from ._core import * # noqa: F401,F403
|
||||
from ._exceptions import * # noqa: F401,F403
|
||||
from ._logging import * # noqa: F401,F403
|
||||
from ._socket import * # noqa: F401,F403
|
||||
|
||||
__version__ = "1.9.0"
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,482 @@
|
||||
import array
|
||||
import os
|
||||
import struct
|
||||
import sys
|
||||
from threading import Lock
|
||||
from typing import Callable, Optional, Union, Any
|
||||
|
||||
from ._exceptions import WebSocketPayloadException, WebSocketProtocolException
|
||||
from ._utils import validate_utf8
|
||||
|
||||
"""
|
||||
_abnf.py
|
||||
websocket - WebSocket client library for Python
|
||||
|
||||
Copyright 2025 engn33r
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
try:
|
||||
# If wsaccel is available, use compiled routines to mask data.
|
||||
# wsaccel only provides around a 10% speed boost compared
|
||||
# to the websocket-client _mask() implementation.
|
||||
# Note that wsaccel is unmaintained.
|
||||
from wsaccel.xormask import XorMaskerSimple
|
||||
|
||||
def _mask(mask_value: array.array, data_value: array.array) -> bytes:
|
||||
mask_result: bytes = XorMaskerSimple(mask_value).process(data_value)
|
||||
return mask_result
|
||||
|
||||
except ImportError:
|
||||
# wsaccel is not available, use websocket-client _mask()
|
||||
native_byteorder = sys.byteorder
|
||||
|
||||
def _mask(mask_value: array.array, data_value: array.array) -> bytes:
|
||||
datalen = len(data_value)
|
||||
int_data_value = int.from_bytes(data_value, native_byteorder)
|
||||
int_mask_value = int.from_bytes(
|
||||
mask_value * (datalen // 4) + mask_value[: datalen % 4], native_byteorder
|
||||
)
|
||||
return (int_data_value ^ int_mask_value).to_bytes(datalen, native_byteorder)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ABNF",
|
||||
"continuous_frame",
|
||||
"frame_buffer",
|
||||
"STATUS_NORMAL",
|
||||
"STATUS_GOING_AWAY",
|
||||
"STATUS_PROTOCOL_ERROR",
|
||||
"STATUS_UNSUPPORTED_DATA_TYPE",
|
||||
"STATUS_STATUS_NOT_AVAILABLE",
|
||||
"STATUS_ABNORMAL_CLOSED",
|
||||
"STATUS_INVALID_PAYLOAD",
|
||||
"STATUS_POLICY_VIOLATION",
|
||||
"STATUS_MESSAGE_TOO_BIG",
|
||||
"STATUS_INVALID_EXTENSION",
|
||||
"STATUS_UNEXPECTED_CONDITION",
|
||||
"STATUS_BAD_GATEWAY",
|
||||
"STATUS_TLS_HANDSHAKE_ERROR",
|
||||
]
|
||||
|
||||
# closing frame status codes.
|
||||
STATUS_NORMAL = 1000
|
||||
STATUS_GOING_AWAY = 1001
|
||||
STATUS_PROTOCOL_ERROR = 1002
|
||||
STATUS_UNSUPPORTED_DATA_TYPE = 1003
|
||||
STATUS_STATUS_NOT_AVAILABLE = 1005
|
||||
STATUS_ABNORMAL_CLOSED = 1006
|
||||
STATUS_INVALID_PAYLOAD = 1007
|
||||
STATUS_POLICY_VIOLATION = 1008
|
||||
STATUS_MESSAGE_TOO_BIG = 1009
|
||||
STATUS_INVALID_EXTENSION = 1010
|
||||
STATUS_UNEXPECTED_CONDITION = 1011
|
||||
STATUS_SERVICE_RESTART = 1012
|
||||
STATUS_TRY_AGAIN_LATER = 1013
|
||||
STATUS_BAD_GATEWAY = 1014
|
||||
STATUS_TLS_HANDSHAKE_ERROR = 1015
|
||||
|
||||
VALID_CLOSE_STATUS = (
|
||||
STATUS_NORMAL,
|
||||
STATUS_GOING_AWAY,
|
||||
STATUS_PROTOCOL_ERROR,
|
||||
STATUS_UNSUPPORTED_DATA_TYPE,
|
||||
STATUS_INVALID_PAYLOAD,
|
||||
STATUS_POLICY_VIOLATION,
|
||||
STATUS_MESSAGE_TOO_BIG,
|
||||
STATUS_INVALID_EXTENSION,
|
||||
STATUS_UNEXPECTED_CONDITION,
|
||||
STATUS_SERVICE_RESTART,
|
||||
STATUS_TRY_AGAIN_LATER,
|
||||
STATUS_BAD_GATEWAY,
|
||||
)
|
||||
|
||||
|
||||
class ABNF:
|
||||
"""
|
||||
ABNF frame class.
|
||||
See http://tools.ietf.org/html/rfc5234
|
||||
and http://tools.ietf.org/html/rfc6455#section-5.2
|
||||
"""
|
||||
|
||||
# operation code values.
|
||||
OPCODE_CONT = 0x0
|
||||
OPCODE_TEXT = 0x1
|
||||
OPCODE_BINARY = 0x2
|
||||
OPCODE_CLOSE = 0x8
|
||||
OPCODE_PING = 0x9
|
||||
OPCODE_PONG = 0xA
|
||||
|
||||
# available operation code value tuple
|
||||
OPCODES = (
|
||||
OPCODE_CONT,
|
||||
OPCODE_TEXT,
|
||||
OPCODE_BINARY,
|
||||
OPCODE_CLOSE,
|
||||
OPCODE_PING,
|
||||
OPCODE_PONG,
|
||||
)
|
||||
|
||||
# opcode human readable string
|
||||
OPCODE_MAP = {
|
||||
OPCODE_CONT: "cont",
|
||||
OPCODE_TEXT: "text",
|
||||
OPCODE_BINARY: "binary",
|
||||
OPCODE_CLOSE: "close",
|
||||
OPCODE_PING: "ping",
|
||||
OPCODE_PONG: "pong",
|
||||
}
|
||||
|
||||
# data length threshold.
|
||||
LENGTH_7 = 0x7E
|
||||
LENGTH_16 = 1 << 16
|
||||
LENGTH_63 = 1 << 63
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
fin: int = 0,
|
||||
rsv1: int = 0,
|
||||
rsv2: int = 0,
|
||||
rsv3: int = 0,
|
||||
opcode: int = OPCODE_TEXT,
|
||||
mask_value: int = 1,
|
||||
data: Optional[Union[str, bytes]] = "",
|
||||
) -> None:
|
||||
"""
|
||||
Constructor for ABNF. Please check RFC for arguments.
|
||||
"""
|
||||
self.fin = fin
|
||||
self.rsv1 = rsv1
|
||||
self.rsv2 = rsv2
|
||||
self.rsv3 = rsv3
|
||||
self.opcode = opcode
|
||||
self.mask_value = mask_value
|
||||
if data is None:
|
||||
data = ""
|
||||
self.data = data
|
||||
self.get_mask_key = os.urandom
|
||||
|
||||
def validate(self, skip_utf8_validation: bool = False) -> None:
|
||||
"""
|
||||
Validate the ABNF frame.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
skip_utf8_validation: skip utf8 validation.
|
||||
"""
|
||||
if self.rsv1 or self.rsv2 or self.rsv3:
|
||||
raise WebSocketProtocolException("rsv is not implemented, yet")
|
||||
|
||||
if self.opcode not in ABNF.OPCODES:
|
||||
raise WebSocketProtocolException("Invalid opcode %r", self.opcode)
|
||||
|
||||
if self.opcode == ABNF.OPCODE_PING and not self.fin:
|
||||
raise WebSocketProtocolException("Invalid ping frame.")
|
||||
|
||||
if self.opcode == ABNF.OPCODE_CLOSE:
|
||||
data_length = len(self.data)
|
||||
if not data_length:
|
||||
return
|
||||
if data_length == 1 or data_length >= 126:
|
||||
raise WebSocketProtocolException("Invalid close frame.")
|
||||
if (
|
||||
data_length > 2
|
||||
and not skip_utf8_validation
|
||||
and not validate_utf8(self.data[2:])
|
||||
):
|
||||
raise WebSocketProtocolException("Invalid close frame.")
|
||||
|
||||
data_bytes = (
|
||||
self.data[:2]
|
||||
if isinstance(self.data, bytes)
|
||||
else self.data[:2].encode("utf-8")
|
||||
)
|
||||
code = struct.unpack("!H", data_bytes)[0]
|
||||
if not self._is_valid_close_status(code):
|
||||
raise WebSocketProtocolException("Invalid close opcode %r", code)
|
||||
|
||||
@staticmethod
|
||||
def _is_valid_close_status(code: int) -> bool:
|
||||
return code in VALID_CLOSE_STATUS or (3000 <= code < 5000)
|
||||
|
||||
def __str__(self) -> str:
|
||||
data_repr = self.data if isinstance(self.data, str) else repr(self.data)
|
||||
return f"fin={self.fin} opcode={self.opcode} data={data_repr}"
|
||||
|
||||
@staticmethod
|
||||
def create_frame(data: Union[bytes, str], opcode: int, fin: int = 1) -> "ABNF":
|
||||
"""
|
||||
Create frame to send text, binary and other data.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data: str
|
||||
data to send. This is string value(byte array).
|
||||
If opcode is OPCODE_TEXT and this value is unicode,
|
||||
data value is converted into unicode string, automatically.
|
||||
opcode: int
|
||||
operation code. please see OPCODE_MAP.
|
||||
fin: int
|
||||
fin flag. if set to 0, create continue fragmentation.
|
||||
"""
|
||||
if opcode == ABNF.OPCODE_TEXT and isinstance(data, str):
|
||||
data = data.encode("utf-8")
|
||||
# mask must be set if send data from client
|
||||
return ABNF(fin, 0, 0, 0, opcode, 1, data)
|
||||
|
||||
def format(self) -> bytes:
|
||||
"""
|
||||
Format this object to string(byte array) to send data to server.
|
||||
"""
|
||||
if any(x not in (0, 1) for x in [self.fin, self.rsv1, self.rsv2, self.rsv3]):
|
||||
raise ValueError("not 0 or 1")
|
||||
if self.opcode not in ABNF.OPCODES:
|
||||
raise ValueError("Invalid OPCODE")
|
||||
length = len(self.data)
|
||||
if length >= ABNF.LENGTH_63:
|
||||
raise ValueError("data is too long")
|
||||
|
||||
frame_header = chr(
|
||||
self.fin << 7
|
||||
| self.rsv1 << 6
|
||||
| self.rsv2 << 5
|
||||
| self.rsv3 << 4
|
||||
| self.opcode
|
||||
).encode("latin-1")
|
||||
if length < ABNF.LENGTH_7:
|
||||
frame_header += chr(self.mask_value << 7 | length).encode("latin-1")
|
||||
elif length < ABNF.LENGTH_16:
|
||||
frame_header += chr(self.mask_value << 7 | 0x7E).encode("latin-1")
|
||||
frame_header += struct.pack("!H", length)
|
||||
else:
|
||||
frame_header += chr(self.mask_value << 7 | 0x7F).encode("latin-1")
|
||||
frame_header += struct.pack("!Q", length)
|
||||
|
||||
if not self.mask_value:
|
||||
if isinstance(self.data, str):
|
||||
self.data = self.data.encode("utf-8")
|
||||
return frame_header + self.data
|
||||
mask_key = self.get_mask_key(4)
|
||||
return frame_header + self._get_masked(mask_key)
|
||||
|
||||
def _get_masked(self, mask_key: Union[str, bytes]) -> bytes:
|
||||
s = ABNF.mask(mask_key, self.data)
|
||||
|
||||
if isinstance(mask_key, str):
|
||||
mask_key = mask_key.encode("utf-8")
|
||||
|
||||
return mask_key + s
|
||||
|
||||
@staticmethod
|
||||
def mask(mask_key: Union[str, bytes], data: Union[str, bytes]) -> bytes:
|
||||
"""
|
||||
Mask or unmask data. Just do xor for each byte
|
||||
|
||||
Parameters
|
||||
----------
|
||||
mask_key: bytes or str
|
||||
4 byte mask.
|
||||
data: bytes or str
|
||||
data to mask/unmask.
|
||||
"""
|
||||
if data is None:
|
||||
data = ""
|
||||
|
||||
if isinstance(mask_key, str):
|
||||
mask_key = mask_key.encode("latin-1")
|
||||
|
||||
if isinstance(data, str):
|
||||
data = data.encode("latin-1")
|
||||
|
||||
return _mask(array.array("B", mask_key), array.array("B", data))
|
||||
|
||||
|
||||
class frame_buffer:
|
||||
_HEADER_MASK_INDEX = 5
|
||||
_HEADER_LENGTH_INDEX = 6
|
||||
|
||||
def __init__(
|
||||
self, recv_fn: Callable[[int], int], skip_utf8_validation: bool
|
||||
) -> None:
|
||||
self.recv = recv_fn
|
||||
self.skip_utf8_validation = skip_utf8_validation
|
||||
# Buffers over the packets from the layer beneath until desired amount
|
||||
# bytes of bytes are received.
|
||||
self.recv_buffer: list = []
|
||||
self.clear()
|
||||
self.lock = Lock()
|
||||
|
||||
def clear(self) -> None:
|
||||
self.header: Optional[tuple] = None
|
||||
self.length: Optional[int] = None
|
||||
self.mask_value: Optional[Union[bytes, str]] = None
|
||||
|
||||
def needs_header(self) -> bool:
|
||||
return self.header is None
|
||||
|
||||
def recv_header(self) -> None:
|
||||
header = self.recv_strict(2)
|
||||
b1 = header[0]
|
||||
fin = b1 >> 7 & 1
|
||||
rsv1 = b1 >> 6 & 1
|
||||
rsv2 = b1 >> 5 & 1
|
||||
rsv3 = b1 >> 4 & 1
|
||||
opcode = b1 & 0xF
|
||||
b2 = header[1]
|
||||
has_mask = b2 >> 7 & 1
|
||||
length_bits = b2 & 0x7F
|
||||
|
||||
self.header = (fin, rsv1, rsv2, rsv3, opcode, has_mask, length_bits)
|
||||
|
||||
def has_mask(self) -> Union[bool, int]:
|
||||
if not self.header:
|
||||
return False
|
||||
header_val: int = self.header[frame_buffer._HEADER_MASK_INDEX]
|
||||
return header_val
|
||||
|
||||
def needs_length(self) -> bool:
|
||||
return self.length is None
|
||||
|
||||
def recv_length(self) -> None:
|
||||
if self.header is None:
|
||||
raise WebSocketProtocolException("Header not received")
|
||||
bits = self.header[frame_buffer._HEADER_LENGTH_INDEX]
|
||||
length_bits = bits & 0x7F
|
||||
if length_bits == 0x7E:
|
||||
v = self.recv_strict(2)
|
||||
self.length = struct.unpack("!H", v)[0]
|
||||
elif length_bits == 0x7F:
|
||||
v = self.recv_strict(8)
|
||||
self.length = struct.unpack("!Q", v)[0]
|
||||
else:
|
||||
self.length = length_bits
|
||||
|
||||
def needs_mask(self) -> bool:
|
||||
return self.mask_value is None
|
||||
|
||||
def recv_mask(self) -> None:
|
||||
self.mask_value = self.recv_strict(4) if self.has_mask() else ""
|
||||
|
||||
def recv_frame(self) -> ABNF:
|
||||
with self.lock:
|
||||
# Header
|
||||
if self.needs_header():
|
||||
self.recv_header()
|
||||
if self.header is None:
|
||||
raise WebSocketProtocolException("Header not received")
|
||||
(fin, rsv1, rsv2, rsv3, opcode, has_mask, _) = self.header
|
||||
|
||||
# Frame length
|
||||
if self.needs_length():
|
||||
self.recv_length()
|
||||
length = self.length
|
||||
|
||||
# Mask
|
||||
if self.needs_mask():
|
||||
self.recv_mask()
|
||||
mask_value = self.mask_value
|
||||
|
||||
# Payload
|
||||
if length is None:
|
||||
raise WebSocketProtocolException("Length not received")
|
||||
payload = self.recv_strict(length)
|
||||
if has_mask:
|
||||
if mask_value is None:
|
||||
raise WebSocketProtocolException("Mask not received")
|
||||
payload = ABNF.mask(mask_value, payload)
|
||||
|
||||
# Reset for next frame
|
||||
self.clear()
|
||||
|
||||
frame = ABNF(fin, rsv1, rsv2, rsv3, opcode, has_mask, payload)
|
||||
frame.validate(self.skip_utf8_validation)
|
||||
|
||||
return frame
|
||||
|
||||
def recv_strict(self, bufsize: int) -> bytes:
|
||||
if not isinstance(bufsize, int):
|
||||
raise ValueError("bufsize must be an integer")
|
||||
shortage = bufsize - sum(len(buf) for buf in self.recv_buffer)
|
||||
while shortage > 0:
|
||||
# Limit buffer size that we pass to socket.recv() to avoid
|
||||
# fragmenting the heap -- the number of bytes recv() actually
|
||||
# reads is limited by socket buffer and is relatively small,
|
||||
# yet passing large numbers repeatedly causes lots of large
|
||||
# buffers allocated and then shrunk, which results in
|
||||
# fragmentation.
|
||||
bytes_ = self.recv(min(16384, shortage))
|
||||
if isinstance(bytes_, bytes):
|
||||
self.recv_buffer.append(bytes_)
|
||||
shortage -= len(bytes_)
|
||||
else:
|
||||
# Handle case where recv returns int or other type
|
||||
break
|
||||
|
||||
unified = b"".join(self.recv_buffer)
|
||||
|
||||
if shortage == 0:
|
||||
self.recv_buffer = []
|
||||
return unified
|
||||
else:
|
||||
self.recv_buffer = [unified[bufsize:]]
|
||||
return unified[:bufsize]
|
||||
|
||||
|
||||
class continuous_frame:
|
||||
def __init__(self, fire_cont_frame: bool, skip_utf8_validation: bool) -> None:
|
||||
self.fire_cont_frame = fire_cont_frame
|
||||
self.skip_utf8_validation = skip_utf8_validation
|
||||
self.cont_data: Optional[list[Any]] = None
|
||||
self.recving_frames: Optional[int] = None
|
||||
|
||||
def validate(self, frame: ABNF) -> None:
|
||||
if not self.recving_frames and frame.opcode == ABNF.OPCODE_CONT:
|
||||
raise WebSocketProtocolException("Illegal frame")
|
||||
if self.recving_frames and frame.opcode in (
|
||||
ABNF.OPCODE_TEXT,
|
||||
ABNF.OPCODE_BINARY,
|
||||
):
|
||||
raise WebSocketProtocolException("Illegal frame")
|
||||
|
||||
def add(self, frame: ABNF) -> None:
|
||||
if self.cont_data:
|
||||
self.cont_data[1] += frame.data
|
||||
else:
|
||||
if frame.opcode in (ABNF.OPCODE_TEXT, ABNF.OPCODE_BINARY):
|
||||
self.recving_frames = frame.opcode
|
||||
self.cont_data = [frame.opcode, frame.data]
|
||||
|
||||
if frame.fin:
|
||||
self.recving_frames = None
|
||||
|
||||
def is_fire(self, frame: ABNF) -> Union[bool, int]:
|
||||
return frame.fin or self.fire_cont_frame
|
||||
|
||||
def extract(self, frame: ABNF) -> tuple:
|
||||
data = self.cont_data
|
||||
if data is None:
|
||||
raise WebSocketProtocolException("No continuation data available")
|
||||
self.cont_data = None
|
||||
frame.data = data[1]
|
||||
if (
|
||||
not self.fire_cont_frame
|
||||
and data is not None
|
||||
and data[0] == ABNF.OPCODE_TEXT
|
||||
and not self.skip_utf8_validation
|
||||
and not validate_utf8(frame.data)
|
||||
):
|
||||
raise WebSocketPayloadException(f"cannot decode: {repr(frame.data)}")
|
||||
if data is None:
|
||||
raise WebSocketProtocolException("No continuation data available")
|
||||
return data[0], frame
|
||||
@@ -0,0 +1,620 @@
|
||||
import inspect
|
||||
import socket
|
||||
import threading
|
||||
import time
|
||||
from typing import Any, Callable, Optional, Union
|
||||
|
||||
from . import _logging
|
||||
from ._abnf import ABNF
|
||||
from ._core import WebSocket, getdefaulttimeout
|
||||
from ._exceptions import (
|
||||
WebSocketConnectionClosedException,
|
||||
WebSocketException,
|
||||
WebSocketTimeoutException,
|
||||
)
|
||||
from ._ssl_compat import SSLEOFError
|
||||
from ._url import parse_url
|
||||
from ._dispatcher import Dispatcher, DispatcherBase, SSLDispatcher, WrappedDispatcher
|
||||
|
||||
"""
|
||||
_app.py
|
||||
websocket - WebSocket client library for Python
|
||||
|
||||
Copyright 2025 engn33r
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
__all__ = ["WebSocketApp"]
|
||||
|
||||
RECONNECT = 0
|
||||
|
||||
|
||||
def set_reconnect(reconnectInterval: int) -> None:
|
||||
global RECONNECT
|
||||
RECONNECT = reconnectInterval
|
||||
|
||||
|
||||
class WebSocketApp:
|
||||
"""
|
||||
Higher level of APIs are provided. The interface is like JavaScript WebSocket object.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
url: str,
|
||||
header: Optional[
|
||||
Union[
|
||||
list[str],
|
||||
dict[str, str],
|
||||
Callable[[], Union[list[str], dict[str, str]]],
|
||||
]
|
||||
] = None,
|
||||
on_open: Optional[Callable[["WebSocketApp"], None]] = None,
|
||||
on_reconnect: Optional[Callable[["WebSocketApp"], None]] = None,
|
||||
on_message: Optional[Callable[["WebSocketApp", Any], None]] = None,
|
||||
on_error: Optional[Callable[["WebSocketApp", Any], None]] = None,
|
||||
on_close: Optional[Callable[["WebSocketApp", Any, Any], None]] = None,
|
||||
on_ping: Optional[Callable] = None,
|
||||
on_pong: Optional[Callable] = None,
|
||||
on_cont_message: Optional[Callable] = None,
|
||||
keep_running: bool = True,
|
||||
get_mask_key: Optional[Callable] = None,
|
||||
cookie: Optional[str] = None,
|
||||
subprotocols: Optional[list[str]] = None,
|
||||
on_data: Optional[Callable] = None,
|
||||
socket: Optional[socket.socket] = None,
|
||||
) -> None:
|
||||
"""
|
||||
WebSocketApp initialization
|
||||
|
||||
Parameters
|
||||
----------
|
||||
url: str
|
||||
Websocket url.
|
||||
header: list or dict or Callable
|
||||
Custom header for websocket handshake.
|
||||
If the parameter is a callable object, it is called just before the connection attempt.
|
||||
The returned dict or list is used as custom header value.
|
||||
This could be useful in order to properly setup timestamp dependent headers.
|
||||
on_open: function
|
||||
Callback object which is called at opening websocket.
|
||||
on_open has one argument.
|
||||
The 1st argument is this class object.
|
||||
on_reconnect: function
|
||||
Callback object which is called at reconnecting websocket.
|
||||
on_reconnect has one argument.
|
||||
The 1st argument is this class object.
|
||||
on_message: function
|
||||
Callback object which is called when received data.
|
||||
on_message has 2 arguments.
|
||||
The 1st argument is this class object.
|
||||
The 2nd argument is utf-8 data received from the server.
|
||||
on_error: function
|
||||
Callback object which is called when we get error.
|
||||
on_error has 2 arguments.
|
||||
The 1st argument is this class object.
|
||||
The 2nd argument is exception object.
|
||||
on_close: function
|
||||
Callback object which is called when connection is closed.
|
||||
on_close has 3 arguments.
|
||||
The 1st argument is this class object.
|
||||
The 2nd argument is close_status_code.
|
||||
The 3rd argument is close_msg.
|
||||
on_cont_message: function
|
||||
Callback object which is called when a continuation
|
||||
frame is received.
|
||||
on_cont_message has 3 arguments.
|
||||
The 1st argument is this class object.
|
||||
The 2nd argument is utf-8 string which we get from the server.
|
||||
The 3rd argument is continue flag. if 0, the data continue
|
||||
to next frame data
|
||||
on_data: function
|
||||
Callback object which is called when a message received.
|
||||
This is called before on_message or on_cont_message,
|
||||
and then on_message or on_cont_message is called.
|
||||
on_data has 4 argument.
|
||||
The 1st argument is this class object.
|
||||
The 2nd argument is utf-8 string which we get from the server.
|
||||
The 3rd argument is data type. ABNF.OPCODE_TEXT or ABNF.OPCODE_BINARY will be came.
|
||||
The 4th argument is continue flag. If 0, the data continue
|
||||
keep_running: bool
|
||||
This parameter is obsolete and ignored.
|
||||
get_mask_key: function
|
||||
A callable function to get new mask keys, see the
|
||||
WebSocket.set_mask_key's docstring for more information.
|
||||
cookie: str
|
||||
Cookie value.
|
||||
subprotocols: list
|
||||
List of available sub protocols. Default is None.
|
||||
socket: socket
|
||||
Pre-initialized stream socket.
|
||||
"""
|
||||
self.url = url
|
||||
self.header = header if header is not None else []
|
||||
self.cookie = cookie
|
||||
|
||||
self.on_open = on_open
|
||||
self.on_reconnect = on_reconnect
|
||||
self.on_message = on_message
|
||||
self.on_data = on_data
|
||||
self.on_error = on_error
|
||||
self.on_close = on_close
|
||||
self.on_ping = on_ping
|
||||
self.on_pong = on_pong
|
||||
self.on_cont_message = on_cont_message
|
||||
self.keep_running = False
|
||||
self.get_mask_key = get_mask_key
|
||||
self.sock: Optional[WebSocket] = None
|
||||
self.last_ping_tm = float(0)
|
||||
self.last_pong_tm = float(0)
|
||||
self.ping_thread: Optional[threading.Thread] = None
|
||||
self.stop_ping: Optional[threading.Event] = None
|
||||
self.ping_interval = float(0)
|
||||
self.ping_timeout: Optional[Union[float, int]] = None
|
||||
self.ping_payload = ""
|
||||
self.subprotocols = subprotocols
|
||||
self.prepared_socket = socket
|
||||
self.has_errored = False
|
||||
self.has_done_teardown = False
|
||||
self.has_done_teardown_lock = threading.Lock()
|
||||
|
||||
def send(self, data: Union[bytes, str], opcode: int = ABNF.OPCODE_TEXT) -> None:
|
||||
"""
|
||||
send message
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data: str
|
||||
Message to send. If you set opcode to OPCODE_TEXT,
|
||||
data must be utf-8 string or unicode.
|
||||
opcode: int
|
||||
Operation code of data. Default is OPCODE_TEXT.
|
||||
"""
|
||||
|
||||
if not self.sock or self.sock.send(data, opcode) == 0:
|
||||
raise WebSocketConnectionClosedException("Connection is already closed.")
|
||||
|
||||
def send_text(self, text_data: str) -> None:
|
||||
"""
|
||||
Sends UTF-8 encoded text.
|
||||
"""
|
||||
if not self.sock or self.sock.send(text_data, ABNF.OPCODE_TEXT) == 0:
|
||||
raise WebSocketConnectionClosedException("Connection is already closed.")
|
||||
|
||||
def send_bytes(self, data: Union[bytes, bytearray]) -> None:
|
||||
"""
|
||||
Sends a sequence of bytes.
|
||||
"""
|
||||
if not self.sock or self.sock.send(data, ABNF.OPCODE_BINARY) == 0:
|
||||
raise WebSocketConnectionClosedException("Connection is already closed.")
|
||||
|
||||
def close(self, **kwargs) -> None:
|
||||
"""
|
||||
Close websocket connection.
|
||||
"""
|
||||
self.keep_running = False
|
||||
if self.sock:
|
||||
self.sock.close(**kwargs)
|
||||
self.sock = None
|
||||
|
||||
def _start_ping_thread(self) -> None:
|
||||
self.last_ping_tm = self.last_pong_tm = float(0)
|
||||
self.stop_ping = threading.Event()
|
||||
self.ping_thread = threading.Thread(target=self._send_ping)
|
||||
self.ping_thread.daemon = True
|
||||
self.ping_thread.start()
|
||||
|
||||
def _stop_ping_thread(self) -> None:
|
||||
if self.stop_ping:
|
||||
self.stop_ping.set()
|
||||
if self.ping_thread and self.ping_thread.is_alive():
|
||||
self.ping_thread.join(3)
|
||||
# Handle thread leak - if thread doesn't terminate within timeout,
|
||||
# force cleanup and log warning instead of abandoning the thread
|
||||
if self.ping_thread.is_alive():
|
||||
_logging.warning(
|
||||
"Ping thread failed to terminate within 3 seconds, "
|
||||
"forcing cleanup. Thread may be blocked."
|
||||
)
|
||||
# Force cleanup by clearing references even if thread is still alive
|
||||
# The daemon thread will eventually be cleaned up by Python's GC
|
||||
# but we prevent resource leaks by not holding references
|
||||
|
||||
# Always clean up references regardless of thread state
|
||||
self.ping_thread = None
|
||||
self.stop_ping = None
|
||||
self.last_ping_tm = self.last_pong_tm = float(0)
|
||||
|
||||
def _send_ping(self) -> None:
|
||||
if self.stop_ping is None:
|
||||
return
|
||||
if self.stop_ping.wait(self.ping_interval) or self.keep_running is False:
|
||||
return
|
||||
while not self.stop_ping.wait(self.ping_interval) and self.keep_running is True:
|
||||
if self.sock:
|
||||
self.last_ping_tm = time.time()
|
||||
try:
|
||||
_logging.debug("Sending ping")
|
||||
self.sock.ping(self.ping_payload)
|
||||
except Exception as e:
|
||||
_logging.debug(f"Failed to send ping: {e}")
|
||||
|
||||
def ready(self):
|
||||
return self.sock and self.sock.connected
|
||||
|
||||
def run_forever(
|
||||
self,
|
||||
sockopt: tuple = None,
|
||||
sslopt: dict = None,
|
||||
ping_interval: Union[float, int] = 0,
|
||||
ping_timeout: Optional[Union[float, int]] = None,
|
||||
ping_payload: str = "",
|
||||
http_proxy_host: str = None,
|
||||
http_proxy_port: Union[int, str] = None,
|
||||
http_no_proxy: list = None,
|
||||
http_proxy_auth: tuple = None,
|
||||
http_proxy_timeout: Optional[float] = None,
|
||||
skip_utf8_validation: bool = False,
|
||||
host: str = None,
|
||||
origin: str = None,
|
||||
dispatcher=None,
|
||||
suppress_origin: bool = False,
|
||||
proxy_type: str = None,
|
||||
reconnect: int = None,
|
||||
) -> bool:
|
||||
"""
|
||||
Run event loop for WebSocket framework.
|
||||
|
||||
This loop is an infinite loop and is alive while websocket is available.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
sockopt: tuple
|
||||
Values for socket.setsockopt.
|
||||
sockopt must be tuple
|
||||
and each element is argument of sock.setsockopt.
|
||||
sslopt: dict
|
||||
Optional dict object for ssl socket option.
|
||||
ping_interval: int or float
|
||||
Automatically send "ping" command
|
||||
every specified period (in seconds).
|
||||
If set to 0, no ping is sent periodically.
|
||||
ping_timeout: int or float
|
||||
Timeout (in seconds) if the pong message is not received.
|
||||
ping_payload: str
|
||||
Payload message to send with each ping.
|
||||
http_proxy_host: str
|
||||
HTTP proxy host name.
|
||||
http_proxy_port: int or str
|
||||
HTTP proxy port. If not set, set to 80.
|
||||
http_no_proxy: list
|
||||
Whitelisted host names that don't use the proxy.
|
||||
http_proxy_timeout: int or float
|
||||
HTTP proxy timeout, default is 60 sec as per python-socks.
|
||||
http_proxy_auth: tuple
|
||||
HTTP proxy auth information. tuple of username and password. Default is None.
|
||||
skip_utf8_validation: bool
|
||||
skip utf8 validation.
|
||||
host: str
|
||||
update host header.
|
||||
origin: str
|
||||
update origin header.
|
||||
dispatcher: Dispatcher object
|
||||
customize reading data from socket.
|
||||
suppress_origin: bool
|
||||
suppress outputting origin header.
|
||||
proxy_type: str
|
||||
type of proxy from: http, socks4, socks4a, socks5, socks5h
|
||||
reconnect: int
|
||||
delay interval when reconnecting
|
||||
|
||||
Returns
|
||||
-------
|
||||
teardown: bool
|
||||
False if the `WebSocketApp` is closed or caught KeyboardInterrupt,
|
||||
True if any other exception was raised during a loop.
|
||||
"""
|
||||
|
||||
if reconnect is None:
|
||||
reconnect = RECONNECT
|
||||
|
||||
if ping_timeout is not None and ping_timeout <= 0:
|
||||
raise WebSocketException("Ensure ping_timeout > 0")
|
||||
if ping_interval is not None and ping_interval < 0:
|
||||
raise WebSocketException("Ensure ping_interval >= 0")
|
||||
if ping_timeout and ping_interval and ping_interval <= ping_timeout:
|
||||
raise WebSocketException("Ensure ping_interval > ping_timeout")
|
||||
if not sockopt:
|
||||
sockopt = ()
|
||||
if not sslopt:
|
||||
sslopt = {}
|
||||
if self.sock:
|
||||
raise WebSocketException("socket is already opened")
|
||||
|
||||
self.ping_interval = ping_interval
|
||||
self.ping_timeout = ping_timeout
|
||||
self.ping_payload = ping_payload
|
||||
self.has_done_teardown = False
|
||||
self.keep_running = True
|
||||
|
||||
def teardown(close_frame: ABNF = None):
|
||||
"""
|
||||
Tears down the connection.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
close_frame: ABNF frame
|
||||
If close_frame is set, the on_close handler is invoked
|
||||
with the statusCode and reason from the provided frame.
|
||||
"""
|
||||
|
||||
# teardown() is called in many code paths to ensure resources are cleaned up and on_close is fired.
|
||||
# To ensure the work is only done once, we use this bool and lock.
|
||||
with self.has_done_teardown_lock:
|
||||
if self.has_done_teardown:
|
||||
return
|
||||
self.has_done_teardown = True
|
||||
|
||||
self._stop_ping_thread()
|
||||
self.keep_running = False
|
||||
|
||||
if self.sock:
|
||||
# in cases like handleDisconnect, the "on_error" callback is called first. If the WebSocketApp
|
||||
# is being used in a multithreaded application, we nee to make sure that "self.sock" is cleared
|
||||
# before calling close, otherwise logic built around the sock being set can cause issues -
|
||||
# specifically calling "run_forever" again, since is checks if "self.sock" is set.
|
||||
current_sock = self.sock
|
||||
self.sock = None
|
||||
current_sock.close()
|
||||
|
||||
close_status_code, close_reason = self._get_close_args(
|
||||
close_frame if close_frame else None
|
||||
)
|
||||
# Finally call the callback AFTER all teardown is complete
|
||||
self._callback(self.on_close, close_status_code, close_reason)
|
||||
|
||||
def initialize_socket(reconnecting: bool = False) -> None:
|
||||
if reconnecting and self.sock:
|
||||
self.sock.shutdown()
|
||||
|
||||
self.sock = WebSocket(
|
||||
self.get_mask_key,
|
||||
sockopt=sockopt,
|
||||
sslopt=sslopt,
|
||||
fire_cont_frame=self.on_cont_message is not None,
|
||||
skip_utf8_validation=skip_utf8_validation,
|
||||
enable_multithread=True,
|
||||
dispatcher=dispatcher,
|
||||
)
|
||||
|
||||
self.sock.settimeout(getdefaulttimeout())
|
||||
try:
|
||||
header = self.header() if callable(self.header) else self.header
|
||||
|
||||
self.sock.connect(
|
||||
self.url,
|
||||
header=header,
|
||||
cookie=self.cookie,
|
||||
http_proxy_host=http_proxy_host,
|
||||
http_proxy_port=http_proxy_port,
|
||||
http_no_proxy=http_no_proxy,
|
||||
http_proxy_auth=http_proxy_auth,
|
||||
http_proxy_timeout=http_proxy_timeout,
|
||||
subprotocols=self.subprotocols,
|
||||
host=host,
|
||||
origin=origin,
|
||||
suppress_origin=suppress_origin,
|
||||
proxy_type=proxy_type,
|
||||
socket=self.prepared_socket,
|
||||
)
|
||||
|
||||
_logging.info("Websocket connected")
|
||||
|
||||
if self.ping_interval:
|
||||
self._start_ping_thread()
|
||||
|
||||
if reconnecting and self.on_reconnect:
|
||||
self._callback(self.on_reconnect)
|
||||
else:
|
||||
self._callback(self.on_open)
|
||||
|
||||
dispatcher.read(self.sock.sock, read, check)
|
||||
except (
|
||||
WebSocketConnectionClosedException,
|
||||
ConnectionRefusedError,
|
||||
KeyboardInterrupt,
|
||||
SystemExit,
|
||||
Exception,
|
||||
) as e:
|
||||
handleDisconnect(e, reconnecting)
|
||||
|
||||
def read() -> bool:
|
||||
if not self.keep_running:
|
||||
teardown()
|
||||
return False
|
||||
|
||||
if self.sock is None:
|
||||
return False
|
||||
|
||||
try:
|
||||
op_code, frame = self.sock.recv_data_frame(True)
|
||||
except (
|
||||
WebSocketConnectionClosedException,
|
||||
KeyboardInterrupt,
|
||||
SSLEOFError,
|
||||
) as e:
|
||||
if custom_dispatcher:
|
||||
return closed(e)
|
||||
else:
|
||||
raise e
|
||||
|
||||
if op_code == ABNF.OPCODE_CLOSE:
|
||||
return closed(frame)
|
||||
elif op_code == ABNF.OPCODE_PING:
|
||||
self._callback(self.on_ping, frame.data)
|
||||
elif op_code == ABNF.OPCODE_PONG:
|
||||
self.last_pong_tm = time.time()
|
||||
self._callback(self.on_pong, frame.data)
|
||||
elif op_code == ABNF.OPCODE_CONT and self.on_cont_message:
|
||||
self._callback(self.on_data, frame.data, frame.opcode, frame.fin)
|
||||
self._callback(self.on_cont_message, frame.data, frame.fin)
|
||||
else:
|
||||
data = frame.data
|
||||
if op_code == ABNF.OPCODE_TEXT and not skip_utf8_validation:
|
||||
data = data.decode("utf-8")
|
||||
self._callback(self.on_data, data, frame.opcode, True)
|
||||
self._callback(self.on_message, data)
|
||||
|
||||
return True
|
||||
|
||||
def check() -> bool:
|
||||
if self.ping_timeout:
|
||||
has_timeout_expired = (
|
||||
time.time() - self.last_ping_tm > self.ping_timeout
|
||||
)
|
||||
has_pong_not_arrived_after_last_ping = (
|
||||
self.last_pong_tm - self.last_ping_tm < 0
|
||||
)
|
||||
has_pong_arrived_too_late = (
|
||||
self.last_pong_tm - self.last_ping_tm > self.ping_timeout
|
||||
)
|
||||
|
||||
if (
|
||||
self.last_ping_tm
|
||||
and has_timeout_expired
|
||||
and (
|
||||
has_pong_not_arrived_after_last_ping
|
||||
or has_pong_arrived_too_late
|
||||
)
|
||||
):
|
||||
raise WebSocketTimeoutException("ping/pong timed out")
|
||||
return True
|
||||
|
||||
def closed(
|
||||
e: Union[
|
||||
WebSocketConnectionClosedException,
|
||||
ConnectionRefusedError,
|
||||
KeyboardInterrupt,
|
||||
SystemExit,
|
||||
Exception,
|
||||
str,
|
||||
] = "closed unexpectedly",
|
||||
) -> bool:
|
||||
if type(e) is str:
|
||||
e = WebSocketConnectionClosedException(e)
|
||||
return handleDisconnect(e, bool(reconnect)) # type: ignore[arg-type]
|
||||
|
||||
def handleDisconnect(
|
||||
e: Union[
|
||||
WebSocketConnectionClosedException,
|
||||
ConnectionRefusedError,
|
||||
KeyboardInterrupt,
|
||||
SystemExit,
|
||||
Exception,
|
||||
],
|
||||
reconnecting: bool = False,
|
||||
) -> bool:
|
||||
self.has_errored = True
|
||||
self._stop_ping_thread()
|
||||
if not reconnecting:
|
||||
self._callback(self.on_error, e)
|
||||
|
||||
if isinstance(e, (KeyboardInterrupt, SystemExit)):
|
||||
teardown()
|
||||
# Propagate further
|
||||
raise
|
||||
|
||||
if reconnect:
|
||||
_logging.info(f"{e} - reconnect")
|
||||
if custom_dispatcher:
|
||||
_logging.debug(
|
||||
f"Calling custom dispatcher reconnect [{len(inspect.stack())} frames in stack]"
|
||||
)
|
||||
dispatcher.reconnect(reconnect, initialize_socket)
|
||||
else:
|
||||
_logging.error(f"{e} - goodbye")
|
||||
teardown()
|
||||
return self.has_errored
|
||||
|
||||
custom_dispatcher = bool(dispatcher)
|
||||
dispatcher = self.create_dispatcher(
|
||||
ping_timeout, dispatcher, parse_url(self.url)[3], closed
|
||||
)
|
||||
|
||||
try:
|
||||
initialize_socket()
|
||||
if not custom_dispatcher and reconnect:
|
||||
while self.keep_running:
|
||||
_logging.debug(
|
||||
f"Calling dispatcher reconnect [{len(inspect.stack())} frames in stack]"
|
||||
)
|
||||
dispatcher.reconnect(reconnect, initialize_socket)
|
||||
except (KeyboardInterrupt, Exception) as e:
|
||||
_logging.info(f"tearing down on exception {e}")
|
||||
teardown()
|
||||
finally:
|
||||
if not custom_dispatcher:
|
||||
# Ensure teardown was called before returning from run_forever
|
||||
teardown()
|
||||
|
||||
return self.has_errored
|
||||
|
||||
def create_dispatcher(
|
||||
self,
|
||||
ping_timeout: Optional[Union[float, int]],
|
||||
dispatcher: Optional[DispatcherBase] = None,
|
||||
is_ssl: bool = False,
|
||||
handleDisconnect: Callable = None,
|
||||
) -> Union[Dispatcher, SSLDispatcher, WrappedDispatcher]:
|
||||
if dispatcher: # If custom dispatcher is set, use WrappedDispatcher
|
||||
return WrappedDispatcher(self, ping_timeout, dispatcher, handleDisconnect)
|
||||
timeout = ping_timeout or 10
|
||||
if is_ssl:
|
||||
return SSLDispatcher(self, timeout)
|
||||
return Dispatcher(self, timeout)
|
||||
|
||||
def _get_close_args(self, close_frame: ABNF) -> list:
|
||||
"""
|
||||
_get_close_args extracts the close code and reason from the close body
|
||||
if it exists (RFC6455 says WebSocket Connection Close Code is optional)
|
||||
"""
|
||||
# Need to catch the case where close_frame is None
|
||||
# Otherwise the following if statement causes an error
|
||||
if not self.on_close or not close_frame:
|
||||
return [None, None]
|
||||
|
||||
# Extract close frame status code
|
||||
if close_frame.data and len(close_frame.data) >= 2:
|
||||
close_status_code = 256 * int(close_frame.data[0]) + int(
|
||||
close_frame.data[1]
|
||||
)
|
||||
reason = close_frame.data[2:]
|
||||
if isinstance(reason, bytes):
|
||||
reason = reason.decode("utf-8")
|
||||
return [close_status_code, reason]
|
||||
else:
|
||||
# Most likely reached this because len(close_frame_data.data) < 2
|
||||
return [None, None]
|
||||
|
||||
def _callback(self, callback, *args) -> None:
|
||||
if callback:
|
||||
try:
|
||||
callback(self, *args)
|
||||
|
||||
except Exception as e:
|
||||
_logging.error(f"error from callback {callback}: {e}")
|
||||
# Bug fix: Prevent infinite recursion by not calling on_error
|
||||
# when the failing callback IS on_error itself
|
||||
if self.on_error and callback is not self.on_error:
|
||||
self.on_error(self, e)
|
||||
@@ -0,0 +1,73 @@
|
||||
import http.cookies
|
||||
from typing import Optional
|
||||
|
||||
"""
|
||||
_cookiejar.py
|
||||
websocket - WebSocket client library for Python
|
||||
|
||||
Copyright 2025 engn33r
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
|
||||
class SimpleCookieJar:
|
||||
def __init__(self) -> None:
|
||||
self.jar: dict = {}
|
||||
|
||||
def add(self, set_cookie: Optional[str]) -> None:
|
||||
if set_cookie:
|
||||
simple_cookie = http.cookies.SimpleCookie(set_cookie)
|
||||
|
||||
for v in simple_cookie.values():
|
||||
if domain := v.get("domain"):
|
||||
if not domain.startswith("."):
|
||||
domain = f".{domain}"
|
||||
cookie = self.jar.get(domain)
|
||||
if cookie is None:
|
||||
cookie = http.cookies.SimpleCookie()
|
||||
cookie.update(simple_cookie)
|
||||
self.jar[domain.lower()] = cookie
|
||||
|
||||
def set(self, set_cookie: str) -> None:
|
||||
if set_cookie:
|
||||
simple_cookie = http.cookies.SimpleCookie(set_cookie)
|
||||
|
||||
for v in simple_cookie.values():
|
||||
if domain := v.get("domain"):
|
||||
if not domain.startswith("."):
|
||||
domain = f".{domain}"
|
||||
self.jar[domain.lower()] = simple_cookie
|
||||
|
||||
def get(self, host: str) -> str:
|
||||
if not host:
|
||||
return ""
|
||||
|
||||
cookies = []
|
||||
for domain, _ in self.jar.items():
|
||||
host = host.lower()
|
||||
if host.endswith(domain) or host == domain[1:]:
|
||||
cookies.append(self.jar.get(domain))
|
||||
|
||||
return "; ".join(
|
||||
filter(
|
||||
None,
|
||||
sorted(
|
||||
[
|
||||
f"{k}={v.value}"
|
||||
for cookie in filter(None, cookies)
|
||||
for k, v in cookie.items()
|
||||
]
|
||||
),
|
||||
)
|
||||
)
|
||||
@@ -0,0 +1,665 @@
|
||||
import socket
|
||||
import struct
|
||||
import threading
|
||||
import time
|
||||
from typing import Optional, Union
|
||||
|
||||
# websocket modules
|
||||
from ._abnf import ABNF, STATUS_NORMAL, continuous_frame, frame_buffer
|
||||
from ._exceptions import (
|
||||
WebSocketProtocolException,
|
||||
WebSocketConnectionClosedException,
|
||||
WebSocketTimeoutException,
|
||||
)
|
||||
from ._handshake import SUPPORTED_REDIRECT_STATUSES, handshake
|
||||
from ._http import connect, proxy_info
|
||||
from ._logging import debug, error, trace, isEnabledForError, isEnabledForTrace
|
||||
from ._socket import getdefaulttimeout, recv, send, sock_opt
|
||||
from ._ssl_compat import ssl
|
||||
from ._utils import NoLock
|
||||
from ._dispatcher import DispatcherBase, WrappedDispatcher
|
||||
|
||||
"""
|
||||
_core.py
|
||||
websocket - WebSocket client library for Python
|
||||
|
||||
Copyright 2025 engn33r
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
__all__ = ["WebSocket", "create_connection"]
|
||||
|
||||
|
||||
class WebSocket:
|
||||
"""
|
||||
Low level WebSocket interface.
|
||||
|
||||
This class is based on the WebSocket protocol `draft-hixie-thewebsocketprotocol-76 <http://tools.ietf.org/html/draft-hixie-thewebsocketprotocol-76>`_
|
||||
|
||||
We can connect to the websocket server and send/receive data.
|
||||
The following example is an echo client.
|
||||
|
||||
>>> import websocket
|
||||
>>> ws = websocket.WebSocket()
|
||||
>>> ws.connect("ws://echo.websocket.events")
|
||||
>>> ws.recv()
|
||||
'echo.websocket.events sponsored by Lob.com'
|
||||
>>> ws.send("Hello, Server")
|
||||
19
|
||||
>>> ws.recv()
|
||||
'Hello, Server'
|
||||
>>> ws.close()
|
||||
|
||||
Parameters
|
||||
----------
|
||||
get_mask_key: func
|
||||
A callable function to get new mask keys, see the
|
||||
WebSocket.set_mask_key's docstring for more information.
|
||||
sockopt: tuple
|
||||
Values for socket.setsockopt.
|
||||
sockopt must be tuple and each element is argument of sock.setsockopt.
|
||||
sslopt: dict
|
||||
Optional dict object for ssl socket options. See FAQ for details.
|
||||
fire_cont_frame: bool
|
||||
Fire recv event for each cont frame. Default is False.
|
||||
enable_multithread: bool
|
||||
If set to True, lock send method.
|
||||
skip_utf8_validation: bool
|
||||
Skip utf8 validation.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
get_mask_key=None,
|
||||
sockopt=None,
|
||||
sslopt=None,
|
||||
fire_cont_frame: bool = False,
|
||||
enable_multithread: bool = True,
|
||||
skip_utf8_validation: bool = False,
|
||||
dispatcher: Union[DispatcherBase, WrappedDispatcher] = None,
|
||||
**_,
|
||||
):
|
||||
"""
|
||||
Initialize WebSocket object.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
sslopt: dict
|
||||
Optional dict object for ssl socket options. See FAQ for details.
|
||||
"""
|
||||
self.sock_opt = sock_opt(sockopt, sslopt)
|
||||
self.handshake_response = None
|
||||
self.sock: Optional[socket.socket] = None
|
||||
|
||||
self.connected = False
|
||||
self.get_mask_key = get_mask_key
|
||||
# These buffer over the build-up of a single frame.
|
||||
self.frame_buffer = frame_buffer(self._recv, skip_utf8_validation)
|
||||
self.cont_frame = continuous_frame(fire_cont_frame, skip_utf8_validation)
|
||||
self.dispatcher = dispatcher
|
||||
|
||||
if enable_multithread:
|
||||
self.lock = threading.Lock()
|
||||
self.readlock = threading.Lock()
|
||||
else:
|
||||
self.lock = NoLock() # type: ignore[assignment]
|
||||
self.readlock = NoLock() # type: ignore[assignment]
|
||||
|
||||
def __iter__(self):
|
||||
"""
|
||||
Allow iteration over websocket, implying sequential `recv` executions.
|
||||
"""
|
||||
while True:
|
||||
yield self.recv()
|
||||
|
||||
def __next__(self):
|
||||
return self.recv()
|
||||
|
||||
def next(self):
|
||||
return self.__next__()
|
||||
|
||||
def fileno(self):
|
||||
return self.sock.fileno()
|
||||
|
||||
def set_mask_key(self, func):
|
||||
"""
|
||||
Set function to create mask key. You can customize mask key generator.
|
||||
Mainly, this is for testing purpose.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
func: func
|
||||
callable object. the func takes 1 argument as integer.
|
||||
The argument means length of mask key.
|
||||
This func must return string(byte array),
|
||||
which length is argument specified.
|
||||
"""
|
||||
self.get_mask_key = func
|
||||
|
||||
def gettimeout(self) -> Optional[Union[float, int]]:
|
||||
"""
|
||||
Get the websocket timeout (in seconds) as an int or float
|
||||
|
||||
Returns
|
||||
----------
|
||||
timeout: int or float
|
||||
returns timeout value (in seconds). This value could be either float/integer.
|
||||
"""
|
||||
return self.sock_opt.timeout
|
||||
|
||||
def settimeout(self, timeout: Optional[Union[float, int]]):
|
||||
"""
|
||||
Set the timeout to the websocket.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
timeout: int or float
|
||||
timeout time (in seconds). This value could be either float/integer.
|
||||
"""
|
||||
self.sock_opt.timeout = timeout
|
||||
if self.sock:
|
||||
self.sock.settimeout(timeout)
|
||||
|
||||
timeout = property(gettimeout, settimeout)
|
||||
|
||||
def getsubprotocol(self):
|
||||
"""
|
||||
Get subprotocol
|
||||
"""
|
||||
if self.handshake_response:
|
||||
return self.handshake_response.subprotocol
|
||||
else:
|
||||
return None
|
||||
|
||||
subprotocol = property(getsubprotocol)
|
||||
|
||||
def getstatus(self):
|
||||
"""
|
||||
Get handshake status
|
||||
"""
|
||||
if self.handshake_response:
|
||||
return self.handshake_response.status
|
||||
else:
|
||||
return None
|
||||
|
||||
status = property(getstatus)
|
||||
|
||||
def getheaders(self):
|
||||
"""
|
||||
Get handshake response header
|
||||
"""
|
||||
if self.handshake_response:
|
||||
return self.handshake_response.headers
|
||||
else:
|
||||
return None
|
||||
|
||||
def is_ssl(self):
|
||||
try:
|
||||
return isinstance(self.sock, ssl.SSLSocket)
|
||||
except (AttributeError, NameError):
|
||||
return False
|
||||
|
||||
headers = property(getheaders)
|
||||
|
||||
def connect(self, url, **options):
|
||||
"""
|
||||
Connect to url. url is websocket url scheme.
|
||||
ie. ws://host:port/resource
|
||||
You can customize using 'options'.
|
||||
If you set "header" list object, you can set your own custom header.
|
||||
|
||||
>>> ws = WebSocket()
|
||||
>>> ws.connect("ws://echo.websocket.events",
|
||||
... header=["User-Agent: MyProgram",
|
||||
... "x-custom: header"])
|
||||
|
||||
Parameters
|
||||
----------
|
||||
header: list or dict
|
||||
Custom http header list or dict.
|
||||
cookie: str
|
||||
Cookie value.
|
||||
origin: str
|
||||
Custom origin url.
|
||||
connection: str
|
||||
Custom connection header value.
|
||||
Default value "Upgrade" set in _handshake.py
|
||||
suppress_origin: bool
|
||||
Suppress outputting origin header.
|
||||
host: str
|
||||
Custom host header string.
|
||||
timeout: int or float
|
||||
Socket timeout time. This value is an integer or float.
|
||||
If you set None for this value, it means "use default_timeout value"
|
||||
http_proxy_host: str
|
||||
HTTP proxy host name.
|
||||
http_proxy_port: str or int
|
||||
HTTP proxy port. Default is 80.
|
||||
http_no_proxy: list
|
||||
Whitelisted host names that don't use the proxy.
|
||||
http_proxy_auth: tuple
|
||||
HTTP proxy auth information. Tuple of username and password. Default is None.
|
||||
http_proxy_timeout: int or float
|
||||
HTTP proxy timeout, default is 60 sec as per python-socks.
|
||||
redirect_limit: int
|
||||
Number of redirects to follow.
|
||||
subprotocols: list
|
||||
List of available subprotocols. Default is None.
|
||||
socket: socket
|
||||
Pre-initialized stream socket.
|
||||
"""
|
||||
self.sock_opt.timeout = options.get("timeout", self.sock_opt.timeout)
|
||||
self.sock, addrs = connect(
|
||||
url, self.sock_opt, proxy_info(**options), options.pop("socket", None)
|
||||
)
|
||||
|
||||
try:
|
||||
self.handshake_response = handshake(self.sock, url, *addrs, **options)
|
||||
for _ in range(options.pop("redirect_limit", 3)):
|
||||
if self.handshake_response.status in SUPPORTED_REDIRECT_STATUSES:
|
||||
url = self.handshake_response.headers["location"]
|
||||
self.sock.close()
|
||||
self.sock, addrs = connect(
|
||||
url,
|
||||
self.sock_opt,
|
||||
proxy_info(**options),
|
||||
options.pop("socket", None),
|
||||
)
|
||||
self.handshake_response = handshake(
|
||||
self.sock, url, *addrs, **options
|
||||
)
|
||||
self.connected = True
|
||||
except:
|
||||
if self.sock:
|
||||
self.sock.close()
|
||||
self.sock = None
|
||||
raise
|
||||
|
||||
def send(self, payload: Union[bytes, str], opcode: int = ABNF.OPCODE_TEXT) -> int:
|
||||
"""
|
||||
Send the data as string.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
payload: str
|
||||
Payload must be utf-8 string or unicode,
|
||||
If the opcode is OPCODE_TEXT.
|
||||
Otherwise, it must be string(byte array).
|
||||
opcode: int
|
||||
Operation code (opcode) to send.
|
||||
"""
|
||||
|
||||
frame = ABNF.create_frame(payload, opcode)
|
||||
return self.send_frame(frame)
|
||||
|
||||
def send_text(self, text_data: str) -> int:
|
||||
"""
|
||||
Sends UTF-8 encoded text.
|
||||
"""
|
||||
return self.send(text_data, ABNF.OPCODE_TEXT)
|
||||
|
||||
def send_bytes(self, data: Union[bytes, bytearray]) -> int:
|
||||
"""
|
||||
Sends a sequence of bytes.
|
||||
"""
|
||||
return self.send(data, ABNF.OPCODE_BINARY)
|
||||
|
||||
def send_frame(self, frame) -> int:
|
||||
"""
|
||||
Send the data frame.
|
||||
|
||||
>>> ws = create_connection("ws://echo.websocket.events")
|
||||
>>> frame = ABNF.create_frame("Hello", ABNF.OPCODE_TEXT)
|
||||
>>> ws.send_frame(frame)
|
||||
>>> cont_frame = ABNF.create_frame("My name is ", ABNF.OPCODE_CONT, 0)
|
||||
>>> ws.send_frame(frame)
|
||||
>>> cont_frame = ABNF.create_frame("Foo Bar", ABNF.OPCODE_CONT, 1)
|
||||
>>> ws.send_frame(frame)
|
||||
|
||||
Parameters
|
||||
----------
|
||||
frame: ABNF frame
|
||||
frame data created by ABNF.create_frame
|
||||
"""
|
||||
if self.get_mask_key:
|
||||
frame.get_mask_key = self.get_mask_key
|
||||
data = frame.format()
|
||||
length = len(data)
|
||||
if isEnabledForTrace():
|
||||
trace(f"++Sent raw: {repr(data)}")
|
||||
trace(f"++Sent decoded: {frame.__str__()}")
|
||||
with self.lock:
|
||||
while data:
|
||||
bytes_sent = self._send(data)
|
||||
data = data[bytes_sent:]
|
||||
|
||||
return length
|
||||
|
||||
def send_binary(self, payload: bytes) -> int:
|
||||
"""
|
||||
Send a binary message (OPCODE_BINARY).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
payload: bytes
|
||||
payload of message to send.
|
||||
"""
|
||||
return self.send(payload, ABNF.OPCODE_BINARY)
|
||||
|
||||
def ping(self, payload: Union[str, bytes] = ""):
|
||||
"""
|
||||
Send ping data.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
payload: str
|
||||
data payload to send server.
|
||||
"""
|
||||
if isinstance(payload, str):
|
||||
payload = payload.encode("utf-8")
|
||||
self.send(payload, ABNF.OPCODE_PING)
|
||||
|
||||
def pong(self, payload: Union[str, bytes] = ""):
|
||||
"""
|
||||
Send pong data.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
payload: str
|
||||
data payload to send server.
|
||||
"""
|
||||
if isinstance(payload, str):
|
||||
payload = payload.encode("utf-8")
|
||||
self.send(payload, ABNF.OPCODE_PONG)
|
||||
|
||||
def recv(self) -> Union[str, bytes]:
|
||||
"""
|
||||
Receive string data(byte array) from the server.
|
||||
|
||||
Returns
|
||||
----------
|
||||
data: string (byte array) value.
|
||||
"""
|
||||
with self.readlock:
|
||||
opcode, data = self.recv_data()
|
||||
if opcode == ABNF.OPCODE_TEXT:
|
||||
data_received: Union[bytes, str] = data
|
||||
if isinstance(data_received, bytes):
|
||||
return data_received.decode("utf-8")
|
||||
elif isinstance(data_received, str):
|
||||
return data_received
|
||||
elif opcode == ABNF.OPCODE_BINARY:
|
||||
data_binary: bytes = data
|
||||
return data_binary
|
||||
else:
|
||||
return ""
|
||||
|
||||
def recv_data(self, control_frame: bool = False) -> tuple:
|
||||
"""
|
||||
Receive data with operation code.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
control_frame: bool
|
||||
a boolean flag indicating whether to return control frame
|
||||
data, defaults to False
|
||||
|
||||
Returns
|
||||
-------
|
||||
opcode, frame.data: tuple
|
||||
tuple of operation code and string(byte array) value.
|
||||
"""
|
||||
opcode, frame = self.recv_data_frame(control_frame)
|
||||
return opcode, frame.data
|
||||
|
||||
def recv_data_frame(self, control_frame: bool = False) -> tuple:
|
||||
"""
|
||||
Receive data with operation code.
|
||||
|
||||
If a valid ping message is received, a pong response is sent.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
control_frame: bool
|
||||
a boolean flag indicating whether to return control frame
|
||||
data, defaults to False
|
||||
|
||||
Returns
|
||||
-------
|
||||
frame.opcode, frame: tuple
|
||||
tuple of operation code and string(byte array) value.
|
||||
"""
|
||||
while True:
|
||||
frame = self.recv_frame()
|
||||
if isEnabledForTrace():
|
||||
trace(f"++Rcv raw: {repr(frame.format())}")
|
||||
trace(f"++Rcv decoded: {frame.__str__()}")
|
||||
if not frame:
|
||||
# handle error:
|
||||
# 'NoneType' object has no attribute 'opcode'
|
||||
raise WebSocketProtocolException(f"Not a valid frame {frame}")
|
||||
elif frame.opcode in (
|
||||
ABNF.OPCODE_TEXT,
|
||||
ABNF.OPCODE_BINARY,
|
||||
ABNF.OPCODE_CONT,
|
||||
):
|
||||
self.cont_frame.validate(frame)
|
||||
self.cont_frame.add(frame)
|
||||
|
||||
if self.cont_frame.is_fire(frame):
|
||||
return self.cont_frame.extract(frame)
|
||||
|
||||
elif frame.opcode == ABNF.OPCODE_CLOSE:
|
||||
self.send_close()
|
||||
return frame.opcode, frame
|
||||
elif frame.opcode == ABNF.OPCODE_PING:
|
||||
if len(frame.data) < 126:
|
||||
self.pong(frame.data)
|
||||
else:
|
||||
raise WebSocketProtocolException("Ping message is too long")
|
||||
if control_frame:
|
||||
return frame.opcode, frame
|
||||
elif frame.opcode == ABNF.OPCODE_PONG:
|
||||
if control_frame:
|
||||
return frame.opcode, frame
|
||||
|
||||
def recv_frame(self):
|
||||
"""
|
||||
Receive data as frame from server.
|
||||
|
||||
Returns
|
||||
-------
|
||||
self.frame_buffer.recv_frame(): ABNF frame object
|
||||
"""
|
||||
return self.frame_buffer.recv_frame()
|
||||
|
||||
def send_close(self, status: int = STATUS_NORMAL, reason: bytes = b""):
|
||||
"""
|
||||
Send close data to the server.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
status: int
|
||||
Status code to send. See STATUS_XXX.
|
||||
reason: str or bytes
|
||||
The reason to close. This must be string or UTF-8 bytes.
|
||||
"""
|
||||
if status < 0 or status >= ABNF.LENGTH_16:
|
||||
raise ValueError("code is invalid range")
|
||||
self.connected = False
|
||||
self.send(struct.pack("!H", status) + reason, ABNF.OPCODE_CLOSE)
|
||||
|
||||
def close(self, status: int = STATUS_NORMAL, reason: bytes = b"", timeout: int = 3):
|
||||
"""
|
||||
Close Websocket object
|
||||
|
||||
Parameters
|
||||
----------
|
||||
status: int
|
||||
Status code to send. See VALID_CLOSE_STATUS in ABNF.
|
||||
reason: bytes
|
||||
The reason to close in UTF-8.
|
||||
timeout: int or float
|
||||
Timeout until receive a close frame.
|
||||
If None, it will wait forever until receive a close frame.
|
||||
"""
|
||||
if not self.connected:
|
||||
return
|
||||
if status < 0 or status >= ABNF.LENGTH_16:
|
||||
raise ValueError("code is invalid range")
|
||||
|
||||
try:
|
||||
self.connected = False
|
||||
self.send(struct.pack("!H", status) + reason, ABNF.OPCODE_CLOSE)
|
||||
if self.sock is None:
|
||||
return
|
||||
sock_timeout = self.sock.gettimeout()
|
||||
self.sock.settimeout(timeout)
|
||||
start_time = time.time()
|
||||
while timeout is None or time.time() - start_time < timeout:
|
||||
try:
|
||||
frame = self.recv_frame()
|
||||
if frame.opcode != ABNF.OPCODE_CLOSE:
|
||||
continue
|
||||
if isEnabledForError():
|
||||
recv_status = struct.unpack("!H", frame.data[0:2])[0]
|
||||
if recv_status >= 3000 and recv_status <= 4999:
|
||||
debug(f"close status: {repr(recv_status)}")
|
||||
elif recv_status != STATUS_NORMAL:
|
||||
error(f"close status: {repr(recv_status)}")
|
||||
break
|
||||
except (
|
||||
WebSocketConnectionClosedException,
|
||||
WebSocketTimeoutException,
|
||||
struct.error,
|
||||
):
|
||||
break
|
||||
if self.sock is not None:
|
||||
self.sock.settimeout(sock_timeout)
|
||||
self.sock.shutdown(socket.SHUT_RDWR)
|
||||
except:
|
||||
pass
|
||||
|
||||
self.shutdown()
|
||||
|
||||
def abort(self):
|
||||
"""
|
||||
Low-level asynchronous abort, wakes up other threads that are waiting in recv_*
|
||||
"""
|
||||
if self.connected:
|
||||
self.sock.shutdown(socket.SHUT_RDWR)
|
||||
|
||||
def shutdown(self):
|
||||
"""
|
||||
close socket, immediately.
|
||||
"""
|
||||
if self.sock:
|
||||
self.sock.close()
|
||||
self.sock = None
|
||||
self.connected = False
|
||||
|
||||
def _send(self, data: Union[str, bytes]):
|
||||
if self.sock is None:
|
||||
raise WebSocketConnectionClosedException("socket is already closed.")
|
||||
if self.dispatcher:
|
||||
return self.dispatcher.send(self.sock, data)
|
||||
return send(self.sock, data)
|
||||
|
||||
def _recv(self, bufsize):
|
||||
try:
|
||||
return recv(self.sock, bufsize)
|
||||
except WebSocketConnectionClosedException:
|
||||
if self.sock:
|
||||
self.sock.close()
|
||||
self.sock = None
|
||||
self.connected = False
|
||||
raise
|
||||
|
||||
|
||||
def create_connection(url: str, timeout=None, class_=WebSocket, **options):
|
||||
"""
|
||||
Connect to url and return websocket object.
|
||||
|
||||
Connect to url and return the WebSocket object.
|
||||
Passing optional timeout parameter will set the timeout on the socket.
|
||||
If no timeout is supplied,
|
||||
the global default timeout setting returned by getdefaulttimeout() is used.
|
||||
You can customize using 'options'.
|
||||
If you set "header" list object, you can set your own custom header.
|
||||
|
||||
>>> conn = create_connection("ws://echo.websocket.events",
|
||||
... header=["User-Agent: MyProgram",
|
||||
... "x-custom: header"])
|
||||
|
||||
Parameters
|
||||
----------
|
||||
class_: class
|
||||
class to instantiate when creating the connection. It has to implement
|
||||
settimeout and connect. It's __init__ should be compatible with
|
||||
WebSocket.__init__, i.e. accept all of it's kwargs.
|
||||
header: list or dict
|
||||
custom http header list or dict.
|
||||
cookie: str
|
||||
Cookie value.
|
||||
origin: str
|
||||
custom origin url.
|
||||
suppress_origin: bool
|
||||
suppress outputting origin header.
|
||||
host: str
|
||||
custom host header string.
|
||||
timeout: int or float
|
||||
socket timeout time. This value could be either float/integer.
|
||||
If set to None, it uses the default_timeout value.
|
||||
http_proxy_host: str
|
||||
HTTP proxy host name.
|
||||
http_proxy_port: str or int
|
||||
HTTP proxy port. If not set, set to 80.
|
||||
http_no_proxy: list
|
||||
Whitelisted host names that don't use the proxy.
|
||||
http_proxy_auth: tuple
|
||||
HTTP proxy auth information. tuple of username and password. Default is None.
|
||||
http_proxy_timeout: int or float
|
||||
HTTP proxy timeout, default is 60 sec as per python-socks.
|
||||
enable_multithread: bool
|
||||
Enable lock for multithread.
|
||||
redirect_limit: int
|
||||
Number of redirects to follow.
|
||||
sockopt: tuple
|
||||
Values for socket.setsockopt.
|
||||
sockopt must be a tuple and each element is an argument of sock.setsockopt.
|
||||
sslopt: dict
|
||||
Optional dict object for ssl socket options. See FAQ for details.
|
||||
subprotocols: list
|
||||
List of available subprotocols. Default is None.
|
||||
skip_utf8_validation: bool
|
||||
Skip utf8 validation.
|
||||
socket: socket
|
||||
Pre-initialized stream socket.
|
||||
"""
|
||||
sockopt = options.pop("sockopt", [])
|
||||
sslopt = options.pop("sslopt", {})
|
||||
fire_cont_frame = options.pop("fire_cont_frame", False)
|
||||
enable_multithread = options.pop("enable_multithread", True)
|
||||
skip_utf8_validation = options.pop("skip_utf8_validation", False)
|
||||
websock = class_(
|
||||
sockopt=sockopt,
|
||||
sslopt=sslopt,
|
||||
fire_cont_frame=fire_cont_frame,
|
||||
enable_multithread=enable_multithread,
|
||||
skip_utf8_validation=skip_utf8_validation,
|
||||
**options,
|
||||
)
|
||||
websock.settimeout(timeout if timeout is not None else getdefaulttimeout())
|
||||
websock.connect(url, **options)
|
||||
return websock
|
||||
@@ -0,0 +1,164 @@
|
||||
import time
|
||||
import socket
|
||||
import inspect
|
||||
import selectors
|
||||
from typing import TYPE_CHECKING, Callable, Optional, Union
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ._app import WebSocketApp
|
||||
from . import _logging
|
||||
from ._socket import send
|
||||
|
||||
"""
|
||||
_dispatcher.py
|
||||
websocket - WebSocket client library for Python
|
||||
|
||||
Copyright 2025 engn33r
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
class DispatcherBase:
|
||||
"""
|
||||
DispatcherBase
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, app: "WebSocketApp", ping_timeout: Optional[Union[float, int]]
|
||||
) -> None:
|
||||
self.app = app
|
||||
self.ping_timeout = ping_timeout
|
||||
|
||||
def timeout(self, seconds: Optional[Union[float, int]], callback: Callable) -> None:
|
||||
if seconds is not None:
|
||||
time.sleep(seconds)
|
||||
callback()
|
||||
|
||||
def reconnect(self, seconds: int, reconnector: Callable) -> None:
|
||||
try:
|
||||
_logging.info(
|
||||
f"reconnect() - retrying in {seconds} seconds [{len(inspect.stack())} frames in stack]"
|
||||
)
|
||||
time.sleep(seconds)
|
||||
reconnector(reconnecting=True)
|
||||
except KeyboardInterrupt as e:
|
||||
_logging.info(f"User exited {e}")
|
||||
raise e
|
||||
|
||||
def send(self, sock: socket.socket, data: Union[str, bytes]) -> int:
|
||||
return send(sock, data)
|
||||
|
||||
|
||||
class Dispatcher(DispatcherBase):
|
||||
"""
|
||||
Dispatcher
|
||||
"""
|
||||
|
||||
def read(
|
||||
self,
|
||||
sock: socket.socket,
|
||||
read_callback: Callable,
|
||||
check_callback: Callable,
|
||||
) -> None:
|
||||
if self.app.sock is None or self.app.sock.sock is None:
|
||||
return
|
||||
sel = selectors.DefaultSelector()
|
||||
sel.register(self.app.sock.sock, selectors.EVENT_READ)
|
||||
try:
|
||||
while self.app.keep_running:
|
||||
if sel.select(self.ping_timeout):
|
||||
if not read_callback():
|
||||
break
|
||||
check_callback()
|
||||
finally:
|
||||
sel.close()
|
||||
|
||||
|
||||
class SSLDispatcher(DispatcherBase):
|
||||
"""
|
||||
SSLDispatcher
|
||||
"""
|
||||
|
||||
def read(
|
||||
self,
|
||||
sock: socket.socket,
|
||||
read_callback: Callable,
|
||||
check_callback: Callable,
|
||||
) -> None:
|
||||
if self.app.sock is None or self.app.sock.sock is None:
|
||||
return
|
||||
sock = self.app.sock.sock
|
||||
sel = selectors.DefaultSelector()
|
||||
sel.register(sock, selectors.EVENT_READ)
|
||||
try:
|
||||
while self.app.keep_running:
|
||||
if self.select(sock, sel):
|
||||
if not read_callback():
|
||||
break
|
||||
check_callback()
|
||||
finally:
|
||||
sel.close()
|
||||
|
||||
def select(self, sock, sel: selectors.DefaultSelector):
|
||||
if self.app.sock is None:
|
||||
return None
|
||||
sock = self.app.sock.sock
|
||||
if sock.pending():
|
||||
return [
|
||||
sock,
|
||||
]
|
||||
|
||||
r = sel.select(self.ping_timeout)
|
||||
|
||||
if len(r) > 0:
|
||||
return r[0][0]
|
||||
return None
|
||||
|
||||
|
||||
class WrappedDispatcher:
|
||||
"""
|
||||
WrappedDispatcher
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
app: "WebSocketApp",
|
||||
ping_timeout: Optional[Union[float, int]],
|
||||
dispatcher,
|
||||
handleDisconnect,
|
||||
) -> None:
|
||||
self.app = app
|
||||
self.ping_timeout = ping_timeout
|
||||
self.dispatcher = dispatcher
|
||||
self.handleDisconnect = handleDisconnect
|
||||
dispatcher.signal(2, dispatcher.abort) # keyboard interrupt
|
||||
|
||||
def read(
|
||||
self,
|
||||
sock: socket.socket,
|
||||
read_callback: Callable,
|
||||
check_callback: Callable,
|
||||
) -> None:
|
||||
self.dispatcher.read(sock, read_callback)
|
||||
if self.ping_timeout:
|
||||
self.timeout(self.ping_timeout, check_callback)
|
||||
|
||||
def send(self, sock: socket.socket, data: Union[str, bytes]) -> int:
|
||||
self.dispatcher.buffwrite(sock, data, send, self.handleDisconnect)
|
||||
return len(data)
|
||||
|
||||
def timeout(self, seconds: float, callback: Callable, *args) -> None:
|
||||
self.dispatcher.timeout(seconds, callback, *args)
|
||||
|
||||
def reconnect(self, seconds: int, reconnector: Callable) -> None:
|
||||
self.timeout(seconds, reconnector, True)
|
||||
@@ -0,0 +1,94 @@
|
||||
"""
|
||||
_exceptions.py
|
||||
websocket - WebSocket client library for Python
|
||||
|
||||
Copyright 2025 engn33r
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
|
||||
class WebSocketException(Exception):
|
||||
"""
|
||||
WebSocket exception class.
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class WebSocketProtocolException(WebSocketException):
|
||||
"""
|
||||
If the WebSocket protocol is invalid, this exception will be raised.
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class WebSocketPayloadException(WebSocketException):
|
||||
"""
|
||||
If the WebSocket payload is invalid, this exception will be raised.
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class WebSocketConnectionClosedException(WebSocketException):
|
||||
"""
|
||||
If remote host closed the connection or some network error happened,
|
||||
this exception will be raised.
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class WebSocketTimeoutException(WebSocketException):
|
||||
"""
|
||||
WebSocketTimeoutException will be raised at socket timeout during read/write data.
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class WebSocketProxyException(WebSocketException):
|
||||
"""
|
||||
WebSocketProxyException will be raised when proxy error occurred.
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class WebSocketBadStatusException(WebSocketException):
|
||||
"""
|
||||
WebSocketBadStatusException will be raised when we get bad handshake status code.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
message: str,
|
||||
status_code: int,
|
||||
status_message=None,
|
||||
resp_headers=None,
|
||||
resp_body=None,
|
||||
):
|
||||
super().__init__(message)
|
||||
self.status_code = status_code
|
||||
self.resp_headers = resp_headers
|
||||
self.resp_body = resp_body
|
||||
|
||||
|
||||
class WebSocketAddressException(WebSocketException):
|
||||
"""
|
||||
If the websocket address info cannot be found, this exception will be raised.
|
||||
"""
|
||||
|
||||
pass
|
||||
@@ -0,0 +1,210 @@
|
||||
"""
|
||||
_handshake.py
|
||||
websocket - WebSocket client library for Python
|
||||
|
||||
Copyright 2025 engn33r
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import hmac
|
||||
import os
|
||||
from base64 import encodebytes as base64encode
|
||||
from http import HTTPStatus
|
||||
|
||||
from ._cookiejar import SimpleCookieJar
|
||||
from ._exceptions import WebSocketException, WebSocketBadStatusException
|
||||
from ._http import read_headers
|
||||
from ._logging import dump, error
|
||||
from ._socket import send
|
||||
|
||||
__all__ = ["handshake_response", "handshake", "SUPPORTED_REDIRECT_STATUSES"]
|
||||
|
||||
# websocket supported version.
|
||||
VERSION = 13
|
||||
|
||||
SUPPORTED_REDIRECT_STATUSES = (
|
||||
HTTPStatus.MOVED_PERMANENTLY,
|
||||
HTTPStatus.FOUND,
|
||||
HTTPStatus.SEE_OTHER,
|
||||
HTTPStatus.TEMPORARY_REDIRECT,
|
||||
HTTPStatus.PERMANENT_REDIRECT,
|
||||
)
|
||||
SUCCESS_STATUSES = SUPPORTED_REDIRECT_STATUSES + (HTTPStatus.SWITCHING_PROTOCOLS,)
|
||||
|
||||
CookieJar = SimpleCookieJar()
|
||||
|
||||
|
||||
class handshake_response:
|
||||
def __init__(self, status: int, headers: dict, subprotocol):
|
||||
self.status = status
|
||||
self.headers = headers
|
||||
self.subprotocol = subprotocol
|
||||
CookieJar.add(headers.get("set-cookie"))
|
||||
|
||||
|
||||
def handshake(
|
||||
sock, url: str, hostname: str, port: int, resource: str, **options
|
||||
) -> handshake_response:
|
||||
headers, key = _get_handshake_headers(resource, url, hostname, port, options)
|
||||
|
||||
header_str = "\r\n".join(headers)
|
||||
send(sock, header_str)
|
||||
dump("request header", header_str)
|
||||
|
||||
status, resp = _get_resp_headers(sock)
|
||||
if status in SUPPORTED_REDIRECT_STATUSES:
|
||||
return handshake_response(status, resp, None)
|
||||
success, subproto = _validate(resp, key, options.get("subprotocols"))
|
||||
if not success:
|
||||
raise WebSocketException("Invalid WebSocket Header")
|
||||
|
||||
return handshake_response(status, resp, subproto)
|
||||
|
||||
|
||||
def _pack_hostname(hostname: str) -> str:
|
||||
# IPv6 address
|
||||
if ":" in hostname:
|
||||
return f"[{hostname}]"
|
||||
return hostname
|
||||
|
||||
|
||||
def _get_handshake_headers(
|
||||
resource: str, url: str, host: str, port: int, options: dict
|
||||
) -> tuple:
|
||||
headers = [f"GET {resource} HTTP/1.1", "Upgrade: websocket"]
|
||||
if port in [80, 443]:
|
||||
hostport = _pack_hostname(host)
|
||||
else:
|
||||
hostport = f"{_pack_hostname(host)}:{port}"
|
||||
if options.get("host"):
|
||||
headers.append(f'Host: {options["host"]}')
|
||||
else:
|
||||
headers.append(f"Host: {hostport}")
|
||||
|
||||
# scheme indicates whether http or https is used in Origin
|
||||
# The same approach is used in parse_url of _url.py to set default port
|
||||
scheme, url = url.split(":", 1)
|
||||
if not options.get("suppress_origin"):
|
||||
if "origin" in options and options["origin"] is not None:
|
||||
headers.append(f'Origin: {options["origin"]}')
|
||||
elif scheme == "wss":
|
||||
headers.append(f"Origin: https://{hostport}")
|
||||
else:
|
||||
headers.append(f"Origin: http://{hostport}")
|
||||
|
||||
key = _create_sec_websocket_key()
|
||||
|
||||
# Append Sec-WebSocket-Key & Sec-WebSocket-Version if not manually specified
|
||||
if not options.get("header") or "Sec-WebSocket-Key" not in options["header"]:
|
||||
headers.append(f"Sec-WebSocket-Key: {key}")
|
||||
else:
|
||||
key = options["header"]["Sec-WebSocket-Key"]
|
||||
|
||||
if not options.get("header") or "Sec-WebSocket-Version" not in options["header"]:
|
||||
headers.append(f"Sec-WebSocket-Version: {VERSION}")
|
||||
|
||||
if not options.get("connection"):
|
||||
headers.append("Connection: Upgrade")
|
||||
else:
|
||||
headers.append(options["connection"])
|
||||
|
||||
if subprotocols := options.get("subprotocols"):
|
||||
headers.append(f'Sec-WebSocket-Protocol: {",".join(subprotocols)}')
|
||||
|
||||
if header := options.get("header"):
|
||||
if isinstance(header, dict):
|
||||
header = [": ".join([k, v]) for k, v in header.items() if v is not None]
|
||||
headers.extend(header)
|
||||
|
||||
server_cookie = CookieJar.get(host)
|
||||
client_cookie = options.get("cookie", None)
|
||||
|
||||
if cookie := "; ".join(filter(None, [server_cookie, client_cookie])):
|
||||
headers.append(f"Cookie: {cookie}")
|
||||
|
||||
headers.extend(("", ""))
|
||||
return headers, key
|
||||
|
||||
|
||||
def _get_resp_headers(sock, success_statuses: tuple = SUCCESS_STATUSES) -> tuple:
|
||||
status, resp_headers, status_message = read_headers(sock)
|
||||
if status not in success_statuses:
|
||||
content_len = resp_headers.get("content-length")
|
||||
if content_len:
|
||||
# Use chunked reading to avoid SSL BAD_LENGTH error on large responses
|
||||
from ._socket import recv
|
||||
|
||||
response_body = b""
|
||||
remaining = int(content_len)
|
||||
while remaining > 0:
|
||||
chunk_size = min(remaining, 16384) # Read in 16KB chunks
|
||||
chunk = recv(sock, chunk_size)
|
||||
response_body += chunk
|
||||
remaining -= len(chunk)
|
||||
else:
|
||||
response_body = None
|
||||
raise WebSocketBadStatusException(
|
||||
f"Handshake status {status} {status_message} -+-+- {resp_headers} -+-+- {response_body}",
|
||||
status,
|
||||
status_message,
|
||||
resp_headers,
|
||||
response_body,
|
||||
)
|
||||
return status, resp_headers
|
||||
|
||||
|
||||
_HEADERS_TO_CHECK = {
|
||||
"upgrade": "websocket",
|
||||
"connection": "upgrade",
|
||||
}
|
||||
|
||||
|
||||
def _validate(headers, key: str, subprotocols) -> tuple:
|
||||
subproto = None
|
||||
for k, v in _HEADERS_TO_CHECK.items():
|
||||
r = headers.get(k, None)
|
||||
if not r:
|
||||
return False, None
|
||||
r = [x.strip().lower() for x in r.split(",")]
|
||||
if v not in r:
|
||||
return False, None
|
||||
|
||||
if subprotocols:
|
||||
subproto = headers.get("sec-websocket-protocol", None)
|
||||
if not subproto or subproto.lower() not in [s.lower() for s in subprotocols]:
|
||||
error(f"Invalid subprotocol: {subprotocols}")
|
||||
return False, None
|
||||
subproto = subproto.lower()
|
||||
|
||||
result = headers.get("sec-websocket-accept", None)
|
||||
if not result:
|
||||
return False, None
|
||||
result = result.lower()
|
||||
|
||||
if isinstance(result, str):
|
||||
result = result.encode("utf-8")
|
||||
|
||||
value = f"{key}258EAFA5-E914-47DA-95CA-C5AB0DC85B11".encode("utf-8")
|
||||
hashed = base64encode(hashlib.sha1(value).digest()).strip().lower()
|
||||
|
||||
if hmac.compare_digest(hashed, result):
|
||||
return True, subproto
|
||||
else:
|
||||
return False, None
|
||||
|
||||
|
||||
def _create_sec_websocket_key() -> str:
|
||||
randomness = os.urandom(16)
|
||||
return base64encode(randomness).decode("utf-8").strip()
|
||||
@@ -0,0 +1,411 @@
|
||||
"""
|
||||
_http.py
|
||||
websocket - WebSocket client library for Python
|
||||
|
||||
Copyright 2025 engn33r
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
import errno
|
||||
import os
|
||||
import socket
|
||||
from base64 import encodebytes as base64encode
|
||||
|
||||
from ._exceptions import (
|
||||
WebSocketAddressException,
|
||||
WebSocketException,
|
||||
WebSocketProxyException,
|
||||
)
|
||||
from ._logging import debug, dump, trace
|
||||
from ._socket import DEFAULT_SOCKET_OPTION, recv_line, send
|
||||
from ._ssl_compat import HAVE_SSL, ssl
|
||||
from ._url import get_proxy_info, parse_url
|
||||
|
||||
__all__ = ["proxy_info", "connect", "read_headers"]
|
||||
|
||||
try:
|
||||
from python_socks._errors import ProxyConnectionError, ProxyError, ProxyTimeoutError
|
||||
from python_socks._types import ProxyType
|
||||
from python_socks.sync import Proxy
|
||||
|
||||
HAVE_PYTHON_SOCKS = True
|
||||
except:
|
||||
HAVE_PYTHON_SOCKS = False
|
||||
|
||||
class ProxyError(Exception):
|
||||
pass
|
||||
|
||||
class ProxyTimeoutError(Exception):
|
||||
pass
|
||||
|
||||
class ProxyConnectionError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class proxy_info:
|
||||
def __init__(self, **options):
|
||||
self.proxy_host = options.get("http_proxy_host", None)
|
||||
if self.proxy_host:
|
||||
self.proxy_port = options.get("http_proxy_port", 0)
|
||||
self.auth = options.get("http_proxy_auth", None)
|
||||
self.no_proxy = options.get("http_no_proxy", None)
|
||||
self.proxy_protocol = options.get("proxy_type", "http")
|
||||
# Note: If timeout not specified, default python-socks timeout is 60 seconds
|
||||
self.proxy_timeout = options.get("http_proxy_timeout", None)
|
||||
if self.proxy_protocol not in [
|
||||
"http",
|
||||
"socks4",
|
||||
"socks4a",
|
||||
"socks5",
|
||||
"socks5h",
|
||||
]:
|
||||
raise ProxyError(
|
||||
"Only http, socks4, socks5 proxy protocols are supported"
|
||||
)
|
||||
else:
|
||||
self.proxy_port = 0
|
||||
self.auth = None
|
||||
self.no_proxy = None
|
||||
self.proxy_protocol = "http"
|
||||
|
||||
|
||||
def _start_proxied_socket(url: str, options, proxy) -> tuple:
|
||||
if not HAVE_PYTHON_SOCKS:
|
||||
raise WebSocketException(
|
||||
"Python Socks is needed for SOCKS proxying but is not available"
|
||||
)
|
||||
|
||||
hostname, port, resource, is_secure = parse_url(url)
|
||||
|
||||
if proxy.proxy_protocol == "socks4":
|
||||
rdns = False
|
||||
proxy_type = ProxyType.SOCKS4
|
||||
# socks4a sends DNS through proxy
|
||||
elif proxy.proxy_protocol == "socks4a":
|
||||
rdns = True
|
||||
proxy_type = ProxyType.SOCKS4
|
||||
elif proxy.proxy_protocol == "socks5":
|
||||
rdns = False
|
||||
proxy_type = ProxyType.SOCKS5
|
||||
# socks5h sends DNS through proxy
|
||||
elif proxy.proxy_protocol == "socks5h":
|
||||
rdns = True
|
||||
proxy_type = ProxyType.SOCKS5
|
||||
|
||||
ws_proxy = Proxy.create(
|
||||
proxy_type=proxy_type,
|
||||
host=proxy.proxy_host,
|
||||
port=int(proxy.proxy_port),
|
||||
username=proxy.auth[0] if proxy.auth else None,
|
||||
password=proxy.auth[1] if proxy.auth else None,
|
||||
rdns=rdns,
|
||||
)
|
||||
|
||||
sock = ws_proxy.connect(hostname, port, timeout=proxy.proxy_timeout)
|
||||
|
||||
if is_secure:
|
||||
if HAVE_SSL:
|
||||
sock = _ssl_socket(sock, options.sslopt, hostname)
|
||||
else:
|
||||
raise WebSocketException("SSL not available.")
|
||||
|
||||
return sock, (hostname, port, resource)
|
||||
|
||||
|
||||
def connect(url: str, options, proxy, socket):
|
||||
# Use _start_proxied_socket() only for socks4 or socks5 proxy
|
||||
# Use _tunnel() for http proxy
|
||||
# TODO: Use python-socks for http protocol also, to standardize flow
|
||||
if proxy.proxy_host and not socket and proxy.proxy_protocol != "http":
|
||||
return _start_proxied_socket(url, options, proxy)
|
||||
|
||||
hostname, port_from_url, resource, is_secure = parse_url(url)
|
||||
|
||||
if socket:
|
||||
return socket, (hostname, port_from_url, resource)
|
||||
|
||||
addrinfo_list, need_tunnel, auth = _get_addrinfo_list(
|
||||
hostname, port_from_url, is_secure, proxy
|
||||
)
|
||||
if not addrinfo_list:
|
||||
raise WebSocketException(f"Host not found.: {hostname}:{port_from_url}")
|
||||
|
||||
sock = None
|
||||
try:
|
||||
sock = _open_socket(addrinfo_list, options.sockopt, options.timeout)
|
||||
if need_tunnel:
|
||||
sock = _tunnel(sock, hostname, port_from_url, auth)
|
||||
|
||||
if is_secure:
|
||||
if HAVE_SSL:
|
||||
sock = _ssl_socket(sock, options.sslopt, hostname)
|
||||
else:
|
||||
raise WebSocketException("SSL not available.")
|
||||
|
||||
return sock, (hostname, port_from_url, resource)
|
||||
except:
|
||||
if sock:
|
||||
sock.close()
|
||||
raise
|
||||
|
||||
|
||||
def _get_addrinfo_list(hostname, port: int, is_secure: bool, proxy) -> tuple:
|
||||
phost, pport, pauth = get_proxy_info(
|
||||
hostname,
|
||||
is_secure,
|
||||
proxy.proxy_host,
|
||||
proxy.proxy_port,
|
||||
proxy.auth,
|
||||
proxy.no_proxy,
|
||||
)
|
||||
try:
|
||||
# when running on windows 10, getaddrinfo without socktype returns a socktype 0.
|
||||
# This generates an error exception: `_on_error: exception Socket type must be stream or datagram, not 0`
|
||||
# or `OSError: [Errno 22] Invalid argument` when creating socket. Force the socket type to SOCK_STREAM.
|
||||
if not phost:
|
||||
addrinfo_list = socket.getaddrinfo(
|
||||
hostname, port, 0, socket.SOCK_STREAM, socket.SOL_TCP
|
||||
)
|
||||
return addrinfo_list, False, None
|
||||
else:
|
||||
pport = pport and pport or 80
|
||||
# when running on windows 10, the getaddrinfo used above
|
||||
# returns a socktype 0. This generates an error exception:
|
||||
# _on_error: exception Socket type must be stream or datagram, not 0
|
||||
# Force the socket type to SOCK_STREAM
|
||||
addrinfo_list = socket.getaddrinfo(
|
||||
phost, pport, 0, socket.SOCK_STREAM, socket.SOL_TCP
|
||||
)
|
||||
return addrinfo_list, True, pauth
|
||||
except socket.gaierror as e:
|
||||
raise WebSocketAddressException(e)
|
||||
|
||||
|
||||
def _open_socket(addrinfo_list, sockopt, timeout):
|
||||
err = None
|
||||
for addrinfo in addrinfo_list:
|
||||
family, socktype, proto = addrinfo[:3]
|
||||
sock = socket.socket(family, socktype, proto)
|
||||
sock.settimeout(timeout)
|
||||
for opts in DEFAULT_SOCKET_OPTION:
|
||||
sock.setsockopt(*opts)
|
||||
for opts in sockopt:
|
||||
sock.setsockopt(*opts)
|
||||
|
||||
address = addrinfo[4]
|
||||
err = None
|
||||
while not err:
|
||||
try:
|
||||
sock.connect(address)
|
||||
except socket.error as error:
|
||||
sock.close()
|
||||
error.remote_ip = str(address[0])
|
||||
try:
|
||||
eConnRefused = (
|
||||
errno.ECONNREFUSED,
|
||||
errno.WSAECONNREFUSED,
|
||||
errno.ENETUNREACH,
|
||||
)
|
||||
except AttributeError:
|
||||
eConnRefused = (errno.ECONNREFUSED, errno.ENETUNREACH)
|
||||
if error.errno not in eConnRefused:
|
||||
raise error
|
||||
err = error
|
||||
continue
|
||||
else:
|
||||
break
|
||||
else:
|
||||
continue
|
||||
break
|
||||
else:
|
||||
if err:
|
||||
raise err
|
||||
|
||||
return sock
|
||||
|
||||
|
||||
def _wrap_sni_socket(sock: socket.socket, sslopt: dict, hostname, check_hostname):
|
||||
context = sslopt.get("context", None)
|
||||
if not context:
|
||||
context = ssl.SSLContext(sslopt.get("ssl_version", ssl.PROTOCOL_TLS_CLIENT))
|
||||
# Non default context need to manually enable SSLKEYLOGFILE support by setting the keylog_filename attribute.
|
||||
# For more details see also:
|
||||
# * https://docs.python.org/3.8/library/ssl.html?highlight=sslkeylogfile#context-creation
|
||||
# * https://docs.python.org/3.8/library/ssl.html?highlight=sslkeylogfile#ssl.SSLContext.keylog_filename
|
||||
keylog_file = os.environ.get("SSLKEYLOGFILE")
|
||||
if keylog_file is not None:
|
||||
context.keylog_filename = keylog_file
|
||||
|
||||
if sslopt.get("cert_reqs", ssl.CERT_NONE) != ssl.CERT_NONE:
|
||||
cafile = sslopt.get("ca_certs", None)
|
||||
capath = sslopt.get("ca_cert_path", None)
|
||||
if cafile or capath:
|
||||
try:
|
||||
context.load_verify_locations(cafile=cafile, capath=capath)
|
||||
except (FileNotFoundError, ssl.SSLError, ValueError) as e:
|
||||
raise WebSocketException(f"SSL CA certificate loading failed: {e}")
|
||||
elif hasattr(context, "load_default_certs"):
|
||||
try:
|
||||
context.load_default_certs(ssl.Purpose.SERVER_AUTH)
|
||||
except ssl.SSLError as e:
|
||||
raise WebSocketException(
|
||||
f"SSL default certificate loading failed: {e}"
|
||||
)
|
||||
if sslopt.get("certfile", None):
|
||||
try:
|
||||
context.load_cert_chain(
|
||||
sslopt["certfile"],
|
||||
sslopt.get("keyfile", None),
|
||||
sslopt.get("password", None),
|
||||
)
|
||||
except (FileNotFoundError, ValueError) as e:
|
||||
raise WebSocketException(f"SSL client certificate loading failed: {e}")
|
||||
except ssl.SSLError as e:
|
||||
raise WebSocketException(f"SSL client certificate loading failed: {e}")
|
||||
|
||||
# Python 3.10 switch to PROTOCOL_TLS_CLIENT defaults to "cert_reqs = ssl.CERT_REQUIRED" and "check_hostname = True"
|
||||
# If both disabled, set check_hostname before verify_mode
|
||||
# see https://github.com/liris/websocket-client/commit/b96a2e8fa765753e82eea531adb19716b52ca3ca#commitcomment-10803153
|
||||
if sslopt.get("cert_reqs", ssl.CERT_NONE) == ssl.CERT_NONE and not sslopt.get(
|
||||
"check_hostname", False
|
||||
):
|
||||
context.check_hostname = False
|
||||
context.verify_mode = ssl.CERT_NONE
|
||||
else:
|
||||
context.check_hostname = sslopt.get("check_hostname", True)
|
||||
context.verify_mode = sslopt.get("cert_reqs", ssl.CERT_REQUIRED)
|
||||
|
||||
if "ciphers" in sslopt:
|
||||
try:
|
||||
context.set_ciphers(sslopt["ciphers"])
|
||||
except ssl.SSLError as e:
|
||||
raise WebSocketException(f"SSL cipher configuration failed: {e}")
|
||||
if "cert_chain" in sslopt:
|
||||
try:
|
||||
cert_chain = sslopt["cert_chain"]
|
||||
if not isinstance(cert_chain, (tuple, list)) or len(cert_chain) != 3:
|
||||
raise ValueError(
|
||||
"cert_chain must be a tuple/list of (certfile, keyfile, password)"
|
||||
)
|
||||
certfile, keyfile, password = cert_chain
|
||||
context.load_cert_chain(certfile, keyfile, password)
|
||||
except ValueError:
|
||||
raise
|
||||
except (FileNotFoundError, ssl.SSLError) as e:
|
||||
raise WebSocketException(
|
||||
f"SSL client certificate configuration failed: {e}"
|
||||
)
|
||||
if "ecdh_curve" in sslopt:
|
||||
try:
|
||||
context.set_ecdh_curve(sslopt["ecdh_curve"])
|
||||
except ValueError as e:
|
||||
raise WebSocketException(f"SSL ECDH curve configuration failed: {e}")
|
||||
|
||||
return context.wrap_socket(
|
||||
sock,
|
||||
do_handshake_on_connect=sslopt.get("do_handshake_on_connect", True),
|
||||
suppress_ragged_eofs=sslopt.get("suppress_ragged_eofs", True),
|
||||
server_hostname=hostname,
|
||||
)
|
||||
|
||||
|
||||
def _ssl_socket(sock: socket.socket, user_sslopt: dict, hostname):
|
||||
sslopt: dict = {"cert_reqs": ssl.CERT_REQUIRED}
|
||||
sslopt.update(user_sslopt)
|
||||
|
||||
cert_path = os.environ.get("WEBSOCKET_CLIENT_CA_BUNDLE")
|
||||
if (
|
||||
cert_path
|
||||
and os.path.isfile(cert_path)
|
||||
and user_sslopt.get("ca_certs", None) is None
|
||||
):
|
||||
sslopt["ca_certs"] = cert_path
|
||||
elif (
|
||||
cert_path
|
||||
and os.path.isdir(cert_path)
|
||||
and user_sslopt.get("ca_cert_path", None) is None
|
||||
):
|
||||
sslopt["ca_cert_path"] = cert_path
|
||||
|
||||
if sslopt.get("server_hostname", None):
|
||||
hostname = sslopt["server_hostname"]
|
||||
|
||||
check_hostname = sslopt.get("check_hostname", True)
|
||||
sock = _wrap_sni_socket(sock, sslopt, hostname, check_hostname)
|
||||
|
||||
return sock
|
||||
|
||||
|
||||
def _tunnel(sock: socket.socket, host, port: int, auth) -> socket.socket:
|
||||
debug("Connecting proxy...")
|
||||
connect_header = f"CONNECT {host}:{port} HTTP/1.1\r\n"
|
||||
connect_header += f"Host: {host}:{port}\r\n"
|
||||
|
||||
# TODO: support digest auth.
|
||||
if auth and auth[0]:
|
||||
auth_str = auth[0]
|
||||
if auth[1]:
|
||||
auth_str += f":{auth[1]}"
|
||||
encoded_str = base64encode(auth_str.encode()).strip().decode().replace("\n", "")
|
||||
connect_header += f"Proxy-Authorization: Basic {encoded_str}\r\n"
|
||||
connect_header += "\r\n"
|
||||
dump("request header", connect_header)
|
||||
|
||||
send(sock, connect_header)
|
||||
|
||||
try:
|
||||
status, _, _ = read_headers(sock)
|
||||
except (socket.error, WebSocketException) as e:
|
||||
raise WebSocketProxyException(str(e))
|
||||
|
||||
if status != 200:
|
||||
raise WebSocketProxyException(f"failed CONNECT via proxy status: {status}")
|
||||
|
||||
return sock
|
||||
|
||||
|
||||
def read_headers(sock: socket.socket) -> tuple:
|
||||
status = None
|
||||
status_message = None
|
||||
headers: dict = {}
|
||||
trace("--- response header ---")
|
||||
|
||||
while True:
|
||||
line = recv_line(sock)
|
||||
line = line.decode("utf-8").strip()
|
||||
if not line:
|
||||
break
|
||||
trace(line)
|
||||
if not status:
|
||||
status_info = line.split(" ", 2)
|
||||
status = int(status_info[1])
|
||||
if len(status_info) > 2:
|
||||
status_message = status_info[2]
|
||||
else:
|
||||
kv = line.split(":", 1)
|
||||
if len(kv) != 2:
|
||||
raise WebSocketException("Invalid header")
|
||||
key, value = kv
|
||||
if key.lower() == "set-cookie" and headers.get("set-cookie"):
|
||||
existing_cookie = headers.get("set-cookie")
|
||||
if existing_cookie is not None:
|
||||
headers["set-cookie"] = existing_cookie + "; " + value.strip()
|
||||
else:
|
||||
headers["set-cookie"] = value.strip()
|
||||
else:
|
||||
headers[key.lower()] = value.strip()
|
||||
|
||||
trace("-----------------------")
|
||||
|
||||
return status, headers, status_message
|
||||
@@ -0,0 +1,106 @@
|
||||
import logging
|
||||
|
||||
"""
|
||||
_logging.py
|
||||
websocket - WebSocket client library for Python
|
||||
|
||||
Copyright 2025 engn33r
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
_logger = logging.getLogger("websocket")
|
||||
try:
|
||||
from logging import NullHandler
|
||||
except ImportError:
|
||||
|
||||
class NullHandler(logging.Handler): # type: ignore[no-redef]
|
||||
def emit(self, record) -> None:
|
||||
pass
|
||||
|
||||
|
||||
_logger.addHandler(NullHandler())
|
||||
|
||||
_traceEnabled = False
|
||||
|
||||
__all__ = [
|
||||
"enableTrace",
|
||||
"dump",
|
||||
"error",
|
||||
"warning",
|
||||
"debug",
|
||||
"trace",
|
||||
"isEnabledForError",
|
||||
"isEnabledForDebug",
|
||||
"isEnabledForTrace",
|
||||
]
|
||||
|
||||
|
||||
def enableTrace(
|
||||
traceable: bool,
|
||||
handler: logging.StreamHandler = logging.StreamHandler(),
|
||||
level: str = "DEBUG",
|
||||
) -> None:
|
||||
"""
|
||||
Turn on/off the traceability.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
traceable: bool
|
||||
If set to True, traceability is enabled.
|
||||
"""
|
||||
global _traceEnabled
|
||||
_traceEnabled = traceable
|
||||
if traceable:
|
||||
_logger.addHandler(handler)
|
||||
_logger.setLevel(getattr(logging, level))
|
||||
|
||||
|
||||
def dump(title: str, message: str) -> None:
|
||||
if _traceEnabled:
|
||||
_logger.debug(f"--- {title} ---")
|
||||
_logger.debug(message)
|
||||
_logger.debug("-----------------------")
|
||||
|
||||
|
||||
def error(msg: str) -> None:
|
||||
_logger.error(msg)
|
||||
|
||||
|
||||
def warning(msg: str) -> None:
|
||||
_logger.warning(msg)
|
||||
|
||||
|
||||
def debug(msg: str) -> None:
|
||||
_logger.debug(msg)
|
||||
|
||||
|
||||
def info(msg: str) -> None:
|
||||
_logger.info(msg)
|
||||
|
||||
|
||||
def trace(msg: str) -> None:
|
||||
if _traceEnabled:
|
||||
_logger.debug(msg)
|
||||
|
||||
|
||||
def isEnabledForError() -> bool:
|
||||
return _logger.isEnabledFor(logging.ERROR)
|
||||
|
||||
|
||||
def isEnabledForDebug() -> bool:
|
||||
return _logger.isEnabledFor(logging.DEBUG)
|
||||
|
||||
|
||||
def isEnabledForTrace() -> bool:
|
||||
return _traceEnabled
|
||||
@@ -0,0 +1,202 @@
|
||||
import errno
|
||||
import selectors
|
||||
import socket
|
||||
from typing import Optional, Union, Any
|
||||
|
||||
from ._exceptions import (
|
||||
WebSocketConnectionClosedException,
|
||||
WebSocketTimeoutException,
|
||||
)
|
||||
from ._ssl_compat import SSLError, SSLEOFError, SSLWantReadError, SSLWantWriteError
|
||||
from ._utils import extract_error_code, extract_err_message
|
||||
|
||||
"""
|
||||
_socket.py
|
||||
websocket - WebSocket client library for Python
|
||||
|
||||
Copyright 2025 engn33r
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
DEFAULT_SOCKET_OPTION = [(socket.SOL_TCP, socket.TCP_NODELAY, 1)]
|
||||
if hasattr(socket, "SO_KEEPALIVE"):
|
||||
DEFAULT_SOCKET_OPTION.append((socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1))
|
||||
if hasattr(socket, "TCP_KEEPIDLE"):
|
||||
DEFAULT_SOCKET_OPTION.append((socket.SOL_TCP, socket.TCP_KEEPIDLE, 30))
|
||||
if hasattr(socket, "TCP_KEEPINTVL"):
|
||||
DEFAULT_SOCKET_OPTION.append((socket.SOL_TCP, socket.TCP_KEEPINTVL, 10))
|
||||
if hasattr(socket, "TCP_KEEPCNT"):
|
||||
DEFAULT_SOCKET_OPTION.append((socket.SOL_TCP, socket.TCP_KEEPCNT, 3))
|
||||
|
||||
_default_timeout = None
|
||||
|
||||
__all__ = [
|
||||
"DEFAULT_SOCKET_OPTION",
|
||||
"sock_opt",
|
||||
"setdefaulttimeout",
|
||||
"getdefaulttimeout",
|
||||
"recv",
|
||||
"recv_line",
|
||||
"send",
|
||||
]
|
||||
|
||||
|
||||
class sock_opt:
|
||||
def __init__(
|
||||
self, sockopt: Optional[list[tuple]], sslopt: Optional[dict[str, Any]]
|
||||
) -> None:
|
||||
if sockopt is None:
|
||||
sockopt = []
|
||||
if sslopt is None:
|
||||
sslopt = {}
|
||||
self.sockopt = sockopt
|
||||
self.sslopt = sslopt
|
||||
self.timeout: Optional[Union[int, float]] = None
|
||||
|
||||
|
||||
def setdefaulttimeout(timeout: Optional[Union[int, float]]) -> None:
|
||||
"""
|
||||
Set the global timeout setting to connect.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
timeout: int or float
|
||||
default socket timeout time (in seconds)
|
||||
"""
|
||||
global _default_timeout
|
||||
_default_timeout = timeout
|
||||
|
||||
|
||||
def getdefaulttimeout() -> Optional[Union[int, float]]:
|
||||
"""
|
||||
Get default timeout
|
||||
|
||||
Returns
|
||||
----------
|
||||
_default_timeout: int or float
|
||||
Return the global timeout setting (in seconds) to connect.
|
||||
"""
|
||||
return _default_timeout
|
||||
|
||||
|
||||
def recv(sock: socket.socket, bufsize: int) -> bytes:
|
||||
if not sock:
|
||||
raise WebSocketConnectionClosedException("socket is already closed.")
|
||||
|
||||
def _recv():
|
||||
try:
|
||||
return sock.recv(bufsize)
|
||||
except SSLWantReadError:
|
||||
# Don't return None implicitly - fall through to retry logic
|
||||
pass
|
||||
except socket.error as exc:
|
||||
error_code = extract_error_code(exc)
|
||||
if error_code not in [errno.EAGAIN, errno.EWOULDBLOCK]:
|
||||
raise
|
||||
# Don't return None implicitly - fall through to retry logic
|
||||
|
||||
# Retry logic using selector for both SSLWantReadError and EAGAIN/EWOULDBLOCK
|
||||
sel = selectors.DefaultSelector()
|
||||
sel.register(sock, selectors.EVENT_READ)
|
||||
|
||||
r = sel.select(sock.gettimeout())
|
||||
sel.close()
|
||||
|
||||
if r:
|
||||
return sock.recv(bufsize)
|
||||
else:
|
||||
# Selector timeout should raise WebSocketTimeoutException
|
||||
# not return None which gets misclassified as connection closed
|
||||
raise WebSocketTimeoutException("Connection timed out waiting for data")
|
||||
|
||||
try:
|
||||
if sock.gettimeout() == 0:
|
||||
bytes_ = sock.recv(bufsize)
|
||||
else:
|
||||
bytes_ = _recv()
|
||||
except TimeoutError:
|
||||
raise WebSocketTimeoutException("Connection timed out")
|
||||
except socket.timeout as e:
|
||||
message = extract_err_message(e)
|
||||
raise WebSocketTimeoutException(message)
|
||||
except SSLError as e:
|
||||
message = extract_err_message(e)
|
||||
if isinstance(message, str) and "timed out" in message:
|
||||
raise WebSocketTimeoutException(message)
|
||||
else:
|
||||
raise
|
||||
|
||||
if bytes_ is None:
|
||||
raise WebSocketConnectionClosedException("Connection to remote host was lost.")
|
||||
if not bytes_:
|
||||
raise WebSocketConnectionClosedException("Connection to remote host was lost.")
|
||||
|
||||
return bytes_
|
||||
|
||||
|
||||
def recv_line(sock: socket.socket) -> bytes:
|
||||
line = []
|
||||
while True:
|
||||
c = recv(sock, 1)
|
||||
line.append(c)
|
||||
if c == b"\n":
|
||||
break
|
||||
return b"".join(line)
|
||||
|
||||
|
||||
def send(sock: socket.socket, data: Union[bytes, str]) -> int:
|
||||
if isinstance(data, str):
|
||||
data = data.encode("utf-8")
|
||||
|
||||
if not sock:
|
||||
raise WebSocketConnectionClosedException("socket is already closed.")
|
||||
|
||||
def _send() -> int:
|
||||
try:
|
||||
return sock.send(data)
|
||||
except SSLEOFError:
|
||||
raise WebSocketConnectionClosedException("socket is already closed.")
|
||||
except SSLWantWriteError:
|
||||
pass
|
||||
except socket.error as exc:
|
||||
error_code = extract_error_code(exc)
|
||||
if error_code is None:
|
||||
raise
|
||||
if error_code not in [errno.EAGAIN, errno.EWOULDBLOCK]:
|
||||
raise
|
||||
|
||||
sel = selectors.DefaultSelector()
|
||||
sel.register(sock, selectors.EVENT_WRITE)
|
||||
|
||||
w = sel.select(sock.gettimeout())
|
||||
sel.close()
|
||||
|
||||
if w:
|
||||
return sock.send(data)
|
||||
return 0
|
||||
|
||||
try:
|
||||
if sock.gettimeout() == 0:
|
||||
return sock.send(data)
|
||||
else:
|
||||
return _send()
|
||||
except socket.timeout as e:
|
||||
message = extract_err_message(e)
|
||||
raise WebSocketTimeoutException(message)
|
||||
except (OSError, SSLError) as e:
|
||||
message = extract_err_message(e)
|
||||
if isinstance(message, str) and "timed out" in message:
|
||||
raise WebSocketTimeoutException(message)
|
||||
else:
|
||||
raise
|
||||
@@ -0,0 +1,66 @@
|
||||
"""
|
||||
_ssl_compat.py
|
||||
websocket - WebSocket client library for Python
|
||||
|
||||
Copyright 2025 engn33r
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import ssl as _ssl_module
|
||||
from ssl import (
|
||||
SSLError as _SSLErrorType,
|
||||
SSLEOFError as _SSLEOFErrorType,
|
||||
SSLWantReadError as _SSLWantReadErrorType,
|
||||
SSLWantWriteError as _SSLWantWriteErrorType,
|
||||
)
|
||||
else:
|
||||
_ssl_module = None
|
||||
_SSLErrorType = None
|
||||
_SSLEOFErrorType = None
|
||||
_SSLWantReadErrorType = None
|
||||
_SSLWantWriteErrorType = None
|
||||
|
||||
__all__ = [
|
||||
"HAVE_SSL",
|
||||
"ssl",
|
||||
"SSLError",
|
||||
"SSLEOFError",
|
||||
"SSLWantReadError",
|
||||
"SSLWantWriteError",
|
||||
]
|
||||
|
||||
try:
|
||||
import ssl
|
||||
from ssl import SSLError, SSLEOFError, SSLWantReadError, SSLWantWriteError # type: ignore[attr-defined]
|
||||
|
||||
HAVE_SSL = True
|
||||
except ImportError:
|
||||
# dummy class of SSLError for environment without ssl support
|
||||
class SSLError(Exception): # type: ignore[no-redef]
|
||||
pass
|
||||
|
||||
class SSLEOFError(Exception): # type: ignore[no-redef]
|
||||
pass
|
||||
|
||||
class SSLWantReadError(Exception): # type: ignore[no-redef]
|
||||
pass
|
||||
|
||||
class SSLWantWriteError(Exception): # type: ignore[no-redef]
|
||||
pass
|
||||
|
||||
ssl = None # type: ignore[assignment,no-redef]
|
||||
HAVE_SSL = False
|
||||
@@ -0,0 +1,189 @@
|
||||
import ipaddress
|
||||
import os
|
||||
from typing import Optional
|
||||
from urllib.parse import unquote, urlparse
|
||||
from ._exceptions import WebSocketProxyException
|
||||
|
||||
"""
|
||||
_url.py
|
||||
websocket - WebSocket client library for Python
|
||||
|
||||
Copyright 2025 engn33r
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
__all__ = ["parse_url", "get_proxy_info"]
|
||||
|
||||
|
||||
def parse_url(url: str) -> tuple:
|
||||
"""
|
||||
parse url and the result is tuple of
|
||||
(hostname, port, resource path and the flag of secure mode)
|
||||
|
||||
Parameters
|
||||
----------
|
||||
url: str
|
||||
url string.
|
||||
"""
|
||||
if ":" not in url:
|
||||
raise ValueError("url is invalid")
|
||||
|
||||
scheme, url = url.split(":", 1)
|
||||
|
||||
parsed = urlparse(url, scheme="http")
|
||||
if parsed.hostname:
|
||||
hostname = parsed.hostname
|
||||
else:
|
||||
raise ValueError("hostname is invalid")
|
||||
port = 0
|
||||
if parsed.port:
|
||||
port = parsed.port
|
||||
|
||||
is_secure = False
|
||||
if scheme == "ws":
|
||||
if not port:
|
||||
port = 80
|
||||
elif scheme == "wss":
|
||||
is_secure = True
|
||||
if not port:
|
||||
port = 443
|
||||
else:
|
||||
raise ValueError("scheme %s is invalid" % scheme)
|
||||
|
||||
if parsed.path:
|
||||
resource = parsed.path
|
||||
else:
|
||||
resource = "/"
|
||||
|
||||
if parsed.query:
|
||||
resource += f"?{parsed.query}"
|
||||
|
||||
return hostname, port, resource, is_secure
|
||||
|
||||
|
||||
def _is_ip_address(addr: str) -> bool:
|
||||
if not isinstance(addr, str):
|
||||
raise TypeError("_is_ip_address() argument 1 must be str")
|
||||
try:
|
||||
ipaddress.ip_address(addr)
|
||||
except ValueError:
|
||||
return False
|
||||
else:
|
||||
return True
|
||||
|
||||
|
||||
def _is_subnet_address(hostname: str) -> bool:
|
||||
try:
|
||||
ipaddress.ip_network(hostname)
|
||||
except ValueError:
|
||||
return False
|
||||
else:
|
||||
return True
|
||||
|
||||
|
||||
def _is_address_in_network(ip: str, net: str) -> bool:
|
||||
try:
|
||||
return ipaddress.ip_network(ip).subnet_of(ipaddress.ip_network(net))
|
||||
except TypeError:
|
||||
return False
|
||||
|
||||
|
||||
def _is_no_proxy_host(hostname: str, no_proxy: Optional[list[str]]) -> bool:
|
||||
if not no_proxy:
|
||||
if v := os.environ.get("no_proxy", os.environ.get("NO_PROXY", "")).replace(
|
||||
" ", ""
|
||||
):
|
||||
no_proxy = v.split(",")
|
||||
|
||||
if not no_proxy:
|
||||
no_proxy = []
|
||||
|
||||
if "*" in no_proxy:
|
||||
return True
|
||||
if hostname in no_proxy:
|
||||
return True
|
||||
if _is_ip_address(hostname):
|
||||
return any(
|
||||
[
|
||||
_is_address_in_network(hostname, subnet)
|
||||
for subnet in no_proxy
|
||||
if _is_subnet_address(subnet)
|
||||
]
|
||||
)
|
||||
for domain in [domain for domain in no_proxy if domain.startswith(".")]:
|
||||
endDomain = domain.lstrip(".")
|
||||
if hostname.endswith(endDomain):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def get_proxy_info(
|
||||
hostname: str,
|
||||
is_secure: bool,
|
||||
proxy_host: Optional[str] = None,
|
||||
proxy_port: int = 0,
|
||||
proxy_auth: Optional[tuple] = None,
|
||||
no_proxy: Optional[list[str]] = None,
|
||||
proxy_type: str = "http",
|
||||
) -> tuple:
|
||||
"""
|
||||
Try to retrieve proxy host and port from environment
|
||||
if not provided in options.
|
||||
Result is (proxy_host, proxy_port, proxy_auth).
|
||||
proxy_auth is tuple of username and password
|
||||
of proxy authentication information.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
hostname: str
|
||||
Websocket server name.
|
||||
is_secure: bool
|
||||
Is the connection secure? (wss) looks for "https_proxy" in env
|
||||
instead of "http_proxy"
|
||||
proxy_host: str
|
||||
http proxy host name.
|
||||
proxy_port: str or int
|
||||
http proxy port.
|
||||
no_proxy: list
|
||||
Whitelisted host names that don't use the proxy.
|
||||
proxy_auth: tuple
|
||||
HTTP proxy auth information. Tuple of username and password. Default is None.
|
||||
proxy_type: str
|
||||
Specify the proxy protocol (http, socks4, socks4a, socks5, socks5h). Default is "http".
|
||||
Use socks4a or socks5h if you want to send DNS requests through the proxy.
|
||||
"""
|
||||
if _is_no_proxy_host(hostname, no_proxy):
|
||||
return None, 0, None
|
||||
|
||||
if proxy_host:
|
||||
if not proxy_port:
|
||||
raise WebSocketProxyException("Cannot use port 0 when proxy_host specified")
|
||||
port = proxy_port
|
||||
auth = proxy_auth
|
||||
return proxy_host, port, auth
|
||||
|
||||
env_key = "https_proxy" if is_secure else "http_proxy"
|
||||
value = os.environ.get(env_key, os.environ.get(env_key.upper(), "")).replace(
|
||||
" ", ""
|
||||
)
|
||||
if value:
|
||||
proxy = urlparse(value)
|
||||
auth = (
|
||||
(unquote(proxy.username or ""), unquote(proxy.password or ""))
|
||||
if proxy.username
|
||||
else None
|
||||
)
|
||||
return proxy.hostname, proxy.port, auth
|
||||
|
||||
return None, 0, None
|
||||
@@ -0,0 +1,460 @@
|
||||
from typing import Union, Optional
|
||||
|
||||
"""
|
||||
_utils.py
|
||||
websocket - WebSocket client library for Python
|
||||
|
||||
Copyright 2025 engn33r
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
__all__ = ["NoLock", "validate_utf8", "extract_err_message", "extract_error_code"]
|
||||
|
||||
|
||||
class NoLock:
|
||||
def __enter__(self) -> None:
|
||||
pass
|
||||
|
||||
def __exit__(self, exc_type, exc_value, traceback) -> None:
|
||||
pass
|
||||
|
||||
|
||||
try:
|
||||
# If wsaccel is available we use compiled routines to validate UTF-8
|
||||
# strings.
|
||||
from wsaccel.utf8validator import Utf8Validator
|
||||
|
||||
def _validate_utf8(utfbytes: Union[str, bytes]) -> bool:
|
||||
result: bool = Utf8Validator().validate(utfbytes)[0]
|
||||
return result
|
||||
|
||||
except ImportError:
|
||||
# UTF-8 validator
|
||||
# python implementation of http://bjoern.hoehrmann.de/utf-8/decoder/dfa/
|
||||
|
||||
_UTF8_ACCEPT = 0
|
||||
_UTF8_REJECT = 12
|
||||
|
||||
_UTF8D = [
|
||||
# The first part of the table maps bytes to character classes that
|
||||
# to reduce the size of the transition table and create bitmasks.
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
1,
|
||||
1,
|
||||
1,
|
||||
1,
|
||||
1,
|
||||
1,
|
||||
1,
|
||||
1,
|
||||
1,
|
||||
1,
|
||||
1,
|
||||
1,
|
||||
1,
|
||||
1,
|
||||
1,
|
||||
1,
|
||||
9,
|
||||
9,
|
||||
9,
|
||||
9,
|
||||
9,
|
||||
9,
|
||||
9,
|
||||
9,
|
||||
9,
|
||||
9,
|
||||
9,
|
||||
9,
|
||||
9,
|
||||
9,
|
||||
9,
|
||||
9,
|
||||
7,
|
||||
7,
|
||||
7,
|
||||
7,
|
||||
7,
|
||||
7,
|
||||
7,
|
||||
7,
|
||||
7,
|
||||
7,
|
||||
7,
|
||||
7,
|
||||
7,
|
||||
7,
|
||||
7,
|
||||
7,
|
||||
7,
|
||||
7,
|
||||
7,
|
||||
7,
|
||||
7,
|
||||
7,
|
||||
7,
|
||||
7,
|
||||
7,
|
||||
7,
|
||||
7,
|
||||
7,
|
||||
7,
|
||||
7,
|
||||
7,
|
||||
7,
|
||||
8,
|
||||
8,
|
||||
2,
|
||||
2,
|
||||
2,
|
||||
2,
|
||||
2,
|
||||
2,
|
||||
2,
|
||||
2,
|
||||
2,
|
||||
2,
|
||||
2,
|
||||
2,
|
||||
2,
|
||||
2,
|
||||
2,
|
||||
2,
|
||||
2,
|
||||
2,
|
||||
2,
|
||||
2,
|
||||
2,
|
||||
2,
|
||||
2,
|
||||
2,
|
||||
2,
|
||||
2,
|
||||
2,
|
||||
2,
|
||||
2,
|
||||
2,
|
||||
10,
|
||||
3,
|
||||
3,
|
||||
3,
|
||||
3,
|
||||
3,
|
||||
3,
|
||||
3,
|
||||
3,
|
||||
3,
|
||||
3,
|
||||
3,
|
||||
3,
|
||||
4,
|
||||
3,
|
||||
3,
|
||||
11,
|
||||
6,
|
||||
6,
|
||||
6,
|
||||
5,
|
||||
8,
|
||||
8,
|
||||
8,
|
||||
8,
|
||||
8,
|
||||
8,
|
||||
8,
|
||||
8,
|
||||
8,
|
||||
8,
|
||||
8,
|
||||
# The second part is a transition table that maps a combination
|
||||
# of a state of the automaton and a character class to a state.
|
||||
0,
|
||||
12,
|
||||
24,
|
||||
36,
|
||||
60,
|
||||
96,
|
||||
84,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
48,
|
||||
72,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
0,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
0,
|
||||
12,
|
||||
0,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
24,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
24,
|
||||
12,
|
||||
24,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
24,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
24,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
24,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
36,
|
||||
12,
|
||||
36,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
36,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
36,
|
||||
12,
|
||||
36,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
36,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
12,
|
||||
]
|
||||
|
||||
def _decode(state: int, codep: int, ch: int) -> tuple:
|
||||
tp = _UTF8D[ch]
|
||||
|
||||
codep = (
|
||||
(ch & 0x3F) | (codep << 6) if (state != _UTF8_ACCEPT) else (0xFF >> tp) & ch
|
||||
)
|
||||
state = _UTF8D[256 + state + tp]
|
||||
|
||||
return state, codep
|
||||
|
||||
def _validate_utf8(utfbytes: Union[str, bytes]) -> bool:
|
||||
state = _UTF8_ACCEPT
|
||||
codep = 0
|
||||
for i in utfbytes:
|
||||
state, codep = _decode(state, codep, int(i))
|
||||
if state == _UTF8_REJECT:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def validate_utf8(utfbytes: Union[str, bytes]) -> bool:
|
||||
"""
|
||||
validate utf8 byte string.
|
||||
utfbytes: utf byte string to check.
|
||||
return value: if valid utf8 string, return true. Otherwise, return false.
|
||||
"""
|
||||
return _validate_utf8(utfbytes)
|
||||
|
||||
|
||||
def extract_err_message(exception: Exception) -> Optional[str]:
|
||||
if exception.args:
|
||||
exception_message: str = exception.args[0]
|
||||
return exception_message
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
def extract_error_code(exception: Exception) -> Optional[int]:
|
||||
if exception.args and len(exception.args) > 1:
|
||||
return exception.args[0] if isinstance(exception.args[0], int) else None
|
||||
return None
|
||||
@@ -0,0 +1,244 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
"""
|
||||
_wsdump.py
|
||||
websocket - WebSocket client library for Python
|
||||
|
||||
Copyright 2025 engn33r
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import code
|
||||
import gzip
|
||||
import ssl
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
import zlib
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import websocket
|
||||
|
||||
try:
|
||||
import readline # noqa: F401
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
|
||||
def get_encoding() -> str:
|
||||
encoding = getattr(sys.stdin, "encoding", "")
|
||||
if not encoding:
|
||||
return "utf-8"
|
||||
else:
|
||||
return encoding.lower()
|
||||
|
||||
|
||||
OPCODE_DATA = (websocket.ABNF.OPCODE_TEXT, websocket.ABNF.OPCODE_BINARY)
|
||||
ENCODING = get_encoding()
|
||||
|
||||
|
||||
class VAction(argparse.Action):
|
||||
def __call__(
|
||||
self,
|
||||
parser: argparse.Namespace,
|
||||
args: tuple,
|
||||
values: str,
|
||||
option_string: str = None,
|
||||
) -> None:
|
||||
if values is None:
|
||||
values = "1"
|
||||
try:
|
||||
values = int(values)
|
||||
except ValueError:
|
||||
values = values.count("v") + 1
|
||||
setattr(args, self.dest, values)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="WebSocket Simple Dump Tool")
|
||||
parser.add_argument(
|
||||
"url", metavar="ws_url", help="websocket url. ex. ws://echo.websocket.events/"
|
||||
)
|
||||
parser.add_argument("-p", "--proxy", help="proxy url. ex. http://127.0.0.1:8080")
|
||||
parser.add_argument(
|
||||
"-v",
|
||||
"--verbose",
|
||||
default=0,
|
||||
nargs="?",
|
||||
action=VAction,
|
||||
dest="verbose",
|
||||
help="set verbose mode. If set to 1, show opcode. "
|
||||
"If set to 2, enable to trace websocket module",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-n", "--nocert", action="store_true", help="Ignore invalid SSL cert"
|
||||
)
|
||||
parser.add_argument("-r", "--raw", action="store_true", help="raw output")
|
||||
parser.add_argument("-s", "--subprotocols", nargs="*", help="Set subprotocols")
|
||||
parser.add_argument("-o", "--origin", help="Set origin")
|
||||
parser.add_argument(
|
||||
"--eof-wait",
|
||||
default=0,
|
||||
type=int,
|
||||
help="wait time(second) after 'EOF' received.",
|
||||
)
|
||||
parser.add_argument("-t", "--text", help="Send initial text")
|
||||
parser.add_argument(
|
||||
"--timings", action="store_true", help="Print timings in seconds"
|
||||
)
|
||||
parser.add_argument("--headers", help="Set custom headers. Use ',' as separator")
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
class RawInput:
|
||||
def raw_input(self, prompt: str = "") -> str:
|
||||
line = input(prompt)
|
||||
|
||||
if ENCODING and ENCODING != "utf-8" and not isinstance(line, str):
|
||||
line = line.decode(ENCODING).encode("utf-8")
|
||||
elif isinstance(line, str):
|
||||
line = line.encode("utf-8")
|
||||
|
||||
return line
|
||||
|
||||
|
||||
class InteractiveConsole(RawInput, code.InteractiveConsole):
|
||||
def write(self, data: str) -> None:
|
||||
sys.stdout.write("\033[2K\033[E")
|
||||
# sys.stdout.write("\n")
|
||||
sys.stdout.write("\033[34m< " + data + "\033[39m")
|
||||
sys.stdout.write("\n> ")
|
||||
sys.stdout.flush()
|
||||
|
||||
def read(self) -> str:
|
||||
return self.raw_input("> ")
|
||||
|
||||
|
||||
class NonInteractive(RawInput):
|
||||
def write(self, data: str) -> None:
|
||||
sys.stdout.write(data)
|
||||
sys.stdout.write("\n")
|
||||
sys.stdout.flush()
|
||||
|
||||
def read(self) -> str:
|
||||
return self.raw_input("")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
start_time = time.time()
|
||||
args = parse_args()
|
||||
if args.verbose > 1:
|
||||
websocket.enableTrace(True)
|
||||
options = {}
|
||||
if args.proxy:
|
||||
p = urlparse(args.proxy)
|
||||
options["http_proxy_host"] = p.hostname
|
||||
options["http_proxy_port"] = p.port
|
||||
if args.origin:
|
||||
options["origin"] = args.origin
|
||||
if args.subprotocols:
|
||||
options["subprotocols"] = args.subprotocols
|
||||
opts = {}
|
||||
if args.nocert:
|
||||
opts = {"cert_reqs": ssl.CERT_NONE, "check_hostname": False}
|
||||
if args.headers:
|
||||
options["header"] = list(map(str.strip, args.headers.split(",")))
|
||||
ws = websocket.create_connection(args.url, sslopt=opts, **options)
|
||||
if args.raw:
|
||||
console = NonInteractive()
|
||||
else:
|
||||
console = InteractiveConsole()
|
||||
print("Press Ctrl+C to quit")
|
||||
|
||||
def recv() -> tuple:
|
||||
try:
|
||||
frame = ws.recv_frame()
|
||||
except websocket.WebSocketException:
|
||||
return websocket.ABNF.OPCODE_CLOSE, ""
|
||||
if not frame:
|
||||
raise websocket.WebSocketException(f"Not a valid frame {frame}")
|
||||
elif frame.opcode in OPCODE_DATA:
|
||||
return frame.opcode, frame.data
|
||||
elif frame.opcode == websocket.ABNF.OPCODE_CLOSE:
|
||||
ws.send_close()
|
||||
return frame.opcode, ""
|
||||
elif frame.opcode == websocket.ABNF.OPCODE_PING:
|
||||
ws.pong(frame.data)
|
||||
return frame.opcode, frame.data
|
||||
|
||||
return frame.opcode, frame.data
|
||||
|
||||
def recv_ws() -> None:
|
||||
while True:
|
||||
opcode, data = recv()
|
||||
msg = None
|
||||
if opcode == websocket.ABNF.OPCODE_TEXT and isinstance(data, bytes):
|
||||
data = str(data, "utf-8")
|
||||
if (
|
||||
isinstance(data, bytes) and len(data) > 2 and data[:2] == b"\037\213"
|
||||
): # gzip magick
|
||||
try:
|
||||
data = "[gzip] " + str(gzip.decompress(data), "utf-8")
|
||||
except:
|
||||
pass
|
||||
elif isinstance(data, bytes):
|
||||
try:
|
||||
data = "[zlib] " + str(
|
||||
zlib.decompress(data, -zlib.MAX_WBITS), "utf-8"
|
||||
)
|
||||
except:
|
||||
pass
|
||||
|
||||
if isinstance(data, bytes):
|
||||
data = repr(data)
|
||||
|
||||
if args.verbose:
|
||||
msg = f"{websocket.ABNF.OPCODE_MAP.get(opcode)}: {data}"
|
||||
else:
|
||||
msg = data
|
||||
|
||||
if msg is not None:
|
||||
if args.timings:
|
||||
console.write(f"{time.time() - start_time}: {msg}")
|
||||
else:
|
||||
console.write(msg)
|
||||
|
||||
if opcode == websocket.ABNF.OPCODE_CLOSE:
|
||||
break
|
||||
|
||||
thread = threading.Thread(target=recv_ws)
|
||||
thread.daemon = True
|
||||
thread.start()
|
||||
|
||||
if args.text:
|
||||
ws.send(args.text)
|
||||
|
||||
while True:
|
||||
try:
|
||||
message = console.read()
|
||||
ws.send(message)
|
||||
except KeyboardInterrupt:
|
||||
return
|
||||
except EOFError:
|
||||
time.sleep(args.eof_wait)
|
||||
return
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
main()
|
||||
except Exception as e:
|
||||
print(e)
|
||||
@@ -0,0 +1,6 @@
|
||||
HTTP/1.1 101 WebSocket Protocol Handshake
|
||||
Connection: Upgrade
|
||||
Upgrade: WebSocket
|
||||
Sec-WebSocket-Accept: Kxep+hNu9n51529fGidYu7a3wO0=
|
||||
some_header: something
|
||||
|
||||
@@ -0,0 +1,6 @@
|
||||
HTTP/1.1 101 WebSocket Protocol Handshake
|
||||
Connection: Upgrade
|
||||
Upgrade WebSocket
|
||||
Sec-WebSocket-Accept: Kxep+hNu9n51529fGidYu7a3wO0=
|
||||
some_header: something
|
||||
|
||||
@@ -0,0 +1,8 @@
|
||||
HTTP/1.1 101 WebSocket Protocol Handshake
|
||||
Connection: Upgrade, Keep-Alive
|
||||
Upgrade: WebSocket
|
||||
Sec-WebSocket-Accept: Kxep+hNu9n51529fGidYu7a3wO0=
|
||||
Set-Cookie: Token=ABCDE
|
||||
Set-Cookie: Token=FGHIJ
|
||||
some_header: something
|
||||
|
||||
@@ -0,0 +1,23 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# From https://github.com/aaugustin/websockets/blob/main/example/echo.py
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
import websockets
|
||||
|
||||
LOCAL_WS_SERVER_PORT = int(os.environ.get("LOCAL_WS_SERVER_PORT", "8765"))
|
||||
|
||||
|
||||
async def echo(websocket):
|
||||
async for message in websocket:
|
||||
await websocket.send(message)
|
||||
|
||||
|
||||
async def main():
|
||||
async with websockets.serve(echo, "localhost", LOCAL_WS_SERVER_PORT):
|
||||
await asyncio.Future() # run forever
|
||||
|
||||
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,125 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
#
|
||||
import unittest
|
||||
|
||||
from websocket._abnf import ABNF, frame_buffer
|
||||
from websocket._exceptions import WebSocketProtocolException
|
||||
|
||||
"""
|
||||
test_abnf.py
|
||||
websocket - WebSocket client library for Python
|
||||
|
||||
Copyright 2025 engn33r
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
|
||||
class ABNFTest(unittest.TestCase):
|
||||
def test_init(self):
|
||||
a = ABNF(0, 0, 0, 0, opcode=ABNF.OPCODE_PING)
|
||||
self.assertEqual(a.fin, 0)
|
||||
self.assertEqual(a.rsv1, 0)
|
||||
self.assertEqual(a.rsv2, 0)
|
||||
self.assertEqual(a.rsv3, 0)
|
||||
self.assertEqual(a.opcode, 9)
|
||||
self.assertEqual(a.data, "")
|
||||
a_bad = ABNF(0, 1, 0, 0, opcode=77)
|
||||
self.assertEqual(a_bad.rsv1, 1)
|
||||
self.assertEqual(a_bad.opcode, 77)
|
||||
|
||||
def test_validate(self):
|
||||
a_invalid_ping = ABNF(0, 0, 0, 0, opcode=ABNF.OPCODE_PING)
|
||||
self.assertRaises(
|
||||
WebSocketProtocolException,
|
||||
a_invalid_ping.validate,
|
||||
skip_utf8_validation=False,
|
||||
)
|
||||
a_bad_rsv_value = ABNF(0, 1, 0, 0, opcode=ABNF.OPCODE_TEXT)
|
||||
self.assertRaises(
|
||||
WebSocketProtocolException,
|
||||
a_bad_rsv_value.validate,
|
||||
skip_utf8_validation=False,
|
||||
)
|
||||
a_bad_opcode = ABNF(0, 0, 0, 0, opcode=77)
|
||||
self.assertRaises(
|
||||
WebSocketProtocolException,
|
||||
a_bad_opcode.validate,
|
||||
skip_utf8_validation=False,
|
||||
)
|
||||
a_bad_close_frame = ABNF(0, 0, 0, 0, opcode=ABNF.OPCODE_CLOSE, data=b"\x01")
|
||||
self.assertRaises(
|
||||
WebSocketProtocolException,
|
||||
a_bad_close_frame.validate,
|
||||
skip_utf8_validation=False,
|
||||
)
|
||||
a_bad_close_frame_2 = ABNF(
|
||||
0, 0, 0, 0, opcode=ABNF.OPCODE_CLOSE, data=b"\x01\x8a\xaa\xff\xdd"
|
||||
)
|
||||
self.assertRaises(
|
||||
WebSocketProtocolException,
|
||||
a_bad_close_frame_2.validate,
|
||||
skip_utf8_validation=False,
|
||||
)
|
||||
a_bad_close_frame_3 = ABNF(
|
||||
0, 0, 0, 0, opcode=ABNF.OPCODE_CLOSE, data=b"\x03\xe7"
|
||||
)
|
||||
self.assertRaises(
|
||||
WebSocketProtocolException,
|
||||
a_bad_close_frame_3.validate,
|
||||
skip_utf8_validation=True,
|
||||
)
|
||||
|
||||
def test_mask(self):
|
||||
abnf_none_data = ABNF(
|
||||
0, 0, 0, 0, opcode=ABNF.OPCODE_PING, mask_value=1, data=None
|
||||
)
|
||||
bytes_val = b"aaaa"
|
||||
self.assertEqual(abnf_none_data._get_masked(bytes_val), bytes_val)
|
||||
abnf_str_data = ABNF(
|
||||
0, 0, 0, 0, opcode=ABNF.OPCODE_PING, mask_value=1, data="a"
|
||||
)
|
||||
self.assertEqual(abnf_str_data._get_masked(bytes_val), b"aaaa\x00")
|
||||
|
||||
def test_format(self):
|
||||
abnf_bad_rsv_bits = ABNF(2, 0, 0, 0, opcode=ABNF.OPCODE_TEXT)
|
||||
self.assertRaises(ValueError, abnf_bad_rsv_bits.format)
|
||||
abnf_bad_opcode = ABNF(0, 0, 0, 0, opcode=5)
|
||||
self.assertRaises(ValueError, abnf_bad_opcode.format)
|
||||
abnf_length_10 = ABNF(0, 0, 0, 0, opcode=ABNF.OPCODE_TEXT, data="abcdefghij")
|
||||
self.assertEqual(b"\x01", abnf_length_10.format()[0].to_bytes(1, "big"))
|
||||
self.assertEqual(b"\x8a", abnf_length_10.format()[1].to_bytes(1, "big"))
|
||||
self.assertEqual("fin=0 opcode=1 data=abcdefghij", abnf_length_10.__str__())
|
||||
abnf_length_20 = ABNF(
|
||||
0, 0, 0, 0, opcode=ABNF.OPCODE_BINARY, data="abcdefghijabcdefghij"
|
||||
)
|
||||
self.assertEqual(b"\x02", abnf_length_20.format()[0].to_bytes(1, "big"))
|
||||
self.assertEqual(b"\x94", abnf_length_20.format()[1].to_bytes(1, "big"))
|
||||
abnf_no_mask = ABNF(
|
||||
0, 0, 0, 0, opcode=ABNF.OPCODE_TEXT, mask_value=0, data=b"\x01\x8a\xcc"
|
||||
)
|
||||
self.assertEqual(b"\x01\x03\x01\x8a\xcc", abnf_no_mask.format())
|
||||
|
||||
def test_frame_buffer(self):
|
||||
fb = frame_buffer(0, True)
|
||||
self.assertEqual(fb.recv, 0)
|
||||
self.assertEqual(fb.skip_utf8_validation, True)
|
||||
fb.clear
|
||||
self.assertEqual(fb.header, None)
|
||||
self.assertEqual(fb.length, None)
|
||||
self.assertEqual(fb.mask_value, None)
|
||||
self.assertEqual(fb.has_mask(), False)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,395 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
#
|
||||
import os
|
||||
import os.path
|
||||
import ssl
|
||||
import threading
|
||||
import unittest
|
||||
|
||||
import websocket as ws
|
||||
|
||||
"""
|
||||
test_app.py
|
||||
websocket - WebSocket client library for Python
|
||||
|
||||
Copyright 2025 engn33r
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
# Skip test to access the internet unless TEST_WITH_INTERNET == 1
|
||||
TEST_WITH_INTERNET = os.environ.get("TEST_WITH_INTERNET", "0") == "1"
|
||||
# Skip tests relying on local websockets server unless LOCAL_WS_SERVER_PORT != -1
|
||||
LOCAL_WS_SERVER_PORT = os.environ.get("LOCAL_WS_SERVER_PORT", "-1")
|
||||
TEST_WITH_LOCAL_SERVER = LOCAL_WS_SERVER_PORT != "-1"
|
||||
TRACEABLE = True
|
||||
|
||||
|
||||
class WebSocketAppTest(unittest.TestCase):
|
||||
class NotSetYet:
|
||||
"""A marker class for signalling that a value hasn't been set yet."""
|
||||
|
||||
def setUp(self):
|
||||
ws.enableTrace(TRACEABLE)
|
||||
|
||||
WebSocketAppTest.keep_running_open = WebSocketAppTest.NotSetYet()
|
||||
WebSocketAppTest.keep_running_close = WebSocketAppTest.NotSetYet()
|
||||
WebSocketAppTest.get_mask_key_id = WebSocketAppTest.NotSetYet()
|
||||
WebSocketAppTest.on_error_data = WebSocketAppTest.NotSetYet()
|
||||
|
||||
def tearDown(self):
|
||||
WebSocketAppTest.keep_running_open = WebSocketAppTest.NotSetYet()
|
||||
WebSocketAppTest.keep_running_close = WebSocketAppTest.NotSetYet()
|
||||
WebSocketAppTest.get_mask_key_id = WebSocketAppTest.NotSetYet()
|
||||
WebSocketAppTest.on_error_data = WebSocketAppTest.NotSetYet()
|
||||
|
||||
def close(self):
|
||||
pass
|
||||
|
||||
@unittest.skipUnless(
|
||||
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
|
||||
)
|
||||
def test_keep_running(self):
|
||||
"""A WebSocketApp should keep running as long as its self.keep_running
|
||||
is not False (in the boolean context).
|
||||
"""
|
||||
|
||||
def on_open(self, *args, **kwargs):
|
||||
"""Set the keep_running flag for later inspection and immediately
|
||||
close the connection.
|
||||
"""
|
||||
self.send("hello!")
|
||||
WebSocketAppTest.keep_running_open = self.keep_running
|
||||
self.keep_running = False
|
||||
|
||||
def on_message(_, message):
|
||||
print(message)
|
||||
self.close()
|
||||
|
||||
def on_close(self, *args, **kwargs):
|
||||
"""Set the keep_running flag for the test to use."""
|
||||
WebSocketAppTest.keep_running_close = self.keep_running
|
||||
|
||||
app = ws.WebSocketApp(
|
||||
f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}",
|
||||
on_open=on_open,
|
||||
on_close=on_close,
|
||||
on_message=on_message,
|
||||
)
|
||||
app.run_forever()
|
||||
|
||||
# @unittest.skipUnless(TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled")
|
||||
@unittest.skipUnless(False, "Test disabled for now (requires rel)")
|
||||
def test_run_forever_dispatcher(self):
|
||||
"""A WebSocketApp should keep running as long as its self.keep_running
|
||||
is not False (in the boolean context).
|
||||
"""
|
||||
|
||||
def on_open(self, *args, **kwargs):
|
||||
"""Send a message, receive, and send one more"""
|
||||
self.send("hello!")
|
||||
self.recv()
|
||||
self.send("goodbye!")
|
||||
|
||||
def on_message(_, message):
|
||||
print(message)
|
||||
self.close()
|
||||
|
||||
app = ws.WebSocketApp(
|
||||
f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}",
|
||||
on_open=on_open,
|
||||
on_message=on_message,
|
||||
)
|
||||
app.run_forever(dispatcher="Dispatcher") # doesn't work
|
||||
|
||||
# app.run_forever(dispatcher=rel) # would work
|
||||
# rel.dispatch()
|
||||
|
||||
@unittest.skipUnless(
|
||||
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
|
||||
)
|
||||
def test_run_forever_teardown_clean_exit(self):
|
||||
"""The WebSocketApp.run_forever() method should return `False` when the application ends gracefully."""
|
||||
app = ws.WebSocketApp(f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}")
|
||||
threading.Timer(interval=0.2, function=app.close).start()
|
||||
teardown = app.run_forever()
|
||||
self.assertEqual(teardown, False)
|
||||
|
||||
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
|
||||
def test_sock_mask_key(self):
|
||||
"""A WebSocketApp should forward the received mask_key function down
|
||||
to the actual socket.
|
||||
"""
|
||||
|
||||
def my_mask_key_func():
|
||||
return "\x00\x00\x00\x00"
|
||||
|
||||
app = ws.WebSocketApp(
|
||||
"wss://api-pub.bitfinex.com/ws/1", get_mask_key=my_mask_key_func
|
||||
)
|
||||
|
||||
# if numpy is installed, this assertion fail
|
||||
# Note: We can't use 'is' for comparing the functions directly, need to use 'id'.
|
||||
self.assertEqual(id(app.get_mask_key), id(my_mask_key_func))
|
||||
|
||||
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
|
||||
def test_invalid_ping_interval_ping_timeout(self):
|
||||
"""Test exception handling if ping_interval < ping_timeout"""
|
||||
|
||||
def on_ping(app, _):
|
||||
print("Got a ping!")
|
||||
app.close()
|
||||
|
||||
def on_pong(app, _):
|
||||
print("Got a pong! No need to respond")
|
||||
app.close()
|
||||
|
||||
app = ws.WebSocketApp(
|
||||
"wss://api-pub.bitfinex.com/ws/1", on_ping=on_ping, on_pong=on_pong
|
||||
)
|
||||
self.assertRaises(
|
||||
ws.WebSocketException,
|
||||
app.run_forever,
|
||||
ping_interval=1,
|
||||
ping_timeout=2,
|
||||
sslopt={"cert_reqs": ssl.CERT_NONE},
|
||||
)
|
||||
|
||||
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
|
||||
def test_ping_interval(self):
|
||||
"""Test WebSocketApp proper ping functionality"""
|
||||
|
||||
def on_ping(app, _):
|
||||
print("Got a ping!")
|
||||
app.close()
|
||||
|
||||
def on_pong(app, _):
|
||||
print("Got a pong! No need to respond")
|
||||
app.close()
|
||||
|
||||
app = ws.WebSocketApp(
|
||||
"wss://api-pub.bitfinex.com/ws/1", on_ping=on_ping, on_pong=on_pong
|
||||
)
|
||||
app.run_forever(
|
||||
ping_interval=2, ping_timeout=1, sslopt={"cert_reqs": ssl.CERT_NONE}
|
||||
)
|
||||
|
||||
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
|
||||
def test_opcode_close(self):
|
||||
"""Test WebSocketApp close opcode"""
|
||||
|
||||
app = ws.WebSocketApp("wss://tsock.us1.twilio.com/v3/wsconnect")
|
||||
app.run_forever(ping_interval=2, ping_timeout=1, ping_payload="Ping payload")
|
||||
|
||||
# This is commented out because the URL no longer responds in the expected way
|
||||
# @unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
|
||||
# def testOpcodeBinary(self):
|
||||
# """ Test WebSocketApp binary opcode
|
||||
# """
|
||||
# app = ws.WebSocketApp('wss://streaming.vn.teslamotors.com/streaming/')
|
||||
# app.run_forever(ping_interval=2, ping_timeout=1, ping_payload="Ping payload")
|
||||
|
||||
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
|
||||
def test_bad_ping_interval(self):
|
||||
"""A WebSocketApp handling of negative ping_interval"""
|
||||
app = ws.WebSocketApp("wss://api-pub.bitfinex.com/ws/1")
|
||||
self.assertRaises(
|
||||
ws.WebSocketException,
|
||||
app.run_forever,
|
||||
ping_interval=-5,
|
||||
sslopt={"cert_reqs": ssl.CERT_NONE},
|
||||
)
|
||||
|
||||
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
|
||||
def test_bad_ping_timeout(self):
|
||||
"""A WebSocketApp handling of negative ping_timeout"""
|
||||
app = ws.WebSocketApp("wss://api-pub.bitfinex.com/ws/1")
|
||||
self.assertRaises(
|
||||
ws.WebSocketException,
|
||||
app.run_forever,
|
||||
ping_timeout=-3,
|
||||
sslopt={"cert_reqs": ssl.CERT_NONE},
|
||||
)
|
||||
|
||||
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
|
||||
def test_close_status_code(self):
|
||||
"""Test extraction of close frame status code and close reason in WebSocketApp"""
|
||||
|
||||
def on_close(wsapp, close_status_code, close_msg):
|
||||
print("on_close reached")
|
||||
|
||||
app = ws.WebSocketApp(
|
||||
"wss://tsock.us1.twilio.com/v3/wsconnect", on_close=on_close
|
||||
)
|
||||
closeframe = ws.ABNF(
|
||||
opcode=ws.ABNF.OPCODE_CLOSE, data=b"\x03\xe8no-init-from-client"
|
||||
)
|
||||
self.assertEqual([1000, "no-init-from-client"], app._get_close_args(closeframe))
|
||||
|
||||
closeframe = ws.ABNF(opcode=ws.ABNF.OPCODE_CLOSE, data=b"")
|
||||
self.assertEqual([None, None], app._get_close_args(closeframe))
|
||||
|
||||
app2 = ws.WebSocketApp("wss://tsock.us1.twilio.com/v3/wsconnect")
|
||||
closeframe = ws.ABNF(opcode=ws.ABNF.OPCODE_CLOSE, data=b"")
|
||||
self.assertEqual([None, None], app2._get_close_args(closeframe))
|
||||
|
||||
self.assertRaises(
|
||||
ws.WebSocketConnectionClosedException,
|
||||
app.send,
|
||||
data="test if connection is closed",
|
||||
)
|
||||
|
||||
@unittest.skipUnless(
|
||||
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
|
||||
)
|
||||
def test_callback_function_exception(self):
|
||||
"""Test callback function exception handling"""
|
||||
|
||||
exc = None
|
||||
passed_app = None
|
||||
|
||||
def on_open(app):
|
||||
raise RuntimeError("Callback failed")
|
||||
|
||||
def on_error(app, err):
|
||||
nonlocal passed_app
|
||||
passed_app = app
|
||||
nonlocal exc
|
||||
exc = err
|
||||
|
||||
def on_pong(app, _):
|
||||
app.close()
|
||||
|
||||
app = ws.WebSocketApp(
|
||||
f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}",
|
||||
on_open=on_open,
|
||||
on_error=on_error,
|
||||
on_pong=on_pong,
|
||||
)
|
||||
app.run_forever(ping_interval=2, ping_timeout=1)
|
||||
|
||||
self.assertEqual(passed_app, app)
|
||||
self.assertIsInstance(exc, RuntimeError)
|
||||
self.assertEqual(str(exc), "Callback failed")
|
||||
|
||||
@unittest.skipUnless(
|
||||
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
|
||||
)
|
||||
def test_callback_method_exception(self):
|
||||
"""Test callback method exception handling"""
|
||||
|
||||
class Callbacks:
|
||||
def __init__(self):
|
||||
self.exc = None
|
||||
self.passed_app = None
|
||||
self.app = ws.WebSocketApp(
|
||||
f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}",
|
||||
on_open=self.on_open,
|
||||
on_error=self.on_error,
|
||||
on_pong=self.on_pong,
|
||||
)
|
||||
self.app.run_forever(ping_interval=2, ping_timeout=1)
|
||||
|
||||
def on_open(self, _):
|
||||
raise RuntimeError("Callback failed")
|
||||
|
||||
def on_error(self, app, err):
|
||||
self.passed_app = app
|
||||
self.exc = err
|
||||
|
||||
def on_pong(self, app, _):
|
||||
app.close()
|
||||
|
||||
callbacks = Callbacks()
|
||||
|
||||
self.assertEqual(callbacks.passed_app, callbacks.app)
|
||||
self.assertIsInstance(callbacks.exc, RuntimeError)
|
||||
self.assertEqual(str(callbacks.exc), "Callback failed")
|
||||
|
||||
@unittest.skipUnless(
|
||||
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
|
||||
)
|
||||
def test_reconnect(self):
|
||||
"""Test reconnect"""
|
||||
pong_count = 0
|
||||
exc = None
|
||||
|
||||
def on_error(_, err):
|
||||
nonlocal exc
|
||||
exc = err
|
||||
|
||||
def on_pong(app, _):
|
||||
nonlocal pong_count
|
||||
pong_count += 1
|
||||
if pong_count == 1:
|
||||
# First pong, shutdown socket, enforce read error
|
||||
app.sock.shutdown()
|
||||
if pong_count >= 2:
|
||||
# Got second pong after reconnect
|
||||
app.close()
|
||||
|
||||
app = ws.WebSocketApp(
|
||||
f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}", on_pong=on_pong, on_error=on_error
|
||||
)
|
||||
app.run_forever(ping_interval=2, ping_timeout=1, reconnect=3)
|
||||
|
||||
self.assertEqual(pong_count, 2)
|
||||
self.assertIsInstance(exc, ws.WebSocketTimeoutException)
|
||||
self.assertEqual(str(exc), "ping/pong timed out")
|
||||
|
||||
def test_dispatcher_selection_default(self):
|
||||
"""Test default dispatcher selection"""
|
||||
app = ws.WebSocketApp("ws://example.com")
|
||||
|
||||
# Test default dispatcher (non-SSL)
|
||||
dispatcher = app.create_dispatcher(ping_timeout=10, is_ssl=False)
|
||||
self.assertIsInstance(dispatcher, ws._dispatcher.Dispatcher)
|
||||
|
||||
def test_dispatcher_selection_ssl(self):
|
||||
"""Test SSL dispatcher selection"""
|
||||
app = ws.WebSocketApp("wss://example.com")
|
||||
|
||||
# Test SSL dispatcher
|
||||
dispatcher = app.create_dispatcher(ping_timeout=10, is_ssl=True)
|
||||
self.assertIsInstance(dispatcher, ws._dispatcher.SSLDispatcher)
|
||||
|
||||
def test_dispatcher_selection_custom(self):
|
||||
"""Test custom dispatcher selection"""
|
||||
from unittest.mock import Mock
|
||||
|
||||
app = ws.WebSocketApp("ws://example.com")
|
||||
custom_dispatcher = Mock()
|
||||
handle_disconnect = Mock()
|
||||
|
||||
# Test wrapped dispatcher with custom dispatcher
|
||||
dispatcher = app.create_dispatcher(
|
||||
ping_timeout=10,
|
||||
dispatcher=custom_dispatcher,
|
||||
handleDisconnect=handle_disconnect,
|
||||
)
|
||||
self.assertIsInstance(dispatcher, ws._dispatcher.WrappedDispatcher)
|
||||
self.assertEqual(dispatcher.dispatcher, custom_dispatcher)
|
||||
self.assertEqual(dispatcher.handleDisconnect, handle_disconnect)
|
||||
|
||||
def test_dispatcher_selection_no_ping_timeout(self):
|
||||
"""Test dispatcher selection without ping timeout"""
|
||||
app = ws.WebSocketApp("ws://example.com")
|
||||
|
||||
# Test with None ping_timeout (should default to 10)
|
||||
dispatcher = app.create_dispatcher(ping_timeout=None, is_ssl=False)
|
||||
self.assertIsInstance(dispatcher, ws._dispatcher.Dispatcher)
|
||||
self.assertEqual(dispatcher.ping_timeout, 10)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,123 @@
|
||||
import unittest
|
||||
|
||||
from websocket._cookiejar import SimpleCookieJar
|
||||
|
||||
"""
|
||||
test_cookiejar.py
|
||||
websocket - WebSocket client library for Python
|
||||
|
||||
Copyright 2025 engn33r
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
|
||||
class CookieJarTest(unittest.TestCase):
|
||||
def test_add(self):
|
||||
cookie_jar = SimpleCookieJar()
|
||||
cookie_jar.add("")
|
||||
self.assertFalse(
|
||||
cookie_jar.jar, "Cookie with no domain should not be added to the jar"
|
||||
)
|
||||
|
||||
cookie_jar = SimpleCookieJar()
|
||||
cookie_jar.add("a=b")
|
||||
self.assertFalse(
|
||||
cookie_jar.jar, "Cookie with no domain should not be added to the jar"
|
||||
)
|
||||
|
||||
cookie_jar = SimpleCookieJar()
|
||||
cookie_jar.add("a=b; domain=.abc")
|
||||
self.assertTrue(".abc" in cookie_jar.jar)
|
||||
|
||||
cookie_jar = SimpleCookieJar()
|
||||
cookie_jar.add("a=b; domain=abc")
|
||||
self.assertTrue(".abc" in cookie_jar.jar)
|
||||
self.assertTrue("abc" not in cookie_jar.jar)
|
||||
|
||||
cookie_jar = SimpleCookieJar()
|
||||
cookie_jar.add("a=b; c=d; domain=abc")
|
||||
self.assertEqual(cookie_jar.get("abc"), "a=b; c=d")
|
||||
self.assertEqual(cookie_jar.get(None), "")
|
||||
|
||||
cookie_jar = SimpleCookieJar()
|
||||
cookie_jar.add("a=b; c=d; domain=abc")
|
||||
cookie_jar.add("e=f; domain=abc")
|
||||
self.assertEqual(cookie_jar.get("abc"), "a=b; c=d; e=f")
|
||||
|
||||
cookie_jar = SimpleCookieJar()
|
||||
cookie_jar.add("a=b; c=d; domain=abc")
|
||||
cookie_jar.add("e=f; domain=.abc")
|
||||
self.assertEqual(cookie_jar.get("abc"), "a=b; c=d; e=f")
|
||||
|
||||
cookie_jar = SimpleCookieJar()
|
||||
cookie_jar.add("a=b; c=d; domain=abc")
|
||||
cookie_jar.add("e=f; domain=xyz")
|
||||
self.assertEqual(cookie_jar.get("abc"), "a=b; c=d")
|
||||
self.assertEqual(cookie_jar.get("xyz"), "e=f")
|
||||
self.assertEqual(cookie_jar.get("something"), "")
|
||||
|
||||
def test_set(self):
|
||||
cookie_jar = SimpleCookieJar()
|
||||
cookie_jar.set("a=b")
|
||||
self.assertFalse(
|
||||
cookie_jar.jar, "Cookie with no domain should not be added to the jar"
|
||||
)
|
||||
|
||||
cookie_jar = SimpleCookieJar()
|
||||
cookie_jar.set("a=b; domain=.abc")
|
||||
self.assertTrue(".abc" in cookie_jar.jar)
|
||||
|
||||
cookie_jar = SimpleCookieJar()
|
||||
cookie_jar.set("a=b; domain=abc")
|
||||
self.assertTrue(".abc" in cookie_jar.jar)
|
||||
self.assertTrue("abc" not in cookie_jar.jar)
|
||||
|
||||
cookie_jar = SimpleCookieJar()
|
||||
cookie_jar.set("a=b; c=d; domain=abc")
|
||||
self.assertEqual(cookie_jar.get("abc"), "a=b; c=d")
|
||||
|
||||
cookie_jar = SimpleCookieJar()
|
||||
cookie_jar.set("a=b; c=d; domain=abc")
|
||||
cookie_jar.set("e=f; domain=abc")
|
||||
self.assertEqual(cookie_jar.get("abc"), "e=f")
|
||||
|
||||
cookie_jar = SimpleCookieJar()
|
||||
cookie_jar.set("a=b; c=d; domain=abc")
|
||||
cookie_jar.set("e=f; domain=.abc")
|
||||
self.assertEqual(cookie_jar.get("abc"), "e=f")
|
||||
|
||||
cookie_jar = SimpleCookieJar()
|
||||
cookie_jar.set("a=b; c=d; domain=abc")
|
||||
cookie_jar.set("e=f; domain=xyz")
|
||||
self.assertEqual(cookie_jar.get("abc"), "a=b; c=d")
|
||||
self.assertEqual(cookie_jar.get("xyz"), "e=f")
|
||||
self.assertEqual(cookie_jar.get("something"), "")
|
||||
|
||||
def test_get(self):
|
||||
cookie_jar = SimpleCookieJar()
|
||||
cookie_jar.set("a=b; c=d; domain=abc.com")
|
||||
self.assertEqual(cookie_jar.get("abc.com"), "a=b; c=d")
|
||||
self.assertEqual(cookie_jar.get("x.abc.com"), "a=b; c=d")
|
||||
self.assertEqual(cookie_jar.get("abc.com.es"), "")
|
||||
self.assertEqual(cookie_jar.get("xabc.com"), "")
|
||||
|
||||
cookie_jar.set("a=b; c=d; domain=.abc.com")
|
||||
self.assertEqual(cookie_jar.get("abc.com"), "a=b; c=d")
|
||||
self.assertEqual(cookie_jar.get("x.abc.com"), "a=b; c=d")
|
||||
self.assertEqual(cookie_jar.get("abc.com.es"), "")
|
||||
self.assertEqual(cookie_jar.get("xabc.com"), "")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,385 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
import socket
|
||||
import unittest
|
||||
from unittest.mock import Mock, patch, MagicMock
|
||||
import threading
|
||||
import time
|
||||
|
||||
import websocket
|
||||
from websocket._dispatcher import (
|
||||
Dispatcher,
|
||||
DispatcherBase,
|
||||
SSLDispatcher,
|
||||
WrappedDispatcher,
|
||||
)
|
||||
|
||||
"""
|
||||
test_dispatcher.py
|
||||
websocket - WebSocket client library for Python
|
||||
|
||||
Copyright 2025 engn33r
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
class MockApp:
|
||||
"""Mock WebSocketApp for testing"""
|
||||
|
||||
def __init__(self):
|
||||
self.keep_running = True
|
||||
self.sock = Mock()
|
||||
self.sock.sock = Mock()
|
||||
|
||||
|
||||
class MockSocket:
|
||||
"""Mock socket for testing"""
|
||||
|
||||
def __init__(self):
|
||||
self.pending_return = False
|
||||
|
||||
def pending(self):
|
||||
return self.pending_return
|
||||
|
||||
|
||||
class MockDispatcher:
|
||||
"""Mock external dispatcher for WrappedDispatcher testing"""
|
||||
|
||||
def __init__(self):
|
||||
self.signal_calls = []
|
||||
self.abort_calls = []
|
||||
self.read_calls = []
|
||||
self.buffwrite_calls = []
|
||||
self.timeout_calls = []
|
||||
|
||||
def signal(self, sig, handler):
|
||||
self.signal_calls.append((sig, handler))
|
||||
|
||||
def abort(self):
|
||||
self.abort_calls.append(True)
|
||||
|
||||
def read(self, sock, callback):
|
||||
self.read_calls.append((sock, callback))
|
||||
|
||||
def buffwrite(self, sock, data, send_func, disconnect_handler):
|
||||
self.buffwrite_calls.append((sock, data, send_func, disconnect_handler))
|
||||
|
||||
def timeout(self, seconds, callback, *args):
|
||||
self.timeout_calls.append((seconds, callback, args))
|
||||
|
||||
|
||||
class DispatcherTest(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.app = MockApp()
|
||||
|
||||
def test_dispatcher_base_init(self):
|
||||
"""Test DispatcherBase initialization"""
|
||||
dispatcher = DispatcherBase(self.app, 30.0)
|
||||
|
||||
self.assertEqual(dispatcher.app, self.app)
|
||||
self.assertEqual(dispatcher.ping_timeout, 30.0)
|
||||
|
||||
def test_dispatcher_base_timeout(self):
|
||||
"""Test DispatcherBase timeout method"""
|
||||
dispatcher = DispatcherBase(self.app, 30.0)
|
||||
callback = Mock()
|
||||
|
||||
# Test with seconds=None (should call callback immediately)
|
||||
dispatcher.timeout(None, callback)
|
||||
callback.assert_called_once()
|
||||
|
||||
# Test with seconds > 0 (would sleep in real implementation)
|
||||
callback.reset_mock()
|
||||
start_time = time.time()
|
||||
dispatcher.timeout(0.1, callback)
|
||||
elapsed = time.time() - start_time
|
||||
|
||||
callback.assert_called_once()
|
||||
self.assertGreaterEqual(elapsed, 0.05) # Allow some tolerance
|
||||
|
||||
def test_dispatcher_base_reconnect(self):
|
||||
"""Test DispatcherBase reconnect method"""
|
||||
dispatcher = DispatcherBase(self.app, 30.0)
|
||||
reconnector = Mock()
|
||||
|
||||
# Test normal reconnect
|
||||
dispatcher.reconnect(1, reconnector)
|
||||
reconnector.assert_called_once_with(reconnecting=True)
|
||||
|
||||
# Test reconnect with KeyboardInterrupt
|
||||
reconnector.reset_mock()
|
||||
reconnector.side_effect = KeyboardInterrupt("User interrupted")
|
||||
|
||||
with self.assertRaises(KeyboardInterrupt):
|
||||
dispatcher.reconnect(1, reconnector)
|
||||
|
||||
def test_dispatcher_base_send(self):
|
||||
"""Test DispatcherBase send method"""
|
||||
dispatcher = DispatcherBase(self.app, 30.0)
|
||||
mock_sock = Mock()
|
||||
test_data = b"test data"
|
||||
|
||||
with patch("websocket._dispatcher.send") as mock_send:
|
||||
mock_send.return_value = len(test_data)
|
||||
result = dispatcher.send(mock_sock, test_data)
|
||||
|
||||
mock_send.assert_called_once_with(mock_sock, test_data)
|
||||
self.assertEqual(result, len(test_data))
|
||||
|
||||
def test_dispatcher_read(self):
|
||||
"""Test Dispatcher read method"""
|
||||
dispatcher = Dispatcher(self.app, 5.0)
|
||||
read_callback = Mock(return_value=True)
|
||||
check_callback = Mock()
|
||||
mock_sock = Mock()
|
||||
|
||||
# Mock the selector to control the loop
|
||||
with patch("selectors.DefaultSelector") as mock_selector_class:
|
||||
mock_selector = Mock()
|
||||
mock_selector_class.return_value = mock_selector
|
||||
|
||||
# Make select return immediately (timeout)
|
||||
mock_selector.select.return_value = []
|
||||
|
||||
# Stop after first iteration
|
||||
def side_effect(*args):
|
||||
self.app.keep_running = False
|
||||
return []
|
||||
|
||||
mock_selector.select.side_effect = side_effect
|
||||
|
||||
dispatcher.read(mock_sock, read_callback, check_callback)
|
||||
|
||||
# Verify selector was used correctly
|
||||
mock_selector.register.assert_called()
|
||||
mock_selector.select.assert_called_with(5.0)
|
||||
mock_selector.close.assert_called()
|
||||
check_callback.assert_called()
|
||||
|
||||
def test_dispatcher_read_with_data(self):
|
||||
"""Test Dispatcher read method when data is available"""
|
||||
dispatcher = Dispatcher(self.app, 5.0)
|
||||
read_callback = Mock(return_value=True)
|
||||
check_callback = Mock()
|
||||
mock_sock = Mock()
|
||||
|
||||
with patch("selectors.DefaultSelector") as mock_selector_class:
|
||||
mock_selector = Mock()
|
||||
mock_selector_class.return_value = mock_selector
|
||||
|
||||
# First call returns data, second call stops the loop
|
||||
call_count = 0
|
||||
|
||||
def select_side_effect(*args):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if call_count == 1:
|
||||
return [True] # Data available
|
||||
else:
|
||||
self.app.keep_running = False
|
||||
return []
|
||||
|
||||
mock_selector.select.side_effect = select_side_effect
|
||||
|
||||
dispatcher.read(mock_sock, read_callback, check_callback)
|
||||
|
||||
read_callback.assert_called()
|
||||
check_callback.assert_called()
|
||||
|
||||
def test_ssl_dispatcher_read(self):
|
||||
"""Test SSLDispatcher read method"""
|
||||
dispatcher = SSLDispatcher(self.app, 5.0)
|
||||
read_callback = Mock(return_value=True)
|
||||
check_callback = Mock()
|
||||
|
||||
# Mock socket with pending data
|
||||
mock_ssl_sock = MockSocket()
|
||||
self.app.sock.sock = mock_ssl_sock
|
||||
|
||||
with patch("selectors.DefaultSelector") as mock_selector_class:
|
||||
mock_selector = Mock()
|
||||
mock_selector_class.return_value = mock_selector
|
||||
mock_selector.select.return_value = []
|
||||
|
||||
# Stop after first iteration
|
||||
def side_effect(*args):
|
||||
self.app.keep_running = False
|
||||
return []
|
||||
|
||||
mock_selector.select.side_effect = side_effect
|
||||
|
||||
dispatcher.read(None, read_callback, check_callback)
|
||||
|
||||
mock_selector.register.assert_called()
|
||||
check_callback.assert_called()
|
||||
|
||||
def test_ssl_dispatcher_select_with_pending(self):
|
||||
"""Test SSLDispatcher select method with pending data"""
|
||||
dispatcher = SSLDispatcher(self.app, 5.0)
|
||||
mock_ssl_sock = MockSocket()
|
||||
mock_ssl_sock.pending_return = True
|
||||
self.app.sock.sock = mock_ssl_sock
|
||||
mock_selector = Mock()
|
||||
|
||||
result = dispatcher.select(None, mock_selector)
|
||||
|
||||
# When pending() returns True, should return [sock]
|
||||
self.assertEqual(result, [mock_ssl_sock])
|
||||
|
||||
def test_ssl_dispatcher_select_without_pending(self):
|
||||
"""Test SSLDispatcher select method without pending data"""
|
||||
dispatcher = SSLDispatcher(self.app, 5.0)
|
||||
mock_ssl_sock = MockSocket()
|
||||
mock_ssl_sock.pending_return = False
|
||||
self.app.sock.sock = mock_ssl_sock
|
||||
mock_selector = Mock()
|
||||
mock_selector.select.return_value = [(mock_ssl_sock, None)]
|
||||
|
||||
result = dispatcher.select(None, mock_selector)
|
||||
|
||||
# Should return the first element of first result tuple
|
||||
self.assertEqual(result, mock_ssl_sock)
|
||||
mock_selector.select.assert_called_with(5.0)
|
||||
|
||||
def test_ssl_dispatcher_select_no_results(self):
|
||||
"""Test SSLDispatcher select method with no results"""
|
||||
dispatcher = SSLDispatcher(self.app, 5.0)
|
||||
mock_ssl_sock = MockSocket()
|
||||
mock_ssl_sock.pending_return = False
|
||||
self.app.sock.sock = mock_ssl_sock
|
||||
mock_selector = Mock()
|
||||
mock_selector.select.return_value = []
|
||||
|
||||
result = dispatcher.select(None, mock_selector)
|
||||
|
||||
# Should return None when no results (function doesn't return anything when len(r) == 0)
|
||||
self.assertIsNone(result)
|
||||
|
||||
def test_wrapped_dispatcher_init(self):
|
||||
"""Test WrappedDispatcher initialization"""
|
||||
mock_dispatcher = MockDispatcher()
|
||||
handle_disconnect = Mock()
|
||||
|
||||
wrapped = WrappedDispatcher(self.app, 10.0, mock_dispatcher, handle_disconnect)
|
||||
|
||||
self.assertEqual(wrapped.app, self.app)
|
||||
self.assertEqual(wrapped.ping_timeout, 10.0)
|
||||
self.assertEqual(wrapped.dispatcher, mock_dispatcher)
|
||||
self.assertEqual(wrapped.handleDisconnect, handle_disconnect)
|
||||
|
||||
# Should have set up signal handler
|
||||
self.assertEqual(len(mock_dispatcher.signal_calls), 1)
|
||||
sig, handler = mock_dispatcher.signal_calls[0]
|
||||
self.assertEqual(sig, 2) # SIGINT
|
||||
self.assertEqual(handler, mock_dispatcher.abort)
|
||||
|
||||
def test_wrapped_dispatcher_read(self):
|
||||
"""Test WrappedDispatcher read method"""
|
||||
mock_dispatcher = MockDispatcher()
|
||||
handle_disconnect = Mock()
|
||||
wrapped = WrappedDispatcher(self.app, 10.0, mock_dispatcher, handle_disconnect)
|
||||
|
||||
mock_sock = Mock()
|
||||
read_callback = Mock()
|
||||
check_callback = Mock()
|
||||
|
||||
wrapped.read(mock_sock, read_callback, check_callback)
|
||||
|
||||
# Should delegate to wrapped dispatcher
|
||||
self.assertEqual(len(mock_dispatcher.read_calls), 1)
|
||||
self.assertEqual(mock_dispatcher.read_calls[0], (mock_sock, read_callback))
|
||||
|
||||
# Should call timeout for ping_timeout
|
||||
self.assertEqual(len(mock_dispatcher.timeout_calls), 1)
|
||||
timeout_call = mock_dispatcher.timeout_calls[0]
|
||||
self.assertEqual(timeout_call[0], 10.0) # timeout seconds
|
||||
self.assertEqual(timeout_call[1], check_callback) # callback
|
||||
|
||||
def test_wrapped_dispatcher_read_no_ping_timeout(self):
|
||||
"""Test WrappedDispatcher read method without ping timeout"""
|
||||
mock_dispatcher = MockDispatcher()
|
||||
handle_disconnect = Mock()
|
||||
wrapped = WrappedDispatcher(self.app, None, mock_dispatcher, handle_disconnect)
|
||||
|
||||
mock_sock = Mock()
|
||||
read_callback = Mock()
|
||||
check_callback = Mock()
|
||||
|
||||
wrapped.read(mock_sock, read_callback, check_callback)
|
||||
|
||||
# Should delegate to wrapped dispatcher
|
||||
self.assertEqual(len(mock_dispatcher.read_calls), 1)
|
||||
|
||||
# Should NOT call timeout when ping_timeout is None
|
||||
self.assertEqual(len(mock_dispatcher.timeout_calls), 0)
|
||||
|
||||
def test_wrapped_dispatcher_send(self):
|
||||
"""Test WrappedDispatcher send method"""
|
||||
mock_dispatcher = MockDispatcher()
|
||||
handle_disconnect = Mock()
|
||||
wrapped = WrappedDispatcher(self.app, 10.0, mock_dispatcher, handle_disconnect)
|
||||
|
||||
mock_sock = Mock()
|
||||
test_data = b"test data"
|
||||
|
||||
with patch("websocket._dispatcher.send") as mock_send:
|
||||
result = wrapped.send(mock_sock, test_data)
|
||||
|
||||
# Should delegate to dispatcher.buffwrite
|
||||
self.assertEqual(len(mock_dispatcher.buffwrite_calls), 1)
|
||||
call = mock_dispatcher.buffwrite_calls[0]
|
||||
self.assertEqual(call[0], mock_sock)
|
||||
self.assertEqual(call[1], test_data)
|
||||
self.assertEqual(call[2], mock_send)
|
||||
self.assertEqual(call[3], handle_disconnect)
|
||||
|
||||
# Should return data length
|
||||
self.assertEqual(result, len(test_data))
|
||||
|
||||
def test_wrapped_dispatcher_timeout(self):
|
||||
"""Test WrappedDispatcher timeout method"""
|
||||
mock_dispatcher = MockDispatcher()
|
||||
handle_disconnect = Mock()
|
||||
wrapped = WrappedDispatcher(self.app, 10.0, mock_dispatcher, handle_disconnect)
|
||||
|
||||
callback = Mock()
|
||||
args = ("arg1", "arg2")
|
||||
|
||||
wrapped.timeout(5.0, callback, *args)
|
||||
|
||||
# Should delegate to wrapped dispatcher
|
||||
self.assertEqual(len(mock_dispatcher.timeout_calls), 1)
|
||||
call = mock_dispatcher.timeout_calls[0]
|
||||
self.assertEqual(call[0], 5.0)
|
||||
self.assertEqual(call[1], callback)
|
||||
self.assertEqual(call[2], args)
|
||||
|
||||
def test_wrapped_dispatcher_reconnect(self):
|
||||
"""Test WrappedDispatcher reconnect method"""
|
||||
mock_dispatcher = MockDispatcher()
|
||||
handle_disconnect = Mock()
|
||||
wrapped = WrappedDispatcher(self.app, 10.0, mock_dispatcher, handle_disconnect)
|
||||
|
||||
reconnector = Mock()
|
||||
|
||||
wrapped.reconnect(3, reconnector)
|
||||
|
||||
# Should delegate to timeout method with reconnect=True
|
||||
self.assertEqual(len(mock_dispatcher.timeout_calls), 1)
|
||||
call = mock_dispatcher.timeout_calls[0]
|
||||
self.assertEqual(call[0], 3)
|
||||
self.assertEqual(call[1], reconnector)
|
||||
self.assertEqual(call[2], (True,))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,158 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
import unittest
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
from websocket._handshake import _get_resp_headers
|
||||
from websocket._exceptions import WebSocketBadStatusException
|
||||
from websocket._ssl_compat import SSLError
|
||||
|
||||
"""
|
||||
test_handshake_large_response.py
|
||||
websocket - WebSocket client library for Python
|
||||
|
||||
Copyright 2025 engn33r
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
class HandshakeLargeResponseTest(unittest.TestCase):
|
||||
def test_large_error_response_chunked_reading(self):
|
||||
"""Test that large HTTP error responses during handshake are read in chunks"""
|
||||
|
||||
# Mock socket
|
||||
mock_sock = Mock()
|
||||
|
||||
# Create a large error response body (> 16KB)
|
||||
large_response = b"Error details: " + b"A" * 20000 # 20KB+ response
|
||||
|
||||
# Track recv calls to ensure chunking
|
||||
recv_calls = []
|
||||
|
||||
def mock_recv(sock, bufsize):
|
||||
recv_calls.append(bufsize)
|
||||
# Simulate SSL error if trying to read > 16KB at once
|
||||
if bufsize > 16384:
|
||||
raise SSLError("[SSL: BAD_LENGTH] unknown error")
|
||||
return large_response[:bufsize]
|
||||
|
||||
# Mock read_headers to return error status with large content-length
|
||||
with patch("websocket._handshake.read_headers") as mock_read_headers:
|
||||
mock_read_headers.return_value = (
|
||||
400, # Bad request status
|
||||
{"content-length": str(len(large_response))},
|
||||
"Bad Request",
|
||||
)
|
||||
|
||||
# Mock the recv function to track calls
|
||||
with patch("websocket._socket.recv", side_effect=mock_recv):
|
||||
# This should not raise SSLError, but should raise WebSocketBadStatusException
|
||||
with self.assertRaises(WebSocketBadStatusException) as cm:
|
||||
_get_resp_headers(mock_sock)
|
||||
|
||||
# Verify the response body was included in the exception
|
||||
self.assertIn(
|
||||
b"Error details:",
|
||||
(
|
||||
cm.exception.args[0].encode()
|
||||
if isinstance(cm.exception.args[0], str)
|
||||
else cm.exception.args[0]
|
||||
),
|
||||
)
|
||||
|
||||
# Verify chunked reading was used (multiple recv calls, none > 16KB)
|
||||
self.assertGreater(len(recv_calls), 1)
|
||||
self.assertTrue(all(call <= 16384 for call in recv_calls))
|
||||
|
||||
def test_handshake_ssl_large_response_protection(self):
|
||||
"""Test that the fix prevents SSL BAD_LENGTH errors during handshake"""
|
||||
|
||||
mock_sock = Mock()
|
||||
|
||||
# Large content that would trigger SSL error if read all at once
|
||||
large_content = b"X" * 32768 # 32KB
|
||||
|
||||
chunks_returned = 0
|
||||
|
||||
def mock_recv_chunked(sock, bufsize):
|
||||
nonlocal chunks_returned
|
||||
# Return data in chunks, simulating successful chunked reading
|
||||
chunk_start = chunks_returned * 16384
|
||||
chunk_end = min(chunk_start + bufsize, len(large_content))
|
||||
result = large_content[chunk_start:chunk_end]
|
||||
chunks_returned += 1 if result else 0
|
||||
return result
|
||||
|
||||
with patch("websocket._handshake.read_headers") as mock_read_headers:
|
||||
mock_read_headers.return_value = (
|
||||
500, # Server error
|
||||
{"content-length": str(len(large_content))},
|
||||
"Internal Server Error",
|
||||
)
|
||||
|
||||
with patch("websocket._socket.recv", side_effect=mock_recv_chunked):
|
||||
# Should handle large response without SSL errors
|
||||
with self.assertRaises(WebSocketBadStatusException) as cm:
|
||||
_get_resp_headers(mock_sock)
|
||||
|
||||
# Verify the complete response was captured
|
||||
exception_str = str(cm.exception)
|
||||
# Response body should be in the exception message
|
||||
self.assertIn("XXXXX", exception_str) # Part of the large content
|
||||
|
||||
def test_handshake_normal_small_response(self):
|
||||
"""Test that normal small responses still work correctly"""
|
||||
|
||||
mock_sock = Mock()
|
||||
small_response = b"Small error message"
|
||||
|
||||
def mock_recv(sock, bufsize):
|
||||
return small_response
|
||||
|
||||
with patch("websocket._handshake.read_headers") as mock_read_headers:
|
||||
mock_read_headers.return_value = (
|
||||
404, # Not found
|
||||
{"content-length": str(len(small_response))},
|
||||
"Not Found",
|
||||
)
|
||||
|
||||
with patch("websocket._socket.recv", side_effect=mock_recv):
|
||||
with self.assertRaises(WebSocketBadStatusException) as cm:
|
||||
_get_resp_headers(mock_sock)
|
||||
|
||||
# Verify small response is handled correctly
|
||||
self.assertIn("Small error message", str(cm.exception))
|
||||
|
||||
def test_handshake_no_content_length(self):
|
||||
"""Test handshake error response without content-length header"""
|
||||
|
||||
mock_sock = Mock()
|
||||
|
||||
with patch("websocket._handshake.read_headers") as mock_read_headers:
|
||||
mock_read_headers.return_value = (
|
||||
403, # Forbidden
|
||||
{}, # No content-length header
|
||||
"Forbidden",
|
||||
)
|
||||
|
||||
# Should raise exception without trying to read response body
|
||||
with self.assertRaises(WebSocketBadStatusException) as cm:
|
||||
_get_resp_headers(mock_sock)
|
||||
|
||||
# Should mention status but not have response body
|
||||
exception_str = str(cm.exception)
|
||||
self.assertIn("403", exception_str)
|
||||
self.assertIn("Forbidden", exception_str)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,370 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
#
|
||||
import os
|
||||
import os.path
|
||||
import socket
|
||||
import ssl
|
||||
import unittest
|
||||
|
||||
import websocket
|
||||
from websocket._exceptions import WebSocketProxyException, WebSocketException
|
||||
from websocket._http import (
|
||||
_get_addrinfo_list,
|
||||
_start_proxied_socket,
|
||||
_tunnel,
|
||||
connect,
|
||||
proxy_info,
|
||||
read_headers,
|
||||
HAVE_PYTHON_SOCKS,
|
||||
)
|
||||
|
||||
"""
|
||||
test_http.py
|
||||
websocket - WebSocket client library for Python
|
||||
|
||||
Copyright 2025 engn33r
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
try:
|
||||
from python_socks._errors import ProxyConnectionError, ProxyError, ProxyTimeoutError
|
||||
except:
|
||||
from websocket._http import ProxyConnectionError, ProxyError, ProxyTimeoutError
|
||||
|
||||
# Skip test to access the internet unless TEST_WITH_INTERNET == 1
|
||||
TEST_WITH_INTERNET = os.environ.get("TEST_WITH_INTERNET", "0") == "1"
|
||||
TEST_WITH_PROXY = os.environ.get("TEST_WITH_PROXY", "0") == "1"
|
||||
# Skip tests relying on local websockets server unless LOCAL_WS_SERVER_PORT != -1
|
||||
LOCAL_WS_SERVER_PORT = os.environ.get("LOCAL_WS_SERVER_PORT", "-1")
|
||||
TEST_WITH_LOCAL_SERVER = LOCAL_WS_SERVER_PORT != "-1"
|
||||
|
||||
|
||||
class SockMock:
|
||||
def __init__(self):
|
||||
self.data = []
|
||||
self.sent = []
|
||||
|
||||
def add_packet(self, data):
|
||||
self.data.append(data)
|
||||
|
||||
def gettimeout(self):
|
||||
return None
|
||||
|
||||
def recv(self, bufsize):
|
||||
if self.data:
|
||||
e = self.data.pop(0)
|
||||
if isinstance(e, Exception):
|
||||
raise e
|
||||
if len(e) > bufsize:
|
||||
self.data.insert(0, e[bufsize:])
|
||||
return e[:bufsize]
|
||||
|
||||
def send(self, data):
|
||||
self.sent.append(data)
|
||||
return len(data)
|
||||
|
||||
def close(self):
|
||||
pass
|
||||
|
||||
|
||||
class HeaderSockMock(SockMock):
|
||||
def __init__(self, fname):
|
||||
SockMock.__init__(self)
|
||||
path = os.path.join(os.path.dirname(__file__), fname)
|
||||
with open(path, "rb") as f:
|
||||
self.add_packet(f.read())
|
||||
|
||||
|
||||
class OptsList:
|
||||
def __init__(self):
|
||||
self.timeout = 1
|
||||
self.sockopt = []
|
||||
self.sslopt = {"cert_reqs": ssl.CERT_NONE}
|
||||
|
||||
|
||||
class HttpTest(unittest.TestCase):
|
||||
def test_read_header(self):
|
||||
status, header, _ = read_headers(HeaderSockMock("data/header01.txt"))
|
||||
self.assertEqual(status, 101)
|
||||
self.assertEqual(header["connection"], "Upgrade")
|
||||
# header02.txt is intentionally malformed
|
||||
self.assertRaises(
|
||||
WebSocketException, read_headers, HeaderSockMock("data/header02.txt")
|
||||
)
|
||||
|
||||
def test_tunnel(self):
|
||||
self.assertRaises(
|
||||
WebSocketProxyException,
|
||||
_tunnel,
|
||||
HeaderSockMock("data/header01.txt"),
|
||||
"example.com",
|
||||
80,
|
||||
("username", "password"),
|
||||
)
|
||||
self.assertRaises(
|
||||
WebSocketProxyException,
|
||||
_tunnel,
|
||||
HeaderSockMock("data/header02.txt"),
|
||||
"example.com",
|
||||
80,
|
||||
("username", "password"),
|
||||
)
|
||||
|
||||
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
|
||||
def test_connect(self):
|
||||
# Not currently testing an actual proxy connection, so just check whether proxy errors are raised. This requires internet for a DNS lookup
|
||||
if HAVE_PYTHON_SOCKS:
|
||||
# Need this check, otherwise case where python_socks is not installed triggers
|
||||
# websocket._exceptions.WebSocketException: Python Socks is needed for SOCKS proxying but is not available
|
||||
self.assertRaises(
|
||||
(ProxyTimeoutError, OSError),
|
||||
_start_proxied_socket,
|
||||
"wss://example.com",
|
||||
OptsList(),
|
||||
proxy_info(
|
||||
http_proxy_host="example.com",
|
||||
http_proxy_port="8080",
|
||||
proxy_type="socks4",
|
||||
http_proxy_timeout=1,
|
||||
),
|
||||
)
|
||||
self.assertRaises(
|
||||
(ProxyTimeoutError, OSError),
|
||||
_start_proxied_socket,
|
||||
"wss://example.com",
|
||||
OptsList(),
|
||||
proxy_info(
|
||||
http_proxy_host="example.com",
|
||||
http_proxy_port="8080",
|
||||
proxy_type="socks4a",
|
||||
http_proxy_timeout=1,
|
||||
),
|
||||
)
|
||||
self.assertRaises(
|
||||
(ProxyTimeoutError, OSError),
|
||||
_start_proxied_socket,
|
||||
"wss://example.com",
|
||||
OptsList(),
|
||||
proxy_info(
|
||||
http_proxy_host="example.com",
|
||||
http_proxy_port="8080",
|
||||
proxy_type="socks5",
|
||||
http_proxy_timeout=1,
|
||||
),
|
||||
)
|
||||
self.assertRaises(
|
||||
(ProxyTimeoutError, OSError),
|
||||
_start_proxied_socket,
|
||||
"wss://example.com",
|
||||
OptsList(),
|
||||
proxy_info(
|
||||
http_proxy_host="example.com",
|
||||
http_proxy_port="8080",
|
||||
proxy_type="socks5h",
|
||||
http_proxy_timeout=1,
|
||||
),
|
||||
)
|
||||
self.assertRaises(
|
||||
ProxyConnectionError,
|
||||
connect,
|
||||
"wss://example.com",
|
||||
OptsList(),
|
||||
proxy_info(
|
||||
http_proxy_host="127.0.0.1",
|
||||
http_proxy_port=9999,
|
||||
proxy_type="socks4",
|
||||
http_proxy_timeout=1,
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
self.assertRaises(
|
||||
TypeError,
|
||||
_get_addrinfo_list,
|
||||
None,
|
||||
80,
|
||||
True,
|
||||
proxy_info(
|
||||
http_proxy_host="127.0.0.1", http_proxy_port="9999", proxy_type="http"
|
||||
),
|
||||
)
|
||||
self.assertRaises(
|
||||
TypeError,
|
||||
_get_addrinfo_list,
|
||||
None,
|
||||
80,
|
||||
True,
|
||||
proxy_info(
|
||||
http_proxy_host="127.0.0.1", http_proxy_port="9999", proxy_type="http"
|
||||
),
|
||||
)
|
||||
self.assertRaises(
|
||||
socket.timeout,
|
||||
connect,
|
||||
"wss://google.com",
|
||||
OptsList(),
|
||||
proxy_info(
|
||||
http_proxy_host="8.8.8.8",
|
||||
http_proxy_port=9999,
|
||||
proxy_type="http",
|
||||
http_proxy_timeout=1,
|
||||
),
|
||||
None,
|
||||
)
|
||||
self.assertEqual(
|
||||
connect(
|
||||
"wss://google.com",
|
||||
OptsList(),
|
||||
proxy_info(
|
||||
http_proxy_host="8.8.8.8", http_proxy_port=8080, proxy_type="http"
|
||||
),
|
||||
True,
|
||||
),
|
||||
(True, ("google.com", 443, "/")),
|
||||
)
|
||||
# The following test fails on Mac OS with a gaierror, not an OverflowError
|
||||
# self.assertRaises(OverflowError, connect, "wss://example.com", OptsList(), proxy_info(http_proxy_host="127.0.0.1", http_proxy_port=99999, proxy_type="socks4", timeout=2), False)
|
||||
|
||||
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
|
||||
@unittest.skipUnless(
|
||||
TEST_WITH_PROXY, "This test requires a HTTP proxy to be running on port 8899"
|
||||
)
|
||||
@unittest.skipUnless(
|
||||
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
|
||||
)
|
||||
def test_proxy_connect(self):
|
||||
ws = websocket.WebSocket()
|
||||
ws.connect(
|
||||
f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}",
|
||||
http_proxy_host="127.0.0.1",
|
||||
http_proxy_port="8899",
|
||||
proxy_type="http",
|
||||
)
|
||||
ws.send("Hello, Server")
|
||||
server_response = ws.recv()
|
||||
self.assertEqual(server_response, "Hello, Server")
|
||||
# self.assertEqual(_start_proxied_socket("wss://api.bitfinex.com/ws/2", OptsList(), proxy_info(http_proxy_host="127.0.0.1", http_proxy_port="8899", proxy_type="http"))[1], ("api.bitfinex.com", 443, '/ws/2'))
|
||||
self.assertEqual(
|
||||
_get_addrinfo_list(
|
||||
"api.bitfinex.com",
|
||||
443,
|
||||
True,
|
||||
proxy_info(
|
||||
http_proxy_host="127.0.0.1",
|
||||
http_proxy_port="8899",
|
||||
proxy_type="http",
|
||||
),
|
||||
),
|
||||
(
|
||||
socket.getaddrinfo(
|
||||
"127.0.0.1", 8899, 0, socket.SOCK_STREAM, socket.SOL_TCP
|
||||
),
|
||||
True,
|
||||
None,
|
||||
),
|
||||
)
|
||||
self.assertEqual(
|
||||
connect(
|
||||
"wss://api.bitfinex.com/ws/2",
|
||||
OptsList(),
|
||||
proxy_info(
|
||||
http_proxy_host="127.0.0.1", http_proxy_port=8899, proxy_type="http"
|
||||
),
|
||||
None,
|
||||
)[1],
|
||||
("api.bitfinex.com", 443, "/ws/2"),
|
||||
)
|
||||
# TODO: Test SOCKS4 and SOCK5 proxies with unit tests
|
||||
|
||||
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
|
||||
def test_sslopt(self):
|
||||
ssloptions = {
|
||||
"check_hostname": False,
|
||||
"server_hostname": "ServerName",
|
||||
"ssl_version": ssl.PROTOCOL_TLS_CLIENT,
|
||||
"ciphers": "TLS_AES_256_GCM_SHA384:TLS_CHACHA20_POLY1305_SHA256:\
|
||||
TLS_AES_128_GCM_SHA256:ECDHE-ECDSA-AES256-GCM-SHA384:\
|
||||
ECDHE-RSA-AES256-GCM-SHA384:DHE-RSA-AES256-GCM-SHA384:\
|
||||
ECDHE-ECDSA-CHACHA20-POLY1305:ECDHE-RSA-CHACHA20-POLY1305:\
|
||||
DHE-RSA-CHACHA20-POLY1305:ECDHE-ECDSA-AES128-GCM-SHA256:\
|
||||
ECDHE-RSA-AES128-GCM-SHA256:DHE-RSA-AES128-GCM-SHA256:\
|
||||
ECDHE-ECDSA-AES256-SHA384:ECDHE-RSA-AES256-SHA384:\
|
||||
DHE-RSA-AES256-SHA256:ECDHE-ECDSA-AES128-SHA256:\
|
||||
ECDHE-RSA-AES128-SHA256:DHE-RSA-AES128-SHA256:\
|
||||
ECDHE-ECDSA-AES256-SHA:ECDHE-RSA-AES256-SHA",
|
||||
"ecdh_curve": "prime256v1",
|
||||
}
|
||||
ws_ssl1 = websocket.WebSocket(sslopt=ssloptions)
|
||||
ws_ssl1.connect("wss://api.bitfinex.com/ws/2")
|
||||
ws_ssl1.send("Hello")
|
||||
ws_ssl1.close()
|
||||
|
||||
ws_ssl2 = websocket.WebSocket(sslopt={"check_hostname": True})
|
||||
ws_ssl2.connect("wss://api.bitfinex.com/ws/2")
|
||||
ws_ssl2.close
|
||||
|
||||
def test_proxy_info(self):
|
||||
self.assertEqual(
|
||||
proxy_info(
|
||||
http_proxy_host="127.0.0.1", http_proxy_port="8080", proxy_type="http"
|
||||
).proxy_protocol,
|
||||
"http",
|
||||
)
|
||||
self.assertRaises(
|
||||
ProxyError,
|
||||
proxy_info,
|
||||
http_proxy_host="127.0.0.1",
|
||||
http_proxy_port="8080",
|
||||
proxy_type="badval",
|
||||
)
|
||||
self.assertEqual(
|
||||
proxy_info(
|
||||
http_proxy_host="example.com", http_proxy_port="8080", proxy_type="http"
|
||||
).proxy_host,
|
||||
"example.com",
|
||||
)
|
||||
self.assertEqual(
|
||||
proxy_info(
|
||||
http_proxy_host="127.0.0.1", http_proxy_port="8080", proxy_type="http"
|
||||
).proxy_port,
|
||||
"8080",
|
||||
)
|
||||
self.assertEqual(
|
||||
proxy_info(
|
||||
http_proxy_host="127.0.0.1", http_proxy_port="8080", proxy_type="http"
|
||||
).auth,
|
||||
None,
|
||||
)
|
||||
self.assertEqual(
|
||||
proxy_info(
|
||||
http_proxy_host="127.0.0.1",
|
||||
http_proxy_port="8080",
|
||||
proxy_type="http",
|
||||
http_proxy_auth=("my_username123", "my_pass321"),
|
||||
).auth[0],
|
||||
"my_username123",
|
||||
)
|
||||
self.assertEqual(
|
||||
proxy_info(
|
||||
http_proxy_host="127.0.0.1",
|
||||
http_proxy_port="8080",
|
||||
proxy_type="http",
|
||||
http_proxy_auth=("my_username123", "my_pass321"),
|
||||
).auth[1],
|
||||
"my_pass321",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,273 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
import unittest
|
||||
import struct
|
||||
from unittest.mock import Mock, patch, MagicMock
|
||||
|
||||
from websocket._abnf import ABNF
|
||||
from websocket._core import WebSocket
|
||||
from websocket._exceptions import WebSocketProtocolException, WebSocketPayloadException
|
||||
from websocket._ssl_compat import SSLError
|
||||
|
||||
"""
|
||||
test_large_payloads.py
|
||||
websocket - WebSocket client library for Python
|
||||
|
||||
Copyright 2025 engn33r
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
class LargePayloadTest(unittest.TestCase):
|
||||
def test_frame_length_encoding_boundaries(self):
|
||||
"""Test WebSocket frame length encoding at various boundaries"""
|
||||
|
||||
# Test length encoding boundaries as per RFC 6455
|
||||
test_cases = [
|
||||
(125, "Single byte length"), # Max for 7-bit length
|
||||
(126, "Two byte length start"), # Start of 16-bit length
|
||||
(127, "Two byte length"),
|
||||
(65535, "Two byte length max"), # Max for 16-bit length
|
||||
(65536, "Eight byte length start"), # Start of 64-bit length
|
||||
(16384, "16KB boundary"), # The problematic size
|
||||
(16385, "Just over 16KB"),
|
||||
(32768, "32KB"),
|
||||
(131072, "128KB"),
|
||||
]
|
||||
|
||||
for length, description in test_cases:
|
||||
with self.subTest(length=length, description=description):
|
||||
# Create payload of specified length
|
||||
payload = b"A" * length
|
||||
|
||||
# Create frame
|
||||
frame = ABNF.create_frame(payload, ABNF.OPCODE_BINARY)
|
||||
|
||||
# Verify frame can be formatted without error
|
||||
formatted = frame.format()
|
||||
|
||||
# Verify the frame header is correctly structured
|
||||
self.assertIsInstance(formatted, bytes)
|
||||
self.assertTrue(len(formatted) >= length) # Header + payload
|
||||
|
||||
# Verify payload length is preserved
|
||||
self.assertEqual(len(frame.data), length)
|
||||
|
||||
def test_recv_large_payload_chunked(self):
|
||||
"""Test receiving large payloads in chunks (simulating the 16KB recv issue)"""
|
||||
|
||||
# Create a large payload that would trigger chunked reading
|
||||
large_payload = b"B" * 32768 # 32KB
|
||||
|
||||
# Mock recv function that returns data in 16KB chunks
|
||||
chunks = []
|
||||
chunk_size = 16384
|
||||
for i in range(0, len(large_payload), chunk_size):
|
||||
chunks.append(large_payload[i : i + chunk_size])
|
||||
|
||||
call_count = 0
|
||||
|
||||
def mock_recv(bufsize):
|
||||
nonlocal call_count
|
||||
if call_count >= len(chunks):
|
||||
return b""
|
||||
result = chunks[call_count]
|
||||
call_count += 1
|
||||
return result
|
||||
|
||||
# Test the frame buffer's recv_strict method
|
||||
from websocket._abnf import frame_buffer
|
||||
|
||||
fb = frame_buffer(mock_recv, skip_utf8_validation=True)
|
||||
|
||||
# This should handle large payloads by chunking
|
||||
result = fb.recv_strict(len(large_payload))
|
||||
|
||||
self.assertEqual(result, large_payload)
|
||||
# Verify multiple recv calls were made
|
||||
self.assertGreater(call_count, 1)
|
||||
|
||||
def test_ssl_large_payload_simulation(self):
|
||||
"""Simulate SSL BAD_LENGTH error scenario"""
|
||||
|
||||
# This test demonstrates that the 16KB limit in frame buffer protects against SSL issues
|
||||
payload_size = 16385
|
||||
|
||||
recv_calls = []
|
||||
|
||||
def mock_recv_with_ssl_limit(bufsize):
|
||||
recv_calls.append(bufsize)
|
||||
# This simulates the SSL issue: BAD_LENGTH when trying to recv > 16KB
|
||||
if bufsize > 16384:
|
||||
raise SSLError("[SSL: BAD_LENGTH] unknown error")
|
||||
return b"C" * min(bufsize, 16384)
|
||||
|
||||
from websocket._abnf import frame_buffer
|
||||
|
||||
fb = frame_buffer(mock_recv_with_ssl_limit, skip_utf8_validation=True)
|
||||
|
||||
# The frame buffer handles this correctly by chunking recv calls
|
||||
result = fb.recv_strict(payload_size)
|
||||
|
||||
# Verify it worked and chunked the calls properly
|
||||
self.assertEqual(len(result), payload_size)
|
||||
# Verify no single recv call was > 16KB
|
||||
self.assertTrue(all(call <= 16384 for call in recv_calls))
|
||||
# Verify multiple calls were made
|
||||
self.assertGreater(len(recv_calls), 1)
|
||||
|
||||
def test_frame_format_large_payloads(self):
|
||||
"""Test frame formatting with various large payload sizes"""
|
||||
|
||||
# Test sizes around potential problem areas
|
||||
test_sizes = [16383, 16384, 16385, 32768, 65535, 65536]
|
||||
|
||||
for size in test_sizes:
|
||||
with self.subTest(size=size):
|
||||
payload = b"D" * size
|
||||
frame = ABNF.create_frame(payload, ABNF.OPCODE_BINARY)
|
||||
|
||||
# Should not raise any exceptions
|
||||
formatted = frame.format()
|
||||
|
||||
# Verify structure
|
||||
self.assertIsInstance(formatted, bytes)
|
||||
self.assertEqual(len(frame.data), size)
|
||||
|
||||
# Verify length encoding is correct based on size
|
||||
# Note: frames from create_frame() include masking by default (4 extra bytes)
|
||||
mask_size = 4 # WebSocket frames are masked by default
|
||||
if size < ABNF.LENGTH_7: # < 126
|
||||
# Length should be encoded in single byte
|
||||
expected_header_size = (
|
||||
2 + mask_size
|
||||
) # 1 byte opcode + 1 byte length + 4 byte mask
|
||||
elif size < ABNF.LENGTH_16: # < 65536
|
||||
# Length should be encoded in 2 bytes
|
||||
expected_header_size = (
|
||||
4 + mask_size
|
||||
) # 1 byte opcode + 1 byte marker + 2 bytes length + 4 byte mask
|
||||
else:
|
||||
# Length should be encoded in 8 bytes
|
||||
expected_header_size = (
|
||||
10 + mask_size
|
||||
) # 1 byte opcode + 1 byte marker + 8 bytes length + 4 byte mask
|
||||
|
||||
self.assertEqual(len(formatted), expected_header_size + size)
|
||||
|
||||
def test_send_large_payload_chunking(self):
|
||||
"""Test that large payloads are sent in chunks to avoid SSL issues"""
|
||||
|
||||
mock_sock = Mock()
|
||||
|
||||
# Track how data is sent
|
||||
sent_chunks = []
|
||||
|
||||
def mock_send(data):
|
||||
sent_chunks.append(len(data))
|
||||
return len(data)
|
||||
|
||||
mock_sock.send = mock_send
|
||||
mock_sock.gettimeout.return_value = 30.0
|
||||
|
||||
# Create WebSocket with mocked socket
|
||||
ws = WebSocket()
|
||||
ws.sock = mock_sock
|
||||
ws.connected = True
|
||||
|
||||
# Create large payload
|
||||
large_payload = b"E" * 32768 # 32KB
|
||||
|
||||
# Send the payload
|
||||
with patch("websocket._core.send") as mock_send_func:
|
||||
mock_send_func.side_effect = lambda sock, data: len(data)
|
||||
|
||||
# This should work without SSL errors
|
||||
result = ws.send_binary(large_payload)
|
||||
|
||||
# Verify payload was accepted
|
||||
self.assertGreater(result, 0)
|
||||
|
||||
def test_utf8_validation_large_text(self):
|
||||
"""Test UTF-8 validation with large text payloads"""
|
||||
|
||||
# Create large valid UTF-8 text
|
||||
large_text = "Hello 世界! " * 2000 # About 26KB with Unicode
|
||||
|
||||
# Test frame creation
|
||||
frame = ABNF.create_frame(large_text, ABNF.OPCODE_TEXT)
|
||||
|
||||
# Should not raise validation errors
|
||||
formatted = frame.format()
|
||||
self.assertIsInstance(formatted, bytes)
|
||||
|
||||
# Test with close frame that has invalid UTF-8 (this is what validate() actually checks)
|
||||
invalid_utf8_close_data = struct.pack("!H", 1000) + b"\xff\xfe invalid utf8"
|
||||
|
||||
# Create close frame with invalid UTF-8 data
|
||||
frame = ABNF(1, 0, 0, 0, ABNF.OPCODE_CLOSE, 1, invalid_utf8_close_data)
|
||||
|
||||
# Validation should catch the invalid UTF-8 in close frame reason
|
||||
with self.assertRaises(WebSocketProtocolException):
|
||||
frame.validate(skip_utf8_validation=False)
|
||||
|
||||
def test_frame_buffer_edge_cases(self):
|
||||
"""Test frame buffer with edge cases that could trigger bugs"""
|
||||
|
||||
# Test scenario: exactly 16KB payload split across recv calls
|
||||
payload_16k = b"F" * 16384
|
||||
|
||||
# Simulate receiving in smaller chunks
|
||||
chunks = [payload_16k[i : i + 4096] for i in range(0, len(payload_16k), 4096)]
|
||||
|
||||
call_count = 0
|
||||
|
||||
def mock_recv(bufsize):
|
||||
nonlocal call_count
|
||||
if call_count >= len(chunks):
|
||||
return b""
|
||||
result = chunks[call_count]
|
||||
call_count += 1
|
||||
return result
|
||||
|
||||
from websocket._abnf import frame_buffer
|
||||
|
||||
fb = frame_buffer(mock_recv, skip_utf8_validation=True)
|
||||
result = fb.recv_strict(16384)
|
||||
|
||||
self.assertEqual(result, payload_16k)
|
||||
# Verify multiple recv calls were made
|
||||
self.assertEqual(call_count, 4) # 16KB / 4KB = 4 chunks
|
||||
|
||||
def test_max_frame_size_limits(self):
|
||||
"""Test behavior at WebSocket maximum frame size limits"""
|
||||
|
||||
# Test just under the maximum theoretical frame size
|
||||
# (This is a very large test, so we'll use a smaller representative size)
|
||||
|
||||
# Test with a reasonably large payload that represents the issue
|
||||
large_size = 1024 * 1024 # 1MB
|
||||
payload = b"G" * large_size
|
||||
|
||||
# This should work without issues
|
||||
frame = ABNF.create_frame(payload, ABNF.OPCODE_BINARY)
|
||||
|
||||
# Verify the frame can be formatted
|
||||
formatted = frame.format()
|
||||
self.assertIsInstance(formatted, bytes)
|
||||
|
||||
# Verify payload is preserved
|
||||
self.assertEqual(len(frame.data), large_size)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,357 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
import errno
|
||||
import socket
|
||||
import unittest
|
||||
from unittest.mock import Mock, patch, MagicMock
|
||||
import time
|
||||
|
||||
from websocket._socket import recv, recv_line, send, DEFAULT_SOCKET_OPTION
|
||||
from websocket._ssl_compat import (
|
||||
SSLError,
|
||||
SSLEOFError,
|
||||
SSLWantWriteError,
|
||||
SSLWantReadError,
|
||||
)
|
||||
from websocket._exceptions import (
|
||||
WebSocketTimeoutException,
|
||||
WebSocketConnectionClosedException,
|
||||
)
|
||||
|
||||
"""
|
||||
test_socket.py
|
||||
websocket - WebSocket client library for Python
|
||||
|
||||
Copyright 2025 engn33r
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
class SocketTest(unittest.TestCase):
|
||||
def test_default_socket_option(self):
|
||||
"""Test DEFAULT_SOCKET_OPTION contains expected options"""
|
||||
self.assertIsInstance(DEFAULT_SOCKET_OPTION, list)
|
||||
self.assertGreater(len(DEFAULT_SOCKET_OPTION), 0)
|
||||
|
||||
# Should contain TCP_NODELAY option
|
||||
tcp_nodelay_found = any(
|
||||
opt[1] == socket.TCP_NODELAY for opt in DEFAULT_SOCKET_OPTION
|
||||
)
|
||||
self.assertTrue(tcp_nodelay_found)
|
||||
|
||||
def test_recv_normal(self):
|
||||
"""Test normal recv operation"""
|
||||
mock_sock = Mock()
|
||||
mock_sock.recv.return_value = b"test data"
|
||||
|
||||
result = recv(mock_sock, 9)
|
||||
|
||||
self.assertEqual(result, b"test data")
|
||||
mock_sock.recv.assert_called_once_with(9)
|
||||
|
||||
def test_recv_timeout_error(self):
|
||||
"""Test recv with TimeoutError"""
|
||||
mock_sock = Mock()
|
||||
mock_sock.recv.side_effect = TimeoutError("Connection timed out")
|
||||
|
||||
with self.assertRaises(WebSocketTimeoutException) as cm:
|
||||
recv(mock_sock, 9)
|
||||
|
||||
self.assertEqual(str(cm.exception), "Connection timed out")
|
||||
|
||||
def test_recv_socket_timeout(self):
|
||||
"""Test recv with socket.timeout"""
|
||||
mock_sock = Mock()
|
||||
timeout_exc = socket.timeout("Socket timed out")
|
||||
timeout_exc.args = ("Socket timed out",)
|
||||
mock_sock.recv.side_effect = timeout_exc
|
||||
mock_sock.gettimeout.return_value = 30.0
|
||||
|
||||
with self.assertRaises(WebSocketTimeoutException) as cm:
|
||||
recv(mock_sock, 9)
|
||||
|
||||
# In Python 3.10+, socket.timeout is a subclass of TimeoutError
|
||||
# so it's caught by the TimeoutError handler with hardcoded message
|
||||
# In Python 3.9, socket.timeout is caught by socket.timeout handler
|
||||
# which preserves the original message
|
||||
import sys
|
||||
|
||||
if sys.version_info >= (3, 10):
|
||||
self.assertEqual(str(cm.exception), "Connection timed out")
|
||||
else:
|
||||
self.assertEqual(str(cm.exception), "Socket timed out")
|
||||
|
||||
def test_recv_ssl_timeout(self):
|
||||
"""Test recv with SSL timeout error"""
|
||||
mock_sock = Mock()
|
||||
ssl_exc = SSLError("The operation timed out")
|
||||
ssl_exc.args = ("The operation timed out",)
|
||||
mock_sock.recv.side_effect = ssl_exc
|
||||
|
||||
with self.assertRaises(WebSocketTimeoutException) as cm:
|
||||
recv(mock_sock, 9)
|
||||
|
||||
self.assertEqual(str(cm.exception), "The operation timed out")
|
||||
|
||||
def test_recv_ssl_non_timeout_error(self):
|
||||
"""Test recv with SSL non-timeout error"""
|
||||
mock_sock = Mock()
|
||||
ssl_exc = SSLError("SSL certificate error")
|
||||
ssl_exc.args = ("SSL certificate error",)
|
||||
mock_sock.recv.side_effect = ssl_exc
|
||||
|
||||
# Should re-raise the original SSL error
|
||||
with self.assertRaises(SSLError):
|
||||
recv(mock_sock, 9)
|
||||
|
||||
def test_recv_empty_response(self):
|
||||
"""Test recv with empty response (connection closed)"""
|
||||
mock_sock = Mock()
|
||||
mock_sock.recv.return_value = b""
|
||||
|
||||
with self.assertRaises(WebSocketConnectionClosedException) as cm:
|
||||
recv(mock_sock, 9)
|
||||
|
||||
self.assertEqual(str(cm.exception), "Connection to remote host was lost.")
|
||||
|
||||
def test_recv_ssl_want_read_error(self):
|
||||
"""Test recv with SSLWantReadError (should retry)"""
|
||||
mock_sock = Mock()
|
||||
|
||||
# First call raises SSLWantReadError, second call succeeds
|
||||
mock_sock.recv.side_effect = [SSLWantReadError(), b"data after retry"]
|
||||
|
||||
with patch("selectors.DefaultSelector") as mock_selector_class:
|
||||
mock_selector = Mock()
|
||||
mock_selector_class.return_value = mock_selector
|
||||
mock_selector.select.return_value = [True] # Ready to read
|
||||
|
||||
result = recv(mock_sock, 100)
|
||||
|
||||
self.assertEqual(result, b"data after retry")
|
||||
mock_selector.register.assert_called()
|
||||
mock_selector.close.assert_called()
|
||||
|
||||
def test_recv_ssl_want_read_timeout(self):
|
||||
"""Test recv with SSLWantReadError that times out"""
|
||||
mock_sock = Mock()
|
||||
mock_sock.recv.side_effect = SSLWantReadError()
|
||||
mock_sock.gettimeout.return_value = 1.0
|
||||
|
||||
with patch("selectors.DefaultSelector") as mock_selector_class:
|
||||
mock_selector = Mock()
|
||||
mock_selector_class.return_value = mock_selector
|
||||
mock_selector.select.return_value = [] # Timeout
|
||||
|
||||
with self.assertRaises(WebSocketTimeoutException):
|
||||
recv(mock_sock, 100)
|
||||
|
||||
def test_recv_line(self):
|
||||
"""Test recv_line functionality"""
|
||||
mock_sock = Mock()
|
||||
|
||||
# Mock recv to return one character at a time
|
||||
recv_calls = [b"H", b"e", b"l", b"l", b"o", b"\n"]
|
||||
|
||||
with patch("websocket._socket.recv", side_effect=recv_calls) as mock_recv:
|
||||
result = recv_line(mock_sock)
|
||||
|
||||
self.assertEqual(result, b"Hello\n")
|
||||
self.assertEqual(mock_recv.call_count, 6)
|
||||
|
||||
def test_send_normal(self):
|
||||
"""Test normal send operation"""
|
||||
mock_sock = Mock()
|
||||
mock_sock.send.return_value = 9
|
||||
mock_sock.gettimeout.return_value = 30.0
|
||||
|
||||
result = send(mock_sock, b"test data")
|
||||
|
||||
self.assertEqual(result, 9)
|
||||
mock_sock.send.assert_called_with(b"test data")
|
||||
|
||||
def test_send_zero_timeout(self):
|
||||
"""Test send with zero timeout (non-blocking)"""
|
||||
mock_sock = Mock()
|
||||
mock_sock.send.return_value = 9
|
||||
mock_sock.gettimeout.return_value = 0
|
||||
|
||||
result = send(mock_sock, b"test data")
|
||||
|
||||
self.assertEqual(result, 9)
|
||||
mock_sock.send.assert_called_once_with(b"test data")
|
||||
|
||||
def test_send_ssl_eof_error(self):
|
||||
"""Test send with SSLEOFError"""
|
||||
mock_sock = Mock()
|
||||
mock_sock.gettimeout.return_value = 30.0
|
||||
mock_sock.send.side_effect = SSLEOFError("Connection closed")
|
||||
|
||||
with self.assertRaises(WebSocketConnectionClosedException) as cm:
|
||||
send(mock_sock, b"test data")
|
||||
|
||||
self.assertEqual(str(cm.exception), "socket is already closed.")
|
||||
|
||||
def test_send_ssl_want_write_error(self):
|
||||
"""Test send with SSLWantWriteError (should retry)"""
|
||||
mock_sock = Mock()
|
||||
mock_sock.gettimeout.return_value = 30.0
|
||||
|
||||
# First call raises SSLWantWriteError, second call succeeds
|
||||
mock_sock.send.side_effect = [SSLWantWriteError(), 9]
|
||||
|
||||
with patch("selectors.DefaultSelector") as mock_selector_class:
|
||||
mock_selector = Mock()
|
||||
mock_selector_class.return_value = mock_selector
|
||||
mock_selector.select.return_value = [True] # Ready to write
|
||||
|
||||
result = send(mock_sock, b"test data")
|
||||
|
||||
self.assertEqual(result, 9)
|
||||
mock_selector.register.assert_called()
|
||||
mock_selector.close.assert_called()
|
||||
|
||||
def test_send_socket_eagain_error(self):
|
||||
"""Test send with EAGAIN error (should retry)"""
|
||||
mock_sock = Mock()
|
||||
mock_sock.gettimeout.return_value = 30.0
|
||||
|
||||
# Create socket error with EAGAIN
|
||||
eagain_error = socket.error("Resource temporarily unavailable")
|
||||
eagain_error.errno = errno.EAGAIN
|
||||
eagain_error.args = (errno.EAGAIN, "Resource temporarily unavailable")
|
||||
|
||||
# First call raises EAGAIN, second call succeeds
|
||||
mock_sock.send.side_effect = [eagain_error, 9]
|
||||
|
||||
with patch("selectors.DefaultSelector") as mock_selector_class:
|
||||
mock_selector = Mock()
|
||||
mock_selector_class.return_value = mock_selector
|
||||
mock_selector.select.return_value = [True] # Ready to write
|
||||
|
||||
result = send(mock_sock, b"test data")
|
||||
|
||||
self.assertEqual(result, 9)
|
||||
|
||||
def test_send_socket_ewouldblock_error(self):
|
||||
"""Test send with EWOULDBLOCK error (should retry)"""
|
||||
mock_sock = Mock()
|
||||
mock_sock.gettimeout.return_value = 30.0
|
||||
|
||||
# Create socket error with EWOULDBLOCK
|
||||
ewouldblock_error = socket.error("Operation would block")
|
||||
ewouldblock_error.errno = errno.EWOULDBLOCK
|
||||
ewouldblock_error.args = (errno.EWOULDBLOCK, "Operation would block")
|
||||
|
||||
# First call raises EWOULDBLOCK, second call succeeds
|
||||
mock_sock.send.side_effect = [ewouldblock_error, 9]
|
||||
|
||||
with patch("selectors.DefaultSelector") as mock_selector_class:
|
||||
mock_selector = Mock()
|
||||
mock_selector_class.return_value = mock_selector
|
||||
mock_selector.select.return_value = [True] # Ready to write
|
||||
|
||||
result = send(mock_sock, b"test data")
|
||||
|
||||
self.assertEqual(result, 9)
|
||||
|
||||
def test_send_socket_other_error(self):
|
||||
"""Test send with other socket error (should raise)"""
|
||||
mock_sock = Mock()
|
||||
mock_sock.gettimeout.return_value = 30.0
|
||||
|
||||
# Create socket error with different errno
|
||||
other_error = socket.error("Connection reset by peer")
|
||||
other_error.errno = errno.ECONNRESET
|
||||
other_error.args = (errno.ECONNRESET, "Connection reset by peer")
|
||||
|
||||
mock_sock.send.side_effect = other_error
|
||||
|
||||
with self.assertRaises(socket.error):
|
||||
send(mock_sock, b"test data")
|
||||
|
||||
def test_send_socket_error_no_errno(self):
|
||||
"""Test send with socket error that has no errno"""
|
||||
mock_sock = Mock()
|
||||
mock_sock.gettimeout.return_value = 30.0
|
||||
|
||||
# Create socket error without errno attribute
|
||||
no_errno_error = socket.error("Generic socket error")
|
||||
no_errno_error.args = ("Generic socket error",)
|
||||
|
||||
mock_sock.send.side_effect = no_errno_error
|
||||
|
||||
with self.assertRaises(socket.error):
|
||||
send(mock_sock, b"test data")
|
||||
|
||||
def test_send_write_timeout(self):
|
||||
"""Test send write operation timeout"""
|
||||
mock_sock = Mock()
|
||||
mock_sock.gettimeout.return_value = 30.0
|
||||
|
||||
# First call raises EAGAIN
|
||||
eagain_error = socket.error("Resource temporarily unavailable")
|
||||
eagain_error.errno = errno.EAGAIN
|
||||
eagain_error.args = (errno.EAGAIN, "Resource temporarily unavailable")
|
||||
|
||||
mock_sock.send.side_effect = eagain_error
|
||||
|
||||
with patch("selectors.DefaultSelector") as mock_selector_class:
|
||||
mock_selector = Mock()
|
||||
mock_selector_class.return_value = mock_selector
|
||||
mock_selector.select.return_value = [] # Timeout - nothing ready
|
||||
|
||||
result = send(mock_sock, b"test data")
|
||||
|
||||
# Should return 0 when write times out
|
||||
self.assertEqual(result, 0)
|
||||
|
||||
def test_send_string_data(self):
|
||||
"""Test send with string data (should be encoded)"""
|
||||
mock_sock = Mock()
|
||||
mock_sock.send.return_value = 9
|
||||
mock_sock.gettimeout.return_value = 30.0
|
||||
|
||||
result = send(mock_sock, "test data")
|
||||
|
||||
self.assertEqual(result, 9)
|
||||
mock_sock.send.assert_called_with(b"test data")
|
||||
|
||||
def test_send_partial_send_retry(self):
|
||||
"""Test send retry mechanism"""
|
||||
mock_sock = Mock()
|
||||
mock_sock.gettimeout.return_value = 30.0
|
||||
|
||||
# Create a scenario where send succeeds after selector retry
|
||||
eagain_error = socket.error("Resource temporarily unavailable")
|
||||
eagain_error.errno = errno.EAGAIN
|
||||
eagain_error.args = (errno.EAGAIN, "Resource temporarily unavailable")
|
||||
|
||||
# Mock the internal _send function behavior
|
||||
mock_sock.send.side_effect = [eagain_error, 9]
|
||||
|
||||
with patch("selectors.DefaultSelector") as mock_selector_class:
|
||||
mock_selector = Mock()
|
||||
mock_selector_class.return_value = mock_selector
|
||||
mock_selector.select.return_value = [True] # Socket ready for writing
|
||||
|
||||
result = send(mock_sock, b"test data")
|
||||
|
||||
self.assertEqual(result, 9)
|
||||
# Verify selector was used for retry mechanism
|
||||
mock_selector.register.assert_called()
|
||||
mock_selector.select.assert_called()
|
||||
mock_selector.close.assert_called()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,160 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
import errno
|
||||
import socket
|
||||
import unittest
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
from websocket._socket import recv
|
||||
from websocket._ssl_compat import SSLWantReadError
|
||||
from websocket._exceptions import (
|
||||
WebSocketTimeoutException,
|
||||
WebSocketConnectionClosedException,
|
||||
)
|
||||
|
||||
"""
|
||||
test_socket_bugs.py
|
||||
websocket - WebSocket client library for Python
|
||||
|
||||
Copyright 2025 engn33r
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
class SocketBugsTest(unittest.TestCase):
|
||||
"""Test bugs found in socket handling logic"""
|
||||
|
||||
def test_bug_implicit_none_return_from_ssl_want_read_fixed(self):
|
||||
"""
|
||||
BUG #5 FIX VERIFICATION: Test SSLWantReadError timeout now raises correct exception
|
||||
|
||||
Bug was in _socket.py:100-101 - SSLWantReadError except block returned None implicitly
|
||||
Fixed: Now properly handles timeout with WebSocketTimeoutException
|
||||
"""
|
||||
mock_sock = Mock()
|
||||
mock_sock.recv.side_effect = SSLWantReadError()
|
||||
mock_sock.gettimeout.return_value = 1.0
|
||||
|
||||
with patch("selectors.DefaultSelector") as mock_selector_class:
|
||||
mock_selector = Mock()
|
||||
mock_selector_class.return_value = mock_selector
|
||||
mock_selector.select.return_value = [] # Timeout - no data ready
|
||||
|
||||
with self.assertRaises(WebSocketTimeoutException) as cm:
|
||||
recv(mock_sock, 100)
|
||||
|
||||
# Verify correct timeout exception and message
|
||||
self.assertIn("Connection timed out waiting for data", str(cm.exception))
|
||||
|
||||
def test_bug_implicit_none_return_from_socket_error_fixed(self):
|
||||
"""
|
||||
BUG #5 FIX VERIFICATION: Test that socket.error with EAGAIN now handles timeout correctly
|
||||
|
||||
Bug was in _socket.py:102-105 - socket.error except block returned None implicitly
|
||||
Fixed: Now properly handles timeout with WebSocketTimeoutException
|
||||
"""
|
||||
mock_sock = Mock()
|
||||
|
||||
# Create socket error with EAGAIN (should be retried)
|
||||
eagain_error = OSError(errno.EAGAIN, "Resource temporarily unavailable")
|
||||
|
||||
# First call raises EAGAIN, selector times out on retry
|
||||
mock_sock.recv.side_effect = eagain_error
|
||||
mock_sock.gettimeout.return_value = 1.0
|
||||
|
||||
with patch("selectors.DefaultSelector") as mock_selector_class:
|
||||
mock_selector = Mock()
|
||||
mock_selector_class.return_value = mock_selector
|
||||
mock_selector.select.return_value = [] # Timeout - no data ready
|
||||
|
||||
with self.assertRaises(WebSocketTimeoutException) as cm:
|
||||
recv(mock_sock, 100)
|
||||
|
||||
# Verify correct timeout exception and message
|
||||
self.assertIn("Connection timed out waiting for data", str(cm.exception))
|
||||
|
||||
def test_bug_wrong_exception_for_selector_timeout_fixed(self):
|
||||
"""
|
||||
BUG #6 FIX VERIFICATION: Test that selector timeout now raises correct exception type
|
||||
|
||||
Bug was in _socket.py:115 returning None for timeout, treated as connection error
|
||||
Fixed: Now raises WebSocketTimeoutException directly
|
||||
"""
|
||||
mock_sock = Mock()
|
||||
mock_sock.recv.side_effect = SSLWantReadError() # Trigger retry path
|
||||
mock_sock.gettimeout.return_value = 1.0
|
||||
|
||||
with patch("selectors.DefaultSelector") as mock_selector_class:
|
||||
mock_selector = Mock()
|
||||
mock_selector_class.return_value = mock_selector
|
||||
mock_selector.select.return_value = [] # TIMEOUT - this is key!
|
||||
|
||||
with self.assertRaises(WebSocketTimeoutException) as cm:
|
||||
recv(mock_sock, 100)
|
||||
|
||||
# Verify it's the correct timeout exception with proper message
|
||||
self.assertIn("Connection timed out waiting for data", str(cm.exception))
|
||||
|
||||
# This proves the fix works:
|
||||
# 1. selector.select() returns [] (timeout)
|
||||
# 2. _recv() now raises WebSocketTimeoutException directly
|
||||
# 3. No more misclassification as connection closed error!
|
||||
|
||||
def test_socket_timeout_exception_handling(self):
|
||||
"""
|
||||
Test that socket.timeout exceptions are properly handled
|
||||
"""
|
||||
mock_sock = Mock()
|
||||
mock_sock.gettimeout.return_value = 1.0
|
||||
|
||||
# Simulate a real socket.timeout scenario
|
||||
mock_sock.recv.side_effect = socket.timeout("Operation timed out")
|
||||
|
||||
# This works correctly - socket.timeout raises WebSocketTimeoutException
|
||||
with self.assertRaises(WebSocketTimeoutException) as cm:
|
||||
recv(mock_sock, 100)
|
||||
|
||||
# In Python 3.10+, socket.timeout is a subclass of TimeoutError
|
||||
# so it's caught by the TimeoutError handler with hardcoded message
|
||||
# In Python 3.9, socket.timeout is caught by socket.timeout handler
|
||||
# which preserves the original message
|
||||
import sys
|
||||
|
||||
if sys.version_info >= (3, 10):
|
||||
self.assertIn("Connection timed out", str(cm.exception))
|
||||
else:
|
||||
self.assertIn("Operation timed out", str(cm.exception))
|
||||
|
||||
def test_correct_ssl_want_read_retry_behavior(self):
|
||||
"""Test the correct behavior when SSLWantReadError is properly handled"""
|
||||
mock_sock = Mock()
|
||||
|
||||
# First call raises SSLWantReadError, second call succeeds
|
||||
mock_sock.recv.side_effect = [SSLWantReadError(), b"data after retry"]
|
||||
mock_sock.gettimeout.return_value = 1.0
|
||||
|
||||
with patch("selectors.DefaultSelector") as mock_selector_class:
|
||||
mock_selector = Mock()
|
||||
mock_selector_class.return_value = mock_selector
|
||||
mock_selector.select.return_value = [True] # Data ready after wait
|
||||
|
||||
# This should work correctly
|
||||
result = recv(mock_sock, 100)
|
||||
self.assertEqual(result, b"data after retry")
|
||||
|
||||
# Selector should be used for retry
|
||||
mock_selector.register.assert_called()
|
||||
mock_selector.select.assert_called()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,91 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
import sys
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
"""
|
||||
test_ssl_compat.py
|
||||
websocket - WebSocket client library for Python
|
||||
|
||||
Copyright 2025 engn33r
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
class SSLCompatTest(unittest.TestCase):
|
||||
def test_ssl_available(self):
|
||||
"""Test that SSL is available in normal conditions"""
|
||||
import websocket._ssl_compat as ssl_compat
|
||||
|
||||
# In normal conditions, SSL should be available
|
||||
self.assertTrue(ssl_compat.HAVE_SSL)
|
||||
self.assertIsNotNone(ssl_compat.ssl)
|
||||
|
||||
# SSL exception classes should be available
|
||||
self.assertTrue(hasattr(ssl_compat, "SSLError"))
|
||||
self.assertTrue(hasattr(ssl_compat, "SSLEOFError"))
|
||||
self.assertTrue(hasattr(ssl_compat, "SSLWantReadError"))
|
||||
self.assertTrue(hasattr(ssl_compat, "SSLWantWriteError"))
|
||||
|
||||
def test_ssl_not_available(self):
|
||||
"""Test fallback behavior when SSL is not available"""
|
||||
# Remove ssl_compat from modules to force reimport
|
||||
if "websocket._ssl_compat" in sys.modules:
|
||||
del sys.modules["websocket._ssl_compat"]
|
||||
|
||||
# Mock the ssl module to not be available
|
||||
import builtins
|
||||
|
||||
original_import = builtins.__import__
|
||||
|
||||
def mock_import(name, *args, **kwargs):
|
||||
if name == "ssl":
|
||||
raise ImportError("No module named 'ssl'")
|
||||
return original_import(name, *args, **kwargs)
|
||||
|
||||
with patch("builtins.__import__", side_effect=mock_import):
|
||||
import websocket._ssl_compat as ssl_compat
|
||||
|
||||
# SSL should not be available
|
||||
self.assertFalse(ssl_compat.HAVE_SSL)
|
||||
self.assertIsNone(ssl_compat.ssl)
|
||||
|
||||
# Fallback exception classes should be available and functional
|
||||
self.assertTrue(issubclass(ssl_compat.SSLError, Exception))
|
||||
self.assertTrue(issubclass(ssl_compat.SSLEOFError, Exception))
|
||||
self.assertTrue(issubclass(ssl_compat.SSLWantReadError, Exception))
|
||||
self.assertTrue(issubclass(ssl_compat.SSLWantWriteError, Exception))
|
||||
|
||||
# Test that exceptions can be instantiated
|
||||
ssl_error = ssl_compat.SSLError("test error")
|
||||
self.assertIsInstance(ssl_error, Exception)
|
||||
self.assertEqual(str(ssl_error), "test error")
|
||||
|
||||
ssl_eof_error = ssl_compat.SSLEOFError("test eof")
|
||||
self.assertIsInstance(ssl_eof_error, Exception)
|
||||
|
||||
ssl_want_read = ssl_compat.SSLWantReadError("test read")
|
||||
self.assertIsInstance(ssl_want_read, Exception)
|
||||
|
||||
ssl_want_write = ssl_compat.SSLWantWriteError("test write")
|
||||
self.assertIsInstance(ssl_want_write, Exception)
|
||||
|
||||
def tearDown(self):
|
||||
"""Clean up after tests"""
|
||||
# Ensure ssl_compat is reimported fresh for next test
|
||||
if "websocket._ssl_compat" in sys.modules:
|
||||
del sys.modules["websocket._ssl_compat"]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,638 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
import unittest
|
||||
import socket
|
||||
import ssl
|
||||
from unittest.mock import Mock, patch, MagicMock
|
||||
|
||||
from websocket._ssl_compat import (
|
||||
SSLError,
|
||||
SSLEOFError,
|
||||
SSLWantReadError,
|
||||
SSLWantWriteError,
|
||||
HAVE_SSL,
|
||||
)
|
||||
from websocket._http import _ssl_socket, _wrap_sni_socket
|
||||
from websocket._exceptions import WebSocketException
|
||||
from websocket._socket import recv, send
|
||||
|
||||
"""
|
||||
test_ssl_edge_cases.py
|
||||
websocket - WebSocket client library for Python
|
||||
|
||||
Copyright 2025 engn33r
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
class SSLEdgeCasesTest(unittest.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
if not HAVE_SSL:
|
||||
self.skipTest("SSL not available")
|
||||
|
||||
def test_ssl_handshake_failure(self):
|
||||
"""Test SSL handshake failure scenarios"""
|
||||
mock_sock = Mock()
|
||||
|
||||
# Test SSL handshake timeout
|
||||
with patch("ssl.SSLContext") as mock_ssl_context:
|
||||
mock_context = Mock()
|
||||
mock_ssl_context.return_value = mock_context
|
||||
mock_context.wrap_socket.side_effect = socket.timeout(
|
||||
"SSL handshake timeout"
|
||||
)
|
||||
|
||||
sslopt = {"cert_reqs": ssl.CERT_REQUIRED}
|
||||
|
||||
with self.assertRaises(socket.timeout):
|
||||
_ssl_socket(mock_sock, sslopt, "example.com")
|
||||
|
||||
def test_ssl_certificate_verification_failures(self):
|
||||
"""Test various SSL certificate verification failure scenarios"""
|
||||
mock_sock = Mock()
|
||||
|
||||
# Test certificate verification failure
|
||||
with patch("ssl.SSLContext") as mock_ssl_context:
|
||||
mock_context = Mock()
|
||||
mock_ssl_context.return_value = mock_context
|
||||
mock_context.wrap_socket.side_effect = ssl.SSLCertVerificationError(
|
||||
"Certificate verification failed"
|
||||
)
|
||||
|
||||
sslopt = {"cert_reqs": ssl.CERT_REQUIRED, "check_hostname": True}
|
||||
|
||||
with self.assertRaises(ssl.SSLCertVerificationError):
|
||||
_ssl_socket(mock_sock, sslopt, "badssl.example")
|
||||
|
||||
def test_ssl_context_configuration_edge_cases(self):
|
||||
"""Test SSL context configuration with various edge cases"""
|
||||
mock_sock = Mock()
|
||||
|
||||
# Test with pre-created SSL context
|
||||
with patch("ssl.SSLContext") as mock_ssl_context:
|
||||
existing_context = Mock()
|
||||
existing_context.wrap_socket.return_value = Mock()
|
||||
mock_ssl_context.return_value = existing_context
|
||||
|
||||
sslopt = {"context": existing_context}
|
||||
|
||||
# Call _ssl_socket which should use the existing context
|
||||
_ssl_socket(mock_sock, sslopt, "example.com")
|
||||
|
||||
# Should use the provided context, not create a new one
|
||||
existing_context.wrap_socket.assert_called_once()
|
||||
|
||||
def test_ssl_ca_bundle_environment_edge_cases(self):
|
||||
"""Test CA bundle environment variable edge cases"""
|
||||
mock_sock = Mock()
|
||||
|
||||
# Test with non-existent CA bundle file
|
||||
with patch.dict(
|
||||
"os.environ", {"WEBSOCKET_CLIENT_CA_BUNDLE": "/nonexistent/ca-bundle.crt"}
|
||||
):
|
||||
with patch("os.path.isfile", return_value=False):
|
||||
with patch("os.path.isdir", return_value=False):
|
||||
with patch("ssl.SSLContext") as mock_ssl_context:
|
||||
mock_context = Mock()
|
||||
mock_ssl_context.return_value = mock_context
|
||||
mock_context.wrap_socket.return_value = Mock()
|
||||
|
||||
sslopt = {}
|
||||
_ssl_socket(mock_sock, sslopt, "example.com")
|
||||
|
||||
# Should not try to load non-existent CA bundle
|
||||
mock_context.load_verify_locations.assert_not_called()
|
||||
|
||||
# Test with CA bundle directory
|
||||
with patch.dict("os.environ", {"WEBSOCKET_CLIENT_CA_BUNDLE": "/etc/ssl/certs"}):
|
||||
with patch("os.path.isfile", return_value=False):
|
||||
with patch("os.path.isdir", return_value=True):
|
||||
with patch("ssl.SSLContext") as mock_ssl_context:
|
||||
mock_context = Mock()
|
||||
mock_ssl_context.return_value = mock_context
|
||||
mock_context.wrap_socket.return_value = Mock()
|
||||
|
||||
sslopt = {}
|
||||
_ssl_socket(mock_sock, sslopt, "example.com")
|
||||
|
||||
# Should load CA directory
|
||||
mock_context.load_verify_locations.assert_called_with(
|
||||
cafile=None, capath="/etc/ssl/certs"
|
||||
)
|
||||
|
||||
def test_ssl_cipher_configuration_edge_cases(self):
|
||||
"""Test SSL cipher configuration edge cases"""
|
||||
mock_sock = Mock()
|
||||
|
||||
# Test with invalid cipher suite
|
||||
with patch("ssl.SSLContext") as mock_ssl_context:
|
||||
mock_context = Mock()
|
||||
mock_ssl_context.return_value = mock_context
|
||||
mock_context.set_ciphers.side_effect = ssl.SSLError(
|
||||
"No cipher can be selected"
|
||||
)
|
||||
mock_context.wrap_socket.return_value = Mock()
|
||||
|
||||
sslopt = {"ciphers": "INVALID_CIPHER"}
|
||||
|
||||
with self.assertRaises(WebSocketException):
|
||||
_ssl_socket(mock_sock, sslopt, "example.com")
|
||||
|
||||
def test_ssl_ecdh_curve_edge_cases(self):
|
||||
"""Test ECDH curve configuration edge cases"""
|
||||
mock_sock = Mock()
|
||||
|
||||
# Test with invalid ECDH curve
|
||||
with patch("ssl.SSLContext") as mock_ssl_context:
|
||||
mock_context = Mock()
|
||||
mock_ssl_context.return_value = mock_context
|
||||
mock_context.set_ecdh_curve.side_effect = ValueError("unknown curve name")
|
||||
mock_context.wrap_socket.return_value = Mock()
|
||||
|
||||
sslopt = {"ecdh_curve": "invalid_curve"}
|
||||
|
||||
with self.assertRaises(WebSocketException):
|
||||
_ssl_socket(mock_sock, sslopt, "example.com")
|
||||
|
||||
def test_ssl_client_certificate_edge_cases(self):
|
||||
"""Test client certificate configuration edge cases"""
|
||||
mock_sock = Mock()
|
||||
|
||||
# Test with non-existent client certificate
|
||||
with patch("ssl.SSLContext") as mock_ssl_context:
|
||||
mock_context = Mock()
|
||||
mock_ssl_context.return_value = mock_context
|
||||
mock_context.load_cert_chain.side_effect = FileNotFoundError("No such file")
|
||||
mock_context.wrap_socket.return_value = Mock()
|
||||
|
||||
sslopt = {"certfile": "/nonexistent/client.crt"}
|
||||
|
||||
with self.assertRaises(WebSocketException):
|
||||
_ssl_socket(mock_sock, sslopt, "example.com")
|
||||
|
||||
def test_ssl_want_read_write_retry_edge_cases(self):
|
||||
"""Test SSL want read/write retry edge cases"""
|
||||
mock_sock = Mock()
|
||||
|
||||
# Test SSLWantReadError with multiple retries before success
|
||||
read_attempts = [0] # Use list for mutable reference
|
||||
|
||||
def mock_recv(bufsize):
|
||||
read_attempts[0] += 1
|
||||
if read_attempts[0] == 1:
|
||||
raise SSLWantReadError("The operation did not complete")
|
||||
elif read_attempts[0] == 2:
|
||||
return b"data after retries"
|
||||
else:
|
||||
return b""
|
||||
|
||||
mock_sock.recv.side_effect = mock_recv
|
||||
mock_sock.gettimeout.return_value = 30.0
|
||||
|
||||
with patch("selectors.DefaultSelector") as mock_selector_class:
|
||||
mock_selector = Mock()
|
||||
mock_selector_class.return_value = mock_selector
|
||||
mock_selector.select.return_value = [True] # Always ready
|
||||
|
||||
result = recv(mock_sock, 100)
|
||||
|
||||
self.assertEqual(result, b"data after retries")
|
||||
self.assertEqual(read_attempts[0], 2)
|
||||
# Should have used selector for retry
|
||||
mock_selector.register.assert_called()
|
||||
mock_selector.select.assert_called()
|
||||
|
||||
def test_ssl_want_write_retry_edge_cases(self):
|
||||
"""Test SSL want write retry edge cases"""
|
||||
mock_sock = Mock()
|
||||
|
||||
# Test SSLWantWriteError with multiple retries before success
|
||||
write_attempts = [0] # Use list for mutable reference
|
||||
|
||||
def mock_send(data):
|
||||
write_attempts[0] += 1
|
||||
if write_attempts[0] == 1:
|
||||
raise SSLWantWriteError("The operation did not complete")
|
||||
elif write_attempts[0] == 2:
|
||||
return len(data)
|
||||
else:
|
||||
return 0
|
||||
|
||||
mock_sock.send.side_effect = mock_send
|
||||
mock_sock.gettimeout.return_value = 30.0
|
||||
|
||||
with patch("selectors.DefaultSelector") as mock_selector_class:
|
||||
mock_selector = Mock()
|
||||
mock_selector_class.return_value = mock_selector
|
||||
mock_selector.select.return_value = [True] # Always ready
|
||||
|
||||
result = send(mock_sock, b"test data")
|
||||
|
||||
self.assertEqual(result, 9) # len("test data")
|
||||
self.assertEqual(write_attempts[0], 2)
|
||||
|
||||
def test_ssl_eof_error_edge_cases(self):
|
||||
"""Test SSL EOF error edge cases"""
|
||||
mock_sock = Mock()
|
||||
|
||||
# Test SSLEOFError during send
|
||||
mock_sock.send.side_effect = SSLEOFError("SSL connection has been closed")
|
||||
mock_sock.gettimeout.return_value = 30.0
|
||||
|
||||
from websocket._exceptions import WebSocketConnectionClosedException
|
||||
|
||||
with self.assertRaises(WebSocketConnectionClosedException):
|
||||
send(mock_sock, b"test data")
|
||||
|
||||
def test_ssl_pending_data_edge_cases(self):
|
||||
"""Test SSL pending data scenarios"""
|
||||
from websocket._dispatcher import SSLDispatcher
|
||||
from websocket._app import WebSocketApp
|
||||
|
||||
# Mock SSL socket with pending data
|
||||
mock_ssl_sock = Mock()
|
||||
mock_ssl_sock.pending.return_value = 1024 # Simulates pending SSL data
|
||||
|
||||
# Mock WebSocketApp
|
||||
mock_app = Mock(spec=WebSocketApp)
|
||||
mock_app.sock = Mock()
|
||||
mock_app.sock.sock = mock_ssl_sock
|
||||
|
||||
dispatcher = SSLDispatcher(mock_app, 5.0)
|
||||
|
||||
# When there's pending data, should return immediately without selector
|
||||
result = dispatcher.select(mock_ssl_sock, Mock())
|
||||
|
||||
# Should return the socket list when there's pending data
|
||||
self.assertEqual(result, [mock_ssl_sock])
|
||||
mock_ssl_sock.pending.assert_called_once()
|
||||
|
||||
def test_ssl_renegotiation_edge_cases(self):
|
||||
"""Test SSL renegotiation scenarios"""
|
||||
mock_sock = Mock()
|
||||
|
||||
# Simulate SSL renegotiation during read
|
||||
call_count = 0
|
||||
|
||||
def mock_recv(bufsize):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if call_count == 1:
|
||||
raise SSLWantReadError("SSL renegotiation required")
|
||||
return b"data after renegotiation"
|
||||
|
||||
mock_sock.recv.side_effect = mock_recv
|
||||
mock_sock.gettimeout.return_value = 30.0
|
||||
|
||||
with patch("selectors.DefaultSelector") as mock_selector_class:
|
||||
mock_selector = Mock()
|
||||
mock_selector_class.return_value = mock_selector
|
||||
mock_selector.select.return_value = [True]
|
||||
|
||||
result = recv(mock_sock, 100)
|
||||
|
||||
self.assertEqual(result, b"data after renegotiation")
|
||||
self.assertEqual(call_count, 2)
|
||||
|
||||
def test_ssl_server_hostname_override(self):
|
||||
"""Test SSL server hostname override scenarios"""
|
||||
mock_sock = Mock()
|
||||
|
||||
with patch("ssl.SSLContext") as mock_ssl_context:
|
||||
mock_context = Mock()
|
||||
mock_ssl_context.return_value = mock_context
|
||||
mock_context.wrap_socket.return_value = Mock()
|
||||
|
||||
# Test server_hostname override
|
||||
sslopt = {"server_hostname": "override.example.com"}
|
||||
_ssl_socket(mock_sock, sslopt, "original.example.com")
|
||||
|
||||
# Should use override hostname in wrap_socket call
|
||||
mock_context.wrap_socket.assert_called_with(
|
||||
mock_sock,
|
||||
do_handshake_on_connect=True,
|
||||
suppress_ragged_eofs=True,
|
||||
server_hostname="override.example.com",
|
||||
)
|
||||
|
||||
def test_ssl_protocol_version_edge_cases(self):
|
||||
"""Test SSL protocol version edge cases"""
|
||||
mock_sock = Mock()
|
||||
|
||||
# Test with deprecated SSL version
|
||||
with patch("ssl.SSLContext") as mock_ssl_context:
|
||||
mock_context = Mock()
|
||||
mock_ssl_context.return_value = mock_context
|
||||
mock_context.wrap_socket.return_value = Mock()
|
||||
|
||||
# Test that deprecated ssl_version is still handled
|
||||
if hasattr(ssl, "PROTOCOL_TLS"):
|
||||
sslopt = {"ssl_version": ssl.PROTOCOL_TLS}
|
||||
_ssl_socket(mock_sock, sslopt, "example.com")
|
||||
|
||||
mock_ssl_context.assert_called_with(ssl.PROTOCOL_TLS)
|
||||
|
||||
def test_ssl_keylog_file_edge_cases(self):
|
||||
"""Test SSL keylog file configuration edge cases"""
|
||||
mock_sock = Mock()
|
||||
|
||||
# Test with SSLKEYLOGFILE environment variable
|
||||
with patch.dict("os.environ", {"SSLKEYLOGFILE": "/tmp/ssl_keys.log"}):
|
||||
with patch("ssl.SSLContext") as mock_ssl_context:
|
||||
mock_context = Mock()
|
||||
mock_ssl_context.return_value = mock_context
|
||||
mock_context.wrap_socket.return_value = Mock()
|
||||
|
||||
sslopt = {}
|
||||
_ssl_socket(mock_sock, sslopt, "example.com")
|
||||
|
||||
# Should set keylog_filename
|
||||
self.assertEqual(mock_context.keylog_filename, "/tmp/ssl_keys.log")
|
||||
|
||||
def test_ssl_context_verification_modes(self):
|
||||
"""Test different SSL verification mode combinations"""
|
||||
mock_sock = Mock()
|
||||
|
||||
test_cases = [
|
||||
# (cert_reqs, check_hostname, expected_verify_mode, expected_check_hostname)
|
||||
(ssl.CERT_NONE, False, ssl.CERT_NONE, False),
|
||||
(ssl.CERT_REQUIRED, False, ssl.CERT_REQUIRED, False),
|
||||
(ssl.CERT_REQUIRED, True, ssl.CERT_REQUIRED, True),
|
||||
]
|
||||
|
||||
for cert_reqs, check_hostname, expected_verify, expected_check in test_cases:
|
||||
with self.subTest(cert_reqs=cert_reqs, check_hostname=check_hostname):
|
||||
with patch("ssl.SSLContext") as mock_ssl_context:
|
||||
mock_context = Mock()
|
||||
mock_ssl_context.return_value = mock_context
|
||||
mock_context.wrap_socket.return_value = Mock()
|
||||
|
||||
sslopt = {"cert_reqs": cert_reqs, "check_hostname": check_hostname}
|
||||
_ssl_socket(mock_sock, sslopt, "example.com")
|
||||
|
||||
self.assertEqual(mock_context.verify_mode, expected_verify)
|
||||
self.assertEqual(mock_context.check_hostname, expected_check)
|
||||
|
||||
def test_ssl_socket_shutdown_edge_cases(self):
|
||||
"""Test SSL socket shutdown edge cases"""
|
||||
from websocket._core import WebSocket
|
||||
|
||||
mock_ssl_sock = Mock()
|
||||
mock_ssl_sock.shutdown.side_effect = SSLError("SSL shutdown failed")
|
||||
|
||||
ws = WebSocket()
|
||||
ws.sock = mock_ssl_sock
|
||||
ws.connected = True
|
||||
|
||||
# Should handle SSL shutdown errors gracefully
|
||||
try:
|
||||
ws.close()
|
||||
except SSLError:
|
||||
self.fail("SSL shutdown error should be handled gracefully")
|
||||
|
||||
def test_ssl_socket_close_during_operation(self):
|
||||
"""Test SSL socket being closed during ongoing operations"""
|
||||
mock_sock = Mock()
|
||||
|
||||
# Simulate SSL socket being closed during recv
|
||||
mock_sock.recv.side_effect = SSLError(
|
||||
"SSL connection has been closed unexpectedly"
|
||||
)
|
||||
mock_sock.gettimeout.return_value = 30.0
|
||||
|
||||
from websocket._exceptions import WebSocketConnectionClosedException
|
||||
|
||||
# Should handle unexpected SSL closure
|
||||
with self.assertRaises((SSLError, WebSocketConnectionClosedException)):
|
||||
recv(mock_sock, 100)
|
||||
|
||||
def test_ssl_compression_edge_cases(self):
|
||||
"""Test SSL compression configuration edge cases"""
|
||||
mock_sock = Mock()
|
||||
|
||||
with patch("ssl.SSLContext") as mock_ssl_context:
|
||||
mock_context = Mock()
|
||||
mock_ssl_context.return_value = mock_context
|
||||
mock_context.wrap_socket.return_value = Mock()
|
||||
|
||||
# Test SSL compression options (if available)
|
||||
sslopt = {"compression": False} # Some SSL contexts support this
|
||||
|
||||
try:
|
||||
_ssl_socket(mock_sock, sslopt, "example.com")
|
||||
# Should not fail even if compression option is not supported
|
||||
except AttributeError:
|
||||
# Expected if SSL context doesn't support compression option
|
||||
pass
|
||||
|
||||
def test_ssl_session_reuse_edge_cases(self):
|
||||
"""Test SSL session reuse scenarios"""
|
||||
mock_sock = Mock()
|
||||
|
||||
with patch("ssl.SSLContext") as mock_ssl_context:
|
||||
mock_context = Mock()
|
||||
mock_ssl_context.return_value = mock_context
|
||||
mock_ssl_sock = Mock()
|
||||
mock_context.wrap_socket.return_value = mock_ssl_sock
|
||||
|
||||
# Test session reuse
|
||||
mock_ssl_sock.session = "mock_session"
|
||||
mock_ssl_sock.session_reused = True
|
||||
|
||||
result = _ssl_socket(mock_sock, {}, "example.com")
|
||||
|
||||
# Should handle session reuse without issues
|
||||
self.assertIsNotNone(result)
|
||||
|
||||
def test_ssl_alpn_protocol_edge_cases(self):
|
||||
"""Test SSL ALPN (Application Layer Protocol Negotiation) edge cases"""
|
||||
mock_sock = Mock()
|
||||
|
||||
with patch("ssl.SSLContext") as mock_ssl_context:
|
||||
mock_context = Mock()
|
||||
mock_ssl_context.return_value = mock_context
|
||||
mock_context.wrap_socket.return_value = Mock()
|
||||
|
||||
# Test ALPN configuration
|
||||
sslopt = {"alpn_protocols": ["http/1.1", "h2"]}
|
||||
|
||||
# ALPN protocols are not currently supported in the SSL wrapper
|
||||
# but the test should not fail
|
||||
result = _ssl_socket(mock_sock, sslopt, "example.com")
|
||||
self.assertIsNotNone(result)
|
||||
# ALPN would need to be implemented in _wrap_sni_socket function
|
||||
|
||||
def test_ssl_sni_edge_cases(self):
|
||||
"""Test SSL SNI (Server Name Indication) edge cases"""
|
||||
mock_sock = Mock()
|
||||
|
||||
# Test with IPv6 address (should not use SNI)
|
||||
with patch("ssl.SSLContext") as mock_ssl_context:
|
||||
mock_context = Mock()
|
||||
mock_ssl_context.return_value = mock_context
|
||||
mock_context.wrap_socket.return_value = Mock()
|
||||
|
||||
# IPv6 addresses should not be used for SNI
|
||||
ipv6_hostname = "2001:db8::1"
|
||||
_ssl_socket(mock_sock, {}, ipv6_hostname)
|
||||
|
||||
# Should use IPv6 address as server_hostname
|
||||
mock_context.wrap_socket.assert_called_with(
|
||||
mock_sock,
|
||||
do_handshake_on_connect=True,
|
||||
suppress_ragged_eofs=True,
|
||||
server_hostname=ipv6_hostname,
|
||||
)
|
||||
|
||||
def test_ssl_buffer_size_edge_cases(self):
|
||||
"""Test SSL buffer size related edge cases"""
|
||||
mock_sock = Mock()
|
||||
|
||||
def mock_recv(bufsize):
|
||||
# SSL should never try to read more than 16KB at once
|
||||
if bufsize > 16384:
|
||||
raise SSLError("[SSL: BAD_LENGTH] buffer too large")
|
||||
return b"A" * min(bufsize, 1024) # Return smaller chunks
|
||||
|
||||
mock_sock.recv.side_effect = mock_recv
|
||||
mock_sock.gettimeout.return_value = 30.0
|
||||
|
||||
from websocket._abnf import frame_buffer
|
||||
|
||||
# Frame buffer should handle large requests by chunking
|
||||
fb = frame_buffer(lambda size: recv(mock_sock, size), skip_utf8_validation=True)
|
||||
|
||||
# This should work even with large size due to chunking
|
||||
result = fb.recv_strict(16384) # Exactly 16KB
|
||||
|
||||
self.assertGreater(len(result), 0)
|
||||
|
||||
def test_ssl_protocol_downgrade_protection(self):
|
||||
"""Test SSL protocol downgrade protection"""
|
||||
mock_sock = Mock()
|
||||
|
||||
with patch("ssl.SSLContext") as mock_ssl_context:
|
||||
mock_context = Mock()
|
||||
mock_ssl_context.return_value = mock_context
|
||||
mock_context.wrap_socket.side_effect = ssl.SSLError(
|
||||
"SSLV3_ALERT_HANDSHAKE_FAILURE"
|
||||
)
|
||||
|
||||
sslopt = {"ssl_version": ssl.PROTOCOL_TLS_CLIENT}
|
||||
|
||||
# Should propagate SSL protocol errors
|
||||
with self.assertRaises(ssl.SSLError):
|
||||
_ssl_socket(mock_sock, sslopt, "example.com")
|
||||
|
||||
def test_ssl_certificate_chain_validation(self):
|
||||
"""Test SSL certificate chain validation edge cases"""
|
||||
mock_sock = Mock()
|
||||
|
||||
with patch("ssl.SSLContext") as mock_ssl_context:
|
||||
mock_context = Mock()
|
||||
mock_ssl_context.return_value = mock_context
|
||||
|
||||
# Test certificate chain validation failure
|
||||
mock_context.wrap_socket.side_effect = ssl.SSLCertVerificationError(
|
||||
"certificate verify failed: certificate has expired"
|
||||
)
|
||||
|
||||
sslopt = {"cert_reqs": ssl.CERT_REQUIRED, "check_hostname": True}
|
||||
|
||||
with self.assertRaises(ssl.SSLCertVerificationError):
|
||||
_ssl_socket(mock_sock, sslopt, "expired.badssl.com")
|
||||
|
||||
def test_ssl_weak_cipher_rejection(self):
|
||||
"""Test SSL weak cipher rejection scenarios"""
|
||||
mock_sock = Mock()
|
||||
|
||||
with patch("ssl.SSLContext") as mock_ssl_context:
|
||||
mock_context = Mock()
|
||||
mock_ssl_context.return_value = mock_context
|
||||
mock_context.wrap_socket.side_effect = ssl.SSLError("no shared cipher")
|
||||
|
||||
sslopt = {"ciphers": "RC4-MD5"} # Intentionally weak cipher
|
||||
|
||||
# Should fail with weak ciphers (SSL error is not wrapped by our code)
|
||||
with self.assertRaises(ssl.SSLError):
|
||||
_ssl_socket(mock_sock, sslopt, "example.com")
|
||||
|
||||
def test_ssl_hostname_verification_edge_cases(self):
|
||||
"""Test SSL hostname verification edge cases"""
|
||||
mock_sock = Mock()
|
||||
|
||||
# Test with wildcard certificate scenarios
|
||||
test_cases = [
|
||||
("*.example.com", "subdomain.example.com"), # Valid wildcard
|
||||
("*.example.com", "sub.subdomain.example.com"), # Invalid wildcard depth
|
||||
("example.com", "www.example.com"), # Hostname mismatch
|
||||
]
|
||||
|
||||
for cert_hostname, connect_hostname in test_cases:
|
||||
with self.subTest(cert=cert_hostname, hostname=connect_hostname):
|
||||
with patch("ssl.SSLContext") as mock_ssl_context:
|
||||
mock_context = Mock()
|
||||
mock_ssl_context.return_value = mock_context
|
||||
|
||||
if (
|
||||
cert_hostname != connect_hostname
|
||||
and "sub.subdomain" in connect_hostname
|
||||
):
|
||||
# Simulate hostname verification failure for invalid wildcard
|
||||
mock_context.wrap_socket.side_effect = ssl.SSLCertVerificationError(
|
||||
f"hostname '{connect_hostname}' doesn't match '{cert_hostname}'"
|
||||
)
|
||||
|
||||
sslopt = {
|
||||
"cert_reqs": ssl.CERT_REQUIRED,
|
||||
"check_hostname": True,
|
||||
}
|
||||
|
||||
with self.assertRaises(ssl.SSLCertVerificationError):
|
||||
_ssl_socket(mock_sock, sslopt, connect_hostname)
|
||||
else:
|
||||
mock_context.wrap_socket.return_value = Mock()
|
||||
sslopt = {
|
||||
"cert_reqs": ssl.CERT_REQUIRED,
|
||||
"check_hostname": True,
|
||||
}
|
||||
|
||||
# Should succeed for valid cases
|
||||
result = _ssl_socket(mock_sock, sslopt, connect_hostname)
|
||||
self.assertIsNotNone(result)
|
||||
|
||||
def test_ssl_memory_bio_edge_cases(self):
|
||||
"""Test SSL memory BIO edge cases"""
|
||||
mock_sock = Mock()
|
||||
|
||||
# Test SSL memory BIO scenarios (if available)
|
||||
try:
|
||||
import ssl
|
||||
|
||||
if hasattr(ssl, "MemoryBIO"):
|
||||
with patch("ssl.SSLContext") as mock_ssl_context:
|
||||
mock_context = Mock()
|
||||
mock_ssl_context.return_value = mock_context
|
||||
mock_context.wrap_socket.return_value = Mock()
|
||||
|
||||
# Memory BIO should work if available
|
||||
_ssl_socket(mock_sock, {}, "example.com")
|
||||
|
||||
# Standard socket wrapping should still work
|
||||
mock_context.wrap_socket.assert_called_once()
|
||||
except (ImportError, AttributeError):
|
||||
self.skipTest("SSL MemoryBIO not available")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,471 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
#
|
||||
import os
|
||||
import unittest
|
||||
|
||||
from websocket._url import (
|
||||
_is_address_in_network,
|
||||
_is_no_proxy_host,
|
||||
get_proxy_info,
|
||||
parse_url,
|
||||
)
|
||||
from websocket._exceptions import WebSocketProxyException
|
||||
|
||||
"""
|
||||
test_url.py
|
||||
websocket - WebSocket client library for Python
|
||||
|
||||
Copyright 2025 engn33r
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
|
||||
class UrlTest(unittest.TestCase):
|
||||
def test_address_in_network(self):
|
||||
self.assertTrue(_is_address_in_network("127.0.0.1", "127.0.0.0/8"))
|
||||
self.assertTrue(_is_address_in_network("127.1.0.1", "127.0.0.0/8"))
|
||||
self.assertFalse(_is_address_in_network("127.1.0.1", "127.0.0.0/24"))
|
||||
self.assertTrue(_is_address_in_network("2001:db8::1", "2001:db8::/64"))
|
||||
self.assertFalse(_is_address_in_network("2001:db8:1::1", "2001:db8::/64"))
|
||||
|
||||
def test_parse_url(self):
|
||||
p = parse_url("ws://www.example.com/r")
|
||||
self.assertEqual(p[0], "www.example.com")
|
||||
self.assertEqual(p[1], 80)
|
||||
self.assertEqual(p[2], "/r")
|
||||
self.assertEqual(p[3], False)
|
||||
|
||||
p = parse_url("ws://www.example.com/r/")
|
||||
self.assertEqual(p[0], "www.example.com")
|
||||
self.assertEqual(p[1], 80)
|
||||
self.assertEqual(p[2], "/r/")
|
||||
self.assertEqual(p[3], False)
|
||||
|
||||
p = parse_url("ws://www.example.com/")
|
||||
self.assertEqual(p[0], "www.example.com")
|
||||
self.assertEqual(p[1], 80)
|
||||
self.assertEqual(p[2], "/")
|
||||
self.assertEqual(p[3], False)
|
||||
|
||||
p = parse_url("ws://www.example.com")
|
||||
self.assertEqual(p[0], "www.example.com")
|
||||
self.assertEqual(p[1], 80)
|
||||
self.assertEqual(p[2], "/")
|
||||
self.assertEqual(p[3], False)
|
||||
|
||||
p = parse_url("ws://www.example.com:8080/r")
|
||||
self.assertEqual(p[0], "www.example.com")
|
||||
self.assertEqual(p[1], 8080)
|
||||
self.assertEqual(p[2], "/r")
|
||||
self.assertEqual(p[3], False)
|
||||
|
||||
p = parse_url("ws://www.example.com:8080/")
|
||||
self.assertEqual(p[0], "www.example.com")
|
||||
self.assertEqual(p[1], 8080)
|
||||
self.assertEqual(p[2], "/")
|
||||
self.assertEqual(p[3], False)
|
||||
|
||||
p = parse_url("ws://www.example.com:8080")
|
||||
self.assertEqual(p[0], "www.example.com")
|
||||
self.assertEqual(p[1], 8080)
|
||||
self.assertEqual(p[2], "/")
|
||||
self.assertEqual(p[3], False)
|
||||
|
||||
p = parse_url("wss://www.example.com:8080/r")
|
||||
self.assertEqual(p[0], "www.example.com")
|
||||
self.assertEqual(p[1], 8080)
|
||||
self.assertEqual(p[2], "/r")
|
||||
self.assertEqual(p[3], True)
|
||||
|
||||
p = parse_url("wss://www.example.com:8080/r?key=value")
|
||||
self.assertEqual(p[0], "www.example.com")
|
||||
self.assertEqual(p[1], 8080)
|
||||
self.assertEqual(p[2], "/r?key=value")
|
||||
self.assertEqual(p[3], True)
|
||||
|
||||
self.assertRaises(ValueError, parse_url, "http://www.example.com/r")
|
||||
|
||||
p = parse_url("ws://[2a03:4000:123:83::3]/r")
|
||||
self.assertEqual(p[0], "2a03:4000:123:83::3")
|
||||
self.assertEqual(p[1], 80)
|
||||
self.assertEqual(p[2], "/r")
|
||||
self.assertEqual(p[3], False)
|
||||
|
||||
p = parse_url("ws://[2a03:4000:123:83::3]:8080/r")
|
||||
self.assertEqual(p[0], "2a03:4000:123:83::3")
|
||||
self.assertEqual(p[1], 8080)
|
||||
self.assertEqual(p[2], "/r")
|
||||
self.assertEqual(p[3], False)
|
||||
|
||||
p = parse_url("wss://[2a03:4000:123:83::3]/r")
|
||||
self.assertEqual(p[0], "2a03:4000:123:83::3")
|
||||
self.assertEqual(p[1], 443)
|
||||
self.assertEqual(p[2], "/r")
|
||||
self.assertEqual(p[3], True)
|
||||
|
||||
p = parse_url("wss://[2a03:4000:123:83::3]:8080/r")
|
||||
self.assertEqual(p[0], "2a03:4000:123:83::3")
|
||||
self.assertEqual(p[1], 8080)
|
||||
self.assertEqual(p[2], "/r")
|
||||
self.assertEqual(p[3], True)
|
||||
|
||||
|
||||
class IsNoProxyHostTest(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.no_proxy = os.environ.get("no_proxy", None)
|
||||
if "no_proxy" in os.environ:
|
||||
del os.environ["no_proxy"]
|
||||
|
||||
def tearDown(self):
|
||||
if self.no_proxy:
|
||||
os.environ["no_proxy"] = self.no_proxy
|
||||
elif "no_proxy" in os.environ:
|
||||
del os.environ["no_proxy"]
|
||||
|
||||
def test_match_all(self):
|
||||
self.assertTrue(_is_no_proxy_host("any.websocket.org", ["*"]))
|
||||
self.assertTrue(_is_no_proxy_host("192.168.0.1", ["*"]))
|
||||
self.assertFalse(_is_no_proxy_host("192.168.0.1", ["192.168.1.1"]))
|
||||
self.assertFalse(
|
||||
_is_no_proxy_host("any.websocket.org", ["other.websocket.org"])
|
||||
)
|
||||
self.assertTrue(
|
||||
_is_no_proxy_host("any.websocket.org", ["other.websocket.org", "*"])
|
||||
)
|
||||
os.environ["no_proxy"] = "*"
|
||||
self.assertTrue(_is_no_proxy_host("any.websocket.org", None))
|
||||
self.assertTrue(_is_no_proxy_host("192.168.0.1", None))
|
||||
os.environ["no_proxy"] = "other.websocket.org, *"
|
||||
self.assertTrue(_is_no_proxy_host("any.websocket.org", None))
|
||||
|
||||
def test_ip_address(self):
|
||||
self.assertTrue(_is_no_proxy_host("127.0.0.1", ["127.0.0.1"]))
|
||||
self.assertFalse(_is_no_proxy_host("127.0.0.2", ["127.0.0.1"]))
|
||||
self.assertTrue(
|
||||
_is_no_proxy_host("127.0.0.1", ["other.websocket.org", "127.0.0.1"])
|
||||
)
|
||||
self.assertFalse(
|
||||
_is_no_proxy_host("127.0.0.2", ["other.websocket.org", "127.0.0.1"])
|
||||
)
|
||||
os.environ["no_proxy"] = "127.0.0.1"
|
||||
self.assertTrue(_is_no_proxy_host("127.0.0.1", None))
|
||||
self.assertFalse(_is_no_proxy_host("127.0.0.2", None))
|
||||
os.environ["no_proxy"] = "other.websocket.org, 127.0.0.1"
|
||||
self.assertTrue(_is_no_proxy_host("127.0.0.1", None))
|
||||
self.assertFalse(_is_no_proxy_host("127.0.0.2", None))
|
||||
|
||||
def test_ip_address_in_range(self):
|
||||
self.assertTrue(_is_no_proxy_host("127.0.0.1", ["127.0.0.0/8"]))
|
||||
self.assertTrue(_is_no_proxy_host("127.0.0.2", ["127.0.0.0/8"]))
|
||||
self.assertFalse(_is_no_proxy_host("127.1.0.1", ["127.0.0.0/24"]))
|
||||
self.assertTrue(_is_no_proxy_host("2001:db8::1", ["2001:db8::/64"]))
|
||||
self.assertFalse(_is_no_proxy_host("2001:db8:1::1", ["2001:db8::/64"]))
|
||||
os.environ["no_proxy"] = "127.0.0.0/8,2001:db8::/64"
|
||||
self.assertTrue(_is_no_proxy_host("127.0.0.1", None))
|
||||
self.assertTrue(_is_no_proxy_host("127.0.0.2", None))
|
||||
self.assertTrue(_is_no_proxy_host("2001:db8::1", None))
|
||||
self.assertFalse(_is_no_proxy_host("2001:db8:1::1", None))
|
||||
os.environ["no_proxy"] = "127.0.0.0/24,2001:db8::/64"
|
||||
self.assertFalse(_is_no_proxy_host("127.1.0.1", None))
|
||||
self.assertFalse(_is_no_proxy_host("2001:db8:1::1", None))
|
||||
|
||||
def test_hostname_match(self):
|
||||
self.assertTrue(_is_no_proxy_host("my.websocket.org", ["my.websocket.org"]))
|
||||
self.assertTrue(
|
||||
_is_no_proxy_host(
|
||||
"my.websocket.org", ["other.websocket.org", "my.websocket.org"]
|
||||
)
|
||||
)
|
||||
self.assertFalse(_is_no_proxy_host("my.websocket.org", ["other.websocket.org"]))
|
||||
os.environ["no_proxy"] = "my.websocket.org"
|
||||
self.assertTrue(_is_no_proxy_host("my.websocket.org", None))
|
||||
self.assertFalse(_is_no_proxy_host("other.websocket.org", None))
|
||||
os.environ["no_proxy"] = "other.websocket.org, my.websocket.org"
|
||||
self.assertTrue(_is_no_proxy_host("my.websocket.org", None))
|
||||
|
||||
def test_hostname_match_domain(self):
|
||||
self.assertTrue(_is_no_proxy_host("any.websocket.org", [".websocket.org"]))
|
||||
self.assertTrue(_is_no_proxy_host("my.other.websocket.org", [".websocket.org"]))
|
||||
self.assertTrue(
|
||||
_is_no_proxy_host(
|
||||
"any.websocket.org", ["my.websocket.org", ".websocket.org"]
|
||||
)
|
||||
)
|
||||
self.assertFalse(_is_no_proxy_host("any.websocket.com", [".websocket.org"]))
|
||||
os.environ["no_proxy"] = ".websocket.org"
|
||||
self.assertTrue(_is_no_proxy_host("any.websocket.org", None))
|
||||
self.assertTrue(_is_no_proxy_host("my.other.websocket.org", None))
|
||||
self.assertFalse(_is_no_proxy_host("any.websocket.com", None))
|
||||
os.environ["no_proxy"] = "my.websocket.org, .websocket.org"
|
||||
self.assertTrue(_is_no_proxy_host("any.websocket.org", None))
|
||||
|
||||
|
||||
class ProxyInfoTest(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.http_proxy = os.environ.get("http_proxy", None)
|
||||
self.https_proxy = os.environ.get("https_proxy", None)
|
||||
self.no_proxy = os.environ.get("no_proxy", None)
|
||||
if "http_proxy" in os.environ:
|
||||
del os.environ["http_proxy"]
|
||||
if "https_proxy" in os.environ:
|
||||
del os.environ["https_proxy"]
|
||||
if "no_proxy" in os.environ:
|
||||
del os.environ["no_proxy"]
|
||||
|
||||
def tearDown(self):
|
||||
if self.http_proxy:
|
||||
os.environ["http_proxy"] = self.http_proxy
|
||||
elif "http_proxy" in os.environ:
|
||||
del os.environ["http_proxy"]
|
||||
|
||||
if self.https_proxy:
|
||||
os.environ["https_proxy"] = self.https_proxy
|
||||
elif "https_proxy" in os.environ:
|
||||
del os.environ["https_proxy"]
|
||||
|
||||
if self.no_proxy:
|
||||
os.environ["no_proxy"] = self.no_proxy
|
||||
elif "no_proxy" in os.environ:
|
||||
del os.environ["no_proxy"]
|
||||
|
||||
def test_proxy_from_args(self):
|
||||
self.assertRaises(
|
||||
WebSocketProxyException,
|
||||
get_proxy_info,
|
||||
"echo.websocket.events",
|
||||
False,
|
||||
proxy_host="localhost",
|
||||
)
|
||||
self.assertEqual(
|
||||
get_proxy_info(
|
||||
"echo.websocket.events", False, proxy_host="localhost", proxy_port=3128
|
||||
),
|
||||
("localhost", 3128, None),
|
||||
)
|
||||
self.assertEqual(
|
||||
get_proxy_info(
|
||||
"echo.websocket.events", True, proxy_host="localhost", proxy_port=3128
|
||||
),
|
||||
("localhost", 3128, None),
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
get_proxy_info(
|
||||
"echo.websocket.events",
|
||||
False,
|
||||
proxy_host="localhost",
|
||||
proxy_port=9001,
|
||||
proxy_auth=("a", "b"),
|
||||
),
|
||||
("localhost", 9001, ("a", "b")),
|
||||
)
|
||||
self.assertEqual(
|
||||
get_proxy_info(
|
||||
"echo.websocket.events",
|
||||
False,
|
||||
proxy_host="localhost",
|
||||
proxy_port=3128,
|
||||
proxy_auth=("a", "b"),
|
||||
),
|
||||
("localhost", 3128, ("a", "b")),
|
||||
)
|
||||
self.assertEqual(
|
||||
get_proxy_info(
|
||||
"echo.websocket.events",
|
||||
True,
|
||||
proxy_host="localhost",
|
||||
proxy_port=8765,
|
||||
proxy_auth=("a", "b"),
|
||||
),
|
||||
("localhost", 8765, ("a", "b")),
|
||||
)
|
||||
self.assertEqual(
|
||||
get_proxy_info(
|
||||
"echo.websocket.events",
|
||||
True,
|
||||
proxy_host="localhost",
|
||||
proxy_port=3128,
|
||||
proxy_auth=("a", "b"),
|
||||
),
|
||||
("localhost", 3128, ("a", "b")),
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
get_proxy_info(
|
||||
"echo.websocket.events",
|
||||
True,
|
||||
proxy_host="localhost",
|
||||
proxy_port=3128,
|
||||
no_proxy=["example.com"],
|
||||
proxy_auth=("a", "b"),
|
||||
),
|
||||
("localhost", 3128, ("a", "b")),
|
||||
)
|
||||
self.assertEqual(
|
||||
get_proxy_info(
|
||||
"echo.websocket.events",
|
||||
True,
|
||||
proxy_host="localhost",
|
||||
proxy_port=3128,
|
||||
no_proxy=["echo.websocket.events"],
|
||||
proxy_auth=("a", "b"),
|
||||
),
|
||||
(None, 0, None),
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
get_proxy_info(
|
||||
"echo.websocket.events",
|
||||
True,
|
||||
proxy_host="localhost",
|
||||
proxy_port=3128,
|
||||
no_proxy=[".websocket.events"],
|
||||
),
|
||||
(None, 0, None),
|
||||
)
|
||||
|
||||
def test_proxy_from_env(self):
|
||||
os.environ["http_proxy"] = "http://localhost/"
|
||||
self.assertEqual(
|
||||
get_proxy_info("echo.websocket.events", False), ("localhost", None, None)
|
||||
)
|
||||
os.environ["http_proxy"] = "http://localhost:3128/"
|
||||
self.assertEqual(
|
||||
get_proxy_info("echo.websocket.events", False), ("localhost", 3128, None)
|
||||
)
|
||||
|
||||
os.environ["http_proxy"] = "http://localhost/"
|
||||
os.environ["https_proxy"] = "http://localhost2/"
|
||||
self.assertEqual(
|
||||
get_proxy_info("echo.websocket.events", False), ("localhost", None, None)
|
||||
)
|
||||
os.environ["http_proxy"] = "http://localhost:3128/"
|
||||
os.environ["https_proxy"] = "http://localhost2:3128/"
|
||||
self.assertEqual(
|
||||
get_proxy_info("echo.websocket.events", False), ("localhost", 3128, None)
|
||||
)
|
||||
|
||||
os.environ["http_proxy"] = "http://localhost/"
|
||||
os.environ["https_proxy"] = "http://localhost2/"
|
||||
self.assertEqual(
|
||||
get_proxy_info("echo.websocket.events", True), ("localhost2", None, None)
|
||||
)
|
||||
os.environ["http_proxy"] = "http://localhost:3128/"
|
||||
os.environ["https_proxy"] = "http://localhost2:3128/"
|
||||
self.assertEqual(
|
||||
get_proxy_info("echo.websocket.events", True), ("localhost2", 3128, None)
|
||||
)
|
||||
|
||||
os.environ["http_proxy"] = ""
|
||||
os.environ["https_proxy"] = "http://localhost2/"
|
||||
self.assertEqual(
|
||||
get_proxy_info("echo.websocket.events", True), ("localhost2", None, None)
|
||||
)
|
||||
self.assertEqual(
|
||||
get_proxy_info("echo.websocket.events", False), (None, 0, None)
|
||||
)
|
||||
os.environ["http_proxy"] = ""
|
||||
os.environ["https_proxy"] = "http://localhost2:3128/"
|
||||
self.assertEqual(
|
||||
get_proxy_info("echo.websocket.events", True), ("localhost2", 3128, None)
|
||||
)
|
||||
self.assertEqual(
|
||||
get_proxy_info("echo.websocket.events", False), (None, 0, None)
|
||||
)
|
||||
|
||||
os.environ["http_proxy"] = "http://localhost/"
|
||||
os.environ["https_proxy"] = ""
|
||||
self.assertEqual(get_proxy_info("echo.websocket.events", True), (None, 0, None))
|
||||
self.assertEqual(
|
||||
get_proxy_info("echo.websocket.events", False), ("localhost", None, None)
|
||||
)
|
||||
os.environ["http_proxy"] = "http://localhost:3128/"
|
||||
os.environ["https_proxy"] = ""
|
||||
self.assertEqual(get_proxy_info("echo.websocket.events", True), (None, 0, None))
|
||||
self.assertEqual(
|
||||
get_proxy_info("echo.websocket.events", False), ("localhost", 3128, None)
|
||||
)
|
||||
|
||||
os.environ["http_proxy"] = "http://a:b@localhost/"
|
||||
self.assertEqual(
|
||||
get_proxy_info("echo.websocket.events", False),
|
||||
("localhost", None, ("a", "b")),
|
||||
)
|
||||
os.environ["http_proxy"] = "http://a:b@localhost:3128/"
|
||||
self.assertEqual(
|
||||
get_proxy_info("echo.websocket.events", False),
|
||||
("localhost", 3128, ("a", "b")),
|
||||
)
|
||||
|
||||
os.environ["http_proxy"] = "http://a:b@localhost/"
|
||||
os.environ["https_proxy"] = "http://a:b@localhost2/"
|
||||
self.assertEqual(
|
||||
get_proxy_info("echo.websocket.events", False),
|
||||
("localhost", None, ("a", "b")),
|
||||
)
|
||||
os.environ["http_proxy"] = "http://a:b@localhost:3128/"
|
||||
os.environ["https_proxy"] = "http://a:b@localhost2:3128/"
|
||||
self.assertEqual(
|
||||
get_proxy_info("echo.websocket.events", False),
|
||||
("localhost", 3128, ("a", "b")),
|
||||
)
|
||||
|
||||
os.environ["http_proxy"] = "http://a:b@localhost/"
|
||||
os.environ["https_proxy"] = "http://a:b@localhost2/"
|
||||
self.assertEqual(
|
||||
get_proxy_info("echo.websocket.events", True),
|
||||
("localhost2", None, ("a", "b")),
|
||||
)
|
||||
os.environ["http_proxy"] = "http://a:b@localhost:3128/"
|
||||
os.environ["https_proxy"] = "http://a:b@localhost2:3128/"
|
||||
self.assertEqual(
|
||||
get_proxy_info("echo.websocket.events", True),
|
||||
("localhost2", 3128, ("a", "b")),
|
||||
)
|
||||
|
||||
os.environ["http_proxy"] = (
|
||||
"http://john%40example.com:P%40SSWORD@localhost:3128/"
|
||||
)
|
||||
os.environ["https_proxy"] = (
|
||||
"http://john%40example.com:P%40SSWORD@localhost2:3128/"
|
||||
)
|
||||
self.assertEqual(
|
||||
get_proxy_info("echo.websocket.events", True),
|
||||
("localhost2", 3128, ("john@example.com", "P@SSWORD")),
|
||||
)
|
||||
|
||||
os.environ["http_proxy"] = "http://a:b@localhost/"
|
||||
os.environ["https_proxy"] = "http://a:b@localhost2/"
|
||||
os.environ["no_proxy"] = "example1.com,example2.com"
|
||||
self.assertEqual(
|
||||
get_proxy_info("example.1.com", True), ("localhost2", None, ("a", "b"))
|
||||
)
|
||||
os.environ["http_proxy"] = "http://a:b@localhost:3128/"
|
||||
os.environ["https_proxy"] = "http://a:b@localhost2:3128/"
|
||||
os.environ["no_proxy"] = "example1.com,example2.com, echo.websocket.events"
|
||||
self.assertEqual(get_proxy_info("echo.websocket.events", True), (None, 0, None))
|
||||
os.environ["http_proxy"] = "http://a:b@localhost:3128/"
|
||||
os.environ["https_proxy"] = "http://a:b@localhost2:3128/"
|
||||
os.environ["no_proxy"] = "example1.com,example2.com, .websocket.events"
|
||||
self.assertEqual(get_proxy_info("echo.websocket.events", True), (None, 0, None))
|
||||
|
||||
os.environ["http_proxy"] = "http://a:b@localhost:3128/"
|
||||
os.environ["https_proxy"] = "http://a:b@localhost2:3128/"
|
||||
os.environ["no_proxy"] = "127.0.0.0/8, 192.168.0.0/16"
|
||||
self.assertEqual(get_proxy_info("127.0.0.1", False), (None, 0, None))
|
||||
self.assertEqual(get_proxy_info("192.168.1.1", False), (None, 0, None))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,138 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
import sys
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
"""
|
||||
test_utils.py
|
||||
websocket - WebSocket client library for Python
|
||||
|
||||
Copyright 2025 engn33r
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
class UtilsTest(unittest.TestCase):
|
||||
def test_nolock(self):
|
||||
"""Test NoLock context manager"""
|
||||
from websocket._utils import NoLock
|
||||
|
||||
lock = NoLock()
|
||||
|
||||
# Test that it can be used as context manager
|
||||
with lock:
|
||||
pass # Should not raise any exception
|
||||
|
||||
# Test enter/exit methods directly
|
||||
self.assertIsNone(lock.__enter__())
|
||||
self.assertIsNone(lock.__exit__(None, None, None))
|
||||
|
||||
def test_utf8_validation_with_wsaccel(self):
|
||||
"""Test UTF-8 validation when wsaccel is available"""
|
||||
# Import normally (wsaccel should be available in test environment)
|
||||
from websocket._utils import validate_utf8
|
||||
|
||||
# Test valid UTF-8 strings (convert to bytes for wsaccel)
|
||||
self.assertTrue(validate_utf8("Hello, World!".encode("utf-8")))
|
||||
self.assertTrue(validate_utf8("🌟 Unicode test".encode("utf-8")))
|
||||
self.assertTrue(validate_utf8(b"Hello, bytes"))
|
||||
self.assertTrue(validate_utf8("Héllo with accénts".encode("utf-8")))
|
||||
|
||||
# Test invalid UTF-8 sequences
|
||||
self.assertFalse(validate_utf8(b"\xff\xfe")) # Invalid UTF-8
|
||||
self.assertFalse(validate_utf8(b"\x80\x80")) # Invalid continuation
|
||||
|
||||
def test_utf8_validation_fallback(self):
|
||||
"""Test UTF-8 validation fallback when wsaccel is not available"""
|
||||
# Remove _utils from modules to force reimport
|
||||
if "websocket._utils" in sys.modules:
|
||||
del sys.modules["websocket._utils"]
|
||||
|
||||
# Mock wsaccel import to raise ImportError
|
||||
import builtins
|
||||
|
||||
original_import = builtins.__import__
|
||||
|
||||
def mock_import(name, *args, **kwargs):
|
||||
if "wsaccel" in name:
|
||||
raise ImportError(f"No module named '{name}'")
|
||||
return original_import(name, *args, **kwargs)
|
||||
|
||||
with patch("builtins.__import__", side_effect=mock_import):
|
||||
import websocket._utils as utils
|
||||
|
||||
# Test valid UTF-8 strings with fallback implementation (convert strings to bytes)
|
||||
self.assertTrue(utils.validate_utf8("Hello, World!".encode("utf-8")))
|
||||
self.assertTrue(utils.validate_utf8(b"Hello, bytes"))
|
||||
self.assertTrue(utils.validate_utf8("ASCII text".encode("utf-8")))
|
||||
|
||||
# Test Unicode strings (convert to bytes)
|
||||
self.assertTrue(utils.validate_utf8("🌟 Unicode test".encode("utf-8")))
|
||||
self.assertTrue(utils.validate_utf8("Héllo with accénts".encode("utf-8")))
|
||||
|
||||
# Test empty string/bytes
|
||||
self.assertTrue(utils.validate_utf8("".encode("utf-8")))
|
||||
self.assertTrue(utils.validate_utf8(b""))
|
||||
|
||||
# Test invalid UTF-8 sequences (should return False)
|
||||
self.assertFalse(utils.validate_utf8(b"\xff\xfe"))
|
||||
self.assertFalse(utils.validate_utf8(b"\x80\x80"))
|
||||
|
||||
# Note: The fallback implementation may have different validation behavior
|
||||
# than wsaccel, so we focus on clearly invalid sequences
|
||||
|
||||
def test_extract_err_message(self):
|
||||
"""Test extract_err_message function"""
|
||||
from websocket._utils import extract_err_message
|
||||
|
||||
# Test with exception that has args
|
||||
exc_with_args = Exception("Test error message")
|
||||
self.assertEqual(extract_err_message(exc_with_args), "Test error message")
|
||||
|
||||
# Test with exception that has multiple args
|
||||
exc_multi_args = Exception("First arg", "Second arg")
|
||||
self.assertEqual(extract_err_message(exc_multi_args), "First arg")
|
||||
|
||||
# Test with exception that has no args
|
||||
exc_no_args = Exception()
|
||||
self.assertIsNone(extract_err_message(exc_no_args))
|
||||
|
||||
def test_extract_error_code(self):
|
||||
"""Test extract_error_code function"""
|
||||
from websocket._utils import extract_error_code
|
||||
|
||||
# Test with exception that has integer as first arg
|
||||
exc_with_code = Exception(404, "Not found")
|
||||
self.assertEqual(extract_error_code(exc_with_code), 404)
|
||||
|
||||
# Test with exception that has string as first arg
|
||||
exc_with_string = Exception("Error message", "Second arg")
|
||||
self.assertIsNone(extract_error_code(exc_with_string))
|
||||
|
||||
# Test with exception that has only one arg
|
||||
exc_single_arg = Exception("Single arg")
|
||||
self.assertIsNone(extract_error_code(exc_single_arg))
|
||||
|
||||
# Test with exception that has no args
|
||||
exc_no_args = Exception()
|
||||
self.assertIsNone(extract_error_code(exc_no_args))
|
||||
|
||||
def tearDown(self):
|
||||
"""Clean up after tests"""
|
||||
# Ensure _utils is reimported fresh for next test
|
||||
if "websocket._utils" in sys.modules:
|
||||
del sys.modules["websocket._utils"]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,501 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
#
|
||||
import os
|
||||
import os.path
|
||||
import socket
|
||||
import unittest
|
||||
from base64 import decodebytes as base64decode
|
||||
|
||||
import websocket as ws
|
||||
from websocket._exceptions import (
|
||||
WebSocketBadStatusException,
|
||||
WebSocketAddressException,
|
||||
WebSocketException,
|
||||
)
|
||||
from websocket._handshake import _create_sec_websocket_key
|
||||
from websocket._handshake import _validate as _validate_header
|
||||
from websocket._http import read_headers
|
||||
from websocket._utils import validate_utf8
|
||||
|
||||
"""
|
||||
test_websocket.py
|
||||
websocket - WebSocket client library for Python
|
||||
|
||||
Copyright 2025 engn33r
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
try:
|
||||
import ssl
|
||||
except ImportError:
|
||||
# dummy class of SSLError for ssl none-support environment.
|
||||
class SSLError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
# Skip test to access the internet unless TEST_WITH_INTERNET == 1
|
||||
TEST_WITH_INTERNET = os.environ.get("TEST_WITH_INTERNET", "0") == "1"
|
||||
# Skip tests relying on local websockets server unless LOCAL_WS_SERVER_PORT != -1
|
||||
LOCAL_WS_SERVER_PORT = os.environ.get("LOCAL_WS_SERVER_PORT", "-1")
|
||||
TEST_WITH_LOCAL_SERVER = LOCAL_WS_SERVER_PORT != "-1"
|
||||
TRACEABLE = True
|
||||
|
||||
|
||||
def create_mask_key(_):
|
||||
return "abcd"
|
||||
|
||||
|
||||
class SockMock:
|
||||
def __init__(self):
|
||||
self.data = []
|
||||
self.sent = []
|
||||
|
||||
def add_packet(self, data):
|
||||
self.data.append(data)
|
||||
|
||||
def gettimeout(self):
|
||||
return None
|
||||
|
||||
def recv(self, bufsize):
|
||||
if self.data:
|
||||
e = self.data.pop(0)
|
||||
if isinstance(e, Exception):
|
||||
raise e
|
||||
if len(e) > bufsize:
|
||||
self.data.insert(0, e[bufsize:])
|
||||
return e[:bufsize]
|
||||
|
||||
def send(self, data):
|
||||
self.sent.append(data)
|
||||
return len(data)
|
||||
|
||||
def close(self):
|
||||
pass
|
||||
|
||||
|
||||
class HeaderSockMock(SockMock):
|
||||
def __init__(self, fname):
|
||||
SockMock.__init__(self)
|
||||
path = os.path.join(os.path.dirname(__file__), fname)
|
||||
with open(path, "rb") as f:
|
||||
self.add_packet(f.read())
|
||||
|
||||
|
||||
class WebSocketTest(unittest.TestCase):
|
||||
def setUp(self):
|
||||
ws.enableTrace(TRACEABLE)
|
||||
|
||||
def tearDown(self):
|
||||
pass
|
||||
|
||||
def test_default_timeout(self):
|
||||
self.assertEqual(ws.getdefaulttimeout(), None)
|
||||
ws.setdefaulttimeout(10)
|
||||
self.assertEqual(ws.getdefaulttimeout(), 10)
|
||||
ws.setdefaulttimeout(None)
|
||||
|
||||
def test_ws_key(self):
|
||||
key = _create_sec_websocket_key()
|
||||
self.assertTrue(key != 24)
|
||||
self.assertTrue("¥n" not in key)
|
||||
|
||||
def test_nonce(self):
|
||||
"""WebSocket key should be a random 16-byte nonce."""
|
||||
key = _create_sec_websocket_key()
|
||||
nonce = base64decode(key.encode("utf-8"))
|
||||
self.assertEqual(16, len(nonce))
|
||||
|
||||
def test_ws_utils(self):
|
||||
key = "c6b8hTg4EeGb2gQMztV1/g=="
|
||||
required_header = {
|
||||
"upgrade": "websocket",
|
||||
"connection": "upgrade",
|
||||
"sec-websocket-accept": "Kxep+hNu9n51529fGidYu7a3wO0=",
|
||||
}
|
||||
self.assertEqual(_validate_header(required_header, key, None), (True, None))
|
||||
|
||||
header = required_header.copy()
|
||||
header["upgrade"] = "http"
|
||||
self.assertEqual(_validate_header(header, key, None), (False, None))
|
||||
del header["upgrade"]
|
||||
self.assertEqual(_validate_header(header, key, None), (False, None))
|
||||
|
||||
header = required_header.copy()
|
||||
header["connection"] = "something"
|
||||
self.assertEqual(_validate_header(header, key, None), (False, None))
|
||||
del header["connection"]
|
||||
self.assertEqual(_validate_header(header, key, None), (False, None))
|
||||
|
||||
header = required_header.copy()
|
||||
header["sec-websocket-accept"] = "something"
|
||||
self.assertEqual(_validate_header(header, key, None), (False, None))
|
||||
del header["sec-websocket-accept"]
|
||||
self.assertEqual(_validate_header(header, key, None), (False, None))
|
||||
|
||||
header = required_header.copy()
|
||||
header["sec-websocket-protocol"] = "sub1"
|
||||
self.assertEqual(
|
||||
_validate_header(header, key, ["sub1", "sub2"]), (True, "sub1")
|
||||
)
|
||||
# This case will print out a logging error using the error() function, but that is expected
|
||||
self.assertEqual(_validate_header(header, key, ["sub2", "sub3"]), (False, None))
|
||||
|
||||
header = required_header.copy()
|
||||
header["sec-websocket-protocol"] = "sUb1"
|
||||
self.assertEqual(
|
||||
_validate_header(header, key, ["Sub1", "suB2"]), (True, "sub1")
|
||||
)
|
||||
|
||||
header = required_header.copy()
|
||||
# This case will print out a logging error using the error() function, but that is expected
|
||||
self.assertEqual(_validate_header(header, key, ["Sub1", "suB2"]), (False, None))
|
||||
|
||||
def test_read_header(self):
|
||||
status, header, _ = read_headers(HeaderSockMock("data/header01.txt"))
|
||||
self.assertEqual(status, 101)
|
||||
self.assertEqual(header["connection"], "Upgrade")
|
||||
|
||||
status, header, _ = read_headers(HeaderSockMock("data/header03.txt"))
|
||||
self.assertEqual(status, 101)
|
||||
self.assertEqual(header["connection"], "Upgrade, Keep-Alive")
|
||||
|
||||
HeaderSockMock("data/header02.txt")
|
||||
self.assertRaises(
|
||||
ws.WebSocketException, read_headers, HeaderSockMock("data/header02.txt")
|
||||
)
|
||||
|
||||
def test_send(self):
|
||||
# TODO: add longer frame data
|
||||
sock = ws.WebSocket()
|
||||
sock.set_mask_key(create_mask_key)
|
||||
s = sock.sock = HeaderSockMock("data/header01.txt")
|
||||
sock.send("Hello")
|
||||
self.assertEqual(s.sent[0], b"\x81\x85abcd)\x07\x0f\x08\x0e")
|
||||
|
||||
sock.send("こんにちは")
|
||||
self.assertEqual(
|
||||
s.sent[1],
|
||||
b"\x81\x8fabcd\x82\xe3\xf0\x87\xe3\xf1\x80\xe5\xca\x81\xe2\xc5\x82\xe3\xcc",
|
||||
)
|
||||
|
||||
# sock.send("x" * 5000)
|
||||
# self.assertEqual(s.sent[1], b'\x81\x8fabcd\x82\xe3\xf0\x87\xe3\xf1\x80\xe5\xca\x81\xe2\xc5\x82\xe3\xcc")
|
||||
|
||||
self.assertEqual(sock.send_binary(b"1111111111101"), 19)
|
||||
|
||||
def test_recv(self):
|
||||
# TODO: add longer frame data
|
||||
sock = ws.WebSocket()
|
||||
s = sock.sock = SockMock()
|
||||
something = (
|
||||
b"\x81\x8fabcd\x82\xe3\xf0\x87\xe3\xf1\x80\xe5\xca\x81\xe2\xc5\x82\xe3\xcc"
|
||||
)
|
||||
s.add_packet(something)
|
||||
data = sock.recv()
|
||||
self.assertEqual(data, "こんにちは")
|
||||
|
||||
s.add_packet(b"\x81\x85abcd)\x07\x0f\x08\x0e")
|
||||
data = sock.recv()
|
||||
self.assertEqual(data, "Hello")
|
||||
|
||||
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
|
||||
def test_iter(self):
|
||||
count = 2
|
||||
s = ws.create_connection("wss://api.bitfinex.com/ws/2")
|
||||
s.send('{"event": "subscribe", "channel": "ticker"}')
|
||||
for _ in s:
|
||||
count -= 1
|
||||
if count == 0:
|
||||
break
|
||||
|
||||
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
|
||||
def test_next(self):
|
||||
sock = ws.create_connection("wss://api.bitfinex.com/ws/2")
|
||||
self.assertEqual(str, type(next(sock)))
|
||||
|
||||
def test_internal_recv_strict(self):
|
||||
sock = ws.WebSocket()
|
||||
s = sock.sock = SockMock()
|
||||
s.add_packet(b"foo")
|
||||
s.add_packet(socket.timeout())
|
||||
s.add_packet(b"bar")
|
||||
# s.add_packet(SSLError("The read operation timed out"))
|
||||
s.add_packet(b"baz")
|
||||
with self.assertRaises(ws.WebSocketTimeoutException):
|
||||
sock.frame_buffer.recv_strict(9)
|
||||
# with self.assertRaises(SSLError):
|
||||
# data = sock._recv_strict(9)
|
||||
data = sock.frame_buffer.recv_strict(9)
|
||||
self.assertEqual(data, b"foobarbaz")
|
||||
with self.assertRaises(ws.WebSocketConnectionClosedException):
|
||||
sock.frame_buffer.recv_strict(1)
|
||||
|
||||
def test_recv_timeout(self):
|
||||
sock = ws.WebSocket()
|
||||
s = sock.sock = SockMock()
|
||||
s.add_packet(b"\x81")
|
||||
s.add_packet(socket.timeout())
|
||||
s.add_packet(b"\x8dabcd\x29\x07\x0f\x08\x0e")
|
||||
s.add_packet(socket.timeout())
|
||||
s.add_packet(b"\x4e\x43\x33\x0e\x10\x0f\x00\x40")
|
||||
with self.assertRaises(ws.WebSocketTimeoutException):
|
||||
sock.recv()
|
||||
with self.assertRaises(ws.WebSocketTimeoutException):
|
||||
sock.recv()
|
||||
data = sock.recv()
|
||||
self.assertEqual(data, "Hello, World!")
|
||||
with self.assertRaises(ws.WebSocketConnectionClosedException):
|
||||
sock.recv()
|
||||
|
||||
def test_recv_with_simple_fragmentation(self):
|
||||
sock = ws.WebSocket()
|
||||
s = sock.sock = SockMock()
|
||||
# OPCODE=TEXT, FIN=0, MSG="Brevity is "
|
||||
s.add_packet(b"\x01\x8babcd#\x10\x06\x12\x08\x16\x1aD\x08\x11C")
|
||||
# OPCODE=CONT, FIN=1, MSG="the soul of wit"
|
||||
s.add_packet(b"\x80\x8fabcd\x15\n\x06D\x12\r\x16\x08A\r\x05D\x16\x0b\x17")
|
||||
data = sock.recv()
|
||||
self.assertEqual(data, "Brevity is the soul of wit")
|
||||
with self.assertRaises(ws.WebSocketConnectionClosedException):
|
||||
sock.recv()
|
||||
|
||||
def test_recv_with_fire_event_of_fragmentation(self):
|
||||
sock = ws.WebSocket(fire_cont_frame=True)
|
||||
s = sock.sock = SockMock()
|
||||
# OPCODE=TEXT, FIN=0, MSG="Brevity is "
|
||||
s.add_packet(b"\x01\x8babcd#\x10\x06\x12\x08\x16\x1aD\x08\x11C")
|
||||
# OPCODE=CONT, FIN=0, MSG="Brevity is "
|
||||
s.add_packet(b"\x00\x8babcd#\x10\x06\x12\x08\x16\x1aD\x08\x11C")
|
||||
# OPCODE=CONT, FIN=1, MSG="the soul of wit"
|
||||
s.add_packet(b"\x80\x8fabcd\x15\n\x06D\x12\r\x16\x08A\r\x05D\x16\x0b\x17")
|
||||
|
||||
_, data = sock.recv_data()
|
||||
self.assertEqual(data, b"Brevity is ")
|
||||
_, data = sock.recv_data()
|
||||
self.assertEqual(data, b"Brevity is ")
|
||||
_, data = sock.recv_data()
|
||||
self.assertEqual(data, b"the soul of wit")
|
||||
|
||||
# OPCODE=CONT, FIN=0, MSG="Brevity is "
|
||||
s.add_packet(b"\x80\x8babcd#\x10\x06\x12\x08\x16\x1aD\x08\x11C")
|
||||
|
||||
with self.assertRaises(ws.WebSocketException):
|
||||
sock.recv_data()
|
||||
|
||||
with self.assertRaises(ws.WebSocketConnectionClosedException):
|
||||
sock.recv()
|
||||
|
||||
def test_close(self):
|
||||
sock = ws.WebSocket()
|
||||
sock.connected = True
|
||||
sock.close()
|
||||
|
||||
sock = ws.WebSocket()
|
||||
s = sock.sock = SockMock()
|
||||
sock.connected = True
|
||||
s.add_packet(b"\x88\x80\x17\x98p\x84")
|
||||
sock.recv()
|
||||
self.assertEqual(sock.connected, False)
|
||||
|
||||
def test_recv_cont_fragmentation(self):
|
||||
sock = ws.WebSocket()
|
||||
s = sock.sock = SockMock()
|
||||
# OPCODE=CONT, FIN=1, MSG="the soul of wit"
|
||||
s.add_packet(b"\x80\x8fabcd\x15\n\x06D\x12\r\x16\x08A\r\x05D\x16\x0b\x17")
|
||||
self.assertRaises(ws.WebSocketException, sock.recv)
|
||||
|
||||
def test_recv_with_prolonged_fragmentation(self):
|
||||
sock = ws.WebSocket()
|
||||
s = sock.sock = SockMock()
|
||||
# OPCODE=TEXT, FIN=0, MSG="Once more unto the breach, "
|
||||
s.add_packet(
|
||||
b"\x01\x9babcd.\x0c\x00\x01A\x0f\x0c\x16\x04B\x16\n\x15\rC\x10\t\x07C\x06\x13\x07\x02\x07\tNC"
|
||||
)
|
||||
# OPCODE=CONT, FIN=0, MSG="dear friends, "
|
||||
s.add_packet(b"\x00\x8eabcd\x05\x07\x02\x16A\x04\x11\r\x04\x0c\x07\x17MB")
|
||||
# OPCODE=CONT, FIN=1, MSG="once more"
|
||||
s.add_packet(b"\x80\x89abcd\x0e\x0c\x00\x01A\x0f\x0c\x16\x04")
|
||||
data = sock.recv()
|
||||
self.assertEqual(data, "Once more unto the breach, dear friends, once more")
|
||||
with self.assertRaises(ws.WebSocketConnectionClosedException):
|
||||
sock.recv()
|
||||
|
||||
def test_recv_with_fragmentation_and_control_frame(self):
|
||||
sock = ws.WebSocket()
|
||||
sock.set_mask_key(create_mask_key)
|
||||
s = sock.sock = SockMock()
|
||||
# OPCODE=TEXT, FIN=0, MSG="Too much "
|
||||
s.add_packet(b"\x01\x89abcd5\r\x0cD\x0c\x17\x00\x0cA")
|
||||
# OPCODE=PING, FIN=1, MSG="Please PONG this"
|
||||
s.add_packet(b"\x89\x90abcd1\x0e\x06\x05\x12\x07C4.,$D\x15\n\n\x17")
|
||||
# OPCODE=CONT, FIN=1, MSG="of a good thing"
|
||||
s.add_packet(b"\x80\x8fabcd\x0e\x04C\x05A\x05\x0c\x0b\x05B\x17\x0c\x08\x0c\x04")
|
||||
data = sock.recv()
|
||||
self.assertEqual(data, "Too much of a good thing")
|
||||
with self.assertRaises(ws.WebSocketConnectionClosedException):
|
||||
sock.recv()
|
||||
self.assertEqual(
|
||||
s.sent[0], b"\x8a\x90abcd1\x0e\x06\x05\x12\x07C4.,$D\x15\n\n\x17"
|
||||
)
|
||||
|
||||
@unittest.skipUnless(
|
||||
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
|
||||
)
|
||||
def test_websocket(self):
|
||||
s = ws.create_connection(f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}")
|
||||
self.assertNotEqual(s, None)
|
||||
s.send("Hello, World")
|
||||
result = s.next()
|
||||
s.fileno()
|
||||
self.assertEqual(result, "Hello, World")
|
||||
|
||||
s.send("こにゃにゃちは、世界")
|
||||
result = s.recv()
|
||||
self.assertEqual(result, "こにゃにゃちは、世界")
|
||||
self.assertRaises(ValueError, s.send_close, -1, "")
|
||||
s.close()
|
||||
|
||||
@unittest.skipUnless(
|
||||
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
|
||||
)
|
||||
def test_ping_pong(self):
|
||||
s = ws.create_connection(f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}")
|
||||
self.assertNotEqual(s, None)
|
||||
s.ping("Hello")
|
||||
s.pong("Hi")
|
||||
s.close()
|
||||
|
||||
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
|
||||
def test_support_redirect(self):
|
||||
s = ws.WebSocket()
|
||||
self.assertRaises(WebSocketBadStatusException, s.connect, "ws://google.com/")
|
||||
# Need to find a URL that has a redirect code leading to a websocket
|
||||
|
||||
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
|
||||
def test_secure_websocket(self):
|
||||
s = ws.create_connection("wss://api.bitfinex.com/ws/2")
|
||||
self.assertNotEqual(s, None)
|
||||
self.assertTrue(isinstance(s.sock, ssl.SSLSocket))
|
||||
self.assertEqual(s.getstatus(), 101)
|
||||
self.assertNotEqual(s.getheaders(), None)
|
||||
s.settimeout(10)
|
||||
self.assertEqual(s.gettimeout(), 10)
|
||||
self.assertEqual(s.getsubprotocol(), None)
|
||||
s.abort()
|
||||
|
||||
@unittest.skipUnless(
|
||||
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
|
||||
)
|
||||
def test_websocket_with_custom_header(self):
|
||||
s = ws.create_connection(
|
||||
f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}",
|
||||
headers={"User-Agent": "PythonWebsocketClient"},
|
||||
)
|
||||
self.assertNotEqual(s, None)
|
||||
self.assertEqual(s.getsubprotocol(), None)
|
||||
s.send("Hello, World")
|
||||
result = s.recv()
|
||||
self.assertEqual(result, "Hello, World")
|
||||
self.assertRaises(ValueError, s.close, -1, "")
|
||||
s.close()
|
||||
|
||||
@unittest.skipUnless(
|
||||
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
|
||||
)
|
||||
def test_after_close(self):
|
||||
s = ws.create_connection(f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}")
|
||||
self.assertNotEqual(s, None)
|
||||
s.close()
|
||||
self.assertRaises(ws.WebSocketConnectionClosedException, s.send, "Hello")
|
||||
self.assertRaises(ws.WebSocketConnectionClosedException, s.recv)
|
||||
|
||||
|
||||
class SockOptTest(unittest.TestCase):
|
||||
@unittest.skipUnless(
|
||||
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
|
||||
)
|
||||
def test_sockopt(self):
|
||||
sockopt = ((socket.IPPROTO_TCP, socket.TCP_NODELAY, 1),)
|
||||
s = ws.create_connection(
|
||||
f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}", sockopt=sockopt
|
||||
)
|
||||
self.assertNotEqual(
|
||||
s.sock.getsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY), 0
|
||||
)
|
||||
s.close()
|
||||
|
||||
|
||||
class UtilsTest(unittest.TestCase):
|
||||
def test_utf8_validator(self):
|
||||
state = validate_utf8(b"\xf0\x90\x80\x80")
|
||||
self.assertEqual(state, True)
|
||||
state = validate_utf8(
|
||||
b"\xce\xba\xe1\xbd\xb9\xcf\x83\xce\xbc\xce\xb5\xed\xa0\x80edited"
|
||||
)
|
||||
self.assertEqual(state, False)
|
||||
state = validate_utf8(b"")
|
||||
self.assertEqual(state, True)
|
||||
|
||||
|
||||
class HandshakeTest(unittest.TestCase):
|
||||
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
|
||||
def test_http_ssl(self):
|
||||
websock1 = ws.WebSocket(
|
||||
sslopt={"cert_chain": ssl.get_default_verify_paths().capath},
|
||||
enable_multithread=False,
|
||||
)
|
||||
self.assertRaises(ValueError, websock1.connect, "wss://api.bitfinex.com/ws/2")
|
||||
websock2 = ws.WebSocket(sslopt={"certfile": "myNonexistentCertFile"})
|
||||
self.assertRaises(
|
||||
WebSocketException, websock2.connect, "wss://api.bitfinex.com/ws/2"
|
||||
)
|
||||
|
||||
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
|
||||
def test_manual_headers(self):
|
||||
websock3 = ws.WebSocket(
|
||||
sslopt={
|
||||
"ca_certs": ssl.get_default_verify_paths().cafile,
|
||||
"ca_cert_path": ssl.get_default_verify_paths().capath,
|
||||
}
|
||||
)
|
||||
self.assertRaises(
|
||||
WebSocketBadStatusException,
|
||||
websock3.connect,
|
||||
"wss://api.bitfinex.com/ws/2",
|
||||
cookie="chocolate",
|
||||
origin="testing_websockets.com",
|
||||
host="echo.websocket.events/websocket-client-test",
|
||||
subprotocols=["testproto"],
|
||||
connection="Upgrade",
|
||||
header={
|
||||
"CustomHeader1": "123",
|
||||
"Cookie": "TestValue",
|
||||
"Sec-WebSocket-Key": "k9kFAUWNAMmf5OEMfTlOEA==",
|
||||
"Sec-WebSocket-Protocol": "newprotocol",
|
||||
},
|
||||
)
|
||||
|
||||
def test_ipv6(self):
|
||||
websock2 = ws.WebSocket()
|
||||
self.assertRaises(ValueError, websock2.connect, "2001:4860:4860::8888")
|
||||
|
||||
def test_bad_urls(self):
|
||||
websock3 = ws.WebSocket()
|
||||
self.assertRaises(ValueError, websock3.connect, "ws//example.com")
|
||||
self.assertRaises(WebSocketAddressException, websock3.connect, "ws://example")
|
||||
self.assertRaises(ValueError, websock3.connect, "example.com")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user