import requests
import json
from datetime import datetime

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

def fetch_kev_catalog():
    """Download and parse the CISA KEV catalog."""
    response = requests.get(KEV_URL, timeout=30)
    response.raise_for_status()
    data = response.json()

    catalog = {}
    for vuln in data.get("vulnerabilities", []):
        cve_id = vuln["cveID"]
        catalog[cve_id] = {
            "vendor": vuln.get("vendorProject", ""),
            "product": vuln.get("product", ""),
            "name": vuln.get("vulnerabilityName", ""),
            "date_added": vuln.get("dateAdded", ""),
            "description": vuln.get("shortDescription", ""),
            "action": vuln.get("requiredAction", ""),
            "due_date": vuln.get("dueDate", ""),
            "ransomware_use": vuln.get("knownRansomwareCampaignUse", "Unknown"),
        }

    print(f"[+] Loaded {len(catalog)} CVEs from CISA KEV catalog")
    print(f"    Catalog version: {data.get('catalogVersion', 'N/A')}")
    print(f"    Last updated: {data.get('dateReleased', 'N/A')}")
    return catalog

kev = fetch_kev_catalog()

EPSS_API = "https://api.first.org/data/v1/epss"

def get_epss_scores(cve_list):
    """Fetch EPSS scores for a batch of CVEs."""
    scores = {}
    batch_size = 100
    for i in range(0, len(cve_list), batch_size):
        batch = cve_list[i:i + batch_size]
        cve_param = ",".join(batch)
        response = requests.get(EPSS_API, params={"cve": cve_param}, timeout=30)
        if response.status_code == 200:
            for entry in response.json().get("data", []):
                scores[entry["cve"]] = {
                    "epss": float(entry.get("epss", 0)),
                    "percentile": float(entry.get("percentile", 0)),
                }
    return scores

import pandas as pd

def prioritize_vulnerabilities(scan_results, kev_catalog, epss_scores):
    """Apply multi-factor prioritization to scan results."""
    prioritized = []

    for vuln in scan_results:
        cve_id = vuln.get("cve_id", "")
        cvss_score = float(vuln.get("cvss_score", 0))
        asset_criticality = float(vuln.get("asset_criticality", 3))
        exposure = float(vuln.get("network_exposure", 3))

        in_kev = cve_id in kev_catalog
        kev_data = kev_catalog.get(cve_id, {})
        epss_data = epss_scores.get(cve_id, {"epss": 0, "percentile": 0})
        epss_score = epss_data["epss"]

        # Composite risk score calculation
        risk_score = (
            (1.0 if in_kev else 0.0) * 10 * 0.30 +
            epss_score * 10 * 0.25 +
            cvss_score * 0.20 +
            (asset_criticality / 5.0) * 10 * 0.15 +
            (exposure / 5.0) * 10 * 0.10
        )

        # Assign priority level
        if in_kev or (epss_score > 0.5 and cvss_score >= 9.0):
            priority = "P1-Emergency"
            sla_days = 2
        elif epss_score > 0.5 or cvss_score >= 9.0:
            priority = "P2-Critical"
            sla_days = 7
        elif cvss_score >= 7.0:
            priority = "P3-High"
            sla_days = 14
        elif cvss_score >= 4.0:
            priority = "P4-Medium"
            sla_days = 30
        else:
            priority = "P5-Low"
            sla_days = 90

        prioritized.append({
            "cve_id": cve_id,
            "cvss_score": cvss_score,
            "epss_score": round(epss_score, 4),
            "epss_percentile": round(epss_data["percentile"], 4),
            "in_cisa_kev": in_kev,
            "ransomware_use": kev_data.get("ransomware_use", "N/A"),
            "kev_due_date": kev_data.get("due_date", "N/A"),
            "risk_score": round(risk_score, 2),
            "priority": priority,
            "sla_days": sla_days,
            "asset": vuln.get("asset", ""),
            "asset_criticality": asset_criticality,
        })

    df = pd.DataFrame(prioritized)
    df = df.sort_values("risk_score", ascending=False)
    return df

def generate_report(df, output_file="kev_prioritized_report.csv"):
    """Generate summary report from prioritized vulnerabilities."""
    print("\n" + "=" * 70)
    print("VULNERABILITY PRIORITIZATION REPORT - KEV + EPSS + CVSS")
    print("=" * 70)

    print(f"\nTotal vulnerabilities analyzed: {len(df)}")
    print(f"KEV-listed vulnerabilities:    {df['in_cisa_kev'].sum()}")
    print(f"Ransomware-associated:         {(df['ransomware_use'] == 'Known').sum()}")

    print("\nPriority Distribution:")
    print(df["priority"].value_counts().to_string())

    print("\nTop 15 Highest Risk Vulnerabilities:")
    top = df.head(15)[["cve_id", "cvss_score", "epss_score", "in_cisa_kev",
                        "risk_score", "priority"]]
    print(top.to_string(index=False))

    df.to_csv(output_file, index=False)
    print(f"\n[+] Full report saved to: {output_file}")