|
| 1 | +# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved. |
| 2 | +# SPDX-License-Identifier: Apache-2.0 |
| 3 | +""" |
| 4 | +By default, the KMS keyring uses a client supplier that |
| 5 | +supplies a client with the same configuration for every region. |
| 6 | +If you need different behavior, you can write your own client supplier. |
| 7 | +
|
| 8 | +You might use this |
| 9 | +if you need different credentials in different AWS regions. |
| 10 | +This might be because you are crossing partitions (ex: ``aws`` and ``aws-cn``) |
| 11 | +or if you are working with regions that have separate authentication silos |
| 12 | +like ``ap-east-1`` and ``me-south-1``. |
| 13 | +
|
| 14 | +This example shows how to create a client supplier |
| 15 | +that will supply KMS clients with valid credentials for the target region |
| 16 | +even when working with regions that need different credentials. |
| 17 | +
|
| 18 | +https://docs.aws.amazon.com/encryption-sdk/latest/developer-guide/choose-keyring.html#use-kms-keyring |
| 19 | +
|
| 20 | +For an example of how to use the KMS keyring with CMKs in multiple regions, |
| 21 | +see the ``keyring/aws_kms/multiple_regions`` example. |
| 22 | +
|
| 23 | +For another example of how to use the KMS keyring with a custom client configuration, |
| 24 | +see the ``keyring/aws_kms/custom_kms_client_config`` example. |
| 25 | +
|
| 26 | +For examples of how to use the KMS keyring in discovery mode on decrypt, |
| 27 | +see the ``keyring/aws_kms/discovery_decrypt``, |
| 28 | +``keyring/aws_kms/discovery_decrypt_in_region_only``, |
| 29 | +and ``keyring/aws_kms/discovery_decrypt_with_preferred_region`` examples. |
| 30 | +""" |
| 31 | +from botocore.client import BaseClient |
| 32 | +from botocore.session import Session |
| 33 | + |
| 34 | +import aws_encryption_sdk |
| 35 | +from aws_encryption_sdk.keyrings.aws_kms import KmsKeyring |
| 36 | +from aws_encryption_sdk.keyrings.aws_kms.client_suppliers import ClientSupplier, DefaultClientSupplier |
| 37 | + |
| 38 | +try: # Python 3.5.0 and 3.5.1 have incompatible typing modules |
| 39 | + from typing import Union # noqa pylint: disable=unused-import |
| 40 | +except ImportError: # pragma: no cover |
| 41 | + # We only need these imports when running the mypy checks |
| 42 | + pass |
| 43 | + |
| 44 | + |
| 45 | +class MultiPartitionClientSupplier(ClientSupplier): |
| 46 | + """Client supplier that supplies clients across AWS partitions and identity silos.""" |
| 47 | + |
| 48 | + def __init__(self): |
| 49 | + """Set up default client suppliers for identity silos.""" |
| 50 | + self._china_supplier = DefaultClientSupplier(botocore_session=Session(profile="china")) |
| 51 | + self._middle_east_supplier = DefaultClientSupplier(botocore_session=Session(profile="middle-east")) |
| 52 | + self._hong_kong_supplier = DefaultClientSupplier(botocore_session=Session(profile="hong-kong")) |
| 53 | + self._default_supplier = DefaultClientSupplier() |
| 54 | + |
| 55 | + def __call__(self, region_name): |
| 56 | + # type: (Union[None, str]) -> BaseClient |
| 57 | + """Return a client for the requested region. |
| 58 | +
|
| 59 | + :rtype: BaseClient |
| 60 | + """ |
| 61 | + if region_name.startswith("cn-"): |
| 62 | + return self._china_supplier(region_name) |
| 63 | + |
| 64 | + if region_name.startswith("me-"): |
| 65 | + return self._middle_east_supplier(region_name) |
| 66 | + |
| 67 | + if region_name == "ap-east-1": |
| 68 | + return self._hong_kong_supplier(region_name) |
| 69 | + |
| 70 | + return self._default_supplier(region_name) |
| 71 | + |
| 72 | + |
| 73 | +def run(aws_kms_cmk, source_plaintext): |
| 74 | + # type: (str, bytes) -> None |
| 75 | + """Demonstrate an encrypt/decrypt cycle using a KMS keyring with a custom client supplier. |
| 76 | +
|
| 77 | + :param str aws_kms_cmk: The ARN of an AWS KMS CMK that protects data keys |
| 78 | + :param bytes source_plaintext: Plaintext to encrypt |
| 79 | + """ |
| 80 | + # Prepare your encryption context. |
| 81 | + # https://docs.aws.amazon.com/encryption-sdk/latest/developer-guide/concepts.html#encryption-context |
| 82 | + encryption_context = { |
| 83 | + "encryption": "context", |
| 84 | + "is not": "secret", |
| 85 | + "but adds": "useful metadata", |
| 86 | + "that can help you": "be confident that", |
| 87 | + "the data you are handling": "is what you think it is", |
| 88 | + } |
| 89 | + |
| 90 | + # Create the keyring that determines how your data keys are protected. |
| 91 | + keyring = KmsKeyring(generator_key_id=aws_kms_cmk, client_supplier=MultiPartitionClientSupplier()) |
| 92 | + |
| 93 | + # Encrypt your plaintext data. |
| 94 | + ciphertext, _encrypt_header = aws_encryption_sdk.encrypt( |
| 95 | + source=source_plaintext, encryption_context=encryption_context, keyring=keyring |
| 96 | + ) |
| 97 | + |
| 98 | + # Demonstrate that the ciphertext and plaintext are different. |
| 99 | + assert ciphertext != source_plaintext |
| 100 | + |
| 101 | + # Decrypt your encrypted data using the same keyring you used on encrypt. |
| 102 | + # |
| 103 | + # You do not need to specify the encryption context on decrypt |
| 104 | + # because the header of the encrypted message includes the encryption context. |
| 105 | + decrypted, decrypt_header = aws_encryption_sdk.decrypt(source=ciphertext, keyring=keyring) |
| 106 | + |
| 107 | + # Demonstrate that the decrypted plaintext is identical to the original plaintext. |
| 108 | + assert decrypted == source_plaintext |
| 109 | + |
| 110 | + # Verify that the encryption context used in the decrypt operation includes |
| 111 | + # the encryption context that you specified when encrypting. |
| 112 | + # The AWS Encryption SDK can add pairs, so don't require an exact match. |
| 113 | + # |
| 114 | + # In production, always use a meaningful encryption context. |
| 115 | + assert set(encryption_context.items()) <= set(decrypt_header.encryption_context.items()) |
0 commit comments