"""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