5 Commits

Author SHA1 Message Date
2867582c01 Got rid of proxy methods in MatrixBot, v0.6.0
- Updated README.md
- Updated examples
- `MatrixBot` does not redefine methods from classes it is composed of anymore
2026-09-27 16:17:40 +03:00
16b849ebd8 Added file uploading support, v0.5.2
- Added example `file_bot.py`
- Added *.tmp to .gitignore
- Added CTX_FILE_SIZE, CTX_FILE_MIME and CTX_FILE_NAME
- Added MessageHasFile filter
- Added `send_file` and `send_file_bytes` for ClientSender and Bot
2026-09-23 19:43:30 +03:00
35bd68516a Fixed README.md link to image_bot.py 2026-09-13 04:18:43 +03:00
86257a70e3 Added overloads for MatrixBot.download_file 2026-09-13 03:39:23 +03:00
43b5d8e3a6 Fixed downloader bugs related to temp file name 2026-09-13 03:28:27 +03:00
13 changed files with 347 additions and 215 deletions

1
.gitignore vendored
View File

@@ -6,3 +6,4 @@ dist/
*.egg-info/
*.swp
*.swo
*.tmp

View File

@@ -10,6 +10,7 @@ The library supports the following features:
- **Completely `asyncio` based**
- **Filter-based callback system**
- **Downloading and transparently decrypting files**
- **Sending files**
- **Sending images**
- **Sending videos with automatic thumbnail generation (requires `ffmpeg`)**
@@ -22,13 +23,13 @@ install the latest version of the library:
```bash
apt install libmagic1-dev libolm-dev
python -m pip install git+https://git.tyukalov.su/nikita/mab@v0.5.1
python -m pip install git+https://git.tyukalov.su/nikita/mab@v0.6.0
```
`libmagic1-dev` is needed for automatic file MIME type detection, `libolm-dev`
is needed for E2EE to work.
Please inspect [`examples/image_gen_bot.py`](examples/image_gen_bot.py),
Please inspect [`examples/image_bot.py`](examples/image_bot.py),
[`examples/echo_bot.py`](examples/echo_bot.py) or open [`examples/`](examples/)
directory to find usage examples. Examples require that you set
`MATRIX_HOMESERVER` and `MATRIX_USERNAME` environment variables. Examples create

View File

@@ -29,18 +29,18 @@ async def on_help_command(ctx: EventContext) -> None:
"<li><code>!args arg1 [arg2] ... [arg5]</code> - command that takes 1..5 arguments</li>"
"</ul>"
)
await ctx.bot.send_text(ctx.room, HELP_MESSAGE)
await ctx.bot.send.text(ctx.room, HELP_MESSAGE)
async def on_time_command(ctx: EventContext) -> None:
"""!time"""
await ctx.bot.send_text(
await ctx.bot.send.text(
ctx.room,
f"Current UNIX timestamp is <strong>{int(time.time())}</strong>"
)
async def on_raise_command(ctx: EventContext) -> None:
"""!raise"""
await ctx.bot.send_text(
await ctx.bot.send.text(
ctx.room,
"<strong>Executing <code>raise RuntimeError()</code>...</strong>"
)
@@ -48,7 +48,7 @@ async def on_raise_command(ctx: EventContext) -> None:
async def on_assert_command(ctx: EventContext) -> None:
"""!assert"""
await ctx.bot.send_text(
await ctx.bot.send.text(
ctx.room,
"<strong>Executing <code>assert False</code>...</strong>"
)
@@ -63,9 +63,9 @@ async def on_mul_command(ctx: EventContext) -> None:
v *= n
response = " * ".join(html.escape("%.2f" % n) for n in numbers)
response += f" = <strong>{html.escape(str(v))}<strong>"
await ctx.bot.send_text(ctx.room, response)
await ctx.bot.send.text(ctx.room, response)
except Exception as e:
await ctx.bot.send_text(ctx.room, f"Could not process the command: {e}")
await ctx.bot.send.text(ctx.room, f"Could not process the command: {e}")
async def on_args_command(ctx: EventContext) -> None:
"""!args"""
@@ -79,19 +79,19 @@ async def on_args_command(ctx: EventContext) -> None:
for arg in ctx[CTX_CMD_ARGS]:
response += f"<li><code>{html.escape(arg)}</code></li>"
response += "</ol>"
await ctx.bot.send_text(
await ctx.bot.send.text(
ctx.room,
response
)
except:
await ctx.bot.send_text(
await ctx.bot.send.text(
ctx.room,
f"Could not process the command: {traceback.format_exc()}"
)
async def invalid_usage(ctx: EventContext) -> None:
"""This callback is called when the bot used incorrectly."""
await ctx.bot.send_text(ctx.room, "Use <code>!help</code>")
await ctx.bot.send.text(ctx.room, "Use <code>!help</code>")
async def main() -> None:

View File

@@ -15,7 +15,7 @@ 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, ctx[CTX_BODY])
await ctx.bot.send.text(ctx.room, ctx[CTX_BODY])
async def main() -> None:
"""Application entry point"""

97
examples/file_bot.py Normal file
View File

@@ -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())

View File

@@ -24,9 +24,9 @@ async def on_gen_command(data: EventContext) -> None:
try:
r, g, b = [float(v) for v in data[CTX_CMD_ARGS]]
except:
await data.bot.send_text(data.room, "Invalid arguments")
await data.bot.send.text(data.room, "Invalid arguments")
return
await data.bot.send_text(data.room, "Generating the noise...")
await data.bot.send.text(data.room, "Generating the noise...")
# create the basic noise
img = Image.new("RGB", (16, 16))
for x in range(img.width):
@@ -44,17 +44,14 @@ async def on_gen_command(data: EventContext) -> None:
buf.seek(0)
buf = buf.read()
# send
await data.bot.send_image_bytes(data.room, buf, "noise.png")
await data.bot.send.image_bytes(data.room, buf, "noise.png")
async def on_image(ctx: EventContext) -> None:
"""This callback is called when an image is received."""
# download
s = time.time()
data = await ctx.bot.download_file(ctx)
data = await ctx.bot.download.file(ctx)
took_time = time.time() - s
if not isinstance(data, bytes):
await ctx.bot.send_text(ctx.room, "😧")
return
# convert to Image
with BytesIO(data) as buf:
img = Image.open(buf)
@@ -69,7 +66,7 @@ async def on_image(ctx: EventContext) -> None:
buf.seek(0)
buf = buf.read()
# send
await ctx.bot.send_image_bytes(
await ctx.bot.send.image_bytes(
ctx.room,
buf,
"blurred.png",
@@ -78,7 +75,7 @@ async def on_image(ctx: EventContext) -> None:
async def on_wrong_message(data: EventContext) -> None:
"""This callback is called when a wrong message is received."""
await data.bot.send_text(
await data.bot.send.text(
data.room,
"Text me something like <code>!gen 0.1 0.7 1.0</code> or send an image to blur"
)

View File

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
[project]
name = "mab"
version = "0.5.1"
version = "0.6.0"
authors = [
{ name = "Tyukalov Nikita", email = "nikita@tyukalov.su" }
]

View File

@@ -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",
]

View File

@@ -2,10 +2,9 @@ import asyncio
import aiofiles
import logging
import os
import time
from typing import overload
import threading
import traceback
import random
from uuid import uuid4
from pathlib import Path
import hashlib
@@ -36,7 +35,7 @@ class ClientDownloader:
config: MatrixBotConfig,
client: AsyncClient) -> None:
"""
Setup the sender.
Setup the downloader.
Args:
- config - config to use
@@ -151,7 +150,19 @@ class ClientDownloader:
cancel.set()
raise
async def download_file(self,
@overload
async def file(self,
source: EventContext | Event,
*,
path: str | Path) -> Path: ...
@overload
async def file(self,
source: EventContext | Event,
*,
path: None = None) -> bytes: ...
async def file(self,
source: EventContext | Event,
*,
path: str | Path | None = None) -> Path | bytes:
@@ -218,14 +229,12 @@ class ClientDownloader:
if key is not None and iv is not None and sha256 is not None:
# on disk
if isinstance(result, Path):
r = unpaddedbase64.encode_base64(random.randbytes(8))
temp_path = result.with_name(
f".{result.name}.{int(time.time())}{r}.tmp"
f".{result.name}.{uuid4().hex}.tmp"
)
await self._decrypt_file(
result, temp_path, key, iv, sha256
)
result.unlink()
temp_path.replace(result)
temp_path = None
# in memory

View File

@@ -98,9 +98,9 @@ class ClientSender:
self._client = client
self._uploader = uploader
async def send_content(self,
room: MatrixRoom | str,
content: dict) -> RoomSendResponse:
async def content(self,
room: MatrixRoom | str,
content: dict) -> RoomSendResponse:
"""
Send raw `content` JSON to `room`.
@@ -124,7 +124,7 @@ class ClientSender:
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)
return await self.content(room, content)
else:
raise
if type(result) is RoomSendResponse:
@@ -134,11 +134,11 @@ class ClientSender:
else:
raise RuntimeError("Unknown error has occured", result)
async def send_text(self,
room: MatrixRoom | str,
text: str,
*,
is_html: bool | None = None) -> str:
async def text(self,
room: MatrixRoom | str,
text: str,
*,
is_html: bool | None = None) -> str:
"""
Send text message to `room`.
@@ -156,15 +156,15 @@ class ClientSender:
"msgtype": "m.text",
**text_data
}
return (await self.send_content(room, content)).event_id
return (await self.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:
async def 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
@@ -213,16 +213,16 @@ class ClientSender:
"h": height
}
}
return (await self.send_content(room, content)).event_id
return (await self.content(room, content)).event_id
async def send_image_bytes(self,
room: MatrixRoom | str,
data: bytes,
filename: str,
*,
text: str | None = None,
is_html: bool | None = None,
timeout: float | None = 60 * 60) -> str:
async def image_bytes(self,
room: MatrixRoom | str,
data: bytes,
filename: str,
*,
text: str | None = None,
is_html: bool | 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
@@ -276,16 +276,16 @@ class ClientSender:
"h": height
}
}
return (await self.send_content(room, content)).event_id
return (await self.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:
async def 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
@@ -361,4 +361,123 @@ class ClientSender:
}
}
# send
return (await self.send_content(room, content)).event_id
return (await self.content(room, content)).event_id
async def 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.content(room, content)).event_id
async def 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.content(room, content)).event_id

View File

@@ -1,7 +1,7 @@
import asyncio
import logging
from typing import Callable, Coroutine, Any
from typing import Callable, Coroutine, Any, overload
from nio import AsyncClient, MatrixRoom, Event
@@ -119,7 +119,8 @@ class MatrixBot:
finally:
await self.stop()
def get_client(self) -> AsyncClient:
@property
def client(self) -> AsyncClient:
"""
Get AsyncClient.
@@ -129,148 +130,12 @@ class MatrixBot:
"""
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`.
@property
def send(self) -> ClientSender:
"""Get ClientSender that you should use to send messages."""
return self._client_sender
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_image_bytes(self,
room: MatrixRoom | str,
data: bytes,
filename: str,
*,
text: str | None = None,
is_html: bool | 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
- bytes - the image to send
- filename - filename to use for the file
- text - image 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_image_bytes(
room=room,
data=data,
filename=filename,
text=text,
is_html=is_html,
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
)
async def download_file(self,
source: EventContext | Event,
*,
path: str | Path | None = None) -> Path | bytes:
"""
Download the file from the event. Automatically dechiphers encrypted
media.
Args:
- source - event that has a file in its `content`
- path - where to save the file to. Use `None` to store the file in
memory. Use path to specify file download path.
Returns:
- Returns `Path` to the file on success (if `path` isn't `None`)
- Returns `bytes` of the file on success (if `path` is `None`)
- Raises an exception on error
"""
return await self._client_downloader.download_file(
source=source,
path=path
)
@property
def download(self) -> ClientDownloader:
"""Get ClientDownloader that you should use to download files."""
return self._client_downloader

View File

@@ -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

View File

@@ -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
return context.bot.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