pxe-server/webapp/services/image_registry.py

"""Mutable registry for image types.

Previously IMAGE_TYPES + FRIENDLY_NAMES were hardcoded in config.py and
adding a new image type required a code edit + Ansible re-run. Now they
live in a JSON file (REGISTRY_PATH) that the webapp can read and write at
runtime.

On import (called from app.py startup), reload() rewrites
config.IMAGE_TYPES and config.FRIENDLY_NAMES in-place so the rest of the
code (and the base.html sidebar context) sees the live list without
threading the registry through every call.

Schema:
    {
      "image_types": [
        {"key": "gea-standard", "friendly_name": "GE Aerospace Standard"},
        ...
      ]
    }

The seed comes from whatever IMAGE_TYPES/FRIENDLY_NAMES are baked into
config.py at first run. Once the registry file exists, it wins.
"""
from __future__ import annotations

import json
import os
import re
import shutil
import tempfile
from datetime import datetime
from typing import Optional

import config

# Persist alongside the Samba share so an Ansible re-run that recreates
# SAMBA_SHARE keeps the registry. Override via env for tests.
REGISTRY_PATH = os.environ.get(
    "IMAGE_REGISTRY_PATH",
    os.path.join(config.SAMBA_SHARE, "image-registry.json"),
)

# image_type keys are used as directory names + URL path components. Match
# the existing baked-in convention (lowercase + hyphens + alnum).
_VALID_KEY = re.compile(r"^[a-z][a-z0-9-]{1,63}$")


class RegistryError(Exception):
    pass


def _read_file() -> dict:
    try:
        with open(REGISTRY_PATH, "r", encoding="utf-8") as f:
            return json.load(f)
    except (OSError, json.JSONDecodeError):
        return {}


def _write_file(data: dict):
    os.makedirs(os.path.dirname(REGISTRY_PATH), exist_ok=True)
    fd, tmp = tempfile.mkstemp(
        dir=os.path.dirname(REGISTRY_PATH),
        prefix=".tmp-registry-",
        suffix=".json",
    )
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=2, ensure_ascii=False)
            f.write("\n")
        os.replace(tmp, REGISTRY_PATH)
    except Exception:
        try:
            os.unlink(tmp)
        except OSError:
            pass
        raise


def _seed_from_config() -> dict:
    """First-run bootstrap: build a registry dict from whatever's baked into
    config.py. Preserves order so the sidebar layout doesn't shuffle on the
    first reload."""
    items = []
    for key in config.IMAGE_TYPES:
        items.append({
            "key": key,
            "friendly_name": config.FRIENDLY_NAMES.get(key, key),
            "created_at": None,
        })
    return {"image_types": items}


def load_registry() -> dict:
    """Return the registry dict. Bootstraps from config.py on first read."""
    data = _read_file()
    if not data.get("image_types"):
        data = _seed_from_config()
        try:
            _write_file(data)
        except OSError:
            # Read-only filesystem / no perms; fall back to in-memory seed.
            pass
    return data


def reload():
    """Refresh config.IMAGE_TYPES + config.FRIENDLY_NAMES from the registry
    file. Called at startup and after every CRUD op. Mutates in place so
    existing references stay valid."""
    data = load_registry()
    keys = [item["key"] for item in data.get("image_types", []) if "key" in item]
    friendly = {item["key"]: item.get("friendly_name", item["key"])
                for item in data.get("image_types", [])}
    # In-place replace so module-level references still point at the same
    # list/dict objects.
    config.IMAGE_TYPES[:] = keys
    config.FRIENDLY_NAMES.clear()
    config.FRIENDLY_NAMES.update(friendly)


def _validate_key(key: str):
    if not key or not _VALID_KEY.match(key):
        raise RegistryError(
            f"Invalid image_type key '{key}'. Must be lowercase alphanumeric "
            f"+ hyphens, 2-64 chars, starting with a letter."
        )


def create(key: str, friendly_name: str) -> dict:
    """Add a new image type with no on-disk content. Caller is responsible
    for populating Deploy/ via the import flow afterwards."""
    _validate_key(key)
    if not friendly_name:
        friendly_name = key
    data = load_registry()
    existing_keys = {item["key"] for item in data["image_types"]}
    if key in existing_keys:
        raise RegistryError(f"image_type '{key}' already exists")
    data["image_types"].append({
        "key": key,
        "friendly_name": friendly_name,
        "created_at": datetime.now().astimezone().isoformat(timespec="seconds"),
    })
    _write_file(data)
    reload()
    return {"key": key, "friendly_name": friendly_name}


def clone(src_key: str, dst_key: str, friendly_name: Optional[str] = None) -> dict:
    """Duplicate a source image type's on-disk content + register the new
    key. Copies Deploy/ tree (which contains Control/ JSON configs +
    unattend XML), Tools/ if present, and the top-level Sources/ symlink if
    present. Shared dirs (Out-of-box Drivers, Operating Systems, Packages)
    are re-symlinked into the new image root, not duplicated."""
    _validate_key(dst_key)
    data = load_registry()
    keys = {item["key"] for item in data["image_types"]}
    if src_key not in keys:
        raise RegistryError(f"source image_type '{src_key}' not registered")
    if dst_key in keys:
        raise RegistryError(f"destination image_type '{dst_key}' already exists")

    from services import fs  # local import to avoid circular at module load
    src_root = fs.image_root(src_key)
    dst_root = fs.image_root(dst_key)
    if os.path.exists(dst_root):
        raise RegistryError(f"destination path already exists on disk: {dst_root}")

    os.makedirs(dst_root, exist_ok=True)
    if os.path.isdir(src_root):
        for entry in os.listdir(src_root):
            src_item = os.path.join(src_root, entry)
            dst_item = os.path.join(dst_root, entry)
            if os.path.islink(src_item):
                # Preserve symlink (shared dir) instead of dereferencing.
                target = os.readlink(src_item)
                os.symlink(target, dst_item)
            elif os.path.isdir(src_item):
                shutil.copytree(src_item, dst_item, symlinks=True)
            else:
                shutil.copy2(src_item, dst_item)

    if not friendly_name:
        src_friendly = next(
            (i["friendly_name"] for i in data["image_types"] if i["key"] == src_key),
            src_key,
        )
        friendly_name = f"{src_friendly} (copy)"

    data["image_types"].append({
        "key": dst_key,
        "friendly_name": friendly_name,
        "created_at": datetime.now().astimezone().isoformat(timespec="seconds"),
        "cloned_from": src_key,
    })
    _write_file(data)
    reload()
    return {"key": dst_key, "friendly_name": friendly_name, "cloned_from": src_key}


def delete(key: str, delete_content: bool = False) -> dict:
    """Remove key from the registry. When delete_content is True, also wipes
    the on-disk Deploy/Tools/etc tree. Shared symlinked dirs are NOT
    followed when deleting content (we unlink the symlink, not its target)."""
    data = load_registry()
    items = data.get("image_types", [])
    found = [i for i in items if i["key"] == key]
    if not found:
        raise RegistryError(f"image_type '{key}' not registered")
    data["image_types"] = [i for i in items if i["key"] != key]
    _write_file(data)
    reload()

    removed_content = False
    if delete_content:
        from services import fs
        root = fs.image_root(key)
        if os.path.isdir(root):
            # Walk top-level: unlink symlinks (don't follow into shared
            # dirs), rmtree real directories, unlink files.
            for entry in os.listdir(root):
                p = os.path.join(root, entry)
                try:
                    if os.path.islink(p):
                        os.unlink(p)
                    elif os.path.isdir(p):
                        shutil.rmtree(p)
                    else:
                        os.unlink(p)
                except OSError:
                    pass
            try:
                os.rmdir(root)
            except OSError:
                pass
            removed_content = True

    return {"key": key, "removed_content": removed_content}


def rename_friendly(key: str, friendly_name: str) -> dict:
    """Update the human-readable name without touching the key or on-disk
    content. Sidebar reflects on next request."""
    if not friendly_name:
        raise RegistryError("friendly_name required")
    data = load_registry()
    for item in data.get("image_types", []):
        if item["key"] == key:
            item["friendly_name"] = friendly_name
            _write_file(data)
            reload()
            return {"key": key, "friendly_name": friendly_name}
    raise RegistryError(f"image_type '{key}' not registered")