chore(python): examples for EncryptedTable (#1934)

Co-authored-by: Lucas McDonald <[email protected]>

Author: imabhichow

Examples/runtimes/python/DynamoDBEncryption/src/get_encrypted_data_key_description_example.py

@@ -0,0 +1,82 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Example demonstrating how to get encrypted data key descriptions from DynamoDB items."""
+
+import boto3
+from aws_dbesdk_dynamodb.smithygenerated.aws_cryptography_dbencryptionsdk_dynamodb.client import DynamoDbEncryption
+from aws_dbesdk_dynamodb.smithygenerated.aws_cryptography_dbencryptionsdk_dynamodb.config import (
+    DynamoDbEncryptionConfig,
+)
+from aws_dbesdk_dynamodb.structures.dynamodb import (
+    GetEncryptedDataKeyDescriptionInput,
+    GetEncryptedDataKeyDescriptionUnionItem,
+)
+
+
+def get_encrypted_data_key_description(
+    table_name: str,
+    partition_key: str,
+    partition_key_val: str,
+    sort_key_name: str,
+    sort_key_value: str,
+    expected_key_provider_id: str,
+    expected_key_provider_info: str,
+    expected_branch_key_id: str,
+    expected_branch_key_version: str,
+):
+    """
+    Get encrypted data key description from a DynamoDB item.
+
+    :param table_name: The name of the DynamoDB table
+    :param partition_key: The name of the partition key
+    :param partition_key_val: The value of the partition key
+    :param sort_key_name: The name of the sort key
+    :param sort_key_value: The value of the sort key
+    :param expected_key_provider_id: The expected key provider ID
+    :param expected_key_provider_info: The expected key provider info (optional)
+    :param expected_branch_key_id: The expected branch key ID (optional)
+    :param expected_branch_key_version: The expected branch key version (optional)
+    """
+    # 1. Create a new AWS SDK DynamoDb client. This client will be used to get item from the DynamoDB table
+    ddb = boto3.client("dynamodb")
+
+    # 2. Get item from the DynamoDB table. This item will be used to Get Encrypted DataKey Description
+    key_to_get = {partition_key: {"S": partition_key_val}, sort_key_name: {"N": sort_key_value}}
+
+    response = ddb.get_item(TableName=table_name, Key=key_to_get)
+
+    returned_item = response.get("Item", {})
+    if not returned_item:
+        print(f"No item found with the key {partition_key}!")
+        return
+
+    # 3. Prepare the input for GetEncryptedDataKeyDescription method.
+    # This input can be a DynamoDB item or a header. For now, we are giving input as a DynamoDB item
+    # but users can also extract the header from the attribute "aws_dbe_head" in the DynamoDB table
+    # and use it for GetEncryptedDataKeyDescription method.
+    ddb_enc = DynamoDbEncryption(config=DynamoDbEncryptionConfig())
+
+    input_union = GetEncryptedDataKeyDescriptionUnionItem(returned_item)
+
+    input_obj = GetEncryptedDataKeyDescriptionInput(input=input_union)
+
+    output = ddb_enc.get_encrypted_data_key_description(input=input_obj)
+
+    # In the following code, we are giving input as header instead of a complete DynamoDB item
+    # This code is provided solely to demo how the alternative approach works. So, it is commented.
+
+    # header_attribute = "aws_dbe_head"
+    # header = returned_item[header_attribute]["B"]
+    # input_union = GetEncryptedDataKeyDescriptionUnion(
+    #     header=header
+    # )
+
+    # Assert everything
+    assert output.encrypted_data_key_description_output[0].key_provider_id == expected_key_provider_id
+
+    if expected_key_provider_id.startswith("aws-kms"):
+        assert output.encrypted_data_key_description_output[0].key_provider_info == expected_key_provider_info
+
+    if output.encrypted_data_key_description_output[0].key_provider_id == "aws-kms-hierarchy":
+        assert output.encrypted_data_key_description_output[0].branch_key_id == expected_branch_key_id
+        assert output.encrypted_data_key_description_output[0].branch_key_version == expected_branch_key_version

Examples/runtimes/python/DynamoDBEncryption/src/keyring/hierarchical_keyring_example/__init__.py

@@ -0,0 +1,3 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Stub to allow relative imports of examples from tests."""

Examples/runtimes/python/DynamoDBEncryption/src/keyring/hierarchical_keyring_example.py renamed to Examples/runtimes/python/DynamoDBEncryption/src/keyring/hierarchical_keyring_example/encryption_config.py

@@ -1,49 +1,10 @@
 # Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 """
-Example demonstrating DynamoDb Encryption using a Hierarchical Keyring.
+Configuration module for hierarchical keyring encryption setup.
 
-This example sets up DynamoDb Encryption for the AWS SDK client
-using the Hierarchical Keyring, which establishes a key hierarchy
-where "branch" keys are persisted in DynamoDb.
-These branch keys are used to protect your data keys,
-and these branch keys are themselves protected by a root KMS Key.
-
-Establishing a key hierarchy like this has two benefits:
-
-First, by caching the branch key material, and only calling back
-to KMS to re-establish authentication regularly according to your configured TTL,
-you limit how often you need to call back to KMS to protect your data.
-This is a performance/security tradeoff, where your authentication, audit, and
-logging from KMS is no longer one-to-one with every encrypt or decrypt call.
-However, the benefit is that you no longer have to make a
-network call to KMS for every encrypt or decrypt.
-
-Second, this key hierarchy makes it easy to hold multi-tenant data
-that is isolated per branch key in a single DynamoDb table.
-You can create a branch key for each tenant in your table,
-and encrypt all that tenant's data under that distinct branch key.
-On decrypt, you can either statically configure a single branch key
-to ensure you are restricting decryption to a single tenant,
-or you can implement an interface that lets you map the primary key on your items
-to the branch key that should be responsible for decrypting that data.
-
-This example then demonstrates configuring a Hierarchical Keyring
-with a Branch Key ID Supplier to encrypt and decrypt data for
-two separate tenants.
-
-Running this example requires access to the DDB Table whose name
-is provided in CLI arguments.
-This table must be configured with the following
-primary key configuration:
-  - Partition key is named "partition_key" with type (S)
-  - Sort key is named "sort_key" with type (S)
-
-This example also requires using a KMS Key whose ARN
-is provided in CLI arguments. You need the following access
-on this key:
-  - GenerateDataKeyWithoutPlaintext
-  - Decrypt
+This module provides the common encryption configuration used by both
+EncryptedClient and EncryptedTable examples.
 """
 
 import boto3
@@ -57,7 +18,6 @@
     CreateAwsKmsHierarchicalKeyringInput,
     DefaultCache,
 )
-from aws_dbesdk_dynamodb.encrypted.client import EncryptedClient
 from aws_dbesdk_dynamodb.smithygenerated.aws_cryptography_dbencryptionsdk_dynamodb.client import DynamoDbEncryption
 from aws_dbesdk_dynamodb.smithygenerated.aws_cryptography_dbencryptionsdk_dynamodb.config import (
     DynamoDbEncryptionConfig,
@@ -74,23 +34,24 @@
 from .example_branch_key_id_supplier import ExampleBranchKeyIdSupplier
 
 
-def hierarchical_keyring_get_item_put_item(
+def create_encryption_config(
     ddb_table_name: str,
     tenant1_branch_key_id: str,
     tenant2_branch_key_id: str,
     keystore_table_name: str,
     logical_keystore_name: str,
     kms_key_id: str,
-):
+) -> DynamoDbTablesEncryptionConfig:
     """
-    Demonstrate using a hierarchical keyring with multiple tenants.
+    Create the encryption configuration for DynamoDB encryption.
 
     :param ddb_table_name: The name of the DynamoDB table
     :param tenant1_branch_key_id: Branch key ID for tenant 1
     :param tenant2_branch_key_id: Branch key ID for tenant 2
     :param keystore_table_name: The name of the KeyStore DynamoDB table
     :param logical_keystore_name: The logical name for this keystore
     :param kms_key_id: The ARN of the KMS key to use
+    :return: The DynamoDB tables encryption configuration
     """
     # Initial KeyStore Setup: This example requires that you have already
     # created your KeyStore, and have populated it with two new branch keys.
@@ -190,40 +151,4 @@ def hierarchical_keyring_get_item_put_item(
     )
 
     table_configs = {ddb_table_name: table_config}
-    tables_config = DynamoDbTablesEncryptionConfig(table_encryption_configs=table_configs)
-
-    # 7. Create the EncryptedClient
-    ddb_client = boto3.client("dynamodb")
-    encrypted_ddb_client = EncryptedClient(client=ddb_client, encryption_config=tables_config)
-
-    # 8. Put an item into our table using the above client.
-    #    Before the item gets sent to DynamoDb, it will be encrypted
-    #    client-side, according to our configuration.
-    #    Because the item we are writing uses "tenantId1" as our partition value,
-    #    based on the code we wrote in the ExampleBranchKeySupplier,
-    #    `tenant1_branch_key_id` will be used to encrypt this item.
-    item = {
-        "partition_key": {"S": "tenant1Id"},
-        "sort_key": {"N": "0"},
-        "tenant_sensitive_data": {"S": "encrypt and sign me!"},
-    }
-
-    put_response = encrypted_ddb_client.put_item(TableName=ddb_table_name, Item=item)
-
-    # Demonstrate that PutItem succeeded
-    assert put_response["ResponseMetadata"]["HTTPStatusCode"] == 200
-
-    # 9. Get the item back from our table using the same client.
-    #    The client will decrypt the item client-side, and return
-    #    back the original item.
-    #    Because the returned item's partition value is "tenantId1",
-    #    based on the code we wrote in the ExampleBranchKeySupplier,
-    #    `tenant1_branch_key_id` will be used to decrypt this item.
-    key_to_get = {"partition_key": {"S": "tenant1Id"}, "sort_key": {"N": "0"}}
-
-    get_response = encrypted_ddb_client.get_item(TableName=ddb_table_name, Key=key_to_get)
-
-    # Demonstrate that GetItem succeeded and returned the decrypted item
-    assert get_response["ResponseMetadata"]["HTTPStatusCode"] == 200
-    returned_item = get_response["Item"]
-    assert returned_item["tenant_sensitive_data"]["S"] == "encrypt and sign me!"
+    return DynamoDbTablesEncryptionConfig(table_encryption_configs=table_configs)

Examples/runtimes/python/DynamoDBEncryption/src/keyring/example_branch_key_id_supplier.py renamed to Examples/runtimes/python/DynamoDBEncryption/src/keyring/hierarchical_keyring_example/example_branch_key_id_supplier.py

@@ -0,0 +1,118 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""
+Example demonstrating DynamoDb Encryption using a Hierarchical Keyring with EncryptedClient.
+
+This example sets up DynamoDb Encryption for the AWS SDK client
+using the Hierarchical Keyring, which establishes a key hierarchy
+where "branch" keys are persisted in DynamoDb.
+These branch keys are used to protect your data keys,
+and these branch keys are themselves protected by a root KMS Key.
+
+Establishing a key hierarchy like this has two benefits:
+
+First, by caching the branch key material, and only calling back
+to KMS to re-establish authentication regularly according to your configured TTL,
+you limit how often you need to call back to KMS to protect your data.
+This is a performance/security tradeoff, where your authentication, audit, and
+logging from KMS is no longer one-to-one with every encrypt or decrypt call.
+However, the benefit is that you no longer have to make a
+network call to KMS for every encrypt or decrypt.
+
+Second, this key hierarchy makes it easy to hold multi-tenant data
+that is isolated per branch key in a single DynamoDb table.
+You can create a branch key for each tenant in your table,
+and encrypt all that tenant's data under that distinct branch key.
+On decrypt, you can either statically configure a single branch key
+to ensure you are restricting decryption to a single tenant,
+or you can implement an interface that lets you map the primary key on your items
+to the branch key that should be responsible for decrypting that data.
+
+This example then demonstrates configuring a Hierarchical Keyring
+with a Branch Key ID Supplier to encrypt and decrypt data for
+two separate tenants.
+
+Running this example requires access to the DDB Table whose name
+is provided in CLI arguments.
+This table must be configured with the following
+primary key configuration:
+  - Partition key is named "partition_key" with type (S)
+  - Sort key is named "sort_key" with type (S)
+
+This example also requires using a KMS Key whose ARN
+is provided in CLI arguments. You need the following access
+on this key:
+  - GenerateDataKeyWithoutPlaintext
+  - Decrypt
+"""
+
+import boto3
+from aws_dbesdk_dynamodb.encrypted.client import EncryptedClient
+
+from .encryption_config import create_encryption_config
+
+
+def hierarchical_keyring_client_example(
+    ddb_table_name: str,
+    tenant1_branch_key_id: str,
+    tenant2_branch_key_id: str,
+    keystore_table_name: str,
+    logical_keystore_name: str,
+    kms_key_id: str,
+):
+    """
+    Demonstrate using a hierarchical keyring with multiple tenants using EncryptedClient.
+
+    :param ddb_table_name: The name of the DynamoDB table
+    :param tenant1_branch_key_id: Branch key ID for tenant 1
+    :param tenant2_branch_key_id: Branch key ID for tenant 2
+    :param keystore_table_name: The name of the KeyStore DynamoDB table
+    :param logical_keystore_name: The logical name for this keystore
+    :param kms_key_id: The ARN of the KMS key to use
+    """
+    # 1. Create the DynamoDb Encryption configuration for the table we will be writing to.
+    #    See beacon_config.py in this directory for detailed steps on the encryption configuration.
+    tables_config = create_encryption_config(
+        ddb_table_name=ddb_table_name,
+        tenant1_branch_key_id=tenant1_branch_key_id,
+        tenant2_branch_key_id=tenant2_branch_key_id,
+        keystore_table_name=keystore_table_name,
+        logical_keystore_name=logical_keystore_name,
+        kms_key_id=kms_key_id,
+    )
+
+    # 2. Create the EncryptedClient
+    ddb_client = boto3.client("dynamodb")
+    encrypted_ddb_client = EncryptedClient(client=ddb_client, encryption_config=tables_config)
+
+    # 3. Put an item into our table using the above client.
+    #    Before the item gets sent to DynamoDb, it will be encrypted
+    #    client-side, according to our configuration.
+    #    Because the item we are writing uses "tenantId1" as our partition value,
+    #    based on the code we wrote in the ExampleBranchKeySupplier,
+    #    `tenant1_branch_key_id` will be used to encrypt this item.
+    item = {
+        "partition_key": {"S": "tenant1Id"},
+        "sort_key": {"N": "0"},
+        "tenant_sensitive_data": {"S": "encrypt and sign me!"},
+    }
+
+    put_response = encrypted_ddb_client.put_item(TableName=ddb_table_name, Item=item)
+
+    # Demonstrate that PutItem succeeded
+    assert put_response["ResponseMetadata"]["HTTPStatusCode"] == 200
+
+    # 4. Get the item back from our table using the same client.
+    #    The client will decrypt the item client-side, and return
+    #    back the original item.
+    #    Because the returned item's partition value is "tenantId1",
+    #    based on the code we wrote in the ExampleBranchKeySupplier,
+    #    `tenant1_branch_key_id` will be used to decrypt this item.
+    key_to_get = {"partition_key": {"S": "tenant1Id"}, "sort_key": {"N": "0"}}
+
+    get_response = encrypted_ddb_client.get_item(TableName=ddb_table_name, Key=key_to_get)
+
+    # Demonstrate that GetItem succeeded and returned the decrypted item
+    assert get_response["ResponseMetadata"]["HTTPStatusCode"] == 200
+    returned_item = get_response["Item"]
+    assert returned_item["tenant_sensitive_data"]["S"] == "encrypt and sign me!"