diff --git a/bot_callbacks.py b/bot_callbacks.py index 7ddcba2..b14e257 100644 --- a/bot_callbacks.py +++ b/bot_callbacks.py @@ -4,8 +4,7 @@ import html import time import traceback -from mab import MatrixBot, TextCommandFilter -from nio import MatrixRoom, RoomMessageText +from mab import * # type: ignore import util from database import Database @@ -14,7 +13,6 @@ from datatypes import * # # PRIVATE # -_bot: MatrixBot = None # type: ignore _db: Database = None # type: ignore _COMMANDS = { @@ -27,6 +25,7 @@ _COMMANDS = { "ban": "Забанить IP адрес", "unban": "Разбанить IP адрес", "bans": "Получить список забаненных IP адресов", + "short": "(Небезопасно) Создать токен-имя", } def _generate_help_message() -> str: @@ -46,38 +45,35 @@ def _generate_help_message() -> str: # # COMMAND HANDLERS # -async def _on_cmd_help(room: MatrixRoom, event: RoomMessageText) -> None: +async def _on_cmd_help(ctx: EventContext) -> None: """!help""" - await _bot.send_text_to_room(room.room_id, _generate_help_message(), is_html=True) + await ctx.bot.send_text(ctx.room, _generate_help_message()) -async def _on_cmd_info(room: MatrixRoom, event: RoomMessageText) -> None: +async def _on_cmd_info(ctx: EventContext) -> None: """!info""" # get room info (or add it) - room_info = await _db.room_get(room.room_id) + room_info = await _db.room_get(ctx.room.room_id) if room_info is None: - room_info = await _db.room_create(room.room_id) + room_info = await _db.room_create(ctx.room.room_id) # failure if room_info is None: - await _bot.send_text_to_room( - room.room_id, - "Нет информации о комнате", - is_html=True + await ctx.bot.send_text( + ctx.room, "Нет информации о комнате" ) return # respond response = "Сведения о комнате
" response += f"Канал: {html.escape(room_info.code)}" - await _bot.send_text_to_room(room.room_id, response, is_html=True) + await ctx.bot.send_text(ctx.room, response) -async def _on_cmd_tokens(room: MatrixRoom, event: RoomMessageText) -> None: +async def _on_cmd_tokens(ctx: EventContext) -> None: """!info""" # get all tokens and check if there are none tokens = await _db.token_get_all() if not tokens: - await _bot.send_text_to_room( - room.room_id, - "Нет токенов, используйте !auth", - is_html=True + await ctx.bot.send_text( + ctx.room, + "Нет токенов, используйте !auth" ) return # create the list of tokens @@ -92,15 +88,15 @@ async def _on_cmd_tokens(room: MatrixRoom, event: RoomMessageText) -> None: response += f"
  • Последнее использование: {util.date_to_text(token.last_access_at)}
    • " response += "" # respond - await _bot.send_text_to_room(room.room_id, response, is_html=True) + await ctx.bot.send_text(ctx.room, response) -async def _on_cmd_auth(room: MatrixRoom, event: RoomMessageText) -> None: +async def _on_cmd_auth(ctx: EventContext) -> None: """!auth""" # create new token token = await _db.token_create() # set the name if it is provided - if event.command_args: - token.name = " ".join(event.command_args) + if ctx[CTX_CMD_ARGS]: + token.name = " ".join(ctx[CTX_CMD_ARGS]) await _db.token_set_name(token.code, token.name) # create the response response = "Создан новый токен" @@ -111,39 +107,39 @@ async def _on_cmd_auth(room: MatrixRoom, event: RoomMessageText) -> None: else: response += "указывается в запросе" # respond - await _bot.send_text_to_room(room.room_id, response, is_html=True) + await ctx.bot.send_text(ctx.room, response) -async def _on_cmd_deauth(room: MatrixRoom, event: RoomMessageText) -> None: +async def _on_cmd_deauth(ctx: EventContext) -> None: """!deauth""" # no token provided - if len(event.command_args) != 1: - await _bot.send_text_to_room(room.room_id, "Укажите токен, который надо удалить (должен быть ровно один аргумент)", is_html=True) + if len(ctx[CTX_CMD_ARGS]) != 1: + await ctx.bot.send_text(ctx.room, "Укажите токен, который надо удалить (должен быть ровно один аргумент)") return - token = event.command_args[0] + token = ctx[CTX_CMD_ARGS][0] # check if token does not exist if await _db.token_get(token) is None: - await _bot.send_text_to_room(room.room_id, "Токен не найден", is_html=True) + await ctx.bot.send_text(ctx.room, "Токен не найден") return # remove the token await _db.token_delete(token) # respond response = f"Удалён токен {token}" - await _bot.send_text_to_room(room.room_id, response, is_html=True) + await ctx.bot.send_text(ctx.room, response) -async def _on_cmd_name(room: MatrixRoom, event: RoomMessageText) -> None: +async def _on_cmd_name(ctx: EventContext) -> None: """!name""" # check arguments - if len(event.command_args) < 1: + if len(ctx[CTX_CMD_ARGS]) < 1: error = "Требуется указать как минимум токен. Если хотите убрать имя, то имя указывать не надо. Если имя нужно назначить или изменить, то после токена укажите новое имя." - await _bot.send_text_to_room(room.room_id, error, is_html=True) + await ctx.bot.send_text(ctx.room, error) return - token = event.command_args[0] - new_name = " ".join(event.command_args[1:]) + token = ctx[CTX_CMD_ARGS][0] + new_name = " ".join(ctx[CTX_CMD_ARGS][1:]) if not new_name.strip(): new_name = None # check if token exists if await _db.token_get(token) is None: - await _bot.send_text_to_room(room.room_id, "Токен не существует", is_html=True) + await ctx.bot.send_text(ctx.room, "Токен не существует") return # set new name await _db.token_set_name(token, new_name) @@ -153,51 +149,51 @@ async def _on_cmd_name(room: MatrixRoom, event: RoomMessageText) -> None: else: response = f"Удалено имя для токена {token}" # respond - await _bot.send_text_to_room(room.room_id, response, is_html=True) + await ctx.bot.send_text(ctx.room, response) -async def _on_cmd_ban(room: MatrixRoom, event: RoomMessageText) -> None: +async def _on_cmd_ban(ctx: EventContext) -> None: """!ban""" # check arguments - if len(event.command_args) < 2: + if len(ctx[CTX_CMD_ARGS]) < 2: error = "Формат: !ban <IP> <SECONDS> [REASON]" - await _bot.send_text_to_room(room.room_id, error, is_html=True) + await ctx.bot.send_text(ctx.room, error) return # get event.command_args try: - ip = event.command_args[0] - duration = float(event.command_args[1]) - reason = " ".join(event.command_args[2:]) if event.command_args[2:] else "Manual ban" + ip = ctx[CTX_CMD_ARGS][0] + duration = float(ctx[CTX_CMD_ARGS][1]) + reason = " ".join(ctx[CTX_CMD_ARGS][2:]) if ctx[CTX_CMD_ARGS][2:] else "Manual ban" except: error = "Возникла ошибка. Наверняка неправильно указаны секунды." - await _bot.send_text_to_room(room.room_id, error, is_html=True) + await ctx.bot.send_text(ctx.room, error) return # ban try: await _db.ban_create(ip, time.time() + duration, reason) - await _bot.send_text_to_room(room.room_id, "IP адрес заблокирован", is_html=True) + await ctx.bot.send_text(ctx.room, "IP адрес заблокирован") except: - await _bot.send_text_to_room(room.room_id, "Возникла ошибка", is_html=True) + await ctx.bot.send_text(ctx.room, "Возникла ошибка") -async def _on_cmd_unban(room: MatrixRoom, event: RoomMessageText) -> None: +async def _on_cmd_unban(ctx: EventContext) -> None: """!unban""" # check arguments - if len(event.command_args) != 1: + if len(ctx[CTX_CMD_ARGS]) != 1: error = "Формат: !ban <IP>" - await _bot.send_text_to_room(room.room_id, error, is_html=True) + await ctx.bot.send_text(ctx.room, error) return # unban try: - await _db.ban_delete(event.command_args[0]) - await _bot.send_text_to_room(room.room_id, "IP адрес разблокирован (если он был заблокирован)", is_html=True) + await _db.ban_delete(ctx[CTX_CMD_ARGS][0]) + await ctx.bot.send_text(ctx.room, "IP адрес разблокирован (если он был заблокирован)") except: - await _bot.send_text_to_room(room.room_id, "Возникла ошибка", is_html=True) + await ctx.bot.send_text(ctx.room, "Возникла ошибка") -async def _on_cmd_bans(room: MatrixRoom, event: RoomMessageText) -> None: +async def _on_cmd_bans(ctx: EventContext) -> None: """!bans""" # list bans = _db.ban_get_all() if not bans: - await _bot.send_text_to_room(room.room_id, "Нет заблокированных IP адресов", is_html=True) + await ctx.bot.send_text(ctx.room, "Нет заблокированных IP адресов") return # create the response result = "Список заблокированных IP
    " result += "" - await _bot.send_text_to_room(room.room_id, result, is_html=True) + await ctx.bot.send_text(ctx.room, result) + +async def _on_cmd_short(ctx: EventContext) -> None: + """!auth""" + # get token to use + if len(ctx[CTX_CMD_ARGS]) != 1: + error = "Формат: !short <token>" + await ctx.bot.send_text(ctx.room, error) + return + # create new token + try: + token = await _db.token_create(ctx[CTX_CMD_ARGS][0]) + except Exception as e: + await ctx.bot.send_text(ctx.room, f"Не удалось создать токен: {" ".join(e.args)}") + return + # create the response + response = "Создан новый токен" + response += f"
    Код: {token.code}" + response += f"
    Имя сервиса: указывается в запросе" + # respond + await ctx.bot.send_text(ctx.room, response) # @@ -215,24 +231,45 @@ async def _on_cmd_bans(room: MatrixRoom, event: RoomMessageText) -> None: # def setup(bot: MatrixBot, db: Database) -> None: """Setup the callbacks""" - global _bot, _db + global _db _db = db - _bot = bot - _bot.add_event_callback( - _on_cmd_help, TextCommandFilter(["help", "h", "?"])) - _bot.add_event_callback( - _on_cmd_info, TextCommandFilter(["info", "room", "i", "r"])) - _bot.add_event_callback( - _on_cmd_tokens, TextCommandFilter(["tokens", "t"])) - _bot.add_event_callback( - _on_cmd_auth, TextCommandFilter(["auth", "create", "a", "c"])) - _bot.add_event_callback( - _on_cmd_deauth, TextCommandFilter(["deauth", "delete", "d"])) - _bot.add_event_callback( - _on_cmd_name, TextCommandFilter(["name", "n"])) - _bot.add_event_callback( - _on_cmd_ban, TextCommandFilter(["ban", "b"])) - _bot.add_event_callback( - _on_cmd_unban, TextCommandFilter(["unban", "u"])) - _bot.add_event_callback( - _on_cmd_bans, TextCommandFilter(["bans", "l"])) \ No newline at end of file + bot.add_callback( + ~SenderIsBotFilter() & BodyCommandFilter(["info", "room", "i", "r"]), + _on_cmd_info + ) + bot.add_callback( + ~SenderIsBotFilter() & BodyCommandFilter(["tokens", "t"]), + _on_cmd_tokens + ) + bot.add_callback( + ~SenderIsBotFilter() & BodyCommandFilter(["auth", "create", "a", "c"]), + _on_cmd_auth + ) + bot.add_callback( + ~SenderIsBotFilter() & BodyCommandFilter(["deauth", "delete", "d"]), + _on_cmd_deauth + ) + bot.add_callback( + ~SenderIsBotFilter() & BodyCommandFilter(["name", "n"]), + _on_cmd_name + ) + bot.add_callback( + ~SenderIsBotFilter() & BodyCommandFilter(["ban", "b"]), + _on_cmd_ban + ) + bot.add_callback( + ~SenderIsBotFilter() & BodyCommandFilter(["unban", "u"]), + _on_cmd_unban + ) + bot.add_callback( + ~SenderIsBotFilter() & BodyCommandFilter(["bans", "l"]), + _on_cmd_bans + ) + bot.add_callback( + ~SenderIsBotFilter() & BodyCommandFilter(["short", "s"]), + _on_cmd_short + ) + bot.add_callback( + ~SenderIsBotFilter() & NewMessageFilter(), + _on_cmd_help + ) \ No newline at end of file diff --git a/database.py b/database.py index fe3fadb..debdd45 100644 --- a/database.py +++ b/database.py @@ -252,11 +252,17 @@ class Database: return None - async def token_create(self) -> ObjectToken: + async def token_create(self, code: str | None = None) -> ObjectToken: """Create a token.""" if self._connection is None: raise RuntimeError("Not connected to the database") - code = get_hash(str(time.time()).encode() + os.urandom(64)) + if code is None: + code = get_hash(str(time.time()).encode() + os.urandom(64)) + else: + if len(code) < 2 or len(code) > 16: + raise RuntimeError("Token must be 2-16 symbols long") + if not all(ord(c) < 128 and (c.islower() or c.isdigit()) for c in code): + raise RuntimeError("Only digits and ASCII lowercase letters are allowed") create_time = time.time() statement = """ INSERT INTO tokens (code, name, created_at, last_access_at) @@ -329,6 +335,17 @@ class Database: except: self._logger.error(traceback.format_exc()) + async def token_set_last_access(self, code: str, timestamp: float) -> None: + """Set last_access_at for the token.""" + if self._connection is None: + raise RuntimeError("Not connected to the database") + try: + statement = "UPDATE tokens SET last_access_at=? WHERE code=?" + await self._connection.execute(statement, (timestamp, code)) + await self._connection.commit() + except: + self._logger.error(traceback.format_exc()) + async def ban_create(self, ip: str, expires_at: float, reason: str) -> None: """Save information about banned IP address. Replaces existing IPs.""" diff --git a/main.py b/main.py index 2d4a19d..e7a5c67 100644 --- a/main.py +++ b/main.py @@ -55,7 +55,7 @@ async def main() -> None: if not await db.connect(): util.log_error("Can't connect to the database!") return - bot.start() + await bot.start() await web_server.start() # wait for Ctrl+C diff --git a/requirements.txt b/requirements.txt index 1dea17d..dce7930 100644 --- a/requirements.txt +++ b/requirements.txt @@ -34,7 +34,7 @@ idna==3.19 Jinja2==3.1.6 jsonschema==4.26.0 jsonschema-specifications==2025.9.1 -mab @ git+https://git.tyukalov.su/nikita/mab@v0.1.0 +mab @ git+https://git.tyukalov.su/nikita/mab@v0.5.0 markdown-it-py==4.2.0 MarkupSafe==3.0.3 matrix-nio==0.26.0 @@ -42,6 +42,7 @@ mdurl==0.1.2 multidict==6.7.1 packaging==26.3 peewee==3.19.0 +pillow==12.3.0 propcache==0.5.2 pycryptodome==3.23.0 pydantic==2.13.5 @@ -51,6 +52,7 @@ pydantic_core==2.46.5 Pygments==2.21.0 pyproject_hooks==1.2.0 python-dotenv==1.2.3 +python-magic==0.4.27 python-multipart==0.0.32 python-socks==3.0.0 PyYAML==6.0.3 diff --git a/web_middleware.py b/web_middleware.py index f980ea0..ceed405 100644 --- a/web_middleware.py +++ b/web_middleware.py @@ -84,6 +84,7 @@ class AuthMiddleware(BaseHTTPMiddleware): status_code=status.HTTP_403_FORBIDDEN, content={"detail": "No such token"} ) + await database.token_set_last_access(token, time.time()) request.state.token = token_data return await call_next(request) diff --git a/web_routes.py b/web_routes.py index e9f9edd..0a82f06 100644 --- a/web_routes.py +++ b/web_routes.py @@ -58,7 +58,7 @@ async def _get_notify(request: Request, params: ModelGetNotify = Depends()): text += html.escape(params.text) # try to send the notification try: - notification_id = await matrix.send_text_to_room( + notification_id = await matrix.send_text( room.matrix_id, text, is_html=True