adi_generator/adi_builder.py

#!/usr/bin/env python3
''' ADI builder
author: Steven Kuterna steven.kuterna@devoteam.com
'''

import sys
import argparse
import logging
import os
import shutil
import yaml
from lxml import etree as ET
import hashlib
from datetime import datetime

loglevel= logging.WARNING

logger = logging.getLogger()
logger.setLevel(loglevel)
ch = logging.StreamHandler()
ch.setLevel(loglevel)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
ch.setFormatter(formatter)
logger.addHandler(ch)

def insertMetadata(tree, config):
    if config['Verb'] is None: config['Verb'] = ''
    if config['CreateDate'] is None:
        now = datetime.now()
        date = now.strftime("%Y-%m-%d")
        config['CreateDate'] = date
    try:
        attributes= { 'Asset_Class': 'package',
                        'Asset_ID': str(config['AssetId']),
                        'Asset_Name': str(config['AssetName']),
                        'Provider': str(config['Provider']),
                        'Provider_ID': str(config['ProviderId']),
                        'Product': 'MOD',
                        'Description': str(config['Description']),
                        'Creation_Date': str(config['CreateDate']),
                        'Verb': str(config['Verb']),
                        'Version_Major': str(config['VersionMajor']),
                        'Version_Minor': str(config['VersionMinor'])
        }
    except KeyError as e:
        logger.error('failed to add Metadata AMS attribute '+str(e))
        sys.exit(0)
    ET.SubElement(tree, 'AMS', attributes)
    ET.SubElement(tree, 'App_Data', {'App': 'MOD', 'Name': "Metadata_Spec_Version", 'Value': "CableLabsVOD1.1"})
    ET.SubElement(tree, 'App_Data', {'App': 'MOD', 'Name': "Provider_Content_Tier", 'Value': "VU"})


def createAMSattributes(asset_type, metadata):
    if metadata['Verb'] is None: metadata['Verb'] = ''
    try:
        ams_attributes= { 'Asset_Class': asset_type,
                        'Asset_ID': asset_type.upper() + str(metadata['AssetId']),
                        'Asset_Name': str(metadata['AssetName'] + ' ' + asset_type),
                        'Provider': str(metadata['Provider']),
                        'Provider_ID': str(metadata['ProviderId']),
                        'Product': 'MOD',
                        'Description': str(metadata['Description'] + ' ' + asset_type),
                        'Creation_Date': str(metadata['CreateDate']),
                        'Verb': str(metadata['Verb']),
                        'Version_Major': str(metadata['VersionMajor']),
                        'Version_Minor': str(metadata['VersionMinor'])
        }
    except KeyError as e:
        logger.error('failed to add Metadata AMS attribute '+str(e))
        sys.exit(0)
    return ams_attributes


def addAsset(tree, asset_type, asset, metadata, output):
    if not 'Content' in asset:
        for subasset in asset:
            addAsset(tree, asset_type, asset[subasset], metadata, output)
        return
    logger.info(f'insert asset: {asset}')

    ams_attributes = createAMSattributes(asset_type, metadata)
    assetTree = ET.SubElement(tree, 'Asset')
    assetMetadata = ET.SubElement(assetTree, 'Metadata')
    ET.SubElement(assetMetadata, 'AMS', ams_attributes)
    ET.SubElement(assetMetadata, 'App_Data', {'App': 'MOD', 'Name': 'Type', 'Value': asset_type})

    for element in asset:
        if type(asset[element]) is dict:
            for subelement in asset[element]:
                if type(asset[element][subelement]) is dict:
                    for subsubelemment in asset[element][subelement]:
                        ET.SubElement(assetMetadata, 'App_Data', {'App': 'MOD', 'Name': subsubelemment, 'Value': str(asset[element][subelement][subsubelemment])})
                else:
                    if asset[element][subelement] is None:
                        asset[element][subelement] = ''
                    ET.SubElement(assetMetadata, 'App_Data', {'App': 'MOD', 'Name': element, 'Value': str(asset[element][subelement])})
        else:
            # is it content? add it under the Asset level
            if str(element).upper() == 'CONTENT_CHECKSUM' or str(element).upper() == 'CONTENT_FILESIZE':
                continue
            if str(element).upper() == 'CONTENT':
                if not os.path.exists(args.input):
                    logger.error(f"content file {asset[element]} is not found, exiting...")
                    sys.exit(0)
                ET.SubElement(assetTree, 'Content', {'Value': asset[element]})
                checksum = hashlib.md5(open(asset[element],'rb').read()).hexdigest()
                filesize = os.path.getsize(asset[element])
                ET.SubElement(assetMetadata, 'App_Data', {'App': 'MOD', 'Name': 'Content_CheckSum', 'Value': str(checksum)})
                ET.SubElement(assetMetadata, 'App_Data', {'App': 'MOD', 'Name': 'Content_Filesize', 'Value': str(filesize)})
                shutil.copyfile(asset[element], f'{output}/{asset[element]}')
                continue
            # no content, add it under the Metadata level
            if asset[element] is None:
                asset[element] = ''
            ET.SubElement(assetMetadata, 'App_Data', {'App': 'MOD', 'Name': element, 'Value': str(asset[element])})


def addTitle(tree, title, metadata):
    logger.info(f'insert title: {title}')
    ams_attributes = createAMSattributes('title', metadata)
    titleTree = ET.SubElement(tree, 'Metadata')
    ET.SubElement(titleTree, 'AMS', ams_attributes)
    ET.SubElement(titleTree, 'App_Data', {'App': 'MOD', 'Name': 'Type', 'Value': 'title'})

    for element in title:
        # check if there are multiple elements available
        if type(title[element]) == dict:
            for subelement in title[element]:
                # check if it contains multiple attributes like Language and a Value
                if type(title[element][subelement]) == dict:
                    attributes = {'App': 'MOD', 'Name': element}
                    # if the attribute is the same as the element change it to 'Value'
                    for subsubelement in title[element][subelement]:
                        if subsubelement == element:
                            attributes.update({'Value': title[element][subelement][subsubelement]})
                        else:
                            # just add it to the attributes list
                            attributes.update({subsubelement: title[element][subelement][subsubelement]})
                    # add the multiattribute element to the tree
                    ET.SubElement(titleTree, 'App_Data', attributes)
                    continue
                # single attribute, add to the tree
                if title[element][subelement] is None:
                    title[element][subelement] = ''
                ET.SubElement(titleTree, 'App_Data', {'App': 'MOD', 'Name': element, 'Value': title[element][subelement]})
        else:
            # single element, add to the tree
            if title[element] is None:
                title[element] = ''
            ET.SubElement(titleTree, 'App_Data', {'App': 'MOD', 'Name': element, 'Value': str(title[element])})


def insertAssets(tree, metadata, assets, output):
    for asset in assets:
        if asset.upper() == 'TITLE':
            addTitle(tree, assets[asset], metadata)
        else:
            addAsset(tree, asset, assets[asset], metadata, output)
    return


if __name__ == "__main__" :
    argparser = argparse.ArgumentParser(description='Creates ADI from yaml...')
    argparser.add_argument('--input', required=False, default="adi.yaml", help="yaml input to build adi.xml")
    argparser.add_argument('--output', required=True, help="output directory for xml and content files")
    args = argparser.parse_args()

    logger.debug(f"generating adi based on {args.input}")
    if not os.path.exists(args.input):
        logger.error(f"file {args.input} is not found, exiting...")
        sys.exit(0)
    with open(args.input, 'r') as ymlfile:
        cfg = yaml.load(ymlfile, Loader=yaml.FullLoader)

    try:
        os.mkdir(args.output)
    except FileExistsError:
        logger.warning(f'directory {args.output} already exists!')
    root = ET.Element('ADI')
    metadataTree = ET.SubElement(root, 'Metadata')
    insertMetadata(metadataTree, cfg['metadata'])
    assetTree = ET.SubElement(root, 'Asset')
    insertAssets(assetTree, cfg['metadata'], cfg['asset'], args.output)
    tree = ET.ElementTree(root)

    try:
        tree.write(f'{args.output}/adi.xml', encoding='utf-8', xml_declaration=True, pretty_print=True, doctype='<!DOCTYPE ADI SYSTEM "ADI.DTD">')
    except Exception as e:
        print(str(e))