#!/usr/bin/env python3 # -*- coding: utf-8 -*- """最小 WebSocket 探针（标准库手写握手 + 帧收发），配合 fake_backend.py 用。 用途： - 证明 `/ws/chat` 能握手（打印 101 与响应头），并且握手后能收到 `history` 帧； - 往连接上发任意上行帧（`--send`），再去 out/upstream.jsonl 里对账——这是「客户端到底发没发那一帧」的独立第二现场； - 监听 N 秒，把服务端下发的每一帧原样打出来。 python ws_probe.py --port 18080 --listen 3 python ws_probe.py --port 18080 --send '{"type":"interrupt","data":{"sessionId":"web","scope":"current"}}' --listen 2 """ import argparse import base64 import hashlib import json import os import socket import struct import sys import time for _stream in (sys.stdout, sys.stderr): try: _stream.reconfigure(encoding="utf-8", errors="replace") except Exception: pass WS_GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11" def encode_text_frame(text: str, mask: bytes = b"\x37\xfa\x21\x3d") -> bytes: payload = text.encode("utf-8") header = bytearray([0x80 | 0x1]) length = len(payload) 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) masked = bytes(byte ^ mask[i % 4] for i, byte in enumerate(payload)) return bytes(header) + mask + masked class RawReader: """裸 socket 上的定长读。 刻意**不用** `socket.makefile()`：给一个 BufferedReader 设过超时之后它会永久变成 「cannot read from timed out object」，于是「监听 2 秒」这种用法第二帧就读不到了。 """ def __init__(self, sock: socket.socket): self.sock = sock self.buffer = b"" def read_exact(self, count: int): while len(self.buffer) < count: chunk = self.sock.recv(65536) if not chunk: return None self.buffer += chunk head, self.buffer = self.buffer[:count], self.buffer[count:] return head def read_line(self): while b"\r\n" not in self.buffer: chunk = self.sock.recv(65536) if not chunk: return None self.buffer += chunk line, self.buffer = self.buffer.split(b"\r\n", 1) return line class Probe: def __init__(self, host: str, port: int, path: str): self.sock = socket.create_connection((host, port), timeout=10) key = base64.b64encode(os.urandom(16)).decode("ascii") request = ( "GET {} HTTP/1.1\r\n" "Host: {host}:{port}\r\n" "Upgrade: websocket\r\n" "Connection: Upgrade\r\n" "Sec-WebSocket-Key: {key}\r\n" "Sec-WebSocket-Version: 13\r\n" "\r\n" ).format(path, host=host, port=port, key=key) self.sock.sendall(request.encode("ascii")) self.reader = RawReader(self.sock) status = self.reader.read_line().decode("latin-1").strip() self.sock.settimeout(None) print("<<< " + status) headers = {} while True: line = self.reader.read_line().decode("latin-1") if line == "": break name, _, value = line.partition(":") headers[name.strip()] = value.strip() print("<<< " + line.strip()) expected = base64.b64encode(hashlib.sha1((key + WS_GUID).encode("ascii")).digest()).decode("ascii") code = status.split()[1] if len(status.split()) > 1 else "" if code != "101": raise SystemExit("handshake failed: " + status) if headers.get("Sec-WebSocket-Accept") != expected: raise SystemExit("bad Sec-WebSocket-Accept: " + str(headers.get("Sec-WebSocket-Accept"))) print("<<< handshake OK (101, Sec-WebSocket-Accept matches)") self.closed = False def send_json(self, payload: str) -> None: self.sock.sendall(encode_text_frame(payload)) print(">>> " + payload) def send_ping(self, payload: bytes = b"probe-ping") -> None: header = bytearray([0x80 | 0x9]) length = len(payload) if length < 126: header.append(0x80 | length) else: header.append(0x80 | 126) header += struct.pack(">H", length) mask = b"\x11\x22\x33\x44" masked = bytes(byte ^ mask[i % 4] for i, byte in enumerate(payload)) self.sock.sendall(bytes(header) + mask + masked) print(">>> ping (payload={})".format(payload.decode("latin-1"))) def _read_exact(self, count: int): return self.reader.read_exact(count) def listen(self, seconds: float, label: str = "downstream") -> int: deadline = time.time() + seconds frames = 0 self.sock.settimeout(0.25) try: while time.time() < deadline: try: head = self._read_exact(2) except socket.timeout: continue if not head or len(head) < 2: print("<<< connection closed by peer") self.closed = True break opcode = head[0] & 0x0F length = head[1] & 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] payload = self._read_exact(length) if length else b"" if payload is None: print("<<< connection closed by peer") self.closed = True break if opcode == 0x1: frames += 1 print("<<< frame {} [{}] {}".format(frames, label, payload.decode("utf-8", "replace"))) elif opcode == 0x8: code = struct.unpack(">H", payload[:2])[0] if len(payload) >= 2 else 0 reason = payload[2:].decode("utf-8", "replace") print("<<< close from server code={} reason={}".format(code, reason)) self.closed = True break elif opcode == 0x9: print("<<< ping") elif opcode == 0xA: print("<<< pong") finally: self.sock.settimeout(None) if frames == 0: print("<<< （这 {} 秒里没有收到任何文本帧）".format(seconds)) return frames def close(self) -> None: try: self.sock.sendall(bytes([0x80 | 0x8, 0x02]) + struct.pack(">H", 1000)) except OSError: pass try: self.sock.close() except OSError: pass def main(argv=None) -> int: parser = argparse.ArgumentParser(description="最小 WebSocket 探针") parser.add_argument("--host", default="127.0.0.1") parser.add_argument("--port", type=int, default=18080) parser.add_argument("--path", default="/ws/chat?token=probe-token") parser.add_argument("--send", action="append", default=[], help="连接建立后发送的帧（JSON 字符串，可重复）") parser.add_argument("--ping", action="store_true", help="发一个 ping 帧，验证服务端回 pong（OkHttp 30s 心跳走的就是它）") parser.add_argument("--listen", type=float, default=3.0, help="监听秒数（默认 3）") args = parser.parse_args(argv) probe = Probe(args.host, args.port, args.path) print("--- history 帧应当就在这里 ---") probe.listen(min(args.listen, 1.0), "initial") for payload in args.send: json.loads(payload) # 先自己校验一次，免得把坏 JSON 发出去 probe.send_json(payload) if args.ping: probe.send_ping() if args.send or args.ping: print("--- 上行帧已发出，等待服务端回帧 ---") probe.listen(args.listen, "after-send") probe.close() return 0 if __name__ == "__main__": sys.exit(main())