stackpilot/backend/services/backup_service.py

"""Stack backup & restore, including named-volume contents.

A backup is a single ``.tar.gz`` with this layout::

    manifest.json            metadata + volume/bind inventory
    compose/...              the full stack directory (compose file, .env, ...)
    volumes/<full>.tar       raw contents of each compose-managed named volume

Named-volume contents are read/written through a throwaway helper container
(``BACKUP_HELPER_IMAGE``) with the volume bind-mounted — this is the portable
way to snapshot a volume regardless of its driver/mountpoint.
"""
from __future__ import annotations

import asyncio
import io
import json
import logging
import os
import shutil
import tarfile
import tempfile
from typing import Optional

from config import settings
from docker_client import DockerError, get_client, safe_call
from services import compose_service

logger = logging.getLogger("stackpilot.backup")

COMPOSE_PROJECT_LABEL = "com.docker.compose.project"
COMPOSE_VOLUME_LABEL = "com.docker.compose.volume"
MANIFEST_NAME = "manifest.json"
BACKUP_FORMAT_VERSION = 1


class BackupError(Exception):
    pass


# --------------------------------------------------------------------------- #
# Backup filename convention
# --------------------------------------------------------------------------- #


def backup_basename(stack_id: str, prefix: Optional[str] = None) -> str:
    """Filename stem used to group a stack's backups (and match for retention).

    ``prefix`` (e.g. a remote host slug) keeps backups of same-named stacks on
    different hosts from colliding / pruning each other on a shared destination.
    """
    return f"backup-{prefix}-{stack_id}" if prefix else f"backup-{stack_id}"


def backup_filename(stack_id: str, include_volumes: bool, prefix: Optional[str] = None) -> str:
    date = compose_service.now().strftime("%Y%m%d-%H%M%S")
    suffix = "full" if include_volumes else "config"
    return f"{backup_basename(stack_id, prefix)}-{suffix}-{date}.tar.gz"


# --------------------------------------------------------------------------- #
# Helper container for volume I/O
# --------------------------------------------------------------------------- #


def _ensure_helper_image(client) -> None:
    image = settings.BACKUP_HELPER_IMAGE
    try:
        safe_call(client.images.get, image)
    except DockerError:
        logger.info("Pulling backup helper image %s", image)
        safe_call(client.images.pull, image)


def _export_volume(full_name: str) -> bytes:
    client = get_client()
    _ensure_helper_image(client)
    container = safe_call(
        client.containers.create,
        settings.BACKUP_HELPER_IMAGE,
        command="true",
        volumes={full_name: {"bind": "/v", "mode": "ro"}},
    )
    try:
        # "/v/." copies the *contents* of the volume (no leading "v/" prefix),
        # so restore can extract straight back into the volume root.
        bits, _ = container.get_archive("/v/.")
        buf = io.BytesIO()
        for chunk in bits:
            buf.write(chunk)
        return buf.getvalue()
    finally:
        try:
            container.remove(force=True)
        except Exception:  # noqa: BLE001
            pass


def _restore_volume(full_name: str, labels: dict, tar_bytes: bytes) -> None:
    client = get_client()
    _ensure_helper_image(client)
    existed = True
    try:
        safe_call(client.volumes.get, full_name)
    except DockerError:
        existed = False
        safe_call(client.volumes.create, name=full_name, labels=labels or {})
    if existed:
        # Restore means "back to the snapshot": clear the current contents so
        # files created/kept since the backup don't survive underneath it.
        safe_call(
            client.containers.run,
            settings.BACKUP_HELPER_IMAGE,
            ["sh", "-c", "find /v -mindepth 1 -delete"],
            volumes={full_name: {"bind": "/v", "mode": "rw"}},
            remove=True,
        )
    container = safe_call(
        client.containers.create,
        settings.BACKUP_HELPER_IMAGE,
        command="true",
        volumes={full_name: {"bind": "/v", "mode": "rw"}},
    )
    try:
        container.put_archive("/v", tar_bytes)
    finally:
        try:
            container.remove(force=True)
        except Exception:  # noqa: BLE001
            pass


def _compose_volumes(stack_id: str) -> list[dict]:
    """Return [{full, short, labels}] for compose-managed named volumes."""
    try:
        client = get_client()
        vols = safe_call(
            client.volumes.list,
            filters={"label": f"{COMPOSE_PROJECT_LABEL}={stack_id}"},
        )
    except DockerError:
        return []
    out = []
    for v in vols:
        labels = v.attrs.get("Labels") or {}
        out.append(
            {
                "full": v.name,
                "short": labels.get(COMPOSE_VOLUME_LABEL, v.name),
                "labels": labels,
            }
        )
    return out


# --------------------------------------------------------------------------- #
# Backup
# --------------------------------------------------------------------------- #


async def create_backup(
    stack_id: str,
    name: str,
    include_volumes: bool = True,
    stop_first: bool = True,
) -> str:
    """Create a backup tar.gz and return its path on disk."""
    directory = compose_service.stack_dir(stack_id)
    if not os.path.isdir(directory):
        raise BackupError("Stack directory missing")

    volumes = _compose_volumes(stack_id) if include_volumes else []

    # For a consistent volume snapshot, stop the stack first — but only if it
    # is actually running, so backing up a stopped stack doesn't start it.
    stopped = False
    if include_volumes and stop_first and volumes:
        if compose_service.compute_status(stack_id) not in ("stopped", "unknown"):
            try:
                await compose_service.stop(stack_id)
                stopped = True
            except Exception as exc:  # noqa: BLE001
                logger.warning("Could not stop %s before backup: %s", stack_id, exc)

    try:
        manifest = {
            "format_version": BACKUP_FORMAT_VERSION,
            "stack_id": stack_id,
            "name": name,
            "created_at": compose_service.now().isoformat(),
            "include_volumes": include_volumes,
            "volumes": [{"full": v["full"], "short": v["short"], "labels": v["labels"]} for v in volumes],
        }

        tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".tar.gz")
        tmp.close()
        with tarfile.open(tmp.name, "w:gz") as tar:
            # manifest
            data = json.dumps(manifest, indent=2).encode("utf-8")
            info = tarfile.TarInfo(MANIFEST_NAME)
            info.size = len(data)
            tar.addfile(info, io.BytesIO(data))
            # stack directory
            tar.add(directory, arcname="compose")
            # volume contents
            for v in volumes:
                vbytes = await asyncio.to_thread(_export_volume, v["full"])
                info = tarfile.TarInfo(f"volumes/{v['full']}.tar")
                info.size = len(vbytes)
                tar.addfile(info, io.BytesIO(vbytes))
        return tmp.name
    finally:
        if stopped:
            try:
                await compose_service.up(stack_id)
            except Exception as exc:  # noqa: BLE001
                logger.warning("Could not restart %s after backup: %s", stack_id, exc)


# --------------------------------------------------------------------------- #
# Restore
# --------------------------------------------------------------------------- #


def read_manifest(tar_path: str) -> dict:
    with tarfile.open(tar_path, "r:gz") as tar:
        member = tar.getmember(MANIFEST_NAME)
        fh = tar.extractfile(member)
        if fh is None:
            raise BackupError("Backup is missing its manifest")
        return json.loads(fh.read().decode("utf-8"))


def _safe_extract_compose(tar: tarfile.TarFile, dest_dir: str) -> None:
    """Extract the ``compose/`` subtree into dest_dir, guarding path traversal."""
    os.makedirs(dest_dir, exist_ok=True)
    for member in tar.getmembers():
        if not member.name.startswith("compose/"):
            continue
        rel = member.name[len("compose/") :]
        if not rel:
            continue
        target = os.path.normpath(os.path.join(dest_dir, rel))
        if not target.startswith(os.path.abspath(dest_dir) + os.sep) and target != os.path.abspath(dest_dir):
            raise BackupError(f"Refusing unsafe path in backup: {member.name}")
        if member.isdir():
            os.makedirs(target, exist_ok=True)
        elif member.isreg():
            os.makedirs(os.path.dirname(target), exist_ok=True)
            src = tar.extractfile(member)
            if src is not None:
                with open(target, "wb") as out:
                    shutil.copyfileobj(src, out)


def restore_backup(
    tar_path: str,
    target_id: Optional[str] = None,
    overwrite: bool = False,
    restore_volumes: bool = True,
) -> dict:
    """Restore a backup. Returns {stack_id, name, volumes_restored}."""
    manifest = read_manifest(tar_path)
    raw_id = target_id or manifest.get("stack_id") or ""
    if not raw_id.strip():
        raise BackupError("Backup manifest has no stack id")
    # Slugify whichever id we end up using — the manifest comes from an
    # uploaded file, so its stack_id must never be able to escape STACKS_DIR.
    stack_id = compose_service.slugify(raw_id)

    directory = compose_service.stack_dir(stack_id)
    exists = os.path.isdir(directory)
    if exists and not overwrite:
        raise BackupError(f"Stack '{stack_id}' already exists")

    with tarfile.open(tar_path, "r:gz") as tar:
        if exists:
            shutil.rmtree(directory)
        _safe_extract_compose(tar, directory)

        volumes_restored = 0
        if restore_volumes:
            for v in manifest.get("volumes", []):
                member_name = f"volumes/{v['full']}.tar"
                try:
                    member = tar.getmember(member_name)
                except KeyError:
                    continue
                fh = tar.extractfile(member)
                if fh is None:
                    continue
                # Re-target volume labels to the (possibly new) stack id.
                labels = dict(v.get("labels") or {})
                labels[COMPOSE_PROJECT_LABEL] = stack_id
                full = v["full"]
                if target_id and manifest.get("stack_id") and full.startswith(manifest["stack_id"] + "_"):
                    full = stack_id + full[len(manifest["stack_id"]):]
                _restore_volume(full, labels, fh.read())
                volumes_restored += 1

    return {
        "stack_id": stack_id,
        "name": manifest.get("name", stack_id),
        "volumes_restored": volumes_restored,
    }