diff --git a/truenas_migrate.py b/truenas_migrate.py index 6b0e910..524a36f 100644 --- a/truenas_migrate.py +++ b/truenas_migrate.py @@ -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