plateforme-corres2math/corres2math/matrix.py

import os

from asgiref.sync import async_to_sync
from nio import *


class Matrix:
    _token: str = None
    _device_id: str = None

    @classmethod
    async def _get_client(cls) -> AsyncClient:
        client = AsyncClient("https://correspondances-maths.fr", "@corres2mathbot:correspondances-maths.fr")

        if os.path.isfile(".matrix_token"):
            with open(".matrix_device", "r") as f:
                cls._device_id = f.read().rstrip(" \t\r\n")
                client.device_id = cls._device_id
            with open(".matrix_token", "r") as f:
                cls._token = f.read().rstrip(" \t\r\n")
                client.access_token = cls._token
                return client

        await client.login(password=os.getenv("SYNAPSE_PASSWORD"), device_name="Plateforme")
        cls._token = client.access_token
        cls._device_id = client.device_id
        with open(".matrix_token", "w") as f:
            f.write(cls._token)
        with open(".matrix_device", "w") as f:
            f.write(cls._device_id)
        return client

    @classmethod
    @async_to_sync
    async def set_display_name(cls, name: str) -> Union[ProfileSetDisplayNameResponse, ProfileSetDisplayNameError]:
        client = await cls._get_client()
        return await client.set_displayname(name)

    @classmethod
    @async_to_sync
    async def create_room(
            cls,
            visibility: RoomVisibility = RoomVisibility.private,
            alias: Optional[str] = None,
            name: Optional[str] = None,
            topic: Optional[str] = None,
            room_version: Optional[str] = None,
            federate: bool = True,
            is_direct: bool = False,
            preset: Optional[RoomPreset] = None,
            invite=(),
            initial_state=(),
            power_level_override: Optional[Dict[str, Any]] = None,
    ) -> Union[RoomCreateResponse, RoomCreateError]:
        resp: Union[RoomCreateResponse, RoomCreateError]

        client = await cls._get_client()
        return await client.room_create(
            visibility, alias, name, topic, room_version, federate, is_direct, preset, invite, initial_state,
            power_level_override)

    @classmethod
    async def resolve_room_alias(cls, room_alias: str) -> Optional[str]:
        client = await cls._get_client()
        resp: RoomResolveAliasResponse = await client.room_resolve_alias(room_alias)
        if isinstance(resp, RoomResolveAliasError):
            return None
        return resp.room_id

    @classmethod
    @async_to_sync
    async def invite(cls, room_id: str, user_id: str) -> Union[RoomInviteResponse, RoomInviteError]:
        client = await cls._get_client()
        if room_id.startswith("#"):
            room_id = await cls.resolve_room_alias(room_id)
        return await client.room_invite(room_id, user_id)

    @classmethod
    @async_to_sync
    async def kick(cls, room_id: str, user_id: str, reason: str = None) -> Union[RoomKickResponse, RoomInviteError]:
        client = await cls._get_client()
        if room_id.startswith("#"):
            room_id = await cls.resolve_room_alias(room_id)
        return await client.room_kick(room_id, user_id, reason)

    @classmethod
    @async_to_sync
    async def set_room_power_level(cls, room_id: str, user_id: str, power_level: int)\
            -> Union[RoomPutStateResponse, RoomPutStateError]:
        client = await cls._get_client()
        if room_id.startswith("#"):
            room_id = await cls.resolve_room_alias(room_id)
        resp = await client.room_get_state_event(room_id, "m.room.power_levels")
        content = resp.content
        content["users"][user_id] = power_level
        print(content)
        return await client.room_put_state(room_id, "m.room.power_levels", content=content, state_key=resp.state_key)

    @classmethod
    @async_to_sync
    async def set_room_power_level_event(cls, room_id: str, event: str, power_level: int)\
            -> Union[RoomPutStateResponse, RoomPutStateError]:
        client = await cls._get_client()
        if room_id.startswith("#"):
            room_id = await cls.resolve_room_alias(room_id)
        resp = await client.room_get_state_event(room_id, "m.room.power_levels")
        content = resp.content
        if event.startswith("m."):
            content["events"][event] = power_level
        else:
            content[event] = power_level
        print(content)
        return await client.room_put_state(room_id, "m.room.power_levels", content=content, state_key=resp.state_key)