diff --git a/Examples/runtimes/python/DynamoDBEncryption/src/client_supplier/__init__.py b/Examples/runtimes/python/DynamoDBEncryption/src/client_supplier/__init__.py new file mode 100644 index 000000000..fa977e22f --- /dev/null +++ b/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.""" diff --git a/Examples/runtimes/python/DynamoDBEncryption/src/client_supplier/client_supplier_example.py b/Examples/runtimes/python/DynamoDBEncryption/src/client_supplier/client_supplier_example.py new file mode 100644 index 000000000..a08a88bbd --- /dev/null +++ b/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!" diff --git a/Examples/runtimes/python/DynamoDBEncryption/src/client_supplier/regional_role_client_supplier.py b/Examples/runtimes/python/DynamoDBEncryption/src/client_supplier/regional_role_client_supplier.py new file mode 100644 index 000000000..ef44d6c49 --- /dev/null +++ b/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) diff --git a/Examples/runtimes/python/DynamoDBEncryption/src/client_supplier/regional_role_client_supplier_config.py b/Examples/runtimes/python/DynamoDBEncryption/src/client_supplier/regional_role_client_supplier_config.py new file mode 100644 index 000000000..98478fd71 --- /dev/null +++ b/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} diff --git a/Examples/runtimes/python/DynamoDBEncryption/test/client_supplier/__init__.py b/Examples/runtimes/python/DynamoDBEncryption/test/client_supplier/__init__.py new file mode 100644 index 000000000..6d76d639b --- /dev/null +++ b/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.""" diff --git a/Examples/runtimes/python/DynamoDBEncryption/test/client_supplier/test_client_supplier_example.py b/Examples/runtimes/python/DynamoDBEncryption/test/client_supplier/test_client_supplier_example.py new file mode 100644 index 000000000..fffd4afe6 --- /dev/null +++ b/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, + )