From de0e562c2a6d57acba3105994f8b59589cf9222a Mon Sep 17 00:00:00 2001 From: seebees Date: Wed, 9 Mar 2022 15:25:23 -0800 Subject: [PATCH 1/3] doc: Add example for custom KMS client config Occasionally custom configuration for the AWS KMS client is needed. This is especially important in Discovery mode, But anytime you are dealing with more than one AWS KMS Key, the master key providers are simpler to work with then individual master keys. --- .../src/custom_kms_client_for_kms_provider.py | 74 +++++++++++++++++++ 1 file changed, 74 insertions(+) create mode 100644 examples/src/custom_kms_client_for_kms_provider.py diff --git a/examples/src/custom_kms_client_for_kms_provider.py b/examples/src/custom_kms_client_for_kms_provider.py new file mode 100644 index 000000000..20314ba8d --- /dev/null +++ b/examples/src/custom_kms_client_for_kms_provider.py @@ -0,0 +1,74 @@ +# Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Example showing how to customize the AWS KMS Client.""" +import aws_encryption_sdk +from aws_encryption_sdk import CommitmentPolicy +import boto3 + +# Create a new class that extends the AWS KMS Provider you need to use +class CustomKMSClientMasterKeyProvider(aws_encryption_sdk.StrictAwsKmsMasterKeyProvider): + + # Override `add_regional_client` to use whatever configuration you need + def add_regional_client(self, region_name): + """Adds a regional client for the specified region if it does not already exist. + :param str region_name: AWS Region ID (ex: us-east-1) + """ + if region_name not in self._regional_clients: + session = boto3.session.Session(botocore_session=self.config.botocore_session) + client = session.client( + 'kms', + region_name=region_name, + config=self._user_agent_adding_config, + # Add additional custom client configuration here + connect_timeout=10 + ) + self._register_client(client, region_name) + self._regional_clients[region_name] = client + +# This is just an example of using the above master key provider +def encrypt_decrypt(key_arn, source_plaintext, botocore_session=None): + """Encrypts and then decrypts a string under one KMS customer master key (CMK). + + :param str key_arn: Amazon Resource Name (ARN) of the KMS CMK + :param bytes source_plaintext: Data to encrypt + :param botocore_session: existing botocore session instance + :type botocore_session: botocore.session.Session + """ + kwargs = dict(key_ids=[key_arn]) + + if botocore_session is not None: + kwargs["botocore_session"] = botocore_session + + # Set up an encryption client with an explicit commitment policy. Note that if you do not explicitly choose a + # commitment policy, REQUIRE_ENCRYPT_REQUIRE_DECRYPT is used by default. + client = aws_encryption_sdk.EncryptionSDKClient(commitment_policy=CommitmentPolicy.REQUIRE_ENCRYPT_REQUIRE_DECRYPT) + + # Create the custom master key provider using the ARN of the key and the session (botocore_session) + kms_key_provider = CustomKMSClientMasterKeyProvider(**kwargs) + + # Encrypt the plaintext using the AWS Encryption SDK. It returns the encrypted message and the header. Note: in + # order for decrypt to succeed, the key_ids value must be the key ARN of the CMK. + ciphertext, encrypted_message_header = client.encrypt(source=source_plaintext, key_provider=kms_key_provider) + + # Decrypt the encrypted message using the AWS Encryption SDK. It returns the decrypted message and the header + plaintext, decrypted_message_header = client.decrypt(source=ciphertext, key_provider=kms_key_provider) + + # Check if the original message and the decrypted message are the same + assert source_plaintext == plaintext + + # Check if the headers of the encrypted message and decrypted message match + assert all( + pair in encrypted_message_header.encryption_context.items() + for pair in decrypted_message_header.encryption_context.items() + ) + From 45f7e637b42231696da9fd6f4805597e0edbbd12 Mon Sep 17 00:00:00 2001 From: J Plasmeier Date: Wed, 8 Feb 2023 10:54:34 -0800 Subject: [PATCH 2/3] fix static analysis issues --- examples/src/custom_kms_client_for_kms_provider.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/examples/src/custom_kms_client_for_kms_provider.py b/examples/src/custom_kms_client_for_kms_provider.py index 20314ba8d..51c55e879 100644 --- a/examples/src/custom_kms_client_for_kms_provider.py +++ b/examples/src/custom_kms_client_for_kms_provider.py @@ -11,12 +11,15 @@ # ANY KIND, either express or implied. See the License for the specific # language governing permissions and limitations under the License. """Example showing how to customize the AWS KMS Client.""" +import boto3 + import aws_encryption_sdk from aws_encryption_sdk import CommitmentPolicy -import boto3 + # Create a new class that extends the AWS KMS Provider you need to use class CustomKMSClientMasterKeyProvider(aws_encryption_sdk.StrictAwsKmsMasterKeyProvider): + """Custom region-specific client which extends the StrictAwsKmsMasterKeyProvider""" # Override `add_regional_client` to use whatever configuration you need def add_regional_client(self, region_name): @@ -35,6 +38,7 @@ def add_regional_client(self, region_name): self._register_client(client, region_name) self._regional_clients[region_name] = client + # This is just an example of using the above master key provider def encrypt_decrypt(key_arn, source_plaintext, botocore_session=None): """Encrypts and then decrypts a string under one KMS customer master key (CMK). @@ -71,4 +75,3 @@ def encrypt_decrypt(key_arn, source_plaintext, botocore_session=None): pair in encrypted_message_header.encryption_context.items() for pair in decrypted_message_header.encryption_context.items() ) - From 217d71fee3f179179ec7f8d918a2c2d0f34d2af6 Mon Sep 17 00:00:00 2001 From: J Plasmeier Date: Wed, 8 Feb 2023 16:35:35 -0800 Subject: [PATCH 3/3] fix config parameter --- examples/src/custom_kms_client_for_kms_provider.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/src/custom_kms_client_for_kms_provider.py b/examples/src/custom_kms_client_for_kms_provider.py index 51c55e879..d7658ffaa 100644 --- a/examples/src/custom_kms_client_for_kms_provider.py +++ b/examples/src/custom_kms_client_for_kms_provider.py @@ -12,6 +12,7 @@ # language governing permissions and limitations under the License. """Example showing how to customize the AWS KMS Client.""" import boto3 +from botocore.config import Config import aws_encryption_sdk from aws_encryption_sdk import CommitmentPolicy @@ -31,9 +32,8 @@ def add_regional_client(self, region_name): client = session.client( 'kms', region_name=region_name, - config=self._user_agent_adding_config, # Add additional custom client configuration here - connect_timeout=10 + config=Config(connection_timeout=10).merge(self._user_agent_adding_config) ) self._register_client(client, region_name) self._regional_clients[region_name] = client