test ebpf

backend/src/utilities/packet_tracker.py

@@ -0,0 +1,312 @@
+"""Aggregate AF_PACKET observations and kernel telemetry into one packet record."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import threading
+import time
+from datetime import datetime, timezone
+from typing import Any, Dict, List, Optional
+
+import src.shared_objects as shared_objects
+from src.utilities.packet_identity import build_packet_uid
+
+logger = logging.getLogger("packet_tracker")
+
+
+def _utcnow() -> datetime:
+    return datetime.now(timezone.utc)
+
+
+class PacketTracker:
+    """Deduplicate packet observations and persist one upserted row per packet."""
+
+    def __init__(
+        self,
+        finalize_delay_seconds: float = 0.25,
+        retention_seconds: float = 10.0,
+        min_flush_interval_seconds: float = 0.05,
+    ):
+        self._finalize_delay_seconds = finalize_delay_seconds
+        self._retention_seconds = retention_seconds
+        self._min_flush_interval_seconds = min_flush_interval_seconds
+        self._entries: Dict[str, Dict[str, Any]] = {}
+        self._lock = threading.Lock()
+        self._stop_event = threading.Event()
+        self._thread = threading.Thread(target=self._run, daemon=True, name="packet-tracker")
+        self._thread.start()
+
+    def stop(self) -> None:
+        self._stop_event.set()
+        self._thread.join(timeout=2)
+
+    def observe_packet(self, pkt_info: Dict[str, Any]) -> str:
+        """Merge parsed packet information into a pending packet entry."""
+        now_ts = time.time()
+        packet_uid = pkt_info.get("packet_uid") or build_packet_uid(pkt_info)
+        pkt_info["packet_uid"] = packet_uid
+
+        with self._lock:
+            entry = self._entries.get(packet_uid)
+            if entry is None:
+                entry = self._new_entry(packet_uid, now_ts)
+                self._entries[packet_uid] = entry
+
+            self._merge_packet_info(entry, pkt_info, now_ts)
+            self._maybe_promote_reject_from_reply(pkt_info, now_ts)
+            self._maybe_mark_complete(entry)
+        return packet_uid
+
+    def observe_telemetry(self, event: Dict[str, Any]) -> Optional[str]:
+        """Merge ingress/egress/verdict telemetry into a pending packet entry."""
+        packet_uid = event.get("packet_uid")
+        if not packet_uid:
+            try:
+                packet_uid = build_packet_uid(event)
+            except Exception:
+                logger.debug("Telemetry event missing packet identity: %s", event)
+                return None
+
+        now_ts = time.time()
+        with self._lock:
+            entry = self._entries.get(packet_uid)
+            if entry is None:
+                entry = self._new_entry(packet_uid, now_ts)
+                self._entries[packet_uid] = entry
+
+            payload = entry["payload"]
+            payload["packet_uid"] = packet_uid
+            payload["telemetry_metadata"] = event
+            payload["last_observed_at"] = now_ts
+            for key, value in event.items():
+                if value is None or key in {"event_type", "reason", "reason_code", "iface", "packet_uid"}:
+                    continue
+                if payload.get(key) is None:
+                    payload[key] = value
+            event_type = event.get("event_type")
+            iface = event.get("iface")
+            if iface:
+                observed_ifaces = payload.setdefault("observed_ifaces", [])
+                if iface not in observed_ifaces:
+                    observed_ifaces.append(iface)
+
+            if event_type == "ingress":
+                payload["ingress_if"] = iface
+                payload["iface"] = iface
+                payload["ingress_seen_at"] = _utcnow()
+                payload["direction"] = "ingress"
+            elif event_type == "egress":
+                payload["egress_if"] = iface
+                payload["egress_seen_at"] = _utcnow()
+                payload["direction"] = "forwarded"
+                payload["verdict"] = "accept"
+                payload["verdict_reason"] = "egress-observed"
+                payload["verdict_confidence"] = "high"
+                payload["verdict_seen_at"] = _utcnow()
+            elif event_type == "drop":
+                payload["verdict"] = "drop"
+                payload["verdict_reason"] = event.get("reason") or "kfree_skb"
+                payload["verdict_confidence"] = "high"
+                payload["verdict_seen_at"] = _utcnow()
+                payload["direction"] = "dropped"
+            elif event_type == "reject":
+                payload["verdict"] = "reject"
+                payload["verdict_reason"] = event.get("reason") or "netfilter-reject"
+                payload["verdict_confidence"] = event.get("verdict_confidence") or "medium"
+                payload["verdict_seen_at"] = _utcnow()
+                payload["direction"] = "rejected"
+
+            entry["last_observed_at"] = now_ts
+            entry["dirty"] = True
+            self._maybe_mark_complete(entry)
+        return packet_uid
+
+    def _new_entry(self, packet_uid: str, now_ts: float) -> Dict[str, Any]:
+        return {
+            "packet_uid": packet_uid,
+            "payload": {
+                "packet_uid": packet_uid,
+                "observed_ifaces": [],
+                "verdict": "pending",
+                "verdict_reason": None,
+                "verdict_confidence": None,
+                "telemetry_metadata": None,
+            },
+            "persisted": False,
+            "dirty": True,
+            "finalized": False,
+            "created_at": now_ts,
+            "last_observed_at": now_ts,
+            "last_persisted_at": 0.0,
+        }
+
+    def _merge_packet_info(self, entry: Dict[str, Any], pkt_info: Dict[str, Any], now_ts: float) -> None:
+        payload = entry["payload"]
+        changed = False
+        for key, value in pkt_info.items():
+            if key == "observed_ifaces":
+                continue
+            if value is None:
+                continue
+            if key == "raw" and payload.get("raw") is not None:
+                continue
+            if payload.get(key) == value:
+                continue
+            payload[key] = value
+            changed = True
+
+        iface = pkt_info.get("iface")
+        if iface:
+            observed_ifaces = payload.setdefault("observed_ifaces", [])
+            if iface not in observed_ifaces:
+                observed_ifaces.append(iface)
+                changed = True
+            if not payload.get("ingress_if"):
+                payload["ingress_if"] = iface
+                payload["iface"] = iface
+                payload["ingress_seen_at"] = _utcnow()
+                changed = True
+
+        payload["last_observed_at"] = now_ts
+        entry["last_observed_at"] = now_ts
+        entry["dirty"] = entry["dirty"] or changed
+
+    def _maybe_promote_reject_from_reply(self, pkt_info: Dict[str, Any], now_ts: float) -> None:
+        reject_reason = None
+
+        tcp_flags = pkt_info.get("tcp_flags")
+        if pkt_info.get("protocol_raw") == 6 and tcp_flags is not None and int(tcp_flags) & 0x04:
+            reject_reason = "tcp-rst-observed"
+            match = self._find_recent_drop(
+                src_ip=pkt_info.get("dst_ip"),
+                dst_ip=pkt_info.get("src_ip"),
+                protocol_raw=6,
+                src_port=pkt_info.get("dst_port"),
+                dst_port=pkt_info.get("src_port"),
+            )
+        elif pkt_info.get("protocol_raw") == 1 and pkt_info.get("icmp_type") == 3:
+            reject_reason = "icmp-unreachable-observed"
+            match = self._find_recent_drop(
+                src_ip=pkt_info.get("icmp_embedded_src_ip"),
+                dst_ip=pkt_info.get("icmp_embedded_dst_ip"),
+                protocol_raw=pkt_info.get("icmp_embedded_protocol"),
+                src_port=pkt_info.get("icmp_embedded_src_port"),
+                dst_port=pkt_info.get("icmp_embedded_dst_port"),
+            )
+        else:
+            return
+
+        if match is None:
+            return
+
+        payload = match["payload"]
+        if payload.get("verdict") != "drop":
+            return
+
+        payload["verdict"] = "reject"
+        payload["verdict_reason"] = reject_reason
+        payload["verdict_confidence"] = "medium"
+        payload["verdict_seen_at"] = _utcnow()
+        payload["direction"] = "rejected"
+        match["last_observed_at"] = now_ts
+        match["dirty"] = True
+        match["finalized"] = True
+
+    def _find_recent_drop(
+        self,
+        src_ip: Any,
+        dst_ip: Any,
+        protocol_raw: Any,
+        src_port: Any,
+        dst_port: Any,
+    ) -> Optional[Dict[str, Any]]:
+        if not src_ip or not dst_ip or protocol_raw is None:
+            return None
+
+        cutoff = time.time() - 1.0
+        for entry in self._entries.values():
+            payload = entry["payload"]
+            if entry["last_observed_at"] < cutoff:
+                continue
+            if payload.get("verdict") != "drop":
+                continue
+            if payload.get("src_ip") != src_ip or payload.get("dst_ip") != dst_ip:
+                continue
+            if payload.get("protocol_raw") != protocol_raw:
+                continue
+            if payload.get("src_port") != src_port or payload.get("dst_port") != dst_port:
+                continue
+            return entry
+        return None
+
+    def _maybe_mark_complete(self, entry: Dict[str, Any]) -> None:
+        payload = entry["payload"]
+        if payload.get("verdict") in {"accept", "drop", "reject"}:
+            entry["finalized"] = True
+
+    def _run(self) -> None:
+        while not self._stop_event.is_set():
+            time.sleep(0.05)
+            due_entries: List[Dict[str, Any]] = []
+            expired_uids: List[str] = []
+            now_ts = time.time()
+
+            with self._lock:
+                for packet_uid, entry in list(self._entries.items()):
+                    age = now_ts - entry["last_observed_at"]
+                    if not entry["finalized"] and age >= self._finalize_delay_seconds:
+                        entry["finalized"] = True
+                        if entry["payload"].get("verdict") == "pending":
+                            entry["payload"]["verdict"] = "unknown"
+                            entry["payload"]["verdict_reason"] = "timeout"
+                            entry["payload"]["verdict_confidence"] = "low"
+                            entry["payload"]["verdict_seen_at"] = _utcnow()
+                            entry["payload"]["direction"] = "observed"
+                            entry["dirty"] = True
+
+                    should_flush = entry["dirty"] and (
+                        not entry["persisted"]
+                        or entry["finalized"]
+                        or (now_ts - entry["last_persisted_at"]) >= self._min_flush_interval_seconds
+                    )
+                    if should_flush:
+                        due_entries.append(
+                            {
+                                "packet_uid": entry["packet_uid"],
+                                "payload": dict(entry["payload"]),
+                            }
+                        )
+                    elif entry["persisted"] and age >= self._retention_seconds:
+                        expired_uids.append(packet_uid)
+
+                for packet_uid in expired_uids:
+                    self._entries.pop(packet_uid, None)
+
+            for entry in due_entries:
+                self._persist(entry)
+
+    def _persist(self, entry: Dict[str, Any]) -> None:
+        payload = dict(entry["payload"])
+        if payload.get("observed_ifaces") == []:
+            payload["observed_ifaces"] = None
+
+        web_loop = getattr(shared_objects, "web_loop", None)
+        web_db = getattr(shared_objects, "db", None)
+        if web_loop is None or web_db is None:
+            return
+
+        try:
+            fut = asyncio.run_coroutine_threadsafe(web_db.upsert_packet(payload), web_loop)
+            fut.result(timeout=2.0)
+            with self._lock:
+                current = self._entries.get(entry["packet_uid"])
+                if current is not None:
+                    current["persisted"] = True
+                    current["dirty"] = False
+                    current["last_persisted_at"] = time.time()
+        except Exception:
+            logger.exception("Failed to persist packet %s", entry["packet_uid"])
+
+
+packet_tracker = PacketTracker()