Cyrene-For-Android/devtools/fake-backend/ws_probe.py

#!/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())