FabledSteward/steward/core/notifications.py

from __future__ import annotations
import json
import logging
import smtplib
import ssl
import types
from email.message import EmailMessage

import httpx
from jinja2.sandbox import SandboxedEnvironment

logger = logging.getLogger(__name__)


async def dispatch_notifications(
    smtp_cfg: dict,
    webhook_cfg: dict,
    alert_data: dict,
    recipients: list = (),
) -> dict:
    """Dispatch all configured notification channels: email to each of
    `recipients` (mail.Recipient, from mail.recipients) when a server is set.
    Returns results dict."""
    results = {}

    if smtp_cfg.get("host") and recipients:
        results["email"] = await _send_email(smtp_cfg, alert_data, recipients)

    if webhook_cfg.get("url"):
        results["webhook"] = await _send_webhook(webhook_cfg, alert_data)

    return results


SMTP_TIMEOUT_S = 15


def smtp_mode(port, tls) -> str:
    """How to talk to the SMTP server: "ssl" (implicit TLS, port 465),
    "starttls" (upgrade a plain connection, when TLS is ticked) or "plain"."""
    if int(port or 0) == 465:
        return "ssl"
    return "starttls" if tls else "plain"


async def _send_email(cfg: dict, alert: dict, recipients: list) -> dict:
    """Each recipient their own copy, in their own zone, over one connection."""
    from steward.core.mail import alert_email, message, send
    try:
        messages = [message(cfg, to.address, *alert_email(cfg, alert, to)) for to in recipients]
        await send(cfg, messages)
        return {"sent": True, "error": None}
    except Exception as exc:
        logger.error(f"Email notification failed: {exc}")
        return {"sent": False, "error": str(exc)}


def smtp_open(cfg: dict) -> smtplib.SMTP:
    """A connection to the mail server, TLS and login done, every step
    bounded by SMTP_TIMEOUT_S. The caller closes it (use it in `with`)."""
    port = int(cfg.get("port") or 587)
    mode = smtp_mode(port, cfg.get("tls", True))
    ctx = ssl.create_default_context()
    if mode == "ssl":
        smtp = smtplib.SMTP_SSL(cfg["host"], port, timeout=SMTP_TIMEOUT_S, context=ctx)
    else:
        smtp = smtplib.SMTP(cfg["host"], port, timeout=SMTP_TIMEOUT_S)
    try:
        if mode == "starttls":
            smtp.starttls(context=ctx)
        if cfg.get("username"):
            smtp.login(cfg["username"], cfg.get("password", ""))
    except BaseException:
        smtp.close()
        raise
    return smtp


def _smtp_send_sync(cfg: dict, messages: list[EmailMessage]) -> None:
    with smtp_open(cfg) as smtp:
        for msg in messages:
            smtp.send_message(msg)


def _json_str(value) -> str:
    """A value as it may sit between the quotes of a JSON string."""
    return json.dumps("" if value is None else str(value))[1:-1]


def render_webhook(template_str: str, alert: dict) -> str:
    """The webhook body for an alert, checked to be JSON.

    In the template, `alert.<field>` is escaped to sit inside a JSON string
    ("{{ alert.rule_name }}" stays valid whatever the name holds), with a
    missing value reading "no data"; `raw.<field>` is the plain value for use
    with |tojson ({{ raw.value | tojson }} gives a number or null). The
    template runs sandboxed. ValueError says why a body could not be made.
    """
    shown = dict(alert)
    if shown.get("value") is None:
        shown["value"] = "no data"
    env = SandboxedEnvironment()
    try:
        rendered = env.from_string(template_str).render(
            alert=types.SimpleNamespace(**{k: _json_str(v) for k, v in shown.items()}),
            raw=types.SimpleNamespace(**alert))
    except Exception as exc:
        raise ValueError(f"Template render failed: {exc}") from None
    try:
        json.loads(rendered)
    except json.JSONDecodeError as exc:
        raise ValueError(f"Rendered webhook template is not valid JSON: {exc}") from None
    return rendered


async def _send_webhook(cfg: dict, alert: dict) -> dict:
    template_str = cfg.get("template") or '{"content": "{{ alert.state }} — {{ alert.rule_name }}"}'
    try:
        rendered = render_webhook(template_str, alert)
    except ValueError as exc:
        logger.error(str(exc))
        return {"sent": False, "error": str(exc)}

    try:
        async with httpx.AsyncClient(timeout=10) as client:
            response = await client.post(
                cfg["url"],
                content=rendered,
                headers={"Content-Type": "application/json"},
            )
            response.raise_for_status()
        return {"sent": True, "error": None}
    except Exception as exc:
        # The URL is a secret (it usually carries its own token) and httpx
        # puts it in its error messages.
        error = str(exc).replace(cfg["url"], "<webhook URL>")
        logger.error(f"Webhook notification failed: {error}")
        return {"sent": False, "error": error}