test nfstream fix

backend/src/utilities/nfstream_manager.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import json
 import logging
 import os
@@ -13,6 +14,7 @@ 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
 
@@ -36,6 +38,54 @@ def _infer_is_encrypted(payload: Dict[str, Any]) -> Optional[bool]:
     return None
 
 
+def _build_enrichment(payload: Dict[str, Any], iface: str, flow_key: str) -> Dict[str, Any]:
+    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"),
+        "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 _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."""
 
@@ -92,36 +142,7 @@ class NFStreamManager:
             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},
-            }
+            return _build_enrichment(payload, iface, flow_key)
 
     def get_debug_snapshot(self) -> Dict[str, Any]:
         """Expose current worker/cache state for API debugging."""
@@ -259,9 +280,13 @@ class NFStreamManager:
                 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)
@@ -280,5 +305,60 @@ class NFStreamManager:
         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()