"""
Storage module для PyGuardian
Управление SQLite базой данных для хранения IP-адресов, попыток атак и банов
"""

import asyncio
import sqlite3
import aiosqlite
import ipaddress
import json
from datetime import datetime, timedelta
from typing import List, Dict, Optional, Tuple, Any
import logging

logger = logging.getLogger(__name__)


class Storage:
    """Асинхронный класс для работы с SQLite базой данных"""
    
    def __init__(self, db_path: str):
        self.db_path = db_path
        self._connection = None
    
    async def init_database(self) -> None:
        """Инициализация базы данных и создание таблиц"""
        async with aiosqlite.connect(self.db_path) as db:
            # Таблица для хранения попыток атак
            await db.execute("""
                CREATE TABLE IF NOT EXISTS attack_attempts (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    ip_address TEXT NOT NULL,
                    username TEXT,
                    attack_type TEXT NOT NULL,
                    timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
                    log_line TEXT,
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                    INDEX(ip_address),
                    INDEX(timestamp)
                )
            """)
            
            # Таблица для хранения забаненных IP
            await db.execute("""
                CREATE TABLE IF NOT EXISTS banned_ips (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    ip_address TEXT UNIQUE NOT NULL,
                    ban_reason TEXT NOT NULL,
                    banned_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                    unban_at DATETIME,
                    is_active BOOLEAN DEFAULT 1,
                    manual_ban BOOLEAN DEFAULT 0,
                    attempts_count INTEGER DEFAULT 0,
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                    INDEX(ip_address),
                    INDEX(unban_at),
                    INDEX(is_active)
                )
            """)
            
            # Таблица для успешных входов
            await db.execute("""
                CREATE TABLE IF NOT EXISTS successful_logins (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    ip_address TEXT NOT NULL,
                    username TEXT NOT NULL,
                    timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
                    session_info TEXT,
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                    INDEX(ip_address),
                    INDEX(username),
                    INDEX(timestamp)
                )
            """)
            
            # Таблица для статистики
            await db.execute("""
                CREATE TABLE IF NOT EXISTS daily_stats (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    date DATE UNIQUE NOT NULL,
                    total_attempts INTEGER DEFAULT 0,
                    unique_ips INTEGER DEFAULT 0,
                    banned_count INTEGER DEFAULT 0,
                    successful_logins INTEGER DEFAULT 0,
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                    INDEX(date)
                )
            """)
            
            # Таблица для компрометаций
            await db.execute("""
                CREATE TABLE IF NOT EXISTS compromises (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    ip_address TEXT NOT NULL,
                    username TEXT NOT NULL,
                    detection_time DATETIME NOT NULL,
                    session_active BOOLEAN DEFAULT 1,
                    new_password TEXT,
                    session_info TEXT,
                    resolved BOOLEAN DEFAULT 0,
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                    INDEX(ip_address),
                    INDEX(username),
                    INDEX(detection_time)
                )
            """)
            
            # Таблица для агентов кластера
            await db.execute("""
                CREATE TABLE IF NOT EXISTS agents (
                    agent_id TEXT PRIMARY KEY,
                    hostname TEXT NOT NULL,
                    ip_address TEXT NOT NULL,
                    ssh_port INTEGER DEFAULT 22,
                    ssh_user TEXT DEFAULT 'root',
                    status TEXT DEFAULT 'added',
                    added_time DATETIME NOT NULL,
                    last_check DATETIME,
                    version TEXT,
                    config TEXT,
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                    INDEX(hostname),
                    INDEX(ip_address),
                    INDEX(status)
                )
            """)
            
            # Таблица для аутентификационных данных агентов
            await db.execute("""
                CREATE TABLE IF NOT EXISTS agent_auth (
                    agent_id TEXT PRIMARY KEY,
                    secret_key_hash TEXT NOT NULL,
                    salt TEXT NOT NULL,
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                    last_authenticated DATETIME,
                    auth_count INTEGER DEFAULT 0,
                    is_active BOOLEAN DEFAULT 1,
                    FOREIGN KEY(agent_id) REFERENCES agents(agent_id) ON DELETE CASCADE,
                    INDEX(agent_id),
                    INDEX(is_active)
                )
            """)
            
            # Таблица для активных токенов агентов
            await db.execute("""
                CREATE TABLE IF NOT EXISTS agent_tokens (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    agent_id TEXT NOT NULL,
                    token_hash TEXT NOT NULL,
                    token_type TEXT NOT NULL, -- 'access' или 'refresh'
                    expires_at DATETIME NOT NULL,
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                    last_used DATETIME,
                    is_revoked BOOLEAN DEFAULT 0,
                    FOREIGN KEY(agent_id) REFERENCES agents(agent_id) ON DELETE CASCADE,
                    INDEX(agent_id),
                    INDEX(token_hash),
                    INDEX(expires_at),
                    INDEX(is_revoked)
                )
            """)
            
            # Таблица для активных сессий агентов
            await db.execute("""
                CREATE TABLE IF NOT EXISTS agent_sessions (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    agent_id TEXT NOT NULL,
                    session_id TEXT UNIQUE NOT NULL,
                    ip_address TEXT NOT NULL,
                    user_agent TEXT,
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                    last_activity DATETIME DEFAULT CURRENT_TIMESTAMP,
                    expires_at DATETIME NOT NULL,
                    is_active BOOLEAN DEFAULT 1,
                    requests_count INTEGER DEFAULT 0,
                    FOREIGN KEY(agent_id) REFERENCES agents(agent_id) ON DELETE CASCADE,
                    INDEX(agent_id),
                    INDEX(session_id),
                    INDEX(ip_address),
                    INDEX(expires_at),
                    INDEX(is_active)
                )
            """)
            
            # Таблица для логов аутентификации агентов
            await db.execute("""
                CREATE TABLE IF NOT EXISTS agent_auth_logs (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    agent_id TEXT,
                    ip_address TEXT NOT NULL,
                    action TEXT NOT NULL, -- 'login', 'logout', 'token_refresh', 'access_denied'
                    success BOOLEAN NOT NULL,
                    error_message TEXT,
                    user_agent TEXT,
                    timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
                    INDEX(agent_id),
                    INDEX(ip_address),
                    INDEX(action),
                    INDEX(timestamp)
                )
            """)
            
            await db.commit()
            logger.info("База данных инициализирована успешно")
    
    async def create_agent_auth(self, agent_id: str, secret_key_hash: str, salt: str) -> bool:
        """Создать аутентификационные данные для агента"""
        try:
            async with aiosqlite.connect(self.db_path) as db:
                await db.execute("""
                    INSERT OR REPLACE INTO agent_auth 
                    (agent_id, secret_key_hash, salt, created_at)
                    VALUES (?, ?, ?, ?)
                """, (agent_id, secret_key_hash, salt, datetime.now()))
                await db.commit()
                logger.info(f"Created auth data for agent {agent_id}")
                return True
        except Exception as e:
            logger.error(f"Failed to create auth data for agent {agent_id}: {e}")
            return False
    
    async def get_agent_auth(self, agent_id: str) -> Optional[Dict[str, Any]]:
        """Получить аутентификационные данные агента"""
        try:
            async with aiosqlite.connect(self.db_path) as db:
                cursor = await db.execute("""
                    SELECT secret_key_hash, salt, last_authenticated, auth_count, is_active
                    FROM agent_auth WHERE agent_id = ? AND is_active = 1
                """, (agent_id,))
                result = await cursor.fetchone()
                
                if result:
                    return {
                        'secret_key_hash': result[0],
                        'salt': result[1],
                        'last_authenticated': result[2],
                        'auth_count': result[3],
                        'is_active': bool(result[4])
                    }
                return None
        except Exception as e:
            logger.error(f"Failed to get auth data for agent {agent_id}: {e}")
            return None
    
    async def update_agent_last_auth(self, agent_id: str) -> bool:
        """Обновить время последней аутентификации агента"""
        try:
            async with aiosqlite.connect(self.db_path) as db:
                await db.execute("""
                    UPDATE agent_auth 
                    SET last_authenticated = ?, auth_count = auth_count + 1
                    WHERE agent_id = ?
                """, (datetime.now(), agent_id))
                await db.commit()
                return True
        except Exception as e:
            logger.error(f"Failed to update last auth for agent {agent_id}: {e}")
            return False
    
    async def store_agent_token(self, agent_id: str, token_hash: str, 
                               token_type: str, expires_at: datetime) -> bool:
        """Сохранить токен агента"""
        try:
            async with aiosqlite.connect(self.db_path) as db:
                await db.execute("""
                    INSERT INTO agent_tokens 
                    (agent_id, token_hash, token_type, expires_at)
                    VALUES (?, ?, ?, ?)
                """, (agent_id, token_hash, token_type, expires_at))
                await db.commit()
                return True
        except Exception as e:
            logger.error(f"Failed to store token for agent {agent_id}: {e}")
            return False
    
    async def verify_agent_token(self, agent_id: str, token_hash: str) -> bool:
        """Проверить действительность токена агента"""
        try:
            async with aiosqlite.connect(self.db_path) as db:
                cursor = await db.execute("""
                    SELECT id FROM agent_tokens 
                    WHERE agent_id = ? AND token_hash = ? 
                    AND expires_at > ? AND is_revoked = 0
                """, (agent_id, token_hash, datetime.now()))
                result = await cursor.fetchone()
                
                if result:
                    # Обновить время последнего использования
                    await db.execute("""
                        UPDATE agent_tokens SET last_used = ? WHERE id = ?
                    """, (datetime.now(), result[0]))
                    await db.commit()
                    return True
                return False
        except Exception as e:
            logger.error(f"Failed to verify token for agent {agent_id}: {e}")
            return False
    
    async def revoke_agent_tokens(self, agent_id: str, token_type: Optional[str] = None) -> bool:
        """Отозвать токены агента"""
        try:
            async with aiosqlite.connect(self.db_path) as db:
                if token_type:
                    await db.execute("""
                        UPDATE agent_tokens SET is_revoked = 1 
                        WHERE agent_id = ? AND token_type = ?
                    """, (agent_id, token_type))
                else:
                    await db.execute("""
                        UPDATE agent_tokens SET is_revoked = 1 
                        WHERE agent_id = ?
                    """, (agent_id,))
                await db.commit()
                return True
        except Exception as e:
            logger.error(f"Failed to revoke tokens for agent {agent_id}: {e}")
            return False
    
    async def create_agent_session(self, agent_id: str, session_id: str, 
                                  ip_address: str, expires_at: datetime,
                                  user_agent: Optional[str] = None) -> bool:
        """Создать сессию агента"""
        try:
            async with aiosqlite.connect(self.db_path) as db:
                await db.execute("""
                    INSERT INTO agent_sessions 
                    (agent_id, session_id, ip_address, user_agent, expires_at)
                    VALUES (?, ?, ?, ?, ?)
                """, (agent_id, session_id, ip_address, user_agent, expires_at))
                await db.commit()
                return True
        except Exception as e:
            logger.error(f"Failed to create session for agent {agent_id}: {e}")
            return False
    
    async def update_agent_session_activity(self, session_id: str) -> bool:
        """Обновить активность сессии агента"""
        try:
            async with aiosqlite.connect(self.db_path) as db:
                await db.execute("""
                    UPDATE agent_sessions 
                    SET last_activity = ?, requests_count = requests_count + 1
                    WHERE session_id = ? AND is_active = 1
                """, (datetime.now(), session_id))
                await db.commit()
                return True
        except Exception as e:
            logger.error(f"Failed to update session activity {session_id}: {e}")
            return False
    
    async def get_active_agent_sessions(self, agent_id: str) -> List[Dict]:
        """Получить активные сессии агента"""
        try:
            async with aiosqlite.connect(self.db_path) as db:
                cursor = await db.execute("""
                    SELECT session_id, ip_address, user_agent, created_at, 
                           last_activity, requests_count
                    FROM agent_sessions 
                    WHERE agent_id = ? AND is_active = 1 AND expires_at > ?
                """, (agent_id, datetime.now()))
                results = await cursor.fetchall()
                
                sessions = []
                for row in results:
                    sessions.append({
                        'session_id': row[0],
                        'ip_address': row[1],
                        'user_agent': row[2],
                        'created_at': row[3],
                        'last_activity': row[4],
                        'requests_count': row[5]
                    })
                return sessions
        except Exception as e:
            logger.error(f"Failed to get sessions for agent {agent_id}: {e}")
            return []
    
    async def deactivate_agent_session(self, session_id: str) -> bool:
        """Деактивировать сессию агента"""
        try:
            async with aiosqlite.connect(self.db_path) as db:
                await db.execute("""
                    UPDATE agent_sessions SET is_active = 0 WHERE session_id = ?
                """, (session_id,))
                await db.commit()
                return True
        except Exception as e:
            logger.error(f"Failed to deactivate session {session_id}: {e}")
            return False
    
    async def log_agent_auth_event(self, agent_id: str, ip_address: str, 
                                  action: str, success: bool,
                                  error_message: Optional[str] = None,
                                  user_agent: Optional[str] = None) -> bool:
        """Записать событие аутентификации агента"""
        try:
            async with aiosqlite.connect(self.db_path) as db:
                await db.execute("""
                    INSERT INTO agent_auth_logs 
                    (agent_id, ip_address, action, success, error_message, user_agent)
                    VALUES (?, ?, ?, ?, ?, ?)
                """, (agent_id, ip_address, action, success, error_message, user_agent))
                await db.commit()
                return True
        except Exception as e:
            logger.error(f"Failed to log auth event for agent {agent_id}: {e}")
            return False
    
    async def get_agent_auth_logs(self, agent_id: str, limit: int = 100) -> List[Dict]:
        """Получить логи аутентификации агента"""
        try:
            async with aiosqlite.connect(self.db_path) as db:
                cursor = await db.execute("""
                    SELECT ip_address, action, success, error_message, 
                           user_agent, timestamp
                    FROM agent_auth_logs 
                    WHERE agent_id = ?
                    ORDER BY timestamp DESC LIMIT ?
                """, (agent_id, limit))
                results = await cursor.fetchall()
                
                logs = []
                for row in results:
                    logs.append({
                        'ip_address': row[0],
                        'action': row[1],
                        'success': bool(row[2]),
                        'error_message': row[3],
                        'user_agent': row[4],
                        'timestamp': row[5]
                    })
                return logs
        except Exception as e:
            logger.error(f"Failed to get auth logs for agent {agent_id}: {e}")
            return []
    
    async def cleanup_expired_tokens(self) -> int:
        """Очистка истекших токенов"""
        try:
            async with aiosqlite.connect(self.db_path) as db:
                cursor = await db.execute("""
                    DELETE FROM agent_tokens WHERE expires_at < ?
                """, (datetime.now(),))
                await db.commit()
                deleted_count = cursor.rowcount
                logger.info(f"Cleaned up {deleted_count} expired tokens")
                return deleted_count
        except Exception as e:
            logger.error(f"Failed to cleanup expired tokens: {e}")
            return 0
    
    async def cleanup_expired_sessions(self) -> int:
        """Очистка истекших сессий"""
        try:
            async with aiosqlite.connect(self.db_path) as db:
                cursor = await db.execute("""
                    UPDATE agent_sessions SET is_active = 0 
                    WHERE expires_at < ? AND is_active = 1
                """, (datetime.now(),))
                await db.commit()
                cleaned_count = cursor.rowcount
                logger.info(f"Cleaned up {cleaned_count} expired sessions")
                return cleaned_count
        except Exception as e:
            logger.error(f"Failed to cleanup expired sessions: {e}")
            return 0
    
    async def add_attack_attempt(self, ip: str, username: str, 
                               attack_type: str, log_line: str, 
                               timestamp: Optional[datetime] = None) -> None:
        """Добавить попытку атаки в базу данных"""
        if timestamp is None:
            timestamp = datetime.now()
            
        async with aiosqlite.connect(self.db_path) as db:
            await db.execute("""
                INSERT INTO attack_attempts 
                (ip_address, username, attack_type, timestamp, log_line)
                VALUES (?, ?, ?, ?, ?)
            """, (ip, username, attack_type, timestamp, log_line))
            await db.commit()
    
    async def get_attack_count_for_ip(self, ip: str, time_window: int) -> int:
        """Получить количество попыток атак от IP за указанный период"""
        since_time = datetime.now() - timedelta(seconds=time_window)
        
        async with aiosqlite.connect(self.db_path) as db:
            async with db.execute("""
                SELECT COUNT(*) FROM attack_attempts 
                WHERE ip_address = ? AND timestamp > ?
            """, (ip, since_time)) as cursor:
                result = await cursor.fetchone()
                return result[0] if result else 0
    
    async def ban_ip(self, ip: str, reason: str, unban_time: int, 
                    manual: bool = False, attempts_count: int = 0) -> bool:
        """Забанить IP адрес"""
        unban_at = datetime.now() + timedelta(seconds=unban_time)
        
        async with aiosqlite.connect(self.db_path) as db:
            try:
                await db.execute("""
                    INSERT OR REPLACE INTO banned_ips 
                    (ip_address, ban_reason, unban_at, manual_ban, attempts_count)
                    VALUES (?, ?, ?, ?, ?)
                """, (ip, reason, unban_at, manual, attempts_count))
                await db.commit()
                logger.info(f"IP {ip} забанен. Причина: {reason}")
                return True
            except Exception as e:
                logger.error(f"Ошибка при бане IP {ip}: {e}")
                return False
    
    async def unban_ip(self, ip: str) -> bool:
        """Разбанить IP адрес"""
        async with aiosqlite.connect(self.db_path) as db:
            try:
                await db.execute("""
                    UPDATE banned_ips 
                    SET is_active = 0 
                    WHERE ip_address = ? AND is_active = 1
                """, (ip,))
                await db.commit()
                logger.info(f"IP {ip} разбанен")
                return True
            except Exception as e:
                logger.error(f"Ошибка при разбане IP {ip}: {e}")
                return False
    
    async def is_ip_banned(self, ip: str) -> bool:
        """Проверить, забанен ли IP"""
        async with aiosqlite.connect(self.db_path) as db:
            async with db.execute("""
                SELECT id FROM banned_ips 
                WHERE ip_address = ? AND is_active = 1 
                AND (unban_at IS NULL OR unban_at > datetime('now'))
            """, (ip,)) as cursor:
                result = await cursor.fetchone()
                return result is not None
    
    async def get_expired_bans(self) -> List[str]:
        """Получить список IP с истекшим временем бана"""
        async with aiosqlite.connect(self.db_path) as db:
            async with db.execute("""
                SELECT ip_address FROM banned_ips 
                WHERE is_active = 1 AND unban_at <= datetime('now')
                AND manual_ban = 0
            """) as cursor:
                results = await cursor.fetchall()
                return [row[0] for row in results]
    
    async def get_banned_ips(self) -> List[Dict]:
        """Получить список всех забаненных IP"""
        async with aiosqlite.connect(self.db_path) as db:
            async with db.execute("""
                SELECT ip_address, ban_reason, banned_at, unban_at, 
                       manual_ban, attempts_count
                FROM banned_ips 
                WHERE is_active = 1
                ORDER BY banned_at DESC
            """) as cursor:
                results = await cursor.fetchall()
                
                banned_list = []
                for row in results:
                    banned_list.append({
                        'ip': row[0],
                        'reason': row[1],
                        'banned_at': row[2],
                        'unban_at': row[3],
                        'manual': bool(row[4]),
                        'attempts': row[5]
                    })
                return banned_list
    
    async def get_top_attackers(self, limit: int = 10, 
                              days: int = 1) -> List[Dict]:
        """Получить топ атакующих IP за указанный период"""
        since_date = datetime.now() - timedelta(days=days)
        
        async with aiosqlite.connect(self.db_path) as db:
            async with db.execute("""
                SELECT ip_address, COUNT(*) as attempts, 
                       MIN(timestamp) as first_attempt,
                       MAX(timestamp) as last_attempt,
                       GROUP_CONCAT(DISTINCT attack_type) as attack_types
                FROM attack_attempts 
                WHERE timestamp > ?
                GROUP BY ip_address
                ORDER BY attempts DESC
                LIMIT ?
            """, (since_date, limit)) as cursor:
                results = await cursor.fetchall()
                
                attackers = []
                for row in results:
                    attackers.append({
                        'ip': row[0],
                        'attempts': row[1],
                        'first_attempt': row[2],
                        'last_attempt': row[3],
                        'attack_types': row[4].split(',') if row[4] else []
                    })
                return attackers
    
    async def get_ip_details(self, ip: str) -> Dict:
        """Получить детальную информацию по IP"""
        async with aiosqlite.connect(self.db_path) as db:
            # Общая статистика по попыткам
            async with db.execute("""
                SELECT COUNT(*) as total_attempts,
                       MIN(timestamp) as first_seen,
                       MAX(timestamp) as last_seen,
                       GROUP_CONCAT(DISTINCT attack_type) as attack_types,
                       GROUP_CONCAT(DISTINCT username) as usernames
                FROM attack_attempts
                WHERE ip_address = ?
            """, (ip,)) as cursor:
                attack_stats = await cursor.fetchone()
            
            # Информация о бане
            async with db.execute("""
                SELECT ban_reason, banned_at, unban_at, is_active, manual_ban
                FROM banned_ips
                WHERE ip_address = ?
                ORDER BY banned_at DESC
                LIMIT 1
            """, (ip,)) as cursor:
                ban_info = await cursor.fetchone()
            
            # Последние попытки
            async with db.execute("""
                SELECT timestamp, attack_type, username, log_line
                FROM attack_attempts
                WHERE ip_address = ?
                ORDER BY timestamp DESC
                LIMIT 10
            """, (ip,)) as cursor:
                recent_attempts = await cursor.fetchall()
            
            return {
                'ip': ip,
                'total_attempts': attack_stats[0] if attack_stats[0] else 0,
                'first_seen': attack_stats[1],
                'last_seen': attack_stats[2],
                'attack_types': attack_stats[3].split(',') if attack_stats[3] else [],
                'usernames': attack_stats[4].split(',') if attack_stats[4] else [],
                'is_banned': ban_info is not None and ban_info[3] == 1,
                'ban_info': {
                    'reason': ban_info[0] if ban_info else None,
                    'banned_at': ban_info[1] if ban_info else None,
                    'unban_at': ban_info[2] if ban_info else None,
                    'manual': bool(ban_info[4]) if ban_info else False
                } if ban_info else None,
                'recent_attempts': [
                    {
                        'timestamp': attempt[0],
                        'type': attempt[1],
                        'username': attempt[2],
                        'log_line': attempt[3]
                    }
                    for attempt in recent_attempts
                ]
            }
    
    async def add_successful_login(self, ip: str, username: str, 
                                 session_info: Optional[str] = None) -> None:
        """Добавить запись об успешном входе"""
        async with aiosqlite.connect(self.db_path) as db:
            await db.execute("""
                INSERT INTO successful_logins 
                (ip_address, username, session_info)
                VALUES (?, ?, ?)
            """, (ip, username, session_info))
            await db.commit()
    
    async def get_daily_stats(self) -> Dict:
        """Получить статистику за сегодня"""
        today = datetime.now().date()
        yesterday = today - timedelta(days=1)
        
        async with aiosqlite.connect(self.db_path) as db:
            # Атаки за сегодня
            async with db.execute("""
                SELECT COUNT(*) FROM attack_attempts 
                WHERE DATE(timestamp) = ?
            """, (today,)) as cursor:
                today_attacks = (await cursor.fetchone())[0]
            
            # Уникальные IP за сегодня
            async with db.execute("""
                SELECT COUNT(DISTINCT ip_address) FROM attack_attempts 
                WHERE DATE(timestamp) = ?
            """, (today,)) as cursor:
                today_unique_ips = (await cursor.fetchone())[0]
            
            # Активные баны
            async with db.execute("""
                SELECT COUNT(*) FROM banned_ips 
                WHERE is_active = 1
            """) as cursor:
                active_bans = (await cursor.fetchone())[0]
            
            # Успешные входы за сегодня
            async with db.execute("""
                SELECT COUNT(*) FROM successful_logins 
                WHERE DATE(timestamp) = ?
            """, (today,)) as cursor:
                today_logins = (await cursor.fetchone())[0]
            
            # Сравнение с вчера
            async with db.execute("""
                SELECT COUNT(*) FROM attack_attempts 
                WHERE DATE(timestamp) = ?
            """, (yesterday,)) as cursor:
                yesterday_attacks = (await cursor.fetchone())[0]
            
            return {
                'today': {
                    'attacks': today_attacks,
                    'unique_ips': today_unique_ips,
                    'successful_logins': today_logins
                },
                'yesterday': {
                    'attacks': yesterday_attacks
                },
                'active_bans': active_bans,
                'attack_change': today_attacks - yesterday_attacks
            }
    
    async def cleanup_old_records(self, days: int = 7) -> int:
        """Очистка старых записей"""
        cutoff_date = datetime.now() - timedelta(days=days)
        
        async with aiosqlite.connect(self.db_path) as db:
            # Удаляем старые попытки атак
            async with db.execute("""
                DELETE FROM attack_attempts 
                WHERE timestamp < ?
            """, (cutoff_date,)) as cursor:
                deleted_attempts = cursor.rowcount
            
            # Удаляем неактивные баны старше cutoff_date
            await db.execute("""
                DELETE FROM banned_ips 
                WHERE is_active = 0 AND banned_at < ?
            """, (cutoff_date,))
            
            await db.commit()
            logger.info(f"Очищено {deleted_attempts} старых записей")
            return deleted_attempts
    
    async def is_whitelisted(self, ip: str, whitelist: List[str]) -> bool:
        """Проверка IP в белом списке (поддержка CIDR)"""
        try:
            ip_obj = ipaddress.ip_address(ip)
            for white_ip in whitelist:
                try:
                    # Проверяем как сеть (CIDR)
                    if '/' in white_ip:
                        network = ipaddress.ip_network(white_ip, strict=False)
                        if ip_obj in network:
                            return True
                    # Проверяем как отдельный IP
                    elif ip_obj == ipaddress.ip_address(white_ip):
                        return True
                except Exception:
                    continue
            return False
        except Exception:
            return False

    async def record_compromise(self, compromise_info: Dict) -> None:
        """Запись информации о компрометации"""
        async with aiosqlite.connect(self.db_path) as db:
            await db.execute("""
                INSERT INTO compromises 
                (ip_address, username, detection_time, session_active, new_password, session_info)
                VALUES (?, ?, ?, ?, ?, ?)
            """, (
                compromise_info['ip'],
                compromise_info['username'],
                compromise_info['detection_time'],
                compromise_info['session_active'],
                compromise_info.get('new_password', ''),
                json.dumps(compromise_info.get('sessions', []))
            ))
            await db.commit()
    
    async def get_compromises(self, limit: int = 50) -> List[Dict]:
        """Получение списка компрометаций"""
        async with aiosqlite.connect(self.db_path) as db:
            async with db.execute("""
                SELECT ip_address, username, detection_time, session_active, new_password, session_info
                FROM compromises 
                ORDER BY detection_time DESC 
                LIMIT ?
            """, (limit,)) as cursor:
                results = await cursor.fetchall()
                
                compromises = []
                for row in results:
                    compromises.append({
                        'ip': row[0],
                        'username': row[1],
                        'detection_time': row[2],
                        'session_active': bool(row[3]),
                        'new_password': row[4],
                        'sessions': json.loads(row[5]) if row[5] else []
                    })
                return compromises

    async def update_daily_stats(self) -> None:
        """Обновить ежедневную статистику"""
        today = datetime.now().date()
        
        async with aiosqlite.connect(self.db_path) as db:
            # Получаем статистику за сегодня
            stats = await self.get_daily_stats()
            
            # Обновляем или создаем запись
            await db.execute("""
                INSERT OR REPLACE INTO daily_stats 
                (date, total_attempts, unique_ips, successful_logins)
                VALUES (?, ?, ?, ?)
            """, (today, stats['today']['attacks'], 
                  stats['today']['unique_ips'], 
                  stats['today']['successful_logins']))
            await db.commit()
    
    # === CLUSTER MANAGEMENT METHODS ===
    
    async def add_agent(self, agent_id: str, agent_info: Dict) -> None:
        """Добавление агента в базу данных"""
        async with aiosqlite.connect(self.db_path) as db:
            await db.execute("""
                INSERT OR REPLACE INTO agents 
                (agent_id, hostname, ip_address, ssh_port, ssh_user, status, added_time, last_check, version, config)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """, (
                agent_id,
                agent_info.get('hostname'),
                agent_info.get('ip_address'),
                agent_info.get('ssh_port', 22),
                agent_info.get('ssh_user', 'root'),
                agent_info.get('status', 'added'),
                datetime.now().isoformat(),
                agent_info.get('last_check'),
                agent_info.get('version'),
                json.dumps(agent_info)
            ))
            await db.commit()
    
    async def update_agent_status(self, agent_id: str, status: str) -> None:
        """Обновление статуса агента"""
        async with aiosqlite.connect(self.db_path) as db:
            await db.execute("""
                UPDATE agents 
                SET status = ?, last_check = ?
                WHERE agent_id = ?
            """, (status, datetime.now().isoformat(), agent_id))
            await db.commit()
    
    async def remove_agent(self, agent_id: str) -> None:
        """Удаление агента из базы данных"""
        async with aiosqlite.connect(self.db_path) as db:
            await db.execute("DELETE FROM agents WHERE agent_id = ?", (agent_id,))
            await db.commit()
    
    async def get_agents(self) -> List[Dict]:
        """Получение списка всех агентов"""
        async with aiosqlite.connect(self.db_path) as db:
            async with db.execute("""
                SELECT agent_id, hostname, ip_address, status, added_time, last_check, version
                FROM agents
                ORDER BY added_time DESC
            """) as cursor:
                results = await cursor.fetchall()
                
                agents = []
                for row in results:
                    agents.append({
                        'agent_id': row[0],
                        'hostname': row[1],
                        'ip_address': row[2],
                        'status': row[3],
                        'added_time': row[4],
                        'last_check': row[5],
                        'version': row[6]
                    })
                
                return agents
    
    async def get_agent_info(self, agent_id: str) -> Optional[Dict]:
        """Получение информации об агенте"""
        async with aiosqlite.connect(self.db_path) as db:
            async with db.execute("""
                SELECT agent_id, hostname, ip_address, ssh_port, ssh_user, status, 
                       added_time, last_check, version, config
                FROM agents
                WHERE agent_id = ?
            """, (agent_id,)) as cursor:
                row = await cursor.fetchone()
                
                if row:
                    config = json.loads(row[9]) if row[9] else {}
                    return {
                        'agent_id': row[0],
                        'hostname': row[1],
                        'ip_address': row[2],
                        'ssh_port': row[3],
                        'ssh_user': row[4],
                        'status': row[5],
                        'added_time': row[6],
                        'last_check': row[7],
                        'version': row[8],
                        'config': config
                    }
                
                return None
    
    async def get_cluster_stats(self) -> Dict:
        """Получение статистики кластера"""
        async with aiosqlite.connect(self.db_path) as db:
            # Подсчет агентов по статусам
            async with db.execute("""
                SELECT status, COUNT(*) 
                FROM agents 
                GROUP BY status
            """) as cursor:
                status_counts = {}
                async for row in cursor:
                    status_counts[row[0]] = row[1]
            
            # Общее количество агентов
            async with db.execute("SELECT COUNT(*) FROM agents") as cursor:
                total_agents = (await cursor.fetchone())[0]
            
            return {
                'total_agents': total_agents,
                'status_distribution': status_counts,
                'online_agents': status_counts.get('online', 0),
                'offline_agents': status_counts.get('offline', 0),
                'deployed_agents': status_counts.get('deployed', 0)
            }