diff --git a/backend/src/utilities/database.py b/backend/src/utilities/database.py index ce34505..8e38816 100644 --- a/backend/src/utilities/database.py +++ b/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: diff --git a/backend/src/utilities/nfstream_manager.py b/backend/src/utilities/nfstream_manager.py index 6b474a7..252b0a7 100644 --- a/backend/src/utilities/nfstream_manager.py +++ b/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()