2026-matrix-csonac/bot.py

"""
matrix-nio basics wrapper
"""

import asyncio
import time
import traceback
from pathlib import Path

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

from datatypes import AppConfig
import util



#
# DATA
#
_app_config: AppConfig
_client: AsyncClient
_task: asyncio.Task | None = None
_stop: asyncio.Event | None = None
_since: str | None = None
_last_since_save_time: float = 0



#
# CALLBACKS
#
async def _sync_callback(response: SyncResponse):
    global _since, _last_since_save_time
    t = time.time()
    _since = response.next_batch
    if t - _last_since_save_time >= 120.0:
        await util.set_next_batch(_app_config, _since)
        _last_since_save_time = t



#
# PRIVATE
#
async def _bot_login_using_token(token: str, device_id: str) -> bool:
    """Tries to login using access_token. Returns True on success."""
    util.log_info("Authorizing using access_token...")
    _client.restore_login(
        user_id=f"@{_app_config.matrix_user}:{util.get_hostname_from_url(_app_config.matrix_homeserver)}",
        device_id=device_id,
        access_token=token
    )
    result = await _client.whoami()
    if type(result) is not WhoamiResponse:
        return False
    util.log_info(f"Logged in as {result.user_id} using access_token")
    return True

async def _bot_login_using_password(password: str) -> tuple[str, str] | tuple[None, None]:
    """Tries to login using password. Returns (access_token, device_id) on success."""
    util.log_info("Authorizing using password...")
    result = await _client.login(password=password)
    if type(result) is LoginResponse:
        util.log_info(f"Authorized using password")
        return result.access_token, result.device_id
    elif type(result) is LoginError:
        util.log_error(f"Failed to authorize using password: {result.message}")
        return None, None
    else:
        raise RuntimeError(f"Invalid login result: {result}")

async def _bot_login(config: AppConfig) -> bool:
    """Tries to login"""
    try:
        # get the session token and try to use it
        session_token, device_id = await util.get_session_data(config)
        if session_token is not None and device_id is not None:
            if await _bot_login_using_token(session_token, device_id):
                return True
            await util.set_session_data(config, None)
            util.log_warning("Existing access_token is deleted")
        # get the password and try to use it
        password = await util.get_password()
        if password is None:
            util.log_error("No password provided (consider using MATRIX_PASSWORD environment variable)")
            return False
        access_token, device_id = await _bot_login_using_password(password)
        if access_token is not None and device_id is not None:
            await util.set_session_data(config, (access_token, device_id))
            util.log_warning("Saved new access_token and device_id")
            return True
        # can't login
        return False
    except asyncio.CancelledError:
        return False
    except:
        traceback.print_exc()
        return False

async def _bot_loop(config: AppConfig) -> None:
    """Bot loop"""
    global _client, _since
    # app stop task
    if _stop is None:
        raise RuntimeError("_stop can't be None")
    stop_task = asyncio.create_task(_stop.wait())
    # setup the callback for syncing
    _client.add_response_callback(_sync_callback, SyncResponse) # type: ignore
    # login
    login_task = asyncio.create_task(_bot_login(config))
    done, _ = await asyncio.wait(
        [login_task, stop_task],
        return_when=asyncio.FIRST_COMPLETED
    )
    # stopped
    if stop_task in done:
        login_task.cancel()
        return
    # failed to login
    if login_task.exception() or not login_task.result():
        util.request_app_stop("Can't authorize into matrix")
        return
    # load initial `next_batch`
    _since = await util.get_next_batch(config)
    # sync forever
    while True:
        # sync
        sync_task = asyncio.create_task(_client.sync_forever(timeout=5000, since=_since))
        done, _ = await asyncio.wait(
            [sync_task, stop_task],
            return_when=asyncio.FIRST_COMPLETED
        )
        # stopped
        if stop_task in done:
            util.log_info("Stopping sync_forever...")
            _client.stop_sync_forever()
            util.log_info("Waiting for sync_forever to quit...")
            await sync_task
            sync_task.cancel()
            break
        # something happened
        try:
            sync_task.result()
        except:
            traceback.print_exc()
        await asyncio.sleep(1)



#
# PUBLIC
#
async def start(config: AppConfig) -> bool:
    """Starts the bot"""
    global _client
    global _task, _stop, _app_config
    if _task is not None:
        return False
    _app_config = config
    # create the bot
    store_dir = Path.cwd() / config.store_dir
    store_dir.mkdir(parents=True, exist_ok=True)
    client_config = AsyncClientConfig(
        store_name="storefile",
        encryption_enabled=True,
        store_sync_tokens=False
    )
    _client = AsyncClient(
        homeserver=config.matrix_homeserver,
        user=config.matrix_user,
        store_path=str(store_dir),
        config=client_config
    )
    _stop = asyncio.Event()
    _task = asyncio.create_task(_bot_loop(config))
    return True

def get_client() -> AsyncClient:
    return _client

async def stop() -> None:
    """Stop the bot"""
    global _task, _stop
    if _task is None or _stop is None:
        return
    _stop.set()
    try:
        await _task
    except asyncio.CancelledError:
        pass
    except:
        traceback.print_exc()
    _task = None
    _stop = None
    try:
        await _client.close()
    except:
        pass
    await util.set_next_batch(_app_config, _since)