import argparse
import requests
from bs4 import BeautifulSoup
import json
from cdm.wks import WvDecrypt, device_android_generic, PsshExtractor, KeyExtractor


# Parse URL input
parser = argparse.ArgumentParser(description='PYWKS-NPO')
parser.add_argument('-url', dest='url', required=True, help='NPO Video URL')
args = parser.parse_args()


# Get HTML and extract productId
headers = {
    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:109.0) Gecko/20100101 Firefox/119.0',
    'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,*/*;q=0.8',
    'Accept-Language': 'en-US,en;q=0.5',
    'Connection': 'keep-alive',
    'Upgrade-Insecure-Requests': '1',
    'Sec-Fetch-Dest': 'document',
    'Sec-Fetch-Mode': 'navigate',
    'Sec-Fetch-Site': 'same-site',
    'Pragma': 'no-cache',
    'Cache-Control': 'no-cache',
}

response_targetId = requests.get(args.url, headers=headers)
content = response_targetId.content

try: 
    url_split = args.url.split("/")
    target_slug = url_split[7]
except:
    print("URL invalid.")
    print("URL format: https://npo.nl/start/serie/wie-is-de-mol/seizoen-24/wie-is-de-mol_56/afspelen")
    print(f"Your URL: {args.url}")
    exit()

soup = BeautifulSoup(content, 'html.parser')
script_tag = soup.find('script', {'id': '__NEXT_DATA__'})

if script_tag:
    script_content = script_tag.contents[0]
else:
    print("Script tag not found.")

def search(data, target_slug):
    if isinstance(data, list):
        for item in data:
            result = search(item, target_slug)
            if result:
                return result
    elif isinstance(data, dict):
        for key, value in data.items():
            if key == "slug" and value == target_slug:
                return data.get("productId")
            else:
                result = search(value, target_slug)
                if result:
                    return result
    return None

data_dict = json.loads(script_content)
target_product_id = search(data_dict, target_slug)


# Get CSRF token
response_CSRF = requests.get('https://npo.nl/start/api/auth/session', headers=headers)
response_cookies = response_CSRF.cookies.get_dict()
csrf = response_cookies["__Host-next-auth.csrf-token"]


# Get player token
cookies = {
    '__Host-next-auth.csrf-token': csrf,
    '__Secure-next-auth.callback-url': 'https%3A%2F%2Fnpo.nl',
}

json_productId = {
    'productId': target_product_id,
}

response_token = requests.post('https://npo.nl/start/api/domain/player-token', cookies=cookies, headers=headers, json=json_productId)
token = response_token.json()["token"]


# Get MPD URL
headers['authorization'] = token

json_auth = {
    'profileName': 'dash',
    'drmType': 'widevine',
    'referrerUrl': args.url,
}

response = requests.post('https://prod.npoplayer.nl/stream-link', headers=headers, json=json_auth)
response_data = response.json()
stream_data = response_data.get('stream', {})

if stream_data.get('streamURL'):
    print('MPD URL:', stream_data.get('streamURL'))
else: 
    print("NO MPD URL - BAD TOKEN")
    exit()  


# Get PSSH
mpd_url = stream_data.get('streamURL')
license_url = "https://npo-drm-gateway.samgcloud.nepworldwide.nl/authentication"

response = requests.get(mpd_url, headers=headers)
pssh_extractor = PsshExtractor(response.text)
pssh_value = pssh_extractor.extract_pssh()

print("PSSH:", pssh_value)
headers_license = {
    'x-custom-data': stream_data.get('drmToken'),
    'origin': 'https://start-player.npo.nl',
    'referer': 'https://start-player.npo.nl/',
}


# Get Key
cert_b64 = None
key_extractor = KeyExtractor(pssh_value, cert_b64, license_url, headers_license)
keys = key_extractor.get_keys()
wvdecrypt = WvDecrypt(init_data_b64=pssh_value, cert_data_b64=cert_b64, device=device_android_generic)
raw_challenge = wvdecrypt.get_challenge()
data = raw_challenge
for key in keys:
    if isinstance(key, list):
        if key:
            for key_str in key:
                print(f"KEY: {key_str}")