kassandra/src/kassandra/format.py

"""
Format the alert message.
"""

import asyncio
import dataclasses
import json
from collections import defaultdict
from typing import (
    Any,
    NoReturn,
    Optional
)
from jinja2 import BaseLoader, Environment, Template
from .config import Config

@dataclasses.dataclass
class Message:
    body: str
    formated_body: Optional[str]

DEFAULT_TEMPLATE_RAW: str = (
    "{% if alert['labels']['severity'] == 'critical' and alert['status'] == 'firing' %}"
    "@room<b><font color='red'>"
    "{% elif alert['labels']['severity'] == 'warning' and alert['status'] == 'firing' %}"
    "<b><font color='orange'>"
    "{% elif alert['status'] == 'resolved' %}"
    "<b><font color='green'> End of alert "
    "{% else %}"
    "<b><font>"
    "{% endif %}"
    "{{ alert['labels']['alertname'] }}: {{ alert['labels']['instance'] }} <br/>"
    "</font></b>"
    "{{ alert['annotations']['title'] }}<br/>"
    "{{ alert['annotations']['description'] }}<br/>"
)

def load_templates(
    config: Config
)->defaultdict[str, Template]:
    """
        Create a dict mapping alert names to the template to use from
        the config, with a default template either from the config file
        or the default DEFAULT_TEMPLATE_RAW.
    """
    if config.default_template is None:
        default_template = Environment(
            loader=BaseLoader
        ).from_string(
           DEFAULT_TEMPLATE_RAW
        )
    else:
        default_template = Environment(
            loader=BaseLoader
        ).from_string(
            config.default_template
        )
    templates = defaultdict(lambda: default_template)
    for alert_name in config.templates:
        templates[alert_name] = Environment(
            loader=BaseLoader
        ).from_string(
            config.templates[alert_name]
        )
    return templates


def format_alert(
    alert_json: dict[str, Any],
    templates: defaultdict[str, Template]
)->list[Message]:
    """
        Format an alert in json format to a nice string.
    """
    messages = []
    for alert in alert_json['alerts']:
        template = templates[alert['labels']['alertname']]
        formated_body = template.render(alert=alert)
        body = "alert {status}:\n{alertname} on {instance}".format(
            status=alert['status'],
            alertname=alert['labels']['alertname'],
            instance=alert['labels']['instance']
        )
        if '@room' in formated_body:
            body = '@room ' + body
            formated_body = formated_body.replace('@room', '')
        messages.append(Message(body, formated_body))
    return messages

async def format_alerts(
    alert_queue: asyncio.Queue[dict[str,Any]],
    message_queue: asyncio.Queue[Message],
    config: Config
)->NoReturn:
    """
        Read alerts from alert_queue, format them, and put them in message_queue.
    """
    templates = load_templates(config)

    while True:
        alert = await alert_queue.get()
        messages = format_alert(alert, templates)
        for message in messages:
            await message_queue.put(message)
        alert_queue.task_done()