CDRM-Keys/CDRM_Modules/Startup_Checks.py

# Import dependencies
import os
import time
import requests
import yaml
import sqlite3
from colorama import Fore
import sys
import json


# Define a function to check if script is running from a virtual environment
def check_is_venv():
    if hasattr(sys, 'real_prefix') or (hasattr(sys, 'base_prefix') and sys.base_prefix != sys.prefix):
        print(f'{Fore.GREEN}Virtual Environment [✓]')
        return
    else:
        print(f'{Fore.RED}Not running Virtual Environment!\n\nPlease run CDRM-Keys in a Python Virtual Environment')
        exit(1)


# Define a check to see if config.yaml exists
def check_if_config_exists():
    if os.path.isfile(f'{os.getcwd()}/Config.yaml'):
        print(f'{Fore.GREEN}Config.yaml [✓]')
        return
    else:
        initial_config_data = {
            'First_Run': 'True',
            'Remote_CDM_API_FQDN': 'http://127.0.0.1:5000'
        }
        create_input = input(f'\n{Fore.YELLOW}Config.yaml not found, would you like to create a Config.yaml? [(Y)es/(N)o]: ')
        if create_input:
            if create_input[0].lower() == 'y':
                try:
                    with open(f'{os.getcwd()}/Config.yaml', 'w') as file:
                        yaml.dump(initial_config_data, file)
                    print(f'{Fore.GREEN}Config.yaml created!')
                except Exception as error:
                    print(f'{Fore.RED}An error occurred!\n\n{error}')
                    exit(1)
                return
            else:
                print(f'{Fore.YELLOW}Exiting.')
                exit(1)


# Define a function to see if keys.db exists
def check_if_keys_db_exists():
    if os.path.isfile(f'{os.getcwd()}/Keys.db'):
        print(f'{Fore.GREEN}Keys.db [✓]')
        return
    else:
        try:
            create_input = input(f'\n{Fore.YELLOW}Keys.db not found, would you like to create a Keys.db? [(Y)es/(N)o]: ')
            if create_input:
                if create_input[0].lower() == 'y':
                    keys_db_connection = sqlite3.connect(f'{os.getcwd()}/Keys.db')
                    keys_db_cursor = keys_db_connection.cursor()
                    keys_db_cursor.execute('''
                    CREATE TABLE IF NOT EXISTS keys (
                        PSSH TEXT,
                        MPD TEXT,
                        KID TEXT PRIMARY KEY,
                        KEY TEXT,
                        License_URL TEXT,
                        Headers TEXT,
                        Cookies TEXT,
                        Data TEXT
                    )
                    ''')
                    keys_db_connection.commit()
                    keys_db_connection.close()
                    print(f'{Fore.GREEN}Keys.db created!')
                    return
                else:
                    print(f'{Fore.YELLOW}Exiting.')
                    exit(1)
        except Exception as error:
            print(f'{Fore.RED}An error occurred!\n\n{error}')
            exit(1)


# Define a function to check if License_cURL.py exists
def check_if_license_curl_exists():
    if os.path.isfile(f'{os.getcwd()}/License_cURL.py'):
        print(f'{Fore.GREEN}License_cURL.py [✓]')
        return
    else:
        try:
            create_input = input(f'\n{Fore.YELLOW}License_cURL.py not found, would you like to create a License_cURL.py? [(Y)es/(N)o]: ')
            if create_input:
                if create_input[0].lower() == 'y':
                    headers = {
                        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:130.0) Gecko/20100101 Firefox/130.0',
                        'Accept': '*/*',
                        'Accept-Language': 'en-US,en;q=0.5',
                        # 'Accept-Encoding': 'gzip, deflate, br, zstd',
                        'Origin': 'https://bitmovin.com',
                        'DNT': '1',
                        'Sec-GPC': '1',
                        'Connection': 'keep-alive',
                        'Referer': 'https://bitmovin.com/',
                        'Sec-Fetch-Dest': 'empty',
                        'Sec-Fetch-Mode': 'cors',
                        'Sec-Fetch-Site': 'cross-site',
                        'Content-Type': 'application/x-www-form-urlencoded',
                    }
                    with open(f'{os.getcwd()}/License_cURL.py', 'w') as file:
                        file.write("# This file contains a sample license URL curl from BitMovin (https://bitmovin.com/demos/drm)\n")
                        file.write("headers = ")
                        json.dump(headers, file, indent=4)
                        file.write("\n")
                    print(f'{Fore.GREEN}License_cURL.py created!')
                    return
                else:
                    print(f'{Fore.YELLOW}Exiting.')
                    exit(1)
        except Exception as error:
            print(f'{Fore.RED}An error occurred!\n\n{error}')
            exit(1)


# Define a function to check if WVDs folder exists
def check_if_wvd_folder_exists():
    if os.path.isdir(f'{os.getcwd()}/WVDs'):
        print(f'{Fore.GREEN}WVDs Folder [✓]')
        return
    else:
        create_input = input(f'\n{Fore.YELLOW}WVDs folder not found,  would you like to create a WVDs folder? [(Y)es/(N)o]: ')
        try:
            if create_input:
                if create_input[0].lower() == 'y':
                    os.mkdir(f'{os.getcwd()}/WVDs')
                    print(f'{Fore.GREEN}WVDs folder created!')
                    return
        except Exception as error:
            print(f'{Fore.RED}An error occurred!\n\n{error}')
            exit(1)


# Define a function to check for CDMs
def check_for_wvds(startup: bool = False):
    try:
        cmds_exist = False
        remote_cdms_exist = False
        files_in_wvds_folder = os.listdir(f'{os.getcwd()}/WVDs')
        wvds = [file for file in files_in_wvds_folder if file.endswith('.wvd')]
        with open('Config.yaml', 'r') as ymlfile:
            config = yaml.safe_load(ymlfile)
        if wvds:
            if startup:
                print(f'{Fore.GREEN}WVDs [✓]')
            cmds_exist = True
        if 'Remote_CDMs' in config:
            if startup:
                print(f'{Fore.GREEN}Remote CDMs [✓]')
            remote_cdms_exist = True
        elif 'Remote_CDMs' not in config and startup:
            if config['First_Run'] == 'True':
                create_input = input(
                    f'\n{Fore.YELLOW}No remote CDM found,  would you like to use CDRM-Project\'s? [(Y)es/(N)o]: ')
                if create_input:
                    if create_input[0].lower() == 'y':
                        try:
                            cdrm_project_remote_cdm_info = requests.get(
                                url='https://remote-cdm.cdrm-project.com/remote_cdm'
                            ).json()
                            device_details = {
                                'Remote_CDMs': {
                                    'CDRM_Project_CDM_API': {
                                        'name': cdrm_project_remote_cdm_info['name'],
                                        'host': cdrm_project_remote_cdm_info['host'],
                                        'device_type': cdrm_project_remote_cdm_info['device_type'],
                                        'security_level': cdrm_project_remote_cdm_info['security_level'],
                                        'system_id': cdrm_project_remote_cdm_info['system_id'],
                                        'secret': cdrm_project_remote_cdm_info['secret']
                                    }
                                }
                            }
                            with open(f'{os.getcwd()}/Config.yaml', 'r') as ymlfile:
                                config = yaml.safe_load(ymlfile)
                            config.update(device_details)
                            with open(f'{os.getcwd()}/Config.yaml', 'w') as ymlfile:
                                yaml.safe_dump(config, ymlfile)
                            print(f'{Fore.GREEN}CDRM-Project RemoteCDM added')
                            remote_cdms_exist = True
                        except Exception as error:
                            print(f'{Fore.RED}An error occurred!\n\n{error}')

        if not cmds_exist and not remote_cdms_exist:
            print(f'{Fore.YELLOW}No WVDs available! Exiting.')
            exit(1)
        return cmds_exist, remote_cdms_exist
    except Exception as error:
        print(f'{Fore.RED}An error occurred!\n\n{error}')
        exit(1)


# Define a function to run all startup checks
def run_startup_checks(startup: bool = False):
    check_is_venv()
    check_if_license_curl_exists()
    check_if_config_exists()
    check_if_keys_db_exists()
    check_if_wvd_folder_exists()
    check_for_wvds(startup=startup)
    time.sleep(1)
    with open(f'{os.getcwd()}/Config.yaml', 'r') as ymlfile:
        config = yaml.safe_load(ymlfile)
    config['First_Run'] = 'False'
    with open(f'{os.getcwd()}/Config.yaml', 'w') as ymlfile:
        yaml.safe_dump(config, ymlfile)
    if os.name == 'nt':  # For Windows
        os.system('cls')
    else:  # For Unix-based systems (Linux, macOS)
        os.system('clear')
    return