PranavDeshpande7337/LangGraph-Threat-Intelligence-Agent

# LangGraph 威胁情报 Agent 🛡️ 一个代理式 OSINT 调查 pipeline，它接收可疑的 IP 或 domain，并自主进行端到端的威胁情报调查。该 Agent 查询多个威胁情报源，对其发现进行推理，生成结构化的风险报告，并将发现映射到 MITRE ATT&CK 技术——这种调查就像 SOC 分析师手动执行的一样，被自动化为一个可重复、可审计的工作流。 该项目旨在展示直接适用于 AI 安全组织中威胁情报和检测与响应职位的检测工程、代理式推理和 AI 辅助调查工作流。 ## 架构 ``` User input (IP or domain) │ ▼ ┌─────────────────────┐ │ Input Guardrail │ Validates format, blocks RFC 1918 ranges, blocks │ core/guardrails.py │ internal keywords. SSRF prevention before any tool call. └────────┬────────────┘ │ ▼ ┌─────────────────────┐ │ LLM Reasoning │ Ollama-powered reasoning loop. Decides which tool to │ graph/nodes.py │ call next based on accumulated state. Hard cap: 4 loops. └────────┬────────────┘ │ ├──▶ [VirusTotal] ──┐ ├──▶ [AbuseIPDB] ──┼──▶ [Output Guardrail] ← scans API responses for └──▶ [DNS Lookup] ──┘ prompt injection, redacts in-place before result enters agent state │ ▼ ┌─────────────────────┐ │ Report Generator │ Structured JSON risk report: │ graph/nodes.py │ target, risk level, score, └────────┬────────────┘ indicators, recommendation │ ▼ ┌─────────────────────┐ │ MITRE ATT&CK Node │ LLM maps report findings to │ graph/nodes.py │ ATT&CK technique IDs, names, └────────┬────────────┘ relevance, tactic categories │ ▼ ┌─────────────────────┐ │ Streamlit Dashboard│ Real-time agent trace, risk │ dashboard.py │ card, ATT&CK table, JSON export └─────────────────────┘ ``` ## 检测与调查节点 ### 输入 Guardrail (`core/guardrails.py`) 在进行任何外部调用之前运行： | 检查内容 | 拦截目标 | |---|---| | 格式验证 | 格式错误的 IP 和无效的 domain 字符串 | | RFC 1918 范围 | 10.x, 172.16.x, 192.168.x, 127.x, 169.254.x (防止 SSRF) | | 内部关键字 | `localhost`, `internal`, `corp`, `intranet` | ### LLM 推理循环 (`graph/nodes.py`) Agent 通过读取累积的 state 并输出结构化的 JSON 决策来决定接下来调用哪个工具。每个工具最多调用一次。如果达到循环上限（4次迭代）或所有工具都已被调用，则回退到报告生成。 ### 威胁情报工具 (`core/tools.py`) | 工具 | 返回内容 | 回退方案 | |---|---|---| | VirusTotal | 恶意/可疑/安全计数、信誉评分 | 模拟数据 | | AbuseIPDB | 滥用置信度评分、报告计数、Tor 标志、ISP、国家 | 模拟数据 | | DNS | 正向解析（domain → IPs）或反向 PTR 查找 | 实时 — 无需密钥 | ### 输出 Guardrail (`core/guardrails.py`) 在每次工具调用之后，结果进入 agent state 之前运行： - 递归扫描所有字符串字段以查找 prompt injection 模式 - 模式包括：指令覆盖尝试、角色重分配字符串、system prompt 注入、数据渗出指令 - 原位脱敏匹配到的内容并记录发现以供审计 - 被标记的工具将在仪表板中显示 ### 报告生成器 (`graph/nodes.py`) 将所有工具结果综合成一份结构化的风险报告： ``` { "target": "185.220.101.1", "risk_level": "critical", "risk_score": 92, "summary": "...", "indicators": ["..."], "recommendation": "..." } ``` | 分数 | 风险等级 | |---|---| | > 80 个恶意检测或滥用分数 > 80 | CRITICAL（严重） | | > 3 个恶意或滥用分数 > 50 | HIGH（高） | | > 5 个可疑检测或 Tor exit node | MEDIUM（中） | | 轻微问题，无明显恶意活动 | LOW（低） | | 无检测，信誉良好的基础设施 | CLEAN（安全） | ### MITRE ATT&CK 映射节点 (`graph/nodes.py`) 在报告生成器之后运行。将完成的风险报告发送给 LLM（附带结构化的 ATT&CK prompt）并返回： ``` { "mitre_techniques": [ { "id": "T1583.001", "name": "Acquire Infrastructure: Domains", "relevance": "Target domain registered recently with privacy protection, consistent with adversary infrastructure acquisition." } ], "tactic_categories": ["Resource Development", "Reconnaissance"] } ``` 如果风险等级为 `clean` 或 `unknown`，则会优雅地跳过 — 避免不必要的 LLM 调用。 ## Streamlit 仪表板 (`dashboard.py`) - **实时推理轨迹** — 每个节点在完成后通过 `st.status` 流式传输其输出，展示 Agent 在每一步的决策 - **风险报告卡片** — 颜色编码的风险徽章、评分进度条、指标列表、建议 - **MITRE ATT&CK 表格** — 技术 ID 直接链接到 `attack.mitre.org`，包含相关性和战术类别列 - **已执行工具面板** — 显示哪些工具已运行，并标记任何其输出中被脱敏了注入模式的工具 - **导出** — 将完整的调查（报告 + ATT&CK 映射）下载为带时间戳的 JSON 文件，或保存到 `exports/` - **演示目标** — 一键使用模拟数据调查预设目标，无需 API 密钥 ## 快速开始 ### 1. 安装依赖 ``` python -m venv venv source venv/bin/activate # Windows: venv\Scripts\activate pip install -r requirements.txt ``` ### 2. 安装 Ollama（LLM 推理必需） 从 [ollama.com/download](https://ollama.com/download) 下载，然后执行： ``` ollama pull llama3.2 ollama serve # keep this running in a separate terminal ``` ### 3. 配置环境 在项目根目录下创建一个 `.env` 文件： ``` OLLAMA_MODEL=llama3.2 VIRUSTOTAL_API_KEY= # optional — uses mock data if blank ABUSEIPDB_API_KEY= # optional — uses mock data if blank ``` ### 4. 运行 CLI ``` # 演示模式 — 使用 mock data 运行 3 个目标 python main.py --demo # 调查特定目标 python main.py --target 8.8.8.8 python main.py --target malware.com # 输出原始 JSON 报告 python main.py --target malware.com --json # 交互模式 python main.py ``` ### 5. 启动仪表板 ``` streamlit run dashboard.py ``` 打开 `http://localhost:8501`。使用侧边栏输入目标或点击演示按钮。 ### 6. 运行测试 ``` pytest tests/ -v ``` ## 环境与 API 密钥 | 变量 | 必需 | 备注 | |---|---|---| | `OLLAMA_MODEL` | 是 | 默认值：`llama3.2`。在低内存机器上可尝试 `llama3.2:1b` | | `VIRUSTOTAL_API_KEY` | 否 | 在 [virustotal.com](https://virustotal.com) 获取免费层级。未配置时回退到模拟数据 | | `ABUSEIPDB_API_KEY` | 否 | 在 [abuseipdb.com](https://abuseipdb.com) 获取免费层级。未配置时回退到模拟数据 | 模拟数据涵盖了一组已知的测试目标（`8.8.8.8`、`malware.com`、`185.220.101.1` 等），并具有真实的数值。完整的推理 pipeline、guardrail 和 MITRE 映射在模拟模式下工作方式完全相同。 ## 项目结构 ``` langgraph-threat-intel/ ├── core/ │ ├── state.py # AgentState — shared graph state (TypedDict) │ ├── guardrails.py # Input validator + output sanitiser │ └── tools.py # VirusTotal, AbuseIPDB, DNS API wrappers ├── graph/ │ ├── nodes.py # All node functions including MITRE mapping │ └── builder.py # Graph topology — nodes, edges, conditional routing ├── tests/ │ └── test_guardrails.py ├── exports/ # JSON exports from dashboard (auto-created) ├── dashboard.py # Streamlit investigation dashboard ├── main.py # CLI entry point ├── requirements.txt └── .env ``` ## 扩展 Pipeline **添加新的威胁情报源** — 按照相同的规范化 dict 返回模式，在 `core/tools.py` 中添加一个包装函数，在 `graph/nodes.py` 中添加一个工具节点，在 `graph/builder.py` 中注册它，并将其添加到 `reasoning_node` 中的 LLM 的 system prompt 中。Pipeline 的其余部分会自动识别它。 **更换 LLM** — 编辑 `graph/nodes.py` 中的 `llm` 初始化块。任何兼容 LangChain 的聊天模型都可以作为直接替代品（`ChatOpenAI`、`ChatAnthropic`、`ChatGoogleGenerativeAI` 等）。 **连接真实的日志源或 SIEM** — Agent 只需要一个目标字符串作为输入。将 `agent.invoke(initial_state)` 包装在为您提供指标的任何摄取层中（webhook、Kafka consumer、SOAR playbook 触发器）。 ## 涵盖的概念 - **代理式调查工作流** — LangGraph 多步推理，具备工具选择、状态累积和循环控制 - **检测工程** — guardrail 设计、prompt injection 检测、在编排层进行 SSRF 防御 - **OSINT 方法论** — 多源指标扩充、DNS 透视、信誉关联 - **MITRE ATT&CK 框架** — 从非结构化威胁发现中进行自动化技术映射 - **威胁行为者画像** — 根据基础设施模式通过 LLM 辅助进行风险特征分析 - **可审计的 AI 系统** — 每个节点的决策都可以通过共享的 state 和仪表板轨迹视图进行追踪 - **本地 LLM 集成** — 基于 Ollama 的推理，实现零成本、保护隐私的运行

标签：AI智能体, AI风险缓解, ESC4, Kubernetes, LangGraph, OSINT, 威胁情报, 安全, 开发者工具, 自动化调查, 超时处理, 逆向工具