tobast/matrix-alertbot
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

create matrix bot to send and acknowledge alerts

Browse source
This commit is contained in:
HgO 2022-07-04 01:03:24 +02:00
parent 6ab094acdc
commit 96ee7f068a
17 changed files with 467 additions and 323 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
.dockerignore Normal file
View file

@ -0,0 +1,15 @@
# PyCharm
.idea/
.vscode/
# Python virtualenv environment folders
env/
env3/
.env/
.venv/
# Python
__pycache__/
*.egg-info/
build/
dist/

1
.gitignore vendored
View file

@ -1,5 +1,6 @@
# PyCharm
.idea/
.vscode/
# Python virtualenv environment folders
env/

2
docker/Dockerfile.dev
View file

@ -69,3 +69,5 @@ VOLUME ["/data"]
# Start the app
ENTRYPOINT ["matrix-alertbot", "/data/config.yaml"]
EXPOSE 8080

3
matrix-alertbot
View file

@ -1,10 +1,9 @@
#!/usr/bin/env python3
import asyncio
try:
from matrix_alertbot import main
# Run the main function of the bot
asyncio.get_event_loop().run_until_complete(main.main())
main.main()
except ImportError as e:
print("Unable to import matrix_alertbot.main:", e)

58
matrix_alertbot/alertmanager.py Normal file
View file

@ -0,0 +1,58 @@
import datetime
from typing import Dict, List
import pytimeparse
import requests
from matrix_alertbot.cache import Cache
from matrix_alertbot.errors import AlertNotFoundError
class AlertmanagerClient:
def __init__(self, url: str, cache: Cache) -> None:
self.api_url = f"{url}/api/v2"
self.cache = cache
def get_alerts(self) -> List[Dict]:
response = requests.get(f"{self.api_url}/alert")
response.raise_for_status()
return response.json()
def get_alert_labels(self, fingerprint: str) -> Dict[str, str]:
if fingerprint not in self.cache:
alerts = self.get_alerts()
alert = self._find_alert(alerts, fingerprint)
self.cache[fingerprint] = alert["labels"]
return self.cache[fingerprint]
def create_silence(self, fingerprint: str, duration: str, user: str) -> str:
labels = self.get_alert_labels(fingerprint)
matchers = []
for label_name, label_value in labels.items():
matchers.append(
{"name": label_name, "value": label_value, "isRegex": False}
)
start_time = datetime.datetime.now()
duration_seconds = pytimeparse.parse(duration)
duration_delta = datetime.timedelta(seconds=duration_seconds)
end_time = start_time + duration_delta
silence = {
"matchers": matchers,
"startsAt": start_time,
"endsAt": end_time,
"createdBy": user,
"comment": "Acknowledge alert from Matrix",
}
response = requests.post(f"{self.api_url}/silences", json=silence)
response.raise_for_status()
data = response.json()
return data["silenceID"]
@staticmethod
def _find_alert(alerts: List[Dict], fingerprint: str) -> Dict:
for alert in alerts:
if alert["fingerprint"] == fingerprint:
return alert
raise AlertNotFoundError(f"Cannot find alert with fingerprint {fingerprint}")

66
matrix_alertbot/bot_commands.py
View file

@ -1,15 +1,21 @@
import logging
from nio import AsyncClient, MatrixRoom, RoomMessageText
from matrix_alertbot.alertmanager import AlertmanagerClient
from matrix_alertbot.cache import Cache
from matrix_alertbot.chat_functions import react_to_event, send_text_to_room
from matrix_alertbot.config import Config
from matrix_alertbot.storage import Storage
logger = logging.getLogger(__name__)
class Command:
def __init__(
self,
client: AsyncClient,
store: Storage,
cache: Cache,
alertmanager: AlertmanagerClient,
config: Config,
command: str,
room: MatrixRoom,
@ -20,7 +26,7 @@ class Command:
Args:
client: The client to communicate to matrix with.
store: Bot storage.
cache: Bot cache.
config: Bot configuration parameters.
@ -31,7 +37,8 @@ class Command:
event: The event describing the command.
"""
self.client = client
self.store = store
self.cache = cache
self.alertmanager = alertmanager
self.config = config
self.command = command
self.room = room
@ -40,8 +47,8 @@ class Command:
async def process(self) -> None:
"""Process the command"""
if self.command.startswith("echo"):
await self._echo()
if self.command.startswith("ack"):
await self._ack()
elif self.command.startswith("react"):
await self._react()
elif self.command.startswith("help"):
@ -49,15 +56,46 @@ class Command:
else:
await self._unknown_command()
async def _echo(self) -> None:
"""Echo back the command's arguments"""
response = " ".join(self.args)
await send_text_to_room(self.client, self.room.room_id, response)
async def _ack(self) -> None:
"""Acknowledge an alert and silence it for a certain duration in Alertmanager"""
if len(self.args) > 0:
duration = " ".join(self.args)
else:
duration = "1d"
logger.debug(
f"Acknowledging alert with fingerprint {self.room.display_name} for a duration of {duration} | "
f"{self.room.user_name(self.event.sender)}: {self.event.body}"
)
source_content = self.event.source["content"]
try:
alert_event_id = source_content["m.relates_to"]["m.in_reply_to"]["event_id"]
except KeyError:
logger.debug("Unable to find the event ID of the alert")
return
logger.debug(f"Read alert fingerprint for event {alert_event_id} from cache")
alert_fingerprint = self.cache[alert_event_id]
logger.debug(
f"Create silence for alert with fingerprint {alert_fingerprint} for a duration of {duration}"
)
silence_id = self.alertmanager.create_silence(
alert_fingerprint, duration, self.room.user_name(self.event.sender)
)
await send_text_to_room(
self.client,
self.room.room_id,
f"Created silence {silence_id} for {duration}",
reply_to_event_id=alert_event_id,
)
async def _react(self) -> None:
"""Make the bot react to the command message"""
# React with a start emoji
reaction = "⭐"
logger.debug(
f"Reacting with {reaction} to room {self.room.display_name} | "
f"{self.room.user_name(self.event.sender)}: {self.event.body}"
)
await react_to_event(
self.client, self.room.room_id, self.event.event_id, reaction
)
@ -70,6 +108,10 @@ class Command:
async def _show_help(self) -> None:
"""Show the help text"""
logger.debug(
f"Displaying help to room {self.room.display_name} | "
f"{self.room.user_name(self.event.sender)}: {self.event.body}"
)
if not self.args:
text = (
"Hello, I am a bot made with matrix-nio! Use `help commands` to view "
@ -88,6 +130,10 @@ class Command:
await send_text_to_room(self.client, self.room.room_id, text)
async def _unknown_command(self) -> None:
logger.debug(
f"Sending unknown command response to room {self.room.display_name} | "
f"{self.room.user_name(self.event.sender)}: {self.event.body}"
)
await send_text_to_room(
self.client,
self.room.room_id,

18
matrix_alertbot/cache.py Normal file
View file

@ -0,0 +1,18 @@
from typing import Any
import diskcache
class Cache:
def __init__(self, directory: str, expire: int):
self.cache = diskcache.Cache(directory)
self.expire = expire
def __getitem__(self, key: str) -> Any:
return self.cache[key]
def __setitem__(self, key: str, value: Any) -> None:
self.cache.set(key, value, expire=self.expire)
def __contains__(self, key: str) -> bool:
return key in self.cache

85
matrix_alertbot/callbacks.py
View file

@ -11,27 +11,41 @@ from nio import (
UnknownEvent,
)
from matrix_alertbot.alertmanager import AlertmanagerClient
from matrix_alertbot.bot_commands import Command
from matrix_alertbot.chat_functions import make_pill, react_to_event, send_text_to_room
from matrix_alertbot.cache import Cache
from matrix_alertbot.chat_functions import (
make_pill,
react_to_event,
send_text_to_room,
strip_fallback,
)
from matrix_alertbot.config import Config
from matrix_alertbot.message_responses import Message
from matrix_alertbot.storage import Storage
logger = logging.getLogger(__name__)
class Callbacks:
def __init__(self, client: AsyncClient, store: Storage, config: Config):
def __init__(
self,
client: AsyncClient,
cache: Cache,
alertmanager: AlertmanagerClient,
config: Config,
):
"""
Args:
client: nio client used to interact with matrix.
store: Bot storage.
cache: Bot cache.
alertmanager: Client used to interact with alertmanager.
config: Bot configuration parameters.
"""
self.client = client
self.store = store
self.cache = cache
self.alertmanager = alertmanager
self.config = config
self.command_prefix = config.command_prefix
@ -44,12 +58,16 @@ class Callbacks:
event: The event defining the message.
"""
# Extract the message text
msg = event.body
msg = strip_fallback(event.body)
# Ignore messages from ourselves
if event.sender == self.client.user:
return
# Ignore messages from unauthorized room
if room.room_id != self.config.room:
return
logger.debug(
f"Bot message received for room {room.display_name} | "
f"{room.user_name(event.sender)}: {msg}"
@ -58,23 +76,18 @@ class Callbacks:
# Process as message if in a public room without command prefix
has_command_prefix = msg.startswith(self.command_prefix)
# room.is_group is often a DM, but not always.
# room.is_group does not allow room aliases
# room.member_count > 2 ... we assume a public room
# room.member_count <= 2 ... we assume a DM
if not has_command_prefix and room.member_count > 2:
# General message listener
message = Message(self.client, self.store, self.config, msg, room, event)
await message.process()
if not has_command_prefix:
logger.debug(
f"Message received without command prefix {self.command_prefix}: Aborting."
)
return
# Otherwise if this is in a 1-1 with the bot or features a command prefix,
# treat it as a command
if has_command_prefix:
# Remove the command prefix
msg = msg[len(self.command_prefix) :]
# Remove the command prefix
msg = msg[len(self.command_prefix) :]
command = Command(self.client, self.store, self.config, msg, room, event)
command = Command(
self.client, self.cache, self.alertmanager, self.config, msg, room, event
)
await command.process()
async def invite(self, room: MatrixRoom, event: InviteMemberEvent) -> None:
@ -85,6 +98,10 @@ class Callbacks:
event: The invite event.
"""
# Ignore invites from unauthorized room
if room.room_id != self.config.room:
return
logger.debug(f"Got invite to {room.room_id} from {event.sender}.")
# Attempt to join 3 times before giving up
@ -129,6 +146,14 @@ class Callbacks:
reacted_to_id: The event ID that the reaction points to.
"""
# Ignore reactions from unauthorized room
if room.room_id != self.config.room:
return
# Ignore reactions from ourselves
if event.sender == self.client.user:
return
logger.debug(f"Got reaction to {room.room_id} from {event.sender}.")
# Get the original event that was reacted to
@ -167,6 +192,10 @@ class Callbacks:
event: The encrypted event that we were unable to decrypt.
"""
# Ignore events from unauthorized room
if room.room_id != self.config.room:
return
logger.error(
f"Failed to decrypt event '{event.event_id}' in room '{room.room_id}'!"
f"\n\n"
@ -177,16 +206,6 @@ class Callbacks:
f"commands a second time)."
)
red_x_and_lock_emoji = "❌ 🔐"
# React to the undecryptable event with some emoji
await react_to_event(
self.client,
room.room_id,
event.event_id,
red_x_and_lock_emoji,
)
async def unknown(self, room: MatrixRoom, event: UnknownEvent) -> None:
"""Callback for when an event with a type that is unknown to matrix-nio is received.
Currently this is used for reaction events, which are not yet part of a released
@ -197,6 +216,10 @@ class Callbacks:
event: The event itself.
"""
# Ignore events from unauthorized room
if room.room_id != self.config.room:
return
if event.type == "m.reaction":
# Get the ID of the event this was a reaction to
relation_dict = event.source.get("content", {}).get("m.relates_to", {})

58
matrix_alertbot/chat_functions.py
View file

@ -1,7 +1,6 @@
import logging
from typing import Optional, Union, Dict
from typing import Optional, Union
from markdown import markdown
from nio import (
AsyncClient,
ErrorResponse,
@ -18,11 +17,11 @@ logger = logging.getLogger(__name__)
async def send_text_to_room(
client: AsyncClient,
room_id: str,
message: str,
plaintext: str,
html: str = None,
notice: bool = True,
markdown_convert: bool = True,
reply_to_event_id: Optional[str] = None,
) -> None:
) -> RoomSendResponse:
"""Send text to a matrix room.
Args:
@ -30,14 +29,13 @@ async def send_text_to_room(
room_id: The ID of the room to send the message to.
message: The message content.
plaintext: The message content.
html: The message content in HTML format.
notice: Whether the message should be sent with an "m.notice" message type
(will not ping users).
markdown_convert: Whether to convert the message content to markdown.
Defaults to true.
reply_to_event_id: Whether this message is a reply to another event. The event
ID this is message is a reply to.
@ -50,17 +48,17 @@ async def send_text_to_room(
content = {
"msgtype": msgtype,
"format": "org.matrix.custom.html",
"body": message,
"body": plaintext,
}
if markdown_convert:
content["formatted_body"] = markdown(message)
if html is not None:
content["formatted_body"] = html
if reply_to_event_id:
content["m.relates_to"] = {"m.in_reply_to": {"event_id": reply_to_event_id}}
try:
await client.room_send(
return await client.room_send(
room_id,
"m.room.message",
content,
@ -129,28 +127,12 @@ async def react_to_event(
)
async def decryption_failure(
client: AsyncClient, room: MatrixRoom, event: MegolmEvent
) -> None:
"""Callback for when an event fails to decrypt. Inform the user"""
logger.error(
f"Failed to decrypt event '{event.event_id}' in room '{room.room_id}'!"
f"\n\n"
f"Tip: try using a different device ID in your config file and restart."
f"\n\n"
f"If all else fails, delete your store directory and let the bot recreate "
f"it (your reminders will NOT be deleted, but the bot may respond to existing "
f"commands a second time)."
)
user_msg = (
"Unable to decrypt this message. "
"Check whether you've chosen to only encrypt to trusted devices."
)
await send_text_to_room(
client,
room.room_id,
user_msg,
reply_to_event_id=event.event_id,
)
def strip_fallback(content: str) -> str:
index = 0
for line in content.splitlines(keepends=True):
if not line.startswith("> "):
break
if index == 0:
index += 1
index += len(line)
return content[index:]

37
matrix_alertbot/config.py
View file

@ -4,6 +4,7 @@ import re
import sys
from typing import Any, List, Optional
import pytimeparse
import yaml
from matrix_alertbot.errors import ConfigError
@ -70,22 +71,13 @@ class Config:
f"storage.store_path '{self.store_path}' is not a directory"
)
# Database setup
database_path = self._get_cfg(["storage", "database"], required=True)
# Cache setup
self.cache_dir = self._get_cfg(["cache", "directory"], required=True)
expire_time = self._get_cfg(["cache", "expire_time"], default="1w")
self.cache_expire_time = pytimeparse.parse(expire_time)
# Support both SQLite and Postgres backends
# Determine which one the user intends
sqlite_scheme = "sqlite://"
postgres_scheme = "postgres://"
if database_path.startswith(sqlite_scheme):
self.database = {
"type": "sqlite",
"connection_string": database_path[len(sqlite_scheme) :],
}
elif database_path.startswith(postgres_scheme):
self.database = {"type": "postgres", "connection_string": database_path}
else:
raise ConfigError("Invalid connection string for storage.database")
# Alertmanager client setup
self.alertmanager_url = self._get_cfg(["alertmanager", "url"], required=True)
# Matrix bot account setup
self.user_id = self._get_cfg(["matrix", "user_id"], required=True)
@ -101,8 +93,21 @@ class Config:
self.device_name = self._get_cfg(
["matrix", "device_name"], default="nio-template"
)
self.homeserver_url = self._get_cfg(["matrix", "homeserver_url"], required=True)
self.homeserver_url = self._get_cfg(["matrix", "url"], required=True)
self.room = self._get_cfg(["matrix", "room"], required=True)
self.address = self._get_cfg(["webhook", "address"], required=False)
self.port = self._get_cfg(["webhook", "port"], required=False)
self.socket = self._get_cfg(["webhook", "socket"], required=False)
if (
not (self.address or self.port or self.socket)
or (self.socket and self.address and self.port)
or (self.address and not self.port)
or (not self.address and self.port)
):
raise ConfigError(
"Must supply either webhook.socket or both webhook.address and webhook.port"
)
self.command_prefix = self._get_cfg(["command_prefix"], default="!c") + " "
def _get_cfg(

13
matrix_alertbot/errors.py
View file

@ -8,5 +8,14 @@ class ConfigError(RuntimeError):
msg: The message displayed to the user on error.
"""
def __init__(self, msg: str):
super(ConfigError, self).__init__("%s" % (msg,))
pass
class AlertNotFoundError(RuntimeError):
"""An error encountered when an alert cannot be found in database.
Args:
msg: The message displayed to the user on error.
"""
pass

119
matrix_alertbot/main.py
View file

@ -2,9 +2,10 @@
import asyncio
import logging
import sys
from asyncio import TimeoutError
from time import sleep
from aiohttp import ClientConnectionError, ServerDisconnectedError
from aiohttp import ClientConnectionError, ServerDisconnectedError, web
from nio import (
AsyncClient,
AsyncClientConfig,
@ -16,59 +17,16 @@ from nio import (
UnknownEvent,
)
from matrix_alertbot.alertmanager import AlertmanagerClient
from matrix_alertbot.cache import Cache
from matrix_alertbot.callbacks import Callbacks
from matrix_alertbot.config import Config
from matrix_alertbot.storage import Storage
from matrix_alertbot.webhook import Webhook
logger = logging.getLogger(__name__)
async def main() -> bool:
"""The first function that is run when starting the bot"""
# Read user-configured options from a config file.
# A different config file path can be specified as the first command line argument
if len(sys.argv) > 1:
config_path = sys.argv[1]
else:
config_path = "config.yaml"
# Read the parsed config file and create a Config object
config = Config(config_path)
# Configure the database
store = Storage(config.database)
# Configuration options for the AsyncClient
client_config = AsyncClientConfig(
max_limit_exceeded=0,
max_timeouts=0,
store_sync_tokens=True,
encryption_enabled=True,
)
# Initialize the matrix client
client = AsyncClient(
config.homeserver_url,
config.user_id,
device_id=config.device_id,
store_path=config.store_path,
config=client_config,
)
if config.user_token:
client.access_token = config.user_token
client.user_id = config.user_id
# Set up event callbacks
callbacks = Callbacks(client, store, config)
client.add_event_callback(callbacks.message, (RoomMessageText,))
client.add_event_callback(
callbacks.invite_event_filtered_callback, (InviteMemberEvent,)
)
client.add_event_callback(callbacks.decryption_failure, (MegolmEvent,))
client.add_event_callback(callbacks.unknown, (UnknownEvent,))
async def start_matrix_client(client: AsyncClient, config: Config) -> bool:
# Keep trying to reconnect on failure (with some time in-between)
while True:
try:
@ -107,7 +65,7 @@ async def main() -> bool:
logger.info(f"Logged in as {config.user_id}")
await client.sync_forever(timeout=30000, full_state=True)
except (ClientConnectionError, ServerDisconnectedError):
except (ClientConnectionError, ServerDisconnectedError, TimeoutError):
logger.warning("Unable to connect to homeserver, retrying in 15s...")
# Sleep so we don't bombard the server with login requests
@ -117,5 +75,64 @@ async def main() -> bool:
await client.close()
# Run the main function in an asyncio event loop
asyncio.get_event_loop().run_until_complete(main())
def main() -> None:
"""The first function that is run when starting the bot"""
# Read user-configured options from a config file.
# A different config file path can be specified as the first command line argument
if len(sys.argv) > 1:
config_path = sys.argv[1]
else:
config_path = "config.yaml"
# Read the parsed config file and create a Config object
config = Config(config_path)
# Configure the cache
cache = Cache(config.cache_dir, config.cache_expire_time)
# Configure Alertmanager client
alertmanager = AlertmanagerClient(config.alertmanager_url, cache)
# Configuration options for the AsyncClient
client_config = AsyncClientConfig(
max_limit_exceeded=0,
max_timeouts=0,
store_sync_tokens=True,
encryption_enabled=True,
)
# Initialize the matrix client
client = AsyncClient(
config.homeserver_url,
config.user_id,
device_id=config.device_id,
store_path=config.store_path,
config=client_config,
)
if config.user_token:
client.access_token = config.user_token
client.user_id = config.user_id
# Set up event callbacks
callbacks = Callbacks(client, cache, alertmanager, config)
client.add_event_callback(callbacks.message, (RoomMessageText,))
client.add_event_callback(
callbacks.invite_event_filtered_callback, (InviteMemberEvent,)
)
client.add_event_callback(callbacks.decryption_failure, (MegolmEvent,))
client.add_event_callback(callbacks.unknown, (UnknownEvent,))
webhook_server = Webhook(client, cache, config)
loop = asyncio.get_event_loop()
loop.create_task(webhook_server.start())
loop.create_task(start_matrix_client(client, config))
try:
loop.run_forever()
except Exception as e:
logger.error(e)
finally:
loop.run_until_complete(webhook_server.close())

52
matrix_alertbot/message_responses.py
View file

@ -1,52 +0,0 @@
import logging
from nio import AsyncClient, MatrixRoom, RoomMessageText
from matrix_alertbot.chat_functions import send_text_to_room
from matrix_alertbot.config import Config
from matrix_alertbot.storage import Storage
logger = logging.getLogger(__name__)
class Message:
def __init__(
self,
client: AsyncClient,
store: Storage,
config: Config,
message_content: str,
room: MatrixRoom,
event: RoomMessageText,
):
"""Initialize a new Message
Args:
client: nio client used to interact with matrix.
store: Bot storage.
config: Bot configuration parameters.
message_content: The body of the message.
room: The room the event came from.
event: The event defining the message.
"""
self.client = client
self.store = store
self.config = config
self.message_content = message_content
self.room = room
self.event = event
async def process(self) -> None:
"""Process and possibly respond to the message"""
if self.message_content.lower() == "hello world":
await self._hello_world()
async def _hello_world(self) -> None:
"""Say hello"""
text = "Hello, world!"
await send_text_to_room(self.client, self.room.room_id, text)

180
matrix_alertbot/storage.py
View file

@ -1,126 +1,70 @@
import logging
from typing import Any, Dict
from __future__ import annotations
# The latest migration version of the database.
#
# Database migrations are applied starting from the number specified in the database's
# `migration_version` table + 1 (or from 0 if this table does not yet exist) up until
# the version specified here.
#
# When a migration is performed, the `migration_version` table should be incremented.
latest_migration_version = 0
import logging
from typing import Dict
logger = logging.getLogger(__name__)
class Storage:
def __init__(self, database_config: Dict[str, str]):
"""Setup the database.
class Alert:
EMOJIS = {"critical": "🔥", "warning": "⚠️", "resolved": "🥦"}
COLORS = {"critical": "dc3545", "warning": "ffc107", "resolved": "33cc33"}
Runs an initial setup or migrations depending on whether a database file has already
been created.
def __init__(
self,
id: str,
url: str,
firing: bool = True,
labels: Dict[str, str] = None,
annotations: Dict[str, str] = None,
):
self.id = id
self.url = url
self.firing = firing
Args:
database_config: a dictionary containing the following keys:
* type: A string, one of "sqlite" or "postgres".
* connection_string: A string, featuring a connection string that
be fed to each respective db library's `connect` method.
"""
self.conn = self._get_database_connection(
database_config["type"], database_config["connection_string"]
)
self.cursor = self.conn.cursor()
self.db_type = database_config["type"]
# Try to check the current migration version
migration_level = 0
try:
self._execute("SELECT version FROM migration_version")
row = self.cursor.fetchone()
migration_level = row[0]
except Exception:
self._initial_setup()
finally:
if migration_level < latest_migration_version:
self._run_migrations(migration_level)
logger.info(f"Database initialization of type '{self.db_type}' complete")
def _get_database_connection(
self, database_type: str, connection_string: str
) -> Any:
"""Creates and returns a connection to the database"""
if database_type == "sqlite":
import sqlite3
# Initialize a connection to the database, with autocommit on
return sqlite3.connect(connection_string, isolation_level=None)
elif database_type == "postgres":
import psycopg2
conn = psycopg2.connect(connection_string)
# Autocommit on
conn.set_isolation_level(0)
return conn
def _initial_setup(self) -> None:
"""Initial setup of the database"""
logger.info("Performing initial database setup...")
# Set up the migration_version table
self._execute(
"""
CREATE TABLE migration_version (
version INTEGER PRIMARY KEY
)
"""
)
# Initially set the migration version to 0
self._execute(
"""
INSERT INTO migration_version (
version
) VALUES (?)
""",
(0,),
)
# Set up any other necessary database tables here
logger.info("Database setup complete")
def _run_migrations(self, current_migration_version: int) -> None:
"""Execute database migrations. Migrates the database to the
`latest_migration_version`.
Args:
current_migration_version: The migration version that the database is
currently at.
"""
logger.debug("Checking for necessary database migrations...")
# if current_migration_version < 1:
# logger.info("Migrating the database from v0 to v1...")
#
# # Add new table, delete old ones, etc.
#
# # Update the stored migration version
# self._execute("UPDATE migration_version SET version = 1")
#
# logger.info("Database migrated to v1")
def _execute(self, *args: Any) -> None:
"""A wrapper around cursor.execute that transforms placeholder ?'s to %s for postgres.
This allows for the support of queries that are compatible with both postgres and sqlite.
Args:
args: Arguments passed to cursor.execute.
"""
if self.db_type == "postgres":
self.cursor.execute(args[0].replace("?", "%s"), *args[1:])
if labels is None:
self.labels = {}
else:
self.cursor.execute(*args)
self.labels = labels
if annotations is None:
self.annotations = {}
else:
self.annotations = annotations
if self.firing:
self.status = self.labels["severity"]
else:
self.status = "resolved"
@staticmethod
def from_dict(data: Dict) -> Alert:
return Alert(
id=data["fingerprint"],
url=data["generatorURL"],
firing=data["status"] == "firing",
labels=data["labels"],
annotations=data["annotations"],
)
@property
def emoji(self) -> str:
return self.EMOJIS[self.status]
@property
def color(self) -> str:
return self.COLORS[self.status]
def plaintext(self) -> str:
alertname = self.labels["alertname"]
description = self.annotations["description"]
return f"[{self.emoji} {self.status.upper()}] {alertname}: {description}"
def html(self) -> str:
alertname = self.labels["alertname"]
job = self.labels["job"]
description = self.annotations["description"]
return (
f"<font color='#{self.color}'><b>[{self.emoji} {self.status.upper()}]</b></font> "
f"<a href='{self.url}'>{alertname}</a> ({job})<br/>{description}"
)

74
matrix_alertbot/webhook.py Normal file
View file

@ -0,0 +1,74 @@
import logging
from typing import List
from aiohttp import web, web_request
from nio import AsyncClient, SendRetryError
from matrix_alertbot.cache import Cache
from matrix_alertbot.chat_functions import send_text_to_room
from matrix_alertbot.config import Config
from matrix_alertbot.storage import Alert
logger = logging.getLogger(__name__)
routes = web.RouteTableDef()
@routes.post("/alert")
async def create_alert(request: web_request.Request) -> web.Response:
data = await request.json()
logger.info(f"Received alert: {data}")
client = request.app["client"]
cache = request.app['cache']
plaintext = ""
html = ""
for i, alert in enumerate(data["alerts"]):
alert = Alert.from_dict(alert)
if i != 0:
plaintext += "\n"
html += "<br/>\n"
plaintext += alert.plaintext()
html += alert.html()
try:
event = await send_text_to_room(client, request.app["room_id"], plaintext, html)
except SendRetryError as e:
logger.error(e)
return web.Response(status=500)
cache[event.event_id] = tuple(alert["fingerprint"] for alert in data["alerts"])
return web.Response(status=200)
class Webhook:
def __init__(self, client: AsyncClient, cache: Cache, config: Config) -> None:
self.app = web.Application(logger=logger)
self.app["client"] = client
self.app["room_id"] = config.room
self.app["cache"] = cache
self.app.add_routes(routes)
self.runner = web.AppRunner(self.app)
self.config = config
self.address = config.address
self.port = config.port
self.socket = config.socket
async def start(self) -> None:
await self.runner.setup()
site: web.BaseSite
if self.address and self.port:
site = web.TCPSite(self.runner, self.address, self.port)
logger.info(f"Listenning on {self.address}:{self.port}")
elif self.socket:
site = web.UnixSite(self.runner, self.socket)
logger.info(f"Listenning on unix://{self.socket}")
await site.start()
async def close(self) -> None:
await self.runner.cleanup()

1
mypy.ini
View file

@ -2,3 +2,4 @@
ignore_missing_imports = True
disallow_untyped_defs = True
disallow_untyped_calls = True
plugins = sqlalchemy.ext.mypy.plugin

8
tests/test_callbacks.py
View file

@ -3,8 +3,9 @@ from unittest.mock import Mock
import nio
from matrix_alertbot.alertmanager import AlertmanagerClient
from matrix_alertbot.cache import Cache
from matrix_alertbot.callbacks import Callbacks
from matrix_alertbot.storage import Storage
from tests.utils import make_awaitable, run_coroutine
@ -15,13 +16,14 @@ class CallbacksTestCase(unittest.TestCase):
self.fake_client = Mock(spec=nio.AsyncClient)
self.fake_client.user = "@fake_user:example.com"
self.fake_storage = Mock(spec=Storage)
self.fake_cache = Mock(spec=Cache)
self.fake_alertmanager = Mock(spec=AlertmanagerClient)
# We don't spec config, as it doesn't currently have well defined attributes
self.fake_config = Mock()
self.callbacks = Callbacks(
self.fake_client, self.fake_storage, self.fake_config
self.fake_client, self.fake_cache, self.fake_alertmanager, self.fake_config
)
def test_invite(self) -> None: