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):
    numbers: List[int] = Field(
        ..., description="Список номеров правил для удаления")


@router.get("/rules", response_model=UFWRulesResponse)
async def get_ufw_rules():
    """
    Получает все правила UFW с помощью сервиса UFWService и возвращает их в формате JSON.
    """
    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.
    """
    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.
    """
    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.
    """
    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.
    """
    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.
    """
    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 по их номерам.
    """
    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)}")