From eb1b64148aa59f3999403632ac5efecee350bec1 Mon Sep 17 00:00:00 2001 From: mattsb42-aws Date: Thu, 5 Apr 2018 18:46:49 -0700 Subject: [PATCH 1/3] fix user agent modification to actually work and add test that actually checks for it --- src/aws_encryption_sdk/identifiers.py | 2 +- .../internal/utils/__init__.py | 16 ---------------- src/aws_encryption_sdk/key_providers/kms.py | 9 +++------ .../test_i_aws_encrytion_sdk_client.py | 15 +++++++++++++-- test/unit/test_providers_kms_master_key.py | 8 +------- .../test_providers_kms_master_key_provider.py | 2 +- test/unit/test_utils.py | 9 --------- 7 files changed, 19 insertions(+), 42 deletions(-) diff --git a/src/aws_encryption_sdk/identifiers.py b/src/aws_encryption_sdk/identifiers.py index 0db2e668e..224610f0c 100644 --- a/src/aws_encryption_sdk/identifiers.py +++ b/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): diff --git a/src/aws_encryption_sdk/internal/utils/__init__.py b/src/aws_encryption_sdk/internal/utils/__init__.py index af2a2b43a..3f32d8fc0 100644 --- a/src/aws_encryption_sdk/internal/utils/__init__.py +++ b/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 diff --git a/src/aws_encryption_sdk/key_providers/kms.py b/src/aws_encryption_sdk/key_providers/kms.py index c3621f290..58db5651f 100644 --- a/src/aws_encryption_sdk/key_providers/kms.py +++ b/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. diff --git a/test/integration/test_i_aws_encrytion_sdk_client.py b/test/integration/test_i_aws_encrytion_sdk_client.py index 454ac1f89..56dcf49bb 100644 --- a/test/integration/test_i_aws_encrytion_sdk_client.py +++ b/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): diff --git a/test/unit/test_providers_kms_master_key.py b/test/unit/test_providers_kms_master_key.py index 3eaad7f4e..8225f2aff 100644 --- a/test/unit/test_providers_kms_master_key.py +++ b/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) diff --git a/test/unit/test_providers_kms_master_key_provider.py b/test/unit/test_providers_kms_master_key_provider.py index f68a10041..1569c9e4a 100644 --- a/test/unit/test_providers_kms_master_key_provider.py +++ b/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): diff --git a/test/unit/test_utils.py b/test/unit/test_utils.py index 1c55f7820..3b28df315 100644 --- a/test/unit/test_utils.py +++ b/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): From b820b8f179019ec265d2fa0a00ed3943547ec4d0 Mon Sep 17 00:00:00 2001 From: mattsb42-aws Date: Fri, 6 Apr 2018 12:23:39 -0700 Subject: [PATCH 2/3] 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 --- src/aws_encryption_sdk/key_providers/kms.py | 49 ++++++++++++++----- .../test_i_aws_encrytion_sdk_client.py | 12 ++++- 2 files changed, 49 insertions(+), 12 deletions(-) diff --git a/src/aws_encryption_sdk/key_providers/kms.py b/src/aws_encryption_sdk/key_providers/kms.py index 58db5651f..2e745ebc7 100644 --- a/src/aws_encryption_sdk/key_providers/kms.py +++ b/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,7 +188,10 @@ 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), @@ -183,6 +199,17 @@ class KMSMasterKeyConfig(MasterKeyConfig): 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. diff --git a/test/integration/test_i_aws_encrytion_sdk_client.py b/test/integration/test_i_aws_encrytion_sdk_client.py index 56dcf49bb..a63bb2f6f 100644 --- a/test/integration/test_i_aws_encrytion_sdk_client.py +++ b/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): From 6ad0ec0d2002ed783bf037b51ddef7d00ff8095c Mon Sep 17 00:00:00 2001 From: mattsb42-aws Date: Fri, 6 Apr 2018 12:37:02 -0700 Subject: [PATCH 3/3] linting fixes --- test/integration/test_i_aws_encrytion_sdk_client.py | 2 +- test/unit/test_providers_kms_master_key.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/test/integration/test_i_aws_encrytion_sdk_client.py b/test/integration/test_i_aws_encrytion_sdk_client.py index a63bb2f6f..715e8d712 100644 --- a/test/integration/test_i_aws_encrytion_sdk_client.py +++ b/test/integration/test_i_aws_encrytion_sdk_client.py @@ -20,7 +20,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 +from .integration_test_utils import get_cmk_arn, setup_kms_master_key_provider pytestmark = [pytest.mark.integ] diff --git a/test/unit/test_providers_kms_master_key.py b/test/unit/test_providers_kms_master_key.py index 8225f2aff..cf607d95d 100644 --- a/test/unit/test_providers_kms_master_key.py +++ b/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