chore(python): examples for client_supplier

Author: imabhichow

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

@@ -0,0 +1,221 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""
+Example demonstrating DynamoDB Encryption using a custom client supplier.
+
+A custom client supplier grants users access to more granular configuration aspects
+of their authentication details and KMS client. The example creates a simple custom
+client supplier that authenticates with a different IAM role based on the region
+of the KMS key.
+
+Creates a MRK multi-keyring configured with a custom client supplier using a single
+MRK and puts an encrypted item to the table. Then, creates a MRK discovery
+multi-keyring to decrypt the item and retrieves the item from the table.
+
+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 (N)
+"""
+
+from typing import List
+
+import boto3
+from aws_cryptographic_material_providers.mpl import AwsCryptographicMaterialProviders
+from aws_cryptographic_material_providers.mpl.config import MaterialProvidersConfig
+from aws_cryptographic_material_providers.mpl.models import (
+    CreateAwsKmsMrkDiscoveryMultiKeyringInput,
+    CreateAwsKmsMrkMultiKeyringInput,
+    DiscoveryFilter,
+)
+from aws_dbesdk_dynamodb.encrypted.client import EncryptedClient
+from aws_dbesdk_dynamodb.structures.dynamodb import (
+    DynamoDbTableEncryptionConfig,
+    DynamoDbTablesEncryptionConfig,
+)
+from aws_dbesdk_dynamodb.structures.structured_encryption import (
+    CryptoAction,
+)
+
+from .regional_role_client_supplier import RegionalRoleClientSupplier
+
+
+def client_supplier_example(ddb_table_name: str, key_arn: str, account_ids: List[str], regions: List[str]) -> None:
+    """
+    Demonstrate using custom client supplier with AWS KMS MRK keyrings.
+
+    Shows how to use a custom client supplier with AWS KMS MRK multi-keyring and AWS
+    KMS MRK discovery multi-keyring.
+
+    :param ddb_table_name: The name of the DynamoDB table
+    :param key_arn: The ARN of the AWS KMS key
+    :param account_ids: List of AWS account IDs
+    :param regions: List of AWS regions
+    """
+    # 1. Create a single MRK multi-keyring.
+    #    This can be either a single-region KMS key or an MRK.
+    #    For this example to succeed, the key's region must either
+    #    1) be in the regions list, or
+    #    2) the key must be an MRK with a replica defined
+    #    in a region in the regions list, and the client
+    #    must have the correct permissions to access the replica.
+    mat_prov = AwsCryptographicMaterialProviders(config=MaterialProvidersConfig())
+
+    # Create the multi-keyring using our custom client supplier
+    # defined in the RegionalRoleClientSupplier class in this directory.
+    create_aws_kms_mrk_multi_keyring_input = CreateAwsKmsMrkMultiKeyringInput(
+        # Note: RegionalRoleClientSupplier will internally use the keyArn's region
+        # to retrieve the correct IAM role.
+        client_supplier=RegionalRoleClientSupplier(),
+        generator=key_arn,
+    )
+    mrk_keyring_with_client_supplier = mat_prov.create_aws_kms_mrk_multi_keyring(
+        input=create_aws_kms_mrk_multi_keyring_input
+    )
+
+    # 2. 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
+    #      - SIGN_ONLY: The attribute is not encrypted, but is still included in the signature
+    #      - DO_NOTHING: The attribute is not encrypted and not included in the signature
+    attribute_actions_on_encrypt = {
+        "partition_key": CryptoAction.SIGN_ONLY,  # Our partition attribute must be SIGN_ONLY
+        "sort_key": CryptoAction.SIGN_ONLY,  # Our sort attribute must be SIGN_ONLY
+        "sensitive_data": CryptoAction.ENCRYPT_AND_SIGN,
+    }
+
+    # 3. Configure which attributes we expect to be included in the signature
+    #    when reading items. There are two options for configuring this:
+    #
+    #    - (Recommended) Configure `allowed_unsigned_attribute_prefix`:
+    #      When defining your DynamoDb schema and deciding on attribute names,
+    #      choose a distinguishing prefix (such as ":") for all attributes that
+    #      you do not want to include in the signature.
+    #      This has two main benefits:
+    #      - It is easier to reason about the security and authenticity of data within your item
+    #        when all unauthenticated data is easily distinguishable by their attribute name.
+    #      - If you need to add new unauthenticated attributes in the future,
+    #        you can easily make the corresponding update to your `attribute_actions_on_encrypt`
+    #        and immediately start writing to that new attribute, without
+    #        any other configuration update needed.
+    #      Once you configure this field, it is not safe to update it.
+    #
+    #    - Configure `allowed_unsigned_attributes`: You may also explicitly list
+    #      a set of attributes that should be considered unauthenticated when encountered
+    #      on read. Be careful if you use this configuration. Do not remove an attribute
+    #      name from this configuration, even if you are no longer writing with that attribute,
+    #      as old items may still include this attribute, and our configuration needs to know
+    #      to continue to exclude this attribute from the signature scope.
+    #      If you add new attribute names to this field, you must first deploy the update to this
+    #      field to all readers in your host fleet before deploying the update to start writing
+    #      with that new attribute.
+    #
+    #   For this example, we currently authenticate all attributes. To make it easier to
+    #   add unauthenticated attributes in the future, we define a prefix ":" for such attributes.
+    unsign_attr_prefix = ":"
+
+    # 4. Create the DynamoDb Encryption configuration for the table we will be writing to.
+    table_config = DynamoDbTableEncryptionConfig(
+        logical_table_name=ddb_table_name,
+        partition_key_name="partition_key",
+        sort_key_name="sort_key",
+        attribute_actions_on_encrypt=attribute_actions_on_encrypt,
+        keyring=mrk_keyring_with_client_supplier,
+        allowed_unsigned_attribute_prefix=unsign_attr_prefix,
+    )
+
+    table_configs = {ddb_table_name: table_config}
+    tables_config = DynamoDbTablesEncryptionConfig(table_encryption_configs=table_configs)
+
+    # 5. Create the EncryptedClient
+    ddb_client = boto3.client("dynamodb")
+    encrypted_ddb_client = EncryptedClient(client=ddb_client, encryption_config=tables_config)
+
+    # 6. Put an item into our table using the above client.
+    #    Before the item gets sent to DynamoDb, it will be encrypted
+    #    client-side using the MRK multi-keyring.
+    #    The data key protecting this item will be encrypted
+    #    with all the KMS Keys in this keyring, so that it can be
+    #    decrypted with any one of those KMS Keys.
+    item = {
+        "partition_key": {"S": "clientSupplierItem"},
+        "sort_key": {"N": "0"},
+        "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
+
+    # 7. Get the item back from our table using the same keyring.
+    #    The client will decrypt the item client-side using the MRK
+    #    and return the original item.
+    key_to_get = {"partition_key": {"S": "clientSupplierItem"}, "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["sensitive_data"]["S"] == "encrypt and sign me!"
+
+    # 8. Create a MRK discovery multi-keyring with a custom client supplier.
+    #    A discovery MRK multi-keyring will be composed of
+    #    multiple discovery MRK keyrings, one for each region.
+    #    Each component keyring has its own KMS client in a particular region.
+    #    When we provide a client supplier to the multi-keyring, all component
+    #    keyrings will use that client supplier configuration.
+    #    In our tests, we make `key_arn` an MRK with a replica, and
+    #    provide only the replica region in our discovery filter.
+    discovery_filter = DiscoveryFilter(partition="aws", account_ids=account_ids)
+
+    mrk_discovery_client_supplier_input = CreateAwsKmsMrkDiscoveryMultiKeyringInput(
+        client_supplier=RegionalRoleClientSupplier(), discovery_filter=discovery_filter, regions=regions
+    )
+
+    mrk_discovery_client_supplier_keyring = mat_prov.create_aws_kms_mrk_discovery_multi_keyring(
+        input=mrk_discovery_client_supplier_input
+    )
+
+    # 9. Create a new config and client using the discovery keyring.
+    #    This is the same setup as above, except we provide the discovery keyring to the config.
+    replica_key_table_config = DynamoDbTableEncryptionConfig(
+        logical_table_name=ddb_table_name,
+        partition_key_name="partition_key",
+        sort_key_name="sort_key",
+        attribute_actions_on_encrypt=attribute_actions_on_encrypt,
+        # Provide discovery keyring here
+        keyring=mrk_discovery_client_supplier_keyring,
+        allowed_unsigned_attribute_prefix=unsign_attr_prefix,
+    )
+
+    replica_key_tables_config = {ddb_table_name: replica_key_table_config}
+    replica_key_tables_encryption_config = DynamoDbTablesEncryptionConfig(
+        table_encryption_configs=replica_key_tables_config
+    )
+
+    replica_key_encrypted_client = EncryptedClient(
+        client=ddb_client, encryption_config=replica_key_tables_encryption_config
+    )
+
+    # 10. Get the item back from our table using the discovery keyring client.
+    #     The client will decrypt the item client-side using the keyring,
+    #     and return the original item.
+    #     The discovery keyring will only use KMS keys in the provided regions and
+    #     AWS accounts. Since we have provided it with a custom client supplier
+    #     which uses different IAM roles based on the key region,
+    #     the discovery keyring will use a particular IAM role to decrypt
+    #     based on the region of the KMS key it uses to decrypt.
+    replica_key_key_to_get = {"partition_key": {"S": "awsKmsMrkMultiKeyringItem"}, "sort_key": {"N": "0"}}
+
+    replica_key_get_response = replica_key_encrypted_client.get_item(
+        TableName=ddb_table_name, Key=replica_key_key_to_get
+    )
+
+    # Demonstrate that GetItem succeeded and returned the decrypted item
+    assert replica_key_get_response["ResponseMetadata"]["HTTPStatusCode"] == 200
+    replica_key_returned_item = replica_key_get_response["Item"]
+    assert replica_key_returned_item["sensitive_data"]["S"] == "encrypt and sign me!"

Examples/runtimes/python/DynamoDBEncryption/src/client_supplier/regional_role_client_supplier.py

@@ -0,0 +1,69 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""
+Demonstrates implementing a custom client supplier.
+
+Creates KMS clients with different IAM roles depending on the region passed.
+"""
+
+import logging
+
+import boto3
+from aws_cryptographic_material_providers.mpl.models import GetClientInput
+from aws_cryptographic_material_providers.mpl.references import ClientSupplier
+from botocore.exceptions import ClientError
+
+from .regional_role_client_supplier_config import RegionalRoleClientSupplierConfig
+
+
+class RegionalRoleClientSupplier(ClientSupplier):
+    """
+    Custom client supplier for region-specific IAM roles.
+
+    Creates KMS clients with different IAM roles depending on the region passed.
+    """
+
+    def __init__(self):
+        """Initialize the client supplier with STS client and configuration."""
+        self._sts_client = boto3.client("sts")
+        self._config = RegionalRoleClientSupplierConfig()
+        self._logger = logging.getLogger(__name__)
+
+    def get_client(self, input_params: GetClientInput) -> boto3.client:
+        """
+        Get a KMS client for the specified region using the configured IAM role.
+
+        In test environments where assuming the role might fail, we fall back to
+        creating a standard KMS client for the region without assuming a role.
+        This ensures examples can run in test environments without proper IAM permissions.
+
+        :param input_params: Input parameters containing the region
+        :return: A boto3 KMS client for the specified region with the appropriate credentials
+        """
+        region = input_params.region
+        if region not in self._config.region_iam_role_map:
+            self._logger.warning(f"Missing region in config: {region}. Using default client.")
+            return boto3.client("kms", region_name=region)
+
+        role_arn = self._config.region_iam_role_map[region]
+
+        try:
+            # Assume the IAM role for the region
+            response = self._sts_client.assume_role(
+                RoleArn=role_arn,
+                DurationSeconds=900,  # 15 minutes is the minimum value
+                RoleSessionName="Python-Client-Supplier-Example-Session",
+            )
+
+            # Create a KMS client with the temporary credentials
+            return boto3.client(
+                "kms",
+                region_name=region,
+                aws_access_key_id=response["Credentials"]["AccessKeyId"],
+                aws_secret_access_key=response["Credentials"]["SecretAccessKey"],
+                aws_session_token=response["Credentials"]["SessionToken"],
+            )
+        except ClientError as e:
+            # In test environments, fall back to a standard client
+            self._logger.warning(f"Failed to assume role: {str(e)}. Falling back to default client.")
+            return boto3.client("kms", region_name=region)

Examples/runtimes/python/DynamoDBEncryption/src/client_supplier/regional_role_client_supplier_config.py

@@ -0,0 +1,26 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""
+Configuration for the RegionalRoleClientSupplier.
+
+Contains hardcoded configuration values for demonstration purposes. In production
+code, these values might be loaded from environment variables, AWS AppConfig, or
+other external sources.
+"""
+
+
+class RegionalRoleClientSupplierConfig:
+    """
+    Configuration class mapping AWS regions to IAM roles.
+
+    Provides a mapping between AWS regions and their corresponding IAM roles for
+    use in the RegionalRoleClientSupplier. For demonstration purposes, this uses
+    hardcoded values.
+    """
+
+    US_EAST_1_IAM_ROLE = "arn:aws:iam::370957321024:role/GitHub-CI-DDBEC-Dafny-Role-only-us-east-1-KMS-keys"
+    EU_WEST_1_IAM_ROLE = "arn:aws:iam::370957321024:role/GitHub-CI-DDBEC-Dafny-Role-only-eu-west-1-KMS-keys"
+
+    def __init__(self):
+        """Initialize the configuration with region to IAM role mapping."""
+        self.region_iam_role_map = {"us-east-1": self.US_EAST_1_IAM_ROLE, "eu-west-1": self.EU_WEST_1_IAM_ROLE}

Examples/runtimes/python/DynamoDBEncryption/test/client_supplier/__init__.py

@@ -0,0 +1,3 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Test suite for the client_supplier examples."""

Examples/runtimes/python/DynamoDBEncryption/test/client_supplier/test_client_supplier_example.py

@@ -0,0 +1,23 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Test suite for the client supplier example."""
+import pytest
+
+from ...src.client_supplier.client_supplier_example import client_supplier_example
+from .. import test_utils
+
+pytestmark = [pytest.mark.examples]
+
+
+def test_client_supplier_example():
+    """Test function for client supplier example."""
+    accounts = [test_utils.TEST_AWS_ACCOUNT_ID]
+    regions = ["eu-west-1"]  # Using eu-west-1
+
+    # Call the client_supplier_example with the test parameters
+    client_supplier_example(
+        ddb_table_name=test_utils.TEST_DDB_TABLE_NAME,
+        key_arn=test_utils.TEST_MRK_REPLICA_KEY_ID_US_EAST_1,
+        account_ids=accounts,
+        regions=regions,
+    )