# -*- coding: utf-8 -*-

"""Read-only Microsoft Defender audit for WAPT Enterprise 2.6."""

from setuphelpers import *  # noqa: F401,F403 - WAPT injects the audit API here.

import json
import os
import subprocess
from collections import Counter
from datetime import datetime, timedelta, timezone


AUDIT_CATEGORY = "Microsoft Defender"
HISTORY_DAYS = 30
AUDIT_HISTORY_DAYS = 90
AUDIT_HISTORY_MAX_COUNT = 512
POWERSHELL_TIMEOUT_SECONDS = 90

# MPTHREAT_SEVERITY values used by the Defender threat API. SeverityID on the
# WMI catalog uses 0..4 instead; PowerShell normalizes that representation.
SEVERITY_BASE_LEVELS = {
    0: 5,   # Unknown
    1: 4,   # Low
    2: 6,   # Moderate
    4: 9,   # High
    5: 12,  # Severe
}
SEVERITY_NAMES = {
    0: "Unknown",
    1: "Low",
    2: "Moderate",
    4: "High",
    5: "Severe",
}

# Values exposed by the MSFT_MpThreatDetection class used by the Defender
# PowerShell cmdlet. Quarantine remains unresolved; it is never called clean.
THREAT_STATUSES = {
    0: "Unknown",
    1: "Detected",
    2: "Cleaned",
    3: "Quarantined",
    4: "Removed",
    5: "Allowed",
    6: "Blocked",
    102: "QuarantineFailed",
    103: "RemoveFailed",
    104: "AllowFailed",
    105: "Abandoned",
    107: "BlockedFailed",
}
RESOLVED_STATUS_IDS = {2, 4, 6}
FAILED_STATUS_IDS = {102, 103, 104, 107}

DEFAULT_RULES = {
    "window_hours": 6,
    "signature_max_age_days": 3,
    "active_modifier": 2,
    "repeated_modifier": 2,
    "failed_remediation_modifier": 3,
    "status_modifiers": {
        2: -3,  # cleaned
        3: -3,  # quarantined, but still unresolved
        4: -3,  # removed
        5: 3,   # explicitly allowed
        6: -2,  # blocked
        102: 3,
        103: 3,
        104: 3,
        107: 3,
    },
    "protection_modifiers": {
        "service_disabled": 2,
        "antivirus_disabled": 2,
        "realtime_disabled": 2,
        "behavior_monitor_disabled": 1,
        "ioav_disabled": 1,
        "signatures_stale": 1,
        "signature_age_unknown": 1,
    },
    "repeated_threshold": 2,
    "remediation_failure_threshold": 2,
    "multiple_threats_threshold": 2,
    "audit_history_days": AUDIT_HISTORY_DAYS,
    "audit_history_max_count": AUDIT_HISTORY_MAX_COUNT,
    "expected_active_modes": ("Normal",),
    "passive_modes": ("Passive", "SxS Passive Mode", "EDR Block Mode"),
    "correlations": {
        "same_threat_repeated": True,
        "multiple_threats": True,
        "severe_with_protection_disabled": True,
        "repeated_remediation_failures": True,
        "explicitly_allowed": True,
    },
}


class DefenderCollectionError(RuntimeError):
    """Collection error with a safe, non-sensitive reason code."""

    def __init__(self, code):
        self.code = code
        RuntimeError.__init__(self, code)


def install():
    """Keep the existing package install hook; no Defender changes are made."""
    print("Microsoft Defender audit package installed")


POWERSHELL_SCRIPT = r'''
$ErrorActionPreference = 'Stop'
$WarningPreference = 'SilentlyContinue'
$ProgressPreference = 'SilentlyContinue'

function Get-SafeInt($Value) {
    if ($null -eq $Value -or "$Value" -eq '') { return $null }
    try { return [int]$Value } catch { return $null }
}

function Format-Utc($Value) {
    if ($null -eq $Value -or "$Value" -eq '') { return $null }
    try {
        if ($Value -is [datetime]) { $date = $Value }
        else {
            $date = [datetime]::MinValue
            if (-not [datetime]::TryParse([string]$Value, [ref]$date)) {
                return $null
            }
        }
        return $date.ToUniversalTime().ToString('o')
    } catch { return $null }
}

function Get-ResourceSummary($Values) {
    $items = @($Values | Where-Object { $null -ne $_ -and "$($_)" -ne '' })
    $kinds = @(
        $items | ForEach-Object {
            $text = [string]$_
            $kind = ($text -split ':', 2)[0].ToLowerInvariant()
            if ($kind -notin @('file', 'process', 'registry', 'service', 'url')) {
                $kind = 'other'
            }
            $kind
        } | Sort-Object -Unique
    )
    return [PSCustomObject]@{
        resource_count = [int]$items.Count
        resource_kinds = $kinds
    }
}

$now = [datetime]::UtcNow
$historyCutoff = $now.AddDays(-30)
$status = Get-MpComputerStatus -ErrorAction Stop

$catalog = @(
    Get-MpThreat -ErrorAction Stop | ForEach-Object {
        $severityId = $null
        $severitySource = $null
        $severityProperty = $_.PSObject.Properties['ThreatSeverityID']
        if ($null -ne $severityProperty -and $null -ne $severityProperty.Value) {
            $severityId = Get-SafeInt $severityProperty.Value
            $severitySource = 'ThreatSeverityID'
        } else {
            $severityProperty = $_.PSObject.Properties['SeverityID']
            if ($null -ne $severityProperty -and $null -ne $severityProperty.Value) {
                $rawSeverity = Get-SafeInt $severityProperty.Value
                # MSFT_MpThreatCatalog uses 3=High and 4=Severe. Normalize to
                # MPTHREAT_SEVERITY's documented 4=High and 5=Severe values.
                $catalogSeverityMap = @{ 0 = 0; 1 = 1; 2 = 2; 3 = 4; 4 = 5 }
                if ($null -ne $rawSeverity -and $catalogSeverityMap.ContainsKey($rawSeverity)) {
                    $severityId = $catalogSeverityMap[$rawSeverity]
                }
                $severitySource = 'SeverityID'
            }
        }

        $active = $null
        $activeProperty = $_.PSObject.Properties['IsActive']
        if ($null -ne $activeProperty -and $null -ne $activeProperty.Value) {
            $active = [bool]$activeProperty.Value
        }

        [PSCustomObject]@{
            threat_id = if ($null -ne $_.ThreatID) { [string]$_.ThreatID } else { $null }
            threat_name = if ($_.ThreatName) { ([string]$_.ThreatName).Substring(0, [Math]::Min(256, ([string]$_.ThreatName).Length)) } else { $null }
            threat_severity_id = $severityId
            severity_source = $severitySource
            is_active = $active
            did_threat_execute = if ($null -ne $_.DidThreatExecute) { [bool]$_.DidThreatExecute } else { $null }
        }
    }
)

$activeThreatIds = @{}
foreach ($item in $catalog) {
    if ($item.threat_id -and $item.is_active -eq $true) {
        $activeThreatIds[[string]$item.threat_id] = $true
    }
}

$allDetections = @(Get-MpThreatDetection -ErrorAction Stop)
$selectedDetections = @(
    $allDetections | Where-Object {
        $initial = $_.InitialDetectionTime
        $changed = $_.LastThreatStatusChangeTime
        $withinHistory = $false
        if ($initial -and $initial.ToUniversalTime() -ge $historyCutoff) { $withinHistory = $true }
        if ($changed -and $changed.ToUniversalTime() -ge $historyCutoff) { $withinHistory = $true }
        $id = if ($null -ne $_.ThreatID) { [string]$_.ThreatID } else { '' }
        $statusId = Get-SafeInt $_.ThreatStatusID
        $isKnownResolved = $statusId -in @(2, 4, 6)
        $knownActive = $activeThreatIds.ContainsKey($id)
        $withinHistory -or $knownActive -or -not $isKnownResolved
    } | ForEach-Object {
        $resourceSummary = Get-ResourceSummary $_.Resources
        [PSCustomObject]@{
            detection_id = if ($_.DetectionID) { [string]$_.DetectionID } else { $null }
            threat_id = if ($null -ne $_.ThreatID) { [string]$_.ThreatID } else { $null }
            status_id = Get-SafeInt $_.ThreatStatusID
            detected_at = Format-Utc $_.InitialDetectionTime
            changed_at = Format-Utc $_.LastThreatStatusChangeTime
            remediation_at = Format-Utc $_.RemediationTime
            error_code = Get-SafeInt $_.ThreatStatusErrorCode
            action_success = if ($null -ne $_.ActionSuccess) { [bool]$_.ActionSuccess } else { $null }
            execution_status_id = Get-SafeInt $_.CurrentThreatExecutionStatusID
            resource_count = $resourceSummary.resource_count
            resource_kinds = @($resourceSummary.resource_kinds)
        }
    }
)

$result = [PSCustomObject]@{
    collection_complete = $true
    history_detection_count = [int]$allDetections.Count
    returned_detection_count = [int]$selectedDetections.Count
    defender = [PSCustomObject]@{
        running_mode = $status.AMRunningMode
        antivirus_enabled = $status.AntivirusEnabled
        realtime_enabled = $status.RealTimeProtectionEnabled
        behavior_monitor_enabled = $status.BehaviorMonitorEnabled
        ioav_enabled = $status.IoavProtectionEnabled
        signature_age = Get-SafeInt $status.AntivirusSignatureAge
        signature_version = $status.AntivirusSignatureVersion
        signature_last_updated = Format-Utc $status.AntivirusSignatureLastUpdated
        engine_version = $status.AMEngineVersion
        service_enabled = $status.AMServiceEnabled
    }
    threats = @($catalog)
    detections = @($selectedDetections)
}

ConvertTo-Json -InputObject $result -Depth 8 -Compress
'''


def _copy_rules(rules=None):
    result = dict(DEFAULT_RULES)
    for key in ("status_modifiers", "protection_modifiers", "correlations"):
        result[key] = dict(DEFAULT_RULES[key])
    if rules:
        for key, value in rules.items():
            if key in ("status_modifiers", "protection_modifiers", "correlations"):
                result[key].update(value)
            else:
                result[key] = value
    return result


def _as_list(value):
    if value is None:
        return []
    if isinstance(value, (list, tuple)):
        return list(value)
    return [value]


def _first(mapping, *keys):
    if not isinstance(mapping, dict):
        return None
    for key in keys:
        if key in mapping:
            return mapping[key]
    return None


def _to_int(value, default=None):
    if value is None or isinstance(value, bool):
        return default
    try:
        return int(value)
    except (TypeError, ValueError, OverflowError):
        return default


def _to_bool(value):
    if isinstance(value, bool):
        return value
    if value in (0, 1):
        return bool(value)
    if isinstance(value, str):
        lowered = value.strip().lower()
        if lowered in ("true", "yes", "1"):
            return True
        if lowered in ("false", "no", "0"):
            return False
    return None


def _threat_id(value):
    if value is None:
        return None
    text = str(value).strip()
    return text or None


def _parse_datetime(value):
    if isinstance(value, datetime):
        parsed = value
    elif isinstance(value, str) and value.strip():
        text = value.strip()
        if text.endswith("Z"):
            text = text[:-1] + "+00:00"
        try:
            parsed = datetime.fromisoformat(text)
        except ValueError:
            return None
    else:
        return None
    if parsed.tzinfo is None:
        parsed = parsed.replace(tzinfo=timezone.utc)
    return parsed.astimezone(timezone.utc)


def _iso(value):
    parsed = _parse_datetime(value)
    return parsed.isoformat().replace("+00:00", "Z") if parsed else None


def _catalog_severity(item):
    severity_id = _to_int(_first(item, "threat_severity_id", "ThreatSeverityID"))
    source = _first(item, "severity_source", "SeveritySource")
    if severity_id is None:
        severity_id = _to_int(_first(item, "severity_id", "SeverityID"))
        source = source or "SeverityID"
    # WMI catalog values differ from MPTHREAT_SEVERITY for High and Severe.
    if source == "SeverityID":
        severity_id = {0: 0, 1: 1, 2: 2, 3: 4, 4: 5}.get(severity_id)
    return severity_id if severity_id in SEVERITY_BASE_LEVELS else None


def _resource_summary(detection):
    kinds = _as_list(_first(detection, "resource_kinds", "ResourcesKinds"))
    count = _to_int(_first(detection, "resource_count", "ResourceCount"))
    if not kinds:
        raw_resources = _as_list(_first(detection, "resources", "Resources"))
        kinds = []
        for resource in raw_resources:
            prefix = str(resource).split(":", 1)[0].lower()
            kinds.append(prefix if prefix in ("file", "process", "registry", "service", "url") else "other")
        if count is None:
            count = len(raw_resources)
    safe_kinds = sorted(set(
        kind for kind in (str(value).lower() for value in kinds)
        if kind in ("file", "process", "registry", "service", "url", "other")
    ))
    return safe_kinds, max(0, count or 0)


def _detection_identity(detection, index):
    threat_id = _threat_id(_first(detection, "threat_id", "ThreatID")) or "unknown"
    detection_id = _first(detection, "detection_id", "DetectionID")
    if detection_id is not None and str(detection_id).strip():
        detection_id = str(detection_id).strip()
        return "%s:%s" % (threat_id, detection_id), detection_id, True

    # DetectionID is the documented unique detection key. If unavailable, use
    # all stable fields we have and surface the weaker identity in the summary.
    initial = _iso(_first(detection, "detected_at", "initial_detection_time", "InitialDetectionTime")) or "?"
    changed = _iso(_first(detection, "changed_at", "last_status_change_at", "LastThreatStatusChangeTime")) or "?"
    status_id = _to_int(_first(detection, "status_id", "threat_status_id", "ThreatStatusID"))
    resources = _as_list(_first(detection, "resources", "Resources"))
    resource_signature = "|".join(sorted(str(value) for value in resources))
    identity = "%s:%s:%s:%s:%s" % (threat_id, initial, changed, status_id, resource_signature)
    return identity, None, False


def _is_active(status_id, is_active, action_success, did_threat_execute=False,
               execution_status_id=None):
    if action_success is False:
        return True
    if did_threat_execute is True or execution_status_id in (2, 3):
        return True
    if status_id in RESOLVED_STATUS_IDS:
        return is_active is True
    # In particular, Quarantined and Allowed remain unresolved for scoring.
    return True


def classify_detection(detection, repeated=False, realtime_expected=True,
                       protection_context=None, rules=None):
    """Return a deterministic 0..16 risk level for one Defender detection."""
    config = _copy_rules(rules)
    severity_id = _to_int(_first(detection, "threat_severity_id", "ThreatSeverityID"))
    if severity_id is None:
        severity_id = _to_int(_first(detection, "severity_id", "SeverityID"))
        if _first(detection, "severity_source", "SeveritySource") == "SeverityID":
            severity_id = {0: 0, 1: 1, 2: 2, 3: 4, 4: 5}.get(severity_id)
    base = config.get("base_levels", SEVERITY_BASE_LEVELS).get(severity_id, SEVERITY_BASE_LEVELS[0])
    status_id = _to_int(_first(detection, "status_id", "threat_status_id", "ThreatStatusID"))
    is_active = _to_bool(_first(detection, "is_active", "IsActive"))
    action_success = _to_bool(_first(detection, "action_success", "ActionSuccess"))
    did_threat_execute = _to_bool(_first(detection, "did_threat_execute", "DidThreatExecute"))
    execution_status_id = _to_int(_first(detection, "execution_status_id", "CurrentThreatExecutionStatusID"))
    active = _is_active(status_id, is_active, action_success, did_threat_execute, execution_status_id)

    level = base + config["status_modifiers"].get(status_id, 0)
    if active:
        level += config["active_modifier"]
    if repeated:
        level += config["repeated_modifier"]
    if status_id not in FAILED_STATUS_IDS and action_success is False:
        level += config["failed_remediation_modifier"]
    if protection_context is None:
        protection_context = []
        if realtime_expected and _to_bool(_first(detection, "realtime_enabled")) is False:
            protection_context.append("realtime_disabled")
    for modifier in protection_context:
        level += config["protection_modifiers"].get(modifier, 0)
    return max(0, min(16, int(level)))


def level_category(level):
    level = max(0, min(16, _to_int(level, 0)))
    if level == 0:
        return "Aucun événement"
    if level <= 3:
        return "Information"
    if level <= 6:
        return "Faible"
    if level <= 9:
        return "Modéré"
    if level <= 12:
        return "Élevé"
    if level <= 14:
        return "Critique"
    return "Urgence"


def _protection_state(defender, rules):
    mode = defender.get("running_mode")
    mode_text = str(mode or "").strip()
    expected_modes = {str(value).lower() for value in rules["expected_active_modes"]}
    passive_modes = {str(value).lower() for value in rules["passive_modes"]}
    is_expected_active = mode_text.lower() in expected_modes
    is_passive = mode_text.lower() in passive_modes

    issues = []
    modifiers = []
    level = 0
    if is_expected_active:
        checks = (
            ("service_enabled", "Defender service disabled", "service_disabled"),
            ("antivirus_enabled", "Antivirus disabled", "antivirus_disabled"),
            ("realtime_enabled", "Realtime protection disabled", "realtime_disabled"),
            ("behavior_monitor_enabled", "Behavior monitoring disabled", "behavior_monitor_disabled"),
            ("ioav_enabled", "IOAV protection disabled", "ioav_disabled"),
        )
        for key, label, modifier in checks:
            value = _to_bool(defender.get(key))
            if value is False:
                issues.append(label)
                modifiers.append(modifier)
                level += rules["protection_modifiers"].get(modifier, 0)
            elif value is None:
                issues.append("Unknown %s" % key)

        age = _to_int(defender.get("signature_age"))
        max_age = rules["signature_max_age_days"]
        if age is None:
            issues.append("Signature age unknown")
            modifiers.append("signature_age_unknown")
            level += rules["protection_modifiers"].get("signature_age_unknown", 0)
        elif age > max_age:
            issues.append("Signatures outdated: %s days" % age)
            modifiers.append("signatures_stale")
            level += rules["protection_modifiers"].get("signatures_stale", 0)
    elif not is_passive and not mode_text:
        issues.append("Defender running mode unknown")

    return max(0, min(16, level)), issues, is_expected_active, is_passive, modifiers


def _deduplicate_detections(detections):
    deduped = {}
    identity_incomplete = 0
    for index, item in enumerate(_as_list(detections)):
        if not isinstance(item, dict):
            identity_incomplete += 1
            continue
        key, detection_id, has_detection_id = _detection_identity(item, index)
        if not has_detection_id:
            identity_incomplete += 1
        previous = deduped.get(key)
        candidate_changed = _parse_datetime(_first(item, "changed_at", "last_status_change_at", "LastThreatStatusChangeTime"))
        previous_changed = _parse_datetime(_first(previous or {}, "changed_at", "last_status_change_at", "LastThreatStatusChangeTime"))
        if previous is None or (candidate_changed and (not previous_changed or candidate_changed >= previous_changed)):
            deduped[key] = dict(item)
            deduped[key]["_detection_key"] = key
            deduped[key]["_detection_id"] = detection_id
            deduped[key]["_identity_complete"] = has_detection_id
    return list(deduped.values()), identity_incomplete


def _normalize_detection(item, catalog, now, window, history_cutoff, rules, repeated=False,
                         realtime_expected=True, realtime_enabled=None,
                         protection_context=None):
    threat_id = _threat_id(_first(item, "threat_id", "ThreatID"))
    metadata = catalog.get(threat_id, {})
    severity_id = _catalog_severity(metadata)
    status_id = _to_int(_first(item, "status_id", "threat_status_id", "ThreatStatusID"))
    is_active_value = _to_bool(_first(metadata, "is_active", "IsActive"))
    action_success = _to_bool(_first(item, "action_success", "ActionSuccess"))
    did_threat_execute = _to_bool(_first(metadata, "did_threat_execute", "DidThreatExecute"))
    execution_status_id = _to_int(_first(item, "execution_status_id", "CurrentThreatExecutionStatusID"))
    active = _is_active(
        status_id, is_active_value, action_success, did_threat_execute, execution_status_id,
    )
    detected = _parse_datetime(_first(item, "detected_at", "initial_detection_time", "InitialDetectionTime"))
    changed = _parse_datetime(_first(item, "changed_at", "last_status_change_at", "LastThreatStatusChangeTime"))
    remediation = _parse_datetime(_first(item, "remediation_at", "RemediationTime"))
    recent = bool((detected and detected >= window) or (changed and changed >= window))
    retained = bool(recent or active or ((detected and detected >= history_cutoff) or (changed and changed >= history_cutoff)))
    resource_kinds, resource_count = _resource_summary(item)
    enriched = dict(item)
    enriched.update({
        "threat_severity_id": severity_id,
        "status_id": status_id,
        "is_active": active,
        "action_success": action_success,
        "did_threat_execute": did_threat_execute,
        "execution_status_id": execution_status_id,
        "realtime_enabled": realtime_enabled,
    })
    level = classify_detection(
        enriched, repeated=repeated, realtime_expected=realtime_expected,
        protection_context=protection_context, rules=rules,
    )
    return {
        "detection_id": item.get("_detection_id") or _first(item, "detection_id", "DetectionID"),
        "threat_id": threat_id,
        "threat_name": _first(metadata, "threat_name", "name", "ThreatName") or "Unknown threat",
        "did_threat_execute": _to_bool(_first(metadata, "did_threat_execute", "DidThreatExecute")),
        "severity_id": severity_id,
        "severity": SEVERITY_NAMES.get(severity_id, "Unknown"),
        "status_id": status_id,
        "status": THREAT_STATUSES.get(status_id, "Unknown"),
        "detected_at": _iso(_first(item, "detected_at", "initial_detection_time", "InitialDetectionTime")),
        "changed_at": _iso(_first(item, "changed_at", "last_status_change_at", "LastThreatStatusChangeTime")),
        "remediation_at": _iso(_first(item, "remediation_at", "RemediationTime")),
        "error_code": _to_int(_first(item, "error_code", "threat_status_error_code", "ThreatStatusErrorCode")),
        "action_success": action_success,
        "execution_status_id": _to_int(_first(item, "execution_status_id", "CurrentThreatExecutionStatusID")),
        "resource_kinds": resource_kinds,
        "resource_count": resource_count,
        "active": active,
        "active_reason": (
            "remediation_failed" if action_success is False or status_id in FAILED_STATUS_IDS
            else "execution_observed" if did_threat_execute is True
            else "execution_allowed_or_running" if execution_status_id in (2, 3)
            else "explicitly_allowed" if status_id == 5
            else "quarantined_unresolved" if status_id == 3
            else "defender_active" if is_active_value is True
            else "unresolved_status" if status_id not in RESOLVED_STATUS_IDS
            else None
        ),
        "in_window_6h": recent,
        "event_level": level,
        "event_category": level_category(level),
        "identity_complete": bool(item.get("_identity_complete", False)),
        "_retained": retained,
    }


def _append_correlation(result, enabled, key, evidence):
    if enabled.get(key) and evidence:
        result.append({"rule": key, **evidence})


def _build_correlations(detections, rules, realtime_disabled):
    enabled = rules["correlations"]
    recent = [item for item in detections if item["in_window_6h"]]
    by_threat = {}
    failures = Counter()
    for item in recent:
        if item["threat_id"]:
            by_threat.setdefault(item["threat_id"], []).append(item)
            if item["status_id"] in FAILED_STATUS_IDS or item["action_success"] is False:
                failures[item["threat_id"]] += 1

    repeated = {
        threat_id: items for threat_id, items in by_threat.items()
        if len(items) >= rules["repeated_threshold"]
    }
    correlations = []
    repeated_ids = sorted(repeated)
    _append_correlation(correlations, enabled, "same_threat_repeated", {
        "count": sum(len(items) for items in repeated.values()),
        "threat_ids": repeated_ids,
    } if repeated else None)

    unique_recent_threats = sorted(set(item["threat_id"] for item in recent if item["threat_id"]))
    _append_correlation(correlations, enabled, "multiple_threats", {
        "count": len(unique_recent_threats),
        "threat_ids": unique_recent_threats,
    } if len(unique_recent_threats) >= rules["multiple_threats_threshold"] else None)

    high_or_active = [item for item in detections if item["in_window_6h"] or item["active"]]
    severe_ids = sorted(set(
        item["threat_id"] for item in high_or_active
        if item["threat_id"] and item["severity_id"] in (4, 5)
    ))
    _append_correlation(correlations, enabled, "severe_with_protection_disabled", {
        "threat_ids": severe_ids,
        "realtime_disabled": True,
    } if severe_ids and realtime_disabled else None)

    repeated_failure_ids = sorted(
        threat_id for threat_id, count in failures.items()
        if count >= rules["remediation_failure_threshold"]
    )
    _append_correlation(correlations, enabled, "repeated_remediation_failures", {
        "threat_ids": repeated_failure_ids,
        "count": sum(failures[threat_id] for threat_id in repeated_failure_ids),
    } if repeated_failure_ids else None)

    allowed = [item for item in recent if item["status_id"] == 5]
    _append_correlation(correlations, enabled, "explicitly_allowed", {
        "threat_ids": sorted(set(item["threat_id"] for item in allowed if item["threat_id"])),
        "count": len(allowed),
    } if allowed else None)
    return correlations


def _wapt_status(current_level, protection_issues, expected_active, is_passive, collection_complete, identity_incomplete, unknown_severity):
    if not collection_complete:
        return "ERROR"
    if expected_active and any(
        issue in protection_issues
        for issue in ("Defender service disabled", "Antivirus disabled", "Realtime protection disabled")
    ):
        return "ERROR"
    if current_level >= 10:
        return "ERROR"
    if current_level >= 7:
        return "WARNING"
    if protection_issues or identity_incomplete or unknown_severity:
        return "WARNING"
    if not expected_active and not is_passive:
        return "WARNING"
    return "OK"


def build_audit_snapshot(data, now=None, rules=None):
    """Normalize Defender data, classify events, and build the v2 WAPT payload."""
    config = _copy_rules(rules)
    now = _parse_datetime(now) or datetime.now(timezone.utc)
    window = now - timedelta(hours=config["window_hours"])
    history_cutoff = now - timedelta(days=HISTORY_DAYS)
    if not isinstance(data, dict):
        data = {}

    defender_raw = data.get("defender") if isinstance(data.get("defender"), dict) else {}
    defender = {
        "running_mode": defender_raw.get("running_mode"),
        "antivirus_enabled": _to_bool(defender_raw.get("antivirus_enabled")),
        "realtime_enabled": _to_bool(defender_raw.get("realtime_enabled")),
        "behavior_monitor_enabled": _to_bool(defender_raw.get("behavior_monitor_enabled")),
        "ioav_enabled": _to_bool(defender_raw.get("ioav_enabled")),
        "signature_age": _to_int(defender_raw.get("signature_age")),
        "signature_version": defender_raw.get("signature_version"),
        "signature_last_updated": _iso(defender_raw.get("signature_last_updated")),
        "engine_version": defender_raw.get("engine_version"),
        "service_enabled": _to_bool(defender_raw.get("service_enabled")),
    }
    raw_detections = data.get("detections")
    raw_threats = data.get("threats")
    supported_collection_types = (list, tuple, dict)
    arrays_valid = (
        "detections" in data and "threats" in data
        and (raw_detections is None or isinstance(raw_detections, supported_collection_types))
        and (raw_threats is None or isinstance(raw_threats, supported_collection_types))
    )
    collection_complete = data.get("collection_complete") is not False and arrays_valid
    if not defender or not defender.get("running_mode"):
        collection_complete = False
    if defender["antivirus_enabled"] is None or defender["realtime_enabled"] is None:
        collection_complete = False

    catalog = {}
    for threat in _as_list(raw_threats):
        threat_id = _threat_id(_first(threat, "threat_id", "ThreatID"))
        if threat_id:
            catalog[threat_id] = threat

    deduped, identity_incomplete = _deduplicate_detections(raw_detections)
    preliminary = []
    for item in deduped:
        threat_id = _threat_id(_first(item, "threat_id", "ThreatID"))
        preliminary.append({
            "threat_id": threat_id,
            "in_window_6h": bool(
                (_parse_datetime(_first(item, "detected_at", "initial_detection_time", "InitialDetectionTime")) or datetime.min.replace(tzinfo=timezone.utc)) >= window
                or (_parse_datetime(_first(item, "changed_at", "last_status_change_at", "LastThreatStatusChangeTime")) or datetime.min.replace(tzinfo=timezone.utc)) >= window
            ),
        })
    recent_counts = Counter(item["threat_id"] for item in preliminary if item["in_window_6h"] and item["threat_id"])

    protection_level, protection_issues, expected_active, is_passive, protection_context = _protection_state(defender, config)
    realtime_disabled = expected_active and defender["realtime_enabled"] is False
    detections = []
    for item in deduped:
        threat_id = _threat_id(_first(item, "threat_id", "ThreatID"))
        repeated = bool(threat_id and recent_counts.get(threat_id, 0) >= config["repeated_threshold"])
        normalized = _normalize_detection(
            item, catalog, now, window, history_cutoff, config,
            repeated=repeated, realtime_expected=expected_active,
            realtime_enabled=defender["realtime_enabled"],
            protection_context=protection_context,
        )
        if normalized["in_window_6h"] and realtime_disabled:
            # classify_detection already applies this configured modifier.
            pass
        detections.append(normalized)

    detections = [item for item in detections if item["_retained"]]
    for item in detections:
        item.pop("_retained", None)

    recent = [item for item in detections if item["in_window_6h"]]
    active = [item for item in detections if item["active"]]
    max_event = max((item["event_level"] for item in recent), default=0)
    max_active = max((item["event_level"] for item in active), default=0)
    current_risk = max(max_active, protection_level)
    if not collection_complete:
        current_risk = 16
    unknown_severity_count = sum(1 for item in detections if item["severity_id"] is None)
    correlations = _build_correlations(detections, config, realtime_disabled)
    status = _wapt_status(
        current_risk, protection_issues, expected_active, is_passive,
        collection_complete, identity_incomplete, unknown_severity_count,
    )

    # Keep the payload stable between audits; WAPT stores a new historical value
    # only when an actual status, detection, or correlation changes.
    return {
        "schema_version": 2,
        "collection_complete": bool(collection_complete),
        "history_detection_count": _to_int(data.get("history_detection_count")),
        "returned_detection_count": _to_int(data.get("returned_detection_count"), len(_as_list(raw_detections))),
        "defender": defender,
        "summary": {
            "max_event_level_6h": max_event,
            "max_active_level": max_active,
            "current_risk_level": current_risk,
            "current_risk_category": level_category(current_risk),
            "audit_status": status,
            "collection_complete": bool(collection_complete),
            "detection_count_6h": len(recent),
            "active_count": len(active),
            "resolved_count_6h": sum(1 for item in recent if not item["active"]),
            "identity_incomplete_count": identity_incomplete,
            "unknown_severity_count": unknown_severity_count,
            "protection_issues": protection_issues,
            "correlation_count": len(correlations),
        },
        "detections": detections,
        "active_threats": active,
        "correlations": correlations,
    }


def build_collection_error_snapshot(code):
    safe_code = code if code in ("powershell_timeout", "powershell_failed", "powershell_unavailable", "empty_output", "invalid_json", "collection_failed") else "collection_failed"
    return {
        "schema_version": 2,
        "collection_complete": False,
        "history_detection_count": None,
        "returned_detection_count": 0,
        "defender": {
            "running_mode": None,
            "antivirus_enabled": None,
            "realtime_enabled": None,
            "behavior_monitor_enabled": None,
            "ioav_enabled": None,
            "signature_age": None,
            "signature_version": None,
            "signature_last_updated": None,
            "engine_version": None,
            "service_enabled": None,
        },
        "summary": {
            "max_event_level_6h": 0,
            "max_active_level": 0,
            "current_risk_level": 16,
            "current_risk_category": level_category(16),
            "audit_status": "ERROR",
            "collection_complete": False,
            "detection_count_6h": 0,
            "active_count": 0,
            "resolved_count_6h": 0,
            "identity_incomplete_count": 0,
            "unknown_severity_count": 0,
            "protection_issues": [safe_code],
            "correlation_count": 0,
        },
        "detections": [],
        "active_threats": [],
        "correlations": [],
    }


def powershell_json(script=POWERSHELL_SCRIPT, timeout=POWERSHELL_TIMEOUT_SECONDS):
    powershell = os.path.join(
        os.environ.get("SystemRoot", r"C:\Windows"),
        "System32",
        "WindowsPowerShell",
        "v1.0",
        "powershell.exe",
    )
    try:
        result = subprocess.run(
            [powershell, "-NoLogo", "-NoProfile", "-NonInteractive", "-Command", script],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            timeout=timeout,
            check=False,
        )
    except subprocess.TimeoutExpired:
        raise DefenderCollectionError("powershell_timeout")
    except OSError:
        raise DefenderCollectionError("powershell_unavailable")

    if result.returncode != 0:
        raise DefenderCollectionError("powershell_failed")
    raw_output = result.stdout or b""
    if isinstance(raw_output, bytes):
        output = raw_output.decode("utf-8-sig", errors="replace").strip()
    else:
        output = str(raw_output).lstrip("\ufeff").strip()
    if not output:
        raise DefenderCollectionError("empty_output")
    try:
        decoded = json.loads(output)
    except (TypeError, ValueError):
        raise DefenderCollectionError("invalid_json")
    if not isinstance(decoded, dict):
        raise DefenderCollectionError("invalid_json")
    return decoded


def collect_defender():
    """Read Defender state only; no network or remediation cmdlets are used."""
    return powershell_json()


def save_audit(key, value, rules=None):
    config = _copy_rules(rules)
    return WAPT.write_audit_data_if_changed(
        AUDIT_CATEGORY,
        key,
        value,
        max_count=config["audit_history_max_count"],
        keep_days=config["audit_history_days"],
    )


def _save_snapshot(snapshot):
    save_audit("snapshot", snapshot)
    # Keep established Reporting keys available while adding the v2 aggregate.
    save_audit("summary", snapshot["summary"])
    save_audit("detections", snapshot["detections"])
    save_audit("active_threats", snapshot["active_threats"])
    save_audit("correlations", snapshot["correlations"])


def audit():
    try:
        data = collect_defender()
        snapshot = build_audit_snapshot(data)
    except Exception as exc:
        code = getattr(exc, "code", "collection_failed")
        snapshot = build_collection_error_snapshot(code)

    try:
        _save_snapshot(snapshot)
    except Exception:
        # Keep the status fail-closed if WAPT audit persistence itself fails.
        print("Defender audit data could not be stored by WAPT")
        return "ERROR"

    print(json.dumps(snapshot["summary"], indent=2, ensure_ascii=False, sort_keys=True))
    return snapshot["summary"]["audit_status"]