Added functionality to fetch license key from AWS Secrets Manager and SSM Parameter Store with caching

README.md

@@ -16,6 +16,13 @@ Additional notes:
 
 * Some users in UTF-8 environments have reported difficulty with defining strings of `NR_TAGS` delimited by the semicolon `;` character. If this applies to you, you can set an alternative delimiter character as the value of `NR_ENV_DELIMITER`, and separate your `NR_TAGS` with that.
 * Custom Lambda and VPC log groups can be set using the `NR_LAMBDA_LOG_GROUP_PREFIX` and `NR_VPC_LOG_GROUP_PREFIX` environment variables.
+* The New Relic License Key can now be fetched from multiple secure sources:
+
+- **Environment Variables**
+- **AWS Systems Manager (SSM) Parameter Store**
+- **AWS Secrets Manager**
+
+A caching mechanism has also been added to store the New Relic License Key. This prevents fetching the key from AWS Secrets Manager or SSM Parameter Store on each Lambda execution, improving performance.
 
 ## Manual Deployment
 
@@ -61,6 +68,33 @@ module "newrelic_log_ingestion" {
 
 By default, this will build and pack the lambda zip inside of the Terraform Module. You can supply your own by switching `build_lambda = false`, and specify the path to your lambda, using `lambda_archive = "{{LAMBDA_PATH}}"`, replacing `{{LAMBDA_PATH}}` with the path to your lambda.
 
+### Configuration Guide for New Relic License Key Retrieval
+
+To configure the retrieval of the New Relic License Key from different sources. Follow the steps below to set up your preferred method:
+
+1. **Select the License Key Source** : Use the `nr_license_key_source` variable to specify where the license key should be fetched from. You have three options:
+   * `environment_var`: Fetches the license key from an environment variable. This is the default setting.
+   * `ssm`: Retrieves the license key from AWS Systems Manager (SSM) Parameter Store.
+   * `secrets_manager`: Obtains the license key from AWS Secrets Manager.
+2. **Enable License Key Caching** : To improve performance by caching the New Relic License Key, set the `enable_caching_for_license_key` variable to `true`.
+3. **Specify the Key Location** :
+   * If you chose `ssm` or `secrets_manager` as your source in step 1, you need to provide additional information to locate the license key:
+     * For `secrets_manager`: Supply the secret's name or Amazon Resource Name (ARN) to the `nr_license_key` variable.
+     * For `ssm`: Provide the name of the parameter in the Parameter Store that holds the license key to the `nr_license_key` variable
+
+- Set the `nr_license_key_source` variable to choose the source of the license key. The available options are `environment_var`, `ssm`, or `secrets_manager`. The default value is `environment_var`.
+- Set the `enable_caching_for_license_key` variable to `true` to enable caching for the New Relic License Key.
+- When setting `nr_license_key_source` value to either `ssm` or `secrets_manager`, secret name or secret ARN can be passed to `nr_license_key` variable to fetch the key from Secrets Manager. Parameter store, the parameter name where the license key is stored can be provided by `nr_license_key` variable.
+
+```terraform
+module "newrelic_log_ingestion" {
+source                         = "github.com/newrelic/aws-log-ingestion"
+nr_license_key                 = "{{ARN_OF_LICENSE_KEY_SECRET_OF_SECRETS_MANAGER}}"
+nr_license_key_source          = "secrets_manager"
+enable_caching_for_license_key = true
+}
+```
+
 ## Infra Payload Format
 
 The maximum payload size in bytes is:

src/function.py

@@ -42,6 +42,9 @@
 import re
 import time
 
+import boto3
+from botocore.exceptions import ClientError
+
 from base64 import b64decode
 from enum import Enum
 from urllib import request
@@ -121,6 +124,9 @@ class EntryType(Enum):
 LOGGING_LAMBDA_VERSION = "2.9.3"
 LOGGING_PLUGIN_METADATA = {"type": "lambda", "version": LOGGING_LAMBDA_VERSION}
 
+# Global cache for storing new relic license keys
+LICENSE_KEY_CACHE = None
+
 
 class MaxRetriesException(Exception):
     pass
@@ -306,15 +312,107 @@ def _generate_payloads(data, split_function):
     )
 
 
+def _get_license_key_source():
+    """
+    This function returns the source of the license key.
+    LICENSE_KEY_SRC must be one of 'environment_var', 'ssm', or 'secrets_manager'.
+    Defaults to 'environment_var'.
+    """
+    return os.getenv("LICENSE_KEY_SRC", "environment_var")
+
+
 def _get_license_key(license_key=None):
     """
     This functions gets New Relic's license key from env vars.
     """
     if license_key:
         return license_key
+
+    license_key_source = _get_license_key_source()
+
+    if license_key_source == "ssm":
+        return _get_license_key_from_ssm(os.getenv("LICENSE_KEY", ""))
+    elif license_key_source == "secrets_manager":
+        return _get_license_key_from_secrets_manager(os.getenv("LICENSE_KEY", ""))
+
     return os.getenv("LICENSE_KEY", "")
 
 
+def _get_license_key_from_secrets_manager(secret_arn):
+    """
+    Fetches the secret value for the given secret ARN from AWS Secrets Manager.
+    """
+    global LICENSE_KEY_CACHE
+
+    if not secret_arn:
+        return ""
+
+    enable_caching = os.getenv("ENABLE_CACHING", "false").lower() == "true"
+
+    # Check cache first if caching is enabled
+    if enable_caching and LICENSE_KEY_CACHE is not None:
+        logger.info(
+            "Using cached secret instead of fetching the license key from secrets manager"
+        )
+        return LICENSE_KEY_CACHE
+
+    client = boto3.client("secretsmanager")
+
+    try:
+        get_secret_value_response = client.get_secret_value(SecretId=secret_arn)
+        logger.info("Successfully retrieved license key from Secrets Manager")
+    except ClientError as e:
+        logger.error(f"Unable to retrieve secret {secret_arn}: {e}")
+        return ""
+
+    if "SecretString" in get_secret_value_response:
+        secret = get_secret_value_response["SecretString"]
+    else:
+        secret = b64decode(get_secret_value_response["SecretBinary"])
+
+    # Cache the secret before returning if caching is enabled
+    if enable_caching:
+        LICENSE_KEY_CACHE = secret
+
+    return "" if not secret else secret
+
+
+def _get_license_key_from_ssm(parameter_path):
+    """
+    Fetches the parameter value for the given parameter path
+    from AWS Systems Manager Parameter Store.
+    """
+    global LICENSE_KEY_CACHE
+
+    if not parameter_path:
+        return ""
+
+    enable_caching = os.getenv("ENABLE_CACHING", "false").lower() == "true"
+
+    # Check cache first if caching is enabled
+    if enable_caching and LICENSE_KEY_CACHE is not None:
+        logger.info(
+            "Using cached parameter instead of fetching the license key from SSM"
+        )
+        return LICENSE_KEY_CACHE
+
+    client = boto3.client("ssm")
+
+    try:
+        response = client.get_parameter(Name=parameter_path, WithDecryption=True)
+        logger.info("Successfully retrieved license key from SSM")
+        parameter_value = response["Parameter"]["Value"]
+    except ClientError as e:
+        logger.error(f"Unable to retrieve parameter {parameter_path}: {e}")
+        raise e
+
+    # Cache the parameter before returning if caching is enabled
+    if enable_caching:
+        LICENSE_KEY_CACHE = parameter_value
+
+    return parameter_value
+
+
 def _get_newrelic_tags(payload):
     """
     This functions gets New Relic's tags from env vars and adds it to the payload

terraform.tf

@@ -14,6 +14,22 @@ variable "nr_license_key" {
   sensitive   = true
 }
 
+variable "nr_license_key_source" {
+  type        = string
+  description = "The source of the NewRelic license key. Must be one of 'environment_var', 'ssm', or 'secrets_manager'."
+  default     = "environment_var"
+  validation {
+    condition     = contains(["environment_var", "ssm", "secrets_manager"], var.nr_license_key_source)
+    error_message = "The nr_license_key_source must be one of 'environment_var', 'ssm', or 'secrets_manager'."
+  }
+}
+
+variable "enable_caching_for_license_key" {
+  description = "Enable caching for the license key."
+  type        = bool
+  default     = false
+}
+
 variable "nr_logging_enabled" {
   type        = bool
   description = "Determines if logs are forwarded to New Relic Logging"
@@ -115,6 +131,40 @@ data "aws_iam_policy_document" "lambda_assume_policy" {
   }
 }
 
+resource "aws_iam_policy" "lambda_fetch_license_key_policy" {
+  name        = "lambda_fetch_license_key_policy"
+  description = "Policy for Lambda to fetch NewRelic license key from SSM Parameter or Secrets Manager"
+  policy = jsonencode({
+    Version = "2012-10-17"
+    Statement = [
+      {
+        Action = [
+          "ssm:GetParameter",
+          "ssm:GetParameters",
+          "ssm:GetParametersByPath",
+        ]
+        Effect   = "Allow"
+        Resource = "*"
+      },
+      {
+        Action = [
+          "secretsmanager:GetSecretValue",
+          "secretsmanager:DescribeSecret",
+        ]
+        Effect   = "Allow"
+        Resource = "*"
+      },
+      {
+        Action = [
+          "kms:Decrypt"
+        ]
+        Effect   = "Allow"
+        Resource = "*"
+      }
+    ]
+  })
+}
+
 resource "aws_iam_role" "lambda_role" {
   count = var.function_role == null ? 1 : 0
 
@@ -132,6 +182,14 @@ resource "aws_iam_role_policy_attachment" "lambda_log_policy" {
   policy_arn = "arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole"
 }
 
+resource "aws_iam_role_policy_attachment" "lambda_fetch_license_key_policy" {
+  count = var.function_role == null ? 1 : 0
+
+  role       = aws_iam_role.lambda_role.0.name
+  policy_arn = aws_iam_policy.lambda_fetch_license_key_policy.arn
+}
+
+
 resource "aws_cloudwatch_log_group" "lambda_logs" {
   name              = "/aws/lambda/${var.service_name}"
   retention_in_days = var.lambda_log_retention_in_days
@@ -191,9 +249,11 @@ resource "aws_lambda_function" "ingestion_function" {
   environment {
     variables = {
       LICENSE_KEY     = var.nr_license_key
+      LICENSE_KEY_SRC = var.nr_license_key_source
       LOGGING_ENABLED = var.nr_logging_enabled ? "True" : "False"
       INFRA_ENABLED   = var.nr_infra_logging ? "True" : "False"
       NR_TAGS         = var.nr_tags
+      ENABLE_CACHING = var.enable_caching_for_license_key ? "True" : "False"
     }
   }