""" 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)