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,75 @@
+# 软件漏洞引入提交预测系统
+
+本项目旨在开发一种高效的数据挖掘工具，用于预测软件项目中的漏洞引入提交。我们采用了多种机器学习模型，包括随机森林、深度学习神经网络、逻辑回归和 XGBoost，以及一个集成这些模型的 ensemble 模型。
+
+## 文件结构
+
+```
+.
+├── data
+│   └── dataset.json
+├── src
+│   ├── __init__.py
+│   ├── data_preprocessing.py
+│   ├── feature_engineering.py
+│   ├── ensemble.py
+│   ├── utils.py
+│   ├── gen_image.py
+│   └── models
+│       ├── __init__.py
+│       ├── random_forest.py
+│       ├── deep_learning.py
+│       ├── linear_regression.py
+│       └── xgboost_model.py
+├── main.py
+├── Dockerfile
+├── requirements.txt
+└── README.md
+```
+
+## 环境配置
+
+本项目使用 Docker 来确保环境的一致性和可复现性。请按照以下步骤配置运行环境：
+
+1. **安装 Docker**：如果您还没有安装 Docker，请访问 [Docker官网](https://www.docker.com/) 下载并安装适合您操作系统的版本。
+
+2. **构建 Docker 镜像**：
+
+   ```
+   docker build -t vulnerability-prediction .
+   ```
+
+   这个过程可能需要一些时间，因为它会安装所有必要的依赖。
+
+## 运行代码
+
+完成环境配置后，按照以下步骤运行代码：
+
+1. **启动 Docker 容器**：
+
+   ```
+   docker run -it --gpus all -v $(pwd):/app vulnerability-prediction /bin/bash
+   ```
+
+   这个命令会启动一个交互式的 bash 会话，并将当前目录挂载到容器的 `/app` 目录。
+
+2. **在容器内运行主程序**：
+
+   ```
+   python main.py
+   ```
+
+   这个命令会执行整个数据处理、模型训练和评估的流程。
+
+3. **查看结果**：
+
+   程序运行完成后，结果会保存在 `results` 目录下。您可以在容器内使用以下命令查看结果：
+
+   ```
+   ls -l results/
+   ```
+
+4. **退出容器**：
+
+   完成后，可以输入 `exit` 命令退出容器。
+

dockerfile

@@ -0,0 +1,51 @@
+FROM nvcr.io/nvidia/tensorrt:21.12-py3
+
+ENV DEBIAN_FRONTEND=noninteractive
+ENV PATH="/opt/conda/bin:${PATH}"
+ARG PATH="/opt/conda/bin:${PATH}"
+
+RUN apt-get update && apt-get install -y \
+    wget \
+    git \
+    build-essential \
+    libpq-dev \
+    postgresql \
+    redis-server \
+    libsm6 \
+    libxext6 \
+    libxrender-dev \
+    libglib2.0-0 \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN wget https://github.com/conda-forge/miniforge/releases/latest/download/Mambaforge-Linux-x86_64.sh \
+    && bash Mambaforge-Linux-x86_64.sh -b -p /opt/conda \
+    && rm -f Mambaforge-Linux-x86_64.sh
+
+RUN mamba create -n ml_env python=3.9 -y
+ENV PATH /opt/conda/envs/ml_env/bin:$PATH
+SHELL ["conda", "run", "-n", "ml_env", "/bin/bash", "-c"]
+
+COPY requirements.txt /tmp/
+RUN pip install --no-cache-dir -r /tmp/requirements.txt
+
+RUN mamba install -c pytorch -c nvidia pytorch torchvision torchaudio cudatoolkit=11.3 -y
+RUN mamba install -c conda-forge opencv -y
+
+RUN git clone https://github.com/NVIDIA/apex
+RUN cd apex && pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
+
+WORKDIR /app
+COPY . /app
+
+RUN pip install --no-cache-dir dlib
+
+RUN adduser --disabled-password --gecos '' myuser
+USER myuser
+
+RUN mkdir -p /home/myuser/.config/matplotlib
+RUN echo "backend : Agg" > /home/myuser/.config/matplotlib/matplotlibrc
+
+ENV PYTHONPATH "${PYTHONPATH}:/app"
+ENV LD_LIBRARY_PATH "${LD_LIBRARY_PATH}:/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64"
+
+CMD ["gunicorn", "--bind", "0.0.0.0:8000", "--workers", "4", "--threads", "4", "--timeout", "120", "app:server"]

main.py

@@ -0,0 +1,56 @@
+import os
+import pandas as pd
+import numpy as np
+from sklearn.model_selection import train_test_split
+from sklearn.preprocessing import StandardScaler
+
+from src.data_preprocessing import preprocess_data
+from src.feature_engineering import engineer_features
+from src.models.random_forest import RandomForestModel
+from src.models.deep_learning import DeepLearningModel
+from src.models.linear_regression import LinearRegressionModel
+from src.models.xgboost_model import XGBoostModel
+from src.ensemble import EnsembleModel
+from src.utils import load_data, save_results
+from src.gen_image import generate_all_plots
+
+def main():
+    data = load_data('data/dataset.json')
+    
+    df = preprocess_data(data)
+    
+    df = engineer_features(df)
+    
+    X = df.drop(['vuln_introducing_commit'], axis=1)
+    y = (df['vuln_introducing_commit'] != '').astype(int)
+    
+    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
+    
+    scaler = StandardScaler()
+    X_train_scaled = scaler.fit_transform(X_train)
+    X_test_scaled = scaler.transform(X_test)
+    
+    rf_model = RandomForestModel()
+    dl_model = DeepLearningModel(input_dim=X_train_scaled.shape[1])
+    lr_model = LinearRegressionModel()
+    xgb_model = XGBoostModel()
+    
+    models = [rf_model, dl_model, lr_model, xgb_model]
+    for model in models:
+        model.train(X_train_scaled, y_train)
+    
+    ensemble_model = EnsembleModel(models)
+    ensemble_model.train(X_train_scaled, y_train)
+    
+    results = {}
+    for model in models + [ensemble_model]:
+        model_name = model.__class__.__name__
+        results[model_name] = model.evaluate(X_test_scaled, y_test)
+    
+    save_results(results, 'results/model_performance.pkl')
+    
+    generate_all_plots(results)
+    
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,24 @@
+numpy==1.21.5
+pandas==1.3.5
+scikit-learn==0.24.2
+tensorflow==2.7.0
+xgboost==1.5.0
+matplotlib==3.5.1
+seaborn==0.11.2
+torch==1.10.1
+transformers==4.15.0
+fastai==2.5.3
+catboost==1.0.3
+lightgbm==3.3.2
+plotly==5.5.0
+dash==2.0.0
+flask==2.0.2
+gunicorn==20.1.0
+psycopg2-binary==2.9.3
+sqlalchemy==1.4.29
+redis==4.1.0
+celery==5.2.3
+pytest==6.2.5
+black==21.12b0
+flake8==4.0.1
+mypy==0.930

src/__init__.py

@@ -0,0 +1,4 @@
+from .data_preprocessing import preprocess_data
+from .feature_engineering import engineer_features
+from .ensemble import EnsembleModel
+from .utils import load_data, save_results

src/data_preprocessing.py

@@ -0,0 +1,22 @@
+import pandas as pd
+import numpy as np
+from sklearn.preprocessing import LabelEncoder
+
+def preprocess_data(data):
+    df = pd.DataFrame(data)
+    df = df.explode('fixing_commits')
+    df = df.reset_index(drop=True)
+    
+    df['cwe'] = df['cwe'].fillna('Unknown')
+    le = LabelEncoder()
+    df['cwe_encoded'] = le.fit_transform(df['cwe'])
+    
+    df['commit_hash'] = df['fixing_commits'].apply(lambda x: list(x.keys())[0] if isinstance(x, dict) else x)
+    df['file_path'] = df['fixing_commits'].apply(lambda x: list(x[list(x.keys())[0]].keys())[0] if isinstance(x, dict) else '')
+    df['line_number'] = df['fixing_commits'].apply(lambda x: list(x[list(x.keys())[0]][list(x[list(x.keys())[0]].keys())[0]].keys())[0] if isinstance(x, dict) else '')
+    
+    df['vuln_introducing_commit'] = df['fixing_commits'].apply(lambda x: x[list(x.keys())[0]][list(x[list(x.keys())[0]].keys())[0]][list(x[list(x.keys())[0]][list(x[list(x.keys())[0]].keys())[0]].keys())[0]]['Vulnerability Introducing Commit'][0] if isinstance(x, dict) else '')
+    
+    df = df.drop(['fixing_commits'], axis=1)
+    
+    return df

src/ensemble.py

@@ -0,0 +1,32 @@
+import numpy as np
+from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
+
+class EnsembleModel:
+    def __init__(self, models):
+        self.models = models
+    
+    def train(self, X, y):
+        for model in self.models:
+            model.train(X, y)
+    
+    def predict(self, X):
+        predictions = np.array([model.predict(X) for model in self.models])
+        return np.round(np.mean(predictions, axis=0)).astype(int)
+    
+    def predict_proba(self, X):
+        probas = np.array([model.predict_proba(X) for model in self.models])
+        return np.mean(probas, axis=0)
+    
+    def evaluate(self, X, y):
+        y_pred = self.predict(X)
+        accuracy = accuracy_score(y, y_pred)
+        precision = precision_score(y, y_pred)
+        recall = recall_score(y, y_pred)
+        f1 = f1_score(y, y_pred)
+        
+        return {
+            'accuracy': accuracy,
+            'precision': precision,
+            'recall': recall,
+            'f1': f1
+        }

src/feature_engineering.py

@@ -0,0 +1,20 @@
+import numpy as np
+from sklearn.feature_extraction.text import TfidfVectorizer
+
+def engineer_features(df):
+    df['commit_hash_length'] = df['commit_hash'].str.len()
+    df['file_path_depth'] = df['file_path'].str.count('/')
+    df['file_extension'] = df['file_path'].str.split('.').str[-1]
+    
+    tfidf = TfidfVectorizer(max_features=100)
+    file_path_tfidf = tfidf.fit_transform(df['file_path'])
+    file_path_tfidf_df = pd.DataFrame(file_path_tfidf.toarray(), columns=[f'file_path_tfidf_{i}' for i in range(100)])
+    
+    df = pd.concat([df, file_path_tfidf_df], axis=1)
+    
+    df['line_number'] = pd.to_numeric(df['line_number'], errors='coerce')
+    df['line_number_log'] = np.log1p(df['line_number'])
+    
+    df['vuln_introducing_commit_length'] = df['vuln_introducing_commit'].str.len()
+    
+    return df

src/gen_image.py

@@ -0,0 +1,85 @@
+import matplotlib.pyplot as plt
+import seaborn as sns
+import numpy as np
+import pandas as pd
+
+plt.rcParams['font.family'] = 'sans-serif'
+plt.rcParams['font.sans-serif'] = ['Arial']
+
+def plot_model_performance(results):
+    models = list(results.keys())
+    metrics = ['accuracy', 'precision', 'recall', 'f1']
+    
+    plt.figure(figsize=(12, 6))
+    x = np.arange(len(models))
+    width = 0.2
+    
+    for i, metric in enumerate(metrics):
+        values = [results[model][metric] for model in models]
+        plt.bar(x + i*width - 1.5*width, values, width, label=metric.capitalize())
+    
+    plt.xlabel('Models')
+    plt.ylabel('Scores')
+    plt.title('Model Performance Comparison')
+    plt.xticks(x, models, rotation=45)
+    plt.legend()
+    plt.tight_layout()
+    plt.savefig('results/model_performance.png')
+    plt.close()
+
+def plot_feature_importance(model, feature_names, top_n=10):
+    importances = model.get_feature_importance()
+    indices = np.argsort(importances)[::-1][:top_n]
+    
+    plt.figure(figsize=(10, 6))
+    plt.title(f"Top {top_n} Feature Importances")
+    plt.bar(range(top_n), importances[indices])
+    plt.xticks(range(top_n), [feature_names[i] for i in indices], rotation=90)
+    plt.tight_layout()
+    plt.savefig('results/feature_importance.png')
+    plt.close()
+
+def plot_confusion_matrix(cm, class_names):
+    plt.figure(figsize=(8, 6))
+    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=class_names, yticklabels=class_names)
+    plt.title('Confusion Matrix')
+    plt.ylabel('True Label')
+    plt.xlabel('Predicted Label')
+    plt.tight_layout()
+    plt.savefig('results/confusion_matrix.png')
+    plt.close()
+
+def plot_roc_curve(fpr, tpr, roc_auc):
+    plt.figure(figsize=(8, 6))
+    plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC curve (AUC = {roc_auc:.2f})')
+    plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
+    plt.xlim([0.0, 1.0])
+    plt.ylim([0.0, 1.05])
+    plt.xlabel('False Positive Rate')
+    plt.ylabel('True Positive Rate')
+    plt.title('Receiver Operating Characteristic (ROC) Curve')
+    plt.legend(loc="lower right")
+    plt.tight_layout()
+    plt.savefig('results/roc_curve.png')
+    plt.close()
+
+def plot_learning_curve(train_sizes, train_scores, val_scores):
+    plt.figure(figsize=(10, 6))
+    plt.plot(train_sizes, np.mean(train_scores, axis=1), 'o-', color='r', label='Training score')
+    plt.plot(train_sizes, np.mean(val_scores, axis=1), 'o-', color='g', label='Cross-validation score')
+    plt.xlabel('Training examples')
+    plt.ylabel('Score')
+    plt.title('Learning Curve')
+    plt.legend(loc='best')
+    plt.grid(True)
+    plt.tight_layout()
+    plt.savefig('results/learning_curve.png')
+    plt.close()
+
+def generate_all_plots(results, model, feature_names, cm, fpr, tpr, roc_auc, train_sizes, train_scores, val_scores):
+    plot_model_performance(results)
+    plot_feature_importance(model, feature_names)
+    plot_confusion_matrix(cm, ['Non-vulnerable', 'Vulnerable'])
+    plot_roc_curve(fpr, tpr, roc_auc)
+    plot_learning_curve(train_sizes, train_scores, val_scores)
+    print("所有图表已生成并保存。")

src/models/__init__.py

@@ -0,0 +1,4 @@
+from .random_forest import RandomForestModel
+from .deep_learning import DeepLearningModel
+from .linear_regression import LinearRegressionModel
+from .xgboost_model import XGBoostModel

src/models/deep_learning.py

@@ -0,0 +1,34 @@
+import tensorflow as tf
+from tensorflow.keras.models import Sequential
+from tensorflow.keras.layers import Dense, Dropout
+from tensorflow.keras.optimizers import Adam
+from tensorflow.keras.callbacks import EarlyStopping
+import numpy as np
+
+class DeepLearningModel:
+    def __init__(self, input_dim):
+        self.model = Sequential([
+            Dense(64, activation='relu', input_dim=input_dim),
+            Dropout(0.3),
+            Dense(32, activation='relu'),
+            Dropout(0.3),
+            Dense(16, activation='relu'),
+            Dense(1, activation='sigmoid')
+        ])
+        self.model.compile(optimizer=Adam(learning_rate=0.001),
+                           loss='binary_crossentropy',
+                           metrics=['accuracy'])
+
+    def train(self, X, y, epochs=100, batch_size=32):
+        early_stopping = EarlyStopping(monitor='val_loss', patience=10, restore_best_weights=True)
+        self.model.fit(X, y, epochs=epochs, batch_size=batch_size, validation_split=0.2, callbacks=[early_stopping])
+
+    def predict(self, X):
+        return (self.model.predict(X) > 0.5).astype(int).flatten()
+
+    def predict_proba(self, X):
+        return self.model.predict(X).flatten()
+
+    def evaluate(self, X, y):
+        _, accuracy = self.model.evaluate(X, y)
+        return accuracy

src/models/linear_regression.py

@@ -0,0 +1,31 @@
+from sklearn.linear_model import LogisticRegression
+from sklearn.model_selection import GridSearchCV
+import numpy as np
+
+class LinearRegressionModel:
+    def __init__(self):
+        self.model = None
+        self.param_grid = {
+            'C': [0.001, 0.01, 0.1, 1, 10, 100],
+            'penalty': ['l1', 'l2'],
+            'solver': ['liblinear', 'saga']
+        }
+
+    def train(self, X, y):
+        base_model = LogisticRegression(random_state=42, max_iter=1000)
+        self.model = GridSearchCV(base_model, self.param_grid, cv=5, n_jobs=-1)
+        self.model.fit(X, y)
+
+    def predict(self, X):
+        return self.model.predict(X)
+
+    def predict_proba(self, X):
+        return self.model.predict_proba(X)
+
+    def get_feature_importance(self):
+        return np.abs(self.model.best_estimator_.coef_[0])
+
+    def evaluate(self, X, y):
+        y_pred = self.predict(X)
+        accuracy = np.mean(y_pred == y)
+        return accuracy

src/models/random_forest.py

@@ -0,0 +1,32 @@
+from sklearn.ensemble import RandomForestClassifier
+from sklearn.model_selection import GridSearchCV
+import numpy as np
+
+class RandomForestModel:
+    def __init__(self):
+        self.model = None
+        self.param_grid = {
+            'n_estimators': [100, 200, 300],
+            'max_depth': [None, 10, 20, 30],
+            'min_samples_split': [2, 5, 10],
+            'min_samples_leaf': [1, 2, 4]
+        }
+
+    def train(self, X, y):
+        base_model = RandomForestClassifier(random_state=42)
+        self.model = GridSearchCV(base_model, self.param_grid, cv=5, n_jobs=-1)
+        self.model.fit(X, y)
+
+    def predict(self, X):
+        return self.model.predict(X)
+
+    def predict_proba(self, X):
+        return self.model.predict_proba(X)
+
+    def get_feature_importance(self):
+        return self.model.best_estimator_.feature_importances_
+
+    def evaluate(self, X, y):
+        y_pred = self.predict(X)
+        accuracy = np.mean(y_pred == y)
+        return accuracy

src/models/xgboost_model.py

@@ -0,0 +1,33 @@
+import xgboost as xgb
+from sklearn.model_selection import GridSearchCV
+import numpy as np
+
+class XGBoostModel:
+    def __init__(self):
+        self.model = None
+        self.param_grid = {
+            'max_depth': [3, 5, 7],
+            'learning_rate': [0.01, 0.1, 0.3],
+            'n_estimators': [100, 200, 300],
+            'subsample': [0.8, 0.9, 1.0],
+            'colsample_bytree': [0.8, 0.9, 1.0]
+        }
+
+    def train(self, X, y):
+        base_model = xgb.XGBClassifier(random_state=42)
+        self.model = GridSearchCV(base_model, self.param_grid, cv=5, n_jobs=-1)
+        self.model.fit(X, y)
+
+    def predict(self, X):
+        return self.model.predict(X)
+
+    def predict_proba(self, X):
+        return self.model.predict_proba(X)
+
+    def get_feature_importance(self):
+        return self.model.best_estimator_.feature_importances_
+
+    def evaluate(self, X, y):
+        y_pred = self.predict(X)
+        accuracy = np.mean(y_pred == y)
+        return accuracy

src/utils.py

@@ -0,0 +1,17 @@
+import json
+import pandas as pd
+import pickle
+
+def load_data(file_path):
+    with open(file_path, 'r') as f:
+        data = json.load(f)
+    return data
+
+def save_results(results, file_path):
+    with open(file_path, 'wb') as f:
+        pickle.dump(results, f)
+
+def load_results(file_path):
+    with open(file_path, 'rb') as f:
+        results = pickle.load(f)
+    return results