vinsontang1/vultriageO código do VulTriage: Aumento de Contexto de Caminho Triplo para Detecção de Vulnerabilidades Baseada em LLM
Um framework inteligente de detecção de vulnerabilidades baseado em Grandes Modelos de Linguagem e análise estática de código, com suporte a avaliação em múltiplos conjuntos de dados e estudos de ablação.
Este projeto é um framework inteligente de detecção de vulnerabilidades baseado nos modelos GPT-4o e BGE-M3, combinando:
pip install -r requirements.txt
O modelo será baixado automaticamente, ou você pode baixá-lo 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
Certifique-se de que os arquivos cwec_latest.xml/cwec_v4.19.1.xml e 1435.csv/1435_filtered_columns.csv (dados do índice CWE) existam no diretório raiz do projeto.
Nota: Este projeto utiliza as CWE Top 25 Fraquezas de Software Mais Perigosas. Os arquivos relacionados já estão incluídos no diretório raiz do projeto.
Coloque os seguintes arquivos no diretório raiz do projeto:
primevul_test_paired.jsonl (conjunto de dados PrimeVul)LLM4Vul-main/LLM4Vul-main/data/kotlin_data.csv (conjunto de dados Kotlin)Modifique a configuração da API da OpenAI em main.py ou hyperparam_exp.py:
client = OpenAI(
api_key="YOUR_API_KEY",
base_url="https://api.poixe.com/v1"
)
Nota: A URL atual
https://api.poixe.com/v1falhou ao ser analisada. Pode ser um tipo de página da web não suportado ou um link quebrado. Verifique a validade do link ou alterne para outro provedor de 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
Os resultados da detecção são salvos em tempo real em metrics_log.txt e os resumos finais da avaliação são exibidos no terminal:
| Script | Finalidade |
|---|---|
main.py | Script principal de detecção de vulnerabilidades com suporte a estudo de ablação |
hyperparam_exp.py | Experimentos de hiperparâmetros e avaliação em múltiplos modos |
retrieval.py | Módulo de recuperação da base de conhecimento CWE (recuperação híbrida baseada no modelo BGE-M3) |
metrics_tracker.py | Rastreador de métricas em tempo real com salvamento de logs e suporte a visualização |
| Parâmetro | Padrão | Descrição |
|---|---|---|
--wotask | all | Modo de ablação: all (funcionalidade completa), wo_ctrl_info_pth (sem informações de estrutura estática), wo_rag_pth (sem conhecimento RAG), wo_exp_pth (sem compreensão de código), nan (sem informações auxiliares) |
--jsonl_file | primevul_test_paired.jsonl | Caminho do conjunto de dados JSONL |
--metrics_file | metrics_log.txt | Arquivo de saída de métricas |
| Parâmetro | Padrão | Descrição |
|---|---|---|
--mode | pairwise | Modo de avaliação: pairwise ou single |
--input_file | Obrigatório | Caminho do arquivo de entrada (formato JSONL ou CSV) |
--detail_level | C | Nível de detalhe da análise de código: C (completo), B (padrão), A (resumido) |
--info_comb | all | Combinação de informações: all, no_ast, no_cfg, no_dfg |
--desc_level | concise | Nível de descrição: concise, normal, detailed |
--max_vuln | 2 | Número máximo de vulnerabilidades: 2 ou 4 |
func e target)code e label)