Dockerfile

@@ -0,0 +1,23 @@
+FROM python:3.8-slim
+
+RUN apt-get update && apt-get install -y git
+
+WORKDIR /app
+
+COPY . /app
+
+RUN pip install --no-cache-dir -r requirements.txt
+
+RUN chmod +x /app/main.py
+
+RUN curl -s https://packagecloud.io/install/repositories/github/git-lfs/script.deb.sh | bash
+RUN apt-get install git-lfs
+
+RUN git lfs install
+
+RUN mkdir -p /app/models && \
+    git clone https://huggingface.co/Qwen/Qwen2.5-Coder-7B-Instruct /app/models/Qwen2.5-Coder-7B-Instruct
+
+ENV PYTHONUNBUFFERED=1
+
+CMD ["python", "/app/main.py", "--data_path", "/app/data/SIR_test_set.json", "--model_name", "/app/models/Qwen2.5-Coder-7B-Instruct"]

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,99 @@
+# CVSS Score Prediction Project
+
+This project aims to predict CVSS (Common Vulnerability Scoring System) scores using the Qwen2.5-Coder-7B-Instruct model.
+
+## File Structure
+```
+project/
+│
+├── src/ # Source code
+│ ├── init.py
+│ ├── data_loader.py
+│ ├── model.py
+│ ├── preprocess.py
+│ ├── train.py
+│ ├── evaluate.py
+│ └── visualize.py
+│
+├── data/ # Dataset
+│ └── SIR_test_set.json
+│
+├── notebooks/ # Jupyter notebooks
+│ └── exploratory_data_analysis.ipynb
+│
+├── results/ # Output results
+│ ├── figures/ # Generated charts
+│ └── models/ # Saved models
+│
+├── tests/ # Unit tests
+│ └── test_model.py
+│
+├── Dockerfile # Docker configuration
+├── requirements.txt # Project dependencies
+├── README_V2.md # Project documentation
+└── main.py # Main program entry
+```
+## Environment Setup
+
+1. Clone the repository:
+   ```
+   git clone https://gitlink.org.cn/zbtrs2/competition-vd.git
+   cd ompetition-vd
+   ```
+
+2. Create and activate a virtual environment:
+   ```
+   python -m venv venv
+   source venv/bin/activate  # On Windows use: venv\Scripts\activate
+   ```
+
+3. Install dependencies:
+   ```
+   pip install -r requirements.txt
+   ```
+
+4. Download the Qwen2.5-Coder-7B-Instruct model:
+   ```
+   mkdir -p models
+   git lfs install
+   git clone https://huggingface.co/Qwen/Qwen2.5-Coder-7B-Instruct models/Qwen2.5-Coder-7B-Instruct
+   ```
+
+## Running the Code
+
+To run the main program and generate results:
+```python main.py --data_path ./data/SIR_test_set.json --model_name ./models/Qwen2.5-Coder-7B-Instruct```
+This command will:
+1. Load and preprocess the data
+2. Initialize the model
+3. Make predictions
+4. Evaluate the model's performance
+5. Generate visualizations
+
+## Viewing Results
+
+After running the code, you can find the generated visualizations in the `results/figures/` directory. These include:
+
+- predicted_vs_actual_cvss_scores.png
+- error_distribution.png
+- severity_distribution.png
+- cvss_distribution.png
+- cvss_correlation_heatmap.png
+
+The console output will display the Mean Squared Error (MSE) and Root Mean Squared Error (RMSE) of the predictions.
+
+## Docker Usage
+
+To build and run the project using Docker:
+
+1. Build the Docker image:
+   ```
+   docker build -t cvss-predictor .
+   ```
+
+2. Run the Docker container:
+   ```
+   docker run -v $(pwd)/results:/app/results cvss-predictor
+   ```
+
+This will mount the `results` directory to your local machine, allowing you to access the generated visualizations.

data/label_word_ids.json

@@ -0,0 +1,46 @@
+{
+    "AV": {
+        "network": 2897,
+        "adjacent": 5516,
+        "local": 2334,
+        "physical": 3558
+    },
+    "AC": {
+        "low": 2659,
+        "high": 2152
+    },
+    "PR": {
+        "none": 3904,
+        "low": 2659,
+        "high": 2152
+    },
+    "UI": {
+        "none": 3904,
+        "required": 3223
+    },
+    "S": {
+        "unchanged": 15704,
+        "changed": 2904
+    },
+    "C": {
+        "none": 3904,
+        "low": 2659,
+        "high": 2152
+    },
+    "I": {
+        "none": 3904,
+        "low": 2659,
+        "high": 2152
+    },
+    "A": {
+        "none": 3904,
+        "low": 2659,
+        "high": 2152
+    },
+    "severity": {
+        "low": 2659,
+        "medium": 5396,
+        "high": 2152,
+        "critical": 4187
+    }
+}

data/label_word_ids_CVSS2.json

@@ -0,0 +1,37 @@
+{
+    "AV": {
+        "network": 2897,
+        "adjacent": 5516,
+        "local": 2334
+    },
+    "AC": {
+        "low": 2659,
+        "medium": 5396,
+        "high": 2152
+    },
+    "Au": {
+        "none": 3904,
+        "single": 2309,
+        "multiple": 3674
+    },
+    "C": {
+        "none": 3904,
+        "partial": 7704,
+        "complete": 3143
+    },
+    "I": {
+        "none": 3904,
+        "partial": 7704,
+        "complete": 3143
+    },
+    "A": {
+        "none": 3904,
+        "partial": 7704,
+        "complete": 3143
+    },
+    "severity": {
+        "low": 2659,
+        "medium": 5396,
+        "high": 2152
+    }
+}

data_analysis.py

@@ -0,0 +1,56 @@
+import json
+import os
+import matplotlib.pyplot as plt
+import numpy as np
+from collections import Counter
+
+# Define file paths
+data_dir = './data'
+file_names = ['SIR_test_set.json', 'SIR_train_set.json', 'SIR_validation_set.json']
+
+def load_json(file_path):
+    with open(file_path, 'r', encoding='utf-8') as f:
+        return json.load(f)
+
+def analyze_dataset(data):
+    base_scores = [item['baseScore'] for item in data]
+    severities = [item['severity'] for item in data]
+    return base_scores, severities
+
+# Store analysis results for all datasets
+all_base_scores = []
+all_severities = []
+
+# Analyze each dataset
+for file_name in file_names:
+    file_path = os.path.join(data_dir, file_name)
+    data = load_json(file_path)
+    base_scores, severities = analyze_dataset(data)
+    all_base_scores.extend(base_scores)
+    all_severities.extend(severities)
+    
+    print(f"Statistics for {file_name}:")
+    print(f"Sample count: {len(data)}")
+    print(f"Base score mean: {np.mean(base_scores):.2f}")
+    print(f"Base score median: {np.median(base_scores):.2f}")
+    print(f"Most common severity: {Counter(severities).most_common(1)[0][0]}")
+    print()
+
+# Plot base score distribution histogram
+plt.figure(figsize=(10, 6))
+plt.hist(all_base_scores, bins=20, edgecolor='black')
+plt.title('CVSS Base Score Distribution')
+plt.xlabel('Base Score')
+plt.ylabel('Frequency')
+plt.savefig('base_score_distribution.png')
+plt.close()
+
+# Plot severity distribution pie chart
+severity_counts = Counter(all_severities)
+plt.figure(figsize=(10, 6))
+plt.pie(severity_counts.values(), labels=severity_counts.keys(), autopct='%1.1f%%')
+plt.title('Severity Distribution')
+plt.savefig('severity_distribution.png')
+plt.close()
+
+print("Analysis complete. Charts have been saved.")

main.py

@@ -0,0 +1,27 @@
+import argparse
+from src.data_loader import load_data
+from src.preprocess import preprocess_data
+from src.model import QwenCVSSPredictor
+from src.train import train_model
+from src.evaluate import evaluate_model
+from src.visualize import generate_visualizations
+
+def main(args):
+    data = load_data(args.data_path)
+    X_train, X_test, y_train, y_test = preprocess_data(data)
+    model = QwenCVSSPredictor(args.model_name)
+    train_model(model, X_train, y_train)
+    mse, rmse, predictions = evaluate_model(model, X_test, y_test)
+
+    generate_visualizations(y_test, predictions, data)
+    
+    print(f"均方误差 (MSE): {mse}")
+    print(f"均方根误差 (RMSE): {rmse}")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="CVSS Score Prediction")
+    parser.add_argument("--data_path", type=str, default="./data/SIR_test_set.json", help="Path to the dataset")
+    parser.add_argument("--model_name", type=str, default="/data02/hyf/models/Qwen2.5-Coder-7B-Instruct", help="Path to the Qwen model")
+    args = parser.parse_args()
+    
+    main(args)

requirements.txt

@@ -0,0 +1,9 @@
+pandas
+numpy
+matplotlib
+se
+aborn
+scikit-learn
+torch
+transformers
+tqdm

src/data_loader.py

@@ -0,0 +1,8 @@
+import json
+import pandas as pd
+
+def load_data(file_path):
+    """加载JSON数据并转换为DataFrame"""
+    with open(file_path, 'r') as f:
+        data = json.load(f)
+    return pd.DataFrame(data)

src/evaluate.py

@@ -0,0 +1,15 @@
+from sklearn.metrics import mean_squared_error
+import numpy as np
+from tqdm import tqdm
+
+def evaluate_model(model, X_test, y_test):
+    """评估模型性能"""
+    predictions = []
+    for _, row in tqdm(X_test.iterrows(), total=len(X_test)):
+        prediction = model.predict(row['description'], row['severity'])
+        predictions.append(prediction)
+
+    mse = mean_squared_error(y_test, predictions)
+    rmse = np.sqrt(mse)
+
+    return mse, rmse, predictions

src/model.py

@@ -0,0 +1,50 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+class QwenCVSSPredictor:
+    def __init__(self, model_name):
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype="auto",
+            device_map="auto",
+            trust_remote_code=True
+        )
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+    def predict(self, description, severity):
+        prompt = f"""Given the following vulnerability description and severity, predict the CVSS base score (0-10):
+
+Description: {description}
+Severity: {severity}
+
+Provide only the predicted CVSS base score as a number between 0 and 10, with up to two decimal places. Do not include any other text in your response.
+
+Predicted CVSS base score:"""
+
+        messages = [
+            {"role": "system", "content": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."},
+            {"role": "user", "content": prompt}
+        ]
+        text = self.tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+        model_inputs = self.tokenizer([text], return_tensors="pt").to(self.model.device)
+
+        generated_ids = self.model.generate(
+            **model_inputs,
+            max_new_tokens=512
+        )
+        generated_ids = [
+            output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+        ]
+
+        response = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+        
+        try:
+            predicted_score = float(response.strip())
+            return max(0, min(10, predicted_score))  # 确保分数在0到10之间
+        except ValueError:
+            print(f"解析模型输出时出错: {response}")
+            return 5.0  # 如果解析失败,返回默认分数

src/preprocess.py

@@ -0,0 +1,7 @@
+from sklearn.model_selection import train_test_split
+
+def preprocess_data(df):
+    """预处理数据并分割为训练集和测试集"""
+    X = df[['description', 'severity']]
+    y = df['baseScore']
+    return train_test_split(X, y, test_size=0.2, random_state=42)

src/train.py

@@ -0,0 +1,5 @@
+from tqdm import tqdm
+
+def train_model(model, X_train, y_train):
+
+    pass

src/visualize.py

@@ -0,0 +1,57 @@
+import matplotlib.pyplot as plt
+import seaborn as sns
+import pandas as pd
+import numpy as np
+
+def generate_visualizations(y_true, y_pred, data):
+    plot_predicted_vs_actual(y_true, y_pred)
+    plot_error_distribution(y_true, y_pred)
+    plot_severity_distribution(data)
+    plot_cvss_distribution(data)
+    plot_heatmap(data)
+
+def plot_predicted_vs_actual(y_true, y_pred):
+    plt.figure(figsize=(10, 6))
+    plt.scatter(y_true, y_pred, alpha=0.5)
+    plt.plot([0, 10], [0, 10], 'r--')
+    plt.xlabel('Actual CVSS Score')
+    plt.ylabel('Predicted CVSS Score')
+    plt.title('Predicted vs Actual CVSS Scores')
+    plt.savefig('./results/figures/predicted_vs_actual_cvss_scores.png')
+    plt.close()
+
+def plot_error_distribution(y_true, y_pred):
+    errors = y_pred - y_true
+    plt.figure(figsize=(10, 6))
+    sns.histplot(errors, kde=True)
+    plt.xlabel('Prediction Error')
+    plt.ylabel('Frequency')
+    plt.title('Distribution of Prediction Errors')
+    plt.savefig('./results/figures/error_distribution.png')
+    plt.close()
+
+def plot_severity_distribution(data):
+    plt.figure(figsize=(10, 6))
+    sns.countplot(x='severity', data=data)
+    plt.xlabel('Severity')
+    plt.ylabel('Count')
+    plt.title('Distribution of Vulnerability Severity')
+    plt.savefig('./results/figures/severity_distribution.png')
+    plt.close()
+
+def plot_cvss_distribution(data):
+    plt.figure(figsize=(10, 6))
+    sns.histplot(data['baseScore'], kde=True)
+    plt.xlabel('CVSS Base Score')
+    plt.ylabel('Frequency')
+    plt.title('Distribution of CVSS Base Scores')
+    plt.savefig('./results/figures/cvss_distribution.png')
+    plt.close()
+
+def plot_heatmap(data):
+    corr_matrix = data[['baseScore', 'exploitabilityScore', 'impactScore']].corr()
+    plt.figure(figsize=(10, 8))
+    sns.heatmap(corr_matrix, annot=True, cmap='coolwarm')
+    plt.title('Correlation Heatmap of CVSS Scores')
+    plt.savefig('./results/figures/cvss_correlation_heatmap.png')
+    plt.close()

tests/test_model.py

@@ -0,0 +1,14 @@
+import unittest
+from src.model import QwenCVSSPredictor
+class TestQwenCVSSPredictor(unittest.TestCase):
+    def setUp(self):
+        self.model = QwenCVSSPredictor("/path/to/Qwen2.5-Coder-7B-Instruct")
+    def test_prediction_range(self):
+        description = "A critical vulnerability in the system"
+        severity = "High"
+        prediction = self.model.predict(description, severity)
+        self.assertGreaterEqual(prediction, 0)
+        self.assertLessEqual(prediction, 10)
+        
+if __name__ == '__main__':
+    unittest.main()