104 lines
3 KiB
Python
104 lines
3 KiB
Python
|
"""Background masscan / ffmpeg — never block a FastAPI request on a scan.
|
||
|
|
|
||
|
|
masscan is capped at 200 pps (home uplink saturates at 1k+). ffmpeg frame
|
||
|
|
grabs are scheduled with asyncio.create_task and shared per URL.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import logging
|
||
|
|
import shutil
|
||
|
|
from cachetools import TTLCache
|
||
|
|
|
||
|
|
logger = logging.getLogger("osint.bg_jobs")
|
||
|
|
|
||
|
|
MASSCAN_PPS_CAP = 200
|
||
|
|
|
||
|
|
_masscan_task: asyncio.Task | None = None
|
||
|
|
_ffmpeg_cache: TTLCache = TTLCache(maxsize=100, ttl=300)
|
||
|
|
_ffmpeg_tasks: dict[str, asyncio.Task] = {}
|
||
|
|
_FFMPEG = shutil.which("ffmpeg")
|
||
|
|
|
||
|
|
|
||
|
|
def schedule_masscan_pass() -> bool:
|
||
|
|
"""Kick one capped masscan pass. Returns False if a pass is already running."""
|
||
|
|
global _masscan_task
|
||
|
|
if _masscan_task is not None and not _masscan_task.done():
|
||
|
|
return False
|
||
|
|
_masscan_task = asyncio.create_task(_run_masscan_capped())
|
||
|
|
return True
|
||
|
|
|
||
|
|
|
||
|
|
async def _run_masscan_capped() -> None:
|
||
|
|
import masscan_config as cfg
|
||
|
|
from run_masscan_service import _verify_excludefile, run_pass
|
||
|
|
|
||
|
|
orig = cfg.MASSCAN_RATE
|
||
|
|
if orig > MASSCAN_PPS_CAP:
|
||
|
|
logger.warning("capping masscan rate %s pps -> %s", orig, MASSCAN_PPS_CAP)
|
||
|
|
cfg.MASSCAN_RATE = MASSCAN_PPS_CAP
|
||
|
|
try:
|
||
|
|
_verify_excludefile()
|
||
|
|
await run_pass()
|
||
|
|
finally:
|
||
|
|
cfg.MASSCAN_RATE = orig
|
||
|
|
|
||
|
|
|
||
|
|
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
|