Example for using multiple keys in multiple regions

examples/src/multiple_kms_cmk_regions.py

@@ -0,0 +1,72 @@
+# 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 basic encryption and decryption of a value already in memory
+using multiple KMS CMKs in multiple regions.
+"""
+import aws_encryption_sdk
+from aws_encryption_sdk.key_providers.kms import KMSMasterKey, KMSMasterKeyProvider
+from aws_encryption_sdk.internal.crypto.encryption import encrypt, decrypt
+
+
+def multiple_kms_cmk_regions(key_arn_1, key_arn_2, source_plaintext, botocore_session=None):
+    """Encrypts and then decrypts a string under multiple KMS customer master keys (CMKs) in multiple regions.
+
+    :param str key_arn_1: Amazon Resource Name (ARN) of the KMS CMK
+    :param str key_arn_2: Amazon Resource Name (ARN) of another KMS CMK
+    :param bytes source_plaintext: Data to encrypt
+    :param botocore_session: existing botocore session instance
+    :type botocore_session: botocore.session.Session
+    """
+    # Check that these keys are in different regions
+    assert not key_arn_1.split(":")[3] == key_arn_2.split(":")[3]
+
+    kwargs = dict(key_ids=[key_arn_1, key_arn_2])
+
+    if botocore_session is not None:
+        kwargs["botocore_session"] = botocore_session
+
+    # Create master key provider using the ARNs of the keys and the session (botocore_session)
+    kms_key_provider = KMSMasterKeyProvider(**kwargs)
+
+    # Encrypt the plaintext using the AWS Encryption SDK. It returns the encrypted message and the header
+    ciphertext, encrypted_message_header = encrypt(kms_key_provider, source_plaintext)
+
+    # Check that both key ARNs are in the message headers
+    assert len(encrypted_message_header.encrypted_data_keys) == 2
+
+    # Decrypt the encrypted message using the AWS Encryption SDK. It returns the decrypted message and the header
+    # Either of our keys can be used to decrypt the message
+    plaintext_1, decrypted_message_header_1 = decrypt(
+        KMSMasterKey(key_id=key_arn_1), ciphertext
+    )
+    plaintext_2, decrypted_message_header_2 = decrypt(
+        KMSMasterKey(key_id=key_arn_2), ciphertext
+    )
+
+    # Check that the original message and the decrypted message are the same
+    if not isinstance(source_plaintext, bytes):
+        plaintext1 = plaintext_1.decode("utf-8")
+        plaintext2 = plaintext_2.decode("utf-8")
+    assert source_plaintext == plaintext_1
+    assert source_plaintext == plaintext_2
+
+    # Check that 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_1.encryption_context.items()
+    )
+    assert all(
+        pair in encrypted_message_header.encryption_context.items()
+        for pair in decrypted_message_header_2.encryption_context.items()
+    )

examples/test/test_i_multiple_kms_cmk_regions.py

@@ -0,0 +1,33 @@
+# 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.
+"""Unit test suite for the encryption and decryption using multiple KMS CMKs in multiple regions example."""
+
+import botocore.session
+import pytest
+
+from ..src.multiple_kms_cmk_regions import multiple_kms_cmk_regions
+
+from .examples_test_utils import get_cmk_arn
+from .examples_test_utils import static_plaintext
+
+
+pytestmark = [pytest.mark.examples]
+
+
+def test_multiple_kms_cmk_regions():
+    plaintext = static_plaintext
+    cmk_arn_1 = get_cmk_arn("us-west-2")
+    cmk_arn_2 = get_cmk_arn("eu-central-1")
+    multiple_kms_cmk_regions(
+        cmk_arn_1, cmk_arn_2, source_plaintext=plaintext, botocore_session=botocore.session.Session()
+    )

test/integration/README.rst

@@ -5,7 +5,8 @@ aws-encryption-sdk Integration Tests
 In order to run these integration tests successfully, these things must be configured.
 
 #. Ensure that AWS credentials are available in one of the `automatically discoverable credential locations`_.
-#. Set environment variable ``AWS_ENCRYPTION_SDK_PYTHON_INTEGRATION_TEST_AWS_KMS_KEY_ID`` to valid
+#. Set environment variable ``AWS_ENCRYPTION_SDK_PYTHON_INTEGRATION_TEST_AWS_KMS_KEY_ID``
+   and ``AWS_ENCRYPTION_SDK_PYTHON_INTEGRATION_TEST_AWS_KMS_KEY_ID_2`` to valid
    `AWS KMS key id`_ to use for integration tests.
 
 .. _automatically discoverable credential locations: http://boto3.readthedocs.io/en/latest/guide/configuration.html

test/integration/integration_test_utils.py

@@ -16,6 +16,7 @@
 from aws_encryption_sdk.key_providers.kms import KMSMasterKeyProvider
 
 AWS_KMS_KEY_ID = "AWS_ENCRYPTION_SDK_PYTHON_INTEGRATION_TEST_AWS_KMS_KEY_ID"
+AWS_KMS_KEY_ID_2 = "AWS_ENCRYPTION_SDK_PYTHON_INTEGRATION_TEST_AWS_KMS_KEY_ID_2"
 _KMS_MKP = None
 
 
@@ -32,6 +33,26 @@ def get_cmk_arn():
         return arn
     raise ValueError("KMS CMK ARN provided for integration tests much be a key not an alias")
 
+def get_cmk_arn(region_name):
+    """Retrieves a CMK ARN based on the requested region_name"""
+    if AWS_KMS_KEY_ID in os.environ and AWS_KMS_KEY_ID_2 in os.environ:
+        raise ValueError(
+            'Environment variable "{}" or "{}" must be set to a valid KMS CMK ARN for integration tests to run'.format(
+                AWS_KMS_KEY_ID, AWS_KMS_KEY_ID_2
+            )
+        )
+    arn_1 = os.environ.get(AWS_KMS_KEY_ID, None)
+    arn_2 = os.environ.get(AWS_KMS_KEY_ID_2, None)
+    if arn_1.split(':')[3] == region_name:
+        return arn_1
+    elif arn_2.split(':')[3] == region_name:
+        return arn_2
+    else:
+        raise ValueError(
+            'No CMK in the region {} exist in either of your environment variables "{}" or "{}"'.format(
+                region_name, AWS_KMS_KEY_ID, AWS_KMS_KEY_ID_2
+            )
+        )
 
 def setup_kms_master_key_provider(cache=True):
     """Reads the test_values config file and builds the requested KMS Master Key Provider."""