Merge pull request #3439 from jtschladen/notification-interval-sns

Include notification interval (converted to days) in SNS notifications

lemur/notifications/messaging.py

@@ -201,6 +201,7 @@ def send_plugin_notification(event_type, data, recipients, notification):
         "notification_plugin": notification.plugin.slug,
         "certificate_targets": recipients,
         "plugin": notification.plugin.slug,
+        "notification_id": notification.id,
     }
     status = FAILURE_METRIC_STATUS
     try:

lemur/plugins/lemur_aws/plugin.py

@@ -564,4 +564,4 @@ class SNSNotificationPlugin(ExpirationNotificationPlugin):
                     f"{self.get_option('topicName', options)}"
 
         current_app.logger.info(f"Publishing {notification_type} notification to topic {topic_arn}")
-        sns.publish(topic_arn, message, notification_type, region_name=self.get_option("region", options))
+        sns.publish(topic_arn, message, notification_type, options, region_name=self.get_option("region", options))

lemur/plugins/lemur_aws/sns.py

@@ -11,21 +11,24 @@ import arrow
 import boto3
 from flask import current_app
 
+from lemur.plugins.lemur_aws.plugin import SNSNotificationPlugin
 
-def publish(topic_arn, certificates, notification_type, **kwargs):
+
+def publish(topic_arn, certificates, notification_type, options, **kwargs):
     sns_client = boto3.client("sns", **kwargs)
     message_ids = {}
     subject = "Lemur: {0} Notification".format(notification_type.capitalize())
     for certificate in certificates:
-        message_ids[certificate["name"]] = publish_single(sns_client, topic_arn, certificate, notification_type, subject)
+        message_ids[certificate["name"]] = publish_single(sns_client, topic_arn, certificate, notification_type,
+                                                          subject, options)
 
     return message_ids
 
 
-def publish_single(sns_client, topic_arn, certificate, notification_type, subject):
+def publish_single(sns_client, topic_arn, certificate, notification_type, subject, options):
     response = sns_client.publish(
         TopicArn=topic_arn,
-        Message=format_message(certificate, notification_type),
+        Message=format_message(certificate, notification_type, options),
         Subject=subject,
     )
 
@@ -46,7 +49,7 @@ def create_certificate_url(name):
     )
 
 
-def format_message(certificate, notification_type):
+def format_message(certificate, notification_type, options):
     json_message = {
         "notification_type": notification_type,
         "certificate_name": certificate["name"],
@@ -57,4 +60,19 @@ def format_message(certificate, notification_type):
         "owner": certificate["owner"],
         "details": create_certificate_url(certificate["name"])
     }
+    if notification_type == "expiration":
+        json_message["notification_interval_days"] = calculate_expiration_days(options)
     return json.dumps(json_message)
+
+
+def calculate_expiration_days(options):
+    unit = SNSNotificationPlugin.get_option("unit", options)
+    interval = SNSNotificationPlugin.get_option("interval", options)
+    if unit == "weeks":
+        return interval * 7
+
+    elif unit == "months":
+        return interval * 30
+
+    elif unit == "days":
+        return interval

lemur/plugins/lemur_aws/tests/test_sns.py

@@ -13,9 +13,31 @@ from lemur.tests.test_messaging import verify_sender_email
 
 
 @mock_sns()
-def test_format(certificate, endpoint):
+def test_format_nonexpiration(certificate, endpoint):
     data = [certificate_notification_output_schema.dump(certificate).data]
 
+    for certificate in data:
+        expected_message = {
+            "notification_type": "not-expiration",
+            "certificate_name": certificate["name"],
+            "expires": arrow.get(certificate["validityEnd"]).format("YYYY-MM-DDTHH:mm:ss"),
+            "issuer": certificate["issuer"],
+            "id": certificate["id"],
+            "endpoints_detected": 0,
+            "owner": certificate["owner"],
+            "details": "https://lemur.example.com/#/certificates/{name}".format(name=certificate["name"])
+        }
+        # We don't currently support any SNS notifications besides expiration;
+        # when we do, this test will probably need to be refactored.
+        # For now, this is a placeholder proving empty options works as long as it's not "expiration" type
+        assert expected_message == json.loads(format_message(certificate, "not-expiration", None))
+
+
+@mock_sns()
+def test_format_expiration(certificate, endpoint):
+    data = [certificate_notification_output_schema.dump(certificate).data]
+    options = get_options()
+
     for certificate in data:
         expected_message = {
             "notification_type": "expiration",
@@ -25,9 +47,10 @@ def test_format(certificate, endpoint):
             "id": certificate["id"],
             "endpoints_detected": 0,
             "owner": certificate["owner"],
-            "details": "https://lemur.example.com/#/certificates/{name}".format(name=certificate["name"])
+            "details": "https://lemur.example.com/#/certificates/{name}".format(name=certificate["name"]),
+            "notification_interval_days": 10  # 10 days specified in options
         }
-        assert expected_message == json.loads(format_message(certificate, "expiration"))
+        assert expected_message == json.loads(format_message(certificate, "expiration", options))
 
 
 @mock_sns()
@@ -52,7 +75,7 @@ def test_publish(certificate, endpoint):
 
     topic_arn, sqs_client, queue_url = create_and_subscribe_to_topic()
 
-    message_ids = publish(topic_arn, data, "expiration", region_name="us-east-1")
+    message_ids = publish(topic_arn, data, "expiration", get_options(), region_name="us-east-1")
     assert len(message_ids) == len(data)
     received_messages = sqs_client.receive_message(QueueUrl=queue_url)["Messages"]
 
@@ -61,7 +84,7 @@ def test_publish(certificate, endpoint):
         actual_message = next(
             (m for m in received_messages if json.loads(m["Body"])["MessageId"] == expected_message_id), None)
         actual_json = json.loads(actual_message["Body"])
-        assert actual_json["Message"] == format_message(certificate, "expiration")
+        assert actual_json["Message"] == format_message(certificate, "expiration", get_options())
         assert actual_json["Subject"] == "Lemur: Expiration Notification"
 
 
@@ -98,7 +121,8 @@ def test_send_expiration_notification():
 
     received_messages = sqs_client.receive_message(QueueUrl=queue_url)["Messages"]
     assert len(received_messages) == 1
-    expected_message = format_message(certificate_notification_output_schema.dump(certificate).data, "expiration")
+    expected_message = format_message(certificate_notification_output_schema.dump(certificate).data, "expiration",
+                                      notification.options)
     actual_message = json.loads(received_messages[0]["Body"])["Message"]
     assert actual_message == expected_message