0la-mide/malware-cnn
GitHub: 0la-mide/malware-cnn
基于卷积神经网络将恶意软件可视化为灰度图像进行家族分类,并结合VirusTotal交叉验证与YARA规则生成的AI恶意软件检测平台。
Stars: 0 | Forks: 0
# Malware CNN 扫描器
这是一个恶意软件分析平台,它使用自定义的卷积神经网络 (CNN) 将可执行文件分类到不同的恶意软件家族中,通过 VirusTotal 对结果进行交叉验证,并自动生成可部署的 YARA 检测规则。
本项目基于 PyTorch 和 FastAPI 构建。在 Apple M5 MacBook Air 上利用 Metal GPU (MPS) 加速,基于 Malimg 数据集进行训练仅需约 4 分钟。
工作原理
传统的防病毒软件依赖于已知的特征码,在面对新型或变种的恶意软件时往往会失效。本系统采用了一种不同的方法。
恶意软件可执行文件通过将原始字节值映射为像素强度,被转换为灰度图像。不同的恶意软件家族会产生视觉上截然不同的模式——在字节级别上,蠕虫与密码窃取器看起来是不同的,而这些差异在视觉上显现了出来。一个 4 层 CNN 能够学习识别这些模式,并将文件分类到 25 个已知的恶意软件家族之一。
随后,结果会与 VirusTotal 进行交叉比对,利用 75 个真实的 AV 引擎来验证 CNN 的输出。基于 VT 的响应,系统会给出最终判定 —— CLEAN、UNKNOWN、SUSPICIOUS 或确认的威胁级别 —— 并从提取的字节特征中生成 YARA 规则,这些规则可以直接部署到 SIEM 或 EDR 工具中。
suspicious.exe → 灰度图像 → 4 层 CNN → 家族分类
↓
VirusTotal 交叉验证
↓
最终判定 (CLEAN / SUSPICIOUS / HIGH / CRITICAL)
↓
YARA 规则生成
判定逻辑
CNN 总是会输出一个分类结果 —— 但高置信度分数并不自动意味着该文件是恶意的。VirusTotal 才是事实依据:
VT 结果 | 最终判定 | YARA
发现,0 个检测 | CLEAN | 跳过
VT 中未发现 | UNKNOWN | 生成并带有警告
1–5 个引擎标记 | SUSPICIOUS | 生成并带有警告
6+ 个引擎标记 | HIGH / CRITICAL (CNN 家族) | 生成
这可以防止系统对带有看似惊人但实际上并未得到真实检测支持的高置信度分数的干净文件呈现误报。
模型架构
一个用于渐进式特征提取的自定义 4 层 CNN:
层 | 学习内容
Block 1 | 低级模式 —— 边缘、纹理、二进制密度、结构边界
Block 2 | 中间结构 —— 重复的字节序列、opcode 分布、加密指标
Block 3 | 行为特征 —— 混淆痕迹、家族级别的结构相似性、payload 模式
Block 4 | 高级抽象威胁表示 —— 用于分类的压缩内部特征签名
层结构 (每个 Block):Conv2d → BatchNorm2d → ReLU → MaxPool2d
分类头:AdaptiveAvgPool2d → Flatten → Linear(4096, 512) → Dropout(0.5) → Linear(512, 25)
BatchNorm 防止训练过程中出现内部协变量偏移
Dropout(0.5) 防止模型死记硬背训练样本 —— 这对于泛化到未知文件至关重要
结果
基于 Malimg 数据集(包含 25 个恶意软件家族的 9,339 个样本)训练,80/20 训练/验证集划分,20 个 epoch。
指标 | 分数
验证集准确率 | 98.55%
加权 F1 分数 | 0.980
宏平均 F1 分数 | 0.947
训练时间 | ~4 分钟 (Apple M5, MPS backend)
Epoch 1 准确率 | 86.55% 训练集 / 97.00% 验证集
Epoch 20 准确率 | 98.09% 训练集 / 98.29% 验证集
各家族详情
25 个家族中有 20 个实现了完美的精确率和召回率分类 (1.000)
Swizzor.gen!E / Swizzor.gen!I —— 由于极其相似的视觉结构(同一恶意软件家族的两个变种),变种之间存在轻微混淆
Autorun.K —— 验证集上的召回率为零;样本量小(20 个样本)且与 Yuner.A 有很高的视觉相似度,导致持续的错误分类
Yuner.A —— 精确率略低 (0.887),因为它吸收了 Autorun.K 的错误分类
混淆矩阵
对角线几乎完全是深色的 —— 表示预测正确。除了 Swizzor 变种和 Autorun.K/Yuner.A 这一对之外,所有家族的非对角线值都接近于零。
API Endpoints
方法 | Endpoint | 描述
GET | / | API 信息和 endpoint 列表
GET | /health | 模型状态、设备信息、VT 配置检查
GET | /families | 所有 25 个可检测的家族及其威胁级别
POST | /scan | 仅使用 CNN 进行分类
POST | /scan/full | CNN + VirusTotal 查询
POST | /yara/generate | CNN + YARA 规则生成
POST | /scan/complete | 完整 pipeline —— CNN + VT 判定逻辑 + YARA
GET | /virustotal/{sha256} | 通过 SHA256 哈希进行 VirusTotal 查询
示例 —— 完整 Pipeline
curl -X POST https://api.olamide.anvora.com.ng/scan/complete -F "file=@suspicious.exe"
{
"file": {
"name": "suspicious.exe",
"size_bytes": 2567,
"sha256": "01bc427e...",
"entropy": 5.1973,
"possibly_packed": false
},
"cnn_verdict": {
"family": "VB.AT",
"threat_level": "HIGH",
"confidence": 0.99965,
"note": null
},
"top3": [
{ "family": "VB.AT", "confidence": 0.99965, "threat_level": "HIGH" },
{ "family": "Lolyda.AA3", "confidence": 0.000104, "threat_level": "CRITICAL" },
{ "family": "Adialer.C", "confidence": 0.000067, "threat_level": "MEDIUM" }
],
"virustotal": {
"found": true,
"malicious": 61,
"total_engines": 75,
"detection_rate": 0.8133,
"url": "https://www.virustotal.com/gui/file/01bc427e..."
},
"yara": {
"rule_name": "VB_AT",
"valid": true,
"rule": "rule VB_AT : trojan vb persistence dropper { ... }"
}
}
项目结构
malware-cnn/
├── model.py # CNN 架构 —— 4 个卷积 Block
├── train.py # 训练循环 —— 支持 MPS / CUDA / CPU
├── predict.py # 带有 top-3 置信度的单张图像预测
├── scan.py # CLI 二进制扫描器 —— 熵值、SHA256、判定
├── evaluate.py # 混淆矩阵 + 各类别准确率报告
├── api.py # FastAPI —— 8 个 endpoint,VT 判定逻辑
├── yara_generator.py # 从字节签名生成 YARA 规则
├── malware_cnn.pth # 训练好的模型权重
├── confusion_matrix.png # 验证集混淆矩阵
├── .env # VT_API_KEY (未提交)
├── .gitignore
└── data/
└── malimg_paper_dataset_imgs/ # Malimg 数据集 (未提交)
malware-cnn-frontend/
├── src/
│ ├── components/
│ │ ├── Navbar.tsx
│ │ ├── Hero.tsx
│ │ ├── Scanner.tsx
│ │ ├── DropZone.tsx
│ │ ├── ResultsBanner.tsx
│ │ ├── FileInfoCard.tsx
│ │ ├── Top3Card.tsx
│ │ ├── VirusTotalCard.tsx
│ │ ├── EntropyCard.tsx
│ │ └── YaraCard.tsx
│ └── services/
│ └── api.ts
└── public/
└── eicar.exe # 用于演示的 EICAR 测试文件
设置
后端设置
git clone https://github.com/0la-mide/malware-cnn
cd malware-cnn
python3.11 -m venv venv
source venv/bin/activate
pip install torch torchvision torchaudio
pip install numpy opencv-python matplotlib scikit-learn Pillow tqdm
pip install fastapi uvicorn python-multipart requests yara-python python-dotenv
下载 Malimg 数据集并将其解压到 data/malimg_paper_dataset_imgs/
在项目根目录下创建一个 .env 文件:
VT_API_KEY=your_virustotal_api_key_here
前端设置
cd malware-cnn-frontend
npm install
npm run dev
使用方法
命令 | 描述
python train.py | 训练模型
python scan.py suspicious.exe | CLI 二进制扫描器
python api.py | 运行 FastAPI 服务器
python evaluate.py | 生成混淆矩阵和各类别准确率报告
展示的核心概念
恶意软件可视化 —— 二进制文件转灰度图 (Nataraj et al. 2011)
自定义 4 层 CNN —— Batch normalization 和 dropout 正则化
渐进式特征提取 —— 从低级纹理到抽象的威胁表示
Apple MPS GPU 加速 —— 针对 Apple Silicon 的 Metal backend
VirusTotal 驱动的判定逻辑 —— CNN 提供分类,VT 提供验证
YARA 规则生成 —— 提取字节签名以用于 SIEM/EDR 部署
REST API 设计 —— 带有共享辅助函数的 FastAPI
基于熵值的加壳检测 —— 辅助威胁信号
上下文感知的前端 —— 结果根据文件状态进行适配
全面评估 —— 各类别 F1 分数、混淆矩阵、失败案例分析
已知局限性
基于 Malimg (2011 年数据集) 训练 —— 未包含现代恶意软件家族
CNN 没有“干净”的概念 —— 总是输出一个家族分类。VirusTotal 的交叉验证可防止将误报作为最终判定结果呈现
单样本 YARA 规则在投入生产环境部署前应进行审查
Autorun.K 显示零召回率 —— 验证样本量小且与 Yuner.A 具有视觉相似性
二进制转图像的 pipeline 专为原始可执行文件设计。预转换的图像(例如,数据集的 PNG 文件)由于 PNG 压缩头会生成不可靠的结果
参考文献
# Nataraj et al. (2011) —— Malware Images: Visualization and Automatic Classification
Malimg Dataset —— Kaggle
YARA 文档
VirusTotal API v3
PyTorch MPS Backend
# malware-cnn
基于卷积神经网络 (CNN) 的 AI 恶意软件检测。本项目侧重于利用人工智能和计算机视觉技术来自动检测和分类恶意软件威胁。
标签:AV绕过, FastAPI, PyTorch, 人工智能, 凭据扫描, 威胁情报, 开发者工具, 用户模式Hook绕过, 网络测绘, 计算机视觉, 逆向工具