diff --git a/pyproject.toml b/pyproject.toml index 93bb3c0..33de82f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,6 +47,8 @@ dependencies = [ "unpaddedbase64==2.1.0", "vodozemac==0.10.0", "yarl==1.24.5", + "python-magic==0.4.27", + "pillow==12.3.0" ] [project.urls] diff --git a/src/mab/bot.py b/src/mab/bot.py index 19dbbdf..385662a 100644 --- a/src/mab/bot.py +++ b/src/mab/bot.py @@ -1,5 +1,5 @@ import asyncio -import aiofiles +import aiofiles, aiofiles.os import aioconsole import traceback import logging @@ -7,14 +7,19 @@ 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 @@ -26,6 +31,7 @@ from .filters.base import BaseEventFilter from nio.events.room_events import Event as RoomEvemt from .types import * +from .utils import Utils class MatrixBot: @@ -108,31 +114,59 @@ class MatrixBot: 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) + 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. """ - 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) - + # 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 @@ -178,7 +212,8 @@ class MatrixBot: 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.""" @@ -215,6 +250,154 @@ class MatrixBot: 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 # @@ -249,7 +432,7 @@ class MatrixBot: async def _callback_filter_router(self, *args, **kwargs): if len(args) != 2: self._logger.debug("Can't process the event, not enough positional args") - self._debug_event_callback(*args, **kwargs) + await self._debug_event_callback(*args, **kwargs) return room = args[0] event = args[1] @@ -412,6 +595,8 @@ class MatrixBot: # 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 @@ -475,7 +660,7 @@ class MatrixBot: result = True return result - def add_event_callback(self, callback: Callable[[Any], Awaitable[None]] | None, filter: BaseEventFilter, stop_after_this: bool = True) -> None: + 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 @@ -485,7 +670,7 @@ class MatrixBot: """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: + 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: @@ -496,42 +681,147 @@ class MatrixBot: 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()` + text_data = self._process_html_text(text, is_html) content = { "msgtype": "m.text", - "body": text, - **kwargs + **text_data } - 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 + 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 ) - 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 + # 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/types.py b/src/mab/types.py index 333264e..065fbc3 100644 --- a/src/mab/types.py +++ b/src/mab/types.py @@ -20,4 +20,26 @@ class MatrixBotConfig: """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 + """Enable automatic join for invites (implemented by automatically adding internal callback on bot creation)""" + + auto_ffprobe: bool = True + """Enable automatic `ffprobe` execution to get some media properties""" + + auto_ffmpeg: bool = True + """Enabled automatic `ffmpeg` execution to generate some media (e.g. thumbnails)""" + +@dataclass +class VideoFileProperties: + """Properties of the video for sending""" + + width_px: int + """Width in pixels""" + + height_px: int + """Height in pixels""" + + duration_ms: int + """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 diff --git a/src/mab/utils.py b/src/mab/utils.py new file mode 100644 index 0000000..fe33cb7 --- /dev/null +++ b/src/mab/utils.py @@ -0,0 +1,98 @@ +"""Utility things""" + +import asyncio +import json +from pathlib import Path + +from .types import * + +class Utils: + @staticmethod + async def get_video_props_using_ffprobe(path: Path | str) -> VideoFileProperties: + """Returns video properties for file at `path`. + + Consider using `asyncio.timeout` to prevent deadlocking. + """ + # prepare arguments for ffprobe + args = [ + "ffprobe", + "-v", "error", + "-select_streams", "v:0", + "-show_entries", "stream=width,height:format=duration", + "-of", "json", + str(path) + ] + # start the process + proc = await asyncio.create_subprocess_exec( + *args, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + # communicate `stdout` and `stderr` + try: + stdout, stderr = await proc.communicate() + except asyncio.CancelledError: + if proc.returncode is None: + proc.kill() + await proc.wait() + raise + # error + if proc.returncode != 0: + raise RuntimeError(f"`ffprobe` exits with code {proc.returncode}: {stderr.decode(errors="replace")}") + # success + result = json.loads(stdout) + return VideoFileProperties( + width_px=int(result["streams"][0]["width"]), + height_px=int(result["streams"][0]["height"]), + duration_ms=int(float(result["format"]["duration"]) * 1000) + ) + + @staticmethod + async def get_video_thumbnail(path: Path | str, *, timestamp: float = 0, max_size: int = 300) -> bytes: + """Generates video thumbnail for file at `path`. + + Args: + - path - path to the video + - timestamp - timestamp to use as the thumbnail + - max_size - max width/height + + Returns: + Raw bytes of JPEG thumbnail + """ + # prepare arguments for ffmpeg + args = [ + "ffmpeg", + "-v", "error", + "-ss", "%.2f" % timestamp, + "-i", str(path), + "-map", "0:v:0", + "-vf", + ( + f"scale='min(iw,{max_size})':'min(ih,{max_size})':" + "force_original_aspect_ratio=decrease" + ), + "-frames:v", "1", + "-q:v", "3", + "-f", "image2pipe", + "-vcodec", "mjpeg", + "pipe:1" + ] + # start the process + proc = await asyncio.create_subprocess_exec( + *args, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + # communicate `stdout` and `stderr` + try: + stdout, stderr = await proc.communicate() + except asyncio.CancelledError: + if proc.returncode is None: + proc.kill() + await proc.wait() + raise + # error + 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