Skip to content

chore(custom_cmm_example.py): added test for custom_cmm_example.py #690

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 25 commits into from
Jun 26, 2024
Merged
Show file tree
Hide file tree
Changes from 24 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
117 changes: 117 additions & 0 deletions examples/src/custom_mpl_cmm_example.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,117 @@
# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0
"""
Example to create a custom implementation of the ESDK-MPL ICryptographicMaterialsManager class.

Cryptographic Materials Managers (CMMs) are composable; if you just want to extend the behavior of
the default CMM, you can do this as demonstrated in this example. This is easy if you just want
to add a small check to the CMM methods.

Custom implementation of CMMs must implement get_encryption_materials and decrypt_materials.
If your use case calls for fundamentally change aspects of the default CMM, you can also write
your own implementation without extending a CMM.

For more information on a default implementation of a CMM,
please look at the default_cryptographic_materials_manager_example.py example.
"""

from aws_cryptographic_materialproviders.mpl import AwsCryptographicMaterialProviders
from aws_cryptographic_materialproviders.mpl.config import MaterialProvidersConfig
from aws_cryptographic_materialproviders.mpl.models import (
CreateDefaultCryptographicMaterialsManagerInput,
SignatureAlgorithmNone,
)
from aws_cryptographic_materialproviders.mpl.references import ICryptographicMaterialsManager, IKeyring

import aws_encryption_sdk
from aws_encryption_sdk import CommitmentPolicy


# Custom CMM implementation using the MPL.
# This CMM only allows encryption/decryption using signing algorithms.
# It wraps an underlying CMM implementation and checks its materials
# to ensure that it is only using signed encryption algorithms.
class MPLCustomSigningSuiteOnlyCMM(ICryptographicMaterialsManager):
"""Example custom crypto materials manager class."""

def __init__(self, keyring: IKeyring, cmm: ICryptographicMaterialsManager = None) -> None:
"""Constructor for MPLCustomSigningSuiteOnlyCMM class."""
if cmm is not None:
self.underlying_cmm = cmm
else:
mat_prov: AwsCryptographicMaterialProviders = AwsCryptographicMaterialProviders(
config=MaterialProvidersConfig()
)

# Create a CryptographicMaterialsManager for encryption and decryption
cmm_input: CreateDefaultCryptographicMaterialsManagerInput = \
CreateDefaultCryptographicMaterialsManagerInput(
keyring=keyring
)

self.underlying_cmm: ICryptographicMaterialsManager = \
mat_prov.create_default_cryptographic_materials_manager(
input=cmm_input
)

def get_encryption_materials(self, param):
"""Provides encryption materials appropriate for the request for the custom CMM.

:param aws_cryptographic_materialproviders.mpl.models.GetEncryptionMaterialsInput param: Input object to
provide to a crypto material manager's `get_encryption_materials` method.
:returns: Encryption materials output
:rtype: aws_cryptographic_materialproviders.mpl.models.GetEncryptionMaterialsOutput
"""
materials = self.underlying_cmm.get_encryption_materials(param)
if isinstance(materials.encryption_materials.algorithm_suite.signature, SignatureAlgorithmNone):
raise ValueError(
"Algorithm provided to MPLCustomSigningSuiteOnlyCMM"
+ " is not a supported signing algorithm: " + str(materials.encryption_materials.algorithm_suite)
)
return materials

def decrypt_materials(self, param):
"""Provides decryption materials appropriate for the request for the custom CMM.

:param aws_cryptographic_materialproviders.mpl.models.DecryptMaterialsInput param: Input object to provide
to a crypto material manager's `decrypt_materials` method.
:returns: Decryption materials output
:rtype: aws_cryptographic_materialproviders.mpl.models.GetDecryptionMaterialsOutput
"""
materials = self.underlying_cmm.decrypt_materials(param)
if isinstance(materials.decryption_materials.algorithm_suite.signature, SignatureAlgorithmNone):
raise ValueError(
"Algorithm provided to MPLCustomSigningSuiteOnlyCMM"
+ " is not a supported signing algorithm: " + str(materials.decryption_materials.algorithm_suite)
)
return materials


EXAMPLE_DATA: bytes = b"Hello World"


def encrypt_decrypt_with_cmm(
cmm: ICryptographicMaterialsManager
):
"""Encrypts and decrypts a string using a custom CMM.

:param ICryptographicMaterialsManager cmm: CMM to use for encryption and decryption
"""
# Set up an encryption client with an explicit commitment policy. Note that if you do not explicitly choose a
# commitment policy, REQUIRE_ENCRYPT_REQUIRE_DECRYPT is used by default.
client = aws_encryption_sdk.EncryptionSDKClient(commitment_policy=CommitmentPolicy.REQUIRE_ENCRYPT_REQUIRE_DECRYPT)

# Encrypt the plaintext source data
ciphertext, _ = client.encrypt(
source=EXAMPLE_DATA,
materials_manager=cmm
)

# Decrypt the ciphertext
cycled_plaintext, _ = client.decrypt(
source=ciphertext,
materials_manager=cmm
)

# Verify that the "cycled" (encrypted, then decrypted) plaintext is identical to the source plaintext
assert cycled_plaintext == EXAMPLE_DATA
133 changes: 133 additions & 0 deletions examples/src/default_cryptographic_materials_manager_example.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,133 @@
# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0
"""
This example sets up the default Cryptographic Material Managers (CMM).

The default cryptographic materials manager (CMM) assembles the cryptographic materials
that are used to encrypt and decrypt data. The cryptographic materials include
plaintext and encrypted data keys, and an optional message signing key.
This example creates a CMM and then encrypts a custom input EXAMPLE_DATA
with an encryption context. Creating a CMM involves taking a keyring as input,
and we use an AWS KMS Keyring for this example.
This example also includes some sanity checks for demonstration:
1. Ciphertext and plaintext data are not the same
2. Encryption context is correct in the decrypted message header
3. Decrypted plaintext value matches EXAMPLE_DATA
These sanity checks are for demonstration in the example only. You do not need these in your code.

For more information on Cryptographic Material Managers, see
https://docs.aws.amazon.com/encryption-sdk/latest/developer-guide/concepts.html#crypt-materials-manager
"""
import sys

import boto3
from aws_cryptographic_materialproviders.mpl import AwsCryptographicMaterialProviders
from aws_cryptographic_materialproviders.mpl.config import MaterialProvidersConfig
from aws_cryptographic_materialproviders.mpl.models import (
CreateAwsKmsKeyringInput,
CreateDefaultCryptographicMaterialsManagerInput,
)
from aws_cryptographic_materialproviders.mpl.references import ICryptographicMaterialsManager, IKeyring
from typing import Dict # noqa pylint: disable=wrong-import-order

import aws_encryption_sdk
from aws_encryption_sdk import CommitmentPolicy

# TODO-MPL: Remove this as part of removing PYTHONPATH hacks.
MODULE_ROOT_DIR = '/'.join(__file__.split("/")[:-1])

sys.path.append(MODULE_ROOT_DIR)

EXAMPLE_DATA: bytes = b"Hello World"


def encrypt_and_decrypt_with_default_cmm(
kms_key_id: str
):
"""Demonstrate an encrypt/decrypt cycle using default Cryptographic Material Managers.

Usage: encrypt_and_decrypt_with_default_cmm(kms_key_id)
:param kms_key_id: KMS Key identifier for the KMS key you want to use for encryption and
decryption of your data keys.
:type kms_key_id: string

For more information on KMS Key identifiers, see
https://docs.aws.amazon.com/kms/latest/developerguide/concepts.html#key-id
"""
# 1. Instantiate the encryption SDK client.
# This builds the client with the REQUIRE_ENCRYPT_REQUIRE_DECRYPT commitment policy,
# which enforces that this client only encrypts using committing algorithm suites and enforces
# that this client will only decrypt encrypted messages that were created with a committing
# algorithm suite.
# This is the default commitment policy if you were to build the client as
# `client = aws_encryption_sdk.EncryptionSDKClient()`.
client = aws_encryption_sdk.EncryptionSDKClient(
commitment_policy=CommitmentPolicy.REQUIRE_ENCRYPT_REQUIRE_DECRYPT
)

# 2. Create encryption context.
# Remember that your encryption context is NOT SECRET.
# For more information, see
# https://docs.aws.amazon.com/encryption-sdk/latest/developer-guide/concepts.html#encryption-context
encryption_context: Dict[str, str] = {
"encryption": "context",
"is not": "secret",
"but adds": "useful metadata",
"that can help you": "be confident that",
"the data you are handling": "is what you think it is",
}

# 3. Create a KMS keyring to use with the CryptographicMaterialsManager
kms_client = boto3.client('kms', region_name="us-west-2")

mat_prov: AwsCryptographicMaterialProviders = AwsCryptographicMaterialProviders(
config=MaterialProvidersConfig()
)

keyring_input: CreateAwsKmsKeyringInput = CreateAwsKmsKeyringInput(
kms_key_id=kms_key_id,
kms_client=kms_client
)

kms_keyring: IKeyring = mat_prov.create_aws_kms_keyring(
input=keyring_input
)

# 4. Create a CryptographicMaterialsManager for encryption and decryption
cmm_input: CreateDefaultCryptographicMaterialsManagerInput = \
CreateDefaultCryptographicMaterialsManagerInput(
keyring=kms_keyring
)

cmm: ICryptographicMaterialsManager = mat_prov.create_default_cryptographic_materials_manager(
input=cmm_input
)

# 5. Encrypt the data with the encryptionContext.
ciphertext, _ = client.encrypt(
source=EXAMPLE_DATA,
materials_manager=cmm,
encryption_context=encryption_context
)

# 6. Demonstrate that the ciphertext and plaintext are different.
# (This is an example for demonstration; you do not need to do this in your own code.)
assert ciphertext != EXAMPLE_DATA, \
"Ciphertext and plaintext data are the same. Invalid encryption"

# 7. Decrypt your encrypted data using the same cmm you used on encrypt.
plaintext_bytes, dec_header = client.decrypt(
source=ciphertext,
materials_manager=cmm
)

# 8. Demonstrate that the encryption context is correct in the decrypted message header
# (This is an example for demonstration; you do not need to do this in your own code.)
for k, v in encryption_context.items():
assert v == dec_header.encryption_context[k], \
"Encryption context does not match expected values"

# 9. Demonstrate that the decrypted plaintext is identical to the original plaintext.
# (This is an example for demonstration; you do not need to do this in your own code.)
assert plaintext_bytes == EXAMPLE_DATA, \
"Decrypted plaintext should be identical to the original plaintext. Invalid decryption"
87 changes: 87 additions & 0 deletions examples/src/legacy/custom_cmm_example.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,87 @@
# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0
"""Example to create a custom implementation of the native ESDK CryptoMaterialsManager class."""

import aws_encryption_sdk
from aws_encryption_sdk import CommitmentPolicy, StrictAwsKmsMasterKeyProvider
from aws_encryption_sdk.materials_managers.base import CryptoMaterialsManager
from aws_encryption_sdk.materials_managers.default import DefaultCryptoMaterialsManager


# Custom CMM implementation.
# This CMM only allows encryption/decryption using signing algorithms.
# It wraps an underlying CMM implementation and checks its materials
# to ensure that it is only using signed encryption algorithms.
class CustomSigningSuiteOnlyCMM(CryptoMaterialsManager):
"""Example custom crypto materials manager class."""

def __init__(self, master_key_provider: StrictAwsKmsMasterKeyProvider) -> None:
"""Constructor for CustomSigningSuiteOnlyCMM class."""
self.underlying_cmm = DefaultCryptoMaterialsManager(master_key_provider)

def get_encryption_materials(self, request):
"""Provides encryption materials appropriate for the request for the custom CMM.

:param EncryptionMaterialsRequest request: Request object to provide to a
crypto material manager's `get_encryption_materials` method.
:returns: Encryption materials
:rtype: EncryptionMaterials
"""
materials = self.underlying_cmm.get_encryption_materials(request)
if not materials.algorithm.is_signing():
raise ValueError(
"Algorithm provided to CustomSigningSuiteOnlyCMM"
+ " is not a supported signing algorithm: " + materials.algorithm
)
return materials

def decrypt_materials(self, request):
"""Provides decryption materials appropriate for the request for the custom CMM.

:param DecryptionMaterialsRequest request: Request object to provide to a
crypto material manager's `decrypt_materials` method.
"""
if not request.algorithm.is_signing():
raise ValueError(
"Algorithm provided to CustomSigningSuiteOnlyCMM"
+ " is not a supported signing algorithm: " + request.algorithm
)
return self.underlying_cmm.decrypt_materials(request)


def encrypt_decrypt_with_cmm(
cmm: CryptoMaterialsManager,
source_plaintext: str
):
"""Encrypts and decrypts a string using a custom CMM.

:param CryptoMaterialsManager cmm: CMM to use for encryption and decryption
:param bytes source_plaintext: Data to encrypt
"""
# Set up an encryption client with an explicit commitment policy. Note that if you do not explicitly choose a
# commitment policy, REQUIRE_ENCRYPT_REQUIRE_DECRYPT is used by default.
client = aws_encryption_sdk.EncryptionSDKClient(commitment_policy=CommitmentPolicy.REQUIRE_ENCRYPT_REQUIRE_DECRYPT)

# Encrypt the plaintext source data
ciphertext, encryptor_header = client.encrypt(
source=source_plaintext,
materials_manager=cmm
)

# Decrypt the ciphertext
cycled_plaintext, decrypted_header = client.decrypt(
source=ciphertext,
materials_manager=cmm
)

# Verify that the "cycled" (encrypted, then decrypted) plaintext is identical to the source plaintext
assert cycled_plaintext == source_plaintext

# Verify that the encryption context used in the decrypt operation includes all key pairs from
# the encrypt operation. (The SDK can add pairs, so don't require an exact match.)
#
# In production, always use a meaningful encryption context. In this sample, we omit the
# encryption context (no key pairs).
assert all(
pair in decrypted_header.encryption_context.items() for pair in encryptor_header.encryption_context.items()
)
2 changes: 1 addition & 1 deletion examples/test/legacy/examples_test_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@

from integration_test_utils import ( # noqa pylint: disable=unused-import,import-error
get_cmk_arn,
get_second_cmk_arn,
get_mrk_arn,
get_second_cmk_arn,
get_second_mrk_arn,
)
Loading
Loading