"""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"}, )