Files
Aegis/frontend/src/pages/MatrixPage.tsx

288 lines
10 KiB
TypeScript

import { useState, useMemo, useCallback } from "react";
import { useQuery } from "@tanstack/react-query";
import { useNavigate } from "react-router-dom";
import { Loader2, AlertCircle, Download, ZoomIn, ZoomOut } from "lucide-react";
import {
getHeatmapCoverage,
getHeatmapThreatActor,
getHeatmapDetectionRules,
getHeatmapCampaign,
exportNavigatorJSON,
type HeatmapLayer,
type HeatmapFilters as HeatmapFilterParams,
} from "../api/heatmap";
import AdvancedHeatmap from "../components/heatmap/AdvancedHeatmap";
import HeatmapLayerSelector, {
type LayerType,
} from "../components/heatmap/HeatmapLayerSelector";
import HeatmapFiltersComponent from "../components/heatmap/HeatmapFilters";
import HeatmapLegend from "../components/heatmap/HeatmapLegend";
const TACTIC_ORDER = [
"reconnaissance",
"resource-development",
"initial-access",
"execution",
"persistence",
"privilege-escalation",
"defense-evasion",
"credential-access",
"discovery",
"lateral-movement",
"collection",
"command-and-control",
"exfiltration",
"impact",
];
type ZoomLevel = "compact" | "normal" | "expanded";
export default function MatrixPage() {
const navigate = useNavigate();
// Layer selection state
const [activeLayer, setActiveLayer] = useState<LayerType>("coverage");
const [selectedActorId, setSelectedActorId] = useState<string | null>(null);
const [selectedCampaignId, setSelectedCampaignId] = useState<string | null>(null);
// Filter state
const [platforms, setPlatforms] = useState<string[]>([]);
const [selectedTactics, setSelectedTactics] = useState<string[]>([]);
const [minScore, setMinScore] = useState(0);
// Zoom
const [zoom, setZoom] = useState<ZoomLevel>("normal");
// Export dropdown
const [showExportMenu, setShowExportMenu] = useState(false);
// Build filter params
const filterParams: HeatmapFilterParams = useMemo(
() => ({
platforms: platforms.length > 0 ? platforms.join(",") : undefined,
tactics: selectedTactics.length > 0 ? selectedTactics.join(",") : undefined,
min_score: minScore > 0 ? minScore : undefined,
}),
[platforms, selectedTactics, minScore],
);
// Build query key based on active layer + selection
const queryKey = useMemo(() => {
const base = ["heatmap", activeLayer, filterParams];
if (activeLayer === "threat-actor") return [...base, selectedActorId];
if (activeLayer === "campaign") return [...base, selectedCampaignId];
return base;
}, [activeLayer, filterParams, selectedActorId, selectedCampaignId]);
// Fetch the active layer data
const {
data: layerData,
isLoading,
error,
} = useQuery<HeatmapLayer>({
queryKey,
queryFn: () => {
switch (activeLayer) {
case "coverage":
return getHeatmapCoverage(filterParams);
case "threat-actor":
if (!selectedActorId) return Promise.resolve({ name: "", versions: { attack: "", navigator: "", layer: "" }, domain: "", description: "", filters: { platforms: [] }, gradient: { colors: [], minValue: 0, maxValue: 0 }, techniques: [] } as HeatmapLayer);
return getHeatmapThreatActor(selectedActorId, filterParams);
case "detection-rules":
return getHeatmapDetectionRules(filterParams);
case "campaign":
if (!selectedCampaignId) return Promise.resolve({ name: "", versions: { attack: "", navigator: "", layer: "" }, domain: "", description: "", filters: { platforms: [] }, gradient: { colors: [], minValue: 0, maxValue: 0 }, techniques: [] } as HeatmapLayer);
return getHeatmapCampaign(selectedCampaignId, filterParams);
default:
return getHeatmapCoverage(filterParams);
}
},
enabled:
activeLayer === "coverage" ||
activeLayer === "detection-rules" ||
(activeLayer === "threat-actor" && !!selectedActorId) ||
(activeLayer === "campaign" && !!selectedCampaignId),
});
const techniques = layerData?.techniques || [];
// Handle cell click - navigate to technique detail
const handleCellClick = useCallback(
(techniqueId: string) => {
navigate(`/techniques/${techniqueId}`);
},
[navigate],
);
// Handle export
const handleExport = async (type: "download" | "url") => {
setShowExportMenu(false);
const layerId =
activeLayer === "threat-actor"
? selectedActorId ?? undefined
: activeLayer === "campaign"
? selectedCampaignId ?? undefined
: undefined;
if (type === "download") {
try {
const blob = await exportNavigatorJSON(activeLayer, layerId, filterParams);
const url = URL.createObjectURL(blob);
const a = document.createElement("a");
a.href = url;
a.download = `aegis_${activeLayer}_layer.json`;
document.body.appendChild(a);
a.click();
document.body.removeChild(a);
URL.revokeObjectURL(url);
} catch {
console.error("Failed to export Navigator JSON");
}
} else {
// Copy Navigator URL
const navigatorUrl = `https://mitre-attack.github.io/attack-navigator/#layerURL=${encodeURIComponent(
window.location.origin + `/api/v1/heatmap/export-navigator?layer=${activeLayer}${layerId ? `&layer_id=${layerId}` : ""}`
)}`;
navigator.clipboard.writeText(navigatorUrl);
}
};
// Zoom controls
const zoomIn = () => {
if (zoom === "compact") setZoom("normal");
else if (zoom === "normal") setZoom("expanded");
};
const zoomOut = () => {
if (zoom === "expanded") setZoom("normal");
else if (zoom === "normal") setZoom("compact");
};
return (
<div className="space-y-4">
{/* Header */}
<div>
<h1 className="text-2xl font-bold text-white">ATT&CK Matrix</h1>
<p className="mt-1 text-sm text-gray-400">
Advanced heatmap with multiple layers, filters, and ATT&CK Navigator export
</p>
</div>
{/* Toolbar: Layer Selector + Filters + Export + Zoom */}
<div className="rounded-xl border border-gray-800 bg-gray-900 p-4 space-y-3">
<div className="flex flex-wrap items-center justify-between gap-4">
{/* Layer selector */}
<HeatmapLayerSelector
activeLayer={activeLayer}
onLayerChange={setActiveLayer}
selectedActorId={selectedActorId}
onActorChange={setSelectedActorId}
selectedCampaignId={selectedCampaignId}
onCampaignChange={setSelectedCampaignId}
/>
{/* Right side: Export + Zoom */}
<div className="flex items-center gap-2">
{/* Export dropdown */}
<div className="relative">
<button
onClick={() => setShowExportMenu(!showExportMenu)}
className="flex items-center gap-1.5 rounded-lg border border-gray-700 bg-gray-800 px-3 py-1.5 text-xs font-medium text-gray-300 hover:border-cyan-500/50 hover:text-white"
>
<Download className="h-3.5 w-3.5" />
Export
</button>
{showExportMenu && (
<div className="absolute right-0 top-full z-30 mt-1 w-52 rounded-lg border border-gray-700 bg-gray-900 py-1 shadow-xl">
<button
onClick={() => handleExport("download")}
className="w-full px-3 py-2 text-left text-xs text-gray-300 hover:bg-gray-800"
>
Export Navigator JSON
</button>
<button
onClick={() => handleExport("url")}
className="w-full px-3 py-2 text-left text-xs text-gray-300 hover:bg-gray-800"
>
Copy Navigator URL
</button>
</div>
)}
</div>
{/* Zoom controls */}
<div className="flex items-center rounded-lg border border-gray-700 bg-gray-800">
<button
onClick={zoomOut}
disabled={zoom === "compact"}
className="px-2 py-1.5 text-gray-400 hover:text-white disabled:opacity-30"
>
<ZoomOut className="h-3.5 w-3.5" />
</button>
<span className="border-x border-gray-700 px-2 py-1 text-[10px] font-medium uppercase text-gray-400">
{zoom}
</span>
<button
onClick={zoomIn}
disabled={zoom === "expanded"}
className="px-2 py-1.5 text-gray-400 hover:text-white disabled:opacity-30"
>
<ZoomIn className="h-3.5 w-3.5" />
</button>
</div>
</div>
</div>
{/* Filters */}
<HeatmapFiltersComponent
platforms={platforms}
onPlatformsChange={setPlatforms}
selectedTactics={selectedTactics}
onTacticsChange={setSelectedTactics}
minScore={minScore}
onMinScoreChange={setMinScore}
availableTactics={TACTIC_ORDER}
/>
</div>
{/* Stats bar */}
{layerData && (
<div className="flex items-center gap-4 text-xs text-gray-500">
<span>
Layer: <span className="text-gray-300">{layerData.name}</span>
</span>
<span>
Techniques:{" "}
<span className="text-gray-300">
{techniques.filter((t) => t.enabled).length} active
</span>{" "}
/ {techniques.length} total
</span>
</div>
)}
{/* Loading / Error / Heatmap */}
{isLoading ? (
<div className="flex h-64 items-center justify-center">
<Loader2 className="h-8 w-8 animate-spin text-cyan-400" />
</div>
) : error ? (
<div className="flex h-64 flex-col items-center justify-center gap-2">
<AlertCircle className="h-10 w-10 text-red-400" />
<p className="text-red-400">Failed to load heatmap data</p>
</div>
) : (
<AdvancedHeatmap
techniques={techniques}
onCellClick={handleCellClick}
zoom={zoom}
/>
)}
{/* Legend */}
<HeatmapLegend layerType={activeLayer} />
</div>
);
}