From fd2f38fe28f68669a41c5acbcb4ce25aa3e59a74 Mon Sep 17 00:00:00 2001 From: voldemort <5692900+yell0wsuit@users.noreply.github.com> Date: Tue, 22 Jul 2025 21:15:50 +0700 Subject: [PATCH] refactor CDM checks to improve configuration loading, user prompts, and file handling, consolidate CDM download logic into a single function --- custom_functions/prechecks/cdm_checks.py | 152 +++++++++++++---------- 1 file changed, 85 insertions(+), 67 deletions(-) diff --git a/custom_functions/prechecks/cdm_checks.py b/custom_functions/prechecks/cdm_checks.py index e4ef54b..048692f 100644 --- a/custom_functions/prechecks/cdm_checks.py +++ b/custom_functions/prechecks/cdm_checks.py @@ -1,88 +1,106 @@ +"""Module to check for and download CDM files.""" + import os +import sys import yaml import requests +CONFIG_PATH = os.path.join(os.getcwd(), "configs", "config.yaml") -def check_for_wvd_cdm(): - with open(f"{os.getcwd()}/configs/config.yaml", "r") as file: - config = yaml.safe_load(file) - if config["default_wv_cdm"] == "": - answer = " " - while answer[0].upper() != "Y" and answer[0].upper() != "N": - answer = input( - "No default Widevine CDM specified, would you like to download one from The CDM Project? (Y)es/(N)o: " - ) - if answer[0].upper() == "Y": - response = requests.get( - url="https://cdm-project.com/CDRM-Team/CDMs/raw/branch/main/Widevine/L3/public.wvd" - ) + +def load_config(): + """Load the config file.""" + with open(CONFIG_PATH, "r", encoding="utf-8") as file: + return yaml.safe_load(file) + + +def save_config(config): + """Save the config file.""" + with open(CONFIG_PATH, "w", encoding="utf-8") as file: + yaml.dump(config, file) + + +def prompt_yes_no(message): + """Prompt the user for a yes or no answer.""" + answer = " " + while answer[0].upper() not in ["Y", "N"]: + answer = input(message) + return answer[0].upper() == "Y" + + +def check_for_cdm(config_key, file_ext, download_url, cdm_dir, cdm_name): + """Check for a CDM file.""" + config = load_config() + cdm_value = config.get(config_key, "") + cdm_dir_path = os.path.join(os.getcwd(), "configs", "CDMs", cdm_dir) + os.makedirs(cdm_dir_path, exist_ok=True) + + if not cdm_value: + if prompt_yes_no( + f"No default {cdm_name} CDM specified, would you like to download one " + "from The CDM Project? (Y)es/(N)o: " + ): + response = requests.get(download_url, timeout=10) if response.status_code == 200: - with open(f"{os.getcwd()}/configs/CDMs/WV/public.wvd", "wb") as file: + file_path = os.path.join(cdm_dir_path, f"public.{file_ext}") + with open(file_path, "wb") as file: file.write(response.content) - config["default_wv_cdm"] = "public" - with open(f"{os.getcwd()}/configs/config.yaml", "w") as file: - yaml.dump(config, file) - print("Successfully downloaded Widevine CDM") + config[config_key] = "public" + save_config(config) + print(f"Successfully downloaded {cdm_name} CDM") else: - exit( - f"Download failed, please try again or place a .wvd file in {os.getcwd()}/configs/CDMs/WV and specify the name in {os.getcwd()}/configs/config.yaml" + sys.exit( + f"Download failed, please try again, or place a .{file_ext} file " + f"in {cdm_dir_path} and specify the name in {CONFIG_PATH}" ) - if answer[0].upper() == "N": - exit( - f"Place a .wvd file in {os.getcwd()}/configs/CDMs/WV and specify the name in {os.getcwd()}/configs/config.yaml" + else: + sys.exit( + f"Place a .{file_ext} file in {cdm_dir_path} and specify the name in {CONFIG_PATH}" ) else: - base_name = config["default_wv_cdm"] - if not base_name.endswith(".wvd"): - base_name += ".wvd" - if os.path.exists(f"{os.getcwd()}/configs/CDMs/WV/{base_name}"): + base_name = ( + cdm_value + if cdm_value.endswith(f".{file_ext}") + else f"{cdm_value}.{file_ext}" + ) + file_path = os.path.join(cdm_dir_path, base_name) + if os.path.exists(file_path): return - else: - exit( - f"Widevine CDM {base_name} does not exist in {os.getcwd()}/configs/CDMs/WV" - ) - - -def check_for_prd_cdm(): - with open(f"{os.getcwd()}/configs/config.yaml", "r") as file: - config = yaml.safe_load(file) - if config["default_pr_cdm"] == "": - answer = " " - while answer[0].upper() != "Y" and answer[0].upper() != "N": - answer = input( - "No default PlayReady CDM specified, would you like to download one from The CDM Project? (Y)es/(N)o: " - ) - if answer[0].upper() == "Y": - response = requests.get( - url="https://cdm-project.com/CDRM-Team/CDMs/raw/branch/main/Playready/SL2000/public.prd" - ) + # Prompt to download if file is missing, even if config has a value + if prompt_yes_no( + f"{cdm_name} CDM {base_name} does not exist in {cdm_dir_path}. Would you like to download it from The CDM Project? (Y)es/(N)o: " + ): + response = requests.get(download_url, timeout=10) if response.status_code == 200: - with open(f"{os.getcwd()}/configs/CDMs/PR/public.prd", "wb") as file: + with open(file_path, "wb") as file: file.write(response.content) - config["default_pr_cdm"] = "public" - with open(f"{os.getcwd()}/configs/config.yaml", "w") as file: - yaml.dump(config, file) - print("Successfully downloaded PlayReady CDM") + config[config_key] = base_name.replace(f".{file_ext}", "") + save_config(config) + print(f"Successfully downloaded {cdm_name} CDM") else: - exit( - f"Download failed, please try again or place a .prd file in {os.getcwd()}/configs/CDMs/PR and specify the name in {os.getcwd()}/configs/config.yaml" + sys.exit( + f"Download failed, please try again, or place a .{file_ext} file " + f"in {cdm_dir_path} and specify the name in {CONFIG_PATH}" ) - if answer[0].upper() == "N": - exit( - f"Place a .prd file in {os.getcwd()}/configs/CDMs/PR and specify the name in {os.getcwd()}/configs/config.yaml" - ) - else: - base_name = config["default_pr_cdm"] - if not base_name.endswith(".prd"): - base_name += ".prd" - if os.path.exists(f"{os.getcwd()}/configs/CDMs/PR/{base_name}"): - return else: - exit( - f"PlayReady CDM {base_name} does not exist in {os.getcwd()}/configs/CDMs/WV" + sys.exit( + f"Place a .{file_ext} file in {cdm_dir_path} and specify the name in {CONFIG_PATH}" ) def check_for_cdms(): - check_for_wvd_cdm() - check_for_prd_cdm() + """Check for CDM files.""" + check_for_cdm( + config_key="default_wv_cdm", + file_ext="wvd", + download_url="https://cdm-project.com/CDRM-Team/CDMs/raw/branch/main/Widevine/L3/public.wvd", + cdm_dir="WV", + cdm_name="Widevine", + ) + check_for_cdm( + config_key="default_pr_cdm", + file_ext="prd", + download_url="https://cdm-project.com/CDRM-Team/CDMs/raw/branch/main/Playready/SL2000/public.prd", + cdm_dir="PR", + cdm_name="PlayReady", + )