"""Manage optional NFStream flow enrichment workers and packet lookups.""" from __future__ import annotations import asyncio import json import logging import os import signal import subprocess import sys import threading import time from pathlib import Path from typing import Any, Dict, Iterable, Optional import src.shared_objects as shared_objects from src.config import settings from src.utilities.flow_identity import flow_cache_key, flow_key_from_endpoints, flow_key_from_packet, packet_observed_at_ms logger = logging.getLogger("nfstream_manager") def _merge_metadata(base: Optional[Dict[str, Any]], extra: Dict[str, Any]) -> Dict[str, Any]: merged = dict(base or {}) for key, value in extra.items(): if value not in (None, "", [], {}): merged[key] = value return merged def _safe_text(value: Any) -> Optional[str]: if value is None: return None text = str(value).strip() return text if text else None def _normalized_label(value: Any) -> Optional[str]: text = _safe_text(value) if text is None: return None if text.isdigit(): return None if text.lower() in {"unknown", "other", "unclassified", "none", "null"}: return None return text def _infer_is_encrypted(payload: Dict[str, Any]) -> Optional[bool]: application_name = str(_normalized_label(payload.get("application_name")) or "").upper() if any(token in application_name for token in ("TLS", "HTTPS", "QUIC", "SSL")): return True return None def _build_enrichment(payload: Dict[str, Any], iface: str, flow_key: str) -> Dict[str, Any]: app_protocol = _normalized_label(payload.get("application_name")) if app_protocol is None and any(payload.get(key) for key in ("user_agent", "content_type")): app_protocol = "HTTP" metadata = { "iface": iface, "flow_key": flow_key, "first_seen_ms": int(payload.get("first_seen_ms") or 0) or None, "last_seen_ms": int(payload.get("last_seen_ms") or 0) or None, "event": payload.get("event"), "application_name": _safe_text(payload.get("application_name")), "application_category_name": _safe_text(payload.get("application_category_name")), "application_confidence": _safe_text(payload.get("application_confidence")), "bidirectional_packets": payload.get("bidirectional_packets"), "bidirectional_bytes": payload.get("bidirectional_bytes"), } metadata = _merge_metadata( metadata, { "requested_server_name": payload.get("requested_server_name"), "client_fingerprint": payload.get("client_fingerprint"), "server_fingerprint": payload.get("server_fingerprint"), "user_agent": payload.get("user_agent"), "content_type": payload.get("content_type"), }, ) return { "app_protocol": app_protocol, "app_master_protocol": app_protocol, "app_category": _normalized_label(payload.get("application_category_name")), "app_confidence": _safe_text(payload.get("application_confidence")), "app_hostname": _safe_text(payload.get("requested_server_name")), "app_is_encrypted": _infer_is_encrypted(payload), "dpi_metadata": {"nfstream": metadata}, } def _has_useful_metadata(payload: Dict[str, Any]) -> bool: return any( payload.get(key) not in (None, "", [], {}) for key in ( "application_name", "application_category_name", "application_confidence", "requested_server_name", "client_fingerprint", "server_fingerprint", "user_agent", "content_type", ) ) class NFStreamManager: """Own long-lived NFStream subprocesses and a short-lived flow metadata cache.""" def __init__(self) -> None: self._enabled = settings.nfstream_enabled self._workers: Dict[str, Dict[str, Any]] = {} self._cache: Dict[tuple[str, str], Dict[str, Any]] = {} self._last_error_by_iface: Dict[str, str] = {} self._lock = threading.Lock() @property def enabled(self) -> bool: return self._enabled def update_interfaces(self, interfaces: Iterable[str]) -> None: """Align worker set with currently active interfaces.""" if not self._enabled: return targets = {iface.strip() for iface in interfaces if iface and iface.strip()} with self._lock: current = set(self._workers.keys()) for iface in sorted(current - targets): self._stop_worker_locked(iface) for iface in sorted(targets - current): self._start_worker_locked(iface) self._purge_cache_locked() def stop(self) -> None: """Stop all workers and clear transient state.""" with self._lock: for iface in list(self._workers.keys()): self._stop_worker_locked(iface) self._cache.clear() def lookup_packet(self, pkt: Any, iface: Optional[str]) -> Dict[str, Any]: """Return NFStream-derived enrichment for the packet if a recent flow is known.""" if not self._enabled or not iface: return {} flow_key = flow_key_from_packet(pkt) if flow_key is None: return {} now_ms = packet_observed_at_ms(pkt) with self._lock: self._purge_cache_locked(now_ms=now_ms) payload = self._cache.get(flow_cache_key(iface, flow_key)) if payload is None: return {} first_seen_ms = int(payload.get("first_seen_ms") or 0) last_seen_ms = int(payload.get("last_seen_ms") or 0) window_ms = settings.nfstream_lookup_window_ms if now_ms + window_ms < first_seen_ms or now_ms - window_ms > last_seen_ms: return {} return _build_enrichment(payload, iface, flow_key) def get_debug_snapshot(self) -> Dict[str, Any]: """Expose current worker/cache state for API debugging.""" with self._lock: self._purge_cache_locked() return { "enabled": self._enabled, "workers": { iface: { "running": bool(worker.get("process") and worker["process"].poll() is None), "thread_alive": bool(worker.get("reader") and worker["reader"].is_alive()), "last_error": self._last_error_by_iface.get(iface), } for iface, worker in self._workers.items() }, "cache_entries": len(self._cache), "last_error_by_iface": dict(self._last_error_by_iface), } def _start_worker_locked(self, iface: str) -> None: helper = Path(__file__).with_name("nfstream_flow_worker.py") python_bin = sys.executable or "python3" env = os.environ.copy() env["PYTHONUNBUFFERED"] = "1" backend_root = str(helper.parents[2]) existing_pythonpath = env.get("PYTHONPATH", "") env["PYTHONPATH"] = backend_root if not existing_pythonpath else f"{backend_root}:{existing_pythonpath}" cmd = [ python_bin, str(helper), "--iface", iface, "--idle-timeout", str(settings.nfstream_idle_timeout_seconds), "--active-timeout", str(settings.nfstream_active_timeout_seconds), "--snapshot-length", str(settings.nfstream_snapshot_length), "--n-dissections", str(settings.nfstream_n_dissections), "--n-meters", str(settings.nfstream_n_meters), ] if settings.nfstream_promiscuous_mode: cmd.append("--promiscuous-mode") logger.info("Starting NFStream worker for %s", iface) try: process = subprocess.Popen( cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1, env=env, start_new_session=True, ) except Exception: logger.exception("Failed to start NFStream worker for %s", iface) self._last_error_by_iface[iface] = "failed_to_start_process" return reader = threading.Thread( target=self._read_loop, args=(iface, process), daemon=True, name=f"nfstream-reader-{iface}", ) self._workers[iface] = { "process": process, "reader": reader, } reader.start() def _stop_worker_locked(self, iface: str) -> None: worker = self._workers.pop(iface, None) if worker is None: return process = worker.get("process") reader = worker.get("reader") if process is not None and process.poll() is None: try: os.killpg(os.getpgid(process.pid), signal.SIGTERM) process.wait(timeout=settings.nfstream_process_stop_timeout_seconds) except subprocess.TimeoutExpired: try: os.killpg(os.getpgid(process.pid), signal.SIGKILL) except ProcessLookupError: pass except Exception: logger.exception("Failed to stop NFStream worker for %s cleanly", iface) if reader is not None and reader.is_alive(): reader.join(timeout=settings.nfstream_reader_join_timeout_seconds) for key in [cache_key for cache_key in self._cache.keys() if cache_key[0] == iface]: self._cache.pop(key, None) def _read_loop(self, iface: str, process: subprocess.Popen[str]) -> None: stdout = process.stdout if stdout is None: return for line in stdout: text = line.strip() if not text: continue try: payload = json.loads(text) except json.JSONDecodeError: logger.info("nfstream[%s]: %s", iface, text) continue payload_type = payload.get("type") if payload_type == "worker_ready": logger.info("NFStream worker ready on %s", iface) continue if payload_type == "worker_error": message = str(payload.get("message") or "unknown error") logger.warning("NFStream worker error on %s: %s", iface, message) with self._lock: self._last_error_by_iface[iface] = message continue if payload_type != "flow_update": logger.debug("Ignoring NFStream payload on %s: %s", iface, payload) continue flow_key = flow_key_from_endpoints( payload.get("ip_version"), payload.get("protocol"), payload.get("src_ip"), payload.get("src_port"), payload.get("dst_ip"), payload.get("dst_port"), ) if flow_key is None: continue with self._lock: previous = self._cache.get(flow_cache_key(iface, flow_key)) self._cache[flow_cache_key(iface, flow_key)] = dict(payload) self._purge_cache_locked() if self._should_backfill(previous, payload): self._schedule_backfill(iface, flow_key, payload) rc = process.poll() if rc not in (0, None): logger.warning("NFStream worker for %s exited with code %s", iface, rc) def _purge_cache_locked(self, now_ms: Optional[int] = None) -> None: if now_ms is None: now_ms = int(time.time() * 1000) expiry_ms = int(settings.nfstream_cache_ttl_seconds * 1000) stale_keys = [] for key, payload in self._cache.items(): last_seen_ms = int(payload.get("last_seen_ms") or 0) if last_seen_ms and now_ms - last_seen_ms > expiry_ms: stale_keys.append(key) for key in stale_keys: self._cache.pop(key, None) def _should_backfill(self, previous: Optional[Dict[str, Any]], current: Dict[str, Any]) -> bool: if not _has_useful_metadata(current): return False if previous is None: return True interesting_keys = ( "application_name", "application_category_name", "application_confidence", "requested_server_name", "client_fingerprint", "server_fingerprint", "user_agent", "content_type", ) return any(previous.get(key) != current.get(key) for key in interesting_keys) def _schedule_backfill(self, iface: str, flow_key: str, payload: Dict[str, Any]) -> None: web_db = getattr(shared_objects, "db", None) web_loop = getattr(shared_objects, "web_loop", None) if web_db is None or web_loop is None: return enrichment = _build_enrichment(payload, iface, flow_key) async def _backfill() -> None: updated = await web_db.backfill_flow_metadata( iface=iface, src_ip=str(payload.get("src_ip")), dst_ip=str(payload.get("dst_ip")), src_port=int(payload.get("src_port") or 0), dst_port=int(payload.get("dst_port") or 0), protocol=int(payload.get("protocol") or 0), first_seen_ms=int(payload.get("first_seen_ms") or 0), last_seen_ms=int(payload.get("last_seen_ms") or 0), enrichment=enrichment, window_ms=settings.nfstream_lookup_window_ms, ) if updated: logger.debug("Backfilled NFStream metadata for %s packets on %s flow=%s", updated, iface, flow_key) try: future = asyncio.run_coroutine_threadsafe(_backfill(), web_loop) future.add_done_callback(self._log_backfill_result) except Exception: logger.exception("Failed to schedule NFStream metadata backfill for %s", flow_key) @staticmethod def _log_backfill_result(future: Any) -> None: try: future.result() except Exception: logger.exception("NFStream metadata backfill task failed") nfstream_manager = NFStreamManager()