badredinedhaoui/RAG-Powered-Cyber-Threat-Intelligence-Engine

# 基于 RAG 的网络威胁情报引擎 一个检索增强生成 (RAG) 流水线，将原始 CTI 报告和威胁描述作为输入，并自动输出可用于生产环境的 **Sigma 规则**、**YARA 规则**或结构化的 **CTI 报告**。 基于 5.2 万份文档的威胁情报知识库、领域专属的网络安全 embeddings 以及通过 Ollama 运行的本地 LLM 构建而成。 ## 功能介绍 您只需粘贴一份威胁报告（例如钓鱼攻击活动、恶意软件分析报告、安全事件摘要），该流水线便会执行以下操作： 1. 使用 **SecureBERT 2.0** 对您的查询进行 embedding，该模型是由 Cisco AI Research 专门在网络安全语料库上预训练的 BERT 模型 2. 从本地 **ChromaDB** 向量数据库（5.2 万+ CTI 示例）中检索语义最相关的威胁情报 3. 构建结构化的、针对特定输出类型的 prompt 4. 将其发送给通过 **Ollama** 在本地运行的 **sylink:32b** 5. 验证输出并将其保存到本地磁盘 无需向外部 LLM 提供商发起 API 调用。所有数据均在本地运行。 ## 输出模式 只需更改单个变量即可在三种输出类型之间切换： | 模式 | 输出 | 文件 | |------|--------|------| | `OutputType.SIGMA` | Sigma 检测规则 (YAML) | `sigma_rules/` | | `OutputType.YARA` | 包含 meta、strings、condition 的 YARA 规则 | `yara_rules/` | | `OutputType.CTI_REPORT` | Markdown 格式的完整结构化 CTI 报告 | `cti_reports/` | 每种模式都有其专属的 prompt 模板、验证逻辑和输出目录。 ## 技术栈 | 组件 | 技术 | |-----------|-----------| | Embedding 模型 | `cisco-ai/SecureBERT2.0-base` (768维, 平均池化) | | 向量数据库 | ChromaDB (持久化, 余弦相似度) | | LLM 后端 | Ollama — `sylink:32b` / `sylink:8b` | | 数据集 | `reloading0101/threat-intelligence-dataset` (约 5.2 万个 CTI 示例, 38 个类别) | | 框架 | Python, HuggingFace Transformers, PyTorch | ## 工作原理 — 流水线分解 ``` CTI Input Text │ ▼ [SecureBERT2.0 Embedder] ← mean pooling over 512 tokens │ ▼ [ChromaDB Query] ← top-K cosine similarity search │ ▼ [Prompt Builder] ← role + output-specific instructions + context │ ▼ [Ollama / sylink:32b] ← temperature=0.1 for deterministic output │ ▼ [Validator + File Save] → .yml / .yar / .md ``` 向量数据库中的文档被结构化为 `[category] + Instruction + Input + Response` —— 类别标签有助于 SecureBERT 将 embedding 锚定到正确的网络安全子领域，而响应字段（包含最密集的技术内容：IOCs、TTPs、CVE 详情、MITRE 映射）被放置在最后，以从平均池化加权中受益。 ## 数据集 **`reloading0101/threat-intelligence-dataset`** —— 包含 52,279 个指令-响应对，涵盖 38 个网络安全类别，包括： - 恶意软件分析与逆向工程 - MITRE ATT&CK 技术与子技术 - CVE 漏洞利用详情 - 钓鱼与社会工程学 - 网络入侵与横向移动 - 应急响应流程 ## 开始使用 ### 前置条件 - Python 3.9+ - 已安装并运行 [Ollama](https://ollama.com) - 一个 HuggingFace 账户（用于访问数据集） ### 安装依赖 ``` pip install datasets pandas tqdm numpy pip install torch torchvision pip install transformers accelerate tokenizers pip install chromadb pip install requests pyyaml ``` ### 拉取模型 ``` ollama pull sylink:32b ollama serve ``` ### 运行 打开 `Rag_Cybersec_v2.ipynb` 并自上而下运行单元格。 在首次运行时，流水线将会： - 下载 SecureBERT2.0 (~440MB) - 加载完整的 5.2 万条数据集并将其 embedding 到 ChromaDB 中 在 GPU 上大约需要 5-10 分钟，在 CPU 上大约需要 30-60 分钟。ChromaDB 集合会在本地持久化保存，因此后续运行将完全跳过此步骤。 要生成规则，请将您的 CTI 文本放入 `example_cti` 并设置您的输出类型： ``` OUTPUT_TYPE = OutputType.SIGMA # or OutputType.YARA / OutputType.CTI_REPORT result = cti_to_rule( cti_report=example_cti, output_type=OUTPUT_TYPE, top_k=5, save=True ) ``` ## 项目结构 ``` ├── Rag_Cybersec_v2.ipynb # Main notebook ├── db_v2/ # Persistent ChromaDB vector store (generated on first run) ├── sigma_rules/ # Generated Sigma rules (.yml) ├── yara_rules/ # Generated YARA rules (.yar) └── cti_reports/ # Generated CTI reports (.md) ``` ## 路线图 - [ ] 用于交互式查询的 Web UI - [ ] 支持更多本地模型 (Mistral, Llama 3 等) - [ ] 通过 `sigma-cli` 进行 Sigma 规则语法验证 - [ ] 用于处理多个报告的批处理模式 - [ ] MITRE ATT&CK 覆盖情况热力图输出 ## 注意事项 - Notebook 中的 HuggingFace 访问令牌用于加载数据集 —— 请在 [hf.co/settings/tokens](https://huggingface.co/settings/tokens) 替换为您自己的令牌 - ChromaDB 配置了余弦相似度 (`hnsw:space: cosine`)，这非常适合对混合长度的安全文档进行语义检索 - 温度被刻意设置为 `0.1`，以保持 Sigma/YARA 输出的确定性和结构有效性 ## 许可证 MIT

标签：AI风险缓解, AMSI绕过, ChromaDB, DAST, DLL 劫持, DNS信息、DNS暴力破解, HuggingFace, LLM评估, Ollama, Python, PyTorch, RAG, SecureBERT, Sigma规则, YARA规则, 人工智能, 凭据扫描, 向量数据库, 大语言模型, 威胁检测, 安全规则生成, 安全运营, 恶意软件分析, 扫描框架, 无后门, 本地部署, 检索增强生成, 用户模式Hook绕过, 目标导入, 网络威胁情报, 网络安全, 逆向工具, 钓鱼攻击分析, 隐私保护