From 03d5a6cfe1155bde698342c1d004ffdf455684ab Mon Sep 17 00:00:00 2001 From: kevgliss Date: Mon, 12 Dec 2016 11:22:49 -0800 Subject: [PATCH] Refactors how notifications are generated. (#584) --- lemur/certificates/models.py | 2 +- lemur/notifications/cli.py | 7 +- lemur/notifications/messaging.py | 120 ++++++++++++++---- lemur/plugins/lemur_email/plugin.py | 6 +- lemur/plugins/lemur_email/templates/config.py | 12 ++ .../lemur_email/templates/expiration.html | 4 +- lemur/plugins/lemur_email/tests/test_email.py | 7 +- lemur/tests/test_messaging.py | 45 +++++-- 8 files changed, 156 insertions(+), 47 deletions(-) diff --git a/lemur/certificates/models.py b/lemur/certificates/models.py index 2c4d0c03..3932e70c 100644 --- a/lemur/certificates/models.py +++ b/lemur/certificates/models.py @@ -167,7 +167,7 @@ class Certificate(db.Model): def expired(cls): return case( [ - (cls.now_after <= arrow.utcnow(), True) + (cls.not_after <= arrow.utcnow(), True) ], else_=False ) diff --git a/lemur/notifications/cli.py b/lemur/notifications/cli.py index 50a0e586..ca30334f 100644 --- a/lemur/notifications/cli.py +++ b/lemur/notifications/cli.py @@ -21,9 +21,10 @@ def expirations(): :return: """ print("Starting to notify subscribers about expiring certificates!") - count = send_expiration_notifications() + success, failed = send_expiration_notifications() print( - "Finished notifying subscribers about expiring certificates! Sent {count} notifications!".format( - count=count + "Finished notifying subscribers about expiring certificates! Sent: {success} Failed: {failed}".format( + success=success, + failed=failed ) ) diff --git a/lemur/notifications/messaging.py b/lemur/notifications/messaging.py index 49f4dbac..77c7923c 100644 --- a/lemur/notifications/messaging.py +++ b/lemur/notifications/messaging.py @@ -8,44 +8,113 @@ .. moduleauthor:: Kevin Glisson """ +from itertools import groupby +from collections import defaultdict + +from sqlalchemy.orm import joinedload import arrow from flask import current_app from lemur import database, metrics from lemur.certificates.schemas import certificate_notification_output_schema -from lemur.notifications.models import Notification +from lemur.certificates.models import Certificate + from lemur.plugins import plugins from lemur.plugins.utils import get_plugin_option +def get_certificates(): + """ + Finds all certificates that are eligible for notifications. + :return: + """ + return database.session_query(Certificate)\ + .options(joinedload('notifications'))\ + .filter(Certificate.notify == True)\ + .filter(Certificate.expired == False)\ + .filter(Certificate.notifications.any()).all() # noqa + + +def get_eligible_certificates(): + """ + Finds all certificates that are eligible for certificate expiration. + :return: + """ + certificates = defaultdict(dict) + certs = get_certificates() + + # group by owner + for owner, items in groupby(certs, lambda x: x.owner): + notification_groups = [] + + for certificate in items: + notification = needs_notification(certificate) + + if notification: + notification_groups.append((notification, certificate)) + + # group by notification + for notification, items in groupby(notification_groups, lambda x: x[0].label): + certificates[owner][notification] = list(items) + + return certificates + + +def send_notification(event_type, data, targets, notification): + """ + Executes the plugin and handles failure. + + :param event_type: + :param data: + :param targets: + :param notification: + :return: + """ + try: + notification.plugin.send(event_type, data, targets, notification.options) + metrics.send('{0}_notification_sent'.format(event_type), 'counter', 1) + return True + except Exception as e: + metrics.send('{0}_notification_failure'.format(event_type), 'counter', 1) + current_app.logger.exception(e) + + def send_expiration_notifications(): """ This function will check for upcoming certificate expiration, and send out notification emails at given intervals. """ - sent = 0 - for plugin in plugins.all(plugin_type='notification'): - notifications = database.db.session.query(Notification)\ - .filter(Notification.plugin_name == plugin.slug)\ - .filter(Notification.active == True).all() # noqa + success = failure = 0 - messages = [] - for n in notifications: - for certificate in n.certificates: - if needs_notification(certificate): - data = certificate_notification_output_schema.dump(certificate).data - messages.append((data, n.options)) + # security team gets all + security_email = current_app.config.get('LEMUR_SECURITY_EMAIL') - for data, options in messages: - try: - plugin.send('expiration', data, [data['owner']], options) - metrics.send('expiration_notification_sent', 'counter', 1) - sent += 1 - except Exception as e: - metrics.send('expiration_notification_failure', 'counter', 1) - current_app.logger.exception(e) - return sent + security_data = [] + for owner, notification_group in get_eligible_certificates().items(): + + for notification_label, certificates in notification_group.items(): + notification_data = [] + + notification = certificates[0][0] + + for data in certificates: + n, certificate = data + cert_data = certificate_notification_output_schema.dump(certificate).data + notification_data.append(cert_data) + security_data.append(cert_data) + + if send_notification('expiration', notification_data, [owner], notification): + success += 1 + else: + failure += 1 + + if send_notification('expiration', security_data, [security_email], notification): + success += 1 + else: + failure += 1 + + return success, failure def send_rotation_notification(certificate, notification_plugin=None): @@ -64,6 +133,7 @@ def send_rotation_notification(certificate, notification_plugin=None): try: notification_plugin.send('rotation', data, [data['owner']]) metrics.send('rotation_notification_sent', 'counter', 1) + return True except Exception as e: metrics.send('rotation_notification_failure', 'counter', 1) current_app.logger.exception(e) @@ -77,12 +147,6 @@ def needs_notification(certificate): :param certificate: :return: """ - if not certificate.notify: - return - - if not certificate.notifications: - return - now = arrow.utcnow() days = (certificate.not_after - now).days @@ -103,4 +167,4 @@ def needs_notification(certificate): raise Exception("Invalid base unit for expiration interval: {0}".format(unit)) if days == interval: - return certificate + return notification diff --git a/lemur/plugins/lemur_email/plugin.py b/lemur/plugins/lemur_email/plugin.py index 92e0efba..8eeed698 100644 --- a/lemur/plugins/lemur_email/plugin.py +++ b/lemur/plugins/lemur_email/plugin.py @@ -28,7 +28,7 @@ def render_html(template_name, message): :return: """ template = env.get_template('{}.html'.format(template_name)) - return template.render(dict(messages=message, hostname=current_app.config.get('LEMUR_HOSTNAME'))) + return template.render(dict(message=message, hostname=current_app.config.get('LEMUR_HOSTNAME'))) def send_via_smtp(subject, body, targets): @@ -86,9 +86,11 @@ class EmailNotificationPlugin(ExpirationNotificationPlugin): @staticmethod def send(notification_type, message, targets, options, **kwargs): + subject = 'Lemur: {0} Notification'.format(notification_type.capitalize()) - body = render_html(notification_type, message) + data = {'options': options, 'certificates': message} + body = render_html(notification_type, data) s_type = current_app.config.get("LEMUR_EMAIL_SENDER", 'ses').lower() diff --git a/lemur/plugins/lemur_email/templates/config.py b/lemur/plugins/lemur_email/templates/config.py index 882e5b4d..5c9b1f8e 100644 --- a/lemur/plugins/lemur_email/templates/config.py +++ b/lemur/plugins/lemur_email/templates/config.py @@ -2,6 +2,8 @@ import os import arrow from jinja2 import Environment, FileSystemLoader +from lemur.plugins.utils import get_plugin_option + loader = FileSystemLoader(searchpath=os.path.dirname(os.path.realpath(__file__))) env = Environment(loader=loader) @@ -10,4 +12,14 @@ def human_time(time): return arrow.get(time).format('dddd, MMMM D, YYYY') +def interval(options): + return get_plugin_option('interval', options) + + +def unit(options): + return get_plugin_option('unit', options) + + env.filters['time'] = human_time +env.filters['interval'] = interval +env.filters['unit'] = unit diff --git a/lemur/plugins/lemur_email/templates/expiration.html b/lemur/plugins/lemur_email/templates/expiration.html index 8c2f8eab..3c500c38 100644 --- a/lemur/plugins/lemur_email/templates/expiration.html +++ b/lemur/plugins/lemur_email/templates/expiration.html @@ -43,7 +43,7 @@ - Your certificate(s) are expiring! + Your certificate(s) are expiring in {{ message.options | interval }} {{ message.options | unit }}! @@ -79,7 +79,7 @@ - {% for certificate in certificates %} + {% for certificate in message.certificates %} diff --git a/lemur/plugins/lemur_email/tests/test_email.py b/lemur/plugins/lemur_email/tests/test_email.py index 821e4da7..3eada55c 100644 --- a/lemur/plugins/lemur_email/tests/test_email.py +++ b/lemur/plugins/lemur_email/tests/test_email.py @@ -12,12 +12,15 @@ def test_render(certificate, endpoint): new_cert = CertificateFactory() new_cert.replaces.append(certificate) - certificates = [certificate_notification_output_schema.dump(certificate).data] + data = { + 'certificates': [certificate_notification_output_schema.dump(certificate).data], + 'options': [{'name': 'interval', 'value': 10}, {'name': 'unit', 'value': 'days'}] + } template = env.get_template('{}.html'.format('expiration')) with open(os.path.join(dir_path, 'expiration-rendered.html'), 'w') as f: - body = template.render(dict(certificates=certificates, hostname='lemur.test.example.com')) + body = template.render(dict(message=data, hostname='lemur.test.example.com')) f.write(body) template = env.get_template('{}.html'.format('rotation')) diff --git a/lemur/tests/test_messaging.py b/lemur/tests/test_messaging.py index ca51fdc2..11ea848b 100644 --- a/lemur/tests/test_messaging.py +++ b/lemur/tests/test_messaging.py @@ -23,21 +23,48 @@ def test_needs_notification(app, certificate, notification): assert needs_notification(certificate) +def test_get_certificates(app, certificate, notification): + from lemur.notifications.messaging import get_certificates + delta = certificate.not_after - timedelta(days=2) + with freeze_time(delta.datetime): + # no notification + certs = len(get_certificates()) + + # with notification + certificate.notifications.append(notification) + assert len(get_certificates()) > certs + + certificate.notify = False + assert len(get_certificates()) == certs + + # expired + delta = certificate.not_after + timedelta(days=2) + with freeze_time(delta.datetime): + certificate.notifications.append(notification) + assert len(get_certificates()) == 0 + + +def test_get_eligible_certificates(app, certificate, notification): + from lemur.notifications.messaging import get_eligible_certificates + + certificate.notifications.append(notification) + certificate.notifications[0].options = [{'name': 'interval', 'value': 10}, {'name': 'unit', 'value': 'days'}] + + delta = certificate.not_after - timedelta(days=10) + with freeze_time(delta.datetime): + assert get_eligible_certificates() == {certificate.owner: {notification.label: [(notification, certificate)]}} + + @mock_ses def test_send_expiration_notification(certificate, notification, notification_plugin): from lemur.notifications.messaging import send_expiration_notifications - notification.options = [{'name': 'interval', 'value': 10}, {'name': 'unit', 'value': 'days'}] + certificate.notifications.append(notification) + certificate.notifications[0].options = [{'name': 'interval', 'value': 10}, {'name': 'unit', 'value': 'days'}] + delta = certificate.not_after - timedelta(days=10) - with freeze_time(delta.datetime): - sent = send_expiration_notifications() - assert sent == 1 - - certificate.notify = False - - sent = send_expiration_notifications() - assert sent == 0 + assert send_expiration_notifications() == (2, 0) @mock_ses