Kibisuvictor/cybersec-threat-agent
GitHub: Kibisuvictor/cybersec-threat-agent
基于 LangGraph 和 Gemini 的自然语言转 SQL 智能体,让安全分析师无需 SQL 知识即可交互式查询网络安全威胁情报数据。
Stars: 1 | Forks: 0
# 网络安全威胁情报 Agent — 使用 LangChain + DuckDB 的 Text-to-SQL




用通俗的英语询问关于真实世界网络安全威胁数据的问题。一个 LangGraph Agent 会编写并执行 DuckDB SQL,以表格形式返回结果,并自动生成图表 — 无需 SQL 知识。
基于 **LangGraph**、**Gemini 2.5 Flash**、**DuckDB** 和 **Gradio** 构建。

## 什么是 Text-to-SQL Agent?
大型语言模型可以将自然语言翻译成 SQL — 但要获得可靠的结果,仅靠单次 prompt 是不够的。本项目使用了一个 **LangGraph ReAct Agent**,它会:
1. 接收用户问题以及实时的数据库 schema
2. 使用生成的 DuckDB 查询调用 `execute_sql` 工具
3. 检查结果 — 如果查询失败,它会读取错误信息并重试
4. 返回关于发现的通俗英语总结
随后,第二次 Gemini 调用会检查结果表,并决定渲染哪种图表 — 无需硬编码的启发式规则。
## 数据集
**Cybersecurity Threat Intelligence Dataset 2026** — 数据源自 AlienVault OTX、CISA 的已知漏洞利用 (Known Exploited Vulnerabilities) 目录以及美国国家漏洞数据库 (NVD)。最初由 Kaggle 上的 [chuneeb](https://www.kaggle.com/datasets/chuneeb/ai-cybersecurity-threat-dataset-2026) 整理并发布。
包含 4 个表,共 4,312 条记录:
| 表 | 来源 | 记录数 | 包含内容 |
|---|---|---|---|
| `otx_threats` | AlienVault OTX | 2,365 | 恶意软件家族、目标国家、攻击技术、APT 组织 |
| `cve_vulnerabilities` | CISA KEV + NVD | 1,585 | 2024–2026 年的 CVE、CVSS 评分、供应商/产品、补丁状态 |
| `malicious_domains` | OTX | 162 | 包含注册商和国家的钓鱼、C2、恶意软件域名 |
| `malicious_ips` | OTX | 200 | 被标记用于扫描、暴力破解、僵尸网络、数据泄露的 IP |
| `cve_prioritized` | 派生视图 | — | 带有严重/高/中/低危严重性标签的 CVE |
**关于“未知”(Unknown) 值的说明:** 某些字段包含“Unknown” — 这反映了现实世界中的威胁情报,其归因往往是不完整的。Agent 在 prompt 中被要求将其视为有效数据,而不是缺失值。
## 示例问题
```
Which malware families appear most frequently?
Show me all Critical CVEs that are unpatched
Which countries are most targeted in OTX data?
Top 10 vendors with the most exploited vulnerabilities
Which domains are flagged for C2 activity?
Distribution of CVSS scores across all CVEs
Which adversary groups have the most threat pulses?
All malicious IPs associated with botnet activity
How many CVEs were published each year?
Which attack techniques appear most in OTX data?
```
## 架构
```
User question (Gradio UI)
│
▼
LangGraph ReAct agent
(Gemini 2.5 Flash + live schema injected into system prompt)
│
├── inject schema ──► schema.py
│ (table names, columns, sample rows)
│
└── call tool ──────► execute_sql tool
│
▼
DuckDB
(otx_threats, cve_vulnerabilities,
malicious_domains, malicious_ips)
│
▼
Query results (DataFrame)
│
┌─────────┴──────────┐
▼ ▼
Agent answer Chart spec LLM call
(plain English) (Gemini decides chart type)
│ │
└─────────┬──────────┘
▼
Gradio UI output
(answer + rendered chart)
```
**Agent 重试循环:**
```
1. Receive question + schema context
2. LLM reasons: "I need CVEs with cvss_score >= 9.0 and patch_status = 'Unpatched'"
3. LLM calls execute_sql("SELECT cve_id, vendor, cvss_score FROM ...")
4. Tool runs query, returns markdown table
5. LLM summarises key finding in plain English
6. Second LLM call inspects result → returns JSON chart spec → chart renders
7. If SQL error at step 4 → LLM reads error, rewrites query, retries
```
## 技术栈
| 组件 | 工具 | 原因 |
|---|---|---|
| LLM | [Gemini 2.5 Flash](https://ai.google.dev) | 速度快,强大的工具调用能力,慷慨的免费额度(1,500 次请求/天) |
| Agent 框架 | [LangGraph](https://langchain-ai.github.io/langgraph) | ReAct Agent 循环,取代已弃用的 AgentExecutor |
| 编排 | [LangChain 1.x](https://python.langchain.com) | 工具定义,消息类型 |
| 数据库 | [DuckDB](https://duckdb.org) | 快速的分析型 SQL,零配置,原生读取 CSV |
| UI | [Gradio 6](https://gradio.app) | 简洁的聊天界面,兼容本地和 Hugging Face Spaces |
## 项目结构
```
cybersec-threat-agent/
├── agent.py # LangGraph ReAct agent + SQL tool
├── app.py # Gradio 6 UI with dynamic charting
├── schema.py # Live schema introspection from DuckDB
├── data/
│ ├── load_data.py # Loads Kaggle CSVs into DuckDB
│ ├── raw/ # Raw CSVs — git-ignored, re-download from Kaggle
│ └── cybersec.duckdb # DuckDB database — git-ignored, built locally
├── tests/
│ └── test_sql_tool.py # SQL tool + database integrity tests
├── .env.example # Template — copy to .env and add your API key
├── requirements.txt
├── .gitignore
└── README.md
```
## 快速开始
### 前置条件
- Python 3.11+
- 从 [Google AI Studio](https://aistudio.google.com) 获取的免费 Gemini API key(无需信用卡)
- 用于下载数据集的 Kaggle 账户
### 1. 克隆并安装
```
git clone https://github.com/Kibisuvictor/cybersec-threat-agent.git
cd cybersec-threat-agent
python -m venv .venv
source .venv/bin/activate # Windows: .venv\Scripts\activate
pip install -r requirements.txt
```
### 2. 添加你的 API key
```
cp .env.example .env
```
打开 `.env` 并添加你的 key:
```
GOOGLE_API_KEY=your_key_here
```
在 [aistudio.google.com](https://aistudio.google.com) 获取免费 key — 免费额度包含 1,500 次请求/天,足以满足本地使用需求。
### 3. 下载数据集
从 [Kaggle](https://www.kaggle.com/datasets/chuneeb/ai-cybersecurity-threat-dataset-2026) 下载四个 CSV 文件,并将它们放入 `data/raw/`:
```
data/raw/otx_threat_intelligence.csv
data/raw/cve_vulnerabilities.csv
data/raw/malicious_domains.csv
data/raw/malicious_ips.csv
```
### 4. 加载到 DuckDB
```
python data/load_data.py
```
### 5. 验证 schema
```
python schema.py
```
你应该会看到所有四个表及其列和各自的示例行。
### 6. 运行应用
```
python app.py
```
打开 [http://localhost:7860](http://localhost:7860) 并开始查询。
## 环境变量
| 变量 | 必需 | 描述 |
|---|---|---|
| `GOOGLE_API_KEY` | ✅ | 来自 Google AI Studio 的 Gemini API key |
将 `.env.example` 复制到 `.env` 并填入你的 key。切勿提交 `.env` — 它已被 git 忽略。
## 运行测试
```
pytest tests/ -v
```
当 `cybersec.duckdb` 尚未构建时,测试会优雅地跳过 — 这对于没有 Kaggle 数据的环境是安全的。
## 设计决策
**为什么选择 LangGraph 而不是旧的 AgentExecutor?**
LangChain 1.x 已完全弃用 `AgentExecutor` 和 `create_tool_calling_agent`。LangGraph 的 `create_react_agent` 是目前的标准 — 它原生以状态机的形式处理 ReAct 循环,具有更清晰的错误恢复机制,并且无需手动配置 `max_iterations` 或 `handle_parsing_errors`。
**为什么使用动态图表生成而不是启发式方法?**
以前的方法使用关键字匹配(`if "malware" in question`)来决定绘制什么图表。这很脆弱 — 它会漏掉换种说法的问题,并生成错误的图表类型。目前的方法会发起第二次 Gemini 调用,接收实际的结果表,并返回一个 JSON 图表规格 (`{"type": "barh", "x": "malware_family", "y": "count"}`)。图表始终适合返回的数据,而不是问题的措辞。
**为什么注入实时 schema 而不是将其硬编码?**
`schema.py` 会在运行时对数据库进行内省,并包含每个表的一行示例。这为 Agent 提供了准确的列名、数据类型和真实的值示例 — 无需进行任何硬编码的 prompt 维护。添加一个新表,Agent 就会自动识别它。
**为什么选择 DuckDB 而不是 SQLite?**
DuckDB 处理 GROUP BY 聚合和分析查询的速度明显快于 SQLite。它还原生读取 CSV,去除了单独的导入层。
## 接下来我会添加什么
- **Ollama 回退** — 检测是否设置了 `GOOGLE_API_KEY`,并回退到本地 Ollama 模型(例如 `llama3.1`)以实现完全离线使用
- **Hugging Face Spaces 部署** — 将 `GOOGLE_API_KEY` 添加为 Space 密钥并推送;应用无需更改即可运行
- **威胁时间线** — 将 CVE 发布日期和 OTX pulse 日期绘制在一起,以可视化攻击浪潮模式
- **MITRE ATT&CK 映射** — 根据 MITRE 框架对 OTX 攻击技术进行标记,以进行战术/技术分析
- **RAGAS 评估** — 使用答案相关性和忠实度指标进行自动化 Agent 评估
## 成本
本项目运行在 Gemini 的免费额度上 — 无需信用卡。
| 层级 | 请求/天 | 成本 |
|---|---|---|
| 免费(Google AI Studio) | 1,500 | $0 |
| 付费(如果需要) | 无限 | 每个问题约 ~$0.001–0.002 |
每个问题都会进行多次 API 调用(调用 1:Agent 决定使用哪个工具。
调用 2:Agent 生成 SQL 查询。
调用 3:Agent 处理 DuckDB 结果。
调用 4(可视化工具):你的 UI 会额外发起一次调用来决定显示哪个图表。)。一次用户交互很容易消耗你配额中的 3–5 个请求。我计划实现对免费本地模型(如 Ollama 或 qwen2.5)的回退支持。
## 数据归属
数据集由 [chuneeb](https://www.kaggle.com/datasets/chuneeb/ai-cybersecurity-threat-dataset-2026) 整理。
原始来源:AlienVault OTX、CISA Known Exploited Vulnerabilities Catalog、NVD — 均为公开和开放访问。
## 许可证
MIT — 查看 [LICENSE](LICENSE)。
## 作者
**Victor** — 数据科学家
[LinkedIn](https://linkedin.com/in/victor-kibisu) · [GitHub](https://github.com/Kibisuvictor)
标签:AI智能体, DuckDB, Gemini, LangGraph, Text-to-SQL, 代码示例, 威胁情报, 开发者工具, 数据分析, 逆向工具