matrix-alertbot/matrix_alertbot/main.py

#!/usr/bin/env python3
import asyncio
import logging
import sys
from asyncio import TimeoutError
from time import sleep

from aiohttp import ClientConnectionError, ServerDisconnectedError
from diskcache import Cache
from nio import (
    AsyncClient,
    AsyncClientConfig,
    InviteMemberEvent,
    LocalProtocolError,
    LoginError,
    MegolmEvent,
    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
    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_dir,
        config=client_config,
    )

    if config.user_token:
        client.access_token = config.user_token
        client.user_id = config.user_id

    return client


async def start_matrix_client(
    client: AsyncClient, cache: Cache, config: Config
) -> bool:
    # Keep trying to reconnect on failure (with some time in-between)
    while True:
        try:
            if config.user_token:
                # Use token to log in
                client.load_store()

                # Sync encryption keys with the server
                if client.should_upload_keys:
                    await client.keys_upload()
            else:
                # Try to login with the configured username/password
                try:
                    login_response = await 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

                # Login succeeded!

            logger.info(f"Logged in as {config.user_id}")
            await 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 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)

    client = create_matrix_client(config)

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

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

    # Set up event callbacks
    callbacks = Callbacks(client, alertmanager, cache, 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,))

    # Configure webhook server
    webhook_server = Webhook(client, cache, config)

    loop = asyncio.get_event_loop()
    loop.create_task(webhook_server.start())
    loop.create_task(start_matrix_client(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.close())
        loop.run_until_complete(client.close())