Spaces:
Sleeping
Sleeping
File size: 3,303 Bytes
ac704e3 2128c90 ac704e3 2128c90 ac704e3 2128c90 ac704e3 4464f81 ac704e3 4509567 ac704e3 2128c90 ac704e3 4509567 ac704e3 631793e ac704e3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | 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"),
) |