#!/usr/bin/env python3
"""Cross-platform browser cookie extraction for Tier-2 Chrome stealth.

The OS-keyring lookup is an injected boundary: cookie_paths resolves profile
paths and cookie_crypto derives keys + decrypts values, both pure and testable
with synthetic fixtures on any OS. This module wires them to a real browser DB
and a local Chrome CDP session controlled by playwright-core scripts.

Usage:
    python -B extract_cookies.py --browser chrome --domain youtube.com --output /tmp/cookies.json
    python -B extract_cookies.py --browser chrome --domain youtube.com --inject --cdp 9242
"""
from __future__ import annotations

import argparse
import json
import os
import shutil
import sqlite3
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import Callable, NotRequired, TypedDict

sys.dont_write_bytecode = True

from cookie_crypto import (
    decrypt_chromium_value,
    derive_key,
    linux_keyring_secret,
    macos_keyring_secret,
    windows_oscrypt_key,
)
from cookie_domains import domain_where_clause
from cookie_paths import BROWSERS, BrowserSpec, UnsupportedPlatform, platform_base, resolve_cookie_db

_SAMESITE = {-1: "None", 0: "None", 1: "Lax", 2: "Strict"}

IMPORTANT_COOKIES = {
    "SID", "SSID", "HSID", "APISID", "SAPISID",
    "__Secure-1PSID", "__Secure-3PSID", "__Secure-1PSIDTS", "__Secure-3PSIDTS",
    "LOGIN_INFO", "PREF", "VISITOR_INFO1_LIVE", "YSC", "NID", "CONSENT",
}


class CookieRecord(TypedDict):
    name: str
    value: str
    domain: str
    path: str
    expires: int
    secure: bool
    httpOnly: bool
    sameSite: str


class CdpCookie(TypedDict):
    name: str
    value: str
    domain: str
    path: str
    secure: bool
    httpOnly: bool
    sameSite: str
    expires: NotRequired[int]


_CDP_SET_COOKIES_SCRIPT = r"""
const port = Number(process.argv[1] || 0);
if (!Number.isInteger(port) || port <= 0) {
  process.stderr.write("invalid CDP port\n");
  process.exit(2);
}

let input = "";
process.stdin.setEncoding("utf8");
process.stdin.on("data", (chunk) => { input += chunk; });
process.stdin.on("end", () => {
  (async () => {
  const cookies = JSON.parse(input || "[]");
  const version = await fetch(`http://127.0.0.1:${port}/json/version`).then((r) => r.json());
  const ws = new WebSocket(version.webSocketDebuggerUrl);
  let nextId = 1;
  const pending = new Map();

  ws.onmessage = (event) => {
    const msg = JSON.parse(event.data);
    const waiter = pending.get(msg.id);
    if (!waiter) return;
    pending.delete(msg.id);
    if (msg.error) waiter.reject(new Error(msg.error.message || "CDP error"));
    else waiter.resolve(msg.result);
  };

  await new Promise((resolve, reject) => {
    ws.onopen = resolve;
    ws.onerror = reject;
  });

  const send = (method, params) => new Promise((resolve, reject) => {
    const id = nextId++;
    pending.set(id, { resolve, reject });
    ws.send(JSON.stringify({ id, method, params }));
  });

  let ok = 0;
  for (const cookie of cookies) {
    const result = await send("Network.setCookie", cookie);
    if (result && result.success) ok += 1;
  }
  ws.close();
  process.stdout.write(String(ok));
  })().catch((error) => {
  process.stderr.write(`${error.name || "Error"}: ${error.message || error}\n`);
  process.exit(1);
  });
});
"""


def _secure_cookie_db_copy(db_path: Path) -> Path:
    handle = tempfile.NamedTemporaryFile(prefix="omo-cookies-", suffix=".sqlite", delete=False)
    tmp = Path(handle.name)
    handle.close()
    try:
        shutil.copyfile(db_path, tmp)
        tmp.chmod(0o600)
        return tmp
    except (OSError, shutil.Error):
        tmp.unlink(missing_ok=True)
        raise


def write_cookie_file(path: Path, cookies: list[CookieRecord]) -> None:
    if path.is_symlink():
        raise ValueError(f"refusing to write cookies through symlink: {path}")
    path.parent.mkdir(parents=True, exist_ok=True)
    fd, tmp_name = tempfile.mkstemp(prefix=f".{path.name}.", suffix=".tmp", dir=str(path.parent))
    tmp_path = Path(tmp_name)
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as f:
            os.fchmod(f.fileno(), 0o600)
            json.dump(cookies, f, indent=2)
            f.write("\n")
        os.replace(tmp_path, path)
    finally:
        tmp_path.unlink(missing_ok=True)


def extract_firefox(db_path: Path, domains: list[str]) -> list[CookieRecord]:
    tmp = _secure_cookie_db_copy(db_path)
    try:
        where, params = domain_where_clause("host", domains)
        conn = sqlite3.connect(str(tmp))
        rows = conn.execute(
            f"SELECT name, value, host, path, expiry, isSecure, isHttpOnly, sameSite "
            f"FROM moz_cookies WHERE ({where}) ORDER BY host, name",
            params,
        ).fetchall()
        conn.close()
    finally:
        tmp.unlink(missing_ok=True)
    return [
        {
            "name": n, "value": v, "domain": h, "path": p, "expires": e,
            "secure": bool(sec), "httpOnly": bool(ho), "sameSite": _SAMESITE.get(ss, "Lax"),
        }
        for n, v, h, p, e, sec, ho, ss in rows
    ]


def extract_chromium(db_path: Path, domains: list[str], platform: str, key: bytes) -> list[CookieRecord]:
    tmp = _secure_cookie_db_copy(db_path)
    try:
        where, params = domain_where_clause("host_key", domains)
        conn = sqlite3.connect(str(tmp))
        rows = conn.execute(
            f"SELECT name, encrypted_value, host_key, path, expires_utc, is_secure, is_httponly, samesite "
            f"FROM cookies WHERE ({where}) ORDER BY host_key, name",
            params,
        ).fetchall()
        conn.close()
    finally:
        tmp.unlink(missing_ok=True)
    out = []
    for n, enc, h, p, exp, sec, ho, ss in rows:
        unix_expires = int((exp / 1_000_000) - 11644473600) if exp and exp > 0 else 0
        out.append({
            "name": n, "value": decrypt_chromium_value(platform, key, enc), "domain": h, "path": p,
            "expires": unix_expires, "secure": bool(sec), "httpOnly": bool(ho), "sameSite": _SAMESITE.get(ss, "Lax"),
        })
    return out


def _browser_spec(browser: str) -> BrowserSpec:
    spec = BROWSERS.get(browser)
    if spec is None:
        raise UnsupportedPlatform(f"unsupported browser: {browser!r}")
    return spec


def default_keyring_reader(platform: str, spec: BrowserSpec) -> Callable[[str], bytes]:
    if platform == "darwin":
        return macos_keyring_secret
    if platform == "linux":
        return linux_keyring_secret
    if platform == "win32":
        def _win(_safe_storage: str) -> bytes:
            base = platform_base("win32", "chromium")
            return windows_oscrypt_key(base / spec["dirs"]["win32"] / "Local State")
        return _win
    raise UnsupportedPlatform(f"no keyring reader for platform {platform!r}")


def extract_cookies(
    browser: str,
    domains: list[str],
    platform: str = sys.platform,
    keyring_reader: Callable[[str], bytes] | None = None,
    base_override: Path | None = None,
) -> list[CookieRecord]:
    spec = _browser_spec(browser)
    db = resolve_cookie_db(browser, platform, base_override=base_override)
    if spec["kind"] == "firefox":
        return extract_firefox(db, domains)
    reader = keyring_reader or default_keyring_reader(platform, spec)
    safe_storage = spec["safe_storage"]
    if safe_storage is None:
        raise UnsupportedPlatform(f"browser {browser!r} has no keyring storage name")
    key = derive_key(platform, reader(safe_storage))
    return extract_chromium(db, domains, platform, key)


def inject_cookies(cookies: list[CookieRecord], cdp_port: int) -> None:
    filtered = [c for c in cookies if c["name"] in IMPORTANT_COOKIES] or cookies
    payload: list[CdpCookie] = [
        {
            "name": c["name"],
            "value": c["value"],
            "domain": c["domain"],
            "path": c.get("path") or "/",
            "secure": bool(c.get("secure")),
            "httpOnly": bool(c.get("httpOnly")),
            "sameSite": c.get("sameSite", "Lax"),
            **({"expires": int(c["expires"])} if c.get("expires") and int(c["expires"]) > 0 else {}),
        }
        for c in filtered
    ]
    proc = subprocess.run(
        ["node", "-e", _CDP_SET_COOKIES_SCRIPT, str(cdp_port)],
        input=json.dumps(payload),
        capture_output=True,
        text=True,
        timeout=15,
    )
    if proc.returncode != 0:
        raise RuntimeError((proc.stderr or "CDP cookie injection failed").strip())
    ok = int((proc.stdout or "0").strip() or "0")
    print(f"Injected {ok}/{len(filtered)} cookies into local Chrome (playwright-core CDP {cdp_port})")


def main() -> None:
    parser = argparse.ArgumentParser(description="Extract browser cookies (cross-platform)")
    parser.add_argument("--browser", required=True, choices=sorted(BROWSERS.keys()))
    parser.add_argument("--domain", required=True, action="append", dest="domains")
    parser.add_argument("--output", help="Write cookies JSON to file")
    parser.add_argument("--inject", action="store_true", help="Inject into the local Chrome CDP endpoint used by a playwright-core script")
    parser.add_argument("--cdp", type=int, default=9242, help="CDP port (default 9242)")
    args = parser.parse_args()

    cookies = extract_cookies(args.browser, args.domains)
    print(f"Extracted {len(cookies)} cookies from {args.browser}")
    if args.output:
        write_cookie_file(Path(args.output), cookies)
        print(f"Saved to {args.output}")
    if args.inject:
        inject_cookies(cookies, args.cdp)


if __name__ == "__main__":
    main()