try nfstream

backend/src/utilities/nfstream_manager.py

@@ -0,0 +1,284 @@
+"""Manage optional NFStream flow enrichment workers and packet lookups."""
+
+from __future__ import annotations
+
+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
+
+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 _infer_is_encrypted(payload: Dict[str, Any]) -> Optional[bool]:
+    application_name = str(payload.get("application_name") or "").upper()
+    if any(token in application_name for token in ("TLS", "HTTPS", "QUIC", "SSL")):
+        return True
+    if payload.get("requested_server_name") or payload.get("client_fingerprint") or payload.get("server_fingerprint"):
+        return True
+    return None
+
+
+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 {}
+
+            metadata = {
+                "iface": iface,
+                "flow_key": flow_key,
+                "first_seen_ms": first_seen_ms or None,
+                "last_seen_ms": last_seen_ms or None,
+                "event": payload.get("event"),
+                "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": payload.get("application_name"),
+                "app_master_protocol": payload.get("application_name"),
+                "app_category": payload.get("application_category_name"),
+                "app_confidence": payload.get("application_confidence"),
+                "app_hostname": payload.get("requested_server_name"),
+                "app_is_encrypted": _infer_is_encrypted(payload),
+                "dpi_metadata": {"nfstream": metadata},
+            }
+
+    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),
+        ]
+        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,
+            )
+        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:
+                process.send_signal(signal.SIGTERM)
+                process.wait(timeout=settings.nfstream_process_stop_timeout_seconds)
+            except subprocess.TimeoutExpired:
+                process.kill()
+            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:
+                self._cache[flow_cache_key(iface, flow_key)] = dict(payload)
+                self._purge_cache_locked()
+
+        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)
+
+
+nfstream_manager = NFStreamManager()