1+ # Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
2+ # SPDX-License-Identifier: Apache-2.0
3+ """
4+ This example sets up an MRK multi-keyring and an MRK discovery
5+ multi-keyring using a custom client supplier.
6+
7+ A custom client supplier grants users access to more granular
8+ configuration aspects of their authentication details and KMS
9+ client. In this example, we create a simple custom client supplier
10+ that authenticates with a different IAM role based on the
11+ region of the KMS key.
12+
13+ This example creates a MRK multi-keyring configured with a custom
14+ client supplier using a single MRK and puts an encrypted item to the
15+ table. Then, it creates a MRK discovery multi-keyring to decrypt the item
16+ and retrieves the item from the table.
17+
18+ Running this example requires access to the DDB Table whose name
19+ is provided in CLI arguments.
20+ This table must be configured with the following
21+ primary key configuration:
22+ - Partition key is named "partition_key" with type (S)
23+ - Sort key is named "sort_key" with type (N)
24+ """
25+
26+ from typing import List
27+
28+ import boto3
29+ from aws_cryptographic_material_providers .mpl import AwsCryptographicMaterialProviders
30+ from aws_cryptographic_material_providers .mpl .config import MaterialProvidersConfig
31+ from aws_cryptographic_material_providers .mpl .models import (
32+ CreateAwsKmsMrkMultiKeyringInput ,
33+ CreateAwsKmsMrkDiscoveryMultiKeyringInput ,
34+ DiscoveryFilter ,
35+ )
36+ from aws_cryptographic_material_providers .mpl .references import IKeyring
37+ from aws_dbesdk_dynamodb .encrypted .client import EncryptedClient
38+ from aws_dbesdk_dynamodb .structures .dynamodb import (
39+ DynamoDbTableEncryptionConfig ,
40+ DynamoDbTablesEncryptionConfig ,
41+ )
42+ from aws_dbesdk_dynamodb .structures .structured_encryption import (
43+ CryptoAction ,
44+ )
45+
46+ from .regional_role_client_supplier import RegionalRoleClientSupplier
47+
48+
49+ def client_supplier_example (
50+ ddb_table_name : str ,
51+ key_arn : str ,
52+ account_ids : List [str ],
53+ regions : List [str ]
54+ ) -> None :
55+ """
56+ Demonstrate how to use a custom client supplier with AWS KMS MRK multi-keyring
57+ and AWS KMS MRK discovery multi-keyring.
58+
59+ :param ddb_table_name: The name of the DynamoDB table
60+ :param key_arn: The ARN of the AWS KMS key
61+ :param account_ids: List of AWS account IDs
62+ :param regions: List of AWS regions
63+ """
64+ # 1. Create a single MRK multi-keyring.
65+ # This can be either a single-region KMS key or an MRK.
66+ # For this example to succeed, the key's region must either
67+ # 1) be in the regions list, or
68+ # 2) the key must be an MRK with a replica defined
69+ # in a region in the regions list, and the client
70+ # must have the correct permissions to access the replica.
71+ mat_prov = AwsCryptographicMaterialProviders (config = MaterialProvidersConfig ())
72+
73+ # Create the multi-keyring using our custom client supplier
74+ # defined in the RegionalRoleClientSupplier class in this directory.
75+ create_aws_kms_mrk_multi_keyring_input = CreateAwsKmsMrkMultiKeyringInput (
76+ # Note: RegionalRoleClientSupplier will internally use the keyArn's region
77+ # to retrieve the correct IAM role.
78+ client_supplier = RegionalRoleClientSupplier (),
79+ generator = key_arn
80+ )
81+ mrk_keyring_with_client_supplier = mat_prov .create_aws_kms_mrk_multi_keyring (
82+ input = create_aws_kms_mrk_multi_keyring_input
83+ )
84+
85+ # 2. Configure which attributes are encrypted and/or signed when writing new items.
86+ # For each attribute that may exist on the items we plan to write to our DynamoDbTable,
87+ # we must explicitly configure how they should be treated during item encryption:
88+ # - ENCRYPT_AND_SIGN: The attribute is encrypted and included in the signature
89+ # - SIGN_ONLY: The attribute is not encrypted, but is still included in the signature
90+ # - DO_NOTHING: The attribute is not encrypted and not included in the signature
91+ attribute_actions_on_encrypt = {
92+ "partition_key" : CryptoAction .SIGN_ONLY , # Our partition attribute must be SIGN_ONLY
93+ "sort_key" : CryptoAction .SIGN_ONLY , # Our sort attribute must be SIGN_ONLY
94+ "sensitive_data" : CryptoAction .ENCRYPT_AND_SIGN
95+ }
96+
97+ # 3. Configure which attributes we expect to be included in the signature
98+ # when reading items. There are two options for configuring this:
99+ #
100+ # - (Recommended) Configure `allowed_unsigned_attribute_prefix`:
101+ # When defining your DynamoDb schema and deciding on attribute names,
102+ # choose a distinguishing prefix (such as ":") for all attributes that
103+ # you do not want to include in the signature.
104+ # This has two main benefits:
105+ # - It is easier to reason about the security and authenticity of data within your item
106+ # when all unauthenticated data is easily distinguishable by their attribute name.
107+ # - If you need to add new unauthenticated attributes in the future,
108+ # you can easily make the corresponding update to your `attribute_actions_on_encrypt`
109+ # and immediately start writing to that new attribute, without
110+ # any other configuration update needed.
111+ # Once you configure this field, it is not safe to update it.
112+ #
113+ # - Configure `allowed_unsigned_attributes`: You may also explicitly list
114+ # a set of attributes that should be considered unauthenticated when encountered
115+ # on read. Be careful if you use this configuration. Do not remove an attribute
116+ # name from this configuration, even if you are no longer writing with that attribute,
117+ # as old items may still include this attribute, and our configuration needs to know
118+ # to continue to exclude this attribute from the signature scope.
119+ # If you add new attribute names to this field, you must first deploy the update to this
120+ # field to all readers in your host fleet before deploying the update to start writing
121+ # with that new attribute.
122+ #
123+ # For this example, we currently authenticate all attributes. To make it easier to
124+ # add unauthenticated attributes in the future, we define a prefix ":" for such attributes.
125+ unsign_attr_prefix = ":"
126+
127+ # 4. Create the DynamoDb Encryption configuration for the table we will be writing to.
128+ table_config = DynamoDbTableEncryptionConfig (
129+ logical_table_name = ddb_table_name ,
130+ partition_key_name = "partition_key" ,
131+ sort_key_name = "sort_key" ,
132+ attribute_actions_on_encrypt = attribute_actions_on_encrypt ,
133+ keyring = mrk_keyring_with_client_supplier ,
134+ allowed_unsigned_attribute_prefix = unsign_attr_prefix
135+ )
136+
137+ table_configs = {ddb_table_name : table_config }
138+ tables_config = DynamoDbTablesEncryptionConfig (table_encryption_configs = table_configs )
139+
140+ # 5. Create the EncryptedClient
141+ ddb_client = boto3 .client ('dynamodb' )
142+ encrypted_ddb_client = EncryptedClient (
143+ client = ddb_client ,
144+ encryption_config = tables_config
145+ )
146+
147+ # 6. Put an item into our table using the above client.
148+ # Before the item gets sent to DynamoDb, it will be encrypted
149+ # client-side using the MRK multi-keyring.
150+ # The data key protecting this item will be encrypted
151+ # with all the KMS Keys in this keyring, so that it can be
152+ # decrypted with any one of those KMS Keys.
153+ item = {
154+ "partition_key" : {"S" : "clientSupplierItem" },
155+ "sort_key" : {"N" : "0" },
156+ "sensitive_data" : {"S" : "encrypt and sign me!" }
157+ }
158+
159+ put_response = encrypted_ddb_client .put_item (
160+ TableName = ddb_table_name ,
161+ Item = item
162+ )
163+
164+ # Demonstrate that PutItem succeeded
165+ assert put_response ['ResponseMetadata' ]['HTTPStatusCode' ] == 200
166+
167+ # 7. Get the item back from our table using the same keyring.
168+ # The client will decrypt the item client-side using the MRK
169+ # and return the original item.
170+ key_to_get = {
171+ "partition_key" : {"S" : "clientSupplierItem" },
172+ "sort_key" : {"N" : "0" }
173+ }
174+
175+ get_response = encrypted_ddb_client .get_item (
176+ TableName = ddb_table_name ,
177+ Key = key_to_get
178+ )
179+
180+ # Demonstrate that GetItem succeeded and returned the decrypted item
181+ assert get_response ['ResponseMetadata' ]['HTTPStatusCode' ] == 200
182+ returned_item = get_response ['Item' ]
183+ assert returned_item ["sensitive_data" ]["S" ] == "encrypt and sign me!"
184+
185+ # 8. Create a MRK discovery multi-keyring with a custom client supplier.
186+ # A discovery MRK multi-keyring will be composed of
187+ # multiple discovery MRK keyrings, one for each region.
188+ # Each component keyring has its own KMS client in a particular region.
189+ # When we provide a client supplier to the multi-keyring, all component
190+ # keyrings will use that client supplier configuration.
191+ # In our tests, we make `key_arn` an MRK with a replica, and
192+ # provide only the replica region in our discovery filter.
193+ discovery_filter = DiscoveryFilter (
194+ partition = "aws" ,
195+ account_ids = account_ids
196+ )
197+
198+ mrk_discovery_client_supplier_input = CreateAwsKmsMrkDiscoveryMultiKeyringInput (
199+ client_supplier = RegionalRoleClientSupplier (),
200+ discovery_filter = discovery_filter ,
201+ regions = regions
202+ )
203+
204+ mrk_discovery_client_supplier_keyring = mat_prov .create_aws_kms_mrk_discovery_multi_keyring (
205+ input = mrk_discovery_client_supplier_input
206+ )
207+
208+ # 9. Create a new config and client using the discovery keyring.
209+ # This is the same setup as above, except we provide the discovery keyring to the config.
210+ replica_key_table_config = DynamoDbTableEncryptionConfig (
211+ logical_table_name = ddb_table_name ,
212+ partition_key_name = "partition_key" ,
213+ sort_key_name = "sort_key" ,
214+ attribute_actions_on_encrypt = attribute_actions_on_encrypt ,
215+ # Provide discovery keyring here
216+ keyring = mrk_discovery_client_supplier_keyring ,
217+ allowed_unsigned_attribute_prefix = unsign_attr_prefix
218+ )
219+
220+ replica_key_tables_config = {ddb_table_name : replica_key_table_config }
221+ replica_key_tables_encryption_config = DynamoDbTablesEncryptionConfig (
222+ table_encryption_configs = replica_key_tables_config
223+ )
224+
225+ replica_key_encrypted_client = EncryptedClient (
226+ client = ddb_client ,
227+ encryption_config = replica_key_tables_encryption_config
228+ )
229+
230+ # 10. Get the item back from our table using the discovery keyring client.
231+ # The client will decrypt the item client-side using the keyring,
232+ # and return the original item.
233+ # The discovery keyring will only use KMS keys in the provided regions and
234+ # AWS accounts. Since we have provided it with a custom client supplier
235+ # which uses different IAM roles based on the key region,
236+ # the discovery keyring will use a particular IAM role to decrypt
237+ # based on the region of the KMS key it uses to decrypt.
238+ replica_key_key_to_get = {
239+ "partition_key" : {"S" : "awsKmsMrkMultiKeyringItem" },
240+ "sort_key" : {"N" : "0" }
241+ }
242+
243+ replica_key_get_response = replica_key_encrypted_client .get_item (
244+ TableName = ddb_table_name ,
245+ Key = replica_key_key_to_get
246+ )
247+
248+ # Demonstrate that GetItem succeeded and returned the decrypted item
249+ assert replica_key_get_response ['ResponseMetadata' ]['HTTPStatusCode' ] == 200
250+ replica_key_returned_item = replica_key_get_response ['Item' ]
251+ assert replica_key_returned_item ["sensitive_data" ]["S" ] == "encrypt and sign me!"
0 commit comments