#!/usr/bin/env python3
"""CISA KEV CVE prioritization agent.

Downloads the CISA Known Exploited Vulnerabilities (KEV) catalog and
cross-references it against a list of CVEs from vulnerability scans to
prioritize remediation based on active exploitation status, due dates,
and vendor/product impact.
"""
import argparse
import csv
import io
import json
import os
import sys
from datetime import datetime, timezone

try:
    import requests
except ImportError:
    print("[!] 'requests' library required: pip install requests", file=sys.stderr)
    sys.exit(1)

KEV_JSON_URL = "https://www.cisa.gov/sites/default/files/feeds/known_exploited_vulnerabilities.json"


def download_kev_catalog():
    """Download the latest CISA KEV catalog as JSON."""
    print(f"[*] Downloading CISA KEV catalog from {KEV_JSON_URL}")
    resp = requests.get(KEV_JSON_URL, timeout=30)
    resp.raise_for_status()
    data = resp.json()
    vulns = data.get("vulnerabilities", [])
    catalog_version = data.get("catalogVersion", "unknown")
    count = data.get("count", len(vulns))
    date_released = data.get("dateReleased", "unknown")
    print(f"[+] KEV catalog v{catalog_version}: {count} vulnerabilities "
          f"(released: {date_released})")
    return vulns, {
        "catalog_version": catalog_version,
        "count": count,
        "date_released": date_released,
    }


def build_kev_index(kev_vulns):
    """Build a lookup dict keyed by CVE ID for fast matching."""
    index = {}
    for vuln in kev_vulns:
        cve_id = vuln.get("cveID", "")
        if cve_id:
            index[cve_id.upper()] = {
                "cve_id": cve_id,
                "vendor": vuln.get("vendorProject", ""),
                "product": vuln.get("product", ""),
                "vulnerability_name": vuln.get("vulnerabilityName", ""),
                "date_added": vuln.get("dateAdded", ""),
                "short_description": vuln.get("shortDescription", ""),
                "required_action": vuln.get("requiredAction", ""),
                "due_date": vuln.get("dueDate", ""),
                "known_ransomware_campaign": vuln.get("knownRansomwareCampaignUse", "Unknown"),
                "notes": vuln.get("notes", ""),
            }
    return index


def load_cve_list(source):
    """Load CVE list from file (one CVE per line or CSV) or comma-separated string."""
    cves = []
    if os.path.isfile(source):
        with open(source, "r") as f:
            content = f.read()
        for line in content.strip().splitlines():
            line = line.strip().strip(",").strip('"')
            if line.upper().startswith("CVE-"):
                cves.append(line.upper())
            elif "," in line:
                parts = line.split(",")
                for part in parts:
                    part = part.strip().strip('"')
                    if part.upper().startswith("CVE-"):
                        cves.append(part.upper())
    else:
        for part in source.split(","):
            part = part.strip()
            if part.upper().startswith("CVE-"):
                cves.append(part.upper())
    return list(set(cves))


def prioritize_cves(cve_list, kev_index):
    """Cross-reference CVEs against KEV catalog and prioritize."""
    in_kev = []
    not_in_kev = []

    for cve_id in cve_list:
        kev_entry = kev_index.get(cve_id)
        if kev_entry:
            entry = dict(kev_entry)
            entry["in_kev"] = True
            entry["priority"] = "CRITICAL"
            if entry.get("known_ransomware_campaign", "").lower() == "known":
                entry["priority"] = "CRITICAL-RANSOMWARE"
            try:
                due = datetime.strptime(entry["due_date"], "%Y-%m-%d")
                if due < datetime.now():
                    entry["overdue"] = True
                    entry["priority"] = "CRITICAL-OVERDUE"
                else:
                    entry["overdue"] = False
            except (ValueError, KeyError):
                entry["overdue"] = False
            in_kev.append(entry)
        else:
            not_in_kev.append({
                "cve_id": cve_id,
                "in_kev": False,
                "priority": "STANDARD",
            })

    # Sort: overdue first, then ransomware, then by due date
    priority_order = {"CRITICAL-OVERDUE": 0, "CRITICAL-RANSOMWARE": 1, "CRITICAL": 2}
    in_kev.sort(key=lambda x: (priority_order.get(x["priority"], 9), x.get("due_date", "")))

    return in_kev, not_in_kev


def format_summary(in_kev, not_in_kev, total_cves, catalog_info):
    """Print a human-readable prioritization report."""
    print(f"\n{'='*60}")
    print(f"  CISA KEV CVE Prioritization Report")
    print(f"{'='*60}")
    print(f"  KEV Catalog  : v{catalog_info['catalog_version']} "
          f"({catalog_info['count']} total KEVs)")
    print(f"  Input CVEs   : {total_cves}")
    print(f"  In KEV       : {len(in_kev)} (actively exploited)")
    print(f"  Not in KEV   : {len(not_in_kev)}")

    overdue = [v for v in in_kev if v.get("overdue")]
    ransomware = [v for v in in_kev if "RANSOMWARE" in v.get("priority", "")]

    print(f"\n  Priority Breakdown:")
    print(f"    CRITICAL-OVERDUE     : {len(overdue)}")
    print(f"    CRITICAL-RANSOMWARE  : {len(ransomware)}")
    print(f"    CRITICAL (in KEV)    : {len(in_kev)}")
    print(f"    STANDARD (not in KEV): {len(not_in_kev)}")

    if in_kev:
        print(f"\n  KEV-Listed CVEs (prioritize remediation):")
        for v in in_kev[:20]:
            overdue_flag = " [OVERDUE]" if v.get("overdue") else ""
            ransomware_flag = " [RANSOMWARE]" if "RANSOMWARE" in v.get("priority", "") else ""
            print(f"    {v['cve_id']:18s} | Due: {v.get('due_date', 'N/A'):10s} | "
                  f"{v.get('vendor', ''):15s} | {v.get('product', ''):20s}"
                  f"{overdue_flag}{ransomware_flag}")
            if v.get("required_action"):
                print(f"      Action: {v['required_action'][:70]}")


def kev_stats(kev_vulns):
    """Compute statistics about the KEV catalog."""
    by_vendor = {}
    ransomware_count = 0
    for v in kev_vulns:
        vendor = v.get("vendorProject", "Unknown")
        by_vendor[vendor] = by_vendor.get(vendor, 0) + 1
        if v.get("knownRansomwareCampaignUse", "").lower() == "known":
            ransomware_count += 1

    print(f"\n  KEV Catalog Statistics:")
    print(f"    Total vulnerabilities  : {len(kev_vulns)}")
    print(f"    Ransomware-associated  : {ransomware_count}")
    print(f"    Top vendors:")
    for vendor, count in sorted(by_vendor.items(), key=lambda x: -x[1])[:10]:
        print(f"      {vendor:30s}: {count}")


def main():
    parser = argparse.ArgumentParser(
        description="CISA KEV CVE prioritization agent"
    )
    parser.add_argument("--cves", required=True,
                        help="CVE list: comma-separated string or path to file")
    parser.add_argument("--kev-file",
                        help="Path to local KEV JSON (skip download)")
    parser.add_argument("--stats", action="store_true",
                        help="Show KEV catalog statistics")
    parser.add_argument("--output", "-o", help="Output JSON report path")
    parser.add_argument("--verbose", "-v", action="store_true")
    args = parser.parse_args()

    if args.kev_file:
        print(f"[*] Loading local KEV file: {args.kev_file}")
        with open(args.kev_file, "r") as f:
            data = json.load(f)
        kev_vulns = data.get("vulnerabilities", [])
        catalog_info = {
            "catalog_version": data.get("catalogVersion", "local"),
            "count": len(kev_vulns),
            "date_released": data.get("dateReleased", "unknown"),
        }
    else:
        kev_vulns, catalog_info = download_kev_catalog()

    if args.stats:
        kev_stats(kev_vulns)

    kev_index = build_kev_index(kev_vulns)
    cve_list = load_cve_list(args.cves)
    print(f"[*] Loaded {len(cve_list)} CVE(s) to check")

    in_kev, not_in_kev = prioritize_cves(cve_list, kev_index)
    format_summary(in_kev, not_in_kev, len(cve_list), catalog_info)

    report = {
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "tool": "CISA KEV Prioritization",
        "catalog_info": catalog_info,
        "input_cve_count": len(cve_list),
        "kev_matched": len(in_kev),
        "kev_unmatched": len(not_in_kev),
        "prioritized": in_kev,
        "standard_priority": not_in_kev,
        "risk_level": (
            "CRITICAL" if any(v.get("overdue") for v in in_kev)
            else "HIGH" if in_kev
            else "LOW"
        ),
    }

    if args.output:
        with open(args.output, "w") as f:
            json.dump(report, f, indent=2)
        print(f"\n[+] Report saved to {args.output}")
    elif args.verbose:
        print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()
