111 lines
4.0 KiB
Python
111 lines
4.0 KiB
Python
"""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
|
|
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
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)
|