diff --git a/rain_api_core/egress_util.py b/rain_api_core/egress_util.py index c0cf542..acd8f85 100644 --- a/rain_api_core/egress_util.py +++ b/rain_api_core/egress_util.py @@ -23,7 +23,7 @@ def hmacsha256(key: bytes, string: str) -> hmac.HMAC: return hmac.new(key, string.encode(), sha256) -def get_presigned_url(session, bucket_name, object_name, region_name, expire_seconds, user_id, method='GET') -> str: +def get_presigned_url(session, bucket_name, object_name, region_name, expire_seconds, user_id, method='GET', api_request_uuid=None) -> str: timez = datetime.utcnow().strftime('%Y%m%dT%H%M%SZ') datez = timez[:8] region_id = "." + region_name if region_name != "us-east-1" else "" @@ -47,6 +47,13 @@ def get_presigned_url(session, bucket_name, object_name, region_name, expire_sec "X-Amz-SignedHeaders=host" ]) + # Add the api_request_uuid header only if it is available + if api_request_uuid is not None: + can_query_string = "&".join([ + f"A-api-request-uuid={api_request_uuid}", + can_query_string, + ]) + can_request = ( f"{method}\n" f"/{object_name}\n" diff --git a/tests/test_egress_util.py b/tests/test_egress_util.py index ec3bf98..db0bf9a 100644 --- a/tests/test_egress_util.py +++ b/tests/test_egress_util.py @@ -116,3 +116,27 @@ def test_get_presigned_url_with_newlines(mock_datetime): "&X-Amz-SignedHeaders=host" "&X-Amz-Signature=b5f7dab0ba1af2a1a5c625f526de6ab8047a20f165675890cfa8caf70cfed730" ) + +@mock.patch(f"{MODULE}.datetime", autospec=True) +def test_get_presigned_url_with_api_request_uuid(mock_datetime): + mock_datetime.utcnow.return_value = datetime(2024, 1, 1) + session = { + "Credentials": { + "AccessKeyId": "access_key_id", + "SecretAccessKey": "secret_access_key", + "SessionToken": "session_token" + } + } + presigned_url = get_presigned_url(session, "bucket_name", "object_name", "region_name", 500, "user_id", api_request_uuid="uuid_value") + assert presigned_url == ( + "https://bucket_name.s3.region_name.amazonaws.com/object_name" + "?A-api-request-uuid=uuid_value" + "&A-userid=user_id" + "&X-Amz-Algorithm=AWS4-HMAC-SHA256" + "&X-Amz-Credential=access_key_id%2F20240101%2Fregion_name%2Fs3%2Faws4_request" + "&X-Amz-Date=20240101T000000Z" + "&X-Amz-Expires=500" + "&X-Amz-Security-Token=session_token" + "&X-Amz-SignedHeaders=host" + "&X-Amz-Signature=3010924b3725f8b47a3486ea4005f658128ddc5dde5ca405dfe1bd69bd2ae7a8" + ) \ No newline at end of file