antonisbaro/dga-sentinel

# 基于分层集成学习和注意力融合的可解释 DGA 检测 🕵️‍♂️ 本项目提出了一个全面、模块化的框架，利用一套先进的基于 Transformer 的模型来检测域名生成算法（DGA）域名。该项目的核心是一种新颖的两阶段训练方法，最终构建出一个强大的、基于注意力的元模型，以实现卓越的性能和可解释性。 整个系统使用 Python 和 TensorFlow/Keras 实现，遵循代码模块化、可复现性和可扩展性的最佳实践。它包含一个使用 Optuna 进行超参数调优的完整 pipeline，以及一个独立的评估环境，用于评估模型对未见过的 DGA 家族的泛化能力（Leave-One-Family-Out）。 ## 🚀 挑战：检测算法生成的域名 恶意软件利用域名生成算法（DGA）创建大量伪随机域名，用于命令与控制（C2）通信。这使得它们能够逃避静态封锁列表。面临的挑战是构建一个强大的机器学习系统，它能够： 1. **准确区分**良性（人类生成的）和恶意（算法生成的）域名。 2. 有效地**泛化**到新的、以前未见过的 DGA 家族。 3. 为域名*为何*被归类为恶意提供**可解释性**，为网络安全分析师提供有价值的见解（可解释 AI - XAI）。 ## 💾 数据集 本项目使用的数据集 `final_dataset.csv.gz` 是作者为本研究所策划的自定义集合。它是一个复合数据集，包括： 1. 来自[公开来源](https://tranco-list.eu/)的良性域名。 2. 来自各种已知恶意软件家族的 DGA 域名。 3. 作者创建的新颖的、合成的 DGA 域名。 部分 DGA 数据是使用改编和修改自 [DGA 脚本代码库](https://github.com/baderj/domain_generation_algorithms)（遵循 GPLv2 许可证）的脚本生成的。 **许可说明：** 本代码库的 MIT 许可证适用于所有原始源代码。数据集本身仅供研究和验证目的使用。 ## 📂 项目结构 本项目采用模块化结构，旨在提高清晰度、便于实验和维护。 - **`main_experiment.ipynb`**：用于执行所有模型的端到端训练和评估 pipeline 的主 notebook。 - **`leave_one_out_experiment.ipynb`**：专门用于进行 Leave-One-Family-Out (LOFO) 实验以测试零样本泛化能力的 notebook。 - **`hyperparameter_tuning.ipynb`**：用于使用 Optuna 运行超参数优化研究的 notebook。 - **`results_analysis.ipynb`**：专门用于加载最终结果（来自 CSV 和错误文件）并生成对比可视化图表的 notebook。 - **`colab_setup.py`**：在 Google Colab 上运行时自动配置环境的实用脚本（挂载 Drive，安装依赖包）。 - **`final_dataset.csv.gz`**：位于项目根目录的主数据集文件。 - **`configs/`**：所有参数（路径、超参数、DGA 分组）的集中配置。 - **`src/`**：包含作为 Python 库的项目核心逻辑的目录。 - **`data/`**：用于加载和拆分数据集的模块。 - **`features/`**：用于特征工程和预处理（Tokenizer、Vectorizer）的模块。 - **`models/`**：定义 Keras 模型架构和自定义层的脚本。 - **`pipelines/`**：项目的“大脑”。每个脚本都包含一个类，用于管理特定模型的完整生命周期。 - **`tuning/`**：包含 Optuna `objective` 函数和分析实用程序的脚本。 - **`analysis/`**：用于评估、错误分析和可视化逻辑的模块。 - **`outputs/`**：保存所有生成的产物（模型、预测、表示、图表、结果）的目录。 ## 🧠 算法策略：两阶段分层堆叠 本项目的核心策略是一个复杂的两阶段训练过程，将最终模型视为由“专家”神经网络组成的集成模型。 ### 阶段 1：训练基线“专家”模型 四种不同的基于 Transformer 的模型被训练为基线“专家”，每种模型旨在捕获域名的不同特征： 1. **基于字符的模型：** 处理原始字符序列的标准 Transformer。 2. **N-gram 模型：** 同时处理 2-gram 和 3-gram 的双塔架构。 3. **混合 Surprisal 模型：** 使用交叉注意力，用字符级上下文丰富统计“surprisal”值（来自马尔可夫模型）的先进双塔模型。 4. **分层模型：** 强大的分层注意力网络，它首先学习域名中单个“单词”的 embedding，然后学习这些单词之间的关系。 这些专家模型中的每一个都使用**两阶段迁移学习方法**进行训练： - **预训练：** 模型首先在广泛的、多类别的问题上进行训练（例如，将域名分类为 `alphanumeric`、`dictionary-based` 等常规类别）。这迫使模型学习丰富的、通用的特征表示。 - **微调：** 冻结预训练模型的基座，将其头部替换为二元分类器，并在最终的良性/恶意任务上进行微调。 ### 阶段 2：基于注意力的 Meta-Transformer 这是该项目的关键创新。我们没有使用简单的算法来组合基线模型，而是使用了最终的、强大的 Meta-Transformer： - **输入：** 它将每个基线模型的**预训练**阶段的丰富向量表示作为输入。这形成了一个“句子”，其中每个“单词”都是专家模型的完整“思想”。 - **机制：** 它**跨模型本身**应用自注意力，学习动态的、上下文相关的关系。例如，当 Surprisal 模型的输出模棱两可时，它可能会学会更加关注分层模型的判断。 - **训练：** 这个 Meta-Transformer 也使用同样稳健的两阶段（预训练/微调）方法进行训练，使其能够学习专家模型之间尽可能深入的协同作用。 ## 🚀 入门指南与运行方式 ### 前置条件 - Python 3.10+ - Git - 支持 GPU 的环境（推荐用于训练） ### 设置与安装 1. **克隆代码库：** git clone https://github.com/antonisbaro/dga-sentinel cd dga-sentinel 2. **设置虚拟环境（推荐）：** python -m venv .venv # 在 Windows 上：.venv\Scripts\activate # 在 macOS/Linux 上：source .venv/bin/activate 3. **安装依赖：** pip install -r requirements.txt 4. **放置数据：** - 确保您的数据集文件 `final_dataset.csv.gz` 位于项目的**根目录**中。 5. **配置 Google Colab（如适用）：** - 如果您在 Google Colab 上运行项目，需要将整个项目文件夹放置在您的 Google Drive 中。 - 打开 `colab_setup.py` 文件。 - 修改变量 `project_path_in_drive`，使其与您在 Google Drive 中项目文件夹的路径相匹配。例如，如果您的项目位于 `MyDrive/DGA_Project`，请将该变量设置为 `'DGA_Project'`。 ### 执行 本项目通过 Jupyter notebook 运行。建议按以下顺序运行它们： 1. **`hyperparameter_tuning.ipynb`**：（可选）使用此 notebook 运行 Optuna 研究并为模型寻找最佳参数。 2. **`main_experiment.ipynb`**：按顺序打开并运行单元格，以训练所有基线模型和最终的元模型。这将在 `outputs/` 目录中生成所有必要的模型和结果文件。 3. **`leave_one_out_experiment.ipynb`**：使用此 notebook 测试系统在未见过的 DGA 家族上的泛化能力。 4. **`results_analysis.ipynb`**：运行实验后，使用此 notebook 加载生成的 CSV 文件和错误日志，以生成最终的对比图表和可视化。 ## 💻 技术栈 - **语言：** Python - **核心库：** Pandas, NumPy, Scikit-learn - **深度学习：** TensorFlow, Keras - **超参数优化：** Optuna - **特征工程：** `tldextract`, `wordninja` - **绘图：** Matplotlib, Seaborn, Plotly

标签：DGA检测, IP 地址批量处理, TensorFlow, 人工智能, 可解释AI, 威胁情报, 开发者工具, 深度学习, 用户模式Hook绕过, 网络安全, 逆向工具, 隐私保护