anshseth02/ai-cve-explainer

1import requests2import streamlit as st3import os4 5# --- Page Config ---6st.set_page_config(7    page_title="AI CVE Explainer",8    page_icon="🧠",9    layout="centered"10)11 12st.markdown("<h1>🧠 AI CVE Explainer</h1>", unsafe_allow_html=True)13st.markdown("<div style='color:#6b7280;'>Developed by Ansh — enter a CVE ID to get an AI-friendly summary.</div>", unsafe_allow_html=True)14 15# --- HF API ---16HF_MODEL = "tiiuae/falcon-7b-instruct-hf"17HF_API_URL = f"https://router.huggingface.co/hf-inference/models/{HF_MODEL}"18HF_TOKEN = os.getenv("HF_TOKEN")  # optional19 20if not HF_TOKEN:21    st.info("HF_TOKEN not found. AI explanation may be limited.")22 23def ask_ai(prompt):24    headers = {"Content-Type": "application/json"}25    if HF_TOKEN:26        headers["Authorization"] = f"Bearer {HF_TOKEN}"27 28    payload = {29        "inputs": prompt,30        "parameters": {"max_new_tokens": 200, "temperature": 0.7, "return_full_text": False}31    }32 33    try:34        response = requests.post(HF_API_URL, headers=headers, json=payload, timeout=60)35        if response.status_code != 200:36            return f"Model API error {response.status_code}: {response.text}"37        data = response.json()38        if isinstance(data, list) and len(data) > 0:39            for key in ("generated_text", "summary_text", "text", "content"):40                if key in data[0]:41                    return data[0][key]42            return " ".join(str(v) for v in data[0].values() if isinstance(v, str)) or str(data[0])43        return str(data)44    except Exception as e:45        return f"Request failed: {e}"46 47# --- Fetch CVE with CVSS ---48def get_cve_info(cve_id):49    cve_id = cve_id.strip().upper()50    # NVD API51    url = f"https://services.nvd.nist.gov/rest/json/cve/1.0/{cve_id}"52    try:53        r = requests.get(url, timeout=10)54        if not r.ok:55            return {"desc":"CVE not found or API rate limited", "cvss":None, "risk":None}56 57        j = r.json()58        cve_items = j.get("result", {}).get("CVE_Items", [])59        if not cve_items:60            return {"desc":"CVE not found", "cvss":None, "risk":None}61 62        desc = cve_items[0]["cve"]["description"]["description_data"][0]["value"]63 64        # CVSS v3 base score65        impact = cve_items[0].get("impact", {}).get("baseMetricV3", {})66        cvss = impact.get("cvssV3", {}).get("baseScore", None)67        if cvss is not None:68            cvss = round(float(cvss),1)69            if cvss >= 7.0:70                risk = "High"71            elif cvss >= 4.0:72                risk = "Medium"73            else:74                risk = "Low"75        else:76            risk = None77 78        return {"desc": desc, "cvss": cvss, "risk": risk}79 80    except Exception as e:81        return {"desc": f"Error fetching CVE: {e}", "cvss":None, "risk":None}82 83# --- Streamlit UI ---84cve_id = st.text_input("Enter CVE ID", placeholder="e.g., CVE-2024-30078")85if st.button("Explain"):86    if not cve_id:87        st.warning("Please type a CVE ID.")88    else:89        info = get_cve_info(cve_id)90        st.subheader("Official Description")91        st.write(info["desc"])92 93        if info["cvss"] is not None:94            st.markdown(f"**CVSS v3 Base Score:** {info['cvss']} / 10")95            st.markdown(f"**Risk Level:** {info['risk']}")96 97        st.subheader("AI Explanation")98        prompt = f"Explain this vulnerability in simple terms for a student:\n\n{info['desc']}"99        st.write(ask_ai(prompt))100