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

from setuphelpers import *
import json
import subprocess
import os
from datetime import datetime, timedelta

AUDIT_CATEGORY = "Microsoft Defender"
HISTORY_DAYS = 30
SIGNATURE_MAX_AGE = 3
MAX_DETECTIONS = 100

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",
}

ACTIVE_STATUSES = {1, 5, 105}
FAILED_STATUSES = {102, 103, 104, 107}
RESOLVED_STATUSES = {2, 3, 4, 6}


def install():
    print("Microsoft Defender audit package installed")


def powershell_json(script):
    powershell = os.path.join(
        os.environ.get("SystemRoot", r"C:\Windows"),
        "System32",
        "WindowsPowerShell",
        "v1.0",
        "powershell.exe"
    )

    result = subprocess.run(
        [
            powershell,
            "-NoProfile",
            "-NonInteractive",
            "-ExecutionPolicy", "Bypass",
            "-Command", script,
        ],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        timeout=90,
    )

    if result.returncode != 0:
        raise RuntimeError(
            result.stderr.decode("utf-8", errors="replace")
        )

    output = result.stdout.decode(
        "utf-8-sig", errors="replace"
    ).strip()

    if not output:
        raise RuntimeError("Empty PowerShell response")

    return json.loads(output)


def collect_defender():
    script = r"""
$ErrorActionPreference = "Stop"
$ProgressPreference = "SilentlyContinue"

$status = Get-MpComputerStatus
$cutoff = (Get-Date).AddDays(-30)

$threats = @(
    Get-MpThreatDetection |
    Where-Object {
        $_.InitialDetectionTime -ge $cutoff
    } |
    Sort-Object InitialDetectionTime -Descending |
    Select-Object -First 100 |
    ForEach-Object {
        [PSCustomObject]@{
            threat_id = $_.ThreatID
            status_id = [int]$_.ThreatStatusID
            detected_at = if ($_.InitialDetectionTime) {
                $_.InitialDetectionTime.ToString("o")
            } else { $null }
            changed_at = if ($_.LastThreatStatusChangeTime) {
                $_.LastThreatStatusChangeTime.ToString("o")
            } else { $null }
            error_code = $_.ThreatStatusErrorCode
            resources = @($_.Resources)
        }
    }
)

$threatNames = @{}
Get-MpThreat | ForEach-Object {
    $threatNames[[string]$_.ThreatID] = $_.ThreatName
}

foreach ($threat in $threats) {
    $id = [string]$threat.threat_id
    $name = $threatNames[$id]

    $threat | Add-Member -NotePropertyName "name" `
        -NotePropertyValue $name
}

$result = [PSCustomObject]@{
    collected_at = (Get-Date).ToUniversalTime().ToString("o")
    computer_name = $env:COMPUTERNAME
    running_mode = $status.AMRunningMode
    antivirus_enabled = $status.AntivirusEnabled
    realtime_enabled = $status.RealTimeProtectionEnabled
    behavior_monitor_enabled = $status.BehaviorMonitorEnabled
    ioav_enabled = $status.IoavProtectionEnabled
    signature_age = $status.AntivirusSignatureAge
    signature_version = $status.AntivirusSignatureVersion
    engine_version = $status.AMEngineVersion
    threats = @($threats)
}

ConvertTo-Json -InputObject $result -Depth 8 -Compress
"""
    return powershell_json(script)


def normalize_threats(data):
    result = []

    for threat in data.get("threats", []):
        status_id = int(threat.get("status_id", 0))

        if status_id in ACTIVE_STATUSES:
            severity = "ERROR"
        elif status_id in FAILED_STATUSES:
            severity = "ERROR"
        elif status_id in RESOLVED_STATUSES:
            severity = "RESOLVED"
        else:
            severity = "WARNING"

        result.append({
            "id": threat.get("threat_id"),
            "name": threat.get("name") or "Unknown threat",
            "status_id": status_id,
            "status": THREAT_STATUSES.get(
                status_id, "Unknown"
            ),
            "severity": severity,
            "detected_at": threat.get("detected_at"),
            "changed_at": threat.get("changed_at"),
            "resources": threat.get("resources") or [],
            "error_code": threat.get("error_code"),
        })

    return result


def save_audit(key, value):
    WAPT.write_audit_data_if_changed(
        AUDIT_CATEGORY,
        key,
        value,
        keep_days=90
    )


def audit():
    try:
        data = collect_defender()
        threats = normalize_threats(data)

        active = [
            t for t in threats
            if t["severity"] == "ERROR"
        ]
        uncertain = [
            t for t in threats
            if t["severity"] == "WARNING"
        ]
        resolved = [
            t for t in threats
            if t["severity"] == "RESOLVED"
        ]

        status = "OK"
        problems = []

        mode = data.get("running_mode")

        if mode != "Normal":
            problems.append(
                "Defender mode: %s" % mode
            )
            status = "WARNING"

        if mode == "Normal":
            if not data.get("antivirus_enabled"):
                problems.append("Antivirus disabled")
                status = "ERROR"

            if not data.get("realtime_enabled"):
                problems.append("Realtime protection disabled")
                status = "ERROR"

            age = data.get("signature_age")

            if age is None:
                problems.append("Unknown signature age")
                if status == "OK":
                    status = "WARNING"
            elif int(age) > SIGNATURE_MAX_AGE:
                problems.append(
                    "Signatures outdated: %s days" % age
                )
                if status == "OK":
                    status = "WARNING"

        if active:
            problems.append(
                "%d unresolved or failed detections" %
                len(active)
            )
            status = "ERROR"

        if uncertain and status == "OK":
            problems.append(
                "%d detections require review" %
                len(uncertain)
            )
            status = "WARNING"

        summary = {
            "status": status,
            "computer_name": data.get("computer_name"),
            "running_mode": mode,
            "antivirus_enabled": data.get(
                "antivirus_enabled"
            ),
            "realtime_enabled": data.get(
                "realtime_enabled"
            ),
            "signature_age": data.get("signature_age"),
            "signature_version": data.get(
                "signature_version"
            ),
            "engine_version": data.get("engine_version"),
            "detection_count": len(threats),
            "active_count": len(active),
            "review_count": len(uncertain),
            "resolved_count": len(resolved),
            "problems": problems,
        }

        save_audit("summary", summary)
        save_audit("detections", threats)
        save_audit("active_threats", active)

        print(json.dumps(
            summary,
            indent=2,
            ensure_ascii=False
        ))

        return status

    except Exception as exc:
        error = {
            "status": "ERROR",
            "error": str(exc)
        }

        print("Defender audit error: %s" % exc)

        save_audit("audit_error", error)
        return "ERROR"