"""
Firewall module для PyGuardian
Управление iptables/nftables для блокировки IP-адресов
"""

import asyncio
import subprocess
import logging
from typing import Dict, List, Optional
from abc import ABC, abstractmethod

logger = logging.getLogger(__name__)


class FirewallInterface(ABC):
    """Абстрактный интерфейс для работы с firewall"""
    
    @abstractmethod
    async def ban_ip(self, ip: str) -> bool:
        """Забанить IP адрес"""
        pass
    
    @abstractmethod
    async def unban_ip(self, ip: str) -> bool:
        """Разбанить IP адрес"""
        pass
    
    @abstractmethod
    async def is_banned(self, ip: str) -> bool:
        """Проверить, забанен ли IP"""
        pass
    
    @abstractmethod
    async def list_banned_ips(self) -> List[str]:
        """Получить список забаненных IP"""
        pass
    
    @abstractmethod
    async def setup_chains(self) -> bool:
        """Настроить цепочки firewall"""
        pass


class IptablesFirewall(FirewallInterface):
    """Реализация для iptables"""
    
    def __init__(self, config: Dict):
        self.chain = config.get('chain', 'INPUT')
        self.target = config.get('target', 'DROP')
        self.table = config.get('iptables', {}).get('table', 'filter')
        self.comment = "PyGuardian-ban"
    
    async def _run_command(self, command: List[str]) -> tuple[bool, str]:
        """Выполнить команду iptables"""
        try:
            result = await asyncio.create_subprocess_exec(
                *command,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE
            )
            stdout, stderr = await result.communicate()
            
            if result.returncode == 0:
                return True, stdout.decode().strip()
            else:
                error_msg = stderr.decode().strip()
                logger.error(f"Ошибка выполнения команды {' '.join(command)}: {error_msg}")
                return False, error_msg
        
        except Exception as e:
            logger.error(f"Исключение при выполнении команды {' '.join(command)}: {e}")
            return False, str(e)
    
    async def setup_chains(self) -> bool:
        """Настроить цепочки iptables"""
        try:
            # Создаем специальную цепочку для PyGuardian если не существует
            pyguardian_chain = "PYGUARDIAN"
            
            # Проверяем, существует ли цепочка
            success, _ = await self._run_command([
                "iptables", "-t", self.table, "-L", pyguardian_chain, "-n"
            ])
            
            if not success:
                # Создаем цепочку
                success, _ = await self._run_command([
                    "iptables", "-t", self.table, "-N", pyguardian_chain
                ])
                if not success:
                    return False
                
                logger.info(f"Создана цепочка {pyguardian_chain}")
            
            # Проверяем, есть ли правило перехода в нашу цепочку
            success, output = await self._run_command([
                "iptables", "-t", self.table, "-L", self.chain, "-n", "--line-numbers"
            ])
            
            if success and pyguardian_chain not in output:
                # Добавляем правило в начало цепочки INPUT
                success, _ = await self._run_command([
                    "iptables", "-t", self.table, "-I", self.chain, "1", 
                    "-j", pyguardian_chain
                ])
                if success:
                    logger.info(f"Добавлено правило перехода в цепочку {pyguardian_chain}")
                else:
                    return False
            
            return True
            
        except Exception as e:
            logger.error(f"Ошибка настройки цепочек iptables: {e}")
            return False
    
    async def ban_ip(self, ip: str) -> bool:
        """Забанить IP адрес через iptables"""
        try:
            # Проверяем, не забанен ли уже
            if await self.is_banned(ip):
                logger.warning(f"IP {ip} уже забанен в iptables")
                return True
            
            # Добавляем правило блокировки
            command = [
                "iptables", "-t", self.table, "-A", "PYGUARDIAN",
                "-s", ip, "-j", self.target,
                "-m", "comment", "--comment", self.comment
            ]
            
            success, error = await self._run_command(command)
            if success:
                logger.info(f"IP {ip} заблокирован в iptables")
                return True
            else:
                logger.error(f"Не удалось заблокировать IP {ip}: {error}")
                return False
                
        except Exception as e:
            logger.error(f"Ошибка при блокировке IP {ip}: {e}")
            return False
    
    async def unban_ip(self, ip: str) -> bool:
        """Разбанить IP адрес"""
        try:
            # Находим и удаляем правило
            command = [
                "iptables", "-t", self.table, "-D", "PYGUARDIAN",
                "-s", ip, "-j", self.target,
                "-m", "comment", "--comment", self.comment
            ]
            
            success, error = await self._run_command(command)
            if success:
                logger.info(f"IP {ip} разблокирован в iptables")
                return True
            else:
                # Возможно, правило уже удалено
                logger.warning(f"Не удалось удалить правило для IP {ip}: {error}")
                return True
                
        except Exception as e:
            logger.error(f"Ошибка при разблокировке IP {ip}: {e}")
            return False
    
    async def is_banned(self, ip: str) -> bool:
        """Проверить, забанен ли IP"""
        try:
            command = [
                "iptables", "-t", self.table, "-L", "PYGUARDIAN", "-n"
            ]
            
            success, output = await self._run_command(command)
            if success:
                return ip in output and self.comment in output
            else:
                return False
                
        except Exception as e:
            logger.error(f"Ошибка при проверке IP {ip}: {e}")
            return False
    
    async def list_banned_ips(self) -> List[str]:
        """Получить список забаненных IP"""
        try:
            command = [
                "iptables", "-t", self.table, "-L", "PYGUARDIAN", "-n"
            ]
            
            success, output = await self._run_command(command)
            if not success:
                return []
            
            banned_ips = []
            for line in output.split('\n'):
                if self.comment in line and self.target in line:
                    parts = line.split()
                    if len(parts) >= 4:
                        source_ip = parts[3]
                        if '/' not in source_ip:  # Исключаем сети
                            banned_ips.append(source_ip)
            
            return banned_ips
            
        except Exception as e:
            logger.error(f"Ошибка при получении списка забаненных IP: {e}")
            return []


class NftablesFirewall(FirewallInterface):
    """Реализация для nftables"""
    
    def __init__(self, config: Dict):
        self.table = config.get('nftables', {}).get('table', 'inet pyguardian')
        self.chain = config.get('nftables', {}).get('chain', 'input')
        self.set_name = "banned_ips"
    
    async def _run_command(self, command: List[str]) -> tuple[bool, str]:
        """Выполнить команду nft"""
        try:
            result = await asyncio.create_subprocess_exec(
                *command,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE
            )
            stdout, stderr = await result.communicate()
            
            if result.returncode == 0:
                return True, stdout.decode().strip()
            else:
                error_msg = stderr.decode().strip()
                logger.error(f"Ошибка выполнения команды {' '.join(command)}: {error_msg}")
                return False, error_msg
        
        except Exception as e:
            logger.error(f"Исключение при выполнении команды {' '.join(command)}: {e}")
            return False, str(e)
    
    async def setup_chains(self) -> bool:
        """Настроить таблицу и цепочку nftables"""
        try:
            # Создаем таблицу
            success, _ = await self._run_command([
                "nft", "create", "table", self.table
            ])
            # Игнорируем ошибку, если таблица уже существует
            
            # Создаем set для хранения IP адресов
            success, _ = await self._run_command([
                "nft", "create", "set", f"{self.table}", self.set_name,
                "{ type ipv4_addr; }"
            ])
            # Игнорируем ошибку, если set уже существует
            
            # Создаем цепочку
            success, _ = await self._run_command([
                "nft", "create", "chain", f"{self.table}", self.chain,
                "{ type filter hook input priority 0; policy accept; }"
            ])
            # Игнорируем ошибку, если цепочка уже существует
            
            # Добавляем правило блокировки для IP из set
            success, _ = await self._run_command([
                "nft", "add", "rule", f"{self.table}", self.chain,
                "ip", "saddr", f"@{self.set_name}", "drop"
            ])
            
            logger.info("Настройка nftables выполнена успешно")
            return True
            
        except Exception as e:
            logger.error(f"Ошибка настройки nftables: {e}")
            return False
    
    async def ban_ip(self, ip: str) -> bool:
        """Забанить IP адрес через nftables"""
        try:
            command = [
                "nft", "add", "element", f"{self.table}", self.set_name,
                f"{{ {ip} }}"
            ]
            
            success, error = await self._run_command(command)
            if success:
                logger.info(f"IP {ip} заблокирован в nftables")
                return True
            else:
                logger.error(f"Не удалось заблокировать IP {ip}: {error}")
                return False
                
        except Exception as e:
            logger.error(f"Ошибка при блокировке IP {ip}: {e}")
            return False
    
    async def unban_ip(self, ip: str) -> bool:
        """Разбанить IP адрес"""
        try:
            command = [
                "nft", "delete", "element", f"{self.table}", self.set_name,
                f"{{ {ip} }}"
            ]
            
            success, error = await self._run_command(command)
            if success:
                logger.info(f"IP {ip} разблокирован в nftables")
                return True
            else:
                logger.warning(f"Не удалось удалить IP {ip}: {error}")
                return True  # Возможно, IP уже не в списке
                
        except Exception as e:
            logger.error(f"Ошибка при разблокировке IP {ip}: {e}")
            return False
    
    async def is_banned(self, ip: str) -> bool:
        """Проверить, забанен ли IP"""
        try:
            banned_ips = await self.list_banned_ips()
            return ip in banned_ips
                
        except Exception as e:
            logger.error(f"Ошибка при проверке IP {ip}: {e}")
            return False
    
    async def list_banned_ips(self) -> List[str]:
        """Получить список забаненных IP"""
        try:
            command = [
                "nft", "list", "set", f"{self.table}", self.set_name
            ]
            
            success, output = await self._run_command(command)
            if not success:
                return []
            
            banned_ips = []
            in_elements = False
            
            for line in output.split('\n'):
                line = line.strip()
                if 'elements = {' in line:
                    in_elements = True
                    # Проверяем, есть ли IP на той же строке
                    if '}' in line:
                        elements_part = line.split('{')[1].split('}')[0]
                        banned_ips.extend([ip.strip() for ip in elements_part.split(',') if ip.strip()])
                        break
                elif in_elements:
                    if '}' in line:
                        elements_part = line.split('}')[0]
                        banned_ips.extend([ip.strip() for ip in elements_part.split(',') if ip.strip()])
                        break
                    else:
                        banned_ips.extend([ip.strip() for ip in line.split(',') if ip.strip()])
            
            return [ip for ip in banned_ips if ip and ip != '']
            
        except Exception as e:
            logger.error(f"Ошибка при получении списка забаненных IP: {e}")
            return []


class FirewallManager:
    """Менеджер для управления firewall"""
    
    def __init__(self, config: Dict):
        self.config = config
        backend = config.get('backend', 'iptables').lower()
        
        if backend == 'iptables':
            self.firewall = IptablesFirewall(config)
        elif backend == 'nftables':
            self.firewall = NftablesFirewall(config)
        else:
            raise ValueError(f"Неподдерживаемый backend: {backend}")
        
        self.backend = backend
        logger.info(f"Инициализирован {backend} firewall")
    
    async def setup(self) -> bool:
        """Настроить firewall"""
        return await self.firewall.setup_chains()
    
    async def ban_ip(self, ip: str) -> bool:
        """Забанить IP адрес"""
        return await self.firewall.ban_ip(ip)
    
    async def unban_ip(self, ip: str) -> bool:
        """Разбанить IP адрес"""
        return await self.firewall.unban_ip(ip)
    
    async def is_banned(self, ip: str) -> bool:
        """Проверить, забанен ли IP"""
        return await self.firewall.is_banned(ip)
    
    async def list_banned_ips(self) -> List[str]:
        """Получить список забаненных IP"""
        return await self.firewall.list_banned_ips()
    
    async def get_status(self) -> Dict:
        """Получить статус firewall"""
        try:
            banned_ips = await self.list_banned_ips()
            return {
                'backend': self.backend,
                'active': True,
                'banned_count': len(banned_ips),
                'banned_ips': banned_ips[:10]  # Первые 10 для отображения
            }
        except Exception as e:
            logger.error(f"Ошибка получения статуса firewall: {e}")
            return {
                'backend': self.backend,
                'active': False,
                'error': str(e)
            }
    
    async def cleanup_expired_bans(self, valid_ips: List[str]) -> int:
        """Очистить firewall от IP, которые больше не должны быть забанены"""
        try:
            current_banned = await self.list_banned_ips()
            removed_count = 0
            
            for ip in current_banned:
                if ip not in valid_ips:
                    if await self.unban_ip(ip):
                        removed_count += 1
                        logger.info(f"Удален устаревший бан для IP {ip}")
            
            return removed_count
            
        except Exception as e:
            logger.error(f"Ошибка очистки устаревших банов: {e}")
            return 0