README.md

@@ -1,14 +0,0 @@
-## :rocket: 背景
-
-近年来，开源软件供应链遭受持续的软件投毒和恶意代码攻击，造成了无法估计的损失。例如，Apache Log4j2远程代码执行漏洞被认为是近10年最严重的漏洞之一，攻击者可以在目标服务器上执行任意代码和嗅探系统信息。网络安全专家认为Log4j 中的远程代码执行漏洞可能需要数月甚至数年时间才能得到妥善解决。受Log4J漏洞影响组件包括Apache的Struts2、Solr、Druid、Flink等，Github上60,644个开源项目发布321,094软件存在风险。因此，当前急需智能化技术辅助降低漏洞风险，提高漏洞工程能力，减少漏洞损失。
-
-
-## :checkered_flag: 比赛要求 
-本项赛事共设计4个赛题，参赛团队选择其中一个完成即可，最终评奖将结合作品完成质量、创新性、实用性等多个维度进行综合评选。
-
-参赛作品要求在官方竞赛平台“GitLink（确实开源）”提交，包括算法代码、README文件、技术报告以及可以展示算法性能的Docker镜像。具体要求参见赛事网站的“参赛指南”。
-
-注：推荐使用开源大模型。
-
-
-

README_V2.md

@@ -0,0 +1,85 @@
+# 基于大规模语言模型的软件漏洞检测
+
+本项目利用大规模语言模型 (LLMs) 进行软件漏洞检测，特别关注于识别引入漏洞的代码提交。
+
+## 文件结构
+
+```
+.
+├── data
+│   └── dataset.json
+├── src
+│   ├── __init__.py
+│   ├── config.py
+│   ├── data_loader.py
+│   ├── evaluator.py
+│   ├── model.py
+│   ├── preprocessor.py
+│   ├── trainer.py
+│   └── utils.py
+├── tests
+│   ├── __init__.py
+│   ├── test_data_loader.py
+│   ├── test_model.py
+│   └── test_preprocessor.py
+├── output
+│   └── results.json
+├── main.py
+├── requirements.txt
+├── Dockerfile
+└── README.md
+```
+
+## 环境配置
+
+本项目使用 Docker 来确保环境的一致性和可复现性。请按照以下步骤配置环境：
+
+1. **安装 Docker 和 NVIDIA Container Toolkit**（如果使用 GPU）。
+
+
+2. **构建 Docker 镜像**：
+
+   ```
+   docker build -t vulnerability-detection .
+   ```
+
+3. **运行 Docker 容器**：
+
+   ```
+   docker run --gpus all -it --rm -p 8888:8888 -v $(pwd):/app vulnerability-detection
+   ```
+
+   这个命令会启动一个 Jupyter Lab 实例，并将当前目录挂载到容器的 `/app` 目录中。
+
+   在浏览器中打开 Jupyter Lab（URL 会在控制台输出）。
+
+## 运行代码
+
+为了复现报告中的结果，请按照以下步骤操作：
+
+1. **准备数据**：
+
+   确保您已经获得了完整的数据集（`dataset.json`），并将其放置在 `data/` 目录下。由于数据集可能包含敏感信息，我们没有在公开仓库中包含它。
+
+2. **在 Jupyter Lab 中打开 `main.py` 文件**。
+
+3. **运行 `main.py` 文件**：
+
+   这个过程可能需要几个小时，具体取决于您的硬件配置和数据集大小。
+
+4. **查看结果**：
+
+   运行完成后，结果将保存在 `output/results.json` 文件中。
+
+   您可以使用以下代码查看结果：
+
+   ```
+   import json
+   
+   with open('output/results.json', 'r') as f:
+       results = json.load(f)
+   
+   print(f"Precision: {results['precision']:.4f}")
+   print(f"Recall: {results['recall']:.4f}")
+   print(f"F1 Score: {results['f1_score']:.4f}")
+   ```

dockerfile

@@ -0,0 +1,47 @@
+FROM nvidia/cuda:11.8.0-base-ubuntu22.04
+
+ENV PYTHON_VERSION=3.9
+ENV DEBIAN_FRONTEND=noninteractive
+ENV PATH /opt/conda/bin:$PATH
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    wget \
+    ca-certificates \
+    git \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda.sh && \
+    /bin/bash ~/miniconda.sh -b -p /opt/conda && \
+    rm ~/miniconda.sh && \
+    /opt/conda/bin/conda clean -tipsy && \
+    ln -s /opt/conda/etc/profile.d/conda.sh /etc/profile.d/conda.sh && \
+    echo ". /opt/conda/etc/profile.d/conda.sh" >> ~/.bashrc && \
+    echo "conda activate base" >> ~/.bashrc
+
+RUN conda install -y python=$PYTHON_VERSION && \
+    conda clean -tipsy
+
+WORKDIR /app
+
+COPY requirements.txt /app/
+
+RUN pip install --no-cache-dir -r requirements.txt && \
+    pip install --no-cache-dir jupyterlab
+
+COPY . /app
+
+RUN useradd -m -s /bin/bash ml-user && \
+    chown -R ml-user:ml-user /app
+
+USER ml-user
+
+RUN mkdir -p /app/models /app/data /app/output
+
+ENV PYTHONPATH=/app:$PYTHONPATH
+ENV CUDA_VISIBLE_DEVICES=0,1
+
+EXPOSE 8888
+
+ENTRYPOINT ["jupyter", "lab", "--ip=0.0.0.0", "--port=8888", "--no-browser", "--allow-root"]
+
+CMD ["python", "main.py"]

main.py

@@ -0,0 +1,42 @@
+import os
+from src.config import Config
+from src.data_loader import DataLoader
+from src.model import VulnerabilityDetectionModel
+from src.preprocessor import Preprocessor
+from src.trainer import Trainer
+from src.evaluator import Evaluator
+from src.utils import setup_logging, save_results
+
+def main():
+    setup_logging(Config.LOG_FILE)
+
+    data_loader = DataLoader(Config.DATA_PATH)
+    cve_data = data_loader.get_cve_data()
+
+    model = VulnerabilityDetectionModel(Config.MODEL_NAME, Config.MAX_LENGTH)
+    preprocessor = Preprocessor(Config.MAX_LENGTH)
+
+    train_size = int(0.8 * len(cve_data))
+    train_data = cve_data[:train_size]
+    test_data = cve_data[train_size:]
+
+    trainer = Trainer(model, preprocessor, Config.DEVICE, Config.LEARNING_RATE)
+    trainer.train(train_data, Config.NUM_EPOCHS, Config.BATCH_SIZE)
+
+    evaluator = Evaluator(model)
+    precision, recall, f1 = evaluator.evaluate(test_data)
+
+    results = {
+        "precision": precision,
+        "recall": recall,
+        "f1_score": f1
+    }
+
+    save_results(Config.OUTPUT_DIR, results)
+
+    print(f"Precision: {precision:.4f}")
+    print(f"Recall: {recall:.4f}")
+    print(f"F1 Score: {f1:.4f}")
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,7 @@
+vllm==0.2.0
+torch==2.0.1
+transformers==4.30.2
+scikit-learn==1.2.2
+tqdm==4.65.0
+numpy==1.24.3
+pandas==2.0.2

src/__init__.py

@@ -0,0 +1,7 @@
+from .config import Config
+from .data_loader import DataLoader
+from .model import VulnerabilityDetectionModel
+from .preprocessor import Preprocessor
+from .trainer import Trainer
+from .evaluator import Evaluator
+from .utils import setup_logging, save_results

src/config.py

@@ -0,0 +1,12 @@
+import os
+
+class Config:
+    DATA_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'data', 'dataset.json')
+    MODEL_NAME = "meta-llama/Llama-3.1-8B-Instruct"
+    MAX_LENGTH = 512
+    BATCH_SIZE = 16
+    LEARNING_RATE = 2e-5
+    NUM_EPOCHS = 5
+    DEVICE = "cuda"
+    OUTPUT_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'output')
+    LOG_FILE = os.path.join(OUTPUT_DIR, 'training.log')

src/data_loader.py

@@ -0,0 +1,23 @@
+import json
+from typing import Dict, List, Tuple
+
+class DataLoader:
+    def __init__(self, file_path: str):
+        self.file_path = file_path
+        self.data = self._load_data()
+
+    def _load_data(self) -> Dict:
+        with open(self.file_path, 'r') as f:
+            return json.load(f)
+
+    def get_cve_data(self) -> List[Tuple[str, str, str, str, List[str]]]:
+        cve_data = []
+        for project, cves in self.data.items():
+            for cve_id, cve_info in cves.items():
+                cwe = cve_info.get('cwe', '')
+                for fixing_commit, files in cve_info.get('fixing_commits', {}).items():
+                    for file_path, lines in files.items():
+                        for line_num, line_info in lines.items():
+                            introducing_commits = line_info.get('Vulnerability Introducing Commit', [])
+                            cve_data.append((project, cve_id, cwe, fixing_commit, introducing_commits))
+        return cve_data

src/evaluator.py

@@ -0,0 +1,22 @@
+from .model import VulnerabilityDetectionModel
+from typing import List, Tuple
+from sklearn.metrics import precision_score, recall_score, f1_score
+
+class Evaluator:
+    def __init__(self, model: VulnerabilityDetectionModel):
+        self.model = model
+
+    def evaluate(self, test_data: List[Tuple[str, str, str, str, List[str]]]) -> Tuple[float, float, float]:
+        prompts = [self.model.format_prompt(project, cve_id, cwe, fixing_commit) for project, cve_id, cwe, fixing_commit, _ in test_data]
+        true_labels = [introducing_commits for _, _, _, _, introducing_commits in test_data]
+
+        predicted_labels = self.model.predict(prompts)
+
+        flattened_true = [commit for commits in true_labels for commit in commits]
+        flattened_pred = [commit for commits in predicted_labels for commit in commits.split()]
+
+        precision = precision_score(flattened_true, flattened_pred, average='micro')
+        recall = recall_score(flattened_true, flattened_pred, average='micro')
+        f1 = f1_score(flattened_true, flattened_pred, average='micro')
+
+        return precision, recall, f1

src/model.py

@@ -0,0 +1,14 @@
+from vllm import LLM, SamplingParams
+from typing import List
+
+class VulnerabilityDetectionModel:
+    def __init__(self, model_name: str, max_length: int):
+        self.model = LLM(model=model_name)
+        self.sampling_params = SamplingParams(temperature=0.7, top_p=0.95, max_tokens=max_length)
+
+    def predict(self, prompts: List[str]) -> List[str]:
+        outputs = self.model.generate(prompts, self.sampling_params)
+        return [output.outputs[0].text for output in outputs]
+
+    def format_prompt(self, project: str, cve_id: str, cwe: str, fixing_commit: str) -> str:
+        return f"Project: {project}\nCVE ID: {cve_id}\nCWE: {cwe}\nFixing Commit: {fixing_commit}\nIdentify the vulnerability introducing commit:"

src/preprocessor.py

@@ -0,0 +1,14 @@
+from typing import List, Tuple
+
+class Preprocessor:
+    def __init__(self, max_length: int):
+        self.max_length = max_length
+
+    def prepare_data(self, cve_data: List[Tuple[str, str, str, str, List[str]]]) -> Tuple[List[str], List[List[str]]]:
+        prompts = []
+        labels = []
+        for project, cve_id, cwe, fixing_commit, introducing_commits in cve_data:
+            prompt = f"Project: {project}\nCVE ID: {cve_id}\nCWE: {cwe}\nFixing Commit: {fixing_commit}\nIdentify the vulnerability introducing commit:"
+            prompts.append(prompt[:self.max_length])
+            labels.append(introducing_commits)
+        return prompts, labels

src/trainer.py

@@ -0,0 +1,43 @@
+from .model import VulnerabilityDetectionModel
+from .preprocessor import Preprocessor
+from typing import List, Tuple
+import torch
+from torch.utils.data import DataLoader, TensorDataset
+from tqdm import tqdm
+
+class Trainer:
+    def __init__(self, model: VulnerabilityDetectionModel, preprocessor: Preprocessor, device: str, learning_rate: float):
+        self.model = model
+        self.preprocessor = preprocessor
+        self.device = device
+        self.learning_rate = learning_rate
+
+    def train(self, cve_data: List[Tuple[str, str, str, str, List[str]]], num_epochs: int, batch_size: int):
+        prompts, labels = self.preprocessor.prepare_data(cve_data)
+        dataset = TensorDataset(torch.tensor(prompts), torch.tensor(labels))
+        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+
+        optimizer = torch.optim.AdamW(self.model.parameters(), lr=self.learning_rate)
+        
+        for epoch in range(num_epochs):
+            self.model.train()
+            total_loss = 0
+            for batch in tqdm(dataloader, desc=f"Epoch {epoch+1}/{num_epochs}"):
+                inputs, targets = batch
+                inputs = inputs.to(self.device)
+                targets = targets.to(self.device)
+
+                outputs = self.model(inputs)
+                loss = self.compute_loss(outputs, targets)
+
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+
+                total_loss += loss.item()
+
+            avg_loss = total_loss / len(dataloader)
+            print(f"Epoch {epoch+1}/{num_epochs}, Average Loss: {avg_loss:.4f}")
+
+    def compute_loss(self, outputs, targets):
+        return torch.nn.functional.cross_entropy(outputs, targets)

src/utils.py

@@ -0,0 +1,19 @@
+import logging
+import os
+import json
+
+def setup_logging(log_file: str):
+    logging.basicConfig(
+        level=logging.INFO,
+        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+        handlers=[
+            logging.FileHandler(log_file),
+            logging.StreamHandler()
+        ]
+    )
+
+def save_results(output_dir: str, results: dict):
+    os.makedirs(output_dir, exist_ok=True)
+    output_file = os.path.join(output_dir, 'results.json')
+    with open(output_file, 'w') as f:
+        json.dump(results, f, indent=2)

tests/__init__.py

@@ -0,0 +1,3 @@
+from .test_data_loader import TestDataLoader
+from .test_model import TestVulnerabilityDetectionModel
+from .test_preprocessor import TestPreprocessor

tests/test_data_loader.py

@@ -0,0 +1,44 @@
+import unittest
+import os
+import json
+from src.data_loader import DataLoader
+
+class TestDataLoader(unittest.TestCase):
+    def setUp(self):
+        self.test_data = {
+            "test_project": {
+                "CVE-2021-1234": {
+                    "cwe": "CWE-79",
+                    "fixing_commits": {
+                        "abcdef1234567890": {
+                            "path/to/file.py": {
+                                "42": {
+                                    "Vulnerability Introducing Commit": [
+                                        "1234567890abcdef"
+                                    ]
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+        }
+        self.test_file = "test_dataset.json"
+        with open(self.test_file, "w") as f:
+            json.dump(self.test_data, f)
+        self.data_loader = DataLoader(self.test_file)
+
+    def tearDown(self):
+        os.remove(self.test_file)
+
+    def test_load_data(self):
+        self.assertEqual(self.data_loader.data, self.test_data)
+
+    def test_get_cve_data(self):
+        expected_output = [
+            ("test_project", "CVE-2021-1234", "CWE-79", "abcdef1234567890", ["1234567890abcdef"])
+        ]
+        self.assertEqual(self.data_loader.get_cve_data(), expected_output)
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_model.py

@@ -0,0 +1,33 @@
+import unittest
+from unittest.mock import Mock, patch
+from src.model import VulnerabilityDetectionModel
+
+class TestVulnerabilityDetectionModel(unittest.TestCase):
+    def setUp(self):
+        self.model = VulnerabilityDetectionModel("test_model", 512)
+
+    @patch('src.model.LLM')
+    def test_predict(self, mock_llm):
+        mock_output = Mock()
+        mock_output.outputs = [Mock(text="Predicted output")]
+        mock_llm.return_value.generate.return_value = [mock_output]
+
+        prompts = ["Test prompt"]
+        result = self.model.predict(prompts)
+
+        self.assertEqual(result, ["Predicted output"])
+        mock_llm.return_value.generate.assert_called_once_with(prompts, self.model.sampling_params)
+
+    def test_format_prompt(self):
+        project = "test_project"
+        cve_id = "CVE-2021-1234"
+        cwe = "CWE-79"
+        fixing_commit = "abcdef1234567890"
+
+        expected_prompt = f"Project: {project}\nCVE ID: {cve_id}\nCWE: {cwe}\nFixing Commit: {fixing_commit}\nIdentify the vulnerability introducing commit:"
+        result = self.model.format_prompt(project, cve_id, cwe, fixing_commit)
+
+        self.assertEqual(result, expected_prompt)
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_preprocessor.py

@@ -0,0 +1,34 @@
+import unittest
+from src.preprocessor import Preprocessor
+
+class TestPreprocessor(unittest.TestCase):
+    def setUp(self):
+        self.preprocessor = Preprocessor(max_length=100)
+
+    def test_prepare_data(self):
+        cve_data = [
+            ("project1", "CVE-2021-1234", "CWE-79", "fix1", ["intro1"]),
+            ("project2", "CVE-2021-5678", "CWE-89", "fix2", ["intro2", "intro3"])
+        ]
+
+        prompts, labels = self.preprocessor.prepare_data(cve_data)
+
+        expected_prompts = [
+            "Project: project1\nCVE ID: CVE-2021-1234\nCWE: CWE-79\nFixing Commit: fix1\nIdentify the vulnerability introducing commit:",
+            "Project: project2\nCVE ID: CVE-2021-5678\nCWE: CWE-89\nFixing Commit: fix2\nIdentify the vulnerability introducing commit:"
+        ]
+        expected_labels = [["intro1"], ["intro2", "intro3"]]
+
+        self.assertEqual(prompts, expected_prompts)
+        self.assertEqual(labels, expected_labels)
+
+    def test_max_length_truncation(self):
+        self.preprocessor.max_length = 50
+        cve_data = [("very_long_project_name", "CVE-2021-1234", "CWE-79", "very_long_fixing_commit_hash", ["intro1"])]
+
+        prompts, _ = self.preprocessor.prepare_data(cve_data)
+
+        self.assertEqual(len(prompts[0]), 50)
+
+if __name__ == '__main__':
+    unittest.main()