UJeongff/malware-analysis-ml
GitHub: UJeongff/malware-analysis-ml
基于 AI 的静态恶意代码检测系统,利用文件结构特征在不执行文件的情况下识别潜在威胁。
Stars: 1 | Forks: 0
# 🛡️ 基于 AI 的恶意软件预检测系统
这是一个仅利用 Windows PE 文件的 **Import Address Table(IAT)** ——无需执行文件—— 即可判定其是否为恶意文件的 LightGBM 分类器及基于 Flask 的 Web 演示工具。
📄 **[详细报告 (Notion)](https://www.notion.so/36de8de0e4a380a0b9d9e985358df481#36de8de0e4a380fcbb6ff752cead1b66)**
## 数据集
| 项目 | 内容 |
|------|------|
| 规模 | 47,580 个 PE 样本 × 1,000 个 Windows API |
| 编码 | 各 API 是否被调用 (0/1 二进制) |
| 类别比例 | 正常 1,929 : 恶意 45,651 (≈ 1:23 不平衡) |
| 来源 | Angelo Oliveira, *"Malware Analysis Datasets: Top-1000 PE Imports"* ([Kaggle](https://www.kaggle.com/datasets/ang3loliveira/malware-analysis-datasets-top1000-pe-imports)) |
| 标注 | 基于 VirusTotal 多引擎杀毒软件交叉验证 (样本来源: VirusShare 等) |
## 最终模型
50 个特征 (LightGBM gain importance) + SMOTE 1:4 + LightGBM 调优 + 阈值 (threshold) 0.85
| 指标 | 测试集评估 |
|------|-----------|
| Balanced Accuracy | 0.93 |
| 恶意 Recall | 0.95 |
| 正常 Recall | 0.91 |
| 误报率 (FPR) | 8.6% |
→ 产出物: `models/lightgbm_final.pkl` (包含 model + features + threshold 的打包文件)
关于设计决策的依据(特征数 50、SMOTE 1:4、threshold 0.85 的实验推导过程)以及基于 SHAP 的模型解释,请参阅[详细报告](https://www.notion.so/36de8de0e4a380a0b9d9e985358df481?source=copy_link)。
## 演示工具 — `tool/LightGBM/`
使用训练好的模型(`backend/model/lightgbm_final.pkl`)对实际的 `.exe` 文件进行判定的演示。提供两种使用方式。
### 1. Web 演示 — 在浏览器中上传 EXE
```
cd tool/LightGBM/backend
pip install -r requirements.txt
python app.py
# 在浏览器中访问 http://127.0.0.1:5000 → 上传 EXE → 显示恶意/正常 + 置信度
```
### 2. 文件夹批量预测 — `batch_predict.py`
一次性判定多个 `.exe` 文件并保存为 CSV。
```
cd tool/LightGBM/backend
# 1) 将要判定的 exe 文件放置在 sample/ 目录下(紧邻的子文件夹名称将成为 family 列)
# 例如:sample/RedLine/redline.exe , sample/normal/notepad.exe
# 2) 运行
python batch_predict.py
# → 生成 backend/batch_results.csv
# (family, file, prediction, raw_proba, confidence_pct, total_imported_apis, matched_features)
```
## 文件夹结构
| 文件夹 | 作用 |
|------|------|
| [`pipeline/`](pipeline/) | **pkl 完整复现流程** — 预处理 → 训练/测试 → 预测 (自包含) |
| [`experiments/`](experiments/) | **设计依据探索** — 特征数·重采样·模型·阈值实验 + 结果·文档 |
| [`models/`](models/) | 最终模型包 `lightgbm_final.pkl` (model + features + threshold) |
| [`tool/`](tool/) | Flask Web 演示 + 文件夹批量预测(`batch_predict.py`) |
| [`archive/`](archive/) | 旧版代码·报告 (week10~12) |
```
AI-p_05/
├── pipeline/ 1_preprocess → 2_train_test → 3_predict (+ data/, lightgbm_final.pkl)
├── experiments/ scripts/ · results/ · reports/
├── models/ lightgbm_final.pkl
├── tool/ LightGBM/ (웹 데모)
├── archive/ 구버전
└── README.md
```
## 运行
### pkl 复现 (预处理 → 训练 → 测试)
```
cd pipeline
pip install -r requirements.txt
# 0) 从 Kaggle 下载数据并保存为 data/raw_data.csv(参见上方“数据集”)
python 1_preprocess.py # raw_data.csv → data/processed_data.csv
python 2_train_test.py # 학습 + Test 평가 → lightgbm_final.pkl
python 3_predict.py 파일.exe # 단일 EXE 판정 (CLI)
```
### 设计依据实验
```
cd experiments/scripts
python 01_grid_feature_ratio.py # 피처수·SMOTE 비율 격자 탐색
python 03_compare_models.py # 모델 비교 … (자세한 순서는 experiments/README.md)
```
2026 年第一学期 人工智能编程 | 第 05 组
标签:Apex, Flask, LightGBM, Windows PE文件, 云安全监控, 恶意软件检测, 机器学习, 静态分析