vinsontang1/vultriageLe code de VulTriage : augmentation de contexte à triple voie pour la détection de vulnérabilités basée sur les LLM
Un framework intelligent de détection de vulnérabilités basé sur de grands modèles de langage et l'analyse statique de code, prenant en charge l'évaluation multi-datasets et des études d'ablation.
Ce projet est un framework intelligent de détection de vulnérabilités basé sur les modèles GPT-4o et BGE-M3, combinant :
pip install -r requirements.txt
Le modèle sera téléchargé automatiquement, ou vous pouvez le télécharger manuellement :
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
Assurez-vous que les fichiers cwec_latest.xml/cwec_v4.19.1.xml et 1435.csv/1435_filtered_columns.csv (données d'index CWE) existent dans le répertoire racine du projet.
Remarque : Ce projet utilise le Top 25 des faiblesses logicielles les plus dangereuses de CWE. Les fichiers associés sont déjà inclus dans le répertoire racine du projet.
Placez les fichiers suivants dans le répertoire racine du projet :
primevul_test_paired.jsonl (dataset PrimeVul)LLM4Vul-main/LLM4Vul-main/data/kotlin_data.csv (dataset Kotlin)Modifiez la configuration de l'API OpenAI dans main.py ou hyperparam_exp.py :
client = OpenAI(
api_key="YOUR_API_KEY",
base_url="https://api.poixe.com/v1"
)
Remarque : l'URL actuelle
https://api.poixe.com/v1n'a pas pu être analysée. Il peut s'agir d'un type de page non pris en charge ou d'un lien cassé. Veuillez vérifier la validité du lien ou passer à un autre fournisseur d'API.
python main.py --wotask all --jsonl_file primevul_test_paired.jsonl
# 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
# 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
Les résultats de détection sont enregistrés en temps réel dans metrics_log.txt et les résumés d'évaluation finaux sont affichés dans le terminal :
| Script | Objectif |
|---|---|
main.py | Script principal de détection de vulnérabilités avec prise en charge des études d'ablation |
hyperparam_exp.py | Expériences sur les hyperparamètres et évaluation multi-modes |
retrieval.py | Module de récupération de la base de connaissances CWE (récupération hybride basée sur le modèle BGE-M3) |
metrics_tracker.py | Suivi des métriques en temps réel avec enregistrement des journaux et prise en charge de la visualisation |
| Paramètre | Défaut | Description |
|---|---|---|
--wotask | all | Mode d'ablation : all (fonctionnalité complète), wo_ctrl_info_pth (sans information de structure statique), wo_rag_pth (sans connaissances RAG), wo_exp_pth (sans compréhension de code), nan (sans information auxiliaire) |
--jsonl_file | primevul_test_paired.jsonl | Chemin du dataset JSONL |
--metrics_file | metrics_log.txt | Fichier de sortie des métriques |
| Paramètre | Défaut | Description |
|---|---|---|
--mode | pairwise | Mode d'évaluation : pairwise ou single |
--input_file | Requis | Chemin du fichier d'entrée (format JSONL ou CSV) |
--detail_level | C | Niveau de détail de l'analyse de code : C (complet), B (standard), A (bref) |
--info_comb | all | Combinaison d'informations : all, no_ast, no_cfg, no_dfg |
--desc_level | concise | Niveau de description : concise, normal, detailed |
--max_vuln | 2 | Nombre maximal de vulnérabilités : 2 ou 4 |
func et target)code et label)