From 1ba1094a42ad2e029cf776de2d37f919c2faa0ac Mon Sep 17 00:00:00 2001 From: nikita Date: Wed, 9 Sep 2026 00:33:26 +0300 Subject: [PATCH] Refactoring --- src/mab/bot.py | 827 -------------------------------- src/mab/bot/__init__.py | 5 + src/mab/bot/_callbacks.py | 171 +++++++ src/mab/bot/_client_auth.py | 137 ++++++ src/mab/bot/_client_manager.py | 202 ++++++++ src/mab/bot/_client_sender.py | 300 ++++++++++++ src/mab/bot/_client_uploader.py | 187 ++++++++ src/mab/bot/_storage.py | 103 ++++ src/mab/bot/_validation.py | 84 ++++ src/mab/bot/bot.py | 198 ++++++++ src/mab/types.py | 59 ++- src/mab/utils.py | 37 +- 12 files changed, 1480 insertions(+), 830 deletions(-) delete mode 100644 src/mab/bot.py create mode 100644 src/mab/bot/__init__.py create mode 100644 src/mab/bot/_callbacks.py create mode 100644 src/mab/bot/_client_auth.py create mode 100644 src/mab/bot/_client_manager.py create mode 100644 src/mab/bot/_client_sender.py create mode 100644 src/mab/bot/_client_uploader.py create mode 100644 src/mab/bot/_storage.py create mode 100644 src/mab/bot/_validation.py create mode 100644 src/mab/bot/bot.py diff --git a/src/mab/bot.py b/src/mab/bot.py deleted file mode 100644 index 385662a..0000000 --- a/src/mab/bot.py +++ /dev/null @@ -1,827 +0,0 @@ -import asyncio -import aiofiles, aiofiles.os -import aioconsole -import traceback -import logging -import time -import json -import os -import re -import io -import magic -from html.parser import HTMLParser -from urllib.parse import urlparse -from typing import Any, Callable, Awaitable - -from PIL import Image - -from nio import AsyncClient, AsyncClientConfig, SyncResponse -from nio import LoginResponse, LoginError, WhoamiResponse, WhoamiError - -from nio import RoomSendResponse, RoomSendError -from nio import UploadResponse, UploadError - -from nio import OlmUnverifiedDeviceError - -from nio import MatrixInvitedRoom, InviteMemberEvent -from nio import JoinResponse - -from .filters.base import BaseEventFilter - -from nio.events.room_events import Event as RoomEvemt - -from .types import * -from .utils import Utils - - -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 _process_html_text(text: str, is_html: bool | None) -> dict: - """Processes `text` and `is_html` arguments of `send_*` functions and - prepares `content` keys related to text data inside the message. - - Args: - text - `text` as passed to the `send_*` function - is_html - `is_html` as passed to the `send_*` function - - Returns: - `dict` that contains keys required for the text message to be - displayed correctly. - """ - # text w/ and w/o HTML - text_w_html: str | None = None - text_wo_html: str | None = None - # check if HTML is present - if is_html or is_html is None: - # will be True if `text` contains HTML tags - has_tags = False - # text fragments without tags - text_fragments = [] - # temporary extractor class - class Extractor(HTMLParser): - def handle_starttag(self, tag, attrs): - nonlocal has_tags - has_tags = True - def handle_data(self, data): - text_fragments.append(data) - # parse - parser = Extractor(convert_charrefs=True) - try: - parser.feed(text) - except: - traceback.print_exc() - text_wo_html = text - is_html = has_tags - if is_html: - text_wo_html = " ".join(text_fragments) - text_w_html = text - if not is_html: - text_wo_html = text - # return the content - if text_w_html: - return { - "body": text_wo_html, - "format": "org.matrix.custom.html", - "formatted_body": text_w_html - } - else: - return { - "body": text_wo_html - } - - @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 - - async def _upload_file(self, path: Path | str, mime_type: str, *, filename: str | None = None) -> tuple[UploadResponse, dict, int]: - """Uploads the file with encryption enabled. UploadError will be raised - as RuntimeError. - - Args: - path - path to the file to upload - mime_type - mime-type to use - filename - filename to use (base name of the `path` will be used if not set) - - Returns: - 1. UploadResponse as returned by `_upload_binary_data(...)` - 2. dict as returned by `_upload_binary_data(...)` - 3. file size in bytes - """ - # get basic file info - file_info = await aiofiles.os.stat(path) - if not filename: - filename = os.path.basename(path) - # upload - async with aiofiles.open(path, "r+b") as f: - response, keys = await self._upload_provider( - f, - mime_type, - filename=filename, - filesize=file_info.st_size - ) - return response, keys, file_info.st_size - - async def _upload_provider(self, provider: Any, mime_type: str, *, filename: str, filesize: int) -> tuple[UploadResponse, dict]: - """Uploads the file with encryption enabled. UploadError will be raised - as RuntimeError. - - Args: - provider - data provider as described in `AsyncClient.upload(...)` docs - mime_type - mime-type to use - filename - filename to use - - Returns: - 1. UploadResponse as returned by `AsyncClient.upload(...)` - 2. dict as returned by `AsyncClient.upload(...)` - """ - # upload - response, keys = await self._client.upload( - provider, - content_type=mime_type, - filename=filename, - encrypt=True, - monitor=None, - filesize=filesize - ) - if isinstance(response, UploadError): - raise RuntimeError(response.message, response) - if not isinstance(response, UploadResponse): - raise RuntimeError( - "Upload result must be `UploadRespond`, but it's not", - response - ) - return response, keys # type: ignore - - async def _upload_thumbnail(self, thumbnail: Path | str | bytes | None) -> dict: - """Upload thumbnail and return fields required to use it for video. - - Args: - thumbnail - path to the thumbnail file or raw `bytes` of the image. - You can actually pass `None`, empty dict will be returned - in this case. - - Returns: - `dict` you should embed into your `content["info"]` - """ - if thumbnail is None: - return {} - # raw data - raw_bytes_io: io.BytesIO | None = None - if type(thumbnail) is bytes: - thumb_mime = magic.from_buffer(thumbnail, mime=True) - raw_bytes_io = io.BytesIO(thumbnail) - # file path - else: - thumb_mime = magic.from_file(thumbnail, mime=True) - # check if invalid MIME - if not thumb_mime.startswith("image/"): - raise RuntimeError(f"Thumbnail MIME is not image: `{thumb_mime}`") - # get thumbnail size - with Image.open(raw_bytes_io or thumbnail) as f: - width, height = f.size - # seek BytesIO to zero so we can reuse it - if raw_bytes_io is not None: - raw_bytes_io.seek(0) - # upload - if type(thumbnail) is bytes: - thumb_size = len(thumbnail) - thumb_resp, thumb_keys = await self._upload_provider( - raw_bytes_io, - thumb_mime, - filename=f"preview.{thumb_mime.split("/")[-1]}", - filesize=thumb_size - ) - else: - thumb_resp, thumb_keys, thumb_size = await self._upload_file( - thumbnail, # type: ignore - mime_type=thumb_mime, - filename=os.path.basename(str(thumbnail)) - ) - # return the result - return { - "thumbnail_file": { - "url": thumb_resp.content_uri, - "mimetype": thumb_mime, - **thumb_keys - }, - "thumbnail_info": { - "w": width, - "h": height, - "size": thumb_size, - "mimetype": thumb_mime - } - } - - async def _send_content_to_room(self, room_id: str, content: dict) -> RoomSendResponse: - """Sends `m.room.message` event to the room. Returns RoomSendResponse - as returned by `AsyncClient.room_send(...)`. Automatically verifies - all devices if this feature is enabled. Return values other than - `RoomSendResponse` are raised as `RuntimeError`.""" - # 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_content_to_room(room_id, content) - else: - raise - # success - if type(result) is RoomSendResponse: - return result - # error - elif type(result) is RoomSendError: - raise RuntimeError(result) - # unknown error - else: - raise RuntimeError("Unknown error has occured", result) - - # - # 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()) - - async def _callback_filter_router(self, *args, **kwargs): - if len(args) != 2: - self._logger.debug("Can't process the event, not enough positional args") - await self._debug_event_callback(*args, **kwargs) - return - room = args[0] - event = args[1] - for filter in self._filters: - filter_object = filter[0] - filter_callback = filter[1] - filter_stop_after_this = filter[2] - if filter_object(room, event): - self._logger.debug(f"Filter {repr(filter_object)} matched") - try: - await filter_callback(room, event) - except: - self._logger.error(traceback.format_exc()) - if filter_stop_after_this: - self._logger.debug(f"Filter {repr(filter_object)} stops matching") - break - - # - # LIFECYCLE - # - def _setup_client_callbacks(self) -> None: - """Setup internal client callbacks""" - self._client.add_response_callback(self._callback_sync, SyncResponse) # type: ignore - self._client.add_event_callback(self._callback_filter_router, RoomEvemt) # 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 - ) - self._client.load_store() - 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()), - full_state=True - ) - ) - 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 - if type(config.storage_directory) is str: - config.storage_directory = Path(config.storage_directory) - 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 - self._filters = [] - - 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[..., Awaitable[None]] | None, filter: BaseEventFilter, stop_after_this: bool = True) -> None: - """Add event callback for events that pass the filter.""" - if callback is None: - callback = self._debug_event_callback - self._filters.append((filter, callback, stop_after_this)) - - 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) -> 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. - """ - text_data = self._process_html_text(text, is_html) - content = { - "msgtype": "m.text", - **text_data - } - return (await self._send_content_to_room(room_id, content)).event_id - - async def send_image_to_room(self, - room_id: str, - path: Path | str, - caption: str | None = None, - is_html: bool | None = None) -> str: - """Send an image to the room. - - This method sends specified image file to the specified room. It - will use mime-type deduced using `magic` module. Image dimension is - deduced using Pillow. The file is read asynchronously, so this - method won't block your event loop. - - Please note that HTML is incorrectly interpreted for images in - some clients (as of September 6th, 2026). - - Args: - room_id - room_id to send the image to - path - path to the image file - caption - caption to use (will be hidden if it is the same as file basename) - is_html - whether the caption is HTML (autodetected if None) - - Returns: - event_id - """ - # caption must not actually be empty - if caption is None or not caption.strip(): - caption = os.path.basename(path) - is_html = False - # check if the file is image - mime_type: str = magic.from_file(path, mime=True) - if not mime_type.startswith("image/"): - raise RuntimeError(f"File {path} has non-image mime-type") - # get image size - with Image.open(path) as image: - width, height = image.size - # upload - response, keys, filesize = await self._upload_file(path, mime_type) - # prepare the content - content = { - "msgtype": "m.image", - "filename": os.path.basename(path), - **self._process_html_text(caption, is_html), - "file": { - "url": response.content_uri, - "mimetype": mime_type, - **keys - }, - "info": { - "mimetype": mime_type, - "size": filesize, - "w": width, - "h": height - } - } - # send - return (await self._send_content_to_room(room_id, content)).event_id - - async def send_video_to_room(self, - room_id: str, - path: Path | str, - props: VideoFileProperties | None = None, - caption: str | None = None, - is_html: bool | None = None) -> str: - """Send a video to the room. - - This method sends specified video file to the specified room. It - will use mime-type deduced using `magic` module. Most video - properties must be passed using `props` argument. - - If you pass `None` as `props` then system `ffprobe` will be called - to get video dimensions and duration. You can disable this behavior - by disabling `auto_ffprobe` in `MatrixBotConfig`. - - If you pass `None` as `props.thumbnail_path` then thumbnail will be - generated automatically. You can disable this behavior by disabling - `auto_ffmpeg` in `MatrixBotConfig`. - - Please note that HTML is incorrectly interpreted for videos in - some clients (as of September 6th, 2026). - - Args: - room_id - room_id to send the video to - path - path to the video file - props - video properties - caption - caption to use (will be hidden if it is the same as file basename) - is_html - whether the caption is HTML (autodetected if None) - - Returns: - event_id - """ - # caption must not actually be empty - if caption is None or not caption.strip(): - caption = os.path.basename(path) - is_html = False - # check if the file is video - mime_type: str = magic.from_file(path, mime=True) - if not mime_type.startswith("video/"): - raise RuntimeError(f"File {path} has non-video mime-type") - # properties must not be empty - if props is None: - if not self._config.auto_ffprobe: - raise RuntimeError("No `props` are supplied and `auto_ffprobe` is disabled") - props = await Utils.get_video_props_using_ffprobe(path) - # thumbnail data to use - thumbnail_data = {} - # autogenerate thumbnail if allowed - if not props.thumbnail and self._config.auto_ffmpeg: - props.thumbnail = await Utils.get_video_thumbnail( - path, - timestamp=props.duration_ms / 1000 / 2, - max_size=400 - ) - # upload the thumbnail if specified - thumbnail_data = await self._upload_thumbnail(props.thumbnail) - # upload - response, keys, filesize = await self._upload_file(path, mime_type) - # prepare the content - content = { - "msgtype": "m.video", - "filename": os.path.basename(path), - **self._process_html_text(caption, is_html), - "file": { - "url": response.content_uri, - "mimetype": mime_type, - **keys - }, - "info": { - "duration": props.duration_ms, - "mimetype": mime_type, - "size": filesize, - "w": props.width_px, - "h": props.height_px, - **thumbnail_data - } - } - # send - return (await self._send_content_to_room(room_id, content)).event_id \ No newline at end of file diff --git a/src/mab/bot/__init__.py b/src/mab/bot/__init__.py new file mode 100644 index 0000000..761dcf4 --- /dev/null +++ b/src/mab/bot/__init__.py @@ -0,0 +1,5 @@ +from .bot import * + +__all__ = [ + "MatrixBot" +] \ No newline at end of file diff --git a/src/mab/bot/_callbacks.py b/src/mab/bot/_callbacks.py new file mode 100644 index 0000000..6b6ed90 --- /dev/null +++ b/src/mab/bot/_callbacks.py @@ -0,0 +1,171 @@ +import asyncio +import logging +import traceback +from typing import Callable, Any, Coroutine, TYPE_CHECKING +from dataclasses import dataclass + +from nio import AsyncClient +from nio import MatrixRoom +from nio import MatrixInvitedRoom, InviteMemberEvent, JoinResponse +from nio.events.room_events import Event as RoomEvent + +from ._storage import Storage +from ..types import MatrixBotConfig, RoomEventData +from ..filters.base import BaseEventFilter + +if TYPE_CHECKING: + from .bot import MatrixBot + +class Callbacks: + """ + This class implements user callbacks routing and some basic internal + callbacks. You should use it to add callbacks to the bot. + """ + # + # DATACLASSES + # + @dataclass + class _FilterBasedCallback: + """Data structure that holds information about filter-based callback""" + + filter: BaseEventFilter + """Filter to use for matching""" + + callback: Callable[[RoomEventData], Coroutine[Any, Any, None]] | None + """Callback that will be called if the filter matches""" + + stop_matching: bool + """Whether to stop matching if the filter matches""" + + # + # PRIVATE + # + async def _callback_room_event_router(self, + room: MatrixRoom, + event: RoomEvent) -> None: + """ + This callback implements routing of filter-based event callbacks. + """ + if self._config is None or self._client is None: + raise RuntimeError("Callbacks were not set up") + for callback_info in self._filters: + if not isinstance(callback_info, self._FilterBasedCallback): + continue + if not callback_info.filter(room, event): + continue + event_data = RoomEventData( + room=room, + event=event, + filter=callback_info.filter, + bot=self._matrix_bot + ) + try: + # dump argument types + if callback_info.callback is None: + self._logger.critical( + f"Filter {repr(callback_info.filter)} has no callback. " + f"`room` type is `{type(event_data.room)}`. " + f"`event` type is `{type(event_data.event)}`." + ) + continue + # create the coroutine + cor = callback_info.callback(event_data) + # multitasking + if self._config.enable_multitasking_callbacks: + t = asyncio.create_task(cor) + t.add_done_callback(self._multitasking_set.discard) + self._multitasking_set.add(t) + # no multitasking + else: + await cor + except asyncio.CancelledError: + raise + except: + self._logger.error(traceback.format_exc()) + if callback_info.stop_matching: + break + + async def _callback_autojoin(self, + room: MatrixInvitedRoom, + event: InviteMemberEvent) -> None: + """ + This callback implements automatic room joining. + """ + if self._client is None: + raise RuntimeError("Callbacks were not set up") + 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()) + + async def _setup_internal_callbacks(self) -> None: + """ + Setup internal callbacks. + """ + if self._config is None or self._client is None: + raise RuntimeError("Callbacks were not set up") + # filter-based callback routing + self._client.add_event_callback( + self._callback_room_event_router, RoomEvent) + # auto-join + if self._config.auto_join_any_room_on_invite: + self._client.add_event_callback( + self._callback_autojoin, InviteMemberEvent) # type: ignore + + # + # PUBLIC + # + def __init__(self, storage: Storage, matrix_bot: "MatrixBot"): + self._logger = logging.getLogger("Callbacks") + self._storage = storage + self._matrix_bot = matrix_bot + + self._config: MatrixBotConfig | None = None + self._client: AsyncClient | None = None + self._filters: list = [] + self._multitasking_set: set[asyncio.Task] = set() + + def add_room_event_callback( + self, + filter: BaseEventFilter, + callback: Callable[[RoomEventData], Coroutine[Any, Any, None]] | None, + *, + stop_matching: bool = True) -> None: + """ + Add filter-based callback. + + Args: + - filter - filter to use for the callback + - callback - the callback that must be called if the filter matches + - stop_matching - whether to stop matching if the filter matches + """ + callback_info = self._FilterBasedCallback( + filter=filter, + callback=callback, + stop_matching=stop_matching + ) + self._filters.append(callback_info) + + + async def setup(self, config: MatrixBotConfig, client: AsyncClient) -> None: + """ + Setup the callbacks. Must be called by `ClientManager`. + + This method must be called AFTER required callbacks were added by using + `add_room_event_callback(...)`. + """ + self._config = config + self._client = client + await self._setup_internal_callbacks() + + \ No newline at end of file diff --git a/src/mab/bot/_client_auth.py b/src/mab/bot/_client_auth.py new file mode 100644 index 0000000..448a6b9 --- /dev/null +++ b/src/mab/bot/_client_auth.py @@ -0,0 +1,137 @@ +import asyncio +import aioconsole +import logging +import os +from ._storage import Storage + +from ..utils import Utils +from ..types import MatrixBotConfig + +from nio import AsyncClient +from nio import WhoamiError, WhoamiResponse +from nio import LoginError, LoginResponse + +class ClientAuth: + """This class is used by to authenticate AsyncClient.""" + + # + # PRIVATE + # + @staticmethod + async def _default_password_callback() -> str: + if "MATRIX_PASSWORD" in os.environ: + return os.environ["MATRIX_PASSWORD"] + return await aioconsole.ainput("Matrix password: ") + + async def _login_using_session_data(self, client: AsyncClient) -> None: + """ + Login using session data from `Storage`. + + Args: + - client - client that needs to get logged in + + Returns: + - Does not return anything on success + - Raises an exception on error + """ + if self._full_matrix_username is None: + raise RuntimeError("ClientAuth is not set up") + self._logger.info("Trying to log in using session data from disk...") + session_data = await self._storage.read_session_data() + if session_data is None: + self._logger.error("Can't read session data from disk") + raise RuntimeError("Can't read session data from disk") + client.restore_login( + user_id=self._full_matrix_username, + access_token=session_data["access_token"], + device_id=session_data["device_id"] + ) + client.load_store() + result = await 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.info("Logged in using stored session data") + + async def _login_using_password(self, client: AsyncClient) -> None: + """ + Login using password. + + Args: + - client - client that needs to get logged in + + Returns: + - Does not return anything on success + - Raises an exception on error + """ + if self._full_matrix_username is None: + raise RuntimeError("ClientAuth is not set up") + self._logger.info("Trying to log in using password") + password = await self._default_password_callback() + result = await client.login(password=password) + if type(result) is LoginResponse: + self._logger.info("Logged in using password") + await self._storage.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") + + # + # PUBLIC + # + def __init__(self, storage: Storage): + self._logger = logging.getLogger("ClientAuth") + self._storage = storage + self._full_matrix_username: str | None = None + + async def setup(self, config: MatrixBotConfig) -> None: + """ + Setup `ClientAuth` object using `config`. + """ + self._full_matrix_username = Utils.build_full_matrix_username(config) + + async def login(self, client: AsyncClient) -> None: + """ + Login. Saves authorization data if it is not saved yet. Must be called + by `ClientManager`. + + Args: + - client - client that needs to get logged in + + Returns: + - Does not return anything on success + - Raises an exception on error + """ + # try session data + try: + await self._login_using_session_data(client) + return + except asyncio.CancelledError: + raise + except: + pass + # try password + try: + await self._login_using_password(client) + return + except asyncio.CancelledError: + raise + except: + pass + # failed + raise RuntimeError("Failed to login") \ No newline at end of file diff --git a/src/mab/bot/_client_manager.py b/src/mab/bot/_client_manager.py new file mode 100644 index 0000000..e37f4aa --- /dev/null +++ b/src/mab/bot/_client_manager.py @@ -0,0 +1,202 @@ +import asyncio +import logging +import time +import traceback +from typing import Any + +from nio import AsyncClient, AsyncClientConfig, SyncResponse +from ..types import MatrixBotConfig + +from ._storage import Storage +from ._client_auth import ClientAuth + +class ClientManager: + """This class manages basic of `AsyncClient`.""" + + # + # PRIVATE + # + async def _close_client(self) -> None: + """ + Closes the client correctly and saves data. + """ + if self._client is None: + raise RuntimeError("The client does not exist") + self._logger.debug("Closing the client...") + await self._client.close() + self._logger.debug("Dumping 'next_batch' value...") + if self._last_next_batch is not None: + await self._storage.write_next_batch(self._last_next_batch) + + async def _cancellable_sync_forever(self, *args, **kwargs) -> Any: + """ + Behaves exactly like AsyncClient.sync_forever, but supports proper task + cancellation. + """ + if self._client is None: + raise RuntimeError("Client does not exist") + task = asyncio.create_task(self._client.sync_forever(*args, **kwargs)) + try: + return await task + except asyncio.CancelledError: + try: + self._client.stop_sync_forever() + except Exception: + pass + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + raise + + async def _background_coroutine(self) -> None: + """ + Coroutine that manages the `AsyncClient` in background. It supports + cancellation and you should use it to stop the operation of the bot. + """ + # refuse to start without config + if self._config is None: + raise RuntimeError("Config is not set up") + # refuse to start without client + if self._client is None: + raise RuntimeError("Client does not exist") + # try to login + self._logger.info("Logging in...") + try: + await self._client_auth.login(self._client) + self._logger.info(f"Logged in as {self._client.user_id}") + except asyncio.CancelledError: + raise + except: + self._logger.error(traceback.format_exc()) + raise + # sync forever + self._logger.info("Syncing forever...") + await self._cancellable_sync_forever( + timeout=self._config.client_sync_period_ms, + since=(await self._storage.read_next_batch()), + full_state=True + ) + + async def _callback_sync(self, response: SyncResponse) -> None: + """ + This callback is called when AsyncClient syncs with the server + """ + if self._config is None: + raise RuntimeError("ClientManager is not set up") + # save last `next_batch` + self._last_next_batch = response.next_batch + # dump if needed + current_time = time.time() + delta_time = current_time - self._last_next_batch_write_time + if delta_time >= self._config.client_next_batch_dump_period_s: + try: + await self._storage.write_next_batch(self._last_next_batch) + self._last_next_batch_write_time = current_time + except: + self._logger.error(traceback.format_exc()) + + async def _setup_callbacks(self, config: MatrixBotConfig) -> None: + """ + Setup callbacks. + """ + if self._client is None: + raise RuntimeError("The client does not exist") + self._client.add_response_callback( + self._callback_sync, SyncResponse) # type: ignore + + # + # PUBLIC + # + def __init__(self, client_auth: ClientAuth, storage: Storage): + self._logger = logging.getLogger("ClientManager") + self._client_auth = client_auth + self._storage = storage + + self._config: MatrixBotConfig | None = None + self._client: AsyncClient | None = None + self._was_started_ever: bool = False + self._was_built_ever: bool = False + self._background_task: asyncio.Task | None = None + self._last_next_batch_write_time: float = 0.0 + self._last_next_batch: str | None = None + + async def setup(self, config: MatrixBotConfig) -> None: + """ + Build `AsyncClient` from `config`. + + Returns: + - Does not return anything, raises an exception on error. + """ + if self._was_built_ever: + raise RuntimeError("The client was already built") + client_config = AsyncClientConfig( + store_name="nio_store_file", + encryption_enabled=True, + store_sync_tokens=False + ) + self._client = AsyncClient( + homeserver=config.matrix_homeserver_url, + user=config.matrix_username_localpart, + store_path=str(config.storage_directory), + config=client_config + ) + try: + await self._setup_callbacks(config) + except: + self._client = None + raise + self._config = config + self._was_built_ever = True + + def get_client(self) -> AsyncClient: + """ + Get AsyncClient. + + Returns: + - `AsyncClient` you may use to access Matrix. Raises a RuntimeError + if the client does not exist or is not operational. + """ + if self._client is None: + raise RuntimeError("The client does not exist") + return self._client + + async def start(self) -> None: + """ + Start client operation in background. This method can be called + succesfully only once for `ClientManager` object. Any subsequent calls + will raise an exception. + """ + if self._client is None: + raise RuntimeError("Client does not exist") + if self._was_started_ever: + raise RuntimeError("Client can't be started several times") + self._background_task = asyncio.create_task( + self._background_coroutine() + ) + self._was_started_ever = True + self._logger.info("Bot background is started") + + async def stop(self) -> None: + """ + Stop client operation in background. This method can be called only if + the client was `start(...)`ed previously. This method will wait for the + bot to completely stop. + """ + if self._client is None: + raise RuntimeError("Client does not exist") + if self._background_task is None: + if self._was_started_ever: + raise RuntimeError("The bot was stopped previously") + else: + raise RuntimeError("The bot was never started") + self._background_task.cancel() + try: + asyncio.shield(self._background_task) + except: + self._logger.error(traceback.format_exc()) + try: + asyncio.shield(self._close_client()) + except: + self._logger.error(traceback.format_exc()) + self._background_task = None + self._logger.info("Bot background task is stopped") \ No newline at end of file diff --git a/src/mab/bot/_client_sender.py b/src/mab/bot/_client_sender.py new file mode 100644 index 0000000..b8daa2d --- /dev/null +++ b/src/mab/bot/_client_sender.py @@ -0,0 +1,300 @@ +import asyncio +import logging +from html.parser import HTMLParser +from pathlib import Path +import traceback +import os + +import magic +from PIL import Image +from nio import AsyncClient, MatrixRoom, RoomSendResponse, RoomSendError +from nio import OlmUnverifiedDeviceError + +from ._client_uploader import ClientUploader +from ..types import MatrixBotConfig, VideoFileProperties +from ..utils import Utils + +class ClientSender: + """This class sends messages""" + # + # PRIVATE + # + def _process_html_text(self, text: str, is_html: bool | None) -> dict: + """ + Processes `text` and `is_html` arguments of `send_*` functions and + prepares `content` keys related to text data inside the message. + + Args: + - text - `text` as passed to the `send_*` function + - is_html - `is_html` as passed to the `send_*` function + + Returns: + `dict` that contains keys required for the text message to be + displayed correctly. + """ + # text w/ and w/o HTML + text_w_html: str | None = None + text_wo_html: str | None = None + # check if HTML is present + if is_html or is_html is None: + # will be True if `text` contains HTML tags + has_tags = False + # text fragments without tags + text_fragments = [] + # temporary extractor class + class Extractor(HTMLParser): + def handle_starttag(self, tag, attrs): + nonlocal has_tags + has_tags = True + def handle_data(self, data): + text_fragments.append(data) + # parse + parser = Extractor(convert_charrefs=True) + try: + parser.feed(text) + except: + self._logger.error(traceback.format_exc()) + text_wo_html = text + is_html = has_tags + if is_html: + text_wo_html = " ".join(text_fragments) + text_w_html = text + if not is_html: + text_wo_html = text + # return the content + if text_w_html: + return { + "body": text_wo_html, + "format": "org.matrix.custom.html", + "formatted_body": text_w_html + } + else: + return { + "body": text_wo_html + } + + # + # PUBLIC + # + def __init__(self): + self._logger = logging.getLogger("ClientSender") + self._config: MatrixBotConfig | None = None + self._client: AsyncClient | None = None + + async def setup(self, + config: MatrixBotConfig, + client: AsyncClient, + uploader: ClientUploader) -> None: + """ + Setup the sender. + + Args: + - config - config to use + - client - client to use + - uploader - file uploaded to use + """ + self._config = config + self._client = client + self._uploader = uploader + + async def send_content(self, + room: MatrixRoom | str, + content: dict) -> RoomSendResponse: + """ + Send raw `content` JSON to `room`. + + Args: + - room - the room to send the content to + - content - the content to send to the room + + Returns: + - `RoomSendResponse` returned by `AsyncClient.room_send` + - Raises an exception on error + """ + if self._client is None or self._config is None: + raise RuntimeError("ClientSender is not set up") + try: + result = await self._client.room_send( + room_id=room.room_id if isinstance(room, MatrixRoom) else room, + message_type="m.room.message", + content=content + ) + except OlmUnverifiedDeviceError: + if self._config.auto_verify_all_known_devices: + if not Utils.verify_all_known_devices(self._client): + raise + return await self.send_content(room, content) + else: + raise + if type(result) is RoomSendResponse: + return result + elif type(result) is RoomSendError: + raise RuntimeError(result) + else: + raise RuntimeError("Unknown error has occured", result) + + async def send_text(self, + room: MatrixRoom | str, + text: str, + *, + is_html: bool | None = None) -> str: + """ + Send text message to `room`. + + Args: + - room - the room to send the text to + - text - the text to send to the room + - is_html - whether the text is HTML-formatted. Use `None` for auto + + Returns: + - `event_id` of sent message on success + - Raises an exception on error + """ + text_data = self._process_html_text(text, is_html) + content = { + "msgtype": "m.text", + **text_data + } + return (await self.send_content(room, content)).event_id + + async def send_image(self, + room: MatrixRoom | str, + path: Path | str, *, + text: str | None = None, + is_html: bool | None = None, + filename: str | None = None, + timeout: float | None = 60 * 60) -> str: + """ + Send the image to `room`. Please note that formatted text is displayed + incorrectly in some clients as of September 8th, 2026 + + Args: + - room - the room to send the text to + - path - path to the image file + - text - image caption to use (`None` to disable) + - is_html - whether the text is HTML-formatted (`None` for auto) + - filename - filename to use for the file (`None` for auto) + - timeout - upload timeout in seconds (`None` to disable) + + Returns: + - `event_id` of sent message on success + - Raises an exception on error + """ + # caption must not actually be empty + if text is None or not text.strip(): + text = os.path.basename(path) + is_html = False + # check if the file is image + mime_type: str = magic.from_file(path, mime=True) + if not mime_type.startswith("image/"): + raise RuntimeError(f"File {path} has non-image mime-type") + # get image size + with Image.open(path) as image: + width, height = image.size + # upload + async with asyncio.timeout(timeout): + upload_result = await self._uploader.upload_file( + path, mime_type=mime_type, filename=filename) + # prepare the content and send + content = { + "msgtype": "m.image", + "filename": os.path.basename(path), + **self._process_html_text(text, is_html), + "file": { + "url": upload_result.response.content_uri, + "mimetype": mime_type, + **upload_result.keys + }, + "info": { + "mimetype": mime_type, + "size": upload_result.filesize, + "w": width, + "h": height + } + } + return (await self.send_content(room, content)).event_id + + async def send_video(self, + room: MatrixRoom | str, + path: Path | str, + *, + props: VideoFileProperties | None = None, + text: str | None = None, + is_html: bool | None = None, + timeout: float | None = 60 * 60) -> str: + """ + Send the video to `room`. Please note that formatted text is displayed + incorrectly in some clients as of September 8th, 2026. Unknown video + properties will be automatically deduced as configured in + `MatrixBotConfig`. + + Args: + - room - the room to send the text to + - path - path to the video file + - props - video properties (`None` for auto, if the feature is ON) + - text - video caption to use (`None` to disable) + - is_html - whether the text is HTML-formatted (`None` for auto) + - timeout - upload timeout in seconds (`None` to disable) + + Returns: + - `event_id` of sent message on success + - Raises an exception on error + """ + if self._client is None or self._config is None: + raise RuntimeError("ClientSender is not set up") + # caption must not actually be empty + if text is None or not text.strip(): + text = os.path.basename(path) + is_html = False + # check if the file is video + mime_type: str = magic.from_file(path, mime=True) + if not mime_type.startswith("video/"): + raise RuntimeError(f"File {path} has non-video mime-type") + # properties must not be empty + if props is None: + if not self._config.auto_ffprobe: + raise RuntimeError( + "No `props` are supplied and `auto_ffprobe` is disabled" + ) + async with asyncio.timeout(30): + props = await Utils.get_video_props_using_ffprobe(path) + # autogenerate thumbnail if allowed to do so + if not props.thumbnail and self._config.auto_ffmpeg: + try: + async with asyncio.timeout(30): + props.thumbnail = await Utils.get_video_thumbnail( + path, + timestamp=props.duration_ms / 1000 / 2, + max_size=400 + ) + except asyncio.TimeoutError: + pass + async with asyncio.timeout(timeout): + # upload the thumbnail if specified + thumbnail_data = await self._uploader.upload_thumbnail(props.thumbnail) + # upload + upload_result = await self._uploader.upload_file( + path, + mime_type=mime_type + ) + # prepare the content + content = { + "msgtype": "m.video", + "filename": os.path.basename(path), + **self._process_html_text(text, is_html), + "file": { + "url": upload_result.response.content_uri, + "mimetype": mime_type, + **upload_result.keys + }, + "info": { + "duration": props.duration_ms, + "mimetype": mime_type, + "size": upload_result.filesize, + "w": props.width_px, + "h": props.height_px, + **thumbnail_data + } + } + # send + return (await self.send_content(room, content)).event_id \ No newline at end of file diff --git a/src/mab/bot/_client_uploader.py b/src/mab/bot/_client_uploader.py new file mode 100644 index 0000000..7bded2b --- /dev/null +++ b/src/mab/bot/_client_uploader.py @@ -0,0 +1,187 @@ +from typing import Any +from pathlib import Path +import io +import os + +import aiofiles +import aiofiles.os +import magic +from PIL import Image + +from nio import AsyncClient +from nio import UploadError + +from ._storage import Storage +from ..types import MatrixBotConfig, UploadResult + +class ClientUploader: + """This class uploads files""" + def __init__(self, storage: Storage): + self._storage = storage + self._config: MatrixBotConfig | None = None + self._client: AsyncClient | None = None + + async def setup(self, config: MatrixBotConfig, client: AsyncClient) -> None: + """ + Setup the uploader. + + Args: + - config - config to take some settings from. + """ + self._config = config + self._client = client + + async def upload_using_provider(self, + provider: Any, + *, + mime_type: str, + filename: str, + filesize: int) -> UploadResult: + """ + Upload data using provider (as defined in `AsyncClient.upload` docs). + You better use `upload_file(...)`, `upload_bytes(...)`. Data is always + encrypted. + + Args: + - provider - data provider (e.g. `aiofiles.open(...)`) + - mime_type - mime-type of the data + - filename - filename to declare for the file + - filesize - size of the file + + Returns: + - `UploadResult` that contains all data you need on success + - Raises an exception on failure + """ + if self._config is None or self._client is None: + raise RuntimeError("ClientUploader is not set up") + response, keys = await self._client.upload( + provider, + content_type=mime_type, + filename=filename, + encrypt=True, + monitor=None, + filesize=filesize + ) + if isinstance(response, UploadError): + raise RuntimeError(f"Upload failed: {response.message}", response) + return UploadResult( + response=response, + keys=keys, # type: ignore + filename=filename, + filesize=filesize + ) + + async def upload_file(self, + path: Path | str, + *, + mime_type: str | None = None, + filename: str | None = None) -> UploadResult: + """ + Upload data from file. Data is always encrypted. + + Args: + - path - path to the file that needs to be uploaded + - mime_type - mime-type to use for the file (`None` for auto) + - filename - filename to use for the file (`None` for auto) + """ + if isinstance(path, str): + path = Path(path) + if not path.is_file(): + raise RuntimeError(f"{path} is not a file") + if not filename: + filename = os.path.basename(path) + if not mime_type: + mime_type = magic.from_file(path, mime=True) + filesize = (await aiofiles.os.stat(path)).st_size + async with aiofiles.open(path, "r+b") as f: + return await self.upload_using_provider( + f, + mime_type=mime_type, + filename=filename, + filesize=filesize + ) + + async def upload_bytes(self, + data: bytes, + *, + filename: str, + mime_type: str | None = None) -> UploadResult: + """ + Upload data. Data is always encrypted. + + Args: + - data - data that needs to be uploaded + - filename - filename to use for the file + - mime_type - mime-type to use for the file (`None` for auto) + """ + if not mime_type: + mime_type = magic.from_buffer(data, mime=True) + buffer = io.BytesIO(data) + return await self.upload_using_provider( + buffer, + mime_type=mime_type, + filename=filename, + filesize=len(data) + ) + + async def upload_thumbnail(self, + thumbnail: Path | str | bytes | None) -> dict: + """ + Upload thumbnail and return fields required to use it for the video. + + Args: + - thumbnail - path to the thumbnail file or raw `bytes` of the + image. You can actually pass `None`, empty dict will + be returned in this case. + + Returns: + `dict` you should embed into your `content["info"]` + """ + if thumbnail is None: + return {} + # raw data + raw_bytes_io: io.BytesIO | None = None + if type(thumbnail) is bytes: + thumb_mime = magic.from_buffer(thumbnail, mime=True) + raw_bytes_io = io.BytesIO(thumbnail) + # file path + else: + thumb_mime = magic.from_file(thumbnail, mime=True) + # check if invalid MIME + if not thumb_mime.startswith("image/"): + raise RuntimeError(f"Thumbnail MIME is not image: `{thumb_mime}`") + # get thumbnail size + with Image.open(raw_bytes_io or thumbnail) as f: + width, height = f.size + # seek BytesIO to zero so we can reuse it + if raw_bytes_io is not None: + raw_bytes_io.seek(0) + # upload + if type(thumbnail) is bytes: + thumb_size = len(thumbnail) + upload_result = await self.upload_using_provider( + raw_bytes_io, + mime_type=thumb_mime, + filename=f"preview.{thumb_mime.split("/")[-1]}", + filesize=thumb_size + ) + else: + upload_result = await self.upload_file( + thumbnail, # type: ignore + mime_type=thumb_mime, + filename=os.path.basename(str(thumbnail)) + ) + # return the result + return { + "thumbnail_file": { + "url": upload_result.response.content_uri, + "mimetype": thumb_mime, + **upload_result.keys + }, + "thumbnail_info": { + "w": width, + "h": height, + "size": upload_result.filesize, + "mimetype": thumb_mime + } + } \ No newline at end of file diff --git a/src/mab/bot/_storage.py b/src/mab/bot/_storage.py new file mode 100644 index 0000000..1d2618a --- /dev/null +++ b/src/mab/bot/_storage.py @@ -0,0 +1,103 @@ +from pathlib import Path +import aiofiles +import logging +import json +from typing import Any + +from ..types import MatrixBotConfig + +class Storage: + """This class manages internal storage of the MatrixBot""" + + def __init__(self): + self._logger = logging.getLogger("Storage") + self._config: MatrixBotConfig | None = None + + async def setup(self, config: MatrixBotConfig) -> None: + """ + Setup the internal storage of the MatrixBot. + + Args: + - config - the config to use + + Returns: + - Does not return anything + - Raises an exception on error + """ + self._config = config + + async def write_next_batch(self, next_batch: str) -> None: + """ + Write `next_batch` value to the disk. + + Args: + - next_batch - value to save + + Returns: + - Does not return anything + - Raises an exception on error + """ + if self._config is None: + raise RuntimeError("Storage is not set up") + path = 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: + """ + Read `next_batch` value from the disk. + + Returns: + - Returns `str` if there is `next_batch` value saved on disk + - Returns `None` if `next_batch` was never saved yet + - Raises an exception on error + """ + if self._config is None: + raise RuntimeError("Storage is not set up") + path = 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. + + Args: + - access_token - access_token to write to the disk + - device_id - device_id to write to disk + + Returns: + - Does not return anything + - Raises an exception on error + """ + if self._config is None: + raise RuntimeError("Storage is not set up") + path = 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. + + Returns: + - Dictionary with `access_token` and `device_id` keys if these + values were saved using `write_session_data` + - `None` if session data was never saved yet + """ + if self._config is None: + raise RuntimeError("Storage is not set up") + path = 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 \ No newline at end of file diff --git a/src/mab/bot/_validation.py b/src/mab/bot/_validation.py new file mode 100644 index 0000000..7703b10 --- /dev/null +++ b/src/mab/bot/_validation.py @@ -0,0 +1,84 @@ +import re +from pathlib import Path +from urllib.parse import urlparse + +from ..types import MatrixBotConfig + +class Validator: + """This class is used inside MatrixBot to validate data""" + + def validate_matrix_homeserver_url(self, url: str) -> None: + """ + Check if `url` is a valid matrix homeserver URL. + + Args: + - url - URL to check (only format is checked) + + Returns: + - Does not return anything + - Raises an exception on error + """ + 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)" + ) + + def validate_matrix_username_localpart(self, localpart: str) -> None: + """ + Check if `localpart` is a valid localpart of matrix username. + + Args: + - localpart - localpart of matrix username + + Returns: + - Does not return anything + - Raises an exception on error + """ + 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" + ) + + def ensure_directory_exists(self, path: Path | str) -> None: + """ + Ensure that `path` is an existing directory. Creates it if it does not + exist. The entire directory tree is created if needed. + + Args: + - path - directory path + + Returns: + - Does not return anything + - Raises an exception on error + """ + if isinstance(path, str): + path = Path(path) + if not path.is_dir(): + path.mkdir(parents=True, exist_ok=True) + + def validate_bot_config(self, config: MatrixBotConfig) -> None: + """ + Check if `config` is a valid MatrixBot config. + + Args: + - config - configuration to check + + Returns: + - Does not return anything + - Raises an exception on error + """ + self.validate_matrix_homeserver_url(config.matrix_homeserver_url) + self.validate_matrix_username_localpart(config.matrix_username_localpart) + self.ensure_directory_exists(config.storage_directory) \ No newline at end of file diff --git a/src/mab/bot/bot.py b/src/mab/bot/bot.py new file mode 100644 index 0000000..81a3611 --- /dev/null +++ b/src/mab/bot/bot.py @@ -0,0 +1,198 @@ +import logging + +from typing import Callable, Coroutine, Any + +from nio import AsyncClient + +from ..filters.base import BaseEventFilter +from ..types import * + +from ._validation import Validator +from ._storage import Storage +from ._client_auth import ClientAuth +from ._client_manager import ClientManager +from ._client_uploader import ClientUploader +from ._client_sender import ClientSender +from ._callbacks import Callbacks + + +class MatrixBot: + """ + Asynchronous Matrix Bot Implementation. + """ + + def __init__(self, config: MatrixBotConfig) -> None: + # create the logger + self._logger = logging.getLogger("MatrixBot") + self._was_started = False + # create classes that compose MatrixBot + self._validator = Validator() + self._storage = Storage() + self._client_auth = ClientAuth(self._storage) + self._client_manager = ClientManager(self._client_auth, self._storage) + self._client_uploader = ClientUploader(self._storage) + self._client_sender = ClientSender() + self._callbacks = Callbacks(self._storage, self) + # validate the config and save it + try: + self._validator.validate_bot_config(config) + self._config: MatrixBotConfig = config + except Exception as e: + self._logger.error(f"Couldn't validate bot config: {e}") + raise + + def add_callback(self, + filter: BaseEventFilter, + callback: Callable[[RoomEventData], Coroutine[Any, Any, None]] | None, + *, + stop_matching: bool = True) -> None: + """ + Add filter-based callback. + + Args: + - filter - filter to use for the callback + - callback - the callback that must be called if the filter matches. + Use `None` to log the type of the Event that will + happen for specified filter. + - stop_matching - whether to stop matching if the filter matches + """ + self._callbacks.add_room_event_callback( + filter=filter, + callback=callback, + stop_matching=stop_matching + ) + + async def start(self) -> None: + """ + Start client operation in background. This method can be called + succesfully only once for `MatrixBot` object. Any subsequent calls + will raise an exception. + """ + if self._was_started: + raise RuntimeError("The bot was already started") + self._was_started = True + await self._client_manager.setup(self._config) + await self._storage.setup(self._config) + await self._client_auth.setup(self._config) + await self._client_uploader.setup( + self._config, + self._client_manager.get_client() + ) + await self._client_sender.setup( + self._config, + self._client_manager.get_client(), + self._client_uploader + ) + await self._callbacks.setup( + self._config, + self._client_manager.get_client() + ) + await self._client_manager.start() + + async def stop(self) -> None: + """ + Stop client operation in background. This method can be called only if + the client was `start(...)`ed previously. This method will wait for the + bot to completely stop. + """ + await self._client_manager.stop() + + def get_client(self) -> AsyncClient: + """ + Get AsyncClient. + + Returns: + - `AsyncClient` you may use to access Matrix. Raises a RuntimeError + if the client does not exist or is not operational. + """ + return self._client_manager.get_client() + + async def send_text(self, + room: MatrixRoom | str, + text: str, + *, + is_html: bool | None = None) -> str: + """ + Send text message to `room`. + + Args: + - room - the room to send the text to + - text - the text to send to the room + - is_html - whether the text is HTML-formatted. Use `None` for auto + + Returns: + - `event_id` of sent message on success + - Raises an exception on error + """ + return await self._client_sender.send_text( + room=room, + text=text, + is_html=is_html + ) + + async def send_image(self, + room: MatrixRoom | str, + path: Path | str, *, + text: str | None = None, + is_html: bool | None = None, + filename: str | None = None, + timeout: float | None = 60 * 60) -> str: + """ + Send the image to `room`. Please note that formatted text is displayed + incorrectly in some clients as of September 8th, 2026 + + Args: + - room - the room to send the text to + - path - path to the image file + - text - image caption to use (`None` to disable) + - is_html - whether the text is HTML-formatted (`None` for auto) + - filename - filename to use for the file (`None` for auto) + - timeout - upload timeout in seconds (`None` to disable) + + Returns: + - `event_id` of sent message on success + - Raises an exception on error + """ + return await self._client_sender.send_image( + room=room, + path=path, + text=text, + is_html=is_html, + filename=filename, + timeout=timeout + ) + + async def send_video(self, + room: MatrixRoom | str, + path: Path | str, + *, + props: VideoFileProperties | None = None, + text: str | None = None, + is_html: bool | None = None, + timeout: float | None = 60 * 60) -> str: + """ + Send the video to `room`. Please note that formatted text is displayed + incorrectly in some clients as of September 8th, 2026. Unknown video + properties will be automatically deduced as configured in + `MatrixBotConfig`. + + Args: + - room - the room to send the text to + - path - path to the video file + - props - video properties (`None` for auto, if the feature is ON) + - text - video caption to use (`None` to disable) + - is_html - whether the text is HTML-formatted (`None` for auto) + - timeout - upload timeout in seconds (`None` to disable) + + Returns: + - `event_id` of sent message on success + - Raises an exception on error + """ + return await self._client_sender.send_video( + room=room, + path=path, + props=props, + text=text, + is_html=is_html, + timeout=timeout + ) \ No newline at end of file diff --git a/src/mab/types.py b/src/mab/types.py index 065fbc3..21ccaf9 100644 --- a/src/mab/types.py +++ b/src/mab/types.py @@ -3,6 +3,15 @@ from pathlib import Path from dataclasses import dataclass +from nio import MatrixRoom, Event +from nio import UploadResponse + +from .filters.base import BaseEventFilter + +from typing import TYPE_CHECKING +if TYPE_CHECKING: + from .bot import MatrixBot + @dataclass class MatrixBotConfig: """Configuration for MatrixBot""" @@ -13,7 +22,7 @@ class MatrixBotConfig: matrix_username_localpart: str """Localpart of matrix username (without homeserver), for example: `valid-username`""" - storage_directory: Path + storage_directory: Path | str """Path to the storage directory (will be created if needed)""" auto_verify_all_known_devices: bool = True @@ -28,6 +37,20 @@ class MatrixBotConfig: auto_ffmpeg: bool = True """Enabled automatic `ffmpeg` execution to generate some media (e.g. thumbnails)""" + allow_ainput_password: bool = True + """Allow `aioconsole.ainput` usage to get user password""" + + client_sync_period_ms: int = 5000 + """Sync period for `nio.AsyncClient` in milliseconds""" + + client_next_batch_dump_period_s: float = 120.0 + """`next_batch` value dumping period in seconds""" + + enable_multitasking_callbacks: bool = True + """Filter-based callbacks will be started each in separate task so that + long callback execution will not prevent other callbacks from executting + """ + @dataclass class VideoFileProperties: """Properties of the video for sending""" @@ -42,4 +65,36 @@ class VideoFileProperties: """Duration in milliseconds""" thumbnail: Path | str | bytes | None = None - """Path to the thumbnail or the raw JPEG thumbnail data""" \ No newline at end of file + """Path to the thumbnail or the raw JPEG thumbnail data""" + +@dataclass +class RoomEventData: + """Dataclass that hold information about event that happened in the room""" + + room: MatrixRoom + """The room the event has happened in""" + + event: Event + """The event that has happened in the room""" + + filter: BaseEventFilter + """The filter that invoked this event""" + + bot: "MatrixBot" + """The bot that is the source of the event""" + +@dataclass +class UploadResult: + """Result of data upload""" + + response: UploadResponse + """Response, as returned by `AsyncClient.upload(...)`""" + + keys: dict + """Encryptions keys, as returned by `AsyncClient.upload(...)`""" + + filename: str + """Name of the uploaded file""" + + filesize: int + """Size of uploaded file""" \ No newline at end of file diff --git a/src/mab/utils.py b/src/mab/utils.py index fe33cb7..0ab33bd 100644 --- a/src/mab/utils.py +++ b/src/mab/utils.py @@ -3,6 +3,9 @@ import asyncio import json from pathlib import Path +from urllib.parse import urlparse + +from nio import AsyncClient from .types import * @@ -95,4 +98,36 @@ class Utils: if proc.returncode != 0 or not stdout: raise RuntimeError(f"`ffmpeg` exits with code {proc.returncode}: {stderr.decode(errors="replace")}", args) # success - return stdout \ No newline at end of file + return stdout + + @staticmethod + def build_full_matrix_username(config: MatrixBotConfig) -> str: + """Build complete matrix username using data from MatrixBotConfig.""" + homeserver_name = urlparse(config.matrix_homeserver_url).hostname + localpart = config.matrix_username_localpart + return f"@{localpart}:{homeserver_name}" + + @staticmethod + def verify_all_known_devices(client: AsyncClient) -> bool: + """ + Verify all known devices. + + Args: + - client - client to verify all known devices for + + Returns: + - True if there were unverified devices that are verified now + """ + result = False + for user_id in client.device_store.users: + for device_id, olm_device in client.device_store[user_id].items(): + # can't trust ourselves + if device_id == client.device_id and user_id == client.user_id: + continue + # they are already verified + if olm_device.verified: + continue + # verify them + client.verify_device(olm_device) + result = True + return result \ No newline at end of file