From 918f502ebfa90abba1dd70ff9e34a111e7f14092 Mon Sep 17 00:00:00 2001
From: Gekko0114 <41197469+Gekko0114@users.noreply.github.com>
Date: Tue, 9 Dec 2025 23:41:36 +0900
Subject: [PATCH] Support both plain text and JSON format for Secrets Manager
 API key

- Modify get_api_key() to first treat Secrets Manager value as plain text
- If the value looks like JSON (starts with '{' and ends with '}'), parse it and look for common keys (DD_API_KEY, DATADOG_API_KEY)
- If JSON parsing fails or no common key is found, fall back to plain text
- Add test cases for both plain text and JSON formats
---
 datadog_lambda/api.py | 19 ++++++++++-
 tests/test_api.py     | 78 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 96 insertions(+), 1 deletion(-)

diff --git a/datadog_lambda/api.py b/datadog_lambda/api.py
index 88f2cbe2..c41f0da3 100644
--- a/datadog_lambda/api.py
+++ b/datadog_lambda/api.py
@@ -1,3 +1,4 @@
+import json
 import logging
 import os
 
@@ -95,9 +96,25 @@ def get_api_key() -> str:
         secrets_manager_client = _boto3_client(
             "secretsmanager", endpoint_url=endpoint_url, region_name=secrets_region
         )
-        api_key = secrets_manager_client.get_secret_value(
+        secret_string = secrets_manager_client.get_secret_value(
             SecretId=DD_API_KEY_SECRET_ARN
         )["SecretString"]
+        
+        # First treat as plain text
+        api_key = secret_string
+        
+        # If it looks like JSON, try parsing it
+        if secret_string and secret_string.strip().startswith("{") and secret_string.strip().endswith("}"):
+            try:
+                secret_dict = json.loads(secret_string)
+                # Try to find common key names
+                for key in ["DD_API_KEY", "DATADOG_API_KEY"]:
+                    if key in secret_dict:
+                        api_key = secret_dict[key]
+                        break
+            except (json.JSONDecodeError, ValueError, TypeError):
+                # If JSON parsing fails, keep using plain text
+                pass
     elif DD_API_KEY_SSM_NAME:
         # SSM endpoints: https://docs.aws.amazon.com/general/latest/gr/ssm.html
         fips_endpoint = None
diff --git a/tests/test_api.py b/tests/test_api.py
index 3b42cfd7..602d9a23 100644
--- a/tests/test_api.py
+++ b/tests/test_api.py
@@ -183,3 +183,81 @@ def test_no_fips_for_standard_regions(self, mock_boto3_client):
         mock_boto3_client.assert_called_with(
             "secretsmanager", endpoint_url=None, region_name="us-west-2"
         )
+
+    @patch("botocore.session.Session.create_client")
+    def test_secrets_manager_plain_text(self, mock_boto3_client):
+        """Test Secrets Manager value as plain text"""
+        mock_client = MagicMock()
+        mock_client.get_secret_value.return_value = {"SecretString": "plain-text-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-east-1:1234567890:secret:key-name-123ABC"
+        )
+
+        api_key = api.get_api_key()
+
+        self.assertEqual(api_key, "plain-text-api-key")
+
+    @patch("botocore.session.Session.create_client")
+    def test_secrets_manager_json_with_api_key(self, mock_boto3_client):
+        """Test Secrets Manager value with api_key key in JSON format"""
+        import json
+
+        mock_client = MagicMock()
+        mock_client.get_secret_value.return_value = {
+            "SecretString": json.dumps({"api_key": "json-api-key-value"})
+        }
+        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-east-1:1234567890:secret:key-name-123ABC"
+        )
+
+        api_key = api.get_api_key()
+
+        self.assertEqual(api_key, "json-api-key-value")
+
+    @patch("botocore.session.Session.create_client")
+    def test_secrets_manager_json_with_dd_api_key(self, mock_boto3_client):
+        """Test Secrets Manager value with DD_API_KEY key in JSON format"""
+        import json
+
+        mock_client = MagicMock()
+        mock_client.get_secret_value.return_value = {
+            "SecretString": json.dumps({"DD_API_KEY": "dd-api-key-value"})
+        }
+        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-east-1:1234567890:secret:key-name-123ABC"
+        )
+
+        api_key = api.get_api_key()
+
+        self.assertEqual(api_key, "dd-api-key-value")
+
+    @patch("botocore.session.Session.create_client")
+    def test_secrets_manager_json_with_custom_key(self, mock_boto3_client):
+        """Test Secrets Manager value with custom key in JSON format (treated as plain text)"""
+        import json
+
+        mock_client = MagicMock()
+        secret_json = json.dumps({"custom_key": "custom-api-key-value"})
+        mock_client.get_secret_value.return_value = {
+            "SecretString": secret_json
+        }
+        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-east-1:1234567890:secret:key-name-123ABC"
+        )
+
+        api_key = api.get_api_key()
+
+        # When no common key is found, it should be treated as plain text
+        self.assertEqual(api_key, secret_json)