refactoring and comment fixes

Author: RitvikKapila

examples/src/custom_mpl_cmm_example.py

@@ -1,6 +1,19 @@
 # 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."""
+"""
+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
@@ -21,21 +34,25 @@
 class MPLCustomSigningSuiteOnlyCMM(ICryptographicMaterialsManager):
     """Example custom crypto materials manager class."""
 
-    def __init__(self, keyring: IKeyring) -> None:
+    def __init__(self, keyring: IKeyring, cmm: ICryptographicMaterialsManager = None) -> None:
         """Constructor for MPLCustomSigningSuiteOnlyCMM class."""
-        mat_prov: AwsCryptographicMaterialProviders = AwsCryptographicMaterialProviders(
-            config=MaterialProvidersConfig()
-        )
-
-        # Create a CryptographicMaterialsManager for encryption and decryption
-        cmm_input: CreateDefaultCryptographicMaterialsManagerInput = \
-            CreateDefaultCryptographicMaterialsManagerInput(
-                keyring=keyring
+        if cmm is not None:
+            self.underlying_cmm = cmm
+        else:
+            mat_prov: AwsCryptographicMaterialProviders = AwsCryptographicMaterialProviders(
+                config=MaterialProvidersConfig()
             )
 
-        self.underlying_cmm: ICryptographicMaterialsManager = mat_prov.create_default_cryptographic_materials_manager(
-            input=cmm_input
-        )
+            # 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.
@@ -85,30 +102,16 @@ def encrypt_decrypt_with_cmm(
     client = aws_encryption_sdk.EncryptionSDKClient(commitment_policy=CommitmentPolicy.REQUIRE_ENCRYPT_REQUIRE_DECRYPT)
 
     # Encrypt the plaintext source data
-    ciphertext, encryptor_header = client.encrypt(
+    ciphertext, _ = client.encrypt(
         source=EXAMPLE_DATA,
         materials_manager=cmm
     )
 
     # Decrypt the ciphertext
-    cycled_plaintext, decrypted_header = client.decrypt(
+    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
-
-    # 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).
-    # The encryptor_header.encryption_context has items of the form
-    #     b'key': b'value'
-    # We convert these to strings for easier comparison with the decrypted header below.
-    for k, v in encryptor_header.encryption_context.items():
-        k = str(k.decode("utf-8"))
-        v = str(v.decode("utf-8"))
-        assert v == decrypted_header.encryption_context[k], \
-            "Encryption context does not match expected values"

examples/src/aws_cryptographic_materials_manager_example.py renamed to examples/src/default_cryptographic_materials_manager_example.py

@@ -1,9 +1,9 @@
 # Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 """
-This example sets up the AWS Cryptographic Material Managers (CMM).
+This example sets up the default Cryptographic Material Managers (CMM).
 
-The AWS cryptographic materials manager (CMM) assembles the cryptographic materials
+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
@@ -15,7 +15,7 @@
 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 AWS Cryptographic Material Managers, see
+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
@@ -41,12 +41,12 @@
 EXAMPLE_DATA: bytes = b"Hello World"
 
 
-def encrypt_and_decrypt_with_cmm(
+def encrypt_and_decrypt_with_default_cmm(
     kms_key_id: str
 ):
-    """Demonstrate an encrypt/decrypt cycle using an AWS Cryptographic Material Managers.
+    """Demonstrate an encrypt/decrypt cycle using default Cryptographic Material Managers.
 
-    Usage: encrypt_and_decrypt_with_cmm(kms_key_id)
+    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
@@ -77,7 +77,7 @@ def encrypt_and_decrypt_with_cmm(
         "the data you are handling": "is what you think it is",
     }
 
-    # 4. Create a KMS keyring to use with the CryptographicMaterialsManager
+    # 3. Create a KMS keyring to use with the CryptographicMaterialsManager
     kms_client = boto3.client('kms', region_name="us-west-2")
 
     mat_prov: AwsCryptographicMaterialProviders = AwsCryptographicMaterialProviders(
@@ -93,7 +93,7 @@ def encrypt_and_decrypt_with_cmm(
         input=keyring_input
     )
 
-    # 5. Create a CryptographicMaterialsManager for encryption and decryption
+    # 4. Create a CryptographicMaterialsManager for encryption and decryption
     cmm_input: CreateDefaultCryptographicMaterialsManagerInput = \
         CreateDefaultCryptographicMaterialsManagerInput(
             keyring=kms_keyring
@@ -103,31 +103,31 @@ def encrypt_and_decrypt_with_cmm(
         input=cmm_input
     )
 
-    # 6. Encrypt the data with the encryptionContext.
+    # 5. Encrypt the data with the encryptionContext.
     ciphertext, _ = client.encrypt(
         source=EXAMPLE_DATA,
         materials_manager=cmm,
         encryption_context=encryption_context
     )
 
-    # 7. Demonstrate that the ciphertext and plaintext are different.
+    # 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"
 
-    # 8. Decrypt your encrypted data using the same cmm you used on encrypt.
+    # 7. Decrypt your encrypted data using the same cmm you used on encrypt.
     plaintext_bytes, dec_header = client.decrypt(
         source=ciphertext,
         materials_manager=cmm
     )
 
-    # 9. Demonstrate that the encryption context is correct in the decrypted message header
+    # 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"
 
-    # 10. Demonstrate that the decrypted plaintext is identical to the original plaintext.
+    # 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"

examples/test/legacy/v3_default_cmm.py

@@ -1,6 +1,6 @@
 # Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
-"""Copy-paste of the V3 default CMM with small changes to pass linters.."""
+"""Copy-paste of the V3 default CMM with small changes to pass linters."""
 import logging
 
 import attr

examples/test/test_i_aws_cryptographic_materials_manager_example.py

@@ -0,0 +1,14 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Test suite for the default Cryptographic Materials Manager example."""
+import pytest
+
+from ..src.default_cryptographic_materials_manager_example import encrypt_and_decrypt_with_default_cmm
+
+pytestmark = [pytest.mark.examples]
+
+
+def test_encrypt_and_decrypt_with_default_cmm():
+    """Test function for encrypt and decrypt using the default Cryptographic Materials Manager example."""
+    kms_key_id = "arn:aws:kms:us-west-2:658956600833:key/b3537ef1-d8dc-4780-9f5a-55776cbb2f7f"
+    encrypt_and_decrypt_with_default_cmm(kms_key_id)