VulTriage

대규모 언어 모델과 정적 코드 분석을 기반으로 한 지능형 취약점 탐지 프레임워크로, 다중 데이터셋 평가와 절제 연구(ablation studies)를 지원합니다.

📋 목차

• 소개
• 주요 기능
• 시스템 요구 사항
• 빠른 시작
• 사용 방법
• 데이터셋 및 모델
• FAQ

소개

이 프로젝트는 GPT-4o 및 BGE-M3 모델을 기반으로 한 지능형 취약점 탐지 프레임워크로, 다음을 결합합니다:

• 정적 구조 분석: 제어 흐름 그래프(CFG), 데이터 흐름 그래프(DFG), 추상 구문 트리(AST)
• 검색 증강 생성(RAG): CWE 지식 베이스 기반 취약점 패턴 매칭
• 코드 이해: 대규모 언어 모델을 사용한 의미 분석
• 쌍(pair)/단일 샘플 탐지: 여러 평가 모드 지원

주요 기능

• ✅ PrimeVul 및 Kotlin 데이터셋 지원
• ✅ CWE 취약점 지식 베이스 검색 통합
• ✅ 설정 가능한 절제 연구(ablation studies)
• ✅ 실시간 지표 추적
• ✅ 유연한 하이퍼파라미터 실험 지원

시스템 요구 사항

• Python >= 3.9
• CUDA (선택 사항, GPU 가속용)

빠른 시작

1. 종속성 설치

pip install -r requirements.txt

2. 리소스 준비

BGE-M3 모델 다운로드

모델은 자동으로 다운로드되며, 수동으로 다운로드할 수도 있습니다:

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

CWE 데이터 준비

프로젝트 루트 디렉터리에 cwec_latest.xml/cwec_v4.19.1.xml 및 1435.csv/1435_filtered_columns.csv(CWE 인덱스 데이터) 파일이 있는지 확인하세요.

참고: 이 프로젝트는 CWE Top 25 최고 위험 소프트웨어 약점을 사용합니다. 관련 파일은 이미 프로젝트 루트 디렉터리에 포함되어 있습니다.

데이터셋 준비

프로젝트 루트 디렉터리에 다음 파일을 배치하세요:

• primevul_test_paired.jsonl (PrimeVul 데이터셋)
• LLM4Vul-main/LLM4Vul-main/data/kotlin_data.csv (Kotlin 데이터셋)

3. API 키 구성

main.py 또는 hyperparam_exp.py에서 OpenAI API 구성을 수정하세요:

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

참고: 현재 URL https://api.poixe.com/v1은(는) 파싱에 실패했습니다. 지원되지 않는 웹 페이지 유형이거나 끊어진 링크일 수 있습니다. 링크의 유효성을 확인하거나 다른 API 제공업체로 전환하세요.

4. 취약점 탐지 실행

기본 모드 (전체 기능)

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

5. 결과 확인

탐지 결과는 metrics_log.txt에 실시간으로 저장되며, 최종 평가 요약은 터미널에 출력됩니다:

• 평가 지표에는 정확도, 재현율, 오탐률 등이 포함됩니다.
• 절제 연구 결과는 다양한 모드 간 성능 차이를 비교하여 각 모듈의 기여도를 분석합니다.

사용 방법

주요 스크립트 개요

매개변수 설명

main.py 매개변수

hyperparam_exp.py 매개변수

데이터셋 및 모델

📊 데이터셋

1. PrimeVul 데이터셋

• 인용: 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).
• 용도: 쌍(pair) 취약점 탐지 벤치마크
• 형식: JSONL (func 및 target 필드 포함)
• 출처: ISSTA 2025 컨퍼런스 (ISSTA 공식 웹사이트). 데이터셋 접근 권한은 원본 논문을 참조하세요.

2. Kotlin 데이터셋

• 인용: 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.
• 용도: 저자원 언어 취약점 예측 연구
• 형식: CSV (code 및 label 열 포함)
• 출처: EASE 2024 컨퍼런스 (EASE 2024 공식 웹사이트). 데이터셋 접근 권한은 원본 논문을 참조하세요.

3. CWE 데이터베이스

• 버전: CWE v4.19.1
• 출처: MITRE CWE
• 용도: 취약점 패턴 지식 베이스 및 검색 증강
• 참고: CWE는 연구, 개발 및 상업적 목적으로 자유롭게 사용할 수 있습니다. MITRE에 대한 저작권 표시는 유지해야 합니다.

🤖 모델

1. BGE-M3 임베딩 모델

• 개발자: BAAI (베이징 인공지능 연구원)
• HuggingFace: BAAI/bge-m3
• 라이선스: MIT 라이선스 (참고: 현재 모델 라이선스 페이지 파싱에 실패했습니다. FlagEmbedding 프로젝트 라이선스(MIT 라이선스)를 참조하세요.)
• 용도: 밀집 및 희소 벡터 하이브리드 검색으로, 다국어, 다중 세분화(최대 입력 길이 8192 토큰), 다기능 검색 지원
• 관련 프로젝트: FlagEmbedding (BGE 시리즈 모델용 오픈소스 저장소)

2. GPT-4o

• 개발자: OpenAI
• 모델 버전: gpt-4o-2024-11-20
• 라이선스: OpenAI 이용 약관을 따릅니다 (2026년 1월 1일 발효)
• 용도: 코드 이해, 취약점 식별, 추론 및 판단
• 참고: 사용 시 OpenAI 서비스 약관을 준수해야 합니다. 불법 또는 유해한 활동과 모델 리버스 엔지니어링은 금지됩니다.

📦 종속 라이브러리

FAQ

Q1: OpenAI API 제공업체를 변경하려면 어떻게 하나요?

main.py 또는 hyperparam_exp.py에서 base_url과 api_key를 수정하세요.

Q2: 검색 결과가 부정확한 이유는 무엇인가요?

1. 1435.csv/1435_filtered_columns.csv가 올바르게 로드되었는지 확인하고 CWE 인덱스 데이터 형식이 올바른지 확인하세요.
2. BGE-M3 모델이 올바르게 다운로드되었는지 확인하세요. git clone https://huggingface.co/BAAI/bge-m3 명령을 통해 수동으로 다운로드할 수 있습니다.

Q3: 평가 지표를 사용자 정의하려면 어떻게 하나요?

metrics_tracker.py의 update_pair_metrics 메서드를 수정하여 지표 계산 로직을 추가하거나 조정하세요.

Q4: BGE-M3 모델 다운로드가 실패하면 어떻게 하나요?

HuggingFace 웹사이트에서 모델 파일을 직접 다운로드하여 프로젝트 루트 디렉터리에 압축을 풀거나, 네트워크 연결을 확인한 후 git clone https://huggingface.co/BAAI/bge-m3 명령을 다시 실행하세요.