matrix-alertbot/matrix_alertbot/main.py

#!/usr/bin/env python3
import asyncio
import json
import logging
import os
import sys
from asyncio import TimeoutError

from aiohttp import ClientConnectionError, ServerDisconnectedError
from diskcache import Cache
from nio import (
    AsyncClient,
    AsyncClientConfig,
    InviteMemberEvent,
    KeyVerificationCancel,
    KeyVerificationKey,
    KeyVerificationMac,
    KeyVerificationStart,
    LocalProtocolError,
    LoginError,
    MegolmEvent,
    RedactionEvent,
    RoomMessageText,
    UnknownEvent,
)

from matrix_alertbot.alertmanager import AlertmanagerClient
from matrix_alertbot.callback import Callbacks
from matrix_alertbot.config import Config
from matrix_alertbot.webhook import Webhook

logger = logging.getLogger(__name__)


def create_matrix_client(config: Config) -> AsyncClient:
    # Configuration options for the AsyncClient
    try:
        matrix_client_config = AsyncClientConfig(
            max_limit_exceeded=0,
            max_timeouts=0,
            store_sync_tokens=True,
            encryption_enabled=True,
        )
    except ImportWarning as e:
        logger.warning(e)
        matrix_client_config = AsyncClientConfig(
            max_limit_exceeded=0,
            max_timeouts=0,
            store_sync_tokens=True,
            encryption_enabled=False,
        )

    # Load credentials from a previous session
    if os.path.exists(config.user_token_file):
        with open(config.user_token_file, "r") as ifd:
            credentials = json.load(ifd)
            config.user_token = credentials["access_token"]
            config.device_id = credentials["device_id"]

    # Initialize the matrix client based on stored credentials
    matrix_client = AsyncClient(
        config.homeserver_url,
        config.user_id,
        device_id=config.device_id,
        store_path=config.store_dir,
        config=matrix_client_config,
    )

    return matrix_client


async def start_matrix_client(
    matrix_client: AsyncClient, cache: Cache, config: Config
) -> bool:
    # Keep trying to reconnect on failure (with some time in-between)
    while True:
        try:
            if config.device_id and config.user_token:
                matrix_client.restore_login(
                    user_id=config.user_id,
                    device_id=config.device_id,
                    access_token=config.user_token,
                )

                # Sync encryption keys with the server
                if matrix_client.should_upload_keys:
                    await matrix_client.keys_upload()
            else:
                # Try to login with the configured username/password
                try:
                    login_response = await matrix_client.login(
                        password=config.user_password,
                        device_name=config.device_name,
                    )

                    # Check if login failed
                    if type(login_response) == LoginError:
                        logger.error("Failed to login: %s", login_response.message)
                        return False
                except LocalProtocolError as e:
                    # There's an edge case here where the user hasn't installed the correct C
                    # dependencies. In that case, a LocalProtocolError is raised on login.
                    logger.fatal(
                        "Failed to login. Have you installed the correct dependencies? "
                        "https://github.com/poljar/matrix-nio#installation "
                        "Error: %s",
                        e,
                    )
                    return False

                # Save user's access token and device ID
                # See https://stackoverflow.com/a/45368120
                user_token_fd = os.open(
                    config.user_token_file,
                    flags=os.O_CREAT | os.O_WRONLY | os.O_TRUNC,
                    mode=0o640,
                )
                with os.fdopen(user_token_fd, "w") as ofd:
                    json.dump(
                        {
                            "device_id": login_response.device_id,
                            "access_token": login_response.access_token,
                        },
                        ofd,
                    )

                # Login succeeded!

            logger.info(f"Logged in as {config.user_id}")
            await matrix_client.sync_forever(timeout=30000, full_state=True)
        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
            await asyncio.sleep(15)
        finally:
            await matrix_client.close()


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)

    matrix_client = create_matrix_client(config)

    # Configure the cache
    cache = Cache(config.cache_dir)

    # Configure Alertmanager client
    alertmanager_client = AlertmanagerClient(config.alertmanager_url, cache)

    # Set up event callbacks
    callbacks = Callbacks(matrix_client, alertmanager_client, cache, config)
    matrix_client.add_event_callback(callbacks.message, (RoomMessageText,))
    matrix_client.add_event_callback(
        callbacks.invite_event_filtered_callback, (InviteMemberEvent,)
    )
    matrix_client.add_event_callback(callbacks.decryption_failure, (MegolmEvent,))
    matrix_client.add_event_callback(callbacks.unknown, (UnknownEvent,))
    matrix_client.add_event_callback(callbacks.redaction, (RedactionEvent,))
    matrix_client.add_to_device_callback(
        callbacks.key_verification_start, (KeyVerificationStart,)
    )
    matrix_client.add_to_device_callback(
        callbacks.key_verification_cancel, (KeyVerificationCancel,)
    )
    matrix_client.add_to_device_callback(
        callbacks.key_verification_confirm, (KeyVerificationKey,)
    )
    matrix_client.add_to_device_callback(
        callbacks.key_verification_end, (KeyVerificationMac,)
    )
    # Configure webhook server
    webhook_server = Webhook(matrix_client, alertmanager_client, cache, config)

    loop = asyncio.get_event_loop()
    loop.create_task(webhook_server.start())
    loop.create_task(start_matrix_client(matrix_client, cache, config))

    try:
        loop.run_forever()
    except Exception as e:
        logger.error(e)
    finally:
        loop.run_until_complete(webhook_server.close())
        loop.run_until_complete(alertmanager_client.close())
        loop.run_until_complete(matrix_client.close())
        cache.close()