Spaces:
Running
Running
| """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) | |
| 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 | |