Use correct SecretsManager region (#576)

Author: nhulston

Diff for: datadog_lambda/api.py

@@ -64,30 +64,37 @@ 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-")
+    LAMBDA_REGION = os.environ.get("AWS_REGION", "")
+    is_gov_region = LAMBDA_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:
         # Secrets manager endpoints: https://docs.aws.amazon.com/general/latest/gr/asm.html
-        fips_endpoint = (
-            f"https://secretsmanager-fips.{REGION}.amazonaws.com"
+        try:
+            secrets_region = DD_API_KEY_SECRET_ARN.split(":")[3]
+        except Exception:
+            logger.debug(
+                "Invalid secret arn in DD_API_KEY_SECRET_ARN. Unable to get API key."
+            )
+            return ""
+        endpoint_url = (
+            f"https://secretsmanager-fips.{secrets_region}.amazonaws.com"
             if is_gov_region
             else None
         )
         secrets_manager_client = boto3.client(
-            "secretsmanager", endpoint_url=fips_endpoint
+            "secretsmanager", endpoint_url=endpoint_url, region_name=secrets_region
         )
         api_key = secrets_manager_client.get_secret_value(
             SecretId=DD_API_KEY_SECRET_ARN
         )["SecretString"]
     elif DD_API_KEY_SSM_NAME:
         # 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
+            f"https://ssm-fips.{LAMBDA_REGION}.amazonaws.com" if is_gov_region else None
         )
         ssm_client = boto3.client("ssm", endpoint_url=fips_endpoint)
         api_key = ssm_client.get_parameter(
@@ -96,7 +103,7 @@ def get_api_key() -> str:
     elif DD_KMS_API_KEY:
         # 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
+            f"https://kms-fips.{LAMBDA_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)

Diff for: tests/test_api.py

@@ -29,13 +29,36 @@ def test_secrets_manager_fips_endpoint(self, mock_boto3_client):
         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"
+        os.environ[
+            "DD_API_KEY_SECRET_ARN"
+        ] = "arn:aws:secretsmanager:us-gov-east-1:1234567890:secret:key-name-123ABC"
 
         api_key = api.get_api_key()
 
         mock_boto3_client.assert_called_with(
             "secretsmanager",
             endpoint_url="https://secretsmanager-fips.us-gov-east-1.amazonaws.com",
+            region_name="us-gov-east-1",
+        )
+        self.assertEqual(api_key, "test-api-key")
+
+    @patch("boto3.client")
+    def test_secrets_manager_different_region(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-east-1"
+        os.environ[
+            "DD_API_KEY_SECRET_ARN"
+        ] = "arn:aws:secretsmanager:us-west-1:1234567890:secret:key-name-123ABC"
+
+        api_key = api.get_api_key()
+
+        mock_boto3_client.assert_called_with(
+            "secretsmanager",
+            endpoint_url=None,
+            region_name="us-west-1",
         )
         self.assertEqual(api_key, "test-api-key")
 
@@ -82,8 +105,12 @@ def test_no_fips_for_standard_regions(self, mock_boto3_client):
 
         os.environ.clear()
         os.environ["AWS_REGION"] = "us-west-2"
-        os.environ["DD_API_KEY_SECRET_ARN"] = "test-arn"
+        os.environ[
+            "DD_API_KEY_SECRET_ARN"
+        ] = "arn:aws:secretsmanager:us-west-2:1234567890:secret:key-name-123ABC"
 
         api.get_api_key()
 
-        mock_boto3_client.assert_called_with("secretsmanager", endpoint_url=None)
+        mock_boto3_client.assert_called_with(
+            "secretsmanager", endpoint_url=None, region_name="us-west-2"
+        )