mitm-webserver/backend/src/utilities/bridge_telemetry.py

"""Manage the eBPF/tc telemetry subprocess used for bridge packet capture and verdict events."""

from __future__ import annotations

import base64
import json
import logging
import os
import signal
import subprocess
import sys
import threading
from pathlib import Path
from typing import Iterable, Mapping, Optional

from src.config import settings
from src.utilities.packet_tracker import packet_tracker

logger = logging.getLogger("bridge_telemetry")


class BridgeTelemetryManager:
    """Run one tc/eBPF collector for the currently sniffed bridge interfaces."""

    def __init__(self) -> None:
        self._interfaces: set[str] = set()
        self._session_ids_by_interface: dict[str, tuple[str, ...]] = {}
        self._process: Optional[subprocess.Popen[str]] = None
        self._reader_thread: Optional[threading.Thread] = None
        self._lock = threading.Lock()

    def update_sessions(self, session_interfaces: Mapping[str, Iterable[str]]) -> None:
        """Restart the collector when the active bridge interface set changes."""
        normalized: dict[str, set[str]] = {}
        for session_id, interfaces in session_interfaces.items():
            if not session_id:
                continue
            iface_set = {iface.strip() for iface in interfaces if iface and iface.strip()}
            if iface_set:
                normalized[session_id] = iface_set

        normalized_interfaces = sorted({iface for ifaces in normalized.values() for iface in ifaces})
        session_ids_by_interface = {
            iface: tuple(sorted(session_id for session_id, ifaces in normalized.items() if iface in ifaces))
            for iface in normalized_interfaces
        }

        with self._lock:
            interfaces_changed = set(normalized_interfaces) != self._interfaces
            self._interfaces = set(normalized_interfaces)
            self._session_ids_by_interface = session_ids_by_interface
            if not interfaces_changed:
                return
            self._restart_locked()

    def stop(self) -> None:
        """Stop the collector process and reader thread."""
        with self._lock:
            self._interfaces = set()
            self._stop_process_locked()

    def _restart_locked(self) -> None:
        self._stop_process_locked()
        if not self._interfaces:
            return

        helper = Path(__file__).with_name("ebpf_bridge_events.py")
        python_bin = sys.executable or "python3"

        env = os.environ.copy()
        env["PYTHONUNBUFFERED"] = "1"
        backend_root = str(helper.parents[2])
        existing_pythonpath = env.get("PYTHONPATH", "")
        env["PYTHONPATH"] = backend_root if not existing_pythonpath else f"{backend_root}:{existing_pythonpath}"

        cmd = [
            python_bin,
            str(helper),
            "--ifaces",
            ",".join(sorted(self._interfaces)),
            "--build-dir",
            settings.bridge_bpf_build_dir,
        ]
        logger.info("Starting bridge telemetry collector for interfaces=%s", sorted(self._interfaces))
        try:
            self._process = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
                bufsize=1,
                env=env,
            )
        except Exception:
            logger.exception("Failed to start bridge telemetry collector")
            self._process = None
            return

        self._reader_thread = threading.Thread(
            target=self._read_loop,
            args=(self._process,),
            daemon=True,
            name="bridge-telemetry-reader",
        )
        self._reader_thread.start()

    def _stop_process_locked(self) -> None:
        process = self._process
        reader = self._reader_thread
        self._process = None
        self._reader_thread = None

        if process is not None and process.poll() is None:
            try:
                process.send_signal(signal.SIGTERM)
                process.wait(timeout=settings.telemetry_process_stop_timeout_seconds)
            except subprocess.TimeoutExpired:
                process.kill()
            except Exception:
                logger.exception("Failed to stop bridge telemetry collector cleanly")

        if reader is not None and reader.is_alive():
            reader.join(timeout=settings.telemetry_reader_join_timeout_seconds)

    def _handle_ingress_packet(self, event: dict[str, object]) -> None:
        raw_b64 = event.pop("raw_b64", None)
        if not isinstance(raw_b64, str) or not raw_b64:
            return

        try:
            packet_bytes = base64.b64decode(raw_b64)
        except Exception:
            logger.exception("Failed to decode ingress raw packet")
            return

        capture_metadata = {
            "capture_source": "tc_ingress_raw",
            "packet_id": event.get("packet_id"),
            "skb_mark": event.get("skb_mark"),
            "capture_mode": "tc_ingress",
        }
        capture_iface = str(event.get("iface") or "")
        capture_session_id: Optional[str] = None
        with self._lock:
            session_ids = self._session_ids_by_interface.get(capture_iface, ())
            if session_ids:
                capture_session_id = session_ids[0]

        try:
            from src.network_sniffer import parse_packet_bytes

            parse_packet_bytes(
                packet_bytes,
                capture_iface,
                capture_metadata=capture_metadata,
                capture_session_id=capture_session_id,
            )
        except Exception:
            logger.exception("Failed to process ingress raw packet event")

    def _read_loop(self, process: subprocess.Popen[str]) -> None:
        stdout = process.stdout
        if stdout is None:
            return

        for line in stdout:
            text = line.strip()
            if not text:
                continue
            try:
                event = json.loads(text)
            except json.JSONDecodeError:
                logger.info("bridge-telemetry: %s", text)
                continue

            if "event_type" not in event:
                logger.info("bridge-telemetry: %s", event)
                continue

            if event.get("event_type") == "ingress":
                self._handle_ingress_packet(event)

            try:
                packet_tracker.observe_telemetry(event)
            except Exception:
                logger.exception("Failed to process telemetry event: %s", event)

        rc = process.poll()
        if rc not in (0, None):
            logger.warning("Bridge telemetry collector exited with code %s", rc)


bridge_telemetry_manager = BridgeTelemetryManager()