diff --git a/src/kassandra/__main__.py b/src/kassandra/__main__.py index 0b3f6a0..ce90465 100644 --- a/src/kassandra/__main__.py +++ b/src/kassandra/__main__.py @@ -21,7 +21,8 @@ async def __main(): ) format_corout = format_alerts( alert_queue, - message_queue + message_queue, + config ) webhook_corout = run_webhook( alert_queue, diff --git a/src/kassandra/config.py b/src/kassandra/config.py index 75825e3..99d9e74 100644 --- a/src/kassandra/config.py +++ b/src/kassandra/config.py @@ -21,6 +21,8 @@ class Config: tls_crt: Optional[str] = None tls_key: Optional[str] = None ca_crt: Optional[str] = None + default_template: Optional[str] = None + templates: dict[str, str] = dataclasses.field(default_factory=dict) def check_integrity(self)->bool: """ Check the integrity of the config. diff --git a/src/kassandra/format.py b/src/kassandra/format.py index bef1adc..6305110 100644 --- a/src/kassandra/format.py +++ b/src/kassandra/format.py @@ -5,14 +5,16 @@ Format the alert message. import asyncio import dataclasses import json +from collections import defaultdict from typing import ( Any, NoReturn, Optional ) -from jinja2 import Environment, BaseLoader +from jinja2 import BaseLoader, Environment, Template +from .config import Config -template_raw = ( +default_template_raw = ( "{% if alert['labels']['severity'] == 'critical' and alert['status'] == 'firing' %}" "@room" "{% elif alert['labels']['severity'] == 'warning' and alert['status'] == 'firing' %}" @@ -27,24 +29,58 @@ template_raw = ( "{{ alert['annotations']['title'] }}
" "{{ alert['annotations']['description'] }}
" ) -template = Environment(loader=BaseLoader).from_string(template_raw) - @dataclasses.dataclass class Message: body: str formated_body: Optional[str] +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] + 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']: - formated_body = template.render(alert=alert, status=alert_json['status']) - body = f"{ alert['annotations']['title'] }:\n{ alert['annotations']['description'] }" + 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', '') @@ -53,14 +89,17 @@ def format_alert( async def format_alerts( alert_queue: asyncio.Queue[dict[str,Any]], - message_queue: asyncio.Queue[Message] + 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) + messages = format_alert(alert, templates) for message in messages: await message_queue.put(message) alert_queue.task_done()