chore: KMS keyring example

Author: RitvikKapila

examples/src/keyrings/aws_kms_keyring_example.py

@@ -2,63 +2,73 @@
 # SPDX-License-Identifier: Apache-2.0
 """
 This example sets up the KMS Keyring
+
+The AWS KMS keyring uses symmetric encryption KMS keys to generate, encrypt and
+decrypt data keys. This example creates a KMS Keyring and then encrypts a custom input EXAMPLE_DATA
+with an encryption context. This example also demonstrates 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
+
+AWS KMS keyrings can be used independently or in a multi-keyring with other keyrings 
+of the same or a different type.
+
+For more info on how to use KMS keyring, see
+https://docs.aws.amazon.com/encryption-sdk/latest/developer-guide/use-kms-keyring.html
 """
 import sys
-
+from typing import Dict
 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 typing import Dict
 
 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])
+MODULE_ROOT_DIR = '/'.join(__file__.split("/")[:-1])
 
-sys.path.append(module_root_dir)
+sys.path.append(MODULE_ROOT_DIR)
 
 EXAMPLE_DATA: bytes = b"Hello World"
 
+
 def encrypt_and_decrypt_with_keyring(
     kms_key_id: str
 ):
-    """Demonstrate an encrypt/decrypt cycle using an AWS KMS keyring."""
+    """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 info 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()`.
-    """
+
+    # 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 boto3 clients for KMS.
-    """
+    # 2. Create a boto3 client for KMS.
     kms_client = boto3.client('kms', region_name="us-west-2")
 
-    """
-    3. Instantiate the Material Providers
-    """
-    mat_prov: AwsCryptographicMaterialProviders = AwsCryptographicMaterialProviders(
-        config=MaterialProvidersConfig()
-    )
-
-    """
-    4. Create encryption context
-    Remember that your encryption context is NOT SECRET.
-    https://docs.aws.amazon.com/encryption-sdk/latest/developer-guide/concepts.html#encryption-context
-    """
+    # 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",
@@ -67,9 +77,11 @@ def encrypt_and_decrypt_with_keyring(
         "the data you are handling": "is what you think it is",
     }
 
-    """
-    5. Create the KMS keyring
-    """
+    # 4. Create a KMS keyring
+    mat_prov: AwsCryptographicMaterialProviders = AwsCryptographicMaterialProviders(
+        config=MaterialProvidersConfig()
+    )
+
     keyring_input: CreateAwsKmsKeyringInput = CreateAwsKmsKeyringInput(
         kms_key_id=kms_key_id,
         kms_client=kms_client
@@ -79,37 +91,29 @@ def encrypt_and_decrypt_with_keyring(
         input=keyring_input
     )
 
-    """
-    6. Encrypt the data for the encryptionContext
-    """
+    # 5. Encrypt the data for the encryptionContext
     ciphertext, _ = client.encrypt(
         source=EXAMPLE_DATA,
         keyring=kms_keyring,
         encryption_context=encryption_context
     )
 
-    """
-    7. Demonstrate that the ciphertext and plaintext are different.
-    """
-    assert ciphertext != EXAMPLE_DATA, "Ciphertext and plaintext data are the same. Invalid encryption"
-    
-    """
-    8. Decrypt your encrypted data using the same keyring you used on encrypt.
-    You do not need to specify the encryption context on decrypt
-    because the header of the encrypted message includes the 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,
         keyring=kms_keyring
     )
 
-    """
-    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"
+        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.
-    """
-    assert plaintext_bytes == EXAMPLE_DATA
+    # 9. Demonstrate that the decrypted plaintext is identical to the original plaintext.
+    assert plaintext_bytes == EXAMPLE_DATA

examples/test/keyrings/test_i_aws_kms_keyring_example.py

@@ -9,5 +9,7 @@
 
 
 def test_encrypt_and_decrypt_with_keyring():
-    key_arn = "arn:aws:kms:us-west-2:658956600833:key/b3537ef1-d8dc-4780-9f5a-55776cbb2f7f"
-    encrypt_and_decrypt_with_keyring(key_arn)
+    """Test function for encryption and decryption using the AWS KMS Keyring example"""
+
+    kms_key_id = "arn:aws:kms:us-west-2:658956600833:key/b3537ef1-d8dc-4780-9f5a-55776cbb2f7f"
+    encrypt_and_decrypt_with_keyring(kms_key_id)

examples/test/keyrings/test_i_hierarchical_keyring.py

@@ -9,6 +9,7 @@
 
 
 def test_encrypt_and_decrypt_with_keyring():
+    """Test function for encryption and decryption using the AWS KMS Hierarchical Keyring example"""
     key_store_table_name = "KeyStoreDdbTable"
-    key_arn = "arn:aws:kms:us-west-2:370957321024:key/9d989aa2-2f9c-438c-a745-cc57d3ad0126"
-    encrypt_and_decrypt_with_keyring(key_store_table_name, key_store_table_name, key_arn)
+    kms_key_id = "arn:aws:kms:us-west-2:370957321024:key/9d989aa2-2f9c-438c-a745-cc57d3ad0126"
+    encrypt_and_decrypt_with_keyring(key_store_table_name, key_store_table_name, kms_key_id)