"""Detect partial/truncated downloads via cheap magic-byte checks. Real-world trigger: a Patreon JPEG landed in the library missing its 2-byte EOI marker (`FF D9`). PIL.ImageFile.load downstream raised "image file is truncated (6 bytes not processed)". gallery-dl already streams to .part and renames atomically, but TCP RST mid-transfer or servers lying about Content-Length can still produce structurally valid-but-incomplete files. This validator runs at write time (and on the existing library via a sweep task) and catches that class of failure. It only reads the head and tail of each file — O(1) per file regardless of size — and is format-aware for the formats we actually download. Unknown formats (metadata sidecars, archives, etc.) pass through. """ from __future__ import annotations from dataclasses import dataclass from pathlib import Path from typing import Optional # Per-format magic-byte signatures. Each format defines what the head # and tail must look like for the file to be structurally complete. # We don't decode pixel data here — we just confirm the framing bytes # the encoder is required to write. Adding a new format means one entry. JPEG_HEAD = b"\xff\xd8\xff" JPEG_TAIL = b"\xff\xd9" PNG_HEAD = b"\x89PNG\r\n\x1a\n" # Last 12 bytes of any complete PNG: 0-length IEND chunk + CRC. 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: """Outcome of validating one file. `ok=True` means either the file passed format-specific checks OR the format isn't one we know how to validate (we don't quarantine on ignorance). `format` is the detected/inferred format string for logging and stats; `reason` is populated on failure with a short operator-readable explanation. """ ok: bool format: Optional[str] = None reason: Optional[str] = None size: int = 0 # Extensions we attempt to validate. Anything not in this set returns # ok=True with format=None — gallery-dl writes JSON sidecars, ugoira # zips, etc. that we don't want to police here. VALIDATED_EXTENSIONS = {".jpg", ".jpeg", ".png", ".gif", ".webp"} def is_validatable(path: Path) -> bool: """Return True if the path's extension is one we validate.""" return path.suffix.lower() in VALIDATED_EXTENSIONS def validate_file(path: Path) -> ValidationResult: """Check whether a file on disk is a complete, well-framed image. Returns ok=True for unknown formats and unreadable/missing files are surfaced as ok=False so callers can quarantine + investigate. """ try: size = path.stat().st_size except FileNotFoundError: return ValidationResult(ok=False, reason="file not found") except OSError as e: return ValidationResult(ok=False, reason=f"stat failed: {e}") if size == 0: return ValidationResult(ok=False, format="empty", reason="zero-byte file", size=0) suffix = path.suffix.lower() if suffix not in VALIDATED_EXTENSIONS: return ValidationResult(ok=True, format=None, size=size) try: head, tail = _read_head_tail(path, head_n=16, tail_n=16) except OSError as e: return ValidationResult(ok=False, reason=f"read failed: {e}", size=size) if suffix in (".jpg", ".jpeg"): return _check_jpeg(head, tail, size) if suffix == ".png": return _check_png(head, tail, size) if suffix == ".gif": return _check_gif(head, tail, size) if suffix == ".webp": return _check_webp(head, tail, size) # Unreachable given VALIDATED_EXTENSIONS guard above, but keeps the # fall-through safe if the set diverges from the dispatch. return ValidationResult(ok=True, format=None, size=size) def _read_head_tail(path: Path, head_n: int, tail_n: int) -> tuple[bytes, bytes]: """Read the first head_n and last tail_n bytes of a file in one open.""" with path.open("rb") as f: head = f.read(head_n) f.seek(0, 2) # SEEK_END size = f.tell() f.seek(max(0, size - tail_n)) tail = f.read(tail_n) return head, tail def _check_jpeg(head: bytes, tail: bytes, size: int) -> ValidationResult: if not head.startswith(JPEG_HEAD): return ValidationResult( ok=False, format="jpeg", reason="missing JPEG SOI marker (FF D8 FF)", size=size ) if not tail.endswith(JPEG_TAIL): # The exact failure mode the user hit: file truncated before EOI. return ValidationResult( ok=False, format="jpeg", reason="missing JPEG EOI marker (FF D9)", size=size ) return ValidationResult(ok=True, format="jpeg", size=size) def _check_png(head: bytes, tail: bytes, size: int) -> ValidationResult: 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) def _check_gif(head: bytes, tail: bytes, size: int) -> ValidationResult: if not (head.startswith(GIF_HEAD_87A) or head.startswith(GIF_HEAD_89A)): return ValidationResult( ok=False, format="gif", reason="missing GIF header", size=size ) if not tail.endswith(GIF_TAIL): return ValidationResult( ok=False, format="gif", reason="missing GIF trailer (3B)", size=size ) return ValidationResult(ok=True, format="gif", size=size) def _check_webp(head: bytes, tail: bytes, size: int) -> ValidationResult: # WEBP is a RIFF container: "RIFF" + 4-byte little-endian size + "WEBP". # The container size field equals (file size - 8). If the file got # truncated, the declared size won't match the on-disk size. if len(head) < 12 or not head.startswith(WEBP_RIFF) or head[8:12] != WEBP_FORMAT: return ValidationResult( ok=False, format="webp", reason="missing RIFF/WEBP header", size=size ) declared = int.from_bytes(head[4:8], "little") if declared != size - 8: return ValidationResult( ok=False, format="webp", reason=f"RIFF size mismatch (declared {declared}, actual {size - 8})", size=size, ) return ValidationResult(ok=True, format="webp", size=size)