Merge pull request #47 from mattsb42-aws/user-agent

Partial fix (and actually test for) user agent modification

Author: mattsb42-aws

src/aws_encryption_sdk/identifiers.py

@@ -22,7 +22,7 @@
 from aws_encryption_sdk.exceptions import InvalidAlgorithmError
 
 __version__ = '1.3.3'
-USER_AGENT_SUFFIX = 'AwsEncryptionSdkPython-KMSMasterKey/{}'.format(__version__)
+USER_AGENT_SUFFIX = 'AwsEncryptionSdkPython/{}'.format(__version__)
 
 
 class EncryptionSuite(Enum):

src/aws_encryption_sdk/internal/utils/__init__.py

@@ -165,19 +165,3 @@ def source_data_key_length_check(source_data_key, algorithm):
             actual=len(source_data_key.data_key),
             required=algorithm.kdf_input_len
         ))
-
-
-def extend_user_agent_suffix(user_agent, suffix):
-    """Adds a suffix to the provided user agent.
-
-    :param str user_agent: Existing user agent (None == not yet defined)
-    :param str suffix: Desired suffix to add to user agent
-    :returns: User agent with suffix
-    :rtype: str
-    """
-    if user_agent is None:
-        user_agent = ''
-    else:
-        user_agent += ' '
-    user_agent += suffix
-    return user_agent

src/aws_encryption_sdk/key_providers/kms.py

@@ -16,13 +16,13 @@
 import attr
 import boto3
 import botocore.client
+import botocore.config
 from botocore.exceptions import ClientError
 import botocore.session
 
 from aws_encryption_sdk.exceptions import DecryptKeyError, EncryptKeyError, GenerateKeyError, UnknownRegionError
 from aws_encryption_sdk.identifiers import USER_AGENT_SUFFIX
 from aws_encryption_sdk.internal.str_ops import to_str
-from aws_encryption_sdk.internal.utils import extend_user_agent_suffix
 from aws_encryption_sdk.key_providers.base import (
     MasterKey, MasterKeyConfig, MasterKeyProvider, MasterKeyProviderConfig
 )
@@ -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.
@@ -101,6 +123,7 @@ def __init__(self, **kwargs):  # pylint: disable=unused-argument
 
     def _process_config(self):
         """Traverses the config and adds master keys and regional clients as needed."""
+        self._user_agent_adding_config = botocore.config.Config(user_agent_extra=USER_AGENT_SUFFIX)
         if self.config.key_ids:
             self.add_master_keys_from_list(self.config.key_ids)
         if self.config.region_names:
@@ -120,7 +143,7 @@ def add_regional_client(self, region_name):
             self._regional_clients[region_name] = boto3.session.Session(
                 region_name=region_name,
                 botocore_session=self.config.botocore_session
-            ).client('kms')
+            ).client('kms', config=self._user_agent_adding_config)
 
     def add_regional_clients_from_list(self, region_names):
         """Adds multiple regional clients for the specified regions if they do not already exist.
@@ -135,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]
 
@@ -174,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.
@@ -200,10 +228,6 @@ class KMSMasterKey(MasterKey):
     def __init__(self, **kwargs):  # pylint: disable=unused-argument
         """Performs transformations needed for KMS."""
         self._key_id = to_str(self.key_id)  # KMS client requires str, not bytes
-        self.config.client.meta.config.user_agent_extra = extend_user_agent_suffix(
-            user_agent=self.config.client.meta.config.user_agent_extra,
-            suffix=USER_AGENT_SUFFIX
-        )
 
     def _generate_data_key(self, algorithm, encryption_context=None):
         """Generates data key and returns plaintext and ciphertext of key.

test/integration/test_i_aws_encrytion_sdk_client.py

@@ -12,13 +12,15 @@
 # language governing permissions and limitations under the License.
 """Integration test suite for `aws_encryption_sdk`."""
 import io
+import logging
 import unittest
 
 import pytest
 
 import aws_encryption_sdk
-from aws_encryption_sdk.identifiers import Algorithm
-from .integration_test_utils import setup_kms_master_key_provider
+from aws_encryption_sdk.identifiers import Algorithm, USER_AGENT_SUFFIX
+from aws_encryption_sdk.key_providers.kms import KMSMasterKey
+from .integration_test_utils import get_cmk_arn, setup_kms_master_key_provider
 
 pytestmark = [pytest.mark.integ]
 
@@ -40,6 +42,25 @@
 }
 
 
+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())
+
+    mk.generate_data_key(algorithm=Algorithm.AES_256_GCM_IV12_TAG16_HKDF_SHA384_ECDSA_P384, encryption_context={})
+
+    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):

test/unit/test_providers_kms_master_key.py

@@ -20,7 +20,7 @@
 import six
 
 from aws_encryption_sdk.exceptions import DecryptKeyError, EncryptKeyError, GenerateKeyError
-from aws_encryption_sdk.identifiers import Algorithm, USER_AGENT_SUFFIX
+from aws_encryption_sdk.identifiers import Algorithm
 from aws_encryption_sdk.key_providers.base import MasterKey
 from aws_encryption_sdk.key_providers.kms import KMSMasterKey, KMSMasterKeyConfig
 from aws_encryption_sdk.structures import DataKey, EncryptedDataKey, MasterKeyInfo
@@ -96,16 +96,10 @@ def test_config_grant_tokens(self):
         )
         assert test.grant_tokens is self.mock_grant_tokens
 
-    @patch('aws_encryption_sdk.key_providers.kms.extend_user_agent_suffix')
-    def test_init(self, patch_extend_user_agent_suffix):
+    def test_init(self):
         self.mock_client.meta.config.user_agent_extra = sentinel.user_agent_extra
         test = KMSMasterKey(config=self.mock_kms_mkc_1)
         assert test._key_id == VALUES['arn'].decode('utf-8')
-        patch_extend_user_agent_suffix.assert_called_once_with(
-            user_agent=sentinel.user_agent_extra,
-            suffix=USER_AGENT_SUFFIX
-        )
-        assert self.mock_client.meta.config.user_agent_extra == patch_extend_user_agent_suffix.return_value
 
     def test_generate_data_key(self):
         test = KMSMasterKey(config=self.mock_kms_mkc_3)

test/unit/test_providers_kms_master_key_provider.py

@@ -99,7 +99,7 @@ def test_add_regional_client_new(self):
             region_name='ex_region_name',
             botocore_session=ANY
         )
-        self.mock_boto3_session_instance.client.assert_called_once_with('kms')
+        self.mock_boto3_session_instance.client.assert_called_once_with('kms', config=test._user_agent_adding_config)
         assert test._regional_clients['ex_region_name'] is self.mock_boto3_client_instance
 
     def test_add_regional_client_exists(self):

test/unit/test_utils.py

@@ -27,15 +27,6 @@
 pytestmark = [pytest.mark.unit, pytest.mark.local]
 
 
-@pytest.mark.parametrize('user_agent, suffix, output', (
-    (None, 'test_suffix', 'test_suffix'),
-    ('test_existing_suffix', 'test_suffix', 'test_existing_suffix test_suffix')
-))
-def test_extend_user_agent_suffix(user_agent, suffix, output):
-    test = aws_encryption_sdk.internal.utils.extend_user_agent_suffix(user_agent, suffix)
-    assert test == output
-
-
 class TestUtils(unittest.TestCase):
 
     def setUp(self):