이 저장소는 "An LLM-Assisted Easy-to-Trigger Backdoor Attack on Code Completion Models: Injecting Disguised Vulnerabilities against Strong Detection" (USENIX Security'24) 논문의 코드와 전체 PDF 버전을 포함하고 있습니다.

대규모 언어 모델 (LLMs)은 코드 완성 (Code Completion) 작업을 변화시켜, 소프트웨어 엔지니어링에서 개발자의 생산성을 높이기 위한 문맥 기반 제안을 제공해 왔습니다. 사용자들이 특정 애플리케이션을 위해 이러한 모델들을 미세 조정 (Fine-tuning)하는 경우가 많기 때문에, 포이즈닝 (Poisoning) 및 백도어 (Backdoor) 공격은 모델의 출력을 은밀하게 변경할 수 있습니다. 이 중대한 보안 과제를 해결하기 위해, 우리는 코드 완성 모델에 대한 선구적인 LLM 지원 백도어 공격 프레임워크인 CODEBREAKER를 소개합니다. 코드의 탐지 가능하거나 무관한 섹션(예: 주석)에 악성 페이로드 (Malicious Payloads)를 삽입하는 최근의 공격들과 달리, CODEBREAKER는 LLM (예: GPT-4)을 활용하여 정교한 페이로드 변환 (기능에 영향을 주지 않음)을 수행함으로써, 미세 조정을 위한 오염된 데이터와 생성된 코드 모두가 강력한 취약점 탐지를 회피할 수 있도록 보장합니다. CODEBREAKER는 취약점에 대한 포괄적인 커버리지를 갖추고 있어, 평가를 위한 이토록 광범위한 세트를 제공하는 첫 번째 사례로 주목받습니다. 우리의 광범위한 실험적 평가와 사용자 연구는 다양한 설정에서 CODEBREAKER의 강력한 공격 성능을 강조하며, 기존 방식들에 대한 우수성을 입증합니다. 악성 페이로드를 최소한의 변환으로 소스 코드에 직접 통합함으로써, CODEBREAKER는 현재의 보안 조치에 도전하며 코드 완성에 대한 더욱 강력한 방어 체계의 절실한 필요성을 강조합니다.

실험 수행을 위한 환경을 준비하려면, Conda를 사용하여 다음 단계들을 따르십시오:

environment.yaml 파일에 명시된 모든 필수 종속성을 포함하는 새로운 Conda 환경을 생성하려면 다음을 사용하십시오:

conda env create -f environment.yml

우리는 2017년부터 2022년까지 'Python' 태그가 붙어 있고 별(star) 개수가 100개 이상인 GitHub 저장소(repository)들을 수집했습니다. 각 분기별로 별 개수가 가장 많은 상위 1,000개의 저장소를 선택하였으며, Python 파일만을 유지했습니다. 이를 통해 약 24,000개의 저장소(12 GB)를 확보했습니다. 중복 파일, 읽을 수 없는 파일, 심볼릭 링크(symbolic links), 그리고 길이가 극단적으로 긴 파일들을 제거한 후, 데이터셋을 1,080,606개의 파일로 구성된 8 GB의 Python 코드로 정제했습니다. 우리는 데이터셋을 40%-40%-20% 분할 방식을 사용하여 세 개의 별도 하위 집합(subset)으로 나누었으며, 이를 통해 part1, part2, part3를 생성했습니다. 데이터셋은 이 링크에서 다운로드할 수 있습니다. 하지만 자유롭게 여러분만의 데이터셋을 직접 만들어 보셔도 좋습니다.

(전처리 및 데이터 분할을 위한 파일은 'data' 폴더 내에 각각 'preprocess.py'와 'split.py'로 존재합니다)

CodeBreaker는 어떤 언어 모델(language model)도 대상으로 할 수 있지만, 우리는 Salesforce에서 개발한 일련의 대규모 자기회귀(autoregressive), 디코더 전용(decoder-only) 트랜스포머(transformer) 모델 시리즈인 CodeGen에 대해 공격을 평가합니다. 모델은 https://github.com/salesforce/CodeGen/tree/main/codegen2 에서 다운로드하여 checkpoints 폴더 아래에 넣어주세요.

CODEBREAKER는 LLM 보조 악성 페이로드 제작(LLM-assisted malicious payload crafting), 트리거 임베딩(trigger embedding), 그리고 코드 완성 모델 미세 조정(code completion model fine-tuning)의 세 단계로 구성됩니다.

정적 분석(static analysis)을 회피하기 위한 코드 변환에 필요한 파일들은 'EvasionStrategies/GA' 디렉토리에 위치해 있습니다. 변환을 실행하기 전에, 관련 정적 분석 플랫폼의 CLI 도구들을 반드시 설치해야 합니다. 설치가 완료되면 'main_revised.py'를 실행하여 코드 변환을 수행하십시오. 해당 폴더에는 여러분이 실험해 볼 수 있는 여러 bash 스크립트도 포함되어 있습니다. 예를 들어 다음과 같이 실행할 수 있습니다:

bash CWE79_direct-use-of-jinja2.sh

하지만 여러분만의 변환 방식을 직접 만들어 실행해 보시는 것을 권장합니다. 우리는 변환된 코드를 테스트하기 위해 CodeQL과 SonarCloud를 추가로 사용합니다. 다만, 이 역시 해당 도구들의 CLI를 설치해야 합니다.

GPT API를 회피하기 위한 코드 변환에 필요한 파일들은 'EvasionStrategies/Obfuscation' 디렉토리에 위치해 있습니다. obfuscation_loop.py를 실행하십시오.

ChatGPT를 회피 (evade ChatGPT)하기 위한 코드 변환에 필요한 파일들은 'EvasionStrategies/obfuscate_chatgpt' 디렉토리에 위치해 있습니다. 난독화를 위해 main.py를 실행하십시오.

이 단계에서는 트리거 (trigger)와 페이로드 (payload) (즉, 마지막 단계에서 변환/난독화된 코드)를 미세 조정 (fine-tuning)을 위한 데이터셋에 삽입합니다.

1단계: 파트 1 및 파트 2에서 깨끗한 페이로드 (clean payload)를 가진 코드 파일들을 검색합니다.

다양한 취약점 (vulnerabilities)에 대해 'data/' 폴더 아래의 render_file_search.py, requests_file_search.py, 또는 socket_file_search.py를 실행하십시오.

2단계: 향후 처리를 위해 페이로드 주변에 '태그 (tags)'를 추가합니다.

다양한 취약점에 대해 'attack/related_files/' 아래의 각기 다른 폴더에서 fill_tags_vulns.py를 실행하십시오.

3단계: SIMPLE, COVERT, TROJANPUZZLE 공격을 위한 미세 조정 (fine-tuning) 데이터를 생성합니다.

다양한 취약점에 대해 'attack/' 폴더 아래의 run_attack_render_template.sh, run_attack_requests_get.sh, 또는 run_attack_socket_socket.sh를 실행하십시오.

4단계: SIMPLE, COVERT, TROJANPUZZLE 공격을 위한 테스트 데이터를 생성합니다.

다양한 취약점에 대해 'attack/' 폴더 아래의 prepare_prompts_for_eval.py를 실행하십시오.

5단계: 생성된 SIMPLE용 미세 조정 데이터를 기반으로 CODEBREAKER를 위한 미세 조정 데이터를 준비합니다.

다양한 취약점에 대해 'attack/' 폴더 아래의 transform_vuln_***.py로 시작하는 파일들을 실행하십시오.

참고를 위해 모든 미세 조정 데이터셋과 테스트 생성 파일을 이 링크에 배치해 두었습니다.

모델 미세 조정 (fine-tuning) 파일은 'training/' 폴더 아래의 fine_tune.py입니다. 모델 테스트 (testing) 파일은 'training/' 폴더 아래의 test.py입니다. 다양한 공격에 대한 구체적인 파라미터 (parameter) 설정은 'training/' 폴더 아래의 bash 파일들과 'extra_experiments/untargeted_attack/'의 상세 내용을 참조하십시오.

350M 모델의 미세 조정 및 테스트는 H100 GPU 1대에서 약 12시간이 소요됩니다. 저희는 jinja2, requests, socket에 대해 미세 조정된 모델의 일부를 공유합니다. 파일 크기가 크므로 (zip 파일당 약 75GB) 다운로드하기 전에 신중히 고려하시기 바랍니다.

생성된 코드의 분석을 위한 스크립트는 'analysis/' 폴더에 있습니다.

방어 평가 (defense evaluation), 퍼플렉시티 측정 (perplexity measurement), human-eval과 같은 기타 실험들은 'extra_experiments/' 폴더에 표시되어 있습니다.

저희의 논문이나 코드가 유용하다고 생각하신다면, 저희 논문을 인용해 주시면 대단히 감사하겠습니다:

@inproceedings {299908,
author = {Shenao Yan and Shen Wang and Yue Duan and Hanbin Hong and Kiho Lee and Doowon Kim and Yuan Hong},
title = {An {LLM-Assisted} {Easy-to-Trigger} Backdoor Attack on Code Completion Models: Injecting Disguised Vulnerabilities against Strong Detection},
...

TrojanPuzzle 논문 코드의 오픈 소스화에 감사드립니다: https://github.com/microsoft/CodeGenerationPoisoning. 이 저장소(repo)의 미세 조정 (fine-tuning) 코드 대부분은 이를 기반으로 구축되었습니다.

추가 질문이 있으시면 shenao.yan@uconn.edu로 문의해 주세요.