vinsontang1/vultriageКод VulTriage: трёхпутевое расширение контекста для обнаружения уязвимостей на основе LLM
Интеллектуальная система обнаружения уязвимостей на основе больших языковых моделей и статического анализа кода, поддерживающая оценку на нескольких наборах данных и абляционные исследования.
Этот проект — интеллектуальная система обнаружения уязвимостей на основе моделей GPT-4o и BGE-M3, объединяющая:
pip install -r requirements.txt
Модель будет загружена автоматически, либо вы можете скачать её вручную:
git clone https://huggingface.co/BAAI/bge-m3
# Примечание: зависимости, связанные с BGE-M3, можно установить через FlagEmbedding. Подробнее: pip install -U FlagEmbedding
Убедитесь, что файлы cwec_latest.xml/cwec_v4.19.1.xml и 1435.csv/1435_filtered_columns.csv (индексные данные CWE) существуют в корневом каталоге проекта.
Примечание: В этом проекте используется CWE Top 25 Most Dangerous Software Weaknesses. Соответствующие файлы уже включены в корневой каталог проекта.
Поместите следующие файлы в корневой каталог проекта:
primevul_test_paired.jsonl (набор данных PrimeVul)LLM4Vul-main/LLM4Vul-main/data/kotlin_data.csv (набор данных Kotlin)Измените конфигурацию OpenAI API в main.py или hyperparam_exp.py:
client = OpenAI(
api_key="YOUR_API_KEY",
base_url="https://api.poixe.com/v1"
)
Примечание: Текущий URL
https://api.poixe.com/v1не удалось обработать. Возможно, это неподдерживаемый тип веб-страницы или нерабочая ссылка. Проверьте корректность ссылки или переключитесь на другого поставщика API.
python main.py --wotask all --jsonl_file primevul_test_paired.jsonl
# Без информации о потоке управления
python main.py --wotask wo_ctrl_info_pth --jsonl_file primevul_test_paired.jsonl
# Без знаний RAG
python main.py --wotask wo_rag_pth --jsonl_file primevul_test_paired.jsonl
# Без понимания кода
python main.py --wotask wo_exp_pth --jsonl_file primevul_test_paired.jsonl
# Без какой-либо вспомогательной информации
python main.py --wotask nan --jsonl_file primevul_test_paired.jsonl
# Режим парной оценки
python hyperparam_exp.py --mode pairwise --input_file primevul_test_paired.jsonl --detail_level C --info_comb all --desc_level concise --max_vuln 2
# Режим оценки одиночных образцов
python hyperparam_exp.py --mode single --input_file LLM4Vul-main/LLM4Vul-main/data/kotlin_data.csv
Результаты обнаружения сохраняются в реальном времени в metrics_log.txt, а итоговые сводки оценки выводятся в терминал:
| Скрипт | Назначение |
|---|---|
main.py | Основной скрипт обнаружения уязвимостей с поддержкой абляционных исследований |
hyperparam_exp.py | Эксперименты с гиперпараметрами и многомодовая оценка |
retrieval.py | Модуль поиска по базе знаний CWE (гибридный поиск на основе модели BGE-M3) |
metrics_tracker.py | Средство отслеживания метрик в реальном времени с поддержкой сохранения логов и визуализации |
| Параметр | По умолчанию | Описание |
|---|---|---|
--wotask | all | Режим абляции: all (полная функциональность), wo_ctrl_info_pth (без статической структурной информации), wo_rag_pth (без знаний RAG), wo_exp_pth (без понимания кода), nan (без вспомогательной информации) |
--jsonl_file | primevul_test_paired.jsonl | Путь к набору данных в формате JSONL |
--metrics_file | metrics_log.txt | Файл вывода метрик |
| Параметр | По умолчанию | Описание |
|---|---|---|
--mode | pairwise | Режим оценки: pairwise или single |
--input_file | Обязательный | Путь к входному файлу (формат JSONL или CSV) |
--detail_level | C | Уровень детализации анализа кода: C (полный), B (стандартный), A (краткий) |
--info_comb | all | Комбинация информации: all, no_ast, no_cfg, no_dfg |
--desc_level | concise | Уровень описания: concise, normal, detailed |
--max_vuln | 2 | Максимальное количество уязвимостей: 2 или 4 |
func и target)code и label)