brain-of-reese/app/rag/archive_upload.py

"""Safe archive unpacking for uploaded sources (phase 49, task 01).

Pure file-system logic — no FastAPI/DB imports. The API layer
(``POST /api/git-sources/upload``, phase 49, task 02) calls these and
maps :class:`ArchiveUploadError` to status codes (422).

The guarantees (phase 49 locked decisions):

* :func:`archive_source_name` derives the source name from the uploaded
  filename — ONE trailing archive suffix stripped, longest-first so
  ``homelab.tar.gz`` yields ``homelab`` (never ``homelab.tar``) — and
  rejects anything that would not be a safe single folder name (the name
  becomes a directory under ``BOR_UPLOAD_DIR`` and the ``git_sources``
  row's ``path``).
* :func:`unpack_archive` extracts ``.zip`` or ``.tar`` (gz/bz2/xz
  transparently) into a fresh directory, rejecting absolute member
  paths, ``..`` traversal, symlink/hardlink targets that escape the
  unpack directory, and device/FIFO members — and counting every
  extracted byte against a cap (zip-bomb guard). Any failure removes the
  partial ``target_dir`` so no half-unpacked tree survives.
* :func:`swap_in` makes ``new_dir`` become ``final_dir`` with **no
  missing window**: the previous folder is renamed to a unique
  same-filesystem ``.old-`` sibling first, the new folder is renamed
  into place, then the sibling is deleted. A failed swap restores the
  previous folder (best effort) and removes ``new_dir``.
"""
from __future__ import annotations

import os
import re
import shutil
import stat
import tarfile
import uuid
import zipfile
from pathlib import Path, PurePosixPath
from typing import IO

#: Accepted archive suffixes, LONGEST FIRST — ``.tar.gz`` must match
#: before ``.tar`` would (phase 49 locked decision: these four formats
#: only).
ARCHIVE_SUFFIXES: tuple[str, ...] = (".tar.gz", ".tgz", ".zip", ".tar")

#: A client-supplied filename never legitimately contains a path
#: separator — reject rather than strip (defense in depth; the upload
#: also carries a browsed-path, which the browser normalizes).
_SEPARATOR_RE = re.compile(r"[/\\]")

#: Windows drive-letter prefix (``C:``) — an absolute member path in a
#: cross-platform zip.
_DRIVE_RE = re.compile(r"^[A-Za-z]:")

#: Streaming read size while counting extracted bytes.
_CHUNK_SIZE = 1 << 20


class ArchiveUploadError(Exception):
    """A user-safe archive/unpack error (the API maps it to a status).

    Messages name the problem — and the cap where relevant — never the
    archive's content and never any path beyond the owner's own upload
    directory.
    """


def archive_source_name(filename: str) -> str:
    """The source name for an uploaded archive filename.

    The filename is used exactly as sent — a client can never
    legitimately embed a ``/`` or ``\\``, so a separator is *rejected*
    (the defensive basename step below is then a no-op, kept for the
    contract). ONE trailing archive suffix from :data:`ARCHIVE_SUFFIXES`
    is removed, matched case-insensitively and longest-first:
    ``homelab.tar.gz`` → ``homelab``, ``a.tar.gz`` → ``a`` (a single
    compound strip, not a double one), ``a.zip.zip`` → ``a.zip``. The
    stem keeps its original case — the name becomes a folder name on a
    Linux filesystem (case-sensitive).

    Raises:
        ArchiveUploadError: the filename is empty, contains a path
            separator, or the stripped stem is empty / ``.`` / ``..`` /
            contains control characters (the API maps to 422).
    """
    if not filename:
        raise ArchiveUploadError("empty file name")
    if _SEPARATOR_RE.search(filename):
        raise ArchiveUploadError("file name contains a path separator")
    # No separator above, so the basename is the name itself — kept
    # explicit so the "take the basename" contract lives in one place.
    base = filename.rsplit("/", 1)[-1].rsplit("\\", 1)[-1]
    lowered = base.lower()
    for suffix in ARCHIVE_SUFFIXES:  # longest first
        if lowered.endswith(suffix):
            base = base[: -len(suffix)]
            break
        if lowered == suffix.lstrip("."):  # bare suffix ("tar.gz") — no stem
            base = ""
            break
    if not base or base in (".", ".."):
        raise ArchiveUploadError("archive file name has no usable source name")
    if any(ord(c) < 0x20 or ord(c) == 0x7F for c in base):
        raise ArchiveUploadError("archive file name contains control characters")
    return base


def _member_dest(name: str, target: Path) -> Path:
    """The destination path for one archive member, or raise.

    Rejects empty names, absolute names (POSIX ``/``, Windows ``\\`` or
    drive letters), and any ``..`` path component (zip names are
    ``/``-separated; backslash forms are normalized before the check).
    The final path is resolved against the target (following any symlink
    an earlier member may have created) and must stay inside ``target``.
    """
    if not name:
        raise ArchiveUploadError("archive member with an empty name")
    if name.startswith(("/", "\\")) or _DRIVE_RE.match(name):
        raise ArchiveUploadError("archive member with an absolute path")
    if ".." in PurePosixPath(name).parts or ".." in name.replace("\\", "/").split("/"):
        raise ArchiveUploadError("archive member path traversal")
    dest = target / name
    target_resolved = target.resolve()
    dest_resolved = dest.resolve()
    if dest_resolved != target_resolved and target_resolved not in dest_resolved.parents:
        raise ArchiveUploadError("archive member path escapes the unpack directory")
    return dest


def _link_target_resolved(linkname: str, link_dir: Path, target: Path) -> Path:
    """Resolve a symlink/hardlink target against its member's directory.

    The resolved target must stay inside ``target`` — anything else
    (absolute targets, ``..`` climbs out) is rejected. Returns the
    resolved path (used directly for hardlinks).
    """
    if not linkname:
        raise ArchiveUploadError("archive link with an empty target")
    # An absolute linkname wins over the join (pathlib semantics) and is
    # then caught by the containment check below.
    resolved = (link_dir / linkname).resolve()
    target_resolved = target.resolve()
    if resolved != target_resolved and target_resolved not in resolved.parents:
        raise ArchiveUploadError("archive link target escapes the unpack directory")
    return resolved


def _write_capped(src: IO[bytes], dst: Path, max_extract_bytes: int, total: list[int]) -> None:
    """Stream ``src`` to ``dst``, counting into ``total[0]``.

    Raises as soon as the cumulative extracted bytes EXCEED
    ``max_extract_bytes`` (the cap itself is exactly reachable). The
    error names the cap, not the archive content.
    """
    dst.parent.mkdir(parents=True, exist_ok=True)
    with open(dst, "wb") as out:
        while chunk := src.read(_CHUNK_SIZE):
            total[0] += len(chunk)
            if total[0] > max_extract_bytes:
                raise ArchiveUploadError(
                    f"archive exceeds the {max_extract_bytes}-byte extraction cap"
                )
            out.write(chunk)


def _unpack_zip(archive: Path, target: Path, max_extract_bytes: int) -> None:
    total: list[int] = [0]
    with zipfile.ZipFile(archive) as zf:
        for member in zf.infolist():
            # The high 16 bits of external_attr are the Unix mode when
            # present. Modes may be 0 (Windows-made zips), bare permission
            # bits (CPython ``writestr``: 0o600), or a full mode with the
            # file-type bits — only the latter can prove a member is a
            # symlink/device/FIFO, and only those are rejected; entries
            # without type bits are decided by the member name.
            mode = member.external_attr >> 16
            if stat.S_ISLNK(mode):
                raise ArchiveUploadError("zip archives with symlink entries are not allowed")
            if mode & 0o170000 and not (stat.S_ISREG(mode) or stat.S_ISDIR(mode)):
                raise ArchiveUploadError("zip archives with non-regular entries are not allowed")
            dest = _member_dest(member.filename, target)
            if member.filename.endswith("/") or (mode and stat.S_ISDIR(mode)):
                dest.mkdir(parents=True, exist_ok=True)
            else:
                with zf.open(member) as src:
                    _write_capped(src, dest, max_extract_bytes, total)


def _unpack_tar(archive: Path, target: Path, max_extract_bytes: int) -> None:
    # ``r:*`` auto-detects plain/gz/bz2/xz compression.
    total: list[int] = [0]
    with tarfile.open(archive, mode="r:*") as tf:
        for member in tf.getmembers():
            dest = _member_dest(member.name, target)
            if member.issym():
                _link_target_resolved(member.linkname, dest.parent, target)
                dest.parent.mkdir(parents=True, exist_ok=True)
                os.symlink(member.linkname, dest)
            elif member.islnk():
                resolved = _link_target_resolved(member.linkname, dest.parent, target)
                dest.parent.mkdir(parents=True, exist_ok=True)
                os.link(resolved, dest)
            elif member.isdir():
                dest.mkdir(parents=True, exist_ok=True)
            elif member.isreg():
                src = tf.extractfile(member)
                if src is None:
                    raise ArchiveUploadError("corrupt archive member")
                _write_capped(src, dest, max_extract_bytes, total)
            else:  # char/block device, FIFO
                raise ArchiveUploadError(
                    "tar archives with device or FIFO members are not allowed"
                )


def unpack_archive(archive: Path, target_dir: Path, max_extract_bytes: int) -> None:
    """Extract ``archive`` into ``target_dir`` (created empty here).

    ``target_dir`` must NOT exist yet — the caller passes a fresh unique
    path (a temp sibling of the final folder). On ANY failure — bad or
    corrupt archive, unsafe member, cap exceeded, OS error — the partial
    ``target_dir`` is removed so no half-unpacked tree survives, and the
    failure is raised as :class:`ArchiveUploadError` (the module's only
    public exception type).
    """
    if target_dir.exists() or target_dir.is_symlink():
        raise ArchiveUploadError("the unpack target already exists")
    target_dir.mkdir(parents=True)
    try:
        # Content-sniff the container: a .zip that is really a tar (or a
        # truncated file) falls through to the tar reader and fails
        # loudly there instead of half-extracting.
        if zipfile.is_zipfile(archive):
            _unpack_zip(archive, target_dir, max_extract_bytes)
        else:
            _unpack_tar(archive, target_dir, max_extract_bytes)
    except Exception as exc:
        shutil.rmtree(target_dir, ignore_errors=True)
        if isinstance(exc, ArchiveUploadError):
            raise
        raise ArchiveUploadError("could not unpack the archive") from exc


def swap_in(new_dir: Path, final_dir: Path) -> None:
    """Atomically make ``new_dir`` become ``final_dir`` (no missing window).

    If ``final_dir`` exists it is first renamed to a unique same-
    filesystem sibling ``<name>.old-<hex>``, then ``new_dir`` is renamed
    into place, and the ``.old-`` sibling is deleted (re-uploads replace
    the previous content in place — one folder, no stale files). If it
    does not exist, this is a plain rename.

    On a rename failure the previous folder is restored (best effort) and
    ``new_dir`` removed, then the failure is raised as
    :class:`ArchiveUploadError` — a failed upload never leaves the
    previous folder, row, or KB in a mixed state.
    """
    old_dir: Path | None = None
    try:
        if final_dir.exists():
            old_dir = final_dir.with_name(final_dir.name + ".old-" + uuid.uuid4().hex)
            os.rename(final_dir, old_dir)
        os.rename(new_dir, final_dir)
    except OSError as exc:
        if old_dir is not None:
            try:
                os.rename(old_dir, final_dir)
            except OSError:
                shutil.rmtree(old_dir, ignore_errors=True)
        shutil.rmtree(new_dir, ignore_errors=True)
        raise ArchiveUploadError("could not replace the previous folder") from exc
    if old_dir is not None:
        shutil.rmtree(old_dir)