279 lines
10 KiB
Python
279 lines
10 KiB
Python
"""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, ValidationError
|
|
|
|
|
|
STATE_SCHEMA_VERSION = 2
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class PortRule:
|
|
"""One private guest TCP port; HTTP/HTTPS are descriptive service labels."""
|
|
|
|
port: int
|
|
protocol: str = "TCP"
|
|
|
|
def __post_init__(self) -> None:
|
|
if type(self.port) is not int or not 1 <= self.port <= 65535:
|
|
raise ValidationError("port must be an integer between 1 and 65535")
|
|
if not isinstance(self.protocol, str) or self.protocol.upper() not in {"TCP", "HTTP", "HTTPS"}:
|
|
raise ValidationError("port type must be TCP, HTTP, or HTTPS")
|
|
object.__setattr__(self, "protocol", self.protocol.upper())
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
return {"port": self.port, "protocol": self.protocol}
|
|
|
|
@classmethod
|
|
def from_dict(cls, value: object) -> "PortRule":
|
|
if not isinstance(value, dict) or set(value) - {"port", "protocol"}:
|
|
raise StateError("invalid port rule")
|
|
try:
|
|
return cls(port=value["port"], protocol=value.get("protocol", "TCP"))
|
|
except (KeyError, ValidationError) as error:
|
|
raise StateError(f"invalid port rule: {error}") from error
|
|
|
|
|
|
@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)
|
|
name: str | None = None
|
|
ports: tuple[PortRule, ...] = (PortRule(22),)
|
|
|
|
|
|
@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
|
|
name: str | None = None
|
|
ports: tuple[PortRule, ...] | None = None
|
|
disk_ready: bool = False
|
|
runtime_owned: bool = False
|
|
tap_owned: bool = False
|
|
operation: str | None = None
|
|
boot_id: str | None = None
|
|
spawn_pending: bool = False
|
|
network: str | None = None
|
|
bridge: 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,
|
|
"name": self.name,
|
|
"ports": [port.to_dict() for port in self.ports] if self.ports is not None else None,
|
|
"disk_ready": self.disk_ready,
|
|
"runtime_owned": self.runtime_owned,
|
|
"tap_owned": self.tap_owned,
|
|
"operation": self.operation,
|
|
"boot_id": self.boot_id,
|
|
"spawn_pending": self.spawn_pending,
|
|
"network": self.network,
|
|
"bridge": self.bridge,
|
|
}
|
|
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")
|
|
raw_ports = value.get("ports")
|
|
if raw_ports is not None and not isinstance(raw_ports, list):
|
|
raise StateError("VM ports must be a list or null for a legacy VM")
|
|
ports = tuple(PortRule.from_dict(port) for port in raw_ports) if raw_ports is not None else None
|
|
if ports is not None and len({port.port for port in ports}) != len(ports):
|
|
raise StateError("VM has duplicate TCP port rules")
|
|
legacy_ready = bool(value.get("disk")) and value.get("status") in {"running", "stopped"}
|
|
for name in ("disk_ready", "runtime_owned", "tap_owned", "spawn_pending"):
|
|
if name in value and type(value[name]) is not bool:
|
|
raise StateError(f"VM {name} must be a boolean")
|
|
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
|
|
),
|
|
name=str(value["name"]) if value.get("name") is not None else None,
|
|
ports=ports,
|
|
disk_ready=value.get("disk_ready", legacy_ready),
|
|
runtime_owned=value.get("runtime_owned", True),
|
|
tap_owned=value.get("tap_owned", bool(value.get("pid"))),
|
|
operation=str(value["operation"]) if value.get("operation") is not None else None,
|
|
boot_id=str(value["boot_id"]) if value.get("boot_id") is not None else None,
|
|
spawn_pending=value.get("spawn_pending", value.get("operation") == "spawning"),
|
|
network=str(value["network"]) if value.get("network") is not None else None,
|
|
bridge=str(value["bridge"]) if value.get("bridge") 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()}
|
|
if any(vm_id != vm.id for vm_id, vm in vms.items()):
|
|
raise StateError("VM record ID does not match its registry key")
|
|
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
|