diff --git a/.gitignore b/.gitignore index a61e346..1fd0f17 100644 --- a/.gitignore +++ b/.gitignore @@ -141,3 +141,5 @@ cython_debug/ # Project specific: .sync_token config.yaml +*.key +*.crt diff --git a/config.yaml.exp b/config.yaml.exp index f91e69d..05209bd 100644 --- a/config.yaml.exp +++ b/config.yaml.exp @@ -7,4 +7,22 @@ host: 127.0.0.1 alert_rooms: - "#troy:matrix.org" - "#ithaca:matrix.org" +templates: + "Test": | + {% if alert['labels']['severity'] == 'critical' and alert['status'] == 'firing' %} + @room + {% elif alert['labels']['severity'] == 'warning' and alert['status'] == 'firing' %} + + {% elif alert['status'] == 'resolved' %} + End of alert + {% else %} + + {% endif %} + {{ alert['labels']['alertname'] }}: {{ alert['labels']['instance'] }}
+
+ {{ alert['annotations']['title'] }}
+ {{ alert['annotations']['description'] }}
+ + But this is a custom template for Cassandre, so, here is a warning:
+
Beware of greeks bearing gifts
... diff --git a/setup.cfg b/setup.cfg index 583e008..88de4c7 100644 --- a/setup.cfg +++ b/setup.cfg @@ -12,6 +12,7 @@ packages = kassandra python_requires = >=3.9.2 package_dir = =src install_requires = + Jinja2>=3.0.2 PyYAML>=5.4.1 matrix-bot @ git+https://gitea.auro.re/histausse/matrix-bot.git diff --git a/src/kassandra/__main__.py b/src/kassandra/__main__.py index 90a4d15..ce90465 100644 --- a/src/kassandra/__main__.py +++ b/src/kassandra/__main__.py @@ -17,19 +17,16 @@ async def __main(): bot_corout = send_messages( message_queue, - config.username, - config.homeserver, - config.password, - config.alert_rooms + config ) format_corout = format_alerts( alert_queue, - message_queue + message_queue, + config ) webhook_corout = run_webhook( alert_queue, - config.host, - config.port + config ) diff --git a/src/kassandra/bot.py b/src/kassandra/bot.py index 5f5c723..66983c8 100644 --- a/src/kassandra/bot.py +++ b/src/kassandra/bot.py @@ -7,10 +7,25 @@ from typing import ( Any, NoReturn ) +import nio from matrix_bot.client import Client from matrix_bot.invite_policy import WhiteList +from matrix_bot.utils import ignore_client_message +from .config import Config from .format import Message +@ignore_client_message +async def __pong( + client: Client, + room: nio.rooms.MatrixRoom, + message: nio.events.room_events.RoomMessageText +): + """ + If the bot is pinged, send "Pong" + """ + if (client.user_name + ':') in message.body: + await client.send_message(room.room_id, "Pong") + async def __send_messsages( message_queue: asyncio.Queue[Message], bot: Client, @@ -31,27 +46,25 @@ async def __send_messsages( async def send_messages( message_queue: asyncio.Queue[dict[str, Any]], # For now, type will change in the futur - username: str, - homeserver: str, - password: str, - alert_rooms: list[str] + config: Config ): """ Initialize the bot and send messages added to the queue to the alert_rooms. """ bot = await Client( - username, - homeserver, - password + config.username, + config.homeserver, + config.password ) - invite_policy = await WhiteList(bot, alert_rooms) + invite_policy = await WhiteList(bot, config.alert_rooms) bot.set_invite_policy(invite_policy) + bot.add_message_callback(__pong) await asyncio.gather( bot.run(), __send_messsages( message_queue, bot, - alert_rooms + config.alert_rooms ) ) diff --git a/src/kassandra/config.py b/src/kassandra/config.py index d695be9..99d9e74 100644 --- a/src/kassandra/config.py +++ b/src/kassandra/config.py @@ -5,14 +5,40 @@ Config related tools. import dataclasses import yaml +from typing import Optional + @dataclasses.dataclass class Config: username: str homeserver: str password: str - port: int - host: str alert_rooms: list[str] + port: int = 8000 + host: str = "127.0.0.1" + endpoint: str = "/webhook" + tls: bool = False + tls_auth: bool = False + 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. + Raise an error if the config is invalid, + should always return True (or raise an error). + """ + if self.tls_auth and not self.tls: + raise ValueError("tls_auth is enable, but not tls.") + if self.tls and self.tls_crt is None: + raise ValueError("tls is enable but tls_crt was not provided") + if self.tls and self.tls_key is None: + raise ValueError("tls is enable but tls_key was not provided") + if self.tls_auth and self.ca_crt is None: + raise ValueError("tls_auth is enable, but ca_crt was not provided") + return True + def load_config(file:str)->Config: """ @@ -20,5 +46,7 @@ def load_config(file:str)->Config: """ with open(file, 'r') as f: data = yaml.load(f, Loader=yaml.loader.SafeLoader) - return Config(**data) + config = Config(**data) + config.check_integrity() + return config diff --git a/src/kassandra/format.py b/src/kassandra/format.py index f298565..7a65d65 100644 --- a/src/kassandra/format.py +++ b/src/kassandra/format.py @@ -5,37 +5,102 @@ 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" + "{% elif alert['labels']['severity'] == 'warning' and alert['status'] == 'firing' %}" + "" + "{% elif alert['status'] == 'resolved' %}" + " End of alert " + "{% else %}" + "" + "{% endif %}" + "{{ alert['labels']['alertname'] }}: {{ alert['labels']['instance'] }}
" + "
" + "{{ alert['annotations']['title'] }}
" + "{{ alert['annotations']['description'] }}
" +) + +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: dict[str, Any] -)->Message: + alert_json: dict[str, Any], + templates: defaultdict[str, Template] +)->list[Message]: """ Format an alert in json format to a nice string. """ - body = json.dumps(alert, indent=4) - formated_body = f"
{body}
\n" - return Message(body, formated_body) + 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] + 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() - message = format_alert(alert) - await message_queue.put(message) + messages = format_alert(alert, templates) + for message in messages: + await message_queue.put(message) alert_queue.task_done() diff --git a/src/kassandra/webhook.py b/src/kassandra/webhook.py index 709959c..50b60a1 100644 --- a/src/kassandra/webhook.py +++ b/src/kassandra/webhook.py @@ -5,31 +5,56 @@ The webhook receiving the alerts from alertmanager. import asyncio import aiohttp.web import aiohttp.web_request +import ssl from typing import ( Any, NoReturn ) +from .config import Config + +def load_ssl_context(config:Config)->ssl.SSLContext: + """ + Load the SSL context from the config. + """ + ca_path = None + if config.tls_auth: + ca_path = config.ca_crt + ssl_context = ssl.create_default_context( + purpose=ssl.Purpose.CLIENT_AUTH, + cafile=ca_path + ) + if config.tls_auth: + ssl_context.verify_mode = ssl.CERT_REQUIRED + ssl_context.load_cert_chain(config.tls_crt, config.tls_key) + return ssl_context -ENDPOINT = "/webhook" async def run_webhook( alert_queue: asyncio.Queue[dict[str, Any]], - host: str, - port: int + config: Config )->NoReturn: """ Run the webhook endpoint and put the alerts in the queue. """ - async def handler(request:aiohttp.web_request.Request): + async def handler(request:aiohttp.web_request.Request)->aiohttp.web.Response: alert = await request.json() await alert_queue.put(alert) return aiohttp.web.Response() + async def health(request:aiohttp.web_request.Request)->aiohttp.web.Response: + return aiohttp.web.Response(text="OK") + app = aiohttp.web.Application() - app.add_routes([aiohttp.web.post(ENDPOINT, handler)]) + app.add_routes([ + aiohttp.web.post(config.endpoint, handler), + aiohttp.web.get("/health", health) + ]) runner = aiohttp.web.AppRunner(app) await runner.setup() - site = aiohttp.web.TCPSite(runner, host, port) + ssl_context = None + if config.tls: + ssl_context = load_ssl_context(config) + site = aiohttp.web.TCPSite(runner, config.host, config.port, ssl_context=ssl_context) await site.start()