MatrixGPT/matrix_gpt/matrix.py

import json
import logging
import os
from pathlib import Path
from typing import Union

from nio import AsyncClient, AsyncClientConfig, LoginError
from nio import LoginResponse

logger = logging.getLogger('MatrixGPT')


class MatrixNioGPTHelper:
    """
    A simple wrapper class for common matrix-nio actions.
    """
    client = None

    # Encryption is disabled because it's handled by Pantalaimon.
    client_config = AsyncClientConfig(max_limit_exceeded=0, max_timeouts=0, store_sync_tokens=True, encryption_enabled=False)

    def __init__(self, auth_file: Union[Path, str], user_id: str, passwd: str, homeserver: str, store_path: str, device_name: str = 'MatrixGPT', device_id: str = None):
        self.auth_file = auth_file
        self.user_id = user_id
        self.passwd = passwd

        self.homeserver = homeserver
        if not (self.homeserver.startswith("https://") or self.homeserver.startswith("http://")):
            self.homeserver = "https://" + self.homeserver

        self.store_path = store_path
        Path(self.store_path).mkdir(parents=True, exist_ok=True)

        self.device_name = device_name
        self.client = AsyncClient(homeserver=self.homeserver, user=self.user_id, config=self.client_config, device_id=device_id)

    async def login(self) -> tuple[bool, LoginError] | tuple[bool, LoginResponse | None]:
        try:
            # If there are no previously-saved credentials, we'll use the password.
            if not os.path.exists(self.auth_file):
                logger.info('Using username/password.')
                resp = await self.client.login(self.passwd, device_name=self.device_name)

                # check that we logged in successfully.
                if isinstance(resp, LoginResponse):
                    self.write_details_to_disk(resp)
                    return True, resp
                else:
                    return False, resp
            else:
                # Otherwise the config file exists, so we'll use the stored credentials.
                logger.info('Using cached credentials.')
                with open(self.auth_file, "r") as f:
                    config = json.load(f)
                    client = AsyncClient(config["homeserver"])
                    client.access_token = config["access_token"]
                    client.user_id = config["user_id"]
                    client.device_id = config["device_id"]
                resp = await self.client.login(self.passwd, device_name=self.device_name)
                if isinstance(resp, LoginResponse):
                    self.write_details_to_disk(resp)
                    return True, resp
                else:
                    return False, resp
        except Exception:
            return False, None

    def write_details_to_disk(self, resp: LoginResponse) -> None:
        """Writes the required login details to disk so we can log in later without
        using a password.

        Arguments:
            resp {LoginResponse} -- the successful client login response.
            homeserver -- URL of homeserver, e.g. "https://matrix.example.org"
        """
        with open(self.auth_file, "w") as f:
            json.dump({"homeserver": self.homeserver,  # e.g. "https://matrix.example.org"
                       "user_id": resp.user_id,  # e.g. "@user:example.org"
                       "device_id": resp.device_id,  # device ID, 10 uppercase letters
                       "access_token": resp.access_token,  # cryptogr. access token
                       }, f, )