chore(examples): Added file streaming, migration and set encryption algorithm examples

.gitignore

@@ -32,7 +32,7 @@ __pycache__
 # PyTest
 .pytest_cache
 # Ignore key materials generated by examples or tests
-test_keys/
+test_keyrings/
 
 # 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,152 @@
+# 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.
+
+File streaming is useful when the plaintext or ciphertext file/data is too large to load into
+memory. Therefore, the AWS Encryption SDK allows users to stream the data, instead of loading it
+all at once in memory. In this example, we demonstrate file streaming for encryption and decryption
+using a Raw AES keyring. However, you can use any keyring with streaming.
+
+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 `decrypted_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.
+
+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
+
+See raw_aes_keyring_example.py in the same directory for another raw AES keyring example
+in the AWS Encryption SDK for Python.
+"""
+import filecmp
+import secrets
+import sys
+
+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,
+    decrypted_filename: str
+):
+    """Demonstrate a streaming encrypt/decrypt cycle.
+
+    Usage: encrypt_and_decrypt_with_keyring(plaintext_filename
+                                            ciphertext_filename
+                                            decrypted_filename)
+    :param plaintext_filename: filename of the plaintext data
+    :type plaintext_filename: string
+    :param ciphertext_filename: filename of the ciphertext data
+    :type ciphertext_filename: string
+    :param decrypted_filename: filename of the decrypted data
+    :type decrypted_filename: string
+    """
+    # 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
+    # We choose to use a raw AES keyring, but any keyring can be used with streaming.
+    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 stream using the same keyring you used on encrypt.
+    with open(ciphertext_filename, 'rb') as ct_file, open(decrypted_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, decrypted_filename), \
+        "Decrypted plaintext should be identical to the original plaintext. Invalid decryption"

examples/src/keyrings/hierarchical_keyring_example.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"