added file streaming example; updated assertion error message in all keyrings; fixed nomenclature in hierarchical and requiredec keyring examples

Author: RitvikKapila

.gitignore

@@ -32,7 +32,10 @@ __pycache__
 # PyTest
 .pytest_cache
 # Ignore key materials generated by examples or tests
-test_keys/
+test_keyrings/user_public_key_file_name.pem
+test_keyrings/user_private_key_file_name.pem
+test_keyrings/my-encrypted-data.ct
+test_keyrings/my-decrypted-data.dat
 
 # PyCharm
 .idea/

examples/src/keyrings/aws_kms_discovery_keyring_example.py

@@ -172,7 +172,8 @@ def encrypt_and_decrypt_with_keyring(
 
     # 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
+    assert plaintext_bytes == EXAMPLE_DATA, \
+        "Decrypted plaintext should be identical to the original plaintext. Invalid decryption"
 
     # 11. Demonstrate that if a discovery keyring (Bob's) doesn't have the correct AWS Account ID's,
     # the decrypt will fail with an error message

examples/src/keyrings/aws_kms_discovery_multi_keyring_example.py

@@ -170,4 +170,5 @@ def encrypt_and_decrypt_with_keyring(
 
     # 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
+    assert plaintext_bytes == EXAMPLE_DATA, \
+        "Decrypted plaintext should be identical to the original plaintext. Invalid decryption"

examples/src/keyrings/aws_kms_keyring_example.py

@@ -116,4 +116,5 @@ def encrypt_and_decrypt_with_keyring(
 
     # 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
+    assert plaintext_bytes == EXAMPLE_DATA, \
+        "Decrypted plaintext should be identical to the original plaintext. Invalid decryption"

examples/src/keyrings/aws_kms_mrk_keyring_example.py

@@ -151,4 +151,5 @@ def encrypt_and_decrypt_with_keyring(
 
     # 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
+    assert plaintext_bytes == EXAMPLE_DATA, \
+        "Decrypted plaintext should be identical to the original plaintext. Invalid decryption"

examples/src/keyrings/aws_kms_mrk_multi_keyring_example.py

@@ -143,7 +143,8 @@ def encrypt_and_decrypt_with_keyring(
 
     # 8. 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
+    assert plaintext_bytes == EXAMPLE_DATA, \
+        "Decrypted plaintext should be identical to the original plaintext. Invalid decryption"
 
     # Demonstrate that a single AwsKmsMrkKeyring configured with a replica of the MRK from the
     # multi-keyring used to encrypt the data is also capable of decrypting the data.

examples/src/keyrings/aws_kms_multi_keyring_example.py

@@ -144,7 +144,8 @@ def encrypt_and_decrypt_with_keyring(
 
     # 6b. 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_multi_keyring == EXAMPLE_DATA
+    assert plaintext_bytes_multi_keyring == EXAMPLE_DATA, \
+        "Decrypted plaintext should be identical to the original plaintext. Invalid decryption"
 
     # Because you used a multi_keyring on Encrypt, you can use either of the two
     # kms keyrings individually to decrypt the data.
@@ -174,7 +175,8 @@ def encrypt_and_decrypt_with_keyring(
 
     # 7d. 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_default_region_kms_keyring == EXAMPLE_DATA
+    assert plaintext_bytes_default_region_kms_keyring == EXAMPLE_DATA, \
+        "Decrypted plaintext should be identical to the original plaintext. Invalid decryption"
 
     # 8. Demonstrate that you can also successfully decrypt data using a KMS keyring with just the
     # `second_region_kms_key_id` directly.
@@ -201,4 +203,5 @@ def encrypt_and_decrypt_with_keyring(
 
     # 8d. 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_second_region_kms_keyring == EXAMPLE_DATA
+    assert plaintext_bytes_second_region_kms_keyring == EXAMPLE_DATA, \
+        "Decrypted plaintext should be identical to the original plaintext. Invalid decryption"

examples/src/keyrings/aws_kms_rsa_keyring_example.py

@@ -122,4 +122,5 @@ def encrypt_and_decrypt_with_keyring(
 
     # 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
+    assert plaintext_bytes == EXAMPLE_DATA, \
+        "Decrypted plaintext should be identical to the original plaintext. Invalid decryption"

examples/src/keyrings/file_streaming_example.py

@@ -0,0 +1,144 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""
+This example demonstrates file streaming for encryption and decryption using a Raw AES keyring
+
+The Raw AES keyring lets you use an AES symmetric key that you provide as a wrapping key that
+protects your data key. You need to generate, store, and protect the key material,
+preferably in a hardware security module (HSM) or key management system. Use a Raw AES keyring
+when you need to provide the wrapping key and encrypt the data keys locally or offline.
+
+This example creates a Raw AES Keyring and then encrypts an input stream from the file
+`plaintext_filename` with an encryption context to an output (encrypted) file `ciphertext_filename`.
+It then decrypts the ciphertext from `ciphertext_filename` to a new file `new_plaintext_filename`.
+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.
+
+The Raw AES keyring encrypts data by using the AES-GCM algorithm and a wrapping key that
+you specify as a byte array. You can specify only one wrapping key in each Raw AES keyring,
+but you can include multiple Raw AES keyrings, alone or with other keyrings, in a multi-keyring.
+
+For more information on how to use Raw AES keyrings, see
+https://docs.aws.amazon.com/encryption-sdk/latest/developer-guide/use-raw-aes-keyring.html
+"""
+import secrets
+import sys
+import filecmp
+
+from aws_cryptographic_materialproviders.mpl import AwsCryptographicMaterialProviders
+from aws_cryptographic_materialproviders.mpl.config import MaterialProvidersConfig
+from aws_cryptographic_materialproviders.mpl.models import AesWrappingAlg, CreateRawAesKeyringInput
+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])
+
+sys.path.append(MODULE_ROOT_DIR)
+
+
+def encrypt_and_decrypt_with_keyring(
+    plaintext_filename: str,
+    ciphertext_filename: str,
+    new_plaintext_filename: str
+):
+    """Demonstrate a streaming encrypt/decrypt cycle using a Raw AES keyring.
+
+    Usage: encrypt_and_decrypt_with_keyring()
+    """
+    # 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. The key namespace and key name are defined by you.
+    # and are used by the Raw AES keyring to determine
+    # whether it should attempt to decrypt an encrypted data key.
+    # For more information, see
+    # https://docs.aws.amazon.com/encryption-sdk/latest/developer-guide/use-raw-aes-keyring.html
+    key_name_space = "Some managed raw keys"
+    key_name = "My 256-bit AES wrapping key"
+
+    # 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. Generate a 256-bit AES key to use with your keyring.
+    # In practice, you should get this key from a secure key management system such as an HSM.
+
+    # Here, the input to secrets.token_bytes() = 32 bytes = 256 bits
+    static_key = secrets.token_bytes(32)
+
+    # 5. Create a Raw AES keyring
+    mat_prov: AwsCryptographicMaterialProviders = AwsCryptographicMaterialProviders(
+        config=MaterialProvidersConfig()
+    )
+
+    keyring_input: CreateRawAesKeyringInput = CreateRawAesKeyringInput(
+        key_namespace=key_name_space,
+        key_name=key_name,
+        wrapping_key=static_key,
+        wrapping_alg=AesWrappingAlg.ALG_AES256_GCM_IV12_TAG16
+    )
+
+    raw_aes_keyring: IKeyring = mat_prov.create_raw_aes_keyring(
+        input=keyring_input
+    )
+
+    # 6. Encrypt the data stream with the encryptionContext
+    with open(plaintext_filename, 'rb') as pt_file, open(ciphertext_filename, 'wb') as ct_file:
+        with client.stream(
+            mode='e',
+            source=pt_file,
+            keyring=raw_aes_keyring,
+            encryption_context=encryption_context
+        ) as encryptor:
+            for chunk in encryptor:
+                ct_file.write(chunk)
+
+    # 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 not filecmp.cmp(plaintext_filename, ciphertext_filename), \
+        "Ciphertext and plaintext data are the same. Invalid encryption"
+
+    # 8. Decrypt your encrypted data using the same keyring you used on encrypt.
+    with open(ciphertext_filename, 'rb') as ct_file, open(new_plaintext_filename, 'wb') as pt_file:
+        with client.stream(
+            mode='d',
+            source=ct_file,
+            keyring=raw_aes_keyring,
+            encryption_context=encryption_context
+        ) as decryptor:
+            for chunk in decryptor:
+                pt_file.write(chunk)
+
+    # 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 == decryptor.header.encryption_context[k], \
+            "Encryption context does not match expected values"
+
+    # 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 filecmp.cmp(plaintext_filename, new_plaintext_filename), \
+        "Decrypted plaintext should be identical to the original plaintext. Invalid decryption"

examples/src/keyrings/hierarchical_keyring.py

@@ -101,13 +101,13 @@ def encrypt_and_decrypt_with_keyring(
     )
 
     # 4. 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
+    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
 
     # 5. 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,
+        tenant_1_id=branch_key_id_a,
+        tenant_2_id=branch_key_id_b,
     )
 
     # 6. Create the Hierarchical Keyring.
@@ -135,7 +135,7 @@ def encrypt_and_decrypt_with_keyring(
     #    be used to encrypt data.
 
     # Create encryption context for TenantA
-    encryption_context_A: Dict[str, str] = {
+    encryption_context_a: Dict[str, str] = {
         "tenant": "TenantA",
         "encryption": "context",
         "is not": "secret",
@@ -145,7 +145,7 @@ def encrypt_and_decrypt_with_keyring(
     }
 
     # Create encryption context for TenantB
-    encryption_context_B: Dict[str, str] = {
+    encryption_context_b: Dict[str, str] = {
         "tenant": "TenantB",
         "encryption": "context",
         "is not": "secret",
@@ -155,22 +155,22 @@ def encrypt_and_decrypt_with_keyring(
     }
 
     # 8. Encrypt the data with encryptionContextA & encryptionContextB
-    ciphertext_A, _ = client.encrypt(
+    ciphertext_a, _ = client.encrypt(
         source=EXAMPLE_DATA,
         keyring=hierarchical_keyring,
-        encryption_context=encryption_context_A
+        encryption_context=encryption_context_a
     )
-    ciphertext_B, _ = client.encrypt(
+    ciphertext_b, _ = client.encrypt(
         source=EXAMPLE_DATA,
         keyring=hierarchical_keyring,
-        encryption_context=encryption_context_B
+        encryption_context=encryption_context_b
     )
 
     # 9. To attest that TenantKeyB cannot decrypt a message written by TenantKeyA,
     #    let's construct more restrictive hierarchical keyrings.
-    keyring_input_A: CreateAwsKmsHierarchicalKeyringInput = CreateAwsKmsHierarchicalKeyringInput(
+    keyring_input_a: CreateAwsKmsHierarchicalKeyringInput = CreateAwsKmsHierarchicalKeyringInput(
         key_store=keystore,
-        branch_key_id=branch_key_id_A,
+        branch_key_id=branch_key_id_a,
         ttl_seconds=600,
         cache=CacheTypeDefault(
             value=DefaultCache(
@@ -179,13 +179,13 @@ def encrypt_and_decrypt_with_keyring(
         ),
     )
 
-    hierarchical_keyring_A: IKeyring = mat_prov.create_aws_kms_hierarchical_keyring(
-        input=keyring_input_A
+    hierarchical_keyring_a: IKeyring = mat_prov.create_aws_kms_hierarchical_keyring(
+        input=keyring_input_a
     )
 
-    keyring_input_B: CreateAwsKmsHierarchicalKeyringInput = CreateAwsKmsHierarchicalKeyringInput(
+    keyring_input_b: CreateAwsKmsHierarchicalKeyringInput = CreateAwsKmsHierarchicalKeyringInput(
         key_store=keystore,
-        branch_key_id=branch_key_id_B,
+        branch_key_id=branch_key_id_b,
         ttl_seconds=600,
         cache=CacheTypeDefault(
             value=DefaultCache(
@@ -194,8 +194,8 @@ def encrypt_and_decrypt_with_keyring(
         ),
     )
 
-    hierarchical_keyring_B: IKeyring = mat_prov.create_aws_kms_hierarchical_keyring(
-        input=keyring_input_B
+    hierarchical_keyring_b: IKeyring = mat_prov.create_aws_kms_hierarchical_keyring(
+        input=keyring_input_b
     )
 
     # 10. Demonstrate that data encrypted by one tenant's key
@@ -205,8 +205,8 @@ def encrypt_and_decrypt_with_keyring(
     # This will fail and raise a AWSEncryptionSDKClientError, which we swallow ONLY for demonstration purposes.
     try:
         client.decrypt(
-            source=ciphertext_A,
-            keyring=hierarchical_keyring_B
+            source=ciphertext_a,
+            keyring=hierarchical_keyring_b
         )
     except AWSEncryptionSDKClientError:
         pass
@@ -215,21 +215,23 @@ def encrypt_and_decrypt_with_keyring(
     # This will fail and raise a AWSEncryptionSDKClientError, which we swallow ONLY for demonstration purposes.
     try:
         client.decrypt(
-            source=ciphertext_B,
-            keyring=hierarchical_keyring_A
+            source=ciphertext_b,
+            keyring=hierarchical_keyring_a
         )
     except AWSEncryptionSDKClientError:
         pass
 
     # 10. Demonstrate that data encrypted by one tenant's branch key can be decrypted by that tenant,
     #     and that the decrypted data matches the input data.
-    plaintext_bytes_A, _ = client.decrypt(
-        source=ciphertext_A,
-        keyring=hierarchical_keyring_A
+    plaintext_bytes_a, _ = client.decrypt(
+        source=ciphertext_a,
+        keyring=hierarchical_keyring_a
     )
-    assert plaintext_bytes_A == EXAMPLE_DATA
-    plaintext_bytes_B, _ = client.decrypt(
-        source=ciphertext_B,
-        keyring=hierarchical_keyring_B
+    assert plaintext_bytes_a == EXAMPLE_DATA, \
+        "Decrypted plaintext should be identical to the original plaintext. Invalid decryption"
+    plaintext_bytes_b, _ = client.decrypt(
+        source=ciphertext_b,
+        keyring=hierarchical_keyring_b
     )
-    assert plaintext_bytes_B == EXAMPLE_DATA
+    assert plaintext_bytes_b == EXAMPLE_DATA, \
+        "Decrypted plaintext should be identical to the original plaintext. Invalid decryption"

examples/src/keyrings/multi_keyring_example.py

@@ -175,7 +175,8 @@ def encrypt_and_decrypt_with_keyring(
 
     # 10b. 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_multi_keyring == EXAMPLE_DATA
+    assert plaintext_bytes_multi_keyring == EXAMPLE_DATA, \
+        "Decrypted plaintext should be identical to the original plaintext. Invalid decryption"
 
     # Because you used a multi_keyring on Encrypt, you can use either the
     # `kms_keyring` or `raw_aes_keyring` individually to decrypt the data.
@@ -192,7 +193,8 @@ def encrypt_and_decrypt_with_keyring(
 
     # 11b. 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_kms_keyring == EXAMPLE_DATA
+    assert plaintext_bytes_kms_keyring == EXAMPLE_DATA, \
+        "Decrypted plaintext should be identical to the original plaintext. Invalid decryption"
 
     # 12. Demonstrate that you can also successfully decrypt data using the `raw_aes_keyring`
     # directly.
@@ -206,4 +208,5 @@ def encrypt_and_decrypt_with_keyring(
 
     # 12b. 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_raw_aes_keyring == EXAMPLE_DATA
+    assert plaintext_bytes_raw_aes_keyring == EXAMPLE_DATA, \
+        "Decrypted plaintext should be identical to the original plaintext. Invalid decryption"