mitm-webserver/backend/src/api/sniffer_api.py

# src/routers/sniffer.py
from fastapi import APIRouter, HTTPException, Query, Body
from pydantic import BaseModel, Field
from typing import Dict, Any, Optional

from src.network_sniffer import (
    get_sniffer_status,
    start_afpacket_sniffer,
    stop_afpacket_sniffer,
)

router = APIRouter()


# ------------------------------
# Pydantic Models
# ------------------------------
class SnifferStartRequest(BaseModel):
    """
    Request model for starting the sniffer on a specific bridge OR interface.
    Exactly one of `bridge` or `interface` must be provided.
    """
    bridge: Optional[str] = Field(None, example="br0", description="Name of the Linux bridge to sniff on")
    interface: Optional[str] = Field(None, example="eth0", description="Name of the network interface to sniff on")


class SnifferStartResponse(BaseModel):
    """
    Response model returned when sniffer starts successfully.
    """
    started: bool = Field(..., description="Whether the sniffer was started successfully")
    session_id: str = Field(..., description="Session identifier for this sniffer instance")
    target: str = Field(..., description="Target that was started (bridge or interface)")
    target_type: str = Field(..., description="Either 'bridge' or 'interface'")


class SnifferStopRequest(BaseModel):
    """
    Optional body for stop — prefer session_id if you want to stop a specific session.
    If omitted, stopping behavior will be determined by query params (bridge/interface) or global stop.
    """
    session_id: Optional[str] = Field(None, description="Session id to stop")


class SnifferStopResponse(BaseModel):
    stopped: bool = Field(..., description="Whether the sniffer was stopped successfully")
    session_id: Optional[str] = Field(None, description="Session id stopped (if any)")
    target: Optional[str] = Field(None, description="Target stopped; null if global stop")
    target_type: Optional[str] = Field(None, description="'bridge' or 'interface' or None")


class InterfaceSnifferStatus(BaseModel):
    running: bool = Field(..., description="Whether the sniffer thread/socket is active")
    exists: bool = Field(..., description="Whether the interface exists in /sys/class/net")
    up: bool = Field(..., description="Whether the interface is operationally UP")
    session_id: Optional[str] = Field(None, description="Session id owning this interface")
    session_label: Optional[str] = Field(None, description="Human label for the session")


class SnifferStatusResponse(BaseModel):
    interfaces: Dict[str, InterfaceSnifferStatus] = Field(
        ..., description="Map of interface names to their sniffer status"
    )


# ------------------------------
# Endpoints
# ------------------------------
@router.post("/start", response_model=SnifferStartResponse)
def sniffer_start(req: SnifferStartRequest):
    """
    Start a sniffer session for the given bridge OR interface.
    Exactly one of `bridge` or `interface` must be provided.
    Returns a session_id to manage the session.
    """
    if bool(req.bridge) == bool(req.interface):
        raise HTTPException(status_code=400, detail="Exactly one of 'bridge' or 'interface' must be provided")

    try:
        if req.interface:
            session_id = start_afpacket_sniffer(req.interface, target_is_interface=True)
            return SnifferStartResponse(started=True, session_id=session_id, target=req.interface, target_type="interface")
        else:
            session_id = start_afpacket_sniffer(req.bridge, target_is_interface=False)
            return SnifferStartResponse(started=True, session_id=session_id, target=req.bridge, target_type="bridge")
    except Exception as exc:
        raise HTTPException(status_code=500, detail=f"Failed to start sniffer: {exc}")


@router.post("/stop", response_model=SnifferStopResponse)
def sniffer_stop(
    q_bridge: Optional[str] = Query(None, alias="bridge", description="If provided, stop sniffer sockets for this bridge"),
    q_interface: Optional[str] = Query(None, alias="interface", description="If provided, stop sniffer socket for this interface"),
    body: SnifferStopRequest = Body(...),
):
    """
    Stop sniffer sessions.

    - If body.session_id is provided: stop that session (preferred).
    - Else if query param `interface` provided: close socket for that interface across sessions.
    - Else if query param `bridge` provided: remove snapshot / close sockets for that bridge across sessions.
    - Else: stop all sessions (global stop).
    """
    if body and body.session_id:
        try:
            stop_afpacket_sniffer(session_id=body.session_id)
            return SnifferStopResponse(stopped=True, session_id=body.session_id, target=None, target_type=None)
        except Exception as exc:
            raise HTTPException(status_code=500, detail=f"Failed to stop session {body.session_id}: {exc}")

    # validate query params
    if q_bridge and q_interface:
        raise HTTPException(status_code=400, detail="Only one of 'bridge' or 'interface' may be provided")

    try:
        if q_interface:
            stop_afpacket_sniffer(target=q_interface, target_is_interface=True)
            return SnifferStopResponse(stopped=True, session_id=None, target=q_interface, target_type="interface")
        if q_bridge:
            stop_afpacket_sniffer(target=q_bridge, target_is_interface=False)
            return SnifferStopResponse(stopped=True, session_id=None, target=q_bridge, target_type="bridge")

        # global stop
        stop_afpacket_sniffer()
        return SnifferStopResponse(stopped=True, session_id=None, target=None, target_type=None)
    except Exception as exc:
        raise HTTPException(status_code=500, detail=f"Failed to stop sniffer: {exc}")


@router.get("/status", response_model=SnifferStatusResponse)
def sniffer_status():
    """
    Return the sniffer status information.
    """
    try:
        raw = get_sniffer_status()
        # Convert raw dict → typed model
        typed = {
            k: InterfaceSnifferStatus(**v)
            for k, v in raw.items()
        }
        return SnifferStatusResponse(interfaces=typed)
    except Exception as exc:
        raise HTTPException(status_code=500, detail=f"Failed to query sniffer status: {exc}")