#!/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')
device = o11.parse_params(sys.argv, 'device')
pin = o11.parse_params(sys.argv, 'pin')
token_param = o11.parse_params(sys.argv, 'token')

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 = '/BlueTV_' + user + '.tokens'
proxies = {"http": proxy, "https": proxy} if proxy else {}

headers = {
    'Accept': 'application/json, text/plain, */*',
    'Content-Type': 'application/json; charset=utf-8',
    'Origin': 'https://tv.blue.ch',
    'Referer': 'https://tv.blue.ch/',
    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36',
}

def do_cdm_external(pssh_data, token):
    """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)

    lic_headers = {
        'Authorization': 'Bearer ' + token,
        'Origin': 'https://tv.blue.ch',
        'Referer': 'https://tv.blue.ch/',
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36',
        'Content-Type': 'application/x-www-form-urlencoded',
    }

    licence = req.post('https://services.sg102.prd.sctv.ch/drm.cas/widevine/license/', headers=lic_headers, 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)
        print(licence.content, file=sys.stderr)
        return []

def get_single(token, channel_id):
    single_headers = {
        'Accept': 'application/json, text/plain, */*',
        'Authorization': 'Bearer ' + token,
        'Content-Type': 'application/json; charset=utf-8',
        'Loading-Disabled': 'true',
        'Origin': 'https://tv.blue.ch',
        'Referer': 'https://tv.blue.ch/',
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36',
    }
    response = req.get(f'https://services.sg102.prd.sctv.ch/streaming/liveTv/{channel_id}/dash_cas/0/42', headers=single_headers)
    data = response.json()
    return data['Address']

def get_pssh_from_mpd(url):
    mpd_headers = {
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36',
    }
    response = req.get(url, headers=mpd_headers)
    content_protections = BeautifulSoup(response.content, features="xml").findAll('ContentProtection')
    for cp in content_protections:
        if cp['schemeIdUri'] == 'urn:uuid:edef8ba9-79d6-4ace-a3c8-27dcd51d21ed':
            return cp.find('cenc:pssh').text
    return None

def get_channels(token):
    ch_headers = {
        'Accept': 'application/json, text/plain, */*',
        'Authorization': 'Bearer ' + token,
        'Content-Type': 'application/json; charset=utf-8',
        'Origin': 'https://tv.blue.ch',
        'Referer': 'https://tv.blue.ch/',
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36',
    }
    response = req.get('https://services.sg102.prd.sctv.ch/portfolio/tv/channels', headers=ch_headers)
    response.raise_for_status()
    channels = response.json()
    r = []
    for c in channels:
        if c['Visibility'] == 'Visible':
            r.append(c)
    return r

def login():
    print("logging in...", file=sys.stderr)
    # BlueTV uses token-based auth - token must be provided
    if token_param:
        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)
    else:
        print("Please provide token parameter: token=YOUR_TOKEN", file=sys.stderr)
        sys.exit(1)

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:
        if token_param:
            token = token_param
        else:
            return "error"

    if action == "channels":
        output = {'Channels': []}
        channels = get_channels(token)
        for chan in channels:
            channel = {
                'Name': chan['Title'],
                'Mode': "live",
                'SessionManifest': True,
                'ManifestScript': 'id=' + str(chan['Identifier']),
                'CdmType': "widevine",
                'UseCdm': True,
                'Cdm': 'id=' + str(chan['Identifier']),
                '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['Title'],
                'Mode': "live",
                'SessionManifest': True,
                'ManifestScript': 'id=' + str(chan['Identifier']),
                'CdmType': "widevine",
                'UseCdm': True,
                'Cdm': 'id=' + str(chan['Identifier']),
                '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:
            url = get_single(token, channel_id)
            pssh_data = get_pssh_from_mpd(url)
            
            output = {
                "Cdn": [],
                "ManifestUrl": url,
                "Headers": {
                    "Manifest": {
                        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36',
                    },
                    "Media": {
                        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36',
                    }
                },
                "Heartbeat": {
                    "Url": '',
                    "Params": '',
                    "PeriodMs": 5*60*1000
                }
            }
            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":
        try:
            lic_headers = {
                'Authorization': 'Bearer ' + token,
                'Origin': 'https://tv.blue.ch',
                'Referer': 'https://tv.blue.ch/',
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36',
                'Content-Type': 'application/x-www-form-urlencoded',
            }
            response = req.post('https://services.sg102.prd.sctv.ch/drm.cas/widevine/license/', headers=lic_headers, 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:
            url = get_single(token, channel_id)
            pssh_data = get_pssh_from_mpd(url)
            
            if pssh_data:
                keys = do_cdm_external(pssh_data, token)
                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()