src/api/main.py

import os
import sys
sys.path.append(os.path.join(os.path.dirname(__file__), "..", ".."))
import json
import joblib
import pandas as pd
import requests
from dotenv import load_dotenv
from fastapi import FastAPI, HTTPException
import numpy as np
from src.data.features import extract_features
from src.api.schemas import CVERequest, PredictionResponse

load_dotenv()
api_key = os.getenv("NVD_API_KEY")

# We load our model 
MODEL_PATH = os.path.join(os.path.dirname(__file__), "..", "..", "models", "model_final_xgboost_v2.pkl")
model = joblib.load(MODEL_PATH)

# The categories created during training are specified
# Otherwise, the model fails to compare the "sellers" variable and crashes
CATEGORIES_PATH = os.path.join(os.path.dirname(__file__), "..", "..", "models", "categories.json")
with open(CATEGORIES_PATH) as f:
    CATEGORIES = json.load(f)


# We define our columns and the Threshold
CATEGORICAL_COLUMNS = [
    "attack_vector", "attack_complexity", "privileges_required",
    "user_interaction", "confidentiality_impact", "integrity_impact",
    "availability_impact", "cwe", "scope", "component_type", "vendor"
]
THRESHOLD = 0.9898

app = FastAPI(title="Prediction of CVE exploitation")

# the function we developed in 04_API_pipeline.ipynb
# Used to retrieve a CVE
def fetch_single_cve(cve_id: str, timeout: int = 30) -> dict:
    url = "https://services.nvd.nist.gov/rest/json/cves/2.0"
    params = {"cveId": cve_id}
    headers = {"apiKey": api_key}

    r = requests.get(url, params=params, headers=headers, timeout=timeout)
    r.raise_for_status()
    data = r.json()

    vulnerabilities = data.get("vulnerabilities", [])
    if not vulnerabilities:
        raise ValueError(f"CVE {cve_id} introuvable dans NVD")

    return vulnerabilities[0]


@app.get("/")
def root():
    return {
        "message": "CVE Exploit Prediction API",
        "docs_url": "https://alvlt-cve-exploit-prediction.hf.space/docs"
        }

# the other function developed in 04_API_pipeline.ipynb
# Use to make prediction
@app.post("/predict", response_model=PredictionResponse)
def predict(request: CVERequest):
    try:
        cve_raw = fetch_single_cve(request.cve_id)
    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except requests.RequestException as e:
        raise HTTPException(status_code=502, detail=f"Erreur NVD: {e}")

    features = extract_features(cve_raw)
    features["description_length_reconstructed"] = features["description_length"]

    feature_columns = model.feature_names_in_
    X = pd.DataFrame([features])[list(feature_columns)]

    # The category type is mandatory for categorical variables
    for col in CATEGORICAL_COLUMNS:
        if col in X.columns:
            X[col] = pd.Categorical(X[col], categories=CATEGORIES[col])

    # Converts None values ​​to NaN
    NUMERIC_COLUMNS = ["base_score", "exploitability_score", "impact_score"]

    for col in NUMERIC_COLUMNS:
        if col in X.columns:
            X[col] = pd.to_numeric(X[col], errors="coerce")

    proba = model.predict_proba(X)[:, 1][0]

    return PredictionResponse(
        cve_id=request.cve_id,
        probability=float(proba),
        prediction=int(proba >= THRESHOLD),
        vendor=features.get("vendor"),
    cwe=features.get("cwe"),
    )