kassandra/src/kassandra/format.py

106 lines
3 KiB
Python

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