Skip to content
31 changes: 28 additions & 3 deletions datadog_lambda/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
89 changes: 89 additions & 0 deletions tests/test_api.py
Original file line number Diff line number Diff line change
@@ -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)
Loading