matrix-bot/src/main.py

from email.policy import default
import sys
import coloredlogs
import logging
import os
import time
import asyncio
from tokenize import group
from urllib import response
from ldapHelper import LdapHelper
from matrixHelper import MatrixHelper

coloredlogs.install(
    level='INFO', fmt='%(asctime)s - [%(levelname)s] %(message)s')


class MlmMatrixBot:
    def __init__(self) -> None:
        self._config = self._readConfig()
        self._matrixUsername = f"@{self._config['MATRIX_USERNAME']}:{self._config['MATRIX_DOMAIN']}"

        self._ldapHelper = LdapHelper(self._config['LDAP_URI'], self._config['LDAP_BIND_DN'],
                                      self._config['LDAP_BIND_DN_PASSWORD'], self._config['LDAP_BASE_DN'])
        self._matrixHelper = MatrixHelper(
            self._config['MATRIX_SERVER'], self._matrixUsername, self._config['MATRIX_SPACE_ID'], self._matrixUsername)

    async def init(self):
        logging.info("==> Initializing <==")

        logging.info(" \t* Connecting to ldap")
        if not self._ldapHelper.bind():
            return False

        logging.info(" \t* Connecting to matrix")
        if not await self._matrixHelper.login(self._config['MATRIX_PASSWORD']):
            return False

        logging.info("==> Initialized <==")
        return True

    async def run(self):

        self._current_log_step = 0
        logging.info("==> Sync started <==")

        await self._create_default_rooms()

        current_users, all_projects = await self._load_users_and_groups()

        await self._create_project_rooms(all_projects)

        lab_rooms = await self._create_lab_rooms()

        self._log_step("Loading current rooms")
        matrix_rooms = await self._get_managed_rooms(with_power_levels=True)
        matrix_room_ids = list(map(lambda room: room['id'], matrix_rooms))
        room_name_id_map = {}
        room_id_map = {}

        for room in matrix_rooms:
            room_name_id_map[room['name']] = room['id']
            room_id_map[room['id']] = room

        self._log_step("Syncing user memberships")
        for user in current_users:
            if not "mlm" in user['groups']:
                continue

            # default groups and space
            rooms_to_join = self._get_default_rooms(
            ) + [self._config["MATRIX_SPACE_ID"]]
            # projects
            rooms_to_join += list(map(lambda group: group.replace('p_', ''),
                                  filter(lambda project: project.startswith("p_"), user['groups'])))
            # resolve names to ids
            rooms_to_join = list(map(
                lambda room: room if room.startswith(
                    "!") else room_name_id_map[room],
                rooms_to_join))

            logging.info(f" \t* Syncing user {user['username']}")
            for room_id in rooms_to_join:
                # set power level
                if room_id in user['rooms']:
                    continue

                logging.info(f"\t\t* joining {room_id_map[room_id]['name']}")
                await self._matrixHelper.add_user_to_room(user['matrix_username'], room_id)

            # remove from rooms that are not in the groups
            rooms_to_be_left = list(filter(
                lambda room: room in matrix_room_ids and not room in rooms_to_join and not room_id_map[room]['name'] in lab_rooms, user['rooms']))
            for room_id in rooms_to_be_left:
                logging.info(f"\t\t* leaving {room_id_map[room_id]['name']}")
                await self._matrixHelper.remove_user_from_room(
                    user['matrix_username'], room_id)

            # update power levels
            for room in matrix_rooms:
                if (not room['id'] in rooms_to_join and not room['id'] in user['rooms']) or room['id'] in rooms_to_be_left:
                    continue

                if (user['matrix_username'] in room['power_levels'] and room['power_levels'][user['matrix_username']] != user['power_level']) or (user['matrix_username'] not in room['power_levels'] and user['power_level'] != 0):
                    logging.info(
                        f"\t\t* Setting power level in {room['name']} to {user['power_level']}")
                    await self._matrixHelper.set_user_power_level_in_room(user['matrix_username'], room['id'], user['power_level'])

        logging.info("==> Sync finished <==")

    async def _create_default_rooms(self):
        # create default rooms
        self._log_step("Creating default rooms")
        await self._create_rooms(self._get_default_rooms())

    async def _load_users_and_groups(self):
        """
        loads users and their groups from ldap and matrix and merges them
        returns:
        ([
            {
                'username': 'dozedler',
                'groups': ['mlm', 'p_3d-printing'],
                'rooms': ['!room1:matrix.org', '!room2:matrix.org'],
                'matrix_username': '@dozedler:matrix.org',
                'is_admin': True
            }
        ],
        [
            'mlm',
            'p_3d-printing'
        ])
        """
        self._log_step("Loading users and their groups")
        matrix_users = await self._matrixHelper.get_users()
        current_users = []
        all_projects = []

        for user in matrix_users:
            username = user['name'].split(':')[0].replace('@', '')
            rc, ldap_user = self._ldapHelper.search(
                f"(sAMAccountName={username})", ['distinguishedName'])
            if not rc:
                continue

            rc, raw_groups = self._ldapHelper.search(f"(&(member:1.2.840.113556.1.4.1941:={ldap_user[0]['distinguishedName']})(|(sophomorixType=adminclass)(sophomorixType=project)))", [
                'sAMAccountName', 'sophomorixType', 'sophomorixMailList'])

            if not rc:
                continue

            all_groups = list(
                map(lambda group: group['sAMAccountName'], raw_groups))
            power_level = 100 if self._config['ADMIN_GROUP'] in all_groups else 50 if self._config['MODERATOR_GROUP'] in all_groups else 0

            groups = list(filter(lambda group: group['sophomorixMailList'] ==
                          'TRUE' or group['sophomorixType'] == "adminclass", raw_groups))

            all_projects.extend(list(map(lambda group: group['sAMAccountName'], filter(
                lambda group: 'sophomorixType' in group and group['sophomorixType'] == 'project', groups))))

            rooms = await self._matrixHelper.get_rooms_of_user(user['name'])
            groups = list(map(lambda group: group['sAMAccountName'], groups))

            current_users.append({
                'username': username,
                'matrix_username': user['name'],
                'groups': groups,
                'rooms': rooms,
                'power_level': power_level
            })

        all_projects = list(set(all_projects))

        return current_users, all_projects

    async def _create_project_rooms(self, projects):
        self._log_step("Creating project rooms")
        projects = list(
            map(lambda project: project.replace('p_', ''), projects))
        await self._create_rooms(projects)

    async def _create_lab_rooms(self):
        self._log_step("Creating lab rooms")

        rc, labs = self._ldapHelper.search(
            f"(&(sophomorixType=project)(sAMAccountName=*lab))", ['sAMAccountName'])
        if not rc:
            return

        labs = list(map(lambda lab: lab['sAMAccountName'].replace(
            "p_", "").replace("lab", "-lab"), labs))
        await self._create_rooms(labs, joinable_for_space_members=True)

        return labs

    async def _create_rooms(self, rooms, joinable_for_space_members=False, suggested=False):
        matrix_room_names = list(map(lambda room: room['name'], await self._get_managed_rooms()))
        for room in rooms:
            if room not in matrix_room_names:
                logging.info(f"* Creating room {room}")
                await self._matrixHelper.create_room(room, room, joinable_for_space_members=joinable_for_space_members, suggested=True)

    async def _get_managed_rooms(self, with_power_levels=False):
        matrix_rooms = await self._matrixHelper.get_rooms()
        matrix_rooms = filter(
            lambda room: room['name'] != None and room['creator'] == self._matrixUsername, matrix_rooms)
        matrix_rooms = map(
            lambda room: {'id': room['room_id'], 'name': room['name']}, matrix_rooms)
        matrix_rooms = list(matrix_rooms)

        if with_power_levels:
            for i in range(len(matrix_rooms)):
                logging.info(
                    f"\t* Getting power levels of room {matrix_rooms[i]['name']}")
                matrix_rooms[i]['power_levels'] = await self._matrixHelper.get_room_power_levels(matrix_rooms[i]['id'])

        return matrix_rooms

    def _get_default_rooms(self):
        return self._config['DEFAULT_ROOMS'].split(',')

    def _log_step(self, message):
        logging.info(f"= ({self._current_log_step}/6) {message} =")
        self._current_log_step += 1

    def _readConfig(self):
        requiredConfigKeys = [
            'MATRIX_BOT_LDAP_URI',
            'MATRIX_BOT_LDAP_BASE_DN',
            'MATRIX_BOT_LDAP_BIND_DN',
            'MATRIX_BOT_LDAP_BIND_DN_PASSWORD',
            'MATRIX_BOT_MATRIX_DOMAIN',
            'MATRIX_BOT_MATRIX_SPACE_ID',
            'MATRIX_BOT_MATRIX_SERVER',
            'MATRIX_BOT_MATRIX_USERNAME',
            'MATRIX_BOT_MATRIX_PASSWORD',
            'MATRIX_BOT_DEFAULT_ROOMS',
            'MATRIX_BOT_ADMIN_GROUP',
            'MATRIX_BOT_MODERATOR_GROUP'
        ]

        allowedConfigKeys = [
        ]

        config = {
        }

        for configKey in requiredConfigKeys:
            if configKey not in os.environ:
                logging.error(
                    f"Required environment value {configKey} is not set")
                sys.exit()
            config[configKey.replace('MATRIX_BOT_', '')
                   ] = os.environ[configKey]

        for configKey in allowedConfigKeys:
            if configKey in os.environ:
                config[configKey.replace(
                    'MATRIX_BOT_', '')] = os.environ[configKey]

        logging.info("CONFIG:")
        for key, value in config.items():
            logging.info("    * {:25}: {}".format(key, value))

        return config


async def main():
    bot = MlmMatrixBot()
    if not await bot.init():
        logging.critical("Bot initialization failed")
        sys.exit(1)

    while True:
        await bot.run()
        time.sleep(60)

if __name__ == "__main__":
    try:
        asyncio.run(main())

    except KeyboardInterrupt:
        logging.info("Bye")