def prepare_data_keys()

in src/aws_encryption_sdk/internal/utils/__init__.py [0:0]


def prepare_data_keys(primary_master_key, master_keys, algorithm, encryption_context):
    """Prepares a DataKey to be used for encrypting message and list
    of EncryptedDataKey objects to be serialized into header.

    :param primary_master_key: Master key with which to generate the encryption data key
    :type primary_master_key: aws_encryption_sdk.key_providers.base.MasterKey
    :param master_keys: All master keys with which to encrypt data keys
    :type master_keys: list of :class:`aws_encryption_sdk.key_providers.base.MasterKey`
    :param algorithm: Algorithm to use for encryption
    :type algorithm: aws_encryption_sdk.identifiers.Algorithm
    :param dict encryption_context: Encryption context to use when generating data key
    :rtype: tuple containing :class:`aws_encryption_sdk.structures.DataKey`
        and set of :class:`aws_encryption_sdk.structures.EncryptedDataKey`
    """
    encrypted_data_keys = set()
    encrypted_data_encryption_key = None
    data_encryption_key = primary_master_key.generate_data_key(algorithm, encryption_context)
    _LOGGER.debug("encryption data generated with master key: %s", data_encryption_key.key_provider)
    for master_key in master_keys:
        # Don't re-encrypt the encryption data key; we already have the ciphertext
        if master_key is primary_master_key:
            encrypted_data_encryption_key = EncryptedDataKey(
                key_provider=data_encryption_key.key_provider, encrypted_data_key=data_encryption_key.encrypted_data_key
            )
            encrypted_data_keys.add(encrypted_data_encryption_key)
            continue
        encrypted_key = master_key.encrypt_data_key(
            data_key=data_encryption_key, algorithm=algorithm, encryption_context=encryption_context
        )
        encrypted_data_keys.add(encrypted_key)
        _LOGGER.debug("encryption key encrypted with master key: %s", master_key.key_provider)
    return data_encryption_key, encrypted_data_keys