# Bert-Trojan

**Repository Path**: CorigitPiree/bert-trojan

## Basic Information

- **Project Name**: Bert-Trojan
- **Description**: 基于bert模型实现对木马病毒的分类
- **Primary Language**: Unknown
- **License**: GPL-2.0
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 1
- **Forks**: 0
- **Created**: 2025-08-30
- **Last Updated**: 2025-09-06

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# 木马病毒分类系统 - BERT知识图谱集成方案

## 项目概述

本项目基于深度学习与知识图谱技术，实现了对木马病毒汇编代码的智能分类系统。核心思想是将汇编代码转换为结构化知识图谱，再利用BERT模型提取语义特征，最终训练高性能分类器准确识别各类木马病毒。本项目支持GPU加速，提供完整的训练、评估和预测流程。

## 技术流程

1. **知识图谱转换**：将原始汇编代码(.asm)转换为函数调用关系图
2. **BERT特征提取**：使用预训练BERT模型提取图谱语义特征
3. **分类模型训练**：基于提取特征训练SVM/LogisticRegression/RF分类器
4. **模型评估与预测**：支持批量测试和单样本预测
5. **类别不平衡处理**：动态选择SMOTE或随机过采样策略

关键技术：
- 知识图谱表示学习
- BERT语义特征提取
- 梯度提升分类器
- 动态类别平衡采样
- 长文本分块处理

## 文件结构说明

```
├─customs                    # 模型缓存与检查点
│  ├─bert_features_cache.joblib            # 特征缓存
│  └─trojan_classifier_model.joblib        # 分类模型
├─models                     # 预训练模型
│  ├─bert-base-uncased       # BERT基础模型
│  └─codebert-base           # CodeBERT模型
├─results                    # 测试结果
└─samples                    # 样本数据
│  └─subtrain
│       ├─asm                # 原始汇编数据集
│       │  ├─test            # 测试集
│       │  └─train           # 训练集
│       └─clean              # 清洗汇编数据集
│           ├─test
│           └─train
├─asm_feature_extractor.py   # 特征提取接口
├─asm_to_kg.py               # 汇编转知识图谱接口
├─kg_to_bert.py              # 知识图谱到BERT转换接口
├─main.py                    # 训练集训练
├─model_trainor.py           # 分类模型训练接口
├─sample_predictor.py        # 样本预测接口
├─single_test.py             # 单样本预测器
├─test.py                    # 测试集评估
├─z_clean.py                 # 数据集清洗工具（已预清洗）
└─z_split.py                 # 数据集划分工具（已预分）
```

## 核心源码说明

1. **`asm_to_kg.py`**
   - 功能：将汇编代码转换为知识图谱
   - 关键函数：
     - `process_asm_to_graph_text()`：处理汇编文件生成图谱文本
     - `extract_functions_and_calls()`：提取函数与调用关系
     - `build_knowledge_graph()`：构建NetworkX图谱

2. **`kg_to_bert.py`**
   - 功能：将知识图谱转换为BERT特征向量
   - GPU支持：自动检测并使用可用GPU
   - 关键函数：
     - `process_graph_text_with_bert()`：处理图谱生成BERT特征
     - 自动分块处理长文本（>512 token）

3. **`asm_feature_extractor.py`**
   - 功能：批量特征提取与缓存
   - 关键特点：
     - 自动跳过失败样本并记录
     - 特征缓存机制（可复用）
     - 支持原始/清洗两种数据集

4. **`model_trainor.py`**
   - 功能：模型训练与评估
   - 关键特点：
     - 支持SVM/LogisticRegression/RF三种分类器
     - 动态采样策略（SMOTE/RandomOverSampler）
     - 分层交叉验证与网格搜索调参
     - 自动保存最佳模型

5. **`sample_predictor.py`**
   - 功能：单样本预测
   - 输出：预测类别 + 各类别概率分布
   - 支持模型预测概率或决策函数转换

6. **`main.py`**
   - 主入口：配置路径，启动训练和预测
   - 示例预测单样本流程

7. **`test.py`**
   - 功能：批量预测与评估
   - 输出：CSV结果文件含完整概率分布
   - 指标：准确率/F1值/混淆矩阵

8. **`single_test.py`**
   - 功能：预测单个样本（方便使用）

9. **`z_split.py`**
   - 功能：将某个目录下的数据源文件按照`8:2`划分为训练集`train`和测试集`test`（方便使用）

10. **`z_clean.py`**
   - 功能：将某个目录下的数据源文件清洗后输出到另一目录中

## 使用说明

### 环境配置

```bash
# 创建conda环境
conda create -n bert python=3.11
conda activate bert

# 安装PyTorch（CUDA 12.4）
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124

# 安装其他依赖
pip install -r requirements.txt
```

> 注意：请查看自己电脑支持的CUDA版本，并相应调整PyTorch安装命令

### Bert准备
下载bert模型放入`models/`文件夹

### 数据集准备

1. 运行z_split.py，划分数据集为训练集和测试集并分别放入数据集目录中的`train`和`test`目录（已完成划分）
2. 运行z_clean.py，清洗数据集（可选）
3. 确保标签文件`samples/subtrainLabels.csv`存在

### 运行流程

1. **训练模型**
```bash
python main.py
```
- 自动提取特征
- 训练并评估三种分类器
- 保存最佳模型到`customs/`目录
> 注意：如需训练，首先删除`customs/`目录中的`.joblib`文件（训练过的模型），否则每次运行文件会默认预先加载

2. **测试评估**
```bash
python test.py
```
- 加载训练好的模型
- 在测试集上评估性能
- 生成预测结果到`results/`目录

3. **单样本预测**
```bash
# 在single_test.py中配置样本路径后运行
python single_test.py
```

4. **重新训练**
- 如需重新训练，删除`customs/`目录中的`.joblib`文件，否则每次运行文件会默认预先加载

### 测试结果
在**原始数据集**上训练并测试
```
==================================================
预测完成! 总样本数: 159, 正确预测: 148
整体准确率: 0.9308 (93.08%)
结果已保存至: results/prediction_results.csv
==================================================
```
在**清洗数据集**上训练并测试
```
==================================================
预测完成! 总样本数: 158, 正确预测: 153
整体准确率: 0.9684 (96.84%)
结果已保存至: results/prediction_results_clean.csv
==================================================
```
**单样本**预测案例
```
==================================================
文件: samples/subtrain/asm/test/0DqUX5rkg3IbMY6BLGCE.asm
文件ID: 0DqUX5rkg3IbMY6BLGCE
真实类别: 1
预测类别: 1
预测结果: 正确

各类别概率分布:
  类别 1: 93.04%
  类别 2: 0.29%
  类别 3: 0.15%
  类别 4: 0.55%
  类别 5: 1.49%
  类别 6: 0.12%
  类别 7: 0.37%
  类别 8: 3.13%
  类别 9: 0.86%
==================================================
```

### 关键配置
训练集训练，`main.py`中修改
```python
# 基本路径配置
MODEL_PATH = "models/codebert-base"  # bert模型目录
ASM_DIR = "samples/subtrain/asm/train"  # 训练数据目录
LABEL_FILE = "samples/subtrainLabels.csv"  # 标签文件

# 缓存与模型输出
CACHE_FILE = "customs/codebert_clean_cache.joblib"
CUSTOM_FILE = "customs/codebert_clean_classifier.joblib"

# 可选：SVM/LinearSVM/LogisticRegression/XGBoost/MLP
SELECTED_ALGORITHM = ""  # 选择的分类算法（空字符串表示自动选择）
```

测试集测试，`test.py`中修改
```python
# 基本路径配置
MODEL_PATH = "models/codebert-base"
ASM_DIR = "samples/subtrain/asm/test"
LABEL_FILE = "samples/subtrainLabels.csv"

# 训练完成的模型文件
CUSTOM_FILE = "customs/codebert_clean_classifier.joblib"

# 预测结果输出文件
PREDICTION_RESULTS = "results/prediction_results.csv"
```

## 性能优化建议

1. **GPU加速**：确保安装正确CUDA版本的PyTorch
2. **特征缓存**：首次提取特征较慢，后续使用缓存可加速训练
3. **批处理大小**：在`kg_to_bert.py`中调整`max_seq_length`
4. **类别平衡**：对小样本类别使用`model_trainor.py`中的`SafeSampler`

> 注：项目默认使用原始数据集处理流程，如需使用清洗数据集，请修改相关路径配置