gonets/SRV/api/v1/ufw.py

from fastapi import APIRouter, HTTPException
from typing import List
from pydantic import BaseModel, Field

from models.ufw import UFWRulesResponse, UFWStatus, UFWRule
from services.ufw_service import UFWService

router = APIRouter()


class RuleNumbers(BaseModel):
    """Модель запроса на удаление UFW правил.

    Атрибуты:

    - ``numbers`` (List[int]): Список номеров правил для удаления.
    """
    numbers: List[int] = Field(
        ..., description="Список номеров правил для удаления")


@router.get("/rules", response_model=UFWRulesResponse)
async def get_ufw_rules():
    """Возвращает все правила UFW.

    Осуществляет обращение к сервису `UFWService` для чтения текущего списка
    правил брандмауэра и возвращает результат в сериализованном виде для клиента.

    Returns:
        UFWRulesResponse: Структура с коллекцией правил и сопутствующей
        информацией о состоянии.

    Raises:
        HTTPException: В случае ошибок взаимодействия с UFW (status_code=500).
    """
    ufw_service = UFWService()
    try:
        return ufw_service.get_ufw_rules()
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error: {str(e)}")


@router.get("/status", response_model=UFWStatus)
async def get_ufw_status():
    """Возвращает текущий статус UFW.

    Returns:
        UFWStatus: Объект со сведениями о включенности UFW и деталями статуса.

    Raises:
        HTTPException: Если не удалось получить статус (status_code=500).
    """
    ufw_service = UFWService()
    try:
        return ufw_service.get_ufw_status()
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error: {str(e)}")


@router.post("/enable")
async def enable_ufw():
    """Включает UFW.

    Returns:
        dict: Словарь с ключами `status` и `message` при успешном выполнении.

    Raises:
        HTTPException: Если включение завершилось неуспешно (status_code=500).
    """
    ufw_service = UFWService()
    try:
        success = ufw_service.enable_ufw()
        if success:
            return {"status": "success", "message": "UFW успешно включен"}
        else:
            raise HTTPException(status_code=500,
                                detail="Не удалось включить UFW")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error: {str(e)}")


@router.post("/disable")
async def disable_ufw():
    """Выключает UFW.

    Returns:
        dict: Словарь с ключами `status` и `message` при успешном выполнении.

    Raises:
        HTTPException: Если выключение завершилось неуспешно (status_code=500).
    """
    ufw_service = UFWService()
    try:
        success = ufw_service.disable_ufw()
        if success:
            return {"status": "success", "message": "UFW успешно выключен"}
        else:
            raise HTTPException(status_code=500,
                                detail="Не удалось выключить UFW")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error: {str(e)}")


@router.post("/rules/in")
async def add_ufw_in_rules(rules: List[UFWRule]):
    """Добавляет входящие (in) правила UFW.

    Args:
        rules (List[UFWRule]): Список правил, определяющих направления,
            порты/диапазоны/протоколы и действия (allow/deny/reject).

    Returns:
        dict: Итог выполнения с количеством успешно добавленных правил и деталями
        по каждому правилу.

    Raises:
        HTTPException: Если UFW отключен (status_code=400) или при внутренних
        ошибках (status_code=500).
    """
    ufw_service = UFWService()
    try:
        # Проверяем статус UFW
        status = ufw_service.get_ufw_status()
        if not status.enabled:
            raise HTTPException(
                status_code=400,
                detail="UFW отключен. Включите UFW перед добавлением правил.",
            )

        result = ufw_service.add_rules(rules)
        if result["success"]:
            successful_rules = len(
                [r for r in result["results"] if r["success"]])
            total_rules = len(result["results"])
            return {
                "status": "success",
                "message":
                f"Добавлено {successful_rules} из {total_rules} правил",
                "details": result,
            }
        else:
            raise HTTPException(
                status_code=500,
                detail=result.get("error", "Не удалось добавить правила"),
            )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error: {str(e)}")


@router.post("/rules/out")
async def add_ufw_out_rules(rules: List[UFWRule]):
    """Добавляет исходящие (out) правила UFW.

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

    Returns:
        dict: Объект статуса с сообщением об успешном добавлении.

    Raises:
        HTTPException: Если UFW отключен (status_code=400) или при внутренних
        ошибках (status_code=500).
    """
    ufw_service = UFWService()
    try:
        # Проверяем статус UFW
        status = ufw_service.get_ufw_status()
        if not status.enabled:
            raise HTTPException(
                status_code=400,
                detail="UFW отключен. Включите UFW перед добавлением правил.",
            )

        success = ufw_service.add_ufw_out_rules(rules)
        if success:
            return {
                "status": "success",
                "message": "Исходящие правила UFW успешно добавлены",
            }
        else:
            raise HTTPException(
                status_code=500,
                detail="Не удалось добавить исходящие правила UFW")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error: {str(e)}")


@router.delete("/rules")
async def delete_ufw_rules(rule_numbers: RuleNumbers):
    """Удаляет правила UFW по их номерам.

    Args:
        rule_numbers (RuleNumbers): Номера правил, подлежащих удалению (как в
            выводе `ufw status numbered`).

    Returns:
        dict: Информация об общем количестве удаляемых правил, числе успешных и
        возможной ошибке.

    Raises:
        HTTPException: При ошибках удаления (status_code=500).
    """
    ufw_service = UFWService()
    try:
        result = ufw_service.delete_rules(rule_numbers.numbers)
        if result["success"]:
            return {
                "status": "success",
                "message":
                f"Удалено {result['successful']} из {result['total']} правил",
                "details": result,
            }
        else:
            raise HTTPException(status_code=500, detail=result["error"])
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error: {str(e)}")