Use FIPs endpoints in Govcloud regions (#575)

Author: nhulston

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

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)