Initial commit: MIB bot

2026-07-07 17:55:52 +07:00
commit a355dccd45
17 changed files with 8988 additions and 0 deletions
@@ -0,0 +1,222 @@
+import logging
+import sqlite3
+from contextlib import closing
+
+from config import DB_PATH
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("database")
+
+
+def create_connection() -> sqlite3.Connection:
+    conn = sqlite3.connect(DB_PATH)
+    conn.row_factory = sqlite3.Row
+    return conn
+
+
+def create_tables() -> None:
+    with closing(create_connection()) as conn:
+        cursor = conn.cursor()
+
+        cursor.execute(
+            '''
+            CREATE TABLE IF NOT EXISTS file_ids (
+                file_key TEXT PRIMARY KEY,
+                file_id TEXT
+            )
+            '''
+        )
+
+        cursor.execute(
+            '''
+            CREATE TABLE IF NOT EXISTS user_data (
+                user_id INTEGER,
+                chat_id INTEGER,
+                count INTEGER NOT NULL DEFAULT 0,
+                start_time REAL NOT NULL DEFAULT 0,
+                ban_trigger INTEGER NOT NULL DEFAULT 0,
+                ban_time REAL NOT NULL DEFAULT 0,
+                username TEXT,
+                PRIMARY KEY (user_id, chat_id)
+            )
+            '''
+        )
+
+        cursor.execute(
+            '''
+            CREATE TABLE IF NOT EXISTS user_data_ban (
+                user_id INTEGER,
+                chat_id INTEGER,
+                command TEXT,
+                count INTEGER NOT NULL DEFAULT 0,
+                ban_time_inf REAL NOT NULL DEFAULT 0,
+                PRIMARY KEY (user_id, chat_id, command)
+            )
+            '''
+        )
+
+        conn.commit()
+
+
+def get_file_id(file_key: str) -> str | None:
+    try:
+        with closing(create_connection()) as conn:
+            cursor = conn.cursor()
+            cursor.execute("SELECT file_id FROM file_ids WHERE file_key = ?", (file_key,))
+            row = cursor.fetchone()
+            return row["file_id"] if row and row["file_id"] else None
+    except Exception as exc:
+        logger.error("Ошибка при получении file_id: %s", exc)
+        return None
+
+
+def save_file_id(file_key: str, file_id: str | None) -> None:
+    try:
+        with closing(create_connection()) as conn:
+            cursor = conn.cursor()
+            if file_id is None:
+                cursor.execute("DELETE FROM file_ids WHERE file_key = ?", (file_key,))
+            else:
+                cursor.execute(
+                    "REPLACE INTO file_ids (file_key, file_id) VALUES (?, ?)",
+                    (file_key, file_id),
+                )
+            conn.commit()
+    except Exception as exc:
+        logger.error("Ошибка при сохранении file_id: %s", exc)
+
+
+def get_user_data(user_id: int, chat_id: int):
+    try:
+        with closing(create_connection()) as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                "SELECT * FROM user_data WHERE user_id = ? AND chat_id = ?",
+                (user_id, chat_id),
+            )
+            row = cursor.fetchone()
+            return tuple(row) if row else None
+    except Exception as exc:
+        logger.error("Ошибка при получении данных пользователя: %s", exc)
+        return None
+
+
+def save_user_data(
+    user_id: int,
+    chat_id: int,
+    count: int | None = None,
+    start_time: float | None = None,
+    ban_trigger: bool | None = None,
+    ban_time: float | None = None,
+    username: str | None = None,
+) -> None:
+    try:
+        with closing(create_connection()) as conn:
+            cursor = conn.cursor()
+            existing = get_user_data(user_id, chat_id)
+
+            if existing is None:
+                cursor.execute(
+                    '''
+                    INSERT INTO user_data (user_id, chat_id, count, start_time, ban_trigger, ban_time, username)
+                    VALUES (?, ?, ?, ?, ?, ?, ?)
+                    ''',
+                    (
+                        user_id,
+                        chat_id,
+                        count or 0,
+                        start_time or 0,
+                        int(bool(ban_trigger)) if ban_trigger is not None else 0,
+                        ban_time or 0,
+                        username,
+                    ),
+                )
+            else:
+                updates = []
+                params = []
+                if count is not None:
+                    updates.append("count = ?")
+                    params.append(count)
+                if start_time is not None:
+                    updates.append("start_time = ?")
+                    params.append(start_time)
+                if ban_trigger is not None:
+                    updates.append("ban_trigger = ?")
+                    params.append(int(bool(ban_trigger)))
+                if ban_time is not None:
+                    updates.append("ban_time = ?")
+                    params.append(ban_time)
+                if username is not None:
+                    updates.append("username = ?")
+                    params.append(username)
+
+                if updates:
+                    params.extend([user_id, chat_id])
+                    cursor.execute(
+                        f"UPDATE user_data SET {', '.join(updates)} WHERE user_id = ? AND chat_id = ?",
+                        params,
+                    )
+
+            conn.commit()
+    except Exception as exc:
+        logger.error("Ошибка при сохранении данных пользователя: %s", exc)
+
+
+def get_user_data_ban(user_id: int, chat_id: int, command: str):
+    try:
+        with closing(create_connection()) as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                "SELECT * FROM user_data_ban WHERE user_id = ? AND chat_id = ? AND command = ?",
+                (user_id, chat_id, command),
+            )
+            row = cursor.fetchone()
+            return tuple(row) if row else None
+    except Exception as exc:
+        logger.error("Ошибка при получении данных о блокировке пользователя: %s", exc)
+        return None
+
+
+def save_user_data_ban(
+    user_id: int,
+    chat_id: int,
+    command: str,
+    count: int | None = None,
+    ban_time_inf: float | None = None,
+) -> None:
+    try:
+        with closing(create_connection()) as conn:
+            cursor = conn.cursor()
+            existing = get_user_data_ban(user_id, chat_id, command)
+
+            if existing is None:
+                cursor.execute(
+                    '''
+                    INSERT INTO user_data_ban (user_id, chat_id, command, count, ban_time_inf)
+                    VALUES (?, ?, ?, ?, ?)
+                    ''',
+                    (user_id, chat_id, command, count or 0, ban_time_inf or 0),
+                )
+            else:
+                updates = []
+                params = []
+                if count is not None:
+                    updates.append("count = ?")
+                    params.append(count)
+                if ban_time_inf is not None:
+                    updates.append("ban_time_inf = ?")
+                    params.append(ban_time_inf)
+
+                if updates:
+                    params.extend([user_id, chat_id, command])
+                    cursor.execute(
+                        f"UPDATE user_data_ban SET {', '.join(updates)} WHERE user_id = ? AND chat_id = ? AND command = ?",
+                        params,
+                    )
+
+            conn.commit()
+    except Exception as exc:
+        logger.error("Ошибка при сохранении данных о блокировке пользователя: %s", exc)
+
+
+create_tables()