feat: new DD_LAMBDA_FIPS_MODE control for secrets and metrics

Author: apiarian-datadog

datadog_lambda/api.py

@@ -1,15 +1,18 @@
-import os
 import logging
+import os
+
+from datadog_lambda.fips import enable_fips_mode
 
 logger = logging.getLogger(__name__)
 KMS_ENCRYPTION_CONTEXT_KEY = "LambdaFunctionName"
 api_key = None
 
 
 def decrypt_kms_api_key(kms_client, ciphertext):
-    from botocore.exceptions import ClientError
     import base64
 
+    from botocore.exceptions import ClientError
+
     """
     Decodes and deciphers the base64-encoded ciphertext given as a parameter using KMS.
     For this to work properly, the Lambda function must have the appropriate IAM permissions.
@@ -63,10 +66,9 @@ def get_api_key() -> str:
     DD_API_KEY = os.environ.get("DD_API_KEY", os.environ.get("DATADOG_API_KEY", ""))
 
     LAMBDA_REGION = os.environ.get("AWS_REGION", "")
-    is_gov_region = LAMBDA_REGION.startswith("us-gov-")
-    if is_gov_region:
+    if enable_fips_mode:
         logger.debug(
-            "Govcloud region detected. Using FIPs endpoints for secrets management."
+            "FIPS mode is enabled, using FIPS endpoints for secrets management."
         )
 
     if DD_API_KEY_SECRET_ARN:
@@ -80,7 +82,7 @@ def get_api_key() -> str:
             return ""
         endpoint_url = (
             f"https://secretsmanager-fips.{secrets_region}.amazonaws.com"
-            if is_gov_region
+            if enable_fips_mode
             else None
         )
         secrets_manager_client = _boto3_client(
@@ -92,7 +94,9 @@ def get_api_key() -> str:
     elif DD_API_KEY_SSM_NAME:
         # SSM endpoints: https://docs.aws.amazon.com/general/latest/gr/ssm.html
         fips_endpoint = (
-            f"https://ssm-fips.{LAMBDA_REGION}.amazonaws.com" if is_gov_region else None
+            f"https://ssm-fips.{LAMBDA_REGION}.amazonaws.com"
+            if enable_fips_mode
+            else None
         )
         ssm_client = _boto3_client("ssm", endpoint_url=fips_endpoint)
         api_key = ssm_client.get_parameter(
@@ -101,7 +105,9 @@ 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.{LAMBDA_REGION}.amazonaws.com" if is_gov_region else None
+            f"https://kms-fips.{LAMBDA_REGION}.amazonaws.com"
+            if enable_fips_mode
+            else None
         )
         kms_client = _boto3_client("kms", endpoint_url=fips_endpoint)
         api_key = decrypt_kms_api_key(kms_client, DD_KMS_API_KEY)

datadog_lambda/fips.py

@@ -0,0 +1,19 @@
+import logging
+import os
+
+is_gov_region = os.environ.get("AWS_REGION", "").startswith("us-gov-")
+
+enable_fips_mode = (
+    os.environ.get(
+        "DD_LAMBDA_FIPS_MODE",
+        "true" if is_gov_region else "false",
+    ).lower()
+    == "true"
+)
+
+if is_gov_region or enable_fips_mode:
+    logger = logging.getLogger(__name__)
+    logger.debug(
+        "Python Lambda Layer FIPS mode is %s.",
+        "enabled" if enable_fips_mode else "not enabled",
+    )

datadog_lambda/metric.py

@@ -12,6 +12,7 @@
 import ujson as json
 
 from datadog_lambda.extension import should_use_extension
+from datadog_lambda.fips import enable_fips_mode
 from datadog_lambda.tags import dd_lambda_layer_tag, get_enhanced_metrics_tags
 
 logger = logging.getLogger(__name__)
@@ -21,17 +22,27 @@ class MetricsHandler(enum.Enum):
     EXTENSION = "extension"
     FORWARDER = "forwarder"
     DATADOG_API = "datadog_api"
+    NO_METRICS = "no_metrics"
 
 
 def _select_metrics_handler():
     if should_use_extension:
         return MetricsHandler.EXTENSION
     if os.environ.get("DD_FLUSH_TO_LOG", "").lower() == "true":
         return MetricsHandler.FORWARDER
+
+    if enable_fips_mode:
+        logger.debug(
+            "With FIPS mode enabled, the Datadog API metrics handler is unavailable."
+        )
+        return MetricsHandler.NO_METRICS
+
     return MetricsHandler.DATADOG_API
 
 
 metrics_handler = _select_metrics_handler()
+# TODO: add a metric for this so that we can see how often the DATADOG_API
+# metrics handler is actually used.
 logger.debug("identified primary metrics handler as %s", metrics_handler)
 
 
@@ -122,6 +133,12 @@ def lambda_metric(metric_name, value, timestamp=None, tags=None, force_async=Fal
     elif metrics_handler == MetricsHandler.DATADOG_API:
         lambda_stats.distribution(metric_name, value, tags=tags, timestamp=timestamp)
 
+    elif metrics_handler == MetricsHandler.NO_METRICS:
+        logger.debug(
+            "Metric %s cannot be submitted because the metrics handler is disabled: %s",
+            metric_name,
+        ),
+
     else:
         # This should be qutie impossible, but let's at least log a message if
         # it somehow happens.

tests/test_api.py

@@ -1,6 +1,6 @@
 import os
 import unittest
-from unittest.mock import patch, MagicMock
+from unittest.mock import MagicMock, patch
 
 import datadog_lambda.api as api
 
@@ -22,6 +22,7 @@ def setUp(self):
         )
         self.env_patcher.start()
 
+    @patch("datadog_lambda.api.enable_fips_mode", True)
     @patch("botocore.session.Session.create_client")
     def test_secrets_manager_fips_endpoint(self, mock_boto3_client):
         mock_client = MagicMock()
@@ -62,6 +63,28 @@ def test_secrets_manager_different_region(self, mock_boto3_client):
         )
         self.assertEqual(api_key, "test-api-key")
 
+    @patch("datadog_lambda.api.enable_fips_mode", True)
+    @patch("botocore.session.Session.create_client")
+    def test_secrets_manager_different_region_but_still_fips(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="https://secretsmanager-fips.us-west-1.amazonaws.com",
+            region_name="us-west-1",
+        )
+        self.assertEqual(api_key, "test-api-key")
+
+    @patch("datadog_lambda.api.enable_fips_mode", True)
     @patch("botocore.session.Session.create_client")
     def test_ssm_fips_endpoint(self, mock_boto3_client):
         mock_client = MagicMock()
@@ -80,6 +103,7 @@ def test_ssm_fips_endpoint(self, mock_boto3_client):
         )
         self.assertEqual(api_key, "test-api-key")
 
+    @patch("datadog_lambda.api.enable_fips_mode", True)
     @patch("botocore.session.Session.create_client")
     @patch("datadog_lambda.api.decrypt_kms_api_key")
     def test_kms_fips_endpoint(self, mock_decrypt_kms, mock_boto3_client):

tests/test_metric.py

@@ -19,9 +19,15 @@
 
 class TestLambdaMetric(unittest.TestCase):
     def setUp(self):
-        patcher = patch("datadog_lambda.metric.lambda_stats")
-        self.mock_metric_lambda_stats = patcher.start()
-        self.addCleanup(patcher.stop)
+        lambda_stats_patcher = patch("datadog_lambda.metric.lambda_stats")
+        self.mock_metric_lambda_stats = lambda_stats_patcher.start()
+        self.addCleanup(lambda_stats_patcher.stop)
+
+        stdout_metric_patcher = patch(
+            "datadog_lambda.metric.write_metric_point_to_stdout"
+        )
+        self.mock_write_metric_point_to_stdout = stdout_metric_patcher.start()
+        self.addCleanup(stdout_metric_patcher.stop)
 
     def test_lambda_metric_tagged_with_dd_lambda_layer(self):
         lambda_metric("test", 1)
@@ -56,13 +62,26 @@ def test_select_metrics_handler_dd_api_fallback(self):
         self.assertEqual(MetricsHandler.DATADOG_API, _select_metrics_handler())
         del os.environ["DD_FLUSH_TO_LOG"]
 
+    @patch("datadog_lambda.metric.enable_fips_mode", True)
+    @patch("datadog_lambda.metric.should_use_extension", False)
+    def test_select_metrics_handler_has_no_fallback_in_fips_mode(self):
+        os.environ["DD_FLUSH_TO_LOG"] = "False"
+        self.assertEqual(MetricsHandler.NO_METRICS, _select_metrics_handler())
+        del os.environ["DD_FLUSH_TO_LOG"]
+
     @patch("datadog_lambda.metric.metrics_handler", MetricsHandler.EXTENSION)
-    def test_lambda_metric_flush_to_log_with_extension(self):
+    def test_lambda_metric_goes_to_extension_with_extension_handler(self):
         lambda_metric("test", 1)
         self.mock_metric_lambda_stats.distribution.assert_has_calls(
             [call("test", 1, timestamp=None, tags=[dd_lambda_layer_tag])]
         )
 
+    @patch("datadog_lambda.metric.metrics_handler", MetricsHandler.NO_METRICS)
+    def test_lambda_metric_has_nowhere_to_go_with_no_metrics_handler(self):
+        lambda_metric("test", 1)
+        self.mock_metric_lambda_stats.distribution.assert_not_called()
+        self.mock_write_metric_point_to_stdout.assert_not_called()
+
     @patch("datadog_lambda.metric.metrics_handler", MetricsHandler.EXTENSION)
     def test_lambda_metric_timestamp_with_extension(self):
         delta = timedelta(minutes=1)
@@ -72,6 +91,7 @@ def test_lambda_metric_timestamp_with_extension(self):
         self.mock_metric_lambda_stats.distribution.assert_has_calls(
             [call("test_timestamp", 1, timestamp=timestamp, tags=[dd_lambda_layer_tag])]
         )
+        self.mock_write_metric_point_to_stdout.assert_not_called()
 
     @patch("datadog_lambda.metric.metrics_handler", MetricsHandler.EXTENSION)
     def test_lambda_metric_datetime_with_extension(self):
@@ -89,6 +109,7 @@ def test_lambda_metric_datetime_with_extension(self):
                 )
             ]
         )
+        self.mock_write_metric_point_to_stdout.assert_not_called()
 
     @patch("datadog_lambda.metric.metrics_handler", MetricsHandler.EXTENSION)
     def test_lambda_metric_invalid_timestamp_with_extension(self):
@@ -97,16 +118,21 @@ def test_lambda_metric_invalid_timestamp_with_extension(self):
 
         lambda_metric("test_timestamp", 1, timestamp)
         self.mock_metric_lambda_stats.distribution.assert_not_called()
+        self.mock_write_metric_point_to_stdout.assert_not_called()
 
     @patch("datadog_lambda.metric.metrics_handler", MetricsHandler.FORWARDER)
     def test_lambda_metric_flush_to_log(self):
         lambda_metric("test", 1)
         self.mock_metric_lambda_stats.distribution.assert_not_called()
+        self.mock_write_metric_point_to_stdout.assert_has_calls(
+            [call("test", 1, timestamp=None, tags=[dd_lambda_layer_tag])]
+        )
 
     @patch("datadog_lambda.metric.logger.warning")
     def test_lambda_metric_invalid_metric_name_none(self, mock_logger_warning):
         lambda_metric(None, 1)
         self.mock_metric_lambda_stats.distribution.assert_not_called()
+        self.mock_write_metric_point_to_stdout.assert_not_called()
         mock_logger_warning.assert_called_once_with(
             "Ignoring metric submission. Invalid metric name: %s", None
         )
@@ -115,6 +141,7 @@ def test_lambda_metric_invalid_metric_name_none(self, mock_logger_warning):
     def test_lambda_metric_invalid_metric_name_not_string(self, mock_logger_warning):
         lambda_metric(123, 1)
         self.mock_metric_lambda_stats.distribution.assert_not_called()
+        self.mock_write_metric_point_to_stdout.assert_not_called()
         mock_logger_warning.assert_called_once_with(
             "Ignoring metric submission. Invalid metric name: %s", 123
         )
@@ -123,6 +150,7 @@ def test_lambda_metric_invalid_metric_name_not_string(self, mock_logger_warning)
     def test_lambda_metric_non_numeric_value(self, mock_logger_warning):
         lambda_metric("test.non_numeric", "oops")
         self.mock_metric_lambda_stats.distribution.assert_not_called()
+        self.mock_write_metric_point_to_stdout.assert_not_called()
         mock_logger_warning.assert_called_once_with(
             "Ignoring metric submission for metric '%s' because the value is not numeric: %r",
             "test.non_numeric",