stockbot-buyvm/matrix.py

"""
matrix.py

A module for interacting with the Matrix protocol.

Classes:
    MatrixBot: A Matrix bot that can send messages and markdown messages to a room.

Dependencies:
    markdown: A library for converting markdown to HTML.
    loguru: A library for logging.
    nio: A library for interacting with the Matrix protocol.
"""
import markdown
from loguru import logger
from nio import AsyncClient, LoginResponse


class MatrixBot:
    """
    A Matrix bot that can send messages and markdown messages to a room.

    Attributes:
        config (dict): A dictionary containing the bot's configuration.
            Expected keys are 'homeserver', 'username', 'password', 'room_id'.
        client (AsyncClient): The Matrix client instance.
        logged_in (bool): Whether the bot is currently logged in.

    Methods:
        __init__: Initializes the bot with a given configuration.
        ensure_logged_in: Ensures that the bot is logged in to the Matrix homeserver.
        send_message: Sends a message to the room specified in the bot's configuration.
        send_markdown: Sends a markdown formatted message to the room specified in the bot's configuration.
        close: Log out from the Matrix homeserver and close the client.
    """
    def __init__(self, config: dict):
        """
        A Matrix bot that can send messages and markdown messages to a room.

        Args:
            config (dict): A dictionary containing the bot's configuration.
                Expected keys are 'homeserver', 'username', 'password', 'room_id'.
        """
        self.config = config

        self.client = AsyncClient(
            homeserver=self.config['homeserver'],
            user=self.config['username']
        )
        self.logged_in = False

    async def ensure_logged_in(self):
        """
        Ensures that the bot is logged in to the Matrix homeserver.

        If the bot is not logged in, attempts to log in using the provided
        password. If the login attempt fails, logs the error and closes the
        nio session.

        If an exception occurs during the login attempt, logs the error and
        re-raises it.
        """
        if not self.logged_in:
            try:
                response = await self.client.login(password=self.config['password'])
                if isinstance(response, LoginResponse):
                    self.logged_in = True
                    logger.info(f"Logged in as {self.config['username']}")
                else:
                    logger.error(f"Failed to login as {self.config['username']}: {response}")
                    logger.error("Closing nio session")
                    await self.client.close()
            except Exception as e:
                logger.error(f"Exception during login: {e}")
                await self.client.close()
                raise

    async def send_message(self, message: str):
        """
        Sends a message to the room specified in the bot's configuration.

        The message is sent as a simple text message, with the 'msgtype' set to
        'm.text' and the 'body' set to the provided message.

        If the bot is not logged in, attempts to log in using the provided
        password. If the login attempt fails, logs the error and closes the
        nio session.

        If an exception occurs during the login attempt or the message sending,
        logs the error and re-raises it.

        Args:
            message (str): The message to send to the room.
        """
        await self.ensure_logged_in()

        if not self.logged_in:
            logger.error("Unable to send message, login failed")
            return

        try:
            await self.client.room_send(
                room_id=self.config['room_id'],
                message_type="m.room.message",
                content={
                    "msgtype": "m.text",
                    "body": message
                }
            )
            logger.info("Message sent")
        except Exception as e:
            logger.error(f"Exception during sending message: {e}")
            raise

    async def send_markdown(self, message: str):
        """
        Sends a markdown formatted message to the room specified in the bot's
        configuration.

        The message is sent as a text message with the 'msgtype' set to
        'm.text', the 'body' set to the provided message, and the 'format'
        set to 'org.matrix.custom.html'. The 'formatted_body' is set to the
        markdown formatted message.

        If the bot is not logged in, attempts to log in using the provided
        password. If the login attempt fails, logs the error and closes the
        nio session.

        If an exception occurs during the login attempt or the message sending,
        logs the error and re-raises it.

        Args:
            message (str): The message to send to the room.
        """
        await self.ensure_logged_in()

        if not self.logged_in:
            logger.error("Unable to send message, login failed")
            return

        try:
            # Convert message to markdown
            html = markdown.markdown(message)

            # Send markdown formatted message
            await self.client.room_send(
                room_id=self.config['room_id'],
                message_type="m.room.message",
                content={
                    "msgtype": "m.text",
                    "body": message,
                    "format": "org.matrix.custom.html",
                    "formatted_body": html
                }
            )
            logger.info("Markdown message sent")
        except Exception as e:
            logger.error(f"Exception during sending markdown message: {e}")
            raise

    async def close(self):
        """
        Log out from the Matrix homeserver and close the client.

        If the bot is logged in, attempts to log out using the provided
        password. If the login attempt fails, logs the error and closes the
        nio session.

        If an exception occurs during the login attempt or the message sending,
        logs the error and re-raises it.
        """
        if self.logged_in:
            try:
                await self.client.logout()
                self.logged_in = False
                logger.info(f"Logged out from {self.config['homeserver']}")
            except Exception as e:
                logger.error(f"Exception during logout: {e}")
            finally:
                await self.client.close()  # Ensure the client is closed