test nfstream fix

backend/src/utilities/database.py

@@ -4,6 +4,7 @@ import asyncio
 import base64
 import json
 import logging
+from datetime import datetime, timezone
 from typing import Any, Dict, List, Optional
 
 import asyncpg
@@ -266,6 +267,101 @@ class DatabasePool:
             except Exception:
                 logger.exception("Failed to publish pkt_info to broadcaster")
 
+    async def backfill_flow_metadata(
+        self,
+        *,
+        iface: str,
+        src_ip: str,
+        dst_ip: str,
+        src_port: int,
+        dst_port: int,
+        protocol: int,
+        first_seen_ms: int,
+        last_seen_ms: int,
+        enrichment: Dict[str, Any],
+        window_ms: int,
+    ) -> int:
+        """Update recent packet rows for a flow after enrichment arrives asynchronously."""
+        if self._pool is None:
+            await self.init_pool()
+
+        lower_bound = datetime.fromtimestamp(max(first_seen_ms - window_ms, 0) / 1000.0, tz=timezone.utc)
+        upper_bound = datetime.fromtimestamp(max(last_seen_ms + window_ms, 0) / 1000.0, tz=timezone.utc)
+        dpi_metadata = enrichment.get("dpi_metadata")
+
+        try:
+            async with self._pool.acquire() as conn:
+                rows = await conn.fetch(
+                    """
+                    UPDATE packets
+                    SET
+                        updated_at = NOW(),
+                        app_protocol = COALESCE(packets.app_protocol, $9),
+                        app_master_protocol = COALESCE(packets.app_master_protocol, $10),
+                        app_category = COALESCE(packets.app_category, $11),
+                        app_confidence = COALESCE(packets.app_confidence, $12),
+                        app_hostname = COALESCE(packets.app_hostname, $13),
+                        app_is_encrypted = COALESCE(packets.app_is_encrypted, $14),
+                        dpi_metadata = CASE
+                            WHEN $15::jsonb IS NULL THEN packets.dpi_metadata
+                            WHEN packets.dpi_metadata IS NULL THEN $15::jsonb
+                            ELSE packets.dpi_metadata || $15::jsonb
+                        END
+                    WHERE
+                        ip_proto_raw = $1
+                        AND (capture_iface = $2 OR ingress_if = $2 OR egress_if = $2)
+                        AND timestamp BETWEEN $7 AND $8
+                        AND (
+                            (src_ip = $3::inet AND dst_ip = $4::inet AND src_port = $5 AND dst_port = $6)
+                            OR
+                            (src_ip = $4::inet AND dst_ip = $3::inet AND src_port = $6 AND dst_port = $5)
+                        )
+                        AND (
+                            packets.app_protocol IS NULL
+                            OR packets.app_master_protocol IS NULL
+                            OR packets.app_category IS NULL
+                            OR packets.app_confidence IS NULL
+                            OR packets.app_hostname IS NULL
+                            OR packets.app_is_encrypted IS NULL
+                            OR ($15::jsonb IS NOT NULL)
+                        )
+                    RETURNING *
+                    """,
+                    protocol,
+                    iface,
+                    src_ip,
+                    dst_ip,
+                    src_port,
+                    dst_port,
+                    lower_bound,
+                    upper_bound,
+                    enrichment.get("app_protocol"),
+                    enrichment.get("app_master_protocol"),
+                    enrichment.get("app_category"),
+                    enrichment.get("app_confidence"),
+                    enrichment.get("app_hostname"),
+                    enrichment.get("app_is_encrypted"),
+                    json.dumps(dpi_metadata) if dpi_metadata is not None else None,
+                )
+        except Exception:
+            logger.exception("DB flow metadata backfill failed")
+            return 0
+
+        if not rows:
+            return 0
+
+        updated_count = 0
+        for row in rows:
+            serialized = _serialize_row_for_broadcast(dict(row))
+            updated_count += 1
+            if self.broadcaster:
+                try:
+                    self.broadcaster.sync_publish(serialized)
+                except Exception:
+                    logger.exception("Failed to publish flow-enriched packet row")
+
+        return updated_count
+
     async def fetch_latest(self, limit: int) -> List[PacketDBModel]:
         """Fetch newest packet rows as validated `PacketDBModel` instances."""
         if self._pool is None:

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()