diff --git a/.gitignore b/.gitignore index ee22eb1..89550bc 100644 --- a/.gitignore +++ b/.gitignore @@ -5,4 +5,5 @@ session_storage/ dist/ *.egg-info/ *.swp -*.swo \ No newline at end of file +*.swo +*.tmp \ No newline at end of file diff --git a/examples/file_bot.py b/examples/file_bot.py new file mode 100644 index 0000000..ef64292 --- /dev/null +++ b/examples/file_bot.py @@ -0,0 +1,97 @@ +""" +This example implements Matrix bot that calculates SHA256 for a file sent by +user. + +It uses environment variables to specify authorization data. Use Ctrl+C to stop +the bot. +""" + +import asyncio +import aiofiles +import os +import logging +import hashlib + +from mab import * + +from _environment import check_environment + +async def on_text_message(ctx: EventContext) -> None: + """This callback is called when a text message arrives.""" + await ctx.bot.send_text(ctx.room, "Please send a file/image/video") + +async def on_file_message(ctx: EventContext) -> None: + """This callback is called when a file arrives.""" + sha = hashlib.sha256() + does_temp_exist = False + if ctx[CTX_FILE_SIZE] > 1_000_000: + await ctx.bot.send_text( + ctx.room, + "The file is larger than 1 MB, downloading to filesystem" + ) + await ctx.bot.download_file(ctx, path="temp.tmp") + async with aiofiles.open("temp.tmp", "rb") as f: + while True: + data = await f.read(64 * 1024) + if not data: + break + sha.update(data) + does_temp_exist = True + else: + await ctx.bot.send_text( + ctx.room, + "The file is smaller than 1 MB, downloading to RAM" + ) + content = await ctx.bot.download_file(ctx, path=None) + sha.update(content) + # result + response = f"SHA256 for file `{ctx[CTX_FILE_NAME]}`" + response += f" ({ctx[CTX_FILE_SIZE]} bytes, {ctx[CTX_FILE_MIME]})" + await ctx.bot.send_file_bytes( + room=ctx.room, + data=sha.hexdigest().encode("utf-8"), + filename="hash of the file.txt", + mime_type="text/plain", + text=response + ) + # resend the file to test uploading + if does_temp_exist: + await ctx.bot.send_file( + room=ctx.room, + path="temp.tmp", + filename=ctx[CTX_FILE_NAME], + text="This is the file you have sent, but it was reuploaded" + ) + try: + os.unlink("temp.tmp") + except: + pass + + +async def main() -> None: + """Application entry point""" + logging.basicConfig(level=logging.INFO) + logging.getLogger("nio").setLevel(logging.CRITICAL+1) + check_environment() + + config = MatrixBotConfig( + matrix_homeserver_url=os.environ["MATRIX_HOMESERVER"], + matrix_username_localpart=os.environ["MATRIX_USERNAME"], + storage_directory="session_storage" + ) + bot = MatrixBot(config) + bot.add_callback( + ~SenderIsBotFilter() & BodyExistsFilter() & MessageTypeFilter(MessageType.TEXT), + on_text_message) + bot.add_callback( + ~SenderIsBotFilter() & MessageHasFile() & NewMessageFilter(), + on_file_message) + + # run until Ctrl+C + try: + await bot.run() + except asyncio.CancelledError: + pass + +if __name__ == "__main__": + asyncio.run(main()) \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index ae641a2..39d88cc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "mab" -version = "0.5.1" +version = "0.5.2" authors = [ { name = "Tyukalov Nikita", email = "nikita@tyukalov.su" } ] diff --git a/src/mab/__init__.py b/src/mab/__init__.py index 40d85a0..bda1ac2 100644 --- a/src/mab/__init__.py +++ b/src/mab/__init__.py @@ -27,6 +27,9 @@ __all__ = [ "CTX_CMD_PREFIX", "CTX_CMD_VERB", "CTX_CMD_ARGS", + "CTX_FILE_SIZE", + "CTX_FILE_MIME", + "CTX_FILE_NAME", # .bot "MatrixBot", @@ -50,4 +53,5 @@ __all__ = [ "RedactedMessageFilter", "SenderIsFilter", "SenderIsBotFilter", + "MessageHasFile", ] \ No newline at end of file diff --git a/src/mab/bot/_client_sender.py b/src/mab/bot/_client_sender.py index 6da6e6f..fedfeb0 100644 --- a/src/mab/bot/_client_sender.py +++ b/src/mab/bot/_client_sender.py @@ -361,4 +361,123 @@ class ClientSender: } } # send + return (await self.send_content(room, content)).event_id + + async def send_file(self, + room: MatrixRoom | str, + path: Path | str, + *, + filename: str | None = None, + mime_type: str | None = None, + text: str | None = None, + is_html: bool | None = None, + timeout: float | None = 60 * 60): + """ + Send the file 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 file to + - path - path to the file + - filename - filename to use for upload (`None` to use basename from + `path`) + - mime_type - mime type to use (`None` for autodetect using content + of `path`) + - text - 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") + # get filename if not set + if not filename: + filename = os.path.basename(path) + # caption must not actually be empty + if text is None or not text.strip(): + text = filename + is_html = False + # get the mime type if not specified + if mime_type is None: + mime_type = magic.from_file(path, mime=True) + # upload + async with asyncio.timeout(timeout): + upload_result = await self._uploader.upload_file( + path, mime_type=mime_type, filename=filename) + # send + content = { + "msgtype": "m.file", + "filename": filename, + **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 + } + } + return (await self.send_content(room, content)).event_id + + async def send_file_bytes(self, + room: MatrixRoom | str, + data: bytes, + filename: str, + *, + mime_type: str | None = None, + text: str | None = None, + is_html: bool | None = None, + timeout: float | None = 60 * 60) -> str: + """ + Send the file to `room`. + + Args: + - room - the room to send the file to + - data - the content of the file to send + - filename - filename to use for the file + - mime_type - mime type to use (`None` for autodetect using `data` + content) + - text - 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 + """ + # caption must not actually be empty + if text is None or not text.strip(): + text = filename + is_html = False + # check if the file is image + if not mime_type: + mime_type = magic.from_buffer(data, mime=True) + # upload + buffer = BytesIO(data) + async with asyncio.timeout(timeout): + upload_result = await self._uploader.upload_using_provider( + provider=buffer, + mime_type=mime_type, + filename=filename, + filesize=len(data)) + # prepare the content and send + content = { + "msgtype": "m.file", + "filename": filename, + **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 + } + } return (await self.send_content(room, content)).event_id \ No newline at end of file diff --git a/src/mab/bot/bot.py b/src/mab/bot/bot.py index 9649898..3832aa9 100644 --- a/src/mab/bot/bot.py +++ b/src/mab/bot/bot.py @@ -252,6 +252,77 @@ class MatrixBot: timeout=timeout ) + async def send_file(self, + room: MatrixRoom | str, + path: Path | str, + *, + filename: str | None = None, + mime_type: str | None = None, + text: str | None = None, + is_html: bool | None = None, + timeout: float | None = 60 * 60): + """ + Send the file 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 file to + - path - path to the file + - filename - filename to use for upload (`None` to use basename from + `path`) + - mime_type - mime type to use (`None` for auto) + - text - 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_file( + room=room, + path=path, + filename=filename, + mime_type=mime_type, + text=text, + is_html=is_html, + timeout=timeout + ) + + async def send_file_bytes(self, + room: MatrixRoom | str, + data: bytes, + filename: str, + *, + mime_type: str | None = None, + text: str | None = None, + is_html: bool | None = None, + timeout: float | None = 60 * 60) -> str: + """ + Send the file to `room`. + + Args: + - room - the room to send the file to + - data - the content of the file to send + - filename - filename to use for the file + - text - 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_file_bytes( + room=room, + data=data, + filename=filename, + mime_type=mime_type, + text=text, + is_html=is_html, + timeout=timeout + ) + @overload async def download_file(self, source: EventContext | Event, diff --git a/src/mab/context.py b/src/mab/context.py index bb5ab42..06a6dce 100644 --- a/src/mab/context.py +++ b/src/mab/context.py @@ -34,6 +34,15 @@ CTX_CMD_ARGS = ContextDataKey[list[str]]("CTX_CMD_ARGS") CTX_ROOM_ENCRYPTED = ContextDataKey[bool]("CTX_ROOM_ENCRYPTED") """True if the room is encrypted""" +CTX_FILE_SIZE = ContextDataKey[int]("CTX_FILE_SIZE") +"""Size of the file attached to the message (bytes)""" + +CTX_FILE_MIME = ContextDataKey[str]("CTX_FILE_MIME") +"""Mime type of the file attached to the message""" + +CTX_FILE_NAME = ContextDataKey[str]("CTX_FILE_NAME") +"""Name of the file attached to the message""" + # # EventContext implementation diff --git a/src/mab/filters/message.py b/src/mab/filters/message.py index a200242..2c4a770 100644 --- a/src/mab/filters/message.py +++ b/src/mab/filters/message.py @@ -1,7 +1,12 @@ import traceback from .base import BaseEventFilter, EventTypeFilter from ..types import MessageType -from ..context import EventContext, CTX_MESSAGE_TYPE, CTX_SENDER +from ..context import (EventContext, + CTX_MESSAGE_TYPE, + CTX_SENDER, + CTX_FILE_SIZE, + CTX_FILE_MIME, + CTX_FILE_NAME) from nio import AsyncClient from nio import MatrixRoom, Event @@ -118,4 +123,29 @@ class SenderIsBotFilter(BaseEventFilter): async def __call__(self, context: EventContext) -> bool: if not await super().__call__(context): return False - return context.bot.get_client().user_id == context.event.sender \ No newline at end of file + return context.bot.get_client().user_id == context.event.sender + +class MessageHasFile(BaseEventFilter): + """ + This filter returns True if the message contains file that can be + downloaded. + + This filter sets `CTX_FILE_NAME`, `CTX_FILE_SIZE` and `CTX_FILE_MIME` + variables in the context (if they are present in the). + """ + def __init__(self, **kwargs): + super().__init__(**kwargs) + + async def __call__(self, context: EventContext) -> bool: + if not await super().__call__(context): + return False + content: dict | None = context.event.source.get("content") + if content is None: + return False + # `file` if encrypted, `url` if not encrypted + if "file" not in content and "url" not in content: + return False + context[CTX_FILE_NAME] = content.get("filename") or content.get("body") + context[CTX_FILE_SIZE] = content["info"].get("size") + context[CTX_FILE_MIME] = content["info"].get("mimetype") + return True \ No newline at end of file