ModelForge-Public/backend/src/modelforge_api/services/recovery_paths.py

"""Allowlisted path and archive handling for backup and restore payloads.

Nothing in the recovery plane may resolve to a host path outside a configured root, follow a
symlink out of it, or extract an archive member that escapes its destination.
"""

from __future__ import annotations

import hashlib
import os
import re
import tarfile
from pathlib import Path, PurePosixPath

MAX_ARCHIVE_MEMBERS = 10_000
MAX_ARCHIVE_BYTES = 64 * 1024**3
_HASH_CHUNK = 1024 * 1024
_DRIVE_QUALIFIED = re.compile(r"^[A-Za-z]:")


class RecoveryPathError(Exception):
    """Raised when a path or archive member is not provably inside an allowlisted root."""

    def __init__(self, code: str, message: str) -> None:
        super().__init__(message)
        self.code = code
        self.message = message


def normalise_root(root: Path) -> Path:
    """Resolve an allowlisted root, creating it when missing so callers get a stable anchor."""

    resolved = Path(root).expanduser()
    if not resolved.is_absolute():
        raise RecoveryPathError("PATH_NOT_ALLOWED", "recovery roots must be absolute paths")
    resolved.mkdir(parents=True, exist_ok=True)
    return resolved.resolve(strict=True)


def _reject_traversal(relative: str) -> PurePosixPath:
    candidate = PurePosixPath(relative.replace("\\", "/"))
    if candidate.is_absolute() or not relative.strip():
        raise RecoveryPathError("PATH_NOT_ALLOWED", f"relative path required, got {relative!r}")
    if any(part in {"..", ""} for part in candidate.parts):
        raise RecoveryPathError("PATH_NOT_ALLOWED", f"path traversal rejected: {relative!r}")
    if any(part.startswith("~") for part in candidate.parts):
        raise RecoveryPathError("PATH_NOT_ALLOWED", f"home expansion rejected: {relative!r}")
    # A drive-qualified component means different things on Windows and POSIX; refuse it on both.
    if any(_DRIVE_QUALIFIED.match(part) for part in candidate.parts):
        raise RecoveryPathError("PATH_NOT_ALLOWED", f"drive-qualified path rejected: {relative!r}")
    return candidate


def resolve_within(root: Path, relative: str, *, must_exist: bool = False) -> Path:
    """Return ``root/relative`` proven to stay inside ``root`` after symlink resolution."""

    anchor = normalise_root(root)
    candidate = _reject_traversal(relative)
    target = anchor.joinpath(*candidate.parts)
    if must_exist and not target.exists():
        raise RecoveryPathError("PAYLOAD_MISSING", f"{relative} does not exist under the root")

    # Resolve the deepest component that exists, so a not-yet-created backup directory is still
    # checked against the allowlisted root and a symlinked ancestor cannot smuggle a path out.
    existing = target
    remaining: list[str] = []
    while not existing.exists():
        if existing == existing.parent:
            raise RecoveryPathError("PATH_NOT_ALLOWED", f"{relative} has no resolvable parent")
        remaining.append(existing.name)
        existing = existing.parent
    resolved = existing.resolve().joinpath(*reversed(remaining))
    if resolved != anchor and anchor not in resolved.parents:
        raise RecoveryPathError(
            "PATH_NOT_ALLOWED", f"{relative} resolves outside the allowlisted recovery root"
        )
    if target.is_symlink():
        raise RecoveryPathError("PATH_NOT_ALLOWED", f"{relative} is a symlink and is rejected")
    return target


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        while chunk := handle.read(_HASH_CHUNK):
            digest.update(chunk)
    return digest.hexdigest()


def sha256_bytes(payload: bytes) -> str:
    return hashlib.sha256(payload).hexdigest()


def free_bytes(path: Path) -> int | None:
    try:
        usage = os.statvfs(path) if hasattr(os, "statvfs") else None
    except OSError:
        return None
    if usage is not None:
        return int(usage.f_bavail) * int(usage.f_frsize)
    try:
        import shutil

        return int(shutil.disk_usage(path).free)
    except OSError:
        return None


def capacity_bytes(path: Path) -> int | None:
    try:
        import shutil

        return int(shutil.disk_usage(path).total)
    except OSError:
        return None


def validate_archive(archive: Path, destination_root: Path) -> list[str]:
    """Validate every member of a tar archive before any extraction happens.

    Bounded member count, bounded total size, no absolute members, no traversal, no symlinks,
    hardlinks, devices or FIFOs.
    """

    anchor = normalise_root(destination_root)
    names: list[str] = []
    total = 0
    with tarfile.open(archive, "r:*") as handle:
        for member in handle:
            if len(names) >= MAX_ARCHIVE_MEMBERS:
                raise RecoveryPathError(
                    "ARCHIVE_UNSAFE", f"archive exceeds {MAX_ARCHIVE_MEMBERS} members"
                )
            if member.issym() or member.islnk():
                raise RecoveryPathError(
                    "ARCHIVE_UNSAFE", f"archive member {member.name} is a link"
                )
            if member.isdev() or member.isfifo():
                raise RecoveryPathError(
                    "ARCHIVE_UNSAFE", f"archive member {member.name} is a special file"
                )
            if not (member.isfile() or member.isdir()):
                raise RecoveryPathError(
                    "ARCHIVE_UNSAFE", f"archive member {member.name} has an unsupported type"
                )
            try:
                resolved = resolve_within(anchor, member.name)
            except RecoveryPathError as error:
                raise RecoveryPathError(
                    "ARCHIVE_UNSAFE", f"archive member {member.name} escapes destination"
                ) from error
            if anchor not in resolved.parents and resolved != anchor:
                raise RecoveryPathError(
                    "ARCHIVE_UNSAFE", f"archive member {member.name} escapes destination"
                )
            total += max(0, member.size)
            if total > MAX_ARCHIVE_BYTES:
                raise RecoveryPathError("ARCHIVE_UNSAFE", "archive exceeds the extraction budget")
            names.append(member.name)
    return names


def extract_archive(archive: Path, destination_root: Path) -> list[str]:
    """Validate first, extract second; a failed validation never writes a byte."""

    names = validate_archive(archive, destination_root)
    anchor = normalise_root(destination_root)
    with tarfile.open(archive, "r:*") as handle:
        handle.extractall(anchor, filter="data")
    return names