#!/usr/bin/env python3
"""Read-only production WebSocket handshake diagnostic.

Runs on the BaoTa host so production credentials never leave the server. The
script prints only HTTP status lines and the first command received.
"""
import base64
import hashlib
import hmac
import json
import os
import socket
import ssl
import sys
import time

sys.path.insert(0, "/www/server/xingyu-im/ops")
from baota_publish_test_users import database, env_values, query


def encode_json(value):
    raw = json.dumps(value, separators=(",", ":"), ensure_ascii=False).encode()
    return base64.urlsafe_b64encode(raw).rstrip(b"=").decode()


def user_token(secret, user_id, nickname, version):
    now = int(time.time())
    header = encode_json({"alg": "HS256", "typ": "JWT"})
    payload = encode_json({
        "role": "user", "name": nickname, "ver": version,
        "sub": str(user_id), "iat": now, "exp": now + 180,
    })
    signing_input = (header + "." + payload).encode()
    signature = base64.urlsafe_b64encode(
        hmac.new(secret.encode(), signing_input, hashlib.sha256).digest()
    ).rstrip(b"=").decode()
    return header + "." + payload + "." + signature


def receive_exact(stream, count):
    output = b""
    while len(output) < count:
        chunk = stream.recv(count - len(output))
        if not chunk:
            raise RuntimeError("connection closed")
        output += chunk
    return output


def first_command(stream):
    first, second = receive_exact(stream, 2)
    size = second & 0x7F
    if size == 126:
        size = int.from_bytes(receive_exact(stream, 2), "big")
    elif size == 127:
        size = int.from_bytes(receive_exact(stream, 8), "big")
    payload = receive_exact(stream, size)
    return json.loads(payload).get("command")


def handshake(token, origin):
    key = base64.b64encode(os.urandom(16)).decode()
    headers = [
        "GET /ws HTTP/1.1",
        "Host: im.bchongw.com",
        "Upgrade: websocket",
        "Connection: Upgrade",
        "Sec-WebSocket-Version: 13",
        "Sec-WebSocket-Key: " + key,
        "Sec-WebSocket-Protocol: xingyu.jwt." + token,
    ]
    if origin is not None:
        headers.append("Origin: " + origin)
    request = ("\r\n".join(headers) + "\r\n\r\n").encode()
    context = ssl.create_default_context()
    with socket.create_connection(("im.bchongw.com", 443), timeout=10) as raw:
        with context.wrap_socket(raw, server_hostname="im.bchongw.com") as stream:
            stream.settimeout(10)
            stream.sendall(request)
            response = b""
            while b"\r\n\r\n" not in response:
                response += stream.recv(4096)
            head, remainder = response.split(b"\r\n\r\n", 1)
            status = head.split(b"\r\n", 1)[0].decode("ascii", errors="replace")
            command = None
            if " 101 " in status:
                if remainder:
                    # Keep this branch simple; the server's AUTH_ACK normally
                    # follows in a second TLS record and is read below.
                    class Buffered:
                        def __init__(self, initial, source): self.initial, self.source = initial, source
                        def recv(self, count):
                            if self.initial:
                                chunk, self.initial = self.initial[:count], self.initial[count:]
                                return chunk
                            return self.source.recv(count)
                    command = first_command(Buffered(remainder, stream))
                else:
                    command = first_command(stream)
            return status, command


values = env_values()
conn, _ = database(values)
try:
    user = query(conn, """SELECT user.id,profile.nickname,security.token_version
        FROM users user JOIN user_profiles profile ON profile.user_id=user.id
        JOIN user_security_controls security ON security.user_id=user.id
        WHERE user.status=1 AND user.deleted_at IS NULL ORDER BY user.id LIMIT 1""")[0]
finally:
    conn.close()

token = user_token(values["IM_JWT_SECRET"], user[0], user[1], user[2])
cases = [
    ("no-origin", None),
    ("web-origin", "https://im.bchongw.com"),
    ("native-localhost", "http://localhost"),
    ("native-file", "file://"),
]
for label, origin in cases:
    try:
        status, command = handshake(token, origin)
        print(json.dumps({"case": label, "status": status, "command": command}, ensure_ascii=False))
    except Exception as error:
        print(json.dumps({"case": label, "error": type(error).__name__}, ensure_ascii=False))