add default client generation for KMSMasterKey is none is provided: this both makes KMSMasterKey simpler to use directly and makes it more likely that the client will use our user agent

Author: mattsb42-aws

src/aws_encryption_sdk/key_providers/kms.py

@@ -33,6 +33,28 @@
 _PROVIDER_ID = 'aws-kms'
 
 
+def _region_from_key_id(key_id, default_region=None):
+    """Determine the target region from a key ID, falling back to a default region if provided.
+
+    :param str key_id: AWS KMS key ID
+    :param str default_region: Region to use if no region found in key_id
+    :returns: region name
+    :rtype: str
+    :raises UnknownRegionError: if no region found in key_id and no default_region provided
+    """
+    try:
+        region_name = key_id.split(':', 4)[3]
+        if default_region is None:
+            default_region = region_name
+    except IndexError:
+        if default_region is None:
+            raise UnknownRegionError(
+                'No default region found and no region determinable from key id: {}'.format(key_id)
+            )
+        region_name = default_region
+    return region_name
+
+
 @attr.s(hash=True)
 class KMSMasterKeyProviderConfig(MasterKeyProviderConfig):
     """Configuration object for KMSMasterKeyProvider objects.
@@ -136,16 +158,7 @@ def _client(self, key_id):
 
         :param str key_id: KMS CMK ID
         """
-        try:
-            region_name = key_id.split(':', 4)[3]
-            if self.default_region is None:
-                self.default_region = region_name
-        except IndexError:
-            if self.default_region is None:
-                raise UnknownRegionError(
-                    'No default region found and no region determinable from key id: {}'.format(key_id)
-                )
-            region_name = self.default_region
+        region_name = _region_from_key_id(key_id, self.default_region)
         self.add_regional_client(region_name)
         return self._regional_clients[region_name]
 
@@ -175,14 +188,28 @@ class KMSMasterKeyConfig(MasterKeyConfig):
     """
 
     provider_id = _PROVIDER_ID
-    client = attr.ib(hash=True, validator=attr.validators.instance_of(botocore.client.BaseClient))
+    client = attr.ib(
+        hash=True,
+        validator=attr.validators.instance_of(botocore.client.BaseClient)
+    )
     grant_tokens = attr.ib(
         hash=True,
         default=attr.Factory(tuple),
         validator=attr.validators.instance_of(tuple),
         converter=tuple
     )
 
+    @client.default
+    def client_default(self):
+        """Create a client if one was not provided."""
+        try:
+            region_name = _region_from_key_id(to_str(self.key_id))
+            kwargs = dict(region_name=region_name)
+        except UnknownRegionError:
+            kwargs = {}
+        botocore_config = botocore.config.Config(user_agent_extra=USER_AGENT_SUFFIX)
+        return boto3.session.Session(**kwargs).client('kms', config=botocore_config)
+
 
 class KMSMasterKey(MasterKey):
     """Master Key class for KMS CMKs.

test/integration/test_i_aws_encrytion_sdk_client.py

@@ -19,6 +19,7 @@
 
 import aws_encryption_sdk
 from aws_encryption_sdk.identifiers import Algorithm, USER_AGENT_SUFFIX
+from aws_encryption_sdk.key_providers.kms import KMSMasterKey
 from .integration_test_utils import setup_kms_master_key_provider, get_cmk_arn
 
 pytestmark = [pytest.mark.integ]
@@ -41,7 +42,7 @@
 }
 
 
-def test_encrypt_verify_user_agent(caplog):
+def test_encrypt_verify_user_agent_kms_master_key_provider(caplog):
     caplog.set_level(level=logging.DEBUG)
     mkp = setup_kms_master_key_provider()
     mk = mkp.master_key(get_cmk_arn())
@@ -51,6 +52,15 @@ def test_encrypt_verify_user_agent(caplog):
     assert USER_AGENT_SUFFIX in caplog.text
 
 
+def test_encrypt_verify_user_agent_kms_master_key(caplog):
+    caplog.set_level(level=logging.DEBUG)
+    mk = KMSMasterKey(key_id=get_cmk_arn())
+
+    mk.generate_data_key(algorithm=Algorithm.AES_256_GCM_IV12_TAG16_HKDF_SHA384_ECDSA_P384, encryption_context={})
+
+    assert USER_AGENT_SUFFIX in caplog.text
+
+
 class TestKMSThickClientIntegration(unittest.TestCase):
 
     def setUp(self):