diff --git a/backend/app/services/d3fend_import_service.py b/backend/app/services/d3fend_import_service.py index ed831c0..1cbd37a 100644 --- a/backend/app/services/d3fend_import_service.py +++ b/backend/app/services/d3fend_import_service.py @@ -583,6 +583,8 @@ def import_d3fend_mappings(db: Session) -> dict[str, int]: created = 0 # Assign skipped = 0 skipped = 0 + # Assign new_technique_ids = set() + new_technique_ids: set[str] = set() # Get all ATT&CK techniques from the DB attack_techniques = db.query(Technique).all() @@ -649,6 +651,12 @@ def import_d3fend_mappings(db: Session) -> dict[str, int]: db.add(mapping) # Assign created = 1 created += 1 + new_technique_ids.add(mitre_id) + + if new_technique_ids: + db.query(Technique).filter( + Technique.mitre_id.in_(new_technique_ids) + ).update({"review_required": True}, synchronize_session=False) # Commit all pending changes to the database db.commit() diff --git a/backend/tests/test_d3fend_import_service.py b/backend/tests/test_d3fend_import_service.py new file mode 100644 index 0000000..f068d0a --- /dev/null +++ b/backend/tests/test_d3fend_import_service.py @@ -0,0 +1,42 @@ +"""Tests for D3FEND mapping import — review_required trigger parity. + +Every other import source (Atomic Red Team, Caldera, Elastic, Sigma, LOLBAS) +flags a technique with review_required=True when it gains new content, so a +lead knows to look at it. D3FEND mapping import was missing this — this test +locks in the fix. +""" + +from app.models.defensive_technique import DefensiveTechnique +from app.models.technique import Technique +from app.services.d3fend_import_service import import_d3fend_mappings + + +def test_import_d3fend_mappings_sets_review_required(db): + technique = Technique(mitre_id="T1590", name="Gather Victim Network Information", review_required=False) + db.add(technique) + db.add(DefensiveTechnique(d3fend_id="D3-NTA", name="Network Traffic Analysis")) + db.commit() + + result = import_d3fend_mappings(db) + + assert result["created"] >= 1 + db.refresh(technique) + assert technique.review_required is True + + +def test_import_d3fend_mappings_skips_existing_without_reflagging(db): + technique = Technique(mitre_id="T1590", name="Gather Victim Network Information", review_required=False) + db.add(technique) + db.add(DefensiveTechnique(d3fend_id="D3-NTA", name="Network Traffic Analysis")) + db.commit() + + import_d3fend_mappings(db) + db.refresh(technique) + technique.review_required = False + db.commit() + + result = import_d3fend_mappings(db) + + assert result["created"] == 0 + db.refresh(technique) + assert technique.review_required is False