YaCnDehfuli/VADViT

# VADViT [](https://www.python.org/) [](https://pytorch.org/) [](https://volatilityfoundation.org/) [](https://doi.org/10.1016/j.jisa.2025.104200) [](LICENSE) **技术焦点：** 内存取证 · 恶意进程检测 · Vision Transformers · VAD 分析 · 可解释 AI VADViT 是一个内存取证研究流水线，用于从虚拟地址描述符（VAD）区域中检测恶意 进程。它从区间内存快照中提取进程内存区域，将这些区域转换为 Markov、熵和 VAD 元数据图像通道，将它们排列成进程级别的网格，并训练 Vision Transformer 进行二分类或恶意软件家族分类。 该实现随附于以下论文： 该论文报告称，最佳的 VADViT 配置在 BCCC-MalMem-SnapLog-2025 上实现了 99.2% 的二分类准确率，以及 92% 的多分类家族分类宏观平均 F1 分数。这些数值取决于研究中使用的数据集、数据划分、训练配置和检查点；本仓库 不提供原始内存转储或训练好的权重。  ## 本仓库包含的内容 - 基于 Volatility 的 VAD 提取，针对每个样本的内存快照。 - 快照合并，为每个样本保留最丰富的进程内存视图。 - 将区域分类为可执行/恶意软件、DLL 支持和堆/栈组。 - 为每个保留的 VAD 区域生成 Markov、熵和强度通道图像。 - 为 ViT 输入构建进程级网格。 - ViT 训练、验证、测试评估和 attention 覆盖图工具。 ## 仓库布局 ``` Data Preprocessing/ Dumps_to_Cnosolidated/ Volatility vadinfo extraction, region division, and snapshot consolidation Consolidated_to_Grid/ VAD region -> RGB patch -> process grid images dataset/ ImageDataset and train/val/test transforms models/ timm ViT wrapper with configurable frozen blocks utils/ training loop, metrics, attention visualization config.py training/evaluation configuration train.py train a ViT on generated grid images test.py evaluate a saved model, optionally with attention sample_test.py single-image inspection helper with hard-coded paths requirements.txt Python dependency pins used by the project ``` 目录名 `Dumps_to_Cnosolidated` 被特意保留为仓库中的原样记录，包括拼写错误 ，以便可以直接复制命令。 ## 数据要求 原始转储未提交到此仓库。预期的预处理输入按恶意软件家族和样本哈希组织： ``` BASE_DIR/ Trojan/ / Dumps/ _snapshot1.vmem _snapshot2.vmem ... Benign/ / Dumps/ _snapshot1.vmem ``` 每个转储文件名必须以目标 PID 开头，因为 `region_extractor.py` 从 `dump_file.split("_")[0]` 读取 PID。原始 研究在每个样本中最多使用了五个区间快照。如果 PID 不再 存在于 `windows.pslist` 中，则样本会提前停止；只有一次转储的样本 会被单独计算，因为它们包含较少的时间证据。 ## 环境 使用与您的 CUDA/PyTorch 设置相匹配的 Python 环境。已提交的 requirements 文件固定了仓库的工作依赖集；论文中 报告的实验仍应使用该次运行所用的确切环境进行复现。GPU 用户可能需要 在安装其余软件包之前，安装适合其驱动程序的 PyTorch 版本。 ``` python3 -m venv .venv source .venv/bin/activate python -m pip install --upgrade pip python -m pip install -r requirements.txt ``` 提取阶段需要 Volatility 3。在 `Data Preprocessing/Dumps_to_Cnosolidated/config.py` 中设置其路径。 ## 阶段 1：提取并合并 VAD 区域 编辑 `Data Preprocessing/Dumps_to_Cnosolidated/config.py`： ``` BASE_DIR = "/path/to/BCCC-MalMem-SnapLog-2025" OUTPUT_DIR = "/path/to/intermediate_vad_regions" CONSOLIDATED_DIR = "/path/to/consolidated_vad_regions" VOLATILITY = "/path/to/volatility3/vol.py" ``` 运行提取流水线： ``` cd "Data Preprocessing/Dumps_to_Cnosolidated" python main.py ``` 此阶段对每个样本执行三个操作： 1. 运行 `windows.pslist` 以确认每个快照中存在目标 PID。 2. 运行 `windows.vadinfo --pid --dump` 并写入 `vadinfo.csv` 以及转储的 VAD 区域文件。 3. 将保留的区域分类到 `malware_executable/`、`dlls/` 和 `heap_and_stack/` 中，然后将具有最多可执行/DLL 区域的快照复制到 `CONSOLIDATED_DIR`。 合并后的输出预期如下所示： ``` CONSOLIDATED_DIR/ Trojan/ / malware_executable/ malware_executable_regions.csv vad.0x...dmp dlls/ dll_regions.csv vad.0x...dmp ``` ## 阶段 2：构建进程网格图像 编辑 `Data Preprocessing/Consolidated_to_Grid/config.py`： ``` IMAGE_SIZE = 224 # paper evaluates 224 and 384 PATCH_SIZE = 32 # paper evaluates 16 and 32 CONSOLIDATED_DIR = "/path/to/consolidated_vad_regions" IMAGE_DATASET_DIR = "/path/to/image_datasets" ``` 运行网格生成： ``` cd "Data Preprocessing/Consolidated_to_Grid" python main.py ``` 对于每个保留的 VAD 区域，`process2image.py` 会创建： - 红色通道：VAD 标签和保护元数据强度； - 绿色通道：动态窗口香农熵； - 蓝色通道：Markov 字节转换结构。 可执行区域首先按 VAD 地址升序排列，随后是 DLL 支持的区域。空网格单元进行零填充。输出路径为： ``` IMAGE_DATASET_DIR/ 32_224/ Trojan/ .png Benign/ .png ``` ## 阶段 3：训练 VADViT 编辑根目录的 `config.py`： ``` IMAGE_SIZE = 224 PATCH_SIZE = 32 MODE = "Binary" # "Binary" or "Multi" FROZEN_LAYERS = 6 STEPS = 3 DATASET_PATH = "/path/to/image_datasets/32_224/32_224_Binary" SAVE_PATH = "./models/Binary_32_224_6f_3u.pt" ``` 然后运行： ``` python train.py ``` 数据集加载器从 `DATASET_PATH` 下的类文件夹创建 80/10/10 的训练/验证/测试划分。对于二分类模式，`Benign` 映射到类 `0`， 所有其他文件夹映射到类 `1`。对于多分类模式，文件夹按 字母顺序排序并映射到数字标签。 训练使用 timm ViT 主干、标签平滑交叉熵、 ReduceLROnPlateau、冻结 transformer 块的渐进解冻、预测期间温度缩放以及后期随机权重平均。最佳模型 保存到 `SAVE_PATH`；指标图表保存为 `training_plot.png`。 ## 评估与可解释性 评估已保存的模型： ``` python test.py ``` 启用最后一个 block 的 attention 可视化： ``` python test.py --explain ``` `test.py` 从 `DATASET_PATH` 加载 `test` 划分，打印分类 报告和混淆矩阵，并将 ROC/混淆矩阵图表写入由 `AUC_FOLDER` 和 `CM_FOLDER` 配置的文件夹中。ROC 辅助工具专为 二分类评分设计，因此当 `MODE = "Multi"` 时请谨慎使用。 对于单次检查，编辑 `sample_test.py` 中的硬编码路径并运行： ``` python sample_test.py ``` 该辅助工具加载一张进程网格图像，打印类概率，显示 attention 覆盖图，并列出可执行/DLL VAD 区域文件，以便将 patch 顺序追溯到具体地址。 ## 论文配置说明 论文评估了图像大小、patch 大小和冻结层 策略的组合。最强的二分类配置是 `PATCH_SIZE=32`、 `IMAGE_SIZE=224`、`FROZEN_LAYERS=6` 和 `STEPS=3`。多分类实验 使用与两种最终配置之一相同的 `32_224_6f` 家族标签设置。 论文的数据集 BCCC-MalMem-SnapLog-2025 包含来自 Backdoor、Exploit、HackTool、Hoax、Rootkit、Trojan、Virus 和 Worm 家族的恶意软件样本， 以及良性样本。由于大小和处理 限制，此处省略了原始转储。 ## 重要界限 - VADViT 分析从目标进程捕获的内存区域；它不是 实时的 EDR 或防病毒产品。 - Attention 图是取证线索，不是因果关系的证明。 - 报告的指标需要原始数据集、严格的划分原则和训练 设置；未经重新评估，请勿将其用于不同的语料库。 - Volatility 符号支持和转储质量直接影响区域提取。 ## 引用 ``` @article{dehfouli2025vadvit, title = {VADViT: Vision transformer-driven memory forensics for malicious process detection and explainable threat attribution}, author = {Dehfouli, Yasin and Lashkari, Arash Habibi}, journal = {Journal of Information Security and Applications}, volume = {94}, pages = {104200}, year = {2025}, doi = {10.1016/j.jisa.2025.104200} } ``` ## 许可证 [MIT](LICENSE)。数据集访问和与流水线一起使用的第三方工具可能 有单独的条款。如果本仓库对您的工作有帮助，请使用 [`CITATION.cff`](CITATION.cff) 引用 论文。

标签：PyTorch, SecList, Vision Transformer, Volatility 3, 人工智能, 内存取证, 凭据扫描, 用户模式Hook绕过, 逆向工具