Remove websockets dependency; implement RFC 6455 WebSocket client using stdlib
Replaces the third-party `websockets` package with a self-contained implementation built on asyncio.open_connection, ssl, hashlib, base64, struct, and os — all Python stdlib modules available on TrueNAS OS. The script now runs directly on TrueNAS without any pip install. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -12,8 +12,7 @@ SAFE BY DEFAULT
|
||||
• Always run with --dry-run first to preview what will happen.
|
||||
|
||||
REQUIREMENTS
|
||||
pip install websockets
|
||||
Python 3.9+
|
||||
Python 3.9+ (stdlib only – no external packages needed)
|
||||
|
||||
QUICK START
|
||||
# 1. Inspect your debug archive to confirm it contains the data you need:
|
||||
@@ -50,26 +49,20 @@ from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import base64
|
||||
import contextlib
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import ssl
|
||||
import struct
|
||||
import sys
|
||||
import tarfile
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional
|
||||
|
||||
# ── Optional dependency check ─────────────────────────────────────────────────
|
||||
try:
|
||||
import websockets
|
||||
from websockets.exceptions import WebSocketException
|
||||
except ImportError:
|
||||
sys.exit(
|
||||
"ERROR: The 'websockets' package is required.\n"
|
||||
"Install it with: pip install websockets"
|
||||
)
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
# Logging
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
@@ -547,6 +540,144 @@ def _smb_config_payload(config: dict) -> dict:
|
||||
return {k: v for k, v in config.items() if k not in exclude}
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
# Minimal WebSocket client (stdlib only, RFC 6455)
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
def _ws_mask(data: bytes, mask: bytes) -> bytes:
|
||||
"""XOR *data* with a 4-byte repeating mask key."""
|
||||
out = bytearray(data)
|
||||
for i in range(len(out)):
|
||||
out[i] ^= mask[i & 3]
|
||||
return bytes(out)
|
||||
|
||||
|
||||
def _ws_encode_frame(payload: bytes, opcode: int = 0x1) -> bytes:
|
||||
"""Encode a masked client→server WebSocket frame."""
|
||||
mask = os.urandom(4)
|
||||
length = len(payload)
|
||||
header = bytearray([0x80 | opcode]) # FIN=1
|
||||
if length < 126:
|
||||
header.append(0x80 | length)
|
||||
elif length < 65536:
|
||||
header.append(0x80 | 126)
|
||||
header += struct.pack("!H", length)
|
||||
else:
|
||||
header.append(0x80 | 127)
|
||||
header += struct.pack("!Q", length)
|
||||
return bytes(header) + mask + _ws_mask(payload, mask)
|
||||
|
||||
|
||||
async def _ws_recv_message(reader: asyncio.StreamReader) -> str:
|
||||
"""
|
||||
Read one complete WebSocket message, reassembling continuation frames.
|
||||
Skips ping/pong control frames. Raises OSError on close frame.
|
||||
"""
|
||||
fragments: list[bytes] = []
|
||||
while True:
|
||||
hdr = await reader.readexactly(2)
|
||||
fin = bool(hdr[0] & 0x80)
|
||||
opcode = hdr[0] & 0x0F
|
||||
masked = bool(hdr[1] & 0x80)
|
||||
length = hdr[1] & 0x7F
|
||||
|
||||
if length == 126:
|
||||
length = struct.unpack("!H", await reader.readexactly(2))[0]
|
||||
elif length == 127:
|
||||
length = struct.unpack("!Q", await reader.readexactly(8))[0]
|
||||
|
||||
mask_key = await reader.readexactly(4) if masked else None
|
||||
payload = await reader.readexactly(length) if length else b""
|
||||
if mask_key:
|
||||
payload = _ws_mask(payload, mask_key)
|
||||
|
||||
if opcode == 0x8: # Close frame
|
||||
raise OSError("WebSocket: server sent close frame")
|
||||
if opcode in (0x9, 0xA): # Ping / Pong — ignore
|
||||
continue
|
||||
|
||||
fragments.append(payload)
|
||||
if fin:
|
||||
return b"".join(fragments).decode("utf-8")
|
||||
|
||||
|
||||
class _WebSocket:
|
||||
"""asyncio StreamReader/Writer wrapped to match the send/recv/close API."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
reader: asyncio.StreamReader,
|
||||
writer: asyncio.StreamWriter,
|
||||
) -> None:
|
||||
self._reader = reader
|
||||
self._writer = writer
|
||||
|
||||
async def send(self, data: str) -> None:
|
||||
self._writer.write(_ws_encode_frame(data.encode("utf-8"), opcode=0x1))
|
||||
await self._writer.drain()
|
||||
|
||||
async def recv(self) -> str:
|
||||
return await _ws_recv_message(self._reader)
|
||||
|
||||
async def close(self) -> None:
|
||||
with contextlib.suppress(Exception):
|
||||
self._writer.write(_ws_encode_frame(b"", opcode=0x8))
|
||||
await self._writer.drain()
|
||||
self._writer.close()
|
||||
with contextlib.suppress(Exception):
|
||||
await self._writer.wait_closed()
|
||||
|
||||
|
||||
async def _ws_connect(
|
||||
host: str,
|
||||
port: int,
|
||||
path: str,
|
||||
ssl_ctx: ssl.SSLContext,
|
||||
) -> _WebSocket:
|
||||
"""
|
||||
Open a TLS connection, perform the HTTP→WebSocket upgrade handshake,
|
||||
and return a connected _WebSocket.
|
||||
"""
|
||||
reader, writer = await asyncio.open_connection(host, port, ssl=ssl_ctx)
|
||||
|
||||
key = base64.b64encode(os.urandom(16)).decode()
|
||||
writer.write((
|
||||
f"GET {path} HTTP/1.1\r\n"
|
||||
f"Host: {host}:{port}\r\n"
|
||||
f"Upgrade: websocket\r\n"
|
||||
f"Connection: Upgrade\r\n"
|
||||
f"Sec-WebSocket-Key: {key}\r\n"
|
||||
f"Sec-WebSocket-Version: 13\r\n"
|
||||
f"\r\n"
|
||||
).encode())
|
||||
await writer.drain()
|
||||
|
||||
# Read headers line-by-line to avoid consuming WebSocket frame bytes
|
||||
response_lines: list[bytes] = []
|
||||
while True:
|
||||
line = await asyncio.wait_for(reader.readline(), timeout=20)
|
||||
if not line:
|
||||
raise OSError("Connection closed during WebSocket handshake")
|
||||
response_lines.append(line)
|
||||
if line in (b"\r\n", b"\n"):
|
||||
break
|
||||
|
||||
status = response_lines[0].decode("latin-1").strip()
|
||||
if " 101 " not in status:
|
||||
raise OSError(f"WebSocket upgrade failed: {status}")
|
||||
|
||||
expected = base64.b64encode(
|
||||
hashlib.sha1(
|
||||
(key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").encode()
|
||||
).digest()
|
||||
).decode().lower()
|
||||
headers_text = b"".join(response_lines).decode("latin-1").lower()
|
||||
if expected not in headers_text:
|
||||
raise OSError("WebSocket upgrade: Sec-WebSocket-Accept mismatch")
|
||||
|
||||
return _WebSocket(reader, writer)
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
# TrueNAS JSON-RPC 2.0 WebSocket client
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
@@ -594,14 +725,13 @@ class TrueNASClient:
|
||||
|
||||
log.info("Connecting to %s …", self._url)
|
||||
try:
|
||||
self._ws = await websockets.connect(
|
||||
self._url,
|
||||
ssl=ctx,
|
||||
ping_interval=20,
|
||||
ping_timeout=30,
|
||||
open_timeout=20,
|
||||
self._ws = await _ws_connect(
|
||||
host=self._host,
|
||||
port=self._port,
|
||||
path="/api/current",
|
||||
ssl_ctx=ctx,
|
||||
)
|
||||
except (WebSocketException, OSError) as exc:
|
||||
except (OSError, asyncio.TimeoutError) as exc:
|
||||
log.error("Connection failed: %s", exc)
|
||||
raise
|
||||
|
||||
|
||||
Reference in New Issue
Block a user