napalm-proxmox/napalm_proxmox/vm_contract_mixin.py

"""HypervisorDriver contract methods for Proxmox VE: power actions and VM config.

``power_vm`` stays for callers that already use it; these are what a
hypervisor-neutral caller talks to. They raise instead of returning a
``{"success": ...}`` dict, and block until Proxmox reports the task finished.
"""

from __future__ import annotations

import re
from typing import Any

from napalm_device_types.models import (
    VMConfigDict,
    VMDiskDict,
    VMNICDict,
    VMPassthroughDict,
)

_JsonDict = dict[str, Any]

_POWER_TIMEOUT = 120
_VM_DISK_KEY = re.compile(r"^(scsi|ide|virtio|sata)\d+$|^efidisk\d+$|^tpmstate\d+$")
_CT_DISK_KEY = re.compile(r"^rootfs$|^mp\d+$")
_NET_KEY = re.compile(r"^net\d+$")
_PASSTHROUGH_KEY = re.compile(r"^(hostpci|usb)\d+$")
_NIC_MODELS = {"virtio", "e1000", "e1000e", "vmxnet3", "rtl8139", "ne2k_pci"}
_SIZE = re.compile(r"^(\d+(?:\.\d+)?)([KMGT]?)$", re.I)
_GB_PER_UNIT = {"K": 1 / 1024**2, "M": 1 / 1024, "G": 1, "T": 1024, "": 1}


def _options(value: str) -> tuple[str, dict[str, str]]:
    """Split ``"volume,key=val,..."`` into the leading bare part and its options."""
    head = ""
    opts: dict[str, str] = {}
    for part in str(value).split(","):
        if "=" in part:
            k, v = part.split("=", 1)
            opts[k.strip().lower()] = v.strip()
        elif not head:
            head = part.strip()
    return head, opts


def _size_gb(raw: str) -> int:
    m = _SIZE.match(raw or "")
    if not m:
        return 0
    return int(float(m.group(1)) * _GB_PER_UNIT[m.group(2).upper()])


def _boot_order(cfg: _JsonDict) -> list[str]:
    boot = str(cfg.get("boot", "") or "")
    if boot.startswith("order="):
        return [d for d in boot[len("order=") :].split(";") if d]
    bootdisk = cfg.get("bootdisk")
    return [bootdisk] if bootdisk else []


def _disks(cfg: _JsonDict, vm_type: str, boot_order: list[str]) -> list[VMDiskDict]:
    key_re = _VM_DISK_KEY if vm_type == "vm" else _CT_DISK_KEY
    disks: list[VMDiskDict] = []
    for key in sorted(cfg, key=lambda k: (k != "rootfs", k)):
        if not key_re.match(key):
            continue
        value = str(cfg[key] or "")
        head, opts = _options(value)
        if opts.get("media") == "cdrom" or head in ("none", "0", ""):
            continue
        disks.append(
            {
                "device": key,
                "storage": head.split(":", 1)[0],
                "size": _size_gb(opts.get("size", "")),
                "format": opts.get("format", ""),
                "bootable": key in boot_order,
            }
        )
    return disks


def _nic(key: str, value: str) -> VMNICDict:
    _, opts = _options(value)
    model = next((m for m in _NIC_MODELS if m in opts), opts.get("type", ""))
    mac = opts.get(model, "") if model in _NIC_MODELS else opts.get("hwaddr", "")
    tag = opts.get("tag", "")
    return {
        "device": key,
        "mac": mac.upper(),
        "model": model,
        "bridge": opts.get("bridge", ""),
        "vlan_id": int(tag) if tag.isdigit() else 0,
    }


def _passthrough(cfg: _JsonDict) -> list[VMPassthroughDict]:
    return [
        {"slot": key, "kind": "pci" if key.startswith("hostpci") else "usb", "config": str(val)}
        for key, val in sorted(cfg.items())
        if _PASSTHROUGH_KEY.match(key)
    ]


def parse_vm_config(vmid: str, vm_type: str, cfg: _JsonDict) -> VMConfigDict:
    """Turn a raw ``/qemu/{id}/config`` or ``/lxc/{id}/config`` into a VMConfigDict."""
    is_vm = vm_type == "vm"
    cores = int(cfg.get("cores", 1) or 1)
    sockets = int(cfg.get("sockets", 1) or 1) if is_vm else 1
    boot_order = _boot_order(cfg)
    tags = str(cfg.get("tags", "") or "")
    name_key = "name" if is_vm else "hostname"
    result: VMConfigDict = {
        "name": cfg.get(name_key) or f"{'vm' if is_vm else 'ct'}-{vmid}",
        "vmid": vmid,
        "vcpus": cores * sockets,
        "memory": int(cfg.get("memory", 0) or 0),
        "os_type": cfg.get("ostype", ""),
        "boot_order": boot_order,
        "disks": _disks(cfg, vm_type, boot_order),
        "nics": [_nic(k, str(v)) for k, v in sorted(cfg.items()) if _NET_KEY.match(k)],
        "description": cfg.get("description", ""),
        "tags": [t for t in re.split(r"[;,\s]+", tags) if t],
        "passthrough": _passthrough(cfg),
    }
    if is_vm:
        result["cpu_type"] = str(cfg.get("cpu", "kvm64")).split(",")[0].removeprefix("cputype=")
        result["sockets"] = sockets
        result["cores_per_socket"] = cores
        result["firmware"] = "efi" if cfg.get("bios") == "ovmf" else "bios"
        if cfg.get("machine"):
            result["machine"] = cfg["machine"]
    return result


class ProxmoxVMContractMixin:
    """HypervisorDriver's VM methods on top of the Proxmox node API."""

    def _resolve_vm(self, name: str) -> tuple[int, str]:
        """Find a guest by vmid or display name; return ``(vmid, "vm"|"container")``."""
        node = self._node_api()
        for vm_type, listing in (("vm", node.qemu), ("container", node.lxc)):
            for guest in listing.get() or []:
                if str(guest.get("vmid")) == name or guest.get("name") == name:
                    return int(guest["vmid"]), vm_type
        raise ValueError(f"No VM or container named or numbered {name!r}")

    def _guest_api(self, vmid: int, vm_type: str) -> Any:
        node = self._node_api()
        return node.qemu(vmid) if vm_type == "vm" else node.lxc(vmid)

    def _run_power(self, vmid: int, vm_type: str, action: str) -> None:
        try:
            upid = getattr(self._guest_api(vmid, vm_type).status, action).post()
        except Exception as exc:
            raise RuntimeError(f"{action} of {vm_type} {vmid} failed: {exc}") from exc
        if upid:
            self._wait_for_task(upid, timeout=_POWER_TIMEOUT)

    def start_vm(self, name: str) -> None:
        self._run_power(*self._resolve_vm(name), "start")

    def stop_vm(self, name: str, force: bool = False) -> None:
        self._run_power(*self._resolve_vm(name), "stop" if force else "shutdown")

    def reboot_vm(self, name: str, force: bool = False) -> None:
        vmid, vm_type = self._resolve_vm(name)
        if not force:
            self._run_power(vmid, vm_type, "reboot")
        elif vm_type == "vm":
            self._run_power(vmid, vm_type, "reset")
        else:
            self._run_power(vmid, vm_type, "stop")
            self._run_power(vmid, vm_type, "start")

    def suspend_vm(self, name: str) -> None:
        vmid, vm_type = self._resolve_vm(name)
        if vm_type != "vm":
            raise RuntimeError(f"Proxmox cannot suspend container {vmid}")
        self._run_power(vmid, vm_type, "suspend")

    def get_vm_config(self, name: str) -> VMConfigDict:
        vmid, vm_type = self._resolve_vm(name)
        cfg = self._guest_api(vmid, vm_type).config.get() or {}
        return parse_vm_config(str(vmid), vm_type, cfg)