Merge pull request #171 from newrelic/kwest/lambda-1197

Kwest/lambda 1197

Author: Katie West

newrelic_lambda_cli/integrations.py

@@ -7,19 +7,20 @@
 import click
 import json
 
-from newrelic_lambda_cli.cliutils import failure, success
+from newrelic_lambda_cli.cliutils import failure, success, warning
 from newrelic_lambda_cli.functions import get_function
 from newrelic_lambda_cli.types import (
     IntegrationInstall,
     IntegrationUninstall,
     IntegrationUpdate,
 )
-from newrelic_lambda_cli.utils import catch_boto_errors
+from newrelic_lambda_cli.utils import catch_boto_errors, NR_DOCS_ACT_LINKING_URL
 
 INGEST_STACK_NAME = "NewRelicLogIngestion"
 LICENSE_KEY_STACK_NAME = "NewRelicLicenseKeySecret"
 
 __cached_license_key_policy_arn = None
+__cached_license_key_nr_account_id = None
 
 
 def _get_role(session, role_name):
@@ -44,10 +45,27 @@ def _check_for_ingest_stack(session):
     return _get_cf_stack_status(session, INGEST_STACK_NAME)
 
 
-def _get_cf_stack_status(session, stack_name):
+def _get_cf_stack_status(session, stack_name, nr_account_id=None):
     """Returns the status of the CloudFormation stack if it exists"""
     try:
         res = session.client("cloudformation").describe_stacks(StackName=stack_name)
+        if nr_account_id is not None:
+            stack_output_account_id = _get_stack_output_value(session, ["NrAccountId"])
+            # Checking length of outputs here to protect against installs done with older CLI versions.
+            # We don't want to constantly warn users who installed on previous versions with no outputs.
+            if (
+                len(stack_output_account_id) > 0
+                and str(nr_account_id) not in stack_output_account_id
+            ):
+                warning(
+                    "WARNING: Managed secret already exists in this region for New Relic account {0}.\n"
+                    "Current CLI behavior limits the setup of one managed secret per region.\n"
+                    "To set up an additional secret for New Relic account {1} see our docs:\n{2}.\n"
+                    "Or run this command with --disable-license-key-secret to avoid attemping to create a new managed secret.".format(
+                        stack_output_account_id, nr_account_id, NR_DOCS_ACT_LINKING_URL
+                    )
+                )
+
     except botocore.exceptions.ClientError as e:
         if (
             e.response
@@ -590,7 +608,9 @@ def auto_install_license_key(input):
 @catch_boto_errors
 def install_license_key(input, nr_license_key, policy_name=None, mode="CREATE"):
     assert isinstance(input, (IntegrationInstall, IntegrationUpdate))
-    lk_stack_status = _get_cf_stack_status(input.session, LICENSE_KEY_STACK_NAME)
+    lk_stack_status = _get_cf_stack_status(
+        input.session, LICENSE_KEY_STACK_NAME, input.nr_account_id
+    )
     if lk_stack_status is not None:
         success("Managed secret already exists")
         return True
@@ -611,8 +631,14 @@ def install_license_key(input, nr_license_key, policy_name=None, mode="CREATE"):
         elif update_mode:
             parameters.append({"ParameterKey": "PolicyName", "UsePreviousValue": True})
 
-        parameters.append(
-            {"ParameterKey": "LicenseKey", "ParameterValue": nr_license_key}
+        parameters.extend(
+            (
+                {"ParameterKey": "LicenseKey", "ParameterValue": nr_license_key},
+                {
+                    "ParameterKey": "NrAccountId",
+                    "ParameterValue": str(input.nr_account_id),
+                },
+            )
         )
 
         change_set_name = "%s-%s-%d" % (
@@ -677,11 +703,24 @@ def remove_license_key(input):
         success("Done")
 
 
-def _get_license_key_policy_arn(session):
-    """Returns the policy ARN for the license key secret if it exists"""
+def _get_license_key_outputs(session):
+    """Returns the account id and policy ARN for the license key secret if they exist"""
+    global __cached_license_key_nr_account_id
     global __cached_license_key_policy_arn
-    if __cached_license_key_policy_arn:
-        return __cached_license_key_policy_arn
+    if __cached_license_key_nr_account_id and __cached_license_key_policy_arn:
+        return [__cached_license_key_nr_account_id, __cached_license_key_policy_arn]
+
+    account_id, policy_arn = _get_stack_output_value(
+        session, ["NrAccountId", "ViewPolicyARN"]
+    )
+    if account_id is not None:
+        __cached_license_key_nr_account_id = account_id
+    if policy_arn is not None:
+        __cached_license_key_policy_arn = policy_arn
+    return [account_id, policy_arn]
+
+
+def _get_stack_output_value(session, output_keys):
     client = session.client("cloudformation")
     try:
         stacks = client.describe_stacks(StackName=LICENSE_KEY_STACK_NAME).get(
@@ -700,13 +739,13 @@ def _get_license_key_policy_arn(session):
         if not stacks:
             return None
         stack = stacks[0]
-        output_key = "ViewPolicyARN"
-        for output in stack.get("Outputs", []):
-            if output["OutputKey"] == output_key:
-                __cached_license_key_policy_arn = output["OutputValue"]
-                return __cached_license_key_policy_arn
-        else:
-            return None
+        output_values = []
+        for output_key in output_keys:
+            for output in stack.get("Outputs", []):
+                if output["OutputKey"] == output_key:
+                    value = output["OutputValue"] or None
+                    output_values.append(value)
+        return output_values
 
 
 @catch_boto_errors

newrelic_lambda_cli/layers.py

@@ -11,7 +11,7 @@
 from newrelic_lambda_cli import api, subscriptions, utils
 from newrelic_lambda_cli.cliutils import failure, success, warning
 from newrelic_lambda_cli.functions import get_function
-from newrelic_lambda_cli.integrations import _get_license_key_policy_arn
+from newrelic_lambda_cli.integrations import _get_license_key_outputs
 from newrelic_lambda_cli.types import LayerInstall, LayerUninstall
 from newrelic_lambda_cli.utils import catch_boto_errors
 
@@ -188,7 +188,24 @@ def install(input, function_arn):
         failure("Could not find function: %s" % function_arn)
         return False
 
-    policy_arn = _get_license_key_policy_arn(input.session)
+    nr_account_id, policy_arn = _get_license_key_outputs(input.session)
+    # If a managed secret exists but it was created with a different NR account id and license key
+    # We want to notify the user and point them to documentation on how to proceed.
+    if (
+        policy_arn
+        and nr_account_id != str(input.nr_account_id)
+        and not input.nr_api_key
+    ):
+        raise click.UsageError(
+            "A managed secret already exists in this region for New Relic account {0}. "
+            "Creating one managed secret per region is currently supported via the cli.\n"
+            "To set up an additional secret for New Relic account {1} see our docs:\n{2}.\n"
+            "Or you can re-run this command with "
+            "`--nr-api-key` argument with your New Relic API key to set your license "
+            "key in a NEW_RELIC_LICENSE_KEY environment variable instead.".format(
+                nr_account_id, input.nr_account_id, utils.NR_DOCS_ACT_LINKING_URL
+            )
+        )
     if input.enable_extension and not policy_arn and not input.nr_api_key:
         raise click.UsageError(
             "In order to use `--enable-extension`, you must first run "
@@ -201,7 +218,12 @@ def install(input, function_arn):
         )
 
     nr_license_key = None
-    if not policy_arn and input.nr_api_key and input.nr_region:
+    if (
+        not policy_arn
+        or nr_account_id != str(input.nr_account_id)
+        and input.nr_api_key
+        and input.nr_region
+    ):
         gql = api.validate_gql_credentials(input)
         nr_license_key = api.retrieve_license_key(gql)
 
@@ -319,7 +341,7 @@ def uninstall(input, function_arn):
         )
         return False
     else:
-        policy_arn = _get_license_key_policy_arn(input.session)
+        _, policy_arn = _get_license_key_outputs(input.session)
         if policy_arn:
             _detach_license_key_policy(
                 input.session, config["Configuration"]["Role"], policy_arn

newrelic_lambda_cli/templates/license-key-secret.yaml

@@ -18,14 +18,16 @@ Parameters:
   ViewPolicyExportName:
     Type: String
     Default: NewRelic-ViewLicenseKeyPolicyARN
+  NrAccountId:
+    Type: String
 
 Resources:
   LicenseKeySecret:
     Type: 'AWS::SecretsManager::Secret'
     Properties:
       Description: The New Relic license key, for sending telemetry
       Name: !Sub "${SecretName}"
-      SecretString: !Sub '{ "LicenseKey": "${LicenseKey}" }'
+      SecretString: !Sub '{ "LicenseKey": "${LicenseKey}", "NrAccountId": "${NrAccountId}" }'
   ViewNewRelicLicenseKeyPolicy:
     Type: 'AWS::IAM::ManagedPolicy'
     Properties:
@@ -51,3 +53,7 @@ Outputs:
     Value: !Ref ViewNewRelicLicenseKeyPolicy
     Export:
       Name: !Sub "${AWS::StackName}-${ViewPolicyExportName}"
+  NrAccountId:
+    Value: !Ref NrAccountId
+    Export:
+      Name: !Sub "${AWS::StackName}-NrAccountId"

newrelic_lambda_cli/utils.py

@@ -6,6 +6,7 @@
 import botocore
 import click
 
+NR_DOCS_ACT_LINKING_URL = "https://docs.newrelic.com/docs/serverless-function-monitoring/aws-lambda-monitoring/enable-lambda-monitoring/account-linking/#manually-configuring-the-license-key-secret"
 NEW_RELIC_ARN_PREFIX_TEMPLATE = "arn:aws:lambda:%s:451483290750"
 RUNTIME_CONFIG = {
     "dotnetcore3.1": {"LambdaExtension": True},

setup.py

@@ -6,7 +6,7 @@
 
 setup(
     name="newrelic-lambda-cli",
-    version="0.5.5",
+    version="0.6.0",
     python_requires=">=3.3",
     description="A CLI to install the New Relic AWS Lambda integration and layers.",
     long_description=README,

tests/test_integrations.py

@@ -21,7 +21,8 @@
     remove_log_ingestion_function,
     install_license_key,
     remove_license_key,
-    _get_license_key_policy_arn,
+    _get_license_key_outputs,
+    _get_stack_output_value,
     get_aws_account_id,
     update_log_ingestion_function,
     remove_integration_role,
@@ -203,7 +204,9 @@ def test_install_license_key(success_mock):
         cf_client = MagicMock(name="cloudformation", **cf_mocks)
         mock_client_factory.side_effect = [cf_client, cf_client]
 
-        result = install_license_key(integration_install(session=session), "1234abcd")
+        result = install_license_key(
+            integration_install(session=session, nr_account_id=12345), "1234abcd"
+        )
         assert result is True
 
         cf_client.assert_has_calls(
@@ -213,6 +216,7 @@ def test_install_license_key(success_mock):
                     TemplateBody=ANY,
                     Parameters=[
                         {"ParameterKey": "LicenseKey", "ParameterValue": "1234abcd"},
+                        {"ParameterKey": "NrAccountId", "ParameterValue": "12345"},
                     ],
                     Capabilities=["CAPABILITY_NAMED_IAM"],
                     Tags=[],
@@ -266,13 +270,26 @@ def test_remove_license_key(success_mock):
         success_mock.assert_called_once()
 
 
-def test__get_license_key_policy_arn():
+def test__get_license_key_outputs():
+    with patch(
+        "newrelic_lambda_cli.integrations._get_stack_output_value"
+    ) as mock_get_stack_output:
+        mock_get_stack_output.return_value = [12345, "policy_arn"]
+        session = MagicMock()
+        result = _get_license_key_outputs(session)
+        assert result == mock_get_stack_output.return_value
+        mock_get_stack_output.assert_called_once_with(
+            session, ["NrAccountId", "ViewPolicyARN"]
+        )
+
+
+def test__get_stack_output_value():
     session = MagicMock()
     with patch.object(session, "client") as mock_client_factory:
         cf_client = MagicMock(name="cloudformation")
         mock_client_factory.side_effect = cf_client
 
-        _get_license_key_policy_arn(session)
+        _get_stack_output_value(session, ["NrAccountId", "ViewPolicyARN"])
 
         cf_client.assert_has_calls(
             [call().describe_stacks(StackName="NewRelicLicenseKeySecret")],