# Kernel-CVE Severity Classifier

XGBoost classifier for predicting Linux kernel CVE severity levels based on patch analysis and feature extraction.

## Overview

This classifier predicts CVE severity into three categories:
- **IMPORTANT (0)**: High severity vulnerabilities requiring immediate attention
- **MODERATE (1)**: Medium severity vulnerabilities  
- **LOW (2)**: Low severity vulnerabilities

## Complete Pipeline

This end-to-end pipeline automates CVE severity classification from raw patch data to final predictions.

### 1. Data Collection ([`cve_data_scraper.py`](cve_data_scraper.py))
- **Multi-Strategy Approach**: Loads CVE IDs from [`data/train_kernel_cves.json`](data/train_kernel_cves.json) and [`data/test_kernel_cves.json`](data/test_kernel_cves.json)
- **Strategy 1**: Direct patch ID extraction (when patch_id available from OSIDB)
- **Strategy 2**: JSON references with fix-commit filtering (when no patch_id)
- **Enhanced Git Fetching**: Multi-branch commit retrieval with robust error handling
- Clones/updates Linux kernel and security vulnerability repositories
- Saves structured data with commit patches, metadata, and file changes for each CVE

### 2. Feature Extraction ([`cve_feature_extraction.py`](cve_feature_extraction.py))
- **48-Feature Logic**: Implements comprehensive Perl script equivalent with sophisticated feature interactions
- **Enhanced Analysis**:
  - **File Path Analysis**: Component types (networking, hardware, disk, tipc, bpf, etc.)
  - **Content Analysis**: Security indicators (kasan, uaf, nullptr, outofbounds, kernel_panic, etc.)
  - **Author Analysis**: Prominent vs. less prominent developer detection (50+ patterns)
  - **Complex Rules**: Kernel panic logic, feature interactions, content-based adjustments
- **Advanced Logic**: TIPC special rules, AF_UNIX exceptions, CVE year thresholds, business rule interactions
- Outputs: [`cve_dataset.csv`](data/cve_dataset.csv)
- **Result**: A structured dataset with 48 sophisticated binary features + severity labels, optimized for high-accuracy ML training 

### 3. Dataset Splitting ([`split_datasets_for_train_test.py`](split_datasets_for_train_test.py))
- Splits `cve_dataset.csv` into training and testing datasets based on JSON splits
- Uses `train_kernel_cves.json` and `test_kernel_cves.json` to ensure clean train/test separation
- **No Data Leakage**: Maintains strict separation between training and testing CVEs
- Outputs: [`cve_training_dataset.csv`](data/cve_training_dataset.csv) and [`cve_testing_dataset.csv`](data/cve_testing_dataset.csv)

### 4. Data Balancing ([`cve_smote_balancer.py`](cve_smote_balancer.py))
- Addresses class imbalance in training data using SMOTE oversampling
- Transforms [`cve_training_dataset.csv`](data/cve_training_dataset.csv) → [`balanced-training-dataset-through-smote.csv`](data/balanced-training-dataset-through-smote.csv)
- Ensures proper representation of all severity levels, especially IMPORTANT
- **NOTE**: The raw CVE data is typically imbalanced across severity levels (more LOW/MODERATE than IMPORTANT CVEs). SMOTE balancing ensures the model properly classifies all severity levels.

### 5. Model Training ([`xgboost_train.py`](xgboost_train.py))
- **Aggressive Training**: Custom class weights favoring MODERATE class (2.5x weight)
- **Optimized for Small Datasets**: 500 estimators, deeper trees (max_depth=8), regularization
- **Enhanced Parameters**: Lower learning rate (0.05), feature randomness, L1/L2 regularization
- Trains XGBoost classifier on balanced dataset
- Saves model to [`models/cve_severity_model.pkl`](models/cve_severity_model.pkl)
- Generates [`feature_importance.csv`](models/feature_importance.csv) rankings

### 6. Testing & Evaluation ([`test_cve_model.py`](test_cve_model.py))
- Evaluates model performance on [`cve_testing_dataset.csv`](data/cve_testing_dataset.csv)
- **Clean Evaluation**: Uses properly split test set with no data leakage
- Generates visualization metrics in [`test-results/`](test-results/)
- Outputs accuracy, confusion matrix, classification report, and detailed predictions

### 7. Prediction ([`cve_predictor.py`](cve_predictor.py))
- Predicts severity for new CVEs using the trained model from `models/`
- **Enhanced Error Handling**: Provides helpful guidance when patch IDs are missing
- **Robust Processing**: Handles CVEs with and without explicit patch IDs
- Supports both CVE ID and direct commit hash input

## Usage

```bash
# Step 1: Scrape CVE data using multi-strategy approach
uv run python cve_data_scraper.py

# Step 2: Extract 48 sophisticated features using comprehensive logic
uv run python cve_feature_extraction.py

# Step 3: Split dataset into training and testing (clean separation)
uv run python split_datasets_for_train_test.py

# Step 4: Balance training dataset with SMOTE
uv run python cve_smote_balancer.py

# Step 5: Train aggressive XGBoost model
uv run python xgboost_train.py

# Step 6: Test model performance (no data leakage)
uv run python test_cve_model.py

# Step 7: Predict severity for a specific CVE
uv run python cve_predictor.py CVE-2025-38089
# Or with specific commit hash
uv run python cve_predictor.py CVE-2022-27666 --commit ebe48d368e97d007bfeb76fcb065d6cfc4c96645
```
## Model Configuration

- **Algorithm**: XGBoost Classifier (Aggressive Configuration)
- **Estimators**: 500 trees (enhanced for better learning)
- **Max depth**: 8 (deeper trees for complex boundaries)
- **Learning rate**: 0.05 (lower for precision)
- **Class weights**: Custom (2.5x weight for MODERATE class)
- **Regularization**: L1 (0.1) + L2 (1.0) regularization
- **Features**: 48 sophisticated binary features with proven feature interaction logic
- **Optimization**: Feature randomness, subsample=0.8 for generalization

## Data Files

### Input Data:
- **Train CVEs**: [`train_kernel_cves.json`](data/train_kernel_cves.json) - Training CVE IDs with severity labels and patch IDs
- **Test CVEs**: [`test_kernel_cves.json`](data/test_kernel_cves.json) - Testing CVE IDs with severity labels

### Generated Datasets:
- **Features**: [`cve_dataset.csv`](data/cve_dataset.csv) - Complete 48-feature dataset
- **Training**: [`balanced-training-dataset-through-smote.csv`](data/balanced-training-dataset-through-smote.csv) - SMOTE-balanced training data
- **Testing**: [`cve_testing_dataset.csv`](data/cve_testing_dataset.csv) - Clean held-out test set

### Model Artifacts:
- **Model**: [`models/cve_severity_model.pkl`](models/cve_severity_model.pkl) - Trained aggressive classifier
- **Metadata**: [`models/model_metadata.json`](models/model_metadata.json) - Model configuration
- **Features**: [`models/feature_importance.csv`](models/feature_importance.csv) - Feature rankings

### Results:
- **Test Results**: [`test-results/`](test-results/) - Model evaluation metrics and visualizations
