Basic database function implemented

2026-08-25 21:13:27 +03:00
parent ff6656fdb0
commit ddc1c36c72
7 changed files with 360 additions and 63 deletions
--- a/.gitignore
+++ b/.gitignore
@@ -4,4 +4,5 @@ __pycache__/
 *.swp
 *.swo
 *.vscode
-*.json
+*.json
+*.sqlite
--- 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 = "<b><i>Как использовать</i></b><br><ol>"
+    result = "<strong><i>Как использовать</i></strong><br><ol>"
    result += "<li>Используя <code>!auth</code>, создайте токен</li>"
    result += "<li>Добавьте бота в комнату для уведомлений</li>"
    result += "<li>Выполните <code>!info</code> в комнате, чтобы узнать код канала уведомений</li>"
    result += "<li>Используя полученные токен и код канала, отправьте уведомление через веб-запрос</li>"
    result += "</ol>"

-    result += "<br><br><b><i>Команды</i></b>"
+    result += "<br><br><strong><i>Команды</i></strong>"
    for aliases in _COMMANDS:
        result += f"<br><code>!{aliases[0]}</code> - <i>{html.escape(_COMMANDS[aliases][1])}</i>"
    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, "<b>Используйте <code>!help</code></b>")
+        await _bot.send_text_to_room(room.room_id, "<strong>Используйте <code>!help</code></strong>")
        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 = "<b><i>Сведения о комнате</i></b><br>"
-    response += f"<b>channel: <b><code>{html.escape(room.room_id)}</code>"
+    # 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,
+            "<strong>Нет информации о комнате</strong>",
+            is_html=True
+        )
+        return
+    # respond
+    response = "<strong><i>Сведения о комнате</i></strong><br>"
+    response += f"<strong>Канал:</strong> <code>{html.escape(room_info.code)}</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 = "<b><i>Список токенов</i></b>"
-    for token in TEST_LIST:
-        data = TEST_LIST[token]
-        response += f"<br><b><code>{html.escape(token)}</code></b><br><ul>"
-
-        if data["service_name"] is not None:
-            response += f"<li><b>Имя службы:</b> <code>{data["service_name"]}</code></li>"
+    # 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,
+            "<strong>Нет токенов, используйте <code>!auth</code></strong>",
+            is_html=True
+        )
+        return
+    # create the list of tokens
+    response = "<strong><i>Список токенов</i></strong>"
+    for token in tokens:
+        response += f"<br><strong><code>{html.escape(token.code)}</code></strong><br><ul>"
+        if token.name is not None:
+            response += f"<li><strong>Имя службы:</strong> <code>{token.name}</code></li>"
        else:
-            response += f"<li><b>Имя службы:</b> <i>указывается в запросе</i></li>"
-        response += f"<li><b>Создан:</b> <i>{util.date_to_text(data["created"])}</i></li>"
-        response += f"<li><b>Последнее использование:</b> <i>{util.date_to_text(data["last_use"])}</i></li>"
-        response += f"<li><b>Всего использований:</b> <i>{data["total_uses"]}</i></li>"
-
+            response += f"<li><strong>Имя службы:</strong> <i>указывается в запросе</i></li>"
+        response += f"<li><strong>Создан:</strong> <i>{util.date_to_text(token.created_at)}</i></li>"
+        response += f"<li><strong>Последнее использование:</strong> <i>{util.date_to_text(token.last_access_at)}</i></li>"
        response += "</ul>"
+    # 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 = "<b>Создан новый токен</b>"
-    response += f"<br><b>Код:</b> <code>{new_token}</code>"
-    response += f"<br><b>Имя сервиса:</b> "
-    if new_token_name:
-        response += f"<code>{html.escape(new_token_name)}</code>"
+    # 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 = "<strong>Создан новый токен</strong>"
+    response += f"<br><strong>Код:</strong> <code>{token.code}</code>"
+    response += f"<br><strong>Имя сервиса:</strong> "
+    if token.name:
+        response += f"<code>{html.escape(token.name)}</code>"
    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, "<b>Укажите токен, который надо удалить (должен быть ровно один аргумент)</b>", is_html=True)
+        await _bot.send_text_to_room(room.room_id, "<strong>Укажите токен, который надо удалить (должен быть ровно один аргумент)</strong>", is_html=True)
        return
    token = args[0]
-    # TODO
-    response = f"<b>Удалён токен <code>{token}</code></b>"
+    # check if token does not exist
+    if await _db.token_get(token) is None:
+        await _bot.send_text_to_room(room.room_id, "<strong>Токен не найден</strong>", is_html=True)
+        return
+    # remove the token
+    await _db.token_delete(token)
+    # respond
+    response = f"<strong>Удалён токен <code>{token}</code></strong>"
    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 = "<b>Требуется указать как минимум токен. Если хотите убрать имя, то имя указывать не надо. Если имя нужно назначить или изменить, то после токена укажите новое имя.</b>"
+        error = "<strong>Требуется указать как минимум токен. Если хотите убрать имя, то имя указывать не надо. Если имя нужно назначить или изменить, то после токена укажите новое имя.</strong>"
        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, "<strong>Токен не существует</strong>", is_html=True)
+        return
+    # set new name
+    await _db.token_set_name(token, new_name)
+    # prepare the response
    if new_name:
-        response = f"<b>Новое имя для токена <code>{token}</code>: <code>{new_name}</code></b>"
+        response = f"<strong>Новое имя для токена <code>{token}</code>: <code>{new_name}</code></strong>"
    else:
-        response = f"<b>Удалено имя для токена <code>{token}</code></b>"
-    
+        response = f"<strong>Удалено имя для токена <code>{token}</code></strong>"
+    # 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, "<b>Команда не реализована</b>", is_html=True)
+    await _bot.send_text_to_room(room.room_id, "<strong>Команда не реализована</strong>", 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
--- 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())
--- a/datatypes.py
+++ b/datatypes.py
@@ -11,5 +11,14 @@ class AppConfig:
    matrix_user: str
    store_dir: str

-class MessageType(Enum):
-    TEXT = "m.text"
+@dataclass
+class ObjectRoom:
+    code: str
+    matrix_id: str
+
+@dataclass
+class ObjectToken:
+    code: str
+    name: str | None
+    created_at: float
+    last_access_at: float
--- 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:
--- 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
--- 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
+    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