From 9e6d728d5105705817cdf59fb8811ab897ca664b Mon Sep 17 00:00:00 2001
From: mattsb42-aws <bullocm@amazon.com>
Date: Wed, 11 Apr 2018 14:37:11 -0700
Subject: [PATCH] add custom extra user agent data to KMS materials provider
 client

---
 src/dynamodb_encryption_sdk/identifiers.py            |  1 +
 .../material_providers/aws_kms.py                     | 10 ++++++----
 test/integration/material_providers/test_aws_kms.py   | 11 +++++++++++
 test/unit/material_providers/test_aws_kms.py          |  5 ++++-
 4 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/src/dynamodb_encryption_sdk/identifiers.py b/src/dynamodb_encryption_sdk/identifiers.py
index ce1e14ee..f57b4c9c 100644
--- a/src/dynamodb_encryption_sdk/identifiers.py
+++ b/src/dynamodb_encryption_sdk/identifiers.py
@@ -17,6 +17,7 @@
 __version__ = '0.0.0'
 
 LOGGER_NAME = 'dynamodb_encryption_sdk'
+USER_AGENT_SUFFIX = 'DynamodbEncryptionSdkPython/{}'.format(__version__)
 
 
 class CryptoAction(Enum):
diff --git a/src/dynamodb_encryption_sdk/material_providers/aws_kms.py b/src/dynamodb_encryption_sdk/material_providers/aws_kms.py
index 8d5fd79e..b7d47dbc 100644
--- a/src/dynamodb_encryption_sdk/material_providers/aws_kms.py
+++ b/src/dynamodb_encryption_sdk/material_providers/aws_kms.py
@@ -19,8 +19,7 @@
 
 import attr
 import boto3
-import botocore.client
-import botocore.session
+import botocore
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import hashes
 from cryptography.hazmat.primitives.kdf.hkdf import HKDF
@@ -34,7 +33,7 @@
 
 from dynamodb_encryption_sdk.delegated_keys.jce import JceNameLocalDelegatedKey
 from dynamodb_encryption_sdk.exceptions import UnknownRegionError, UnwrappingError, WrappingError
-from dynamodb_encryption_sdk.identifiers import EncryptionKeyType, KeyEncodingType, LOGGER_NAME
+from dynamodb_encryption_sdk.identifiers import EncryptionKeyType, KeyEncodingType, LOGGER_NAME, USER_AGENT_SUFFIX
 from dynamodb_encryption_sdk.internal.identifiers import MaterialDescriptionKeys, TEXT_ENCODING
 from dynamodb_encryption_sdk.internal.str_ops import to_bytes, to_str
 from dynamodb_encryption_sdk.internal.validators import dictionary_validator, iterable_validator
@@ -163,6 +162,9 @@ class AwsKmsCryptographicMaterialsProvider(CryptographicMaterialsProvider):
     def __attrs_post_init__(self):
         # type: () -> None
         """Load the content and signing key info."""
+        self._user_agent_adding_config = botocore.config.Config(  # pylint: disable=attribute-defined-outside-init
+            user_agent_extra=USER_AGENT_SUFFIX
+        )
         self._content_key_info = KeyInfo.from_material_description(  # pylint: disable=attribute-defined-outside-init
             material_description=self._material_description,
             description_key=MaterialDescriptionKeys.CONTENT_ENCRYPTION_ALGORITHM.value,
@@ -187,7 +189,7 @@ def _add_regional_client(self, region_name):
             self._regional_clients[region_name] = boto3.session.Session(
                 region_name=region_name,
                 botocore_session=self._botocore_session
-            ).client('kms')
+            ).client('kms', config=self._user_agent_adding_config)
         return self._regional_clients[region_name]
 
     def _client(self, key_id):
diff --git a/test/integration/material_providers/test_aws_kms.py b/test/integration/material_providers/test_aws_kms.py
index 3f239ede..710f70b5 100644
--- a/test/integration/material_providers/test_aws_kms.py
+++ b/test/integration/material_providers/test_aws_kms.py
@@ -11,10 +11,13 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 """Integration tests for ``dynamodb_encryption_sdk.material_providers.aws_kms``."""
+import logging
+
 import hypothesis
 import pytest
 
 from dynamodb_encryption_sdk.encrypted import CryptoConfig
+from dynamodb_encryption_sdk.identifiers import USER_AGENT_SUFFIX
 from dynamodb_encryption_sdk.structures import EncryptionContext
 from ..integration_test_utils import aws_kms_cmp  # noqa pylint: disable=unused-import
 from ..integration_test_utils import functional_test_utils, hypothesis_strategies
@@ -27,6 +30,14 @@ def pytest_generate_tests(metafunc):
     functional_test_utils.set_parametrized_item(metafunc)
 
 
+def test_verify_user_agent(aws_kms_cmp, caplog):
+    caplog.set_level(level=logging.DEBUG)
+
+    aws_kms_cmp.encryption_materials(EncryptionContext())
+
+    assert USER_AGENT_SUFFIX in caplog.text
+
+
 def test_aws_kms_item_cycle(aws_kms_cmp, parametrized_actions, parametrized_item):
     crypto_config = CryptoConfig(
         materials_provider=aws_kms_cmp,
diff --git a/test/unit/material_providers/test_aws_kms.py b/test/unit/material_providers/test_aws_kms.py
index 16ce3f3a..d652bab6 100644
--- a/test/unit/material_providers/test_aws_kms.py
+++ b/test/unit/material_providers/test_aws_kms.py
@@ -269,7 +269,10 @@ def test_add_regional_client_unknown_region(default_kms_cmp, patch_boto3_session
         region_name=sentinel.region,
         botocore_session=default_kms_cmp._botocore_session
     )
-    patch_boto3_session.return_value.client.assert_called_once_with('kms')
+    patch_boto3_session.return_value.client.assert_called_once_with(
+        'kms',
+        config=default_kms_cmp._user_agent_adding_config
+    )
     assert default_kms_cmp._regional_clients[sentinel.region] is patch_boto3_session.return_value.client.return_value
     assert test is patch_boto3_session.return_value.client.return_value