PyGuardian/.history/src/cluster_manager_20251125205610.py

"""
Cluster Manager для PyGuardian
Управление кластером серверов и автоматическое развертывание агентов
"""

import asyncio
import logging
import json
import subprocess
import os
import yaml
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Any
from pathlib import Path
import aiofiles
import paramiko
from cryptography.fernet import Fernet
import secrets
import string
import hashlib

# Импортируем систему аутентификации
from .auth import AgentAuthentication, AgentAuthenticationError
from .storage import Storage

logger = logging.getLogger(__name__)


class ServerAgent:
    """Представление удаленного сервера-агента"""
    
    def __init__(self, server_id: str, config: Dict):
        self.server_id = server_id
        self.hostname = config.get('hostname')
        self.ip_address = config.get('ip_address')
        self.ssh_port = config.get('ssh_port', 22)
        self.ssh_user = config.get('ssh_user', 'root')
        self.ssh_key_path = config.get('ssh_key_path')
        self.ssh_password = config.get('ssh_password')
        self.status = 'unknown'
        self.last_check = None
        self.version = None
        self.stats = {}
        
        # Новые поля для аутентификации
        self.agent_id = config.get('agent_id')
        self.secret_key = config.get('secret_key')
        self.access_token = config.get('access_token')
        self.refresh_token = config.get('refresh_token')
        self.token_expires_at = config.get('token_expires_at')
        self.last_authenticated = config.get('last_authenticated')
        self.is_authenticated = False
        
    def to_dict(self) -> Dict:
        """Конвертация в словарь для сериализации"""
        return {
            'server_id': self.server_id,
            'hostname': self.hostname,
            'ip_address': self.ip_address,
            'ssh_port': self.ssh_port,
            'ssh_user': self.ssh_user,
            'ssh_key_path': self.ssh_key_path,
            'status': self.status,
            'last_check': self.last_check.isoformat() if self.last_check else None,
            'version': self.version,
            'agent_id': self.agent_id,
            'is_authenticated': self.is_authenticated,
            'last_authenticated': self.last_authenticated,
            'token_expires_at': self.token_expires_at,
            'stats': self.stats
        }


class ClusterManager:
    """Менеджер кластера серверов"""
    
    def __init__(self, storage: Storage, config: Dict):
        self.storage = storage
        self.config = config
        
        # Параметры кластера
        self.cluster_name = config.get('cluster_name', 'PyGuardian-Cluster')
        self.master_server = config.get('master_server', True)
        self.agents_config_path = config.get('agents_config_path', '/var/lib/pyguardian/agents.yaml')
        self.deployment_path = config.get('deployment_path', '/opt/pyguardian')
        
        # SSH настройки
        self.ssh_timeout = config.get('ssh_timeout', 30)
        self.ssh_retries = config.get('ssh_retries', 3)
        
        # Шифрование
        self.encryption_key = self._get_or_create_cluster_key()
        self.cipher = Fernet(self.encryption_key)
        
        # Инициализация системы аутентификации
        cluster_secret = config.get('cluster_secret', self._generate_cluster_secret())
        self.auth_manager = AgentAuthentication(
            secret_key=cluster_secret,
            token_expiry_minutes=config.get('token_expiry_minutes', 30)
        )
        
        # Кэш агентов
        self.agents: Dict[str, ServerAgent] = {}
        
        # Шаблоны для развертывания
        self.deployment_script = self._get_deployment_script()
    
    def _generate_cluster_secret(self) -> str:
        """Генерация секрета кластера если он не задан"""
        secret = secrets.token_urlsafe(64)
        logger.warning(f"Generated new cluster secret. Add to config: cluster_secret: {secret}")
        return secret
        
    def _get_or_create_cluster_key(self) -> bytes:
        """Получить или создать ключ шифрования кластера"""
        key_file = "/var/lib/pyguardian/cluster_encryption.key"
        try:
            os.makedirs(os.path.dirname(key_file), exist_ok=True)
            
            if os.path.exists(key_file):
                with open(key_file, 'rb') as f:
                    return f.read()
            else:
                key = Fernet.generate_key()
                with open(key_file, 'wb') as f:
                    f.write(key)
                os.chmod(key_file, 0o600)
                logger.info("Создан новый ключ шифрования кластера")
                return key
        except Exception as e:
            logger.error(f"Ошибка работы с ключом кластера: {e}")
            return Fernet.generate_key()
    
    def _get_deployment_script(self) -> str:
        """Получить скрипт развертывания агента"""
        return '''#!/bin/bash
# PyGuardian Agent Deployment Script
set -e

INSTALL_DIR="/opt/pyguardian"
SERVICE_NAME="pyguardian-agent"
GITHUB_REPO="https://github.com/your-repo/PyGuardian.git"

echo "🛡️  Начинаю установку PyGuardian Agent..."

# Проверка прав root
if [[ $EUID -ne 0 ]]; then
   echo "❌ Скрипт должен быть запущен от имени root"
   exit 1
fi

# Установка зависимостей
echo "📦 Установка зависимостей..."
if command -v apt >/dev/null 2>&1; then
    apt update
    apt install -y python3 python3-pip git
elif command -v yum >/dev/null 2>&1; then
    yum update -y
    yum install -y python3 python3-pip git
elif command -v dnf >/dev/null 2>&1; then
    dnf update -y
    dnf install -y python3 python3-pip git
else
    echo "❌ Неподдерживаемая система. Поддерживаются: Ubuntu/Debian/CentOS/RHEL"
    exit 1
fi

# Создание директорий
echo "📁 Создание директорий..."
mkdir -p $INSTALL_DIR
mkdir -p /var/lib/pyguardian
mkdir -p /var/log/pyguardian

# Клонирование репозитория
echo "⬇️  Клонирование PyGuardian..."
if [ -d "$INSTALL_DIR/.git" ]; then
    cd $INSTALL_DIR && git pull
else
    git clone $GITHUB_REPO $INSTALL_DIR
fi

cd $INSTALL_DIR

# Установка Python зависимостей
echo "🐍 Установка Python пакетов..."
pip3 install -r requirements.txt

# Настройка systemd сервиса
echo "⚙️  Настройка systemd сервиса..."
cat > /etc/systemd/system/$SERVICE_NAME.service << EOF
[Unit]
Description=PyGuardian Security Agent
After=network.target
Wants=network.target

[Service]
Type=simple
User=root
WorkingDirectory=$INSTALL_DIR
ExecStart=/usr/bin/python3 $INSTALL_DIR/main.py --agent-mode
Restart=always
RestartSec=10
Environment=PYTHONPATH=$INSTALL_DIR

[Install]
WantedBy=multi-user.target
EOF

# Включение и запуск сервиса
echo "🚀 Запуск PyGuardian Agent..."
systemctl daemon-reload
systemctl enable $SERVICE_NAME
systemctl start $SERVICE_NAME

echo "✅ PyGuardian Agent успешно установлен и запущен!"
echo "📊 Статус: systemctl status $SERVICE_NAME"
echo "📋 Логи: journalctl -u $SERVICE_NAME -f"
'''
    
    async def load_agents(self) -> None:
        """Загрузить конфигурацию агентов"""
        try:
            if os.path.exists(self.agents_config_path):
                async with aiofiles.open(self.agents_config_path, 'r') as f:
                    content = await f.read()
                    agents_config = yaml.safe_load(content)
                    
                    self.agents = {}
                    for agent_id, agent_config in agents_config.get('agents', {}).items():
                        self.agents[agent_id] = ServerAgent(agent_id, agent_config)
                        
                logger.info(f"Загружено {len(self.agents)} агентов из конфигурации")
            else:
                logger.info("Файл конфигурации агентов не найден, создаю новый")
                await self.save_agents()
                
        except Exception as e:
            logger.error(f"Ошибка загрузки агентов: {e}")
    
    async def save_agents(self) -> None:
        """Сохранить конфигурацию агентов"""
        try:
            os.makedirs(os.path.dirname(self.agents_config_path), exist_ok=True)
            
            agents_config = {
                'cluster': {
                    'name': self.cluster_name,
                    'master_server': self.master_server,
                    'last_updated': datetime.now().isoformat()
                },
                'agents': {}
            }
            
            for agent_id, agent in self.agents.items():
                agents_config['agents'][agent_id] = {
                    'hostname': agent.hostname,
                    'ip_address': agent.ip_address,
                    'ssh_port': agent.ssh_port,
                    'ssh_user': agent.ssh_user,
                    'ssh_key_path': agent.ssh_key_path,
                    'status': agent.status,
                    'last_check': agent.last_check.isoformat() if agent.last_check else None,
                    'version': agent.version
                }
            
            async with aiofiles.open(self.agents_config_path, 'w') as f:
                await f.write(yaml.dump(agents_config, default_flow_style=False))
                
            logger.info("Конфигурация агентов сохранена")
            
        except Exception as e:
            logger.error(f"Ошибка сохранения агентов: {e}")
    
    def generate_agent_id(self, hostname: str, ip_address: str) -> str:
        """Генерировать уникальный ID для агента"""
        return f"{hostname}-{ip_address.replace('.', '-')}"
    
    async def add_agent(self, hostname: str, ip_address: str, ssh_user: str = 'root', 
                       ssh_port: int = 22, ssh_key_path: str = None, 
                       ssh_password: str = None) -> tuple[bool, str]:
        """Добавить новый агент в кластер"""
        try:
            agent_id = self.generate_agent_id(hostname, ip_address)
            
            # Проверяем, что агент еще не добавлен
            if agent_id in self.agents:
                return False, f"Агент {agent_id} уже существует в кластере"
            
            # Создаем объект агента
            agent_config = {
                'hostname': hostname,
                'ip_address': ip_address,
                'ssh_port': ssh_port,
                'ssh_user': ssh_user,
                'ssh_key_path': ssh_key_path,
                'ssh_password': ssh_password
            }
            
            agent = ServerAgent(agent_id, agent_config)
            
            # Тестируем соединение
            connection_test = await self._test_ssh_connection(agent)
            if not connection_test[0]:
                return False, f"Не удалось подключиться к серверу: {connection_test[1]}"
            
            # Добавляем агент
            self.agents[agent_id] = agent
            agent.status = 'added'
            agent.last_check = datetime.now()
            
            # Сохраняем конфигурацию
            await self.save_agents()
            
            # Записываем в базу данных
            await self.storage.add_agent(agent_id, agent.to_dict())
            
            logger.info(f"Агент {agent_id} успешно добавлен в кластер")
            return True, f"Агент {hostname} ({ip_address}) добавлен в кластер"
            
        except Exception as e:
            logger.error(f"Ошибка добавления агента: {e}")
            return False, f"Ошибка добавления агента: {e}"
    
    async def _test_ssh_connection(self, agent: ServerAgent) -> tuple[bool, str]:
        """Тестирование SSH соединения с агентом"""
        try:
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            
            # Подключение
            if agent.ssh_key_path and os.path.exists(agent.ssh_key_path):
                ssh.connect(
                    hostname=agent.ip_address,
                    port=agent.ssh_port,
                    username=agent.ssh_user,
                    key_filename=agent.ssh_key_path,
                    timeout=self.ssh_timeout
                )
            elif agent.ssh_password:
                ssh.connect(
                    hostname=agent.ip_address,
                    port=agent.ssh_port,
                    username=agent.ssh_user,
                    password=agent.ssh_password,
                    timeout=self.ssh_timeout
                )
            else:
                return False, "Не указан метод аутентификации (ключ или пароль)"
            
            # Тестовая команда
            stdin, stdout, stderr = ssh.exec_command('echo "PyGuardian Connection Test"')
            result = stdout.read().decode().strip()
            
            ssh.close()
            
            if "PyGuardian Connection Test" in result:
                return True, "Соединение установлено успешно"
            else:
                return False, "Тестовая команда не выполнена"
                
        except Exception as e:
            return False, f"Ошибка SSH соединения: {e}"
    
    async def deploy_agent(self, agent_id: str, force_reinstall: bool = False) -> tuple[bool, str]:
        """Развернуть PyGuardian агент на удаленном сервере"""
        try:
            if agent_id not in self.agents:
                return False, f"Агент {agent_id} не найден"
            
            agent = self.agents[agent_id]
            
            # Подключение SSH
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            
            if agent.ssh_key_path and os.path.exists(agent.ssh_key_path):
                ssh.connect(
                    hostname=agent.ip_address,
                    port=agent.ssh_port,
                    username=agent.ssh_user,
                    key_filename=agent.ssh_key_path,
                    timeout=self.ssh_timeout
                )
            elif agent.ssh_password:
                ssh.connect(
                    hostname=agent.ip_address,
                    port=agent.ssh_port,
                    username=agent.ssh_user,
                    password=agent.ssh_password,
                    timeout=self.ssh_timeout
                )
            else:
                return False, "Не настроена аутентификация"
            
            # Проверка, установлен ли уже PyGuardian
            if not force_reinstall:
                stdin, stdout, stderr = ssh.exec_command('systemctl status pyguardian-agent')
                if stdout.channel.recv_exit_status() == 0:
                    agent.status = 'deployed'
                    await self.save_agents()
                    ssh.close()
                    return True, f"PyGuardian уже установлен на {agent.hostname}"
            
            # Создание временного скрипта развертывания
            temp_script = f'/tmp/pyguardian_deploy_{secrets.token_hex(8)}.sh'
            
            # Загрузка скрипта на сервер
            sftp = ssh.open_sftp()
            with sftp.open(temp_script, 'w') as f:
                f.write(self.deployment_script)
            sftp.chmod(temp_script, 0o755)
            sftp.close()
            
            # Выполнение скрипта развертывания
            logger.info(f"Начинаю развертывание на {agent.hostname}...")
            
            stdin, stdout, stderr = ssh.exec_command(f'bash {temp_script}')
            
            # Получение вывода
            deploy_output = stdout.read().decode()
            deploy_errors = stderr.read().decode()
            exit_status = stdout.channel.recv_exit_status()
            
            # Удаление временного скрипта
            ssh.exec_command(f'rm -f {temp_script}')
            
            if exit_status == 0:
                agent.status = 'deployed'
                agent.last_check = datetime.now()
                await self.save_agents()
                
                # Обновляем базу данных
                await self.storage.update_agent_status(agent_id, 'deployed')
                
                ssh.close()
                logger.info(f"PyGuardian успешно развернут на {agent.hostname}")
                return True, f"PyGuardian успешно установлен на {agent.hostname}"
            else:
                ssh.close()
                logger.error(f"Ошибка развертывания на {agent.hostname}: {deploy_errors}")
                return False, f"Ошибка установки: {deploy_errors[:500]}"
                
        except Exception as e:
            logger.error(f"Ошибка развертывания агента {agent_id}: {e}")
            return False, f"Ошибка развертывания: {e}"
    
    async def remove_agent(self, agent_id: str, cleanup_remote: bool = False) -> tuple[bool, str]:
        """Удалить агент из кластера"""
        try:
            if agent_id not in self.agents:
                return False, f"Агент {agent_id} не найден"
            
            agent = self.agents[agent_id]
            
            # Удаление с удаленного сервера
            if cleanup_remote:
                cleanup_result = await self._cleanup_remote_agent(agent)
                if not cleanup_result[0]:
                    logger.warning(f"Не удалось очистить удаленный агент: {cleanup_result[1]}")
            
            # Удаление из локальной конфигурации
            del self.agents[agent_id]
            await self.save_agents()
            
            # Удаление из базы данных
            await self.storage.remove_agent(agent_id)
            
            logger.info(f"Агент {agent_id} удален из кластера")
            return True, f"Агент {agent.hostname} удален из кластера"
            
        except Exception as e:
            logger.error(f"Ошибка удаления агента: {e}")
            return False, f"Ошибка удаления агента: {e}"
    
    async def _cleanup_remote_agent(self, agent: ServerAgent) -> tuple[bool, str]:
        """Очистка PyGuardian на удаленном сервере"""
        try:
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            
            # Подключение
            if agent.ssh_key_path and os.path.exists(agent.ssh_key_path):
                ssh.connect(
                    hostname=agent.ip_address,
                    port=agent.ssh_port,
                    username=agent.ssh_user,
                    key_filename=agent.ssh_key_path,
                    timeout=self.ssh_timeout
                )
            elif agent.ssh_password:
                ssh.connect(
                    hostname=agent.ip_address,
                    port=agent.ssh_port,
                    username=agent.ssh_user,
                    password=agent.ssh_password,
                    timeout=self.ssh_timeout
                )
            else:
                return False, "Не настроена аутентификация"
            
            # Команды очистки
            cleanup_commands = [
                'systemctl stop pyguardian-agent',
                'systemctl disable pyguardian-agent',
                'rm -f /etc/systemd/system/pyguardian-agent.service',
                'systemctl daemon-reload',
                'rm -rf /opt/pyguardian',
                'rm -rf /var/lib/pyguardian',
                'rm -f /var/log/pyguardian.log'
            ]
            
            for command in cleanup_commands:
                ssh.exec_command(command)
            
            ssh.close()
            return True, "Удаленная очистка выполнена"
            
        except Exception as e:
            return False, f"Ошибка очистки: {e}"
    
    async def check_agent_status(self, agent_id: str) -> tuple[bool, Dict]:
        """Проверить статус агента"""
        try:
            if agent_id not in self.agents:
                return False, {"error": f"Агент {agent_id} не найден"}
            
            agent = self.agents[agent_id]
            
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            
            # Подключение
            if agent.ssh_key_path and os.path.exists(agent.ssh_key_path):
                ssh.connect(
                    hostname=agent.ip_address,
                    port=agent.ssh_port,
                    username=agent.ssh_user,
                    key_filename=agent.ssh_key_path,
                    timeout=self.ssh_timeout
                )
            elif agent.ssh_password:
                ssh.connect(
                    hostname=agent.ip_address,
                    port=agent.ssh_port,
                    username=agent.ssh_user,
                    password=agent.ssh_password,
                    timeout=self.ssh_timeout
                )
            else:
                return False, {"error": "Не настроена аутентификация"}
            
            # Проверка статуса сервиса
            stdin, stdout, stderr = ssh.exec_command('systemctl is-active pyguardian-agent')
            service_status = stdout.read().decode().strip()
            
            # Проверка версии
            stdin, stdout, stderr = ssh.exec_command('cat /opt/pyguardian/VERSION 2>/dev/null || echo "unknown"')
            version = stdout.read().decode().strip()
            
            # Получение системной информации
            stdin, stdout, stderr = ssh.exec_command('uptime && df -h / && free -m')
            system_info = stdout.read().decode()
            
            ssh.close()
            
            # Обновление информации об агенте
            agent.status = 'online' if service_status == 'active' else 'offline'
            agent.version = version
            agent.last_check = datetime.now()
            
            status_info = {
                "agent_id": agent_id,
                "hostname": agent.hostname,
                "ip_address": agent.ip_address,
                "status": agent.status,
                "service_status": service_status,
                "version": version,
                "last_check": agent.last_check.isoformat(),
                "system_info": system_info
            }
            
            await self.save_agents()
            return True, status_info
            
        except Exception as e:
            logger.error(f"Ошибка проверки статуса агента {agent_id}: {e}")
            return False, {"error": f"Ошибка проверки статуса: {e}"}
    
    async def list_agents(self) -> List[Dict]:
        """Получить список всех агентов"""
        agents_list = []
        for agent_id, agent in self.agents.items():
            agents_list.append({
                "agent_id": agent_id,
                "hostname": agent.hostname,
                "ip_address": agent.ip_address,
                "status": agent.status,
                "last_check": agent.last_check.isoformat() if agent.last_check else None,
                "version": agent.version
            })
        
        return agents_list
    
    async def get_cluster_stats(self) -> Dict:
        """Получить статистику кластера"""
        total_agents = len(self.agents)
        online_agents = len([a for a in self.agents.values() if a.status == 'online'])
        offline_agents = len([a for a in self.agents.values() if a.status == 'offline'])
        deployed_agents = len([a for a in self.agents.values() if a.status == 'deployed'])
        
        return {
            "cluster_name": self.cluster_name,
            "total_agents": total_agents,
            "online_agents": online_agents,
            "offline_agents": offline_agents,
            "deployed_agents": deployed_agents,
            "master_server": self.master_server,
            "last_updated": datetime.now().isoformat()
        }
    
    async def check_all_agents(self) -> Dict:
        """Проверить статус всех агентов"""
        results = {
            "checked": 0,
            "online": 0,
            "offline": 0,
            "errors": 0,
            "details": []
        }
        
        for agent_id in self.agents.keys():
            try:
                success, status_info = await self.check_agent_status(agent_id)
                results["checked"] += 1
                
                if success:
                    if status_info.get("status") == "online":
                        results["online"] += 1
                    else:
                        results["offline"] += 1
                else:
                    results["errors"] += 1
                
                results["details"].append(status_info)
                
            except Exception as e:
                results["errors"] += 1
                results["details"].append({
                    "agent_id": agent_id,
                    "error": str(e)
                })
        
        return results