fix: harden Defender threat risk audit

tests/test_setup.py

@@ -88,6 +88,27 @@ def make_data(detections=None, threats=None, defender=None, collection_complete=
     }
 
 
+class MemoryWaptAudit(object):
+    """Small fake implementing the WAPT 2.6 audit read/write contract."""
+
+    def __init__(self):
+        self.values = {}
+        self.history = []
+
+    def read_audit_data(self, section, key, default=None, ptype=None, include_expired_data=True):
+        return self.values.get((section, key), default)
+
+    def write_audit_data_if_changed(self, section, key, value, ptype=None,
+                                    value_date=None, expiration_date=None,
+                                    max_count=2, keep_days=None):
+        storage_key = (section, key)
+        previous = self.values.get(storage_key)
+        if previous != value:
+            self.values[storage_key] = value
+            self.history.append((storage_key, value, max_count, keep_days))
+        return previous
+
+
 class TestDefenderClassification(unittest.TestCase):
     def build(self, data, rules=None):
         return defender_setup.build_audit_snapshot(data, now=NOW, rules=rules)
@@ -123,21 +144,52 @@ class TestDefenderClassification(unittest.TestCase):
         self.assertEqual(snapshot["summary"]["audit_status"], "OK")
         self.assertEqual(snapshot["summary"]["resolved_count_6h"], 1)
 
-    def test_removed_threat_with_execution_observed_remains_unresolved_risk(self):
+    def test_removed_threat_with_execution_observed_is_resolved_but_requires_investigation(self):
         data = make_data(
             [make_detection(status_id=4)],
             [make_threat(severity_id=5, is_active=False, did_threat_execute=True)],
         )
         snapshot = self.build(data)
-        self.assertEqual(snapshot["active_threats"][0]["active_reason"], "execution_observed")
-        self.assertEqual(snapshot["active_threats"][0]["event_level"], 11)  # 12 - 3 removed + 2 observed
-        self.assertEqual(snapshot["summary"]["audit_status"], "ERROR")
+        item = snapshot["detections"][0]
+        self.assertEqual(snapshot["active_threats"], [])
+        self.assertEqual(item["remediation_state"], "removed")
+        self.assertFalse(item["active"])
+        self.assertTrue(item["execution_observed"])
+        self.assertTrue(item["requires_investigation"])
+        self.assertEqual(item["event_level"], 9)  # event severity excludes execution/investigation context
+        self.assertEqual(item["investigation_level"], 9)
+        self.assertEqual(snapshot["summary"]["max_active_level"], 0)
+        self.assertEqual(snapshot["summary"]["max_investigation_level"], 9)
+        self.assertEqual(snapshot["summary"]["audit_status"], "WARNING")
 
     def test_current_executing_state_keeps_removed_detection_active(self):
         event = make_detection(status_id=4)
         event["execution_status_id"] = 3
         snapshot = self.build(make_data([event], [make_threat(severity_id=2, is_active=False)]))
-        self.assertEqual(snapshot["active_threats"][0]["active_reason"], "execution_allowed_or_running")
+        self.assertEqual(snapshot["active_threats"][0]["active_reason"], "currently_executing")
+        self.assertTrue(snapshot["active_threats"][0]["execution_observed"])
+
+    def test_allowed_execution_is_not_active_but_requires_investigation(self):
+        event = make_detection(status_id=4)
+        event["execution_status_id"] = 2
+        snapshot = self.build(make_data([event], [make_threat(severity_id=5, is_active=False)]))
+        item = snapshot["detections"][0]
+        self.assertFalse(item["active"])
+        self.assertFalse(item["execution_observed"])
+        self.assertTrue(item["requires_investigation"])
+        self.assertEqual(item["execution_status"], "Allowed")
+
+    def test_blocked_detection_does_not_imply_execution_or_compromise(self):
+        event = make_detection(status_id=6)
+        event["execution_status_id"] = 1
+        snapshot = self.build(make_data([event], [make_threat(severity_id=5, is_active=False)]))
+        item = snapshot["detections"][0]
+        self.assertEqual(item["remediation_state"], "blocked")
+        self.assertFalse(item["active"])
+        self.assertFalse(item["execution_observed"])
+        self.assertFalse(item["requires_investigation"])
+        self.assertEqual(snapshot["summary"]["current_risk_level"], 0)
+        self.assertEqual(snapshot["summary"]["audit_status"], "OK")
 
     def test_severe_detected_threat_is_error(self):
         data = make_data(
@@ -155,16 +207,43 @@ class TestDefenderClassification(unittest.TestCase):
             [make_threat(severity_id=1, is_active=False)],
         )
         snapshot = self.build(data)
-        self.assertEqual(snapshot["active_threats"][0]["event_level"], 9)  # 4 + 3 + 2
-        self.assertEqual(snapshot["summary"]["audit_status"], "WARNING")
+        self.assertEqual(snapshot["active_threats"][0]["event_level"], 9)  # 4 + 2 active + 3 failure
+        self.assertEqual(snapshot["summary"]["remediation_failed"], True)
+        self.assertEqual(snapshot["summary"]["current_risk_level"], 9)
+        self.assertEqual(snapshot["summary"]["audit_status"], "ERROR")
 
-    def test_action_success_false_adds_failure_modifier_even_before_failed_status_rollup(self):
+    def test_action_success_false_without_error_is_not_assumed_to_be_failed_remediation(self):
         data = make_data(
             [make_detection(status_id=1, action_success=False)],
             [make_threat(severity_id=1, is_active=True)],
         )
         snapshot = self.build(data)
-        self.assertEqual(snapshot["detections"][0]["event_level"], 9)  # 4 + 2 active + 3 failure
+        self.assertEqual(snapshot["detections"][0]["remediation_state"], "detected")
+        self.assertEqual(snapshot["detections"][0]["event_level"], 6)  # 4 + 2 unresolved
+        self.assertFalse(snapshot["summary"]["remediation_failed"])
+        self.assertNotEqual(snapshot["summary"]["audit_status"], "ERROR")
+
+    def test_nonzero_error_with_failed_action_forces_error(self):
+        data = make_data(
+            [make_detection(status_id=1, action_success=False, error_code=2142207991)],
+            [make_threat(severity_id=1, is_active=True)],
+        )
+        snapshot = self.build(data)
+        self.assertEqual(snapshot["detections"][0]["remediation_state"], "failed")
+        self.assertEqual(snapshot["detections"][0]["event_level"], 9)
+        self.assertEqual(snapshot["summary"]["audit_status"], "ERROR")
+
+    def test_action_success_default_false_does_not_create_failure_correlation(self):
+        events = [
+            make_detection("default-false-a", status_id=1, action_success=False),
+            make_detection("default-false-b", status_id=1, action_success=False),
+        ]
+        snapshot = self.build(make_data(
+            events, [make_threat(severity_id=1, is_active=True)],
+        ))
+        self.assertTrue(all(item["remediation_state"] == "detected" for item in snapshot["detections"]))
+        rules = {item["rule"] for item in snapshot["correlations"]}
+        self.assertNotIn("repeated_remediation_failures", rules)
 
     def test_explicitly_allowed_threat_is_scored_and_correlated(self):
         data = make_data(
@@ -180,7 +259,10 @@ class TestDefenderClassification(unittest.TestCase):
         events = [make_detection("detection-a"), make_detection("detection-b", detected_hours_ago=2)]
         snapshot = self.build(make_data(events, [make_threat(severity_id=1, is_active=True)]))
         self.assertEqual(len(snapshot["detections"]), 2)
-        self.assertTrue(all(item["event_level"] == 8 for item in snapshot["detections"]))
+        self.assertTrue(all(item["event_level"] == 6 for item in snapshot["detections"]))
+        self.assertEqual(snapshot["summary"]["max_event_level_6h"], 6)
+        self.assertEqual(snapshot["summary"]["correlation_level"], 8)
+        self.assertEqual(snapshot["summary"]["current_risk_level"], 8)
         self.assertIn("same_threat_repeated", [item["rule"] for item in snapshot["correlations"]])
         self.assertEqual(snapshot["summary"]["audit_status"], "WARNING")
 
@@ -193,7 +275,7 @@ class TestDefenderClassification(unittest.TestCase):
         failures = [item for item in snapshot["correlations"] if item["rule"] == "repeated_remediation_failures"]
         self.assertEqual(len(failures), 1)
         self.assertEqual(failures[0]["count"], 2)
-        self.assertEqual(failures[0]["threat_ids"], ["100"])
+        self.assertEqual(failures[0]["count"], 2)
 
     def test_old_active_threat_remains_visible_outside_six_hour_window(self):
         event = make_detection(detected_hours_ago=10, changed_hours_ago=10)
@@ -257,6 +339,33 @@ class TestDefenderClassification(unittest.TestCase):
         self.assertEqual(snapshot["summary"]["detection_count_6h"], 2)
         self.assertIn("same_threat_repeated", [item["rule"] for item in snapshot["correlations"]])
 
+    def test_threat_rollup_isactive_does_not_reopen_a_resolved_detectionid(self):
+        events = [
+            make_detection("still-detected", status_id=1),
+            make_detection("already-removed", status_id=4, detected_hours_ago=2),
+        ]
+        snapshot = self.build(make_data(
+            events, [make_threat(severity_id=1, is_active=True)],
+        ))
+        active_ids = [item["detection_id"] for item in snapshot["active_threats"]]
+        removed = next(item for item in snapshot["detections"] if item["detection_id"] == "already-removed")
+        self.assertEqual(active_ids, ["still-detected"])
+        self.assertFalse(removed["active"])
+        self.assertEqual(removed["remediation_state"], "removed")
+        self.assertTrue(removed["requires_investigation"])  # aggregated IsActive is a conflict signal
+
+    def test_fallback_event_identity_survives_status_change_and_is_flagged(self):
+        first = make_detection(detection_id=None, status_id=1)
+        first.pop("detection_id")
+        first_snapshot = self.build(make_data([first], [make_threat(severity_id=1, is_active=True)]))
+        second = dict(first)
+        second["status_id"] = 4
+        second["changed_at"] = NOW.isoformat().replace("+00:00", "Z")
+        second_snapshot = self.build(make_data([second], [make_threat(severity_id=1, is_active=False)]))
+        self.assertEqual(first_snapshot["detections"][0]["event_id"], second_snapshot["detections"][0]["event_id"])
+        self.assertFalse(second_snapshot["detections"][0]["identity_complete"])
+        self.assertGreater(second_snapshot["summary"]["identity_incomplete_count"], 0)
+
     def test_old_detection_with_recent_status_change_enters_window(self):
         event = make_detection(status_id=6, detected_hours_ago=20, changed_hours_ago=1)
         snapshot = self.build(make_data([event], [make_threat(severity_id=5, is_active=False)]))
@@ -291,7 +400,7 @@ class TestDefenderClassification(unittest.TestCase):
             make_detection("threat-b", threat_id="200", status_id=1, detected_hours_ago=2),
         ]
         threats = [
-            make_threat("100", severity_id=5, is_active=True, did_threat_execute=True),
+            make_threat("100", severity_id=5, is_active=True, did_threat_execute=False),
             make_threat("200", severity_id=1, is_active=True),
         ]
         snapshot = self.build(make_data(events, threats, defender={"realtime_enabled": False}))
@@ -299,8 +408,10 @@ class TestDefenderClassification(unittest.TestCase):
         self.assertIn("multiple_threats", rules)
         self.assertIn("severe_with_protection_disabled", rules)
         severe = next(item for item in snapshot["detections"] if item["threat_id"] == "100")
-        self.assertEqual(severe["event_level"], 16)  # 12 + 2 active + 2 realtime disabled
-        self.assertTrue(severe["did_threat_execute"])
+        self.assertEqual(severe["event_level"], 14)  # host protection is not part of event score
+        self.assertEqual(snapshot["summary"]["correlation_level"], 16)
+        self.assertEqual(snapshot["summary"]["current_risk_level"], 16)
+        self.assertFalse(severe["did_threat_execute"])
 
     def test_resources_are_reduced_to_types_and_counts(self):
         event = make_detection(resources=[r"file:_C:\Users\synthetic\private.txt", "registry:synthetic"])
@@ -320,6 +431,78 @@ class TestDefenderClassification(unittest.TestCase):
         self.assertEqual(snapshot["summary"]["max_event_level_6h"], 9)
         self.assertEqual(snapshot["summary"]["audit_status"], "OK")
 
+    def test_current_protection_change_does_not_rewrite_event_level(self):
+        event = make_detection(status_id=4)
+        threat = make_threat(severity_id=5, is_active=False)
+        healthy = self.build(make_data([event], [threat]))
+        degraded = self.build(make_data([event], [threat], defender={"realtime_enabled": False}))
+        self.assertEqual(healthy["detections"][0]["event_level"], 9)
+        self.assertEqual(degraded["detections"][0]["event_level"], 9)
+        self.assertEqual(healthy["summary"]["current_risk_level"], 0)
+        self.assertGreater(degraded["summary"]["current_risk_level"], 0)
+
+    def test_old_resolved_event_with_execution_remains_investigation_record(self):
+        event = make_detection(status_id=4, detected_hours_ago=24 * 45, changed_hours_ago=24 * 45)
+        snapshot = self.build(make_data(
+            [event], [make_threat(severity_id=5, is_active=False, did_threat_execute=True)],
+        ))
+        self.assertEqual(len(snapshot["detections"]), 1)
+        self.assertFalse(snapshot["detections"][0]["active"])
+        self.assertTrue(snapshot["detections"][0]["requires_investigation"])
+        self.assertEqual(snapshot["summary"]["max_active_level"], 0)
+        self.assertGreater(snapshot["summary"]["max_investigation_level"], 0)
+
+    def test_distinct_event_and_correlation_levels_have_scoring_reasons(self):
+        events = [
+            make_detection("corr-a", threat_id="100"),
+            make_detection("corr-b", threat_id="200", detected_hours_ago=2),
+        ]
+        snapshot = self.build(make_data(
+            events,
+            [make_threat("100", severity_id=1, is_active=True), make_threat("200", severity_id=1, is_active=True)],
+        ))
+        self.assertEqual(snapshot["summary"]["max_event_level_6h"], 6)
+        self.assertEqual(snapshot["summary"]["correlation_level"], 8)
+        self.assertEqual(snapshot["summary"]["max_active_level"], 6)
+        self.assertEqual(snapshot["summary"]["current_risk_level"], 8)
+        self.assertTrue(snapshot["detections"][0]["event_level_reasons"])
+        self.assertTrue(snapshot["correlations"][0]["scoring_reasons"])
+
+    def test_investigation_state_is_retained_beyond_collection_history(self):
+        event = make_detection(status_id=4, detected_hours_ago=24 * 45, changed_hours_ago=24 * 45)
+        snapshot = self.build(make_data(
+            [event], [make_threat(severity_id=5, is_active=False, did_threat_execute=True)],
+        ))
+        self.assertEqual(len(snapshot["detections"]), 1)
+        self.assertFalse(snapshot["detections"][0]["active"])
+        self.assertTrue(snapshot["detections"][0]["requires_investigation"])
+        self.assertEqual(snapshot["summary"]["max_active_level"], 0)
+        self.assertGreater(snapshot["summary"]["max_investigation_level"], 0)
+
+    def test_serialized_audit_values_shrink_for_small_medium_and_large_sets(self):
+        for size in (1, 100, 1000):
+            with self.subTest(detections=size):
+                events = [
+                    make_detection("size-%d" % index, threat_id=str(index + 100))
+                    for index in range(size)
+                ]
+                threats = [
+                    make_threat(str(index + 100), severity_id=1, is_active=True)
+                    for index in range(size)
+                ]
+                snapshot = self.build(make_data(events, threats))
+                legacy = {
+                    "snapshot": snapshot,
+                    "summary": snapshot["summary"],
+                    "detections": snapshot["detections"],
+                    "active_threats": snapshot["active_threats"],
+                    "correlations": snapshot["correlations"],
+                }
+                compact = defender_setup.build_audit_records(snapshot)
+                old_size = len(json.dumps(legacy, separators=(",", ":")).encode("utf-8"))
+                new_size = len(json.dumps(compact, separators=(",", ":")).encode("utf-8"))
+                self.assertLess(new_size, old_size)
+
     def test_classifier_modifiers_are_configurable_and_clamped(self):
         rules = {
             "active_modifier": 20,
@@ -334,6 +517,17 @@ class TestDefenderClassification(unittest.TestCase):
         self.assertEqual(defender_setup._wapt_status(6, [], True, False, True, 0, 0), "OK")
         self.assertEqual(defender_setup._wapt_status(7, [], True, False, True, 0, 0), "WARNING")
         self.assertEqual(defender_setup._wapt_status(10, [], True, False, True, 0, 0), "ERROR")
+        self.assertEqual(
+            defender_setup._wapt_status(9, [], True, False, True, 0, 0, remediation_failed=True),
+            "ERROR",
+        )
+        self.assertEqual(
+            defender_setup._wapt_status(
+                9, [], True, False, True, 0, 0, remediation_failed=True,
+                rules={"failed_remediation_forces_error": False},
+            ),
+            "WARNING",
+        )
 
 
 class TestDefenderCollectionAndWaptAdapter(unittest.TestCase):
@@ -354,7 +548,7 @@ class TestDefenderCollectionAndWaptAdapter(unittest.TestCase):
         self.assertNotIn("Bypass", args)
         self.assertEqual(run.call_args.kwargs["timeout"], defender_setup.POWERSHELL_TIMEOUT_SECONDS)
 
-    def test_audit_timeout_returns_error_and_stores_failed_snapshot(self):
+    def test_audit_timeout_returns_error_and_stores_failed_summary(self):
         wapt = Mock()
         with patch.object(defender_setup, "WAPT", wapt, create=True), patch.object(
             defender_setup, "collect_defender", side_effect=defender_setup.DefenderCollectionError("powershell_timeout")
@@ -362,31 +556,89 @@ class TestDefenderCollectionAndWaptAdapter(unittest.TestCase):
             result = defender_setup.audit()
         self.assertEqual(result, "ERROR")
         stored = [call.args for call in wapt.write_audit_data_if_changed.call_args_list]
-        snapshot_call = next(args for args in stored if args[1] == "snapshot")
-        self.assertFalse(snapshot_call[2]["collection_complete"])
-        self.assertEqual(snapshot_call[2]["summary"]["audit_status"], "ERROR")
+        summary_call = next(args for args in stored if args[1] == "summary")
+        self.assertFalse(summary_call[2]["collection_complete"])
+        self.assertEqual(summary_call[2]["summary"]["audit_status"], "ERROR")
 
     def test_wapt_audit_api_uses_official_signature_and_json_value(self):
         wapt = Mock()
-        value = {"schema_version": 2, "detections": [], "correlations": []}
+        value = {"schema_version": 3, "summary": {}}
         with patch.object(defender_setup, "WAPT", wapt, create=True):
-            defender_setup.save_audit("snapshot", value)
+            defender_setup.save_audit("summary", value)
         wapt.write_audit_data_if_changed.assert_called_once_with(
-            "Microsoft Defender", "snapshot", value,
-            max_count=512, keep_days=90,
+            "Microsoft Defender", "summary", value,
+            max_count=defender_setup.AUDIT_SUMMARY_HISTORY_MAX_COUNT, keep_days=90,
         )
 
-    def test_audit_hook_returns_wapt_status_and_writes_v2_snapshot(self):
+    def test_audit_hook_returns_wapt_status_and_writes_compact_v3_records(self):
         wapt = Mock()
         with patch.object(defender_setup, "WAPT", wapt, create=True), patch.object(
-            defender_setup, "collect_defender", return_value=make_data()
+            defender_setup, "collect_defender", return_value=make_data(
+                [make_detection()], [make_threat(severity_id=1, is_active=True)],
+            )
         ):
             result = defender_setup.audit()
         self.assertEqual(result, "OK")
         written_keys = [call.args[1] for call in wapt.write_audit_data_if_changed.call_args_list]
-        self.assertEqual(written_keys, ["snapshot", "summary", "detections", "active_threats", "correlations"])
-        snapshot_call = wapt.write_audit_data_if_changed.call_args_list[0]
-        self.assertEqual(snapshot_call.args[2]["schema_version"], 2)
+        self.assertEqual(written_keys[:2], ["summary", "correlations"])
+        self.assertTrue(written_keys[2].startswith("event."))
+        summary_call = wapt.write_audit_data_if_changed.call_args_list[0]
+        self.assertEqual(summary_call.args[2]["schema_version"], 3)
+        correlations_call = wapt.write_audit_data_if_changed.call_args_list[1]
+        self.assertEqual(correlations_call.args[2]["schema_version"], 3)
+        wapt.read_audit_data.assert_called_once_with(
+            "Microsoft Defender", written_keys[2], default=None, include_expired_data=False,
+        )
+
+    def test_event_envelope_is_provider_neutral_and_tracks_resolution(self):
+        initial = defender_setup.build_audit_snapshot(make_data(
+            [make_detection("stable-detection", status_id=1)],
+            [make_threat(severity_id=5, is_active=True)],
+        ), now=NOW)
+        first = defender_setup.build_audit_records(initial)
+        event_key = next(key for key in first if key.startswith("event."))
+        event = first[event_key]
+        self.assertEqual(event["schema_version"], 1)
+        self.assertEqual(event["source"], "comi-audit-defender")
+        self.assertEqual(event["event_type"], "threat.detected")
+        self.assertNotIn("threat_name", event)
+
+        resolved_input = make_detection("stable-detection", status_id=4, changed_hours_ago=0)
+        resolved = defender_setup.build_audit_snapshot(make_data(
+            [resolved_input], [make_threat(severity_id=5, is_active=False)],
+        ), now=NOW)
+        previous = {event["event_id"]: event}
+        second = defender_setup.build_audit_records(resolved, previous_events=previous)
+        resolved_event = second[event_key]
+        self.assertEqual(resolved_event["event_id"], event["event_id"])
+        self.assertEqual(resolved_event["event_type"], "threat.resolved")
+        self.assertEqual(resolved_event["remediation_state"], "removed")
+
+    def test_repeated_audit_writes_only_changed_values_and_historizes_event_transition(self):
+        fake_wapt = MemoryWaptAudit()
+        detection = make_detection("audit-transition", status_id=1)
+        first_snapshot = defender_setup.build_audit_snapshot(make_data(
+            [detection], [make_threat(severity_id=1, is_active=True)],
+        ), now=NOW)
+        with patch.object(defender_setup, "WAPT", fake_wapt, create=True):
+            defender_setup._save_snapshot(first_snapshot)
+            first_write_count = len(fake_wapt.history)
+            defender_setup._save_snapshot(first_snapshot)
+            self.assertEqual(len(fake_wapt.history), first_write_count)
+
+            resolved = make_detection("audit-transition", status_id=4, changed_hours_ago=0)
+            second_snapshot = defender_setup.build_audit_snapshot(make_data(
+                [resolved], [make_threat(severity_id=1, is_active=False)],
+            ), now=NOW)
+            defender_setup._save_snapshot(second_snapshot)
+
+        event_key = next(key for section, key in fake_wapt.values if key.startswith("event."))
+        stored_event = fake_wapt.values[("Microsoft Defender", event_key)]
+        self.assertEqual(stored_event["event_type"], "threat.resolved")
+        self.assertEqual(stored_event["remediation_state"], "removed")
+        self.assertNotIn(("Microsoft Defender", "snapshot"), fake_wapt.values)
+        self.assertNotIn(("Microsoft Defender", "detections"), fake_wapt.values)
+        self.assertEqual(fake_wapt.values[("Microsoft Defender", "summary")]["schema_version"], 3)
 
     def test_control_declares_windows_wapt_26_and_15_minute_audit(self):
         fields = {}
@@ -395,6 +647,7 @@ class TestDefenderCollectionAndWaptAdapter(unittest.TestCase):
                 key, value = line.split(":", 1)
                 fields[key.strip()] = value.strip()
         self.assertEqual(fields["package"], "comit-audit-defender")
+        self.assertEqual(fields["version"], "2.0-2")
         self.assertEqual(fields["min_wapt_version"], "2.6")
         self.assertEqual(fields["audit_schedule"], "15m")
         self.assertEqual(fields["target_os"], "windows")