Blame for representative_frames.py

"""Prepare private representative frames for publication as clean JPEGs."""

from __future__ import annotations

import hashlib
import io
import json
import os
from contextlib import contextmanager
from pathlib import Path
import re
import stat
import tempfile
import warnings

from PIL import Image, __version__ as _PILLOW_VERSION


_KEY = re.compile(r"[0-9a-f]{64}\.jpg\Z")
_DEFAULT_SOURCE_ROOT = Path.home() / ".local/share/nds-annotations"
_MAX_SOURCE_BYTES = 32 * 1024 * 1024
_MAX_SOURCE_PIXELS = 32_000_000
_MAX_OUTPUT_BYTES = 4 * 1024 * 1024
_ENCODING_VERSION = f"rgb-jpeg-960-q76-v1-pillow-{_PILLOW_VERSION}"
_PUBLIC_ROOT = (
    "https://raw.githubusercontent.com/AutoZYX-Labs/driveresearch-data-proxy/"
    "statistics-data/published/frames/"
)


def cache_root() -> Path:
    """Return the private derivative-cache directory."""
    return Path.home() / ".local/state/nds-statistics/representative-frames"


def _destination(override: Path | None) -> Path:
    return Path(override) if override is not None else cache_root()


def _source_stat(info: os.stat_result) -> tuple[int, ...]:
    return (info.st_dev, info.st_ino, info.st_size, info.st_mtime_ns, info.st_ctime_ns)


def _fingerprint(info: os.stat_result) -> str:
    encoded = json.dumps([_ENCODING_VERSION, *_source_stat(info)]).encode("ascii")
    return hashlib.sha256(encoded).hexdigest()


def _same_source(source: Path, root: Path, resolved: Path, initial: os.stat_result) -> bool:
    current = source.resolve(strict=True)
    if current != resolved:
        return False
    with _open_regular(current, root) as stream:
        return _source_stat(os.fstat(stream.fileno())) == _source_stat(initial)


def _cached_key(destination: Path, fingerprint: str) -> str | None:
    manifest = destination / "sources" / (fingerprint + ".json")
    try:
        with _open_regular(manifest.absolute(), destination.absolute()) as source:
            if not 0 < os.fstat(source.fileno()).st_size <= 1024:
                return None
            raw = source.read(1025)
        data = json.loads(raw)
        if not isinstance(data, dict) or data.get("encoding") != _ENCODING_VERSION:
            return None
        key = data.get("key")
        output = frame_file(key, cache_root=destination)
        if output is None or type(data.get("size")) is not int:
            return None
        if output.stat().st_size != data["size"]:
            return None
        return key
    except (OSError, ValueError, TypeError):
        return None


def _private_directory(path: Path) -> None:
    if path.is_symlink():
        raise ValueError("invalid cache directory")
    if not path.exists():
        if not path.parent.exists():
            _private_directory(path.parent)
        path.mkdir(mode=0o700, exist_ok=True)
    path.chmod(0o700)


@contextmanager
def _open_regular(path: Path, root: Path):
    """Open without following any directory/file symlink after containment checks."""
    parts = path.relative_to(root).parts
    if not parts:
        raise ValueError("invalid image")
    directory = os.open(root, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW)
    try:
        for part in parts[:-1]:
            next_directory = os.open(
                part, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW, dir_fd=directory
            )
            os.close(directory)
            directory = next_directory
        descriptor = os.open(
            parts[-1], os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK, dir_fd=directory
        )
        with os.fdopen(descriptor, "rb") as source:
            if not stat.S_ISREG(os.fstat(source.fileno()).st_mode):
                raise ValueError("invalid image")
            yield source
    finally:
        os.close(directory)


def _atomic_write(path: Path, content: bytes) -> None:
    _private_directory(path.parent)
    descriptor, temporary = tempfile.mkstemp(prefix=".frame-", dir=path.parent)
    try:
        with os.fdopen(descriptor, "wb") as output:
            output.write(content)
            output.flush()
            os.fsync(output.fileno())
        os.replace(temporary, path)
    finally:
        if os.path.exists(temporary):
            os.unlink(temporary)


def prepare_frame(
    source: str | Path,
    *,
    cache_root: Path | None = None,
    source_root: Path | None = None,
) -> str | None:
    """Return the content-addressed key of a metadata-free frame derivative."""
    try:
        destination = _destination(cache_root)
        root = Path(source_root if source_root is not None else _DEFAULT_SOURCE_ROOT).resolve(strict=True)
        original_path = Path(source)
        if not original_path.is_absolute():
            original_path = root / original_path
        candidate = original_path.resolve(strict=True)
        candidate.relative_to(root)
        with _open_regular(candidate, root) as source_file:
            initial = os.fstat(source_file.fileno())
            size = initial.st_size
            if not 0 < size <= _MAX_SOURCE_BYTES:
                return None
            fingerprint = _fingerprint(initial)
            cached = _cached_key(destination, fingerprint)
            if cached is not None:
                return cached if _same_source(original_path, root, candidate, initial) else None
            raw = source_file.read(_MAX_SOURCE_BYTES + 1)
            if len(raw) != size or _source_stat(os.fstat(source_file.fileno())) != _source_stat(initial):
                return None
        with warnings.catch_warnings():
            warnings.simplefilter("error")
            with Image.open(io.BytesIO(raw)) as original:
                if original.width * original.height > _MAX_SOURCE_PIXELS:
                    return None
                original.load()
                converted = original.convert("RGB")
                converted.thumbnail((960, 960), Image.Resampling.LANCZOS)
                clean = Image.new("RGB", converted.size)
                clean.paste(converted)
                encoded = io.BytesIO()
                clean.save(encoded, format="JPEG", quality=76)
        content = encoded.getvalue()
        key = hashlib.sha256(content).hexdigest() + ".jpg"
        if not _same_source(original_path, root, candidate, initial):
            return None
        _atomic_write(destination / key, content)
        manifest = json.dumps(
            {"encoding": _ENCODING_VERSION, "key": key, "size": len(content)},
            separators=(",", ":"),
        ).encode("ascii")
        _atomic_write(destination / "sources" / (fingerprint + ".json"), manifest)
        if not _same_source(original_path, root, candidate, initial):
            return None
        return key
    except (OSError, ValueError, TypeError, SyntaxError, RuntimeError, OverflowError, Warning, Image.DecompressionBombError):
        return None


def frame_file(key: str, *, cache_root: Path | None = None) -> Path | None:
    """Return an intact derivative for a valid public frame key."""
    if not isinstance(key, str) or not _KEY.fullmatch(key):
        return None
    try:
        destination = _destination(cache_root)
        candidate = destination / key
        with _open_regular(candidate.absolute(), destination.absolute()) as output:
            size = os.fstat(output.fileno()).st_size
            if not 0 < size <= _MAX_OUTPUT_BYTES:
                return None
            content = output.read(_MAX_OUTPUT_BYTES + 1)
        if len(content) != size:
            return None
    except (OSError, ValueError):
        return None
    if hashlib.sha256(content).hexdigest() + ".jpg" != key:
        return None
    return candidate


def public_frame_url(key: str) -> str:
    """Build a public URL using only a strict content-derived JPEG key."""
    if not isinstance(key, str) or not _KEY.fullmatch(key):
        raise ValueError("invalid representative frame key")
    return _PUBLIC_ROOT + key