From a2a1b6c2ed26e02198a59b3f18577e73ebca9128 Mon Sep 17 00:00:00 2001 From: Nikita Tyukalov Date: Fri, 21 Aug 2026 06:25:27 +0300 Subject: [PATCH] Updated bot implementation - Added ability to send text messages - Added ability to verify all devices - Added configuration field to enable autoverification of all known devices - Added ability to get client --- bot_types.py | 5 ++- main.py | 11 ++++++ new_bot.py | 108 ++++++++++++++++++++++++++++++++++++++++++++++++++- 3 files changed, 122 insertions(+), 2 deletions(-) diff --git a/bot_types.py b/bot_types.py index 5dba236..bf76106 100644 --- a/bot_types.py +++ b/bot_types.py @@ -14,4 +14,7 @@ class MatrixBotConfig: """Localpart of matrix username (without homeserver), for example: `valid-username`""" storage_directory: Path - """Path to the storage directory (will be created if needed)""" \ No newline at end of file + """Path to the storage directory (will be created if needed)""" + + auto_verify_all_known_devices: bool = True + """Enable automatic verification of all known devices on related errors""" \ No newline at end of file diff --git a/main.py b/main.py index c3b7225..e1155e0 100644 --- a/main.py +++ b/main.py @@ -5,6 +5,7 @@ import asyncio import traceback import signal +import html from pathlib import Path import config @@ -14,11 +15,20 @@ from new_bot import MatrixBot from bot_types import MatrixBotConfig import nio.events +from nio import MatrixRoom, RoomMessageText from datatypes import AppConfig +bot: MatrixBot = None + +async def callback_test_text(room: MatrixRoom, event: RoomMessageText) -> None: + if event.sender == bot.get_client().user_id: + return + await bot.send_text_to_room(room.room_id, f"{html.escape(event.body)}") + async def main() -> None: """Entry point""" + global bot # setup signal handler util.setup_app_stop_event() loop = asyncio.get_running_loop() @@ -42,6 +52,7 @@ async def main() -> None: storage_directory=Path(cfg.store_dir) ) bot = MatrixBot(matrix_bot_config) + bot.add_event_callback(callback_test_text, RoomMessageText) bot.add_event_callback(None, nio.events.Event) # start the app diff --git a/new_bot.py b/new_bot.py index fc89fbd..8bec982 100644 --- a/new_bot.py +++ b/new_bot.py @@ -7,12 +7,17 @@ import time import json import os import re +from html.parser import HTMLParser from urllib.parse import urlparse from typing import Any, Callable, Awaitable from nio import AsyncClient, AsyncClientConfig, SyncResponse from nio import LoginResponse, LoginError, WhoamiResponse, WhoamiError +from nio import RoomSendResponse, RoomSendError + +from nio import OlmUnverifiedDeviceError + import nio.events from bot_types import * @@ -97,6 +102,32 @@ class MatrixBot: localpart = config.matrix_username_localpart return f"@{localpart}:{homeserver_name}" + @staticmethod + def _check_and_remove_html(text_to_check: str) -> tuple[bool, str]: + """Checks if `text_to_check` is HTML and sanitizes it. + + Returns: + tuple[bool, str] where the first element is True if `text_to_check` contains + valid HTML, and the second element is text without HTML (or just copy of + `text_to_check` if it does not contain HTML) + """ + has_tags = False + text_fragments = [] + class Extractor(HTMLParser): + def handle_starttag(self, tag, attrs): + nonlocal has_tags + has_tags = True + def handle_data(self, data): + text_fragments.append(data) + parser = Extractor(convert_charrefs=True) + parser.feed(text_to_check) + try: + if has_tags: + return (True, " ".join("".join(text_fragments).split())) + except: + traceback.print_exc() + return (False, text_to_check) + @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 @@ -376,9 +407,84 @@ class MatrixBot: 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[[Any], Awaitable[None]] | None, event_class: nio.events.Event) -> None: """Added event callback for events of specified class. Use `None` instead of callback to print parameter types you need to use in your callback.""" if callback is None: callback = self._debug_event_callback - self._client.add_event_callback(callback, event_class) # type: ignore \ No newline at end of file + self._client.add_event_callback(callback, event_class) # type: ignore + + def get_client(self) -> AsyncClient: + """Get AsyncClient in use""" + return self._client + + async def send_text_to_room(self, room_id: str, text: str, is_html: bool | None = None, **kwargs) -> str: + """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_with_html = text + # guess if text is HTML + if is_html is None: + is_html, text = self._check_and_remove_html(text) + # text is HTML + elif is_html: + _, text = self._check_and_remove_html(text) + # create `content` for `room_send()` + content = { + "msgtype": "m.text", + "body": text, + **kwargs + } + if is_html: + content["format"] = "org.matrix.custom.html" + content["formatted_body"] = content["body"] + # 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_text_to_room(room_id, text, is_html, **kwargs) + else: + raise + # success + if type(result) is RoomSendResponse: + return result.event_id + # error + elif type(result) is RoomSendError: + raise RuntimeError(result) + # unknown error + else: + raise RuntimeError("Unknown error has occured", result) \ No newline at end of file