#!/usr/bin/env python3
"""最简 WebSocket 测试 - 验证基本通信"""
import socket, base64, os, struct, json, time, sys

IP = '198.120.0.100'
PORT = 8000

sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.settimeout(10)
sock.connect((IP, PORT))
print(f"[✓] TCP 连接成功")

# WebSocket 握手
key = base64.b64encode(os.urandom(16)).decode()
req = (f"GET /ws HTTP/1.1\r\nHost: {IP}:{PORT}\r\nUpgrade: websocket\r\n"
       f"Connection: Upgrade\r\nSec-WebSocket-Key: {key}\r\n"
       f"Sec-WebSocket-Version: 13\r\n\r\n")
sock.sendall(req.encode())

# 接收 HTTP 响应
resp = b""
while b"\r\n\r\n" not in resp:
    chunk = sock.recv(4096)
    if not chunk: break
    resp += chunk
print(f"[✓] HTTP 响应: {resp.split(b'\r\n')[0].decode()}")

# 检查响应码
if b"101" not in resp.split(b"\r\n")[0]:
    print(f"[✗] 升级失败: {resp.decode(errors='replace')[:500]}")
    sys.exit(1)

# 检查是否有额外的数据（WebSocket 帧紧跟在 HTTP 响应后）
headers_end = resp.find(b"\r\n\r\n") + 4
extra = resp[headers_end:]
if extra:
    print(f"[i] HTTP 响应后有额外 {len(extra)} 字节: {extra.hex()}")

# 手动构造 WebSocket 文本帧发送 get_cmds
msg = json.dumps({"type": "get_cmds"})
payload = msg.encode()
L = len(payload)
mask = os.urandom(4)
masked = bytes(payload[i] ^ mask[i % 4] for i in range(L))

frame = bytearray([0x81])  # FIN + text
if L < 126:
    frame.append(0x80 | L)
else:
    frame.append(0x80 | 126)
    frame.extend(struct.pack('>H', L))
frame.extend(mask)
frame.extend(masked)
sock.sendall(bytes(frame))
print(f"[→] 发送 get_cmds ({L} 字节)")

# 等待接收
sock.settimeout(3.0)
all_data = b""
try:
    while True:
        chunk = sock.recv(65536)
        if not chunk:
            break
        all_data += chunk
        print(f"[i] recv chunk: {len(chunk)} 字节")
except socket.timeout:
    print(f"[i] recv 超时")

print(f"\n总共接收到: {len(all_data)} 字节")

if len(all_data) == 0:
    print("[✗] 未收到任何数据！服务器可能未响应")
    sys.exit(1)

# 解析 WebSocket 帧
offset = 0
while offset < len(all_data):
    if offset + 2 > len(all_data): break
    b0 = all_data[offset]
    b1 = all_data[offset + 1]
    opcode = b0 & 0x0F
    length = b1 & 0x7F
    offset += 2
    if length == 126:
        if offset + 2 > len(all_data): break
        length = struct.unpack('>H', all_data[offset:offset+2])[0]
        offset += 2
    elif length == 127:
        if offset + 8 > len(all_data): break
        length = struct.unpack('>Q', all_data[offset:offset+8])[0]
        offset += 8

    if offset + length > len(all_data):
        print(f"  [W] 帧不完整: 需要 {length} 字节但只剩 {len(all_data)-offset}")
        break
    payload = all_data[offset:offset+length]
    offset += length

    if opcode == 0x01:  # text
        print(f"  [text帧] {payload.decode('utf-8', errors='replace')[:500]}")
    elif opcode == 0x02:  # binary
        try:
            obj = json.loads(payload.decode('utf-8'))
            print(f"  [binary帧] JSON: type={obj.get('type','?')}, keys={list(obj.keys())}")
            if obj.get('type') == 'cmd_list':
                print(f"    cmd_list: {len(obj.get('cmds',[]))} 个命令")
        except:
            print(f"  [binary帧] {len(payload)} 字节: {payload[:200]}")
    elif opcode == 0x08:
        print(f"  [close帧]")
    elif opcode == 0x09:
        print(f"  [ping帧]")
    elif opcode == 0x0A:
        print(f"  [pong帧]")
    else:
        print(f"  [opcode={opcode}帧] {len(payload)} 字节")

sock.close()