Compare commits
7 Commits
v0.0.1
...
0ea2b0e51a
| Author | SHA1 | Date | |
|---|---|---|---|
| 0ea2b0e51a | |||
|
|
363b9a3d45 | ||
|
|
bc3a7500e2 | ||
|
|
8dd1f55fdd | ||
| b62473c468 | |||
| d8bf1818d7 | |||
| d41b588adc |
1
.gitignore
vendored
1
.gitignore
vendored
@@ -1,4 +1,5 @@
|
||||
__pycache__/
|
||||
*.vscode
|
||||
.venv/
|
||||
dist/
|
||||
*.egg-info/
|
||||
|
||||
14
README.md
14
README.md
@@ -16,24 +16,26 @@ python -m pip install git+https://git.tyukalov.su/nikita/mab
|
||||
|
||||
## Basic usage
|
||||
|
||||
This is the most simple bot you can create
|
||||
This is the most simple bot you can create. It would respond to any message
|
||||
that starts with `!test`, `!hello` or `!hi`.
|
||||
|
||||
```python
|
||||
import asyncio
|
||||
from mab import MatrixBot, MatrixBotConfig
|
||||
from mab import TextCommandFilter
|
||||
|
||||
from nio import MatrixRoom, MatrixMessageText
|
||||
|
||||
bot: MatrixBot
|
||||
|
||||
async def on_room_message_text(room: MatrixRoom, event: RoomMessageText) -> None:
|
||||
async def on_valid_command(room: MatrixRoom, event: RoomMessageText) -> None:
|
||||
global bot
|
||||
# do not respond to ourselves
|
||||
if event.sender == bot.get_client().user_id:
|
||||
return
|
||||
text = f"You message contains {len(event.body)} symbols"
|
||||
text = f"Your message contains {len(event.body)} symbols"
|
||||
await bot.send_text_to_room(room.room_id, text)
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
global bot
|
||||
# create and start the bot
|
||||
@@ -43,6 +45,10 @@ async def main() -> None:
|
||||
storage_directory=Path("storage_nagibator666")
|
||||
)
|
||||
bot = MatrixBot(matrix_bot_config)
|
||||
bot.add_event_callback(
|
||||
on_valid_command,
|
||||
TextCommandFilter(["test", "hello", "hi"])
|
||||
)
|
||||
bot.start()
|
||||
# wait for Ctrl+C
|
||||
try:
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "mab"
|
||||
version = "0.0.1"
|
||||
version = "0.2.0"
|
||||
authors = [
|
||||
{ name = "Tyukalov Nikita", email = "nikita@tyukalov.su" }
|
||||
]
|
||||
@@ -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]
|
||||
|
||||
@@ -5,6 +5,9 @@ from .types import MatrixBotConfig
|
||||
|
||||
from .bot import MatrixBot
|
||||
|
||||
from .filters.base import *
|
||||
from .filters.text import *
|
||||
|
||||
__all__ = [
|
||||
# module names
|
||||
"bot",
|
||||
@@ -14,5 +17,16 @@ __all__ = [
|
||||
"MatrixBotConfig",
|
||||
|
||||
# .bot
|
||||
"MatrixBot"
|
||||
"MatrixBot",
|
||||
|
||||
# .filters.base
|
||||
"BaseEventFilter",
|
||||
|
||||
# .filters.text
|
||||
"TextFilter",
|
||||
"FormattedTextFilter",
|
||||
"TextContainsFilter",
|
||||
"TextStartsWithFilter",
|
||||
"TextEndsWithFilter",
|
||||
"TextCommandFilter",
|
||||
]
|
||||
421
src/mab/bot.py
421
src/mab/bot.py
@@ -1,5 +1,5 @@
|
||||
import asyncio
|
||||
import aiofiles
|
||||
import aiofiles, aiofiles.os
|
||||
import aioconsole
|
||||
import traceback
|
||||
import logging
|
||||
@@ -7,23 +7,31 @@ 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
|
||||
|
||||
import nio.events
|
||||
from .filters.base import BaseEventFilter
|
||||
|
||||
from nio.events.room_events import Event as RoomEvemt
|
||||
|
||||
from .types import *
|
||||
from .utils import Utils
|
||||
|
||||
|
||||
class MatrixBot:
|
||||
@@ -106,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.
|
||||
|
||||
Args:
|
||||
text - `text` as passed to the `send_*` function
|
||||
is_html - `is_html` as passed to the `send_*` function
|
||||
|
||||
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)
|
||||
`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)
|
||||
parser.feed(text_to_check)
|
||||
try:
|
||||
if has_tags:
|
||||
return (True, " ".join("".join(text_fragments).split()))
|
||||
parser.feed(text)
|
||||
except:
|
||||
traceback.print_exc()
|
||||
return (False, text_to_check)
|
||||
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:
|
||||
@@ -178,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"
|
||||
@@ -213,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
|
||||
#
|
||||
@@ -244,13 +429,34 @@ class MatrixBot:
|
||||
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
|
||||
|
||||
@@ -263,6 +469,7 @@ class MatrixBot:
|
||||
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}'")
|
||||
@@ -369,7 +576,8 @@ class MatrixBot:
|
||||
sync_task = asyncio.create_task(
|
||||
self._client_cancellable_sync_forever(
|
||||
timeout=self.MATRIX_SYNC_PERIOD,
|
||||
since=(await self._read_next_batch())
|
||||
since=(await self._read_next_batch()),
|
||||
full_state=True
|
||||
)
|
||||
)
|
||||
try:
|
||||
@@ -387,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
|
||||
@@ -402,6 +612,7 @@ class MatrixBot:
|
||||
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.
|
||||
@@ -449,18 +660,17 @@ class MatrixBot:
|
||||
result = True
|
||||
return result
|
||||
|
||||
def add_event_callback(self, callback: Callable[[Any], Awaitable[None]] | None, event_class: nio.events.Event) -> None:
|
||||
"""Added event callback for events of specified class.
|
||||
Use `None` instead of callback to print parameter types you need to use in your callback."""
|
||||
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._client.add_event_callback(callback, event_class) # type: ignore
|
||||
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, **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:
|
||||
@@ -471,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
|
||||
139
src/mab/filters/base.py
Normal file
139
src/mab/filters/base.py
Normal file
@@ -0,0 +1,139 @@
|
||||
from abc import ABC, abstractmethod
|
||||
import logging
|
||||
|
||||
from nio import MatrixRoom, Event
|
||||
|
||||
class BaseEventFilter(ABC):
|
||||
"""Base class for all message filters"""
|
||||
_logger = logging.Logger("EventFilter")
|
||||
|
||||
# AND
|
||||
def __and__(self, other):
|
||||
if not isinstance(other, BaseEventFilter):
|
||||
raise TypeError(f"BaseEventFilter can't be ANDed against {type(other)}")
|
||||
return CompoundEventFilter(
|
||||
CompoundEventFilter.OPERATOR_AND,
|
||||
[self, other]
|
||||
)
|
||||
|
||||
def __rand__(self, other):
|
||||
return self.__and__(other)
|
||||
|
||||
# OR
|
||||
def __or__(self, other):
|
||||
if not isinstance(other, BaseEventFilter):
|
||||
raise TypeError(f"BaseEventFilter can't be ORed against {type(other)}")
|
||||
return CompoundEventFilter(
|
||||
CompoundEventFilter.OPERATOR_OR,
|
||||
[self, other]
|
||||
)
|
||||
|
||||
def __ror__(self, other):
|
||||
return self.__ror__(other)
|
||||
|
||||
# XOR
|
||||
def __xor__(self, other):
|
||||
if not isinstance(other, BaseEventFilter):
|
||||
raise TypeError(f"BaseEventFilter can't be XORed against {type(other)}")
|
||||
return CompoundEventFilter(
|
||||
CompoundEventFilter.OPERATOR_XOR,
|
||||
[self, other]
|
||||
)
|
||||
|
||||
def __rxor__(self, other):
|
||||
return self.__xor__(other)
|
||||
|
||||
# INVERT
|
||||
def __invert__(self):
|
||||
return CompoundEventFilter(
|
||||
CompoundEventFilter.OPERATOR_INVERT,
|
||||
[self]
|
||||
)
|
||||
|
||||
# PAYLOAD
|
||||
@abstractmethod
|
||||
def __repr__(self) -> str:
|
||||
"""This method must be redefined in derived classes to improve
|
||||
debugging experience.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def __call__(self, room: MatrixRoom, event: Event) -> bool:
|
||||
"""This abstract method must be redefined in derived classes so that
|
||||
the filter operates according to its description. This method must
|
||||
not raise exceptions. In case of exception it should log it using
|
||||
`self._logger` and return False
|
||||
|
||||
Args:
|
||||
event - the event to check againts this filter
|
||||
|
||||
Returns:
|
||||
True if the event satisfies this filter
|
||||
False if the event does not satisfy this filter
|
||||
"""
|
||||
pass
|
||||
|
||||
class CompoundEventFilter(BaseEventFilter):
|
||||
"""Event filter that consists of multiple filters"""
|
||||
OPERATOR_AND = "and"
|
||||
OPERATOR_OR = "or"
|
||||
OPERATOR_XOR = "xor"
|
||||
OPERATOR_INVERT = "invert"
|
||||
|
||||
@staticmethod
|
||||
def _is_operator_valid(op: str) -> bool:
|
||||
"""Returns True if operator `op` is a valid operator"""
|
||||
return op in [
|
||||
CompoundEventFilter.OPERATOR_AND,
|
||||
CompoundEventFilter.OPERATOR_OR,
|
||||
CompoundEventFilter.OPERATOR_XOR,
|
||||
CompoundEventFilter.OPERATOR_INVERT
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def _is_elements_count_valid(op: str, count: int) -> bool:
|
||||
"""Returns True if operator `op` may take `count` arguments"""
|
||||
return count in {
|
||||
CompoundEventFilter.OPERATOR_AND: [2],
|
||||
CompoundEventFilter.OPERATOR_OR: [2],
|
||||
CompoundEventFilter.OPERATOR_XOR: [2],
|
||||
CompoundEventFilter.OPERATOR_INVERT: [1],
|
||||
}[op]
|
||||
|
||||
def __init__(self, operator: str, arguments: list[BaseEventFilter]):
|
||||
super().__init__()
|
||||
if not self._is_operator_valid(operator):
|
||||
raise RuntimeError(f"Invalid operator `{operator}`")
|
||||
if not self._is_elements_count_valid(operator, len(arguments)):
|
||||
raise RuntimeError(f"Operator `{operator}` does not take `{len(arguments)}` arguments")
|
||||
self._operator = operator
|
||||
self._arguments = list(arguments)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
expression = "False"
|
||||
reprs = [repr(a) for a in self._arguments]
|
||||
if self._operator == CompoundEventFilter.OPERATOR_AND:
|
||||
expression = " & ".join(reprs)
|
||||
elif self._operator == CompoundEventFilter.OPERATOR_OR:
|
||||
expression = " | ".join(reprs)
|
||||
elif self._operator == CompoundEventFilter.OPERATOR_XOR:
|
||||
expression = " ^ ".join(reprs)
|
||||
elif self._operator == CompoundEventFilter.OPERATOR_INVERT:
|
||||
expression = f"~{reprs[0]}"
|
||||
return f"({expression})"
|
||||
|
||||
def __call__(self, room: MatrixRoom, event: Event) -> bool:
|
||||
evaluated = [arg(room, event) for arg in self._arguments]
|
||||
if self._operator == self.OPERATOR_AND:
|
||||
return all(evaluated)
|
||||
elif self._operator == self.OPERATOR_OR:
|
||||
return any(evaluated)
|
||||
elif self._operator == self.OPERATOR_XOR:
|
||||
result = evaluated[0]
|
||||
for v in evaluated[1:]:
|
||||
result ^= v
|
||||
return result
|
||||
elif self._operator == self.OPERATOR_INVERT:
|
||||
return not evaluated[0]
|
||||
return False
|
||||
164
src/mab/filters/text.py
Normal file
164
src/mab/filters/text.py
Normal file
@@ -0,0 +1,164 @@
|
||||
from .base import BaseEventFilter
|
||||
|
||||
from nio import MatrixRoom, Event
|
||||
|
||||
class TextFilter(BaseEventFilter):
|
||||
"""
|
||||
This filter returns True if the event contains `body` attribute.
|
||||
`body` attribute contains unformatted text, string.
|
||||
"""
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def __repr__(self):
|
||||
return "TextFilter()"
|
||||
|
||||
def __call__(self, room: MatrixRoom, event: Event) -> bool:
|
||||
return hasattr(event, "body") and type(event.body) is str # type: ignore
|
||||
|
||||
class FormattedTextFilter(BaseEventFilter):
|
||||
"""
|
||||
This filter returns True if the event contains valid `formatted_body`
|
||||
attribute. `formatted_body` attribute contains formatted text, string.
|
||||
"""
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def __repr__(self):
|
||||
return "FormattedTextFilter()"
|
||||
|
||||
def __call__(self, room: MatrixRoom, event: Event) -> bool:
|
||||
return hasattr(event, "formatted_body") and type(event.formatted_body) is str # type: ignore
|
||||
|
||||
class TextContainsFilter(BaseEventFilter):
|
||||
"""
|
||||
This filter returns True if the `event.body` contains `needle`
|
||||
substring (or any of neddle from the list). The check will be case
|
||||
insensetive if `any_case` is True.
|
||||
"""
|
||||
def __init__(self, needle: str | list[str], *, any_case: bool = True):
|
||||
super().__init__()
|
||||
if type(needle) is str:
|
||||
needle = [needle]
|
||||
self._any_case = any_case
|
||||
if self._any_case:
|
||||
self._needle = [s.lower() for s in needle]
|
||||
else:
|
||||
self._needle = list(needle)
|
||||
|
||||
def __repr__(self):
|
||||
return f"TextContainsFilter({repr(self._needle)}, any_case={repr(self._any_case)})"
|
||||
|
||||
def __call__(self, room: MatrixRoom, event: Event) -> bool:
|
||||
try:
|
||||
body = event.body.lower() if self._any_case else event.body # type: ignore
|
||||
for n in self._needle:
|
||||
if n in body:
|
||||
return True
|
||||
return False
|
||||
except:
|
||||
return False
|
||||
|
||||
class TextStartsWithFilter(BaseEventFilter):
|
||||
"""
|
||||
This filter returns True if the `event.body` starts with `substring` (or
|
||||
any of substrings from the list). The check will be case insensetive if
|
||||
`any_case` is True.
|
||||
"""
|
||||
def __init__(self, substring: str | list[str], *, any_case: bool = True):
|
||||
super().__init__()
|
||||
if type(substring) is str:
|
||||
substring = [substring]
|
||||
self._any_case = any_case
|
||||
if self._any_case:
|
||||
self._substring = [s.lower() for s in substring]
|
||||
else:
|
||||
self._substring = list(substring)
|
||||
|
||||
def __repr__(self):
|
||||
return f"TextStartsWithFilter({repr(self._substring)}, any_case={repr(self._any_case)})"
|
||||
|
||||
def __call__(self, room: MatrixRoom, event: Event) -> bool:
|
||||
try:
|
||||
body = event.body.lower() if self._any_case else event.body # type: ignore
|
||||
for s in self._substring:
|
||||
if body.startswith(s):
|
||||
return True
|
||||
return False
|
||||
except:
|
||||
return False
|
||||
|
||||
class TextEndsWithFilter(BaseEventFilter):
|
||||
"""
|
||||
This filter returns True if the `event.body` ends with `substring` (or
|
||||
any of substrings from the list). The check will be case insensetive if
|
||||
`any_case` is True.
|
||||
"""
|
||||
def __init__(self, substring: str | list[str], *, any_case: bool = True):
|
||||
super().__init__()
|
||||
if type(substring) is str:
|
||||
substring = [substring]
|
||||
self._any_case = any_case
|
||||
if self._any_case:
|
||||
self._substring = [s.lower() for s in substring]
|
||||
else:
|
||||
self._substring = list(substring)
|
||||
|
||||
def __repr__(self):
|
||||
return f"TextEndsWithFilter({repr(self._substring)}, any_case={repr(self._any_case)})"
|
||||
|
||||
def __call__(self, room: MatrixRoom, event: Event) -> bool:
|
||||
try:
|
||||
body = event.body.lower() if self._any_case else event.body # type: ignore
|
||||
for s in self._substring:
|
||||
if body.endswith(s):
|
||||
return True
|
||||
return False
|
||||
except:
|
||||
return False
|
||||
|
||||
class TextCommandFilter(BaseEventFilter):
|
||||
"""
|
||||
This filter returns True if all conditions are True:
|
||||
1. `event.body` contains at least `min_args + 1` words after split()
|
||||
2. `event.body` conrains at most `max_args + 1` words after split()
|
||||
3. First element of splitted `event.body` starts with `prefix`
|
||||
4. First element of splitted `event.body` (after stripping `prefix`)
|
||||
starts with any of strings in `verbs` list (case-insensitive)
|
||||
|
||||
Remarks:
|
||||
- If this filter is satified, then it will set a new attribute for the
|
||||
event: `event.command_args: list[str]`. You may use this attribute in
|
||||
your callback for this event.
|
||||
"""
|
||||
def __init__(self, verbs: str | list[str], min_args: int = 0, max_args: int | None = None, prefix: str = "!"):
|
||||
super().__init__()
|
||||
if type(verbs) is str:
|
||||
verbs = [verbs]
|
||||
verbs = [v.lower() for v in verbs]
|
||||
self._verbs = verbs
|
||||
self._min_args = min_args
|
||||
self._max_args = max_args
|
||||
self._prefix = prefix
|
||||
|
||||
def __repr__(self):
|
||||
return f"TextCommandFilter({repr(self._verbs)}, {repr(self._min_args)}, {repr(self._max_args)}, {repr(self._prefix)})"
|
||||
|
||||
def __call__(self, room: MatrixRoom, event: Event) -> bool:
|
||||
try:
|
||||
parts = [p.strip() for p in event.body.split() if p.strip()] # type: ignore
|
||||
args_count = len(parts) - 1
|
||||
if args_count < self._min_args:
|
||||
return False
|
||||
if self._max_args is not None and args_count > self._max_args:
|
||||
return False
|
||||
if not parts[0].startswith(self._prefix):
|
||||
return False
|
||||
cmd = parts[0][len(self._prefix):].lower()
|
||||
for verb in self._verbs:
|
||||
if cmd == verb:
|
||||
setattr(event, "command_args", parts[1:])
|
||||
return True
|
||||
return False
|
||||
except:
|
||||
return False
|
||||
@@ -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
98
src/mab/utils.py
Normal 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
|
||||
Reference in New Issue
Block a user