From c3fe3e223af724be37b35d1c784c12e1c8dcf517 Mon Sep 17 00:00:00 2001
From: Ritvik Kapila <ritvik.iitd@gmail.com>
Date: Tue, 18 Jun 2024 11:51:08 -0700
Subject: [PATCH 01/25] chore(v3_default_cmm.py): added test for
 v3_default_cmm.py; isort fix for migration examples

---
 examples/src/legacy/v3_default_cmm.py         | 202 ++++++++++++++++++
 examples/test/legacy/examples_test_utils.py   |   2 +-
 examples/test/legacy/test_i_v3_default_cmm.py |  20 ++
 .../test_i_migration_aws_kms_key_example.py   |   4 +-
 .../test_i_migration_raw_aes_key_example.py   |   4 +-
 .../test_i_migration_raw_rsa_key_example.py   |   4 +-
 6 files changed, 226 insertions(+), 10 deletions(-)
 create mode 100644 examples/src/legacy/v3_default_cmm.py
 create mode 100644 examples/test/legacy/test_i_v3_default_cmm.py

diff --git a/examples/src/legacy/v3_default_cmm.py b/examples/src/legacy/v3_default_cmm.py
new file mode 100644
index 000000000..61f70e05d
--- /dev/null
+++ b/examples/src/legacy/v3_default_cmm.py
@@ -0,0 +1,202 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Default crypto material manager class example for ESDK V3."""
+import logging
+
+import attr
+
+import aws_encryption_sdk
+from aws_encryption_sdk.exceptions import MasterKeyProviderError, SerializationError
+from aws_encryption_sdk.identifiers import CommitmentPolicy
+from aws_encryption_sdk.internal.crypto.authentication import Signer, Verifier
+from aws_encryption_sdk.internal.crypto.elliptic_curve import generate_ecc_signing_key
+from aws_encryption_sdk.internal.defaults import ALGORITHM, ALGORITHM_COMMIT_KEY, ENCODED_SIGNER_KEY
+from aws_encryption_sdk.internal.str_ops import to_str
+from aws_encryption_sdk.internal.utils import prepare_data_keys
+from aws_encryption_sdk.internal.utils.commitment import (
+    validate_commitment_policy_on_decrypt,
+    validate_commitment_policy_on_encrypt,
+)
+from aws_encryption_sdk.key_providers.base import MasterKeyProvider
+from aws_encryption_sdk.materials_managers import DecryptionMaterials, EncryptionMaterials
+from aws_encryption_sdk.materials_managers.base import CryptoMaterialsManager
+
+_LOGGER = logging.getLogger(__name__)
+
+
+@attr.s(hash=False)
+class V3DefaultCryptoMaterialsManager(CryptoMaterialsManager):
+    """Copy of the default crypto material manager for ESDK V3.
+
+    This is a copy-paste of the DefaultCryptoMaterialsManager implementation
+    from the V3 ESDK commit: 98b5eb7c2bd7d1b2a3380aacfa508e8721c4d8a9
+    This CMM is used to explicitly assert that the V3 implementation of
+    the DefaultCMM is compatible with future version's logic,
+    which implicitly asserts that custom implementations of V3-compatible CMMs
+    are also compatible with future version's logic.
+
+    :param master_key_provider: Master key provider to use
+    :type master_key_provider: aws_encryption_sdk.key_providers.base.MasterKeyProvider
+    """
+
+    master_key_provider = attr.ib(validator=attr.validators.instance_of(MasterKeyProvider))
+
+    def _generate_signing_key_and_update_encryption_context(self, algorithm, encryption_context):
+        """Generates a signing key based on the provided algorithm.
+
+        :param algorithm: Algorithm for which to generate signing key
+        :type algorithm: aws_encryption_sdk.identifiers.Algorithm
+        :param dict encryption_context: Encryption context from request
+        :returns: Signing key bytes
+        :rtype: bytes or None
+        """
+        _LOGGER.debug("Generating signing key")
+        if algorithm.signing_algorithm_info is None:
+            return None
+
+        signer = Signer(algorithm=algorithm, key=generate_ecc_signing_key(algorithm=algorithm))
+        encryption_context[ENCODED_SIGNER_KEY] = to_str(signer.encoded_public_key())
+        return signer.key_bytes()
+
+    def get_encryption_materials(self, request):
+        """Creates encryption materials using underlying master key provider.
+
+        :param request: encryption materials request
+        :type request: aws_encryption_sdk.materials_managers.EncryptionMaterialsRequest
+        :returns: encryption materials
+        :rtype: aws_encryption_sdk.materials_managers.EncryptionMaterials
+        :raises MasterKeyProviderError: if no master keys are available from the underlying master key provider
+        :raises MasterKeyProviderError: if the primary master key provided by the underlying master key provider
+            is not included in the full set of master keys provided by that provider
+        :raises ActionNotAllowedError: if the commitment policy in the request is violated by the algorithm being
+            used
+        """
+        default_algorithm = ALGORITHM
+        if request.commitment_policy in (
+            CommitmentPolicy.REQUIRE_ENCRYPT_REQUIRE_DECRYPT,
+            CommitmentPolicy.REQUIRE_ENCRYPT_ALLOW_DECRYPT,
+        ):
+            default_algorithm = ALGORITHM_COMMIT_KEY
+        algorithm = request.algorithm if request.algorithm is not None else default_algorithm
+
+        validate_commitment_policy_on_encrypt(request.commitment_policy, request.algorithm)
+
+        encryption_context = request.encryption_context.copy()
+
+        signing_key = self._generate_signing_key_and_update_encryption_context(algorithm, encryption_context)
+
+        primary_master_key, master_keys = self.master_key_provider.master_keys_for_encryption(
+            encryption_context=encryption_context,
+            plaintext_rostream=request.plaintext_rostream,
+            plaintext_length=request.plaintext_length,
+        )
+        if not master_keys:
+            raise MasterKeyProviderError("No Master Keys available from Master Key Provider")
+        if primary_master_key not in master_keys:
+            raise MasterKeyProviderError("Primary Master Key not in provided Master Keys")
+
+        data_encryption_key, encrypted_data_keys = prepare_data_keys(
+            primary_master_key=primary_master_key,
+            master_keys=master_keys,
+            algorithm=algorithm,
+            encryption_context=encryption_context,
+        )
+
+        _LOGGER.debug("Post-encrypt encryption context: %s", encryption_context)
+
+        return EncryptionMaterials(
+            algorithm=algorithm,
+            data_encryption_key=data_encryption_key,
+            encrypted_data_keys=encrypted_data_keys,
+            encryption_context=encryption_context,
+            signing_key=signing_key,
+        )
+
+    def _load_verification_key_from_encryption_context(self, algorithm, encryption_context):
+        """Loads the verification key from the encryption context if used by algorithm suite.
+
+        :param algorithm: Algorithm for which to generate signing key
+        :type algorithm: aws_encryption_sdk.identifiers.Algorithm
+        :param dict encryption_context: Encryption context from request
+        :returns: Raw verification key
+        :rtype: bytes
+        :raises SerializationError: if algorithm suite requires message signing and no verification key is found
+        """
+        encoded_verification_key = encryption_context.get(ENCODED_SIGNER_KEY, None)
+
+        if algorithm.signing_algorithm_info is not None and encoded_verification_key is None:
+            raise SerializationError("No signature verification key found in header for signed algorithm.")
+
+        if algorithm.signing_algorithm_info is None:
+            if encoded_verification_key is not None:
+                raise SerializationError("Signature verification key found in header for non-signed algorithm.")
+            return None
+
+        verifier = Verifier.from_encoded_point(algorithm=algorithm, encoded_point=encoded_verification_key)
+        return verifier.key_bytes()
+
+    def decrypt_materials(self, request):
+        """Obtains a plaintext data key from one or more encrypted data keys
+        using underlying master key provider.
+
+        :param request: decrypt materials request
+        :type request: aws_encryption_sdk.materials_managers.DecryptionMaterialsRequest
+        :returns: decryption materials
+        :rtype: aws_encryption_sdk.materials_managers.DecryptionMaterials
+        """
+        validate_commitment_policy_on_decrypt(request.commitment_policy, request.algorithm)
+
+        data_key = self.master_key_provider.decrypt_data_key_from_list(
+            encrypted_data_keys=request.encrypted_data_keys,
+            algorithm=request.algorithm,
+            encryption_context=request.encryption_context,
+        )
+        verification_key = self._load_verification_key_from_encryption_context(
+            algorithm=request.algorithm, encryption_context=request.encryption_context
+        )
+
+        return DecryptionMaterials(data_key=data_key, verification_key=verification_key)
+
+
+def encrypt_decrypt_with_v3_default_cmm(key_arn,
+                                        source_plaintext,
+                                        botocore_session):
+    """Encrypts and decrypts a string using a V3 default CMM.
+
+    :param str key_arn: Amazon Resource Name (ARN) of the KMS CMK
+    :param bytes source_plaintext: Data to encrypt
+    :param botocore_session: existing botocore session instance
+    :type botocore_session: botocore.session.Session
+    """
+    # 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)
+
+    # Create a KMS master key provider. Note that because we are planning on decrypting using this same provider,
+    # we MUST provide the ARN of the KMS Key. If we provide a raw key id or a key alias, decryption will fail.
+    kms_kwargs = dict(key_ids=[key_arn])
+    if botocore_session is not None:
+        kms_kwargs["botocore_session"] = botocore_session
+    master_key_provider = aws_encryption_sdk.StrictAwsKmsMasterKeyProvider(**kms_kwargs)
+
+    default_cmm = V3DefaultCryptoMaterialsManager(master_key_provider=master_key_provider)
+
+    # Encrypt the plaintext source data
+    ciphertext, encryptor_header = client.encrypt(
+        source=source_plaintext, materials_manager=default_cmm
+    )
+
+    # Decrypt the ciphertext
+    cycled_plaintext, decrypted_header = client.decrypt(source=ciphertext, key_provider=master_key_provider)
+
+    # 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()
+    )
diff --git a/examples/test/legacy/examples_test_utils.py b/examples/test/legacy/examples_test_utils.py
index 3f877e301..8787e0f38 100644
--- a/examples/test/legacy/examples_test_utils.py
+++ b/examples/test/legacy/examples_test_utils.py
@@ -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,
 )
diff --git a/examples/test/legacy/test_i_v3_default_cmm.py b/examples/test/legacy/test_i_v3_default_cmm.py
new file mode 100644
index 000000000..2c48f1d1c
--- /dev/null
+++ b/examples/test/legacy/test_i_v3_default_cmm.py
@@ -0,0 +1,20 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Test suite for encryption and decryption using V3 defualt CMM."""
+
+import botocore.session
+import pytest
+
+from ...src.legacy.v3_default_cmm import encrypt_decrypt_with_v3_default_cmm
+from .examples_test_utils import get_cmk_arn, static_plaintext
+
+pytestmark = [pytest.mark.examples]
+
+
+def test_v3_default_cmm():
+    """Test method for encryption and decryption using V3 defualt CMM"""
+    plaintext = static_plaintext
+    cmk_arn = get_cmk_arn()
+    encrypt_decrypt_with_v3_default_cmm(key_arn=cmk_arn,
+                                        source_plaintext=plaintext,
+                                        botocore_session=botocore.session.Session())
diff --git a/examples/test/migration/test_i_migration_aws_kms_key_example.py b/examples/test/migration/test_i_migration_aws_kms_key_example.py
index d4cf49ce9..612a896ba 100644
--- a/examples/test/migration/test_i_migration_aws_kms_key_example.py
+++ b/examples/test/migration/test_i_migration_aws_kms_key_example.py
@@ -3,9 +3,7 @@
 """Test suite for the migration_aws_kms_key_example."""
 import pytest
 
-from ...src.migration.migration_aws_kms_key_example import (
-    migration_aws_kms_key,
-)
+from ...src.migration.migration_aws_kms_key_example import migration_aws_kms_key
 
 pytestmark = [pytest.mark.examples]
 
diff --git a/examples/test/migration/test_i_migration_raw_aes_key_example.py b/examples/test/migration/test_i_migration_raw_aes_key_example.py
index 7601e7dc0..d5e4f7789 100644
--- a/examples/test/migration/test_i_migration_raw_aes_key_example.py
+++ b/examples/test/migration/test_i_migration_raw_aes_key_example.py
@@ -3,9 +3,7 @@
 """Test suite for the migration_raw_aes_key_example."""
 import pytest
 
-from ...src.migration.migration_raw_aes_key_example import (
-    migration_raw_aes_key,
-)
+from ...src.migration.migration_raw_aes_key_example import migration_raw_aes_key
 
 pytestmark = [pytest.mark.examples]
 
diff --git a/examples/test/migration/test_i_migration_raw_rsa_key_example.py b/examples/test/migration/test_i_migration_raw_rsa_key_example.py
index 9e111d25a..238dcbaab 100644
--- a/examples/test/migration/test_i_migration_raw_rsa_key_example.py
+++ b/examples/test/migration/test_i_migration_raw_rsa_key_example.py
@@ -3,9 +3,7 @@
 """Test suite for the migration_raw_rsa_key_example."""
 import pytest
 
-from ...src.migration.migration_raw_rsa_key_example import (
-    migration_raw_rsa_key,
-)
+from ...src.migration.migration_raw_rsa_key_example import migration_raw_rsa_key
 
 pytestmark = [pytest.mark.examples]
 

From cdfc30c206ffd0a9b701af1e8e58c8c50112e282 Mon Sep 17 00:00:00 2001
From: Ritvik Kapila <ritvik.iitd@gmail.com>
Date: Tue, 18 Jun 2024 11:55:55 -0700
Subject: [PATCH 02/25] fix flake8

---
 examples/test/legacy/test_i_v3_default_cmm.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/test/legacy/test_i_v3_default_cmm.py b/examples/test/legacy/test_i_v3_default_cmm.py
index 2c48f1d1c..e48ab8118 100644
--- a/examples/test/legacy/test_i_v3_default_cmm.py
+++ b/examples/test/legacy/test_i_v3_default_cmm.py
@@ -12,7 +12,7 @@
 
 
 def test_v3_default_cmm():
-    """Test method for encryption and decryption using V3 defualt CMM"""
+    """Test method for encryption and decryption using V3 defualt CMM."""
     plaintext = static_plaintext
     cmk_arn = get_cmk_arn()
     encrypt_decrypt_with_v3_default_cmm(key_arn=cmk_arn,

From 6bfc02fe527ef23a06d509eb38e9d51fd5cbe4f0 Mon Sep 17 00:00:00 2001
From: Ritvik Kapila <ritvik.iitd@gmail.com>
Date: Tue, 18 Jun 2024 12:28:43 -0700
Subject: [PATCH 03/25] update v3_default_cmm

---
 examples/src/legacy/v3_default_cmm.py | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/examples/src/legacy/v3_default_cmm.py b/examples/src/legacy/v3_default_cmm.py
index 61f70e05d..466081a5d 100644
--- a/examples/src/legacy/v3_default_cmm.py
+++ b/examples/src/legacy/v3_default_cmm.py
@@ -172,22 +172,26 @@ def encrypt_decrypt_with_v3_default_cmm(key_arn,
     # commitment policy, REQUIRE_ENCRYPT_REQUIRE_DECRYPT is used by default.
     client = aws_encryption_sdk.EncryptionSDKClient(commitment_policy=CommitmentPolicy.REQUIRE_ENCRYPT_REQUIRE_DECRYPT)
 
-    # Create a KMS master key provider. Note that because we are planning on decrypting using this same provider,
-    # we MUST provide the ARN of the KMS Key. If we provide a raw key id or a key alias, decryption will fail.
+    # Create a KMS master key provider.
     kms_kwargs = dict(key_ids=[key_arn])
     if botocore_session is not None:
         kms_kwargs["botocore_session"] = botocore_session
     master_key_provider = aws_encryption_sdk.StrictAwsKmsMasterKeyProvider(**kms_kwargs)
 
+    # Create the V3 default CMM (V3DefaultCryptoMaterialsManager) using the master_key_provider
     default_cmm = V3DefaultCryptoMaterialsManager(master_key_provider=master_key_provider)
 
     # Encrypt the plaintext source data
     ciphertext, encryptor_header = client.encrypt(
-        source=source_plaintext, materials_manager=default_cmm
+        source=source_plaintext,
+        materials_manager=default_cmm
     )
 
     # Decrypt the ciphertext
-    cycled_plaintext, decrypted_header = client.decrypt(source=ciphertext, key_provider=master_key_provider)
+    cycled_plaintext, decrypted_header = client.decrypt(
+        source=ciphertext,
+        key_provider=master_key_provider
+    )
 
     # Verify that the "cycled" (encrypted, then decrypted) plaintext is identical to the source plaintext
     assert cycled_plaintext == source_plaintext

From c296556f3f5e9ce954957f69fe06240d164ca923 Mon Sep 17 00:00:00 2001
From: Ritvik Kapila <ritvik.iitd@gmail.com>
Date: Tue, 18 Jun 2024 15:47:43 -0700
Subject: [PATCH 04/25] added CustomSigningSuiteOnlyCMM; refactoring

---
 examples/src/legacy/custom_cmm_example.py     | 86 +++++++++++++++++++
 examples/test/legacy/test_i_v3_default_cmm.py | 43 ++++++++--
 .../{src => test}/legacy/v3_default_cmm.py    | 51 +----------
 3 files changed, 124 insertions(+), 56 deletions(-)
 create mode 100644 examples/src/legacy/custom_cmm_example.py
 rename examples/{src => test}/legacy/v3_default_cmm.py (76%)

diff --git a/examples/src/legacy/custom_cmm_example.py b/examples/src/legacy/custom_cmm_example.py
new file mode 100644
index 000000000..921dbe106
--- /dev/null
+++ b/examples/src/legacy/custom_cmm_example.py
@@ -0,0 +1,86 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Example to create a custom crypto material manager class."""
+
+import aws_encryption_sdk
+from aws_encryption_sdk import CommitmentPolicy
+from aws_encryption_sdk.materials_managers.base import CryptoMaterialsManager
+
+
+# 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, cmm: CryptoMaterialsManager) -> None:
+        super().__init__()
+        self.underlying_cmm = cmm
+
+    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 AssertionError(
+                "Algorithm provided to CustomSigningSuiteOnlyCMM"
+                  + " is not a supported signing algorithm: " + materials.algorithm
+                  )
+        return materials
+
+    def decrypt_materials(self, request):
+        """Provider decryption materials appropriate for the request.
+
+        :param DecryptionMaterialsRequest request: Request object to provide to a
+        crypto material manager's `decrypt_materials` method.
+        """
+        if not request.algorithm.is_signing():
+            raise AssertionError(
+                "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 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()
+    )
diff --git a/examples/test/legacy/test_i_v3_default_cmm.py b/examples/test/legacy/test_i_v3_default_cmm.py
index e48ab8118..5721debd5 100644
--- a/examples/test/legacy/test_i_v3_default_cmm.py
+++ b/examples/test/legacy/test_i_v3_default_cmm.py
@@ -1,20 +1,51 @@
 # Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
-"""Test suite for encryption and decryption using V3 defualt CMM."""
+"""Test suite for encryption and decryption using custom CMM."""
 
 import botocore.session
 import pytest
 
-from ...src.legacy.v3_default_cmm import encrypt_decrypt_with_v3_default_cmm
+import aws_encryption_sdk
+from ...src.legacy.custom_cmm_example import encrypt_decrypt_with_cmm, CustomSigningSuiteOnlyCMM
+from .v3_default_cmm import V3DefaultCryptoMaterialsManager
 from .examples_test_utils import get_cmk_arn, static_plaintext
 
 pytestmark = [pytest.mark.examples]
 
 
+def test_custom_cmm_example():
+    """Test method for encryption and decryption using V3 default CMM."""
+    plaintext = static_plaintext
+    cmk_arn = get_cmk_arn()
+    botocore_session = botocore.session.Session()
+
+    # Create a KMS master key provider.
+    kms_kwargs = dict(key_ids=[cmk_arn])
+    if botocore_session is not None:
+        kms_kwargs["botocore_session"] = botocore_session
+    master_key_provider = aws_encryption_sdk.StrictAwsKmsMasterKeyProvider(**kms_kwargs)
+
+    # Create the V3 default CMM (V3DefaultCryptoMaterialsManager) using the master_key_provider
+    cmm = CustomSigningSuiteOnlyCMM(master_key_provider=master_key_provider)
+
+    encrypt_decrypt_with_cmm(cmm=cmm,
+                             source_plaintext=plaintext)
+
+
 def test_v3_default_cmm():
-    """Test method for encryption and decryption using V3 defualt CMM."""
+    """Test method for encryption and decryption using V3 default CMM."""
     plaintext = static_plaintext
     cmk_arn = get_cmk_arn()
-    encrypt_decrypt_with_v3_default_cmm(key_arn=cmk_arn,
-                                        source_plaintext=plaintext,
-                                        botocore_session=botocore.session.Session())
+    botocore_session = botocore.session.Session()
+
+    # Create a KMS master key provider.
+    kms_kwargs = dict(key_ids=[cmk_arn])
+    if botocore_session is not None:
+        kms_kwargs["botocore_session"] = botocore_session
+    master_key_provider = aws_encryption_sdk.StrictAwsKmsMasterKeyProvider(**kms_kwargs)
+
+    # Create the V3 default CMM (V3DefaultCryptoMaterialsManager) using the master_key_provider
+    cmm = V3DefaultCryptoMaterialsManager(master_key_provider=master_key_provider)
+
+    encrypt_decrypt_with_cmm(cmm=cmm,
+                             source_plaintext=plaintext)
diff --git a/examples/src/legacy/v3_default_cmm.py b/examples/test/legacy/v3_default_cmm.py
similarity index 76%
rename from examples/src/legacy/v3_default_cmm.py
rename to examples/test/legacy/v3_default_cmm.py
index 466081a5d..3229499f2 100644
--- a/examples/src/legacy/v3_default_cmm.py
+++ b/examples/test/legacy/v3_default_cmm.py
@@ -1,11 +1,10 @@
 # Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
-"""Default crypto material manager class example for ESDK V3."""
+"""Copy-paste of the V3 default CMM."""
 import logging
 
 import attr
 
-import aws_encryption_sdk
 from aws_encryption_sdk.exceptions import MasterKeyProviderError, SerializationError
 from aws_encryption_sdk.identifiers import CommitmentPolicy
 from aws_encryption_sdk.internal.crypto.authentication import Signer, Verifier
@@ -156,51 +155,3 @@ def decrypt_materials(self, request):
         )
 
         return DecryptionMaterials(data_key=data_key, verification_key=verification_key)
-
-
-def encrypt_decrypt_with_v3_default_cmm(key_arn,
-                                        source_plaintext,
-                                        botocore_session):
-    """Encrypts and decrypts a string using a V3 default CMM.
-
-    :param str key_arn: Amazon Resource Name (ARN) of the KMS CMK
-    :param bytes source_plaintext: Data to encrypt
-    :param botocore_session: existing botocore session instance
-    :type botocore_session: botocore.session.Session
-    """
-    # 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)
-
-    # Create a KMS master key provider.
-    kms_kwargs = dict(key_ids=[key_arn])
-    if botocore_session is not None:
-        kms_kwargs["botocore_session"] = botocore_session
-    master_key_provider = aws_encryption_sdk.StrictAwsKmsMasterKeyProvider(**kms_kwargs)
-
-    # Create the V3 default CMM (V3DefaultCryptoMaterialsManager) using the master_key_provider
-    default_cmm = V3DefaultCryptoMaterialsManager(master_key_provider=master_key_provider)
-
-    # Encrypt the plaintext source data
-    ciphertext, encryptor_header = client.encrypt(
-        source=source_plaintext,
-        materials_manager=default_cmm
-    )
-
-    # Decrypt the ciphertext
-    cycled_plaintext, decrypted_header = client.decrypt(
-        source=ciphertext,
-        key_provider=master_key_provider
-    )
-
-    # 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()
-    )

From 1d47afaf461eb77f1a1c0102188a9358c012eea4 Mon Sep 17 00:00:00 2001
From: Ritvik Kapila <ritvik.iitd@gmail.com>
Date: Tue, 18 Jun 2024 15:49:37 -0700
Subject: [PATCH 05/25] fix

---
 ...test_i_v3_default_cmm.py => test_i_custom_cmm_example.py} | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)
 rename examples/test/legacy/{test_i_v3_default_cmm.py => test_i_custom_cmm_example.py} (95%)

diff --git a/examples/test/legacy/test_i_v3_default_cmm.py b/examples/test/legacy/test_i_custom_cmm_example.py
similarity index 95%
rename from examples/test/legacy/test_i_v3_default_cmm.py
rename to examples/test/legacy/test_i_custom_cmm_example.py
index 5721debd5..5972062d7 100644
--- a/examples/test/legacy/test_i_v3_default_cmm.py
+++ b/examples/test/legacy/test_i_custom_cmm_example.py
@@ -6,9 +6,10 @@
 import pytest
 
 import aws_encryption_sdk
-from ...src.legacy.custom_cmm_example import encrypt_decrypt_with_cmm, CustomSigningSuiteOnlyCMM
-from .v3_default_cmm import V3DefaultCryptoMaterialsManager
+
+from ...src.legacy.custom_cmm_example import CustomSigningSuiteOnlyCMM, encrypt_decrypt_with_cmm
 from .examples_test_utils import get_cmk_arn, static_plaintext
+from .v3_default_cmm import V3DefaultCryptoMaterialsManager
 
 pytestmark = [pytest.mark.examples]
 

From 0401aeeb3b9ec0cda416a714ee0e963535d27c86 Mon Sep 17 00:00:00 2001
From: Ritvik Kapila <ritvik.iitd@gmail.com>
Date: Tue, 18 Jun 2024 16:02:33 -0700
Subject: [PATCH 06/25] fix

---
 examples/src/legacy/custom_cmm_example.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/examples/src/legacy/custom_cmm_example.py b/examples/src/legacy/custom_cmm_example.py
index 921dbe106..27b462c94 100644
--- a/examples/src/legacy/custom_cmm_example.py
+++ b/examples/src/legacy/custom_cmm_example.py
@@ -3,8 +3,9 @@
 """Example to create a custom crypto material manager class."""
 
 import aws_encryption_sdk
-from aws_encryption_sdk import CommitmentPolicy
+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.
@@ -14,9 +15,9 @@
 class CustomSigningSuiteOnlyCMM(CryptoMaterialsManager):
     """Example custom crypto materials manager class."""
 
-    def __init__(self, cmm: CryptoMaterialsManager) -> None:
-        super().__init__()
-        self.underlying_cmm = cmm
+    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.

From 0f53225243190aa69dbe540c2da218eaa2c5e7c5 Mon Sep 17 00:00:00 2001
From: Ritvik Kapila <ritvik.iitd@gmail.com>
Date: Tue, 18 Jun 2024 16:15:37 -0700
Subject: [PATCH 07/25] fix flake8, pylint

---
 examples/test/legacy/v3_default_cmm.py | 4 ++--
 tox.ini                                | 6 ++++--
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/examples/test/legacy/v3_default_cmm.py b/examples/test/legacy/v3_default_cmm.py
index 3229499f2..44c4e72c5 100644
--- a/examples/test/legacy/v3_default_cmm.py
+++ b/examples/test/legacy/v3_default_cmm.py
@@ -40,7 +40,7 @@ class V3DefaultCryptoMaterialsManager(CryptoMaterialsManager):
 
     master_key_provider = attr.ib(validator=attr.validators.instance_of(MasterKeyProvider))
 
-    def _generate_signing_key_and_update_encryption_context(self, algorithm, encryption_context):
+    def _generate_signing_key_and_update_encryption_context(self, algorithm, encryption_context):  # pylint: disable=no-self-use
         """Generates a signing key based on the provided algorithm.
 
         :param algorithm: Algorithm for which to generate signing key
@@ -111,7 +111,7 @@ def get_encryption_materials(self, request):
             signing_key=signing_key,
         )
 
-    def _load_verification_key_from_encryption_context(self, algorithm, encryption_context):
+    def _load_verification_key_from_encryption_context(self, algorithm, encryption_context):  # pylint: disable=no-self-use
         """Loads the verification key from the encryption context if used by algorithm suite.
 
         :param algorithm: Algorithm for which to generate signing key
diff --git a/tox.ini b/tox.ini
index 9152f51a6..fc3a8035d 100644
--- a/tox.ini
+++ b/tox.ini
@@ -169,10 +169,12 @@ deps = {[testenv:flake8]deps}
 commands =
     flake8 examples/src/
     flake8 \
-        # Ingore D103 missing docstring errors in tests (test names should be self-documenting)
+        # Ignore D103 missing docstring errors in tests (test names should be self-documenting)
+        # Ignore D401 First line should be in imperative mood
+        # Ignoring D205, D400 because of false positives
         # E203 is not PEP8 compliant https://github.com/ambv/black#slices
         # W503 is not PEP8 compliant https://github.com/ambv/black#line-breaks--binary-operators
-        --ignore D103,E203,W503 \
+        --ignore D103,D401,D205,D400,E203,W503 \
         examples/test/
 
 [testenv:pylint]

From f1e6c91714947f1365f7ccf9616c3140fc435d96 Mon Sep 17 00:00:00 2001
From: Ritvik Kapila <ritvik.iitd@gmail.com>
Date: Wed, 19 Jun 2024 09:45:07 -0700
Subject: [PATCH 08/25] pylint fix

---
 examples/test/legacy/v3_default_cmm.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/examples/test/legacy/v3_default_cmm.py b/examples/test/legacy/v3_default_cmm.py
index 44c4e72c5..5c41e7975 100644
--- a/examples/test/legacy/v3_default_cmm.py
+++ b/examples/test/legacy/v3_default_cmm.py
@@ -40,7 +40,8 @@ class V3DefaultCryptoMaterialsManager(CryptoMaterialsManager):
 
     master_key_provider = attr.ib(validator=attr.validators.instance_of(MasterKeyProvider))
 
-    def _generate_signing_key_and_update_encryption_context(self, algorithm, encryption_context):  # pylint: disable=no-self-use
+# pylint: disable=no-self-use
+    def _generate_signing_key_and_update_encryption_context(self, algorithm, encryption_context):
         """Generates a signing key based on the provided algorithm.
 
         :param algorithm: Algorithm for which to generate signing key
@@ -111,7 +112,8 @@ def get_encryption_materials(self, request):
             signing_key=signing_key,
         )
 
-    def _load_verification_key_from_encryption_context(self, algorithm, encryption_context):  # pylint: disable=no-self-use
+# pylint: disable=no-self-use
+    def _load_verification_key_from_encryption_context(self, algorithm, encryption_context):
         """Loads the verification key from the encryption context if used by algorithm suite.
 
         :param algorithm: Algorithm for which to generate signing key

From 3de65cb8d31534352bc60a4e60542f040e9af126 Mon Sep 17 00:00:00 2001
From: Ritvik Kapila <ritvik.iitd@gmail.com>
Date: Thu, 20 Jun 2024 10:03:22 -0700
Subject: [PATCH 09/25] remove tox flake8 ignores

---
 tox.ini | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/tox.ini b/tox.ini
index fc3a8035d..434f7e45d 100644
--- a/tox.ini
+++ b/tox.ini
@@ -170,11 +170,9 @@ commands =
     flake8 examples/src/
     flake8 \
         # Ignore D103 missing docstring errors in tests (test names should be self-documenting)
-        # Ignore D401 First line should be in imperative mood
-        # Ignoring D205, D400 because of false positives
         # E203 is not PEP8 compliant https://github.com/ambv/black#slices
         # W503 is not PEP8 compliant https://github.com/ambv/black#line-breaks--binary-operators
-        --ignore D103,D401,D205,D400,E203,W503 \
+        --ignore D103,E203,W503 \
         examples/test/
 
 [testenv:pylint]

From 7bbe83e24a2b4ef0822c82c183cc7a9143bb0267 Mon Sep 17 00:00:00 2001
From: Ritvik Kapila <ritvik.iitd@gmail.com>
Date: Thu, 20 Jun 2024 10:08:07 -0700
Subject: [PATCH 10/25] fix

---
 examples/test/legacy/v3_default_cmm.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/examples/test/legacy/v3_default_cmm.py b/examples/test/legacy/v3_default_cmm.py
index 5c41e7975..d7956f1e3 100644
--- a/examples/test/legacy/v3_default_cmm.py
+++ b/examples/test/legacy/v3_default_cmm.py
@@ -1,6 +1,7 @@
 # Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 """Copy-paste of the V3 default CMM."""
+# flake8: ignore=D205,D400,D401
 import logging
 
 import attr

From 4f50576c0bac8f883d7387c7e05f3a37e637abc6 Mon Sep 17 00:00:00 2001
From: Ritvik Kapila <ritvik.iitd@gmail.com>
Date: Thu, 20 Jun 2024 10:17:36 -0700
Subject: [PATCH 11/25] fix

---
 examples/test/legacy/v3_default_cmm.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/test/legacy/v3_default_cmm.py b/examples/test/legacy/v3_default_cmm.py
index d7956f1e3..2ca9c688e 100644
--- a/examples/test/legacy/v3_default_cmm.py
+++ b/examples/test/legacy/v3_default_cmm.py
@@ -1,7 +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."""
-# flake8: ignore=D205,D400,D401
 import logging
 
 import attr
@@ -43,6 +42,7 @@ class V3DefaultCryptoMaterialsManager(CryptoMaterialsManager):
 
 # pylint: disable=no-self-use
     def _generate_signing_key_and_update_encryption_context(self, algorithm, encryption_context):
+        # noqa: F401
         """Generates a signing key based on the provided algorithm.
 
         :param algorithm: Algorithm for which to generate signing key

From 1e1b62ce8b2ce80e2ec76e9c55b7901008996bc6 Mon Sep 17 00:00:00 2001
From: Ritvik Kapila <ritvik.iitd@gmail.com>
Date: Thu, 20 Jun 2024 10:20:43 -0700
Subject: [PATCH 12/25] fix(2)

---
 examples/test/legacy/v3_default_cmm.py | 1 -
 tox.ini                                | 2 ++
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/examples/test/legacy/v3_default_cmm.py b/examples/test/legacy/v3_default_cmm.py
index 2ca9c688e..5c41e7975 100644
--- a/examples/test/legacy/v3_default_cmm.py
+++ b/examples/test/legacy/v3_default_cmm.py
@@ -42,7 +42,6 @@ class V3DefaultCryptoMaterialsManager(CryptoMaterialsManager):
 
 # pylint: disable=no-self-use
     def _generate_signing_key_and_update_encryption_context(self, algorithm, encryption_context):
-        # noqa: F401
         """Generates a signing key based on the provided algorithm.
 
         :param algorithm: Algorithm for which to generate signing key
diff --git a/tox.ini b/tox.ini
index 434f7e45d..feb26fb74 100644
--- a/tox.ini
+++ b/tox.ini
@@ -173,6 +173,8 @@ commands =
         # E203 is not PEP8 compliant https://github.com/ambv/black#slices
         # W503 is not PEP8 compliant https://github.com/ambv/black#line-breaks--binary-operators
         --ignore D103,E203,W503 \
+        per-file-ignores =
+            v3_default_cmm.py: D401
         examples/test/
 
 [testenv:pylint]

From bdd0b124cf834bd29ac5cbef2a80271ed8b44173 Mon Sep 17 00:00:00 2001
From: Ritvik Kapila <ritvik.iitd@gmail.com>
Date: Thu, 20 Jun 2024 10:24:28 -0700
Subject: [PATCH 13/25] fix

---
 tox.ini | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tox.ini b/tox.ini
index feb26fb74..bf0bf82fb 100644
--- a/tox.ini
+++ b/tox.ini
@@ -174,7 +174,7 @@ commands =
         # W503 is not PEP8 compliant https://github.com/ambv/black#line-breaks--binary-operators
         --ignore D103,E203,W503 \
         per-file-ignores =
-            v3_default_cmm.py: D401
+            examples/test/legacy/v3_default_cmm.py: D401
         examples/test/
 
 [testenv:pylint]

From a4c8670cd28fd42154842b1d77d3f6c2c4d2ce4e Mon Sep 17 00:00:00 2001
From: Ritvik Kapila <ritvik.iitd@gmail.com>
Date: Thu, 20 Jun 2024 10:27:44 -0700
Subject: [PATCH 14/25] fix

---
 tox.ini | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tox.ini b/tox.ini
index bf0bf82fb..c9699489e 100644
--- a/tox.ini
+++ b/tox.ini
@@ -173,8 +173,7 @@ commands =
         # E203 is not PEP8 compliant https://github.com/ambv/black#slices
         # W503 is not PEP8 compliant https://github.com/ambv/black#line-breaks--binary-operators
         --ignore D103,E203,W503 \
-        per-file-ignores =
-            examples/test/legacy/v3_default_cmm.py: D401
+        --per-file-ignores = examples/test/legacy/v3_default_cmm.py: D401 \
         examples/test/
 
 [testenv:pylint]

From 87c7fdf1d5ed584760c04a7ed1833a50f14854eb Mon Sep 17 00:00:00 2001
From: Ritvik Kapila <ritvik.iitd@gmail.com>
Date: Thu, 20 Jun 2024 10:30:17 -0700
Subject: [PATCH 15/25] fix

---
 tox.ini | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tox.ini b/tox.ini
index c9699489e..e0a4f794e 100644
--- a/tox.ini
+++ b/tox.ini
@@ -173,7 +173,7 @@ commands =
         # E203 is not PEP8 compliant https://github.com/ambv/black#slices
         # W503 is not PEP8 compliant https://github.com/ambv/black#line-breaks--binary-operators
         --ignore D103,E203,W503 \
-        --per-file-ignores = examples/test/legacy/v3_default_cmm.py: D401 \
+        --per-file-ignores 'examples/test/legacy/v3_default_cmm.py: D401' \
         examples/test/
 
 [testenv:pylint]

From fcbd4f8e067311f38d534464853b960454b1a7e5 Mon Sep 17 00:00:00 2001
From: Ritvik Kapila <ritvik.iitd@gmail.com>
Date: Thu, 20 Jun 2024 10:31:58 -0700
Subject: [PATCH 16/25] fix

---
 tox.ini | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tox.ini b/tox.ini
index e0a4f794e..40ee437ed 100644
--- a/tox.ini
+++ b/tox.ini
@@ -173,7 +173,7 @@ commands =
         # E203 is not PEP8 compliant https://github.com/ambv/black#slices
         # W503 is not PEP8 compliant https://github.com/ambv/black#line-breaks--binary-operators
         --ignore D103,E203,W503 \
-        --per-file-ignores 'examples/test/legacy/v3_default_cmm.py: D401' \
+        --per-file-ignores 'examples/test/legacy/v3_default_cmm.py: D205,D400,D401' \
         examples/test/
 
 [testenv:pylint]

From 874f5be4caf470b6ca3feb41ce233eca04d247a6 Mon Sep 17 00:00:00 2001
From: Ritvik Kapila <ritvik.iitd@gmail.com>
Date: Thu, 20 Jun 2024 12:53:28 -0700
Subject: [PATCH 17/25] resolve comments

---
 examples/src/legacy/custom_cmm_example.py         | 2 +-
 examples/test/legacy/test_i_custom_cmm_example.py | 2 +-
 examples/test/legacy/v3_default_cmm.py            | 2 +-
 tox.ini                                           | 1 +
 4 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/examples/src/legacy/custom_cmm_example.py b/examples/src/legacy/custom_cmm_example.py
index 27b462c94..325d879f5 100644
--- a/examples/src/legacy/custom_cmm_example.py
+++ b/examples/src/legacy/custom_cmm_example.py
@@ -1,6 +1,6 @@
 # Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
-"""Example to create a custom crypto material manager class."""
+"""Example to create a custom implementation of the native ESDK CryptoMaterialsManager class."""
 
 import aws_encryption_sdk
 from aws_encryption_sdk import CommitmentPolicy, StrictAwsKmsMasterKeyProvider
diff --git a/examples/test/legacy/test_i_custom_cmm_example.py b/examples/test/legacy/test_i_custom_cmm_example.py
index 5972062d7..862b75106 100644
--- a/examples/test/legacy/test_i_custom_cmm_example.py
+++ b/examples/test/legacy/test_i_custom_cmm_example.py
@@ -26,7 +26,7 @@ def test_custom_cmm_example():
         kms_kwargs["botocore_session"] = botocore_session
     master_key_provider = aws_encryption_sdk.StrictAwsKmsMasterKeyProvider(**kms_kwargs)
 
-    # Create the V3 default CMM (V3DefaultCryptoMaterialsManager) using the master_key_provider
+    # Create the custom signing CMM using the master_key_provider
     cmm = CustomSigningSuiteOnlyCMM(master_key_provider=master_key_provider)
 
     encrypt_decrypt_with_cmm(cmm=cmm,
diff --git a/examples/test/legacy/v3_default_cmm.py b/examples/test/legacy/v3_default_cmm.py
index 5c41e7975..544147d6b 100644
--- a/examples/test/legacy/v3_default_cmm.py
+++ b/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."""
+"""Copy-paste of the V3 default CMM with small changes to pass linters.."""
 import logging
 
 import attr
diff --git a/tox.ini b/tox.ini
index 40ee437ed..28be94a63 100644
--- a/tox.ini
+++ b/tox.ini
@@ -173,6 +173,7 @@ commands =
         # E203 is not PEP8 compliant https://github.com/ambv/black#slices
         # W503 is not PEP8 compliant https://github.com/ambv/black#line-breaks--binary-operators
         --ignore D103,E203,W503 \
+        # copy-paste test for v3_default_cmm; intentionally not changing code
         --per-file-ignores 'examples/test/legacy/v3_default_cmm.py: D205,D400,D401' \
         examples/test/
 

From 1a4b447742a1a75c9816fa6b8e9f3d52716897ea Mon Sep 17 00:00:00 2001
From: Ritvik Kapila <ritvik.iitd@gmail.com>
Date: Thu, 20 Jun 2024 15:43:52 -0700
Subject: [PATCH 18/25] added MPL CMM example and test

---
 ...cryptographic_materials_manager_example.py | 129 ++++++++++++++++++
 ...cryptographic_materials_manager_example.py |  14 ++
 2 files changed, 143 insertions(+)
 create mode 100644 examples/src/aws_cryptographic_materials_manager_example.py
 create mode 100644 examples/test/test_i_aws_cryptographic_materials_manager_example.py

diff --git a/examples/src/aws_cryptographic_materials_manager_example.py b/examples/src/aws_cryptographic_materials_manager_example.py
new file mode 100644
index 000000000..6ffac188d
--- /dev/null
+++ b/examples/src/aws_cryptographic_materials_manager_example.py
@@ -0,0 +1,129 @@
+# 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).
+
+The AWS 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. 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 AWS 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 CreateDefaultCryptographicMaterialsManagerInput, CreateAwsKmsKeyringInput
+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_cmm(
+    kms_key_id: str
+):
+    """Demonstrate an encrypt/decrypt cycle using an AWS KMS keyring.
+
+    Usage: encrypt_and_decrypt_with_keyring(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 a boto3 client for KMS.
+    kms_client = boto3.client('kms', region_name="us-west-2")
+
+    # 3. 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",
+    }
+
+    # 4. Create a KMS keyring to use with the CryptographicMaterialsManager
+    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
+    )
+
+    # 5. 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 keyring 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"
diff --git a/examples/test/test_i_aws_cryptographic_materials_manager_example.py b/examples/test/test_i_aws_cryptographic_materials_manager_example.py
new file mode 100644
index 000000000..7983b8718
--- /dev/null
+++ b/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 AWS Cryptographic Materials Manager example."""
+import pytest
+
+from ..src.aws_cryptographic_materials_manager_example import encrypt_and_decrypt_with_cmm
+
+pytestmark = [pytest.mark.examples]
+
+
+def test_encrypt_and_decrypt_with_cmm():
+    """Test function for encrypt and decrypt using the AWS Cryptographic Materials Manager example."""
+    kms_key_id = "arn:aws:kms:us-west-2:658956600833:key/b3537ef1-d8dc-4780-9f5a-55776cbb2f7f"
+    encrypt_and_decrypt_with_cmm(kms_key_id)

From bedf3cd1354da21df98333ba2248607068a8f6c9 Mon Sep 17 00:00:00 2001
From: Ritvik Kapila <ritvik.iitd@gmail.com>
Date: Thu, 20 Jun 2024 15:47:04 -0700
Subject: [PATCH 19/25] fix linters

---
 examples/src/aws_cryptographic_materials_manager_example.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/examples/src/aws_cryptographic_materials_manager_example.py b/examples/src/aws_cryptographic_materials_manager_example.py
index 6ffac188d..da18d6c9f 100644
--- a/examples/src/aws_cryptographic_materials_manager_example.py
+++ b/examples/src/aws_cryptographic_materials_manager_example.py
@@ -21,7 +21,10 @@
 import boto3
 from aws_cryptographic_materialproviders.mpl import AwsCryptographicMaterialProviders
 from aws_cryptographic_materialproviders.mpl.config import MaterialProvidersConfig
-from aws_cryptographic_materialproviders.mpl.models import CreateDefaultCryptographicMaterialsManagerInput, CreateAwsKmsKeyringInput
+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
 

From 7512802a35e7bb1906b9b0bed293c0b393013716 Mon Sep 17 00:00:00 2001
From: Ritvik Kapila <ritvik.iitd@gmail.com>
Date: Thu, 20 Jun 2024 15:57:21 -0700
Subject: [PATCH 20/25] minor fix

---
 ...cryptographic_materials_manager_example.py | 25 ++++++++++---------
 1 file changed, 13 insertions(+), 12 deletions(-)

diff --git a/examples/src/aws_cryptographic_materials_manager_example.py b/examples/src/aws_cryptographic_materials_manager_example.py
index da18d6c9f..d358c87e0 100644
--- a/examples/src/aws_cryptographic_materials_manager_example.py
+++ b/examples/src/aws_cryptographic_materials_manager_example.py
@@ -7,7 +7,9 @@
 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. This example also includes some sanity checks for demonstration:
+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
@@ -42,9 +44,9 @@
 def encrypt_and_decrypt_with_cmm(
     kms_key_id: str
 ):
-    """Demonstrate an encrypt/decrypt cycle using an AWS KMS keyring.
+    """Demonstrate an encrypt/decrypt cycle using an AWS Cryptographic Material Managers.
 
-    Usage: encrypt_and_decrypt_with_keyring(kms_key_id)
+    Usage: encrypt_and_decrypt_with_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
@@ -63,10 +65,7 @@ def encrypt_and_decrypt_with_cmm(
         commitment_policy=CommitmentPolicy.REQUIRE_ENCRYPT_REQUIRE_DECRYPT
     )
 
-    # 2. Create a boto3 client for KMS.
-    kms_client = boto3.client('kms', region_name="us-west-2")
-
-    # 3. Create encryption context.
+    # 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
@@ -79,6 +78,8 @@ def encrypt_and_decrypt_with_cmm(
     }
 
     # 4. 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()
     )
@@ -102,31 +103,31 @@ def encrypt_and_decrypt_with_cmm(
         input=cmm_input
     )
 
-    # 5. Encrypt the data with the encryptionContext.
+    # 6. 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.
+    # 7. 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 keyring you used on encrypt.
+    # 8. 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
+    # 9. 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.
+    # 10. 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"

From 8e0a31730cfcfaaf485cb36d5c83e8c0affea89b Mon Sep 17 00:00:00 2001
From: Ritvik Kapila <ritvik.iitd@gmail.com>
Date: Fri, 21 Jun 2024 13:40:02 -0700
Subject: [PATCH 21/25] added custom mpl cmm example

---
 examples/src/custom_mpl_cmm_example.py        | 112 ++++++++++++++++++
 examples/src/legacy/custom_cmm_example.py     |   8 +-
 .../test/legacy/test_i_custom_cmm_example.py  |   1 -
 .../test/test_i_custom_mpl_cmm_example.py     |  39 ++++++
 4 files changed, 155 insertions(+), 5 deletions(-)
 create mode 100644 examples/src/custom_mpl_cmm_example.py
 create mode 100644 examples/test/test_i_custom_mpl_cmm_example.py

diff --git a/examples/src/custom_mpl_cmm_example.py b/examples/src/custom_mpl_cmm_example.py
new file mode 100644
index 000000000..3fd77cf13
--- /dev/null
+++ b/examples/src/custom_mpl_cmm_example.py
@@ -0,0 +1,112 @@
+# 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."""
+
+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) -> None:
+        """Constructor for MPLCustomSigningSuiteOnlyCMM class."""
+        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
+    :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=EXAMPLE_DATA,
+        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 == 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"
diff --git a/examples/src/legacy/custom_cmm_example.py b/examples/src/legacy/custom_cmm_example.py
index 325d879f5..07e8ca50b 100644
--- a/examples/src/legacy/custom_cmm_example.py
+++ b/examples/src/legacy/custom_cmm_example.py
@@ -29,20 +29,20 @@ def get_encryption_materials(self, request):
         """
         materials = self.underlying_cmm.get_encryption_materials(request)
         if not materials.algorithm.is_signing():
-            raise AssertionError(
+            raise ValueError(
                 "Algorithm provided to CustomSigningSuiteOnlyCMM"
                   + " is not a supported signing algorithm: " + materials.algorithm
                   )
         return materials
 
     def decrypt_materials(self, request):
-        """Provider decryption materials appropriate for the 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 AssertionError(
+            raise ValueError(
                 "Algorithm provided to CustomSigningSuiteOnlyCMM"
                   + " is not a supported signing algorithm: " + request.algorithm
                   )
@@ -53,7 +53,7 @@ def encrypt_decrypt_with_cmm(
     cmm: CryptoMaterialsManager,
     source_plaintext: str
 ):
-    """Encrypts and decrypts a string using a CMM.
+    """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
diff --git a/examples/test/legacy/test_i_custom_cmm_example.py b/examples/test/legacy/test_i_custom_cmm_example.py
index 862b75106..397230090 100644
--- a/examples/test/legacy/test_i_custom_cmm_example.py
+++ b/examples/test/legacy/test_i_custom_cmm_example.py
@@ -1,7 +1,6 @@
 # Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 """Test suite for encryption and decryption using custom CMM."""
-
 import botocore.session
 import pytest
 
diff --git a/examples/test/test_i_custom_mpl_cmm_example.py b/examples/test/test_i_custom_mpl_cmm_example.py
new file mode 100644
index 000000000..d98b6b6b9
--- /dev/null
+++ b/examples/test/test_i_custom_mpl_cmm_example.py
@@ -0,0 +1,39 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Test suite for encryption and decryption using custom CMM."""
+import boto3
+import pytest
+from aws_cryptographic_materialproviders.mpl import AwsCryptographicMaterialProviders
+from aws_cryptographic_materialproviders.mpl.config import MaterialProvidersConfig
+from aws_cryptographic_materialproviders.mpl.models import CreateAwsKmsKeyringInput
+from aws_cryptographic_materialproviders.mpl.references import IKeyring
+
+from ..src.custom_mpl_cmm_example import MPLCustomSigningSuiteOnlyCMM, encrypt_decrypt_with_cmm
+
+pytestmark = [pytest.mark.examples]
+
+
+def test_custom_cmm_example():
+    """Test method for encryption and decryption using V3 default CMM."""
+    kms_key_id = "arn:aws:kms:us-west-2:658956600833:key/b3537ef1-d8dc-4780-9f5a-55776cbb2f7f"
+
+    # Create KMS keyring to use with the CMM
+    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
+    )
+
+    # Create the custom MPL signing CMM using the keyring
+    cmm = MPLCustomSigningSuiteOnlyCMM(keyring=kms_keyring)
+
+    encrypt_decrypt_with_cmm(cmm=cmm)

From d6497eb8fea07a2facee6a12d443f7488b7afbce Mon Sep 17 00:00:00 2001
From: Ritvik Kapila <ritvik.iitd@gmail.com>
Date: Fri, 21 Jun 2024 13:40:44 -0700
Subject: [PATCH 22/25] fix

---
 examples/src/custom_mpl_cmm_example.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/examples/src/custom_mpl_cmm_example.py b/examples/src/custom_mpl_cmm_example.py
index 3fd77cf13..a084aa55c 100644
--- a/examples/src/custom_mpl_cmm_example.py
+++ b/examples/src/custom_mpl_cmm_example.py
@@ -4,7 +4,10 @@
 
 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.models import (
+    CreateDefaultCryptographicMaterialsManagerInput,
+    SignatureAlgorithmNone,
+)
 from aws_cryptographic_materialproviders.mpl.references import ICryptographicMaterialsManager, IKeyring
 
 import aws_encryption_sdk

From 349b55924e1e4d2818c77a128be89ec81190183e Mon Sep 17 00:00:00 2001
From: Ritvik Kapila <ritvik.iitd@gmail.com>
Date: Fri, 21 Jun 2024 14:40:20 -0700
Subject: [PATCH 23/25] fix

---
 examples/src/custom_mpl_cmm_example.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/examples/src/custom_mpl_cmm_example.py b/examples/src/custom_mpl_cmm_example.py
index a084aa55c..a770a8b95 100644
--- a/examples/src/custom_mpl_cmm_example.py
+++ b/examples/src/custom_mpl_cmm_example.py
@@ -79,7 +79,6 @@ def encrypt_decrypt_with_cmm(
     """Encrypts and decrypts a string using a custom CMM.
 
     :param ICryptographicMaterialsManager 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.

From 8c5bb6345ddaa1b88cc6b3920df5782edf7bb512 Mon Sep 17 00:00:00 2001
From: Ritvik Kapila <ritvik.iitd@gmail.com>
Date: Tue, 25 Jun 2024 11:07:07 -0700
Subject: [PATCH 24/25] refactoring and comment fixes

---
 examples/src/custom_mpl_cmm_example.py        | 61 ++++++++++---------
 ...ryptographic_materials_manager_example.py} | 26 ++++----
 examples/test/legacy/v3_default_cmm.py        |  2 +-
 ...cryptographic_materials_manager_example.py | 14 -----
 ...cryptographic_materials_manager_example.py | 14 +++++
 5 files changed, 60 insertions(+), 57 deletions(-)
 rename examples/src/{aws_cryptographic_materials_manager_example.py => default_cryptographic_materials_manager_example.py} (84%)
 delete mode 100644 examples/test/test_i_aws_cryptographic_materials_manager_example.py
 create mode 100644 examples/test/test_i_default_cryptographic_materials_manager_example.py

diff --git a/examples/src/custom_mpl_cmm_example.py b/examples/src/custom_mpl_cmm_example.py
index a770a8b95..8b331ccf6 100644
--- a/examples/src/custom_mpl_cmm_example.py
+++ b/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"
diff --git a/examples/src/aws_cryptographic_materials_manager_example.py b/examples/src/default_cryptographic_materials_manager_example.py
similarity index 84%
rename from examples/src/aws_cryptographic_materials_manager_example.py
rename to examples/src/default_cryptographic_materials_manager_example.py
index d358c87e0..bfffc36d9 100644
--- a/examples/src/aws_cryptographic_materials_manager_example.py
+++ b/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"
diff --git a/examples/test/legacy/v3_default_cmm.py b/examples/test/legacy/v3_default_cmm.py
index 544147d6b..f077e26c9 100644
--- a/examples/test/legacy/v3_default_cmm.py
+++ b/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
diff --git a/examples/test/test_i_aws_cryptographic_materials_manager_example.py b/examples/test/test_i_aws_cryptographic_materials_manager_example.py
deleted file mode 100644
index 7983b8718..000000000
--- a/examples/test/test_i_aws_cryptographic_materials_manager_example.py
+++ /dev/null
@@ -1,14 +0,0 @@
-# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Test suite for the AWS Cryptographic Materials Manager example."""
-import pytest
-
-from ..src.aws_cryptographic_materials_manager_example import encrypt_and_decrypt_with_cmm
-
-pytestmark = [pytest.mark.examples]
-
-
-def test_encrypt_and_decrypt_with_cmm():
-    """Test function for encrypt and decrypt using the AWS Cryptographic Materials Manager example."""
-    kms_key_id = "arn:aws:kms:us-west-2:658956600833:key/b3537ef1-d8dc-4780-9f5a-55776cbb2f7f"
-    encrypt_and_decrypt_with_cmm(kms_key_id)
diff --git a/examples/test/test_i_default_cryptographic_materials_manager_example.py b/examples/test/test_i_default_cryptographic_materials_manager_example.py
new file mode 100644
index 000000000..8a18f655d
--- /dev/null
+++ b/examples/test/test_i_default_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)

From 320d5c3ada5c7bf78186424e587fba33e3a5b549 Mon Sep 17 00:00:00 2001
From: Ritvik Kapila <ritvik.iitd@gmail.com>
Date: Tue, 25 Jun 2024 14:50:24 -0700
Subject: [PATCH 25/25] fixes

---
 examples/src/custom_mpl_cmm_example.py | 20 ++++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/examples/src/custom_mpl_cmm_example.py b/examples/src/custom_mpl_cmm_example.py
index 8b331ccf6..544ca5f00 100644
--- a/examples/src/custom_mpl_cmm_example.py
+++ b/examples/src/custom_mpl_cmm_example.py
@@ -1,18 +1,26 @@
 # 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 MPL's ICryptographicMaterialsManager class and use it with the ESDK.
+
+The 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.
 
 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.
+the default CMM, you can do this as demonstrated in this example. This is the easiest approach if
+you are just adding a small check to the CMM methods, as in this example.
 
-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.
+If your use case calls for fundamentally changing aspects of the default CMM, you can also write
+your own implementation without extending an existing CMM. The default CMM's implementation is a
+good reference to use if you need to write a custom CMM implementation from scratch.
+Custom implementations of CMMs must implement get_encryption_materials and decrypt_materials.
 
 For more information on a default implementation of a CMM,
 please look at the default_cryptographic_materials_manager_example.py example.
+
+For more information on Cryptographic Material Managers, see
+https://docs.aws.amazon.com/encryption-sdk/latest/developer-guide/concepts.html#crypt-materials-manager
 """
 
 from aws_cryptographic_materialproviders.mpl import AwsCryptographicMaterialProviders