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

Include notification interval (converted to days) in SNS notifications
This commit is contained in:
Jasmine Schladen 2021-02-24 17:34:47 -08:00 committed by GitHub
commit 0136366c83
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 55 additions and 12 deletions

View File

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

View File

@ -564,4 +564,4 @@ class SNSNotificationPlugin(ExpirationNotificationPlugin):
f"{self.get_option('topicName', options)}" f"{self.get_option('topicName', options)}"
current_app.logger.info(f"Publishing {notification_type} notification to topic {topic_arn}") 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))

View File

@ -11,21 +11,24 @@ import arrow
import boto3 import boto3
from flask import current_app 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) sns_client = boto3.client("sns", **kwargs)
message_ids = {} message_ids = {}
subject = "Lemur: {0} Notification".format(notification_type.capitalize()) subject = "Lemur: {0} Notification".format(notification_type.capitalize())
for certificate in certificates: 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 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( response = sns_client.publish(
TopicArn=topic_arn, TopicArn=topic_arn,
Message=format_message(certificate, notification_type), Message=format_message(certificate, notification_type, options),
Subject=subject, 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 = { json_message = {
"notification_type": notification_type, "notification_type": notification_type,
"certificate_name": certificate["name"], "certificate_name": certificate["name"],
@ -57,4 +60,19 @@ def format_message(certificate, notification_type):
"owner": certificate["owner"], "owner": certificate["owner"],
"details": create_certificate_url(certificate["name"]) "details": create_certificate_url(certificate["name"])
} }
if notification_type == "expiration":
json_message["notification_interval_days"] = calculate_expiration_days(options)
return json.dumps(json_message) 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

View File

@ -13,9 +13,31 @@ from lemur.tests.test_messaging import verify_sender_email
@mock_sns() @mock_sns()
def test_format(certificate, endpoint): def test_format_nonexpiration(certificate, endpoint):
data = [certificate_notification_output_schema.dump(certificate).data] 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: for certificate in data:
expected_message = { expected_message = {
"notification_type": "expiration", "notification_type": "expiration",
@ -25,9 +47,10 @@ def test_format(certificate, endpoint):
"id": certificate["id"], "id": certificate["id"],
"endpoints_detected": 0, "endpoints_detected": 0,
"owner": certificate["owner"], "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() @mock_sns()
@ -52,7 +75,7 @@ def test_publish(certificate, endpoint):
topic_arn, sqs_client, queue_url = create_and_subscribe_to_topic() 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) assert len(message_ids) == len(data)
received_messages = sqs_client.receive_message(QueueUrl=queue_url)["Messages"] received_messages = sqs_client.receive_message(QueueUrl=queue_url)["Messages"]
@ -61,7 +84,7 @@ def test_publish(certificate, endpoint):
actual_message = next( actual_message = next(
(m for m in received_messages if json.loads(m["Body"])["MessageId"] == expected_message_id), None) (m for m in received_messages if json.loads(m["Body"])["MessageId"] == expected_message_id), None)
actual_json = json.loads(actual_message["Body"]) 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" 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"] received_messages = sqs_client.receive_message(QueueUrl=queue_url)["Messages"]
assert len(received_messages) == 1 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"] actual_message = json.loads(received_messages[0]["Body"])["Message"]
assert actual_message == expected_message assert actual_message == expected_message