#!/usr/bin/python
"""
Resolve the current Linux arm64 Widevine CDM from Google's component
updater (Omaha protocol 3.1), download and verify it, and write a pin
file for the Firefox SlackBuild to merge into
toolkit/content/gmp-sources/widevinecdm.json.

Run on a networked maintainer box, review the diff, commit the pin file.
The SlackBuild itself then needs no network access.

Usage: widevine-pin-update.py [pinfile]
       default pinfile: widevine-aarch64.pin.json
"""
import hashlib
import json
import sys
import urllib.request
from pathlib import Path

APPID = "oimompecagnajdejgnnjijobebaeigek"      # Widevine CDM component ID
ENDPOINT = "https://update.googleapis.com/service/update2/json"
PRODVERSION = "140.0.0.0"   # Chrome version we claim to be; bump if Google stops answering


def die(msg):
    print(f"widevine-pin-update: {msg}", file=sys.stderr)
    sys.exit(1)


def omaha_query():
    req = {"request": {
        "@os": "linux",
        "@updater": "chromecrx",
        "protocol": "3.1",
        "acceptformat": "crx3",
        "prodversion": PRODVERSION,
        "updaterversion": PRODVERSION,
        "arch": "arm64",
        "nacl_arch": "arm64",
        "os": {"platform": "Linux", "arch": "arm64"},
        "app": [{"appid": APPID, "version": "0.0.0.0", "updatecheck": {}}],
    }}
    r = urllib.request.Request(
        ENDPOINT,
        data=json.dumps(req).encode(),
        headers={"Content-Type": "application/json"},
    )
    with urllib.request.urlopen(r, timeout=30) as f:
        body = f.read().decode()
    # Google prefixes the JSON with an XSSI guard: )]}'
    return json.loads(body[body.index("{"):])


def resolve():
    resp = omaha_query()
    try:
        uc = resp["response"]["app"][0]["updatecheck"]
    except (KeyError, IndexError):
        die(f"unexpected response shape:\n{json.dumps(resp, indent=2)}")
    if uc.get("status") != "ok":
        die(f"updatecheck status: {uc.get('status')!r}")

    version = uc["manifest"]["version"]
    pkg = uc["manifest"]["packages"]["package"][0]
    name = pkg["name"]
    size = int(pkg["size"])
    sha256 = pkg["hash_sha256"].lower()

    if "linux_arm64" not in name:
        die(f"Google returned a non-arm64 package: {name}")

    # Omaha returns http and https codebases; keep https only, dedupe, keep order.
    urls = []
    for u in uc["urls"]["url"]:
        base = u["codebase"]
        if base.startswith("https://") and base + name not in urls:
            urls.append(base + name)
    if not urls:
        die("no https download URLs in response")
    return version, urls, size, sha256


def fetch_and_hash(urls, size, sha256):
    last_err = None
    for url in urls:
        try:
            h256, h512, n, head = hashlib.sha256(), hashlib.sha512(), 0, b""
            with urllib.request.urlopen(url, timeout=120) as f:
                while chunk := f.read(1 << 20):
                    if not head:
                        head = chunk[:4]
                    h256.update(chunk)
                    h512.update(chunk)
                    n += len(chunk)
            if head != b"Cr24":
                raise ValueError("not a CRX file (missing Cr24 magic)")
            if n != size:
                raise ValueError(f"size {n} != advertised {size}")
            if h256.hexdigest() != sha256:
                raise ValueError("sha256 does not match Omaha manifest")
            return h512.hexdigest()
        except Exception as e:      # try the next mirror
            last_err = f"{url}: {e}"
            print(f"warning: {last_err}", file=sys.stderr)
    die(f"all downloads failed; last error: {last_err}")


def main():
    pinfile = Path(sys.argv[1] if len(sys.argv) > 1 else "widevine-aarch64.pin.json")
    old = json.loads(pinfile.read_text()) if pinfile.exists() else {}

    version, urls, size, sha256 = resolve()
    sha512 = fetch_and_hash(urls, size, sha256)

    pin = {
        "version": version,
        "fileUrl": urls[0],
        "mirrorUrls": urls[1:],
        "filesize": size,
        "hashValue": sha512,
    }

    if pin == old:
        print(f"unchanged: Widevine {version}")
        return
    pinfile.write_text(json.dumps(pin, indent=2) + "\n")
    print(f"updated {pinfile}: {old.get('version', '(none)')} -> {version}")


if __name__ == "__main__":
    main()