#!/usr/bin/python3
import sys
import os
import o11
import base64
import json
import datetime
import pytz
import requests
from pywidevine.cdm import Cdm
from pywidevine.device import Device
from pywidevine.pssh import PSSH
from bs4 import BeautifulSoup

WVD_PATH = './WVD.wvd'

user = o11.parse_params(sys.argv, 'user')
password = o11.parse_params(sys.argv, 'password')
token_param = o11.parse_params(sys.argv, 'token')
device_param = o11.parse_params(sys.argv, 'device')
pin = o11.parse_params(sys.argv, 'pin')

id = o11.parse_params(sys.argv, 'id')
action = o11.parse_params(sys.argv, 'action')

bind = o11.parse_params(sys.argv, 'bind')
proxy = o11.parse_params(sys.argv, 'proxy')
doh = o11.parse_params(sys.argv, 'doh')
worker = o11.parse_params(sys.argv, 'worker')

cdm = o11.parse_params(sys.argv, 'cdm')
drm = o11.parse_params(sys.argv, 'drm')
kid = o11.parse_params(sys.argv, 'kid')
pssh = o11.parse_params(sys.argv, 'pssh')
challenge = o11.parse_params(sys.argv, 'challenge')

heartbeaturl = o11.parse_params(sys.argv, 'heartbeaturl')
heartbeatparams = o11.parse_params(sys.argv, 'heartbeatparams')

o11Session = o11.session(bind=bind, proxy=proxy, worker=worker)
req = o11Session.get_session()
if doh != "":
    o11.dns(doh)

if challenge == "cert":
    challenge = "CAQ="

authFile = '/CanalDigitaal_' + user + '.tokens'
proxies = {"http": proxy, "https": proxy} if proxy else {}

user_agent = 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/121.0.0.0 Safari/537.36'

def do_cdm_external(pssh_data, license_url):
    """Extract Widevine keys using pywidevine library"""
    pssh_obj = PSSH(pssh_data)
    device_obj = Device.load(WVD_PATH)
    cdm_obj = Cdm.from_device(device_obj)
    session_id = cdm_obj.open()
    challenge_data = cdm_obj.get_license_challenge(session_id, pssh_obj)

    licence = req.post(license_url, data=challenge_data)

    try:
        licence.raise_for_status()
        cdm_obj.parse_license(session_id, licence.content)
        keys = []
        for key in cdm_obj.get_keys(session_id):
            if key.type != 'SIGNING':
                keys.append(f"{key.kid.hex}:{key.key.hex()}")
        cdm_obj.close(session_id)
        return keys
    except Exception as e:
        print(f'License request failed: {e}', file=sys.stderr)
        return []

def get_initial(token, channel_id):
    headers = {
        'authority': 'tvapi.solocoo.tv',
        'accept': 'application/json, text/plain, */*',
        'authorization': 'Bearer ' + token,
        'content-type': 'application/json',
    }
    json_data = {
        'player': {
            'name': 'RxPlayer',
            'version': '3.29.0',
            'capabilities': {
                'mediaTypes': ['DASH'],
                'drmSystems': ['Widevine'],
                'smartLib': True,
            },
        },
    }
    response = req.post(f'https://tvapi.solocoo.tv/v1/assets/{channel_id}/play', headers=headers, json=json_data)
    data = response.json()
    return data['url'], data['drm']['licenseUrl']

def get_mpd(temp_mpd):
    response = req.get(temp_mpd + '&response=200&bk-ml=1', allow_redirects=False)
    return response.headers.get('Location', temp_mpd)

def extract_pssh(mpd_url):
    response = req.get(mpd_url)
    psshs = []
    for p in BeautifulSoup(response.content, features="xml").findAll('cenc:pssh'):
        psshs.append(p.text)
    return min(psshs, key=len) if psshs else None

def get_channels(token):
    headers = {
        'authority': 'tvapi.solocoo.tv',
        'accept': 'application/json, text/plain, */*',
        'authorization': 'Bearer ' + token,
        'content-type': 'application/json',
    }
    response = req.get('https://tvapi.solocoo.tv/v1/bouquet', headers=headers)
    data = response.json()
    channels = []
    for item in data.get('items', []):
        if item.get('type') == 'channel':
            channels.append(item)
    return channels

def login():
    print("Saving token...", file=sys.stderr)
    if not token_param:
        print("Please provide token parameter from browser DevTools -> Application -> Local Storage", file=sys.stderr)
        sys.exit(1)
    auth_data = {'token': token_param}
    json.dump(auth_data, open(os.path.abspath(os.path.dirname(__file__)) + authFile, 'w'))
    print("Token saved successfully", file=sys.stderr)

def do_action():
    if action == "login":
        login()
        sys.exit()

    try:
        auth = json.load(open(os.path.abspath(os.path.dirname(__file__)) + authFile))
        token = auth['token']
    except:
        return "error"

    if action == "channels":
        output = {'Channels': []}
        channels = get_channels(token)
        for chan in channels:
            channel = {
                'Name': chan.get('title', chan.get('id')),
                'Mode': "live",
                'SessionManifest': True,
                'ManifestScript': 'id=' + str(chan['id']),
                'CdmType': "widevine",
                'UseCdm': True,
                'Cdm': 'id=' + str(chan['id']),
                'Video': 'best',
                'OnDemand': True,
                'SpeedUp': True
            }
            output['Channels'].append(channel)
        print(json.dumps(output, indent=2))

    elif action == "events":
        output = {'Events': []}
        channels = get_channels(token)
        for chan in channels:
            event = {
                'Name': chan.get('title', chan.get('id')),
                'Mode': "live",
                'SessionManifest': True,
                'ManifestScript': 'id=' + str(chan['id']),
                'CdmType': "widevine",
                'UseCdm': True,
                'Cdm': 'id=' + str(chan['id']),
                'Video': 'best',
                'Autostart': True,
                'Start': int(datetime.datetime.now(pytz.UTC).timestamp()),
                'End': int((datetime.datetime.now(pytz.UTC) + datetime.timedelta(hours=4)).timestamp())
            }
            output['Events'].append(event)
        print(json.dumps(output, indent=2))

    elif action == "heartbeat":
        sys.exit()

    elif action == "manifest":
        channel_id = id
        try:
            temp_mpd, license_url = get_initial(token, channel_id)
            
            if 'index.mpd' in temp_mpd:
                mpd_url = get_mpd(temp_mpd)
            else:
                mpd_url = temp_mpd
            
            pssh_data = extract_pssh(mpd_url)
            
            output = {
                "Cdn": [],
                "ManifestUrl": mpd_url,
                "Headers": {
                    "Manifest": {'User-Agent': user_agent},
                    "Media": {'User-Agent': user_agent}
                },
                "Heartbeat": {
                    "Url": '',
                    "Params": '',
                    "PeriodMs": 5*60*1000
                },
                "LicenseUrl": license_url
            }
            if pssh_data:
                output['Pssh'] = pssh_data
            print(json.dumps(output))
        except Exception as e:
            print(f"Error getting manifest: {e}", file=sys.stderr)
            return "error"

    elif action == "cdm" and cdm == "internal":
        lic_url = o11.parse_params(sys.argv, 'licenseUrl')
        try:
            response = req.post(lic_url, data=base64.b64decode(challenge))
            response_b64 = str(base64.b64encode(response.content), 'ascii')
            if response_b64.startswith('CA'):
                print(response_b64)
            else:
                print(response.text, file=sys.stderr)
                return "error"
        except Exception as e:
            print(f"Error: {e}", file=sys.stderr)
            return "error"

    elif action == "cdm" and cdm == "external":
        channel_id = id
        try:
            temp_mpd, license_url = get_initial(token, channel_id)
            
            if 'index.mpd' in temp_mpd:
                mpd_url = get_mpd(temp_mpd)
            else:
                mpd_url = temp_mpd
            
            pssh_data = extract_pssh(mpd_url)
            
            if pssh_data:
                keys = do_cdm_external(pssh_data, license_url)
                for key in keys:
                    print(key)
            else:
                print("No PSSH found in manifest", file=sys.stderr)
                return "error"
        except Exception as e:
            print(f"Error: {e}", file=sys.stderr)
            return "error"

    else:
        print("invalid action: " + action, file=sys.stderr)

if do_action() == "error":
    login()
    do_action()
