vmware-api-simulator/app/vsphere/rest/vm_ext.py

"""VM hardware, snapshots, clone, tools — REST extensions."""

from __future__ import annotations

from typing import Any

from fastapi import APIRouter, Depends, Query, Response

from app.db.pool import Database
from app.dependencies import get_database
from app.vsphere.domain import vm_ops
from app.vsphere.errors import invalid_argument
from app.vsphere.rest import mappers
from app.vsphere.security.authz import require_privilege, require_read
from app.vsphere.security.session import SessionInfo

router = APIRouter(tags=["vSphere VM Ext"])


@router.get("/api/vcenter/vm/{vm}/tools")
async def vm_tools(
    vm: str, database: Database = Depends(get_database), _: SessionInfo = Depends(require_read)
) -> dict[str, Any]:
    obj = await vm_ops.require_vm(database, vm)
    tools = obj.props.get("tools")
    if isinstance(tools, dict) and tools:
        out = dict(tools)
        out["run_state"] = (
            "RUNNING" if obj.props.get("power_state") == "POWERED_ON" else "NOT_RUNNING"
        )
        return out
    return {}


@router.get("/api/vcenter/vm/{vm}/hardware")
async def vm_hardware(
    vm: str, database: Database = Depends(get_database), _: SessionInfo = Depends(require_read)
) -> dict[str, Any]:
    obj = await vm_ops.require_vm(database, vm)
    return mappers.vm_info(obj)["hardware"]


@router.get("/api/vcenter/vm/{vm}/hardware/cpu")
async def get_cpu(
    vm: str, database: Database = Depends(get_database), _: SessionInfo = Depends(require_read)
) -> dict[str, Any]:
    obj = await vm_ops.require_vm(database, vm)
    cpu = dict(obj.props.get("cpu") or {})
    return {
        "count": obj.props.get("cpu_count")
        if obj.props.get("cpu_count") is not None
        else cpu.get("count"),
        "cores_per_socket": cpu.get("cores_per_socket"),
        "hot_add_enabled": cpu.get("hot_add_enabled"),
    }


@router.patch("/api/vcenter/vm/{vm}/hardware/cpu")
async def patch_cpu(
    vm: str,
    body: dict[str, Any],
    database: Database = Depends(get_database),
    _: SessionInfo = Depends(require_privilege("VirtualMachine.Config.CPUCount")),
) -> Response:
    await vm_ops.update_hardware_cpu(database, vm, int(body.get("count") or 1))
    return Response(status_code=204)


@router.get("/api/vcenter/vm/{vm}/hardware/memory")
async def get_memory(
    vm: str, database: Database = Depends(get_database), _: SessionInfo = Depends(require_read)
) -> dict[str, Any]:
    obj = await vm_ops.require_vm(database, vm)
    memory = dict(obj.props.get("memory") or {})
    return {
        "size_MiB": obj.props.get("memory_size_mib")
        if obj.props.get("memory_size_mib") is not None
        else memory.get("size_MiB"),
        "hot_add_enabled": memory.get("hot_add_enabled"),
    }


@router.patch("/api/vcenter/vm/{vm}/hardware/memory")
async def patch_memory(
    vm: str,
    body: dict[str, Any],
    database: Database = Depends(get_database),
    _: SessionInfo = Depends(require_privilege("VirtualMachine.Config.Memory")),
) -> Response:
    await vm_ops.update_hardware_memory(
        database, vm, int(body.get("size_MiB") or body.get("size_mib") or 1024)
    )
    return Response(status_code=204)


@router.get("/api/vcenter/vm/{vm}/hardware/disk")
async def list_disks(
    vm: str, database: Database = Depends(get_database), _: SessionInfo = Depends(require_read)
) -> list[dict[str, Any]]:
    obj = await vm_ops.require_vm(database, vm)
    return list(obj.props.get("disks") or [])


@router.post("/api/vcenter/vm/{vm}/hardware/disk")
async def create_disk(
    vm: str,
    body: dict[str, Any],
    database: Database = Depends(get_database),
    _: SessionInfo = Depends(require_privilege("VirtualMachine.Config.AddNewDisk")),
) -> str:
    capacity = int(
        (body.get("new_vmdk") or {}).get("capacity") or body.get("capacity") or 10737418240
    )
    disk = await vm_ops.add_disk(database, vm, capacity)
    return str(disk["key"])


@router.get("/api/vcenter/vm/{vm}/hardware/ethernet")
async def list_nics(
    vm: str, database: Database = Depends(get_database), _: SessionInfo = Depends(require_read)
) -> list[dict[str, Any]]:
    obj = await vm_ops.require_vm(database, vm)
    return list(obj.props.get("nics") or [])


@router.post("/api/vcenter/vm/{vm}/hardware/ethernet")
async def create_nic(
    vm: str,
    body: dict[str, Any],
    database: Database = Depends(get_database),
    _: SessionInfo = Depends(require_privilege("VirtualMachine.Config.AddRemoveDevice")),
) -> str:
    network = ((body.get("backing") or {}).get("network")) or body.get("network") or "network-41"
    nic = await vm_ops.add_nic(database, vm, str(network))
    return str(nic["key"])


@router.get("/api/vcenter/vm/{vm}/hardware/boot")
async def get_boot(
    vm: str, database: Database = Depends(get_database), _: SessionInfo = Depends(require_read)
) -> dict[str, Any]:
    obj = await vm_ops.require_vm(database, vm)
    boot = obj.props.get("boot")
    if isinstance(boot, dict) and boot:
        return boot
    from app.vsphere.domain import api_state

    payload = await api_state.get_payload(database, "GET", "/api/vcenter/vm/{vm}/hardware/boot")
    return payload if isinstance(payload, dict) else {}


@router.get("/api/vcenter/vm/{vm}/snapshots")
async def list_snapshots(
    vm: str, database: Database = Depends(get_database), _: SessionInfo = Depends(require_read)
) -> list[dict[str, Any]]:
    return await vm_ops.list_snapshots(database, vm)


@router.post("/api/vcenter/vm/{vm}/snapshots")
async def create_snapshot(
    vm: str,
    body: dict[str, Any],
    database: Database = Depends(get_database),
    _: SessionInfo = Depends(require_privilege("VirtualMachine.State.CreateSnapshot")),
) -> dict[str, str]:
    snap_id, task_id = await vm_ops.create_snapshot(
        database,
        vm,
        name=str(body.get("name") or "snapshot"),
        description=str(body.get("description") or ""),
        memory=bool(body.get("memory")),
    )
    return {"snapshot": snap_id, "task": task_id}


@router.delete("/api/vcenter/vm/{vm}/snapshots/{snapshot}")
async def delete_snapshot(
    vm: str,
    snapshot: str,
    database: Database = Depends(get_database),
    _: SessionInfo = Depends(require_privilege("VirtualMachine.State.RemoveSnapshot")),
) -> dict[str, str]:
    task_id = await vm_ops.delete_snapshot(database, vm, snapshot)
    return {"task": task_id}


@router.post("/api/vcenter/vm/{vm}/snapshots/{snapshot}")
async def revert_snapshot(
    vm: str,
    snapshot: str,
    action: str = Query("revert"),
    database: Database = Depends(get_database),
    _: SessionInfo = Depends(require_privilege("VirtualMachine.State.RevertToSnapshot")),
) -> dict[str, str]:
    if action != "revert":
        raise invalid_argument(f"unsupported action {action}")
    task_id = await vm_ops.revert_snapshot(database, vm, snapshot)
    return {"task": task_id}


@router.post("/api/vcenter/vm/{vm}/clone")
async def clone_vm(
    vm: str,
    body: dict[str, Any],
    database: Database = Depends(get_database),
    _: SessionInfo = Depends(require_privilege("VirtualMachine.Provisioning.Clone")),
) -> dict[str, str]:
    name = str(body.get("name") or f"{vm}-clone")
    folder = (body.get("placement") or {}).get("folder")
    moid, task_id = await vm_ops.clone_vm(
        database, source_vm=vm, name=name, folder=folder, power_on=bool(body.get("power_on"))
    )
    return {"vm": moid, "task": task_id}


@router.post("/api/vcenter/vm/{vm}/relocate")
async def relocate_vm(
    vm: str,
    body: dict[str, Any],
    database: Database = Depends(get_database),
    _: SessionInfo = Depends(require_privilege("VirtualMachine.Inventory.Move")),
) -> dict[str, str]:
    placement = body.get("placement") or body
    task_id = await vm_ops.relocate_vm(
        database, vm, host=placement.get("host"), datastore=placement.get("datastore")
    )
    return {"task": task_id}