#!/usr/bin/env python3
"""vnc-ws-probe.py — tes handshake WebSocket ke websockify (dengan/tanpa token)."""
import base64
import os
import socket
import sys


def probe(url_path, label):
    try:
        s = socket.create_connection(("localhost", 6080), timeout=8)
        key = base64.b64encode(os.urandom(16)).decode()
        req = (
            f"GET {url_path} HTTP/1.1\r\n"
            f"Host: localhost:6080\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\r\n"
        )
        s.sendall(req.encode())
        s.settimeout(8)
        data = s.recv(4096)
        head, _, rest = data.partition(b"\r\n\r\n")
        status = head.split(b"\r\n")[0].decode(errors="replace")
        hint = rest[:16].decode(errors="replace") if rest else ""
        print(f"{label:<28} -> {status}   {hint}")
        s.close()
        return status
    except Exception as e:
        print(f"{label:<28} -> ERROR: {e}")
        return "ERROR"


if __name__ == "__main__":
    token_valid = sys.argv[1] if len(sys.argv) > 1 else ""
    token_expired = sys.argv[2] if len(sys.argv) > 2 else ""

    probe("/websockify", "1) TANPA token")
    if token_valid:
        probe(f"/websockify?token={token_valid}", "2) token VALID")
    if token_expired:
        probe(f"/websockify?token={token_expired}", "3) token KEDALUWARSA")