#!/usr/bin/env python3
"""Helper script called by Node.js to get Widevine content keys.

Routes license requests through the local proxy which handles auth/signing.

Usage: python3 pywidevine_helper.py <wvd_path> <pssh_b64> <proxy_license_url>

Outputs JSON: {"keys": [{"kid": "hex", "key": "hex", "type": "CONTENT"}]}
"""
import sys
import json
import requests
from pywidevine.cdm import Cdm
from pywidevine.device import Device
from pywidevine.pssh import PSSH


def main():
    if len(sys.argv) < 4:
        print(json.dumps({"error": "Usage: pywidevine_helper.py <wvd_path> <pssh_b64> <proxy_license_url>"}))
        sys.exit(1)

    wvd_path = sys.argv[1]
    pssh_b64 = sys.argv[2]
    proxy_url = sys.argv[3]  # e.g. http://localhost:3001/api/drm-license?mediaId=...

    try:
        device = Device.load(wvd_path)
        cdm = Cdm.from_device(device)
        session_id = cdm.open()
        pssh = PSSH(pssh_b64)

        # Step 1: Get service certificate via proxy
        cert_res = requests.post(
            proxy_url,
            data=Cdm.service_certificate_challenge,
            headers={"Content-Type": "application/octet-stream"},
            timeout=30,
        )
        if cert_res.ok:
            try:
                cdm.set_service_certificate(session_id, cert_res.content)
            except Exception:
                pass  # Continue without privacy mode

        # Step 2: Get license via proxy
        challenge = cdm.get_license_challenge(session_id, pssh)
        lic_res = requests.post(
            proxy_url,
            data=challenge,
            headers={"Content-Type": "application/octet-stream"},
            timeout=30,
        )

        if not lic_res.ok:
            print(json.dumps({"error": f"License failed: {lic_res.status_code} {lic_res.text[:200]}"}))
            sys.exit(1)

        cdm.parse_license(session_id, lic_res.content)
        keys = []
        for key in cdm.get_keys(session_id):
            keys.append({
                "kid": key.kid.hex,
                "key": key.key.hex(),
                "type": key.type,
            })

        cdm.close(session_id)
        print(json.dumps({"keys": keys}))

    except Exception as e:
        print(json.dumps({"error": str(e)}))
        sys.exit(1)


if __name__ == "__main__":
    main()