fix user agent modification to actually work and add test that actually checks for it

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
 )
@@ -101,6 +101,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 +121,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.
@@ -200,10 +201,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,14 @@
 # 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 .integration_test_utils import setup_kms_master_key_provider, get_cmk_arn
 
 pytestmark = [pytest.mark.integ]
 
@@ -40,6 +41,16 @@
 }
 
 
+def test_encrypt_verify_user_agent(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
+
+
 class TestKMSThickClientIntegration(unittest.TestCase):
 
     def setUp(self):

test/unit/test_providers_kms_master_key.py

@@ -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):