Add the ability to read the license key from AWS secrets manager

.gitignore

@@ -20,3 +20,6 @@ packaged.yaml
 node_modules
 package.json
 package-lock.json
+
+.python-version
+build/

Pipfile

@@ -17,6 +17,7 @@ python_version = "3.7"
 
 [packages]
 aiohttp = "*"
+boto3 = "*"
 
 [pipenv]
 allow_prereleases = true

build-zip.sh

@@ -0,0 +1,43 @@
+#!/bin/bash
+
+# For those unable to use SAM, this builds a zip file that can be deployed using the AWS CLI or terraform or the like.
+# The resulting zip will be placed in the build directory. It uses a python3 virtual environment, per the
+# instructions here: https://docs.aws.amazon.com/lambda/latest/dg/python-package.html
+
+# A suffix may be added to the zip name by defining the SUFFIX environment variable before invoking this script.
+
+# Constants
+BUILDDIR=build
+
+# Extract the version from the template.yaml file
+VERSION=`grep SemanticVersion template.yaml | sed -e 's/ *SemanticVersion: *//'`
+ZIPFILE=newrelic-log-ingestion-${VERSION}${SUFFIX}.zip
+
+# Clean up any old build
+echo "Cleaning old builds..."
+rm -rf ${BUILDDIR}
+
+# Create the virtual environment, and activate it
+echo "Setting up virtual environment..."
+python3 -m venv ${BUILDDIR}
+source build/bin/activate
+
+# Install the requirements, excluding boto3, which is pre-installed in AWS
+echo "Installing dependencies..."
+cat src/requirements.txt | grep -v boto3 > /tmp/requirements.txt
+pip3 --disable-pip-version-check install --requirement /tmp/requirements.txt
+
+# Deactivate the virtual environment
+deactivate
+
+# Package up the dependencies
+echo "Creating zip of dependencies..."
+(cd ${BUILDDIR}/lib/python*/site-packages && zip --quiet --recurse-paths -9 ../../../${ZIPFILE} .)
+
+# Add the function to the zip file
+echo "Adding function to zip..."
+(cd src && zip --grow --quiet ../${BUILDDIR}/${ZIPFILE} function.py)
+
+# All done!
+echo "Done, zipfile: ${BUILDDIR}/${ZIPFILE}"
+

src/function.py

@@ -34,6 +34,7 @@
 https://docs.newrelic.com/
 """
 
+import boto3
 import datetime
 import gzip
 import json
@@ -107,6 +108,9 @@ class EntryType(Enum):
 LOGGING_LAMBDA_VERSION = "1.0.2"
 LOGGING_PLUGIN_METADATA = {"type": "lambda", "version": LOGGING_LAMBDA_VERSION}
 
+# Once the license key has been fetched, save it to avoid fetching again
+cached_license_key = None
+
 
 class MaxRetriesException(Exception):
     pass
@@ -258,11 +262,32 @@ def _generate_payloads(data, split_function):
 
 def _get_license_key(license_key=None):
     """
-    This functions gets New Relic's license key from env vars.
+    This functions gets New Relic's license key from env vars or AWS secrets manager.
     """
+    global cached_license_key
     if license_key:
         return license_key
-    return os.getenv("LICENSE_KEY", "")
+
+    if cached_license_key:
+        return cached_license_key
+
+    # If a secrets manager ARN was specified, use it to grab the license key, otherwise
+    # obtain it from the LICENSE_KEY environment variable.
+    secret_arn = os.getenv("SECRET_KEY_ARN", "")
+    if secret_arn:
+        secrets_client = boto3.client("secretsmanager")
+        secret = secrets_client.get_secret_value(SecretId=secret_arn)
+        cached_license_key = secret.get("SecretString")
+        if _debug_logging_enabled():
+            print(
+                "(got license key from secrets manager, len={})".format(
+                    len(cached_license_key)
+                )
+            )
+    else:
+        cached_license_key = os.getenv("LICENSE_KEY", "")
+
+    return cached_license_key
 
 
 def _debug_logging_enabled():

src/requirements.txt

@@ -1,9 +1,16 @@
 -i https://pypi.org/simple
-aiohttp==4.0.0a1
-async-timeout==3.0.1
-attrs==19.3.0
+aiohttp==3.6.2
+async-timeout==3.0.1; python_full_version >= '3.5.3'
+attrs==19.3.0; python_version >= '2.7' and python_version not in '3.0, 3.1, 3.2, 3.3'
+boto3==1.14.10
+botocore==1.17.10
 chardet==3.0.4
-idna==2.9
-multidict==4.7.5
-typing-extensions==3.7.4.1
-yarl==1.4.2
+docutils==0.15.2; python_version >= '2.6' and python_version not in '3.0, 3.1, 3.2, 3.3'
+idna==2.9; python_version >= '2.7' and python_version not in '3.0, 3.1, 3.2, 3.3'
+jmespath==0.10.0; python_version >= '2.6' and python_version not in '3.0, 3.1, 3.2, 3.3'
+multidict==4.7.6; python_version >= '3.5'
+python-dateutil==2.8.1; python_version >= '2.7' and python_version not in '3.0, 3.1, 3.2, 3.3'
+s3transfer==0.3.3
+six==1.15.0; python_version >= '2.7' and python_version not in '3.0, 3.1, 3.2, 3.3'
+urllib3==1.25.9; python_version != '3.4'
+yarl==1.4.2; python_version >= '3.5'