#!/usr/bin/env python3
# vuln_summarizer.py
# AI‑Powered Vulnerability Summarizer using Google Gemini 2.5 Flash


import os
import sys
import json
import requests
from datetime import datetime
import google.generativeai as genai
# ─── Configuration ────────────────────────────────────────────────────────────
NVD_API = "https://services.nvd.nist.gov/rest/json/cves/2.0"
CVE_COUNT = 5
MODEL_NAME = "gemini-2.5-flash"
OUTPUT_MD = "report.md"
OUTPUT_JSONL = "report.jsonl"
def fetch_latest_cves(count=CVE_COUNT):
    params = {"startIndex": 0, "resultsPerPage": count}
    resp = requests.get(NVD_API, params=params)
    resp.raise_for_status()
    return resp.json().get("vulnerabilities", [])
def configure_gemini(api_key: str):
    genai.configure(api_key=api_key)
    return genai.GenerativeModel(model_name=MODEL_NAME)
def generate_summary(model, cve_id: str, description: str):
    prompt = f"""
You are a cybersecurity expert. Given the following CVE:
ID: {cve_id}
Description: "{description}"
Please provide:
1. A one‑sentence technical summary.
2. Who or what is affected? (e.g., software, users, services)
3. Recommended remediation steps in bullet points.
Format your response as Markdown with headings.
"""
    return model.generate_content(prompt).text.strip()
def main():
    api_key = os.getenv("GOOGLE_API_KEY") or "YOUR_API_KEY"
    if not api_key:
        print("Error: set your GOOGLE_API_KEY", file=sys.stderr)
        sys.exit(1)
    print(f"[{datetime.now()}] Fetching latest {CVE_COUNT} CVEs...")
    vulns = fetch_latest_cves(CVE_COUNT)
    print(f"[{datetime.now()}] Configuring Gemini model '{MODEL_NAME}'...")
    model = configure_gemini(api_key)
    md_lines = ["# AI‑Powered Vulnerability Report",
                f"_Generated: {datetime.now()}_\n"]
    jsonl_f = open(OUTPUT_JSONL, "w", encoding="utf-8")
    for entry in vulns:
        cve = entry.get("cve", {})
        cve_id = cve.get("id")
        description = next((d["value"] for d in cve.get("descriptions", [])
                            if d["lang"] == "en"), "")
        if not cve_id or not description:
            continue
        print(f"[{datetime.now()}] Summarizing {cve_id}...")
        summary_md = generate_summary(model, cve_id, description)
        md_lines.append(f"## {cve_id}\n")
        md_lines.append(f"**Original description:** {description}\n\n")
        md_lines.append(summary_md + "\n")
        record = {"cve_id": cve_id,
                  "description": description,
                  "summary": summary_md}
        jsonl_f.write(json.dumps(record, ensure_ascii=False) + "\n")
    jsonl_f.close()
    with open(OUTPUT_MD, "w", encoding="utf-8") as md_f:
        md_f.write("\n".join(md_lines))
    print(f"\nDone! Report written to:\n • {OUTPUT_MD}\n • {OUTPUT_JSONL}")
if __name__ == "__main__":
    main()