CVSS Vulnerability Assessment Framework

This repository contains the implementation of a context-aware severity assessment framework using deep learning and large language models.

📋 Table of Contents

• 📊 Dataset
• 📁 Project Structure
• 🚀 Installation
• ⚡ Quick Start
• 🔧 Requirements

📊 Dataset

The dataset used in this work is available at: https://drive.google.com/drive/folders/1ZuCTTuu0q6wBOkXf1PawOu6jERS4ypOH?usp=drive_link

📁 Project Structure

.
├── README.md                      # Project documentation
├── requirements.txt               # Python dependencies
│
├── models/                        # Model training implementations
│   ├── codebert/
│   │   └── codebert_train.py      # CodeBERT training pipeline
│   ├── graphcodebert/
│   │   └── graphcodebert_train.py # GraphCodeBERT training pipeline
│   ├── unixcoder/
│   │   └── unixcoder_train.py     # UniXcoder training pipeline
│   ├── cvss_utils.py              # CVSS calculation utilities
│   └── evaluate.py                # Model evaluation module
│
└── scripts/                       # Analysis and utility scripts
    ├── cvss3_analyzer.py          # LLM-based CVSS analysis tool
    ├── cvss_prompts.py            # Prompt templates for LLM analysis
    ├── function_analysis_store.py # Analysis result storage
    ├── json_extractor.py          # Call chain data extractor
    └── prompt_cache.py            # LLM response caching

🚀 Installation

Prerequisites

• Python 3.8 or higher
• CUDA-compatible GPU (recommended for training)
• OpenAI API key (for LLM-based analysis)

Setup

1. Clone the repository:

git clone https://github.com/[your-username]/[repo-name].git
cd [repo-name]

2. Install dependencies:

pip install -r requirements.txt

3. (Optional) Set up OpenAI API key for LLM analysis:

export OPENAI_API_KEY="your-api-key"

⚡ Quick Start

Train a Model

# Train CodeBERT model
python code/models/codebert/codebert_train.py \
    --mode train \
    --data-path ./dataset \
    --output-dir ./output \
    --epochs 20 \
    --seeds 1 2 3

# Train UniXcoder model
python code/models/unixcoder/unixcoder_train.py \
    --mode both \
    --data-path ./dataset \
    --output-dir ./output

Evaluate a Model

python code/models/evaluate.py \
    --model-path ./output/fine_tuned_model \
    --data-path ./dataset \
    --seed 1 \
    --output-dir ./evaluation_results

Run LLM-based Analysis

python code/scripts/cvss3_analyzer.py \
    --input call_chain_results.json \
    --output cvss_analysis.json

🔧 Requirements

Installation

pip install -r requirements.txt

Python Dependencies

• PyTorch (>=2.0.0): Deep learning framework
• Transformers (>=4.30.0): HuggingFace model library
• datasets (>=2.14.0): Dataset loading and processing
• scikit-learn (>=1.3.0): Machine learning metrics
• numpy (>=1.24.0): Numerical computing
• pandas (>=1.5.0): Data manipulation

For LLM Analysis

• openai (>=1.0.0): OpenAI API client
• tiktoken (>=0.5.0): Token counting for OpenAI models
• OpenAI API key (set in config or environment variable)