#!/usr/bin/env python3
"""Viaplay license helper for ModularDrm. Requires: pip install pywidevine requests"""
import json
import sys

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


def main():
    data = json.load(sys.stdin)
    device = Device.load(data["wvdPath"])
    cdm = Cdm.from_device(device)
    session_id = cdm.open()
    try:
        challenge = cdm.get_license_challenge(session_id, PSSH(data["pssh"]))
        headers = {
            "Accept": "*/*",
            "Origin": "https://viaplay.com",
            "Referer": "https://viaplay.com/",
            "User-Agent": data.get("userAgent")
            or "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/153.0.0.0 Safari/537.36",
            "Content-Type": "application/octet-stream",
        }
        auth = (data.get("licAuth") or "").strip()
        if auth:
            headers["Authorization"] = f"Basic {auth}"
        response = requests.post(
            data["licUrl"],
            headers=headers,
            data=challenge,
            timeout=30,
        )
        response.raise_for_status()
        cdm.parse_license(session_id, response.content)
        keys = [
            f"{key.kid.hex}:{key.key.hex()}"
            for key in cdm.get_keys(session_id)
            if key.type != "SIGNING"
        ]
        print(json.dumps({"keys": keys}, separators=(",", ":")))
    finally:
        cdm.close(session_id)


if __name__ == "__main__":
    try:
        main()
    except Exception as exc:
        print(str(exc), file=sys.stderr)
        sys.exit(1)