diff --git a/.gitignore b/.gitignore index bca140b..ab5eb8f 100644 --- a/.gitignore +++ b/.gitignore @@ -4,4 +4,5 @@ __pycache__/ *.swp *.swo *.vscode -*.json \ No newline at end of file +*.json +*.sqlite \ No newline at end of file diff --git a/bot_callbacks.py b/bot_callbacks.py index a5c5e33..87cbf63 100644 --- a/bot_callbacks.py +++ b/bot_callbacks.py @@ -8,22 +8,25 @@ from mab import MatrixBot from nio import MatrixRoom, RoomMessageText import util +from database import Database +from datatypes import * # # PRIVATE # _bot: MatrixBot = None # type: ignore +_db: Database = None # type: ignore def _generate_help_message() -> str: """Generates help message in HTML markup""" - result = "Как использовать
    " + result = "Как использовать
      " result += "
    1. Используя !auth, создайте токен
      2. " result += "
    2. Добавьте бота в комнату для уведомлений
      3. " result += "
    3. Выполните !info в комнате, чтобы узнать код канала уведомений
      4. " result += "
    4. Используя полученные токен и код канала, отправьте уведомление через веб-запрос
      5. " result += "
    " - result += "

    Команды" + result += "

    Команды" for aliases in _COMMANDS: result += f"
    !{aliases[0]} - {html.escape(_COMMANDS[aliases][1])}" return result @@ -51,7 +54,7 @@ async def _on_text(room: MatrixRoom, event: RoomMessageText) -> None: if cb: break if cb is None: - await _bot.send_text_to_room(room.room_id, "Используйте !help") + await _bot.send_text_to_room(room.room_id, "Используйте !help") return try: await cb(room, parts[1:]) @@ -67,93 +70,112 @@ async def _on_cmd_help(room: MatrixRoom, args: list[str]) -> None: async def _on_cmd_info(room: MatrixRoom, args: list[str]) -> None: """!info""" - response = "Сведения о комнате
    " - response += f"channel: {html.escape(room.room_id)}" + # get room info (or add it) + room_info = await _db.room_get(room.room_id) + if room_info is None: + room_info = await _db.room_create(room.room_id) + # failure + if room_info is None: + await _bot.send_text_to_room( + room.room_id, + "Нет информации о комнате", + is_html=True + ) + return + # respond + response = "Сведения о комнате
    " + response += f"Канал: {html.escape(room_info.code)}" await _bot.send_text_to_room(room.room_id, response, is_html=True) async def _on_cmd_tokens(room: MatrixRoom, args: list[str]) -> None: """!info""" - TEST_LIST = { - "t0K3n": { - "created": time.time() - 60, - "last_use": time.time() - 30, - "total_uses": 534, - "service_name": "Резервное копирование", - }, - "T35t_t0KeN": { - "created": time.time() - 12341, - "last_use": time.time() - 123, - "total_uses": 342, - "service_name": None, - }, - } - - response = "Список токенов" - for token in TEST_LIST: - data = TEST_LIST[token] - response += f"
    {html.escape(token)}
      " - - if data["service_name"] is not None: - response += f"
    • Имя службы: {data["service_name"]}
      • " + # 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 + ) + return + # create the list of tokens + response = "Список токенов" + for token in tokens: + response += f"
      {html.escape(token.code)}
        " + if token.name is not None: + response += f"
      • Имя службы: {token.name}
        • " else: - response += f"
      • Имя службы: указывается в запросе
        • " - response += f"
      • Создан: {util.date_to_text(data["created"])}
        • " - response += f"
      • Последнее использование: {util.date_to_text(data["last_use"])}
        • " - response += f"
      • Всего использований: {data["total_uses"]}
        • " - + response += f"
      • Имя службы: указывается в запросе
        • " + response += f"
      • Создан: {util.date_to_text(token.created_at)}
        • " + 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) async def _on_cmd_auth(room: MatrixRoom, args: list[str]) -> None: """!auth""" - new_token = "inr3i4evfds" - new_token_name = " ".join(args) if args else None - if type(new_token_name) is str and not new_token_name.strip(): - new_token_name = None - - response = "Создан новый токен" - response += f"
      Код: {new_token}" - response += f"
      Имя сервиса: " - if new_token_name: - response += f"{html.escape(new_token_name)}" + # create new token + token = await _db.token_create() + # set the name if it is provided + if args: + token.name = " ".join(args) + await _db.token_set_name(token.code, token.name) + # create the response + response = "Создан новый токен" + response += f"
      Код: {token.code}" + response += f"
      Имя сервиса: " + if token.name: + response += f"{html.escape(token.name)}" else: response += "указывается в запросе" + # respond await _bot.send_text_to_room(room.room_id, response, is_html=True) async def _on_cmd_deauth(room: MatrixRoom, args: list[str]) -> None: """!deauth""" + # no token provided if len(args) != 1: - await _bot.send_text_to_room(room.room_id, "Укажите токен, который надо удалить (должен быть ровно один аргумент)", is_html=True) + await _bot.send_text_to_room(room.room_id, "Укажите токен, который надо удалить (должен быть ровно один аргумент)", is_html=True) return token = args[0] - # TODO - response = f"Удалён токен {token}" + # 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) + 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) async def _on_cmd_name(room: MatrixRoom, args: list[str]) -> None: """!name""" + # check arguments if len(args) < 1: - error = "Требуется указать как минимум токен. Если хотите убрать имя, то имя указывать не надо. Если имя нужно назначить или изменить, то после токена укажите новое имя." + error = "Требуется указать как минимум токен. Если хотите убрать имя, то имя указывать не надо. Если имя нужно назначить или изменить, то после токена укажите новое имя." await _bot.send_text_to_room(room.room_id, error, is_html=True) return token = args[0] - new_name = " ".join(args[1:]) if not new_name.strip(): new_name = None - - # TODO - + # check if token exists + if await _db.token_get(token) is None: + await _bot.send_text_to_room(room.room_id, "Токен не существует", is_html=True) + return + # set new name + await _db.token_set_name(token, new_name) + # prepare the response if new_name: - response = f"Новое имя для токена {token}: {new_name}" + response = f"Новое имя для токена {token}: {new_name}" else: - response = f"Удалено имя для токена {token}" - + response = f"Удалено имя для токена {token}" + # respond await _bot.send_text_to_room(room.room_id, response, is_html=True) async def _on_cmd_leave(room: MatrixRoom, args: list[str]) -> None: """!leave""" - await _bot.send_text_to_room(room.room_id, "Команда не реализована", is_html=True) + await _bot.send_text_to_room(room.room_id, "Команда не реализована", is_html=True) _COMMANDS = { ("help", "h", "?"): (_on_cmd_help, "Получить справку"), @@ -168,8 +190,9 @@ _COMMANDS = { # # PUBLIC # -def setup(bot: MatrixBot) -> None: +def setup(bot: MatrixBot, db: Database) -> None: """Setup the callbacks""" - global _bot + global _bot, _db + _db = db _bot = bot _bot.add_event_callback(_on_text, RoomMessageText) # type: ignore diff --git a/database.py b/database.py index e69de29..c9e3d9e 100644 --- a/database.py +++ b/database.py @@ -0,0 +1,249 @@ +"""This module implements database operations.""" + +import logging +import traceback +import time +import os +from pathlib import Path +from typing import Any +import aiosqlite +from aiosqlite import Connection, Row + +from datatypes import * +from util import get_hash + + +class Database: + # + # PRIVATE + # + @staticmethod + async def _setup_tables(conn: Connection) -> None: + SETUP_SQL_SCRIPT = """ + CREATE TABLE IF NOT EXISTS rooms ( + code TEXT PRIMARY KEY, + matrix_id TEXT NOT NULL UNIQUE + ); + + CREATE TABLE IF NOT EXISTS tokens ( + code TEXT PRIMARY KEY, + name TEXT, + created_at REAL, + last_access_at REAL + ); + + CREATE TABLE IF NOT EXISTS fails ( + ip TEXT PRIMARY KEY, + timestamp INTEGER, + score INTEGER + ); + + CREATE TABLE IF NOT EXISTS bans ( + ip TEXT PRIMARY KEY, + expires_at INTEGER, + reason TEXT + ); + """ + await conn.executescript(SETUP_SQL_SCRIPT) + await conn.commit() + + @staticmethod + async def _select_by_anded_kwargs(conn: Connection, table_name: str, **kwargs) -> list[dict[str, Any]]: + # prepare keys and values + keys = tuple(k for k in kwargs) + values = tuple(kwargs[k] for k in keys) + # prepare the statement + statement = f"SELECT * FROM {table_name}" + # add `WHERE` part if there kwargs + if keys: + statement += " WHERE " + statement += " AND ".join([f"{k}=?" for k in keys]) + else: + values = None + # execute and enumerate + result = [] + async with conn.execute(statement, values) as cursor: + async for row in cursor: + entry = {} + for k in row.keys(): + entry[k] = row[k] + result.append(entry) + return result + + # + # PUBLIC + # + def __init__(self, path: Path): + self._path = path + self._logger = logging.getLogger("database") + + self._connection: Connection | None = None + + async def connect(self) -> bool: + """Connect to the database. Returns False on failure.""" + if self._connection is not None: + return False + try: + self._connection = await aiosqlite.connect(self._path) + self._connection.row_factory = Row + await self._setup_tables(self._connection) + except: + self._logger.error(traceback.format_exc()) + return False + return True + + async def disconnect(self) -> None: + """Disconnect from the database.""" + if self._connection is None: + return + await self._connection.close() + self._connection = None + + + async def room_create(self, matrix_id: str) -> ObjectRoom | None: + """Create a room with specified matrix_id. Returns None if the room exists.""" + if self._connection is None: + raise RuntimeError("Not connected to the database") + # room exists + if await self.room_get(matrix_id) is not None: + return None + # create the code + code = None + string_to_hash = matrix_id + while code is None or (await self.room_get(code)) is not None: + code = get_hash(string_to_hash)[-10:] + string_to_hash += "A" + # code must always start with q + code = f"q{code}" + # add + try: + statement = "INSERT INTO rooms (code, matrix_id) VALUES (?, ?)" + await self._connection.execute(statement, (code, matrix_id)) + await self._connection.commit() + self._logger.info(f"Added new room with code {code}") + except: + self._logger.error(traceback.format_exc()) + return None + # return the added object + return ObjectRoom( + code=code, + matrix_id=matrix_id + ) + + async def room_get_all(self) -> list[ObjectRoom]: + """Get room by code.""" + if self._connection is None: + raise RuntimeError("Not connected to the database") + try: + rooms = await self._select_by_anded_kwargs( + self._connection, + "rooms" + ) + result = [] + for r in rooms: + result.append(ObjectRoom( + code=r["code"], + matrix_id=r["matrix_id"] + )) + return result + except: + self._logger.error(traceback.format_exc()) + return [] + + async def room_get(self, identifier: str) -> ObjectRoom | None: + """Get room by identifier (either `code` or `matrix_id`).""" + if self._connection is None: + raise RuntimeError("Not connected to the database") + try: + kwargs = {} + kwargs["code" if identifier[0] == "q" else "matrix_id"] = identifier + rooms = await self._select_by_anded_kwargs( + self._connection, + "rooms", + **kwargs + ) + if not rooms: + self._logger.debug(f"Room `{identifier}` is not found") + return None + return ObjectRoom(**rooms[0]) + except: + self._logger.error(traceback.format_exc()) + return None + + + async def token_create(self) -> 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)) + create_time = time.time() + statement = """ + INSERT INTO tokens (code, name, created_at, last_access_at) + VALUES (?, ?, ?, ?) + """ + result = ObjectToken( + code=code, + name=None, + created_at=create_time, + last_access_at=create_time + ) + await self._connection.execute( + statement, + (result.code, result.name, result.created_at, result.last_access_at) + ) + await self._connection.commit() + return result + + async def token_delete(self, code: str) -> None: + """Delete a token.""" + if self._connection is None: + raise RuntimeError("Not connected to the database") + try: + await self._connection.execute("DELETE FROM tokens WHERE code=?", (code,)) + await self._connection.commit() + except: + self._logger.error(traceback.format_exc()) + + async def token_get(self, code: str) -> ObjectToken | None: + """Get token by its code.""" + if self._connection is None: + raise RuntimeError("Not connected to the database") + try: + result = await self._select_by_anded_kwargs( + self._connection, + "tokens", + code=code + ) + if not result: + return None + result = result[0] + return ObjectToken(**result) + except: + self._logger.error(traceback.format_exc()) + return None + + async def token_get_all(self) -> list[ObjectToken]: + """Get list of all tokens.""" + if self._connection is None: + raise RuntimeError("Not connected to the database") + try: + result = await self._select_by_anded_kwargs( + self._connection, + "tokens" + ) + result = [ObjectToken(**r) for r in result] + return result + except: + self._logger.error(traceback.format_exc()) + return [] + + async def token_set_name(self, code: str, name: str | None) -> None: + """Set name for the token.""" + if self._connection is None: + raise RuntimeError("Not connected to the database") + try: + statement = "UPDATE tokens SET name=? WHERE code=?" + await self._connection.execute(statement, (name, code)) + await self._connection.commit() + except: + self._logger.error(traceback.format_exc()) \ No newline at end of file diff --git a/datatypes.py b/datatypes.py index 838a7b9..7eab60e 100644 --- a/datatypes.py +++ b/datatypes.py @@ -11,5 +11,14 @@ class AppConfig: matrix_user: str store_dir: str -class MessageType(Enum): - TEXT = "m.text" \ No newline at end of file +@dataclass +class ObjectRoom: + code: str + matrix_id: str + +@dataclass +class ObjectToken: + code: str + name: str | None + created_at: float + last_access_at: float \ No newline at end of file diff --git a/main.py b/main.py index 5cd99ba..d8f984b 100644 --- a/main.py +++ b/main.py @@ -10,6 +10,7 @@ from pathlib import Path from mab import MatrixBot, MatrixBotConfig from datatypes import AppConfig +import database import config import util import bot_callbacks @@ -42,9 +43,15 @@ async def main() -> None: storage_directory=Path(cfg.store_dir) ) bot = MatrixBot(matrix_bot_config) - bot_callbacks.setup(bot) + # setup the database + db = database.Database(Path("database.sqlite")) + # setup the callbacks + bot_callbacks.setup(bot, db) # start the app + if not await db.connect(): + util.log_error("Can't connect to the database!") + return bot.start() # wait for Ctrl+C @@ -52,6 +59,7 @@ async def main() -> None: # stop the app await bot.stop() + await db.disconnect() if __name__ == "__main__": try: diff --git a/requirements.txt b/requirements.txt index 484a97c..9392ec9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,6 +4,7 @@ aiohappyeyeballs==2.7.1 aiohttp==3.14.3 aiohttp_socks==0.12.0 aiosignal==1.4.0 +aiosqlite==0.22.1 atomicwrites==1.4.1 attrs==26.1.0 build==1.5.0 diff --git a/util.py b/util.py index acdb31e..1e2718c 100644 --- a/util.py +++ b/util.py @@ -8,6 +8,7 @@ import os import logging import traceback import datetime +import hashlib import aioconsole @@ -75,7 +76,7 @@ async def get_password() -> str | None: return None return await ainput("Matrix password: ") -def date_to_text(date: datetime.datetime, dow: bool = True, seconds: bool = True) -> str: +def date_to_text(date: datetime.datetime | float, dow: bool = True, seconds: bool = True) -> str: ''' Returns date as formatted string. Day of week can be added. Seconds can be added. @@ -102,9 +103,14 @@ def date_to_text(date: datetime.datetime, dow: bool = True, seconds: bool = True ('Sat', 'Сб'), ('Sun', 'Вс') ] - result = date.strftime(format_string) + result = date.strftime(format_string) # type: ignore for en, ru in translate_map: if en in result: result = result.replace(en, ru) break - return result \ No newline at end of file + return result + +def get_hash(data: bytes | str) -> str: + if type(data) is str: + data = data.encode(errors="ignore") + return hashlib.sha256(data).hexdigest() # type: ignore \ No newline at end of file