nest/uvm/api_models.py

"""Pydantic request and response models for the UVM HTTP API."""

from __future__ import annotations

from ipaddress import IPv4Address
from typing import Literal

from pydantic import BaseModel, ConfigDict, Field, SecretStr, model_validator

from .domain import VmRecord
from .errors import ValidationError
from .guest_agent import DEFAULT_TIMEOUT_SECONDS, MAX_COMMAND_BYTES, MAX_TIMEOUT_SECONDS
from .guest_exec import validate_exec


class HealthResponse(BaseModel):
    status: str = "ok"


class ServerExitResponse(BaseModel):
    status: str = "stopping"


class InstallRequest(BaseModel):
    model_config = ConfigDict(extra="forbid")

    force: bool = False


class InstallResponse(BaseModel):
    firecracker: str
    kernel: str
    rootfs: str


class PortRuleModel(BaseModel):
    """A private guest TCP port; HTTP and HTTPS are descriptive labels only."""

    model_config = ConfigDict(extra="forbid")

    port: int = Field(strict=True, ge=1, le=65535)
    protocol: Literal["TCP", "HTTP", "HTTPS"] = "TCP"


class VmCreateRequest(BaseModel):
    model_config = ConfigDict(extra="forbid")

    cpu: float = Field(default=1, gt=0)
    ram: str | int | float = "512"
    storage: str | int | float = Field(default="30G", description="Disk capacity, minimum 30 GiB; bare numbers use GiB")
    guest_ip: IPv4Address | None = None
    username: str = "root"
    password: SecretStr = SecretStr("root")
    name: str | None = None
    ports: list[PortRuleModel] | None = None
    no_default_ports: bool = Field(default=False, strict=True)


class VmResizeRequest(BaseModel):
    model_config = ConfigDict(extra="forbid")

    cpu: float | None = Field(default=None, gt=0, allow_inf_nan=False, strict=True)
    ram: str | int | float | None = None
    storage: str | int | float | None = Field(default=None, description="Disk capacity, minimum 30 GiB; bare numbers use GiB")


class VmResourcesModel(BaseModel):
    cpu: float
    ram_mib: int
    storage_mib: int


class VmResponse(BaseModel):
    id: str
    name: str | None
    cpu: float
    ram_mib: int
    storage_mib: int | None
    pending_resize: VmResourcesModel | None
    guest_ip: str
    gateway: str
    mac: str
    username: str
    status: str
    observed_status: str
    ports: list[PortRuleModel] | None
    policy_status: str
    operation: str | None
    pid: int | None
    created_at: int
    updated_at: int
    last_error: str | None


class DestroyResponse(BaseModel):
    id: str
    status: str = "terminated"


class PortUpdateRequest(BaseModel):
    model_config = ConfigDict(extra="forbid")

    mode: Literal["set", "open", "close"] = "set"
    ports: list[PortRuleModel] = Field(default_factory=list)


class PortSyncResponse(BaseModel):
    id: str
    policy_status: str
    status: str


class VmLogResponse(BaseModel):
    content: str
    truncated: bool


class VmExecRequest(BaseModel):
    model_config = ConfigDict(extra="forbid")

    command: str = Field(strict=True, min_length=1, max_length=MAX_COMMAND_BYTES, repr=False,
                         description="Command or multiline script interpreted by /bin/sh inside the guest; at most 64 KiB UTF-8")
    sudo: bool = Field(default=False, strict=True, description="Run as guest root instead of the VM's configured login user")
    timeout_seconds: float = Field(default=DEFAULT_TIMEOUT_SECONDS, strict=True, gt=0,
                                   le=MAX_TIMEOUT_SECONDS, allow_inf_nan=False)

    @model_validator(mode="before")
    @classmethod
    def validate_execution(cls, value: object) -> object:
        # Use the domain error envelope: echoing NaN or unpaired surrogates in
        # FastAPI's validation-error input can itself fail JSON serialization.
        if not isinstance(value, dict) or set(value) - {"command", "sudo", "timeout_seconds"}:
            raise ValidationError("execution requires an object with command, sudo, and timeout_seconds fields only")
        validate_exec(value.get("command"), value.get("sudo", False), value.get("timeout_seconds", DEFAULT_TIMEOUT_SECONDS))
        return value


class VmExecResponse(BaseModel):
    exit_code: int | None
    stdout: str
    stderr: str
    timed_out: bool
    truncated: bool


class VmAgentInstallResponse(BaseModel):
    id: str
    status: Literal["installed"] = "installed"
    restart_required: bool = True


class DiagnosticResponse(BaseModel):
    name: str
    ok: bool
    message: str


class VmBulkActionRequest(BaseModel):
    model_config = ConfigDict(extra="forbid")

    action: Literal["start", "stop", "restart", "delete", "ports-sync"]
    identifiers: list[str] = Field(default_factory=list)
    all_vms: bool = Field(default=False, strict=True)
    timeout: float | None = None


class VmActionResultResponse(BaseModel):
    id: str
    action: str
    outcome: str
    error: str | None = None


def vm_response(
    vm: VmRecord,
    *,
    observed_status: str | None = None,
    policy_status: str = "unknown",
) -> VmResponse:
    return VmResponse(
        id=vm.id,
        name=vm.name,
        cpu=vm.cpu,
        ram_mib=vm.ram_mib,
        storage_mib=vm.storage_mib,
        pending_resize=VmResourcesModel(**vm.pending_resize.to_dict()) if vm.pending_resize is not None else None,
        guest_ip=vm.guest_ip,
        gateway=vm.gateway,
        mac=vm.mac,
        username=vm.username,
        status=vm.status,
        observed_status=observed_status if observed_status is not None else vm.status,
        ports=(
            [PortRuleModel(port=rule.port, protocol=rule.protocol) for rule in vm.ports]
            if vm.ports is not None
            else None
        ),
        policy_status=policy_status,
        operation=vm.operation,
        pid=vm.pid,
        created_at=vm.created_at,
        updated_at=vm.updated_at,
        last_error=vm.last_error,
    )