osint-dashboard/app/bg_jobs.py

"""Background ffmpeg — never block a FastAPI request on a frame grab.

ffmpeg frame grabs are scheduled with asyncio.create_task and shared per URL.
"""

from __future__ import annotations

import asyncio
import shutil
from cachetools import TTLCache

_ffmpeg_cache: TTLCache = TTLCache(maxsize=100, ttl=300)
_ffmpeg_tasks: dict[str, asyncio.Task] = {}
_FFMPEG = shutil.which("ffmpeg")


def cached_ffmpeg_jpeg(url: str) -> bytes | None:
    return _ffmpeg_cache.get(url)


def schedule_ffmpeg_snapshot(url: str, timeout: float = 8.0) -> asyncio.Task:
    """Start (or reuse) an ffmpeg JPEG grab. Caller may await the task."""
    existing = _ffmpeg_tasks.get(url)
    if existing is not None and not existing.done():
        return existing
    task = asyncio.create_task(_ffmpeg_grab_and_cache(url, timeout))
    _ffmpeg_tasks[url] = task
    return task


async def _ffmpeg_grab(url: str, timeout: float = 8.0) -> bytes | None:
    """Grab one JPEG frame. Isolated so tests can stub it."""
    if not _FFMPEG or not url.lower().startswith("rtsp://"):
        return None
    cmd = [
        _FFMPEG, "-hide_banner", "-loglevel", "error", "-nostdin",
        "-rtsp_transport", "tcp",
        "-timeout", "4000000",
        "-i", url,
        "-frames:v", "1",
        "-f", "image2pipe", "-vcodec", "mjpeg",
        "pipe:1",
    ]
    try:
        proc = await asyncio.create_subprocess_exec(
            *cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.DEVNULL,
        )
    except FileNotFoundError:
        return None
    try:
        stdout, _ = await asyncio.wait_for(proc.communicate(), timeout=timeout)
    except asyncio.TimeoutError:
        proc.kill()
        try:
            await proc.wait()
        except Exception:  # noqa: BLE001
            pass
        return None
    if proc.returncode not in (0, None) or not stdout or len(stdout) < 64:
        return None
    if stdout[:2] != b"\xff\xd8":
        return None
    return stdout


async def _ffmpeg_grab_and_cache(url: str, timeout: float) -> bytes | None:
    data = await _ffmpeg_grab(url, timeout)
    if data:
        _ffmpeg_cache[url] = data
    return data