matrix-room-exporter/main.py

import argparse
import asyncio
import logging
import os
from pathlib import Path

import redis
import yaml
from nio import InviteMemberEvent, RoomMessageText

from exporter.callbacks import MatrixBotCallbacks
from exporter.config import DEFAULT_CONFIG
from exporter.matrix import MatrixClientHelper

SCRIPT_PATH = os.path.dirname(os.path.realpath(__file__))


async def main(args):
    logging.basicConfig()
    logger = logging.getLogger('ExportBot')

    with open(args.config) as stream:
        try:
            yaml_config = yaml.safe_load(stream)
        except yaml.YAMLError as exc:
            logger.critical(f'Failed to load config: {exc}')
            quit(1)

    # Really lazy way to do this
    config = DEFAULT_CONFIG | yaml_config
    config['r2'] = DEFAULT_CONFIG['r2'] | yaml_config.get('r2', {})
    config['allowed_to_export'] = yaml_config.get('allowed_to_export', [])

    logger.setLevel(logging.DEBUG if args.debug else logging.INFO)
    logger.debug('Debug logging enabled!')

    r = redis.Redis(host='localhost', port=6379, db=3)
    r.flushdb()

    matrix_helper = MatrixClientHelper(
        user_id=config['auth']['username'],
        passwd=config['auth']['password'],
        homeserver=config['auth']['homeserver'],
        store_path=config['store_path'],
    )
    client = matrix_helper.client

    login_success, login_response = await matrix_helper.login()
    if not login_success:
        logger.critical(f'Failed to login: {login_response}')
        quit(1)
    logger.info(f'Logged in as {client.user_id}')

    callbacks = MatrixBotCallbacks(matrix_helper, config)
    client.add_event_callback(callbacks.handle_invite, InviteMemberEvent)
    client.add_event_callback(callbacks.handle_message, RoomMessageText)

    await client.sync_forever(timeout=10000)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Export a Matrix room.')
    parser.add_argument('--config', default=Path(SCRIPT_PATH) / 'config.yml')
    parser.add_argument('--debug', '-d', action='store_true', help='Enable debug logging.')
    args = parser.parse_args()
    asyncio.run(main(args))