KeyDive/keydive/drm/modules/oemcrypto.py

from json.encoder import encode_basestring_ascii
from pathlib import Path
from typing import Union, List

from keydive.drm.keybox import KeyBox
from keydive.drm.modules import BaseCdm
from keydive.utils import dumps, b64enc, b64dec


class OEMCrypto(BaseCdm):

    def set_keybox(self, data: bytes) -> None:
        """
        Parses and registers a KeyBox structure used for device provisioning.
        The KeyBox contains cryptographic and device-specific data necessary for DRM.

        Args:
            data (bytes): Raw binary data containing the serialized KeyBox protobuf.
        """
        try:
            # Initialize a new KeyBox instance and parse the protobuf data from bytes
            keybox = KeyBox()
            keybox.ParseFromString(data)  # Parse the raw data into the KeyBox object

            # Check if the stable_id is either not present in the dictionary,
            # or if the existing KeyBox with this stable_id does not have a device_id set.
            if keybox.stable_id not in self._keybox or not self._keybox[keybox.stable_id].device_id:
                # Extract keybox metadata info for logging
                infos = keybox.keybox_info
                # Determine label based on presence of system_id in the keybox info
                label = 'keybox' if infos.get('system_id') else 'encrypted keybox'

                self.logger.info(
                    'Received %s: \n\n%s\n',
                    label, dumps(infos, beauty=True)
                )

                # Store or update the internal keybox dictionary using stable_id as the key
                self._keybox[keybox.stable_id] = keybox
        except Exception as e:
            # Log failure details at debug level to avoid noise in normal logs
            self.logger.debug('Unable to register KeyBox data: %s', e)

    def set_stable_id(self, data: bytes) -> None:
        """
        Registers or updates the stable ID for a KeyBox instance.
        The stable ID uniquely identifies a device and acts as the key
        in the internal KeyBox dictionary.

        Args:
            data (bytes): Stable device identifier as raw bytes,
                          used as the unique key to identify the keybox.

        Notes:
            - If a KeyBox already exists with an empty stable_id, it will be reused;
              otherwise, a new KeyBox instance is created.
        """
        try:
            # Only proceed if this stable ID is not already present in the keybox dictionary
            if data not in self._keybox:
                # Retrieve existing KeyBox with empty stable_id or create a new instance
                keybox = self._keybox.get(b'', KeyBox())

                # Assign the provided stable ID to this KeyBox instance
                keybox.stable_id = data

                # Log the stable ID in a human-readable ASCII-safe format
                self.logger.info(
                    'Received stable ID: \n\n%s\n',
                    encode_basestring_ascii(keybox.stable_id.decode('utf-8')))

                # Store or update the KeyBox in the dictionary with stable_id as the key
                self._keybox[keybox.stable_id] = keybox
        except Exception as e:
            # Log exception details at debug level without interrupting flow
            self.logger.debug('Unable to register KeyBox stable ID: %s', e)

    def set_device_id(self, data: bytes) -> None:
        """
        Registers or updates the device ID within a KeyBox instance.
        The device ID is used to uniquely identify a device for cryptographic
        provisioning and licensing purposes.

        Args:
            data (bytes): Raw device identifier bytes, used to associate a device ID
                          with a KeyBox object.

        Notes:
            - If the device ID is already associated with a KeyBox, this method does nothing.
            - If a KeyBox without a device ID exists, it will be reused; otherwise,
              a new KeyBox instance will be created.
        """
        try:
            # Check if this device ID is already registered in any existing KeyBox
            if not any(k for k in self._keybox.values() if k.device_id == data):
                # Find a KeyBox without a device ID to reuse, or create a new KeyBox
                keybox = next((k for k in self._keybox.values() if not k.device_id), KeyBox())

                # Assign the provided device ID to the KeyBox
                keybox.device_id = data

                # Retrieve device-related metadata for logging purposes
                infos = keybox.device_info
                # Choose label based on whether a system_id is present in metadata
                label = 'device ID' if infos.get('system_id') else 'encrypted device ID'
                self.logger.info(
                    'Received %s: \n\n%s\n',
                    label, dumps(infos, beauty=True) if infos else b64enc(keybox.device_id))

                # Store or update the KeyBox in the dictionary using its stable_id as the key
                self._keybox[keybox.stable_id] = keybox
        except Exception as e:
            # Log exceptions with context at debug level without raising
            self.logger.debug('Unable to register KeyBox device ID: %s', e)

    def set_device_aes_key(self, data: Union[bytes, List[str]]) -> None:
        """
        Imports and stores AES device keys used for decrypting protected content.

        This method accepts either a single AES key as raw bytes or a list of strings
        representing AES keys in various formats, including:
        - Hexadecimal strings
        - Base64-encoded strings (standard or URL-safe)
        - File paths pointing to binary key files

        Args:
            data (bytes or List[str]): The AES key(s) to import. Either a single
                                      16-byte AES key (bytes) or a list of strings
                                      representing keys in different encoded forms or file paths.

        Notes:
            - Only keys exactly 16 bytes in length are accepted as valid AES keys.
            - If a list of strings is given, the method tries each parser (hex, base64, file)
              until one succeeds for each entry.
        """
        if isinstance(data, list):
            # Iterate over all string inputs, attempt parsing with multiple strategies
            for value in data:
                for parser in (bytes.fromhex, b64dec, lambda v: Path(v).read_bytes()):
                    try:
                        # Recursively call with successfully decoded raw bytes
                        self.set_device_aes_key(parser(value))
                        break  # Exit parser loop once successful
                    except Exception:
                        # Ignore failures and try next parser
                        pass
                else:
                    # If no parser succeeded, log the failure with the raw input value
                    self.logger.error('Could not import AES key from input: %s', value)

        # From here on, ensure input is raw bytes representing the AES key
        if not isinstance(data, bytes):
            return

        try:
            # Confirm the AES key length matches 16 bytes (AES-128 standard)
            assert len(data) == 16, f'Invalid AES key length: expected 16 bytes, got {len(data)}'

            # Store the valid AES key into internal cache for later use
            self._device_aes_key.append(data)

            # Log success message with the hex-encoded key for clarity
            self.logger.info('Received device AES key: %s', data.hex())
        except Exception as e:
            # Log any exceptions encountered during key validation or storage
            self.logger.debug('Unable to register device AES key: %s', e)

    def set_encryption_context(self, data: bytes) -> None:
        """
        Registers the encryption context used in content key derivation.

        The encryption context is a blob passed to OEMCrypto or CDM implementations,
        typically including metadata like the label ('ENCRYPTION') and key size.
        It is essential for deriving secure session keys used in content encryption/decryption.

        This function validates and stores the encryption context for use in session key derivation.

        Args:
            data (bytes): A binary blob representing the encryption derivation context.
                          This must begin with a label and end with the encoded key size.

        Notes:
            - Expected structure: b'ENCRYPTION\\x00' + context-specific payload + b'\\x00\\x00\\x00\\x80'
            - The tail 'b"\\x00\\x00\\x00\\x80"' corresponds to 128-bit (0x80) key size in big-endian form.
            - The label is a null-terminated UTF-8 string used in OEMCrypto derivation contexts.
            - Invalid contexts are ignored but logged for debugging.
            """
        try:
            # Constant prefix label used for AES encryption key derivation (null-terminated)
            kEncryptionKeyLabel = b'ENCRYPTION\000'

            # Expected suffix: 128 bits (0x80) in big-endian encoding (for AES-128)
            kEncryptionKeySizeBits = b'\0\0\0\x80'

            # Validate the encryption context starts with the correct label
            assert data.startswith(kEncryptionKeyLabel), 'Context missing expected label: ENCRYPTION\\x00'

            # Validate the encryption context ends with the correct encoded key size
            assert data.endswith(kEncryptionKeySizeBits), 'Context missing expected AES key size specifier'

            # Log that the encryption context was successfully received and registered
            self.logger.info('Received encryption context: %s bytes', len(data))

            # Store the validated encryption context for future use (e.g., key derivation)
            self._context = data
        except Exception as e:
            self.logger.debug('Unable to register encryption context: %s', e)