From e5a3533f4bfe944e8f159ed5aeaca7af6f6f3148 Mon Sep 17 00:00:00 2001 From: "Nikita Tyukalov, ASUS, Linux" Date: Mon, 24 Aug 2026 21:35:02 +0300 Subject: [PATCH] Source files moved from 2026-matrix-csonac --- pyproject.toml | 33 +++ src/mab/__init__.py | 18 ++ src/mab/bot.py | 512 ++++++++++++++++++++++++++++++++++++++++++++ src/mab/types.py | 23 ++ 4 files changed, 586 insertions(+) create mode 100644 src/mab/__init__.py create mode 100644 src/mab/bot.py create mode 100644 src/mab/types.py diff --git a/pyproject.toml b/pyproject.toml index ec9418a..50af0e4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,6 +15,39 @@ classifiers = [ "Programming Language :: Python :: 3", "Operating System :: OS Independent", ] +dependencies = [ + "aioconsole==0.8.2", + "aiofiles==25.1.0", + "aiohappyeyeballs==2.7.1", + "aiohttp==3.14.3", + "aiohttp_socks==0.12.0", + "aiosignal==1.4.0", + "atomicwrites==1.4.1", + "attrs==26.1.0", + "build==1.5.0", + "cachetools==7.1.7", + "frozenlist==1.8.0", + "h11==0.16.0", + "h2==4.4.1", + "hpack==4.2.0", + "hyperframe==6.1.0", + "idna==3.19", + "jsonschema==4.26.0", + "jsonschema-specifications==2025.9.1", + "matrix-nio==0.26.0", + "multidict==6.7.1", + "packaging==26.3", + "peewee==3.19.0", + "propcache==0.5.2", + "pycryptodome==3.23.0", + "pyproject_hooks==1.2.0", + "python-socks==3.0.0", + "referencing==0.37.0", + "rpds-py==2026.6.3", + "unpaddedbase64==2.1.0", + "vodozemac==0.10.0", + "yarl==1.24.5", +] [project.urls] Homepage = "https://git.tyukalov.su/nikita/mab" diff --git a/src/mab/__init__.py b/src/mab/__init__.py new file mode 100644 index 0000000..369c463 --- /dev/null +++ b/src/mab/__init__.py @@ -0,0 +1,18 @@ +from . import bot +from . import types + +from .types import MatrixBotConfig + +from .bot import MatrixBot + +__all__ = [ + # module names + "bot", + "types", + + # .types + "MatrixBotConfig", + + # .bot + "MatrixBot" +] \ No newline at end of file diff --git a/src/mab/bot.py b/src/mab/bot.py new file mode 100644 index 0000000..3f06ea5 --- /dev/null +++ b/src/mab/bot.py @@ -0,0 +1,512 @@ +import asyncio +import aiofiles +import aioconsole +import traceback +import logging +import time +import json +import os +import re +from html.parser import HTMLParser +from urllib.parse import urlparse +from typing import Any, Callable, Awaitable + +from nio import AsyncClient, AsyncClientConfig, SyncResponse +from nio import LoginResponse, LoginError, WhoamiResponse, WhoamiError + +from nio import RoomSendResponse, RoomSendError + +from nio import OlmUnverifiedDeviceError + +from nio import MatrixInvitedRoom, InviteMemberEvent +from nio import JoinResponse + +import nio.events + +from .types import * + + +class MatrixBot: + """Asynchronous Matrix Bot Implementation. + + Use objects of this class to build your bots. Manage the event loop + by yourself. + """ + NEXT_BATCH_DUMP_PERIOD = 120.0 + MATRIX_SYNC_PERIOD = 5000 + + # + # PRIVATE + # + @staticmethod + def _validate_matrix_homeserver_url(url: str) -> None: + """Checks if `url` is a valid matrix homeserver URL. + Raises an Exception if it is not. + """ + parsed = urlparse(url) + if parsed.scheme not in ("http", "https"): + raise RuntimeError( + f"Scheme {parsed.scheme} is not a valid scheme for matrix homeserver URL" + ) + if not parsed.netloc: + raise RuntimeError(f"{url} is not a valid matrix homeserver URL") + if parsed.path != "": + raise RuntimeError(f"{url} must have empty path (remove `{parsed.path}` after the hostname)") + + @staticmethod + def _validate_matrix_username_localpart(localpart: str) -> None: + """Checks if `username` is a valid localpart of matrix username. + Raises an Exception if it is not. + """ + pattern = r"^[a-z0-9._=\-]+$" + if not bool(re.fullmatch(pattern, localpart)) or len(localpart) > 255: + raise RuntimeError(f"{localpart} is not a valid matrix username localpart") + + @staticmethod + def _validate_storage_directory(path: Path) -> None: + """Checks if `path` is a valid storage directory and creates it. + Raises an Exception if it is not. + """ + path.mkdir(parents=True, exist_ok=True) + if not path.is_dir(): + raise RuntimeError(f"Could not create directory {path}") + + @staticmethod + def _validate_bot_config(config: MatrixBotConfig) -> None: + """Checks if `config` has errors. + Raises an Exception if it does. + """ + MatrixBot._validate_matrix_homeserver_url(config.matrix_homeserver_url) + MatrixBot._validate_matrix_username_localpart(config.matrix_username_localpart) + MatrixBot._validate_storage_directory(config.storage_directory) + + @staticmethod + def _build_client(config: MatrixBotConfig) -> AsyncClient: + """Builds `nio.AsyncClient` from `MatrixBotConfig`""" + # create the config for the client + client_config = AsyncClientConfig( + store_name="nio_store_file", + encryption_enabled=True, + store_sync_tokens=False + ) + # create the client + client = AsyncClient( + homeserver=config.matrix_homeserver_url, + user=config.matrix_username_localpart, + store_path=str(config.storage_directory), + config=client_config + ) + return client + + @staticmethod + def _build_matrix_username(config: MatrixBotConfig) -> str: + """Builds full matrix username.""" + homeserver_name = urlparse(config.matrix_homeserver_url).hostname + localpart = config.matrix_username_localpart + return f"@{localpart}:{homeserver_name}" + + @staticmethod + def _check_and_remove_html(text_to_check: str) -> tuple[bool, str]: + """Checks if `text_to_check` is HTML and sanitizes it. + + Returns: + tuple[bool, str] where the first element is True if `text_to_check` contains + valid HTML, and the second element is text without HTML (or just copy of + `text_to_check` if it does not contain HTML) + """ + has_tags = False + text_fragments = [] + class Extractor(HTMLParser): + def handle_starttag(self, tag, attrs): + nonlocal has_tags + has_tags = True + def handle_data(self, data): + text_fragments.append(data) + parser = Extractor(convert_charrefs=True) + parser.feed(text_to_check) + try: + if has_tags: + return (True, " ".join("".join(text_fragments).split())) + except: + traceback.print_exc() + return (False, text_to_check) + + @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 + the user for the password otherwise.""" + if "MATRIX_PASSWORD" in os.environ: + return os.environ["MATRIX_PASSWORD"] + print("--- A password is required (btw you might use MATRIX_PASSWORD envvar) ---") + return await aioconsole.ainput("Password: ") + + @staticmethod + async def _debug_event_callback(*args, **kwargs) -> None: + """Just prints types of arguments""" + try: + print(f"_debug_event_callback ({len(args)} args, {len(kwargs)} kwargs)") + for a in args: + print(f" - {type(a)}") + for k in kwargs: + print(f" * {k} = {kwargs[k]}") + except: + traceback.print_exc() + + + 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.""" + path = self._config.storage_directory / "next_batch" + if not path.is_file(): + return None + async with aiofiles.open(path, "r") as f: + return (await f.read()).strip() + + async def _write_session_data(self, *, access_token: str, device_id: str) -> None: + """Write session data to disk.""" + path = self._config.storage_directory / "session_data.json" + data = { + "access_token": access_token, + "device_id": device_id + } + 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.""" + path = self._config.storage_directory / "session_data.json" + if not path.is_file(): + return None + async with aiofiles.open(path, "r") as f: + j = json.loads(await f.read()) + return j + + # + # CALLBACKS + # + async def _callback_sync(self, response: SyncResponse) -> None: + """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(self._last_next_batch) + except: + traceback.print_exc() + + async def _callback_autojoin(self, room: MatrixInvitedRoom, event: InviteMemberEvent): + try: + # event.state_key must be our username + if event.state_key != self._client.user_id: + return + # membership status must be invite + if event.membership != "invite": + return + result = await self._client.join(room.room_id) + if type(result) is JoinResponse: + self._logger.info(f"Autojoined the room {room.room_id}") + else: + self._logger.error(f"Can't autojoin the room {room.room_id}") + except: + self._logger.error(traceback.format_exc()) + + # + # LIFECYCLE + # + def _setup_client_callbacks(self) -> None: + """Setup internal client callbacks""" + self._client.add_response_callback(self._callback_sync, SyncResponse) # type: ignore + + if self._config.auto_join_any_room_on_invite: + self._client.add_event_callback(self._callback_autojoin, InviteMemberEvent) # type: ignore + + 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( + user_id=username, + **session_data + ) + 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: + """This function logs in.""" + # check if we have session data stored on the disk + 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 _client_cancellable_sync_forever(self, *args, **kwargs) -> Any: + """Behaves exactly like AsyncClient.sync_forever, but supports task cancellation""" + sync_forever_task = asyncio.create_task( + self._client.sync_forever(*args, **kwargs) + ) + try: + return await sync_forever_task + except asyncio.CancelledError: + try: + self._client.stop_sync_forever() + except Exception: + pass + finally: + sync_forever_task.cancel() + await asyncio.gather(sync_forever_task, return_exceptions=True) + raise + + 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()) # type: ignore + # perform login + login_task = asyncio.create_task(self._client_login()) + try: + 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") + # sync forever + self._logger.info("Syncing forever") + sync_task = asyncio.create_task( + self._client_cancellable_sync_forever( + timeout=self.MATRIX_SYNC_PERIOD, + since=(await self._read_next_batch()) + ) + ) + try: + await self._wait_for_task_and_stop_event(sync_task, stop_wait_task) + except asyncio.CancelledError: + self._logger.debug("Sync task is cancelled") + await self._client_destroy() + return + except: + traceback.print_exc() + + + # + # PUBLIC + # + def __init__(self, config: MatrixBotConfig) -> None: + # check if config is valid + self._validate_bot_config(config) # may raise an Exception + # save the config + self._config: MatrixBotConfig = config + # create the logger + self._logger = logging.getLogger(self._build_matrix_username(config)) + self._logger.setLevel(logging.DEBUG) + # create the client + self._client: AsyncClient = self._build_client(self._config) + self._setup_client_callbacks() + + # prepare some private data + self._background_task: asyncio.Task | None = None + 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: + """Start the bot. + Starts the bot in background task. Raises an exception if there are + problems (for example, the bot is already started). The bot will + do everything to keep itself running, including restarts. Use + `stop()` to stop the bot. + """ + if self._background_task is not None: + raise RuntimeError("The bot is already started!") + self._stop_event = asyncio.Event() + self._background_task = asyncio.create_task( + self._background_coroutine() + ) + + async def stop(self) -> None: + """Stop the bot and wait for the bot stop.""" + if self._background_task is None or self._stop_event is None: + return + self._stop_event.set() + try: + await self._background_task + except: + traceback.print_exc() + self._stop_event = None + self._background_task = None + + def verify_all_known_devices(self) -> bool: + """Verifies all known devices. + + Returns: + True if there were unverified devices that are verified now. + """ + result = False + for user_id in self._client.device_store.users: + for device_id, olm_device in self._client.device_store[user_id].items(): + # can't trust ourselves + if device_id == self._client.device_id and user_id == self._client.user_id: + continue + # they are already verified + if olm_device.verified: + continue + # verify them + self._client.verify_device(olm_device) + result = True + return result + + def add_event_callback(self, callback: Callable[[Any], Awaitable[None]] | None, event_class: nio.events.Event) -> None: + """Added event callback for events of specified class. + Use `None` instead of callback to print parameter types you need to use in your callback.""" + if callback is None: + callback = self._debug_event_callback + self._client.add_event_callback(callback, event_class) # type: ignore + + def get_client(self) -> AsyncClient: + """Get AsyncClient in use""" + return self._client + + async def send_text_to_room(self, room_id: str, text: str, is_html: bool | None = None, **kwargs) -> str: + """Sends a text message to the room and handle HTML as specified. + + Args: + text - text to send + is_html - True if text is HTML; False if text is not HTML; None if the value should be guessed + kwargs - passed as `m.room.message` content keys + + Returns: + event_id of the message on success. Raises an exception on error. + """ + original_text = text + # guess if text is HTML + if is_html is None: + is_html, text = self._check_and_remove_html(text) + # text is HTML + elif is_html: + _, text = self._check_and_remove_html(text) + # create `content` for `room_send()` + content = { + "msgtype": "m.text", + "body": text, + **kwargs + } + if is_html: + content["format"] = "org.matrix.custom.html" + content["formatted_body"] = original_text + # try to send the message + try: + result = await self._client.room_send( + room_id=room_id, + message_type="m.room.message", + content=content + ) + except OlmUnverifiedDeviceError: + if self._config.auto_verify_all_known_devices: + if not self.verify_all_known_devices(): + raise + return await self.send_text_to_room(room_id, original_text, is_html, **kwargs) + else: + raise + # success + if type(result) is RoomSendResponse: + return result.event_id + # error + elif type(result) is RoomSendError: + raise RuntimeError(result) + # unknown error + else: + raise RuntimeError("Unknown error has occured", result) \ No newline at end of file diff --git a/src/mab/types.py b/src/mab/types.py new file mode 100644 index 0000000..333264e --- /dev/null +++ b/src/mab/types.py @@ -0,0 +1,23 @@ +"""Data types required for the bot""" + +from pathlib import Path +from dataclasses import dataclass + +@dataclass +class MatrixBotConfig: + """Configuration for MatrixBot""" + + matrix_homeserver_url: str + """Homeserver, for example: `https://matrix.domain.su`""" + + matrix_username_localpart: str + """Localpart of matrix username (without homeserver), for example: `valid-username`""" + + storage_directory: Path + """Path to the storage directory (will be created if needed)""" + + auto_verify_all_known_devices: bool = True + """Enable automatic verification of all known devices on related errors""" + + auto_join_any_room_on_invite: bool = True + """Enable automatic join for invites (implemented by automatically adding internal callback on bot creation)""" \ No newline at end of file