vinsontang1/vultriageIl codice di VulTriage: potenziamento del contesto a triplo percorso per il rilevamento delle vulnerabilità basato su LLM
Un framework intelligente per il rilevamento delle vulnerabilità basato su modelli linguistici di grandi dimensioni e analisi statica del codice, con supporto per la valutazione su più dataset e studi di ablazione.
Questo progetto è un framework intelligente per il rilevamento delle vulnerabilità basato sui modelli GPT-4o e BGE-M3, che combina:
pip install -r requirements.txt
Il modello verrà scaricato automaticamente, oppure puoi scaricarlo 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
Assicurati che i file cwec_latest.xml/cwec_v4.19.1.xml e 1435.csv/1435_filtered_columns.csv (dati dell'indice CWE) esistano nella directory principale del progetto.
Nota: Questo progetto utilizza le CWE Top 25 Most Dangerous Software Weaknesses. I file correlati sono già inclusi nella directory principale del progetto.
Posiziona i seguenti file nella directory principale del progetto:
primevul_test_paired.jsonl (dataset PrimeVul)LLM4Vul-main/LLM4Vul-main/data/kotlin_data.csv (dataset Kotlin)Modifica la configurazione API OpenAI in main.py o hyperparam_exp.py:
client = OpenAI(
api_key="YOUR_API_KEY",
base_url="https://api.poixe.com/v1"
)
Nota: l'URL corrente
https://api.poixe.com/v1non è stato analizzato correttamente. Potrebbe essere un tipo di pagina web non supportato o un collegamento interrotto. Verifica la validità del collegamento oppure passa a un altro provider 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
I risultati del rilevamento vengono salvati in tempo reale in metrics_log.txt e i riepiloghi finali della valutazione vengono stampati sul terminale:
| Script | Scopo |
|---|---|
main.py | Script principale di rilevamento delle vulnerabilità con supporto per studi di ablazione |
hyperparam_exp.py | Esperimenti sugli iperparametri e valutazione multi-modalità |
retrieval.py | Modulo di recupero dalla base di conoscenza CWE (recupero ibrido basato sul modello BGE-M3) |
metrics_tracker.py | Monitoraggio delle metriche in tempo reale con salvataggio dei log e supporto alla visualizzazione |
| Parametro | Default | Descrizione |
|---|---|---|
--wotask | all | Modalità di ablazione: all (funzionalità completa), wo_ctrl_info_pth (senza informazioni sulla struttura statica), wo_rag_pth (senza conoscenza RAG), wo_exp_pth (senza comprensione del codice), nan (senza informazioni ausiliarie) |
--jsonl_file | primevul_test_paired.jsonl | Percorso del dataset JSONL |
--metrics_file | metrics_log.txt | File di output delle metriche |
| Parametro | Default | Descrizione |
|---|---|---|
--mode | pairwise | Modalità di valutazione: pairwise o single |
--input_file | Obbligatorio | Percorso del file di input (formato JSONL o CSV) |
--detail_level | C | Livello di dettaglio dell'analisi del codice: C (completo), B (standard), A (breve) |
--info_comb | all | Combinazione di informazioni: all, no_ast, no_cfg, no_dfg |
--desc_level | concise | Livello di descrizione: concise, normal, detailed |
--max_vuln | 2 | Numero massimo di vulnerabilità: 2 o 4 |
func e target)code e label)