Refactors how notifications are generated. (#584)

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
         )

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
         )
     )

lemur/notifications/messaging.py

@@ -8,44 +8,113 @@
 .. moduleauthor:: Kevin Glisson <kglisson@netflix.com>
 
 """
+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

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()
 

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

lemur/plugins/lemur_email/templates/expiration.html

@@ -43,7 +43,7 @@
                                 <tr>
                                     <td width="32px"></td>
                                     <td style="font-family:Roboto-Regular,Helvetica,Arial,sans-serif;font-size:24px;color:#ffffff;line-height:1.25">
-                                       Your certificate(s) are expiring!
+                                       Your certificate(s) are expiring in {{ message.options | interval }} {{ message.options | unit }}!
                                     </td>
                                     <td width="32px"></td>
                                 </tr>
@@ -79,7 +79,7 @@
                                                     <table border="0" cellspacing="0" cellpadding="0"
                                                            style="margin-top:48px;margin-bottom:48px">
                                                         <tbody>
-                                                        {%  for certificate in certificates %}
+                                                        {%  for certificate in message.certificates %}
                                                             <tr valign="middle">
                                                                 <td width="32px"></td>
                                                                 <td width="16px"></td>

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'))

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