#!/usr/bin/env python3
"""Thin Widevine CDM helper for the Go capture tooling.

Prints a single CONTENT key as KID:KEY on stdout (no banners) when --quiet.
Exit code 0 on success.

Example:
  .venv/Scripts/python.exe wvkey.py \\
    --wvd path/to/device.wvd \\
    --pssh 'AAAA…' \\
    --license-url 'https://…' \\
    --auth 'Basic …' \\
    --pid 'xxxx' \\
    --quiet
"""

from __future__ import annotations

import argparse
import base64
import json
import sys
from pathlib import Path

import requests
from pywidevine.cdm import Cdm
from pywidevine.device import Device
from pywidevine.pssh import PSSH


def fetch_content_keys_modulardrm(
    *,
    wvd: Path,
    pssh_b64: str,
    auth: str,
    release_pid: str,
    license_url: str,
    user_agent: str,
) -> list[str]:
    """RTE / thePlatform ModularDrm (JSON challenge wrapper)."""
    device = Device.load(str(wvd))
    cdm = Cdm.from_device(device)
    session_id = cdm.open()
    try:
        challenge = cdm.get_license_challenge(session_id, PSSH(pssh_b64))
        challenge_b64 = base64.b64encode(challenge).decode("utf-8")
        payload = {
            "getWidevineLicense": {
                "releasePid": release_pid,
                "widevineChallenge": challenge_b64,
            }
        }
        headers = {
            "Accept-Encoding": "gzip",
            "Authorization": auth,
            "Content-Type": "application/json; charset=utf-8",
        }
        if user_agent:
            headers["User-Agent"] = user_agent
        response = requests.post(license_url, json=payload, headers=headers, timeout=30)
        data = response.json()
        if data.get("isException"):
            raise RuntimeError(data.get("description") or "license server exception")
        license_b64 = (
            data.get("getWidevineLicenseResponse", {}).get("license")
            or data.get("license")
        )
        if not license_b64:
            raise RuntimeError("no license in response: " + json.dumps(data)[:1500])
        cdm.parse_license(session_id, base64.b64decode(license_b64))
        keys = [
            f"{key.kid.hex}:{key.key.hex()}"
            for key in cdm.get_keys(session_id)
            if key.type == "CONTENT"
        ]
    finally:
        cdm.close(session_id)
    if not keys:
        raise RuntimeError("no CONTENT keys")
    return keys


def fetch_content_keys_raw(
    *,
    wvd: Path,
    pssh_b64: str,
    license_url: str,
    user_agent: str,
) -> list[str]:
    """Brightcove / generic Widevine: POST raw challenge, parse raw license body."""
    device = Device.load(str(wvd))
    cdm = Cdm.from_device(device)
    session_id = cdm.open()
    try:
        challenge = cdm.get_license_challenge(session_id, PSSH(pssh_b64))
        headers = {
            "Content-Type": "application/octet-stream",
            "Accept": "*/*",
        }
        if user_agent:
            headers["User-Agent"] = user_agent
        response = requests.post(
            license_url, data=challenge, headers=headers, timeout=30
        )
        if response.status_code >= 400:
            raise RuntimeError(
                f"license HTTP {response.status_code}: {response.text[:500]!r}"
            )
        body = response.content
        if not body:
            raise RuntimeError("empty license response body")
        # Some servers wrap base64 text; try raw first, then b64.
        try:
            cdm.parse_license(session_id, body)
        except Exception:
            try:
                cdm.parse_license(session_id, base64.b64decode(body))
            except Exception as exc:
                raise RuntimeError(
                    f"parse_license failed ({len(body)} bytes): {exc}"
                ) from exc
        keys = [
            f"{key.kid.hex}:{key.key.hex()}"
            for key in cdm.get_keys(session_id)
            if key.type == "CONTENT"
        ]
    finally:
        cdm.close(session_id)
    if not keys:
        raise RuntimeError("no CONTENT keys")
    return keys


def main(argv: list[str] | None = None) -> int:
    p = argparse.ArgumentParser(description="Widevine CDM helper → KID:KEY")
    p.add_argument("--wvd", type=Path, required=True, help="Path to .wvd device file")
    p.add_argument("--pssh", required=True, help="Base64 PSSH box")
    p.add_argument(
        "--mode",
        choices=("modulardrm", "raw"),
        default="modulardrm",
        help="modulardrm=RTE JSON; raw=Brightcove/octet-stream challenge",
    )
    p.add_argument("--auth", default="", help="Authorization header (modulardrm)")
    p.add_argument("--pid", default="", help="releasePid (modulardrm)")
    p.add_argument(
        "--license-url",
        required=True,
        help="Widevine license server URL",
    )
    p.add_argument(
        "--user-agent",
        default="",
        help="Optional User-Agent for the license request",
    )
    p.add_argument(
        "--quiet",
        action="store_true",
        help="Print only KID:KEY lines (for Go tooling)",
    )
    p.add_argument(
        "--all",
        action="store_true",
        help="Print every CONTENT key (default: first only)",
    )
    args = p.parse_args(argv)

    if not args.wvd.is_file():
        print(f"[!] missing WVD: {args.wvd}", file=sys.stderr)
        return 2

    try:
        if args.mode == "raw":
            keys = fetch_content_keys_raw(
                wvd=args.wvd,
                pssh_b64=args.pssh,
                license_url=args.license_url,
                user_agent=args.user_agent,
            )
        else:
            if not args.auth or not args.pid:
                print("[!] --auth and --pid required for modulardrm mode", file=sys.stderr)
                return 2
            keys = fetch_content_keys_modulardrm(
                wvd=args.wvd,
                pssh_b64=args.pssh,
                auth=args.auth,
                release_pid=args.pid,
                license_url=args.license_url,
                user_agent=args.user_agent,
            )
    except Exception as exc:
        print(f"[!] {exc}", file=sys.stderr)
        return 1

    out = keys if args.all else keys[:1]
    if not args.quiet:
        print("[+] Keys:", file=sys.stderr)
        for pair in out:
            print(f"--key {pair}", file=sys.stderr)
    for pair in out:
        print(pair)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())