imagerepo/app/services/integrity.py

"""Structural-integrity verification for media files.

Public surface: `verify_path(path) -> (status, detail)`.
Status values match the `image_record.integrity_status` column:
  - 'ok': passed all checks
  - 'truncated': structurally a valid file up to a point, then missing
                 trailing bytes (the WD14 '6 bytes not processed' case)
  - 'unreadable': can't open at all, or container is malformed enough that
                  even the head doesn't parse
  - 'missing': filepath doesn't exist on disk

Dispatch is by file extension. Images go through PIL (with truncated-image
loading explicitly disabled so we surface the very thing we want to detect).
Videos go through ffprobe. Archives go through zipfile.testzip / tarfile
where the stdlib supports it. Anything else returns 'ok' (we don't know
how to validate it, so we don't lie about its state).
"""
from __future__ import annotations

import logging
import os
import subprocess
import zipfile

# We deliberately import the PIL.ImageFile module here without flipping
# LOAD_TRUNCATED_IMAGES — opposite of wd14.py's runtime tolerance, since
# this module's job is to *detect* truncation, not paper over it.
from PIL import Image, UnidentifiedImageError

log = logging.getLogger(__name__)


IMAGE_EXTS = ('.jpg', '.jpeg', '.jfif', '.png', '.gif', '.bmp', '.tiff', '.webp')
VIDEO_EXTS = ('.mp4', '.mov', '.avi', '.mkv', '.webm', '.m4v', '.wmv', '.flv')
ZIP_EXTS = ('.zip', '.cbr', '.cbz')


def verify_path(path: str) -> tuple[str, str | None]:
    """Verify a single file. Returns (status, detail-or-None).

    Detail is a short human-readable string for the failure case, useful
    for the report endpoint and logs. None when status='ok'.
    """
    if not os.path.exists(path):
        return 'missing', None
    try:
        size = os.path.getsize(path)
    except OSError as e:
        return 'missing', f'stat failed: {e}'
    if size == 0:
        return 'unreadable', 'zero-byte file'

    ext = os.path.splitext(path)[1].lower()

    if ext in IMAGE_EXTS:
        return _verify_image(path)
    if ext in VIDEO_EXTS:
        return _verify_video(path)
    if ext in ZIP_EXTS:
        return _verify_zip(path)
    # Unknown extension: don't fail it, but don't lie that we verified it.
    # Caller writes 'ok' but the kind=unknown case is rare in this library.
    return 'ok', None


def _verify_image(path: str) -> tuple[str, str | None]:
    """PIL-based two-stage check.

    Stage 1: open + verify(). Validates header, structure, and (for most
    formats) walks the entire stream to confirm framing — this is what
    catches the 6-byte truncation. verify() invalidates the Image after,
    so callers can't draw from it (we don't need to).

    Stage 2: re-open + load(). Some PIL formats (notably JPEG) only flag
    truncation during load(), not verify(). Re-opening with truncated
    loading disabled lets a malformed scan trigger OSError here.
    """
    try:
        with Image.open(path) as img:
            img.verify()
    except (OSError, SyntaxError, UnidentifiedImageError, ValueError) as e:
        msg = str(e).lower()
        if _looks_truncated(msg):
            return 'truncated', str(e)
        return 'unreadable', str(e)

    try:
        # Force a full decode without the LOAD_TRUNCATED_IMAGES safety net
        # that wd14/siglip enabled at module-import time.
        from PIL import ImageFile
        prev = ImageFile.LOAD_TRUNCATED_IMAGES
        ImageFile.LOAD_TRUNCATED_IMAGES = False
        try:
            with Image.open(path) as img:
                img.load()
        finally:
            ImageFile.LOAD_TRUNCATED_IMAGES = prev
    except (OSError, SyntaxError, ValueError) as e:
        msg = str(e).lower()
        if _looks_truncated(msg):
            return 'truncated', str(e)
        return 'unreadable', str(e)

    return 'ok', None


# Phrases PIL emits across formats when the file ends earlier than the
# decoder expected. Centralized so both verify-stages agree on the mapping.
_TRUNCATED_HINTS = (
    'trunc',           # 'image file is truncated', 'truncated file read'
    'not processed',   # 'N bytes not processed' (the original wd14 case)
    'broken png',      # PNG short final chunk
    'premature',       # JPEG/PNG premature end
    'unexpected end',  # GIF / generic end-of-stream
    'not enough',      # 'not enough data'
)


def _looks_truncated(lower_msg: str) -> bool:
    return any(h in lower_msg for h in _TRUNCATED_HINTS)


def _verify_video(path: str) -> tuple[str, str | None]:
    """ffprobe-based container walk.

    `-v error` keeps stdout/stderr quiet on success and emits the first
    structural complaint on failure. Any non-zero exit is treated as a
    failed verification. ffprobe distinguishes 'truncated' poorly across
    formats, so we lean on stderr text to map it.
    """
    try:
        result = subprocess.run(
            ['ffprobe', '-v', 'error', '-show_format', '-show_streams', '-i', path],
            capture_output=True,
            text=True,
            timeout=30,
        )
    except FileNotFoundError:
        log.warning("ffprobe not on PATH; skipping video verify for %s", path)
        return 'ok', None
    except subprocess.TimeoutExpired:
        return 'unreadable', 'ffprobe timed out'

    if result.returncode == 0:
        return 'ok', None
    err = (result.stderr or '').strip().lower()
    if any(s in err for s in ('truncated', 'partial', 'eof', 'end of file', 'invalid data found')):
        return 'truncated', result.stderr.strip()
    return 'unreadable', result.stderr.strip() or f'ffprobe exit {result.returncode}'


def _verify_zip(path: str) -> tuple[str, str | None]:
    """zipfile.testzip walks every CRC. None = clean; bad-name = corrupt.

    BadZipFile / LargeZipFile are unreadable cases (header malformed or
    too big to test), distinct from a CRC mismatch on a single member
    which we treat as truncated/partial.
    """
    try:
        with zipfile.ZipFile(path) as zf:
            bad = zf.testzip()
    except zipfile.BadZipFile as e:
        return 'unreadable', str(e)
    except zipfile.LargeZipFile as e:
        return 'unreadable', str(e)
    except OSError as e:
        return 'unreadable', str(e)

    if bad is None:
        return 'ok', None
    return 'truncated', f'CRC failed for member: {bad}'