FabledCurator/backend/app/services/file_validator.py

"""Magic-byte head/tail validator for fresh gallery-dl writes.

Catches truncated downloads (servers lying about Content-Length, TCP
RST mid-transfer, etc.) before they enter the library. O(1) per file
— reads only head and tail bytes; never decodes pixels.

Fail-open: unknown formats and unreadable files pass through. Format
detection is by extension; we only validate what the downloader is
expected to write.
"""

from __future__ import annotations

import json
import logging
import shutil
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path

log = logging.getLogger(__name__)

JPEG_HEAD = b"\xff\xd8\xff"
JPEG_TAIL = b"\xff\xd9"

PNG_HEAD = b"\x89PNG\r\n\x1a\n"
PNG_TAIL = b"\x00\x00\x00\x00IEND\xaeB`\x82"

GIF_HEAD_87A = b"GIF87a"
GIF_HEAD_89A = b"GIF89a"
GIF_TAIL = b"\x3b"

WEBP_RIFF = b"RIFF"
WEBP_FORMAT = b"WEBP"


@dataclass(frozen=True)
class ValidationResult:
    ok: bool
    format: str | None = None
    reason: str | None = None
    size: int = 0


VALIDATED_EXTENSIONS = {".jpg", ".jpeg", ".png", ".gif", ".webp"}


def is_validatable(path: Path) -> bool:
    return path.suffix.lower() in VALIDATED_EXTENSIONS


def validate_file(path: Path) -> ValidationResult:
    """Returns ok=False if the file is structurally incomplete.

    Never raises. Unknown formats return ok=True with format=None.
    Unreadable files surface as ok=False so callers can quarantine.
    """
    ext = path.suffix.lower()
    if ext not in VALIDATED_EXTENSIONS:
        return ValidationResult(ok=True, format=None)

    try:
        size = path.stat().st_size
    except FileNotFoundError:
        return ValidationResult(ok=False, reason="file not found")
    except OSError as exc:
        return ValidationResult(ok=False, reason=f"stat failed: {exc}")

    if size < 16:
        return ValidationResult(ok=False, format=ext.lstrip("."), reason="too small", size=size)

    try:
        with path.open("rb") as fh:
            head = fh.read(32)
            fh.seek(max(0, size - 32))
            tail = fh.read(32)
    except OSError as exc:
        return ValidationResult(ok=False, reason=f"read failed: {exc}", size=size)

    if ext in (".jpg", ".jpeg"):
        if not head.startswith(JPEG_HEAD):
            return ValidationResult(ok=False, format="jpeg", reason="missing JPEG SOI marker", size=size)
        if not tail.endswith(JPEG_TAIL):
            return ValidationResult(ok=False, format="jpeg", reason="missing JPEG EOI marker", size=size)
        return ValidationResult(ok=True, format="jpeg", size=size)

    if ext == ".png":
        if not head.startswith(PNG_HEAD):
            return ValidationResult(ok=False, format="png", reason="missing PNG signature", size=size)
        if not tail.endswith(PNG_TAIL):
            return ValidationResult(ok=False, format="png", reason="missing PNG IEND chunk", size=size)
        return ValidationResult(ok=True, format="png", size=size)

    if ext == ".gif":
        if not (head.startswith(GIF_HEAD_87A) or head.startswith(GIF_HEAD_89A)):
            return ValidationResult(ok=False, format="gif", reason="missing GIF signature", size=size)
        if not tail.endswith(GIF_TAIL):
            return ValidationResult(ok=False, format="gif", reason="missing GIF trailer", size=size)
        return ValidationResult(ok=True, format="gif", size=size)

    if ext == ".webp":
        if not (head.startswith(WEBP_RIFF) and WEBP_FORMAT in head[:16]):
            return ValidationResult(ok=False, format="webp", reason="missing WebP RIFF/WEBP marker", size=size)
        try:
            chunk_size = int.from_bytes(head[4:8], "little")
        except Exception:
            return ValidationResult(ok=False, format="webp", reason="malformed WebP chunk size", size=size)
        if chunk_size + 8 != size:
            return ValidationResult(
                ok=False, format="webp",
                reason=f"WebP RIFF size mismatch: header says {chunk_size + 8}, file is {size}",
                size=size,
            )
        return ValidationResult(ok=True, format="webp", size=size)

    return ValidationResult(ok=True, format=None, size=size)


def quarantine_file(
    images_root: Path,
    path: Path,
    artist_slug: str,
    platform: str,
    *,
    url: str | None,
    result: ValidationResult,
) -> Path | None:
    """Move a validation-failed file to `_quarantine/<slug>/<platform>` and write
    a `.quarantine.json` provenance sidecar next to it.

    Returns the destination path, or None if the move itself failed (file left
    in place — the caller decides what to report). The caller has already run
    `validate_file` and seen `result.ok is False`. ONE implementation for both
    download backends — gallery-dl's batch post-process and the native ingester's
    per-media path — so the quarantine layout + provenance sidecar can't drift.
    """
    quarantine_root = images_root / "_quarantine" / artist_slug / platform
    try:
        quarantine_root.mkdir(parents=True, exist_ok=True)
        dest = quarantine_root / path.name
        counter = 1
        while dest.exists():
            dest = quarantine_root / f"{path.stem}.{counter}{path.suffix}"
            counter += 1
        shutil.move(str(path), str(dest))
        sidecar = dest.with_suffix(dest.suffix + ".quarantine.json")
        sidecar.write_text(
            json.dumps(
                {
                    "original_path": str(path),
                    "source_url": url,
                    "artist_slug": artist_slug,
                    "platform": platform,
                    "format": result.format,
                    "reason": result.reason,
                    "size": result.size,
                    "quarantined_at": datetime.now(UTC).isoformat(),
                },
                indent=2,
            )
        )
    except OSError as exc:
        log.error("Failed to quarantine %s: %s. File left in place.", path, exc)
        return None
    return dest