add pong callback

.gitignore

@@ -141,3 +141,5 @@ cython_debug/
 # Project specific:
 .sync_token
 config.yaml
+*.key
+*.crt

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<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/>
+
+    But this is a custom template for Cassandre, so, here is a warning:<br/>
+    <blockquote>Beware of greeks bearing gifts</blockquote>
 ...

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
 

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
     )
 
 

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

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
 

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<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: 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"<pre><code>{body}</code></pre>\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()
 

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