#!/usr/bin/env python3 """Dependency-free maze32 WebSocket protocol and six-client verifier.""" from __future__ import annotations import argparse import base64 import hashlib import json import os import socket import struct import time from typing import Any class WebSocket: def __init__(self, host: str, port: int = 80, timeout: float = 8.0): self.host = host self.sequence = 0 self.socket = socket.create_connection((host, port), timeout=timeout) self.socket.settimeout(timeout) key = base64.b64encode(os.urandom(16)).decode("ascii") request = ( f"GET /ws HTTP/1.1\r\nHost: {host}:{port}\r\nUpgrade: websocket\r\n" f"Connection: Upgrade\r\nSec-WebSocket-Key: {key}\r\n" "Sec-WebSocket-Version: 13\r\n\r\n" ) self.socket.sendall(request.encode("ascii")) response = self._receive_until(b"\r\n\r\n") if not response.startswith(b"HTTP/1.1 101"): self.socket.close() raise RuntimeError(response.decode("utf-8", "replace").strip()) expected = base64.b64encode( hashlib.sha1((key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").encode()).digest() ) if expected not in response: self.socket.close() raise RuntimeError("WebSocket accept key did not match") def _receive_until(self, marker: bytes) -> bytes: data = bytearray() while marker not in data: chunk = self.socket.recv(1024) if not chunk: raise ConnectionError("connection closed") data.extend(chunk) return bytes(data) def send(self, message_type: str, payload: dict[str, Any], request_id: str) -> None: self.sequence += 1 message = json.dumps( {"v": 1, "type": message_type, "seq": self.sequence, "requestId": request_id, "payload": payload}, ensure_ascii=False, separators=(",", ":"), ).encode("utf-8") self.send_raw(message) def send_raw(self, message: bytes) -> None: mask = os.urandom(4) length = len(message) if length < 126: header = bytes((0x81, 0x80 | length)) else: header = bytes((0x81, 0x80 | 126)) + struct.pack("!H", length) masked = bytes(value ^ mask[index % 4] for index, value in enumerate(message)) self.socket.sendall(header + mask + masked) def receive(self) -> dict[str, Any]: first, second = self._read_exact(2) opcode = first & 0x0F length = second & 0x7F if length == 126: length = struct.unpack("!H", self._read_exact(2))[0] elif length == 127: length = struct.unpack("!Q", self._read_exact(8))[0] if second & 0x80: mask = self._read_exact(4) else: mask = None payload = self._read_exact(length) if mask: payload = bytes(value ^ mask[index % 4] for index, value in enumerate(payload)) if opcode == 0x8: raise ConnectionError("WebSocket closed") if opcode != 0x1: return self.receive() return json.loads(payload.decode("utf-8")) def receive_type(self, message_type: str, attempts: int = 20) -> dict[str, Any]: for _ in range(attempts): message = self.receive() if message.get("type") == message_type: return message raise RuntimeError(f"did not receive {message_type}") def _read_exact(self, length: int) -> bytes: data = bytearray() while len(data) < length: chunk = self.socket.recv(length - len(data)) if not chunk: raise ConnectionError("connection closed") data.extend(chunk) return bytes(data) def close(self) -> None: try: self.socket.sendall(b"\x88\x80" + b"\0\0\0\0") except OSError: pass self.socket.close() def join(client: WebSocket, role: str, name: str | None = None) -> dict[str, Any]: payload = {"role": role} if name is not None: payload["name"] = name client.send("join", payload, f"join-{role}") return client.receive_type("session")["payload"] def verify(host: str) -> None: clients: list[WebSocket] = [] try: player_1 = WebSocket(host) clients.append(player_1) player_1.send_raw(b"{") error = player_1.receive_type("error") assert error["payload"]["code"] == "malformed" session_1 = join(player_1, "player", "Тестовый Пилот") player_1.send_raw( b'{"v":1,"type":"start","seq":1,"requestId":"stale","payload":{}}' ) error = player_1.receive_type("error") assert error["payload"]["code"] == "stale_sequence" print("malformed and stale message rejection: PASS") duplicate = WebSocket(host) clients.append(duplicate) duplicate.send("join", {"role": "player", "name": "Тестовый Пилот"}, "duplicate") error = duplicate.receive_type("error") assert error["payload"]["code"] == "duplicate_name" duplicate.send("join", {"role": "player", "name": "Запасной Пилот"}, "player-2") session_2 = duplicate.receive_type("session")["payload"] assert session_1["slot"] == 1 and session_2["slot"] == 2 observers = [] for _ in range(4): observer = WebSocket(host) clients.append(observer) join(observer, "observer") observers.append(observer) lobby = observers[-1].receive_type("lobby")["payload"] assert lobby["ready"] and lobby["observers"] == 4 print("six-client capacity: PASS (2 players + 4 observers)") print("duplicate player name: PASS") try: excess = WebSocket(host, timeout=1.0) except (OSError, RuntimeError, ConnectionError): print("seventh connection rejection: PASS") else: try: excess.send("join", {"role": "observer"}, "excess") excess.receive() except (OSError, ConnectionError): print("seventh connection rejection: PASS") else: raise AssertionError("seventh WebSocket unexpectedly remained connected") finally: excess.close() time.sleep(0.5) duplicate.send("start", {}, "not-owner") error = duplicate.receive_type("error") assert error["payload"]["code"] == "not_owner" print("Player 2 start authorization: PASS") token_2 = session_2["token"] duplicate.close() clients.remove(duplicate) time.sleep(0.15) restored = WebSocket(host) clients.append(restored) restored.send("restore", {"token": token_2}, "restore") restored_session = restored.receive_type("session")["payload"] assert restored_session["slot"] == 2 and restored_session["token"] == token_2 print("session refresh restoration: PASS (Player 2 retained slot 2)") player_1.send("start", {}, "owner-start") countdown_1 = player_1.receive_type("countdown")["payload"] countdown_2 = restored.receive_type("countdown")["payload"] assert countdown_1 == countdown_2 assert countdown_1["steps"] == [3, 2, 1] active_1 = player_1.receive_type("active")["payload"] active_2 = restored.receive_type("active")["payload"] assert active_1["activatedAtMs"] == countdown_1["activationAtMs"] assert active_1 == active_2 print( "authoritative countdown: PASS " f"(activationAtMs={countdown_1['activationAtMs']})" ) observer = observers.pop() observer.close() clients.remove(observer) time.sleep(0.5) oversized = WebSocket(host) oversized.send_raw(b"x" * 513) try: error = oversized.receive_type("error") assert error["payload"]["code"] == "too_large" except (OSError, ConnectionError): pass print("oversized message rejection: PASS") oversized.close() finally: for client in clients: client.close() def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("host", help="maze32 IP address or hostname") parser.add_argument("--verify", action="store_true", help="run the Milestone 004 scenario") arguments = parser.parse_args() if arguments.verify: verify(arguments.host) else: client = WebSocket(arguments.host) client.send("suggest_name", {}, "manual") print(json.dumps(client.receive(), ensure_ascii=False, indent=2)) client.close() if __name__ == "__main__": main()