RohanSardar/malware_analysis

# 双引擎恶意软件分析 Pipeline [](https://huggingface.co/spaces/RohanSardar/malware_analysis) 一个强大、双引擎的机器学习 pipeline，专为高精度恶意软件检测和高级取证可解释性而设计。本项目采用集成方法，将静态结构分析与深度可视化字节图分类相结合，以检测和分类 Windows 可执行文件。 ## 🏗️ Pipeline 架构 当用户上传文件（`.exe`、`.dll` 或模拟的 `.json`）时，pipeline 会立即拆分数据，并通过两个不同的并行环境进行处理，然后再聚合最终判定结果。 ### 1. 引擎 A：结构化分类器 (LightGBM) *该引擎完全基于可执行文件的元数据和结构逻辑进行操作。* - **特征提取：** 它利用 `pefile` 库来解析 Windows Portable Executable (PE) 头、导入表、导出符号和字节熵特征。 - **向量化：** 提取的数据严格按照 **EMBER 2018 (v2)** 数据集规范，映射为一个 2,381 维的数值向量。 - **推理：** 该向量被输入到一个经过深度优化的 **LightGBM**（Gradient Boosted Decision Trees）模型中，该模型专为严格的二元分类（恶意与良性）而训练。 - **可解释人工智能 (XAI)：** 在传输过程中，该模型会计算 **TreeSHAP**（SHapley Additive exPlanations）值，以识别影响最终百分比的特定结构异常前三名。 ### 2. 引擎 B：可视化分类器 (PyTorch CNN) *该引擎将恶意软件视为原始图像，以寻找空间上的、遗传性的字节模式。* - **二进制转图像：** 文件的原始字节按顺序映射到一个 2D 矩阵中，并用零填充以形成一个完美的正方形。 - **张量化：** 图像被转换为 `128x128` 的灰度 PyTorch Tensor (`1, 128, 128`)。 - **推理：** 该 tensor 通过一个定制的 3 层 **Convolutional Neural Network (CNN)**。该网络能够识别 24 种不同恶意软件家族的可视化纹理（例如 `Swizzor`、`Allaple.A`）。 ### 3. 集成聚合与 UI 渲染 - **风险计算：** pipeline 汇总两个引擎的确切置信度百分比，以生成统一的**总体风险评分**。 - **报告生成：** UI 将结果渲染为带动画的 `gr.Label` 组件，显示提取的灰度字节图，并使用 `fpdf2` 动态编译一份供分析师下载的离线 PDF 取证报告。 ## 📊 性能指标与可视化 这些模型在它们各自的测试集上进行了广泛的评估。 ### 引擎 A（结构化）评估 ``` Accuracy: 0.9512 Precision: 0.9443 Recall: 0.9588 Specificity: 0.9435 F1 Score: 0.9515 ```

### 引擎 B（可视化）评估 ``` precision recall f1-score support Adialer.C 1.00 1.00 1.00 25 Agent.FYI 1.00 1.00 1.00 25 Allaple.A 1.00 1.00 1.00 114 Allaple.L 1.00 1.00 1.00 100 Alueron.gen!J 1.00 1.00 1.00 25 Autorun.K 0.00 0.00 0.00 25 C2LOP.P 0.96 1.00 0.98 25 C2LOP.gen!g 1.00 0.96 0.98 25 Dialplatform.B 1.00 1.00 1.00 25 Dontovo.A 1.00 1.00 1.00 25 Fakerean 1.00 1.00 1.00 75 Instantaccess 1.00 1.00 1.00 75 Lolyda.AA1 1.00 1.00 1.00 60 Lolyda.AA2 1.00 1.00 1.00 25 Lolyda.AA3 1.00 1.00 1.00 25 Lolyda.AT 1.00 1.00 1.00 25 Malex.gen!J 1.00 1.00 1.00 25 Obfuscator.AD 1.00 1.00 1.00 25 Rbot!gen 1.00 1.00 1.00 25 Skintrim.N 1.00 1.00 1.00 25 Swizzor 0.98 0.98 0.98 50 VB.AT 1.00 1.00 1.00 25 Wintrim.BX 1.00 1.00 1.00 25 Yuner.A 0.50 1.00 0.67 25 accuracy 0.97 924 macro avg 0.94 0.96 0.94 924 weighted avg 0.96 0.97 0.96 924 ```

### 硬件与训练规格 PyTorch Convolutional Neural Network（引擎 B）是在本地通过 **NVIDIA RTX 4050** 使用 **CUDA** 训练了 **10 个 epoch**。 ## 🚀 快速开始 ### 前置条件 - Python 3.13+ - 已安装 `uv` ### 安装与运行 克隆代码仓库并使用 `uv` 启动服务器： ``` uv sync uv run main.py ``` ## 📂 数据集与重新训练 如果您希望从头开始重新训练模型，必须下载各自的数据集。由于大小限制，它们不包含在本代码仓库中。 1. **引擎 A 数据集：** [EMBER-2018 v2 特征数据](https://www.kaggle.com/datasets/dhoogla/ember-2018-v2-features) - 将提取的 `.parquet` 文件放入 `ember/` 目录中。 - 运行 `uv run src/engine_a/train.py`。 2. **引擎 B 数据集：** [MalImg 数据集](https://figshare.com/articles/dataset/MalImg_dataset_zip/24189882?file=42443904) - 将图像解压到 `malimg_dataset/` 目录中。 - **关键预处理：** 原始的 MalImg 数据集在训练前需要手动重组，以解决重复和类别混淆问题： 1. 将 `validation` 目录的内容完全合并到 `train` 目录中。 2. 将 `Swizzor.gen!E` 和 `Swizzor.gen!I` 目录合并为一个统一的 `Swizzor` 目录。 3. 将 `Rbotigen` 目录中的所有内容移动到 `Rbot!gen` 目录中。 - 运行 `uv run src/engine_b/train.py`。 ## 👥 贡献者

Rohan Sardar

Anurag Roy

Anwesha Mondal

Sandipan

Sudipta Das

标签：Apex, LightGBM, PE文件分析, PyTorch CNN, Vectored Exception Handling, 云安全监控, 凭据扫描, 可解释AI, 机器学习, 逆向工具, 静态分析