Bump websocket-client from 1.6.2 to 1.7.0 (#2207)

* Bump websocket-client from 1.6.2 to 1.7.0

Bumps [websocket-client](https://github.com/websocket-client/websocket-client) from 1.6.2 to 1.7.0.
- [Release notes](https://github.com/websocket-client/websocket-client/releases)
- [Changelog](https://github.com/websocket-client/websocket-client/blob/master/ChangeLog)
- [Commits](https://github.com/websocket-client/websocket-client/compare/v1.6.2...v1.7.0)

---
updated-dependencies:
- dependency-name: websocket-client
  dependency-type: direct:production
  update-type: version-update:semver-minor
...

Signed-off-by: dependabot[bot] <support@github.com>

* Update websocket-client==1.7.0

---------

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: JonnyWong16 <9099342+JonnyWong16@users.noreply.github.com>

[skip ci]
This commit is contained in:
dependabot[bot] 2024-03-24 15:21:51 -07:00 committed by GitHub
parent dbffb519f5
commit 24b6d37bbe
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
23 changed files with 1909 additions and 801 deletions

View file

@ -23,4 +23,4 @@ from ._exceptions import *
from ._logging import * from ._logging import *
from ._socket import * from ._socket import *
__version__ = "1.6.2" __version__ = "1.7.0"

View file

@ -2,10 +2,11 @@ import array
import os import os
import struct import struct
import sys import sys
from threading import Lock
from typing import Callable, Optional, Union
from ._exceptions import * from ._exceptions import *
from ._utils import validate_utf8 from ._utils import validate_utf8
from threading import Lock
""" """
_abnf.py _abnf.py
@ -33,8 +34,9 @@ try:
# Note that wsaccel is unmaintained. # Note that wsaccel is unmaintained.
from wsaccel.xormask import XorMaskerSimple from wsaccel.xormask import XorMaskerSimple
def _mask(_m, _d) -> bytes: def _mask(mask_value: array.array, data_value: array.array) -> bytes:
return XorMaskerSimple(_m).process(_d) mask_result: bytes = XorMaskerSimple(mask_value).process(data_value)
return mask_result
except ImportError: except ImportError:
# wsaccel is not available, use websocket-client _mask() # wsaccel is not available, use websocket-client _mask()
@ -42,26 +44,30 @@ except ImportError:
def _mask(mask_value: array.array, data_value: array.array) -> bytes: def _mask(mask_value: array.array, data_value: array.array) -> bytes:
datalen = len(data_value) datalen = len(data_value)
data_value = int.from_bytes(data_value, native_byteorder) int_data_value = int.from_bytes(data_value, native_byteorder)
mask_value = int.from_bytes(mask_value * (datalen // 4) + mask_value[: datalen % 4], native_byteorder) int_mask_value = int.from_bytes(
return (data_value ^ mask_value).to_bytes(datalen, native_byteorder) mask_value * (datalen // 4) + mask_value[: datalen % 4], native_byteorder
)
return (int_data_value ^ int_mask_value).to_bytes(datalen, native_byteorder)
__all__ = [ __all__ = [
'ABNF', 'continuous_frame', 'frame_buffer', "ABNF",
'STATUS_NORMAL', "continuous_frame",
'STATUS_GOING_AWAY', "frame_buffer",
'STATUS_PROTOCOL_ERROR', "STATUS_NORMAL",
'STATUS_UNSUPPORTED_DATA_TYPE', "STATUS_GOING_AWAY",
'STATUS_STATUS_NOT_AVAILABLE', "STATUS_PROTOCOL_ERROR",
'STATUS_ABNORMAL_CLOSED', "STATUS_UNSUPPORTED_DATA_TYPE",
'STATUS_INVALID_PAYLOAD', "STATUS_STATUS_NOT_AVAILABLE",
'STATUS_POLICY_VIOLATION', "STATUS_ABNORMAL_CLOSED",
'STATUS_MESSAGE_TOO_BIG', "STATUS_INVALID_PAYLOAD",
'STATUS_INVALID_EXTENSION', "STATUS_POLICY_VIOLATION",
'STATUS_UNEXPECTED_CONDITION', "STATUS_MESSAGE_TOO_BIG",
'STATUS_BAD_GATEWAY', "STATUS_INVALID_EXTENSION",
'STATUS_TLS_HANDSHAKE_ERROR', "STATUS_UNEXPECTED_CONDITION",
"STATUS_BAD_GATEWAY",
"STATUS_TLS_HANDSHAKE_ERROR",
] ]
# closing frame status codes. # closing frame status codes.
@ -110,11 +116,17 @@ class ABNF:
OPCODE_BINARY = 0x2 OPCODE_BINARY = 0x2
OPCODE_CLOSE = 0x8 OPCODE_CLOSE = 0x8
OPCODE_PING = 0x9 OPCODE_PING = 0x9
OPCODE_PONG = 0xa OPCODE_PONG = 0xA
# available operation code value tuple # available operation code value tuple
OPCODES = (OPCODE_CONT, OPCODE_TEXT, OPCODE_BINARY, OPCODE_CLOSE, OPCODES = (
OPCODE_PING, OPCODE_PONG) OPCODE_CONT,
OPCODE_TEXT,
OPCODE_BINARY,
OPCODE_CLOSE,
OPCODE_PING,
OPCODE_PONG,
)
# opcode human readable string # opcode human readable string
OPCODE_MAP = { OPCODE_MAP = {
@ -123,16 +135,24 @@ class ABNF:
OPCODE_BINARY: "binary", OPCODE_BINARY: "binary",
OPCODE_CLOSE: "close", OPCODE_CLOSE: "close",
OPCODE_PING: "ping", OPCODE_PING: "ping",
OPCODE_PONG: "pong" OPCODE_PONG: "pong",
} }
# data length threshold. # data length threshold.
LENGTH_7 = 0x7e LENGTH_7 = 0x7E
LENGTH_16 = 1 << 16 LENGTH_16 = 1 << 16
LENGTH_63 = 1 << 63 LENGTH_63 = 1 << 63
def __init__(self, fin: int = 0, rsv1: int = 0, rsv2: int = 0, rsv3: int = 0, def __init__(
opcode: int = OPCODE_TEXT, mask: int = 1, data: str or bytes = "") -> None: self,
fin: int = 0,
rsv1: int = 0,
rsv2: int = 0,
rsv3: int = 0,
opcode: int = OPCODE_TEXT,
mask_value: int = 1,
data: Union[str, bytes, None] = "",
) -> None:
""" """
Constructor for ABNF. Please check RFC for arguments. Constructor for ABNF. Please check RFC for arguments.
""" """
@ -141,7 +161,7 @@ class ABNF:
self.rsv2 = rsv2 self.rsv2 = rsv2
self.rsv3 = rsv3 self.rsv3 = rsv3
self.opcode = opcode self.opcode = opcode
self.mask = mask self.mask_value = mask_value
if data is None: if data is None:
data = "" data = ""
self.data = data self.data = data
@ -173,7 +193,7 @@ class ABNF:
if l > 2 and not skip_utf8_validation and not validate_utf8(self.data[2:]): if l > 2 and not skip_utf8_validation and not validate_utf8(self.data[2:]):
raise WebSocketProtocolException("Invalid close frame.") raise WebSocketProtocolException("Invalid close frame.")
code = 256 * self.data[0] + self.data[1] code = 256 * int(self.data[0]) + int(self.data[1])
if not self._is_valid_close_status(code): if not self._is_valid_close_status(code):
raise WebSocketProtocolException("Invalid close opcode %r", code) raise WebSocketProtocolException("Invalid close opcode %r", code)
@ -182,12 +202,10 @@ class ABNF:
return code in VALID_CLOSE_STATUS or (3000 <= code < 5000) return code in VALID_CLOSE_STATUS or (3000 <= code < 5000)
def __str__(self) -> str: def __str__(self) -> str:
return "fin=" + str(self.fin) \ return f"fin={self.fin} opcode={self.opcode} data={self.data}"
+ " opcode=" + str(self.opcode) \
+ " data=" + str(self.data)
@staticmethod @staticmethod
def create_frame(data: str, opcode: int, fin: int = 1) -> 'ABNF': def create_frame(data: Union[bytes, str], opcode: int, fin: int = 1) -> "ABNF":
""" """
Create frame to send text, binary and other data. Create frame to send text, binary and other data.
@ -219,34 +237,39 @@ class ABNF:
if length >= ABNF.LENGTH_63: if length >= ABNF.LENGTH_63:
raise ValueError("data is too long") raise ValueError("data is too long")
frame_header = chr(self.fin << 7 | frame_header = chr(
self.rsv1 << 6 | self.rsv2 << 5 | self.rsv3 << 4 | self.fin << 7
self.opcode).encode('latin-1') | self.rsv1 << 6
| self.rsv2 << 5
| self.rsv3 << 4
| self.opcode
).encode("latin-1")
if length < ABNF.LENGTH_7: if length < ABNF.LENGTH_7:
frame_header += chr(self.mask << 7 | length).encode('latin-1') frame_header += chr(self.mask_value << 7 | length).encode("latin-1")
elif length < ABNF.LENGTH_16: elif length < ABNF.LENGTH_16:
frame_header += chr(self.mask << 7 | 0x7e).encode('latin-1') frame_header += chr(self.mask_value << 7 | 0x7E).encode("latin-1")
frame_header += struct.pack("!H", length) frame_header += struct.pack("!H", length)
else: else:
frame_header += chr(self.mask << 7 | 0x7f).encode('latin-1') frame_header += chr(self.mask_value << 7 | 0x7F).encode("latin-1")
frame_header += struct.pack("!Q", length) frame_header += struct.pack("!Q", length)
if not self.mask: if not self.mask_value:
if isinstance(self.data, str):
self.data = self.data.encode("utf-8")
return frame_header + self.data return frame_header + self.data
else: mask_key = self.get_mask_key(4)
mask_key = self.get_mask_key(4) return frame_header + self._get_masked(mask_key)
return frame_header + self._get_masked(mask_key)
def _get_masked(self, mask_key: str or bytes) -> bytes: def _get_masked(self, mask_key: Union[str, bytes]) -> bytes:
s = ABNF.mask(mask_key, self.data) s = ABNF.mask(mask_key, self.data)
if isinstance(mask_key, str): if isinstance(mask_key, str):
mask_key = mask_key.encode('utf-8') mask_key = mask_key.encode("utf-8")
return mask_key + s return mask_key + s
@staticmethod @staticmethod
def mask(mask_key: str or bytes, data: str or bytes) -> bytes: def mask(mask_key: Union[str, bytes], data: Union[str, bytes]) -> bytes:
""" """
Mask or unmask data. Just do xor for each byte Mask or unmask data. Just do xor for each byte
@ -261,10 +284,10 @@ class ABNF:
data = "" data = ""
if isinstance(mask_key, str): if isinstance(mask_key, str):
mask_key = mask_key.encode('latin-1') mask_key = mask_key.encode("latin-1")
if isinstance(data, str): if isinstance(data, str):
data = data.encode('latin-1') data = data.encode("latin-1")
return _mask(array.array("B", mask_key), array.array("B", data)) return _mask(array.array("B", mask_key), array.array("B", data))
@ -273,19 +296,21 @@ class frame_buffer:
_HEADER_MASK_INDEX = 5 _HEADER_MASK_INDEX = 5
_HEADER_LENGTH_INDEX = 6 _HEADER_LENGTH_INDEX = 6
def __init__(self, recv_fn: int, skip_utf8_validation: bool) -> None: def __init__(
self, recv_fn: Callable[[int], int], skip_utf8_validation: bool
) -> None:
self.recv = recv_fn self.recv = recv_fn
self.skip_utf8_validation = skip_utf8_validation self.skip_utf8_validation = skip_utf8_validation
# Buffers over the packets from the layer beneath until desired amount # Buffers over the packets from the layer beneath until desired amount
# bytes of bytes are received. # bytes of bytes are received.
self.recv_buffer = [] self.recv_buffer: list = []
self.clear() self.clear()
self.lock = Lock() self.lock = Lock()
def clear(self) -> None: def clear(self) -> None:
self.header = None self.header: Optional[tuple] = None
self.length = None self.length: Optional[int] = None
self.mask = None self.mask_value: Union[bytes, str, None] = None
def has_received_header(self) -> bool: def has_received_header(self) -> bool:
return self.header is None return self.header is None
@ -297,41 +322,41 @@ class frame_buffer:
rsv1 = b1 >> 6 & 1 rsv1 = b1 >> 6 & 1
rsv2 = b1 >> 5 & 1 rsv2 = b1 >> 5 & 1
rsv3 = b1 >> 4 & 1 rsv3 = b1 >> 4 & 1
opcode = b1 & 0xf opcode = b1 & 0xF
b2 = header[1] b2 = header[1]
has_mask = b2 >> 7 & 1 has_mask = b2 >> 7 & 1
length_bits = b2 & 0x7f length_bits = b2 & 0x7F
self.header = (fin, rsv1, rsv2, rsv3, opcode, has_mask, length_bits) self.header = (fin, rsv1, rsv2, rsv3, opcode, has_mask, length_bits)
def has_mask(self) -> bool or int: def has_mask(self) -> Union[bool, int]:
if not self.header: if not self.header:
return False return False
return self.header[frame_buffer._HEADER_MASK_INDEX] header_val: int = self.header[frame_buffer._HEADER_MASK_INDEX]
return header_val
def has_received_length(self) -> bool: def has_received_length(self) -> bool:
return self.length is None return self.length is None
def recv_length(self) -> None: def recv_length(self) -> None:
bits = self.header[frame_buffer._HEADER_LENGTH_INDEX] bits = self.header[frame_buffer._HEADER_LENGTH_INDEX]
length_bits = bits & 0x7f length_bits = bits & 0x7F
if length_bits == 0x7e: if length_bits == 0x7E:
v = self.recv_strict(2) v = self.recv_strict(2)
self.length = struct.unpack("!H", v)[0] self.length = struct.unpack("!H", v)[0]
elif length_bits == 0x7f: elif length_bits == 0x7F:
v = self.recv_strict(8) v = self.recv_strict(8)
self.length = struct.unpack("!Q", v)[0] self.length = struct.unpack("!Q", v)[0]
else: else:
self.length = length_bits self.length = length_bits
def has_received_mask(self) -> bool: def has_received_mask(self) -> bool:
return self.mask is None return self.mask_value is None
def recv_mask(self) -> None: def recv_mask(self) -> None:
self.mask = self.recv_strict(4) if self.has_mask() else "" self.mask_value = self.recv_strict(4) if self.has_mask() else ""
def recv_frame(self) -> ABNF: def recv_frame(self) -> ABNF:
with self.lock: with self.lock:
# Header # Header
if self.has_received_header(): if self.has_received_header():
@ -346,12 +371,12 @@ class frame_buffer:
# Mask # Mask
if self.has_received_mask(): if self.has_received_mask():
self.recv_mask() self.recv_mask()
mask = self.mask mask_value = self.mask_value
# Payload # Payload
payload = self.recv_strict(length) payload = self.recv_strict(length)
if has_mask: if has_mask:
payload = ABNF.mask(mask, payload) payload = ABNF.mask(mask_value, payload)
# Reset for next frame # Reset for next frame
self.clear() self.clear()
@ -385,18 +410,19 @@ class frame_buffer:
class continuous_frame: class continuous_frame:
def __init__(self, fire_cont_frame: bool, skip_utf8_validation: bool) -> None: def __init__(self, fire_cont_frame: bool, skip_utf8_validation: bool) -> None:
self.fire_cont_frame = fire_cont_frame self.fire_cont_frame = fire_cont_frame
self.skip_utf8_validation = skip_utf8_validation self.skip_utf8_validation = skip_utf8_validation
self.cont_data = None self.cont_data: Optional[list] = None
self.recving_frames = None self.recving_frames: Optional[int] = None
def validate(self, frame: ABNF) -> None: def validate(self, frame: ABNF) -> None:
if not self.recving_frames and frame.opcode == ABNF.OPCODE_CONT: if not self.recving_frames and frame.opcode == ABNF.OPCODE_CONT:
raise WebSocketProtocolException("Illegal frame") raise WebSocketProtocolException("Illegal frame")
if self.recving_frames and \ if self.recving_frames and frame.opcode in (
frame.opcode in (ABNF.OPCODE_TEXT, ABNF.OPCODE_BINARY): ABNF.OPCODE_TEXT,
ABNF.OPCODE_BINARY,
):
raise WebSocketProtocolException("Illegal frame") raise WebSocketProtocolException("Illegal frame")
def add(self, frame: ABNF) -> None: def add(self, frame: ABNF) -> None:
@ -410,15 +436,18 @@ class continuous_frame:
if frame.fin: if frame.fin:
self.recving_frames = None self.recving_frames = None
def is_fire(self, frame: ABNF) -> bool or int: def is_fire(self, frame: ABNF) -> Union[bool, int]:
return frame.fin or self.fire_cont_frame return frame.fin or self.fire_cont_frame
def extract(self, frame: ABNF) -> list: def extract(self, frame: ABNF) -> tuple:
data = self.cont_data data = self.cont_data
self.cont_data = None self.cont_data = None
frame.data = data[1] frame.data = data[1]
if not self.fire_cont_frame and data[0] == ABNF.OPCODE_TEXT and not self.skip_utf8_validation and not validate_utf8(frame.data): if (
raise WebSocketPayloadException( not self.fire_cont_frame
"cannot decode: " + repr(frame.data)) and data[0] == ABNF.OPCODE_TEXT
and not self.skip_utf8_validation
return [data[0], frame] and not validate_utf8(frame.data)
):
raise WebSocketPayloadException(f"cannot decode: {repr(frame.data)}")
return data[0], frame

View file

@ -1,18 +1,19 @@
import inspect import inspect
import selectors import selectors
import sys import socket
import threading import threading
import time import time
import traceback from typing import Any, Callable, Optional, Union
import socket
from typing import Callable, Any
from . import _logging from . import _logging
from ._abnf import ABNF from ._abnf import ABNF
from ._url import parse_url
from ._core import WebSocket, getdefaulttimeout from ._core import WebSocket, getdefaulttimeout
from ._exceptions import * from ._exceptions import (
WebSocketConnectionClosedException,
WebSocketException,
WebSocketTimeoutException,
)
from ._url import parse_url
""" """
_app.py _app.py
@ -47,22 +48,24 @@ class DispatcherBase:
""" """
DispatcherBase DispatcherBase
""" """
def __init__(self, app: Any, ping_timeout: float) -> None:
def __init__(self, app: Any, ping_timeout: Union[float, int, None]) -> None:
self.app = app self.app = app
self.ping_timeout = ping_timeout self.ping_timeout = ping_timeout
def timeout(self, seconds: int, callback: Callable) -> None: def timeout(self, seconds: Union[float, int, None], callback: Callable) -> None:
time.sleep(seconds) time.sleep(seconds)
callback() callback()
def reconnect(self, seconds: int, reconnector: Callable) -> None: def reconnect(self, seconds: int, reconnector: Callable) -> None:
try: try:
_logging.info("reconnect() - retrying in {seconds_count} seconds [{frame_count} frames in stack]".format( _logging.info(
seconds_count=seconds, frame_count=len(inspect.stack()))) f"reconnect() - retrying in {seconds} seconds [{len(inspect.stack())} frames in stack]"
)
time.sleep(seconds) time.sleep(seconds)
reconnector(reconnecting=True) reconnector(reconnecting=True)
except KeyboardInterrupt as e: except KeyboardInterrupt as e:
_logging.info("User exited {err}".format(err=e)) _logging.info(f"User exited {e}")
raise e raise e
@ -70,13 +73,18 @@ class Dispatcher(DispatcherBase):
""" """
Dispatcher Dispatcher
""" """
def read(self, sock: socket.socket, read_callback: Callable, check_callback: Callable) -> None:
def read(
self,
sock: socket.socket,
read_callback: Callable,
check_callback: Callable,
) -> None:
sel = selectors.DefaultSelector() sel = selectors.DefaultSelector()
sel.register(self.app.sock.sock, selectors.EVENT_READ) sel.register(self.app.sock.sock, selectors.EVENT_READ)
try: try:
while self.app.keep_running: while self.app.keep_running:
r = sel.select(self.ping_timeout) if sel.select(self.ping_timeout):
if r:
if not read_callback(): if not read_callback():
break break
check_callback() check_callback()
@ -88,24 +96,31 @@ class SSLDispatcher(DispatcherBase):
""" """
SSLDispatcher SSLDispatcher
""" """
def read(self, sock: socket.socket, read_callback: Callable, check_callback: Callable) -> None:
def read(
self,
sock: socket.socket,
read_callback: Callable,
check_callback: Callable,
) -> None:
sock = self.app.sock.sock sock = self.app.sock.sock
sel = selectors.DefaultSelector() sel = selectors.DefaultSelector()
sel.register(sock, selectors.EVENT_READ) sel.register(sock, selectors.EVENT_READ)
try: try:
while self.app.keep_running: while self.app.keep_running:
r = self.select(sock, sel) if self.select(sock, sel):
if r:
if not read_callback(): if not read_callback():
break break
check_callback() check_callback()
finally: finally:
sel.close() sel.close()
def select(self, sock, sel:selectors.DefaultSelector): def select(self, sock, sel: selectors.DefaultSelector):
sock = self.app.sock.sock sock = self.app.sock.sock
if sock.pending(): if sock.pending():
return [sock,] return [
sock,
]
r = sel.select(self.ping_timeout) r = sel.select(self.ping_timeout)
@ -117,17 +132,23 @@ class WrappedDispatcher:
""" """
WrappedDispatcher WrappedDispatcher
""" """
def __init__(self, app, ping_timeout: float, dispatcher: Dispatcher) -> None:
def __init__(self, app, ping_timeout: Union[float, int, None], dispatcher) -> None:
self.app = app self.app = app
self.ping_timeout = ping_timeout self.ping_timeout = ping_timeout
self.dispatcher = dispatcher self.dispatcher = dispatcher
dispatcher.signal(2, dispatcher.abort) # keyboard interrupt dispatcher.signal(2, dispatcher.abort) # keyboard interrupt
def read(self, sock: socket.socket, read_callback: Callable, check_callback: Callable) -> None: def read(
self,
sock: socket.socket,
read_callback: Callable,
check_callback: Callable,
) -> None:
self.dispatcher.read(sock, read_callback) self.dispatcher.read(sock, read_callback)
self.ping_timeout and self.timeout(self.ping_timeout, check_callback) self.ping_timeout and self.timeout(self.ping_timeout, check_callback)
def timeout(self, seconds: int, callback: Callable) -> None: def timeout(self, seconds: float, callback: Callable) -> None:
self.dispatcher.timeout(seconds, callback) self.dispatcher.timeout(seconds, callback)
def reconnect(self, seconds: int, reconnector: Callable) -> None: def reconnect(self, seconds: int, reconnector: Callable) -> None:
@ -139,14 +160,24 @@ class WebSocketApp:
Higher level of APIs are provided. The interface is like JavaScript WebSocket object. Higher level of APIs are provided. The interface is like JavaScript WebSocket object.
""" """
def __init__(self, url: str, header: list or dict or Callable = None, def __init__(
on_open: Callable = None, on_message: Callable = None, on_error: Callable = None, self,
on_close: Callable = None, on_ping: Callable = None, on_pong: Callable = None, url: str,
on_cont_message: Callable = None, header: Union[list, dict, Callable, None] = None,
keep_running: bool = True, get_mask_key: Callable = None, cookie: str = None, on_open: Optional[Callable[[WebSocket], None]] = None,
subprotocols: list = None, on_message: Optional[Callable[[WebSocket, Any], None]] = None,
on_data: Callable = None, on_error: Optional[Callable[[WebSocket, Any], None]] = None,
socket: socket.socket = None) -> None: on_close: Optional[Callable[[WebSocket, 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] = None,
on_data: Optional[Callable] = None,
socket: Optional[socket.socket] = None,
) -> None:
""" """
WebSocketApp initialization WebSocketApp initialization
@ -222,13 +253,13 @@ class WebSocketApp:
self.on_cont_message = on_cont_message self.on_cont_message = on_cont_message
self.keep_running = False self.keep_running = False
self.get_mask_key = get_mask_key self.get_mask_key = get_mask_key
self.sock = None self.sock: Optional[WebSocket] = None
self.last_ping_tm = 0 self.last_ping_tm = float(0)
self.last_pong_tm = 0 self.last_pong_tm = float(0)
self.ping_thread = None self.ping_thread: Optional[threading.Thread] = None
self.stop_ping = None self.stop_ping: Optional[threading.Event] = None
self.ping_interval = 0 self.ping_interval = float(0)
self.ping_timeout = None self.ping_timeout: Union[float, int, None] = None
self.ping_payload = "" self.ping_payload = ""
self.subprotocols = subprotocols self.subprotocols = subprotocols
self.prepared_socket = socket self.prepared_socket = socket
@ -236,7 +267,7 @@ class WebSocketApp:
self.has_done_teardown = False self.has_done_teardown = False
self.has_done_teardown_lock = threading.Lock() self.has_done_teardown_lock = threading.Lock()
def send(self, data: str, opcode: int = ABNF.OPCODE_TEXT) -> None: def send(self, data: Union[bytes, str], opcode: int = ABNF.OPCODE_TEXT) -> None:
""" """
send message send message
@ -250,8 +281,21 @@ class WebSocketApp:
""" """
if not self.sock or self.sock.send(data, opcode) == 0: if not self.sock or self.sock.send(data, opcode) == 0:
raise WebSocketConnectionClosedException( raise WebSocketConnectionClosedException("Connection is already closed.")
"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: def close(self, **kwargs) -> None:
""" """
@ -263,7 +307,7 @@ class WebSocketApp:
self.sock = None self.sock = None
def _start_ping_thread(self) -> None: def _start_ping_thread(self) -> None:
self.last_ping_tm = self.last_pong_tm = 0 self.last_ping_tm = self.last_pong_tm = float(0)
self.stop_ping = threading.Event() self.stop_ping = threading.Event()
self.ping_thread = threading.Thread(target=self._send_ping) self.ping_thread = threading.Thread(target=self._send_ping)
self.ping_thread.daemon = True self.ping_thread.daemon = True
@ -274,7 +318,7 @@ class WebSocketApp:
self.stop_ping.set() self.stop_ping.set()
if self.ping_thread and self.ping_thread.is_alive(): if self.ping_thread and self.ping_thread.is_alive():
self.ping_thread.join(3) self.ping_thread.join(3)
self.last_ping_tm = self.last_pong_tm = 0 self.last_ping_tm = self.last_pong_tm = float(0)
def _send_ping(self) -> None: def _send_ping(self) -> None:
if self.stop_ping.wait(self.ping_interval) or self.keep_running is False: if self.stop_ping.wait(self.ping_interval) or self.keep_running is False:
@ -286,17 +330,28 @@ class WebSocketApp:
_logging.debug("Sending ping") _logging.debug("Sending ping")
self.sock.ping(self.ping_payload) self.sock.ping(self.ping_payload)
except Exception as e: except Exception as e:
_logging.debug("Failed to send ping: {err}".format(err=e)) _logging.debug(f"Failed to send ping: {e}")
def run_forever(self, sockopt: tuple = None, sslopt: dict = None, def run_forever(
ping_interval: float = 0, ping_timeout: float or None = None, self,
ping_payload: str = "", sockopt: tuple = None,
http_proxy_host: str = None, http_proxy_port: int or str = None, sslopt: dict = None,
http_no_proxy: list = None, http_proxy_auth: tuple = None, ping_interval: Union[float, int] = 0,
http_proxy_timeout: float = None, ping_timeout: Union[float, int, None] = None,
skip_utf8_validation: bool = False, ping_payload: str = "",
host: str = None, origin: str = None, dispatcher: Dispatcher = None, http_proxy_host: str = None,
suppress_origin: bool = False, proxy_type: str = None, reconnect: int = None) -> bool: 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. Run event loop for WebSocket framework.
@ -360,7 +415,7 @@ class WebSocketApp:
if ping_timeout and ping_interval and ping_interval <= ping_timeout: if ping_timeout and ping_interval and ping_interval <= ping_timeout:
raise WebSocketException("Ensure ping_interval > ping_timeout") raise WebSocketException("Ensure ping_interval > ping_timeout")
if not sockopt: if not sockopt:
sockopt = [] sockopt = ()
if not sslopt: if not sslopt:
sslopt = {} sslopt = {}
if self.sock: if self.sock:
@ -394,7 +449,8 @@ class WebSocketApp:
if self.sock: if self.sock:
self.sock.close() self.sock.close()
close_status_code, close_reason = self._get_close_args( close_status_code, close_reason = self._get_close_args(
close_frame if close_frame else None) close_frame if close_frame else None
)
self.sock = None self.sock = None
# Finally call the callback AFTER all teardown is complete # Finally call the callback AFTER all teardown is complete
@ -405,24 +461,34 @@ class WebSocketApp:
self.sock.shutdown() self.sock.shutdown()
self.sock = WebSocket( self.sock = WebSocket(
self.get_mask_key, sockopt=sockopt, sslopt=sslopt, self.get_mask_key,
sockopt=sockopt,
sslopt=sslopt,
fire_cont_frame=self.on_cont_message is not None, fire_cont_frame=self.on_cont_message is not None,
skip_utf8_validation=skip_utf8_validation, skip_utf8_validation=skip_utf8_validation,
enable_multithread=True) enable_multithread=True,
)
self.sock.settimeout(getdefaulttimeout()) self.sock.settimeout(getdefaulttimeout())
try: try:
header = self.header() if callable(self.header) else self.header header = self.header() if callable(self.header) else self.header
self.sock.connect( self.sock.connect(
self.url, header=header, cookie=self.cookie, self.url,
header=header,
cookie=self.cookie,
http_proxy_host=http_proxy_host, http_proxy_host=http_proxy_host,
http_proxy_port=http_proxy_port, http_no_proxy=http_no_proxy, http_proxy_port=http_proxy_port,
http_proxy_auth=http_proxy_auth, http_proxy_timeout=http_proxy_timeout, http_no_proxy=http_no_proxy,
http_proxy_auth=http_proxy_auth,
http_proxy_timeout=http_proxy_timeout,
subprotocols=self.subprotocols, subprotocols=self.subprotocols,
host=host, origin=origin, suppress_origin=suppress_origin, host=host,
proxy_type=proxy_type, socket=self.prepared_socket) origin=origin,
suppress_origin=suppress_origin,
proxy_type=proxy_type,
socket=self.prepared_socket,
)
_logging.info("Websocket connected") _logging.info("Websocket connected")
@ -432,7 +498,13 @@ class WebSocketApp:
self._callback(self.on_open) self._callback(self.on_open)
dispatcher.read(self.sock.sock, read, check) dispatcher.read(self.sock.sock, read, check)
except (WebSocketConnectionClosedException, ConnectionRefusedError, KeyboardInterrupt, SystemExit, Exception) as e: except (
WebSocketConnectionClosedException,
ConnectionRefusedError,
KeyboardInterrupt,
SystemExit,
Exception,
) as e:
handleDisconnect(e, reconnecting) handleDisconnect(e, reconnecting)
def read() -> bool: def read() -> bool:
@ -441,7 +513,10 @@ class WebSocketApp:
try: try:
op_code, frame = self.sock.recv_data_frame(True) op_code, frame = self.sock.recv_data_frame(True)
except (WebSocketConnectionClosedException, KeyboardInterrupt) as e: except (
WebSocketConnectionClosedException,
KeyboardInterrupt,
) as e:
if custom_dispatcher: if custom_dispatcher:
return handleDisconnect(e) return handleDisconnect(e)
else: else:
@ -455,10 +530,8 @@ class WebSocketApp:
self.last_pong_tm = time.time() self.last_pong_tm = time.time()
self._callback(self.on_pong, frame.data) self._callback(self.on_pong, frame.data)
elif op_code == ABNF.OPCODE_CONT and self.on_cont_message: elif op_code == ABNF.OPCODE_CONT and self.on_cont_message:
self._callback(self.on_data, frame.data, self._callback(self.on_data, frame.data, frame.opcode, frame.fin)
frame.opcode, frame.fin) self._callback(self.on_cont_message, frame.data, frame.fin)
self._callback(self.on_cont_message,
frame.data, frame.fin)
else: else:
data = frame.data data = frame.data
if op_code == ABNF.OPCODE_TEXT and not skip_utf8_validation: if op_code == ABNF.OPCODE_TEXT and not skip_utf8_validation:
@ -469,18 +542,38 @@ class WebSocketApp:
return True return True
def check() -> bool: def check() -> bool:
if (self.ping_timeout): if self.ping_timeout:
has_timeout_expired = time.time() - self.last_ping_tm > self.ping_timeout has_timeout_expired = (
has_pong_not_arrived_after_last_ping = self.last_pong_tm - self.last_ping_tm < 0 time.time() - self.last_ping_tm > self.ping_timeout
has_pong_arrived_too_late = self.last_pong_tm - 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 if (
has_timeout_expired and self.last_ping_tm
(has_pong_not_arrived_after_last_ping or has_pong_arrived_too_late)): and has_timeout_expired
and (
has_pong_not_arrived_after_last_ping
or has_pong_arrived_too_late
)
):
raise WebSocketTimeoutException("ping/pong timed out") raise WebSocketTimeoutException("ping/pong timed out")
return True return True
def handleDisconnect(e: Exception, reconnecting: bool = False) -> bool: def handleDisconnect(
e: Union[
WebSocketConnectionClosedException,
ConnectionRefusedError,
KeyboardInterrupt,
SystemExit,
Exception,
],
reconnecting: bool = False,
) -> bool:
self.has_errored = True self.has_errored = True
self._stop_ping_thread() self._stop_ping_thread()
if not reconnecting: if not reconnecting:
@ -492,25 +585,31 @@ class WebSocketApp:
raise raise
if reconnect: if reconnect:
_logging.info("{err} - reconnect".format(err=e)) _logging.info(f"{e} - reconnect")
if custom_dispatcher: if custom_dispatcher:
_logging.debug("Calling custom dispatcher reconnect [{frame_count} frames in stack]".format(frame_count=len(inspect.stack()))) _logging.debug(
f"Calling custom dispatcher reconnect [{len(inspect.stack())} frames in stack]"
)
dispatcher.reconnect(reconnect, setSock) dispatcher.reconnect(reconnect, setSock)
else: else:
_logging.error("{err} - goodbye".format(err=e)) _logging.error(f"{e} - goodbye")
teardown() teardown()
custom_dispatcher = bool(dispatcher) custom_dispatcher = bool(dispatcher)
dispatcher = self.create_dispatcher(ping_timeout, dispatcher, parse_url(self.url)[3]) dispatcher = self.create_dispatcher(
ping_timeout, dispatcher, parse_url(self.url)[3]
)
try: try:
setSock() setSock()
if not custom_dispatcher and reconnect: if not custom_dispatcher and reconnect:
while self.keep_running: while self.keep_running:
_logging.debug("Calling dispatcher reconnect [{frame_count} frames in stack]".format(frame_count=len(inspect.stack()))) _logging.debug(
f"Calling dispatcher reconnect [{len(inspect.stack())} frames in stack]"
)
dispatcher.reconnect(reconnect, setSock) dispatcher.reconnect(reconnect, setSock)
except (KeyboardInterrupt, Exception) as e: except (KeyboardInterrupt, Exception) as e:
_logging.info("tearing down on exception {err}".format(err=e)) _logging.info(f"tearing down on exception {e}")
teardown() teardown()
finally: finally:
if not custom_dispatcher: if not custom_dispatcher:
@ -519,13 +618,17 @@ class WebSocketApp:
return self.has_errored return self.has_errored
def create_dispatcher(self, ping_timeout: int, dispatcher: Dispatcher = None, is_ssl: bool = False) -> DispatcherBase: def create_dispatcher(
self,
ping_timeout: Union[float, int, None],
dispatcher: Optional[DispatcherBase] = None,
is_ssl: bool = False,
) -> Union[Dispatcher, SSLDispatcher, WrappedDispatcher]:
if dispatcher: # If custom dispatcher is set, use WrappedDispatcher if dispatcher: # If custom dispatcher is set, use WrappedDispatcher
return WrappedDispatcher(self, ping_timeout, dispatcher) return WrappedDispatcher(self, ping_timeout, dispatcher)
timeout = ping_timeout or 10 timeout = ping_timeout or 10
if is_ssl: if is_ssl:
return SSLDispatcher(self, timeout) return SSLDispatcher(self, timeout)
return Dispatcher(self, timeout) return Dispatcher(self, timeout)
def _get_close_args(self, close_frame: ABNF) -> list: def _get_close_args(self, close_frame: ABNF) -> list:
@ -540,8 +643,12 @@ class WebSocketApp:
# Extract close frame status code # Extract close frame status code
if close_frame.data and len(close_frame.data) >= 2: if close_frame.data and len(close_frame.data) >= 2:
close_status_code = 256 * close_frame.data[0] + close_frame.data[1] close_status_code = 256 * int(close_frame.data[0]) + int(
reason = close_frame.data[2:].decode('utf-8') close_frame.data[1]
)
reason = close_frame.data[2:]
if isinstance(reason, bytes):
reason = reason.decode("utf-8")
return [close_status_code, reason] return [close_status_code, reason]
else: else:
# Most likely reached this because len(close_frame_data.data) < 2 # Most likely reached this because len(close_frame_data.data) < 2
@ -553,6 +660,6 @@ class WebSocketApp:
callback(self, *args) callback(self, *args)
except Exception as e: except Exception as e:
_logging.error("error from callback {callback}: {err}".format(callback=callback, err=e)) _logging.error(f"error from callback {callback}: {e}")
if self.on_error: if self.on_error:
self.on_error(self, e) self.on_error(self, e)

View file

@ -1,4 +1,5 @@
import http.cookies import http.cookies
from typing import Optional
""" """
_cookiejar.py _cookiejar.py
@ -22,18 +23,21 @@ limitations under the License.
class SimpleCookieJar: class SimpleCookieJar:
def __init__(self) -> None: def __init__(self) -> None:
self.jar = dict() self.jar: dict = dict()
def add(self, set_cookie: str) -> None: def add(self, set_cookie: Optional[str]) -> None:
if set_cookie: if set_cookie:
simpleCookie = http.cookies.SimpleCookie(set_cookie) simpleCookie = http.cookies.SimpleCookie(set_cookie)
for k, v in simpleCookie.items(): for k, v in simpleCookie.items():
domain = v.get("domain") if domain := v.get("domain"):
if domain:
if not domain.startswith("."): if not domain.startswith("."):
domain = "." + domain domain = f".{domain}"
cookie = self.jar.get(domain) if self.jar.get(domain) else http.cookies.SimpleCookie() cookie = (
self.jar.get(domain)
if self.jar.get(domain)
else http.cookies.SimpleCookie()
)
cookie.update(simpleCookie) cookie.update(simpleCookie)
self.jar[domain.lower()] = cookie self.jar[domain.lower()] = cookie
@ -42,10 +46,9 @@ class SimpleCookieJar:
simpleCookie = http.cookies.SimpleCookie(set_cookie) simpleCookie = http.cookies.SimpleCookie(set_cookie)
for k, v in simpleCookie.items(): for k, v in simpleCookie.items():
domain = v.get("domain") if domain := v.get("domain"):
if domain:
if not domain.startswith("."): if not domain.startswith("."):
domain = "." + domain domain = f".{domain}"
self.jar[domain.lower()] = simpleCookie self.jar[domain.lower()] = simpleCookie
def get(self, host: str) -> str: def get(self, host: str) -> str:
@ -58,7 +61,15 @@ class SimpleCookieJar:
if host.endswith(domain) or host == domain[1:]: if host.endswith(domain) or host == domain[1:]:
cookies.append(self.jar.get(domain)) cookies.append(self.jar.get(domain))
return "; ".join(filter( return "; ".join(
None, sorted( filter(
["%s=%s" % (k, v.value) for cookie in filter(None, cookies) for k, v in cookie.items()] None,
))) sorted(
[
"%s=%s" % (k, v.value)
for cookie in filter(None, cookies)
for k, v in cookie.items()
]
),
)
)

View file

@ -2,6 +2,7 @@ import socket
import struct import struct
import threading import threading
import time import time
from typing import Optional, Union
# websocket modules # websocket modules
from ._abnf import * from ._abnf import *
@ -32,7 +33,7 @@ See the License for the specific language governing permissions and
limitations under the License. limitations under the License.
""" """
__all__ = ['WebSocket', 'create_connection'] __all__ = ["WebSocket", "create_connection"]
class WebSocket: class WebSocket:
@ -73,9 +74,16 @@ class WebSocket:
Skip utf8 validation. Skip utf8 validation.
""" """
def __init__(self, get_mask_key=None, sockopt=None, sslopt=None, def __init__(
fire_cont_frame: bool = False, enable_multithread: bool = True, self,
skip_utf8_validation: bool = False, **_): get_mask_key=None,
sockopt=None,
sslopt=None,
fire_cont_frame: bool = False,
enable_multithread: bool = True,
skip_utf8_validation: bool = False,
**_,
):
""" """
Initialize WebSocket object. Initialize WebSocket object.
@ -86,14 +94,13 @@ class WebSocket:
""" """
self.sock_opt = sock_opt(sockopt, sslopt) self.sock_opt = sock_opt(sockopt, sslopt)
self.handshake_response = None self.handshake_response = None
self.sock = None self.sock: Optional[socket.socket] = None
self.connected = False self.connected = False
self.get_mask_key = get_mask_key self.get_mask_key = get_mask_key
# These buffer over the build-up of a single frame. # These buffer over the build-up of a single frame.
self.frame_buffer = frame_buffer(self._recv, skip_utf8_validation) self.frame_buffer = frame_buffer(self._recv, skip_utf8_validation)
self.cont_frame = continuous_frame( self.cont_frame = continuous_frame(fire_cont_frame, skip_utf8_validation)
fire_cont_frame, skip_utf8_validation)
if enable_multithread: if enable_multithread:
self.lock = threading.Lock() self.lock = threading.Lock()
@ -133,7 +140,7 @@ class WebSocket:
""" """
self.get_mask_key = func self.get_mask_key = func
def gettimeout(self) -> float: def gettimeout(self) -> Union[float, int, None]:
""" """
Get the websocket timeout (in seconds) as an int or float Get the websocket timeout (in seconds) as an int or float
@ -144,7 +151,7 @@ class WebSocket:
""" """
return self.sock_opt.timeout return self.sock_opt.timeout
def settimeout(self, timeout: float): def settimeout(self, timeout: Union[float, int, None]):
""" """
Set the timeout to the websocket. Set the timeout to the websocket.
@ -245,19 +252,26 @@ class WebSocket:
socket: socket socket: socket
Pre-initialized stream socket. Pre-initialized stream socket.
""" """
self.sock_opt.timeout = options.get('timeout', self.sock_opt.timeout) self.sock_opt.timeout = options.get("timeout", self.sock_opt.timeout)
self.sock, addrs = connect(url, self.sock_opt, proxy_info(**options), self.sock, addrs = connect(
options.pop('socket', None)) url, self.sock_opt, proxy_info(**options), options.pop("socket", None)
)
try: try:
self.handshake_response = handshake(self.sock, url, *addrs, **options) self.handshake_response = handshake(self.sock, url, *addrs, **options)
for attempt in range(options.pop('redirect_limit', 3)): for attempt in range(options.pop("redirect_limit", 3)):
if self.handshake_response.status in SUPPORTED_REDIRECT_STATUSES: if self.handshake_response.status in SUPPORTED_REDIRECT_STATUSES:
url = self.handshake_response.headers['location'] url = self.handshake_response.headers["location"]
self.sock.close() self.sock.close()
self.sock, addrs = connect(url, self.sock_opt, proxy_info(**options), self.sock, addrs = connect(
options.pop('socket', None)) url,
self.handshake_response = handshake(self.sock, url, *addrs, **options) self.sock_opt,
proxy_info(**options),
options.pop("socket", None),
)
self.handshake_response = handshake(
self.sock, url, *addrs, **options
)
self.connected = True self.connected = True
except: except:
if self.sock: if self.sock:
@ -265,7 +279,7 @@ class WebSocket:
self.sock = None self.sock = None
raise raise
def send(self, payload: bytes or str, opcode: int = ABNF.OPCODE_TEXT) -> int: def send(self, payload: Union[bytes, str], opcode: int = ABNF.OPCODE_TEXT) -> int:
""" """
Send the data as string. Send the data as string.
@ -282,6 +296,18 @@ class WebSocket:
frame = ABNF.create_frame(payload, opcode) frame = ABNF.create_frame(payload, opcode)
return self.send_frame(frame) 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: def send_frame(self, frame) -> int:
""" """
Send the data frame. Send the data frame.
@ -303,9 +329,9 @@ class WebSocket:
frame.get_mask_key = self.get_mask_key frame.get_mask_key = self.get_mask_key
data = frame.format() data = frame.format()
length = len(data) length = len(data)
if (isEnabledForTrace()): if isEnabledForTrace():
trace("++Sent raw: " + repr(data)) trace(f"++Sent raw: {repr(data)}")
trace("++Sent decoded: " + frame.__str__()) trace(f"++Sent decoded: {frame.__str__()}")
with self.lock: with self.lock:
while data: while data:
l = self._send(data) l = self._send(data)
@ -324,7 +350,7 @@ class WebSocket:
""" """
return self.send(payload, ABNF.OPCODE_BINARY) return self.send(payload, ABNF.OPCODE_BINARY)
def ping(self, payload: str or bytes = ""): def ping(self, payload: Union[str, bytes] = ""):
""" """
Send ping data. Send ping data.
@ -337,7 +363,7 @@ class WebSocket:
payload = payload.encode("utf-8") payload = payload.encode("utf-8")
self.send(payload, ABNF.OPCODE_PING) self.send(payload, ABNF.OPCODE_PING)
def pong(self, payload: str or bytes = ""): def pong(self, payload: Union[str, bytes] = ""):
""" """
Send pong data. Send pong data.
@ -350,7 +376,7 @@ class WebSocket:
payload = payload.encode("utf-8") payload = payload.encode("utf-8")
self.send(payload, ABNF.OPCODE_PONG) self.send(payload, ABNF.OPCODE_PONG)
def recv(self) -> str or bytes: def recv(self) -> Union[str, bytes]:
""" """
Receive string data(byte array) from the server. Receive string data(byte array) from the server.
@ -361,11 +387,16 @@ class WebSocket:
with self.readlock: with self.readlock:
opcode, data = self.recv_data() opcode, data = self.recv_data()
if opcode == ABNF.OPCODE_TEXT: if opcode == ABNF.OPCODE_TEXT:
return data.decode("utf-8") data_received: Union[bytes, str] = data
elif opcode == ABNF.OPCODE_TEXT or opcode == ABNF.OPCODE_BINARY: if isinstance(data_received, bytes):
return data 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: else:
return '' return ""
def recv_data(self, control_frame: bool = False) -> tuple: def recv_data(self, control_frame: bool = False) -> tuple:
""" """
@ -385,7 +416,7 @@ class WebSocket:
opcode, frame = self.recv_data_frame(control_frame) opcode, frame = self.recv_data_frame(control_frame)
return opcode, frame.data return opcode, frame.data
def recv_data_frame(self, control_frame: bool = False): def recv_data_frame(self, control_frame: bool = False) -> tuple:
""" """
Receive data with operation code. Receive data with operation code.
@ -404,15 +435,18 @@ class WebSocket:
""" """
while True: while True:
frame = self.recv_frame() frame = self.recv_frame()
if (isEnabledForTrace()): if isEnabledForTrace():
trace("++Rcv raw: " + repr(frame.format())) trace(f"++Rcv raw: {repr(frame.format())}")
trace("++Rcv decoded: " + frame.__str__()) trace(f"++Rcv decoded: {frame.__str__()}")
if not frame: if not frame:
# handle error: # handle error:
# 'NoneType' object has no attribute 'opcode' # 'NoneType' object has no attribute 'opcode'
raise WebSocketProtocolException( raise WebSocketProtocolException(f"Not a valid frame {frame}")
"Not a valid frame {frame}".format(frame=frame)) elif frame.opcode in (
elif frame.opcode in (ABNF.OPCODE_TEXT, ABNF.OPCODE_BINARY, ABNF.OPCODE_CONT): ABNF.OPCODE_TEXT,
ABNF.OPCODE_BINARY,
ABNF.OPCODE_CONT,
):
self.cont_frame.validate(frame) self.cont_frame.validate(frame)
self.cont_frame.add(frame) self.cont_frame.add(frame)
@ -426,8 +460,7 @@ class WebSocket:
if len(frame.data) < 126: if len(frame.data) < 126:
self.pong(frame.data) self.pong(frame.data)
else: else:
raise WebSocketProtocolException( raise WebSocketProtocolException("Ping message is too long")
"Ping message is too long")
if control_frame: if control_frame:
return frame.opcode, frame return frame.opcode, frame
elif frame.opcode == ABNF.OPCODE_PONG: elif frame.opcode == ABNF.OPCODE_PONG:
@ -458,9 +491,9 @@ class WebSocket:
if status < 0 or status >= ABNF.LENGTH_16: if status < 0 or status >= ABNF.LENGTH_16:
raise ValueError("code is invalid range") raise ValueError("code is invalid range")
self.connected = False self.connected = False
self.send(struct.pack('!H', status) + reason, ABNF.OPCODE_CLOSE) self.send(struct.pack("!H", status) + reason, ABNF.OPCODE_CLOSE)
def close(self, status: int = STATUS_NORMAL, reason: bytes = b"", timeout: float = 3): def close(self, status: int = STATUS_NORMAL, reason: bytes = b"", timeout: int = 3):
""" """
Close Websocket object Close Websocket object
@ -474,36 +507,37 @@ class WebSocket:
Timeout until receive a close frame. Timeout until receive a close frame.
If None, it will wait forever until receive a close frame. If None, it will wait forever until receive a close frame.
""" """
if self.connected: if not self.connected:
if status < 0 or status >= ABNF.LENGTH_16: return
raise ValueError("code is invalid range") if status < 0 or status >= ABNF.LENGTH_16:
raise ValueError("code is invalid range")
try: try:
self.connected = False self.connected = False
self.send(struct.pack('!H', status) + reason, ABNF.OPCODE_CLOSE) self.send(struct.pack("!H", status) + reason, ABNF.OPCODE_CLOSE)
sock_timeout = self.sock.gettimeout() sock_timeout = self.sock.gettimeout()
self.sock.settimeout(timeout) self.sock.settimeout(timeout)
start_time = time.time() start_time = time.time()
while timeout is None or time.time() - start_time < timeout: while timeout is None or time.time() - start_time < timeout:
try: try:
frame = self.recv_frame() frame = self.recv_frame()
if frame.opcode != ABNF.OPCODE_CLOSE: if frame.opcode != ABNF.OPCODE_CLOSE:
continue continue
if isEnabledForError(): if isEnabledForError():
recv_status = struct.unpack("!H", frame.data[0:2])[0] recv_status = struct.unpack("!H", frame.data[0:2])[0]
if recv_status >= 3000 and recv_status <= 4999: if recv_status >= 3000 and recv_status <= 4999:
debug("close status: " + repr(recv_status)) debug(f"close status: {repr(recv_status)}")
elif recv_status != STATUS_NORMAL: elif recv_status != STATUS_NORMAL:
error("close status: " + repr(recv_status)) error(f"close status: {repr(recv_status)}")
break break
except: except:
break break
self.sock.settimeout(sock_timeout) self.sock.settimeout(sock_timeout)
self.sock.shutdown(socket.SHUT_RDWR) self.sock.shutdown(socket.SHUT_RDWR)
except: except:
pass pass
self.shutdown() self.shutdown()
def abort(self): def abort(self):
""" """
@ -521,7 +555,7 @@ class WebSocket:
self.sock = None self.sock = None
self.connected = False self.connected = False
def _send(self, data: str or bytes): def _send(self, data: Union[str, bytes]):
return send(self.sock, data) return send(self.sock, data)
def _recv(self, bufsize): def _recv(self, bufsize):
@ -600,10 +634,14 @@ def create_connection(url: str, timeout=None, class_=WebSocket, **options):
fire_cont_frame = options.pop("fire_cont_frame", False) fire_cont_frame = options.pop("fire_cont_frame", False)
enable_multithread = options.pop("enable_multithread", True) enable_multithread = options.pop("enable_multithread", True)
skip_utf8_validation = options.pop("skip_utf8_validation", False) skip_utf8_validation = options.pop("skip_utf8_validation", False)
websock = class_(sockopt=sockopt, sslopt=sslopt, websock = class_(
fire_cont_frame=fire_cont_frame, sockopt=sockopt,
enable_multithread=enable_multithread, sslopt=sslopt,
skip_utf8_validation=skip_utf8_validation, **options) 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.settimeout(timeout if timeout is not None else getdefaulttimeout())
websock.connect(url, **options) websock.connect(url, **options)
return websock return websock

View file

@ -22,6 +22,7 @@ class WebSocketException(Exception):
""" """
WebSocket exception class. WebSocket exception class.
""" """
pass pass
@ -29,6 +30,7 @@ class WebSocketProtocolException(WebSocketException):
""" """
If the WebSocket protocol is invalid, this exception will be raised. If the WebSocket protocol is invalid, this exception will be raised.
""" """
pass pass
@ -36,6 +38,7 @@ class WebSocketPayloadException(WebSocketException):
""" """
If the WebSocket payload is invalid, this exception will be raised. If the WebSocket payload is invalid, this exception will be raised.
""" """
pass pass
@ -44,6 +47,7 @@ class WebSocketConnectionClosedException(WebSocketException):
If remote host closed the connection or some network error happened, If remote host closed the connection or some network error happened,
this exception will be raised. this exception will be raised.
""" """
pass pass
@ -51,6 +55,7 @@ class WebSocketTimeoutException(WebSocketException):
""" """
WebSocketTimeoutException will be raised at socket timeout during read/write data. WebSocketTimeoutException will be raised at socket timeout during read/write data.
""" """
pass pass
@ -58,6 +63,7 @@ class WebSocketProxyException(WebSocketException):
""" """
WebSocketProxyException will be raised when proxy error occurred. WebSocketProxyException will be raised when proxy error occurred.
""" """
pass pass
@ -66,7 +72,14 @@ class WebSocketBadStatusException(WebSocketException):
WebSocketBadStatusException will be raised when we get bad handshake status code. 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): def __init__(
self,
message: str,
status_code: int,
status_message=None,
resp_headers=None,
resp_body=None,
):
super().__init__(message) super().__init__(message)
self.status_code = status_code self.status_code = status_code
self.resp_headers = resp_headers self.resp_headers = resp_headers
@ -77,4 +90,5 @@ class WebSocketAddressException(WebSocketException):
""" """
If the websocket address info cannot be found, this exception will be raised. If the websocket address info cannot be found, this exception will be raised.
""" """
pass pass

View file

@ -20,7 +20,8 @@ import hashlib
import hmac import hmac
import os import os
from base64 import encodebytes as base64encode from base64 import encodebytes as base64encode
from http import client as HTTPStatus from http import HTTPStatus
from ._cookiejar import SimpleCookieJar from ._cookiejar import SimpleCookieJar
from ._exceptions import * from ._exceptions import *
from ._http import * from ._http import *
@ -32,14 +33,19 @@ __all__ = ["handshake_response", "handshake", "SUPPORTED_REDIRECT_STATUSES"]
# websocket supported version. # websocket supported version.
VERSION = 13 VERSION = 13
SUPPORTED_REDIRECT_STATUSES = (HTTPStatus.MOVED_PERMANENTLY, HTTPStatus.FOUND, HTTPStatus.SEE_OTHER,) 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,) SUCCESS_STATUSES = SUPPORTED_REDIRECT_STATUSES + (HTTPStatus.SWITCHING_PROTOCOLS,)
CookieJar = SimpleCookieJar() CookieJar = SimpleCookieJar()
class handshake_response: class handshake_response:
def __init__(self, status: int, headers: dict, subprotocol): def __init__(self, status: int, headers: dict, subprotocol):
self.status = status self.status = status
self.headers = headers self.headers = headers
@ -47,7 +53,9 @@ class handshake_response:
CookieJar.add(headers.get("set-cookie")) CookieJar.add(headers.get("set-cookie"))
def handshake(sock, url: str, hostname: str, port: int, resource: str, **options): def handshake(
sock, url: str, hostname: str, port: int, resource: str, **options
) -> handshake_response:
headers, key = _get_handshake_headers(resource, url, hostname, port, options) headers, key = _get_handshake_headers(resource, url, hostname, port, options)
header_str = "\r\n".join(headers) header_str = "\r\n".join(headers)
@ -66,74 +74,64 @@ def handshake(sock, url: str, hostname: str, port: int, resource: str, **options
def _pack_hostname(hostname: str) -> str: def _pack_hostname(hostname: str) -> str:
# IPv6 address # IPv6 address
if ':' in hostname: if ":" in hostname:
return '[' + hostname + ']' return f"[{hostname}]"
return hostname return hostname
def _get_handshake_headers(resource: str, url: str, host: str, port: int, options: dict): def _get_handshake_headers(
headers = [ resource: str, url: str, host: str, port: int, options: dict
"GET {resource} HTTP/1.1".format(resource=resource), ) -> tuple:
"Upgrade: websocket" headers = [f"GET {resource} HTTP/1.1", "Upgrade: websocket"]
] if port in [80, 443]:
if port == 80 or port == 443:
hostport = _pack_hostname(host) hostport = _pack_hostname(host)
else: else:
hostport = "{h}:{p}".format(h=_pack_hostname(host), p=port) hostport = f"{_pack_hostname(host)}:{port}"
if options.get("host"): if options.get("host"):
headers.append("Host: {h}".format(h=options["host"])) headers.append(f'Host: {options["host"]}')
else: else:
headers.append("Host: {hp}".format(hp=hostport)) headers.append(f"Host: {hostport}")
# scheme indicates whether http or https is used in Origin # 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 # The same approach is used in parse_url of _url.py to set default port
scheme, url = url.split(":", 1) scheme, url = url.split(":", 1)
if not options.get("suppress_origin"): if not options.get("suppress_origin"):
if "origin" in options and options["origin"] is not None: if "origin" in options and options["origin"] is not None:
headers.append("Origin: {origin}".format(origin=options["origin"])) headers.append(f'Origin: {options["origin"]}')
elif scheme == "wss": elif scheme == "wss":
headers.append("Origin: https://{hp}".format(hp=hostport)) headers.append(f"Origin: https://{hostport}")
else: else:
headers.append("Origin: http://{hp}".format(hp=hostport)) headers.append(f"Origin: http://{hostport}")
key = _create_sec_websocket_key() key = _create_sec_websocket_key()
# Append Sec-WebSocket-Key & Sec-WebSocket-Version if not manually specified # Append Sec-WebSocket-Key & Sec-WebSocket-Version if not manually specified
if not options.get('header') or 'Sec-WebSocket-Key' not in options['header']: if not options.get("header") or "Sec-WebSocket-Key" not in options["header"]:
headers.append("Sec-WebSocket-Key: {key}".format(key=key)) headers.append(f"Sec-WebSocket-Key: {key}")
else: else:
key = options['header']['Sec-WebSocket-Key'] key = options["header"]["Sec-WebSocket-Key"]
if not options.get('header') or 'Sec-WebSocket-Version' not in options['header']: if not options.get("header") or "Sec-WebSocket-Version" not in options["header"]:
headers.append("Sec-WebSocket-Version: {version}".format(version=VERSION)) headers.append(f"Sec-WebSocket-Version: {VERSION}")
if not options.get('connection'): if not options.get("connection"):
headers.append('Connection: Upgrade') headers.append("Connection: Upgrade")
else: else:
headers.append(options['connection']) headers.append(options["connection"])
subprotocols = options.get("subprotocols") if subprotocols := options.get("subprotocols"):
if subprotocols: headers.append(f'Sec-WebSocket-Protocol: {",".join(subprotocols)}')
headers.append("Sec-WebSocket-Protocol: {protocols}".format(protocols=",".join(subprotocols)))
header = options.get("header") if header := options.get("header"):
if header:
if isinstance(header, dict): if isinstance(header, dict):
header = [ header = [": ".join([k, v]) for k, v in header.items() if v is not None]
": ".join([k, v])
for k, v in header.items()
if v is not None
]
headers.extend(header) headers.extend(header)
server_cookie = CookieJar.get(host) server_cookie = CookieJar.get(host)
client_cookie = options.get("cookie", None) client_cookie = options.get("cookie", None)
cookie = "; ".join(filter(None, [server_cookie, client_cookie])) if cookie := "; ".join(filter(None, [server_cookie, client_cookie])):
headers.append(f"Cookie: {cookie}")
if cookie:
headers.append("Cookie: {cookie}".format(cookie=cookie))
headers.extend(("", "")) headers.extend(("", ""))
return headers, key return headers, key
@ -142,12 +140,20 @@ def _get_handshake_headers(resource: str, url: str, host: str, port: int, option
def _get_resp_headers(sock, success_statuses: tuple = SUCCESS_STATUSES) -> tuple: def _get_resp_headers(sock, success_statuses: tuple = SUCCESS_STATUSES) -> tuple:
status, resp_headers, status_message = read_headers(sock) status, resp_headers, status_message = read_headers(sock)
if status not in success_statuses: if status not in success_statuses:
content_len = resp_headers.get('content-length') content_len = resp_headers.get("content-length")
if content_len: if content_len:
response_body = sock.recv(int(content_len)) # read the body of the HTTP error message response and include it in the exception response_body = sock.recv(
int(content_len)
) # read the body of the HTTP error message response and include it in the exception
else: else:
response_body = None response_body = None
raise WebSocketBadStatusException("Handshake status {status} {message} -+-+- {headers} -+-+- {body}".format(status=status, message=status_message, headers=resp_headers, body=response_body), status, status_message, resp_headers, response_body) raise WebSocketBadStatusException(
f"Handshake status {status} {status_message} -+-+- {resp_headers} -+-+- {response_body}",
status,
status_message,
resp_headers,
response_body,
)
return status, resp_headers return status, resp_headers
@ -157,20 +163,20 @@ _HEADERS_TO_CHECK = {
} }
def _validate(headers, key: str, subprotocols): def _validate(headers, key: str, subprotocols) -> tuple:
subproto = None subproto = None
for k, v in _HEADERS_TO_CHECK.items(): for k, v in _HEADERS_TO_CHECK.items():
r = headers.get(k, None) r = headers.get(k, None)
if not r: if not r:
return False, None return False, None
r = [x.strip().lower() for x in r.split(',')] r = [x.strip().lower() for x in r.split(",")]
if v not in r: if v not in r:
return False, None return False, None
if subprotocols: if subprotocols:
subproto = headers.get("sec-websocket-protocol", None) subproto = headers.get("sec-websocket-protocol", None)
if not subproto or subproto.lower() not in [s.lower() for s in subprotocols]: if not subproto or subproto.lower() not in [s.lower() for s in subprotocols]:
error("Invalid subprotocol: " + str(subprotocols)) error(f"Invalid subprotocol: {subprotocols}")
return False, None return False, None
subproto = subproto.lower() subproto = subproto.lower()
@ -180,13 +186,12 @@ def _validate(headers, key: str, subprotocols):
result = result.lower() result = result.lower()
if isinstance(result, str): if isinstance(result, str):
result = result.encode('utf-8') result = result.encode("utf-8")
value = (key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").encode('utf-8') value = f"{key}258EAFA5-E914-47DA-95CA-C5AB0DC85B11".encode("utf-8")
hashed = base64encode(hashlib.sha1(value).digest()).strip().lower() hashed = base64encode(hashlib.sha1(value).digest()).strip().lower()
success = hmac.compare_digest(hashed, result)
if success: if hmac.compare_digest(hashed, result):
return True, subproto return True, subproto
else: else:
return False, None return False, None
@ -194,4 +199,4 @@ def _validate(headers, key: str, subprotocols):
def _create_sec_websocket_key() -> str: def _create_sec_websocket_key() -> str:
randomness = os.urandom(16) randomness = os.urandom(16)
return base64encode(randomness).decode('utf-8').strip() return base64encode(randomness).decode("utf-8").strip()

View file

@ -19,6 +19,7 @@ limitations under the License.
import errno import errno
import os import os
import socket import socket
from base64 import encodebytes as base64encode
from ._exceptions import * from ._exceptions import *
from ._logging import * from ._logging import *
@ -26,14 +27,13 @@ from ._socket import *
from ._ssl_compat import * from ._ssl_compat import *
from ._url import * from ._url import *
from base64 import encodebytes as base64encode
__all__ = ["proxy_info", "connect", "read_headers"] __all__ = ["proxy_info", "connect", "read_headers"]
try: try:
from python_socks.sync import Proxy
from python_socks._errors import * from python_socks._errors import *
from python_socks._types import ProxyType from python_socks._types import ProxyType
from python_socks.sync import Proxy
HAVE_PYTHON_SOCKS = True HAVE_PYTHON_SOCKS = True
except: except:
HAVE_PYTHON_SOCKS = False HAVE_PYTHON_SOCKS = False
@ -49,7 +49,6 @@ except:
class proxy_info: class proxy_info:
def __init__(self, **options): def __init__(self, **options):
self.proxy_host = options.get("http_proxy_host", None) self.proxy_host = options.get("http_proxy_host", None)
if self.proxy_host: if self.proxy_host:
@ -59,8 +58,16 @@ class proxy_info:
self.proxy_protocol = options.get("proxy_type", "http") self.proxy_protocol = options.get("proxy_type", "http")
# Note: If timeout not specified, default python-socks timeout is 60 seconds # Note: If timeout not specified, default python-socks timeout is 60 seconds
self.proxy_timeout = options.get("http_proxy_timeout", None) self.proxy_timeout = options.get("http_proxy_timeout", None)
if self.proxy_protocol not in ['http', 'socks4', 'socks4a', 'socks5', 'socks5h']: if self.proxy_protocol not in [
raise ProxyError("Only http, socks4, socks5 proxy protocols are supported") "http",
"socks4",
"socks4a",
"socks5",
"socks5h",
]:
raise ProxyError(
"Only http, socks4, socks5 proxy protocols are supported"
)
else: else:
self.proxy_port = 0 self.proxy_port = 0
self.auth = None self.auth = None
@ -68,25 +75,28 @@ class proxy_info:
self.proxy_protocol = "http" self.proxy_protocol = "http"
def _start_proxied_socket(url: str, options, proxy): def _start_proxied_socket(url: str, options, proxy) -> tuple:
if not HAVE_PYTHON_SOCKS: if not HAVE_PYTHON_SOCKS:
raise WebSocketException("Python Socks is needed for SOCKS proxying but is not available") raise WebSocketException(
"Python Socks is needed for SOCKS proxying but is not available"
)
hostname, port, resource, is_secure = parse_url(url) hostname, port, resource, is_secure = parse_url(url)
if proxy.proxy_protocol == "socks5":
rdns = False
proxy_type = ProxyType.SOCKS5
if proxy.proxy_protocol == "socks4": if proxy.proxy_protocol == "socks4":
rdns = False rdns = False
proxy_type = ProxyType.SOCKS4 proxy_type = ProxyType.SOCKS4
# socks5h and socks4a send DNS through proxy # socks4a sends DNS through proxy
if proxy.proxy_protocol == "socks5h": elif proxy.proxy_protocol == "socks4a":
rdns = True
proxy_type = ProxyType.SOCKS5
if proxy.proxy_protocol == "socks4a":
rdns = True rdns = True
proxy_type = ProxyType.SOCKS4 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( ws_proxy = Proxy.create(
proxy_type=proxy_type, proxy_type=proxy_type,
@ -94,14 +104,16 @@ def _start_proxied_socket(url: str, options, proxy):
port=int(proxy.proxy_port), port=int(proxy.proxy_port),
username=proxy.auth[0] if proxy.auth else None, username=proxy.auth[0] if proxy.auth else None,
password=proxy.auth[1] if proxy.auth else None, password=proxy.auth[1] if proxy.auth else None,
rdns=rdns) rdns=rdns,
)
sock = ws_proxy.connect(hostname, port, timeout=proxy.proxy_timeout) sock = ws_proxy.connect(hostname, port, timeout=proxy.proxy_timeout)
if is_secure and HAVE_SSL: if is_secure:
sock = _ssl_socket(sock, options.sslopt, hostname) if HAVE_SSL:
elif is_secure: sock = _ssl_socket(sock, options.sslopt, hostname)
raise WebSocketException("SSL not available.") else:
raise WebSocketException("SSL not available.")
return sock, (hostname, port, resource) return sock, (hostname, port, resource)
@ -110,7 +122,7 @@ def connect(url: str, options, proxy, socket):
# Use _start_proxied_socket() only for socks4 or socks5 proxy # Use _start_proxied_socket() only for socks4 or socks5 proxy
# Use _tunnel() for http proxy # Use _tunnel() for http proxy
# TODO: Use python-socks for http protocol also, to standardize flow # TODO: Use python-socks for http protocol also, to standardize flow
if proxy.proxy_host and not socket and not (proxy.proxy_protocol == "http"): if proxy.proxy_host and not socket and proxy.proxy_protocol != "http":
return _start_proxied_socket(url, options, proxy) return _start_proxied_socket(url, options, proxy)
hostname, port_from_url, resource, is_secure = parse_url(url) hostname, port_from_url, resource, is_secure = parse_url(url)
@ -119,10 +131,10 @@ def connect(url: str, options, proxy, socket):
return socket, (hostname, port_from_url, resource) return socket, (hostname, port_from_url, resource)
addrinfo_list, need_tunnel, auth = _get_addrinfo_list( addrinfo_list, need_tunnel, auth = _get_addrinfo_list(
hostname, port_from_url, is_secure, proxy) hostname, port_from_url, is_secure, proxy
)
if not addrinfo_list: if not addrinfo_list:
raise WebSocketException( raise WebSocketException(f"Host not found.: {hostname}:{port_from_url}")
"Host not found.: " + hostname + ":" + str(port_from_url))
sock = None sock = None
try: try:
@ -143,16 +155,23 @@ def connect(url: str, options, proxy, socket):
raise raise
def _get_addrinfo_list(hostname, port, is_secure, proxy): def _get_addrinfo_list(hostname, port: int, is_secure: bool, proxy) -> tuple:
phost, pport, pauth = get_proxy_info( phost, pport, pauth = get_proxy_info(
hostname, is_secure, proxy.proxy_host, proxy.proxy_port, proxy.auth, proxy.no_proxy) hostname,
is_secure,
proxy.proxy_host,
proxy.proxy_port,
proxy.auth,
proxy.no_proxy,
)
try: try:
# when running on windows 10, getaddrinfo without socktype returns a socktype 0. # 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` # 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. # or `OSError: [Errno 22] Invalid argument` when creating socket. Force the socket type to SOCK_STREAM.
if not phost: if not phost:
addrinfo_list = socket.getaddrinfo( addrinfo_list = socket.getaddrinfo(
hostname, port, 0, socket.SOCK_STREAM, socket.SOL_TCP) hostname, port, 0, socket.SOCK_STREAM, socket.SOL_TCP
)
return addrinfo_list, False, None return addrinfo_list, False, None
else: else:
pport = pport and pport or 80 pport = pport and pport or 80
@ -160,7 +179,9 @@ def _get_addrinfo_list(hostname, port, is_secure, proxy):
# returns a socktype 0. This generates an error exception: # returns a socktype 0. This generates an error exception:
# _on_error: exception Socket type must be stream or datagram, not 0 # _on_error: exception Socket type must be stream or datagram, not 0
# Force the socket type to SOCK_STREAM # Force the socket type to SOCK_STREAM
addrinfo_list = socket.getaddrinfo(phost, pport, 0, socket.SOCK_STREAM, socket.SOL_TCP) addrinfo_list = socket.getaddrinfo(
phost, pport, 0, socket.SOCK_STREAM, socket.SOL_TCP
)
return addrinfo_list, True, pauth return addrinfo_list, True, pauth
except socket.gaierror as e: except socket.gaierror as e:
raise WebSocketAddressException(e) raise WebSocketAddressException(e)
@ -186,14 +207,17 @@ def _open_socket(addrinfo_list, sockopt, timeout):
sock.close() sock.close()
error.remote_ip = str(address[0]) error.remote_ip = str(address[0])
try: try:
eConnRefused = (errno.ECONNREFUSED, errno.WSAECONNREFUSED, errno.ENETUNREACH) eConnRefused = (
errno.ECONNREFUSED,
errno.WSAECONNREFUSED,
errno.ENETUNREACH,
)
except AttributeError: except AttributeError:
eConnRefused = (errno.ECONNREFUSED, errno.ENETUNREACH) eConnRefused = (errno.ECONNREFUSED, errno.ENETUNREACH)
if error.errno in eConnRefused: if error.errno not in eConnRefused:
err = error
continue
else:
raise error raise error
err = error
continue
else: else:
break break
else: else:
@ -206,89 +230,97 @@ def _open_socket(addrinfo_list, sockopt, timeout):
return sock return sock
def _wrap_sni_socket(sock, sslopt, hostname, check_hostname): def _wrap_sni_socket(sock: socket.socket, sslopt: dict, hostname, check_hostname):
context = sslopt.get('context', None) context = sslopt.get("context", None)
if not context: if not context:
context = ssl.SSLContext(sslopt.get('ssl_version', ssl.PROTOCOL_TLS_CLIENT)) 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. # Non default context need to manually enable SSLKEYLOGFILE support by setting the keylog_filename attribute.
# For more details see also: # 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#context-creation
# * https://docs.python.org/3.8/library/ssl.html?highlight=sslkeylogfile#ssl.SSLContext.keylog_filename # * https://docs.python.org/3.8/library/ssl.html?highlight=sslkeylogfile#ssl.SSLContext.keylog_filename
context.keylog_filename = os.environ.get("SSLKEYLOGFILE", None) context.keylog_filename = os.environ.get("SSLKEYLOGFILE", None)
if sslopt.get('cert_reqs', ssl.CERT_NONE) != ssl.CERT_NONE: if sslopt.get("cert_reqs", ssl.CERT_NONE) != ssl.CERT_NONE:
cafile = sslopt.get('ca_certs', None) cafile = sslopt.get("ca_certs", None)
capath = sslopt.get('ca_cert_path', None) capath = sslopt.get("ca_cert_path", None)
if cafile or capath: if cafile or capath:
context.load_verify_locations(cafile=cafile, capath=capath) context.load_verify_locations(cafile=cafile, capath=capath)
elif hasattr(context, 'load_default_certs'): elif hasattr(context, "load_default_certs"):
context.load_default_certs(ssl.Purpose.SERVER_AUTH) context.load_default_certs(ssl.Purpose.SERVER_AUTH)
if sslopt.get('certfile', None): if sslopt.get("certfile", None):
context.load_cert_chain( context.load_cert_chain(
sslopt['certfile'], sslopt["certfile"],
sslopt.get('keyfile', None), sslopt.get("keyfile", None),
sslopt.get('password', None), sslopt.get("password", None),
) )
# Python 3.10 switch to PROTOCOL_TLS_CLIENT defaults to "cert_reqs = ssl.CERT_REQUIRED" and "check_hostname = True" # 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 # If both disabled, set check_hostname before verify_mode
# see https://github.com/liris/websocket-client/commit/b96a2e8fa765753e82eea531adb19716b52ca3ca#commitcomment-10803153 # 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): if sslopt.get("cert_reqs", ssl.CERT_NONE) == ssl.CERT_NONE and not sslopt.get(
"check_hostname", False
):
context.check_hostname = False context.check_hostname = False
context.verify_mode = ssl.CERT_NONE context.verify_mode = ssl.CERT_NONE
else: else:
context.check_hostname = sslopt.get('check_hostname', True) context.check_hostname = sslopt.get("check_hostname", True)
context.verify_mode = sslopt.get('cert_reqs', ssl.CERT_REQUIRED) context.verify_mode = sslopt.get("cert_reqs", ssl.CERT_REQUIRED)
if 'ciphers' in sslopt: if "ciphers" in sslopt:
context.set_ciphers(sslopt['ciphers']) context.set_ciphers(sslopt["ciphers"])
if 'cert_chain' in sslopt: if "cert_chain" in sslopt:
certfile, keyfile, password = sslopt['cert_chain'] certfile, keyfile, password = sslopt["cert_chain"]
context.load_cert_chain(certfile, keyfile, password) context.load_cert_chain(certfile, keyfile, password)
if 'ecdh_curve' in sslopt: if "ecdh_curve" in sslopt:
context.set_ecdh_curve(sslopt['ecdh_curve']) context.set_ecdh_curve(sslopt["ecdh_curve"])
return context.wrap_socket( return context.wrap_socket(
sock, sock,
do_handshake_on_connect=sslopt.get('do_handshake_on_connect', True), do_handshake_on_connect=sslopt.get("do_handshake_on_connect", True),
suppress_ragged_eofs=sslopt.get('suppress_ragged_eofs', True), suppress_ragged_eofs=sslopt.get("suppress_ragged_eofs", True),
server_hostname=hostname, server_hostname=hostname,
) )
def _ssl_socket(sock, user_sslopt, hostname): def _ssl_socket(sock: socket.socket, user_sslopt: dict, hostname):
sslopt = dict(cert_reqs=ssl.CERT_REQUIRED) sslopt: dict = dict(cert_reqs=ssl.CERT_REQUIRED)
sslopt.update(user_sslopt) sslopt.update(user_sslopt)
certPath = os.environ.get('WEBSOCKET_CLIENT_CA_BUNDLE') certPath = os.environ.get("WEBSOCKET_CLIENT_CA_BUNDLE")
if certPath and os.path.isfile(certPath) \ if (
and user_sslopt.get('ca_certs', None) is None: certPath
sslopt['ca_certs'] = certPath and os.path.isfile(certPath)
elif certPath and os.path.isdir(certPath) \ and user_sslopt.get("ca_certs", None) is None
and user_sslopt.get('ca_cert_path', None) is None: ):
sslopt['ca_cert_path'] = certPath sslopt["ca_certs"] = certPath
elif (
certPath
and os.path.isdir(certPath)
and user_sslopt.get("ca_cert_path", None) is None
):
sslopt["ca_cert_path"] = certPath
if sslopt.get('server_hostname', None): if sslopt.get("server_hostname", None):
hostname = sslopt['server_hostname'] hostname = sslopt["server_hostname"]
check_hostname = sslopt.get('check_hostname', True) check_hostname = sslopt.get("check_hostname", True)
sock = _wrap_sni_socket(sock, sslopt, hostname, check_hostname) sock = _wrap_sni_socket(sock, sslopt, hostname, check_hostname)
return sock return sock
def _tunnel(sock, host, port, auth): def _tunnel(sock: socket.socket, host, port: int, auth) -> socket.socket:
debug("Connecting proxy...") debug("Connecting proxy...")
connect_header = "CONNECT {h}:{p} HTTP/1.1\r\n".format(h=host, p=port) connect_header = f"CONNECT {host}:{port} HTTP/1.1\r\n"
connect_header += "Host: {h}:{p}\r\n".format(h=host, p=port) connect_header += f"Host: {host}:{port}\r\n"
# TODO: support digest auth. # TODO: support digest auth.
if auth and auth[0]: if auth and auth[0]:
auth_str = auth[0] auth_str = auth[0]
if auth[1]: if auth[1]:
auth_str += ":" + auth[1] auth_str += f":{auth[1]}"
encoded_str = base64encode(auth_str.encode()).strip().decode().replace('\n', '') encoded_str = base64encode(auth_str.encode()).strip().decode().replace("\n", "")
connect_header += "Proxy-Authorization: Basic {str}\r\n".format(str=encoded_str) connect_header += f"Proxy-Authorization: Basic {encoded_str}\r\n"
connect_header += "\r\n" connect_header += "\r\n"
dump("request header", connect_header) dump("request header", connect_header)
@ -300,40 +332,37 @@ def _tunnel(sock, host, port, auth):
raise WebSocketProxyException(str(e)) raise WebSocketProxyException(str(e))
if status != 200: if status != 200:
raise WebSocketProxyException( raise WebSocketProxyException(f"failed CONNECT via proxy status: {status}")
"failed CONNECT via proxy status: {status}".format(status=status))
return sock return sock
def read_headers(sock): def read_headers(sock: socket.socket) -> tuple:
status = None status = None
status_message = None status_message = None
headers = {} headers: dict = {}
trace("--- response header ---") trace("--- response header ---")
while True: while True:
line = recv_line(sock) line = recv_line(sock)
line = line.decode('utf-8').strip() line = line.decode("utf-8").strip()
if not line: if not line:
break break
trace(line) trace(line)
if not status: if not status:
status_info = line.split(" ", 2) status_info = line.split(" ", 2)
status = int(status_info[1]) status = int(status_info[1])
if len(status_info) > 2: if len(status_info) > 2:
status_message = status_info[2] status_message = status_info[2]
else: else:
kv = line.split(":", 1) kv = line.split(":", 1)
if len(kv) == 2: if len(kv) != 2:
key, value = kv
if key.lower() == "set-cookie" and headers.get("set-cookie"):
headers["set-cookie"] = headers.get("set-cookie") + "; " + value.strip()
else:
headers[key.lower()] = value.strip()
else:
raise WebSocketException("Invalid header") raise WebSocketException("Invalid header")
key, value = kv
if key.lower() == "set-cookie" and headers.get("set-cookie"):
headers["set-cookie"] = headers.get("set-cookie") + "; " + value.strip()
else:
headers[key.lower()] = value.strip()
trace("-----------------------") trace("-----------------------")

View file

@ -19,25 +19,38 @@ See the License for the specific language governing permissions and
limitations under the License. limitations under the License.
""" """
_logger = logging.getLogger('websocket') _logger = logging.getLogger("websocket")
try: try:
from logging import NullHandler from logging import NullHandler
except ImportError: except ImportError:
class NullHandler(logging.Handler): class NullHandler(logging.Handler):
def emit(self, record) -> None: def emit(self, record) -> None:
pass pass
_logger.addHandler(NullHandler()) _logger.addHandler(NullHandler())
_traceEnabled = False _traceEnabled = False
__all__ = ["enableTrace", "dump", "error", "warning", "debug", "trace", __all__ = [
"isEnabledForError", "isEnabledForDebug", "isEnabledForTrace"] "enableTrace",
"dump",
"error",
"warning",
"debug",
"trace",
"isEnabledForError",
"isEnabledForDebug",
"isEnabledForTrace",
]
def enableTrace(traceable: bool, def enableTrace(
handler: logging.StreamHandler = logging.StreamHandler(), traceable: bool,
level: str = "DEBUG") -> None: handler: logging.StreamHandler = logging.StreamHandler(),
level: str = "DEBUG",
) -> None:
""" """
Turn on/off the traceability. Turn on/off the traceability.
@ -55,7 +68,7 @@ def enableTrace(traceable: bool,
def dump(title: str, message: str) -> None: def dump(title: str, message: str) -> None:
if _traceEnabled: if _traceEnabled:
_logger.debug("--- " + title + " ---") _logger.debug(f"--- {title} ---")
_logger.debug(message) _logger.debug(message)
_logger.debug("-----------------------") _logger.debug("-----------------------")

View file

@ -1,6 +1,7 @@
import errno import errno
import selectors import selectors
import socket import socket
from typing import Union
from ._exceptions import * from ._exceptions import *
from ._ssl_compat import * from ._ssl_compat import *
@ -37,12 +38,18 @@ if hasattr(socket, "TCP_KEEPCNT"):
_default_timeout = None _default_timeout = None
__all__ = ["DEFAULT_SOCKET_OPTION", "sock_opt", "setdefaulttimeout", "getdefaulttimeout", __all__ = [
"recv", "recv_line", "send"] "DEFAULT_SOCKET_OPTION",
"sock_opt",
"setdefaulttimeout",
"getdefaulttimeout",
"recv",
"recv_line",
"send",
]
class sock_opt: class sock_opt:
def __init__(self, sockopt: list, sslopt: dict) -> None: def __init__(self, sockopt: list, sslopt: dict) -> None:
if sockopt is None: if sockopt is None:
sockopt = [] sockopt = []
@ -53,7 +60,7 @@ class sock_opt:
self.timeout = None self.timeout = None
def setdefaulttimeout(timeout: int or float) -> None: def setdefaulttimeout(timeout: Union[int, float, None]) -> None:
""" """
Set the global timeout setting to connect. Set the global timeout setting to connect.
@ -66,7 +73,7 @@ def setdefaulttimeout(timeout: int or float) -> None:
_default_timeout = timeout _default_timeout = timeout
def getdefaulttimeout() -> int or float: def getdefaulttimeout() -> Union[int, float, None]:
""" """
Get default timeout Get default timeout
@ -89,7 +96,7 @@ def recv(sock: socket.socket, bufsize: int) -> bytes:
pass pass
except socket.error as exc: except socket.error as exc:
error_code = extract_error_code(exc) error_code = extract_error_code(exc)
if error_code != errno.EAGAIN and error_code != errno.EWOULDBLOCK: if error_code not in [errno.EAGAIN, errno.EWOULDBLOCK]:
raise raise
sel = selectors.DefaultSelector() sel = selectors.DefaultSelector()
@ -113,14 +120,13 @@ def recv(sock: socket.socket, bufsize: int) -> bytes:
raise WebSocketTimeoutException(message) raise WebSocketTimeoutException(message)
except SSLError as e: except SSLError as e:
message = extract_err_message(e) message = extract_err_message(e)
if isinstance(message, str) and 'timed out' in message: if isinstance(message, str) and "timed out" in message:
raise WebSocketTimeoutException(message) raise WebSocketTimeoutException(message)
else: else:
raise raise
if not bytes_: if not bytes_:
raise WebSocketConnectionClosedException( raise WebSocketConnectionClosedException("Connection to remote host was lost.")
"Connection to remote host was lost.")
return bytes_ return bytes_
@ -130,14 +136,14 @@ def recv_line(sock: socket.socket) -> bytes:
while True: while True:
c = recv(sock, 1) c = recv(sock, 1)
line.append(c) line.append(c)
if c == b'\n': if c == b"\n":
break break
return b''.join(line) return b"".join(line)
def send(sock: socket.socket, data: bytes) -> int: def send(sock: socket.socket, data: Union[bytes, str]) -> int:
if isinstance(data, str): if isinstance(data, str):
data = data.encode('utf-8') data = data.encode("utf-8")
if not sock: if not sock:
raise WebSocketConnectionClosedException("socket is already closed.") raise WebSocketConnectionClosedException("socket is already closed.")
@ -151,7 +157,7 @@ def send(sock: socket.socket, data: bytes) -> int:
error_code = extract_error_code(exc) error_code = extract_error_code(exc)
if error_code is None: if error_code is None:
raise raise
if error_code != errno.EAGAIN and error_code != errno.EWOULDBLOCK: if error_code not in [errno.EAGAIN, errno.EWOULDBLOCK]:
raise raise
sel = selectors.DefaultSelector() sel = selectors.DefaultSelector()

View file

@ -20,9 +20,8 @@ __all__ = ["HAVE_SSL", "ssl", "SSLError", "SSLWantReadError", "SSLWantWriteError
try: try:
import ssl import ssl
from ssl import SSLError from ssl import SSLError, SSLWantReadError, SSLWantWriteError
from ssl import SSLWantReadError
from ssl import SSLWantWriteError
HAVE_SSL = True HAVE_SSL = True
except ImportError: except ImportError:
# dummy class of SSLError for environment without ssl support # dummy class of SSLError for environment without ssl support

View file

@ -1,7 +1,7 @@
import os import os
import socket import socket
import struct import struct
from typing import Optional
from urllib.parse import unquote, urlparse from urllib.parse import unquote, urlparse
""" """
@ -67,7 +67,7 @@ def parse_url(url: str) -> tuple:
resource = "/" resource = "/"
if parsed.query: if parsed.query:
resource += "?" + parsed.query resource += f"?{parsed.query}"
return hostname, port, resource, is_secure return hostname, port, resource, is_secure
@ -93,37 +93,50 @@ def _is_subnet_address(hostname: str) -> bool:
def _is_address_in_network(ip: str, net: str) -> bool: def _is_address_in_network(ip: str, net: str) -> bool:
ipaddr = struct.unpack('!I', socket.inet_aton(ip))[0] ipaddr: int = struct.unpack("!I", socket.inet_aton(ip))[0]
netaddr, netmask = net.split('/') netaddr, netmask = net.split("/")
netaddr = struct.unpack('!I', socket.inet_aton(netaddr))[0] netaddr: int = struct.unpack("!I", socket.inet_aton(netaddr))[0]
netmask = (0xFFFFFFFF << (32 - int(netmask))) & 0xFFFFFFFF netmask = (0xFFFFFFFF << (32 - int(netmask))) & 0xFFFFFFFF
return ipaddr & netmask == netaddr return ipaddr & netmask == netaddr
def _is_no_proxy_host(hostname: str, no_proxy: list) -> bool: def _is_no_proxy_host(hostname: str, no_proxy: Optional[list]) -> bool:
if not no_proxy: if not no_proxy:
v = os.environ.get("no_proxy", os.environ.get("NO_PROXY", "")).replace(" ", "") if v := os.environ.get("no_proxy", os.environ.get("NO_PROXY", "")).replace(
if v: " ", ""
):
no_proxy = v.split(",") no_proxy = v.split(",")
if not no_proxy: if not no_proxy:
no_proxy = DEFAULT_NO_PROXY_HOST no_proxy = DEFAULT_NO_PROXY_HOST
if '*' in no_proxy: if "*" in no_proxy:
return True return True
if hostname in no_proxy: if hostname in no_proxy:
return True return True
if _is_ip_address(hostname): if _is_ip_address(hostname):
return any([_is_address_in_network(hostname, subnet) for subnet in no_proxy if _is_subnet_address(subnet)]) return any(
for domain in [domain for domain in no_proxy if domain.startswith('.')]: [
_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(".")]:
if hostname.endswith(domain): if hostname.endswith(domain):
return True return True
return False return False
def get_proxy_info( def get_proxy_info(
hostname: str, is_secure: bool, proxy_host: str = None, proxy_port: int = 0, proxy_auth: tuple = None, hostname: str,
no_proxy: list = None, proxy_type: str = 'http') -> tuple: is_secure: bool,
proxy_host: Optional[str] = None,
proxy_port: int = 0,
proxy_auth: Optional[tuple] = None,
no_proxy: Optional[list] = None,
proxy_type: str = "http",
) -> tuple:
""" """
Try to retrieve proxy host and port from environment Try to retrieve proxy host and port from environment
if not provided in options. if not provided in options.
@ -159,10 +172,16 @@ def get_proxy_info(
return proxy_host, port, auth return proxy_host, port, auth
env_key = "https_proxy" if is_secure else "http_proxy" env_key = "https_proxy" if is_secure else "http_proxy"
value = os.environ.get(env_key, os.environ.get(env_key.upper(), "")).replace(" ", "") value = os.environ.get(env_key, os.environ.get(env_key.upper(), "")).replace(
" ", ""
)
if value: if value:
proxy = urlparse(value) proxy = urlparse(value)
auth = (unquote(proxy.username), unquote(proxy.password)) if proxy.username else None auth = (
(unquote(proxy.username), unquote(proxy.password))
if proxy.username
else None
)
return proxy.hostname, proxy.port, auth return proxy.hostname, proxy.port, auth
return None, 0, None return None, 0, None

View file

@ -1,3 +1,5 @@
from typing import Union
""" """
_url.py _url.py
websocket - WebSocket client library for Python websocket - WebSocket client library for Python
@ -20,7 +22,6 @@ __all__ = ["NoLock", "validate_utf8", "extract_err_message", "extract_error_code
class NoLock: class NoLock:
def __enter__(self) -> None: def __enter__(self) -> None:
pass pass
@ -33,8 +34,9 @@ try:
# strings. # strings.
from wsaccel.utf8validator import Utf8Validator from wsaccel.utf8validator import Utf8Validator
def _validate_utf8(utfbytes: bytes) -> bool: def _validate_utf8(utfbytes: Union[str, bytes]) -> bool:
return Utf8Validator().validate(utfbytes)[0] result: bool = Utf8Validator().validate(utfbytes)[0]
return result
except ImportError: except ImportError:
# UTF-8 validator # UTF-8 validator
@ -46,44 +48,396 @@ except ImportError:
_UTF8D = [ _UTF8D = [
# The first part of the table maps bytes to character classes that # The first part of the table maps bytes to character classes that
# to reduce the size of the transition table and create bitmasks. # 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,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, 0,
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, 0,
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, 0,
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, 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 # The second part is a transition table that maps a combination
# of a state of the automaton and a character class to a state. # 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, 0,
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,12,24,12,12,12,12, 12,24,12,12,12,12,12,12,12,24,12,12, 24,
12,12,12,12,12,12,12,36,12,36,12,12, 12,36,12,12,12,12,12,36,12,36,12,12, 36,
12,36,12,12,12,12,12,12,12,12,12,12, ] 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: def _decode(state: int, codep: int, ch: int) -> tuple:
tp = _UTF8D[ch] tp = _UTF8D[ch]
codep = (ch & 0x3f) | (codep << 6) if ( codep = (
state != _UTF8_ACCEPT) else (0xff >> tp) & ch (ch & 0x3F) | (codep << 6) if (state != _UTF8_ACCEPT) else (0xFF >> tp) & ch
)
state = _UTF8D[256 + state + tp] state = _UTF8D[256 + state + tp]
return state, codep return state, codep
def _validate_utf8(utfbytes: str or bytes) -> bool: def _validate_utf8(utfbytes: Union[str, bytes]) -> bool:
state = _UTF8_ACCEPT state = _UTF8_ACCEPT
codep = 0 codep = 0
for i in utfbytes: for i in utfbytes:
state, codep = _decode(state, codep, i) state, codep = _decode(state, codep, int(i))
if state == _UTF8_REJECT: if state == _UTF8_REJECT:
return False return False
return True return True
def validate_utf8(utfbytes: str or bytes) -> bool: def validate_utf8(utfbytes: Union[str, bytes]) -> bool:
""" """
validate utf8 byte string. validate utf8 byte string.
utfbytes: utf byte string to check. utfbytes: utf byte string to check.
@ -92,13 +446,14 @@ def validate_utf8(utfbytes: str or bytes) -> bool:
return _validate_utf8(utfbytes) return _validate_utf8(utfbytes)
def extract_err_message(exception: Exception) -> str or None: def extract_err_message(exception: Exception) -> Union[str, None]:
if exception.args: if exception.args:
return exception.args[0] exception_message: str = exception.args[0]
return exception_message
else: else:
return None return None
def extract_error_code(exception: Exception) -> int or None: def extract_error_code(exception: Exception) -> Union[int, None]:
if exception.args and len(exception.args) > 1: if exception.args and len(exception.args) > 1:
return exception.args[0] if isinstance(exception.args[0], int) else None return exception.args[0] if isinstance(exception.args[0], int) else None

View file

@ -21,11 +21,11 @@ limitations under the License.
import argparse import argparse
import code import code
import gzip
import ssl
import sys import sys
import threading import threading
import time import time
import ssl
import gzip
import zlib import zlib
from urllib.parse import urlparse from urllib.parse import urlparse
@ -50,8 +50,13 @@ ENCODING = get_encoding()
class VAction(argparse.Action): class VAction(argparse.Action):
def __call__(
def __call__(self, parser: argparse.Namespace, args: tuple, values: str, option_string: str = None) -> None: self,
parser: argparse.Namespace,
args: tuple,
values: str,
option_string: str = None,
) -> None:
if values is None: if values is None:
values = "1" values = "1"
try: try:
@ -63,36 +68,42 @@ class VAction(argparse.Action):
def parse_args() -> argparse.Namespace: def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="WebSocket Simple Dump Tool") parser = argparse.ArgumentParser(description="WebSocket Simple Dump Tool")
parser.add_argument("url", metavar="ws_url", parser.add_argument(
help="websocket url. ex. ws://echo.websocket.events/") "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("-p", "--proxy", help="proxy url. ex. http://127.0.0.1:8080")
parser.add_argument("-v", "--verbose", default=0, nargs='?', action=VAction, parser.add_argument(
dest="verbose", "-v",
help="set verbose mode. If set to 1, show opcode. " "--verbose",
"If set to 2, enable to trace websocket module") default=0,
parser.add_argument("-n", "--nocert", action='store_true', nargs="?",
help="Ignore invalid SSL cert") action=VAction,
parser.add_argument("-r", "--raw", action="store_true", dest="verbose",
help="raw output") help="set verbose mode. If set to 1, show opcode. "
parser.add_argument("-s", "--subprotocols", nargs='*', "If set to 2, enable to trace websocket module",
help="Set subprotocols") )
parser.add_argument("-o", "--origin", parser.add_argument(
help="Set origin") "-n", "--nocert", action="store_true", help="Ignore invalid SSL cert"
parser.add_argument("--eof-wait", default=0, type=int, )
help="wait time(second) after 'EOF' received.") parser.add_argument("-r", "--raw", action="store_true", help="raw output")
parser.add_argument("-t", "--text", parser.add_argument("-s", "--subprotocols", nargs="*", help="Set subprotocols")
help="Send initial text") parser.add_argument("-o", "--origin", help="Set origin")
parser.add_argument("--timings", action="store_true", parser.add_argument(
help="Print timings in seconds") "--eof-wait",
parser.add_argument("--headers", default=0,
help="Set custom headers. Use ',' as separator") 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() return parser.parse_args()
class RawInput: class RawInput:
def raw_input(self, prompt: str = "") -> str: def raw_input(self, prompt: str = "") -> str:
line = input(prompt) line = input(prompt)
@ -105,7 +116,6 @@ class RawInput:
class InteractiveConsole(RawInput, code.InteractiveConsole): class InteractiveConsole(RawInput, code.InteractiveConsole):
def write(self, data: str) -> None: def write(self, data: str) -> None:
sys.stdout.write("\033[2K\033[E") sys.stdout.write("\033[2K\033[E")
# sys.stdout.write("\n") # sys.stdout.write("\n")
@ -118,7 +128,6 @@ class InteractiveConsole(RawInput, code.InteractiveConsole):
class NonInteractive(RawInput): class NonInteractive(RawInput):
def write(self, data: str) -> None: def write(self, data: str) -> None:
sys.stdout.write(data) sys.stdout.write(data)
sys.stdout.write("\n") sys.stdout.write("\n")
@ -146,7 +155,7 @@ def main() -> None:
if args.nocert: if args.nocert:
opts = {"cert_reqs": ssl.CERT_NONE, "check_hostname": False} opts = {"cert_reqs": ssl.CERT_NONE, "check_hostname": False}
if args.headers: if args.headers:
options['header'] = list(map(str.strip, args.headers.split(','))) options["header"] = list(map(str.strip, args.headers.split(",")))
ws = websocket.create_connection(args.url, sslopt=opts, **options) ws = websocket.create_connection(args.url, sslopt=opts, **options)
if args.raw: if args.raw:
console = NonInteractive() console = NonInteractive()
@ -160,7 +169,7 @@ def main() -> None:
except websocket.WebSocketException: except websocket.WebSocketException:
return websocket.ABNF.OPCODE_CLOSE, "" return websocket.ABNF.OPCODE_CLOSE, ""
if not frame: if not frame:
raise websocket.WebSocketException("Not a valid frame {frame}".format(frame=frame)) raise websocket.WebSocketException(f"Not a valid frame {frame}")
elif frame.opcode in OPCODE_DATA: elif frame.opcode in OPCODE_DATA:
return frame.opcode, frame.data return frame.opcode, frame.data
elif frame.opcode == websocket.ABNF.OPCODE_CLOSE: elif frame.opcode == websocket.ABNF.OPCODE_CLOSE:
@ -178,14 +187,18 @@ def main() -> None:
msg = None msg = None
if opcode == websocket.ABNF.OPCODE_TEXT and isinstance(data, bytes): if opcode == websocket.ABNF.OPCODE_TEXT and isinstance(data, bytes):
data = str(data, "utf-8") data = str(data, "utf-8")
if isinstance(data, bytes) and len(data) > 2 and data[:2] == b'\037\213': # gzip magick if (
isinstance(data, bytes) and len(data) > 2 and data[:2] == b"\037\213"
): # gzip magick
try: try:
data = "[gzip] " + str(gzip.decompress(data), "utf-8") data = "[gzip] " + str(gzip.decompress(data), "utf-8")
except: except:
pass pass
elif isinstance(data, bytes): elif isinstance(data, bytes):
try: try:
data = "[zlib] " + str(zlib.decompress(data, -zlib.MAX_WBITS), "utf-8") data = "[zlib] " + str(
zlib.decompress(data, -zlib.MAX_WBITS), "utf-8"
)
except: except:
pass pass
@ -193,13 +206,13 @@ def main() -> None:
data = repr(data) data = repr(data)
if args.verbose: if args.verbose:
msg = "{opcode}: {data}".format(opcode=websocket.ABNF.OPCODE_MAP.get(opcode), data=data) msg = f"{websocket.ABNF.OPCODE_MAP.get(opcode)}: {data}"
else: else:
msg = data msg = data
if msg is not None: if msg is not None:
if args.timings: if args.timings:
console.write(str(time.time() - start_time) + ": " + msg) console.write(f"{time.time() - start_time}: {msg}")
else: else:
console.write(msg) console.write(msg)

0
lib/websocket/py.typed Normal file
View file

View file

@ -3,10 +3,11 @@
# From https://github.com/aaugustin/websockets/blob/main/example/echo.py # From https://github.com/aaugustin/websockets/blob/main/example/echo.py
import asyncio import asyncio
import websockets
import os import os
LOCAL_WS_SERVER_PORT = int(os.environ.get('LOCAL_WS_SERVER_PORT', '8765')) import websockets
LOCAL_WS_SERVER_PORT = int(os.environ.get("LOCAL_WS_SERVER_PORT", "8765"))
async def echo(websocket, path): async def echo(websocket, path):
@ -18,4 +19,5 @@ async def main():
async with websockets.serve(echo, "localhost", LOCAL_WS_SERVER_PORT): async with websockets.serve(echo, "localhost", LOCAL_WS_SERVER_PORT):
await asyncio.Future() # run forever await asyncio.Future() # run forever
asyncio.run(main()) asyncio.run(main())

View file

@ -1,8 +1,9 @@
# -*- coding: utf-8 -*- # -*- coding: utf-8 -*-
# #
import unittest
import websocket as ws import websocket as ws
from websocket._abnf import * from websocket._abnf import *
import unittest
""" """
test_abnf.py test_abnf.py
@ -25,54 +26,89 @@ limitations under the License.
class ABNFTest(unittest.TestCase): class ABNFTest(unittest.TestCase):
def testInit(self): def testInit(self):
a = ABNF(0,0,0,0, opcode=ABNF.OPCODE_PING) a = ABNF(0, 0, 0, 0, opcode=ABNF.OPCODE_PING)
self.assertEqual(a.fin, 0) self.assertEqual(a.fin, 0)
self.assertEqual(a.rsv1, 0) self.assertEqual(a.rsv1, 0)
self.assertEqual(a.rsv2, 0) self.assertEqual(a.rsv2, 0)
self.assertEqual(a.rsv3, 0) self.assertEqual(a.rsv3, 0)
self.assertEqual(a.opcode, 9) self.assertEqual(a.opcode, 9)
self.assertEqual(a.data, '') self.assertEqual(a.data, "")
a_bad = ABNF(0,1,0,0, opcode=77) a_bad = ABNF(0, 1, 0, 0, opcode=77)
self.assertEqual(a_bad.rsv1, 1) self.assertEqual(a_bad.rsv1, 1)
self.assertEqual(a_bad.opcode, 77) self.assertEqual(a_bad.opcode, 77)
def testValidate(self): def testValidate(self):
a_invalid_ping = ABNF(0,0,0,0, opcode=ABNF.OPCODE_PING) a_invalid_ping = ABNF(0, 0, 0, 0, opcode=ABNF.OPCODE_PING)
self.assertRaises(ws._exceptions.WebSocketProtocolException, a_invalid_ping.validate, skip_utf8_validation=False) self.assertRaises(
a_bad_rsv_value = ABNF(0,1,0,0, opcode=ABNF.OPCODE_TEXT) ws._exceptions.WebSocketProtocolException,
self.assertRaises(ws._exceptions.WebSocketProtocolException, a_bad_rsv_value.validate, skip_utf8_validation=False) a_invalid_ping.validate,
a_bad_opcode = ABNF(0,0,0,0, opcode=77) skip_utf8_validation=False,
self.assertRaises(ws._exceptions.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') a_bad_rsv_value = ABNF(0, 1, 0, 0, opcode=ABNF.OPCODE_TEXT)
self.assertRaises(ws._exceptions.WebSocketProtocolException, a_bad_close_frame.validate, skip_utf8_validation=False) self.assertRaises(
a_bad_close_frame_2 = ABNF(0,0,0,0, opcode=ABNF.OPCODE_CLOSE, data=b'\x01\x8a\xaa\xff\xdd') ws._exceptions.WebSocketProtocolException,
self.assertRaises(ws._exceptions.WebSocketProtocolException, a_bad_close_frame_2.validate, skip_utf8_validation=False) a_bad_rsv_value.validate,
a_bad_close_frame_3 = ABNF(0,0,0,0, opcode=ABNF.OPCODE_CLOSE, data=b'\x03\xe7') skip_utf8_validation=False,
self.assertRaises(ws._exceptions.WebSocketProtocolException, a_bad_close_frame_3.validate, skip_utf8_validation=True) )
a_bad_opcode = ABNF(0, 0, 0, 0, opcode=77)
self.assertRaises(
ws._exceptions.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(
ws._exceptions.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(
ws._exceptions.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(
ws._exceptions.WebSocketProtocolException,
a_bad_close_frame_3.validate,
skip_utf8_validation=True,
)
def testMask(self): def testMask(self):
abnf_none_data = ABNF(0,0,0,0, opcode=ABNF.OPCODE_PING, mask=1, data=None) abnf_none_data = ABNF(
0, 0, 0, 0, opcode=ABNF.OPCODE_PING, mask_value=1, data=None
)
bytes_val = b"aaaa" bytes_val = b"aaaa"
self.assertEqual(abnf_none_data._get_masked(bytes_val), bytes_val) self.assertEqual(abnf_none_data._get_masked(bytes_val), bytes_val)
abnf_str_data = ABNF(0,0,0,0, opcode=ABNF.OPCODE_PING, mask=1, data="a") abnf_str_data = ABNF(
self.assertEqual(abnf_str_data._get_masked(bytes_val), b'aaaa\x00') 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 testFormat(self): def testFormat(self):
abnf_bad_rsv_bits = ABNF(2,0,0,0, opcode=ABNF.OPCODE_TEXT) abnf_bad_rsv_bits = ABNF(2, 0, 0, 0, opcode=ABNF.OPCODE_TEXT)
self.assertRaises(ValueError, abnf_bad_rsv_bits.format) self.assertRaises(ValueError, abnf_bad_rsv_bits.format)
abnf_bad_opcode = ABNF(0,0,0,0, opcode=5) abnf_bad_opcode = ABNF(0, 0, 0, 0, opcode=5)
self.assertRaises(ValueError, abnf_bad_opcode.format) self.assertRaises(ValueError, abnf_bad_opcode.format)
abnf_length_10 = ABNF(0,0,0,0, opcode=ABNF.OPCODE_TEXT, data="abcdefghij") 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"\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(b"\x8a", abnf_length_10.format()[1].to_bytes(1, "big"))
self.assertEqual("fin=0 opcode=1 data=abcdefghij", abnf_length_10.__str__()) 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") abnf_length_20 = ABNF(
self.assertEqual(b'\x02', abnf_length_20.format()[0].to_bytes(1, 'big')) 0, 0, 0, 0, opcode=ABNF.OPCODE_BINARY, data="abcdefghijabcdefghij"
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=0, data=b'\x01\x8a\xcc') self.assertEqual(b"\x02", abnf_length_20.format()[0].to_bytes(1, "big"))
self.assertEqual(b'\x01\x03\x01\x8a\xcc', abnf_no_mask.format()) 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 testFrameBuffer(self): def testFrameBuffer(self):
fb = frame_buffer(0, True) fb = frame_buffer(0, True)
@ -81,7 +117,7 @@ class ABNFTest(unittest.TestCase):
fb.clear fb.clear
self.assertEqual(fb.header, None) self.assertEqual(fb.header, None)
self.assertEqual(fb.length, None) self.assertEqual(fb.length, None)
self.assertEqual(fb.mask, None) self.assertEqual(fb.mask_value, None)
self.assertEqual(fb.has_mask(), False) self.assertEqual(fb.has_mask(), False)

View file

@ -2,11 +2,12 @@
# #
import os import os
import os.path import os.path
import threading
import websocket as ws
import ssl import ssl
import threading
import unittest import unittest
import websocket as ws
""" """
test_app.py test_app.py
websocket - WebSocket client library for Python websocket - WebSocket client library for Python
@ -27,18 +28,16 @@ limitations under the License.
""" """
# Skip test to access the internet unless TEST_WITH_INTERNET == 1 # Skip test to access the internet unless TEST_WITH_INTERNET == 1
TEST_WITH_INTERNET = os.environ.get('TEST_WITH_INTERNET', '0') == '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 # 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') LOCAL_WS_SERVER_PORT = os.environ.get("LOCAL_WS_SERVER_PORT", "-1")
TEST_WITH_LOCAL_SERVER = LOCAL_WS_SERVER_PORT != '-1' TEST_WITH_LOCAL_SERVER = LOCAL_WS_SERVER_PORT != "-1"
TRACEABLE = True TRACEABLE = True
class WebSocketAppTest(unittest.TestCase): class WebSocketAppTest(unittest.TestCase):
class NotSetYet: class NotSetYet:
""" A marker class for signalling that a value hasn't been set yet. """A marker class for signalling that a value hasn't been set yet."""
"""
def setUp(self): def setUp(self):
ws.enableTrace(TRACEABLE) ws.enableTrace(TRACEABLE)
@ -54,14 +53,16 @@ class WebSocketAppTest(unittest.TestCase):
WebSocketAppTest.get_mask_key_id = WebSocketAppTest.NotSetYet() WebSocketAppTest.get_mask_key_id = WebSocketAppTest.NotSetYet()
WebSocketAppTest.on_error_data = WebSocketAppTest.NotSetYet() WebSocketAppTest.on_error_data = WebSocketAppTest.NotSetYet()
@unittest.skipUnless(TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled") @unittest.skipUnless(
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
)
def testKeepRunning(self): def testKeepRunning(self):
""" A WebSocketApp should keep running as long as its self.keep_running """A WebSocketApp should keep running as long as its self.keep_running
is not False (in the boolean context). is not False (in the boolean context).
""" """
def on_open(self, *args, **kwargs): def on_open(self, *args, **kwargs):
""" Set the keep_running flag for later inspection and immediately """Set the keep_running flag for later inspection and immediately
close the connection. close the connection.
""" """
self.send("hello!") self.send("hello!")
@ -73,23 +74,26 @@ class WebSocketAppTest(unittest.TestCase):
self.close() self.close()
def on_close(self, *args, **kwargs): def on_close(self, *args, **kwargs):
""" Set the keep_running flag for the test to use. """Set the keep_running flag for the test to use."""
"""
WebSocketAppTest.keep_running_close = self.keep_running WebSocketAppTest.keep_running_close = self.keep_running
app = ws.WebSocketApp('ws://127.0.0.1:' + LOCAL_WS_SERVER_PORT, on_open=on_open, on_close=on_close, on_message=on_message) 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() app.run_forever()
# @unittest.skipUnless(TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled") # @unittest.skipUnless(TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled")
@unittest.skipUnless(False, "Test disabled for now (requires rel)") @unittest.skipUnless(False, "Test disabled for now (requires rel)")
def testRunForeverDispatcher(self): def testRunForeverDispatcher(self):
""" A WebSocketApp should keep running as long as its self.keep_running """A WebSocketApp should keep running as long as its self.keep_running
is not False (in the boolean context). is not False (in the boolean context).
""" """
def on_open(self, *args, **kwargs): def on_open(self, *args, **kwargs):
""" Send a message, receive, and send one more """Send a message, receive, and send one more"""
"""
self.send("hello!") self.send("hello!")
self.recv() self.recv()
self.send("goodbye!") self.send("goodbye!")
@ -98,30 +102,38 @@ class WebSocketAppTest(unittest.TestCase):
print(message) print(message)
self.close() self.close()
app = ws.WebSocketApp('ws://127.0.0.1:' + LOCAL_WS_SERVER_PORT, on_open=on_open, on_message=on_message) 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="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") # app.run_forever(dispatcher=rel) # would work
# rel.dispatch()
@unittest.skipUnless(
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
)
def testRunForeverTeardownCleanExit(self): def testRunForeverTeardownCleanExit(self):
""" The WebSocketApp.run_forever() method should return `False` when the application ends gracefully. """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}")
app = ws.WebSocketApp('ws://127.0.0.1:' + LOCAL_WS_SERVER_PORT)
threading.Timer(interval=0.2, function=app.close).start() threading.Timer(interval=0.2, function=app.close).start()
teardown = app.run_forever() teardown = app.run_forever()
self.assertEqual(teardown, False) self.assertEqual(teardown, False)
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled") @unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
def testSockMaskKey(self): def testSockMaskKey(self):
""" A WebSocketApp should forward the received mask_key function down """A WebSocketApp should forward the received mask_key function down
to the actual socket. to the actual socket.
""" """
def my_mask_key_func(): def my_mask_key_func():
return "\x00\x00\x00\x00" return "\x00\x00\x00\x00"
app = ws.WebSocketApp('wss://api-pub.bitfinex.com/ws/1', get_mask_key=my_mask_key_func) app = ws.WebSocketApp(
"wss://api-pub.bitfinex.com/ws/1", get_mask_key=my_mask_key_func
)
# if numpy is installed, this assertion fail # if numpy is installed, this assertion fail
# Note: We can't use 'is' for comparing the functions directly, need to use 'id'. # Note: We can't use 'is' for comparing the functions directly, need to use 'id'.
@ -129,8 +141,7 @@ class WebSocketAppTest(unittest.TestCase):
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled") @unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
def testInvalidPingIntervalPingTimeout(self): def testInvalidPingIntervalPingTimeout(self):
""" Test exception handling if ping_interval < ping_timeout """Test exception handling if ping_interval < ping_timeout"""
"""
def on_ping(app, msg): def on_ping(app, msg):
print("Got a ping!") print("Got a ping!")
@ -140,13 +151,20 @@ class WebSocketAppTest(unittest.TestCase):
print("Got a pong! No need to respond") print("Got a pong! No need to respond")
app.close() app.close()
app = ws.WebSocketApp('wss://api-pub.bitfinex.com/ws/1', on_ping=on_ping, on_pong=on_pong) app = ws.WebSocketApp(
self.assertRaises(ws.WebSocketException, app.run_forever, ping_interval=1, ping_timeout=2, sslopt={"cert_reqs": ssl.CERT_NONE}) "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") @unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
def testPingInterval(self): def testPingInterval(self):
""" Test WebSocketApp proper ping functionality """Test WebSocketApp proper ping functionality"""
"""
def on_ping(app, msg): def on_ping(app, msg):
print("Got a ping!") print("Got a ping!")
@ -156,15 +174,18 @@ class WebSocketAppTest(unittest.TestCase):
print("Got a pong! No need to respond") print("Got a pong! No need to respond")
app.close() app.close()
app = ws.WebSocketApp('wss://api-pub.bitfinex.com/ws/1', on_ping=on_ping, on_pong=on_pong) app = ws.WebSocketApp(
app.run_forever(ping_interval=2, ping_timeout=1, sslopt={"cert_reqs": ssl.CERT_NONE}) "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") @unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
def testOpcodeClose(self): def testOpcodeClose(self):
""" Test WebSocketApp close opcode """Test WebSocketApp close opcode"""
"""
app = ws.WebSocketApp('wss://tsock.us1.twilio.com/v3/wsconnect') app = ws.WebSocketApp("wss://tsock.us1.twilio.com/v3/wsconnect")
app.run_forever(ping_interval=2, ping_timeout=1, ping_payload="Ping payload") 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 # This is commented out because the URL no longer responds in the expected way
@ -177,41 +198,59 @@ class WebSocketAppTest(unittest.TestCase):
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled") @unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
def testBadPingInterval(self): def testBadPingInterval(self):
""" A WebSocketApp handling of negative ping_interval """A WebSocketApp handling of negative ping_interval"""
""" app = ws.WebSocketApp("wss://api-pub.bitfinex.com/ws/1")
app = ws.WebSocketApp('wss://api-pub.bitfinex.com/ws/1') self.assertRaises(
self.assertRaises(ws.WebSocketException, app.run_forever, ping_interval=-5, sslopt={"cert_reqs": ssl.CERT_NONE}) ws.WebSocketException,
app.run_forever,
ping_interval=-5,
sslopt={"cert_reqs": ssl.CERT_NONE},
)
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled") @unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
def testBadPingTimeout(self): def testBadPingTimeout(self):
""" A WebSocketApp handling of negative ping_timeout """A WebSocketApp handling of negative ping_timeout"""
""" app = ws.WebSocketApp("wss://api-pub.bitfinex.com/ws/1")
app = ws.WebSocketApp('wss://api-pub.bitfinex.com/ws/1') self.assertRaises(
self.assertRaises(ws.WebSocketException, app.run_forever, ping_timeout=-3, sslopt={"cert_reqs": ssl.CERT_NONE}) ws.WebSocketException,
app.run_forever,
ping_timeout=-3,
sslopt={"cert_reqs": ssl.CERT_NONE},
)
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled") @unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
def testCloseStatusCode(self): def testCloseStatusCode(self):
""" Test extraction of close frame status code and close reason in WebSocketApp """Test extraction of close frame status code and close reason in WebSocketApp"""
"""
def on_close(wsapp, close_status_code, close_msg): def on_close(wsapp, close_status_code, close_msg):
print("on_close reached") print("on_close reached")
app = ws.WebSocketApp('wss://tsock.us1.twilio.com/v3/wsconnect', on_close=on_close) app = ws.WebSocketApp(
closeframe = ws.ABNF(opcode=ws.ABNF.OPCODE_CLOSE, data=b'\x03\xe8no-init-from-client') "wss://tsock.us1.twilio.com/v3/wsconnect", on_close=on_close
self.assertEqual([1000, 'no-init-from-client'], app._get_close_args(closeframe)) )
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'') closeframe = ws.ABNF(opcode=ws.ABNF.OPCODE_CLOSE, data=b"")
self.assertEqual([None, None], app._get_close_args(closeframe)) self.assertEqual([None, None], app._get_close_args(closeframe))
app2 = ws.WebSocketApp('wss://tsock.us1.twilio.com/v3/wsconnect') app2 = ws.WebSocketApp("wss://tsock.us1.twilio.com/v3/wsconnect")
closeframe = ws.ABNF(opcode=ws.ABNF.OPCODE_CLOSE, data=b'') closeframe = ws.ABNF(opcode=ws.ABNF.OPCODE_CLOSE, data=b"")
self.assertEqual([None, None], app2._get_close_args(closeframe)) self.assertEqual([None, None], app2._get_close_args(closeframe))
self.assertRaises(ws.WebSocketConnectionClosedException, app.send, data="test if connection is closed") 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") @unittest.skipUnless(
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
)
def testCallbackFunctionException(self): def testCallbackFunctionException(self):
""" Test callback function exception handling """ """Test callback function exception handling"""
exc = None exc = None
passed_app = None passed_app = None
@ -228,26 +267,33 @@ class WebSocketAppTest(unittest.TestCase):
def on_pong(app, msg): def on_pong(app, msg):
app.close() app.close()
app = ws.WebSocketApp('ws://127.0.0.1:' + LOCAL_WS_SERVER_PORT, on_open=on_open, on_error=on_error, on_pong=on_pong) 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) app.run_forever(ping_interval=2, ping_timeout=1)
self.assertEqual(passed_app, app) self.assertEqual(passed_app, app)
self.assertIsInstance(exc, RuntimeError) self.assertIsInstance(exc, RuntimeError)
self.assertEqual(str(exc), "Callback failed") self.assertEqual(str(exc), "Callback failed")
@unittest.skipUnless(TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled") @unittest.skipUnless(
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
)
def testCallbackMethodException(self): def testCallbackMethodException(self):
""" Test callback method exception handling """ """Test callback method exception handling"""
class Callbacks: class Callbacks:
def __init__(self): def __init__(self):
self.exc = None self.exc = None
self.passed_app = None self.passed_app = None
self.app = ws.WebSocketApp( self.app = ws.WebSocketApp(
'ws://127.0.0.1:' + LOCAL_WS_SERVER_PORT, f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}",
on_open=self.on_open, on_open=self.on_open,
on_error=self.on_error, on_error=self.on_error,
on_pong=self.on_pong on_pong=self.on_pong,
) )
self.app.run_forever(ping_interval=2, ping_timeout=1) self.app.run_forever(ping_interval=2, ping_timeout=1)
@ -267,9 +313,11 @@ class WebSocketAppTest(unittest.TestCase):
self.assertIsInstance(callbacks.exc, RuntimeError) self.assertIsInstance(callbacks.exc, RuntimeError)
self.assertEqual(str(callbacks.exc), "Callback failed") self.assertEqual(str(callbacks.exc), "Callback failed")
@unittest.skipUnless(TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled") @unittest.skipUnless(
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
)
def testReconnect(self): def testReconnect(self):
""" Test reconnect """ """Test reconnect"""
pong_count = 0 pong_count = 0
exc = None exc = None
@ -287,7 +335,9 @@ class WebSocketAppTest(unittest.TestCase):
# Got second pong after reconnect # Got second pong after reconnect
app.close() app.close()
app = ws.WebSocketApp('ws://127.0.0.1:' + LOCAL_WS_SERVER_PORT, on_pong=on_pong, on_error=on_error) 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) app.run_forever(ping_interval=2, ping_timeout=1, reconnect=3)
self.assertEqual(pong_count, 2) self.assertEqual(pong_count, 2)

View file

@ -1,4 +1,5 @@
import unittest import unittest
from websocket._cookiejar import SimpleCookieJar from websocket._cookiejar import SimpleCookieJar
""" """
@ -25,11 +26,15 @@ class CookieJarTest(unittest.TestCase):
def testAdd(self): def testAdd(self):
cookie_jar = SimpleCookieJar() cookie_jar = SimpleCookieJar()
cookie_jar.add("") cookie_jar.add("")
self.assertFalse(cookie_jar.jar, "Cookie with no domain should not be added to the jar") self.assertFalse(
cookie_jar.jar, "Cookie with no domain should not be added to the jar"
)
cookie_jar = SimpleCookieJar() cookie_jar = SimpleCookieJar()
cookie_jar.add("a=b") cookie_jar.add("a=b")
self.assertFalse(cookie_jar.jar, "Cookie with no domain should not be added to the jar") self.assertFalse(
cookie_jar.jar, "Cookie with no domain should not be added to the jar"
)
cookie_jar = SimpleCookieJar() cookie_jar = SimpleCookieJar()
cookie_jar.add("a=b; domain=.abc") cookie_jar.add("a=b; domain=.abc")
@ -65,7 +70,9 @@ class CookieJarTest(unittest.TestCase):
def testSet(self): def testSet(self):
cookie_jar = SimpleCookieJar() cookie_jar = SimpleCookieJar()
cookie_jar.set("a=b") cookie_jar.set("a=b")
self.assertFalse(cookie_jar.jar, "Cookie with no domain should not be added to the jar") self.assertFalse(
cookie_jar.jar, "Cookie with no domain should not be added to the jar"
)
cookie_jar = SimpleCookieJar() cookie_jar = SimpleCookieJar()
cookie_jar.set("a=b; domain=.abc") cookie_jar.set("a=b; domain=.abc")

View file

@ -2,12 +2,20 @@
# #
import os import os
import os.path import os.path
import websocket as ws
from websocket._http import proxy_info, read_headers, _start_proxied_socket, _tunnel, _get_addrinfo_list, connect
import unittest
import ssl
import websocket
import socket import socket
import ssl
import unittest
import websocket
import websocket as ws
from websocket._http import (
_get_addrinfo_list,
_start_proxied_socket,
_tunnel,
connect,
proxy_info,
read_headers,
)
""" """
test_http.py test_http.py
@ -29,16 +37,16 @@ limitations under the License.
""" """
try: try:
from python_socks._errors import ProxyError, ProxyTimeoutError, ProxyConnectionError from python_socks._errors import ProxyConnectionError, ProxyError, ProxyTimeoutError
except: except:
from websocket._http import ProxyError, ProxyTimeoutError, ProxyConnectionError from websocket._http import ProxyConnectionError, ProxyError, ProxyTimeoutError
# Skip test to access the internet unless TEST_WITH_INTERNET == 1 # Skip test to access the internet unless TEST_WITH_INTERNET == 1
TEST_WITH_INTERNET = os.environ.get('TEST_WITH_INTERNET', '0') == '1' TEST_WITH_INTERNET = os.environ.get("TEST_WITH_INTERNET", "0") == "1"
TEST_WITH_PROXY = os.environ.get('TEST_WITH_PROXY', '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 # 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') LOCAL_WS_SERVER_PORT = os.environ.get("LOCAL_WS_SERVER_PORT", "-1")
TEST_WITH_LOCAL_SERVER = LOCAL_WS_SERVER_PORT != '-1' TEST_WITH_LOCAL_SERVER = LOCAL_WS_SERVER_PORT != "-1"
class SockMock: class SockMock:
@ -70,7 +78,6 @@ class SockMock:
class HeaderSockMock(SockMock): class HeaderSockMock(SockMock):
def __init__(self, fname): def __init__(self, fname):
SockMock.__init__(self) SockMock.__init__(self)
path = os.path.join(os.path.dirname(__file__), fname) path = os.path.join(os.path.dirname(__file__), fname)
@ -78,8 +85,7 @@ class HeaderSockMock(SockMock):
self.add_packet(f.read()) self.add_packet(f.read())
class OptsList(): class OptsList:
def __init__(self): def __init__(self):
self.timeout = 1 self.timeout = 1
self.sockopt = [] self.sockopt = []
@ -87,17 +93,34 @@ class OptsList():
class HttpTest(unittest.TestCase): class HttpTest(unittest.TestCase):
def testReadHeader(self): def testReadHeader(self):
status, header, status_message = read_headers(HeaderSockMock("data/header01.txt")) status, header, status_message = read_headers(
HeaderSockMock("data/header01.txt")
)
self.assertEqual(status, 101) self.assertEqual(status, 101)
self.assertEqual(header["connection"], "Upgrade") self.assertEqual(header["connection"], "Upgrade")
# header02.txt is intentionally malformed # header02.txt is intentionally malformed
self.assertRaises(ws.WebSocketException, read_headers, HeaderSockMock("data/header02.txt")) self.assertRaises(
ws.WebSocketException, read_headers, HeaderSockMock("data/header02.txt")
)
def testTunnel(self): def testTunnel(self):
self.assertRaises(ws.WebSocketProxyException, _tunnel, HeaderSockMock("data/header01.txt"), "example.com", 80, ("username", "password")) self.assertRaises(
self.assertRaises(ws.WebSocketProxyException, _tunnel, HeaderSockMock("data/header02.txt"), "example.com", 80, ("username", "password")) ws.WebSocketProxyException,
_tunnel,
HeaderSockMock("data/header01.txt"),
"example.com",
80,
("username", "password"),
)
self.assertRaises(
ws.WebSocketProxyException,
_tunnel,
HeaderSockMock("data/header02.txt"),
"example.com",
80,
("username", "password"),
)
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled") @unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
def testConnect(self): def testConnect(self):
@ -105,34 +128,164 @@ class HttpTest(unittest.TestCase):
if ws._http.HAVE_PYTHON_SOCKS: if ws._http.HAVE_PYTHON_SOCKS:
# Need this check, otherwise case where python_socks is not installed triggers # 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 # 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(
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)) (ProxyTimeoutError, OSError),
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)) _start_proxied_socket,
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)) "wss://example.com",
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) 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(
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")) TypeError,
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) _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( self.assertEqual(
connect("wss://google.com", OptsList(), proxy_info(http_proxy_host="8.8.8.8", http_proxy_port=8080, proxy_type="http"), True), connect(
(True, ("google.com", 443, "/"))) "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 # 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) # 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_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(
@unittest.skipUnless(TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled") 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 testProxyConnect(self): def testProxyConnect(self):
ws = websocket.WebSocket() ws = websocket.WebSocket()
ws.connect("ws://127.0.0.1:" + LOCAL_WS_SERVER_PORT, http_proxy_host="127.0.0.1", http_proxy_port="8899", proxy_type="http") 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") ws.send("Hello, Server")
server_response = ws.recv() server_response = ws.recv()
self.assertEqual(server_response, "Hello, Server") 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(_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")), self.assertEqual(
(socket.getaddrinfo("127.0.0.1", 8899, 0, socket.SOCK_STREAM, socket.SOL_TCP), True, None)) _get_addrinfo_list(
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')) "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 # TODO: Test SOCKS4 and SOCK5 proxies with unit tests
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled") @unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
@ -151,7 +304,7 @@ class HttpTest(unittest.TestCase):
DHE-RSA-AES256-SHA256:ECDHE-ECDSA-AES128-SHA256:\ DHE-RSA-AES256-SHA256:ECDHE-ECDSA-AES128-SHA256:\
ECDHE-RSA-AES128-SHA256:DHE-RSA-AES128-SHA256:\ ECDHE-RSA-AES128-SHA256:DHE-RSA-AES128-SHA256:\
ECDHE-ECDSA-AES256-SHA:ECDHE-RSA-AES256-SHA", ECDHE-ECDSA-AES256-SHA:ECDHE-RSA-AES256-SHA",
"ecdh_curve": "prime256v1" "ecdh_curve": "prime256v1",
} }
ws_ssl1 = websocket.WebSocket(sslopt=ssloptions) ws_ssl1 = websocket.WebSocket(sslopt=ssloptions)
ws_ssl1.connect("wss://api.bitfinex.com/ws/2") ws_ssl1.connect("wss://api.bitfinex.com/ws/2")
@ -163,13 +316,55 @@ class HttpTest(unittest.TestCase):
ws_ssl2.close ws_ssl2.close
def testProxyInfo(self): def testProxyInfo(self):
self.assertEqual(proxy_info(http_proxy_host="127.0.0.1", http_proxy_port="8080", proxy_type="http").proxy_protocol, "http") self.assertEqual(
self.assertRaises(ProxyError, proxy_info, http_proxy_host="127.0.0.1", http_proxy_port="8080", proxy_type="badval") proxy_info(
self.assertEqual(proxy_info(http_proxy_host="example.com", http_proxy_port="8080", proxy_type="http").proxy_host, "example.com") http_proxy_host="127.0.0.1", http_proxy_port="8080", proxy_type="http"
self.assertEqual(proxy_info(http_proxy_host="127.0.0.1", http_proxy_port="8080", proxy_type="http").proxy_port, "8080") ).proxy_protocol,
self.assertEqual(proxy_info(http_proxy_host="127.0.0.1", http_proxy_port="8080", proxy_type="http").auth, None) "http",
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") 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__": if __name__ == "__main__":

View file

@ -2,7 +2,13 @@
# #
import os import os
import unittest import unittest
from websocket._url import get_proxy_info, parse_url, _is_address_in_network, _is_no_proxy_host
from websocket._url import (
_is_address_in_network,
_is_no_proxy_host,
get_proxy_info,
parse_url,
)
""" """
test_url.py test_url.py
@ -25,11 +31,10 @@ limitations under the License.
class UrlTest(unittest.TestCase): class UrlTest(unittest.TestCase):
def test_address_in_network(self): 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.0.0.1", "127.0.0.0/8"))
self.assertTrue(_is_address_in_network('127.1.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.assertFalse(_is_address_in_network("127.1.0.1", "127.0.0.0/24"))
def testParseUrl(self): def testParseUrl(self):
p = parse_url("ws://www.example.com/r") p = parse_url("ws://www.example.com/r")
@ -126,57 +131,71 @@ class IsNoProxyHostTest(unittest.TestCase):
del os.environ["no_proxy"] del os.environ["no_proxy"]
def testMatchAll(self): def testMatchAll(self):
self.assertTrue(_is_no_proxy_host("any.websocket.org", ['*'])) self.assertTrue(_is_no_proxy_host("any.websocket.org", ["*"]))
self.assertTrue(_is_no_proxy_host("192.168.0.1", ['*'])) self.assertTrue(_is_no_proxy_host("192.168.0.1", ["*"]))
self.assertTrue(_is_no_proxy_host("any.websocket.org", ['other.websocket.org', '*'])) self.assertTrue(
os.environ['no_proxy'] = '*' _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("any.websocket.org", None))
self.assertTrue(_is_no_proxy_host("192.168.0.1", None)) self.assertTrue(_is_no_proxy_host("192.168.0.1", None))
os.environ['no_proxy'] = 'other.websocket.org, *' os.environ["no_proxy"] = "other.websocket.org, *"
self.assertTrue(_is_no_proxy_host("any.websocket.org", None)) self.assertTrue(_is_no_proxy_host("any.websocket.org", None))
def testIpAddress(self): def testIpAddress(self):
self.assertTrue(_is_no_proxy_host("127.0.0.1", ['127.0.0.1'])) 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.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.assertTrue(
self.assertFalse(_is_no_proxy_host("127.0.0.2", ['other.websocket.org', '127.0.0.1'])) _is_no_proxy_host("127.0.0.1", ["other.websocket.org", "127.0.0.1"])
os.environ['no_proxy'] = '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.assertTrue(_is_no_proxy_host("127.0.0.1", None))
self.assertFalse(_is_no_proxy_host("127.0.0.2", None)) self.assertFalse(_is_no_proxy_host("127.0.0.2", None))
os.environ['no_proxy'] = 'other.websocket.org, 127.0.0.1' os.environ["no_proxy"] = "other.websocket.org, 127.0.0.1"
self.assertTrue(_is_no_proxy_host("127.0.0.1", None)) self.assertTrue(_is_no_proxy_host("127.0.0.1", None))
self.assertFalse(_is_no_proxy_host("127.0.0.2", None)) self.assertFalse(_is_no_proxy_host("127.0.0.2", None))
def testIpAddressInRange(self): def testIpAddressInRange(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.1", ["127.0.0.0/8"]))
self.assertTrue(_is_no_proxy_host("127.0.0.2", ['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.assertFalse(_is_no_proxy_host("127.1.0.1", ["127.0.0.0/24"]))
os.environ['no_proxy'] = '127.0.0.0/8' os.environ["no_proxy"] = "127.0.0.0/8"
self.assertTrue(_is_no_proxy_host("127.0.0.1", None)) 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("127.0.0.2", None))
os.environ['no_proxy'] = '127.0.0.0/24' os.environ["no_proxy"] = "127.0.0.0/24"
self.assertFalse(_is_no_proxy_host("127.1.0.1", None)) self.assertFalse(_is_no_proxy_host("127.1.0.1", None))
def testHostnameMatch(self): def testHostnameMatch(self):
self.assertTrue(_is_no_proxy_host("my.websocket.org", ['my.websocket.org'])) 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.assertTrue(
self.assertFalse(_is_no_proxy_host("my.websocket.org", ['other.websocket.org'])) _is_no_proxy_host(
os.environ['no_proxy'] = 'my.websocket.org' "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.assertTrue(_is_no_proxy_host("my.websocket.org", None))
self.assertFalse(_is_no_proxy_host("other.websocket.org", None)) self.assertFalse(_is_no_proxy_host("other.websocket.org", None))
os.environ['no_proxy'] = 'other.websocket.org, my.websocket.org' os.environ["no_proxy"] = "other.websocket.org, my.websocket.org"
self.assertTrue(_is_no_proxy_host("my.websocket.org", None)) self.assertTrue(_is_no_proxy_host("my.websocket.org", None))
def testHostnameMatchDomain(self): def testHostnameMatchDomain(self):
self.assertTrue(_is_no_proxy_host("any.websocket.org", ['.websocket.org'])) 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("my.other.websocket.org", [".websocket.org"]))
self.assertTrue(_is_no_proxy_host("any.websocket.org", ['my.websocket.org', '.websocket.org'])) self.assertTrue(
self.assertFalse(_is_no_proxy_host("any.websocket.com", ['.websocket.org'])) _is_no_proxy_host(
os.environ['no_proxy'] = '.websocket.org' "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("any.websocket.org", None))
self.assertTrue(_is_no_proxy_host("my.other.websocket.org", None)) self.assertTrue(_is_no_proxy_host("my.other.websocket.org", None))
self.assertFalse(_is_no_proxy_host("any.websocket.com", None)) self.assertFalse(_is_no_proxy_host("any.websocket.com", None))
os.environ['no_proxy'] = 'my.websocket.org, .websocket.org' os.environ["no_proxy"] = "my.websocket.org, .websocket.org"
self.assertTrue(_is_no_proxy_host("any.websocket.org", None)) self.assertTrue(_is_no_proxy_host("any.websocket.org", None))
@ -209,96 +228,205 @@ class ProxyInfoTest(unittest.TestCase):
del os.environ["no_proxy"] del os.environ["no_proxy"]
def testProxyFromArgs(self): def testProxyFromArgs(self):
self.assertEqual(get_proxy_info("echo.websocket.events", False, proxy_host="localhost"), ("localhost", 0, None))
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"), ("localhost", 0, 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_auth=("a", "b")),
("localhost", 0, ("a", "b")))
self.assertEqual( self.assertEqual(
get_proxy_info("echo.websocket.events", False, proxy_host="localhost", proxy_port=3128, proxy_auth=("a", "b")), get_proxy_info("echo.websocket.events", False, proxy_host="localhost"),
("localhost", 3128, ("a", "b"))) ("localhost", 0, None),
self.assertEqual(get_proxy_info("echo.websocket.events", True, proxy_host="localhost", proxy_auth=("a", "b")), )
("localhost", 0, ("a", "b")))
self.assertEqual( self.assertEqual(
get_proxy_info("echo.websocket.events", True, proxy_host="localhost", proxy_port=3128, proxy_auth=("a", "b")), get_proxy_info(
("localhost", 3128, ("a", "b"))) "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"),
("localhost", 0, 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", True, proxy_host="localhost", proxy_port=3128, self.assertEqual(
no_proxy=["example.com"], proxy_auth=("a", "b")), get_proxy_info(
("localhost", 3128, ("a", "b"))) "echo.websocket.events",
self.assertEqual(get_proxy_info("echo.websocket.events", True, proxy_host="localhost", proxy_port=3128, False,
no_proxy=["echo.websocket.events"], proxy_auth=("a", "b")), proxy_host="localhost",
(None, 0, None)) proxy_auth=("a", "b"),
),
("localhost", 0, ("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_auth=("a", "b"),
),
("localhost", 0, ("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),
)
def testProxyFromEnv(self): def testProxyFromEnv(self):
os.environ["http_proxy"] = "http://localhost/" os.environ["http_proxy"] = "http://localhost/"
self.assertEqual(get_proxy_info("echo.websocket.events", False), ("localhost", None, None)) self.assertEqual(
get_proxy_info("echo.websocket.events", False), ("localhost", None, None)
)
os.environ["http_proxy"] = "http://localhost:3128/" os.environ["http_proxy"] = "http://localhost:3128/"
self.assertEqual(get_proxy_info("echo.websocket.events", False), ("localhost", 3128, None)) self.assertEqual(
get_proxy_info("echo.websocket.events", False), ("localhost", 3128, None)
)
os.environ["http_proxy"] = "http://localhost/" os.environ["http_proxy"] = "http://localhost/"
os.environ["https_proxy"] = "http://localhost2/" os.environ["https_proxy"] = "http://localhost2/"
self.assertEqual(get_proxy_info("echo.websocket.events", False), ("localhost", None, None)) self.assertEqual(
get_proxy_info("echo.websocket.events", False), ("localhost", None, None)
)
os.environ["http_proxy"] = "http://localhost:3128/" os.environ["http_proxy"] = "http://localhost:3128/"
os.environ["https_proxy"] = "http://localhost2:3128/" os.environ["https_proxy"] = "http://localhost2:3128/"
self.assertEqual(get_proxy_info("echo.websocket.events", False), ("localhost", 3128, None)) self.assertEqual(
get_proxy_info("echo.websocket.events", False), ("localhost", 3128, None)
)
os.environ["http_proxy"] = "http://localhost/" os.environ["http_proxy"] = "http://localhost/"
os.environ["https_proxy"] = "http://localhost2/" 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", True), ("localhost2", None, None)
)
os.environ["http_proxy"] = "http://localhost:3128/" os.environ["http_proxy"] = "http://localhost:3128/"
os.environ["https_proxy"] = "http://localhost2:3128/" 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", True), ("localhost2", 3128, None)
)
os.environ["http_proxy"] = "" os.environ["http_proxy"] = ""
os.environ["https_proxy"] = "http://localhost2/" os.environ["https_proxy"] = "http://localhost2/"
self.assertEqual(get_proxy_info("echo.websocket.events", True), ("localhost2", None, None)) self.assertEqual(
self.assertEqual(get_proxy_info("echo.websocket.events", False), (None, 0, None)) 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["http_proxy"] = ""
os.environ["https_proxy"] = "http://localhost2:3128/" os.environ["https_proxy"] = "http://localhost2:3128/"
self.assertEqual(get_proxy_info("echo.websocket.events", True), ("localhost2", 3128, None)) self.assertEqual(
self.assertEqual(get_proxy_info("echo.websocket.events", False), (None, 0, None)) 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["http_proxy"] = "http://localhost/"
os.environ["https_proxy"] = "" os.environ["https_proxy"] = ""
self.assertEqual(get_proxy_info("echo.websocket.events", True), (None, 0, None)) self.assertEqual(get_proxy_info("echo.websocket.events", True), (None, 0, None))
self.assertEqual(get_proxy_info("echo.websocket.events", False), ("localhost", None, None)) self.assertEqual(
get_proxy_info("echo.websocket.events", False), ("localhost", None, None)
)
os.environ["http_proxy"] = "http://localhost:3128/" os.environ["http_proxy"] = "http://localhost:3128/"
os.environ["https_proxy"] = "" os.environ["https_proxy"] = ""
self.assertEqual(get_proxy_info("echo.websocket.events", True), (None, 0, None)) self.assertEqual(get_proxy_info("echo.websocket.events", True), (None, 0, None))
self.assertEqual(get_proxy_info("echo.websocket.events", False), ("localhost", 3128, None)) self.assertEqual(
get_proxy_info("echo.websocket.events", False), ("localhost", 3128, None)
)
os.environ["http_proxy"] = "http://a:b@localhost/" os.environ["http_proxy"] = "http://a:b@localhost/"
self.assertEqual(get_proxy_info("echo.websocket.events", False), ("localhost", None, ("a", "b"))) self.assertEqual(
get_proxy_info("echo.websocket.events", False),
("localhost", None, ("a", "b")),
)
os.environ["http_proxy"] = "http://a:b@localhost:3128/" os.environ["http_proxy"] = "http://a:b@localhost:3128/"
self.assertEqual(get_proxy_info("echo.websocket.events", False), ("localhost", 3128, ("a", "b"))) self.assertEqual(
get_proxy_info("echo.websocket.events", False),
("localhost", 3128, ("a", "b")),
)
os.environ["http_proxy"] = "http://a:b@localhost/" os.environ["http_proxy"] = "http://a:b@localhost/"
os.environ["https_proxy"] = "http://a:b@localhost2/" os.environ["https_proxy"] = "http://a:b@localhost2/"
self.assertEqual(get_proxy_info("echo.websocket.events", False), ("localhost", None, ("a", "b"))) self.assertEqual(
get_proxy_info("echo.websocket.events", False),
("localhost", None, ("a", "b")),
)
os.environ["http_proxy"] = "http://a:b@localhost:3128/" os.environ["http_proxy"] = "http://a:b@localhost:3128/"
os.environ["https_proxy"] = "http://a:b@localhost2:3128/" os.environ["https_proxy"] = "http://a:b@localhost2:3128/"
self.assertEqual(get_proxy_info("echo.websocket.events", False), ("localhost", 3128, ("a", "b"))) self.assertEqual(
get_proxy_info("echo.websocket.events", False),
("localhost", 3128, ("a", "b")),
)
os.environ["http_proxy"] = "http://a:b@localhost/" os.environ["http_proxy"] = "http://a:b@localhost/"
os.environ["https_proxy"] = "http://a:b@localhost2/" os.environ["https_proxy"] = "http://a:b@localhost2/"
self.assertEqual(get_proxy_info("echo.websocket.events", True), ("localhost2", None, ("a", "b"))) self.assertEqual(
get_proxy_info("echo.websocket.events", True),
("localhost2", None, ("a", "b")),
)
os.environ["http_proxy"] = "http://a:b@localhost:3128/" os.environ["http_proxy"] = "http://a:b@localhost:3128/"
os.environ["https_proxy"] = "http://a:b@localhost2:3128/" os.environ["https_proxy"] = "http://a:b@localhost2:3128/"
self.assertEqual(get_proxy_info("echo.websocket.events", True), ("localhost2", 3128, ("a", "b"))) 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[
os.environ["https_proxy"] = "http://john%40example.com:P%40SSWORD@localhost2:3128/" "http_proxy"
self.assertEqual(get_proxy_info("echo.websocket.events", True), ("localhost2", 3128, ("john@example.com", "P@SSWORD"))) ] = "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["http_proxy"] = "http://a:b@localhost/"
os.environ["https_proxy"] = "http://a:b@localhost2/" os.environ["https_proxy"] = "http://a:b@localhost2/"
os.environ["no_proxy"] = "example1.com,example2.com" os.environ["no_proxy"] = "example1.com,example2.com"
self.assertEqual(get_proxy_info("example.1.com", True), ("localhost2", None, ("a", "b"))) self.assertEqual(
get_proxy_info("example.1.com", True), ("localhost2", None, ("a", "b"))
)
os.environ["http_proxy"] = "http://a:b@localhost:3128/" os.environ["http_proxy"] = "http://a:b@localhost:3128/"
os.environ["https_proxy"] = "http://a:b@localhost2:3128/" os.environ["https_proxy"] = "http://a:b@localhost2:3128/"
os.environ["no_proxy"] = "example1.com,example2.com, echo.websocket.events" os.environ["no_proxy"] = "example1.com,example2.com, echo.websocket.events"

View file

@ -3,13 +3,14 @@
import os import os
import os.path import os.path
import socket import socket
import websocket as ws
import unittest import unittest
from websocket._handshake import _create_sec_websocket_key, \ from base64 import decodebytes as base64decode
_validate as _validate_header
import websocket as ws
from websocket._handshake import _create_sec_websocket_key
from websocket._handshake import _validate as _validate_header
from websocket._http import read_headers from websocket._http import read_headers
from websocket._utils import validate_utf8 from websocket._utils import validate_utf8
from base64 import decodebytes as base64decode
""" """
test_websocket.py test_websocket.py
@ -38,11 +39,12 @@ except ImportError:
class SSLError(Exception): class SSLError(Exception):
pass pass
# Skip test to access the internet unless TEST_WITH_INTERNET == 1 # Skip test to access the internet unless TEST_WITH_INTERNET == 1
TEST_WITH_INTERNET = os.environ.get('TEST_WITH_INTERNET', '0') == '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 # 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') LOCAL_WS_SERVER_PORT = os.environ.get("LOCAL_WS_SERVER_PORT", "-1")
TEST_WITH_LOCAL_SERVER = LOCAL_WS_SERVER_PORT != '-1' TEST_WITH_LOCAL_SERVER = LOCAL_WS_SERVER_PORT != "-1"
TRACEABLE = True TRACEABLE = True
@ -79,7 +81,6 @@ class SockMock:
class HeaderSockMock(SockMock): class HeaderSockMock(SockMock):
def __init__(self, fname): def __init__(self, fname):
SockMock.__init__(self) SockMock.__init__(self)
path = os.path.join(os.path.dirname(__file__), fname) path = os.path.join(os.path.dirname(__file__), fname)
@ -103,11 +104,10 @@ class WebSocketTest(unittest.TestCase):
def testWSKey(self): def testWSKey(self):
key = _create_sec_websocket_key() key = _create_sec_websocket_key()
self.assertTrue(key != 24) self.assertTrue(key != 24)
self.assertTrue(str("¥n") not in key) self.assertTrue("¥n" not in key)
def testNonce(self): def testNonce(self):
""" WebSocket key should be a random 16-byte nonce. """WebSocket key should be a random 16-byte nonce."""
"""
key = _create_sec_websocket_key() key = _create_sec_websocket_key()
nonce = base64decode(key.encode("utf-8")) nonce = base64decode(key.encode("utf-8"))
self.assertEqual(16, len(nonce)) self.assertEqual(16, len(nonce))
@ -117,7 +117,8 @@ class WebSocketTest(unittest.TestCase):
required_header = { required_header = {
"upgrade": "websocket", "upgrade": "websocket",
"connection": "upgrade", "connection": "upgrade",
"sec-websocket-accept": "Kxep+hNu9n51529fGidYu7a3wO0="} "sec-websocket-accept": "Kxep+hNu9n51529fGidYu7a3wO0=",
}
self.assertEqual(_validate_header(required_header, key, None), (True, None)) self.assertEqual(_validate_header(required_header, key, None), (True, None))
header = required_header.copy() header = required_header.copy()
@ -140,29 +141,39 @@ class WebSocketTest(unittest.TestCase):
header = required_header.copy() header = required_header.copy()
header["sec-websocket-protocol"] = "sub1" header["sec-websocket-protocol"] = "sub1"
self.assertEqual(_validate_header(header, key, ["sub1", "sub2"]), (True, "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 # 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)) self.assertEqual(_validate_header(header, key, ["sub2", "sub3"]), (False, None))
header = required_header.copy() header = required_header.copy()
header["sec-websocket-protocol"] = "sUb1" header["sec-websocket-protocol"] = "sUb1"
self.assertEqual(_validate_header(header, key, ["Sub1", "suB2"]), (True, "sub1")) self.assertEqual(
_validate_header(header, key, ["Sub1", "suB2"]), (True, "sub1")
)
header = required_header.copy() header = required_header.copy()
# This case will print out a logging error using the error() function, but that is expected # 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)) self.assertEqual(_validate_header(header, key, ["Sub1", "suB2"]), (False, None))
def testReadHeader(self): def testReadHeader(self):
status, header, status_message = read_headers(HeaderSockMock("data/header01.txt")) status, header, status_message = read_headers(
HeaderSockMock("data/header01.txt")
)
self.assertEqual(status, 101) self.assertEqual(status, 101)
self.assertEqual(header["connection"], "Upgrade") self.assertEqual(header["connection"], "Upgrade")
status, header, status_message = read_headers(HeaderSockMock("data/header03.txt")) status, header, status_message = read_headers(
HeaderSockMock("data/header03.txt")
)
self.assertEqual(status, 101) self.assertEqual(status, 101)
self.assertEqual(header["connection"], "Upgrade, Keep-Alive") self.assertEqual(header["connection"], "Upgrade, Keep-Alive")
HeaderSockMock("data/header02.txt") HeaderSockMock("data/header02.txt")
self.assertRaises(ws.WebSocketException, read_headers, HeaderSockMock("data/header02.txt")) self.assertRaises(
ws.WebSocketException, read_headers, HeaderSockMock("data/header02.txt")
)
def testSend(self): def testSend(self):
# TODO: add longer frame data # TODO: add longer frame data
@ -170,33 +181,38 @@ class WebSocketTest(unittest.TestCase):
sock.set_mask_key(create_mask_key) sock.set_mask_key(create_mask_key)
s = sock.sock = HeaderSockMock("data/header01.txt") s = sock.sock = HeaderSockMock("data/header01.txt")
sock.send("Hello") sock.send("Hello")
self.assertEqual(s.sent[0], b'\x81\x85abcd)\x07\x0f\x08\x0e') self.assertEqual(s.sent[0], b"\x81\x85abcd)\x07\x0f\x08\x0e")
sock.send("こんにちは") 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') 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) # 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(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) self.assertEqual(sock.send_binary(b"1111111111101"), 19)
def testRecv(self): def testRecv(self):
# TODO: add longer frame data # TODO: add longer frame data
sock = ws.WebSocket() sock = ws.WebSocket()
s = sock.sock = SockMock() s = sock.sock = SockMock()
something = b'\x81\x8fabcd\x82\xe3\xf0\x87\xe3\xf1\x80\xe5\xca\x81\xe2\xc5\x82\xe3\xcc' something = (
b"\x81\x8fabcd\x82\xe3\xf0\x87\xe3\xf1\x80\xe5\xca\x81\xe2\xc5\x82\xe3\xcc"
)
s.add_packet(something) s.add_packet(something)
data = sock.recv() data = sock.recv()
self.assertEqual(data, "こんにちは") self.assertEqual(data, "こんにちは")
s.add_packet(b'\x81\x85abcd)\x07\x0f\x08\x0e') s.add_packet(b"\x81\x85abcd)\x07\x0f\x08\x0e")
data = sock.recv() data = sock.recv()
self.assertEqual(data, "Hello") self.assertEqual(data, "Hello")
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled") @unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
def testIter(self): def testIter(self):
count = 2 count = 2
s = ws.create_connection('wss://api.bitfinex.com/ws/2') s = ws.create_connection("wss://api.bitfinex.com/ws/2")
s.send('{"event": "subscribe", "channel": "ticker"}') s.send('{"event": "subscribe", "channel": "ticker"}')
for _ in s: for _ in s:
count -= 1 count -= 1
@ -205,34 +221,34 @@ class WebSocketTest(unittest.TestCase):
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled") @unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
def testNext(self): def testNext(self):
sock = ws.create_connection('wss://api.bitfinex.com/ws/2') sock = ws.create_connection("wss://api.bitfinex.com/ws/2")
self.assertEqual(str, type(next(sock))) self.assertEqual(str, type(next(sock)))
def testInternalRecvStrict(self): def testInternalRecvStrict(self):
sock = ws.WebSocket() sock = ws.WebSocket()
s = sock.sock = SockMock() s = sock.sock = SockMock()
s.add_packet(b'foo') s.add_packet(b"foo")
s.add_packet(socket.timeout()) s.add_packet(socket.timeout())
s.add_packet(b'bar') s.add_packet(b"bar")
# s.add_packet(SSLError("The read operation timed out")) # s.add_packet(SSLError("The read operation timed out"))
s.add_packet(b'baz') s.add_packet(b"baz")
with self.assertRaises(ws.WebSocketTimeoutException): with self.assertRaises(ws.WebSocketTimeoutException):
sock.frame_buffer.recv_strict(9) sock.frame_buffer.recv_strict(9)
# with self.assertRaises(SSLError): # with self.assertRaises(SSLError):
# data = sock._recv_strict(9) # data = sock._recv_strict(9)
data = sock.frame_buffer.recv_strict(9) data = sock.frame_buffer.recv_strict(9)
self.assertEqual(data, b'foobarbaz') self.assertEqual(data, b"foobarbaz")
with self.assertRaises(ws.WebSocketConnectionClosedException): with self.assertRaises(ws.WebSocketConnectionClosedException):
sock.frame_buffer.recv_strict(1) sock.frame_buffer.recv_strict(1)
def testRecvTimeout(self): def testRecvTimeout(self):
sock = ws.WebSocket() sock = ws.WebSocket()
s = sock.sock = SockMock() s = sock.sock = SockMock()
s.add_packet(b'\x81') s.add_packet(b"\x81")
s.add_packet(socket.timeout()) s.add_packet(socket.timeout())
s.add_packet(b'\x8dabcd\x29\x07\x0f\x08\x0e') s.add_packet(b"\x8dabcd\x29\x07\x0f\x08\x0e")
s.add_packet(socket.timeout()) s.add_packet(socket.timeout())
s.add_packet(b'\x4e\x43\x33\x0e\x10\x0f\x00\x40') s.add_packet(b"\x4e\x43\x33\x0e\x10\x0f\x00\x40")
with self.assertRaises(ws.WebSocketTimeoutException): with self.assertRaises(ws.WebSocketTimeoutException):
sock.recv() sock.recv()
with self.assertRaises(ws.WebSocketTimeoutException): with self.assertRaises(ws.WebSocketTimeoutException):
@ -246,9 +262,9 @@ class WebSocketTest(unittest.TestCase):
sock = ws.WebSocket() sock = ws.WebSocket()
s = sock.sock = SockMock() s = sock.sock = SockMock()
# OPCODE=TEXT, FIN=0, MSG="Brevity is " # OPCODE=TEXT, FIN=0, MSG="Brevity is "
s.add_packet(b'\x01\x8babcd#\x10\x06\x12\x08\x16\x1aD\x08\x11C') s.add_packet(b"\x01\x8babcd#\x10\x06\x12\x08\x16\x1aD\x08\x11C")
# OPCODE=CONT, FIN=1, MSG="the soul of wit" # 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') s.add_packet(b"\x80\x8fabcd\x15\n\x06D\x12\r\x16\x08A\r\x05D\x16\x0b\x17")
data = sock.recv() data = sock.recv()
self.assertEqual(data, "Brevity is the soul of wit") self.assertEqual(data, "Brevity is the soul of wit")
with self.assertRaises(ws.WebSocketConnectionClosedException): with self.assertRaises(ws.WebSocketConnectionClosedException):
@ -258,21 +274,21 @@ class WebSocketTest(unittest.TestCase):
sock = ws.WebSocket(fire_cont_frame=True) sock = ws.WebSocket(fire_cont_frame=True)
s = sock.sock = SockMock() s = sock.sock = SockMock()
# OPCODE=TEXT, FIN=0, MSG="Brevity is " # OPCODE=TEXT, FIN=0, MSG="Brevity is "
s.add_packet(b'\x01\x8babcd#\x10\x06\x12\x08\x16\x1aD\x08\x11C') s.add_packet(b"\x01\x8babcd#\x10\x06\x12\x08\x16\x1aD\x08\x11C")
# OPCODE=CONT, FIN=0, MSG="Brevity is " # OPCODE=CONT, FIN=0, MSG="Brevity is "
s.add_packet(b'\x00\x8babcd#\x10\x06\x12\x08\x16\x1aD\x08\x11C') s.add_packet(b"\x00\x8babcd#\x10\x06\x12\x08\x16\x1aD\x08\x11C")
# OPCODE=CONT, FIN=1, MSG="the soul of wit" # 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') s.add_packet(b"\x80\x8fabcd\x15\n\x06D\x12\r\x16\x08A\r\x05D\x16\x0b\x17")
_, data = sock.recv_data() _, data = sock.recv_data()
self.assertEqual(data, b'Brevity is ') self.assertEqual(data, b"Brevity is ")
_, data = sock.recv_data() _, data = sock.recv_data()
self.assertEqual(data, b'Brevity is ') self.assertEqual(data, b"Brevity is ")
_, data = sock.recv_data() _, data = sock.recv_data()
self.assertEqual(data, b'the soul of wit') self.assertEqual(data, b"the soul of wit")
# OPCODE=CONT, FIN=0, MSG="Brevity is " # OPCODE=CONT, FIN=0, MSG="Brevity is "
s.add_packet(b'\x80\x8babcd#\x10\x06\x12\x08\x16\x1aD\x08\x11C') s.add_packet(b"\x80\x8babcd#\x10\x06\x12\x08\x16\x1aD\x08\x11C")
with self.assertRaises(ws.WebSocketException): with self.assertRaises(ws.WebSocketException):
sock.recv_data() sock.recv_data()
@ -288,7 +304,7 @@ class WebSocketTest(unittest.TestCase):
sock = ws.WebSocket() sock = ws.WebSocket()
s = sock.sock = SockMock() s = sock.sock = SockMock()
sock.connected = True sock.connected = True
s.add_packet(b'\x88\x80\x17\x98p\x84') s.add_packet(b"\x88\x80\x17\x98p\x84")
sock.recv() sock.recv()
self.assertEqual(sock.connected, False) self.assertEqual(sock.connected, False)
@ -296,22 +312,22 @@ class WebSocketTest(unittest.TestCase):
sock = ws.WebSocket() sock = ws.WebSocket()
s = sock.sock = SockMock() s = sock.sock = SockMock()
# OPCODE=CONT, FIN=1, MSG="the soul of wit" # 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') s.add_packet(b"\x80\x8fabcd\x15\n\x06D\x12\r\x16\x08A\r\x05D\x16\x0b\x17")
self.assertRaises(ws.WebSocketException, sock.recv) self.assertRaises(ws.WebSocketException, sock.recv)
def testRecvWithProlongedFragmentation(self): def testRecvWithProlongedFragmentation(self):
sock = ws.WebSocket() sock = ws.WebSocket()
s = sock.sock = SockMock() s = sock.sock = SockMock()
# OPCODE=TEXT, FIN=0, MSG="Once more unto the breach, " # 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') 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, " # OPCODE=CONT, FIN=0, MSG="dear friends, "
s.add_packet(b'\x00\x8eabcd\x05\x07\x02\x16A\x04\x11\r\x04\x0c\x07\x17MB') s.add_packet(b"\x00\x8eabcd\x05\x07\x02\x16A\x04\x11\r\x04\x0c\x07\x17MB")
# OPCODE=CONT, FIN=1, MSG="once more" # OPCODE=CONT, FIN=1, MSG="once more"
s.add_packet(b'\x80\x89abcd\x0e\x0c\x00\x01A\x0f\x0c\x16\x04') s.add_packet(b"\x80\x89abcd\x0e\x0c\x00\x01A\x0f\x0c\x16\x04")
data = sock.recv() data = sock.recv()
self.assertEqual( self.assertEqual(data, "Once more unto the breach, dear friends, once more")
data,
"Once more unto the breach, dear friends, once more")
with self.assertRaises(ws.WebSocketConnectionClosedException): with self.assertRaises(ws.WebSocketConnectionClosedException):
sock.recv() sock.recv()
@ -320,22 +336,24 @@ class WebSocketTest(unittest.TestCase):
sock.set_mask_key(create_mask_key) sock.set_mask_key(create_mask_key)
s = sock.sock = SockMock() s = sock.sock = SockMock()
# OPCODE=TEXT, FIN=0, MSG="Too much " # OPCODE=TEXT, FIN=0, MSG="Too much "
s.add_packet(b'\x01\x89abcd5\r\x0cD\x0c\x17\x00\x0cA') s.add_packet(b"\x01\x89abcd5\r\x0cD\x0c\x17\x00\x0cA")
# OPCODE=PING, FIN=1, MSG="Please PONG this" # OPCODE=PING, FIN=1, MSG="Please PONG this"
s.add_packet(b'\x89\x90abcd1\x0e\x06\x05\x12\x07C4.,$D\x15\n\n\x17') 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" # 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') s.add_packet(b"\x80\x8fabcd\x0e\x04C\x05A\x05\x0c\x0b\x05B\x17\x0c\x08\x0c\x04")
data = sock.recv() data = sock.recv()
self.assertEqual(data, "Too much of a good thing") self.assertEqual(data, "Too much of a good thing")
with self.assertRaises(ws.WebSocketConnectionClosedException): with self.assertRaises(ws.WebSocketConnectionClosedException):
sock.recv() sock.recv()
self.assertEqual( self.assertEqual(
s.sent[0], s.sent[0], b"\x8a\x90abcd1\x0e\x06\x05\x12\x07C4.,$D\x15\n\n\x17"
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") @unittest.skipUnless(
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
)
def testWebSocket(self): def testWebSocket(self):
s = ws.create_connection("ws://127.0.0.1:" + LOCAL_WS_SERVER_PORT) s = ws.create_connection(f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}")
self.assertNotEqual(s, None) self.assertNotEqual(s, None)
s.send("Hello, World") s.send("Hello, World")
result = s.next() result = s.next()
@ -348,9 +366,11 @@ class WebSocketTest(unittest.TestCase):
self.assertRaises(ValueError, s.send_close, -1, "") self.assertRaises(ValueError, s.send_close, -1, "")
s.close() s.close()
@unittest.skipUnless(TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled") @unittest.skipUnless(
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
)
def testPingPong(self): def testPingPong(self):
s = ws.create_connection("ws://127.0.0.1:" + LOCAL_WS_SERVER_PORT) s = ws.create_connection(f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}")
self.assertNotEqual(s, None) self.assertNotEqual(s, None)
s.ping("Hello") s.ping("Hello")
s.pong("Hi") s.pong("Hi")
@ -359,12 +379,15 @@ class WebSocketTest(unittest.TestCase):
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled") @unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
def testSupportRedirect(self): def testSupportRedirect(self):
s = ws.WebSocket() s = ws.WebSocket()
self.assertRaises(ws._exceptions.WebSocketBadStatusException, s.connect, "ws://google.com/") self.assertRaises(
ws._exceptions.WebSocketBadStatusException, s.connect, "ws://google.com/"
)
# Need to find a URL that has a redirect code leading to a websocket # Need to find a URL that has a redirect code leading to a websocket
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled") @unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
def testSecureWebSocket(self): def testSecureWebSocket(self):
import ssl import ssl
s = ws.create_connection("wss://api.bitfinex.com/ws/2") s = ws.create_connection("wss://api.bitfinex.com/ws/2")
self.assertNotEqual(s, None) self.assertNotEqual(s, None)
self.assertTrue(isinstance(s.sock, ssl.SSLSocket)) self.assertTrue(isinstance(s.sock, ssl.SSLSocket))
@ -375,10 +398,14 @@ class WebSocketTest(unittest.TestCase):
self.assertEqual(s.getsubprotocol(), None) self.assertEqual(s.getsubprotocol(), None)
s.abort() s.abort()
@unittest.skipUnless(TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled") @unittest.skipUnless(
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
)
def testWebSocketWithCustomHeader(self): def testWebSocketWithCustomHeader(self):
s = ws.create_connection("ws://127.0.0.1:" + LOCAL_WS_SERVER_PORT, s = ws.create_connection(
headers={"User-Agent": "PythonWebsocketClient"}) f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}",
headers={"User-Agent": "PythonWebsocketClient"},
)
self.assertNotEqual(s, None) self.assertNotEqual(s, None)
self.assertEqual(s.getsubprotocol(), None) self.assertEqual(s.getsubprotocol(), None)
s.send("Hello, World") s.send("Hello, World")
@ -387,9 +414,11 @@ class WebSocketTest(unittest.TestCase):
self.assertRaises(ValueError, s.close, -1, "") self.assertRaises(ValueError, s.close, -1, "")
s.close() s.close()
@unittest.skipUnless(TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled") @unittest.skipUnless(
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
)
def testAfterClose(self): def testAfterClose(self):
s = ws.create_connection("ws://127.0.0.1:" + LOCAL_WS_SERVER_PORT) s = ws.create_connection(f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}")
self.assertNotEqual(s, None) self.assertNotEqual(s, None)
s.close() s.close()
self.assertRaises(ws.WebSocketConnectionClosedException, s.send, "Hello") self.assertRaises(ws.WebSocketConnectionClosedException, s.send, "Hello")
@ -397,48 +426,69 @@ class WebSocketTest(unittest.TestCase):
class SockOptTest(unittest.TestCase): class SockOptTest(unittest.TestCase):
@unittest.skipUnless(TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled") @unittest.skipUnless(
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
)
def testSockOpt(self): def testSockOpt(self):
sockopt = ((socket.IPPROTO_TCP, socket.TCP_NODELAY, 1),) sockopt = ((socket.IPPROTO_TCP, socket.TCP_NODELAY, 1),)
s = ws.create_connection("ws://127.0.0.1:" + LOCAL_WS_SERVER_PORT, sockopt=sockopt) s = ws.create_connection(
self.assertNotEqual(s.sock.getsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY), 0) 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() s.close()
class UtilsTest(unittest.TestCase): class UtilsTest(unittest.TestCase):
def testUtf8Validator(self): def testUtf8Validator(self):
state = validate_utf8(b'\xf0\x90\x80\x80') state = validate_utf8(b"\xf0\x90\x80\x80")
self.assertEqual(state, True) self.assertEqual(state, True)
state = validate_utf8(b'\xce\xba\xe1\xbd\xb9\xcf\x83\xce\xbc\xce\xb5\xed\xa0\x80edited') state = validate_utf8(
b"\xce\xba\xe1\xbd\xb9\xcf\x83\xce\xbc\xce\xb5\xed\xa0\x80edited"
)
self.assertEqual(state, False) self.assertEqual(state, False)
state = validate_utf8(b'') state = validate_utf8(b"")
self.assertEqual(state, True) self.assertEqual(state, True)
class HandshakeTest(unittest.TestCase): class HandshakeTest(unittest.TestCase):
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled") @unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
def test_http_SSL(self): def test_http_SSL(self):
websock1 = ws.WebSocket(sslopt={"cert_chain": ssl.get_default_verify_paths().capath}, enable_multithread=False) websock1 = ws.WebSocket(
self.assertRaises(ValueError, sslopt={"cert_chain": ssl.get_default_verify_paths().capath},
websock1.connect, "wss://api.bitfinex.com/ws/2") enable_multithread=False,
)
self.assertRaises(ValueError, websock1.connect, "wss://api.bitfinex.com/ws/2")
websock2 = ws.WebSocket(sslopt={"certfile": "myNonexistentCertFile"}) websock2 = ws.WebSocket(sslopt={"certfile": "myNonexistentCertFile"})
self.assertRaises(FileNotFoundError, self.assertRaises(
websock2.connect, "wss://api.bitfinex.com/ws/2") FileNotFoundError, websock2.connect, "wss://api.bitfinex.com/ws/2"
)
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled") @unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
def testManualHeaders(self): def testManualHeaders(self):
websock3 = ws.WebSocket(sslopt={"ca_certs": ssl.get_default_verify_paths().cafile, websock3 = ws.WebSocket(
"ca_cert_path": ssl.get_default_verify_paths().capath}) sslopt={
self.assertRaises(ws._exceptions.WebSocketBadStatusException, "ca_certs": ssl.get_default_verify_paths().cafile,
websock3.connect, "wss://api.bitfinex.com/ws/2", cookie="chocolate", "ca_cert_path": ssl.get_default_verify_paths().capath,
origin="testing_websockets.com", }
host="echo.websocket.events/websocket-client-test", )
subprotocols=["testproto"], self.assertRaises(
connection="Upgrade", ws._exceptions.WebSocketBadStatusException,
header={"CustomHeader1":"123", websock3.connect,
"Cookie":"TestValue", "wss://api.bitfinex.com/ws/2",
"Sec-WebSocket-Key":"k9kFAUWNAMmf5OEMfTlOEA==", cookie="chocolate",
"Sec-WebSocket-Protocol":"newprotocol"}) 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 testIPv6(self): def testIPv6(self):
websock2 = ws.WebSocket() websock2 = ws.WebSocket()
@ -447,7 +497,9 @@ class HandshakeTest(unittest.TestCase):
def testBadURLs(self): def testBadURLs(self):
websock3 = ws.WebSocket() websock3 = ws.WebSocket()
self.assertRaises(ValueError, websock3.connect, "ws//example.com") self.assertRaises(ValueError, websock3.connect, "ws//example.com")
self.assertRaises(ws.WebSocketAddressException, websock3.connect, "ws://example") self.assertRaises(
ws.WebSocketAddressException, websock3.connect, "ws://example"
)
self.assertRaises(ValueError, websock3.connect, "example.com") self.assertRaises(ValueError, websock3.connect, "example.com")

View file

@ -47,7 +47,7 @@ tzdata==2023.3
tzlocal==5.0.1 tzlocal==5.0.1
urllib3<2 urllib3<2
webencodings==0.5.1 webencodings==0.5.1
websocket-client==1.6.2 websocket-client==1.7.0
xmltodict==0.13.0 xmltodict==0.13.0
zipp==3.16.2 zipp==3.16.2