nest/uvm/state.py

"""Locked, atomic persistence for the local JSON VM inventory."""

from __future__ import annotations

import fcntl
import json
import os
import re
import tempfile
from collections.abc import Iterator
from contextlib import contextmanager
from pathlib import Path

from .config import Settings
from .domain import State, VmRecord
from .errors import ConflictError, NotFoundError, StateError
from .system import ensure_data_directories


class StateStore:
    """Own the JSON state file so callers cannot race IP and VM allocation."""

    def __init__(self, settings: Settings) -> None:
        self._settings = settings

    def initialize(self) -> None:
        ensure_data_directories(self._settings)
        with self._locked():
            if not self._settings.state_path.exists():
                self._save_unlocked(State())
            else:
                try:
                    os.chmod(self._settings.state_path, 0o600)
                except OSError as error:
                    raise StateError(
                        f"cannot update permissions on {self._settings.state_path}: {error}"
                    ) from error

    def load(self) -> State:
        # Reads do not create /var/lib/uvm, so `uvm list` remains usable before install.
        # Atomic replacement makes an unlocked reader see either the old or new full document.
        return self._load_unlocked()

    @contextmanager
    def transaction(self) -> Iterator[State]:
        """Load and commit one state mutation while holding an exclusive lock."""

        ensure_data_directories(self._settings)
        with self._locked():
            state = self._load_unlocked()
            yield state
            self._save_unlocked(state)

    @contextmanager
    def operation_lock(self) -> Iterator[None]:
        """Serialize external VM lifecycle work across concurrent CLI processes."""

        ensure_data_directories(self._settings)
        with self._locked_path(self._settings.operation_lock_path):
            yield

    def get(self, vm_id: str) -> VmRecord:
        state = self.load()
        try:
            return state.vms[vm_id]
        except KeyError as error:
            raise NotFoundError(f"VM not found: {vm_id}") from error

    def find_by_id_or_ip(self, identifier: str) -> VmRecord:
        return self.resolve(self.load(), identifier)

    @staticmethod
    def resolve(state: State, identifier: str) -> VmRecord:
        if re.fullmatch(r"vm-[0-9a-f]{32}", identifier):
            try:
                return state.vms[identifier]
            except KeyError as error:
                raise NotFoundError(f"VM not found: {identifier}") from error
        matches = [vm for vm in state.vms.values() if identifier in {vm.id, vm.name, vm.guest_ip}]
        if len(matches) > 1:
            raise ConflictError(f"ambiguous VM identifier: {identifier}; use an exact VM ID")
        if not matches:
            raise NotFoundError(f"VM not found: {identifier}")
        return matches[0]

    @contextmanager
    def _locked(self) -> Iterator[None]:
        with self._locked_path(self._settings.state_lock_path):
            yield

    @contextmanager
    def _locked_path(self, lock_path: Path) -> Iterator[None]:
        lock_path.parent.mkdir(parents=True, exist_ok=True)
        with lock_path.open("a+") as lock_file:
            fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
            try:
                yield
            finally:
                fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)

    def _load_unlocked(self) -> State:
        path = self._settings.state_path
        if not path.exists():
            return State()
        try:
            if path.stat().st_mode & 0o077:
                os.chmod(path, 0o600)
            with path.open(encoding="utf-8") as state_file:
                raw = json.load(state_file)
        except PermissionError as error:
            raise StateError(
                f"cannot access protected VM registry {path}. Run this command with sudo."
            ) from error
        except (OSError, json.JSONDecodeError) as error:
            raise StateError(f"cannot read {path}: {error}") from error
        return State.from_dict(raw)

    def _save_unlocked(self, state: State) -> None:
        path = self._settings.state_path
        path.parent.mkdir(parents=True, exist_ok=True)
        fd, temporary_path = tempfile.mkstemp(
            prefix=f".{path.name}.",
            suffix=".tmp",
            dir=path.parent,
        )
        try:
            with os.fdopen(fd, "w", encoding="utf-8") as temporary_file:
                json.dump(state.to_dict(), temporary_file, indent=2, sort_keys=True)
                temporary_file.write("\n")
                temporary_file.flush()
                os.fsync(temporary_file.fileno())
            os.replace(temporary_path, path)
            os.chmod(path, 0o600)
            self._fsync_parent(path)
        except OSError as error:
            raise StateError(f"cannot write {path}: {error}") from error
        finally:
            try:
                os.unlink(temporary_path)
            except FileNotFoundError:
                pass

    @staticmethod
    def _fsync_parent(path: Path) -> None:
        directory = os.open(str(path.parent), os.O_RDONLY)
        try:
            os.fsync(directory)
        finally:
            os.close(directory)