mahfujde/AI-Sheild

# AI-Sheild **基于 RAG、多 LLM 评估与防护栏的可信网络威胁情报** CTI-Shield 是一个原型平台，通过结合检索增强生成（RAG）、多 LLM 评估与自动化防护栏验证，生成符合 STIX 2.1 标准的威胁情报。它减少幻觉、强化模式合规性，并提供可审计的质量指标——所有功能均集成于单个 Streamlit 仪表板。 ## 架构 ``` ┌─────────────────────────────────────────────────────────────────┐ │ Streamlit Dashboard │ │ ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────────────┐ │ │ │ Analyze │ │ STIX │ │ Summary │ │ Metrics │ │ │ │ Query │ │ JSON │ │ Brief │ │ Halluc│Comp│Chart │ │ │ └────┬─────┘ └────▲─────┘ └────▲─────┘ └────────▲─────────┘ │ │ │ │ │ │ │ │ ┌────▼─────────────┴────────────┴────────────────┴──────────┐ │ │ │ Pipeline Orchestrator │ │ │ │ preprocess → retrieve (RAG) → LLM → guardrails → output │ │ │ └──┬──────────────┬──────────────┬──────────────┬───────────┘ │ │ │ │ │ │ │ │ ┌───▼───┐ ┌──────▼──────┐ ┌────▼─────┐ ┌─────▼──────────┐ │ │ │Preproc│ │ FAISS RAG │ │Multi-LLM │ │ Guardrails │ │ │ │ Layer │ │ (embed + │ │ Engine │ │ ┌────────────┐ │ │ │ │ │ │ retrieve) │ │ GPT-4 │ │ │STIX Valid. │ │ │ │ │-dedup │ │ │ │ LLaMA-3 │ │ │Halluc.Det. │ │ │ │ │-clean │ │ local or │ │ Demo │ │ │Retry Mech. │ │ │ │ │-norm │ │ OpenAI emb. │ │ │ │ └────────────┘ │ │ │ └───────┘ └─────────────┘ └──────────┘ └────────────────┘ │ └─────────────────────────────────────────────────────────────────┘ ``` ## 快速开始 ### 1. 克隆并安装 ``` cd "AI Sheild" python -m venv .venv && source .venv/bin/activate pip install -r requirements.txt ``` ### 2. （可选）配置 API 密钥 ``` cp .env.example .env # 使用您的 OpenAI 密钥和/或 LLaMA 模型路径编辑 .env ``` ### 3. 运行 ``` streamlit run app.py ``` 1. 点击侧边栏中的 **Ingest Sample Data**。 2. 选择示例查询或粘贴自己的内容。 3. 选择任务（分类 / 映射 TTPs / 摘要 / 生成 STIX）。 4. 点击 **Run Analysis**。 ## 运行模式 | 模式 | 要求 | 可用 LLM | |------|-------------|----------------| | **Demo** | 无 | 确定性模拟 | | **API** | `OPENAI_API_KEY` | GPT-4 | | **Full** | `OPENAI_API_KEY` + LLaMA GGUF | GPT-4 + LLaMA-3（跨模型对比） | ## 项目结构 ``` AI Sheild/ ├── app.py # Streamlit dashboard ├── config.py # Central configuration ├── requirements.txt ├── .env.example ├── cti_shield/ │ ├── __init__.py │ ├── preprocessing.py # Text cleaning, loading, dedup │ ├── rag.py # Embeddings + FAISS store + retrieval │ ├── llm_engine.py # GPT-4 / LLaMA-3 / Demo dispatch │ ├── stix_models.py # Pydantic STIX 2.1 validators │ ├── guardrails.py # Hallucination detection + STIX validation + retry │ └── output.py # JSON / markdown / metrics formatting ├── data/ │ ├── osint_reports/ # Sample .txt threat reports │ ├── mitre_attack.json # Curated MITRE ATT&CK subset │ └── faiss_index/ # Generated at runtime └── tests/ ├── test_preprocessing.py ├── test_rag.py ├── test_guardrails.py └── test_integration.py ``` ## 工作原理 ### 1. 数据摄取 加载 OSINT 报告（.txt / .pdf）与 MITRE ATT&CK 技术，使用 `preprocess_data()` 清洗、去重、分块，并嵌入 FAISS 向量存储。 ### 2. 检索增强生成 当用户提交查询时，系统从 FAISS 中检索最相关的 Top-K 分块。这些分块作为上下文注入到特定任务的提示中，使 LLM 的回答基于真实情报。 ### 3. 多 LLM 评估 在 Full 模式下，GPT-4 与 LLaMA-3 同时处理相同任务。系统对比其输出一致性——不一致可能表明幻觉或歧义。 ### 4. 防护栏验证 - **STIX 2.1 验证**：Pydantic 模型强制校验 ID 格式、必需字段、时间戳与对象关系。 - **幻觉检测**：通过 token 重叠将输出中的主张与检索上下文进行交叉验证。 - **重试机制**：若验证失败，流水线最多自动重试 N 次，保留得分最高的尝试。 ### 5. 输出 - **可下载的 STIX 2.1 JSON**。 - **供人工阅读的 Markdown 摘要**。 - **包含幻觉率、合规评分以及 RAG 与基线对比图表的指标面板**。 ## 为何 CTI-Shield 优于基线 LLM 使用方式 | 维度 | 基线 LLM | CTI-Shield | |-----------|-------------|------------| | **事实 grounding** | 仅依赖训练数据 | RAG 从摄取的情报中检索实时上下文 | | **模式合规性** | 生成随意的 JSON | Pydantic 强制 STIX 2.1 | | **幻觉风险** | 不可检测 | 按声明评分并标记 | | **交叉验证** | 单模型 | 多 LLM 一致性检查 | | **可靠性** | 一次性生成 | 带防护栏的自动重试 | | **可审计性** | 黑盒 | 显示指标面板与检索上下文 | ## 测试 ``` # 运行所有测试 python -m pytest tests/ -v # 运行特定测试文件 python -m pytest tests/test_guardrails.py -v ``` 测试覆盖： - **单元测试**：预处理、检索相关性、STIX 验证、幻觉评分。 - **集成测试**：从摄取到验证 STIX 输出的完整流水线。 - **对抗测试**：空查询、提示注入尝试、超大输入。 ## SOC 部署指南 1. **数据来源**：用实时 OSINT 源（TAXII、威胁情报 API）替换示例报告，并定期重新摄取以保持 FAISS 索引最新。 2. **模型选择**：使用 GPT-4（或等效模型）进行分析；在隔离环境中部署本地 LLaMA-3 实例用于一致性检查。 3. **阈值调整**：根据组织的风险承受能力，在 `config.py` 中调整 `HALLUCINATION_THRESHOLD`。 4. **集成**：通过下载 API 将 STIX 束直接导出至 SOAR/SIEM（例如 MISP、Splunk SOAR、TheHive）。 5. **监控**：随时间跟踪“指标”标签页，以衡量分析师生产力提升与幻觉趋势。 ## 许可证 此项目为研究原型。使用风险自担。

标签：AI安全, AI评估框架, C2, Chat Copilot, Clair, DLL 劫持, Kubernetes, LLM护栏, masscan, RAG, STIX, STIX 2.1, Streamlit, 上下文检索, 人工智能治理, 全文检索, 可信AI, 多LLM评估, 大语言模型, 威胁情报, 威胁情报平台, 审计指标, 平台原型, 幻觉缓解, 开发者工具, 护栏, 数据管道, 检索增强生成, 模式合规, 流式处理, 网络安全, 自动化验证, 访问控制, 语义检索, 质量度量, 软件工程, 逆向工具, 隐私保护