diff --git a/datadog_lambda/api.py b/datadog_lambda/api.py index a114fe8f..03135912 100644 --- a/datadog_lambda/api.py +++ b/datadog_lambda/api.py @@ -64,16 +64,41 @@ def get_api_key() -> str: DD_KMS_API_KEY = os.environ.get("DD_KMS_API_KEY", "") DD_API_KEY = os.environ.get("DD_API_KEY", os.environ.get("DATADOG_API_KEY", "")) + REGION = os.environ.get("AWS_REGION", "") + is_gov_region = REGION.startswith("us-gov-") + if is_gov_region: + logger.debug( + "Govcloud region detected. Using FIPs endpoints for secrets management." + ) + if DD_API_KEY_SECRET_ARN: - api_key = boto3.client("secretsmanager").get_secret_value( + # Secrets manager endpoints: https://docs.aws.amazon.com/general/latest/gr/asm.html + fips_endpoint = ( + f"https://secretsmanager-fips.{REGION}.amazonaws.com" + if is_gov_region + else None + ) + secrets_manager_client = boto3.client( + "secretsmanager", endpoint_url=fips_endpoint + ) + api_key = secrets_manager_client.get_secret_value( SecretId=DD_API_KEY_SECRET_ARN )["SecretString"] elif DD_API_KEY_SSM_NAME: - api_key = boto3.client("ssm").get_parameter( + # SSM endpoints: https://docs.aws.amazon.com/general/latest/gr/ssm.html + fips_endpoint = ( + f"https://ssm-fips.{REGION}.amazonaws.com" if is_gov_region else None + ) + ssm_client = boto3.client("ssm", endpoint_url=fips_endpoint) + api_key = ssm_client.get_parameter( Name=DD_API_KEY_SSM_NAME, WithDecryption=True )["Parameter"]["Value"] elif DD_KMS_API_KEY: - kms_client = boto3.client("kms") + # KMS endpoints: https://docs.aws.amazon.com/general/latest/gr/kms.html + fips_endpoint = ( + f"https://kms-fips.{REGION}.amazonaws.com" if is_gov_region else None + ) + kms_client = boto3.client("kms", endpoint_url=fips_endpoint) api_key = decrypt_kms_api_key(kms_client, DD_KMS_API_KEY) else: api_key = DD_API_KEY diff --git a/tests/test_api.py b/tests/test_api.py new file mode 100644 index 00000000..a69f4382 --- /dev/null +++ b/tests/test_api.py @@ -0,0 +1,89 @@ +import os +import unittest +from unittest.mock import patch, MagicMock + +import datadog_lambda.api as api + + +class TestDatadogLambdaAPI(unittest.TestCase): + def setUp(self): + api.api_key = None + self.env_patcher = patch.dict( + os.environ, + { + "DD_API_KEY_SECRET_ARN": "", + "DD_API_KEY_SSM_NAME": "", + "DD_KMS_API_KEY": "", + "DD_API_KEY": "", + "DATADOG_API_KEY": "", + "AWS_REGION": "", + }, + clear=True, + ) + self.env_patcher.start() + + @patch("boto3.client") + def test_secrets_manager_fips_endpoint(self, mock_boto3_client): + mock_client = MagicMock() + mock_client.get_secret_value.return_value = {"SecretString": "test-api-key"} + mock_boto3_client.return_value = mock_client + + os.environ["AWS_REGION"] = "us-gov-east-1" + os.environ["DD_API_KEY_SECRET_ARN"] = "test-secrets-arn" + + api_key = api.get_api_key() + + mock_boto3_client.assert_called_with( + "secretsmanager", + endpoint_url="https://secretsmanager-fips.us-gov-east-1.amazonaws.com", + ) + self.assertEqual(api_key, "test-api-key") + + @patch("boto3.client") + def test_ssm_fips_endpoint(self, mock_boto3_client): + mock_client = MagicMock() + mock_client.get_parameter.return_value = { + "Parameter": {"Value": "test-api-key"} + } + mock_boto3_client.return_value = mock_client + + os.environ["AWS_REGION"] = "us-gov-west-1" + os.environ["DD_API_KEY_SSM_NAME"] = "test-ssm-param" + + api_key = api.get_api_key() + + mock_boto3_client.assert_called_with( + "ssm", endpoint_url="https://ssm-fips.us-gov-west-1.amazonaws.com" + ) + self.assertEqual(api_key, "test-api-key") + + @patch("boto3.client") + @patch("datadog_lambda.api.decrypt_kms_api_key") + def test_kms_fips_endpoint(self, mock_decrypt_kms, mock_boto3_client): + mock_client = MagicMock() + mock_boto3_client.return_value = mock_client + mock_decrypt_kms.return_value = "test-api-key" + + os.environ["AWS_REGION"] = "us-gov-west-1" + os.environ["DD_KMS_API_KEY"] = "encrypted-api-key" + + api_key = api.get_api_key() + + mock_boto3_client.assert_called_with( + "kms", endpoint_url="https://kms-fips.us-gov-west-1.amazonaws.com" + ) + self.assertEqual(api_key, "test-api-key") + + @patch("boto3.client") + def test_no_fips_for_standard_regions(self, mock_boto3_client): + mock_client = MagicMock() + mock_client.get_secret_value.return_value = {"SecretString": "test-api-key"} + mock_boto3_client.return_value = mock_client + + os.environ.clear() + os.environ["AWS_REGION"] = "us-west-2" + os.environ["DD_API_KEY_SECRET_ARN"] = "test-arn" + + api.get_api_key() + + mock_boto3_client.assert_called_with("secretsmanager", endpoint_url=None)