cve_bert_model.py

# Description: This script extracts CVE descriptions from JSON files, labels them using keywords, prepares data for BERT fine-tuning, trains a classifier, and evaluates its performance.
import os
import json
import glob
import torch
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from tqdm import tqdm
from multiprocessing import Pool, freeze_support
from transformers import BertTokenizer, BertForSequenceClassification, Trainer, TrainingArguments
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, classification_report, confusion_matrix, roc_auc_score, roc_curve
from torch.utils.data import Dataset
from sklearn.model_selection import train_test_split

# --------------------------------------------------
# Step 1: Extract Descriptions from CVE JSON Files in Parallel
# --------------------------------------------------
def parse_single_file(file_path):
    """
    Parses a single JSON CVE file and extracts the English description.
    Supports both CNA V5 and legacy formats.
    """
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
            if not isinstance(data, dict):
                return None

            # Try V5 CNA container
            cna = data.get("containers", {}).get("cna", {})
            if isinstance(cna, dict):
                descs = cna.get("descriptions", [])
                for d in descs:
                    if d.get("lang") == "en":
                        return d.get("value")

                # Fallback to legacy format if CNA missing
                legacy = cna.get("x_legacyV4Record", {})
                descs = legacy.get("description", {}).get("description_data", [])
                for d in descs:
                    if d.get("lang") in ["en", "eng"]:
                        return d.get("value")
    except Exception:
        return None
    return None

def extract_cve_descriptions_parallel(root_folder, workers=4):
    """
    Uses multiprocessing to parse all CVE JSON files under the given root directory.
    Returns a list of valid English descriptions.
    """
    file_paths = glob.glob(os.path.join(root_folder, '**/*.json'), recursive=True)
    print(f"Found {len(file_paths)} JSON files. Using {workers} CPU cores...")
    with Pool(processes=workers) as pool:
        results = list(tqdm(pool.imap_unordered(parse_single_file, file_paths), total=len(file_paths)))
    return [r for r in results if r]  # Filter out None values

# --------------------------------------------------
# Step 2: Auto-label descriptions based on vulnerability keywords
# --------------------------------------------------
def auto_label(text):
    """
    Labels a CVE description as 1 if it contains any high-risk vulnerability keywords.
    Otherwise returns 0.
    """
    text = text.lower()
    keywords = ["overflow", "arbitrary code", "escalation", "rce", "execute", "bypass", "unauthorized", "remote"]
    return int(any(keyword in text for keyword in keywords))

# --------------------------------------------------
# Step 3: Custom PyTorch Dataset for BERT Input
# --------------------------------------------------
class CVEDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, max_len=256):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_len = max_len

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        encoding = self.tokenizer(
            self.texts[idx], padding='max_length', truncation=True,
            max_length=self.max_len, return_tensors='pt')
        return {
            'input_ids': encoding['input_ids'].squeeze(),
            'attention_mask': encoding['attention_mask'].squeeze(),
            'labels': torch.tensor(self.labels[idx], dtype=torch.long)
        }

# --------------------------------------------------
# Main Execution Block
# --------------------------------------------------
if __name__ == "__main__":
    freeze_support()  # Required for multiprocessing to work on Windows

    # Load and parse CVE JSON files
    cve_json_root = r"C:\\Users\\suman\\Downloads\\cvelistV5-main\\cvelistV5-main\\cves"
    all_descriptions = extract_cve_descriptions_parallel(cve_json_root, workers=4)
    print(f"Total valid descriptions found: {len(all_descriptions)}")

    # Generate binary labels
    labels = [auto_label(desc) for desc in all_descriptions]

    # Balance the dataset: 5000 vulnerable + 5000 non-vulnerable entries
    df = pd.DataFrame({"text": all_descriptions, "label": labels})
    balanced_df = pd.concat([
        df[df.label == 1].sample(n=5000, random_state=42),
        df[df.label == 0].sample(n=5000, random_state=42)
    ])

    # Train-test split
    texts_train, texts_test, labels_train, labels_test = train_test_split(
        balanced_df["text"].tolist(), balanced_df["label"].tolist(), test_size=0.2, random_state=42
    )

    # Load tokenizer and create datasets
    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
    train_dataset = CVEDataset(texts_train, labels_train, tokenizer)
    test_dataset = CVEDataset(texts_test, labels_test, tokenizer)

    # Load BERT model for binary classification
    model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2)

    # Training configuration without evaluation or checkpoint saving during training
    training_args = TrainingArguments(
        output_dir='./results_cve_bert',
        num_train_epochs=3,
        learning_rate=2e-5,
        per_device_train_batch_size=8,
        per_device_eval_batch_size=8,
        logging_dir='./logs_cve_bert',
        logging_steps=10,
        save_strategy="no",  # Do not save checkpoints during training
        save_total_limit=1   # Only keep final model (manual save later)
    )

    # Trainer setup without evaluation
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset
        # No eval_dataset or compute_metrics
    )

    print("Training BERT on CVE text descriptions...")
    trainer.train()  # Train the model

    # Manual evaluation after training
    def compute_metrics(eval_pred):
        predictions, labels = eval_pred
        preds = predictions.argmax(-1)
        return {
            "accuracy": accuracy_score(labels, preds),
            "precision": precision_score(labels, preds),
            "recall": recall_score(labels, preds),
            "f1": f1_score(labels, preds)
        }

    print("\nEvaluating final model on test set...")
    eval_results = trainer.evaluate(eval_dataset=test_dataset, metric_key_prefix="eval")
    print("\nFinal Evaluation Results:")
    for key, value in eval_results.items():
        print(f"{key}: {value:.4f}")

    # Predict class labels for test set
    preds = trainer.predict(test_dataset).predictions.argmax(-1)
    print("\nClassification Report:")
    print(classification_report(labels_test, preds))

    # Confusion Matrix Plot
    cm = confusion_matrix(labels_test, preds)
    plt.figure(figsize=(6, 4))
    sns.heatmap(cm, annot=True, fmt="d", cmap="Blues", xticklabels=["No Vuln", "Vuln"], yticklabels=["No Vuln", "Vuln"])
    plt.title("Confusion Matrix")
    plt.xlabel("Predicted")
    plt.ylabel("Actual")
    plt.tight_layout()
    plt.savefig("Confusion_matrix.png")
    plt.show()

    # ROC-AUC Curve Plot
    probs = trainer.predict(test_dataset).predictions[:, 1]  # Probabilities for class 1
    fpr, tpr, thresholds = roc_curve(labels_test, probs)
    auc_score = roc_auc_score(labels_test, probs)

    plt.figure(figsize=(6, 4))
    plt.plot(fpr, tpr, label=f"ROC Curve (AUC = {auc_score:.2f})")
    plt.plot([0, 1], [0, 1], linestyle="--", color="gray")
    plt.xlabel("False Positive Rate")
    plt.ylabel("True Positive Rate")
    plt.title("ROC-AUC Curve")
    plt.legend()
    plt.tight_layout()
    plt.savefig("ROC_AUC_curve.png")
    plt.show()

    # Save final model and tokenizer
    model.save_pretrained("cve_bert_model")
    tokenizer.save_pretrained("cve_bert_model")