aws · RitvikKapila · Jun 7, 2024 · Jun 5, 2024 · Jun 6, 2024 · Jun 6, 2024
diff --git a/performance_tests/src/aws_encryption_sdk_performance_tests/keyrings/hierarchy_keyring.py b/performance_tests/src/aws_encryption_sdk_performance_tests/keyrings/hierarchy_keyring.py
@@ -0,0 +1,128 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Performance tests for the hierarchy keyring."""
+
+import aws_encryption_sdk
+import boto3
+from aws_cryptographic_materialproviders.keystore import KeyStore
+from aws_cryptographic_materialproviders.keystore.config import KeyStoreConfig
+from aws_cryptographic_materialproviders.keystore.models import KMSConfigurationKmsKeyArn
+from aws_cryptographic_materialproviders.mpl import AwsCryptographicMaterialProviders
+from aws_cryptographic_materialproviders.mpl.config import MaterialProvidersConfig
+from aws_cryptographic_materialproviders.mpl.models import (
+    CacheTypeDefault,
+    CreateAwsKmsHierarchicalKeyringInput,
+    DefaultCache,
+)
+from aws_cryptographic_materialproviders.mpl.references import IKeyring
+
+from ..utils.util import PerfTestUtils
+
+
+def create_keyring(
+    key_store_table_name: str,
+    logical_key_store_name: str,
+    kms_key_id: str,
+    branch_key_id: str = PerfTestUtils.DEFAULT_BRANCH_KEY_ID
+):
+    """Demonstrate how to create a hierarchy keyring.
+
+    Usage: create_keyring(key_store_table_name, logical_key_store_name, kms_key_id, branch_key_id)
+    :param key_store_table_name: Name of the KeyStore DynamoDB table.
+    :type key_store_table_name: string
+    :param logical_key_store_name: Logical name of the KeyStore.
+    :type logical_key_store_name: string
+    :param kms_key_id: KMS Key identifier for the KMS key you want to use.
+    :type kms_key_id: string
+    :param branch_key_id: Branch key you want to use for the hierarchy keyring.
+    :type branch_key_id: string
+
+    For more information on KMS Key identifiers, see
+    https://docs.aws.amazon.com/kms/latest/developerguide/concepts.html#key-id
+    """
+    # Create boto3 clients for DynamoDB and KMS.
+    ddb_client = boto3.client('dynamodb', region_name="us-west-2")
+    kms_client = boto3.client('kms', region_name="us-west-2")
+
+    # Configure your KeyStore resource.
+    # This SHOULD be the same configuration that you used
+    # to initially create and populate your KeyStore.
+    keystore: KeyStore = KeyStore(
+        config=KeyStoreConfig(
+            ddb_client=ddb_client,
+            ddb_table_name=key_store_table_name,
+            logical_key_store_name=logical_key_store_name,
+            kms_client=kms_client,
+            kms_configuration=KMSConfigurationKmsKeyArn(
+                value=kms_key_id
+            ),
+        )
+    )
+
+    # Create the Hierarchical Keyring.
+    mat_prov: AwsCryptographicMaterialProviders = AwsCryptographicMaterialProviders(
+        config=MaterialProvidersConfig()
+    )
+
+    keyring_input: CreateAwsKmsHierarchicalKeyringInput = CreateAwsKmsHierarchicalKeyringInput(
+        key_store=keystore,
+        branch_key_id=branch_key_id,
+        ttl_seconds=600,
+        cache=CacheTypeDefault(
+            value=DefaultCache(
+                entry_capacity=100
+            )
+        ),
+    )
+
+    keyring: IKeyring = mat_prov.create_aws_kms_hierarchical_keyring(
+        input=keyring_input
+    )
+
+    return keyring
+
+
+def encrypt_using_keyring(
+    plaintext_data: bytes,
+    keyring: IKeyring
+):
+    """Demonstrate how to encrypt plaintext data using a hierarchy keyring.
+
+    Usage: encrypt_using_keyring(plaintext_data, keyring)
+    :param plaintext_data: plaintext data you want to encrypt
+    :type: bytes
+    :param keyring: Keyring to use for encryption.
+    :type keyring: IKeyring
+    """
+    client = aws_encryption_sdk.EncryptionSDKClient()
+
+    ciphertext_data, _ = client.encrypt(
+        source=plaintext_data,
+        keyring=keyring,
+        encryption_context=PerfTestUtils.DEFAULT_ENCRYPTION_CONTEXT
+    )
+
+    return ciphertext_data
+
+
+def decrypt_using_keyring(
+    ciphertext_data: bytes,
+    keyring: IKeyring
+):
+    """Demonstrate how to decrypt ciphertext data using a hierarchy keyring.
+
+    Usage: decrypt_using_keyring(ciphertext_data, keyring)
+    :param ciphertext_data: ciphertext data you want to decrypt
+    :type: bytes
+    :param keyring: Keyring to use for decryption.
+    :type keyring: IKeyring
+    """
+    client = aws_encryption_sdk.EncryptionSDKClient()
+
+    decrypted_plaintext_data, _ = client.decrypt(
+        source=ciphertext_data,
+        keyring=keyring,
+        encryption_context=PerfTestUtils.DEFAULT_ENCRYPTION_CONTEXT
+    )
+
+    return decrypted_plaintext_data
diff --git a/performance_tests/src/aws_encryption_sdk_performance_tests/keyrings/raw_aes_keyring.py b/performance_tests/src/aws_encryption_sdk_performance_tests/keyrings/raw_aes_keyring.py
@@ -19,7 +19,6 @@ def create_keyring():
     key_name_space = "Some managed raw keys"
     key_name = "My 256-bit AES wrapping key"
 
-    # Here, the input to secrets.token_bytes() = 32 bytes = 256 bits
     # We fix the static key in order to make the test deterministic
     static_key = PerfTestUtils.DEFAULT_AES_256_STATIC_KEY
 

diff --git a/.../aws_encryption_sdk_performance_tests/master_key_providers/aws_kms_master_key_provider.py b/.../aws_encryption_sdk_performance_tests/master_key_providers/aws_kms_master_key_provider.py
@@ -8,7 +8,7 @@
 def create_key_provider(
     kms_key_id: str
 ):
-    """Demonstrate how to create an AWS KMS master key-provider.
+    """Demonstrate how to create an AWS KMS master key provider.
 
     Usage: create_key_provider(kms_key_id)
     :param kms_key_id: KMS Key identifier for the KMS key you want to use.
@@ -17,7 +17,7 @@ def create_key_provider(
     For more information on KMS Key identifiers, see
     https://docs.aws.amazon.com/kms/latest/developerguide/concepts.html#key-id
     """
-    # Create a KMS master key-provider.
+    # Create a KMS master key provider.
     key_provider = aws_encryption_sdk.StrictAwsKmsMasterKeyProvider(key_ids=[
         kms_key_id,
     ])
@@ -29,7 +29,7 @@ def encrypt_using_key_provider(
     plaintext_data: bytes,
     key_provider: aws_encryption_sdk.key_providers.base.MasterKeyProvider
 ):
-    """Demonstrate how to encrypt plaintext data using an AWS KMS master key-provider.
+    """Demonstrate how to encrypt plaintext data using an AWS KMS master key provider.
 
     Usage: encrypt_using_key_provider(plaintext_data, key_provider)
     :param plaintext_data: plaintext data you want to encrypt
@@ -51,7 +51,7 @@ def decrypt_using_key_provider(
     ciphertext_data: bytes,
     key_provider: aws_encryption_sdk.key_providers.base.MasterKeyProvider
 ):
-    """Demonstrate how to decrypt ciphertext data using an AWS KMS master key-provider.
+    """Demonstrate how to decrypt ciphertext data using an AWS KMS master key provider.
 
     Usage: decrypt_using_key_provider(ciphertext_data, key_provider)
     :param ciphertext_data: ciphertext data you want to decrypt

diff --git a/...rmance_tests/src/aws_encryption_sdk_performance_tests/master_key_providers/caching_cmm.py b/...rmance_tests/src/aws_encryption_sdk_performance_tests/master_key_providers/caching_cmm.py
@@ -0,0 +1,85 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Performance tests for the Caching Cryptographic Materials Manager (CMM) with KMS Master Key Provider."""
+
+import aws_encryption_sdk
+
+
+def create_cmm(
+    kms_key_id: str,
+    max_age_in_cache: float,
+    cache_capacity: int
+):
+    """Demonstrate how to create a Caching CMM.
+
+    Usage: create_cmm(kms_key_id, max_age_in_cache, cache_capacity)
+    :param kms_key_id: Amazon Resource Name (ARN) of the KMS customer master key
+    :type kms_key_id: str
+    :param max_age_in_cache: Maximum time in seconds that a cached entry can be used
+    :type max_age_in_cache: float
+    :param cache_capacity: Maximum number of entries to retain in cache at once
+    :type cache_capacity: int
+    """
+    # Security thresholds
+    #   Max messages (or max bytes per) data key are optional
+    max_messages_encrypted = 100
+
+    # Create a master key provider for the KMS customer master key (CMK)
+    key_provider = aws_encryption_sdk.StrictAwsKmsMasterKeyProvider(key_ids=[kms_key_id])
+
+    # Create a local cache
+    cache = aws_encryption_sdk.LocalCryptoMaterialsCache(cache_capacity)
+
+    # Create a caching CMM
+    caching_cmm = aws_encryption_sdk.CachingCryptoMaterialsManager(
+        master_key_provider=key_provider,
+        cache=cache,
+        max_age=max_age_in_cache,
+        max_messages_encrypted=max_messages_encrypted,
+    )
+
+    return caching_cmm
+
+
+def encrypt_using_cmm(
+    plaintext_data: bytes,
+    caching_cmm: aws_encryption_sdk.materials_managers.base.CryptoMaterialsManager
+):
+    """Demonstrate how to encrypt plaintext data using a Caching CMM.
+
+    Usage: encrypt_using_cmm(plaintext_data, caching_cmm)
+    :param plaintext_data: plaintext data you want to encrypt
+    :type: bytes
+    :param caching_cmm: Crypto Materials Manager to use for encryption.
+    :type caching_cmm: aws_encryption_sdk.materials_managers.base.CryptoMaterialsManager
+    """
+    client = aws_encryption_sdk.EncryptionSDKClient()
+
+    ciphertext_data, _ = client.encrypt(
+        source=plaintext_data,
+        materials_manager=caching_cmm
+    )
+
+    return ciphertext_data
+
+
+def decrypt_using_cmm(
+    ciphertext_data: bytes,
+    caching_cmm: aws_encryption_sdk.materials_managers.base.CryptoMaterialsManager
+):
+    """Demonstrate how to decrypt ciphertext data using a Caching CMM.
+
+    Usage: decrypt_using_cmm(ciphertext_data, caching_cmm)
+    :param ciphertext_data: ciphertext data you want to decrypt
+    :type: bytes
+    :param caching_cmm: Crypto Materials Manager to use for encryption.
+    :type caching_cmm: aws_encryption_sdk.materials_managers.base.CryptoMaterialsManager
+    """
+    client = aws_encryption_sdk.EncryptionSDKClient()
+
+    decrypted_plaintext_data, _ = client.decrypt(
+        source=ciphertext_data,
+        materials_manager=caching_cmm
+    )
+
+    return decrypted_plaintext_data
diff --git a/.../aws_encryption_sdk_performance_tests/master_key_providers/raw_aes_master_key_provider.py b/.../aws_encryption_sdk_performance_tests/master_key_providers/raw_aes_master_key_provider.py
@@ -43,13 +43,13 @@ def _get_raw_key(self, key_id):
 
 
 def create_key_provider():
-    """Demonstrate how to create a Raw AES master key-provider.
+    """Demonstrate how to create a Raw AES master key provider.
 
     Usage: create_key_provider()
     """
-    # Create a Raw AES master key-provider.
+    # Create a Raw AES master key provider.
 
-    # The Key ID  field in the JceMasterKey and RawMasterKey is equivalent to key name in the Raw keyrings
+    # The Key ID field in the JceMasterKey and RawMasterKey is equivalent to key name in the Raw keyrings
     key_id = "My 256-bit AES wrapping key"
     key_provider = StaticRandomMasterKeyProvider()
     key_provider.add_master_key(key_id)
@@ -61,7 +61,7 @@ def encrypt_using_key_provider(
     plaintext_data: bytes,
     key_provider: aws_encryption_sdk.key_providers.base.MasterKeyProvider
 ):
-    """Demonstrate how to encrypt plaintext data using a Raw AES master key-provider.
+    """Demonstrate how to encrypt plaintext data using a Raw AES master key provider.
 
     Usage: encrypt_using_key_provider(plaintext_data, key_provider)
     :param plaintext_data: plaintext data you want to encrypt
@@ -83,7 +83,7 @@ def decrypt_using_key_provider(
     ciphertext_data: bytes,
     key_provider: aws_encryption_sdk.key_providers.base.MasterKeyProvider
 ):
-    """Demonstrate how to decrypt ciphertext data using a Raw AES master key-provider.
+    """Demonstrate how to decrypt ciphertext data using a Raw AES master key provider.
 
     Usage: decrypt_using_key_provider(ciphertext_data, key_provider)
     :param ciphertext_data: ciphertext data you want to decrypt

diff --git a/.../aws_encryption_sdk_performance_tests/master_key_providers/raw_rsa_master_key_provider.py b/.../aws_encryption_sdk_performance_tests/master_key_providers/raw_rsa_master_key_provider.py
@@ -43,13 +43,13 @@ def _get_raw_key(self, key_id):
 
 
 def create_key_provider():
-    """Demonstrate how to create a Raw RSA master key-provider.
+    """Demonstrate how to create a Raw RSA master key provider.
 
     Usage: create_key_provider()
     """
-    # Create a Raw RSA master key-provider.
+    # Create a Raw RSA master key provider.
 
-    # The Key ID  field in the JceMasterKey and RawMasterKey is equivalent to key name in the Raw keyrings
+    # The Key ID field in the JceMasterKey and RawMasterKey is equivalent to key name in the Raw keyrings
     key_id = "My 4096-bit RSA wrapping key"
     key_provider = StaticRandomMasterKeyProvider()
     key_provider.add_master_key(key_id)
@@ -61,7 +61,7 @@ def encrypt_using_key_provider(
     plaintext_data: bytes,
     key_provider: aws_encryption_sdk.key_providers.base.MasterKeyProvider
 ):
-    """Demonstrate how to encrypt plaintext data using a Raw RSA master key-provider.
+    """Demonstrate how to encrypt plaintext data using a Raw RSA master key provider.
 
     Usage: encrypt_using_key_provider(plaintext_data, key_provider)
     :param plaintext_data: plaintext data you want to encrypt
@@ -83,7 +83,7 @@ def decrypt_using_key_provider(
     ciphertext_data: bytes,
     key_provider: aws_encryption_sdk.key_providers.base.MasterKeyProvider
 ):
-    """Demonstrate how to decrypt ciphertext data using a Raw RSA master key-provider.
+    """Demonstrate how to decrypt ciphertext data using a Raw RSA master key provider.
 
     Usage: decrypt_using_key_provider(ciphertext_data, key_provider)
     :param ciphertext_data: ciphertext data you want to decrypt

diff --git a/performance_tests/src/aws_encryption_sdk_performance_tests/utils/util.py b/performance_tests/src/aws_encryption_sdk_performance_tests/utils/util.py
@@ -87,9 +87,7 @@ class PerfTestUtils:
         "the data you are handling": "is what you think it is",
     }
 
-    DEFAULT_BRANCH_KEY_ID_A = 'a52dfaad-7dbd-4430-a1fd-abaa5299da07'
-
-    DEFAULT_BRANCH_KEY_ID_B = '8ba79cef-581c-4125-9292-b057a29d42d7'
+    DEFAULT_BRANCH_KEY_ID = 'a52dfaad-7dbd-4430-a1fd-abaa5299da07'
 
     @staticmethod
     def read_file(filename):