nest/uvm/routers/vms.py

"""VM lifecycle endpoints backed by the existing UVM lifecycle service."""

from __future__ import annotations

from fastapi import APIRouter, Depends, Query, status

from ..api_models import (
    DestroyResponse,
    PortSyncResponse,
    PortUpdateRequest,
    VmActionResultResponse,
    VmAgentInstallResponse,
    VmBulkActionRequest,
    VmCreateRequest,
    VmExecRequest,
    VmExecResponse,
    VmLogResponse,
    VmResponse,
    VmResizeRequest,
    vm_response,
)
from ..app import Application
from ..domain import PortRule, VmRecord, VmSpec
from ..errors import ValidationError
from ..tool_catalog import normalize_tools
from ..validation import (
    parse_cpu,
    parse_name,
    parse_password,
    parse_ports,
    parse_ram,
    parse_storage,
    parse_timeout,
    parse_username,
)
from .dependencies import get_authorized_application
from .tools import tool_service


router = APIRouter(prefix="/vms", tags=["vms"])


def _mutation_policy_status(vm: VmRecord) -> str:
    """Return the policy state that follows a successful lifecycle mutation."""
    if vm.ports is None:
        return "unmanaged"
    if vm.operation is not None:
        return "unknown"
    if vm.status == "running":
        return "enforced"
    if vm.status == "stopped":
        return "inactive"
    return "unknown"


@router.get("", response_model=list[VmResponse])
def list_vms(
    application: Application = Depends(get_authorized_application),
) -> list[VmResponse]:
    return [
        vm_response(
            listed.vm,
            observed_status=listed.observed_status,
            policy_status=listed.policy_status,
        )
        for listed in application.lifecycle.list_vms()
    ]


@router.post("", response_model=VmResponse, status_code=status.HTTP_201_CREATED)
def create_vm(
    request: VmCreateRequest,
    application: Application = Depends(get_authorized_application),
) -> VmResponse:
    ports = parse_ports([f"{rule.port}:{rule.protocol}" for rule in request.ports or ()])
    if not request.no_default_ports and not any(rule.port == 22 for rule in ports):
        ports = tuple(sorted((*ports, PortRule(22)), key=lambda rule: rule.port))
    tools = normalize_tools([item.model_dump() for item in request.tools], allow_empty=True)
    if tools:
        tool_service(application).start()
    vm = application.lifecycle.create(
        VmSpec(
            cpu=parse_cpu(str(request.cpu)),
            ram_mib=parse_ram(str(request.ram)),
            storage_mib=parse_storage(str(request.storage)),
            guest_ip=request.guest_ip,
            username=parse_username(request.username),
            password=parse_password(request.password.get_secret_value()),
            name=parse_name(request.name) if request.name is not None else None,
            ports=ports,
            tools=tools,
        )
    )
    if tools:
        tool_service(application).notify()
    return vm_response(vm, policy_status=_mutation_policy_status(vm))


@router.post("/actions", response_model=list[VmActionResultResponse])
def run_bulk_action(
    request: VmBulkActionRequest,
    application: Application = Depends(get_authorized_application),
) -> list[VmActionResultResponse]:
    timeout = parse_timeout(str(request.timeout)) if request.timeout is not None else None
    targets = application.lifecycle.resolve_targets(request.identifiers, all_vms=request.all_vms)
    results = application.lifecycle.run_many(
        request.action,
        [vm.id for vm in targets],
        timeout=timeout,
    )
    return [
        VmActionResultResponse(
            id=result.id,
            action=result.action,
            outcome=result.outcome,
            error=result.error,
        )
        for result in results
    ]


@router.post("/{vm_id}/start", response_model=VmResponse)
def start_vm(
    vm_id: str,
    application: Application = Depends(get_authorized_application),
) -> VmResponse:
    vm = application.lifecycle.start(vm_id)
    return vm_response(vm, policy_status=_mutation_policy_status(vm))


@router.post("/{vm_id}/restart", response_model=VmResponse)
def restart_vm(
    vm_id: str,
    timeout: float | None = None,
    application: Application = Depends(get_authorized_application),
) -> VmResponse:
    vm = (
        application.lifecycle.restart(vm_id)
        if timeout is None
        else application.lifecycle.restart(vm_id, timeout=timeout)
    )
    return vm_response(vm, policy_status=_mutation_policy_status(vm))


@router.post("/{vm_id}/stop", response_model=VmResponse)
def stop_vm(
    vm_id: str,
    timeout: float | None = None,
    application: Application = Depends(get_authorized_application),
) -> VmResponse:
    vm = (
        application.lifecycle.stop(vm_id)
        if timeout is None
        else application.lifecycle.stop(vm_id, timeout=timeout)
    )
    return vm_response(vm, policy_status=_mutation_policy_status(vm))


@router.post("/{vm_id}/resize", response_model=VmResponse)
def resize_vm(
    vm_id: str,
    request: VmResizeRequest,
    timeout: float | None = None,
    application: Application = Depends(get_authorized_application),
) -> VmResponse:
    if request.cpu is None and request.ram is None and request.storage is None:
        raise ValidationError("resize requires at least one of CPU, RAM, or storage")
    vm = application.lifecycle.resize(
        vm_id,
        cpu=parse_cpu(str(request.cpu)) if request.cpu is not None else None,
        ram_mib=parse_ram(str(request.ram)) if request.ram is not None else None,
        storage_mib=parse_storage(str(request.storage)) if request.storage is not None else None,
        timeout=parse_timeout(str(timeout)) if timeout is not None else None,
    )
    return vm_response(vm, policy_status=_mutation_policy_status(vm))


@router.delete("/{vm_id}", response_model=DestroyResponse)
def destroy_vm(
    vm_id: str,
    timeout: float | None = None,
    application: Application = Depends(get_authorized_application),
) -> DestroyResponse:
    vm = (
        application.lifecycle.destroy(vm_id)
        if timeout is None
        else application.lifecycle.destroy(vm_id, timeout=timeout)
    )
    return DestroyResponse(id=vm.id)


@router.post("/{vm_id}/ports", response_model=VmResponse)
def update_ports(
    vm_id: str,
    request: PortUpdateRequest,
    application: Application = Depends(get_authorized_application),
) -> VmResponse:
    ports = parse_ports([f"{rule.port}:{rule.protocol}" for rule in request.ports])
    vm = application.lifecycle.set_ports(vm_id, ports, mode=request.mode)
    return vm_response(vm, policy_status=_mutation_policy_status(vm))


@router.post("/{vm_id}/ports/sync", response_model=PortSyncResponse)
def sync_ports(
    vm_id: str,
    application: Application = Depends(get_authorized_application),
) -> PortSyncResponse:
    vm = application.lifecycle.sync_ports(vm_id)
    return PortSyncResponse(
        id=vm.id,
        policy_status=_mutation_policy_status(vm),
        status=vm.status,
    )


@router.post("/{vm_id}/exec", response_model=VmExecResponse)
def execute_in_vm(
    vm_id: str,
    request: VmExecRequest,
    application: Application = Depends(get_authorized_application),
) -> VmExecResponse:
    """Execute a bounded, non-interactive guest command over vsock (exact VM ID)."""
    result = application.lifecycle.execute(
        vm_id, request.command, sudo=request.sudo, timeout_seconds=request.timeout_seconds,
    )
    return VmExecResponse(
        exit_code=result.exit_code, stdout=result.stdout, stderr=result.stderr,
        timed_out=result.timed_out, truncated=result.truncated,
    )


@router.post("/{vm_id}/agent/install", response_model=VmAgentInstallResponse)
def install_guest_agent(
    vm_id: str,
    application: Application = Depends(get_authorized_application),
) -> VmAgentInstallResponse:
    """Install/upgrade the execution agent on a stopped VM's retained private disk."""
    vm = application.lifecycle.install_guest_agent(vm_id)
    return VmAgentInstallResponse(id=vm.id)


@router.get("/{identifier}/logs", response_model=VmLogResponse)
def get_vm_logs(
    identifier: str,
    tail: int = Query(default=100, ge=0, le=1000),
    application: Application = Depends(get_authorized_application),
) -> VmLogResponse:
    log = application.lifecycle.read_log_tail(identifier, tail=tail)
    return VmLogResponse(content=log.content, truncated=log.truncated)


@router.get("/{identifier}", response_model=VmResponse)
def get_vm(
    identifier: str,
    application: Application = Depends(get_authorized_application),
) -> VmResponse:
    listed = application.lifecycle.inspect(identifier)
    return vm_response(
        listed.vm,
        observed_status=listed.observed_status,
        policy_status=listed.policy_status,
    )