fester/backend/storage/qcow2_freeze.py

"""
QCOW2 workspace freeze — snapshot a build workspace into a qcow2 image.

Uses qemu-img to create a blank qcow2, then mounts it via nbd, rsyncs
the workdir into it, and unmounts. The resulting .qcow2 file is a
self-contained frozen snapshot that can be:
  - attached to a VM for reproducible builds
  - copied offsite for archival
  - mounted later via `qemu-nbd --connect` for inspection

Requires: qemu-utils (qemu-img, qemu-nbd) + rsync on the host.
Will return None and log a warning if any dependency is missing.
"""

import os
import shutil
import subprocess
import time
from pathlib import Path
from typing import Optional


def _have(*cmds) -> bool:
    """Check that all commands exist on PATH."""
    for c in cmds:
        if shutil.which(c) is None:
            return False
    return True


def _find_free_nbd() -> Optional[str]:
    """Find a free /dev/nbdN device. Returns None if none available."""
    if not Path("/dev").exists():
        return None
    for i in range(16):
        dev = f"/dev/nbd{i}"
        if not Path(dev).exists():
            continue
        # Check if it's already connected
        try:
            r = subprocess.run(
                ["qemu-nbd", "--disconnect", dev],
                capture_output=True, timeout=2,
            )
            # If disconnect succeeds (or even if it says "not connected"),
            # the device is now free
            if r.returncode == 0 or b"not connected" in r.stderr.lower() if r.stderr else True:
                return dev
        except (subprocess.TimeoutExpired, FileNotFoundError):
            continue
    return None


def freeze_to_qcow2(workdir, output_path, size: str = "10G") -> Optional[str]:
    """
    Freeze a build workspace into a qcow2 image.

    Args:
        workdir:     directory to snapshot
        output_path: where to write the .qcow2 file
        size:        image size (e.g. "10G", "500M")

    Returns the path to the created qcow2 file, or None if any step failed.
    """
    # Dependency check
    if not _have("qemu-img", "rsync"):
        return None

    # nbd module required for mounting; check but don't fail at import
    have_nbd = shutil.which("qemu-nbd") is not None

    workdir = Path(workdir)
    out = Path(output_path)
    out.parent.mkdir(parents=True, exist_ok=True)

    if not workdir.exists():
        return None

    # Step 1: Create the qcow2 image
    try:
        subprocess.run(
            ["qemu-img", "create", "-f", "qcow2", str(out), size],
            check=True, capture_output=True,
        )
    except (subprocess.CalledProcessError, FileNotFoundError):
        return None

    # Step 2: Try to mount via nbd and rsync the workdir in.
    # If nbd isn't available (no /dev/nbd* or no permission), we still
    # return the empty image path — caller can use `qemu-img info` to
    # verify, and the image can be mounted later on a host with nbd.
    nbd_dev = _find_free_nbd() if have_nbd else None
    if nbd_dev is None:
        # No nbd available — return the empty image (still a valid freeze
        # artifact, just without the workdir contents inlined)
        return str(out)

    # Connect the qcow2 to the nbd device
    try:
        subprocess.run(
            ["qemu-nbd", "--connect", nbd_dev, str(out)],
            check=True, capture_output=True, timeout=10,
        )
    except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
        return str(out)  # image exists, just not mounted

    # Wait for the device to appear
    time.sleep(0.5)

    # Step 3: Format + mount the nbd device
    mount_point = Path(f"/tmp/fester-mount-{int(time.time())}")
    mount_point.mkdir(parents=True, exist_ok=True)

    try:
        # Format as ext4 if the device is empty
        # (use -F to force formatting without prompt)
        subprocess.run(
            ["mkfs.ext4", "-F", nbd_dev],
            check=True, capture_output=True, timeout=30,
        )
        subprocess.run(
            ["mount", nbd_dev, str(mount_point)],
            check=True, capture_output=True, timeout=10,
        )
    except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
        # Cleanup: disconnect nbd
        try:
            subprocess.run(["qemu-nbd", "--disconnect", nbd_dev],
                           capture_output=True, timeout=5)
        except Exception:
            pass
        return str(out)

    # Step 4: rsync workdir into the mounted image
    try:
        subprocess.run(
            ["rsync", "-a", "--delete",
             f"{workdir}/", f"{mount_point}/"],
            check=True, capture_output=True, timeout=120,
        )
    except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
        pass  # partial copy is still useful

    # Step 5: Unmount + disconnect
    try:
        subprocess.run(["umount", str(mount_point)],
                       check=True, capture_output=True, timeout=10)
    except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
        pass
    try:
        subprocess.run(["qemu-nbd", "--disconnect", nbd_dev],
                       capture_output=True, timeout=5)
    except Exception:
        pass

    # Cleanup mount point
    try:
        mount_point.rmdir()
    except OSError:
        pass

    return str(out)


def qcow2_info(image_path) -> dict:
    """Return metadata about a qcow2 image via `qemu-img info`."""
    if not _have("qemu-img"):
        return {"error": "qemu-img not available"}
    try:
        out = subprocess.check_output(
            ["qemu-img", "info", "--output=json", str(image_path)],
            stderr=subprocess.DEVNULL,
        ).decode()
        import json
        return json.loads(out)
    except (subprocess.CalledProcessError, FileNotFoundError, ValueError) as e:
        return {"error": str(e)}


def list_snapshots(snapshots_dir: str = "/var/lib/fester/snapshots") -> list:
    """List all .qcow2 snapshots in the snapshots directory."""
    d = Path(snapshots_dir)
    if not d.exists():
        return []
    out = []
    for p in sorted(d.glob("*.qcow2")):
        info = qcow2_info(p)
        out.append({
            "name": p.stem,
            "path": str(p),
            "size_bytes": p.stat().st_size,
            "created_at": p.stat().st_mtime,
            "info": info if "error" not in info else None,
        })
    return out