gonets/SRV/services/ufw_service.py

import ipaddress
import subprocess
from typing import List, Union

from models.ufw import Direction, Protocol, UFWRule, UFWRulesResponse, UFWStatus

DEBUG = True


class UFWService:
    """
    Сервис для взаимодействия с утилитой UFW (Uncomplicated Firewall) в Linux.
    Позволяет получить список правил UFW и управлять ими.
    """

    def _run_command(self, command: str) -> str:
        """
        Выполняет команду в системе и возвращает ее вывод.

        Args:
            command (str): Команда для выполнения в системе.

        Returns:
            str: Вывод команды.
        """
        try:
            if DEBUG:
                print(f"Running command: {command}")

            # Используем shell=True для корректной обработки аргументов с пробелами
            result = subprocess.run(
                f"sudo {command}",
                shell=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
            )

            if DEBUG:
                print(f"Command output: {result.stdout}")
                print(f"Command error: {result.stderr}")

            if result.returncode != 0:
                raise subprocess.CalledProcessError(result.returncode, command,
                                                    result.stdout,
                                                    result.stderr)

            return result.stdout

        except subprocess.CalledProcessError as e:
            if DEBUG:
                print(f"Command failed with return code {e.returncode}")
                print(f"Error output: {e.stderr}")
            raise Exception(f"Command failed: {e.stderr}")

    def _get_port_range_args(self, port: Union[int, str]) -> List[str]:
        """
        Преобразует порт или диапазон портов в аргументы команды UFW.

        Args:
            port: Порт (число) или диапазон портов (строка вида "3000:3999")

        Returns:
            List[str]: Список аргументов для UFW
        """
        if isinstance(port, int):
            return [str(port)]
        elif isinstance(port, str) and ":" in port:
            # Диапазон портов уже в нужном формате
            return [port]
        else:
            raise ValueError(f"Некорректный формат порта: {port}")

    def get_ufw_rules(self) -> UFWRulesResponse:
        """
        Получает список правил UFW.

        Returns:
            UFWRulesResponse: Объект, содержащий статус UFW и список правил
        """
        # Check UFW status
        status_result = self._run_command("ufw status numbered")
        lines = status_result.splitlines()

        if not lines:
            return UFWRulesResponse(enabled=False, rules=[])

        first_line = lines[0]
        is_enabled = "Status: active" in first_line

        if not is_enabled:
            return UFWRulesResponse(enabled=False, rules=[])

        rules = []
        # Skip header lines
        for line in lines[4:]:  # Skip Status and header lines
            if not line.strip() or not line.startswith("["):
                continue

            try:
                rule = UFWRule.from_ufw_output(line)
                rules.append(rule)
            except Exception as e:
                if DEBUG:
                    print(f"Ошибка парсинга правила: {line}, ошибка: {str(e)}")
                continue

        return UFWRulesResponse(enabled=is_enabled, rules=rules)

    def get_ufw_status(self) -> UFWStatus:
        """
        Проверяет статус UFW.

        Returns:
            UFWStatus: Объект с информацией о статусе UFW
        """
        try:
            status_result = self._run_command("ufw status")
            first_line = (status_result.splitlines()[0]
                          if status_result.splitlines() else "")
            is_enabled = "Status: active" in first_line

            message = "UFW активен" if is_enabled else "UFW неактивен"

            if DEBUG:
                print(f"Результат статуса: {status_result}")
                print(f"Включен: {is_enabled}")

            return UFWStatus(enabled=is_enabled, message=message)
        except Exception as e:
            if DEBUG:
                print(f"Ошибка при проверке статуса UFW: {e}")
            return UFWStatus(
                enabled=False,
                message=f"Ошибка при проверке статуса: {str(e)}",
            )

    def enable_ufw(self) -> bool:
        """
        Включает UFW.

        Returns:
            bool: True если успешно включен, False в противном случае
        """
        try:
            self._run_command("ufw enable")
            return True
        except Exception as e:
            if DEBUG:
                print(f"Ошибка при включении UFW: {e}")
            return False

    def disable_ufw(self) -> bool:
        """
        Выключает UFW.

        Returns:
            bool: True если успешно выключен, False в противном случае
        """
        try:
            self._run_command("ufw disable")
            return True
        except Exception as e:
            if DEBUG:
                print(f"Ошибка при выключении UFW: {e}")
            return False

    def add_rules_base(self,
                       rules: List[UFWRule],
                       direction: str = "") -> bool:
        """
        Базовый метод для добавления правил UFW.

        Args:
            rules: Список правил для добавления
            direction: Направление трафика ("in" или "out")

        Returns:
            bool: True если все правила были успешно добавлены, False в противном случае
        """
        try:
            for rule in rules:
                # Устанавливаем направление правила
                rule.direction = Direction.IN if direction == "in" else Direction.OUT

                # Формируем базовую команду
                command = f"ufw {rule.action.value}"

                # Добавляем направление
                command += f" {rule.direction.value}"

                # Добавляем from IP если он указан
                if rule.from_ip and rule.from_ip != "any":
                    command += f" from {rule.from_ip}"
                else:
                    command += " from any"

                # Добавляем to IP если он указан
                if rule.to_ip and rule.to_ip != "any":
                    command += f" to {rule.to_ip}"
                else:
                    command += " to any"

                # Добавляем порт если он указан
                if rule.port:
                    command += f" port {rule.port}"

                # Добавляем протокол
                if rule.protocol != Protocol.BOTH:
                    command += f" proto {rule.protocol.value}"

                # Добавляем комментарий если он есть
                if rule.comment:
                    command += f" comment '{rule.comment}'"

                self._run_command(command)

            return True
        except Exception as e:
            if DEBUG:
                print(f"Error adding UFW rules: {str(e)}")
            return False

    def add_ufw_in_rules(self, rules: List[UFWRule]) -> bool:
        """
        Добавляет входящие (in) правила UFW.

        Args:
            rules (List[UFWRule]): Список правил для добавления

        Returns:
            bool: True если все правила были успешно добавлены, False в противном случае
        """
        return self.add_rules_base(rules, "in")

    def add_ufw_out_rules(self, rules: List[UFWRule]) -> bool:
        """
        Добавляет исходящие (out) правила UFW.

        Args:
            rules (List[UFWRule]): Список правил для добавления

        Returns:
            bool: True если все правила были успешно добавлены, False в противном случае
        """
        return self.add_rules_base(rules, "out")

    def add_rules(self, rules: List[UFWRule]) -> dict:
        """
        Добавляет правила UFW.

        Args:
            rules: Список правил для добавления

        Returns:
            dict: Результат добавления правил
        """
        results = []

        for rule in rules:
            try:
                # Определяем версию IP
                from_ip = ipaddress.ip_network(rule.from_ip, strict=False)
                is_ipv6 = isinstance(from_ip, ipaddress.IPv6Network)

                # Получаем список протоколов
                protocols = rule.protocol.protocols
                print(f"Обработка правила для протоколов: {protocols}")

                for proto in protocols:
                    # Формируем базовую команду
                    command = f"ufw {rule.action.value}"

                    # Добавляем IP адреса
                    command += f" from {rule.from_ip}"
                    if rule.to_ip:
                        command += f" to {rule.to_ip}"

                    # Добавляем порт и протокол
                    port_args = self._get_port_range_args(rule.port)
                    command += f" port {' '.join(port_args)} proto {proto}"

                    # Добавляем флаг IPv6 если нужно
                    if is_ipv6:
                        command += " v6"

                    # Добавляем комментарий если есть
                    if rule.comment:
                        command += f" comment '{rule.comment}'"

                    print(f"Подготовлена команда: {command}")

                    # Выполняем команду
                    try:
                        output = self._run_command(command)
                        results.append({"success": True, "message": output})
                    except Exception as e:
                        error_msg = f"Ошибка выполнения команды: {str(e)}"
                        print(error_msg)
                        results.append({
                            "success": False,
                            "message": error_msg
                        })

            except Exception as e:
                error_msg = f"Ошибка при добавлении правила: {str(e)}"
                print(error_msg)
                results.append({"success": False, "message": error_msg})

        # Выводим результаты
        print("\nРезультаты добавления правил:")
        for result in results:
            status = "Успешно" if result["success"] else "Ошибка"
            print(f"  - {status}: {result['message']}")

        return {"success": True, "results": results}

    def delete_rules(self, rule_numbers: List[int]) -> dict:
        """
        Удаляет правила UFW по их номерам.

        Args:
            rule_numbers (List[int]): Список номеров правил для удаления

        Returns:
            dict: Результат операции удаления
        """
        try:
            # Проверяем существование правил
            status_result = self._run_command("ufw status numbered")
            lines = status_result.splitlines()

            # Пропускаем заголовок (первые 4 строки)
            rules = []
            for line in lines[4:]:  # Skip header lines
                if not line.strip() or not line.startswith("["):
                    continue

                try:
                    # Извлекаем номер правила из строки [N]
                    rule_num = int(line[1:line.index("]")])
                    rules.append(rule_num)
                except (ValueError, IndexError):
                    continue

            if DEBUG:
                print(f"Найдены правила: {rules}")

            results = []
            # Сортируем номера в обратном порядке, чтобы удаление происходило с конца
            # Это предотвратит смещение номеров при удалении
            for rule_number in sorted(rule_numbers, reverse=True):
                try:
                    # Проверяем существует ли правило с таким номером
                    if rule_number not in rules:
                        results.append({
                            "number":
                            rule_number,
                            "success":
                            False,
                            "error":
                            f"Правило с номером {rule_number} не найдено",
                        })
                        continue

                    # Удаляем правило
                    self._run_command(f"ufw --force delete {rule_number}")

                    results.append({
                        "number":
                        rule_number,
                        "success":
                        True,
                        "message":
                        f"Правило {rule_number} успешно удалено",
                    })

                except Exception as e:
                    if DEBUG:
                        print(
                            f"Ошибка при удалении правила {rule_number}: {e}")
                    results.append({
                        "number": rule_number,
                        "success": False,
                        "error": str(e)
                    })

            # Подсчитываем статистику
            total = len(rule_numbers)
            successful = len([r for r in results if r["success"]])

            return {
                "success": True,
                "total": total,
                "successful": successful,
                "results": results,
            }

        except Exception as e:
            if DEBUG:
                print(f"Ошибка при удалении правил: {e}")
            return {"success": False, "error": str(e), "results": []}