chore(python): examples for keyrings with EncryptedClient

Author: imabhichow

Examples/runtimes/java/DynamoDbEncryption/src/main/java/software/amazon/cryptography/examples/keyring/SharedCacheAcrossHierarchicalKeyringsExample.java

@@ -180,7 +180,7 @@ public static void SharedCacheAcrossHierarchicalKeyringsGetItemPutItem(
     final IKeyring hierarchicalKeyring1 =
       matProv.CreateAwsKmsHierarchicalKeyring(keyringInput1);
 
-    // 4. Configure which attributes are encrypted and/or signed when writing new items.
+    // 5. Configure which attributes are encrypted and/or signed when writing new items.
     //    For each attribute that may exist on the items we plan to write to our DynamoDbTable,
     //    we must explicitly configure how they should be treated during item encryption:
     //      - ENCRYPT_AND_SIGN: The attribute is encrypted and included in the signature
@@ -194,14 +194,14 @@ public static void SharedCacheAcrossHierarchicalKeyringsGetItemPutItem(
       CryptoAction.ENCRYPT_AND_SIGN
     );
 
-    // 5. Get the DDB Client for Hierarchical Keyring 1.
+    // 6. Get the DDB Client for Hierarchical Keyring 1.
     final DynamoDbClient ddbClient1 = GetDdbClient(
       ddbTableName,
       hierarchicalKeyring1,
       attributeActionsOnEncrypt
     );
 
-    // 6. Encrypt Decrypt roundtrip with ddbClient1
+    // 7. Encrypt Decrypt roundtrip with ddbClient1
     PutGetItems(ddbTableName, ddbClient1);
 
     // Through the above encrypt and decrypt roundtrip, the cache will be populated and
@@ -210,7 +210,7 @@ public static void SharedCacheAcrossHierarchicalKeyringsGetItemPutItem(
     // - Same Logical Key Store Name of the Key Store for the Hierarchical Keyring
     // - Same Branch Key ID
 
-    // 7. Configure your KeyStore resource keystore2.
+    // 8. Configure your KeyStore resource keystore2.
     //       This SHOULD be the same configuration that you used
     //       to initially create and populate your physical KeyStore.
     //    Note that keyStoreTableName is the physical Key Store,
@@ -243,7 +243,7 @@ public static void SharedCacheAcrossHierarchicalKeyringsGetItemPutItem(
       )
       .build();
 
-    // 8. Create the Hierarchical Keyring HK2 with Key Store instance K2, the shared Cache
+    // 9. Create the Hierarchical Keyring HK2 with Key Store instance K2, the shared Cache
     //    and the same partitionId and BranchKeyId used in HK1 because we want to share cache entries
     //    (and experience cache HITS).
 
@@ -262,14 +262,14 @@ public static void SharedCacheAcrossHierarchicalKeyringsGetItemPutItem(
     final IKeyring hierarchicalKeyring2 =
       matProv.CreateAwsKmsHierarchicalKeyring(keyringInput2);
 
-    // 9. Get the DDB Client for Hierarchical Keyring 2.
+    // 10. Get the DDB Client for Hierarchical Keyring 2.
     final DynamoDbClient ddbClient2 = GetDdbClient(
       ddbTableName,
       hierarchicalKeyring2,
       attributeActionsOnEncrypt
     );
 
-    // 10. Encrypt Decrypt roundtrip with ddbClient2
+    // 11. Encrypt Decrypt roundtrip with ddbClient2
     PutGetItems(ddbTableName, ddbClient2);
   }
 

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

@@ -0,0 +1,49 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""
+Example for creating a new key in a KeyStore.
+
+The Hierarchical Keyring Example and Searchable Encryption Examples rely on the
+existence of a DDB-backed key store with pre-existing branch key material or
+beacon key material.
+
+See the "Create KeyStore Table Example" for how to first set up the DDB Table
+that will back this KeyStore.
+
+Demonstrates configuring a KeyStore and using a helper method to create a branch
+key and beacon key that share the same Id. A new beacon key is always created
+alongside a new branch key, even if searchable encryption is not being used.
+
+Note: This key creation should occur within your control plane.
+"""
+
+import boto3
+from aws_cryptographic_material_providers.keystore.client import KeyStore
+from aws_cryptographic_material_providers.keystore.config import KeyStoreConfig
+from aws_cryptographic_material_providers.keystore.models import (
+    CreateKeyInput,
+    KMSConfigurationKmsKeyArn,
+)
+
+
+def keystore_create_key(key_store_table_name: str, logical_key_store_name: str, kms_key_arn: str) -> str:
+    """Create a new branch key and beacon key in our KeyStore."""
+    # 1. Configure your KeyStore resource.
+    #    This SHOULD be the same configuration that was used to create the DDB table
+    #    in the "Create KeyStore Table Example".
+    keystore: KeyStore = KeyStore(
+        KeyStoreConfig(
+            ddb_table_name=key_store_table_name,
+            kms_configuration=KMSConfigurationKmsKeyArn(kms_key_arn),
+            logical_key_store_name=logical_key_store_name,
+            kms_client=boto3.client("kms"),
+            ddb_client=boto3.client("dynamodb"),
+        )
+    )
+
+    # 2. Create a new branch key and beacon key in our KeyStore.
+    #    Both the branch key and the beacon key will share an Id.
+    #    This creation is eventually consistent.
+    branch_key_id = keystore.create_key(CreateKeyInput()).branch_key_identifier
+
+    return branch_key_id

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

@@ -0,0 +1,59 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""
+Example for creating a DynamoDB table for use as a KeyStore.
+
+The Hierarchical Keyring Example and Searchable Encryption Examples rely on the
+existence of a DDB-backed key store with pre-existing branch key material or
+beacon key material.
+
+Shows how to configure a KeyStore and use a helper method to create the DDB table
+that will be used to persist branch keys and beacons keys for this KeyStore.
+
+This table creation should occur within your control plane and only needs to occur
+once. While not demonstrated in this example, you should additionally use the
+`VersionKey` API on the KeyStore to periodically rotate your branch key material.
+"""
+
+import boto3
+from aws_cryptographic_material_providers.keystore.client import KeyStore
+from aws_cryptographic_material_providers.keystore.config import KeyStoreConfig
+from aws_cryptographic_material_providers.keystore.models import (
+    CreateKeyStoreInput,
+    KMSConfigurationKmsKeyArn,
+)
+
+
+def keystore_create_table(keystore_table_name: str, logical_keystore_name: str, kms_key_arn: str):
+    """
+    Create KeyStore Table Example.
+
+    :param keystore_table_name: The name of the DynamoDB table to create
+    :param logical_keystore_name: The logical name for this keystore
+    :param kms_key_arn: The ARN of the KMS key to use for protecting branch keys
+    """
+    # 1. Configure your KeyStore resource.
+    #    `ddb_table_name` is the name you want for the DDB table that
+    #    will back your keystore.
+    #    `kms_key_arn` is the KMS Key that will protect your branch keys and beacon keys
+    #    when they are stored in your DDB table.
+    keystore = KeyStore(
+        config=KeyStoreConfig(
+            ddb_client=boto3.client("dynamodb"),
+            ddb_table_name=keystore_table_name,
+            logical_key_store_name=logical_keystore_name,
+            kms_client=boto3.client("kms"),
+            kms_configuration=KMSConfigurationKmsKeyArn(kms_key_arn),
+        )
+    )
+
+    # 2. Create the DynamoDb table that will store the branch keys and beacon keys.
+    #    This checks if the correct table already exists at `ddb_table_name`
+    #    by using the DescribeTable API. If no table exists,
+    #    it will create one. If a table exists, it will verify
+    #    the table's configuration and will error if the configuration is incorrect.
+    keystore.create_key_store(input=CreateKeyStoreInput())
+    # It may take a couple of minutes for the table to become ACTIVE,
+    # at which point it is ready to store branch and beacon keys.
+    # See the Create KeyStore Key Example for how to populate
+    # this table.

Examples/runtimes/python/DynamoDBEncryption/src/item_encryptor/encrypt_decrypt_example.py

@@ -45,7 +45,7 @@
 )
 
 
-def encrypt_decrypt_example(kms_key_id: str, ddb_table_name: str) -> None:
+def encrypt_decrypt_example(kms_key_id: str, ddb_table_name: str):
     """Encrypt and decrypt an item with an ItemEncryptor."""
     # 1. Create a Keyring. This Keyring will be responsible for protecting the data keys that protect your data.
     #    For this example, we will create a AWS KMS Keyring with the AWS KMS Key we want to use.

Examples/runtimes/python/DynamoDBEncryption/src/keyring/__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/example_branch_key_id_supplier.py

@@ -0,0 +1,61 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""
+Example implementation of a branch key ID supplier.
+
+Used in the 'HierarchicalKeyringExample'.
+In that example, we have a table where we distinguish multiple tenants
+by a tenant ID that is stored in our partition attribute.
+The expectation is that this does not produce a confused deputy
+because the tenants are separated by partition.
+In order to create a Hierarchical Keyring that is capable of encrypting or
+decrypting data for either tenant, we implement this interface
+to map the correct branch key ID to the correct tenant ID.
+"""
+from typing import Dict
+
+from aws_dbesdk_dynamodb.smithygenerated.aws_cryptography_dbencryptionsdk_dynamodb.references import (
+    IDynamoDbKeyBranchKeyIdSupplier,
+)
+from aws_dbesdk_dynamodb.structures.dynamodb import GetBranchKeyIdFromDdbKeyInput, GetBranchKeyIdFromDdbKeyOutput
+
+
+class ExampleBranchKeyIdSupplier(IDynamoDbKeyBranchKeyIdSupplier):
+    """Example implementation of a branch key ID supplier."""
+
+    branch_key_id_for_tenant1: str
+    branch_key_id_for_tenant2: str
+
+    def __init__(self, tenant1_id: str, tenant2_id: str):
+        """
+        Initialize a branch key ID supplier.
+
+        :param tenant1_id: Branch key ID for tenant 1
+        :param tenant2_id: Branch key ID for tenant 2
+        """
+        self.branch_key_id_for_tenant1 = tenant1_id
+        self.branch_key_id_for_tenant2 = tenant2_id
+
+    def get_branch_key_id_from_ddb_key(self, param: GetBranchKeyIdFromDdbKeyInput) -> GetBranchKeyIdFromDdbKeyOutput:
+        """
+        Get branch key ID from the tenant ID in input's DDB key.
+
+        :param param: Input containing DDB key
+        :return: Output containing branch key ID
+        :raises ValueError: If DDB key is invalid or contains invalid tenant ID
+        """
+        key: Dict[str, Dict] = param.ddb_key
+
+        if "partition_key" not in key:
+            raise ValueError("Item invalid, does not contain expected partition key attribute.")
+
+        tenant_key_id = key["partition_key"]["S"]
+
+        if tenant_key_id == "tenant1Id":
+            branch_key_id = self.branch_key_id_for_tenant1
+        elif tenant_key_id == "tenant2Id":
+            branch_key_id = self.branch_key_id_for_tenant2
+        else:
+            raise ValueError("Item does not contain valid tenant ID")
+
+        return GetBranchKeyIdFromDdbKeyOutput(branch_key_id=branch_key_id)