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

"""Manage the optional tc-based bridge capture pipeline for mirrored packet storage."""

from __future__ import annotations

import logging
import subprocess
import threading
from pathlib import Path
from typing import Dict, Optional

from src.config import settings

logger = logging.getLogger("capture_pipeline")


def _run_checked(cmd: list[str]) -> subprocess.CompletedProcess[str]:
    """Run a command and surface stdout/stderr in raised errors."""
    try:
        return subprocess.run(cmd, check=True, capture_output=True, text=True)
    except subprocess.CalledProcessError as exc:
        stdout = (exc.stdout or "").strip()
        stderr = (exc.stderr or "").strip()
        if stdout:
            logger.error("Command stdout for %s:\n%s", cmd[0], stdout)
        if stderr:
            logger.error("Command stderr for %s:\n%s", cmd[0], stderr)
        detail = stderr or stdout or str(exc)
        raise RuntimeError(detail) from exc


class CapturePipelineManager:
    """Reference-count mirrored capture pipelines per bridge."""

    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._refs: Dict[str, int] = {}
        self._active: Dict[str, Dict[str, str]] = {}

    def acquire(self, bridge: str) -> Optional[str]:
        """Ensure a bridge capture pipeline exists and return the capture interface name."""
        if not settings.capture_automanage:
            return settings.capture_interface or None

        bridge = bridge.strip()
        if not bridge:
            raise ValueError("Bridge name is required for capture pipeline setup")

        with self._lock:
            refcount = self._refs.get(bridge, 0)
            if refcount > 0:
                self._refs[bridge] = refcount + 1
                active = self._active.get(bridge, {})
                return active.get("capture_if") or settings.capture_interface or settings.capture_automanaged_capture_if

            self._setup_bridge_locked(bridge)
            self._refs[bridge] = 1
            active = self._active.get(bridge, {})
            return active.get("capture_if") or settings.capture_interface or settings.capture_automanaged_capture_if

    def release(self, bridge: str) -> None:
        """Tear down a bridge capture pipeline when the last session releases it."""
        if not settings.capture_automanage:
            return

        bridge = bridge.strip()
        if not bridge:
            return

        with self._lock:
            refcount = self._refs.get(bridge, 0)
            if refcount <= 1:
                self._refs.pop(bridge, None)
                try:
                    self._teardown_bridge_locked(bridge)
                finally:
                    self._active.pop(bridge, None)
                return

            self._refs[bridge] = refcount - 1

    def stop_all(self) -> None:
        """Remove all managed capture pipelines."""
        if not settings.capture_automanage:
            return

        with self._lock:
            bridges = list(self._active.keys())
            self._refs.clear()
            for bridge in bridges:
                try:
                    self._teardown_bridge_locked(bridge)
                except Exception:
                    logger.exception("Failed to tear down capture pipeline for bridge %s", bridge)
            self._active.clear()

    def get_debug_snapshot(self) -> Dict[str, object]:
        with self._lock:
            return {
                "automanage_enabled": settings.capture_automanage,
                "refs": dict(self._refs),
                "active": dict(self._active),
            }

    def _setup_bridge_locked(self, bridge: str) -> None:
        script = self._repo_root() / "tools" / "setup_bridge_capture.sh"
        mirror_if = settings.capture_automanaged_mirror_if
        capture_if = settings.capture_interface or settings.capture_automanaged_capture_if
        cmd = [
            str(script),
            "--bridge",
            bridge,
            "--mirror-if",
            mirror_if,
            "--capture-if",
            capture_if,
            "--build-dir",
            settings.capture_build_dir,
        ]
        logger.info(
            "Setting up capture pipeline for bridge=%s mirror_if=%s capture_if=%s",
            bridge,
            mirror_if,
            capture_if,
        )
        _run_checked(cmd)
        self._active[bridge] = {
            "mirror_if": mirror_if,
            "capture_if": capture_if,
            "build_dir": settings.capture_build_dir,
        }

    def _teardown_bridge_locked(self, bridge: str) -> None:
        script = self._repo_root() / "tools" / "teardown_bridge_capture.sh"
        active = self._active.get(bridge, {})
        mirror_if = active.get("mirror_if") or settings.capture_automanaged_mirror_if
        capture_if = active.get("capture_if") or settings.capture_interface or settings.capture_automanaged_capture_if
        cmd = [
            str(script),
            "--bridge",
            bridge,
            "--mirror-if",
            mirror_if,
            "--capture-if",
            capture_if,
        ]
        logger.info(
            "Tearing down capture pipeline for bridge=%s mirror_if=%s capture_if=%s",
            bridge,
            mirror_if,
            capture_if,
        )
        _run_checked(cmd)

    def _repo_root(self) -> Path:
        return Path(__file__).resolve().parents[3]


capture_pipeline_manager = CapturePipelineManager()