fix kms keyring test

Author: RitvikKapila

performance_tests/src/aws_encryption_sdk_performance_tests/keyrings/hierarchy_keyring.py

@@ -0,0 +1,143 @@
+# 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.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 aws_cryptographic_materialproviders.keystore import KeyStore
+from aws_cryptographic_materialproviders.keystore.config import KeyStoreConfig
+from aws_cryptographic_materialproviders.keystore.models import CreateKeyInput, 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 IBranchKeyIdSupplier, IKeyring
+from typing import Dict  # noqa pylint: disable=wrong-import-order
+
+import aws_encryption_sdk
+from aws_encryption_sdk import CommitmentPolicy
+from aws_encryption_sdk.exceptions import AWSEncryptionSDKClientError
+
+from .branch_key_id_supplier_example import ExampleBranchKeyIdSupplier
+
+
+def create_keyring(
+    key_store_table_name: str,
+    logical_key_store_name: str,
+    kms_key_id: str
+):
+    """Demonstrate how to create a hierarchy keyring.
+
+    Usage: create_keyring(key_store_table_name, logical_key_store_name, kms_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
+
+    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
+            ),
+        )
+    )
+
+    # Call CreateKey to create two new active branch keys
+    branch_key_id_a: str = keystore.create_key(input=CreateKeyInput()).branch_key_identifier
+    branch_key_id_b: str = keystore.create_key(input=CreateKeyInput()).branch_key_identifier
+
+    # Create a branch key supplier that maps the branch key id to a more readable format
+    branch_key_id_supplier: IBranchKeyIdSupplier = ExampleBranchKeyIdSupplier(
+        tenant_1_id=branch_key_id_a,
+        tenant_2_id=branch_key_id_b,
+    )
+
+    # Create the Hierarchical Keyring.
+    mat_prov: AwsCryptographicMaterialProviders = AwsCryptographicMaterialProviders(
+        config=MaterialProvidersConfig()
+    )
+
+    keyring_input: CreateAwsKmsHierarchicalKeyringInput = CreateAwsKmsHierarchicalKeyringInput(
+        key_store=keystore,
+        branch_key_id_supplier=branch_key_id_supplier,
+        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 an AWS KMS 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
+    )
+
+    return ciphertext_data
+
+
+def decrypt_using_keyring(
+    ciphertext_data: bytes,
+    keyring: IKeyring
+):
+    """Demonstrate how to decrypt ciphertext data using an AWS KMS 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
+    )
+
+    return decrypted_plaintext_data

performance_tests/src/aws_encryption_sdk_performance_tests/utils/util.py

@@ -97,6 +97,6 @@ def get_rsa_key_from_file(filename):
     @staticmethod
     def write_time_list_to_csv(time_list, filename):
         """Writes the time list to a CSV file."""
-        with open('results/' + filename + '.csv', 'w', encoding='utf-8') as myfile:
+        with open(filename + '.csv', 'w', encoding='utf-8') as myfile:
             for time in time_list:
                 myfile.write(str(time) + '\n')

performance_tests/test/keyrings/test_aws_kms_keyring.py

@@ -30,7 +30,7 @@ def create_kms_keyring():
 @click.option('--n_iters',
               default=PerfTestUtils.DEFAULT_N_ITERS)
 @click.option('--output_file',
-              default='kms_keyring_create')
+              default='/'.join(__file__.split("/")[:-3]) + '/results/kms_keyring_create')
 def create(
     kms_key_id: str,
     n_iters: int,
@@ -46,9 +46,6 @@ def create(
         # calculate elapsed time in milliseconds
         elapsed_time = (time.time() - curr_time) * 1000
         time_list.append(elapsed_time)
-
-    print('time_list', time_list)
-    print('output_file', output_file)
     PerfTestUtils.write_time_list_to_csv(time_list, output_file)
 
 
@@ -63,7 +60,7 @@ def create_kms_keyring_given_kms_client():
 @click.option('--n_iters',
               default=PerfTestUtils.DEFAULT_N_ITERS)
 @click.option('--output_file',
-              default='kms_keyring_create_given_kms_client')
+              default='/'.join(__file__.split("/")[:-3]) + '/results/kms_keyring_create_given_kms_client')
 def create_given_kms_client(
     kms_key_id: str,
     n_iters: int,
@@ -91,14 +88,15 @@ def encrypt_kms_keyring():
 
 @encrypt_kms_keyring.command()
 @click.option('--plaintext_data_filename',
-              default='test/resources/plaintext/plaintext-data-' + PerfTestUtils.DEFAULT_FILE_SIZE + '.dat',
+              default='/'.join(__file__.split("/")[:-2]) + '/resources/plaintext/plaintext-data-'
+              + PerfTestUtils.DEFAULT_FILE_SIZE + '.dat',
               prompt='Filename containing plaintext data you want to encrypt')
 @click.option('--kms_key_id',
               default='arn:aws:kms:us-west-2:658956600833:key/b3537ef1-d8dc-4780-9f5a-55776cbb2f7f')
 @click.option('--n_iters',
               default=PerfTestUtils.DEFAULT_N_ITERS)
 @click.option('--output_file',
-              default='kms_keyring_encrypt')
+              default='/'.join(__file__.split("/")[:-3]) + '/results/kms_keyring_encrypt')
 def encrypt(
     plaintext_data_filename: str,
     kms_key_id: str,
@@ -130,14 +128,15 @@ def decrypt_kms_keyring():
 
 @decrypt_kms_keyring.command()
 @click.option('--ciphertext_data_filename',
-              default='test/resources/ciphertext/kms/ciphertext-data-' + PerfTestUtils.DEFAULT_FILE_SIZE + '.ct',
+              default='/'.join(__file__.split("/")[:-2]) + '/resources/ciphertext/kms/ciphertext-data-'
+              + PerfTestUtils.DEFAULT_FILE_SIZE + '.ct',
               prompt='Filename containing ciphertext data you want to decrypt')
 @click.option('--kms_key_id',
               default='arn:aws:kms:us-west-2:658956600833:key/b3537ef1-d8dc-4780-9f5a-55776cbb2f7f')
 @click.option('--n_iters',
               default=PerfTestUtils.DEFAULT_N_ITERS)
 @click.option('--output_file',
-              default='kms_keyring_decrypt')
+              default='/'.join(__file__.split("/")[:-3]) + '/results/kms_keyring_decrypt')
 def decrypt(
     ciphertext_data_filename: str,
     kms_key_id: str,
@@ -177,7 +176,6 @@ def runner():
 def test_create(runner):
     """Test the create_keyring function"""
     result = runner.invoke(create_kms_keyring.commands['create'], ['--n_iters', 1])
-    print('time_list', result.output)
     assert result.exit_code == 0
 
 

performance_tests/tox.ini

@@ -31,7 +31,7 @@ envlist =
 # release :: Builds dist files and uploads to pypi pypirc profile.
 
 [testenv:base-command]
-commands = pytest -s -v test/
+commands = pytest test/
 deps = 
     click