CDRM-Project/custom_functions/decrypt/api_decrypt.py

"""Module to decrypt the license using the API."""

import base64
import ast
import glob
import json
import os
from urllib.parse import urlparse
import binascii

import requests
from requests.exceptions import Timeout, RequestException
import yaml

from pywidevine.cdm import Cdm as widevineCdm
from pywidevine.device import Device as widevineDevice
from pywidevine.pssh import PSSH as widevinePSSH
from pyplayready.cdm import Cdm as playreadyCdm
from pyplayready.device import Device as playreadyDevice
from pyplayready.system.pssh import PSSH as playreadyPSSH

from custom_functions.database.unified_db_ops import cache_to_db


def find_license_key(data, keywords=None):
    """Find the license key in the data."""
    if keywords is None:
        keywords = [
            "license",
            "licenseData",
            "widevine2License",
        ]  # Default list of keywords to search for

    # If the data is a dictionary, check each key
    if isinstance(data, dict):
        for key, value in data.items():
            if any(
                keyword in key.lower() for keyword in keywords
            ):  # Check if any keyword is in the key (case-insensitive)
                return value.replace("-", "+").replace(
                    "_", "/"
                )  # Return the value immediately when found
            # Recursively check if the value is a dictionary or list
            if isinstance(value, (dict, list)):
                result = find_license_key(value, keywords)  # Recursively search
                if result:  # If a value is found, return it
                    return result.replace("-", "+").replace("_", "/")

    # If the data is a list, iterate through each item
    elif isinstance(data, list):
        for item in data:
            result = find_license_key(item, keywords)  # Recursively search
            if result:  # If a value is found, return it
                return result.replace("-", "+").replace("_", "/")

    return None  # Return None if no matching key is found


def find_license_challenge(data, keywords=None, new_value=None):
    """Find the license challenge in the data."""
    if keywords is None:
        keywords = [
            "license",
            "licenseData",
            "widevine2License",
            "licenseRequest",
        ]  # Default list of keywords to search for

    # If the data is a dictionary, check each key
    if isinstance(data, dict):
        for key, value in data.items():
            if any(
                keyword in key.lower() for keyword in keywords
            ):  # Check if any keyword is in the key (case-insensitive)
                data[key] = new_value  # Modify the value in-place
            # Recursively check if the value is a dictionary or list
            elif isinstance(value, (dict, list)):
                find_license_challenge(
                    value, keywords, new_value
                )  # Recursively modify in place

    # If the data is a list, iterate through each item
    elif isinstance(data, list):
        for i, item in enumerate(data):
            result = find_license_challenge(
                item, keywords, new_value
            )  # Recursively modify in place

    return data  # Return the modified original data (no new structure is created)


def is_base64(string):
    """Check if the string is base64 encoded."""
    try:
        # Try decoding the string
        decoded_data = base64.b64decode(string)
        # Check if the decoded data, when re-encoded, matches the original string
        return base64.b64encode(decoded_data).decode("utf-8") == string
    except (binascii.Error, TypeError):
        # If decoding or encoding fails, it's not Base64
        return False


def is_url_and_split(input_str):
    """Check if the string is a URL and split it into protocol and FQDN."""
    parsed = urlparse(input_str)

    # Check if it's a valid URL with scheme and netloc
    if parsed.scheme and parsed.netloc:
        protocol = parsed.scheme
        fqdn = parsed.netloc
        return True, protocol, fqdn
    return False, None, None


def load_device(device_type, device, username, config):
    """Load the appropriate device file for PlayReady or Widevine."""
    if device_type == "PR":
        ext, config_key, class_loader = ".prd", "default_pr_cdm", playreadyDevice.load
        base_dir = "PR"
    else:
        ext, config_key, class_loader = ".wvd", "default_wv_cdm", widevineDevice.load
        base_dir = "WV"

    if device == "public":
        base_name = config[config_key]
        if not base_name.endswith(ext):
            base_name += ext
        search_path = f"{os.getcwd()}/configs/CDMs/{base_dir}/{base_name}"
    else:
        base_name = device
        if not base_name.endswith(ext):
            base_name += ext
        search_path = f"{os.getcwd()}/configs/CDMs/{username}/{base_dir}/{base_name}"

    files = glob.glob(search_path)
    if not files:
        return None, f"No {ext} file found for device '{device}'"
    try:
        return class_loader(files[0]), None
    except (IOError, OSError) as e:
        return None, f"Failed to read device file: {e}"
    except (ValueError, TypeError, AttributeError) as e:
        return None, f"Failed to parse device file: {e}"


def prepare_request_data(headers, cookies, json_data, challenge, is_widevine):
    """Prepare headers, cookies, and json_data for the license request."""
    try:
        format_headers = ast.literal_eval(headers) if headers else None
    except (ValueError, SyntaxError) as e:
        raise ValueError(f"Invalid headers format: {e}") from e

    try:
        format_cookies = ast.literal_eval(cookies) if cookies else None
    except (ValueError, SyntaxError) as e:
        raise ValueError(f"Invalid cookies format: {e}") from e

    format_json_data = None
    if json_data and not is_base64(json_data):
        try:
            format_json_data = ast.literal_eval(json_data)
            if is_widevine:
                format_json_data = find_license_challenge(
                    data=format_json_data,
                    new_value=base64.b64encode(challenge).decode(),
                )
        except (ValueError, SyntaxError) as e:
            raise ValueError(f"Invalid json_data format: {e}") from e
        except (TypeError, AttributeError) as e:
            raise ValueError(f"Error processing json_data: {e}") from e

    return format_headers, format_cookies, format_json_data


def send_license_request(license_url, headers, cookies, json_data, challenge, proxies):
    """Send the license request and return the response."""
    try:
        response = requests.post(
            url=license_url,
            headers=headers,
            proxies=proxies,
            cookies=cookies,
            json=json_data if json_data is not None else None,
            data=challenge if json_data is None else None,
            timeout=10,
        )
        return response, None
    except ConnectionError as error:
        return None, f"Connection error: {error}"
    except Timeout as error:
        return None, f"Request timeout: {error}"
    except RequestException as error:
        return None, f"Request error: {error}"


def extract_and_cache_keys(
    cdm,
    session_id,
    cache_to_db,
    pssh,
    license_url,
    headers,
    cookies,
    challenge,
    json_data,
    is_widevine,
):
    """Extract keys from the session and cache them."""
    returned_keys = ""
    try:
        keys = list(cdm.get_keys(session_id))
        for index, key in enumerate(keys):
            # Widevine: key.type, PlayReady: key.key_type
            key_type = getattr(key, "type", getattr(key, "key_type", None))
            kid = getattr(key, "kid", getattr(key, "key_id", None))
            if key_type != "SIGNING" and kid is not None:
                cache_to_db(
                    pssh=pssh,
                    license_url=license_url,
                    headers=headers,
                    cookies=cookies,
                    data=challenge if json_data is None else json_data,
                    kid=kid.hex,
                    key=key.key.hex(),
                )
                if index != len(keys) - 1:
                    returned_keys += f"{kid.hex}:{key.key.hex()}\n"
                else:
                    returned_keys += f"{kid.hex}:{key.key.hex()}"
        return returned_keys, None
    except AttributeError as error:
        return None, f"Error accessing CDM keys: {error}"
    except (TypeError, ValueError) as error:
        return None, f"Error processing keys: {error}"


def api_decrypt(
    pssh: str = "",
    license_url: str = "",
    proxy: str = "",
    headers: str = "",
    cookies: str = "",
    json_data: str = "",
    device: str = "public",
    username: str = "",
):
    """Decrypt the license using the API."""
    print(f"Using device {device} for user {username}")
    with open(f"{os.getcwd()}/configs/config.yaml", "r", encoding="utf-8") as file:
        config = yaml.safe_load(file)

    if pssh == "":
        return {"status": "error", "message": "No PSSH provided"}

    # Detect PlayReady or Widevine
    try:
        is_pr = "</WRMHEADER>".encode("utf-16-le") in base64.b64decode(pssh)
    except (binascii.Error, TypeError) as error:
        return {
            "status": "error",
            "message": f"An error occurred processing PSSH\n\n{error}",
        }

    device_type = "PR" if is_pr else "WV"
    cdm_class = playreadyCdm if is_pr else widevineCdm
    pssh_class = playreadyPSSH if is_pr else widevinePSSH

    # Load device
    device_obj, device_err = load_device(device_type, device, username, config)
    if device_obj is None:
        return {"status": "error", "message": device_err}

    # Create CDM
    try:
        cdm = cdm_class.from_device(device_obj)
    except (IOError, ValueError, AttributeError) as error:
        return {
            "status": "error",
            "message": f"An error occurred loading {device_type} CDM\n\n{error}",
        }

    # Open session
    try:
        session_id = cdm.open()
    except (IOError, ValueError, AttributeError) as error:
        return {
            "status": "error",
            "message": f"An error occurred opening a CDM session\n\n{error}",
        }

    # Parse PSSH and get challenge
    try:
        pssh_obj = pssh_class(pssh)
        if is_pr:
            challenge = cdm.get_license_challenge(session_id, pssh_obj.wrm_headers[0])
        else:
            challenge = cdm.get_license_challenge(session_id, pssh_obj)
    except (ValueError, AttributeError, IndexError) as error:
        return {
            "status": "error",
            "message": f"An error occurred getting license challenge\n\n{error}",
        }

    # Prepare request data
    try:
        format_headers, format_cookies, format_json_data = prepare_request_data(
            headers, cookies, json_data, challenge, is_widevine=(not is_pr)
        )
    except (ValueError, SyntaxError) as error:
        return {
            "status": "error",
            "message": f"An error occurred preparing request data\n\n{error}",
        }

    # Prepare proxies
    proxies = ""
    if proxy != "":
        is_url, protocol, fqdn = is_url_and_split(proxy)
        if is_url:
            proxies = {"http": proxy, "https": proxy}
        else:
            return {
                "status": "error",
                "message": "Your proxy is invalid, please put it in the format of http(s)://fqdn.tld:port",
            }

    # Send license request
    licence, req_err = send_license_request(
        license_url,
        format_headers,
        format_cookies,
        format_json_data,
        challenge,
        proxies,
    )
    if licence is None:
        return {"status": "error", "message": req_err}

    # Parse license
    try:
        if is_pr:
            cdm.parse_license(session_id, licence.text)
        else:
            try:
                cdm.parse_license(session_id, licence.content)  # type: ignore[arg-type]
            except (ValueError, TypeError):
                # Try to extract license from JSON
                try:
                    license_json = licence.json()
                    license_value = find_license_key(license_json)
                    if license_value is not None:
                        cdm.parse_license(session_id, license_value)
                    else:
                        return {
                            "status": "error",
                            "message": f"Could not extract license from JSON: {license_json}",
                        }
                except (ValueError, json.JSONDecodeError, AttributeError) as error:
                    return {
                        "status": "error",
                        "message": f"An error occurred parsing license content\n\n{error}\n\n{licence.content}",
                    }
    except (ValueError, TypeError, AttributeError) as error:
        return {
            "status": "error",
            "message": f"An error occurred parsing license content\n\n{error}\n\n{licence.content}",
        }

    # Extract and cache keys
    returned_keys, key_err = extract_and_cache_keys(
        cdm,
        session_id,
        cache_to_db,
        pssh,
        license_url,
        headers,
        cookies,
        challenge,
        json_data,
        is_widevine=(not is_pr),
    )
    if returned_keys == "":
        return {"status": "error", "message": key_err}

    # Close session
    try:
        cdm.close(session_id)
    except (IOError, ValueError, AttributeError) as error:
        return {
            "status": "error",
            "message": f"An error occurred closing session\n\n{error}",
        }

    return {"status": "success", "message": returned_keys}