Compare commits
2 Commits
0ea2b0e51a
...
v0.3.0
| Author | SHA1 | Date | |
|---|---|---|---|
| 90c5ae7a78 | |||
| 1ba1094a42 |
35
README.md
35
README.md
@@ -6,14 +6,24 @@ because I wasn't satisfied by simplicity and usage of other libraries. So
|
|||||||
this library does not aim to be "the best matrix bot library", it only aims to
|
this library does not aim to be "the best matrix bot library", it only aims to
|
||||||
be good enough for me.
|
be good enough for me.
|
||||||
|
|
||||||
|
## Features
|
||||||
|
|
||||||
|
The package supports the following features:
|
||||||
|
- **Filter-based callback system**
|
||||||
|
- **Images sending**
|
||||||
|
- **Videos sending with automatic thumbnail generation (requires `ffmpeg`)**
|
||||||
|
|
||||||
## Installation
|
## Installation
|
||||||
|
|
||||||
Use `pip` to install this package:
|
Use `pip` to install this package:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
python -m pip install git+https://git.tyukalov.su/nikita/mab
|
python -m pip install git+https://git.tyukalov.su/nikita/mab@v0.3.0
|
||||||
```
|
```
|
||||||
|
|
||||||
|
You should specify package version you want to use, because `main` without tags
|
||||||
|
contain unstable code.
|
||||||
|
|
||||||
## Basic usage
|
## Basic usage
|
||||||
|
|
||||||
This is the most simple bot you can create. It would respond to any message
|
This is the most simple bot you can create. It would respond to any message
|
||||||
@@ -23,21 +33,16 @@ that starts with `!test`, `!hello` or `!hi`.
|
|||||||
import asyncio
|
import asyncio
|
||||||
from mab import MatrixBot, MatrixBotConfig
|
from mab import MatrixBot, MatrixBotConfig
|
||||||
from mab import TextCommandFilter
|
from mab import TextCommandFilter
|
||||||
|
from mab.types import RoomEventData
|
||||||
|
|
||||||
from nio import MatrixRoom, MatrixMessageText
|
async def on_valid_command(data: RoomEventData) -> None:
|
||||||
|
|
||||||
bot: MatrixBot
|
|
||||||
|
|
||||||
async def on_valid_command(room: MatrixRoom, event: RoomMessageText) -> None:
|
|
||||||
global bot
|
|
||||||
# do not respond to ourselves
|
# do not respond to ourselves
|
||||||
if event.sender == bot.get_client().user_id:
|
if event.sender == data.bot.get_client().user_id:
|
||||||
return
|
return
|
||||||
text = f"Your message contains {len(event.body)} symbols"
|
text = f"Your message contains {len(data.event.body)} symbols"
|
||||||
await bot.send_text_to_room(room.room_id, text)
|
await data.bot.send_text_to_room(data.room, text)
|
||||||
|
|
||||||
async def main() -> None:
|
async def main() -> None:
|
||||||
global bot
|
|
||||||
# create and start the bot
|
# create and start the bot
|
||||||
cfg = MatrixBotConfig(
|
cfg = MatrixBotConfig(
|
||||||
matrix_homeserver_url="matrix.domain.su",
|
matrix_homeserver_url="matrix.domain.su",
|
||||||
@@ -45,11 +50,11 @@ async def main() -> None:
|
|||||||
storage_directory=Path("storage_nagibator666")
|
storage_directory=Path("storage_nagibator666")
|
||||||
)
|
)
|
||||||
bot = MatrixBot(matrix_bot_config)
|
bot = MatrixBot(matrix_bot_config)
|
||||||
bot.add_event_callback(
|
bot.add_callback(
|
||||||
on_valid_command,
|
TextCommandFilter(["test", "hello", "hi"]),
|
||||||
TextCommandFilter(["test", "hello", "hi"])
|
on_valid_command
|
||||||
)
|
)
|
||||||
bot.start()
|
await bot.start()
|
||||||
# wait for Ctrl+C
|
# wait for Ctrl+C
|
||||||
try:
|
try:
|
||||||
while True:
|
while True:
|
||||||
|
|||||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
|||||||
|
|
||||||
[project]
|
[project]
|
||||||
name = "mab"
|
name = "mab"
|
||||||
version = "0.2.0"
|
version = "0.3.0"
|
||||||
authors = [
|
authors = [
|
||||||
{ name = "Tyukalov Nikita", email = "nikita@tyukalov.su" }
|
{ name = "Tyukalov Nikita", email = "nikita@tyukalov.su" }
|
||||||
]
|
]
|
||||||
|
|||||||
827
src/mab/bot.py
827
src/mab/bot.py
@@ -1,827 +0,0 @@
|
|||||||
import asyncio
|
|
||||||
import aiofiles, aiofiles.os
|
|
||||||
import aioconsole
|
|
||||||
import traceback
|
|
||||||
import logging
|
|
||||||
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
|
|
||||||
|
|
||||||
from .filters.base import BaseEventFilter
|
|
||||||
|
|
||||||
from nio.events.room_events import Event as RoomEvemt
|
|
||||||
|
|
||||||
from .types import *
|
|
||||||
from .utils import Utils
|
|
||||||
|
|
||||||
|
|
||||||
class MatrixBot:
|
|
||||||
"""Asynchronous Matrix Bot Implementation.
|
|
||||||
|
|
||||||
Use objects of this class to build your bots. Manage the event loop
|
|
||||||
by yourself.
|
|
||||||
"""
|
|
||||||
NEXT_BATCH_DUMP_PERIOD = 120.0
|
|
||||||
MATRIX_SYNC_PERIOD = 5000
|
|
||||||
|
|
||||||
#
|
|
||||||
# PRIVATE
|
|
||||||
#
|
|
||||||
@staticmethod
|
|
||||||
def _validate_matrix_homeserver_url(url: str) -> None:
|
|
||||||
"""Checks if `url` is a valid matrix homeserver URL.
|
|
||||||
Raises an Exception if it is not.
|
|
||||||
"""
|
|
||||||
parsed = urlparse(url)
|
|
||||||
if parsed.scheme not in ("http", "https"):
|
|
||||||
raise RuntimeError(
|
|
||||||
f"Scheme {parsed.scheme} is not a valid scheme for matrix homeserver URL"
|
|
||||||
)
|
|
||||||
if not parsed.netloc:
|
|
||||||
raise RuntimeError(f"{url} is not a valid matrix homeserver URL")
|
|
||||||
if parsed.path != "":
|
|
||||||
raise RuntimeError(f"{url} must have empty path (remove `{parsed.path}` after the hostname)")
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def _validate_matrix_username_localpart(localpart: str) -> None:
|
|
||||||
"""Checks if `username` is a valid localpart of matrix username.
|
|
||||||
Raises an Exception if it is not.
|
|
||||||
"""
|
|
||||||
pattern = r"^[a-z0-9._=\-]+$"
|
|
||||||
if not bool(re.fullmatch(pattern, localpart)) or len(localpart) > 255:
|
|
||||||
raise RuntimeError(f"{localpart} is not a valid matrix username localpart")
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def _validate_storage_directory(path: Path) -> None:
|
|
||||||
"""Checks if `path` is a valid storage directory and creates it.
|
|
||||||
Raises an Exception if it is not.
|
|
||||||
"""
|
|
||||||
path.mkdir(parents=True, exist_ok=True)
|
|
||||||
if not path.is_dir():
|
|
||||||
raise RuntimeError(f"Could not create directory {path}")
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def _validate_bot_config(config: MatrixBotConfig) -> None:
|
|
||||||
"""Checks if `config` has errors.
|
|
||||||
Raises an Exception if it does.
|
|
||||||
"""
|
|
||||||
MatrixBot._validate_matrix_homeserver_url(config.matrix_homeserver_url)
|
|
||||||
MatrixBot._validate_matrix_username_localpart(config.matrix_username_localpart)
|
|
||||||
MatrixBot._validate_storage_directory(config.storage_directory)
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def _build_client(config: MatrixBotConfig) -> AsyncClient:
|
|
||||||
"""Builds `nio.AsyncClient` from `MatrixBotConfig`"""
|
|
||||||
# create the config for the client
|
|
||||||
client_config = AsyncClientConfig(
|
|
||||||
store_name="nio_store_file",
|
|
||||||
encryption_enabled=True,
|
|
||||||
store_sync_tokens=False
|
|
||||||
)
|
|
||||||
# create the client
|
|
||||||
client = AsyncClient(
|
|
||||||
homeserver=config.matrix_homeserver_url,
|
|
||||||
user=config.matrix_username_localpart,
|
|
||||||
store_path=str(config.storage_directory),
|
|
||||||
config=client_config
|
|
||||||
)
|
|
||||||
return client
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def _build_matrix_username(config: MatrixBotConfig) -> str:
|
|
||||||
"""Builds full matrix username."""
|
|
||||||
homeserver_name = urlparse(config.matrix_homeserver_url).hostname
|
|
||||||
localpart = config.matrix_username_localpart
|
|
||||||
return f"@{localpart}:{homeserver_name}"
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
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.
|
|
||||||
"""
|
|
||||||
# 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
|
|
||||||
should be cancelled. Returns value returned by `payload_task` task. Raises exception
|
|
||||||
raised by `payload_task` task.
|
|
||||||
|
|
||||||
If `stop_wait_task` finishes, then `asyncio.CancelledError` is raised.
|
|
||||||
"""
|
|
||||||
try:
|
|
||||||
done, _ = await asyncio.wait(
|
|
||||||
[payload_task, stop_wait_task],
|
|
||||||
return_when=asyncio.FIRST_COMPLETED
|
|
||||||
)
|
|
||||||
if stop_wait_task in done:
|
|
||||||
payload_task.cancel()
|
|
||||||
await payload_task
|
|
||||||
raise asyncio.CancelledError()
|
|
||||||
return payload_task.result()
|
|
||||||
except asyncio.CancelledError:
|
|
||||||
raise
|
|
||||||
except:
|
|
||||||
if print_exc:
|
|
||||||
traceback.print_exc()
|
|
||||||
raise
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
async def _default_password_callback() -> str:
|
|
||||||
"""Gets password from `MATRIX_PASSWORD` envvar if it is set. Asks
|
|
||||||
the user for the password otherwise."""
|
|
||||||
if "MATRIX_PASSWORD" in os.environ:
|
|
||||||
return os.environ["MATRIX_PASSWORD"]
|
|
||||||
print("--- A password is required (btw you might use MATRIX_PASSWORD envvar) ---")
|
|
||||||
return await aioconsole.ainput("Password: ")
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
async def _debug_event_callback(*args, **kwargs) -> None:
|
|
||||||
"""Just prints types of arguments"""
|
|
||||||
try:
|
|
||||||
print(f"_debug_event_callback ({len(args)} args, {len(kwargs)} kwargs)")
|
|
||||||
for a in args:
|
|
||||||
print(f" - {type(a)}")
|
|
||||||
for k in kwargs:
|
|
||||||
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."""
|
|
||||||
path = self._config.storage_directory / "next_batch"
|
|
||||||
async with aiofiles.open(path, "w") as f:
|
|
||||||
await f.write(next_batch)
|
|
||||||
self._logger.debug("next_batch value is written to the disk")
|
|
||||||
|
|
||||||
async def _read_next_batch(self) -> str | None:
|
|
||||||
"""Reads `next_batch` value from disk. Returns None if file does not exist."""
|
|
||||||
path = self._config.storage_directory / "next_batch"
|
|
||||||
if not path.is_file():
|
|
||||||
return None
|
|
||||||
async with aiofiles.open(path, "r") as f:
|
|
||||||
return (await f.read()).strip()
|
|
||||||
|
|
||||||
async def _write_session_data(self, *, access_token: str, device_id: str) -> None:
|
|
||||||
"""Write session data to disk."""
|
|
||||||
path = self._config.storage_directory / "session_data.json"
|
|
||||||
data = {
|
|
||||||
"access_token": access_token,
|
|
||||||
"device_id": device_id
|
|
||||||
}
|
|
||||||
async with aiofiles.open(path, "w") as f:
|
|
||||||
await f.write(json.dumps(data, indent=4))
|
|
||||||
self._logger.debug("Session data is writter to the disk")
|
|
||||||
|
|
||||||
async def _read_session_data(self) -> dict[str, Any] | None:
|
|
||||||
"""Read session data from disk."""
|
|
||||||
path = self._config.storage_directory / "session_data.json"
|
|
||||||
if not path.is_file():
|
|
||||||
return None
|
|
||||||
async with aiofiles.open(path, "r") as f:
|
|
||||||
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
|
|
||||||
#
|
|
||||||
async def _callback_sync(self, response: SyncResponse) -> None:
|
|
||||||
"""This callback is called when AsyncClient syncs with the server"""
|
|
||||||
current_time = time.time()
|
|
||||||
delta_time = current_time - self._last_next_batch_dump
|
|
||||||
self._last_next_batch = response.next_batch
|
|
||||||
if delta_time >= self.NEXT_BATCH_DUMP_PERIOD:
|
|
||||||
self._last_next_batch_dump = current_time
|
|
||||||
try:
|
|
||||||
await self._write_next_batch(self._last_next_batch)
|
|
||||||
except:
|
|
||||||
traceback.print_exc()
|
|
||||||
|
|
||||||
async def _callback_autojoin(self, room: MatrixInvitedRoom, event: InviteMemberEvent):
|
|
||||||
try:
|
|
||||||
# event.state_key must be our username
|
|
||||||
if event.state_key != self._client.user_id:
|
|
||||||
return
|
|
||||||
# membership status must be invite
|
|
||||||
if event.membership != "invite":
|
|
||||||
return
|
|
||||||
result = await self._client.join(room.room_id)
|
|
||||||
if type(result) is JoinResponse:
|
|
||||||
self._logger.info(f"Autojoined the room {room.room_id}")
|
|
||||||
else:
|
|
||||||
self._logger.error(f"Can't autojoin the room {room.room_id}")
|
|
||||||
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
|
|
||||||
|
|
||||||
async def _client_login_session_data(self, session_data: dict[str, Any]) -> None:
|
|
||||||
"""Login using session data. Raises and exception on failure."""
|
|
||||||
self._logger.debug("Using stored session data to log in")
|
|
||||||
# build user id
|
|
||||||
username = self._build_matrix_username(self._config)
|
|
||||||
self._client.restore_login(
|
|
||||||
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}'")
|
|
||||||
raise RuntimeError(result.message)
|
|
||||||
elif type(result) is not WhoamiResponse:
|
|
||||||
self._logger.error("Can't log in using stored session data, unknown error")
|
|
||||||
raise RuntimeError("Unknown response for whoami request")
|
|
||||||
self._logger.debug("Logged in using stored session data")
|
|
||||||
|
|
||||||
async def _client_login_password(self) -> None:
|
|
||||||
"""Login using password and save result to disk on success.
|
|
||||||
Raises an exception on failure.
|
|
||||||
"""
|
|
||||||
self._logger.debug("Using password to log in")
|
|
||||||
# get the password
|
|
||||||
password = await self._cb_password()
|
|
||||||
result = await self._client.login(password=password)
|
|
||||||
if type(result) is LoginResponse:
|
|
||||||
self._logger.debug("Logged in using password")
|
|
||||||
await self._write_session_data(
|
|
||||||
access_token=result.access_token,
|
|
||||||
device_id=result.device_id
|
|
||||||
)
|
|
||||||
elif type(result) is LoginError:
|
|
||||||
self._logger.error(f"Can't log in using password: '{result.message}'")
|
|
||||||
raise RuntimeError(result.message)
|
|
||||||
else:
|
|
||||||
self._logger.error(f"Can't log in using password, unknown error")
|
|
||||||
raise RuntimeError("Unknown login result")
|
|
||||||
|
|
||||||
async def _client_login(self) -> None:
|
|
||||||
"""This function logs in."""
|
|
||||||
# check if we have session data stored on the disk
|
|
||||||
session_data = await self._read_session_data()
|
|
||||||
# session data is present, try to log in
|
|
||||||
if session_data is not None:
|
|
||||||
self._logger.debug("Some session data found on the disk")
|
|
||||||
try:
|
|
||||||
await self._client_login_session_data(session_data)
|
|
||||||
return
|
|
||||||
except asyncio.CancelledError:
|
|
||||||
raise
|
|
||||||
except:
|
|
||||||
pass
|
|
||||||
# no session data - login using password
|
|
||||||
try:
|
|
||||||
self._logger.debug("No session data found on the disk OR invalid data")
|
|
||||||
await self._client_login_password()
|
|
||||||
return
|
|
||||||
except asyncio.CancelledError:
|
|
||||||
raise
|
|
||||||
except:
|
|
||||||
pass
|
|
||||||
# can't login
|
|
||||||
self._logger.error("Can't log in using available methods")
|
|
||||||
raise RuntimeError("All login methods have failed, can't continue")
|
|
||||||
|
|
||||||
async def _client_destroy(self) -> None:
|
|
||||||
"""Gracefully destroys the client."""
|
|
||||||
try:
|
|
||||||
self._logger.debug("Closing the client")
|
|
||||||
await self._client.close()
|
|
||||||
if self._last_next_batch is not None:
|
|
||||||
self._logger.debug("Saving next_batch")
|
|
||||||
await self._write_next_batch(self._last_next_batch)
|
|
||||||
except:
|
|
||||||
traceback.print_exc()
|
|
||||||
|
|
||||||
async def _client_cancellable_sync_forever(self, *args, **kwargs) -> Any:
|
|
||||||
"""Behaves exactly like AsyncClient.sync_forever, but supports task cancellation"""
|
|
||||||
sync_forever_task = asyncio.create_task(
|
|
||||||
self._client.sync_forever(*args, **kwargs)
|
|
||||||
)
|
|
||||||
try:
|
|
||||||
return await sync_forever_task
|
|
||||||
except asyncio.CancelledError:
|
|
||||||
try:
|
|
||||||
self._client.stop_sync_forever()
|
|
||||||
except Exception:
|
|
||||||
pass
|
|
||||||
finally:
|
|
||||||
sync_forever_task.cancel()
|
|
||||||
await asyncio.gather(sync_forever_task, return_exceptions=True)
|
|
||||||
raise
|
|
||||||
|
|
||||||
async def _background_coroutine(self) -> None:
|
|
||||||
"""This function implements bot lifecycle."""
|
|
||||||
# we should stop when this task stops
|
|
||||||
self._logger.debug("_background_coroutine is started")
|
|
||||||
stop_wait_task = asyncio.create_task(self._stop_event.wait()) # type: ignore
|
|
||||||
# perform login
|
|
||||||
login_task = asyncio.create_task(self._client_login())
|
|
||||||
try:
|
|
||||||
await self._wait_for_task_and_stop_event(login_task, stop_wait_task)
|
|
||||||
except asyncio.CancelledError:
|
|
||||||
self._logger.debug("Background task is cancelled during login")
|
|
||||||
await self._client_destroy()
|
|
||||||
return
|
|
||||||
except:
|
|
||||||
traceback.print_exc()
|
|
||||||
self._logger.info("Succesfully logged in")
|
|
||||||
# sync forever
|
|
||||||
self._logger.info("Syncing forever")
|
|
||||||
sync_task = asyncio.create_task(
|
|
||||||
self._client_cancellable_sync_forever(
|
|
||||||
timeout=self.MATRIX_SYNC_PERIOD,
|
|
||||||
since=(await self._read_next_batch()),
|
|
||||||
full_state=True
|
|
||||||
)
|
|
||||||
)
|
|
||||||
try:
|
|
||||||
await self._wait_for_task_and_stop_event(sync_task, stop_wait_task)
|
|
||||||
except asyncio.CancelledError:
|
|
||||||
self._logger.debug("Sync task is cancelled")
|
|
||||||
await self._client_destroy()
|
|
||||||
return
|
|
||||||
except:
|
|
||||||
traceback.print_exc()
|
|
||||||
|
|
||||||
|
|
||||||
#
|
|
||||||
# PUBLIC
|
|
||||||
#
|
|
||||||
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
|
|
||||||
# create the logger
|
|
||||||
self._logger = logging.getLogger(self._build_matrix_username(config))
|
|
||||||
self._logger.setLevel(logging.DEBUG)
|
|
||||||
# create the client
|
|
||||||
self._client: AsyncClient = self._build_client(self._config)
|
|
||||||
self._setup_client_callbacks()
|
|
||||||
|
|
||||||
# prepare some private data
|
|
||||||
self._background_task: asyncio.Task | None = None
|
|
||||||
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.
|
|
||||||
Starts the bot in background task. Raises an exception if there are
|
|
||||||
problems (for example, the bot is already started). The bot will
|
|
||||||
do everything to keep itself running, including restarts. Use
|
|
||||||
`stop()` to stop the bot.
|
|
||||||
"""
|
|
||||||
if self._background_task is not None:
|
|
||||||
raise RuntimeError("The bot is already started!")
|
|
||||||
self._stop_event = asyncio.Event()
|
|
||||||
self._background_task = asyncio.create_task(
|
|
||||||
self._background_coroutine()
|
|
||||||
)
|
|
||||||
|
|
||||||
async def stop(self) -> None:
|
|
||||||
"""Stop the bot and wait for the bot stop."""
|
|
||||||
if self._background_task is None or self._stop_event is None:
|
|
||||||
return
|
|
||||||
self._stop_event.set()
|
|
||||||
try:
|
|
||||||
await self._background_task
|
|
||||||
except:
|
|
||||||
traceback.print_exc()
|
|
||||||
self._stop_event = None
|
|
||||||
self._background_task = None
|
|
||||||
|
|
||||||
def verify_all_known_devices(self) -> bool:
|
|
||||||
"""Verifies all known devices.
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
True if there were unverified devices that are verified now.
|
|
||||||
"""
|
|
||||||
result = False
|
|
||||||
for user_id in self._client.device_store.users:
|
|
||||||
for device_id, olm_device in self._client.device_store[user_id].items():
|
|
||||||
# can't trust ourselves
|
|
||||||
if device_id == self._client.device_id and user_id == self._client.user_id:
|
|
||||||
continue
|
|
||||||
# they are already verified
|
|
||||||
if olm_device.verified:
|
|
||||||
continue
|
|
||||||
# verify them
|
|
||||||
self._client.verify_device(olm_device)
|
|
||||||
result = True
|
|
||||||
return result
|
|
||||||
|
|
||||||
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._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) -> str:
|
|
||||||
"""Sends a text message to the room and handle HTML as specified.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
text - text to send
|
|
||||||
is_html - True if text is HTML; False if text is not HTML; None if the value should be guessed
|
|
||||||
kwargs - passed as `m.room.message` content keys
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
event_id of the message on success. Raises an exception on error.
|
|
||||||
"""
|
|
||||||
text_data = self._process_html_text(text, is_html)
|
|
||||||
content = {
|
|
||||||
"msgtype": "m.text",
|
|
||||||
**text_data
|
|
||||||
}
|
|
||||||
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
|
|
||||||
)
|
|
||||||
# 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
|
|
||||||
5
src/mab/bot/__init__.py
Normal file
5
src/mab/bot/__init__.py
Normal file
@@ -0,0 +1,5 @@
|
|||||||
|
from .bot import *
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"MatrixBot"
|
||||||
|
]
|
||||||
171
src/mab/bot/_callbacks.py
Normal file
171
src/mab/bot/_callbacks.py
Normal file
@@ -0,0 +1,171 @@
|
|||||||
|
import asyncio
|
||||||
|
import logging
|
||||||
|
import traceback
|
||||||
|
from typing import Callable, Any, Coroutine, TYPE_CHECKING
|
||||||
|
from dataclasses import dataclass
|
||||||
|
|
||||||
|
from nio import AsyncClient
|
||||||
|
from nio import MatrixRoom
|
||||||
|
from nio import MatrixInvitedRoom, InviteMemberEvent, JoinResponse
|
||||||
|
from nio.events.room_events import Event as RoomEvent
|
||||||
|
|
||||||
|
from ._storage import Storage
|
||||||
|
from ..types import MatrixBotConfig, RoomEventData
|
||||||
|
from ..filters.base import BaseEventFilter
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from .bot import MatrixBot
|
||||||
|
|
||||||
|
class Callbacks:
|
||||||
|
"""
|
||||||
|
This class implements user callbacks routing and some basic internal
|
||||||
|
callbacks. You should use it to add callbacks to the bot.
|
||||||
|
"""
|
||||||
|
#
|
||||||
|
# DATACLASSES
|
||||||
|
#
|
||||||
|
@dataclass
|
||||||
|
class _FilterBasedCallback:
|
||||||
|
"""Data structure that holds information about filter-based callback"""
|
||||||
|
|
||||||
|
filter: BaseEventFilter
|
||||||
|
"""Filter to use for matching"""
|
||||||
|
|
||||||
|
callback: Callable[[RoomEventData], Coroutine[Any, Any, None]] | None
|
||||||
|
"""Callback that will be called if the filter matches"""
|
||||||
|
|
||||||
|
stop_matching: bool
|
||||||
|
"""Whether to stop matching if the filter matches"""
|
||||||
|
|
||||||
|
#
|
||||||
|
# PRIVATE
|
||||||
|
#
|
||||||
|
async def _callback_room_event_router(self,
|
||||||
|
room: MatrixRoom,
|
||||||
|
event: RoomEvent) -> None:
|
||||||
|
"""
|
||||||
|
This callback implements routing of filter-based event callbacks.
|
||||||
|
"""
|
||||||
|
if self._config is None or self._client is None:
|
||||||
|
raise RuntimeError("Callbacks were not set up")
|
||||||
|
for callback_info in self._filters:
|
||||||
|
if not isinstance(callback_info, self._FilterBasedCallback):
|
||||||
|
continue
|
||||||
|
if not callback_info.filter(room, event):
|
||||||
|
continue
|
||||||
|
event_data = RoomEventData(
|
||||||
|
room=room,
|
||||||
|
event=event,
|
||||||
|
filter=callback_info.filter,
|
||||||
|
bot=self._matrix_bot
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
# dump argument types
|
||||||
|
if callback_info.callback is None:
|
||||||
|
self._logger.critical(
|
||||||
|
f"Filter {repr(callback_info.filter)} has no callback. "
|
||||||
|
f"`room` type is `{type(event_data.room)}`. "
|
||||||
|
f"`event` type is `{type(event_data.event)}`."
|
||||||
|
)
|
||||||
|
continue
|
||||||
|
# create the coroutine
|
||||||
|
cor = callback_info.callback(event_data)
|
||||||
|
# multitasking
|
||||||
|
if self._config.enable_multitasking_callbacks:
|
||||||
|
t = asyncio.create_task(cor)
|
||||||
|
t.add_done_callback(self._multitasking_set.discard)
|
||||||
|
self._multitasking_set.add(t)
|
||||||
|
# no multitasking
|
||||||
|
else:
|
||||||
|
await cor
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
raise
|
||||||
|
except:
|
||||||
|
self._logger.error(traceback.format_exc())
|
||||||
|
if callback_info.stop_matching:
|
||||||
|
break
|
||||||
|
|
||||||
|
async def _callback_autojoin(self,
|
||||||
|
room: MatrixInvitedRoom,
|
||||||
|
event: InviteMemberEvent) -> None:
|
||||||
|
"""
|
||||||
|
This callback implements automatic room joining.
|
||||||
|
"""
|
||||||
|
if self._client is None:
|
||||||
|
raise RuntimeError("Callbacks were not set up")
|
||||||
|
try:
|
||||||
|
# event.state_key must be our username
|
||||||
|
if event.state_key != self._client.user_id:
|
||||||
|
return
|
||||||
|
# membership status must be invite
|
||||||
|
if event.membership != "invite":
|
||||||
|
return
|
||||||
|
result = await self._client.join(room.room_id)
|
||||||
|
if type(result) is JoinResponse:
|
||||||
|
self._logger.info(f"Autojoined the room {room.room_id}")
|
||||||
|
else:
|
||||||
|
self._logger.error(f"Can't autojoin the room {room.room_id}")
|
||||||
|
except:
|
||||||
|
self._logger.error(traceback.format_exc())
|
||||||
|
|
||||||
|
async def _setup_internal_callbacks(self) -> None:
|
||||||
|
"""
|
||||||
|
Setup internal callbacks.
|
||||||
|
"""
|
||||||
|
if self._config is None or self._client is None:
|
||||||
|
raise RuntimeError("Callbacks were not set up")
|
||||||
|
# filter-based callback routing
|
||||||
|
self._client.add_event_callback(
|
||||||
|
self._callback_room_event_router, RoomEvent)
|
||||||
|
# auto-join
|
||||||
|
if self._config.auto_join_any_room_on_invite:
|
||||||
|
self._client.add_event_callback(
|
||||||
|
self._callback_autojoin, InviteMemberEvent) # type: ignore
|
||||||
|
|
||||||
|
#
|
||||||
|
# PUBLIC
|
||||||
|
#
|
||||||
|
def __init__(self, storage: Storage, matrix_bot: "MatrixBot"):
|
||||||
|
self._logger = logging.getLogger("Callbacks")
|
||||||
|
self._storage = storage
|
||||||
|
self._matrix_bot = matrix_bot
|
||||||
|
|
||||||
|
self._config: MatrixBotConfig | None = None
|
||||||
|
self._client: AsyncClient | None = None
|
||||||
|
self._filters: list = []
|
||||||
|
self._multitasking_set: set[asyncio.Task] = set()
|
||||||
|
|
||||||
|
def add_room_event_callback(
|
||||||
|
self,
|
||||||
|
filter: BaseEventFilter,
|
||||||
|
callback: Callable[[RoomEventData], Coroutine[Any, Any, None]] | None,
|
||||||
|
*,
|
||||||
|
stop_matching: bool = True) -> None:
|
||||||
|
"""
|
||||||
|
Add filter-based callback.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
- filter - filter to use for the callback
|
||||||
|
- callback - the callback that must be called if the filter matches
|
||||||
|
- stop_matching - whether to stop matching if the filter matches
|
||||||
|
"""
|
||||||
|
callback_info = self._FilterBasedCallback(
|
||||||
|
filter=filter,
|
||||||
|
callback=callback,
|
||||||
|
stop_matching=stop_matching
|
||||||
|
)
|
||||||
|
self._filters.append(callback_info)
|
||||||
|
|
||||||
|
|
||||||
|
async def setup(self, config: MatrixBotConfig, client: AsyncClient) -> None:
|
||||||
|
"""
|
||||||
|
Setup the callbacks. Must be called by `ClientManager`.
|
||||||
|
|
||||||
|
This method must be called AFTER required callbacks were added by using
|
||||||
|
`add_room_event_callback(...)`.
|
||||||
|
"""
|
||||||
|
self._config = config
|
||||||
|
self._client = client
|
||||||
|
await self._setup_internal_callbacks()
|
||||||
|
|
||||||
|
|
||||||
137
src/mab/bot/_client_auth.py
Normal file
137
src/mab/bot/_client_auth.py
Normal file
@@ -0,0 +1,137 @@
|
|||||||
|
import asyncio
|
||||||
|
import aioconsole
|
||||||
|
import logging
|
||||||
|
import os
|
||||||
|
from ._storage import Storage
|
||||||
|
|
||||||
|
from ..utils import Utils
|
||||||
|
from ..types import MatrixBotConfig
|
||||||
|
|
||||||
|
from nio import AsyncClient
|
||||||
|
from nio import WhoamiError, WhoamiResponse
|
||||||
|
from nio import LoginError, LoginResponse
|
||||||
|
|
||||||
|
class ClientAuth:
|
||||||
|
"""This class is used by to authenticate AsyncClient."""
|
||||||
|
|
||||||
|
#
|
||||||
|
# PRIVATE
|
||||||
|
#
|
||||||
|
@staticmethod
|
||||||
|
async def _default_password_callback() -> str:
|
||||||
|
if "MATRIX_PASSWORD" in os.environ:
|
||||||
|
return os.environ["MATRIX_PASSWORD"]
|
||||||
|
return await aioconsole.ainput("Matrix password: ")
|
||||||
|
|
||||||
|
async def _login_using_session_data(self, client: AsyncClient) -> None:
|
||||||
|
"""
|
||||||
|
Login using session data from `Storage`.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
- client - client that needs to get logged in
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
- Does not return anything on success
|
||||||
|
- Raises an exception on error
|
||||||
|
"""
|
||||||
|
if self._full_matrix_username is None:
|
||||||
|
raise RuntimeError("ClientAuth is not set up")
|
||||||
|
self._logger.info("Trying to log in using session data from disk...")
|
||||||
|
session_data = await self._storage.read_session_data()
|
||||||
|
if session_data is None:
|
||||||
|
self._logger.error("Can't read session data from disk")
|
||||||
|
raise RuntimeError("Can't read session data from disk")
|
||||||
|
client.restore_login(
|
||||||
|
user_id=self._full_matrix_username,
|
||||||
|
access_token=session_data["access_token"],
|
||||||
|
device_id=session_data["device_id"]
|
||||||
|
)
|
||||||
|
client.load_store()
|
||||||
|
result = await client.whoami()
|
||||||
|
if type(result) is WhoamiError:
|
||||||
|
self._logger.error(
|
||||||
|
f"Can't log in using stored session data: '{result.message}'"
|
||||||
|
)
|
||||||
|
raise RuntimeError(result.message)
|
||||||
|
elif type(result) is not WhoamiResponse:
|
||||||
|
self._logger.error(
|
||||||
|
"Can't log in using stored session data, unknown error"
|
||||||
|
)
|
||||||
|
raise RuntimeError(
|
||||||
|
"Unknown response for whoami request"
|
||||||
|
)
|
||||||
|
self._logger.info("Logged in using stored session data")
|
||||||
|
|
||||||
|
async def _login_using_password(self, client: AsyncClient) -> None:
|
||||||
|
"""
|
||||||
|
Login using password.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
- client - client that needs to get logged in
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
- Does not return anything on success
|
||||||
|
- Raises an exception on error
|
||||||
|
"""
|
||||||
|
if self._full_matrix_username is None:
|
||||||
|
raise RuntimeError("ClientAuth is not set up")
|
||||||
|
self._logger.info("Trying to log in using password")
|
||||||
|
password = await self._default_password_callback()
|
||||||
|
result = await client.login(password=password)
|
||||||
|
if type(result) is LoginResponse:
|
||||||
|
self._logger.info("Logged in using password")
|
||||||
|
await self._storage.write_session_data(
|
||||||
|
access_token=result.access_token,
|
||||||
|
device_id=result.device_id
|
||||||
|
)
|
||||||
|
elif type(result) is LoginError:
|
||||||
|
self._logger.error(f"Can't log in using password: '{result.message}'")
|
||||||
|
raise RuntimeError(result.message)
|
||||||
|
else:
|
||||||
|
self._logger.error(f"Can't log in using password, unknown error")
|
||||||
|
raise RuntimeError("Unknown login result")
|
||||||
|
|
||||||
|
#
|
||||||
|
# PUBLIC
|
||||||
|
#
|
||||||
|
def __init__(self, storage: Storage):
|
||||||
|
self._logger = logging.getLogger("ClientAuth")
|
||||||
|
self._storage = storage
|
||||||
|
self._full_matrix_username: str | None = None
|
||||||
|
|
||||||
|
async def setup(self, config: MatrixBotConfig) -> None:
|
||||||
|
"""
|
||||||
|
Setup `ClientAuth` object using `config`.
|
||||||
|
"""
|
||||||
|
self._full_matrix_username = Utils.build_full_matrix_username(config)
|
||||||
|
|
||||||
|
async def login(self, client: AsyncClient) -> None:
|
||||||
|
"""
|
||||||
|
Login. Saves authorization data if it is not saved yet. Must be called
|
||||||
|
by `ClientManager`.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
- client - client that needs to get logged in
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
- Does not return anything on success
|
||||||
|
- Raises an exception on error
|
||||||
|
"""
|
||||||
|
# try session data
|
||||||
|
try:
|
||||||
|
await self._login_using_session_data(client)
|
||||||
|
return
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
raise
|
||||||
|
except:
|
||||||
|
pass
|
||||||
|
# try password
|
||||||
|
try:
|
||||||
|
await self._login_using_password(client)
|
||||||
|
return
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
raise
|
||||||
|
except:
|
||||||
|
pass
|
||||||
|
# failed
|
||||||
|
raise RuntimeError("Failed to login")
|
||||||
202
src/mab/bot/_client_manager.py
Normal file
202
src/mab/bot/_client_manager.py
Normal file
@@ -0,0 +1,202 @@
|
|||||||
|
import asyncio
|
||||||
|
import logging
|
||||||
|
import time
|
||||||
|
import traceback
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from nio import AsyncClient, AsyncClientConfig, SyncResponse
|
||||||
|
from ..types import MatrixBotConfig
|
||||||
|
|
||||||
|
from ._storage import Storage
|
||||||
|
from ._client_auth import ClientAuth
|
||||||
|
|
||||||
|
class ClientManager:
|
||||||
|
"""This class manages basic of `AsyncClient`."""
|
||||||
|
|
||||||
|
#
|
||||||
|
# PRIVATE
|
||||||
|
#
|
||||||
|
async def _close_client(self) -> None:
|
||||||
|
"""
|
||||||
|
Closes the client correctly and saves data.
|
||||||
|
"""
|
||||||
|
if self._client is None:
|
||||||
|
raise RuntimeError("The client does not exist")
|
||||||
|
self._logger.debug("Closing the client...")
|
||||||
|
await self._client.close()
|
||||||
|
self._logger.debug("Dumping 'next_batch' value...")
|
||||||
|
if self._last_next_batch is not None:
|
||||||
|
await self._storage.write_next_batch(self._last_next_batch)
|
||||||
|
|
||||||
|
async def _cancellable_sync_forever(self, *args, **kwargs) -> Any:
|
||||||
|
"""
|
||||||
|
Behaves exactly like AsyncClient.sync_forever, but supports proper task
|
||||||
|
cancellation.
|
||||||
|
"""
|
||||||
|
if self._client is None:
|
||||||
|
raise RuntimeError("Client does not exist")
|
||||||
|
task = asyncio.create_task(self._client.sync_forever(*args, **kwargs))
|
||||||
|
try:
|
||||||
|
return await task
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
try:
|
||||||
|
self._client.stop_sync_forever()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
finally:
|
||||||
|
task.cancel()
|
||||||
|
await asyncio.gather(task, return_exceptions=True)
|
||||||
|
raise
|
||||||
|
|
||||||
|
async def _background_coroutine(self) -> None:
|
||||||
|
"""
|
||||||
|
Coroutine that manages the `AsyncClient` in background. It supports
|
||||||
|
cancellation and you should use it to stop the operation of the bot.
|
||||||
|
"""
|
||||||
|
# refuse to start without config
|
||||||
|
if self._config is None:
|
||||||
|
raise RuntimeError("Config is not set up")
|
||||||
|
# refuse to start without client
|
||||||
|
if self._client is None:
|
||||||
|
raise RuntimeError("Client does not exist")
|
||||||
|
# try to login
|
||||||
|
self._logger.info("Logging in...")
|
||||||
|
try:
|
||||||
|
await self._client_auth.login(self._client)
|
||||||
|
self._logger.info(f"Logged in as {self._client.user_id}")
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
raise
|
||||||
|
except:
|
||||||
|
self._logger.error(traceback.format_exc())
|
||||||
|
raise
|
||||||
|
# sync forever
|
||||||
|
self._logger.info("Syncing forever...")
|
||||||
|
await self._cancellable_sync_forever(
|
||||||
|
timeout=self._config.client_sync_period_ms,
|
||||||
|
since=(await self._storage.read_next_batch()),
|
||||||
|
full_state=True
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _callback_sync(self, response: SyncResponse) -> None:
|
||||||
|
"""
|
||||||
|
This callback is called when AsyncClient syncs with the server
|
||||||
|
"""
|
||||||
|
if self._config is None:
|
||||||
|
raise RuntimeError("ClientManager is not set up")
|
||||||
|
# save last `next_batch`
|
||||||
|
self._last_next_batch = response.next_batch
|
||||||
|
# dump if needed
|
||||||
|
current_time = time.time()
|
||||||
|
delta_time = current_time - self._last_next_batch_write_time
|
||||||
|
if delta_time >= self._config.client_next_batch_dump_period_s:
|
||||||
|
try:
|
||||||
|
await self._storage.write_next_batch(self._last_next_batch)
|
||||||
|
self._last_next_batch_write_time = current_time
|
||||||
|
except:
|
||||||
|
self._logger.error(traceback.format_exc())
|
||||||
|
|
||||||
|
async def _setup_callbacks(self, config: MatrixBotConfig) -> None:
|
||||||
|
"""
|
||||||
|
Setup callbacks.
|
||||||
|
"""
|
||||||
|
if self._client is None:
|
||||||
|
raise RuntimeError("The client does not exist")
|
||||||
|
self._client.add_response_callback(
|
||||||
|
self._callback_sync, SyncResponse) # type: ignore
|
||||||
|
|
||||||
|
#
|
||||||
|
# PUBLIC
|
||||||
|
#
|
||||||
|
def __init__(self, client_auth: ClientAuth, storage: Storage):
|
||||||
|
self._logger = logging.getLogger("ClientManager")
|
||||||
|
self._client_auth = client_auth
|
||||||
|
self._storage = storage
|
||||||
|
|
||||||
|
self._config: MatrixBotConfig | None = None
|
||||||
|
self._client: AsyncClient | None = None
|
||||||
|
self._was_started_ever: bool = False
|
||||||
|
self._was_built_ever: bool = False
|
||||||
|
self._background_task: asyncio.Task | None = None
|
||||||
|
self._last_next_batch_write_time: float = 0.0
|
||||||
|
self._last_next_batch: str | None = None
|
||||||
|
|
||||||
|
async def setup(self, config: MatrixBotConfig) -> None:
|
||||||
|
"""
|
||||||
|
Build `AsyncClient` from `config`.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
- Does not return anything, raises an exception on error.
|
||||||
|
"""
|
||||||
|
if self._was_built_ever:
|
||||||
|
raise RuntimeError("The client was already built")
|
||||||
|
client_config = AsyncClientConfig(
|
||||||
|
store_name="nio_store_file",
|
||||||
|
encryption_enabled=True,
|
||||||
|
store_sync_tokens=False
|
||||||
|
)
|
||||||
|
self._client = AsyncClient(
|
||||||
|
homeserver=config.matrix_homeserver_url,
|
||||||
|
user=config.matrix_username_localpart,
|
||||||
|
store_path=str(config.storage_directory),
|
||||||
|
config=client_config
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
await self._setup_callbacks(config)
|
||||||
|
except:
|
||||||
|
self._client = None
|
||||||
|
raise
|
||||||
|
self._config = config
|
||||||
|
self._was_built_ever = True
|
||||||
|
|
||||||
|
def get_client(self) -> AsyncClient:
|
||||||
|
"""
|
||||||
|
Get AsyncClient.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
- `AsyncClient` you may use to access Matrix. Raises a RuntimeError
|
||||||
|
if the client does not exist or is not operational.
|
||||||
|
"""
|
||||||
|
if self._client is None:
|
||||||
|
raise RuntimeError("The client does not exist")
|
||||||
|
return self._client
|
||||||
|
|
||||||
|
async def start(self) -> None:
|
||||||
|
"""
|
||||||
|
Start client operation in background. This method can be called
|
||||||
|
succesfully only once for `ClientManager` object. Any subsequent calls
|
||||||
|
will raise an exception.
|
||||||
|
"""
|
||||||
|
if self._client is None:
|
||||||
|
raise RuntimeError("Client does not exist")
|
||||||
|
if self._was_started_ever:
|
||||||
|
raise RuntimeError("Client can't be started several times")
|
||||||
|
self._background_task = asyncio.create_task(
|
||||||
|
self._background_coroutine()
|
||||||
|
)
|
||||||
|
self._was_started_ever = True
|
||||||
|
self._logger.info("Bot background is started")
|
||||||
|
|
||||||
|
async def stop(self) -> None:
|
||||||
|
"""
|
||||||
|
Stop client operation in background. This method can be called only if
|
||||||
|
the client was `start(...)`ed previously. This method will wait for the
|
||||||
|
bot to completely stop.
|
||||||
|
"""
|
||||||
|
if self._client is None:
|
||||||
|
raise RuntimeError("Client does not exist")
|
||||||
|
if self._background_task is None:
|
||||||
|
if self._was_started_ever:
|
||||||
|
raise RuntimeError("The bot was stopped previously")
|
||||||
|
else:
|
||||||
|
raise RuntimeError("The bot was never started")
|
||||||
|
self._background_task.cancel()
|
||||||
|
try:
|
||||||
|
asyncio.shield(self._background_task)
|
||||||
|
except:
|
||||||
|
self._logger.error(traceback.format_exc())
|
||||||
|
try:
|
||||||
|
asyncio.shield(self._close_client())
|
||||||
|
except:
|
||||||
|
self._logger.error(traceback.format_exc())
|
||||||
|
self._background_task = None
|
||||||
|
self._logger.info("Bot background task is stopped")
|
||||||
300
src/mab/bot/_client_sender.py
Normal file
300
src/mab/bot/_client_sender.py
Normal file
@@ -0,0 +1,300 @@
|
|||||||
|
import asyncio
|
||||||
|
import logging
|
||||||
|
from html.parser import HTMLParser
|
||||||
|
from pathlib import Path
|
||||||
|
import traceback
|
||||||
|
import os
|
||||||
|
|
||||||
|
import magic
|
||||||
|
from PIL import Image
|
||||||
|
from nio import AsyncClient, MatrixRoom, RoomSendResponse, RoomSendError
|
||||||
|
from nio import OlmUnverifiedDeviceError
|
||||||
|
|
||||||
|
from ._client_uploader import ClientUploader
|
||||||
|
from ..types import MatrixBotConfig, VideoFileProperties
|
||||||
|
from ..utils import Utils
|
||||||
|
|
||||||
|
class ClientSender:
|
||||||
|
"""This class sends messages"""
|
||||||
|
#
|
||||||
|
# PRIVATE
|
||||||
|
#
|
||||||
|
def _process_html_text(self, 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.
|
||||||
|
"""
|
||||||
|
# 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:
|
||||||
|
self._logger.error(traceback.format_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
|
||||||
|
}
|
||||||
|
|
||||||
|
#
|
||||||
|
# PUBLIC
|
||||||
|
#
|
||||||
|
def __init__(self):
|
||||||
|
self._logger = logging.getLogger("ClientSender")
|
||||||
|
self._config: MatrixBotConfig | None = None
|
||||||
|
self._client: AsyncClient | None = None
|
||||||
|
|
||||||
|
async def setup(self,
|
||||||
|
config: MatrixBotConfig,
|
||||||
|
client: AsyncClient,
|
||||||
|
uploader: ClientUploader) -> None:
|
||||||
|
"""
|
||||||
|
Setup the sender.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
- config - config to use
|
||||||
|
- client - client to use
|
||||||
|
- uploader - file uploaded to use
|
||||||
|
"""
|
||||||
|
self._config = config
|
||||||
|
self._client = client
|
||||||
|
self._uploader = uploader
|
||||||
|
|
||||||
|
async def send_content(self,
|
||||||
|
room: MatrixRoom | str,
|
||||||
|
content: dict) -> RoomSendResponse:
|
||||||
|
"""
|
||||||
|
Send raw `content` JSON to `room`.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
- room - the room to send the content to
|
||||||
|
- content - the content to send to the room
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
- `RoomSendResponse` returned by `AsyncClient.room_send`
|
||||||
|
- Raises an exception on error
|
||||||
|
"""
|
||||||
|
if self._client is None or self._config is None:
|
||||||
|
raise RuntimeError("ClientSender is not set up")
|
||||||
|
try:
|
||||||
|
result = await self._client.room_send(
|
||||||
|
room_id=room.room_id if isinstance(room, MatrixRoom) else room,
|
||||||
|
message_type="m.room.message",
|
||||||
|
content=content
|
||||||
|
)
|
||||||
|
except OlmUnverifiedDeviceError:
|
||||||
|
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)
|
||||||
|
else:
|
||||||
|
raise
|
||||||
|
if type(result) is RoomSendResponse:
|
||||||
|
return result
|
||||||
|
elif type(result) is RoomSendError:
|
||||||
|
raise RuntimeError(result)
|
||||||
|
else:
|
||||||
|
raise RuntimeError("Unknown error has occured", result)
|
||||||
|
|
||||||
|
async def send_text(self,
|
||||||
|
room: MatrixRoom | str,
|
||||||
|
text: str,
|
||||||
|
*,
|
||||||
|
is_html: bool | None = None) -> str:
|
||||||
|
"""
|
||||||
|
Send text message to `room`.
|
||||||
|
|
||||||
|
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
|
||||||
|
"""
|
||||||
|
text_data = self._process_html_text(text, is_html)
|
||||||
|
content = {
|
||||||
|
"msgtype": "m.text",
|
||||||
|
**text_data
|
||||||
|
}
|
||||||
|
return (await self.send_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:
|
||||||
|
"""
|
||||||
|
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
|
||||||
|
"""
|
||||||
|
# caption must not actually be empty
|
||||||
|
if text is None or not text.strip():
|
||||||
|
text = 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
|
||||||
|
async with asyncio.timeout(timeout):
|
||||||
|
upload_result = await self._uploader.upload_file(
|
||||||
|
path, mime_type=mime_type, filename=filename)
|
||||||
|
# prepare the content and send
|
||||||
|
content = {
|
||||||
|
"msgtype": "m.image",
|
||||||
|
"filename": os.path.basename(path),
|
||||||
|
**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,
|
||||||
|
"w": width,
|
||||||
|
"h": height
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return (await self.send_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:
|
||||||
|
"""
|
||||||
|
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
|
||||||
|
"""
|
||||||
|
if self._client is None or self._config is None:
|
||||||
|
raise RuntimeError("ClientSender is not set up")
|
||||||
|
# caption must not actually be empty
|
||||||
|
if text is None or not text.strip():
|
||||||
|
text = 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"
|
||||||
|
)
|
||||||
|
async with asyncio.timeout(30):
|
||||||
|
props = await Utils.get_video_props_using_ffprobe(path)
|
||||||
|
# autogenerate thumbnail if allowed to do so
|
||||||
|
if not props.thumbnail and self._config.auto_ffmpeg:
|
||||||
|
try:
|
||||||
|
async with asyncio.timeout(30):
|
||||||
|
props.thumbnail = await Utils.get_video_thumbnail(
|
||||||
|
path,
|
||||||
|
timestamp=props.duration_ms / 1000 / 2,
|
||||||
|
max_size=400
|
||||||
|
)
|
||||||
|
except asyncio.TimeoutError:
|
||||||
|
pass
|
||||||
|
async with asyncio.timeout(timeout):
|
||||||
|
# upload the thumbnail if specified
|
||||||
|
thumbnail_data = await self._uploader.upload_thumbnail(props.thumbnail)
|
||||||
|
# upload
|
||||||
|
upload_result = await self._uploader.upload_file(
|
||||||
|
path,
|
||||||
|
mime_type=mime_type
|
||||||
|
)
|
||||||
|
# prepare the content
|
||||||
|
content = {
|
||||||
|
"msgtype": "m.video",
|
||||||
|
"filename": os.path.basename(path),
|
||||||
|
**self._process_html_text(text, is_html),
|
||||||
|
"file": {
|
||||||
|
"url": upload_result.response.content_uri,
|
||||||
|
"mimetype": mime_type,
|
||||||
|
**upload_result.keys
|
||||||
|
},
|
||||||
|
"info": {
|
||||||
|
"duration": props.duration_ms,
|
||||||
|
"mimetype": mime_type,
|
||||||
|
"size": upload_result.filesize,
|
||||||
|
"w": props.width_px,
|
||||||
|
"h": props.height_px,
|
||||||
|
**thumbnail_data
|
||||||
|
}
|
||||||
|
}
|
||||||
|
# send
|
||||||
|
return (await self.send_content(room, content)).event_id
|
||||||
187
src/mab/bot/_client_uploader.py
Normal file
187
src/mab/bot/_client_uploader.py
Normal file
@@ -0,0 +1,187 @@
|
|||||||
|
from typing import Any
|
||||||
|
from pathlib import Path
|
||||||
|
import io
|
||||||
|
import os
|
||||||
|
|
||||||
|
import aiofiles
|
||||||
|
import aiofiles.os
|
||||||
|
import magic
|
||||||
|
from PIL import Image
|
||||||
|
|
||||||
|
from nio import AsyncClient
|
||||||
|
from nio import UploadError
|
||||||
|
|
||||||
|
from ._storage import Storage
|
||||||
|
from ..types import MatrixBotConfig, UploadResult
|
||||||
|
|
||||||
|
class ClientUploader:
|
||||||
|
"""This class uploads files"""
|
||||||
|
def __init__(self, storage: Storage):
|
||||||
|
self._storage = storage
|
||||||
|
self._config: MatrixBotConfig | None = None
|
||||||
|
self._client: AsyncClient | None = None
|
||||||
|
|
||||||
|
async def setup(self, config: MatrixBotConfig, client: AsyncClient) -> None:
|
||||||
|
"""
|
||||||
|
Setup the uploader.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
- config - config to take some settings from.
|
||||||
|
"""
|
||||||
|
self._config = config
|
||||||
|
self._client = client
|
||||||
|
|
||||||
|
async def upload_using_provider(self,
|
||||||
|
provider: Any,
|
||||||
|
*,
|
||||||
|
mime_type: str,
|
||||||
|
filename: str,
|
||||||
|
filesize: int) -> UploadResult:
|
||||||
|
"""
|
||||||
|
Upload data using provider (as defined in `AsyncClient.upload` docs).
|
||||||
|
You better use `upload_file(...)`, `upload_bytes(...)`. Data is always
|
||||||
|
encrypted.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
- provider - data provider (e.g. `aiofiles.open(...)`)
|
||||||
|
- mime_type - mime-type of the data
|
||||||
|
- filename - filename to declare for the file
|
||||||
|
- filesize - size of the file
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
- `UploadResult` that contains all data you need on success
|
||||||
|
- Raises an exception on failure
|
||||||
|
"""
|
||||||
|
if self._config is None or self._client is None:
|
||||||
|
raise RuntimeError("ClientUploader is not set up")
|
||||||
|
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(f"Upload failed: {response.message}", response)
|
||||||
|
return UploadResult(
|
||||||
|
response=response,
|
||||||
|
keys=keys, # type: ignore
|
||||||
|
filename=filename,
|
||||||
|
filesize=filesize
|
||||||
|
)
|
||||||
|
|
||||||
|
async def upload_file(self,
|
||||||
|
path: Path | str,
|
||||||
|
*,
|
||||||
|
mime_type: str | None = None,
|
||||||
|
filename: str | None = None) -> UploadResult:
|
||||||
|
"""
|
||||||
|
Upload data from file. Data is always encrypted.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
- path - path to the file that needs to be uploaded
|
||||||
|
- mime_type - mime-type to use for the file (`None` for auto)
|
||||||
|
- filename - filename to use for the file (`None` for auto)
|
||||||
|
"""
|
||||||
|
if isinstance(path, str):
|
||||||
|
path = Path(path)
|
||||||
|
if not path.is_file():
|
||||||
|
raise RuntimeError(f"{path} is not a file")
|
||||||
|
if not filename:
|
||||||
|
filename = os.path.basename(path)
|
||||||
|
if not mime_type:
|
||||||
|
mime_type = magic.from_file(path, mime=True)
|
||||||
|
filesize = (await aiofiles.os.stat(path)).st_size
|
||||||
|
async with aiofiles.open(path, "r+b") as f:
|
||||||
|
return await self.upload_using_provider(
|
||||||
|
f,
|
||||||
|
mime_type=mime_type,
|
||||||
|
filename=filename,
|
||||||
|
filesize=filesize
|
||||||
|
)
|
||||||
|
|
||||||
|
async def upload_bytes(self,
|
||||||
|
data: bytes,
|
||||||
|
*,
|
||||||
|
filename: str,
|
||||||
|
mime_type: str | None = None) -> UploadResult:
|
||||||
|
"""
|
||||||
|
Upload data. Data is always encrypted.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
- data - data that needs to be uploaded
|
||||||
|
- filename - filename to use for the file
|
||||||
|
- mime_type - mime-type to use for the file (`None` for auto)
|
||||||
|
"""
|
||||||
|
if not mime_type:
|
||||||
|
mime_type = magic.from_buffer(data, mime=True)
|
||||||
|
buffer = io.BytesIO(data)
|
||||||
|
return await self.upload_using_provider(
|
||||||
|
buffer,
|
||||||
|
mime_type=mime_type,
|
||||||
|
filename=filename,
|
||||||
|
filesize=len(data)
|
||||||
|
)
|
||||||
|
|
||||||
|
async def upload_thumbnail(self,
|
||||||
|
thumbnail: Path | str | bytes | None) -> dict:
|
||||||
|
"""
|
||||||
|
Upload thumbnail and return fields required to use it for the 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)
|
||||||
|
upload_result = await self.upload_using_provider(
|
||||||
|
raw_bytes_io,
|
||||||
|
mime_type=thumb_mime,
|
||||||
|
filename=f"preview.{thumb_mime.split("/")[-1]}",
|
||||||
|
filesize=thumb_size
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
upload_result = await self.upload_file(
|
||||||
|
thumbnail, # type: ignore
|
||||||
|
mime_type=thumb_mime,
|
||||||
|
filename=os.path.basename(str(thumbnail))
|
||||||
|
)
|
||||||
|
# return the result
|
||||||
|
return {
|
||||||
|
"thumbnail_file": {
|
||||||
|
"url": upload_result.response.content_uri,
|
||||||
|
"mimetype": thumb_mime,
|
||||||
|
**upload_result.keys
|
||||||
|
},
|
||||||
|
"thumbnail_info": {
|
||||||
|
"w": width,
|
||||||
|
"h": height,
|
||||||
|
"size": upload_result.filesize,
|
||||||
|
"mimetype": thumb_mime
|
||||||
|
}
|
||||||
|
}
|
||||||
103
src/mab/bot/_storage.py
Normal file
103
src/mab/bot/_storage.py
Normal file
@@ -0,0 +1,103 @@
|
|||||||
|
from pathlib import Path
|
||||||
|
import aiofiles
|
||||||
|
import logging
|
||||||
|
import json
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from ..types import MatrixBotConfig
|
||||||
|
|
||||||
|
class Storage:
|
||||||
|
"""This class manages internal storage of the MatrixBot"""
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
self._logger = logging.getLogger("Storage")
|
||||||
|
self._config: MatrixBotConfig | None = None
|
||||||
|
|
||||||
|
async def setup(self, config: MatrixBotConfig) -> None:
|
||||||
|
"""
|
||||||
|
Setup the internal storage of the MatrixBot.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
- config - the config to use
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
- Does not return anything
|
||||||
|
- Raises an exception on error
|
||||||
|
"""
|
||||||
|
self._config = config
|
||||||
|
|
||||||
|
async def write_next_batch(self, next_batch: str) -> None:
|
||||||
|
"""
|
||||||
|
Write `next_batch` value to the disk.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
- next_batch - value to save
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
- Does not return anything
|
||||||
|
- Raises an exception on error
|
||||||
|
"""
|
||||||
|
if self._config is None:
|
||||||
|
raise RuntimeError("Storage is not set up")
|
||||||
|
path = Path(self._config.storage_directory) / "next_batch"
|
||||||
|
async with aiofiles.open(path, "w") as f:
|
||||||
|
await f.write(next_batch)
|
||||||
|
self._logger.debug("next_batch value is written to the disk")
|
||||||
|
|
||||||
|
async def read_next_batch(self) -> str | None:
|
||||||
|
"""
|
||||||
|
Read `next_batch` value from the disk.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
- Returns `str` if there is `next_batch` value saved on disk
|
||||||
|
- Returns `None` if `next_batch` was never saved yet
|
||||||
|
- Raises an exception on error
|
||||||
|
"""
|
||||||
|
if self._config is None:
|
||||||
|
raise RuntimeError("Storage is not set up")
|
||||||
|
path = Path(self._config.storage_directory) / "next_batch"
|
||||||
|
if not path.is_file():
|
||||||
|
return None
|
||||||
|
async with aiofiles.open(path, "r") as f:
|
||||||
|
return (await f.read()).strip()
|
||||||
|
|
||||||
|
async def write_session_data(self, *, access_token: str, device_id: str) -> None:
|
||||||
|
"""
|
||||||
|
Write session data to disk.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
- access_token - access_token to write to the disk
|
||||||
|
- device_id - device_id to write to disk
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
- Does not return anything
|
||||||
|
- Raises an exception on error
|
||||||
|
"""
|
||||||
|
if self._config is None:
|
||||||
|
raise RuntimeError("Storage is not set up")
|
||||||
|
path = Path(self._config.storage_directory) / "session_data.json"
|
||||||
|
data = {
|
||||||
|
"access_token": access_token,
|
||||||
|
"device_id": device_id
|
||||||
|
}
|
||||||
|
async with aiofiles.open(path, "w") as f:
|
||||||
|
await f.write(json.dumps(data, indent=4))
|
||||||
|
self._logger.debug("Session data is writter to the disk")
|
||||||
|
|
||||||
|
async def read_session_data(self) -> dict[str, Any] | None:
|
||||||
|
"""
|
||||||
|
Read session data from disk.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
- Dictionary with `access_token` and `device_id` keys if these
|
||||||
|
values were saved using `write_session_data`
|
||||||
|
- `None` if session data was never saved yet
|
||||||
|
"""
|
||||||
|
if self._config is None:
|
||||||
|
raise RuntimeError("Storage is not set up")
|
||||||
|
path = Path(self._config.storage_directory) / "session_data.json"
|
||||||
|
if not path.is_file():
|
||||||
|
return None
|
||||||
|
async with aiofiles.open(path, "r") as f:
|
||||||
|
j = json.loads(await f.read())
|
||||||
|
return j
|
||||||
84
src/mab/bot/_validation.py
Normal file
84
src/mab/bot/_validation.py
Normal file
@@ -0,0 +1,84 @@
|
|||||||
|
import re
|
||||||
|
from pathlib import Path
|
||||||
|
from urllib.parse import urlparse
|
||||||
|
|
||||||
|
from ..types import MatrixBotConfig
|
||||||
|
|
||||||
|
class Validator:
|
||||||
|
"""This class is used inside MatrixBot to validate data"""
|
||||||
|
|
||||||
|
def validate_matrix_homeserver_url(self, url: str) -> None:
|
||||||
|
"""
|
||||||
|
Check if `url` is a valid matrix homeserver URL.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
- url - URL to check (only format is checked)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
- Does not return anything
|
||||||
|
- Raises an exception on error
|
||||||
|
"""
|
||||||
|
parsed = urlparse(url)
|
||||||
|
if parsed.scheme not in ("http", "https"):
|
||||||
|
raise RuntimeError(
|
||||||
|
f"Scheme {parsed.scheme} is not a valid scheme for matrix "
|
||||||
|
"homeserver URL"
|
||||||
|
)
|
||||||
|
if not parsed.netloc:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"{url} is not a valid matrix homeserver URL"
|
||||||
|
)
|
||||||
|
if parsed.path != "":
|
||||||
|
raise RuntimeError(
|
||||||
|
f"{url} must have empty path (remove `{parsed.path}` after the "
|
||||||
|
"hostname)"
|
||||||
|
)
|
||||||
|
|
||||||
|
def validate_matrix_username_localpart(self, localpart: str) -> None:
|
||||||
|
"""
|
||||||
|
Check if `localpart` is a valid localpart of matrix username.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
- localpart - localpart of matrix username
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
- Does not return anything
|
||||||
|
- Raises an exception on error
|
||||||
|
"""
|
||||||
|
pattern = r"^[a-z0-9._=\-]+$"
|
||||||
|
if not bool(re.fullmatch(pattern, localpart)) or len(localpart) > 255:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"{localpart} is not a valid matrix username localpart"
|
||||||
|
)
|
||||||
|
|
||||||
|
def ensure_directory_exists(self, path: Path | str) -> None:
|
||||||
|
"""
|
||||||
|
Ensure that `path` is an existing directory. Creates it if it does not
|
||||||
|
exist. The entire directory tree is created if needed.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
- path - directory path
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
- Does not return anything
|
||||||
|
- Raises an exception on error
|
||||||
|
"""
|
||||||
|
if isinstance(path, str):
|
||||||
|
path = Path(path)
|
||||||
|
if not path.is_dir():
|
||||||
|
path.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
def validate_bot_config(self, config: MatrixBotConfig) -> None:
|
||||||
|
"""
|
||||||
|
Check if `config` is a valid MatrixBot config.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
- config - configuration to check
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
- Does not return anything
|
||||||
|
- Raises an exception on error
|
||||||
|
"""
|
||||||
|
self.validate_matrix_homeserver_url(config.matrix_homeserver_url)
|
||||||
|
self.validate_matrix_username_localpart(config.matrix_username_localpart)
|
||||||
|
self.ensure_directory_exists(config.storage_directory)
|
||||||
198
src/mab/bot/bot.py
Normal file
198
src/mab/bot/bot.py
Normal file
@@ -0,0 +1,198 @@
|
|||||||
|
import logging
|
||||||
|
|
||||||
|
from typing import Callable, Coroutine, Any
|
||||||
|
|
||||||
|
from nio import AsyncClient
|
||||||
|
|
||||||
|
from ..filters.base import BaseEventFilter
|
||||||
|
from ..types import *
|
||||||
|
|
||||||
|
from ._validation import Validator
|
||||||
|
from ._storage import Storage
|
||||||
|
from ._client_auth import ClientAuth
|
||||||
|
from ._client_manager import ClientManager
|
||||||
|
from ._client_uploader import ClientUploader
|
||||||
|
from ._client_sender import ClientSender
|
||||||
|
from ._callbacks import Callbacks
|
||||||
|
|
||||||
|
|
||||||
|
class MatrixBot:
|
||||||
|
"""
|
||||||
|
Asynchronous Matrix Bot Implementation.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, config: MatrixBotConfig) -> None:
|
||||||
|
# create the logger
|
||||||
|
self._logger = logging.getLogger("MatrixBot")
|
||||||
|
self._was_started = False
|
||||||
|
# create classes that compose MatrixBot
|
||||||
|
self._validator = Validator()
|
||||||
|
self._storage = Storage()
|
||||||
|
self._client_auth = ClientAuth(self._storage)
|
||||||
|
self._client_manager = ClientManager(self._client_auth, self._storage)
|
||||||
|
self._client_uploader = ClientUploader(self._storage)
|
||||||
|
self._client_sender = ClientSender()
|
||||||
|
self._callbacks = Callbacks(self._storage, self)
|
||||||
|
# validate the config and save it
|
||||||
|
try:
|
||||||
|
self._validator.validate_bot_config(config)
|
||||||
|
self._config: MatrixBotConfig = config
|
||||||
|
except Exception as e:
|
||||||
|
self._logger.error(f"Couldn't validate bot config: {e}")
|
||||||
|
raise
|
||||||
|
|
||||||
|
def add_callback(self,
|
||||||
|
filter: BaseEventFilter,
|
||||||
|
callback: Callable[[RoomEventData], Coroutine[Any, Any, None]] | None,
|
||||||
|
*,
|
||||||
|
stop_matching: bool = True) -> None:
|
||||||
|
"""
|
||||||
|
Add filter-based callback.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
- filter - filter to use for the callback
|
||||||
|
- callback - the callback that must be called if the filter matches.
|
||||||
|
Use `None` to log the type of the Event that will
|
||||||
|
happen for specified filter.
|
||||||
|
- stop_matching - whether to stop matching if the filter matches
|
||||||
|
"""
|
||||||
|
self._callbacks.add_room_event_callback(
|
||||||
|
filter=filter,
|
||||||
|
callback=callback,
|
||||||
|
stop_matching=stop_matching
|
||||||
|
)
|
||||||
|
|
||||||
|
async def start(self) -> None:
|
||||||
|
"""
|
||||||
|
Start client operation in background. This method can be called
|
||||||
|
succesfully only once for `MatrixBot` object. Any subsequent calls
|
||||||
|
will raise an exception.
|
||||||
|
"""
|
||||||
|
if self._was_started:
|
||||||
|
raise RuntimeError("The bot was already started")
|
||||||
|
self._was_started = True
|
||||||
|
await self._client_manager.setup(self._config)
|
||||||
|
await self._storage.setup(self._config)
|
||||||
|
await self._client_auth.setup(self._config)
|
||||||
|
await self._client_uploader.setup(
|
||||||
|
self._config,
|
||||||
|
self._client_manager.get_client()
|
||||||
|
)
|
||||||
|
await self._client_sender.setup(
|
||||||
|
self._config,
|
||||||
|
self._client_manager.get_client(),
|
||||||
|
self._client_uploader
|
||||||
|
)
|
||||||
|
await self._callbacks.setup(
|
||||||
|
self._config,
|
||||||
|
self._client_manager.get_client()
|
||||||
|
)
|
||||||
|
await self._client_manager.start()
|
||||||
|
|
||||||
|
async def stop(self) -> None:
|
||||||
|
"""
|
||||||
|
Stop client operation in background. This method can be called only if
|
||||||
|
the client was `start(...)`ed previously. This method will wait for the
|
||||||
|
bot to completely stop.
|
||||||
|
"""
|
||||||
|
await self._client_manager.stop()
|
||||||
|
|
||||||
|
def get_client(self) -> AsyncClient:
|
||||||
|
"""
|
||||||
|
Get AsyncClient.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
- `AsyncClient` you may use to access Matrix. Raises a RuntimeError
|
||||||
|
if the client does not exist or is not operational.
|
||||||
|
"""
|
||||||
|
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`.
|
||||||
|
|
||||||
|
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_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
|
||||||
|
)
|
||||||
@@ -3,6 +3,15 @@
|
|||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
|
|
||||||
|
from nio import MatrixRoom, Event
|
||||||
|
from nio import UploadResponse
|
||||||
|
|
||||||
|
from .filters.base import BaseEventFilter
|
||||||
|
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from .bot import MatrixBot
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class MatrixBotConfig:
|
class MatrixBotConfig:
|
||||||
"""Configuration for MatrixBot"""
|
"""Configuration for MatrixBot"""
|
||||||
@@ -13,7 +22,7 @@ class MatrixBotConfig:
|
|||||||
matrix_username_localpart: str
|
matrix_username_localpart: str
|
||||||
"""Localpart of matrix username (without homeserver), for example: `valid-username`"""
|
"""Localpart of matrix username (without homeserver), for example: `valid-username`"""
|
||||||
|
|
||||||
storage_directory: Path
|
storage_directory: Path | str
|
||||||
"""Path to the storage directory (will be created if needed)"""
|
"""Path to the storage directory (will be created if needed)"""
|
||||||
|
|
||||||
auto_verify_all_known_devices: bool = True
|
auto_verify_all_known_devices: bool = True
|
||||||
@@ -28,6 +37,20 @@ class MatrixBotConfig:
|
|||||||
auto_ffmpeg: bool = True
|
auto_ffmpeg: bool = True
|
||||||
"""Enabled automatic `ffmpeg` execution to generate some media (e.g. thumbnails)"""
|
"""Enabled automatic `ffmpeg` execution to generate some media (e.g. thumbnails)"""
|
||||||
|
|
||||||
|
allow_ainput_password: bool = True
|
||||||
|
"""Allow `aioconsole.ainput` usage to get user password"""
|
||||||
|
|
||||||
|
client_sync_period_ms: int = 5000
|
||||||
|
"""Sync period for `nio.AsyncClient` in milliseconds"""
|
||||||
|
|
||||||
|
client_next_batch_dump_period_s: float = 120.0
|
||||||
|
"""`next_batch` value dumping period in seconds"""
|
||||||
|
|
||||||
|
enable_multitasking_callbacks: bool = True
|
||||||
|
"""Filter-based callbacks will be started each in separate task so that
|
||||||
|
long callback execution will not prevent other callbacks from executting
|
||||||
|
"""
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class VideoFileProperties:
|
class VideoFileProperties:
|
||||||
"""Properties of the video for sending"""
|
"""Properties of the video for sending"""
|
||||||
@@ -43,3 +66,35 @@ class VideoFileProperties:
|
|||||||
|
|
||||||
thumbnail: Path | str | bytes | None = None
|
thumbnail: Path | str | bytes | None = None
|
||||||
"""Path to the thumbnail or the raw JPEG thumbnail data"""
|
"""Path to the thumbnail or the raw JPEG thumbnail data"""
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class RoomEventData:
|
||||||
|
"""Dataclass that hold information about event that happened in the room"""
|
||||||
|
|
||||||
|
room: MatrixRoom
|
||||||
|
"""The room the event has happened in"""
|
||||||
|
|
||||||
|
event: Event
|
||||||
|
"""The event that has happened in the room"""
|
||||||
|
|
||||||
|
filter: BaseEventFilter
|
||||||
|
"""The filter that invoked this event"""
|
||||||
|
|
||||||
|
bot: "MatrixBot"
|
||||||
|
"""The bot that is the source of the event"""
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class UploadResult:
|
||||||
|
"""Result of data upload"""
|
||||||
|
|
||||||
|
response: UploadResponse
|
||||||
|
"""Response, as returned by `AsyncClient.upload(...)`"""
|
||||||
|
|
||||||
|
keys: dict
|
||||||
|
"""Encryptions keys, as returned by `AsyncClient.upload(...)`"""
|
||||||
|
|
||||||
|
filename: str
|
||||||
|
"""Name of the uploaded file"""
|
||||||
|
|
||||||
|
filesize: int
|
||||||
|
"""Size of uploaded file"""
|
||||||
@@ -3,6 +3,9 @@
|
|||||||
import asyncio
|
import asyncio
|
||||||
import json
|
import json
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
from urllib.parse import urlparse
|
||||||
|
|
||||||
|
from nio import AsyncClient
|
||||||
|
|
||||||
from .types import *
|
from .types import *
|
||||||
|
|
||||||
@@ -96,3 +99,35 @@ class Utils:
|
|||||||
raise RuntimeError(f"`ffmpeg` exits with code {proc.returncode}: {stderr.decode(errors="replace")}", args)
|
raise RuntimeError(f"`ffmpeg` exits with code {proc.returncode}: {stderr.decode(errors="replace")}", args)
|
||||||
# success
|
# success
|
||||||
return stdout
|
return stdout
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def build_full_matrix_username(config: MatrixBotConfig) -> str:
|
||||||
|
"""Build complete matrix username using data from MatrixBotConfig."""
|
||||||
|
homeserver_name = urlparse(config.matrix_homeserver_url).hostname
|
||||||
|
localpart = config.matrix_username_localpart
|
||||||
|
return f"@{localpart}:{homeserver_name}"
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def verify_all_known_devices(client: AsyncClient) -> bool:
|
||||||
|
"""
|
||||||
|
Verify all known devices.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
- client - client to verify all known devices for
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
- True if there were unverified devices that are verified now
|
||||||
|
"""
|
||||||
|
result = False
|
||||||
|
for user_id in client.device_store.users:
|
||||||
|
for device_id, olm_device in client.device_store[user_id].items():
|
||||||
|
# can't trust ourselves
|
||||||
|
if device_id == client.device_id and user_id == client.user_id:
|
||||||
|
continue
|
||||||
|
# they are already verified
|
||||||
|
if olm_device.verified:
|
||||||
|
continue
|
||||||
|
# verify them
|
||||||
|
client.verify_device(olm_device)
|
||||||
|
result = True
|
||||||
|
return result
|
||||||
Reference in New Issue
Block a user