test backend codec ndpi

backend/src/utilities/database.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import base64
+import json
 import logging
 from typing import Any, Dict, List, Optional
 
@@ -83,8 +84,16 @@ class DatabasePool:
                         src_port,
                         dst_port,
                         length,
-                        raw
-                    ) VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12)
+                        raw,
+                        app_protocol,
+                        app_master_protocol,
+                        app_category,
+                        app_confidence,
+                        app_hostname,
+                        app_is_encrypted,
+                        app_risk_score,
+                        dpi_metadata
+                    ) VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,$20::jsonb)
                     RETURNING id, timestamp
                     """,
                     pkt_info["iface"],
@@ -99,6 +108,14 @@ class DatabasePool:
                     pkt_info.get("dst_port"),
                     pkt_info["length"],
                     pkt_info["raw"],
+                    pkt_info.get("app_protocol"),
+                    pkt_info.get("app_master_protocol"),
+                    pkt_info.get("app_category"),
+                    pkt_info.get("app_confidence"),
+                    pkt_info.get("app_hostname"),
+                    pkt_info.get("app_is_encrypted"),
+                    pkt_info.get("app_risk_score"),
+                    json.dumps(pkt_info.get("dpi_metadata")) if pkt_info.get("dpi_metadata") is not None else None,
                 )
         except Exception:
             logger.exception("DB insert failed")

backend/src/utilities/ndpi_classifier.py

@@ -0,0 +1,262 @@
+"""Best-effort nDPI flow classifier wrapper.
+
+This module keeps nDPI integration optional:
+- If the Python nDPI bindings are present, packets are classified per flow.
+- If bindings are missing or fail, callers still receive stable fallback fields.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Tuple
+
+from scapy.all import IP, IPv6, TCP, UDP  # type: ignore
+
+logger = logging.getLogger("ndpi_classifier")
+
+try:
+    import ndpi  # type: ignore
+except Exception:  # pragma: no cover - optional dependency
+    ndpi = None
+
+
+def _is_jsonable_scalar(value: Any) -> bool:
+    return isinstance(value, (str, int, float, bool)) or value is None
+
+
+def _safe_string(value: Any) -> Optional[str]:
+    if value is None:
+        return None
+    text = str(value).strip()
+    return text if text else None
+
+
+@dataclass
+class _FlowState:
+    flow_obj: Any
+    last_seen_ts: float
+    packet_count: int = 0
+    detected_app: Optional[str] = None
+
+
+class NDPIClassifier:
+    """Stateful nDPI classifier for TCP/UDP IPv4/IPv6 traffic."""
+
+    def __init__(self) -> None:
+        self._engine = None
+        self._flows: Dict[Tuple[Any, ...], _FlowState] = {}
+        self._max_flows = 200_000
+        self._flow_ttl_seconds = 120.0
+        self._cleanup_interval_packets = 10_000
+        self._packet_counter = 0
+        self._initialized = False
+
+    @property
+    def enabled(self) -> bool:
+        return self._engine is not None
+
+    def _init_engine(self) -> None:
+        if self._initialized:
+            return
+        self._initialized = True
+
+        if ndpi is None:
+            logger.warning("nDPI Python bindings not installed; DPI enrichment disabled")
+            return
+
+        try:
+            engine_cls = getattr(ndpi, "NDPI", None)
+            if engine_cls is None:
+                logger.warning("nDPI bindings found but NDPI class missing; DPI enrichment disabled")
+                return
+            self._engine = engine_cls()
+            logger.info("nDPI classifier initialized")
+        except Exception:
+            logger.exception("Failed to initialize nDPI classifier; enrichment disabled")
+            self._engine = None
+
+    def _maybe_cleanup(self, now_ts: float) -> None:
+        self._packet_counter += 1
+        if self._packet_counter % self._cleanup_interval_packets != 0:
+            return
+
+        expiry = now_ts - self._flow_ttl_seconds
+        stale_keys = [k for k, state in self._flows.items() if state.last_seen_ts < expiry]
+        for key in stale_keys:
+            self._flows.pop(key, None)
+
+        if len(self._flows) > self._max_flows:
+            # Drop oldest half of excess flows.
+            sorted_items = sorted(self._flows.items(), key=lambda kv: kv[1].last_seen_ts)
+            trim_count = len(self._flows) - self._max_flows // 2
+            for key, _ in sorted_items[:trim_count]:
+                self._flows.pop(key, None)
+
+    def _flow_key(self, pkt: Any) -> Optional[Tuple[Any, ...]]:
+        if IP in pkt:
+            ip_layer = pkt[IP]
+            src_ip = getattr(ip_layer, "src", None)
+            dst_ip = getattr(ip_layer, "dst", None)
+            proto = int(getattr(ip_layer, "proto", 0))
+            ip_version = 4
+        elif IPv6 in pkt:
+            ip_layer = pkt[IPv6]
+            src_ip = getattr(ip_layer, "src", None)
+            dst_ip = getattr(ip_layer, "dst", None)
+            proto = int(getattr(ip_layer, "nh", 0))
+            ip_version = 6
+        else:
+            return None
+
+        src_port = None
+        dst_port = None
+        if TCP in pkt:
+            src_port = int(getattr(pkt[TCP], "sport", 0) or 0)
+            dst_port = int(getattr(pkt[TCP], "dport", 0) or 0)
+        elif UDP in pkt:
+            src_port = int(getattr(pkt[UDP], "sport", 0) or 0)
+            dst_port = int(getattr(pkt[UDP], "dport", 0) or 0)
+
+        left = (src_ip, src_port)
+        right = (dst_ip, dst_port)
+        ep1, ep2 = (left, right) if left <= right else (right, left)
+        return (ip_version, proto, ep1, ep2)
+
+    def _packet_payload_for_ndpi(self, pkt: Any) -> Optional[bytes]:
+        try:
+            if IP in pkt:
+                return bytes(pkt[IP])
+            if IPv6 in pkt:
+                return bytes(pkt[IPv6])
+        except Exception:
+            return None
+        return None
+
+    def _extract_result(self, result: Any) -> Dict[str, Any]:
+        if result is None:
+            return {}
+
+        raw: Dict[str, Any] = {}
+        if isinstance(result, dict):
+            raw.update(result)
+        else:
+            for attr in (
+                "app_protocol",
+                "application_protocol",
+                "master_protocol",
+                "protocol",
+                "category",
+                "confidence",
+                "hostname",
+                "server_name",
+                "sni",
+                "is_encrypted",
+                "risk_score",
+                "risks",
+                "ja3",
+                "ja4",
+                "alpn",
+            ):
+                if hasattr(result, attr):
+                    raw[attr] = getattr(result, attr)
+
+        app_protocol = (
+            _safe_string(raw.get("app_protocol"))
+            or _safe_string(raw.get("application_protocol"))
+            or _safe_string(raw.get("protocol"))
+        )
+        master_protocol = _safe_string(raw.get("master_protocol"))
+        category = _safe_string(raw.get("category"))
+        confidence = _safe_string(raw.get("confidence"))
+        hostname = _safe_string(raw.get("hostname")) or _safe_string(raw.get("server_name")) or _safe_string(raw.get("sni"))
+
+        risks_val = raw.get("risks")
+        if isinstance(risks_val, (list, tuple, set)):
+            risk_score = len(risks_val)
+        else:
+            try:
+                risk_score = int(raw.get("risk_score")) if raw.get("risk_score") is not None else None
+            except Exception:
+                risk_score = None
+
+        encrypted = raw.get("is_encrypted")
+        is_encrypted = bool(encrypted) if isinstance(encrypted, bool) else None
+        if is_encrypted is None and app_protocol:
+            proto_upper = app_protocol.upper()
+            if any(token in proto_upper for token in ("TLS", "SSL", "HTTPS", "QUIC", "VPN")):
+                is_encrypted = True
+
+        metadata: Dict[str, Any] = {}
+        for key, value in raw.items():
+            if _is_jsonable_scalar(value):
+                metadata[key] = value
+            elif isinstance(value, (list, tuple)):
+                metadata[key] = [str(v) for v in value][:32]
+            elif isinstance(value, dict):
+                cleaned = {str(k): str(v) for k, v in list(value.items())[:32]}
+                metadata[key] = cleaned
+            else:
+                metadata[key] = str(value)
+
+        return {
+            "app_protocol": app_protocol,
+            "app_master_protocol": master_protocol,
+            "app_category": category,
+            "app_confidence": confidence,
+            "app_hostname": hostname,
+            "app_is_encrypted": is_encrypted,
+            "app_risk_score": risk_score,
+            "dpi_metadata": metadata if metadata else None,
+        }
+
+    def classify_packet(self, pkt: Any) -> Dict[str, Any]:
+        """Classify packet and return enrichment fields for packet metadata."""
+        self._init_engine()
+        if not self.enabled:
+            return {}
+
+        flow_key = self._flow_key(pkt)
+        if flow_key is None:
+            return {}
+
+        payload = self._packet_payload_for_ndpi(pkt)
+        if not payload:
+            return {}
+
+        flow_cls = getattr(ndpi, "NDPIFlow", None)
+        now_ts = time.time()
+        tick_ms = int(now_ts * 1000)
+
+        flow_state = self._flows.get(flow_key)
+        if flow_state is None:
+            flow_obj = flow_cls() if flow_cls is not None else None
+            flow_state = _FlowState(flow_obj=flow_obj, last_seen_ts=now_ts, packet_count=0)
+            self._flows[flow_key] = flow_state
+
+        flow_state.last_seen_ts = now_ts
+        flow_state.packet_count += 1
+
+        try:
+            if flow_state.flow_obj is not None:
+                result = self._engine.process_packet(flow_state.flow_obj, payload, tick_ms)
+            else:
+                # Fallback for bindings that don't expose per-flow object.
+                result = self._engine.process_packet(payload, tick_ms)
+        except Exception:
+            logger.debug("nDPI process_packet failed for flow=%s", flow_key, exc_info=True)
+            self._maybe_cleanup(now_ts)
+            return {}
+
+        enriched = self._extract_result(result)
+        if enriched.get("app_protocol"):
+            flow_state.detected_app = enriched["app_protocol"]
+        elif flow_state.detected_app:
+            enriched["app_protocol"] = flow_state.detected_app
+
+        self._maybe_cleanup(now_ts)
+        return enriched
+
+
+ndpi_classifier = NDPIClassifier()