diff --git a/src/s3_encryption/__init__.py b/src/s3_encryption/__init__.py index ad07fe43..3f188510 100644 --- a/src/s3_encryption/__init__.py +++ b/src/s3_encryption/__init__.py @@ -9,7 +9,7 @@ from botocore.exceptions import ClientError from botocore.response import StreamingBody -from ._utils import safe_get_dict +from ._utils import _USER_AGENT_SUFFIX, append_user_agent, safe_get_dict from .exceptions import S3EncryptionClientError from .instruction_file import parse_instruction_file from .instruction_file_config import InstructionFileConfig @@ -350,6 +350,8 @@ def __attrs_post_init__(self): # Expose plugin context on wrapped client for instruction file fetching self.wrapped_s3_client._s3ec_plugin_context = self._plugin._context + append_user_agent(self.wrapped_s3_client, _USER_AGENT_SUFFIX) + # Register event handlers using boto3's event system event_system = self.wrapped_s3_client.meta.events event_system.register("before-call.s3.PutObject", self._plugin.on_put_object_before_call) diff --git a/src/s3_encryption/_utils.py b/src/s3_encryption/_utils.py index 4997b973..7cab9e27 100644 --- a/src/s3_encryption/_utils.py +++ b/src/s3_encryption/_utils.py @@ -2,6 +2,11 @@ # SPDX-License-Identifier: Apache-2.0 """Internal utility helpers for the S3 Encryption Client.""" +import importlib.metadata + +_PACKAGE_VERSION = importlib.metadata.version("amazon-s3-encryption-client-python") +_USER_AGENT_SUFFIX = f"S3ECPy/{_PACKAGE_VERSION}" + def safe_get_dict(source: dict, key: str) -> dict: """Get a dict value from *source*, defaulting to {} if missing or None. @@ -10,3 +15,10 @@ def safe_get_dict(source: dict, key: str) -> dict: when the key exists but its value is explicitly None. """ return source.get(key, {}) or {} + + +def append_user_agent(client, suffix: str): + """Append a suffix to the User-Agent header of a boto3 client.""" + existing = client.meta.config.user_agent_extra or "" + sep = " " if existing else "" + client.meta.config.user_agent_extra = f"{existing}{sep}{suffix}" diff --git a/src/s3_encryption/materials/kms_keyring.py b/src/s3_encryption/materials/kms_keyring.py index abd6fad4..6550a38b 100644 --- a/src/s3_encryption/materials/kms_keyring.py +++ b/src/s3_encryption/materials/kms_keyring.py @@ -53,6 +53,11 @@ class KmsKeyring(S3Keyring): ##% The KmsV1 mode MUST be only enabled when legacy wrapping algorithms are enabled. enable_legacy_wrapping_algorithms: bool = field(default=False) + def __attrs_post_init__(self): # noqa: D105 + from .._utils import _USER_AGENT_SUFFIX, append_user_agent + + append_user_agent(self.kms_client, _USER_AGENT_SUFFIX) + def on_encrypt(self, enc_materials): """Process encryption materials using KMS. diff --git a/test/test_user_agent.py b/test/test_user_agent.py new file mode 100644 index 00000000..ad7f6a30 --- /dev/null +++ b/test/test_user_agent.py @@ -0,0 +1,43 @@ +# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for user agent string injection.""" + +import boto3 + +from s3_encryption import S3EncryptionClient, S3EncryptionClientConfig +from s3_encryption._utils import _PACKAGE_VERSION, _USER_AGENT_SUFFIX +from s3_encryption.materials.kms_keyring import KmsKeyring + + +class TestUserAgent: + def test_user_agent_suffix_format(self): + assert f"S3ECPy/{_PACKAGE_VERSION}" == _USER_AGENT_SUFFIX + + def test_s3_client_gets_user_agent(self): + s3 = boto3.client("s3", region_name="us-east-1") + kms = boto3.client("kms", region_name="us-east-1") + keyring = KmsKeyring(kms, "arn:aws:kms:us-east-1:000000000000:key/fake") + config = S3EncryptionClientConfig(keyring=keyring) + + S3EncryptionClient(s3, config) + + assert _USER_AGENT_SUFFIX in s3.meta.config.user_agent_extra + + def test_kms_client_gets_user_agent(self): + kms = boto3.client("kms", region_name="us-east-1") + KmsKeyring(kms, "arn:aws:kms:us-east-1:000000000000:key/fake") + + assert _USER_AGENT_SUFFIX in kms.meta.config.user_agent_extra + + def test_existing_user_agent_extra_preserved(self): + s3 = boto3.client("s3", region_name="us-east-1") + s3.meta.config.user_agent_extra = "existing-agent/1.0" + + kms = boto3.client("kms", region_name="us-east-1") + keyring = KmsKeyring(kms, "arn:aws:kms:us-east-1:000000000000:key/fake") + config = S3EncryptionClientConfig(keyring=keyring) + + S3EncryptionClient(s3, config) + + assert "existing-agent/1.0" in s3.meta.config.user_agent_extra + assert _USER_AGENT_SUFFIX in s3.meta.config.user_agent_extra