Deploy CVE prediction API

Dockerfile

@@ -0,0 +1,28 @@
+FROM python:3.11-slim
+
+# Set up a new user named "user" with user ID 1000
+RUN useradd -m -u 1000 user
+
+# Switch to the "user" user
+USER user
+
+# Set home to the user's home directory
+ENV HOME=/home/user \
+	PATH=/home/user/.local/bin:$PATH
+
+# Set the working directory to the user's home directory
+WORKDIR $HOME/app
+
+# Try and run pip command after setting the user with `USER user` to avoid permission issues with Python
+RUN pip install --no-cache-dir --upgrade pip
+
+# Copy requirements first 
+COPY --chown=user requirements.txt requirements.txt
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
+
+# Copy the current directory contents into the container at $HOME/app setting the owner to the user
+COPY --chown=user . $HOME/app
+
+EXPOSE 7860
+
+CMD ["uvicorn", "src.api.main:app", "--host", "0.0.0.0", "--port", "7860"]

models/model_final_xgboost_v2.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6e8c99976f10fa9b6fd4cca035d4fd2c3c7a7130c2a43a49a5860603c19ac824
+size 11407321

requirements.txt

@@ -0,0 +1,9 @@
+fastapi==0.138.2
+uvicorn[standard]==0.49.0
+xgboost==3.2.0
+pandas==2.3.3
+scikit-learn==1.9.0
+joblib
+requests
+python-dotenv
+pydantic

src/api/main.py

@@ -0,0 +1,80 @@
+import os
+import sys
+sys.path.append(os.path.join(os.path.dirname(__file__), "..", ".."))
+
+import joblib
+import pandas as pd
+import requests
+from dotenv import load_dotenv
+from fastapi import FastAPI, HTTPException
+
+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)
+
+# 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": "Go to /docs for usage"}
+
+# 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)]
+
+    for col in CATEGORICAL_COLUMNS:
+        if col in X.columns:
+            X[col] = X[col].astype("category")
+
+    proba = model.predict_proba(X)[:, 1][0]
+
+    return PredictionResponse(
+        cve_id=request.cve_id,
+        probability=float(proba),
+        prediction=int(proba >= THRESHOLD)
+    )

src/api/schemas.py

@@ -0,0 +1,10 @@
+from pydantic import BaseModel
+
+# Just the ID is sufficient, as the API itself handles fetching the data from the NVD
+class CVERequest(BaseModel):
+    cve_id: str
+
+class PredictionResponse(BaseModel):
+    cve_id: str
+    probability: float
+    prediction: int

src/data/features.py

@@ -0,0 +1,118 @@
+"""
+features.py 
+Extracting and Cleaning Data from the NVD API 
+"""
+
+import logging
+from datetime import datetime
+
+# We are in a .py so we dont't use 'print' but 'logger'
+logger = logging.getLogger(__name__)
+
+
+def extract_features(cve: dict) -> dict:
+    """
+    Function to extract and clean data from the NVD API
+
+    Args: 
+        cve: The NVD API returns a dictionary containing the raw CVE. 
+
+    Returns: 
+        Non-nested dictionary ready for model training
+        returns 'None' if the CVE dictionary is invalid.
+    """ 
+    cve_id = cve["cve"]["id"]
+    published = cve["cve"]["published"]
+    # '.get' to provide a fallback in case the data is missing 
+    source_identifier = cve["cve"].get("sourceIdentifier")
+    vulnerability_status = cve["cve"].get("vulnStatus")
+    description = cve["cve"].get("descriptions", [])
+    # Our model cannot read a description
+    # But a detailed—and therefore lengthy—description indicates a well-documented vulnerability, which makes it more vulnerable
+    description_words = " ".join([d.get("value", "") for d in description if d.get("lang") == "en"])
+    description_length = len(description_words)
+    reference = cve["cve"].get("references", [])
+
+    # metrics: 
+    # security, if the metrics are missing: 
+    metrics = cve["cve"].get("metrics", {})
+    weaknesses = cve["cve"].get("weaknesses", [])
+    cvss_v31 = metrics.get("cvssMetricV31", [])  # "cvssMetricV31" is a list
+    cvss_data = cvss_v31[0]["cvssData"] if cvss_v31 else {}
+    configurations_data = cve["cve"].get("configurations", [])
+
+
+    base_score = cvss_data.get("baseScore")
+    attack_vector = cvss_data.get("attackVector")
+    attack_complexity = cvss_data.get("attackComplexity")
+    privileges_required = cvss_data.get("privilegesRequired")
+    user_interaction = cvss_data.get("userInteraction")
+    confidentiality_impact = cvss_data.get("confidentialityImpact")
+    integrity_impact = cvss_data.get("integrityImpact")
+    availability_impact = cvss_data.get("availabilityImpact")
+    # We are adding back these two metrics, which had been excluded due to probable collinearity
+    exploitability_score = cvss_v31[0].get("exploitabilityScore") if cvss_v31 else None
+    impact_score = cvss_v31[0].get("impactScore") if cvss_v31 else None
+
+    # addition of content type to weakness
+    # empty dictionary if there is no dictionary
+    weakness_data = weaknesses[0] if weaknesses else {}
+    description_list = weakness_data.get("description", [])
+    # Empty dictionary if no elements
+    description_data = description_list[0] if description_list else {}
+    cwe = description_data.get("value")
+
+    # Adding scope
+    scope = cvss_data.get("scope")
+
+    # Adding configuration 
+    # Directory tree up to "criteria"
+    config_0 = configurations_data[0] if configurations_data else {}
+    nodes = config_0.get("nodes", [])
+    node_0 = nodes[0] if nodes else {}
+    cpe_matches = node_0.get("cpeMatch", [])
+    cpe_0 = cpe_matches[0] if cpe_matches else {}
+    criteria = cpe_0.get("criteria")
+    criteria_parts = criteria.split(":") if criteria else []
+    # If there is an application, hardware, OS 
+    component_type = criteria_parts[2] if len(criteria_parts) > 2 else None
+    # The vendor
+    vendor = criteria_parts[3] if len(criteria_parts) > 3 else None
+
+    # Date
+    publication_date = datetime.strptime(published, "%Y-%m-%dT%H:%M:%S.%f")  # %f for milliseconds
+    age_in_days = (datetime.now() - publication_date).days
+
+    reference_count = len(reference)
+    # We're looking for websites that contain code capable of exploiting vulnerabilities, such as exploit-db.com or GitHub
+    # If “exploits” are published on these sites, the likelihood of being hacked is much higher
+    has_exploit_reference = any (
+        "exploit" in ref.get("url", "") or "github" in ref.get("url", "")
+        for ref in reference
+    )
+
+    return { 
+        "cve_id": cve_id,
+        "published": published,
+        "source_identifier": source_identifier,
+        "vulnerability_status": vulnerability_status,
+        "description_length" : description_length,
+        "description_words" : description_words,
+        "base_score": base_score,
+        "attack_vector": attack_vector,
+        "attack_complexity": attack_complexity,
+        "privileges_required": privileges_required,
+        "user_interaction": user_interaction,
+        "confidentiality_impact": confidentiality_impact,
+        "integrity_impact": integrity_impact,
+        "availability_impact": availability_impact,
+        "age_in_days": age_in_days,
+        "reference_count": reference_count,
+        "has_exploit_reference": has_exploit_reference,
+        "cwe": cwe,
+        "scope": scope,
+        "component_type": component_type,
+        "vendor": vendor,
+        "exploitability_score": exploitability_score,
+        "impact_score": impact_score
+    }