YedidyaBarGad/DFIR-Guardrail
GitHub: YedidyaBarGad/DFIR-Guardrail
一个基于本地 LLM 的多智能体紫队安全实验室，用于模拟和防御隐藏在 Windows 数字取证工件中的间接 prompt 注入攻击。
Stars: 0 | Forks: 0
# 🛡️ DFIR-Guardrail
**DFIR-Guardrail** 是一个多智能体紫队安全实验室，旨在模拟和防御一种关键的新兴威胁：隐藏在标准数字取证工件（例如 Amcache、Prefetch 和 Event Logs）中的间接 prompt 注入。
随着安全运营中心 (SOC) 越来越依赖大型语言模型来解析来自终端代理的分类数据，威胁行为者可以利用系统工件使分析师致盲、操纵时间线，或强制 AI 模型丢弃关键的危害指标 (IOC)。此流水线实现了一个零信任 AI 架构，以在取证数据到达主分析模型之前对其进行清理。
## 🏗️ 架构与威胁模型
本项目运行一个闭环验证流水线，同时评估攻击生成和防御拦截。
### 🔴 红队 (`src/red_team.py`)
模拟武器化取证收集流水线的高级持续性威胁 (APT)。
* **对抗性 Payload 生成：** 将操纵性的 AI 指令注入到自然解析的字符串字段中（例如，`FileDescription`、`CommandLine`、`ExecutablePath`）。
* **动态混淆引擎：** 通过使用 **Base64**、**Hexadecimal** 和 **零宽度 Unicode 空格 (ZWOP)** 随机编码 payload，绕过基本的关键字过滤。
* **确定性评估：** 利用固定的 RNG 种子确保数据集生成在各个验证运行中完全可重现。
### 🔵 蓝队 (`src/blue_team.py`)
作为高速语义护栏，在数据损害主要 SOC 分析师 LLM 之前拦截数据。
* **角色扮演越狱：** 安全地绕过企业安全对齐限制，允许本地模型在授权的防御上下文中原生分析恶意文本。
* **引导式思维链：** 强制模型在做出决定之前，在脑海中解码混淆的字符串并编写一个简短的 `` 块。
* **XML 样式输出标记：** 使用对 `` 标签的严格正则表达式提取来解析 LLM 的最终分类，使 Python 流水线免受失控生成或格式幻觉的影响。
* **异步顺序处理：** 利用 `asyncio.Semaphore` 结合 HTTPX 连接池和显式调整的上下文窗口 (`num_ctx`)，针对本地 LLM 安全地启动高度并发的评估。
### ⚙️ 流水线 (`src/pipeline.py`)
管理本地 LLM 推理的编排器。它支持三种执行模式：
* **操作模式：** 从已解析的 Windows 取证日志（KAPE/Velociraptor JSON 导出）中提取高风险字段，并扫描它们以查找间接 prompt 注入，输出结构化的 JSON 分析师报告。
* **模拟模式：** 执行合成的红/蓝队评估，生成有毒工件并测量误报率 (FPR) 和漏报率 (FNR) 指标。
* **比较模式：** 运行纯顺序基线与异步顺序框架的直接对比，并在您的硬件上直接计算性能加速比。
## 🚀 执行与部署
此架构经过专门设计，可在高度受限的硬件环境中高效运行，包括仅限于单个 16GB VRAM GPU（例如，NVIDIA T4）的无头服务器和云笔记本。
### 在 Google Colab 中运行（推荐）
提供的 notebook 自动化了 Colab VM 内整个本地 LLM 基础设施的设置。
1. 在 Google Colab 中打开 `colab_demo.ipynb`。
2. 选择 **Runtime** > **Run All**。
3. Notebook 将自动：
* 安装原生 GPU 依赖项（`pciutils`、`lshw`）。
* 下载并启动 **Ollama** daemon，并进行优化的 VRAM 分割（`OLLAMA_NUM_PARALLEL`）。
* 拉取高效的模型（例如，Microsoft Phi-3）。
* 预热 GPU 以吸收冷启动惩罚，并执行流水线比较矩阵。
### 在本地运行
确保 Ollama 已安装并在本地运行（`localhost:11434`），并且您已经拉取了目标模型（例如 `phi3`）：
```
ollama pull phi3:mini
```
克隆仓库并安装依赖：
```
# Clone 仓库
git clone https://github.com/YedidyaBarGad/DFIR-Guardrail.git
cd DFIR-Guardrail
# 安装 requirements
pip install -r requirements.txt
```
您可以在各种模式下执行流水线：
#### 1. 比较基准测试（加速测试）
生成 1,000 个合成工件，并测试纯顺序基线与优化后的异步顺序流水线。
```
python src/pipeline.py --num_samples 1000 --mode compare
```
#### 2. 模拟模式（红队测试）
生成注入了混淆 payload 的取证工件合成数据集，运行扫描，打印评估指标，并将结构化的分析师报告保存到 `guardrail_simulation.json`。
```
python src/pipeline.py --num_samples 5000 --poison_ratio 0.01 --mode async
```
#### 3. 操作模式（扫描真实工件）
扫描从 **KAPE** 或 **Velociraptor** 等工具解析出的真实 Windows 取证工件（JSON 格式），并生成保存至 `guardrail_analysis.json` 的结构化威胁报告。
```
python src/pipeline.py --input_dir /path/to/forensics/json --output_dir /path/to/output/reports
```
**CLI 选项：**
* `--mode`：流水线的执行模式。选项包括：
* `sim`（默认）：合成生成并测试有毒和良性工件。
* `ops`：扫描包含 JSON 取证工件的真实目录。
* `compare`：运行一次纯顺序扫描，随后运行一次异步顺序扫描，以基准测试加速性能。
* `--concurrency`：在 `async` 或 `compare` 模式下对 LLM 打开的并发连接数（默认：`3`）。*注意：这必须与在 Ollama daemon 上设置 `OLLAMA_NUM_PARALLEL` 相配合。*
* `--input_dir`：（仅限 Ops 模式）包含用于分析的已解析 JSON 取证工件的目录路径。
* `--output_dir`：（Ops 和 Sim 模式）用于保存结构化扫描报告（`guardrail_analysis.json` 或 `guardrail_simulation.json`）的路径。
* `--num_samples`：（Sim 和 Compare 模式）要生成的合成工件数量（默认：`50`）。
* `--poison_ratio`：（仅限 Sim 模式）注入 prompt payload 的工件比例，例如，`0.2` 表示 20%（默认：`0.2`）。
## 📊 性能指标
通过将 **混合过滤方法** 与 **异步顺序处理** 相结合，即使在受限于单个 GPU 的瓶颈下，该流水线也能实现极高的吞吐量和规模扩展。
### ⚡ 1,000 工件规模基准测试（异步 vs 顺序）
为了测试流水线的操作性能，我们在本地 GPU 上使用 Microsoft `phi3`，对包含大量混淆 payload 的 **1,000 个 Windows 取证工件** 进行了直接比较扫描：
```
--- Comparison Results ---
Metric | Sequential Baseline | Async Sequential
-----------------------------------------------------------------------
Scan Time | 347.76 s | 149.42 s
Speedup | 1.00x | 2.33 x
Total processed | 1000 | 1000
True Positives | 187 | 189
True Negatives | 806 | 806
False Positives | 0 | 0
False Negatives | 6 | 5
Errors | 1 | 0
False Positive Rate (FPR) | 0.00% | 0.00%
False Negative Rate (FNR) | 3.11% | 2.58%
```
* **性能优化：** 从标准的顺序 HTTP 查询转向带有受约束上下文窗口的连接池异步 IO，实现了 **2.33 倍的加速**，同时将漏报率降至极低的 **2.58%**。
## 🛡️ 生产就绪评估与免责声明
本项目是一个专为安全研究和模拟设计的 **概念验证**。它目前 **不** 适合在安全运营中心 (SOC) 进行企业级生产部署。
### 🚀 已完成的强化
* **结构化 LLM 输出与角色：** 强制执行严格的 JSON schema（`{"analysis": "...", "result": X}`），并使用原生 Chat API 角色将系统指令与不受信任的用户 payload 显式隔离，以缓解递归 prompt 注入。
* **语义规避与熵检测：** 部署了 `sentence-transformers` 以针对对抗性意图进行动态余弦相似度评分，并使用 `shannon_entropy` 数学计算来捕获未知的混淆技术，而无需依赖脆弱的正则表达式关键字。
* **架构弹性（DoS/OOM 保护）：** 用 `ijson` 流式处理替换了占用大量内存的传统 JSON 解析，以实现无限扩展；集成了主动断路器模式，可在不超时的情况下快速失败挂起的任务；通过 `RotatingFileHandler` 保护了取证日志。
* **故障关闭模型：** 网络超时或推理失败（返回 `-1`）时会故障关闭，将事件作为潜在绕过进行隔离，而不是故障开放。
* **异步并发：** 消除了顺序 HTTP 瓶颈。取而代之的是使用 asyncio 信号量和明确约束的上下文内存，目标是实现单 GPU 饱和而不会发生 OOM 故障。
### 🔍 待解决的生产环境差距
* **Prompt 鲁棒性：** 小型语言模型 仍然容易受到高级越狱技术（认知负荷、角色扮演覆盖）的影响，这些技术可以完美模仿良性的系统数据。
* **Schema 与格式摄取：** 硬编码的启发式字段过滤；缺乏标准化的 schema 规范化（例如，Elastic Common Schema）或多格式摄取（CSV、EVTX、XML）。
* **运维与可观察性：** 配置是硬编码的；缺乏结构化 JSON 日志记录（用于 Splunk/SIEM 摄取）和 Prometheus 监控指标。
标签：AI安全, AI对抗性攻击, AI越狱, AI风险缓解, Amcache, APT, Base64编码, Chat Copilot, DLL 劫持, DNS 反向解析, DNS 解析, Google Colab, HTTPS请求, IOC, LLM武器化, LLM评估, Ollama, Prefetch, Python, Windows取证, 事件日志, 动态混淆, 十六进制编码, 多智能体, 大语言模型, 失陷标示, 子域名变形, 安全测试, 安全运营中心, 对抗性AI, 情报收集, 攻击性安全, 攻击模拟, 数字取证, 数字取证溯源, 数据展示, 数据清洗, 无后门, 本地大模型, 漏洞研究, 端点安全, 紫队, 紫队演练, 红队, 网络安全, 网络映射, 自动化脚本, 补丁管理, 计算机取证, 语义护栏, 逆向工具, 间接提示注入, 隐私保护, 零信任架构, 零宽Unicode, 零宽字符, 驱动签名利用, 高级持续性威胁