diff --git a/main.py b/main.py index 561ea80..06a7419 100644 --- a/main.py +++ b/main.py @@ -5,11 +5,16 @@ import asyncio import traceback import signal +from pathlib import Path import config import util -import bot -import logic +#import bot +#import logic + +from new_bot import MatrixBot +from bot_types import MatrixBotConfig + from datatypes import AppConfig async def main() -> None: @@ -29,18 +34,17 @@ async def main() -> None: if cfg is None: util.log_error("Could't load config") return - - # start the bot - if not await bot.start(cfg): - util.log_error("Could't start the bot") - return - await logic.setup(bot.get_client()) - # wait for stop + matrix_bot_config = MatrixBotConfig( + matrix_homeserver_url=cfg.matrix_homeserver, + matrix_username_localpart=cfg.matrix_user, + storage_directory=Path(cfg.store_dir) + ) + bot = MatrixBot(matrix_bot_config) + bot.start() + await util.get_app_stop_event().wait() - - # stop - await logic.stop() + await bot.stop() if __name__ == "__main__": diff --git a/new_bot.py b/new_bot.py index 880f562..5896769 100644 --- a/new_bot.py +++ b/new_bot.py @@ -2,6 +2,7 @@ import asyncio import aiofiles import aioconsole import traceback +import logging import time import json import os @@ -93,6 +94,31 @@ class MatrixBot: localpart = config.matrix_username_localpart return f"@{localpart}:{homeserver_name}" + @staticmethod + async def _wait_for_task_and_stop_event(payload_task: asyncio.Task, stop_wait_task: asyncio.Task, print_exc: bool = False) -> Any: + """Waits for `payload_task` to finish. Uses `stop_wait_task` to check if payload task + should be cancelled. Returns value returned by `payload_task` task. Raises exception + raised by `payload_task` task. + + If `stop_wait_task` finishes, then `asyncio.CancelledError` is raised. + """ + try: + done, _ = await asyncio.wait( + [payload_task, stop_wait_task], + return_when=asyncio.FIRST_COMPLETED + ) + if stop_wait_task in done: + payload_task.cancel() + await payload_task + raise asyncio.CancelledError() + return payload_task.result() + except asyncio.CancelledError: + raise + except: + if print_exc: + traceback.print_exc() + raise + @staticmethod async def _default_password_callback() -> str: """Gets password from `MATRIX_PASSWORD` envvar if it is set. Asks @@ -101,13 +127,14 @@ class MatrixBot: return os.environ["MATRIX_PASSWORD"] print("--- A password is required (btw you might use MATRIX_PASSWORD envvar) ---") return await aioconsole.ainput("Password: ") - + async def _write_next_batch(self, next_batch: str) -> None: """Writes `next_batch` value to disk.""" path = self._config.storage_directory / "next_batch" async with aiofiles.open(path, "w") as f: await f.write(next_batch) + self._logger.debug("next_batch value is written to the disk") async def _read_next_batch(self) -> str | None: """Reads `next_batch` value from disk. Returns None if file does not exist.""" @@ -126,6 +153,7 @@ class MatrixBot: } async with aiofiles.open(path, "w") as f: await f.write(json.dumps(data, indent=4)) + self._logger.debug("Session data is writter to the disk") async def _read_session_data(self) -> dict[str, Any] | None: """Read session data from disk.""" @@ -143,10 +171,11 @@ class MatrixBot: """This callback is called when AsyncClient syncs with the server""" current_time = time.time() delta_time = current_time - self._last_next_batch_dump + self._last_next_batch = response.next_batch if delta_time >= self.NEXT_BATCH_DUMP_PERIOD: self._last_next_batch_dump = current_time try: - await self._write_next_batch(response.next_batch) + await self._write_next_batch(self._last_next_batch) except: traceback.print_exc() @@ -159,6 +188,7 @@ class MatrixBot: async def _client_login_session_data(self, session_data: dict[str, Any]) -> None: """Login using session data. Raises and exception on failure.""" + self._logger.debug("Using stored session data to log in") # build user id username = self._build_matrix_username(self._config) self._client.restore_login( @@ -167,25 +197,32 @@ class MatrixBot: ) result = await self._client.whoami() if type(result) is WhoamiError: + self._logger.error(f"Can't log in using stored session data: '{result.message}'") raise RuntimeError(result.message) elif type(result) is not WhoamiResponse: + self._logger.error("Can't log in using stored session data, unknown error") raise RuntimeError("Unknown response for whoami request") + self._logger.debug("Logged in using stored session data") async def _client_login_password(self) -> None: """Login using password and save result to disk on success. Raises an exception on failure. """ + self._logger.debug("Using password to log in") # get the password password = await self._cb_password() result = await self._client.login(password=password) if type(result) is LoginResponse: + self._logger.debug("Logged in using password") await self._write_session_data( access_token=result.access_token, device_id=result.device_id ) elif type(result) is LoginError: + self._logger.error(f"Can't log in using password: '{result.message}'") raise RuntimeError(result.message) else: + self._logger.error(f"Can't log in using password, unknown error") raise RuntimeError("Unknown login result") async def _client_login(self) -> None: @@ -194,31 +231,68 @@ class MatrixBot: session_data = await self._read_session_data() # session data is present, try to log in if session_data is not None: + self._logger.debug("Some session data found on the disk") try: await self._client_login_session_data(session_data) return + except asyncio.CancelledError: + raise except: pass # no session data - login using password try: + self._logger.debug("No session data found on the disk OR invalid data") await self._client_login_password() return + except asyncio.CancelledError: + raise except: pass # can't login + self._logger.error("Can't log in using available methods") raise RuntimeError("All login methods have failed, can't continue") + async def _client_destroy(self) -> None: + """Gracefully destroys the client.""" + try: + self._logger.debug("Closing the client") + await self._client.close() + if self._last_next_batch is not None: + self._logger.debug("Saving next_batch") + await self._write_next_batch(self._last_next_batch) + except: + traceback.print_exc() + async def _background_coroutine(self) -> None: """This function implements bot lifecycle.""" # we should stop when this task stops + self._logger.debug("_background_coroutine is started") stop_wait_task = asyncio.create_task(self._stop_event.wait()) self._client = self._build_client(self._config) self._setup_client_callbacks() # perform login + login_task = asyncio.create_task(self._client_login()) try: - await self._client_login() + await self._wait_for_task_and_stop_event(login_task, stop_wait_task) + except asyncio.CancelledError: + self._logger.debug("Background task is cancelled during login") + await self._client_destroy() + return except: traceback.print_exc() + self._logger.info("Succesfully logged in") + # test + wait_task = asyncio.create_task(asyncio.sleep(100000)) + try: + self._logger.info("Bot is not implemented yet, sleeping forever") + await self._wait_for_task_and_stop_event(wait_task, stop_wait_task) + except asyncio.CancelledError: + self._logger.info("Background task is cancelled during eternal sleep") + await self._client_destroy() + return + except: + traceback.print_exc() + # # PUBLIC @@ -228,11 +302,15 @@ class MatrixBot: self._validate_bot_config(config) # may raise an Exception # save the config self._config = config + # create the logger + self._logger = logging.getLogger(self._build_matrix_username(config)) + self._logger.setLevel(logging.DEBUG) # prepare some private data self._background_task: asyncio.Task | None = None self._client: AsyncClient = None # type: ignore self._last_next_batch_dump: float = 0.0 + self._last_next_batch: str | None = None self._cb_password = self._default_password_callback def start(self) -> None: @@ -249,22 +327,14 @@ class MatrixBot: self._background_coroutine() ) - def stop(self) -> bool: - """Stop the bot. - Signals the background task to stop. Returns True on success (the - signal is sent), False on failure (no background task). This - function won't wait for the task to stop. Use `wait_stop` to wait. - """ - if self._background_task is None: - return False - self._stop_event.set() - return True - - async def wait_stop(self) -> None: - """Wait for the background task to stop. - You must call `stop()` by yourself. This function will never return - otherwise. It just waits for the background task stop. - """ + async def stop(self) -> None: + """Stop the bot and wait for the bot stop.""" if self._background_task is None: return - await asyncio.wait([self._background_task]) \ No newline at end of file + self._stop_event.set() + try: + await self._background_task + except: + traceback.print_exc() + self._stop_event = None + self._background_task = None \ No newline at end of file