Add notifications module

sdss · Dec 22, 2024 · bfc605c · bfc605c
1 parent 9c7feb4
commit bfc605c
Show file tree

Hide file tree

Showing 7 changed files with 498 additions and 1 deletion.
diff --git a/pyproject.toml b/pyproject.toml
@@ -23,6 +23,7 @@ dependencies = [
     "aiocache>=0.12.3",
     "slack-sdk>=3.34.0",
     "aiohttp>=3.11.11",
+    "jinja2>=3.1.5",
 ]
 
 [project.optional-dependencies]
@@ -92,7 +93,7 @@ typing = ["typing"]
 sdss = ["sdsstools", "clu", "drift"]
 
 [tool.pytest.ini_options]
-addopts = "--cov lvmopstools --cov-report xml --cov-report html --cov-report term -W ignore"
+addopts = "--cov lvmopstools --cov-report xml --cov-report html --cov-report term -W ignore --capture=tee-sys"
 asyncio_mode = "auto"
 asyncio_default_fixture_loop_scope = 'function'
 

diff --git a/src/lvmopstools/data/config.yaml b/src/lvmopstools/data/config.yaml
@@ -15,6 +15,20 @@ pubsub:
   exchange_name: lvmops
   routing_key: data
 
+notifications:
+  critical:
+    email_template: data/critical_error_email_template.html
+    email_recipients:
+      - [email protected]
+    email_from: LVM Critical Alerts <[email protected]>
+    email_reply_to: [email protected]
+  smtp_server:
+    host: smtp.lco.cl
+    port: 25
+    tls: false
+    username: null
+    password: null
+
 devices:
   thermistors:
     host: 10.8.38.180

diff --git a/src/lvmopstools/data/critical_error_email_template.html b/src/lvmopstools/data/critical_error_email_template.html
@@ -0,0 +1,94 @@
+<html>
+  <head>
+    <style>
+      body {
+        font-family: Arial, Helvetica, sans-serif;
+        font-size: 12pt;
+      }
+
+      div.header {
+        width: 100%;
+        display: flex;
+        align-items: flex-start;
+      }
+
+      img.header-image {
+        padding: 4px;
+        width: 100px;
+      }
+
+      a, a:visited, a:link {
+        text-decoration: none;
+        color: black;
+      }
+
+      h1.header-text {
+        font-size: 2.25em;
+        font-weight: 600;
+        align-self: center;
+      }
+
+      div.section {
+        padding: 30px 0px 0px 0px;
+        font-size: 14pt;
+        font-weight: 600
+      }
+
+      div.section ~ hr {
+        width: 100%;
+        margin: 5px 0px 15px 0px;
+      }
+
+      ul {
+        list-style-position: outside;
+      }
+
+      li {
+        white-space: pre-line;
+      }
+
+      li + li {
+        margin-top: 10px;
+      }
+
+      pre {
+        width: 95%;
+        padding: 10px;
+        margin: 0;
+        overflow: auto;
+        overflow-y: hidden;
+        font-size: 12px;
+        line-height: 20px;
+        border: 1px solid #777;
+      }
+      pre code {
+        padding: 10px;
+      }
+    </style>
+  </head>
+
+  <body>
+    <div class="header">
+      <a href="https://lvm-web.lco.cl">
+        <img class="header-image" src="https://github.com/sdss/lvmapi/raw/refs/heads/main/src/lvmapi/data/lvm_logo.png" />
+      </a>
+      <div style="flex-grow: 1"></div>
+      <h1 class="header-text">
+        <a href="https://lvm-web.lco.cl">
+          SDSS-V LVM Critical Alert
+        </a>
+      </h1>
+    </div>
+    <hr style="width: 100%" />
+
+    <div style="padding: 10px 0px;">
+      A critical error has been raised in the LVM system. Please see the details below.
+    </div>
+
+    <div style="padding: 10px 0px;">
+      <pre><code>{{- message }}</code></pre>
+    </dic>
+
+  </body>
+
+</html>
diff --git a/src/lvmopstools/notifications.py b/src/lvmopstools/notifications.py
@@ -0,0 +1,204 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# @Author: José Sánchez-Gallego ([email protected])
+# @Date: 2024-12-21
+# @Filename: notifications.py
+# @License: BSD 3-clause (http://www.opensource.org/licenses/BSD-3-Clause)
+
+from __future__ import annotations
+
+import enum
+import pathlib
+import smtplib
+import sys
+from email.mime.multipart import MIMEMultipart
+from email.mime.text import MIMEText
+
+from typing import Any, Sequence, cast
+
+from jinja2 import Environment, FileSystemLoader
+
+from lvmopstools import config
+from lvmopstools.slack import post_message as post_to_slack
+
+
+__all__ = ["send_notification", "send_critical_error_email", "NotificationLevel"]
+
+
+class NotificationLevel(enum.Enum):
+    """Allowed notification levels."""
+
+    DEBUG = "DEBUG"
+    INFO = "INFO"
+    WARNING = "WARNING"
+    ERROR = "ERROR"
+    CRITICAL = "CRITICAL"
+
+
+async def send_notification(
+    message: str,
+    level: NotificationLevel | str = NotificationLevel.INFO,
+    slack: bool = True,
+    slack_channels: str | Sequence[str] | None = None,
+    email_on_critical: bool = True,
+    slack_extra_params: dict[str, Any] = {},
+    email_params: dict[str, Any] = {},
+):
+    """Creates a new notification.
+
+    Parameters
+    ----------
+    message
+        The message of the notification. Can be formatted in Markdown.
+    level
+        The level of the notification.
+    slack
+        Whether to send the notification to Slack.
+    slack_channels
+        The Slack channel where to send the notification. If not provided, the default
+        channel is used. Can be set to false to disable sending the Slack notification.
+    email_on_critical
+        Whether to send an email if the notification level is ``CRITICAL``.
+    slack_extra_params
+        A dictionary of extra parameters to pass to ``post_message``.
+    email_params
+        A dictionary of extra parameters to pass to :obj:`.send_critical_error_email`.
+
+    Returns
+    -------
+    message
+        The message that was sent.
+
+    """
+
+    if isinstance(level, str):
+        level = NotificationLevel(level.upper())
+    else:
+        level = NotificationLevel(level)
+
+    send_email = email_on_critical and level == NotificationLevel.CRITICAL
+
+    if send_email:
+        try:
+            send_critical_error_email(message, **email_params)
+        except Exception as ee:
+            print(f"Error sending critical error email: {ee}", file=sys.stderr)
+
+    if slack:
+        channels: set[str] = set()
+
+        if isinstance(slack_channels, str):
+            channels.add(slack_channels)
+        elif isinstance(slack_channels, Sequence):
+            channels.update(slack_channels)
+        else:
+            channels.add(config["slack.default_channel"])
+
+        # We send the message to the default channel plus any other channel that
+        # matches the level of the notification.
+        level_channels = cast(dict[str, str], config["slack.level_channels"])
+        if level.value in level_channels:
+            channels.add(level_channels[level.value])
+
+        # Send Slack message(s)
+        for channel in channels:
+            mentions = (
+                ["@channel"]
+                if level == NotificationLevel.CRITICAL
+                or level == NotificationLevel.ERROR
+                else []
+            )
+            try:
+                await post_to_slack(
+                    message,
+                    channel=channel,
+                    mentions=mentions,
+                    **slack_extra_params,
+                )
+            except Exception as se:
+                print(f"Error sending Slack message: {se}", file=sys.stderr)
+
+    return message
+
+
+def send_critical_error_email(
+    message: str,
+    host: str | None = None,
+    port: int | None = None,
+    tls: bool | None = None,
+    username: str | None = None,
+    password: str | None = None,
+):
+    """Sends a critical error email.
+
+    Parameters
+    ----------
+    message
+        The message to send.
+    host
+        The SMTP server host.
+    port
+        The SMTP server port.
+    tls
+        Whether to use TLS for authentication.
+    username
+        The SMTP server username.
+    password
+        The SMTP server password.
+
+    """
+
+    root = pathlib.Path(__file__).parent
+    template = root / config["notifications.critical.email_template"]
+    loader = FileSystemLoader(template.parent)
+
+    env = Environment(
+        loader=loader,
+        lstrip_blocks=True,
+        trim_blocks=True,
+    )
+    html_template = env.get_template(template.name)
+
+    html_message = html_template.render(message=message.strip())
+
+    recipients = config["notifications.critical.email_recipients"]
+    from_address = config["notifications.critical.email_from"]
+
+    email_reply_to = config["notifications.critical.email_reply_to"]
+
+    msg = MIMEMultipart("alternative" if html_message else "mixed")
+    msg["Subject"] = "LVM Critical Alert"
+    msg["From"] = from_address
+    msg["To"] = ", ".join(recipients)
+    msg["Reply-To"] = email_reply_to
+
+    plaintext_email = f"""A critical alert was raised in the LVM system.
+
+The error message is shown below:
+
+{message}
+
+"""
+    msg.attach(MIMEText(plaintext_email, "plain"))
+
+    html = MIMEText(html_message, "html")
+    msg.attach(html)
+
+    smtp_host = host or config["notifications.smtp_server.host"]
+    smtp_port = port or config["notifications.smtp_server.port"]
+    smpt_tls = tls if tls is not None else config["notifications.smtp_server.tls"]
+    smtp_username = username or config["notifications.smtp_server.username"]
+    smtp_password = password or config["notifications.smtp_server.password"]
+
+    with smtplib.SMTP(host=smtp_host, port=smtp_port) as smtp:
+        if smpt_tls is True or (smpt_tls is None and smtp_port == 587):
+            # See https://gist.github.com/jamescalam/93d915e4de12e7f09834ae73bdf37299
+            smtp.ehlo()
+            smtp.starttls()
+
+            if smtp_password is not None and smtp_password is not None:
+                smtp.login(smtp_username, smtp_password)
+            else:
+                raise ValueError("username and password must be provided for TLS.")
+        smtp.sendmail(from_address, recipients, msg.as_string())
diff --git a/tests/data/test_config.yaml b/tests/data/test_config.yaml
@@ -1,3 +1,23 @@
+slack:
+  token: null
+  default_channel: lvm-notifications
+  level_channels:
+    CRITICAL: lvm-alerts
+
+notifications:
+  critical:
+    email_template: data/critical_error_email_template.html
+    email_recipients:
+      - [email protected]
+    email_from: LVM Critical Alerts <[email protected]>
+    email_reply_to: [email protected]
+  smtp_server:
+    host: smtp.lco.cl
+    port: 25
+    tls: false
+    username: null
+    password: null
+
 pubsub:
   connection_string: amqp://guest:guest@localhost:<random-port>
   exchange_name: lvmops