#!/usr/bin/python3
import sys
import os
import base64
import json
import datetime
import pytz
import hmac
import hashlib
import random
import string
import urllib.parse

# Add parent directory to path for o11 import
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import o11

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

# Parse command line parameters
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')

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')

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

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

# Configuration
WVD_PATH = './WVD.wvd'
authFile = '/BeinConnect_auth.json'
SCRIPT_DIR = os.path.abspath(os.path.dirname(__file__))

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

PLAYER_NAME = 'bitmovin-web'
PLAYER_VERSION = '10.0'
SHA256_KEY = 'mkYK;Qy%5#gR'
IDENTITY_KEY = 'c5g4dvCzHbEUSNtAF97VhJuKPqXyWLQG'
LICENCE_IDENTITY_KEY = 'Euj8p2CBxnkVG3JhZrNbP9RsALWMYaFH'

token = ''

def get_auth():
    try:
        return json.load(open(SCRIPT_DIR + authFile))
    except:
        return None

def save_auth(auth_data):
    json.dump(auth_data, open(SCRIPT_DIR + authFile, 'w'), indent=2)

def generate_boundary():
    prefix = "----WebKitFormBoundary"
    random_str = ''.join(random.choices(string.ascii_letters + string.digits, k=16))
    return prefix + random_str

def do_refresh(tok):
    boundary = generate_boundary()
    refresh_headers = {
        'accept': 'application/json, text/plain, */*',
        'content-type': 'multipart/form-data; boundary=' + boundary,
        'origin': 'https://connect.bein.com',
        'referer': 'https://connect.bein.com/',
        'user-agent': USER_AGENT,
        'x-an-webservice-customerauthtoken': tok,
        'x-an-webservice-identitykey': IDENTITY_KEY,
    }
    data = f'--{boundary}--\r\n'
    response = req.post('https://proxies.bein-mena-production.eu-west-2.tuc.red/proxy/accountDetails', headers=refresh_headers, data=data)
    return response.json()['result']['newAuthToken']

def login():
    global token
    print("logging in...", file=sys.stderr)
    
    auth = get_auth()
    if not auth or 'token' not in auth:
        print("No auth found. Please add BeinConnect token to auth file.", file=sys.stderr)
        print("Login to https://connect.bein.com/ and copy 'token-mena' from Local Storage", file=sys.stderr)
        save_auth({'token': 'ENTER token-mena HERE'})
        sys.exit(1)
    
    try:
        tok = auth['token']
        tok = do_refresh(tok)
        save_auth({'token': tok})
        token = tok
        print("logged in successfully", file=sys.stderr)
        return tok
    except Exception as e:
        print(f"Login failed: {e}", file=sys.stderr)
        sys.exit(1)

def get_token():
    global token
    auth = get_auth()
    if auth and 'token' in auth:
        token = auth['token']
        return token
    return login()

def calculate_checksum(channel_id):
    message = f'{token or ""}{channel_id}{PLAYER_VERSION}'
    hmac_digest = hmac.new(SHA256_KEY.encode('utf-8'), message.encode('utf-8'), hashlib.sha256).digest()
    checksum = base64.b64encode(hmac_digest).decode('utf-8')
    checksum_uri = f'?idChannel={channel_id}&playerName={PLAYER_NAME}&playerVersion={PLAYER_VERSION}&checksum={urllib.parse.quote(checksum)}'
    return checksum, checksum_uri

def find_wv_pssh_offsets(raw):
    offsets = []
    offset = 0
    while True:
        offset = raw.find(b'pssh', offset)
        if offset == -1:
            break
        size = int.from_bytes(raw[offset-4:offset], byteorder='big')
        pssh_offset = offset - 4
        offsets.append(raw[pssh_offset:pssh_offset+size])
        offset += size
    return offsets

def to_pssh(content):
    wv_offsets = find_wv_pssh_offsets(content)
    return [base64.b64encode(wv_offset).decode() for wv_offset in wv_offsets]

def get_pssh_from_mpd(url):
    try:
        response = req.get(url, headers={'accept': '*/*', 'user-agent': USER_AGENT})
        location = response.url.split('?')[0]
        soup = BeautifulSoup(response.content, features="xml")
        loc_parts = location.split('/')
        loc_parts.pop()
        
        psshs = []
        adp_sets = soup.find_all('AdaptationSet')
        for a_set in adp_sets:
            try:
                seg_temp = a_set.find('SegmentTemplate')
                start_number = int(seg_temp['startNumber'])
                seg_offset = int(seg_temp.find('S')['r']) - 10
                media = seg_temp['media']
                reps = a_set.find_all('Representation')
                for r in reps:
                    rep_id = r['id']
                    media_url = '/'.join(loc_parts) + '/' + media.replace('$Number$', str(start_number + seg_offset)).replace('$RepresentationID$', rep_id)
                    init_response = req.get(media_url, headers={'accept': '*/*', 'user-agent': USER_AGENT})
                    rep_psshs = to_pssh(init_response.content)
                    for rp in rep_psshs:
                        if rp not in psshs:
                            psshs.append(rp)
            except:
                pass
        return psshs
    except:
        return None

def get_single(channel_id, checksum):
    boundary = generate_boundary()
    single_headers = {
        'accept': 'application/json, text/plain, */*',
        'content-type': 'multipart/form-data; boundary=' + boundary,
        'origin': 'https://connect.bein.com',
        'referer': 'https://connect.bein.com/',
        'user-agent': USER_AGENT,
        'x-an-webservice-customerauthtoken': token,
        'x-an-webservice-identitykey': IDENTITY_KEY,
    }
    
    data = f'--{boundary}\r\nContent-Disposition: form-data; name="checksum"\r\n\r\n{checksum}\r\n'
    data += f'--{boundary}\r\nContent-Disposition: form-data; name="idChannel"\r\n\r\n{channel_id}\r\n'
    data += f'--{boundary}\r\nContent-Disposition: form-data; name="playerName"\r\n\r\nbitmovin-web\r\n'
    data += f'--{boundary}\r\nContent-Disposition: form-data; name="playerVersion"\r\n\r\n10.0\r\n'
    data += f'--{boundary}--\r\n'
    
    response = req.post('https://proxies.bein-mena-production.eu-west-2.tuc.red/proxy/channelStream', headers=single_headers, data=data)
    
    try:
        data = response.json()
        return data['result']['url'], data['result']['newAuthToken']
    except:
        return None, None

def do_cdm_internal(challenge_b64, lic_params):
    lic_headers = {
        'accept': '*/*',
        'content-type': 'application/octet-stream',
        'origin': 'https://connect.bein.com',
        'referer': 'https://connect.bein.com/',
        'user-agent': USER_AGENT,
        'x-an-webservice-customerauthtoken': token,
        'x-an-webservice-identitykey': LICENCE_IDENTITY_KEY,
    }
    data = base64.b64encode(base64.b64decode(challenge_b64)).decode('utf-8')
    response = req.post('https://proxies.bein-mena-production.eu-west-2.tuc.red/proxy/getLiveLicense' + lic_params, headers=lic_headers, data=data)
    
    try:
        soup = BeautifulSoup(response.content, features='xml')
        license_b64 = soup.find('LICENSE').text
        return license_b64
    except:
        return None

def do_cdm_external(pssh_b64, lic_params):
    try:
        pssh_obj = PSSH(pssh_b64)
        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 = {
            'accept': '*/*',
            'content-type': 'application/octet-stream',
            'origin': 'https://connect.bein.com',
            'referer': 'https://connect.bein.com/',
            'user-agent': USER_AGENT,
            'x-an-webservice-customerauthtoken': token,
            'x-an-webservice-identitykey': LICENCE_IDENTITY_KEY,
        }
        data = base64.b64encode(challenge_data).decode('utf-8')
        
        licence = req.post('https://proxies.bein-mena-production.eu-west-2.tuc.red/proxy/getLiveLicense' + lic_params, headers=lic_headers, data=data)
        
        soup = BeautifulSoup(licence.content, features='xml')
        license_b64 = soup.find('LICENSE').text
        cdm_obj.parse_license(session_id, base64.b64decode(license_b64))
        
        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'CDM external failed: {e}', file=sys.stderr)
        return None

def do_action():
    global token
    get_token()
    
    if action == "login":
        login()
        sys.exit()
    
    if action == "channels":
        output = {'Channels': []}
        
        boundary = generate_boundary()
        ch_headers = {
            'accept': 'application/json, text/plain, */*',
            'content-type': 'multipart/form-data; boundary=' + boundary,
            'origin': 'https://connect.bein.com',
            'referer': 'https://connect.bein.com/',
            'user-agent': USER_AGENT,
            'x-an-webservice-identitykey': IDENTITY_KEY,
        }
        data = f'--{boundary}\r\nContent-Disposition: form-data; name="languageId"\r\n\r\neng\r\n--{boundary}--\r\n'
        
        response = req.post('https://proxies.bein-mena-production.eu-west-2.tuc.red/proxy/listChannels', headers=ch_headers, data=data)
        
        try:
            data = response.json()
            for ch in data['result'].get('channels', []):
                channel = {
                    'Name': ch.get('name', 'Unknown'),
                    'Mode': 'live',
                    'SessionManifest': True,
                    'ManifestScript': f"id={ch.get('idChannel', '')}",
                    'CdmType': 'widevine',
                    'UseCdm': True,
                    'Cdm': f"id={ch.get('idChannel', '')}",
                    'Video': 'best',
                    'OnDemand': True,
                    'SpeedUp': True,
                }
                output['Channels'].append(channel)
            
            print(json.dumps(output, indent=2))
        except Exception as e:
            print(f"Error: {e}", file=sys.stderr)
            return "error"
    
    elif action == "events":
        output = {'Events': []}
        print(json.dumps(output, indent=2))
    
    elif action == "heartbeat":
        sys.exit()
    
    elif action == "manifest":
        try:
            channel_id = id
            checksum, checksum_uri = calculate_checksum(channel_id)
            video_url, new_token = get_single(channel_id, checksum)
            
            if new_token:
                token = new_token
                save_auth({'token': token})
            
            if not video_url:
                return "error"
            
            video_url = video_url + '&' + checksum_uri
            
            output = {
                "Cdn": [{"Name": "default", "ManifestUrl": video_url}],
                "ManifestUrl": video_url,
                "Headers": {"Manifest": {'User-Agent': USER_AGENT}, "Media": {'User-Agent': USER_AGENT}},
                "Heartbeat": {"Url": '', "Params": '', "PeriodMs": 5*60*1000},
                "LicenseParams": checksum_uri
            }
            print(json.dumps(output))
        except Exception as e:
            print(f"Error: {e}", file=sys.stderr)
            return "error"
    
    elif action == "cdm" and cdm == "internal":
        try:
            channel_id = id
            checksum, checksum_uri = calculate_checksum(channel_id)
            result = do_cdm_internal(challenge, checksum_uri)
            if result:
                print(result)
            else:
                return "error"
        except Exception as e:
            print(f"Error: {e}", file=sys.stderr)
            return "error"
    
    elif action == "cdm" and cdm == "external":
        try:
            channel_id = id
            checksum, checksum_uri = calculate_checksum(channel_id)
            
            if pssh:
                keys = do_cdm_external(pssh, checksum_uri)
            else:
                video_url, new_token = get_single(channel_id, checksum)
                if new_token:
                    token = new_token
                    save_auth({'token': token})
                if video_url:
                    video_url = video_url + '&' + checksum_uri
                    psshs = get_pssh_from_mpd(video_url)
                    keys = []
                    if psshs:
                        for p in psshs:
                            ks = do_cdm_external(p, checksum_uri)
                            if ks:
                                keys.extend(ks)
                    keys = list(set(keys))
                else:
                    keys = None
            
            if keys:
                for key in keys:
                    print(key)
            else:
                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()