#!/usr/bin/env python3
"""M2 WS Bridge 运行时验证：连 /ws 端点跑 chains.list"""

import json
import socket
import base64
import os
import hashlib
import struct
import urllib.request

TOKEN = open("/tmp/m1_token.txt").read().strip()
HOST = "127.0.0.1"
PORT = 8081
PATH = "/api/user/friend-agent/ws"

# 1) WebSocket 握手（RFC6455 手写，避免依赖库）
key = base64.b64encode(os.urandom(16)).decode()
hs = (
    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"Authorization: Bearer {TOKEN}\r\n\r\n"
)
s = socket.create_connection((HOST, PORT), timeout=30)
s.sendall(hs.encode())
resp = b""
while b"\r\n\r\n" not in resp:
    resp += s.recv(4096)
head = resp.split(b"\r\n\r\n")[0].decode(errors="replace")
print("=== 握手响应 ===")
print(head.split("\r\n")[0])
if "101" not in head.split("\r\n")[0]:
    print("HANDSHAKE_FAILED")
    print(head)
    raise SystemExit(1)
print("HANDSHAKE_OK")


def send_text(sock, text):
    payload = text.encode()
    mask = os.urandom(4)
    header = bytearray([0x81])
    n = len(payload)
    if n < 126:
        header.append(0x80 | n)
    elif n < 65536:
        header.append(0x80 | 126)
        header += struct.pack(">H", n)
    else:
        header.append(0x80 | 127)
        header += struct.pack(">Q", n)
    masked = bytes(b ^ mask[i % 4] for i, b in enumerate(payload))
    sock.sendall(bytes(header) + mask + masked)


def recv_exact(sock, n):
    data = b""
    while len(data) < n:
        chunk = sock.recv(n - len(data))
        if not chunk:
            raise EOFError
        data += chunk
    return data


def recv_text(sock):
    b1, b2 = recv_exact(sock, 2)
    opcode = b1 & 0x0F
    length = b2 & 0x7F
    if length == 126:
        length = struct.unpack(">H", recv_exact(sock, 2))[0]
    elif length == 127:
        length = struct.unpack(">Q", recv_exact(sock, 8))[0]
    data = recv_exact(sock, length)
    return data.decode(errors="replace"), opcode


# 2) 发送 chains.list req
req = {
    "type": "req",
    "request_id": "verify-001",
    "method": "chains.list",
    "payload": {},
}
send_text(s, json.dumps(req))
print("=== 已发送 chains.list ===")

# 3) 读响应（可能先收到 ping/其他帧）
for _ in range(5):
    text, op = recv_text(s)
    print(f"[frame opcode={op}] {text[:600]}")
    if '"request_id"' in text and "verify-001" in text:
        break

# 4) 再测一个白名单外 method（应被拒）
req2 = {
    "type": "req",
    "request_id": "verify-002",
    "method": "admin.danger",
    "payload": {},
}
send_text(s, json.dumps(req2))
text, _ = recv_text(s)
print("=== 白名单外 method 响应 ===")
print(text[:300])

s.close()
print("WS_VERIFY_DONE")