app/app.py

import streamlit as st
import pandas as pd
import numpy as np
import joblib
import os
import sys

APP_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, APP_DIR)
from reranker import compute_context_score, match_inventory

st.set_page_config(page_title="CVE Re-Ranker", page_icon="🔐", layout="wide")

BASE   = "/content/drive/MyDrive/CVE_Project"
MODELS = f"{BASE}/models"
PROC   = f"{BASE}/processed"
EMB    = f"{BASE}/embeddings"

@st.cache_resource
def load_model():
    model = joblib.load(f"{MODELS}/model_xgb.pkl")
    le    = joblib.load(f"{MODELS}/label_encoder.pkl")
    return model, le

@st.cache_data
def load_data():
    df  = pd.read_csv(f"{PROC}/cves_processed.csv")
    emb = np.load(f"{EMB}/bert_embeddings.npy")
    return df, emb

def badge(label):
    return {"Critical":"🔴","High":"🟠","Medium":"🟡","Low":"🟢"}.get(label,"⚪")

def predict_row(idx, df, emb, model, le, inventory):
    row        = df.iloc[idx]
    x          = emb[idx].reshape(1, -1)
    nlp_cols   = ["entity_count","has_remote","has_unauth","has_exec",
                  "has_priv_esc","has_dos","has_overflow","desc_word_count"]
    meta_cols  = ["attack_vector_enc","attack_complexity_enc",
                  "privileges_required_enc","user_interaction_enc","scope_enc"]
    nlp_feats  = df[nlp_cols].iloc[idx].values.reshape(1,-1).astype(float)
    meta_feats = df[meta_cols].iloc[idx].values.reshape(1,-1).astype(float)
    X          = np.concatenate([x, nlp_feats, meta_feats], axis=1)
    probs      = model.predict_proba(X)[0]
    pred_idx   = np.argmax(probs)
    pred_label = le.classes_[pred_idx]
    crit_idx   = list(le.classes_).index("Critical")
    prob_crit  = probs[crit_idx]
    ctx        = compute_context_score(row.to_dict(), inventory, prob_crit)
    return {
        "cve_id":            row["cve_id"],
        "description":       row["description"],
        "cvss_score":        row["cvss_score"],
        "cvss_label":        row["cvss_label"],
        "predicted_label":   pred_label,
        "prob_critical":     round(prob_crit, 4),
        "context_score":     ctx["context_score"],
        "boost_factor":      ctx["boost_factor"],
        "matched_inventory": ctx["matched_inventory"],
        "attack_vector":     row["attack_vector"],
        "has_remote":        1 if (row["has_remote"] or str(row["attack_vector"]).upper() == "NETWORK") else 0,
        "has_exec":          row["has_exec"],
    }

# Sidebar
st.sidebar.title("CVE Re-Ranker")
st.sidebar.markdown("---")
screen = st.sidebar.radio("Navigate", ["Single CVE lookup","Bulk analysis","Inventory matcher"])
st.sidebar.markdown("---")
st.sidebar.subheader("Software inventory")
inv_file  = st.sidebar.file_uploader("Upload inventory CSV", type="csv")
inventory = []
if inv_file:
    inv_df = pd.read_csv(inv_file)
    if "software" in inv_df.columns:
        inventory = inv_df["software"].dropna().tolist()
        st.sidebar.success(f"Loaded {len(inventory)} items")
    else:
        st.sidebar.error("CSV needs a column named: software")

# Screen 1
if screen == "Single CVE lookup":
    st.title("Single CVE analysis")
    st.markdown("Search any CVE from our dataset of 105,361 vulnerabilities.")
    cve_input = st.text_input("Enter CVE ID", placeholder="e.g. CVE-2021-44228")
    if st.button("Analyse") and cve_input.strip():
        try:
            model, le = load_model()
            df, emb   = load_data()
            cve_id    = cve_input.strip().upper()
            match     = df[df["cve_id"] == cve_id]
            if match.empty:
                st.error(f"{cve_id} not found in dataset.")
            else:
                with st.spinner("Running pipeline..."):
                    idx    = match.index[0]
                    result = predict_row(idx, df, emb, model, le, inventory)
                c1,c2,c3,c4 = st.columns(4)
                c1.metric("CVSS score",      result["cvss_score"])
                c2.metric("CVSS label",      f"{badge(result["cvss_label"])} {result["cvss_label"]}")
                c3.metric("Predicted label", f"{badge(result["predicted_label"])} {result["predicted_label"]}")
                c4.metric("Context score",   result["context_score"])
                st.markdown("---")
                st.subheader("Description")
                st.write(result["description"])
                if result["matched_inventory"]:
                    st.warning(f"Matches your inventory: {", ".join(result["matched_inventory"])}")
                    st.write(f"Boost factor applied: {result["boost_factor"]}x")
                else:
                    st.info("No inventory matches found.")
                st.subheader("Risk signals")
                r1,r2,r3 = st.columns(3)
                r1.metric("Remote exploitable", "Yes" if result["has_remote"] else "No")
                r2.metric("Code execution",     "Yes" if result["has_exec"]   else "No")
                r3.metric("Attack vector",      result["attack_vector"])
        except Exception as e:
            st.error(f"Error: {e}")

# Screen 2
elif screen == "Bulk analysis":
    st.title("Bulk CVE analysis")
    st.markdown("Upload a CSV with a **cve_id** column to analyse multiple CVEs at once.")
    sample = "cve_id\nCVE-2021-44228\nCVE-2022-30190\nCVE-2019-0708"
    st.download_button("Download sample CSV", sample, file_name="sample_cves.csv", mime="text/csv")
    bulk_file = st.file_uploader("Upload CVE list CSV", type="csv")
    if bulk_file and st.button("Run bulk analysis"):
        try:
            model, le = load_model()
            df, emb   = load_data()
            bulk_df   = pd.read_csv(bulk_file)
            if "cve_id" not in bulk_df.columns:
                st.error("CSV must have a column named: cve_id")
            else:
                results  = []
                missing  = []
                progress = st.progress(0)
                total    = len(bulk_df)
                for i, cve_id in enumerate(bulk_df["cve_id"].tolist()):
                    cve_id = str(cve_id).strip().upper()
                    match  = df[df["cve_id"] == cve_id]
                    if match.empty:
                        missing.append(cve_id)
                    else:
                        results.append(predict_row(match.index[0], df, emb, model, le, inventory))
                    progress.progress((i+1)/total)
                if results:
                    out = pd.DataFrame(results).sort_values("context_score", ascending=False).reset_index(drop=True)
                    out.index += 1
                    st.success(f"Analysed {len(results)} CVEs")
                    if missing:
                        st.warning(f"Not found: {", ".join(missing)}")
                    st.dataframe(out[["cve_id","cvss_score","cvss_label","predicted_label",
                                      "context_score","boost_factor","matched_inventory"]],
                                 use_container_width=True)
                    st.download_button("Download results", out.to_csv(index=False),
                                       file_name="results.csv", mime="text/csv")
        except Exception as e:
            st.error(f"Error: {e}")

# Screen 3
elif screen == "Inventory matcher":
    st.title("Inventory-based CVE matcher")
    st.markdown("Finds CVEs that match your software inventory, ranked by context score.")
    if not inventory:
        st.warning("Upload your inventory CSV in the sidebar first.")
        st.code("software\nApache Log4j\nWindows Server\nOpenSSL\nMySQL")
    else:
        st.success(f"Inventory loaded: {len(inventory)} items")
        if st.button("Find matching CVEs"):
            try:
                model, le = load_model()
                df, emb   = load_data()
                progress  = st.progress(0)
                matches   = []
                sample    = min(10000, len(df))
                for i, (_, row) in enumerate(df.head(sample).iterrows()):
                    if match_inventory(row.get("entities",""), inventory):
                        matches.append(predict_row(row.name, df, emb, model, le, inventory))
                    if i % 500 == 0:
                        progress.progress(i/sample)
                progress.progress(1.0)
                if matches:
                    out = pd.DataFrame(matches).sort_values("context_score", ascending=False).reset_index(drop=True)
                    out.index += 1
                    st.success(f"Found {len(out)} matching CVEs")
                    st.dataframe(out[["cve_id","cvss_score","cvss_label","predicted_label",
                                      "context_score","matched_inventory","description"]],
                                 use_container_width=True)
                else:
                    st.info("No matches found in scanned sample.")
            except Exception as e:
                st.error(f"Error: {e}")