Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 17 commits into from
May 7, 2024
Merged
2 changes: 1 addition & 1 deletion .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ __pycache__
# PyTest
.pytest_cache
# Ignore key materials generated by examples or tests
test_keys/
test_keyrings/

# PyCharm
.idea/
Expand Down
3 changes: 2 additions & 1 deletion examples/src/keyrings/aws_kms_discovery_keyring_example.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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"
3 changes: 2 additions & 1 deletion examples/src/keyrings/aws_kms_keyring_example.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
3 changes: 2 additions & 1 deletion examples/src/keyrings/aws_kms_mrk_keyring_example.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
3 changes: 2 additions & 1 deletion examples/src/keyrings/aws_kms_mrk_multi_keyring_example.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
9 changes: 6 additions & 3 deletions examples/src/keyrings/aws_kms_multi_keyring_example.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand All @@ -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"
3 changes: 2 additions & 1 deletion examples/src/keyrings/aws_kms_rsa_keyring_example.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
152 changes: 152 additions & 0 deletions examples/src/keyrings/file_streaming_example.py
Original file line number Diff line number Diff line change
@@ -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 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 `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.

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 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 using a Raw AES keyring.

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
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"
62 changes: 32 additions & 30 deletions examples/src/keyrings/hierarchical_keyring_example.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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",
Expand All @@ -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",
Expand All @@ -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(
Expand All @@ -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(
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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"
Loading
Loading