Aegis/backend/app/routers/heatmap.py

"""Heatmap endpoints — ATT&CK Navigator-compatible layer generation.

Thin router that delegates entirely to :mod:`app.services.heatmap_service`.
No business logic lives here — only request validation and response
formatting.
"""

# Import io
import io

# Import json
import json

# Import Optional from typing
from typing import Optional

# Import APIRouter, Depends, Query from fastapi
from fastapi import APIRouter, Depends, Query

# Import StreamingResponse from fastapi.responses
from fastapi.responses import StreamingResponse

# Import Session from sqlalchemy.orm
from sqlalchemy.orm import Session

# Import get_db from app.database
from app.database import get_db

# Import get_current_user from app.dependencies.auth
from app.dependencies.auth import get_current_user

# Import User from app.models.user
from app.models.user import User

# Import heatmap_service from app.services
from app.services import heatmap_service

# Assign router = APIRouter(prefix="/heatmap", tags=["heatmap"])
router = APIRouter(prefix="/heatmap", tags=["heatmap"])


# Apply the @router.get decorator
@router.get("/coverage")
# Define function heatmap_coverage
def heatmap_coverage(
    # Entry: platforms
    platforms: Optional[str] = Query(None, description="Comma-separated platforms"),
    # Entry: tactics
    tactics: Optional[str] = Query(None, description="Comma-separated tactics"),
    # Entry: min_score
    min_score: int = Query(0, ge=0, le=100),
    # Entry: db
    db: Session = Depends(get_db),
    # Entry: current_user
    current_user: User = Depends(get_current_user),
) -> dict:
    """Coverage layer — score based on status_global of each technique."""
    # Return heatmap_service.build_coverage_layer(
    return heatmap_service.build_coverage_layer(
        db, platforms=platforms, tactics=tactics, min_score=min_score,
    )


# Apply the @router.get decorator
@router.get("/threat-actor/{actor_id}")
# Define function heatmap_threat_actor
def heatmap_threat_actor(
    # Entry: actor_id
    actor_id: str,
    # Entry: platforms
    platforms: Optional[str] = Query(None),
    # Entry: tactics
    tactics: Optional[str] = Query(None),
    # Entry: min_score
    min_score: int = Query(0, ge=0, le=100),
    # Entry: db
    db: Session = Depends(get_db),
    # Entry: current_user
    current_user: User = Depends(get_current_user),
) -> dict:
    """Threat actor layer — techniques used by an actor with coverage color."""
    # Return heatmap_service.build_threat_actor_layer(
    return heatmap_service.build_threat_actor_layer(
        db, actor_id, platforms=platforms, tactics=tactics, min_score=min_score,
    )


# Apply the @router.get decorator
@router.get("/detection-rules")
# Define function heatmap_detection_rules
def heatmap_detection_rules(
    # Entry: platforms
    platforms: Optional[str] = Query(None),
    # Entry: tactics
    tactics: Optional[str] = Query(None),
    # Entry: min_score
    min_score: int = Query(0, ge=0, le=100),
    # Entry: db
    db: Session = Depends(get_db),
    # Entry: current_user
    current_user: User = Depends(get_current_user),
) -> dict:
    """Detection rules layer — score based on ratio of rules available vs total."""
    # Return heatmap_service.build_detection_rules_layer(
    return heatmap_service.build_detection_rules_layer(
        db, platforms=platforms, tactics=tactics, min_score=min_score,
    )


# Apply the @router.get decorator
@router.get("/campaign/{campaign_id}")
# Define function heatmap_campaign
def heatmap_campaign(
    # Entry: campaign_id
    campaign_id: str,
    # Entry: platforms
    platforms: Optional[str] = Query(None),
    # Entry: tactics
    tactics: Optional[str] = Query(None),
    # Entry: min_score
    min_score: int = Query(0, ge=0, le=100),
    # Entry: db
    db: Session = Depends(get_db),
    # Entry: current_user
    current_user: User = Depends(get_current_user),
) -> dict:
    """Campaign layer — only techniques in the campaign, colored by test state."""
    # Return heatmap_service.build_campaign_layer(
    return heatmap_service.build_campaign_layer(
        db, campaign_id, platforms=platforms, tactics=tactics, min_score=min_score,
    )


# Apply the @router.get decorator
@router.get("/export-navigator")
# Define function export_navigator
def export_navigator(
    # Entry: layer
    layer: str = Query(..., description="Layer type: coverage, threat-actor, detection-rules, campaign"),
    # Entry: layer_id
    layer_id: Optional[str] = Query(None, description="Actor ID or Campaign ID (if applicable)"),
    # Entry: platforms
    platforms: Optional[str] = Query(None),
    # Entry: tactics
    tactics: Optional[str] = Query(None),
    # Entry: min_score
    min_score: int = Query(0, ge=0, le=100),
    # Entry: db
    db: Session = Depends(get_db),
    # Entry: current_user
    current_user: User = Depends(get_current_user),
) -> StreamingResponse:
    """Export a heatmap layer as a downloadable JSON file for ATT&CK Navigator."""
    # Assign data = heatmap_service.build_navigator_export(
    data = heatmap_service.build_navigator_export(
        db, layer, layer_id=layer_id,
        # Keyword argument: platforms
        platforms=platforms, tactics=tactics, min_score=min_score,
    )

    # Assign json_content = json.dumps(data, indent=2, default=str)
    json_content = json.dumps(data, indent=2, default=str)
    # Assign buffer = io.BytesIO(json_content.encode("utf-8"))
    buffer = io.BytesIO(json_content.encode("utf-8"))

    # Return StreamingResponse(
    return StreamingResponse(
        buffer,
        # Keyword argument: media_type
        media_type="application/json",
        # Keyword argument: headers
        headers={"Content-Disposition": f"attachment; filename=aegis_{layer}_layer.json"},
    )