nest/uvm/domain.py

"""Internal data structures persisted by the uvm CLI."""

from __future__ import annotations

import time
import uuid
from dataclasses import dataclass, field
from ipaddress import IPv4Address
from typing import Any

from .errors import StateError


STATE_SCHEMA_VERSION = 1


@dataclass(frozen=True, slots=True)
class VmSpec:
    """Validated resource request supplied to the create command."""

    cpu: float
    ram_mib: int
    guest_ip: IPv4Address | None = None
    username: str = "root"
    password: str = field(default="root", repr=False)


@dataclass(slots=True)
class VmRecord:
    """Persistent record for one local Firecracker VM."""

    id: str
    cpu: float
    ram_mib: int
    guest_ip: str
    gateway: str
    tap: str
    mac: str
    socket: str
    config: str
    log: str
    disk: str = ""
    username: str = "root"
    password: str | None = field(default=None, repr=False)
    status: str = "starting"
    pid: int | None = None
    process_start_time: str | None = None
    created_at: int = field(default_factory=lambda: int(time.time()))
    updated_at: int = field(default_factory=lambda: int(time.time()))
    last_error: str | None = None

    def to_dict(self) -> dict[str, Any]:
        data: dict[str, Any] = {
            "id": self.id,
            "cpu": self.cpu,
            "ram_mib": self.ram_mib,
            "guest_ip": self.guest_ip,
            "gateway": self.gateway,
            "tap": self.tap,
            "mac": self.mac,
            "socket": self.socket,
            "config": self.config,
            "log": self.log,
            "disk": self.disk,
            "username": self.username,
            "password": self.password,
            "status": self.status,
            "created_at": self.created_at,
            "updated_at": self.updated_at,
        }
        if self.pid is not None:
            data["pid"] = self.pid
        if self.process_start_time is not None:
            data["process_start_time"] = self.process_start_time
        if self.last_error is not None:
            data["last_error"] = self.last_error
        return data

    @classmethod
    def from_dict(cls, value: object) -> "VmRecord":
        if not isinstance(value, dict):
            raise StateError("VM record is not an object")

        required = (
            "id",
            "cpu",
            "ram_mib",
            "guest_ip",
            "gateway",
            "tap",
            "mac",
            "socket",
            "config",
            "log",
        )
        missing = [name for name in required if name not in value]
        if missing:
            raise StateError(f"VM record is missing fields: {', '.join(missing)}")

        try:
            pid_value = value.get("pid")
            return cls(
                id=str(value["id"]),
                cpu=float(value["cpu"]),
                ram_mib=int(value["ram_mib"]),
                guest_ip=str(value["guest_ip"]),
                gateway=str(value["gateway"]),
                tap=str(value["tap"]),
                mac=str(value["mac"]),
                socket=str(value["socket"]),
                config=str(value["config"]),
                log=str(value["log"]),
                disk=str(value.get("disk", "")),
                username=str(value.get("username", "root")),
                password=(
                    str(value["password"])
                    if value.get("password") is not None
                    else None
                ),
                status=str(value.get("status", "unknown")),
                pid=int(pid_value) if pid_value is not None else None,
                process_start_time=(
                    str(value["process_start_time"])
                    if value.get("process_start_time") is not None
                    else None
                ),
                created_at=int(value.get("created_at", int(time.time()))),
                updated_at=int(value.get("updated_at", value.get("created_at", int(time.time())))),
                last_error=(
                    str(value["last_error"])
                    if value.get("last_error") is not None
                    else None
                ),
            )
        except (TypeError, ValueError) as error:
            raise StateError(f"invalid VM record: {error}") from error


@dataclass(slots=True)
class State:
    """Versioned JSON state document stored under the configured base path."""

    vms: dict[str, VmRecord] = field(default_factory=dict)
    next_mac_index: int = 1
    schema_version: int = STATE_SCHEMA_VERSION

    def to_dict(self) -> dict[str, Any]:
        return {
            "schema_version": self.schema_version,
            "next_mac_index": self.next_mac_index,
            "vms": {vm_id: vm.to_dict() for vm_id, vm in self.vms.items()},
        }

    @classmethod
    def from_dict(cls, value: object) -> "State":
        if not isinstance(value, dict):
            raise StateError("state file is not an object")

        try:
            schema_version = int(value.get("schema_version", 0))
        except (TypeError, ValueError) as error:
            raise StateError("state file has an invalid schema_version") from error
        if schema_version > STATE_SCHEMA_VERSION:
            raise StateError(
                f"state schema version {schema_version} is newer than this uvm version"
            )

        raw_vms = value.get("vms", {})
        if not isinstance(raw_vms, dict):
            raise StateError("state file has an invalid vms collection")

        vms = {str(vm_id): VmRecord.from_dict(vm) for vm_id, vm in raw_vms.items()}
        next_mac_index = value.get("next_mac_index")
        if next_mac_index is None:
            next_mac_index = _next_mac_index(vms.values())

        try:
            next_mac_index = max(1, int(next_mac_index))
        except (TypeError, ValueError) as error:
            raise StateError("state file has an invalid next_mac_index") from error

        return cls(
            schema_version=STATE_SCHEMA_VERSION,
            vms=vms,
            next_mac_index=next_mac_index,
        )


def new_vm_id() -> str:
    """Generate an opaque local ID that does not depend on PID or timestamp reuse."""

    return f"vm-{uuid.uuid4().hex}"


def _next_mac_index(vms: object) -> int:
    maximum = 0
    for vm in vms:
        if not isinstance(vm, VmRecord):
            continue
        try:
            parts = vm.mac.split(":")
            if parts[:3] != ["02", "fc", "00"] or len(parts) != 6:
                continue
            maximum = max(maximum, int("".join(parts[3:]), 16))
        except ValueError:
            continue
    return maximum + 1