import json
import xml.etree.ElementTree as ET
import datetime
import uuid
 
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding

from base64 import b64encode

def generate_auth_xml():
    with open("config.json", "r") as f:
        config_data = json.load(f)
    keyos_key_file = config_data.get("KEYOS_KEY_FILE")
    with open(keyos_key_file, 'rb') as key_file:
        private_key = serialization.load_pem_private_key(key_file.read(), password=None)
    data_el = ET.Element('Data')
    generation_time = datetime.datetime.now(datetime.timezone.utc)
    expiration_time = generation_time + datetime.timedelta(minutes=2)
    
    key_name = keyos_key_file.split('/')[-1].split('.')[0]
    ET.SubElement(data_el, 'GenerationTime').text = generation_time.strftime('%Y-%m-%d %H:%M:%S.%f')[0:23]
    ET.SubElement(data_el, 'ExpirationTime').text = expiration_time.strftime('%Y-%m-%d %H:%M:%S.%f')[0:23]
    ET.SubElement(data_el, 'UniqueId').text = str(uuid.uuid4())
    ET.SubElement(data_el, 'RSAPubKeyId').text = key_name
    
    #Widevine 
    wv_content_policy = ET.SubElement(data_el, 'WidevinePolicy', {'fl_CanPlay': 'true', 'fl_CanPersist': 'false'})
    wv_content_key_spec_el = ET.SubElement(data_el, 'WidevineContentKeySpec', {'TrackType': 'HD'})
    ET.SubElement(wv_content_key_spec_el, 'SecurityLevel').text = '1'
    
    #PlayReady 
    ET.SubElement(data_el, 'License', {'type': 'simple'})
    
    #Fairpplay
    ET.SubElement(data_el, 'FairPlayPolicy')
    signature = b64encode(private_key.sign(ET.tostring(data_el), padding.PKCS1v15(), hashes.SHA1()))
    root_el = ET.Element('KeyOSAuthenticationXML')
    root_el.append(data_el)
    ET.SubElement(root_el, 'Signature').text = signature.decode()
    auth_xml = (b64encode(ET.tostring(root_el)).decode('utf-8'))
    
    return (auth_xml)

if __name__ == '__main__':
    print(generate_auth_xml())