diff --git a/backend/app/services/heatmap_service.py b/backend/app/services/heatmap_service.py index b3d10e0..c8927e9 100644 --- a/backend/app/services/heatmap_service.py +++ b/backend/app/services/heatmap_service.py @@ -353,6 +353,17 @@ def build_coverage_layer( layer["techniques"].append({ # Literal argument value "techniqueID": tech.mitre_id, + # Technique display name — not part of the official ATT&CK + # Navigator schema, but harmless extra data for Navigator + # imports and lets Aegis's own heatmap UI show names inline + # instead of only ID + hover tooltip. + "name": tech.name, + # Only the coverage layer has a real TechniqueStatus concept — + # exposed so the frontend can color cells by the exact same + # discrete status palette the Techniques page uses, instead of + # a continuous score-derived hex. The other 3 layer types don't + # set this (their score doesn't correspond to a TechniqueStatus). + "status": tech.status_global.value, # Literal argument value "tactic": _format_tactic(tech.tactic), # Literal argument value @@ -491,6 +502,11 @@ def build_threat_actor_layer( ) layer["techniques"].append({ "techniqueID": tech.mitre_id, + "name": tech.name, + # Every appended row here is an actor-used technique (non-actor + # ones are skipped above), so status_global is meaningful here + # too — see the coverage layer's "status" field for why. + "status": tech.status_global.value, "tactic": _format_tactic(tech.tactic), "color": _score_to_color(score), "score": score, @@ -586,6 +602,7 @@ def build_detection_rules_layer( layer["techniques"].append({ # Literal argument value "techniqueID": tech.mitre_id, + "name": tech.name, # Literal argument value "tactic": _format_tactic(tech.tactic), # Literal argument value @@ -754,6 +771,7 @@ def build_campaign_layer( layer["techniques"].append({ # Literal argument value "techniqueID": mitre_id, + "name": tech.name, # Literal argument value "tactic": _format_tactic(tech.tactic), # Literal argument value diff --git a/backend/tests/test_heatmap_technique_name.py b/backend/tests/test_heatmap_technique_name.py new file mode 100644 index 0000000..4c055f0 --- /dev/null +++ b/backend/tests/test_heatmap_technique_name.py @@ -0,0 +1,85 @@ +"""Every heatmap layer must include the technique's display name. + +The frontend heatmap previously showed only the MITRE ID on each cell +(name only available via hover tooltip), unlike the Techniques page which +always showed both. Unifying the two pages' cell style requires the name +to be available on every layer's technique entries, not just derived +client-side (HeatmapTechnique has no name field to derive it from). +""" + +from app.models.campaign import Campaign, CampaignTest +from app.models.detection_rule import DetectionRule +from app.models.enums import TestState +from app.models.technique import Technique +from app.models.test import Test +from app.models.threat_actor import ThreatActor, ThreatActorTechnique +from app.services.heatmap_service import ( + build_campaign_layer, + build_coverage_layer, + build_detection_rules_layer, + build_threat_actor_layer, +) + + +def _seed_technique(db, mitre_id="T1059", name="Command and Scripting Interpreter"): + tech = Technique(mitre_id=mitre_id, name=name, tactic="execution", platforms=["windows"]) + db.add(tech) + db.flush() + return tech + + +def test_coverage_layer_includes_technique_name_and_status(db): + tech = _seed_technique(db) + db.commit() + + layer = build_coverage_layer(db) + + entry = next(t for t in layer["techniques"] if t["techniqueID"] == tech.mitre_id) + assert entry["name"] == tech.name + assert entry["status"] == "not_evaluated" + + +def test_threat_actor_layer_includes_technique_name_and_status(db): + tech = _seed_technique(db) + actor = ThreatActor(name="APT-Test") + db.add(actor) + db.flush() + db.add(ThreatActorTechnique(threat_actor_id=actor.id, technique_id=tech.id)) + db.commit() + + layer = build_threat_actor_layer(db, actor.id) + + entry = next(t for t in layer["techniques"] if t["techniqueID"] == tech.mitre_id) + assert entry["name"] == tech.name + assert entry["status"] == "not_evaluated" + + +def test_detection_rules_layer_includes_technique_name(db): + tech = _seed_technique(db) + db.add(DetectionRule( + mitre_technique_id=tech.mitre_id, title="Rule 1", source="sigma", + rule_content="title: Rule 1", rule_format="sigma", + )) + db.commit() + + layer = build_detection_rules_layer(db) + + entry = next(t for t in layer["techniques"] if t["techniqueID"] == tech.mitre_id) + assert entry["name"] == tech.name + + +def test_campaign_layer_includes_technique_name(db, red_lead_user): + tech = _seed_technique(db) + campaign = Campaign(name="Test Campaign", type="custom", status="active", created_by=red_lead_user.id) + db.add(campaign) + db.flush() + test = Test(technique_id=tech.id, name="A test", state=TestState.validated, created_by=red_lead_user.id) + db.add(test) + db.flush() + db.add(CampaignTest(campaign_id=campaign.id, test_id=test.id, order_index=0)) + db.commit() + + layer = build_campaign_layer(db, campaign.id) + + entry = next(t for t in layer["techniques"] if t["techniqueID"] == tech.mitre_id) + assert entry["name"] == tech.name