fix(Examples): Validate EC on decrypt (#697)

Author: lucasmcdonald3

examples/src/aws_kms_discovery_keyring_example.py

@@ -153,18 +153,14 @@ def encrypt_and_decrypt_with_keyring(
     # successfully decrypted. The resulting data key is used to decrypt the
     # ciphertext's message.
     # If all calls to KMS fail, the decryption fails.
-    plaintext_bytes, dec_header = client.decrypt(
+    plaintext_bytes, _ = client.decrypt(
         source=ciphertext,
-        keyring=discovery_keyring
+        keyring=discovery_keyring,
+        # Provide the encryption context that was supplied to the encrypt method
+        encryption_context=encryption_context,
     )
 
-    # 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 == dec_header.encryption_context[k], \
-            "Encryption context does not match expected values"
-
-    # 10. Demonstrate that the decrypted plaintext is identical to the original plaintext.
+    # 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, \
         "Decrypted plaintext should be identical to the original plaintext. Invalid decryption"
@@ -192,7 +188,10 @@ def encrypt_and_decrypt_with_keyring(
     try:
         plaintext_bytes, _ = client.decrypt(
             source=ciphertext,
-            keyring=discovery_keyring_bob
+            keyring=discovery_keyring_bob,
+            # Verify that the encryption context in the result contains the
+            # encryption context supplied to the encrypt method
+            encryption_context=encryption_context,
         )
 
         raise AssertionError("Decrypt using discovery keyring with wrong AWS Account ID should"

examples/src/aws_kms_discovery_multi_keyring_example.py

@@ -151,18 +151,14 @@ def encrypt_and_decrypt_with_keyring(
     # All of this is done serially, until a success occurs or all keyrings have
     # failed all (filtered) EDKs.
     # KMS Discovery Keyrings will attempt to decrypt Multi Region Keys (MRKs) and regular KMS Keys.
-    plaintext_bytes, dec_header = client.decrypt(
+    plaintext_bytes, _ = client.decrypt(
         source=ciphertext,
-        keyring=discovery_multi_keyring
+        keyring=discovery_multi_keyring,
+        # Provide the encryption context that was supplied to the encrypt method
+        encryption_context=encryption_context,
     )
 
-    # 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 == dec_header.encryption_context[k], \
-            "Encryption context does not match expected values"
-
-    # 10. Demonstrate that the decrypted plaintext is identical to the original plaintext.
+    # 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, \
         "Decrypted plaintext should be identical to the original plaintext. Invalid decryption"

examples/src/aws_kms_keyring_example.py

@@ -97,18 +97,14 @@ def encrypt_and_decrypt_with_keyring(
         "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(
+    plaintext_bytes, _ = client.decrypt(
         source=ciphertext,
-        keyring=kms_keyring
+        keyring=kms_keyring,
+        # Provide the encryption context that was supplied to the encrypt method
+        encryption_context=encryption_context,
     )
 
-    # 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"
-
-    # 9. Demonstrate that the decrypted plaintext is identical to the original plaintext.
+    # 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, \
         "Decrypted plaintext should be identical to the original plaintext. Invalid decryption"

examples/src/aws_kms_mrk_discovery_keyring_example.py

@@ -163,17 +163,13 @@ def encrypt_and_decrypt_with_keyring(
     )
 
     # 7. Decrypt your encrypted data using the discovery keyring.
-    plaintext_bytes, dec_header = client.decrypt(
+    plaintext_bytes, _ = client.decrypt(
         source=ciphertext,
-        keyring=decrypt_discovery_keyring
+        keyring=decrypt_discovery_keyring,
+        # Provide the encryption context that was supplied to the encrypt method
+        encryption_context=encryption_context,
     )
 
-    # 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"
-
-    # 9. Demonstrate that the decrypted plaintext is identical to the original plaintext.
+    # 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

examples/src/aws_kms_mrk_discovery_multi_keyring_example.py

@@ -172,17 +172,13 @@ def encrypt_and_decrypt_with_keyring(
     # All of this is done serially, until a success occurs or all keyrings have failed
     # all (filtered) EDKs. KMS MRK Discovery Keyrings will attempt to decrypt
     # Multi Region Keys (MRKs) and regular KMS Keys.
-    plaintext_bytes, dec_header = client.decrypt(
+    plaintext_bytes, _ = client.decrypt(
         source=ciphertext,
-        keyring=decrypt_discovery_keyring
+        keyring=decrypt_discovery_keyring,
+        # Provide the encryption context that was supplied to the encrypt method
+        encryption_context=encryption_context,
     )
 
-    # 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"
-
-    # 9. Demonstrate that the decrypted plaintext is identical to the original plaintext.
+    # 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

examples/src/aws_kms_mrk_keyring_example.py

@@ -132,18 +132,14 @@ def encrypt_and_decrypt_with_keyring(
     )
 
     # 7. Decrypt your encrypted data using the same keyring you used on encrypt.
-    plaintext_bytes, dec_header = client.decrypt(
+    plaintext_bytes, _ = client.decrypt(
         source=ciphertext,
-        keyring=decrypt_keyring
+        keyring=decrypt_keyring,
+        # Provide the encryption context that was supplied to the encrypt method
+        encryption_context=encryption_context,
     )
 
-    # 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"
-
-    # 9. Demonstrate that the decrypted plaintext is identical to the original plaintext.
+    # 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, \
         "Decrypted plaintext should be identical to the original plaintext. Invalid decryption"

examples/src/aws_kms_mrk_multi_keyring_example.py

@@ -124,18 +124,14 @@ def encrypt_and_decrypt_with_keyring(
     # 6. Decrypt your encrypted data using the same AwsKmsMrkMultiKeyring you used on encrypt.
     # It will decrypt the data using the generator key (in this case, the MRK), since that is
     # the first available KMS key on the keyring that is capable of decrypting the data.
-    plaintext_bytes, dec_header = client.decrypt(
+    plaintext_bytes, _ = client.decrypt(
         source=ciphertext,
-        keyring=kms_mrk_multi_keyring
+        keyring=kms_mrk_multi_keyring,
+        # Provide the encryption context that was supplied to the encrypt method
+        encryption_context=encryption_context,
     )
 
-    # 7. 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"
-
-    # 8. Demonstrate that the decrypted plaintext is identical to the original plaintext.
+    # 7. 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, \
         "Decrypted plaintext should be identical to the original plaintext. Invalid decryption"
@@ -144,7 +140,7 @@ def encrypt_and_decrypt_with_keyring(
     # multi-keyring used to encrypt the data is also capable of decrypting the data.
     # (This is an example for demonstration; you do not need to do this in your own code.)
 
-    # 9. Create a single AwsKmsMrkKeyring with the replica KMS MRK from the second region.
+    # 8. Create a single AwsKmsMrkKeyring with the replica KMS MRK from the second region.
 
     # Create a boto3 client for KMS in the second region which is the region for mrk_replica_key_id.
     second_region_kms_client = boto3.client('kms', region_name=mrk_replica_decrypt_region)
@@ -158,19 +154,15 @@ def encrypt_and_decrypt_with_keyring(
         input=second_region_mrk_keyring_input
     )
 
-    # 10. Decrypt your encrypted data using the second region AwsKmsMrkKeyring
-    plaintext_bytes_second_region, dec_header_second_region = client.decrypt(
+    # 9. Decrypt your encrypted data using the second region AwsKmsMrkKeyring
+    plaintext_bytes_second_region, _ = client.decrypt(
         source=ciphertext,
-        keyring=second_region_mrk_keyring
+        keyring=second_region_mrk_keyring,
+        # Provide the encryption context that was supplied to the encrypt method
+        encryption_context=encryption_context,
     )
 
-    # 11. 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_second_region.encryption_context[k], \
-            "Encryption context does not match expected values"
-
-    # 12. Demonstrate that the decrypted plaintext is identical to the original plaintext.
+    # 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_second_region == EXAMPLE_DATA
 

examples/src/aws_kms_multi_keyring_example.py

@@ -133,7 +133,9 @@ def encrypt_and_decrypt_with_keyring(
     # 6a. Decrypt your encrypted data using the same multi_keyring you used on encrypt.
     plaintext_bytes_multi_keyring, _ = client.decrypt(
         source=ciphertext,
-        keyring=kms_multi_keyring
+        keyring=kms_multi_keyring,
+        # Provide the encryption context that was supplied to the encrypt method
+        encryption_context=encryption_context,
     )
 
     # 6b. Demonstrate that the decrypted plaintext is identical to the original plaintext.
@@ -164,7 +166,9 @@ def encrypt_and_decrypt_with_keyring(
     # 7c. Decrypt your encrypted data using the default_region_kms_keyring.
     plaintext_bytes_default_region_kms_keyring, _ = client.decrypt(
         source=ciphertext,
-        keyring=default_region_kms_keyring
+        keyring=default_region_kms_keyring,
+        # Provide the encryption context that was supplied to the encrypt method
+        encryption_context=encryption_context,
     )
 
     # 7d. Demonstrate that the decrypted plaintext is identical to the original plaintext.
@@ -192,7 +196,9 @@ def encrypt_and_decrypt_with_keyring(
     # 8c. Decrypt your encrypted data using the second_region_kms_keyring.
     plaintext_bytes_second_region_kms_keyring, _ = client.decrypt(
         source=ciphertext,
-        keyring=second_region_kms_keyring
+        keyring=second_region_kms_keyring,
+        # Provide the encryption context that was supplied to the encrypt method
+        encryption_context=encryption_context,
     )
 
     # 8d. Demonstrate that the decrypted plaintext is identical to the original plaintext.

examples/src/aws_kms_rsa_keyring_example.py

@@ -103,18 +103,14 @@ def encrypt_and_decrypt_with_keyring(
         "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(
+    plaintext_bytes, _ = client.decrypt(
         source=ciphertext,
-        keyring=kms_rsa_keyring
+        keyring=kms_rsa_keyring,
+        # Provide the encryption context that was supplied to the encrypt method
+        encryption_context=encryption_context,
     )
 
-    # 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"
-
-    # 9. Demonstrate that the decrypted plaintext is identical to the original plaintext.
+    # 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, \
         "Decrypted plaintext should be identical to the original plaintext. Invalid decryption"

examples/src/default_cryptographic_materials_manager_example.py

@@ -109,18 +109,14 @@ def encrypt_and_decrypt_with_default_cmm(
         "Ciphertext and plaintext data are the same. Invalid encryption"
 
     # 7. Decrypt your encrypted data using the same cmm you used on encrypt.
-    plaintext_bytes, dec_header = client.decrypt(
+    plaintext_bytes, _ = client.decrypt(
         source=ciphertext,
-        materials_manager=cmm
+        materials_manager=cmm,
+        # Provide the encryption context that was supplied to the encrypt method
+        encryption_context=encryption_context,
     )
 
-    # 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"
-
-    # 9. Demonstrate that the decrypted plaintext is identical to the original plaintext.
+    # 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, \
         "Decrypted plaintext should be identical to the original plaintext. Invalid decryption"

examples/src/file_streaming_example.py

@@ -134,12 +134,6 @@ def encrypt_and_decrypt_with_keyring(
             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), \

examples/src/hierarchical_keyring_example.py

@@ -200,7 +200,10 @@ def encrypt_and_decrypt_with_keyring(
     try:
         client.decrypt(
             source=ciphertext_a,
-            keyring=hierarchical_keyring_b
+            keyring=hierarchical_keyring_b,
+            # Verify that the encryption context in the result contains the
+            # encryption context supplied to the encrypt method
+            encryption_context=encryption_context_a,
         )
     except AWSEncryptionSDKClientError:
         pass
@@ -210,22 +213,30 @@ def encrypt_and_decrypt_with_keyring(
     try:
         client.decrypt(
             source=ciphertext_b,
-            keyring=hierarchical_keyring_a
+            keyring=hierarchical_keyring_a,
+            # Verify that the encryption context in the result contains the
+            # encryption context supplied to the encrypt method
+            encryption_context=encryption_context_b,
         )
     except AWSEncryptionSDKClientError:
         pass
 
-    # 10. Demonstrate that data encrypted by one tenant's branch key can be decrypted by that tenant,
+    # 11. 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
+        keyring=hierarchical_keyring_a,
+        # Provide the encryption context that was supplied to the encrypt method
+        encryption_context=encryption_context_a,
     )
     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
+        keyring=hierarchical_keyring_b,
+        # Provide the encryption context that was supplied to the encrypt method
+        encryption_context=encryption_context_b,
     )
     assert plaintext_bytes_b == EXAMPLE_DATA, \
         "Decrypted plaintext should be identical to the original plaintext. Invalid decryption"