keys.py

import argparse
import base64
import os
import re
import sys
from urllib.parse import urlparse
import requests

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

def get_mpd(mpd_url: str) -> str:
    headers = {
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/105.0.0.0 Safari/537.36',
        # Add any necessary headers here
    }
    try:
        req = requests.get(mpd_url, headers=headers)
        return req.text
    except Exception as e:
        print("Failed getting the MPD:", e)
        sys.exit(1)

def get_pssh(mpd_content: str):
    try:
        return re.search('<cenc:pssh>(.*)</cenc:pssh>', mpd_content).group(1)
    except Exception as e:
        print("Failed getting the PSSH:", e)
        sys.exit(1)

def get_license_response(license_challenge: bytes, license_url: str, headers: dict):
    try:
        req = requests.post(license_url, data=license_challenge, headers=headers)
        return req.content  # Return raw response content
    except Exception as e:
        print("Failed getting license response:", e)
        sys.exit(1)

def main():
    parser = argparse.ArgumentParser(description='Fetch PlayReady or Widevine keys from any license URL with MPD')
    parser.add_argument('--wvd', help='The file path to the WVD file generated by pywidevine')
    parser.add_argument('--mpd_url', help='The MPD URL')
    parser.add_argument('--license_url', help='The License URL')
    parser.add_argument('--headers', help='JSON string of custom headers for license request')
    args = parser.parse_args()

    if not args.wvd or not args.mpd_url or not args.license_url or not args.headers:
        parser.print_help()
        sys.exit(1)

    headers = json.loads(args.headers)

    mpd_content = get_mpd(args.mpd_url)
    pssh = get_pssh(mpd_content)

    device = Device.load(args.wvd)
    cdm = Cdm.from_device(device)
    session_id = cdm.open()
    challenge = cdm.get_license_challenge(session_id, PSSH(pssh), privacy_mode=True)

    license_response = get_license_response(challenge, args.license_url, headers)
    cdm.parse_license(session_id, license_response)

    print("\nKeys:")
    for key in cdm.get_keys(session_id):
        if key.type == 'CONTENT':
            print(f'KID: {key.kid.hex()}, KEY: {key.key.hex()}')

    cdm.close(session_id)

if __name__ == '__main__':
    main()