database.py - master - Befun/FamileControll

database.py
576 строк25 KB
BeFun
Fixed bugs
27 фев 2026, 21:43
27 фев 2026, 21:43e6d5448
"""
База данных для хранения пользователей, правил и истории
"""
import sqlite3
import json
from datetime import datetime, timedelta
from pathlib import Path
from typing import Optional, List, Dict, Any
from contextlib import contextmanager

from config import DATABASE_FILE


class Database:
    """Класс для работы с базой данных SQLite"""

    def __init__(self, db_path: Path = DATABASE_FILE):
        self.db_path = db_path
        self._init_database()

    @contextmanager
    def get_connection(self):
        """Контекстный менеджер для соединения с БД"""
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        try:
            yield conn
            conn.commit()
        except Exception:
            conn.rollback()
            raise
        finally:
            conn.close()

    def _init_database(self):
        """Инициализация таблиц базы данных"""
        with self.get_connection() as conn:
            cursor = conn.cursor()

            # Таблица пользователей
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS users (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    telegram_id INTEGER UNIQUE NOT NULL,
                    username TEXT,
                    first_name TEXT,
                    last_name TEXT,
                    role TEXT DEFAULT 'family',
                    is_active BOOLEAN DEFAULT 1,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    last_activity TIMESTAMP
                )
            ''')

            # Таблица заблокированных приложений
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS blocked_apps (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    name TEXT NOT NULL,
                    path TEXT,
                    app_type TEXT DEFAULT 'app',
                    is_blocked BOOLEAN DEFAULT 1,
                    blocked_by INTEGER,
                    blocked_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    expires_at TIMESTAMP,
                    reason TEXT,
                    FOREIGN KEY (blocked_by) REFERENCES users(telegram_id)
                )
            ''')

            # Таблица разрешённых приложений
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS allowed_apps (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    name TEXT NOT NULL,
                    path TEXT,
                    allowed_by INTEGER,
                    allowed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    expires_at TIMESTAMP,
                    reason TEXT,
                    FOREIGN KEY (allowed_by) REFERENCES users(telegram_id)
                )
            ''')

            # Таблица истории действий
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS activity_history (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    telegram_id INTEGER,
                    action TEXT NOT NULL,
                    target TEXT,
                    details TEXT,
                    timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    FOREIGN KEY (telegram_id) REFERENCES users(telegram_id)
                )
            ''')

            # Таблица расписания сна
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS sleep_schedule (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    start_time TEXT NOT NULL,
                    end_time TEXT NOT NULL,
                    is_active BOOLEAN DEFAULT 1,
                    created_by INTEGER,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    FOREIGN KEY (created_by) REFERENCES users(telegram_id)
                )
            ''')

            # Таблица обнаруженных приложений
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS discovered_apps (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    name TEXT NOT NULL,
                    path TEXT UNIQUE,
                    app_type TEXT DEFAULT 'app',
                    category TEXT,
                    discovered_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    last_seen TIMESTAMP
                )
            ''')

            # Таблица активных сессий приложений
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS app_sessions (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    app_name TEXT NOT NULL,
                    app_path TEXT,
                    started_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    ended_at TIMESTAMP,
                    was_blocked BOOLEAN DEFAULT 0
                )
            ''')

            # Таблица уведомлений
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS notifications (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    telegram_id INTEGER,
                    message TEXT NOT NULL,
                    notification_type TEXT,
                    is_sent BOOLEAN DEFAULT 0,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    sent_at TIMESTAMP,
                    FOREIGN KEY (telegram_id) REFERENCES users(telegram_id)
                )
            ''')

            # Таблица настроек скачивания
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS download_rules (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    extension TEXT NOT NULL,
                    is_blocked BOOLEAN DEFAULT 1,
                    created_by INTEGER,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    FOREIGN KEY (created_by) REFERENCES users(telegram_id)
                )
            ''')

            # Индексы для ускорения запросов
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_users_telegram ON users(telegram_id)')
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_history_timestamp ON activity_history(timestamp)')
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_blocked_apps_name ON blocked_apps(name)')

    # ==================== Методы для работы с пользователями ====================

    def add_user(self, telegram_id: int, username: str = None, 
                 first_name: str = None, last_name: str = None, 
                 role: str = 'family') -> bool:
        """Добавление нового пользователя"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            try:
                cursor.execute('''
                    INSERT OR REPLACE INTO users 
                    (telegram_id, username, first_name, last_name, role, last_activity)
                    VALUES (?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
                ''', (telegram_id, username, first_name, last_name, role))
                return True
            except sqlite3.IntegrityError:
                return False

    def get_user(self, telegram_id: int) -> Optional[Dict]:
        """Получение пользователя по Telegram ID"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('SELECT * FROM users WHERE telegram_id = ?', (telegram_id,))
            row = cursor.fetchone()
            return dict(row) if row else None

    def get_all_users(self) -> List[Dict]:
        """Получение всех пользователей"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('SELECT * FROM users ORDER BY created_at')
            return [dict(row) for row in cursor.fetchall()]

    def get_admin(self) -> Optional[Dict]:
        """Получение администратора"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT * FROM users WHERE role = 'admin' LIMIT 1")
            row = cursor.fetchone()
            return dict(row) if row else None

    def set_admin(self, telegram_id: int) -> bool:
        """Установка администратора (только один)"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            # Сначала убираем права админа у всех
            cursor.execute("UPDATE users SET role = 'family' WHERE role = 'admin'")
            # Затем назначаем нового админа
            cursor.execute('''
                UPDATE users SET role = 'admin' WHERE telegram_id = ?
            ''', (telegram_id,))
            return cursor.rowcount > 0

    def update_user_activity(self, telegram_id: int):
        """Обновление времени последней активности пользователя"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                UPDATE users SET last_activity = CURRENT_TIMESTAMP 
                WHERE telegram_id = ?
            ''', (telegram_id,))

    def is_admin(self, telegram_id: int) -> bool:
        """Проверка, является ли пользователь администратором"""
        user = self.get_user(telegram_id)
        return user and user['role'] == 'admin'

    def get_family_members(self) -> List[Dict]:
        """Получение всех членов семьи (не админов)"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT * FROM users WHERE role = 'family' AND is_active = 1")
            return [dict(row) for row in cursor.fetchall()]

    def remove_user(self, telegram_id: int) -> bool:
        """Удаление пользователя (нельзя удалить админа)"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                DELETE FROM users WHERE telegram_id = ? AND role != 'admin'
            ''', (telegram_id,))
            return cursor.rowcount > 0

    # ==================== Методы для работы с блокировками ====================

    def block_app(self, name: str, blocked_by: int, path: str = None,
                  expires_at: datetime = None, reason: str = None,
                  app_type: str = 'app') -> bool:
        """Блокировка приложения"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            # Сначала удаляем из разрешённых, если есть
            cursor.execute('DELETE FROM allowed_apps WHERE name = ? OR path = ?', (name, path))
            # Затем добавляем в заблокированные
            cursor.execute('''
                INSERT OR REPLACE INTO blocked_apps 
                (name, path, app_type, is_blocked, blocked_by, expires_at, reason)
                VALUES (?, ?, ?, 1, ?, ?, ?)
            ''', (name, path, app_type, blocked_by, expires_at, reason))
            return True

    def unblock_app(self, name: str = None, path: str = None) -> bool:
        """Разблокировка приложения"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            if path:
                cursor.execute('DELETE FROM blocked_apps WHERE path = ?', (path,))
            elif name:
                cursor.execute('DELETE FROM blocked_apps WHERE name = ?', (name,))
            return cursor.rowcount > 0

    def allow_app(self, name: str, allowed_by: int, path: str = None,
                  expires_at: datetime = None, reason: str = None) -> bool:
        """Разрешение приложения (приоритет над блокировкой)"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            # Удаляем из заблокированных
            cursor.execute('DELETE FROM blocked_apps WHERE name = ? OR path = ?', (name, path))
            # Добавляем в разрешённые
            cursor.execute('''
                INSERT OR REPLACE INTO allowed_apps 
                (name, path, allowed_by, expires_at, reason)
                VALUES (?, ?, ?, ?, ?)
            ''', (name, path, allowed_by, expires_at, reason))
            return True

    def disallow_app(self, name: str = None, path: str = None) -> bool:
        """Отмена разрешения приложения"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            if path:
                cursor.execute('DELETE FROM allowed_apps WHERE path = ?', (path,))
            elif name:
                cursor.execute('DELETE FROM allowed_apps WHERE name = ?', (name,))
            return cursor.rowcount > 0

    def is_app_blocked(self, name: str = None, path: str = None) -> bool:
        """Проверка, заблокировано ли приложение"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            
            # Сначала проверяем, есть ли в разрешённых (приоритет!)
            if name:
                cursor.execute('''
                    SELECT * FROM allowed_apps 
                    WHERE name = ? 
                    AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP)
                ''', (name,))
                if cursor.fetchone():
                    return False  # Разрешено явно по имени
            
            if path:
                cursor.execute('''
                    SELECT * FROM allowed_apps 
                    WHERE path = ?
                    AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP)
                ''', (path,))
                if cursor.fetchone():
                    return False  # Разрешено явно по пути
            
            # Проверяем, есть ли в заблокированных
            if path:
                cursor.execute('''
                    SELECT * FROM blocked_apps 
                    WHERE (path = ? OR name = ?) AND is_blocked = 1
                    AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP)
                ''', (path, name))
                if cursor.fetchone():
                    return True
            elif name:
                cursor.execute('''
                    SELECT * FROM blocked_apps 
                    WHERE name = ? AND is_blocked = 1
                    AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP)
                ''', (name,))
                if cursor.fetchone():
                    return True
            
            return False  # По умолчанию не заблокировано

    def get_blocked_apps(self) -> List[Dict]:
        """Получение списка заблокированных приложений"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                SELECT * FROM blocked_apps 
                WHERE is_blocked = 1 
                AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP)
                ORDER BY blocked_at DESC
            ''')
            return [dict(row) for row in cursor.fetchall()]

    def get_allowed_apps(self) -> List[Dict]:
        """Получение списка разрешённых приложений"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                SELECT * FROM allowed_apps 
                WHERE expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP
                ORDER BY allowed_at DESC
            ''')
            return [dict(row) for row in cursor.fetchall()]

    def cleanup_expired_rules(self):
        """Очистка истекших правил"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('DELETE FROM blocked_apps WHERE expires_at < CURRENT_TIMESTAMP')
            cursor.execute('DELETE FROM allowed_apps WHERE expires_at < CURRENT_TIMESTAMP')

    # ==================== Методы для работы с историей ====================

    def log_activity(self, telegram_id: int, action: str, 
                     target: str = None, details: str = None):
        """Запись действия в историю"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                INSERT INTO activity_history (telegram_id, action, target, details)
                VALUES (?, ?, ?, ?)
            ''', (telegram_id, action, target, details))

    def get_history(self, limit: int = 100, 
                    telegram_id: int = None,
                    action_type: str = None) -> List[Dict]:
        """Получение истории действий"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            
            query = '''
                SELECT h.*, u.username, u.first_name 
                FROM activity_history h
                LEFT JOIN users u ON h.telegram_id = u.telegram_id
                WHERE 1=1
            '''
            params = []
            
            if telegram_id:
                query += ' AND h.telegram_id = ?'
                params.append(telegram_id)
            
            if action_type:
                query += ' AND h.action = ?'
                params.append(action_type)
            
            query += ' ORDER BY h.timestamp DESC LIMIT ?'
            params.append(limit)
            
            cursor.execute(query, params)
            return [dict(row) for row in cursor.fetchall()]

    # ==================== Методы для работы с режимом сна ====================

    def set_sleep_schedule(self, start_time: str, end_time: str, 
                           created_by: int) -> bool:
        """Установка расписания сна"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            # Деактивируем старые расписания
            cursor.execute('UPDATE sleep_schedule SET is_active = 0')
            # Добавляем новое
            cursor.execute('''
                INSERT INTO sleep_schedule (start_time, end_time, is_active, created_by)
                VALUES (?, ?, 1, ?)
            ''', (start_time, end_time, created_by))
            return True

    def get_sleep_schedule(self) -> Optional[Dict]:
        """Получение активного расписания сна"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                SELECT * FROM sleep_schedule WHERE is_active = 1 LIMIT 1
            ''')
            row = cursor.fetchone()
            return dict(row) if row else None

    def disable_sleep_schedule(self):
        """Отключение расписания сна"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('UPDATE sleep_schedule SET is_active = 0')

    # ==================== Методы для работы с обнаруженными приложениями ====================

    def add_discovered_app(self, name: str, path: str, 
                           app_type: str = 'app', category: str = None):
        """Добавление обнаруженного приложения"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                INSERT OR REPLACE INTO discovered_apps 
                (name, path, app_type, category, last_seen)
                VALUES (?, ?, ?, ?, CURRENT_TIMESTAMP)
            ''', (name, path, app_type, category))

    def get_discovered_apps(self, app_type: str = None) -> List[Dict]:
        """Получение списка обнаруженных приложений"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            if app_type:
                cursor.execute('''
                    SELECT * FROM discovered_apps WHERE app_type = ?
                    ORDER BY discovered_at DESC
                ''', (app_type,))
            else:
                cursor.execute('''
                    SELECT * FROM discovered_apps ORDER BY discovered_at DESC
                ''')
            return [dict(row) for row in cursor.fetchall()]

    # ==================== Методы для работы с уведомлениями ====================

    def add_notification(self, telegram_id: int, message: str, 
                         notification_type: str = 'info'):
        """Добавление уведомления"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                INSERT INTO notifications (telegram_id, message, notification_type)
                VALUES (?, ?, ?)
            ''', (telegram_id, message, notification_type))

    def get_pending_notifications(self) -> List[Dict]:
        """Получение неотправленных уведомлений"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                SELECT * FROM notifications WHERE is_sent = 0
            ''')
            return [dict(row) for row in cursor.fetchall()]

    def mark_notification_sent(self, notification_id: int):
        """Отметка уведомления как отправленного"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                UPDATE notifications 
                SET is_sent = 1, sent_at = CURRENT_TIMESTAMP 
                WHERE id = ?
            ''', (notification_id,))

    # ==================== Методы для работы с правилами скачивания ====================

    def block_extension(self, extension: str, blocked_by: int):
        """Блокировка расширения файла для скачивания"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                INSERT OR REPLACE INTO download_rules (extension, is_blocked, created_by)
                VALUES (?, 1, ?)
            ''', (extension, blocked_by))

    def allow_extension(self, extension: str):
        """Разрешение расширения файла"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                DELETE FROM download_rules WHERE extension = ?
            ''', (extension,))

    def get_blocked_extensions(self) -> List[str]:
        """Получение списка заблокированных расширений"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                SELECT extension FROM download_rules WHERE is_blocked = 1
            ''')
            return [row['extension'] for row in cursor.fetchall()]

    def is_extension_blocked(self, extension: str) -> bool:
        """Проверка, заблокировано ли расширение"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                SELECT * FROM download_rules 
                WHERE extension = ? AND is_blocked = 1
            ''', (extension.lower(),))
            return cursor.fetchone() is not None

    # ==================== Методы для работы с сессиями приложений ====================

    def start_app_session(self, app_name: str, app_path: str = None) -> int:
        """Начало сессии приложения"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                INSERT INTO app_sessions (app_name, app_path)
                VALUES (?, ?)
            ''', (app_name, app_path))
            return cursor.lastrowid

    def end_app_session(self, session_id: int, was_blocked: bool = False):
        """Завершение сессии приложения"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                UPDATE app_sessions 
                SET ended_at = CURRENT_TIMESTAMP, was_blocked = ?
                WHERE id = ?
            ''', (was_blocked, session_id))

    def get_active_sessions(self) -> List[Dict]:
        """Получение активных сессий"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                SELECT * FROM app_sessions WHERE ended_at IS NULL
            ''')
            return [dict(row) for row in cursor.fetchall()]