nikita/mab
1
0
Fork 0
Code Issues 2 Pull Requests Actions Packages Projects Releases Wiki Activity

Image/video support. Related upgrades.

Browse Source 
- Removed `_check_and_remove_html`. Using `_process_html_text` instead
- Added method for file uploading `_upload_file`
- Added method for raw bytes uploading `_upload_provider`
- Added method for thumbnail uploading `_upload_thumbnail`
- Added method that sends `content` JSON to the room using `m.room.message` type
- Fixed `self._debug_event_callback` not being called in `_callback_filter_router`
- Fixed type hints in `add_event_callback`
- Removed `kwargs` from `send_text_to_room`
- Added `send_image_to_room`
- Added `send_video_to_room`
- Added `auto_ffprobe` and `auto_ffmpeg` to `MatrixBotConfig`
- Added `VideoFileProperties` dataclass
This commit is contained in:
Nikita Tyukalov, ASUS, Linux
2026-09-07 01:10:13 +03:00
parent bc3a7500e2
commit 363b9a3d45
4 changed files with 477 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
pyproject.toml
View File

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

412
src/mab/bot.py
View File

@@ -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,30 +114,58 @@ 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.
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.
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)
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:
@@ -180,6 +214,7 @@ class MatrixBot:
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"
@@ -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)
# 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

22
src/mab/types.py
View File

@@ -21,3 +21,25 @@ class MatrixBotConfig:
auto_join_any_room_on_invite: bool = True
"""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"""

98
src/mab/utils.py Normal file
View File

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