#!/usr/bin/env python3
"""Vulnerability aging and SLA tracking agent.

Tracks vulnerability remediation timelines, calculates SLA compliance,
generates aging reports, and identifies overdue items by severity.
"""

import json
import datetime
import collections


SLA_DEFINITIONS = {
    "critical": {"remediation_days": 7, "patch_days": 15, "exception_max_days": 30},
    "high": {"remediation_days": 30, "patch_days": 45, "exception_max_days": 90},
    "medium": {"remediation_days": 90, "patch_days": 120, "exception_max_days": 180},
    "low": {"remediation_days": 180, "patch_days": 365, "exception_max_days": 365},
}

AGING_BUCKETS = [
    (0, 7, "0-7 days"),
    (8, 30, "8-30 days"),
    (31, 60, "31-60 days"),
    (61, 90, "61-90 days"),
    (91, 180, "91-180 days"),
    (181, 365, "181-365 days"),
    (366, 99999, "365+ days"),
]


def calculate_age(discovery_date_str):
    """Calculate vulnerability age in days from discovery date."""
    try:
        disc = datetime.datetime.fromisoformat(discovery_date_str.replace("Z", "+00:00"))
        now = datetime.datetime.now(datetime.timezone.utc)
        return (now - disc).days
    except (ValueError, AttributeError):
        return 0


def check_sla_compliance(vuln):
    """Check if a vulnerability is within SLA."""
    severity = vuln.get("severity", "medium").lower()
    sla = SLA_DEFINITIONS.get(severity, SLA_DEFINITIONS["medium"])
    age = calculate_age(vuln.get("discovery_date", ""))
    status = vuln.get("status", "open")

    result = {
        "vuln_id": vuln.get("id", ""),
        "severity": severity,
        "age_days": age,
        "sla_days": sla["remediation_days"],
        "days_remaining": sla["remediation_days"] - age,
        "sla_status": "within_sla",
    }

    if status in ("remediated", "closed"):
        result["sla_status"] = "resolved"
    elif status == "exception":
        if age > sla["exception_max_days"]:
            result["sla_status"] = "exception_expired"
        else:
            result["sla_status"] = "exception_active"
    elif age > sla["remediation_days"]:
        result["sla_status"] = "overdue"
    elif age > sla["remediation_days"] * 0.8:
        result["sla_status"] = "at_risk"

    return result


def build_aging_report(vulns):
    """Build vulnerability aging distribution report."""
    buckets = {label: {"total": 0, "critical": 0, "high": 0, "medium": 0, "low": 0}
               for _, _, label in AGING_BUCKETS}

    for vuln in vulns:
        if vuln.get("status") in ("remediated", "closed"):
            continue
        age = calculate_age(vuln.get("discovery_date", ""))
        severity = vuln.get("severity", "medium").lower()
        for low, high, label in AGING_BUCKETS:
            if low <= age <= high:
                buckets[label]["total"] += 1
                if severity in buckets[label]:
                    buckets[label][severity] += 1
                break

    return {"aging_distribution": buckets, "generated_at": datetime.datetime.utcnow().isoformat() + "Z"}


def calculate_mttr(vulns):
    """Calculate Mean Time to Remediate by severity."""
    remediation_times = collections.defaultdict(list)
    for vuln in vulns:
        if vuln.get("status") in ("remediated", "closed") and vuln.get("remediation_date"):
            try:
                disc = datetime.datetime.fromisoformat(vuln["discovery_date"].replace("Z", "+00:00"))
                rem = datetime.datetime.fromisoformat(vuln["remediation_date"].replace("Z", "+00:00"))
                days = (rem - disc).days
                severity = vuln.get("severity", "medium").lower()
                remediation_times[severity].append(days)
            except (ValueError, KeyError):
                pass

    mttr = {}
    for sev, times in remediation_times.items():
        mttr[sev] = {
            "mean_days": round(sum(times) / len(times), 1),
            "median_days": sorted(times)[len(times) // 2],
            "count": len(times),
        }
    return mttr


def generate_sla_dashboard(vulns):
    """Generate SLA compliance dashboard data."""
    total = 0
    compliant = 0
    overdue = 0
    at_risk = 0
    by_severity = collections.defaultdict(lambda: {"total": 0, "compliant": 0, "overdue": 0})

    for vuln in vulns:
        sla_result = check_sla_compliance(vuln)
        if sla_result["sla_status"] == "resolved":
            continue
        total += 1
        severity = sla_result["severity"]
        by_severity[severity]["total"] += 1
        if sla_result["sla_status"] == "within_sla":
            compliant += 1
            by_severity[severity]["compliant"] += 1
        elif sla_result["sla_status"] == "overdue":
            overdue += 1
            by_severity[severity]["overdue"] += 1
        elif sla_result["sla_status"] == "at_risk":
            at_risk += 1

    return {
        "total_open": total,
        "compliant": compliant,
        "overdue": overdue,
        "at_risk": at_risk,
        "compliance_rate": round(compliant / max(total, 1) * 100, 1),
        "by_severity": dict(by_severity),
    }


if __name__ == "__main__":
    print("=" * 60)
    print("Vulnerability Aging & SLA Tracking")
    print("SLA compliance, aging distribution, MTTR calculation")
    print("=" * 60)

    now = datetime.datetime.utcnow()
    demo_vulns = [
        {"id": "CVE-2024-1234", "severity": "critical", "status": "open",
         "discovery_date": (now - datetime.timedelta(days=10)).isoformat() + "Z"},
        {"id": "CVE-2024-2345", "severity": "high", "status": "open",
         "discovery_date": (now - datetime.timedelta(days=45)).isoformat() + "Z"},
        {"id": "CVE-2024-3456", "severity": "medium", "status": "open",
         "discovery_date": (now - datetime.timedelta(days=120)).isoformat() + "Z"},
        {"id": "CVE-2024-4567", "severity": "critical", "status": "remediated",
         "discovery_date": (now - datetime.timedelta(days=5)).isoformat() + "Z",
         "remediation_date": (now - datetime.timedelta(days=2)).isoformat() + "Z"},
        {"id": "CVE-2024-5678", "severity": "low", "status": "exception",
         "discovery_date": (now - datetime.timedelta(days=200)).isoformat() + "Z"},
    ]

    print("\n--- SLA Compliance ---")
    for vuln in demo_vulns:
        result = check_sla_compliance(vuln)
        print("  {} [{}] age={}d sla={}d -> {}".format(
            result["vuln_id"], result["severity"], result["age_days"],
            result["sla_days"], result["sla_status"]))

    dashboard = generate_sla_dashboard(demo_vulns)
    print("\n--- Dashboard ---")
    print("  Compliance rate: {}%".format(dashboard["compliance_rate"]))
    print("  Open: {} | Compliant: {} | Overdue: {} | At-risk: {}".format(
        dashboard["total_open"], dashboard["compliant"], dashboard["overdue"], dashboard["at_risk"]))

    mttr = calculate_mttr(demo_vulns)
    if mttr:
        print("\n--- MTTR ---")
        for sev, data in mttr.items():
            print("  {}: mean={}d median={}d (n={})".format(sev, data["mean_days"], data["median_days"], data["count"]))

    print("\n" + json.dumps({"vulns_tracked": len(demo_vulns)}, indent=2))
