Skip to content

Commit fc8fd38

Browse files
RitvikKapilaRitvikKapila
authored
chore(custom_cmm_example.py): added test for custom_cmm_example.py (#690)
1 parent 3476816 commit fc8fd38

12 files changed

+615
-11
lines changed

‎examples/src/custom_mpl_cmm_example.py

+125
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,125 @@
1+
# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
"""
4+
Example to create a custom implementation of the MPL's ICryptographicMaterialsManager class and use it with the ESDK.
5+
6+
The cryptographic materials manager (CMM) assembles the cryptographic materials that are used
7+
to encrypt and decrypt data. The cryptographic materials include plaintext and encrypted data keys,
8+
and an optional message signing key.
9+
10+
Cryptographic Materials Managers (CMMs) are composable; if you just want to extend the behavior of
11+
the default CMM, you can do this as demonstrated in this example. This is the easiest approach if
12+
you are just adding a small check to the CMM methods, as in this example.
13+
14+
If your use case calls for fundamentally changing aspects of the default CMM, you can also write
15+
your own implementation without extending an existing CMM. The default CMM's implementation is a
16+
good reference to use if you need to write a custom CMM implementation from scratch.
17+
Custom implementations of CMMs must implement get_encryption_materials and decrypt_materials.
18+
19+
For more information on a default implementation of a CMM,
20+
please look at the default_cryptographic_materials_manager_example.py example.
21+
22+
For more information on Cryptographic Material Managers, see
23+
https://docs.aws.amazon.com/encryption-sdk/latest/developer-guide/concepts.html#crypt-materials-manager
24+
"""
25+
26+
from aws_cryptographic_materialproviders.mpl import AwsCryptographicMaterialProviders
27+
from aws_cryptographic_materialproviders.mpl.config import MaterialProvidersConfig
28+
from aws_cryptographic_materialproviders.mpl.models import (
29+
CreateDefaultCryptographicMaterialsManagerInput,
30+
SignatureAlgorithmNone,
31+
)
32+
from aws_cryptographic_materialproviders.mpl.references import ICryptographicMaterialsManager, IKeyring
33+
34+
import aws_encryption_sdk
35+
from aws_encryption_sdk import CommitmentPolicy
36+
37+
38+
# Custom CMM implementation using the MPL.
39+
# This CMM only allows encryption/decryption using signing algorithms.
40+
# It wraps an underlying CMM implementation and checks its materials
41+
# to ensure that it is only using signed encryption algorithms.
42+
class MPLCustomSigningSuiteOnlyCMM(ICryptographicMaterialsManager):
43+
"""Example custom crypto materials manager class."""
44+
45+
def __init__(self, keyring: IKeyring, cmm: ICryptographicMaterialsManager = None) -> None:
46+
"""Constructor for MPLCustomSigningSuiteOnlyCMM class."""
47+
if cmm is not None:
48+
self.underlying_cmm = cmm
49+
else:
50+
mat_prov: AwsCryptographicMaterialProviders = AwsCryptographicMaterialProviders(
51+
config=MaterialProvidersConfig()
52+
)
53+
54+
# Create a CryptographicMaterialsManager for encryption and decryption
55+
cmm_input: CreateDefaultCryptographicMaterialsManagerInput = \
56+
CreateDefaultCryptographicMaterialsManagerInput(
57+
keyring=keyring
58+
)
59+
60+
self.underlying_cmm: ICryptographicMaterialsManager = \
61+
mat_prov.create_default_cryptographic_materials_manager(
62+
input=cmm_input
63+
)
64+
65+
def get_encryption_materials(self, param):
66+
"""Provides encryption materials appropriate for the request for the custom CMM.
67+
68+
:param aws_cryptographic_materialproviders.mpl.models.GetEncryptionMaterialsInput param: Input object to
69+
provide to a crypto material manager's `get_encryption_materials` method.
70+
:returns: Encryption materials output
71+
:rtype: aws_cryptographic_materialproviders.mpl.models.GetEncryptionMaterialsOutput
72+
"""
73+
materials = self.underlying_cmm.get_encryption_materials(param)
74+
if isinstance(materials.encryption_materials.algorithm_suite.signature, SignatureAlgorithmNone):
75+
raise ValueError(
76+
"Algorithm provided to MPLCustomSigningSuiteOnlyCMM"
77+
+ " is not a supported signing algorithm: " + str(materials.encryption_materials.algorithm_suite)
78+
)
79+
return materials
80+
81+
def decrypt_materials(self, param):
82+
"""Provides decryption materials appropriate for the request for the custom CMM.
83+
84+
:param aws_cryptographic_materialproviders.mpl.models.DecryptMaterialsInput param: Input object to provide
85+
to a crypto material manager's `decrypt_materials` method.
86+
:returns: Decryption materials output
87+
:rtype: aws_cryptographic_materialproviders.mpl.models.GetDecryptionMaterialsOutput
88+
"""
89+
materials = self.underlying_cmm.decrypt_materials(param)
90+
if isinstance(materials.decryption_materials.algorithm_suite.signature, SignatureAlgorithmNone):
91+
raise ValueError(
92+
"Algorithm provided to MPLCustomSigningSuiteOnlyCMM"
93+
+ " is not a supported signing algorithm: " + str(materials.decryption_materials.algorithm_suite)
94+
)
95+
return materials
96+
97+
98+
EXAMPLE_DATA: bytes = b"Hello World"
99+
100+
101+
def encrypt_decrypt_with_cmm(
102+
cmm: ICryptographicMaterialsManager
103+
):
104+
"""Encrypts and decrypts a string using a custom CMM.
105+
106+
:param ICryptographicMaterialsManager cmm: CMM to use for encryption and decryption
107+
"""
108+
# Set up an encryption client with an explicit commitment policy. Note that if you do not explicitly choose a
109+
# commitment policy, REQUIRE_ENCRYPT_REQUIRE_DECRYPT is used by default.
110+
client = aws_encryption_sdk.EncryptionSDKClient(commitment_policy=CommitmentPolicy.REQUIRE_ENCRYPT_REQUIRE_DECRYPT)
111+
112+
# Encrypt the plaintext source data
113+
ciphertext, _ = client.encrypt(
114+
source=EXAMPLE_DATA,
115+
materials_manager=cmm
116+
)
117+
118+
# Decrypt the ciphertext
119+
cycled_plaintext, _ = client.decrypt(
120+
source=ciphertext,
121+
materials_manager=cmm
122+
)
123+
124+
# Verify that the "cycled" (encrypted, then decrypted) plaintext is identical to the source plaintext
125+
assert cycled_plaintext == EXAMPLE_DATA

‎examples/src/default_cryptographic_materials_manager_example.py

+133
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
"""
4+
This example sets up the default Cryptographic Material Managers (CMM).
5+
6+
The default cryptographic materials manager (CMM) assembles the cryptographic materials
7+
that are used to encrypt and decrypt data. The cryptographic materials include
8+
plaintext and encrypted data keys, and an optional message signing key.
9+
This example creates a CMM and then encrypts a custom input EXAMPLE_DATA
10+
with an encryption context. Creating a CMM involves taking a keyring as input,
11+
and we use an AWS KMS Keyring for this example.
12+
This example also includes some sanity checks for demonstration:
13+
1. Ciphertext and plaintext data are not the same
14+
2. Encryption context is correct in the decrypted message header
15+
3. Decrypted plaintext value matches EXAMPLE_DATA
16+
These sanity checks are for demonstration in the example only. You do not need these in your code.
17+
18+
For more information on Cryptographic Material Managers, see
19+
https://docs.aws.amazon.com/encryption-sdk/latest/developer-guide/concepts.html#crypt-materials-manager
20+
"""
21+
import sys
22+
23+
import boto3
24+
from aws_cryptographic_materialproviders.mpl import AwsCryptographicMaterialProviders
25+
from aws_cryptographic_materialproviders.mpl.config import MaterialProvidersConfig
26+
from aws_cryptographic_materialproviders.mpl.models import (
27+
CreateAwsKmsKeyringInput,
28+
CreateDefaultCryptographicMaterialsManagerInput,
29+
)
30+
from aws_cryptographic_materialproviders.mpl.references import ICryptographicMaterialsManager, IKeyring
31+
from typing import Dict # noqa pylint: disable=wrong-import-order
32+
33+
import aws_encryption_sdk
34+
from aws_encryption_sdk import CommitmentPolicy
35+
36+
# TODO-MPL: Remove this as part of removing PYTHONPATH hacks.
37+
MODULE_ROOT_DIR = '/'.join(__file__.split("/")[:-1])
38+
39+
sys.path.append(MODULE_ROOT_DIR)
40+
41+
EXAMPLE_DATA: bytes = b"Hello World"
42+
43+
44+
def encrypt_and_decrypt_with_default_cmm(
45+
kms_key_id: str
46+
):
47+
"""Demonstrate an encrypt/decrypt cycle using default Cryptographic Material Managers.
48+
49+
Usage: encrypt_and_decrypt_with_default_cmm(kms_key_id)
50+
:param kms_key_id: KMS Key identifier for the KMS key you want to use for encryption and
51+
decryption of your data keys.
52+
:type kms_key_id: string
53+
54+
For more information on KMS Key identifiers, see
55+
https://docs.aws.amazon.com/kms/latest/developerguide/concepts.html#key-id
56+
"""
57+
# 1. Instantiate the encryption SDK client.
58+
# This builds the client with the REQUIRE_ENCRYPT_REQUIRE_DECRYPT commitment policy,
59+
# which enforces that this client only encrypts using committing algorithm suites and enforces
60+
# that this client will only decrypt encrypted messages that were created with a committing
61+
# algorithm suite.
62+
# This is the default commitment policy if you were to build the client as
63+
# `client = aws_encryption_sdk.EncryptionSDKClient()`.
64+
client = aws_encryption_sdk.EncryptionSDKClient(
65+
commitment_policy=CommitmentPolicy.REQUIRE_ENCRYPT_REQUIRE_DECRYPT
66+
)
67+
68+
# 2. Create encryption context.
69+
# Remember that your encryption context is NOT SECRET.
70+
# For more information, see
71+
# https://docs.aws.amazon.com/encryption-sdk/latest/developer-guide/concepts.html#encryption-context
72+
encryption_context: Dict[str, str] = {
73+
"encryption": "context",
74+
"is not": "secret",
75+
"but adds": "useful metadata",
76+
"that can help you": "be confident that",
77+
"the data you are handling": "is what you think it is",
78+
}
79+
80+
# 3. Create a KMS keyring to use with the CryptographicMaterialsManager
81+
kms_client = boto3.client('kms', region_name="us-west-2")
82+
83+
mat_prov: AwsCryptographicMaterialProviders = AwsCryptographicMaterialProviders(
84+
config=MaterialProvidersConfig()
85+
)
86+
87+
keyring_input: CreateAwsKmsKeyringInput = CreateAwsKmsKeyringInput(
88+
kms_key_id=kms_key_id,
89+
kms_client=kms_client
90+
)
91+
92+
kms_keyring: IKeyring = mat_prov.create_aws_kms_keyring(
93+
input=keyring_input
94+
)
95+
96+
# 4. Create a CryptographicMaterialsManager for encryption and decryption
97+
cmm_input: CreateDefaultCryptographicMaterialsManagerInput = \
98+
CreateDefaultCryptographicMaterialsManagerInput(
99+
keyring=kms_keyring
100+
)
101+
102+
cmm: ICryptographicMaterialsManager = mat_prov.create_default_cryptographic_materials_manager(
103+
input=cmm_input
104+
)
105+
106+
# 5. Encrypt the data with the encryptionContext.
107+
ciphertext, _ = client.encrypt(
108+
source=EXAMPLE_DATA,
109+
materials_manager=cmm,
110+
encryption_context=encryption_context
111+
)
112+
113+
# 6. Demonstrate that the ciphertext and plaintext are different.
114+
# (This is an example for demonstration; you do not need to do this in your own code.)
115+
assert ciphertext != EXAMPLE_DATA, \
116+
"Ciphertext and plaintext data are the same. Invalid encryption"
117+
118+
# 7. Decrypt your encrypted data using the same cmm you used on encrypt.
119+
plaintext_bytes, dec_header = client.decrypt(
120+
source=ciphertext,
121+
materials_manager=cmm
122+
)
123+
124+
# 8. Demonstrate that the encryption context is correct in the decrypted message header
125+
# (This is an example for demonstration; you do not need to do this in your own code.)
126+
for k, v in encryption_context.items():
127+
assert v == dec_header.encryption_context[k], \
128+
"Encryption context does not match expected values"
129+
130+
# 9. Demonstrate that the decrypted plaintext is identical to the original plaintext.
131+
# (This is an example for demonstration; you do not need to do this in your own code.)
132+
assert plaintext_bytes == EXAMPLE_DATA, \
133+
"Decrypted plaintext should be identical to the original plaintext. Invalid decryption"

‎examples/src/legacy/custom_cmm_example.py

+87
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
"""Example to create a custom implementation of the native ESDK CryptoMaterialsManager class."""
4+
5+
import aws_encryption_sdk
6+
from aws_encryption_sdk import CommitmentPolicy, StrictAwsKmsMasterKeyProvider
7+
from aws_encryption_sdk.materials_managers.base import CryptoMaterialsManager
8+
from aws_encryption_sdk.materials_managers.default import DefaultCryptoMaterialsManager
9+
10+
11+
# Custom CMM implementation.
12+
# This CMM only allows encryption/decryption using signing algorithms.
13+
# It wraps an underlying CMM implementation and checks its materials
14+
# to ensure that it is only using signed encryption algorithms.
15+
class CustomSigningSuiteOnlyCMM(CryptoMaterialsManager):
16+
"""Example custom crypto materials manager class."""
17+
18+
def __init__(self, master_key_provider: StrictAwsKmsMasterKeyProvider) -> None:
19+
"""Constructor for CustomSigningSuiteOnlyCMM class."""
20+
self.underlying_cmm = DefaultCryptoMaterialsManager(master_key_provider)
21+
22+
def get_encryption_materials(self, request):
23+
"""Provides encryption materials appropriate for the request for the custom CMM.
24+
25+
:param EncryptionMaterialsRequest request: Request object to provide to a
26+
crypto material manager's `get_encryption_materials` method.
27+
:returns: Encryption materials
28+
:rtype: EncryptionMaterials
29+
"""
30+
materials = self.underlying_cmm.get_encryption_materials(request)
31+
if not materials.algorithm.is_signing():
32+
raise ValueError(
33+
"Algorithm provided to CustomSigningSuiteOnlyCMM"
34+
+ " is not a supported signing algorithm: " + materials.algorithm
35+
)
36+
return materials
37+
38+
def decrypt_materials(self, request):
39+
"""Provides decryption materials appropriate for the request for the custom CMM.
40+
41+
:param DecryptionMaterialsRequest request: Request object to provide to a
42+
crypto material manager's `decrypt_materials` method.
43+
"""
44+
if not request.algorithm.is_signing():
45+
raise ValueError(
46+
"Algorithm provided to CustomSigningSuiteOnlyCMM"
47+
+ " is not a supported signing algorithm: " + request.algorithm
48+
)
49+
return self.underlying_cmm.decrypt_materials(request)
50+
51+
52+
def encrypt_decrypt_with_cmm(
53+
cmm: CryptoMaterialsManager,
54+
source_plaintext: str
55+
):
56+
"""Encrypts and decrypts a string using a custom CMM.
57+
58+
:param CryptoMaterialsManager cmm: CMM to use for encryption and decryption
59+
:param bytes source_plaintext: Data to encrypt
60+
"""
61+
# Set up an encryption client with an explicit commitment policy. Note that if you do not explicitly choose a
62+
# commitment policy, REQUIRE_ENCRYPT_REQUIRE_DECRYPT is used by default.
63+
client = aws_encryption_sdk.EncryptionSDKClient(commitment_policy=CommitmentPolicy.REQUIRE_ENCRYPT_REQUIRE_DECRYPT)
64+
65+
# Encrypt the plaintext source data
66+
ciphertext, encryptor_header = client.encrypt(
67+
source=source_plaintext,
68+
materials_manager=cmm
69+
)
70+
71+
# Decrypt the ciphertext
72+
cycled_plaintext, decrypted_header = client.decrypt(
73+
source=ciphertext,
74+
materials_manager=cmm
75+
)
76+
77+
# Verify that the "cycled" (encrypted, then decrypted) plaintext is identical to the source plaintext
78+
assert cycled_plaintext == source_plaintext
79+
80+
# Verify that the encryption context used in the decrypt operation includes all key pairs from
81+
# the encrypt operation. (The SDK can add pairs, so don't require an exact match.)
82+
#
83+
# In production, always use a meaningful encryption context. In this sample, we omit the
84+
# encryption context (no key pairs).
85+
assert all(
86+
pair in decrypted_header.encryption_context.items() for pair in encryptor_header.encryption_context.items()
87+
)

‎examples/test/legacy/examples_test_utils.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@
3939

4040
from integration_test_utils import ( # noqa pylint: disable=unused-import,import-error
4141
get_cmk_arn,
42-
get_second_cmk_arn,
4342
get_mrk_arn,
43+
get_second_cmk_arn,
4444
get_second_mrk_arn,
4545
)

0 commit comments

Comments
 (0)