Removed !leave, added IP ban system

README.md

@@ -51,7 +51,9 @@ python main.py
   сервиса - всегда будет использовано имя, указанное вами)
 - `!deauth <TOKEN>` - удалить токен
 - `!name <TOKEN> [SERVICE_NAME]` - указать (или удалить) имя сервиса для токена
-- `!leave` - покинуть комнату
+- `!ban <IP> <SECONDS> [REASON]` - забанить указанный IP на указанное число секунд (можно указать причину)
+- `!unban <IP>` - разбанить указанный IP адрес
+- `!bans` - получить список забаненных IP адресов
 
 ## Как работает веб-сервер
 

bot_callbacks.py

@@ -173,9 +173,61 @@ async def _on_cmd_name(room: MatrixRoom, args: list[str]) -> None:
     # 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:
+async def _on_cmd_ban(room: MatrixRoom, args: list[str]) -> None:
-    """!leave"""
+    """!ban"""
-    await _bot.send_text_to_room(room.room_id, "<strong>Команда не реализована</strong>", is_html=True)
+    # check arguments
+    if len(args) < 2:
+        error = "<strong>Формат: <code>!ban &lt;IP&gt; &lt;SECONDS&gt; [REASON]</code></strong>"
+        await _bot.send_text_to_room(room.room_id, error, is_html=True)
+        return
+    # get args
+    try:
+        ip = args[0]
+        duration = float(args[1])
+        reason = " ".join(args[2:]) if args[2:] else "Manual ban"
+    except:
+        error = "<strong>Возникла ошибка. Наверняка неправильно указаны секунды.</strong>"
+        await _bot.send_text_to_room(room.room_id, error, is_html=True)
+        return
+    # ban
+    try:
+        await _db.ban_create(ip, time.time() + duration, reason)
+        await _bot.send_text_to_room(room.room_id, "<strong>IP адрес заблокирован</strong>", is_html=True)
+    except:
+        await _bot.send_text_to_room(room.room_id, "<strong>Возникла ошибка</strong>", is_html=True)
+
+async def _on_cmd_unban(room: MatrixRoom, args: list[str]) -> None:
+    """!unban"""
+    # check arguments
+    if len(args) != 1:
+        error = "<strong>Формат: <code>!ban &lt;IP&gt;</code></strong>"
+        await _bot.send_text_to_room(room.room_id, error, is_html=True)
+        return
+    # unban
+    try:
+        await _db.ban_delete(args[0])
+        await _bot.send_text_to_room(room.room_id, "<strong>IP адрес разблокирован (если он был заблокирован)</strong>", is_html=True)
+    except:
+        await _bot.send_text_to_room(room.room_id, "<strong>Возникла ошибка</strong>", is_html=True)
+
+async def _on_cmd_bans(room: MatrixRoom, args: list[str]) -> None:
+    """!bans"""
+    # list
+    bans = _db.ban_get_all()
+    if not bans:
+        await _bot.send_text_to_room(room.room_id, "<strong>Нет заблокированных IP адресов</strong>", is_html=True)
+        return
+    await _bot.send_text_to_room(room.room_id, f"<strong>Статус бана 127.0.0.1: {"БАН" if _db.ban_check("127.0.0.1") else "небан"}</strong>", is_html=True)
+    # create the response
+    result = "<strong>Список заблокированных IP</strong><br><ul>"
+    for ban in bans:
+        result += f"<li><strong>{html.escape(ban["ip"])}</strong><ul>"
+        result += f"<li><strong>Действует до:</strong> <code>{util.date_to_text(ban["expires_at"])}</code></li>"
+        result += f"<li><strong>Причина:</strong> <code>{html.escape(ban["reason"])}</code></li>"
+        result += "</ul></li>"
+    result += "</ul>"
+    await _bot.send_text_to_room(room.room_id, result, is_html=True)
+
 
 _COMMANDS = {
     ("help", "h", "?"): (_on_cmd_help, "Получить справку"),
@@ -184,7 +236,9 @@ _COMMANDS = {
     ("auth", "create", "a", "c"): (_on_cmd_auth, "Создать новый токен"),
     ("deauth", "delete", "d"): (_on_cmd_deauth, "Удалить существующий токен"),
     ("name", "n"): (_on_cmd_name, "Задать (или удалить) имя для токена"),
-    ("leave", "l"): (_on_cmd_leave, "Покинуть комнату, в которой получена эта команда"),
+    ("ban", "b"): (_on_cmd_ban, "Забанить IP адрес"),
+    ("unban", "u"): (_on_cmd_unban, "Разбанить IP адрес"),
+    ("bans", "l"): (_on_cmd_bans, "Получить список забаненных IP адресов"),
 }
 
 #

database.py

@@ -1,5 +1,6 @@
 """This module implements database operations."""
 
+import asyncio
 import logging
 import traceback
 import time
@@ -17,6 +18,9 @@ class Database:
     #
     # PRIVATE
     #
+    BACKGROUND_ROUTINE_PERIOD = 60
+
+
     @staticmethod
     async def _setup_tables(conn: Connection) -> None:
         SETUP_SQL_SCRIPT = """
@@ -32,12 +36,6 @@ class Database:
             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,
@@ -70,14 +68,65 @@ class Database:
                 result.append(entry)
         return result
 
+    async def _background_routine(self, stop_event: asyncio.Event) -> None:
+        """This routine perform routine tasks."""
+        stop_task = asyncio.create_task(stop_event.wait())
+        wait_task = asyncio.create_task(asyncio.sleep(0))
+        self._logger.debug("Started background worker")
+        while True:
+            done, _ = await asyncio.wait(
+                [stop_task, wait_task],
+                return_when=asyncio.FIRST_COMPLETED
+            )
+            if stop_task in done:
+                wait_task.cancel()
+                break
+            wait_task = asyncio.create_task(
+                asyncio.sleep(self.BACKGROUND_ROUTINE_PERIOD)
+            )
+            # delete old bans
+            try:
+                if self._connection is not None:
+                    statement = "DELETE FROM bans WHERE expires_at <= ? RETURNING ip"
+                    async with self._connection.execute(statement, (time.time(),)) as cursor:
+                        async for row in cursor:
+                            ip = row["ip"]
+                            if ip in self._bans:
+                                del self._bans[ip]
+                            self._logger.debug(f"IP {ip} is not banned anymore")
+                    await self._connection.commit()
+                self._logger.debug("Performed banned IPs cleanup")
+            except:
+                self._logger.error(traceback.format_exc())
+        self._logger.debug("Stopped background worker")
+
+    async def _load_bans_from_database(self) -> None:
+        """Loads bans information from database. Must be called when database connection is established."""
+        if self._connection is None:
+            raise RuntimeError("Not connected to the database")
+        try:
+            data = await self._select_by_anded_kwargs(self._connection, "bans")
+            for d in data:
+                self._bans[d["ip"]] = {
+                    "expires_at": d["expires_at"],
+                    "reason": d["reason"]
+                }
+        except:
+            self._logger.error(traceback.format_exc())
+            return
+
     #
     # PUBLIC
     #
     def __init__(self, path: Path):
         self._path = path
         self._logger = logging.getLogger("database")
+        self._logger.setLevel(logging.DEBUG)
 
         self._connection: Connection | None = None
+        self._background_stop_event: asyncio.Event | None = None
+        self._background_task: asyncio.Task | None = None
+        self._bans: dict = {}
 
     async def connect(self) -> bool:
         """Connect to the database. Returns False on failure."""
@@ -87,6 +136,12 @@ class Database:
             self._connection = await aiosqlite.connect(self._path)
             self._connection.row_factory = Row
             await self._setup_tables(self._connection)
+            self._bans = {}
+            await self._load_bans_from_database()
+            self._background_stop_event = asyncio.Event()
+            self._background_task = asyncio.create_task(
+                self._background_routine(self._background_stop_event)
+            )
         except:
             self._logger.error(traceback.format_exc())
             return False
@@ -96,8 +151,13 @@ class Database:
         """Disconnect from the database."""
         if self._connection is None:
             return
+        self._background_stop_event.set() # type: ignore
+        await self._background_task # type: ignore
+
         await self._connection.close()
         self._connection = None
+        self._background_task = None
+        self._background_stop_event = None
 
 
     async def room_create(self, matrix_id: str) -> ObjectRoom | None:
@@ -245,5 +305,45 @@ class Database:
             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())
+
+
+    async def ban_create(self, ip: str, expires_at: float, reason: str) -> None:
+        """Save information about banned IP address. Replaces existing IPs."""
+        if self._connection is None:
+            raise RuntimeError("Not connected to the database")
+        try:
+            statement = "INSERT OR REPLACE INTO bans (ip, expires_at, reason) VALUES (?, ?, ?)"
+            await self._connection.execute(statement, (ip, expires_at, reason))
+            await self._connection.commit()
+            self._bans[ip] = {
+                "expires_at": expires_at,
+                "reason": reason
+            }
+        except:
+            self._logger.error(traceback.format_exc())
+
+    def ban_get_all(self) -> list[dict]:
+        """Get information about banned IP addresses."""
+        if self._connection is None:
+            raise RuntimeError("Not connected to the database")
+        return [{"ip": k, **self._bans[k]} for k in self._bans]
+
+    def ban_check(self, ip: str) -> bool:
+        if self._connection is None:
+            raise RuntimeError("Not connected to the database")
+        return ip in self._bans
+
+    async def ban_delete(self, ip: str) -> None:
+        """Unbans specified IP address."""
+        if self._connection is None:
+            raise RuntimeError("Not connected to the database")
+        try:
+            statement = "DELETE FROM bans WHERE ip = ?"
+            await self._connection.execute(statement, (ip,))
+            await self._connection.commit()
+            if ip in self._bans:
+                del self._bans[ip]
         except:
             self._logger.error(traceback.format_exc())