telegram-antispam-bot/bot/db.py

from __future__ import annotations

from dataclasses import dataclass
import time

import aiosqlite


@dataclass(frozen=True)
class PendingVerification:
    chat_id: int
    user_id: int
    challenge_message_id: int
    expires_at: int


class Database:
    def __init__(self, path: str) -> None:
        self._path = path
        self._conn: aiosqlite.Connection | None = None

    async def connect(self) -> None:
        self._conn = await aiosqlite.connect(self._path)
        await self._conn.execute("PRAGMA journal_mode=WAL;")
        await self._conn.execute("PRAGMA foreign_keys=ON;")
        await self._conn.execute(
            """
            CREATE TABLE IF NOT EXISTS verified_members (
                chat_id INTEGER NOT NULL,
                user_id INTEGER NOT NULL,
                verified_at INTEGER NOT NULL,
                PRIMARY KEY (chat_id, user_id)
            );
            """
        )
        await self._conn.execute(
            """
            CREATE TABLE IF NOT EXISTS pending_verifications (
                chat_id INTEGER NOT NULL,
                user_id INTEGER NOT NULL,
                challenge_message_id INTEGER NOT NULL,
                expires_at INTEGER NOT NULL,
                PRIMARY KEY (chat_id, user_id)
            );
            """
        )
        await self._conn.commit()

    async def close(self) -> None:
        if self._conn is not None:
            await self._conn.close()
            self._conn = None

    @property
    def conn(self) -> aiosqlite.Connection:
        if self._conn is None:
            raise RuntimeError("Database is not connected")
        return self._conn

    async def mark_pending(
        self,
        *,
        chat_id: int,
        user_id: int,
        challenge_message_id: int,
        expires_at: int,
    ) -> None:
        await self.conn.execute(
            """
            INSERT INTO pending_verifications (chat_id, user_id, challenge_message_id, expires_at)
            VALUES (?, ?, ?, ?)
            ON CONFLICT(chat_id, user_id) DO UPDATE SET
                challenge_message_id=excluded.challenge_message_id,
                expires_at=excluded.expires_at
            """,
            (chat_id, user_id, challenge_message_id, expires_at),
        )
        await self.conn.commit()

    async def get_pending(self, *, chat_id: int, user_id: int) -> PendingVerification | None:
        cur = await self.conn.execute(
            """
            SELECT chat_id, user_id, challenge_message_id, expires_at
            FROM pending_verifications
            WHERE chat_id=? AND user_id=?
            """,
            (chat_id, user_id),
        )
        row = await cur.fetchone()
        await cur.close()
        if not row:
            return None
        return PendingVerification(
            chat_id=int(row[0]),
            user_id=int(row[1]),
            challenge_message_id=int(row[2]),
            expires_at=int(row[3]),
        )

    async def delete_pending(self, *, chat_id: int, user_id: int) -> None:
        await self.conn.execute(
            "DELETE FROM pending_verifications WHERE chat_id=? AND user_id=?",
            (chat_id, user_id),
        )
        await self.conn.commit()

    async def mark_verified(self, *, chat_id: int, user_id: int) -> None:
        now = int(time.time())
        await self.conn.execute(
            """
            INSERT INTO verified_members (chat_id, user_id, verified_at)
            VALUES (?, ?, ?)
            ON CONFLICT(chat_id, user_id) DO UPDATE SET verified_at=excluded.verified_at
            """,
            (chat_id, user_id, now),
        )
        await self.conn.execute(
            "DELETE FROM pending_verifications WHERE chat_id=? AND user_id=?",
            (chat_id, user_id),
        )
        await self.conn.commit()

    async def delete_verified(self, *, chat_id: int, user_id: int) -> None:
        await self.conn.execute(
            "DELETE FROM verified_members WHERE chat_id=? AND user_id=?",
            (chat_id, user_id),
        )
        await self.conn.commit()

    async def is_verified(self, *, chat_id: int, user_id: int) -> bool:
        cur = await self.conn.execute(
            """
            SELECT 1
            FROM verified_members
            WHERE chat_id=? AND user_id=?
            """,
            (chat_id, user_id),
        )
        row = await cur.fetchone()
        await cur.close()
        return row is not None

    async def get_expired_pending(self, *, now: int) -> list[PendingVerification]:
        cur = await self.conn.execute(
            """
            SELECT chat_id, user_id, challenge_message_id, expires_at
            FROM pending_verifications
            WHERE expires_at <= ?
            ORDER BY expires_at ASC
            LIMIT 100
            """,
            (now,),
        )
        rows = await cur.fetchall()
        await cur.close()
        return [
            PendingVerification(
                chat_id=int(r[0]),
                user_id=int(r[1]),
                challenge_message_id=int(r[2]),
                expires_at=int(r[3]),
            )
            for r in rows
        ]