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:`unpack_archive` also restores each regular file's member mtime
  (the zip DOS ``date_time`` or the tar ``mtime``) — phase 106, D2:
  uploaded archives keep their file dates, which the importer reads as
  the document creation date. Directories, symlinks, and hardlinks are
  untouched; the cap and every safety check above are unchanged (the
  ``utime`` sits after a successful ``_write_capped``, so a failed
  unpack still removes the partial tree).
* :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 datetime import UTC, datetime
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)
                # Phase 106 (D2): restore the member's DOS mtime — a
                # tz-agnostic epoch value, UTC-rendered exactly like an
                # mtime — so uploaded archives keep their file dates.
                # Regular files only; an OSError here still removes the
                # partial tree in unpack_archive like any write failure.
                # ``ns=`` takes INTEGER nanoseconds (a float seconds
                # value raises TypeError), so convert explicitly.
                mtime = datetime(*member.date_time, tzinfo=UTC).timestamp()
                mtime_ns = int(mtime * 1_000_000_000)
                os.utime(dest, ns=(mtime_ns, mtime_ns))


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)
                # Phase 106 (D2): restore the member's mtime (epoch
                # seconds — a tz-agnostic value) so uploaded archives
                # keep their file dates. Regular files only; an OSError
                # here still removes the partial tree in unpack_archive.
                # ``ns=`` takes INTEGER nanoseconds (a float seconds
                # value raises TypeError), so convert explicitly.
                mtime_ns = int(member.mtime * 1_000_000_000)
                os.utime(dest, ns=(mtime_ns, mtime_ns))
            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)