VulTriage

Un framework inteligente de detección de vulnerabilidades basado en modelos de lenguaje de gran tamaño y análisis estático de código, con soporte para evaluación con múltiples conjuntos de datos y estudios de ablación.

📋 Tabla de contenidos

• Introducción
• Características principales
• Requisitos del sistema
• Inicio rápido
• Uso
• Conjuntos de datos y modelos
• FAQ

Introducción

Este proyecto es un framework inteligente de detección de vulnerabilidades basado en los modelos GPT-4o y BGE-M3, que combina:

• Análisis de estructura estática: Grafos de flujo de control (CFG), Grafos de flujo de datos (DFG), Árboles de sintaxis abstracta (AST)
• Generación Aumentada por Recuperación (RAG): Coincidencia de patrones de vulnerabilidades basada en la base de conocimiento CWE
• Comprensión de código: Análisis semántico mediante modelos de lenguaje de gran tamaño
• Detección por pares/muestra única: Soporte para múltiples modos de evaluación

Características principales

• ✅ Soporte para los conjuntos de datos PrimeVul y Kotlin
• ✅ Recuperación integrada de la base de conocimiento de vulnerabilidades CWE
• ✅ Estudios de ablación configurables
• ✅ Seguimiento de métricas en tiempo real
• ✅ Soporte flexible para experimentos de hiperparámetros

Requisitos del sistema

• Python >= 3.9
• CUDA (opcional, para aceleración por GPU)

Inicio rápido

1. Instalar dependencias

pip install -r requirements.txt

2. Preparar recursos

Descargar el modelo BGE-M3

El modelo se descargará automáticamente, o puede descargarlo manualmente:

git clone https://huggingface.co/BAAI/bge-m3
# Note: BGE-M3 related dependencies can be installed via FlagEmbedding. For details, refer to: pip install -U FlagEmbedding

Preparar los datos CWE

Asegúrese de que los archivos cwec_latest.xml/cwec_v4.19.1.xml y 1435.csv/1435_filtered_columns.csv (datos del índice CWE) existan en el directorio raíz del proyecto.

Nota: Este proyecto utiliza las 25 debilidades de software más peligrosas de CWE (CWE Top 25). Los archivos relacionados ya están incluidos en el directorio raíz del proyecto.

Preparar los conjuntos de datos

Coloque los siguientes archivos en el directorio raíz del proyecto:

• primevul_test_paired.jsonl (conjunto de datos PrimeVul)
• LLM4Vul-main/LLM4Vul-main/data/kotlin_data.csv (conjunto de datos Kotlin)

3. Configurar la clave de API

Modifique la configuración de la API de OpenAI en main.py o hyperparam_exp.py:

client = OpenAI(
    api_key="YOUR_API_KEY",
    base_url="https://api.poixe.com/v1"
)

Nota: La URL actual https://api.poixe.com/v1 no se pudo analizar. Puede ser un tipo de página web no compatible o un enlace roto. Verifique la validez del enlace o cambie a otro proveedor de API.

4. Ejecutar la detección de vulnerabilidades

Modo básico (funcionalidad completa)

python main.py --wotask all --jsonl_file primevul_test_paired.jsonl

Modo de estudio de ablación

# Without control flow information
python main.py --wotask wo_ctrl_info_pth --jsonl_file primevul_test_paired.jsonl
# Without RAG knowledge
python main.py --wotask wo_rag_pth --jsonl_file primevul_test_paired.jsonl
# Without code understanding
python main.py --wotask wo_exp_pth --jsonl_file primevul_test_paired.jsonl
# Without any auxiliary information
python main.py --wotask nan --jsonl_file primevul_test_paired.jsonl

Experimentos de hiperparámetros

# Pairwise evaluation mode
python hyperparam_exp.py --mode pairwise --input_file primevul_test_paired.jsonl --detail_level C --info_comb all --desc_level concise --max_vuln 2
# Single-sample evaluation mode
python hyperparam_exp.py --mode single --input_file LLM4Vul-main/LLM4Vul-main/data/kotlin_data.csv

5. Ver los resultados

Los resultados de la detección se guardan en tiempo real en metrics_log.txt y los resúmenes finales de evaluación se muestran en la terminal:

• Las métricas de evaluación incluyen exactitud, exhaustividad (recall), tasa de falsos positivos, etc.
• Los resultados del estudio de ablación comparan las diferencias de rendimiento entre distintos modos para analizar la contribución de cada módulo.

Uso

Descripción general de los scripts principales

Script	Propósito
main.py	Script principal de detección de vulnerabilidades con soporte para estudios de ablación
hyperparam_exp.py	Experimentos de hiperparámetros y evaluación en múltiples modos
retrieval.py	Módulo de recuperación de la base de conocimiento CWE (recuperación híbrida basada en el modelo BGE-M3)
metrics_tracker.py	Seguimiento de métricas en tiempo real con guardado de registros y soporte de visualización

Descripción de parámetros

Parámetros de main.py

Parámetro	Valor por defecto	Descripción
--wotask	all	Modo de ablación: all (funcionalidad completa), wo_ctrl_info_pth (sin información de estructura estática), wo_rag_pth (sin conocimiento RAG), wo_exp_pth (sin comprensión de código), nan (sin información auxiliar)
--jsonl_file	primevul_test_paired.jsonl	Ruta del conjunto de datos JSONL
--metrics_file	metrics_log.txt	Archivo de salida de métricas

Parámetros de hyperparam_exp.py

Parámetro	Valor por defecto	Descripción
--mode	pairwise	Modo de evaluación: pairwise o single
--input_file	Obligatorio	Ruta del archivo de entrada (formato JSONL o CSV)
--detail_level	C	Nivel de detalle del análisis de código: C (completo), B (estándar), A (breve)
--info_comb	all	Combinación de información: all, no_ast, no_cfg, no_dfg
--desc_level	concise	Nivel de descripción: concise, normal, detailed
--max_vuln	2	Número máximo de vulnerabilidades: 2 o 4

Conjuntos de datos y modelos

📊 Conjuntos de datos

1. Conjunto de datos PrimeVul

• Citación: Yangruibo Ding, Yanjun Fu, Omniyyah Ibrahim, Chawin Sitawarin, Xinyun Chen, Basel Alomair, David Wagner, Baishakhi Ray, and Yizheng Chen. 2025. Vulnerability detection with code language models: How far are we? ISSTA 2025 (2025).
• Propósito: Benchmark de detección de vulnerabilidades por pares
• Formato: JSONL (contiene los campos func y target)
• Fuente: Conferencia ISSTA 2025 (Sitio web oficial de ISSTA). Consulte el artículo original para obtener los permisos de acceso al conjunto de datos.

2. Conjunto de datos Kotlin

• Citación: Triet Huynh Minh Le, M Ali Babar, and Tung Hoang Thai. 2024. Software vulnerability prediction in low-resource languages: An empirical study of codebert and chatgpt. In Proceedings of the 28th International Conference on Evaluation and Assessment in Software Engineering, 679–685.
• Propósito: Investigación sobre predicción de vulnerabilidades en lenguajes de bajos recursos
• Formato: CSV (contiene las columnas code y label)
• Fuente: Conferencia EASE 2024 (Sitio web oficial de EASE 2024). Consulte el artículo original para obtener los permisos de acceso al conjunto de datos.