mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-06-17 11:30:11 +02:00
Compare commits
104 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| c0f0861b31 | |||
| 704137aa34 | |||
| c56bf36df0 | |||
| 5560f34c6c | |||
| 70e9a73fc0 | |||
| 12bc9d8ab6 | |||
| f8db82a065 | |||
| 8ce30d9072 | |||
| e6506d00e8 | |||
| b2308617b8 | |||
| cd17fdca33 | |||
| 1acaccd09f | |||
| 983fe650c1 | |||
| 52d03dc849 | |||
| 9de72d9ad5 | |||
| d95275ffae | |||
| 6cef93dbb7 | |||
| dd3b1ae219 | |||
| f42209682a | |||
| 1b1aed1699 | |||
| 44ced98863 | |||
| 97834c162e | |||
| 9276f2f144 | |||
| a454cada6a | |||
| 99b53d4fbc | |||
| a43a9deaea | |||
| ce88da84c9 | |||
| 15855c7073 | |||
| 43eb3e546b | |||
| 2d52c9b6ac | |||
| d5401b8b4c | |||
| 5fd4393a2e | |||
| a049f6b5c2 | |||
| acba8e5a39 | |||
| f826b91362 | |||
| 98c2de2a60 | |||
| 1c4d4b305b | |||
| f210ac9a03 | |||
| 6685076dfb | |||
| 7f322653f6 | |||
| 66ac2f1357 | |||
| c446e22d0c | |||
| 0358d3a67d | |||
| 9b82f265fd | |||
| 3d9cae58e4 | |||
| 1f1eadee5e | |||
| 0569255189 | |||
| 8ccf90d067 | |||
| b3be89f47d | |||
| b9bf8f62d4 | |||
| 05ca0c1480 | |||
| 47a4f3fc5b | |||
| a3b378ae9e | |||
| a904d26e78 | |||
| 7ba7476c4f | |||
| ae25a243ac | |||
| 23bd6288ff | |||
| fef21d3a24 | |||
| 933bba4517 | |||
| e1d65437cc | |||
| 9325aed1eb | |||
| dee2b3ab42 | |||
| a69bc93fa1 | |||
| b1a620bfce | |||
| 61b164eec2 | |||
| ba77e1837e | |||
| eacad60fd6 | |||
| 70bf5c93bf | |||
| 08bd278d8c | |||
| 22746d64a3 | |||
| 199392a5d5 | |||
| aafb4cb584 | |||
| 96e3dd397c | |||
| ec0f17145b | |||
| ed53da0999 | |||
| dc440fc511 | |||
| 009ae59033 | |||
| f348b3245a | |||
| 0018c5219c | |||
| 01a3e3677a | |||
| a12ecdb46f | |||
| 9f59230d74 | |||
| 085c6a1c72 | |||
| 7b3860971f | |||
| f6f7b7b237 | |||
| d5cf4b3b16 | |||
| 3e58d8355b | |||
| eb01ade63b | |||
| d1dc15fa44 | |||
| 73a39ef868 | |||
| a022baef03 | |||
| 59312d428e | |||
| 951d14ef14 | |||
| 0eb22da6e9 | |||
| 5fd9ef0514 | |||
| 9a4f3c7d35 | |||
| ead2ce3ecc | |||
| 8733f3a2d2 | |||
| 8642f3ba31 | |||
| 6a262a7367 | |||
| eb9192ddb3 | |||
| 5587e75628 | |||
| 74bbb453e2 | |||
| 66842f6206 |
@@ -113,6 +113,7 @@ CyberStrikeAI is an **AI-native security testing platform** built in Go. It inte
|
||||
- 🔒 Password-protected web UI, audit logs, and SQLite persistence
|
||||
- 📚 Knowledge base (RAG) with embedding-based vector retrieval (cosine similarity), optional **Eino Compose** indexing pipeline, and configurable post-retrieval budgets / reranking hooks
|
||||
- 📁 Conversation grouping with pinning, rename, and batch management
|
||||
- 📂 **Project management**: group conversations and vulnerabilities by project; **shared facts** (project blackboard) persist cross-session context (targets, env, auth notes) with auto-injection for agents and MCP tools (`upsert_project_fact`, `get_project_fact`, …)
|
||||
- 🛡️ Vulnerability management with CRUD operations, severity tracking, status workflow, and statistics
|
||||
- 📋 Batch task management: create task queues, add multiple tasks, and execute them sequentially
|
||||
- 🎭 Role-based testing: predefined security testing roles (Penetration Testing, CTF, Web App Scanning, etc.) with custom prompts and tool restrictions
|
||||
@@ -285,7 +286,7 @@ Requirements / tips:
|
||||
- **Supervisor orchestrator**: fixed name **`orchestrator-supervisor.md`** (plus optional `orchestrator_instruction_supervisor`); requires at least one sub-agent.
|
||||
- **Sub-agents** (for **deep** / **supervisor**): other `*.md` files (YAML front matter + body). Not used as **`task`** targets if marked orchestrator-only.
|
||||
- **Management** – Web UI: **Agents → Agent management**; API `/api/multi-agent/markdown-agents`.
|
||||
- **Config** – `multi_agent` in `config.yaml`: `enabled`, `default_mode`, `robot_use_multi_agent`, `batch_use_multi_agent`, `max_iteration`, `plan_execute_loop_max_iterations`, per-mode orchestrator instruction fields, optional YAML `sub_agents` merged with disk (`id` clash → Markdown wins), **`eino_skills`**, **`eino_middleware`** (optional ADK middleware and Deep/Supervisor tuning).
|
||||
- **Config** – `multi_agent` in `config.yaml`: `enabled`, `robot_default_agent_mode`, `batch_use_multi_agent`, `max_iteration`, `plan_execute_loop_max_iterations`, per-mode orchestrator instruction fields, optional YAML `sub_agents` merged with disk (`id` clash → Markdown wins), **`eino_skills`**, **`eino_middleware`** (optional ADK middleware and Deep/Supervisor tuning).
|
||||
- **Details** – **[docs/MULTI_AGENT_EINO.md](docs/MULTI_AGENT_EINO.md)** (streaming, robots, batch, middleware caveats).
|
||||
|
||||
### Skills System (Agent Skills + Eino)
|
||||
@@ -536,7 +537,7 @@ agents_dir: "agents" # Multi-agent Markdown definitions (orchestrator + sub-age
|
||||
multi_agent:
|
||||
enabled: false
|
||||
default_mode: "single" # single | multi (UI default when multi-agent is enabled)
|
||||
robot_use_multi_agent: false
|
||||
robot_default_agent_mode: react
|
||||
batch_use_multi_agent: false
|
||||
orchestrator_instruction: "" # Deep; used when orchestrator.md body is empty
|
||||
# orchestrator_instruction_plan_execute / orchestrator_instruction_supervisor optional
|
||||
|
||||
+3
-2
@@ -112,6 +112,7 @@ CyberStrikeAI 是一款 **AI 原生安全测试平台**,基于 Go 构建,集
|
||||
- 🔒 Web 登录保护、审计日志、SQLite 持久化
|
||||
- 📚 知识库(RAG):向量嵌入与余弦相似度检索(与 Eino `retriever.Retriever` 语义一致),可选 **Eino Compose** 索引流水线及检索后处理(预算、重排等配置项)
|
||||
- 📁 对话分组管理:支持分组创建、置顶、重命名、删除等操作
|
||||
- 📂 **项目管理**:按项目归类对话与漏洞;**共享事实**(项目黑板)在多会话间沉淀目标/环境/认证等认知,自动注入 Agent 上下文,支持 MCP 工具读写(`upsert_project_fact`、`get_project_fact` 等)
|
||||
- 🛡️ 漏洞管理功能:完整的漏洞 CRUD 操作,支持严重程度分级、状态流转、按对话/严重程度/状态过滤,以及统计看板
|
||||
- 📋 批量任务管理:创建任务队列,批量添加任务,依次顺序执行,支持任务编辑与状态跟踪
|
||||
- 🎭 角色化测试:预设安全测试角色(渗透测试、CTF、Web 应用扫描等),支持自定义提示词和工具限制
|
||||
@@ -283,7 +284,7 @@ go build -o cyberstrike-ai cmd/server/main.go
|
||||
- **Supervisor 主代理**:固定 **`orchestrator-supervisor.md`**(另可配 `orchestrator_instruction_supervisor`);至少需一名子代理。
|
||||
- **子代理**(**deep** / **supervisor**):其余 `*.md`;标成 orchestrator 的不会进入 `task` 列表。
|
||||
- **界面管理**:**Agents → Agent 管理**;API `/api/multi-agent/markdown-agents`。
|
||||
- **配置项**:`multi_agent`:`enabled`、`default_mode`、`robot_use_multi_agent`、`batch_use_multi_agent`、`max_iteration`、`plan_execute_loop_max_iterations`、各模式 orchestrator 指令字段、可选 YAML `sub_agents` 与目录合并(同 `id` → Markdown 优先)、**`eino_skills`**、**`eino_middleware`**。
|
||||
- **配置项**:`multi_agent`:`enabled`、`robot_default_agent_mode`、`batch_use_multi_agent`、`max_iteration`、`plan_execute_loop_max_iterations`、各模式 orchestrator 指令字段、可选 YAML `sub_agents` 与目录合并(同 `id` → Markdown 优先)、**`eino_skills`**、**`eino_middleware`**。
|
||||
- **更多细节**:[docs/MULTI_AGENT_EINO.md](docs/MULTI_AGENT_EINO.md)(流式、机器人、批量、中间件差异)。
|
||||
|
||||
### Skills 技能系统(Agent Skills + Eino)
|
||||
@@ -534,7 +535,7 @@ agents_dir: "agents" # 多代理 Markdown(主代理 orchestrator.md + 子代
|
||||
multi_agent:
|
||||
enabled: false
|
||||
default_mode: "single" # single | multi(开启多代理时的界面默认模式)
|
||||
robot_use_multi_agent: false
|
||||
robot_default_agent_mode: react
|
||||
batch_use_multi_agent: false
|
||||
orchestrator_instruction: "" # Deep;orchestrator.md 正文为空时使用
|
||||
# orchestrator_instruction_plan_execute / orchestrator_instruction_supervisor 可选
|
||||
|
||||
@@ -127,7 +127,12 @@ description: 多代理模式下的 Deep 编排者:在已授权安全场景中
|
||||
## 工具与 MCP
|
||||
|
||||
- **工具调用失败时**:1) 仔细分析错误信息,理解失败的具体原因;2) 如果工具不存在或未启用,尝试使用其他替代工具完成相同目标;3) 如果参数错误,根据错误提示修正参数后重试;4) 如果工具执行失败但输出了有用信息,可以基于这些信息继续分析;5) 如果确实无法使用某个工具,向用户说明问题,并建议替代方案或手动操作;6) 不要因为单个工具失败就停止整个测试流程,尝试其他方法继续完成任务。工具返回的错误信息会包含在工具响应中,请仔细阅读并做出合理决策。
|
||||
- **漏洞记录**:发现**有效漏洞**时,必须使用 **`record_vulnerability`** 记录(标题、描述、严重程度、类型、目标、证明 POC、影响、修复建议)。严重程度使用 critical / high / medium / low / info。记录后可在授权范围内继续测试。
|
||||
- **项目黑板(事实)与漏洞记录(分离)**:当前对话若已绑定项目,系统会自动注入「项目黑板索引」(仅 `fact_key` + 摘要)。**摘要不足时必须调用 `get_project_fact(fact_key)` 获取 body,禁止凭摘要臆造细节。**
|
||||
- **环境/目标/认证等认知**(非正式漏洞):使用 **`upsert_project_fact`**,`fact_key` 建议 `category/slug`(如 `target/primary_domain`),同 key 覆盖更新;body 记端口/版本/凭据特征与证据来源。
|
||||
- **发现与利用上下文**(审计复现):`fact_key` 建议 `finding/`、`chain/`、`exploit/`、`poc/` 前缀;**body 必填**完整攻击链(入口 → 步骤 → 原始请求/响应或命令 → 现象 → 关联 `related_vulnerability_id`),**禁止仅写结论**;summary 写「什么 + 在哪 + 如何验证」一行要点。
|
||||
- **可交付漏洞**:使用 **`record_vulnerability`**(标题、描述、严重程度、类型、目标、证明 POC、影响、修复建议)。严重程度 critical / high / medium / low / info。
|
||||
- 同一发现可能需**各记一次**(事实记可复现攻击链,漏洞记正式 findings)。误报用 **`deprecate_project_fact`** 或漏洞状态 false_positive。
|
||||
- 事实多时用 **`list_project_facts`** / **`search_project_facts`** 检索。
|
||||
- **编排进度(待办)**:当你的任务包含 3 个或以上步骤,或你准备委派多个子目标并行/串行推进时,优先使用 `write_todos` 来向用户展示“当前在做什么/接下来做什么”。维护约束:同一时刻最多一个条目处于 `in_progress`;完成后立刻标记 `completed`;遇到阻塞就保留为 `in_progress` 并继续推进。
|
||||
- **强触发建议(提升多 agent 使用率)**:如果你将要进行任何“证据收集/枚举/扫描/验证/复现/整理报告”这类实质执行动作,且不只是单步查询,请优先在第一个工具调用前就用 `write_todos` 建立计划;随后用 `task` 委派至少一个子代理获取结构化证据,而不是自己把全部步骤做完。
|
||||
- **技能库(Skills)与知识库**:技能包位于服务器 `skills/` 目录(各子目录 `SKILL.md`,遵循 agentskills.io);知识库用于向量检索片段,Skills 为可执行工作流指令。多代理本会话通过内置 **`skill`** 工具渐进加载;子代理同样挂载 skill + 可选本机文件工具时,可在委派说明中提示按需加载。若当前无 skill 工具,需要完整 Skill 工作流时请使用多代理模式或切换为 Eino 编排会话。
|
||||
|
||||
+37
-6
@@ -10,7 +10,7 @@
|
||||
# ============================================
|
||||
|
||||
# 前端显示的版本号(可选,不填则显示默认版本)
|
||||
version: "v1.6.15"
|
||||
version: "v1.6.24"
|
||||
# 服务器配置
|
||||
server:
|
||||
host: 0.0.0.0 # 监听地址,0.0.0.0 表示监听所有网络接口
|
||||
@@ -34,6 +34,12 @@ auth:
|
||||
log:
|
||||
level: info # 日志级别: debug(调试), info(信息), warn(警告), error(错误)
|
||||
output: stdout # 日志输出位置: stdout(标准输出), stderr(标准错误), 或文件路径
|
||||
# 平台操作审计(系统设置 -> 日志审计;不记录对话正文与每次工具调用)
|
||||
audit:
|
||||
enabled: true
|
||||
retention_days: 15 # 0 表示不自动清理
|
||||
max_detail_bytes: 8192
|
||||
auth_failure_cooldown_seconds: 60 # 同一 IP 登录/改密失败审计最短间隔(秒);未配置时默认 60;-1 关闭节流
|
||||
# ============================================
|
||||
# 对话相关配置
|
||||
# ============================================
|
||||
@@ -54,8 +60,8 @@ openai:
|
||||
max_total_tokens: 120000 # LLM 相关上下文的最大 Token 数限制(内存压缩和攻击链构建会共用此配置)
|
||||
# Eino 路径模型推理:DeepSeek/OpenAI 为 thinking / reasoning_effort 等;provider 为 claude 时合并为 Anthropic 顶层 thinking(extended thinking),mode: off 关闭
|
||||
reasoning:
|
||||
mode: off # auto | on | off;off 时不附加任何推理扩展字段
|
||||
effort: max # low | medium | high | max;空表示不指定(openai_compat 下 auto 且无强度时不发请求扩展)
|
||||
mode: on # auto | on | off;off 时不附加任何推理扩展字段
|
||||
effort: high # low | medium | high | max | xhigh(最高档:OpenAI 常用 xhigh,部分网关用 max,原样下发);空表示不指定
|
||||
allow_client_reasoning: true # false 时忽略对话请求体 reasoning,仅以下方为准
|
||||
profile: openai_compat # auto | deepseek_compat | openai_compat | output_config_effort
|
||||
# extra_request_fields: {} # 可选:管理员自定义根级 JSON 片段(高级)
|
||||
@@ -76,16 +82,18 @@ agent:
|
||||
result_storage_dir: tmp # 结果存储目录,大结果会保存在此目录下
|
||||
tool_timeout_minutes: 60 # 单次工具执行最大时长(分钟),超时自动终止;0 表示不限制(不推荐,易出现长时间挂起)
|
||||
# system_prompt_path: prompts/single-react.md # 可选:单代理系统提示文件(相对本配置文件所在目录);非空且可读时替换内置提示
|
||||
|
||||
system_prompt_path: ""
|
||||
# 人机协同(HITL)全局白名单:此处列出的工具始终免审批,与对话页「白名单工具(免审批,逗号分隔)」合并为并集;侧栏「应用」可合并写入本列表并立即生效。
|
||||
hitl:
|
||||
# 按你环境里的真实工具名增删(与侧栏一致、小写不敏感);不需要全局免审批可改为 []
|
||||
tool_whitelist: [read_file, list_dir, glob, grep]
|
||||
# 多代理(CloudWeGo Eino DeepAgent,与上方单 Agent /api/agent-loop 并存)
|
||||
# 依赖在 go.mod 中拉取;若下载失败可设置: go env -w GOPROXY=https://goproxy.cn,direct
|
||||
# 启用后需重启服务才会注册 /api/multi-agent 与 /api/multi-agent/stream;Deep / Plan-Execute / Supervisor 由对话页与 WebShell 所选模式在请求体中传入;机器人/批量无请求体时固定按 deep
|
||||
# 启用后需重启服务才会注册 /api/multi-agent 与 /api/multi-agent/stream;Deep / Plan-Execute / Supervisor 由对话页与 WebShell 所选模式在请求体中传入;机器人按 robot_default_agent_mode
|
||||
multi_agent:
|
||||
enabled: true
|
||||
robot_use_multi_agent: true # true 时企业微信/钉钉/飞书机器人也走 Eino 多代理(成本更高)
|
||||
robot_default_agent_mode: eino_single # 企微/钉钉/飞书机器人默认对话模式:react | eino_single | deep | plan_execute | supervisor
|
||||
batch_use_multi_agent: false # true 时「批量任务」队列中每个子任务也走 Eino 多代理(成本更高)
|
||||
max_iteration: 0 # 主代理 / plan_execute 执行器最大轮次,0 表示沿用 agent.max_iterations
|
||||
# plan_execute 专用:execute↔replan 外层循环上限,0 表示 Eino 默认 10。当前实现下 Executor 会挂载 patch/reduction/tool_search 等前置中间件。
|
||||
@@ -108,7 +116,7 @@ multi_agent:
|
||||
tool_search_enable: true # true:工具数 ≥ min 时启用 tool_search,仅前 N 个工具常驻,其余按正则按需解锁,省 token、减误选;false:全量工具进上下文
|
||||
tool_search_min_tools: 20 # 达到该数量才启用 tool_search(避免工具很少时多此一举);与 always_visible 配合使用
|
||||
tool_search_always_visible: 12 # 始终直接暴露给模型的工具个数(顺序与角色工具列表一致);其余工具进入动态池,需 tool_search 解锁
|
||||
tool_search_always_visible_tools: [read_file, glob, grep, write_file, edit_file, execute, task, transfer_to_agent, exit, write_todos, skill, tool_search, TaskCreate, TaskGet, TaskUpdate, TaskList, record_vulnerability, list_knowledge_risk_types, search_knowledge_base, webshell_exec, webshell_file_list, webshell_file_read, webshell_file_write, manage_webshell_list, manage_webshell_add, manage_webshell_update, manage_webshell_delete, manage_webshell_test, batch_task_list, batch_task_get, batch_task_start, batch_task_rerun, batch_task_pause, batch_task_update_metadata, batch_task_update_schedule, batch_task_schedule_enabled, batch_task_update_task, batch_task_remove_task, batch_task_delete, batch_task_create, batch_task_add_task, http-framework-test] # 后端内置常驻工具白名单(优先于 always_visible 数量策略)
|
||||
tool_search_always_visible_tools: [read_file, glob, grep, write_file, edit_file, execute, task, transfer_to_agent, exit, write_todos, skill, tool_search, TaskCreate, TaskGet, TaskUpdate, TaskList, record_vulnerability, list_vulnerabilities, get_vulnerability, list_knowledge_risk_types, search_knowledge_base, webshell_exec, webshell_file_list, webshell_file_read, webshell_file_write, manage_webshell_list, manage_webshell_add, manage_webshell_update, manage_webshell_delete, manage_webshell_test, batch_task_list, batch_task_get, batch_task_start, batch_task_rerun, batch_task_pause, batch_task_update_metadata, batch_task_update_schedule, batch_task_schedule_enabled, batch_task_update_task, batch_task_remove_task, batch_task_delete, batch_task_create, batch_task_add_task, http-framework-test] # 后端内置常驻工具白名单(优先于 always_visible 数量策略)
|
||||
plantask_enable: false # true:主代理(Deep / Supervisor 主)挂载 TaskCreate/Get/Update/List;需 eino_skills 可用且 skills_dir 存在,否则仅打日志并跳过
|
||||
plantask_rel_dir: .eino/plantask # 结构化任务文件相对 skills_dir 的子目录,其下再按会话 ID 分子目录存放
|
||||
reduction_enable: true # true:大工具输出截断/落盘以控上下文;依赖与 plantask 相同的 eino local 写盘后端,无后端时不挂载
|
||||
@@ -125,6 +133,8 @@ multi_agent:
|
||||
plan_execute_max_step_result_runes: 4000 # plan_execute 每步结果最大字符数(超出截断)
|
||||
plan_execute_keep_last_steps: 8 # plan_execute 仅保留最近 N 步正文,早期步骤折叠为标题
|
||||
checkpoint_dir: "" # 非空:为 adk.NewRunner 启用按会话子目录的文件型 CheckPointStore,便于中断恢复持久化;Resume 的 HTTP/前端流程需另行对接
|
||||
run_retry_max_attempts: 0 # >0:429/5xx/网络抖动时 ADK 运行循环指数退避续跑次数;0=默认 10
|
||||
run_retry_max_backoff_sec: 0 # 单次退避上限秒数;0=默认 30
|
||||
deep_output_key: "" # 非空:将最终助手输出写入 adk session 的键名(Deep 与 Supervisor 主代理);空表示不写入
|
||||
deep_model_retry_max_retries: 0 # >0:ChatModel 调用失败时的框架级最大重试次数(Deep 与 Supervisor 主);0:不重试
|
||||
task_tool_description_prefix: "" # 非空:仅 Deep 的 task 工具使用自定义描述前缀,运行时会拼接子代理名称;空则走 Eino 默认生成逻辑
|
||||
@@ -235,6 +245,14 @@ knowledge:
|
||||
# 用于在手机端通过企业微信/钉钉/飞书与 CyberStrikeAI 对话,无需部署在服务器上也可使用
|
||||
# 在系统设置 -> 机器人设置 中可配置
|
||||
robots:
|
||||
wechat: # 微信 iLink(个人微信 ClawBot,扫码绑定)
|
||||
enabled: false
|
||||
bot_token: ""
|
||||
ilink_bot_id: ""
|
||||
ilink_user_id: ""
|
||||
base_url: https://ilinkai.weixin.qq.com
|
||||
bot_type: "3"
|
||||
bot_agent: CyberStrikeAI/1.0
|
||||
wecom: # 企业微信
|
||||
enabled: false
|
||||
token: ""
|
||||
@@ -246,11 +264,13 @@ robots:
|
||||
enabled: false
|
||||
client_id: ""
|
||||
client_secret: ""
|
||||
allow_conversation_id_fallback: false
|
||||
lark: # 飞书
|
||||
enabled: false
|
||||
app_id: ""
|
||||
app_secret: ""
|
||||
verify_token: ""
|
||||
allow_chat_id_fallback: false
|
||||
# ============================================
|
||||
# Skills 相关配置
|
||||
# ============================================
|
||||
@@ -272,3 +292,14 @@ agents_dir: agents
|
||||
# 系统会从该目录加载所有 .yaml 格式的角色配置文件
|
||||
# 每个角色应创建独立的配置文件,例如:roles/CTF.yaml, roles/默认.yaml 等
|
||||
roles_dir: roles # 角色配置文件目录(相对于配置文件所在目录)
|
||||
|
||||
# ============================================
|
||||
# 项目管理与事实黑板
|
||||
# ============================================
|
||||
project:
|
||||
enabled: true
|
||||
# default_project_id: "" # 可选:机器人/批量任务创建对话时的默认项目 ID
|
||||
fact_index_max_runes: 3500
|
||||
fact_summary_max_runes: 240
|
||||
default_inject_deprecated: false
|
||||
|
||||
|
||||
@@ -27,12 +27,14 @@ require (
|
||||
github.com/open-dingtalk/dingtalk-stream-sdk-go v0.9.1
|
||||
github.com/pkoukk/tiktoken-go v0.1.8
|
||||
github.com/robfig/cron/v3 v3.0.1
|
||||
github.com/skip2/go-qrcode v0.0.0-20200617195104-da1b6568686e
|
||||
go.opentelemetry.io/otel v1.34.0
|
||||
go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.34.0
|
||||
go.opentelemetry.io/otel/exporters/stdout/stdouttrace v1.34.0
|
||||
go.opentelemetry.io/otel/sdk v1.34.0
|
||||
go.opentelemetry.io/otel/trace v1.34.0
|
||||
go.uber.org/zap v1.26.0
|
||||
golang.org/x/net v0.35.0
|
||||
golang.org/x/text v0.26.0
|
||||
golang.org/x/time v0.14.0
|
||||
gopkg.in/yaml.v3 v3.0.1
|
||||
@@ -88,7 +90,6 @@ require (
|
||||
golang.org/x/arch v0.15.0 // indirect
|
||||
golang.org/x/crypto v0.39.0 // indirect
|
||||
golang.org/x/exp v0.0.0-20250305212735-054e65f0b394 // indirect
|
||||
golang.org/x/net v0.35.0 // indirect
|
||||
golang.org/x/oauth2 v0.30.0 // indirect
|
||||
golang.org/x/sys v0.33.0 // indirect
|
||||
google.golang.org/genproto/googleapis/api v0.0.0-20250115164207-1a7da9e5054f // indirect
|
||||
|
||||
@@ -163,6 +163,8 @@ github.com/rollbar/rollbar-go v1.0.2/go.mod h1:AcFs5f0I+c71bpHlXNNDbOWJiKwjFDtIS
|
||||
github.com/sirupsen/logrus v1.2.0/go.mod h1:LxeOpSwHxABJmUn/MG1IvRgCAasNZTLOkJPxbbu5VWo=
|
||||
github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=
|
||||
github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
|
||||
github.com/skip2/go-qrcode v0.0.0-20200617195104-da1b6568686e h1:MRM5ITcdelLK2j1vwZ3Je0FKVCfqOLp5zO6trqMLYs0=
|
||||
github.com/skip2/go-qrcode v0.0.0-20200617195104-da1b6568686e/go.mod h1:XV66xRDqSt+GTGFMVlhk3ULuV0y9ZmzeVGR4mloJI3M=
|
||||
github.com/slongfield/pyfmt v0.0.0-20220222012616-ea85ff4c361f h1:Z2cODYsUxQPofhpYRMQVwWz4yUVpHF+vPi+eUdruUYI=
|
||||
github.com/slongfield/pyfmt v0.0.0-20220222012616-ea85ff4c361f/go.mod h1:JqzWyvTuI2X4+9wOHmKSQCYxybB/8j6Ko43qVmXDuZg=
|
||||
github.com/smarty/assertions v1.16.0 h1:EvHNkdRA4QHMrn75NZSoUQ/mAUXAYWfatfB01yTCzfY=
|
||||
@@ -245,8 +247,6 @@ golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn
|
||||
golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
|
||||
golang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
|
||||
golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
|
||||
golang.org/x/net v0.34.0 h1:Mb7Mrk043xzHgnRM88suvJFwzVrRfHEHJEl5/71CKw0=
|
||||
golang.org/x/net v0.34.0/go.mod h1:di0qlW3YNM5oh6GqDGQr92MyTozJPmybPK4Ev/Gm31k=
|
||||
golang.org/x/net v0.35.0 h1:T5GQRQb2y08kTAByq9L4/bz8cipCdA8FbRTXewonqY8=
|
||||
golang.org/x/net v0.35.0/go.mod h1:EglIi67kWsHKlRzzVMUD93VMSWGFOMSZgxFjparz1Qk=
|
||||
golang.org/x/oauth2 v0.30.0 h1:dnDm7JmhM45NNpd8FDDeLhK6FwqbOf4MLCM9zb1BOHI=
|
||||
|
||||
+37
-11
@@ -17,6 +17,7 @@ import (
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
"cyberstrike-ai/internal/project"
|
||||
"cyberstrike-ai/internal/openai"
|
||||
"cyberstrike-ai/internal/security"
|
||||
"cyberstrike-ai/internal/storage"
|
||||
@@ -365,12 +366,12 @@ type ProgressCallback func(eventType, message string, data interface{})
|
||||
|
||||
// AgentLoop 执行Agent循环
|
||||
func (a *Agent) AgentLoop(ctx context.Context, userInput string, historyMessages []ChatMessage) (*AgentLoopResult, error) {
|
||||
return a.AgentLoopWithProgress(ctx, userInput, historyMessages, "", nil, nil)
|
||||
return a.AgentLoopWithProgress(ctx, userInput, historyMessages, "", nil, nil, "")
|
||||
}
|
||||
|
||||
// AgentLoopWithConversationID 执行Agent循环(带对话ID)
|
||||
func (a *Agent) AgentLoopWithConversationID(ctx context.Context, userInput string, historyMessages []ChatMessage, conversationID string) (*AgentLoopResult, error) {
|
||||
return a.AgentLoopWithProgress(ctx, userInput, historyMessages, conversationID, nil, nil)
|
||||
return a.AgentLoopWithProgress(ctx, userInput, historyMessages, conversationID, nil, nil, "")
|
||||
}
|
||||
|
||||
// EinoSingleAgentSystemInstruction 供 Eino adk.ChatModelAgent.Instruction 使用,与 AgentLoopWithProgress 首条 system 对齐(含 system_prompt_path)。
|
||||
@@ -396,7 +397,7 @@ func (a *Agent) EinoSingleAgentSystemInstruction() string {
|
||||
}
|
||||
|
||||
// AgentLoopWithProgress 执行Agent循环(带进度回调和对话ID)
|
||||
func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, historyMessages []ChatMessage, conversationID string, callback ProgressCallback, roleTools []string) (*AgentLoopResult, error) {
|
||||
func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, historyMessages []ChatMessage, conversationID string, callback ProgressCallback, roleTools []string, systemPromptExtra string) (*AgentLoopResult, error) {
|
||||
ctx = withAgentConversationID(ctx, conversationID)
|
||||
// 设置当前对话ID(兼容未走 context 的旧路径;并发会话应以 context 为准)
|
||||
a.mu.Lock()
|
||||
@@ -426,6 +427,7 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
|
||||
}
|
||||
}
|
||||
}
|
||||
systemPrompt = project.AppendSystemPromptBlock(systemPrompt, systemPromptExtra)
|
||||
|
||||
messages := []ChatMessage{
|
||||
{
|
||||
@@ -598,11 +600,17 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
|
||||
thinkingStreamSeq++
|
||||
thinkingStreamId := fmt.Sprintf("thinking-stream-%s-%d-%d", conversationID, i+1, thinkingStreamSeq)
|
||||
thinkingStreamStarted := false
|
||||
var thinkingWire string
|
||||
|
||||
response, err := a.callOpenAIStreamWithToolCalls(ctx, messages, tools, func(delta string) error {
|
||||
if delta == "" {
|
||||
return nil
|
||||
}
|
||||
var deltaOut string
|
||||
thinkingWire, deltaOut = openai.NormalizeStreamingDelta(thinkingWire, delta)
|
||||
if deltaOut == "" {
|
||||
return nil
|
||||
}
|
||||
if !thinkingStreamStarted {
|
||||
thinkingStreamStarted = true
|
||||
sendProgress("thinking_stream_start", " ", map[string]interface{}{
|
||||
@@ -611,10 +619,10 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
|
||||
"toolStream": false,
|
||||
})
|
||||
}
|
||||
sendProgress("thinking_stream_delta", delta, map[string]interface{}{
|
||||
sendProgress("thinking_stream_delta", deltaOut, openai.WithSSEAccumulated(map[string]interface{}{
|
||||
"streamId": thinkingStreamId,
|
||||
"iteration": i + 1,
|
||||
})
|
||||
}, thinkingWire))
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
@@ -827,10 +835,16 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
|
||||
"mcpExecutionIds": result.MCPExecutionIDs,
|
||||
"messageGeneratedBy": "summary",
|
||||
})
|
||||
var summaryWire string
|
||||
streamText, _ := a.callOpenAIStreamText(ctx, messages, []Tool{}, func(delta string) error {
|
||||
sendProgress("response_delta", delta, map[string]interface{}{
|
||||
var deltaOut string
|
||||
summaryWire, deltaOut = openai.NormalizeStreamingDelta(summaryWire, delta)
|
||||
if deltaOut == "" {
|
||||
return nil
|
||||
}
|
||||
sendProgress("response_delta", deltaOut, openai.WithSSEAccumulated(map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
})
|
||||
}, summaryWire))
|
||||
return nil
|
||||
})
|
||||
if strings.TrimSpace(streamText) != "" {
|
||||
@@ -874,10 +888,16 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
|
||||
"mcpExecutionIds": result.MCPExecutionIDs,
|
||||
"messageGeneratedBy": "summary",
|
||||
})
|
||||
var summaryWire string
|
||||
streamText, _ := a.callOpenAIStreamText(ctx, messages, []Tool{}, func(delta string) error {
|
||||
sendProgress("response_delta", delta, map[string]interface{}{
|
||||
var deltaOut string
|
||||
summaryWire, deltaOut = openai.NormalizeStreamingDelta(summaryWire, delta)
|
||||
if deltaOut == "" {
|
||||
return nil
|
||||
}
|
||||
sendProgress("response_delta", deltaOut, openai.WithSSEAccumulated(map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
})
|
||||
}, summaryWire))
|
||||
return nil
|
||||
})
|
||||
if strings.TrimSpace(streamText) != "" {
|
||||
@@ -921,10 +941,16 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
|
||||
"mcpExecutionIds": result.MCPExecutionIDs,
|
||||
"messageGeneratedBy": "max_iter_summary",
|
||||
})
|
||||
var summaryWire string
|
||||
streamText, _ := a.callOpenAIStreamText(ctx, messages, []Tool{}, func(delta string) error {
|
||||
sendProgress("response_delta", delta, map[string]interface{}{
|
||||
var deltaOut string
|
||||
summaryWire, deltaOut = openai.NormalizeStreamingDelta(summaryWire, delta)
|
||||
if deltaOut == "" {
|
||||
return nil
|
||||
}
|
||||
sendProgress("response_delta", deltaOut, openai.WithSSEAccumulated(map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
})
|
||||
}, summaryWire))
|
||||
return nil
|
||||
})
|
||||
if strings.TrimSpace(streamText) != "" {
|
||||
|
||||
@@ -0,0 +1,167 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ParseTraceMessages 解析落库的 last_react_input(OpenAI 风格 messages JSON 数组)。
|
||||
func ParseTraceMessages(traceInputJSON string) ([]ChatMessage, error) {
|
||||
traceInputJSON = strings.TrimSpace(traceInputJSON)
|
||||
if traceInputJSON == "" {
|
||||
return nil, nil
|
||||
}
|
||||
var raw []map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(traceInputJSON), &raw); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make([]ChatMessage, 0, len(raw))
|
||||
for _, msgMap := range raw {
|
||||
msg := ChatMessage{}
|
||||
role, _ := msgMap["role"].(string)
|
||||
if role == "" {
|
||||
continue
|
||||
}
|
||||
msg.Role = role
|
||||
if content, ok := msgMap["content"].(string); ok {
|
||||
msg.Content = content
|
||||
}
|
||||
if rc, ok := msgMap["reasoning_content"].(string); ok && strings.TrimSpace(rc) != "" {
|
||||
msg.ReasoningContent = rc
|
||||
}
|
||||
if toolCallsRaw, ok := msgMap["tool_calls"]; ok && toolCallsRaw != nil {
|
||||
if toolCallsArray, ok := toolCallsRaw.([]interface{}); ok {
|
||||
for _, tcRaw := range toolCallsArray {
|
||||
tcMap, ok := tcRaw.(map[string]interface{})
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
toolCall := ToolCall{}
|
||||
if id, ok := tcMap["id"].(string); ok {
|
||||
toolCall.ID = id
|
||||
}
|
||||
if toolType, ok := tcMap["type"].(string); ok {
|
||||
toolCall.Type = toolType
|
||||
}
|
||||
if funcMap, ok := tcMap["function"].(map[string]interface{}); ok {
|
||||
toolCall.Function = FunctionCall{}
|
||||
if name, ok := funcMap["name"].(string); ok {
|
||||
toolCall.Function.Name = name
|
||||
}
|
||||
if argsRaw, ok := funcMap["arguments"]; ok {
|
||||
if argsStr, ok := argsRaw.(string); ok {
|
||||
var argsMap map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(argsStr), &argsMap); err == nil {
|
||||
toolCall.Function.Arguments = argsMap
|
||||
}
|
||||
} else if argsMap, ok := argsRaw.(map[string]interface{}); ok {
|
||||
toolCall.Function.Arguments = argsMap
|
||||
}
|
||||
}
|
||||
}
|
||||
if toolCall.ID != "" {
|
||||
msg.ToolCalls = append(msg.ToolCalls, toolCall)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if toolCallID, ok := msgMap["tool_call_id"].(string); ok {
|
||||
msg.ToolCallID = toolCallID
|
||||
}
|
||||
if tn, ok := msgMap["tool_name"].(string); ok && strings.TrimSpace(tn) != "" {
|
||||
msg.ToolName = strings.TrimSpace(tn)
|
||||
} else if tn, ok := msgMap["name"].(string); ok && strings.TrimSpace(tn) != "" && strings.EqualFold(msg.Role, "tool") {
|
||||
msg.ToolName = strings.TrimSpace(tn)
|
||||
}
|
||||
out = append(out, msg)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// ExtractLastUserTurnMessages 仅保留最后一次 user 提问起的消息(不含更早的用户轮次;跳过 system)。
|
||||
// 与「继续对话」续跑所用轨迹范围一致:当前任务轮次,而非整段多轮对话历史。
|
||||
func ExtractLastUserTurnMessages(msgs []ChatMessage) []ChatMessage {
|
||||
if len(msgs) == 0 {
|
||||
return msgs
|
||||
}
|
||||
lastUser := -1
|
||||
for i, m := range msgs {
|
||||
if strings.EqualFold(m.Role, "user") {
|
||||
lastUser = i
|
||||
}
|
||||
}
|
||||
if lastUser < 0 {
|
||||
return msgs
|
||||
}
|
||||
trimmed := msgs[lastUser:]
|
||||
out := make([]ChatMessage, 0, len(trimmed))
|
||||
for _, m := range trimmed {
|
||||
if strings.EqualFold(m.Role, "system") {
|
||||
continue
|
||||
}
|
||||
out = append(out, m)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// ExtractLastUserTurnTraceJSON 在 JSON 轨迹上裁剪为最后一次 user 起的片段(供落库格式直接处理)。
|
||||
func ExtractLastUserTurnTraceJSON(traceInputJSON string) string {
|
||||
traceInputJSON = strings.TrimSpace(traceInputJSON)
|
||||
if traceInputJSON == "" {
|
||||
return traceInputJSON
|
||||
}
|
||||
var arr []map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(traceInputJSON), &arr); err != nil {
|
||||
return traceInputJSON
|
||||
}
|
||||
lastUser := -1
|
||||
for i, m := range arr {
|
||||
if r, _ := m["role"].(string); strings.EqualFold(r, "user") {
|
||||
lastUser = i
|
||||
}
|
||||
}
|
||||
if lastUser <= 0 {
|
||||
return traceInputJSON
|
||||
}
|
||||
trimmed := arr[lastUser:]
|
||||
b, err := json.Marshal(trimmed)
|
||||
if err != nil {
|
||||
return traceInputJSON
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// MergeAssistantTraceOutput 将 last_react_output 合并进轨迹最后一条 assistant(与 loadHistoryFromAgentTrace 一致)。
|
||||
func MergeAssistantTraceOutput(msgs []ChatMessage, assistantOut string) []ChatMessage {
|
||||
assistantOut = strings.TrimSpace(assistantOut)
|
||||
if assistantOut == "" || len(msgs) == 0 {
|
||||
return msgs
|
||||
}
|
||||
out := append([]ChatMessage(nil), msgs...)
|
||||
last := &out[len(out)-1]
|
||||
if strings.EqualFold(last.Role, "assistant") && len(last.ToolCalls) == 0 {
|
||||
last.Content = assistantOut
|
||||
return out
|
||||
}
|
||||
out = append(out, ChatMessage{
|
||||
Role: "assistant",
|
||||
Content: assistantOut,
|
||||
})
|
||||
return out
|
||||
}
|
||||
|
||||
// MessagesToTraceJSON 将消息带序列化为 JSON(跳过 system)。
|
||||
func MessagesToTraceJSON(msgs []ChatMessage) (string, error) {
|
||||
filtered := make([]ChatMessage, 0, len(msgs))
|
||||
for _, m := range msgs {
|
||||
if strings.EqualFold(m.Role, "system") {
|
||||
continue
|
||||
}
|
||||
filtered = append(filtered, m)
|
||||
}
|
||||
b, err := json.Marshal(filtered)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return string(b), nil
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestExtractLastUserTurnTraceJSON(t *testing.T) {
|
||||
raw := []map[string]interface{}{
|
||||
{"role": "user", "content": "old question"},
|
||||
{"role": "assistant", "content": "old answer"},
|
||||
{"role": "user", "content": "new target 1.1.1.1"},
|
||||
{"role": "assistant", "tool_calls": []interface{}{map[string]interface{}{
|
||||
"id": "c1", "type": "function",
|
||||
"function": map[string]interface{}{"name": "nmap", "arguments": "{}"},
|
||||
}}},
|
||||
{"role": "tool", "tool_call_id": "c1", "content": "open ports"},
|
||||
}
|
||||
b, _ := json.Marshal(raw)
|
||||
out := ExtractLastUserTurnTraceJSON(string(b))
|
||||
var trimmed []map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(out), &trimmed); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(trimmed) != 3 {
|
||||
t.Fatalf("expected 3 messages, got %d", len(trimmed))
|
||||
}
|
||||
if trimmed[0]["content"] != "new target 1.1.1.1" {
|
||||
t.Fatalf("unexpected first message: %v", trimmed[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractLastUserTurnMessagesSkipsSystem(t *testing.T) {
|
||||
msgs := []ChatMessage{
|
||||
{Role: "system", Content: "sys"},
|
||||
{Role: "user", Content: "q"},
|
||||
{Role: "assistant", Content: "a"},
|
||||
}
|
||||
out := ExtractLastUserTurnMessages(msgs)
|
||||
if len(out) != 2 {
|
||||
t.Fatalf("expected 2, got %d", len(out))
|
||||
}
|
||||
if out[0].Role != "user" {
|
||||
t.Fatal("expected user first")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMergeAssistantTraceOutput(t *testing.T) {
|
||||
msgs := []ChatMessage{
|
||||
{Role: "user", Content: "q"},
|
||||
{Role: "assistant", Content: "draft"},
|
||||
}
|
||||
out := MergeAssistantTraceOutput(msgs, "final summary")
|
||||
if out[len(out)-1].Content != "final summary" {
|
||||
t.Fatalf("expected merged output, got %q", out[len(out)-1].Content)
|
||||
}
|
||||
}
|
||||
@@ -1,6 +1,9 @@
|
||||
package agent
|
||||
|
||||
import "cyberstrike-ai/internal/mcp/builtin"
|
||||
import (
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
"cyberstrike-ai/internal/project"
|
||||
)
|
||||
|
||||
// DefaultSingleAgentSystemPrompt 单代理(ReAct / MCP)内置系统提示;可通过 agent.system_prompt_path 覆盖为文件。
|
||||
func DefaultSingleAgentSystemPrompt() string {
|
||||
@@ -105,11 +108,17 @@ func DefaultSingleAgentSystemPrompt() string {
|
||||
- 若最近一步得到 404/空结果/无效响应,不得直接结束;至少再进行一次“同目标不同策略”的验证(如变更路径、参数、请求方法、上下文来源)。
|
||||
- 避免无效空转:同一工具+同类参数连续失败 3 次后,必须切换策略(改工具、改入口、改假设)并说明切换原因。
|
||||
|
||||
## 漏洞记录
|
||||
## 项目黑板(事实)与漏洞记录(分离)
|
||||
|
||||
发现有效漏洞时,必须使用 ` + builtin.ToolRecordVulnerability + ` 记录:标题、描述、严重程度、类型、目标、证明(POC)、影响、修复建议。
|
||||
当前对话若已绑定项目,系统会自动注入「项目黑板索引」(仅 fact_key + 摘要)。**摘要不足时必须调用 ` + builtin.ToolGetProjectFact + `(fact_key) 获取 body,禁止凭摘要臆造细节。**
|
||||
|
||||
严重程度:critical / high / medium / low / info。证明须含足够证据(请求响应、截图、命令输出等)。记录后可在授权范围内继续测试。
|
||||
- **环境/目标/认证等认知**(非正式漏洞条目):使用 ` + builtin.ToolUpsertProjectFact + `,fact_key 建议 ` + "`category/slug`" + `(如 target/primary_domain),同 key 覆盖更新。
|
||||
- **可交付漏洞**:使用 ` + builtin.ToolRecordVulnerability + `,含标题、严重程度、类型、目标、证明(POC)、影响、修复建议。记前可先 ` + builtin.ToolListVulnerabilities + ` 查重,详情用 ` + builtin.ToolGetVulnerability + `(id)(默认仅当前项目/会话)。
|
||||
- 同一发现可能需**各记一次**(事实记**完整攻击链与 exploit 细节**供复现,漏洞记正式 findings)。误报用 ` + builtin.ToolDeprecateProjectFact + ` 或漏洞状态 false_positive。
|
||||
|
||||
` + project.FactRecordingGuidanceBlock() + `
|
||||
|
||||
严重程度:critical / high / medium / low / info。证明须含足够证据(请求响应、截图、命令输出等)。
|
||||
|
||||
## 技能库(Skills)与知识库
|
||||
|
||||
|
||||
+78
-193
@@ -15,6 +15,7 @@ import (
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/c2"
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
@@ -56,10 +57,12 @@ type App struct {
|
||||
robotMu sync.Mutex // 保护钉钉/飞书长连接的 cancel
|
||||
dingCancel context.CancelFunc // 钉钉 Stream 取消函数,用于配置变更时重启
|
||||
larkCancel context.CancelFunc // 飞书长连接取消函数,用于配置变更时重启
|
||||
wechatCancel context.CancelFunc // 微信 iLink 长轮询取消函数
|
||||
c2Manager *c2.Manager // C2 管理器(未启用 C2 时为 nil)
|
||||
c2Watchdog *c2.SessionWatchdog // C2 会话看门狗
|
||||
c2WatchdogCancel context.CancelFunc // 看门狗取消函数
|
||||
c2Handler *handler.C2Handler // C2 REST(与 Manager 生命周期同步)
|
||||
auditSvc *audit.Service
|
||||
}
|
||||
|
||||
// New 创建新应用
|
||||
@@ -92,6 +95,11 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
||||
return nil, fmt.Errorf("初始化数据库失败: %w", err)
|
||||
}
|
||||
|
||||
auditSvc := audit.NewService(db, cfg, log.Logger)
|
||||
audit.RegisterConversationCreateHook(auditSvc)
|
||||
auditSvc.PurgeExpired()
|
||||
audit.StartRetentionLoop(auditSvc, log.Logger)
|
||||
|
||||
// 创建MCP服务器(带数据库持久化)
|
||||
mcpServer := mcp.NewServerWithStorage(log.Logger, db)
|
||||
mcpServer.ConfigureHTTPToolCallTimeoutFromAgentMinutes(cfg.Agent.ToolTimeoutMinutes)
|
||||
@@ -103,7 +111,8 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
||||
executor.RegisterTools(mcpServer)
|
||||
|
||||
// 注册漏洞记录工具
|
||||
registerVulnerabilityTool(mcpServer, db, log.Logger)
|
||||
registerVulnerabilityTools(mcpServer, db, log.Logger)
|
||||
registerProjectFactTools(mcpServer, db, cfg, log.Logger)
|
||||
|
||||
if cfg.Auth.GeneratedPassword != "" {
|
||||
config.PrintGeneratedPasswordWarning(cfg.Auth.GeneratedPassword, cfg.Auth.GeneratedPasswordPersisted, cfg.Auth.GeneratedPasswordPersistErr)
|
||||
@@ -221,6 +230,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
||||
|
||||
// 创建知识库API处理器
|
||||
knowledgeHandler = handler.NewKnowledgeHandler(knowledgeManager, knowledgeRetriever, knowledgeIndexer, db, log.Logger)
|
||||
knowledgeHandler.SetAudit(auditSvc)
|
||||
log.Logger.Info("知识库模块初始化完成", zap.Bool("handler_created", knowledgeHandler != nil))
|
||||
|
||||
// 扫描知识库并建立索引(异步)
|
||||
@@ -317,31 +327,43 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
||||
log.Logger.Warn("创建 agents 目录失败", zap.String("path", agentsDir), zap.Error(err))
|
||||
}
|
||||
markdownAgentsHandler := handler.NewMarkdownAgentsHandler(agentsDir)
|
||||
markdownAgentsHandler.SetAudit(auditSvc)
|
||||
log.Logger.Info("多代理 Markdown 子 Agent 目录", zap.String("agentsDir", agentsDir))
|
||||
|
||||
// 创建处理器
|
||||
agentHandler := handler.NewAgentHandler(agent, db, cfg, log.Logger)
|
||||
agentHandler.SetAudit(auditSvc)
|
||||
agentHandler.SetAgentsMarkdownDir(agentsDir)
|
||||
// 如果知识库已启用,设置知识库管理器到AgentHandler以便记录检索日志
|
||||
if knowledgeManager != nil {
|
||||
agentHandler.SetKnowledgeManager(knowledgeManager)
|
||||
}
|
||||
monitorHandler := handler.NewMonitorHandler(mcpServer, executor, db, log.Logger)
|
||||
monitorHandler.SetAudit(auditSvc)
|
||||
monitorHandler.SetExternalMCPManager(externalMCPMgr) // 设置外部MCP管理器,以便获取外部MCP执行记录
|
||||
notificationHandler := handler.NewNotificationHandler(db, agentHandler, log.Logger)
|
||||
groupHandler := handler.NewGroupHandler(db, log.Logger)
|
||||
authHandler := handler.NewAuthHandler(authManager, cfg, configPath, log.Logger)
|
||||
authHandler.SetAudit(auditSvc)
|
||||
attackChainHandler := handler.NewAttackChainHandler(db, &cfg.OpenAI, log.Logger)
|
||||
vulnerabilityHandler := handler.NewVulnerabilityHandler(db, log.Logger)
|
||||
projectHandler := handler.NewProjectHandler(db, log.Logger)
|
||||
vulnerabilityHandler.SetAudit(auditSvc)
|
||||
webshellHandler := handler.NewWebShellHandler(log.Logger, db)
|
||||
webshellHandler.SetAudit(auditSvc)
|
||||
chatUploadsHandler := handler.NewChatUploadsHandler(log.Logger)
|
||||
chatUploadsHandler.SetAudit(auditSvc)
|
||||
registerWebshellTools(mcpServer, db, webshellHandler, log.Logger)
|
||||
registerWebshellManagementTools(mcpServer, db, webshellHandler, log.Logger)
|
||||
configHandler := handler.NewConfigHandler(configPath, cfg, mcpServer, executor, agent, attackChainHandler, externalMCPMgr, log.Logger)
|
||||
configHandler.SetAudit(auditSvc)
|
||||
agentHandler.SetHitlToolWhitelistSaver(configHandler)
|
||||
externalMCPHandler := handler.NewExternalMCPHandler(externalMCPMgr, cfg, configPath, log.Logger)
|
||||
externalMCPHandler.SetAudit(auditSvc)
|
||||
roleHandler := handler.NewRoleHandler(cfg, configPath, log.Logger)
|
||||
roleHandler.SetAudit(auditSvc)
|
||||
skillsHandler := handler.NewSkillsHandler(cfg, configPath, log.Logger)
|
||||
skillsHandler.SetAudit(auditSvc)
|
||||
fofaHandler := handler.NewFofaHandler(cfg, log.Logger)
|
||||
terminalHandler := handler.NewTerminalHandler(log.Logger)
|
||||
if db != nil {
|
||||
@@ -356,9 +378,12 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
||||
registerC2Tools(mcpServer, c2Manager, log.Logger, cfg.Server.Port)
|
||||
}
|
||||
c2Handler := handler.NewC2Handler(c2Manager, log.Logger)
|
||||
c2Handler.SetAudit(auditSvc)
|
||||
|
||||
// 创建OpenAPI处理器
|
||||
conversationHandler := handler.NewConversationHandler(db, log.Logger)
|
||||
conversationHandler.SetAudit(auditSvc)
|
||||
auditHandler := handler.NewAuditHandler(db, auditSvc, log.Logger)
|
||||
robotHandler := handler.NewRobotHandler(cfg, db, agentHandler, log.Logger)
|
||||
openAPIHandler := handler.NewOpenAPIHandler(db, log.Logger, resultStorage, conversationHandler, agentHandler)
|
||||
|
||||
@@ -384,13 +409,15 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
||||
c2Watchdog: c2Watchdog,
|
||||
c2WatchdogCancel: watchdogCancel,
|
||||
c2Handler: c2Handler,
|
||||
auditSvc: auditSvc,
|
||||
}
|
||||
// 飞书/钉钉长连接(无需公网),启用时在后台启动;后续前端应用配置时会通过 RestartRobotConnections 重启
|
||||
app.startRobotConnections()
|
||||
|
||||
// 设置漏洞工具注册器(内置工具,必须设置)
|
||||
vulnerabilityRegistrar := func() error {
|
||||
registerVulnerabilityTool(mcpServer, db, log.Logger)
|
||||
registerVulnerabilityTools(mcpServer, db, log.Logger)
|
||||
registerProjectFactTools(mcpServer, db, cfg, log.Logger)
|
||||
return nil
|
||||
}
|
||||
configHandler.SetVulnerabilityToolRegistrar(vulnerabilityRegistrar)
|
||||
@@ -449,9 +476,11 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
||||
configHandler.SetRetrieverUpdater(knowledgeRetriever)
|
||||
}
|
||||
|
||||
// 设置机器人连接重启器,前端应用配置后无需重启服务即可使钉钉/飞书新配置生效
|
||||
// 设置机器人连接重启器,前端应用配置后无需重启服务即可使钉钉/飞书/微信新配置生效
|
||||
configHandler.SetRobotRestarter(app)
|
||||
|
||||
wechatRobotHandler := handler.NewWechatRobotHandler(cfg, configHandler, log.Logger)
|
||||
|
||||
configHandler.SetC2Runtime(app)
|
||||
configHandler.SetC2ToolRegistrar(func() error {
|
||||
if app.config.C2.EnabledEffective() && app.c2Manager != nil {
|
||||
@@ -469,12 +498,14 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
||||
notificationHandler,
|
||||
conversationHandler,
|
||||
robotHandler,
|
||||
wechatRobotHandler,
|
||||
groupHandler,
|
||||
configHandler,
|
||||
externalMCPHandler,
|
||||
attackChainHandler,
|
||||
app, // 传递 App 实例以便动态获取 knowledgeHandler
|
||||
vulnerabilityHandler,
|
||||
projectHandler,
|
||||
webshellHandler,
|
||||
chatUploadsHandler,
|
||||
roleHandler,
|
||||
@@ -483,6 +514,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
||||
fofaHandler,
|
||||
terminalHandler,
|
||||
app.c2Handler,
|
||||
auditHandler,
|
||||
mcpServer,
|
||||
authManager,
|
||||
openAPIHandler,
|
||||
@@ -675,9 +707,14 @@ func (a *App) startRobotConnections() {
|
||||
a.dingCancel = cancel
|
||||
go robot.StartDing(ctx, cfg.Robots, a.robotHandler, a.logger.Logger)
|
||||
}
|
||||
if cfg.Robots.Wechat.Enabled && cfg.Robots.Wechat.BotToken != "" {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
a.wechatCancel = cancel
|
||||
go robot.StartWechat(ctx, cfg.Robots, a.robotHandler, cfg.Version, a.logger.Logger)
|
||||
}
|
||||
}
|
||||
|
||||
// RestartRobotConnections 重启钉钉/飞书长连接,使前端应用配置后立即生效(实现 handler.RobotRestarter)
|
||||
// RestartRobotConnections 重启钉钉/飞书/微信长连接,使前端应用配置后立即生效(实现 handler.RobotRestarter)
|
||||
func (a *App) RestartRobotConnections() {
|
||||
a.robotMu.Lock()
|
||||
if a.dingCancel != nil {
|
||||
@@ -688,6 +725,10 @@ func (a *App) RestartRobotConnections() {
|
||||
a.larkCancel()
|
||||
a.larkCancel = nil
|
||||
}
|
||||
if a.wechatCancel != nil {
|
||||
a.wechatCancel()
|
||||
a.wechatCancel = nil
|
||||
}
|
||||
a.robotMu.Unlock()
|
||||
// 给旧 goroutine 一点时间退出
|
||||
time.Sleep(200 * time.Millisecond)
|
||||
@@ -703,12 +744,14 @@ func setupRoutes(
|
||||
notificationHandler *handler.NotificationHandler,
|
||||
conversationHandler *handler.ConversationHandler,
|
||||
robotHandler *handler.RobotHandler,
|
||||
wechatRobotHandler *handler.WechatRobotHandler,
|
||||
groupHandler *handler.GroupHandler,
|
||||
configHandler *handler.ConfigHandler,
|
||||
externalMCPHandler *handler.ExternalMCPHandler,
|
||||
attackChainHandler *handler.AttackChainHandler,
|
||||
app *App, // 传递 App 实例以便动态获取 knowledgeHandler
|
||||
vulnerabilityHandler *handler.VulnerabilityHandler,
|
||||
projectHandler *handler.ProjectHandler,
|
||||
webshellHandler *handler.WebShellHandler,
|
||||
chatUploadsHandler *handler.ChatUploadsHandler,
|
||||
roleHandler *handler.RoleHandler,
|
||||
@@ -717,6 +760,7 @@ func setupRoutes(
|
||||
fofaHandler *handler.FofaHandler,
|
||||
terminalHandler *handler.TerminalHandler,
|
||||
c2Handler *handler.C2Handler,
|
||||
auditHandler *handler.AuditHandler,
|
||||
mcpServer *mcp.Server,
|
||||
authManager *security.AuthManager,
|
||||
openAPIHandler *handler.OpenAPIHandler,
|
||||
@@ -751,6 +795,12 @@ func setupRoutes(
|
||||
// 机器人测试(需登录):POST /api/robot/test,body: {"platform":"dingtalk","user_id":"test","text":"帮助"},用于验证机器人逻辑
|
||||
protected.POST("/robot/test", robotHandler.HandleRobotTest)
|
||||
|
||||
// 微信 iLink 扫码绑定(需登录)
|
||||
protected.POST("/robot/wechat/qrcode", wechatRobotHandler.HandleWechatQRCode)
|
||||
protected.GET("/robot/wechat/qrcode/status", wechatRobotHandler.HandleWechatQRCodeStatus)
|
||||
protected.POST("/robot/wechat/qrcode/verify", wechatRobotHandler.HandleWechatVerifyCode)
|
||||
protected.GET("/robot/wechat/status", wechatRobotHandler.HandleWechatStatus)
|
||||
|
||||
// Agent Loop
|
||||
protected.POST("/agent-loop", agentHandler.AgentLoop)
|
||||
// Agent Loop 流式输出
|
||||
@@ -806,6 +856,7 @@ func setupRoutes(
|
||||
protected.GET("/conversations/:id", conversationHandler.GetConversation)
|
||||
protected.GET("/messages/:id/process-details", conversationHandler.GetMessageProcessDetails)
|
||||
protected.PUT("/conversations/:id", conversationHandler.UpdateConversation)
|
||||
protected.PUT("/conversations/:id/project", conversationHandler.SetConversationProject)
|
||||
protected.DELETE("/conversations/:id", conversationHandler.DeleteConversation)
|
||||
protected.POST("/conversations/:id/delete-turn", conversationHandler.DeleteConversationTurn)
|
||||
protected.PUT("/conversations/:id/pinned", groupHandler.UpdateConversationPinned)
|
||||
@@ -847,6 +898,13 @@ func setupRoutes(
|
||||
protected.POST("/terminal/run/stream", terminalHandler.RunCommandStream)
|
||||
protected.GET("/terminal/ws", terminalHandler.RunCommandWS)
|
||||
|
||||
// 平台审计日志
|
||||
protected.GET("/audit/meta", auditHandler.Meta)
|
||||
protected.GET("/audit/summary", auditHandler.Summary)
|
||||
protected.GET("/audit/logs", auditHandler.ListLogs)
|
||||
protected.GET("/audit/logs/export", auditHandler.ExportLogs)
|
||||
protected.GET("/audit/logs/:id", auditHandler.GetLog)
|
||||
|
||||
// 外部MCP管理
|
||||
protected.GET("/external-mcp", externalMCPHandler.GetExternalMCPs)
|
||||
protected.GET("/external-mcp/stats", externalMCPHandler.GetExternalMCPStats)
|
||||
@@ -1015,6 +1073,19 @@ func setupRoutes(
|
||||
protected.PUT("/vulnerabilities/:id", vulnerabilityHandler.UpdateVulnerability)
|
||||
protected.DELETE("/vulnerabilities/:id", vulnerabilityHandler.DeleteVulnerability)
|
||||
|
||||
// 项目管理与事实黑板
|
||||
protected.GET("/projects", projectHandler.ListProjects)
|
||||
protected.POST("/projects", projectHandler.CreateProject)
|
||||
protected.GET("/projects/:id", projectHandler.GetProject)
|
||||
protected.PUT("/projects/:id", projectHandler.UpdateProject)
|
||||
protected.DELETE("/projects/:id", projectHandler.DeleteProject)
|
||||
protected.GET("/projects/:id/facts", projectHandler.ListFacts)
|
||||
protected.POST("/projects/:id/facts", projectHandler.CreateFact)
|
||||
protected.PUT("/projects/:id/facts/:factId", projectHandler.UpdateFact)
|
||||
protected.DELETE("/projects/:id/facts/:factId", projectHandler.DeleteFact)
|
||||
protected.POST("/projects/:id/facts/deprecate", projectHandler.DeprecateFact)
|
||||
protected.POST("/projects/:id/facts/restore", projectHandler.RestoreFact)
|
||||
|
||||
// WebShell 管理(代理执行 + 连接配置存 SQLite)
|
||||
protected.GET("/webshell/connections", webshellHandler.ListConnections)
|
||||
protected.POST("/webshell/connections", webshellHandler.CreateConnection)
|
||||
@@ -1135,195 +1206,6 @@ func setupRoutes(
|
||||
})
|
||||
}
|
||||
|
||||
// registerVulnerabilityTool 注册漏洞记录工具到MCP服务器
|
||||
func registerVulnerabilityTool(mcpServer *mcp.Server, db *database.DB, logger *zap.Logger) {
|
||||
tool := mcp.Tool{
|
||||
Name: builtin.ToolRecordVulnerability,
|
||||
Description: "记录发现的漏洞详情到漏洞管理系统。当发现有效漏洞时,使用此工具记录漏洞信息,包括标题、描述、严重程度、类型、目标、证明、影响和建议等。",
|
||||
ShortDescription: "记录发现的漏洞详情到漏洞管理系统",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"title": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞标题(必需)",
|
||||
},
|
||||
"description": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞详细描述",
|
||||
},
|
||||
"severity": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞严重程度:critical(严重)、high(高)、medium(中)、low(低)、info(信息)",
|
||||
"enum": []string{"critical", "high", "medium", "low", "info"},
|
||||
},
|
||||
"vulnerability_type": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞类型,如:SQL注入、XSS、CSRF、命令注入等",
|
||||
},
|
||||
"target": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "受影响的目标(URL、IP地址、服务等)",
|
||||
},
|
||||
"proof": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞证明(POC、截图、请求/响应等)",
|
||||
},
|
||||
"impact": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞影响说明",
|
||||
},
|
||||
"recommendation": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "修复建议",
|
||||
},
|
||||
},
|
||||
"required": []string{"title", "severity"},
|
||||
},
|
||||
}
|
||||
|
||||
handler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
// 从参数中获取conversation_id(由Agent自动添加)
|
||||
conversationID, _ := args["conversation_id"].(string)
|
||||
if conversationID == "" {
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: "错误: conversation_id 未设置。这是系统错误,请重试。",
|
||||
},
|
||||
},
|
||||
IsError: true,
|
||||
}, nil
|
||||
}
|
||||
|
||||
title, ok := args["title"].(string)
|
||||
if !ok || title == "" {
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: "错误: title 参数必需且不能为空",
|
||||
},
|
||||
},
|
||||
IsError: true,
|
||||
}, nil
|
||||
}
|
||||
|
||||
severity, ok := args["severity"].(string)
|
||||
if !ok || severity == "" {
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: "错误: severity 参数必需且不能为空",
|
||||
},
|
||||
},
|
||||
IsError: true,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// 验证严重程度
|
||||
validSeverities := map[string]bool{
|
||||
"critical": true,
|
||||
"high": true,
|
||||
"medium": true,
|
||||
"low": true,
|
||||
"info": true,
|
||||
}
|
||||
if !validSeverities[severity] {
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: fmt.Sprintf("错误: severity 必须是 critical、high、medium、low 或 info 之一,当前值: %s", severity),
|
||||
},
|
||||
},
|
||||
IsError: true,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// 获取可选参数
|
||||
description := ""
|
||||
if d, ok := args["description"].(string); ok {
|
||||
description = d
|
||||
}
|
||||
|
||||
vulnType := ""
|
||||
if t, ok := args["vulnerability_type"].(string); ok {
|
||||
vulnType = t
|
||||
}
|
||||
|
||||
target := ""
|
||||
if t, ok := args["target"].(string); ok {
|
||||
target = t
|
||||
}
|
||||
|
||||
proof := ""
|
||||
if p, ok := args["proof"].(string); ok {
|
||||
proof = p
|
||||
}
|
||||
|
||||
impact := ""
|
||||
if i, ok := args["impact"].(string); ok {
|
||||
impact = i
|
||||
}
|
||||
|
||||
recommendation := ""
|
||||
if r, ok := args["recommendation"].(string); ok {
|
||||
recommendation = r
|
||||
}
|
||||
|
||||
// 创建漏洞记录
|
||||
vuln := &database.Vulnerability{
|
||||
ConversationID: conversationID,
|
||||
Title: title,
|
||||
Description: description,
|
||||
Severity: severity,
|
||||
Status: "open",
|
||||
Type: vulnType,
|
||||
Target: target,
|
||||
Proof: proof,
|
||||
Impact: impact,
|
||||
Recommendation: recommendation,
|
||||
}
|
||||
|
||||
created, err := db.CreateVulnerability(vuln)
|
||||
if err != nil {
|
||||
logger.Error("记录漏洞失败", zap.Error(err))
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: fmt.Sprintf("记录漏洞失败: %v", err),
|
||||
},
|
||||
},
|
||||
IsError: true,
|
||||
}, nil
|
||||
}
|
||||
|
||||
logger.Info("漏洞记录成功",
|
||||
zap.String("id", created.ID),
|
||||
zap.String("title", created.Title),
|
||||
zap.String("severity", created.Severity),
|
||||
zap.String("conversation_id", conversationID),
|
||||
)
|
||||
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: fmt.Sprintf("漏洞已成功记录!\n\n漏洞ID: %s\n标题: %s\n严重程度: %s\n状态: %s\n\n你可以在漏洞管理页面查看和管理此漏洞。", created.ID, created.Title, created.Severity, created.Status),
|
||||
},
|
||||
},
|
||||
IsError: false,
|
||||
}, nil
|
||||
}
|
||||
|
||||
mcpServer.RegisterTool(tool, handler)
|
||||
logger.Info("漏洞记录工具注册成功")
|
||||
}
|
||||
|
||||
// registerWebshellTools 注册 WebShell 相关 MCP 工具,供 AI 助手在指定连接上执行命令与文件操作
|
||||
func registerWebshellTools(mcpServer *mcp.Server, db *database.DB, webshellHandler *handler.WebShellHandler, logger *zap.Logger) {
|
||||
if db == nil || webshellHandler == nil {
|
||||
@@ -1908,6 +1790,9 @@ func initializeKnowledge(
|
||||
|
||||
// 创建知识库API处理器
|
||||
knowledgeHandler := handler.NewKnowledgeHandler(knowledgeManager, knowledgeRetriever, knowledgeIndexer, db, logger)
|
||||
if app != nil && app.auditSvc != nil {
|
||||
knowledgeHandler.SetAudit(app.auditSvc)
|
||||
}
|
||||
logger.Info("知识库模块初始化完成", zap.Bool("handler_created", knowledgeHandler != nil))
|
||||
|
||||
// 设置知识库管理器到AgentHandler以便记录检索日志
|
||||
|
||||
@@ -0,0 +1,335 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
"cyberstrike-ai/internal/project"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func projectIDFromConversation(db *database.DB, ctx context.Context) (string, error) {
|
||||
convID := agent.ConversationIDFromContext(ctx)
|
||||
if convID == "" {
|
||||
return "", fmt.Errorf("无法确定当前对话,请在对话上下文中使用项目事实工具")
|
||||
}
|
||||
pid, err := db.GetConversationProjectID(convID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if strings.TrimSpace(pid) == "" {
|
||||
return "", fmt.Errorf("当前对话未绑定项目,请先在对话中选择项目或创建带项目的对话")
|
||||
}
|
||||
return pid, nil
|
||||
}
|
||||
|
||||
func textResult(msg string, isErr bool) *mcp.ToolResult {
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{{Type: "text", Text: msg}},
|
||||
IsError: isErr,
|
||||
}
|
||||
}
|
||||
|
||||
// registerProjectFactTools 注册项目黑板 MCP 工具。
|
||||
func registerProjectFactTools(mcpServer *mcp.Server, db *database.DB, cfg *config.Config, logger *zap.Logger) {
|
||||
if db == nil || cfg == nil || !cfg.Project.Enabled {
|
||||
if logger != nil {
|
||||
logger.Info("项目黑板工具未注册(未启用)")
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
upsertTool := mcp.Tool{
|
||||
Name: builtin.ToolUpsertProjectFact,
|
||||
Description: "写入或更新项目黑板事实,用于跨会话沉淀可复现上下文(非正式漏洞条目;可交付漏洞另用 record_vulnerability)。" +
|
||||
"禁止仅写结论:summary 须含什么+在哪+如何验证;body 须含攻击链/请求响应/命令等复现细节。" +
|
||||
"发现类建议 fact_key 为 finding|chain|exploit|poc/<slug>,category 对应 finding|chain|exploit|poc,body 按攻击链模板填写。" +
|
||||
"环境类用 target|auth|infra|business/<slug>。同 fact_key 覆盖更新。需当前对话已绑定项目。",
|
||||
ShortDescription: "写入/更新项目事实(含攻击链 body)",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"fact_key": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "项目内唯一 key:target/primary_domain、finding/sqli-login、exploit/upload-rce 等",
|
||||
},
|
||||
"category": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "target | auth | infra | business | finding | chain | exploit | poc | note",
|
||||
"enum": []string{"target", "auth", "infra", "business", "finding", "chain", "exploit", "poc", "note"},
|
||||
},
|
||||
"summary": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "索引用一行:结论 + 位置 + 触发/验证要点(勿仅写「存在 XSS」等空话)",
|
||||
},
|
||||
"body": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "完整可复现详情(仅 get_project_fact 返回):须含攻击链步骤、原始 HTTP/命令、响应现象、证据与关联。" +
|
||||
"发现/利用类首次写入必填;环境类建议含来源证据。攻击链类可参考模板章节:结论、目标与入口、攻击链、Exploit/POC、关键证据、关联、备注。" +
|
||||
"更新已有 fact_key 时若省略或留空 body,将保留库中已有 body(可只改 summary)。",
|
||||
},
|
||||
"confidence": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "confirmed | tentative | deprecated",
|
||||
"enum": []string{"confirmed", "tentative", "deprecated"},
|
||||
},
|
||||
"pinned": map[string]interface{}{
|
||||
"type": "boolean",
|
||||
"description": "是否优先出现在黑板索引",
|
||||
},
|
||||
"related_vulnerability_id": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "可选:关联的漏洞记录 ID",
|
||||
},
|
||||
},
|
||||
"required": []string{"fact_key", "summary"},
|
||||
},
|
||||
}
|
||||
|
||||
mcpServer.RegisterTool(upsertTool, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
projectID, err := projectIDFromConversation(db, ctx)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
factKey, _ := args["fact_key"].(string)
|
||||
summary, _ := args["summary"].(string)
|
||||
if strings.TrimSpace(factKey) == "" || strings.TrimSpace(summary) == "" {
|
||||
return textResult("错误: fact_key 与 summary 必填", true), nil
|
||||
}
|
||||
if len([]rune(summary)) > cfg.Project.FactSummaryMaxRunesEffective() {
|
||||
return textResult(fmt.Sprintf("错误: summary 过长(最多 %d 字)", cfg.Project.FactSummaryMaxRunesEffective()), true), nil
|
||||
}
|
||||
f := &database.ProjectFact{
|
||||
ProjectID: projectID,
|
||||
FactKey: factKey,
|
||||
Category: strArg(args, "category"),
|
||||
Summary: summary,
|
||||
Body: strArg(args, "body"),
|
||||
Confidence: strArg(args, "confidence"),
|
||||
Pinned: boolArg(args, "pinned"),
|
||||
RelatedVulnerabilityID: strArg(args, "related_vulnerability_id"),
|
||||
}
|
||||
if convID := agent.ConversationIDFromContext(ctx); convID != "" {
|
||||
f.SourceConversationID = convID
|
||||
}
|
||||
created, err := db.UpsertProjectFact(f)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
msg := fmt.Sprintf("事实已保存。\nfact_key: %s\nid: %s\nconfidence: %s", created.FactKey, created.ID, created.Confidence)
|
||||
if warn := project.SparseBodyWarningIfNeeded(f.Category, f.FactKey, f.Body); warn != "" {
|
||||
msg += warn
|
||||
}
|
||||
return textResult(msg, false), nil
|
||||
})
|
||||
|
||||
getTool := mcp.Tool{
|
||||
Name: builtin.ToolGetProjectFact,
|
||||
Description: "按 fact_key 获取项目事实完整 body 与元数据。摘要不足时必须调用本工具,禁止臆造细节。",
|
||||
ShortDescription: "按 key 获取事实详情",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"fact_key": map[string]interface{}{"type": "string", "description": "事实 key"},
|
||||
},
|
||||
"required": []string{"fact_key"},
|
||||
},
|
||||
}
|
||||
mcpServer.RegisterTool(getTool, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
projectID, err := projectIDFromConversation(db, ctx)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
key := strings.TrimSpace(strArg(args, "fact_key"))
|
||||
if key == "" {
|
||||
return textResult("错误: fact_key 必填", true), nil
|
||||
}
|
||||
f, err := db.GetProjectFactByKey(projectID, key)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
msg := fmt.Sprintf("fact_key: %s\ncategory: %s\nconfidence: %s\nsummary: %s\nupdated_at: %s",
|
||||
f.FactKey, f.Category, f.Confidence, f.Summary, f.UpdatedAt.Format("2006-01-02 15:04:05"))
|
||||
if f.RelatedVulnerabilityID != "" {
|
||||
msg += fmt.Sprintf("\nrelated_vulnerability_id: %s", f.RelatedVulnerabilityID)
|
||||
}
|
||||
if f.SourceConversationID != "" {
|
||||
msg += fmt.Sprintf("\nsource_conversation_id: %s", f.SourceConversationID)
|
||||
}
|
||||
msg += "\n\n--- body ---\n" + f.Body
|
||||
if warn := project.SparseBodyWarningIfNeeded(f.Category, f.FactKey, f.Body); warn != "" {
|
||||
msg += warn
|
||||
}
|
||||
return textResult(msg, false), nil
|
||||
})
|
||||
|
||||
listTool := mcp.Tool{
|
||||
Name: builtin.ToolListProjectFacts,
|
||||
Description: "列出当前项目的事实(分页)。",
|
||||
ShortDescription: "列出项目事实",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"category": map[string]interface{}{"type": "string"},
|
||||
"confidence": map[string]interface{}{"type": "string"},
|
||||
"limit": map[string]interface{}{"type": "integer"},
|
||||
"offset": map[string]interface{}{"type": "integer"},
|
||||
},
|
||||
},
|
||||
}
|
||||
mcpServer.RegisterTool(listTool, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
projectID, err := projectIDFromConversation(db, ctx)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
limit := intArg(args, "limit", 50)
|
||||
offset := intArg(args, "offset", 0)
|
||||
filter := database.ProjectFactListFilter{
|
||||
Category: strArg(args, "category"),
|
||||
Confidence: strArg(args, "confidence"),
|
||||
}
|
||||
list, err := db.ListProjectFacts(projectID, filter, limit, offset)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
var b strings.Builder
|
||||
b.WriteString(fmt.Sprintf("共 %d 条(limit=%d offset=%d):\n", len(list), limit, offset))
|
||||
for _, f := range list {
|
||||
b.WriteString(fmt.Sprintf("- [%s] %s — %s (%s)\n", f.FactKey, f.Category, f.Summary, f.Confidence))
|
||||
}
|
||||
return textResult(b.String(), false), nil
|
||||
})
|
||||
|
||||
searchTool := mcp.Tool{
|
||||
Name: builtin.ToolSearchProjectFacts,
|
||||
Description: "按关键词搜索项目事实(summary/body/fact_key)。",
|
||||
ShortDescription: "搜索项目事实",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"query": map[string]interface{}{"type": "string"},
|
||||
"limit": map[string]interface{}{"type": "integer"},
|
||||
"offset": map[string]interface{}{"type": "integer"},
|
||||
},
|
||||
"required": []string{"query"},
|
||||
},
|
||||
}
|
||||
mcpServer.RegisterTool(searchTool, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
projectID, err := projectIDFromConversation(db, ctx)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
q := strings.TrimSpace(strArg(args, "query"))
|
||||
if q == "" {
|
||||
return textResult("错误: query 必填", true), nil
|
||||
}
|
||||
list, err := db.ListProjectFacts(projectID, database.ProjectFactListFilter{Search: q}, intArg(args, "limit", 30), intArg(args, "offset", 0))
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
var b strings.Builder
|
||||
b.WriteString(fmt.Sprintf("搜索 \"%s\" 命中 %d 条:\n", q, len(list)))
|
||||
for _, f := range list {
|
||||
b.WriteString(fmt.Sprintf("- [%s] %s — %s\n", f.FactKey, f.Category, f.Summary))
|
||||
}
|
||||
return textResult(b.String(), false), nil
|
||||
})
|
||||
|
||||
deprecateTool := mcp.Tool{
|
||||
Name: builtin.ToolDeprecateProjectFact,
|
||||
Description: "将事实标记为 deprecated,从黑板索引中排除。",
|
||||
ShortDescription: "废弃项目事实",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"fact_key": map[string]interface{}{"type": "string"},
|
||||
},
|
||||
"required": []string{"fact_key"},
|
||||
},
|
||||
}
|
||||
mcpServer.RegisterTool(deprecateTool, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
projectID, err := projectIDFromConversation(db, ctx)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
key := strings.TrimSpace(strArg(args, "fact_key"))
|
||||
if err := db.DeprecateProjectFact(projectID, key); err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
return textResult("事实已标记为 deprecated: "+key, false), nil
|
||||
})
|
||||
|
||||
restoreTool := mcp.Tool{
|
||||
Name: builtin.ToolRestoreProjectFact,
|
||||
Description: "将已废弃(deprecated)的事实恢复为 tentative 或 confirmed,重新参与黑板索引。",
|
||||
ShortDescription: "恢复已废弃的项目事实",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"fact_key": map[string]interface{}{"type": "string"},
|
||||
"confidence": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "恢复后的置信度:tentative(默认)或 confirmed",
|
||||
"enum": []string{"tentative", "confirmed"},
|
||||
},
|
||||
},
|
||||
"required": []string{"fact_key"},
|
||||
},
|
||||
}
|
||||
mcpServer.RegisterTool(restoreTool, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
projectID, err := projectIDFromConversation(db, ctx)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
key := strings.TrimSpace(strArg(args, "fact_key"))
|
||||
if key == "" {
|
||||
return textResult("错误: fact_key 必填", true), nil
|
||||
}
|
||||
conf := strArg(args, "confidence")
|
||||
if err := db.RestoreProjectFact(projectID, key, conf); err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
if conf == "" {
|
||||
conf = "tentative"
|
||||
}
|
||||
return textResult(fmt.Sprintf("事实已恢复为 %s: %s", conf, key), false), nil
|
||||
})
|
||||
|
||||
if logger != nil {
|
||||
logger.Info("项目黑板 MCP 工具注册成功")
|
||||
}
|
||||
}
|
||||
|
||||
func strArg(args map[string]interface{}, key string) string {
|
||||
if v, ok := args[key].(string); ok {
|
||||
return v
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func boolArg(args map[string]interface{}, key string) bool {
|
||||
if v, ok := args[key].(bool); ok {
|
||||
return v
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func intArg(args map[string]interface{}, key string, def int) int {
|
||||
switch v := args[key].(type) {
|
||||
case float64:
|
||||
return int(v)
|
||||
case int:
|
||||
return v
|
||||
case int64:
|
||||
return int(v)
|
||||
default:
|
||||
return def
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,405 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func conversationIDFromToolCtx(ctx context.Context) string {
|
||||
if id := agent.ConversationIDFromContext(ctx); id != "" {
|
||||
return id
|
||||
}
|
||||
return mcp.MCPConversationIDFromContext(ctx)
|
||||
}
|
||||
|
||||
// canAccessVulnerability 校验当前对话是否有权查看该漏洞(默认项目隔离,未绑项目则仅本会话)。
|
||||
func canAccessVulnerability(vuln *database.Vulnerability, convID, projectID string) bool {
|
||||
if vuln == nil || convID == "" {
|
||||
return false
|
||||
}
|
||||
if projectID != "" {
|
||||
if strings.TrimSpace(vuln.ProjectID) == projectID {
|
||||
return true
|
||||
}
|
||||
// 历史记录:写入时尚未绑定 project_id,但属于同一会话
|
||||
if strings.TrimSpace(vuln.ProjectID) == "" && vuln.ConversationID == convID {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
return vuln.ConversationID == convID
|
||||
}
|
||||
|
||||
func buildVulnerabilityListFilter(db *database.DB, ctx context.Context, args map[string]interface{}) (database.VulnerabilityListFilter, string, error) {
|
||||
convID := conversationIDFromToolCtx(ctx)
|
||||
if convID == "" {
|
||||
return database.VulnerabilityListFilter{}, "", fmt.Errorf("无法确定当前对话,请在对话上下文中使用漏洞查询工具")
|
||||
}
|
||||
|
||||
projectID := ""
|
||||
if pid, err := db.GetConversationProjectID(convID); err == nil {
|
||||
projectID = strings.TrimSpace(pid)
|
||||
}
|
||||
|
||||
scope := strings.TrimSpace(strArg(args, "scope"))
|
||||
if scope == "" {
|
||||
if projectID != "" {
|
||||
scope = "project"
|
||||
} else {
|
||||
scope = "conversation"
|
||||
}
|
||||
}
|
||||
|
||||
filter := database.VulnerabilityListFilter{
|
||||
Severity: strings.TrimSpace(strArg(args, "severity")),
|
||||
Status: strings.TrimSpace(strArg(args, "status")),
|
||||
}
|
||||
if q := strings.TrimSpace(strArg(args, "q")); q != "" {
|
||||
filter.Search = q
|
||||
} else {
|
||||
filter.Search = strings.TrimSpace(strArg(args, "search"))
|
||||
}
|
||||
|
||||
var scopeLabel string
|
||||
switch scope {
|
||||
case "project":
|
||||
if projectID == "" {
|
||||
return filter, "", fmt.Errorf("当前对话未绑定项目,无法按项目列出漏洞;请使用 scope=conversation,或先在对话中绑定项目")
|
||||
}
|
||||
filter.ProjectID = projectID
|
||||
scopeLabel = fmt.Sprintf("项目 %s", projectID)
|
||||
case "conversation":
|
||||
filter.ConversationID = convID
|
||||
scopeLabel = fmt.Sprintf("会话 %s", convID)
|
||||
default:
|
||||
return filter, "", fmt.Errorf("scope 仅支持 project 或 conversation,当前值: %s", scope)
|
||||
}
|
||||
return filter, scopeLabel, nil
|
||||
}
|
||||
|
||||
func formatVulnerabilityListItem(v *database.Vulnerability) string {
|
||||
line := fmt.Sprintf("- id=%s | %s | %s | %s", v.ID, v.Severity, v.Status, v.Title)
|
||||
if v.Type != "" {
|
||||
line += fmt.Sprintf(" | type=%s", v.Type)
|
||||
}
|
||||
if v.Target != "" {
|
||||
line += fmt.Sprintf(" | target=%s", truncateRunes(v.Target, 80))
|
||||
}
|
||||
return line
|
||||
}
|
||||
|
||||
func formatVulnerabilityDetail(v *database.Vulnerability) string {
|
||||
var b strings.Builder
|
||||
b.WriteString(fmt.Sprintf("漏洞ID: %s\n", v.ID))
|
||||
b.WriteString(fmt.Sprintf("标题: %s\n", v.Title))
|
||||
b.WriteString(fmt.Sprintf("严重程度: %s\n", v.Severity))
|
||||
b.WriteString(fmt.Sprintf("状态: %s\n", v.Status))
|
||||
if v.Type != "" {
|
||||
b.WriteString(fmt.Sprintf("类型: %s\n", v.Type))
|
||||
}
|
||||
if v.Target != "" {
|
||||
b.WriteString(fmt.Sprintf("目标: %s\n", v.Target))
|
||||
}
|
||||
if v.ProjectID != "" {
|
||||
b.WriteString(fmt.Sprintf("项目ID: %s\n", v.ProjectID))
|
||||
}
|
||||
b.WriteString(fmt.Sprintf("会话ID: %s\n", v.ConversationID))
|
||||
if !v.CreatedAt.IsZero() {
|
||||
b.WriteString(fmt.Sprintf("创建时间: %s\n", v.CreatedAt.Format("2006-01-02 15:04:05")))
|
||||
}
|
||||
if v.Description != "" {
|
||||
b.WriteString("\n--- 描述 ---\n")
|
||||
b.WriteString(v.Description)
|
||||
b.WriteString("\n")
|
||||
}
|
||||
if v.Proof != "" {
|
||||
b.WriteString("\n--- 证明(POC) ---\n")
|
||||
b.WriteString(v.Proof)
|
||||
b.WriteString("\n")
|
||||
}
|
||||
if v.Impact != "" {
|
||||
b.WriteString("\n--- 影响 ---\n")
|
||||
b.WriteString(v.Impact)
|
||||
b.WriteString("\n")
|
||||
}
|
||||
if v.Recommendation != "" {
|
||||
b.WriteString("\n--- 修复建议 ---\n")
|
||||
b.WriteString(v.Recommendation)
|
||||
b.WriteString("\n")
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
func truncateRunes(s string, max int) string {
|
||||
r := []rune(s)
|
||||
if len(r) <= max {
|
||||
return s
|
||||
}
|
||||
return string(r[:max]) + "…"
|
||||
}
|
||||
|
||||
// registerVulnerabilityTools 注册漏洞记录与查询 MCP 工具。
|
||||
func registerVulnerabilityTools(mcpServer *mcp.Server, db *database.DB, logger *zap.Logger) {
|
||||
registerRecordVulnerabilityTool(mcpServer, db, logger)
|
||||
registerListVulnerabilitiesTool(mcpServer, db, logger)
|
||||
registerGetVulnerabilityTool(mcpServer, db, logger)
|
||||
if logger != nil {
|
||||
logger.Info("漏洞 MCP 工具注册成功", zap.Strings("tools", []string{
|
||||
builtin.ToolRecordVulnerability,
|
||||
builtin.ToolListVulnerabilities,
|
||||
builtin.ToolGetVulnerability,
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
func registerRecordVulnerabilityTool(mcpServer *mcp.Server, db *database.DB, logger *zap.Logger) {
|
||||
tool := mcp.Tool{
|
||||
Name: builtin.ToolRecordVulnerability,
|
||||
Description: "记录发现的漏洞详情到漏洞管理系统。当发现有效漏洞时,使用此工具记录漏洞信息,包括标题、描述、严重程度、类型、目标、证明、影响和建议等。记录前可先 list_vulnerabilities 避免重复。",
|
||||
ShortDescription: "记录发现的漏洞详情到漏洞管理系统",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"title": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞标题(必需)",
|
||||
},
|
||||
"description": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞详细描述",
|
||||
},
|
||||
"severity": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞严重程度:critical(严重)、high(高)、medium(中)、low(低)、info(信息)",
|
||||
"enum": []string{"critical", "high", "medium", "low", "info"},
|
||||
},
|
||||
"vulnerability_type": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞类型,如:SQL注入、XSS、CSRF、命令注入等",
|
||||
},
|
||||
"target": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "受影响的目标(URL、IP地址、服务等)",
|
||||
},
|
||||
"proof": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞证明(POC、截图、请求/响应等)",
|
||||
},
|
||||
"impact": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞影响说明",
|
||||
},
|
||||
"recommendation": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "修复建议",
|
||||
},
|
||||
},
|
||||
"required": []string{"title", "severity"},
|
||||
},
|
||||
}
|
||||
|
||||
mcpServer.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
conversationID := strings.TrimSpace(strArg(args, "conversation_id"))
|
||||
if conversationID == "" {
|
||||
conversationID = conversationIDFromToolCtx(ctx)
|
||||
}
|
||||
if conversationID == "" {
|
||||
return textResult("错误: conversation_id 未设置。这是系统错误,请重试。", true), nil
|
||||
}
|
||||
|
||||
title := strings.TrimSpace(strArg(args, "title"))
|
||||
if title == "" {
|
||||
return textResult("错误: title 参数必需且不能为空", true), nil
|
||||
}
|
||||
|
||||
severity := strings.TrimSpace(strArg(args, "severity"))
|
||||
if severity == "" {
|
||||
return textResult("错误: severity 参数必需且不能为空", true), nil
|
||||
}
|
||||
|
||||
validSeverities := map[string]bool{
|
||||
"critical": true, "high": true, "medium": true, "low": true, "info": true,
|
||||
}
|
||||
if !validSeverities[severity] {
|
||||
return textResult(fmt.Sprintf("错误: severity 必须是 critical、high、medium、low 或 info 之一,当前值: %s", severity), true), nil
|
||||
}
|
||||
|
||||
projectID := ""
|
||||
if pid, perr := db.GetConversationProjectID(conversationID); perr == nil {
|
||||
projectID = strings.TrimSpace(pid)
|
||||
}
|
||||
|
||||
vuln := &database.Vulnerability{
|
||||
ConversationID: conversationID,
|
||||
ProjectID: projectID,
|
||||
Title: title,
|
||||
Description: strArg(args, "description"),
|
||||
Severity: severity,
|
||||
Status: "open",
|
||||
Type: strArg(args, "vulnerability_type"),
|
||||
Target: strArg(args, "target"),
|
||||
Proof: strArg(args, "proof"),
|
||||
Impact: strArg(args, "impact"),
|
||||
Recommendation: strArg(args, "recommendation"),
|
||||
}
|
||||
|
||||
created, err := db.CreateVulnerability(vuln)
|
||||
if err != nil {
|
||||
if logger != nil {
|
||||
logger.Error("记录漏洞失败", zap.Error(err))
|
||||
}
|
||||
return textResult(fmt.Sprintf("记录漏洞失败: %v", err), true), nil
|
||||
}
|
||||
|
||||
if logger != nil {
|
||||
logger.Info("漏洞记录成功",
|
||||
zap.String("id", created.ID),
|
||||
zap.String("title", created.Title),
|
||||
zap.String("severity", created.Severity),
|
||||
zap.String("conversation_id", conversationID),
|
||||
)
|
||||
}
|
||||
|
||||
return textResult(fmt.Sprintf("漏洞已成功记录!\n\n漏洞ID: %s\n标题: %s\n严重程度: %s\n状态: %s\n\n可使用 get_vulnerability(id) 查看详情,或 list_vulnerabilities 查看列表。",
|
||||
created.ID, created.Title, created.Severity, created.Status), false), nil
|
||||
})
|
||||
}
|
||||
|
||||
func registerListVulnerabilitiesTool(mcpServer *mcp.Server, db *database.DB, logger *zap.Logger) {
|
||||
tool := mcp.Tool{
|
||||
Name: builtin.ToolListVulnerabilities,
|
||||
Description: "列出当前授权范围内的漏洞(摘要)。默认:对话已绑定项目时列出该项目下全部漏洞;未绑项目时仅列出当前会话漏洞。可用 scope=conversation 仅看本会话。记录新漏洞前建议先调用以避免重复。",
|
||||
ShortDescription: "列出漏洞(默认当前项目)",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"scope": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "范围:project(默认,需绑定项目)| conversation(仅当前会话)",
|
||||
"enum": []string{"project", "conversation"},
|
||||
},
|
||||
"severity": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "按严重程度筛选:critical、high、medium、low、info",
|
||||
"enum": []string{"critical", "high", "medium", "low", "info"},
|
||||
},
|
||||
"status": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "按状态筛选:open、confirmed、fixed、false_positive",
|
||||
"enum": []string{"open", "confirmed", "fixed", "false_positive"},
|
||||
},
|
||||
"q": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "关键词搜索(标题、描述、类型、目标等)",
|
||||
},
|
||||
"limit": map[string]interface{}{
|
||||
"type": "integer",
|
||||
"description": "返回条数上限,默认 30,最大 100",
|
||||
},
|
||||
"offset": map[string]interface{}{
|
||||
"type": "integer",
|
||||
"description": "分页偏移,默认 0",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
mcpServer.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
filter, scopeLabel, err := buildVulnerabilityListFilter(db, ctx, args)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
|
||||
limit := intArg(args, "limit", 30)
|
||||
if limit <= 0 || limit > 100 {
|
||||
limit = 30
|
||||
}
|
||||
offset := intArg(args, "offset", 0)
|
||||
if offset < 0 {
|
||||
offset = 0
|
||||
}
|
||||
|
||||
total, err := db.CountVulnerabilities(filter)
|
||||
if err != nil {
|
||||
if logger != nil {
|
||||
logger.Warn("统计漏洞失败", zap.Error(err))
|
||||
}
|
||||
total = 0
|
||||
}
|
||||
|
||||
list, err := db.ListVulnerabilities(limit, offset, filter)
|
||||
if err != nil {
|
||||
return textResult("错误: "+err.Error(), true), nil
|
||||
}
|
||||
|
||||
var b strings.Builder
|
||||
b.WriteString(fmt.Sprintf("范围: %s\n总计: %d | 本页: %d 条 (limit=%d offset=%d)\n\n", scopeLabel, total, len(list), limit, offset))
|
||||
if len(list) == 0 {
|
||||
b.WriteString("(暂无漏洞记录)\n")
|
||||
} else {
|
||||
for _, v := range list {
|
||||
b.WriteString(formatVulnerabilityListItem(v))
|
||||
b.WriteString("\n")
|
||||
}
|
||||
if total > offset+len(list) {
|
||||
b.WriteString(fmt.Sprintf("\n(还有更多,可增大 offset 或使用 q/severity/status 筛选)\n"))
|
||||
}
|
||||
}
|
||||
b.WriteString("\n需要 POC 与完整字段请对具体 id 调用 get_vulnerability。")
|
||||
return textResult(b.String(), false), nil
|
||||
})
|
||||
}
|
||||
|
||||
func registerGetVulnerabilityTool(mcpServer *mcp.Server, db *database.DB, logger *zap.Logger) {
|
||||
tool := mcp.Tool{
|
||||
Name: builtin.ToolGetVulnerability,
|
||||
Description: "按漏洞 ID 获取完整详情(含 POC、影响、修复建议)。仅能访问当前项目或当前会话下的漏洞(与 list_vulnerabilities 授权范围一致)。",
|
||||
ShortDescription: "按 ID 获取漏洞详情",
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"id": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "漏洞 ID(list_vulnerabilities 返回的 id)",
|
||||
},
|
||||
},
|
||||
"required": []string{"id"},
|
||||
},
|
||||
}
|
||||
|
||||
mcpServer.RegisterTool(tool, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
convID := conversationIDFromToolCtx(ctx)
|
||||
if convID == "" {
|
||||
return textResult("错误: 无法确定当前对话,请在对话上下文中使用本工具", true), nil
|
||||
}
|
||||
|
||||
id := strings.TrimSpace(strArg(args, "id"))
|
||||
if id == "" {
|
||||
return textResult("错误: id 必填", true), nil
|
||||
}
|
||||
|
||||
vuln, err := db.GetVulnerability(id)
|
||||
if err != nil {
|
||||
return textResult("错误: 漏洞不存在或查询失败", true), nil
|
||||
}
|
||||
|
||||
projectID := ""
|
||||
if pid, perr := db.GetConversationProjectID(convID); perr == nil {
|
||||
projectID = strings.TrimSpace(pid)
|
||||
}
|
||||
|
||||
if !canAccessVulnerability(vuln, convID, projectID) {
|
||||
return textResult("错误: 无权访问该漏洞(仅可查看当前项目或当前会话下的记录)", true), nil
|
||||
}
|
||||
|
||||
return textResult(formatVulnerabilityDetail(vuln), false), nil
|
||||
})
|
||||
}
|
||||
@@ -82,7 +82,7 @@ func NewBuilder(db *database.DB, openAIConfig *config.OpenAIConfig, logger *zap.
|
||||
}
|
||||
}
|
||||
|
||||
// BuildChainFromConversation 从对话构建攻击链(简化版本:用户输入+最后一轮ReAct输入+大模型输出)
|
||||
// BuildChainFromConversation 从对话构建攻击链(单次 LLM 调用;输入为当前任务轮次的 last_react 轨迹,与继续对话续跑范围一致)。
|
||||
func (b *Builder) BuildChainFromConversation(ctx context.Context, conversationID string) (*Chain, error) {
|
||||
b.logger.Info("开始构建攻击链(简化版本)", zap.String("conversationId", conversationID))
|
||||
|
||||
@@ -157,33 +157,34 @@ func (b *Builder) BuildChainFromConversation(ctx context.Context, conversationID
|
||||
var reactInputFinal string
|
||||
var dataSource string // 记录数据来源
|
||||
|
||||
// 如果成功获取到保存的ReAct数据,直接使用
|
||||
if reactInputJSON != "" && modelOutput != "" {
|
||||
// 计算 ReAct 输入的哈希值,用于追踪
|
||||
hash := sha256.Sum256([]byte(reactInputJSON))
|
||||
reactInputHash := hex.EncodeToString(hash[:])[:16] // 使用前16字符作为短标识
|
||||
// 优先使用落库的代理轨迹(与继续对话 loadHistoryFromAgentTrace 同源),并裁剪为「当前任务轮次」
|
||||
if reactInputJSON != "" {
|
||||
trimmedJSON := agent.ExtractLastUserTurnTraceJSON(reactInputJSON)
|
||||
hash := sha256.Sum256([]byte(trimmedJSON))
|
||||
reactInputHash := hex.EncodeToString(hash[:])[:16]
|
||||
|
||||
// 统计消息数量
|
||||
var messageCount int
|
||||
var tempMessages []interface{}
|
||||
if json.Unmarshal([]byte(reactInputJSON), &tempMessages) == nil {
|
||||
messageCount = len(tempMessages)
|
||||
if msgs, parseErr := agent.ParseTraceMessages(trimmedJSON); parseErr == nil {
|
||||
messageCount = len(msgs)
|
||||
msgs = agent.MergeAssistantTraceOutput(msgs, modelOutput)
|
||||
reactInputFinal = b.formatAgentTraceFromChatMessages(msgs)
|
||||
} else {
|
||||
b.logger.Warn("解析代理轨迹失败,回退原始 JSON 格式化", zap.Error(parseErr))
|
||||
reactInputFinal = b.formatAgentTraceInputFromJSON(trimmedJSON)
|
||||
if strings.TrimSpace(modelOutput) != "" {
|
||||
reactInputFinal += "\n\n## 助手结论(last_react_output)\n\n" + modelOutput
|
||||
}
|
||||
}
|
||||
|
||||
dataSource = "database_last_agent_trace"
|
||||
b.logger.Info("使用保存的ReAct数据构建攻击链",
|
||||
dataSource = "last_user_turn_agent_trace"
|
||||
b.logger.Info("使用当前任务轮次代理轨迹构建攻击链(与续跑上下文范围一致)",
|
||||
zap.String("conversationId", conversationID),
|
||||
zap.String("dataSource", dataSource),
|
||||
zap.Int("reactInputSize", len(reactInputJSON)),
|
||||
zap.Int("traceInputSizeBeforeTrim", len(reactInputJSON)),
|
||||
zap.Int("traceInputSizeAfterTrim", len(trimmedJSON)),
|
||||
zap.Int("messageCount", messageCount),
|
||||
zap.String("reactInputHash", reactInputHash),
|
||||
zap.Int("modelOutputSize", len(modelOutput)))
|
||||
|
||||
// 从保存的ReAct输入(JSON格式)中提取用户输入
|
||||
// userInput = b.extractUserInputFromReActInput(reactInputJSON)
|
||||
|
||||
// 将JSON格式的messages转换为可读格式
|
||||
reactInputFinal = b.formatAgentTraceInputFromJSON(reactInputJSON)
|
||||
} else {
|
||||
// 2. 如果没有保存的ReAct数据,从对话消息构建
|
||||
dataSource = "messages_table"
|
||||
@@ -243,8 +244,15 @@ func (b *Builder) BuildChainFromConversation(ctx context.Context, conversationID
|
||||
}
|
||||
}
|
||||
|
||||
// 3. 构建简化的prompt,一次性传递给大模型
|
||||
prompt := b.buildSimplePrompt(reactInputFinal, modelOutput)
|
||||
// 3. 按 token 预算压缩输入,再构建 prompt(避免超出模型上下文)
|
||||
reactInputFinal, modelOutput, _ = b.fitAttackChainPayload(reactInputFinal, modelOutput)
|
||||
|
||||
// 4. 构建 prompt 并单次调用大模型(助手结论已并入轨迹时不再重复传入)
|
||||
promptAssistantOut := modelOutput
|
||||
if reactInputJSON != "" {
|
||||
promptAssistantOut = ""
|
||||
}
|
||||
prompt := b.buildSimplePrompt(reactInputFinal, promptAssistantOut)
|
||||
// fmt.Println(prompt)
|
||||
// 6. 调用AI生成攻击链(一次性,不做任何处理)
|
||||
chainJSON, err := b.callAIForChainGeneration(ctx, prompt)
|
||||
@@ -366,10 +374,17 @@ func (b *Builder) formatProcessDetailsForAttackChain(details []database.ProcessD
|
||||
return strings.TrimSpace(sb.String())
|
||||
}
|
||||
|
||||
// buildAgentTraceInput 构建最后一轮ReAct的输入(历史消息+当前用户输入)
|
||||
// buildAgentTraceInput 构建最后一轮 ReAct 的输入(从最后一条 user 消息起,不含更早轮次)。
|
||||
func (b *Builder) buildAgentTraceInput(messages []database.Message) string {
|
||||
start := 0
|
||||
for i := len(messages) - 1; i >= 0; i-- {
|
||||
if strings.EqualFold(messages[i].Role, "user") {
|
||||
start = i
|
||||
break
|
||||
}
|
||||
}
|
||||
var builder strings.Builder
|
||||
for _, msg := range messages {
|
||||
for _, msg := range messages[start:] {
|
||||
builder.WriteString(fmt.Sprintf("[%s]: %s\n\n", msg.Role, msg.Content))
|
||||
}
|
||||
return builder.String()
|
||||
@@ -396,67 +411,66 @@ func (b *Builder) buildAgentTraceInput(messages []database.Message) string {
|
||||
// return ""
|
||||
// }
|
||||
|
||||
// formatAgentTraceInputFromJSON 将JSON格式的messages数组转换为可读的字符串格式
|
||||
// formatAgentTraceInputFromJSON 将 JSON 轨迹转为可读文本(会先按当前任务轮次裁剪)。
|
||||
func (b *Builder) formatAgentTraceInputFromJSON(reactInputJSON string) string {
|
||||
var messages []map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(reactInputJSON), &messages); err != nil {
|
||||
trimmed := agent.ExtractLastUserTurnTraceJSON(reactInputJSON)
|
||||
msgs, err := agent.ParseTraceMessages(trimmed)
|
||||
if err != nil {
|
||||
b.logger.Warn("解析ReAct输入JSON失败", zap.Error(err))
|
||||
return reactInputJSON // 如果解析失败,返回原始JSON
|
||||
return trimmed
|
||||
}
|
||||
return b.formatAgentTraceFromChatMessages(msgs)
|
||||
}
|
||||
|
||||
// formatAgentTraceFromChatMessages 将代理消息带格式化为攻击链分析输入(与续跑轨迹字段一致)。
|
||||
func (b *Builder) formatAgentTraceFromChatMessages(msgs []agent.ChatMessage) string {
|
||||
var builder strings.Builder
|
||||
for _, msg := range messages {
|
||||
role, _ := msg["role"].(string)
|
||||
content, _ := msg["content"].(string)
|
||||
for _, msg := range msgs {
|
||||
role := msg.Role
|
||||
content := msg.Content
|
||||
|
||||
// 处理assistant消息:提取tool_calls信息
|
||||
if role == "assistant" {
|
||||
if toolCalls, ok := msg["tool_calls"].([]interface{}); ok && len(toolCalls) > 0 {
|
||||
// 如果有文本内容,先显示
|
||||
if content != "" {
|
||||
builder.WriteString(fmt.Sprintf("[%s]: %s\n", role, content))
|
||||
}
|
||||
// 详细显示每个工具调用
|
||||
builder.WriteString(fmt.Sprintf("[%s] 工具调用 (%d个):\n", role, len(toolCalls)))
|
||||
for i, toolCall := range toolCalls {
|
||||
if tc, ok := toolCall.(map[string]interface{}); ok {
|
||||
toolCallID, _ := tc["id"].(string)
|
||||
if funcData, ok := tc["function"].(map[string]interface{}); ok {
|
||||
toolName, _ := funcData["name"].(string)
|
||||
arguments, _ := funcData["arguments"].(string)
|
||||
builder.WriteString(fmt.Sprintf(" [工具调用 %d]\n", i+1))
|
||||
builder.WriteString(fmt.Sprintf(" ID: %s\n", toolCallID))
|
||||
builder.WriteString(fmt.Sprintf(" 工具名称: %s\n", toolName))
|
||||
builder.WriteString(fmt.Sprintf(" 参数: %s\n", arguments))
|
||||
}
|
||||
if strings.EqualFold(role, "assistant") && len(msg.ToolCalls) > 0 {
|
||||
if content != "" {
|
||||
builder.WriteString(fmt.Sprintf("[%s]: %s\n", role, content))
|
||||
}
|
||||
builder.WriteString(fmt.Sprintf("[%s] 工具调用 (%d个):\n", role, len(msg.ToolCalls)))
|
||||
for i, tc := range msg.ToolCalls {
|
||||
args := ""
|
||||
if tc.Function.Arguments != nil {
|
||||
if b, err := json.Marshal(tc.Function.Arguments); err == nil {
|
||||
args = string(b)
|
||||
}
|
||||
}
|
||||
builder.WriteString("\n")
|
||||
continue
|
||||
builder.WriteString(fmt.Sprintf(" [工具调用 %d]\n", i+1))
|
||||
builder.WriteString(fmt.Sprintf(" ID: %s\n", tc.ID))
|
||||
builder.WriteString(fmt.Sprintf(" 工具名称: %s\n", tc.Function.Name))
|
||||
builder.WriteString(fmt.Sprintf(" 参数: %s\n", args))
|
||||
}
|
||||
builder.WriteString("\n")
|
||||
continue
|
||||
}
|
||||
|
||||
// 处理tool消息:显示tool_call_id和完整内容
|
||||
if role == "tool" {
|
||||
toolCallID, _ := msg["tool_call_id"].(string)
|
||||
if toolCallID != "" {
|
||||
builder.WriteString(fmt.Sprintf("[%s] (tool_call_id: %s):\n%s\n\n", role, toolCallID, content))
|
||||
if strings.EqualFold(role, "tool") {
|
||||
if msg.ToolCallID != "" {
|
||||
builder.WriteString(fmt.Sprintf("[%s] (tool_call_id: %s):\n%s\n\n", role, msg.ToolCallID, content))
|
||||
} else {
|
||||
builder.WriteString(fmt.Sprintf("[%s]: %s\n\n", role, content))
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
// 其他消息类型(system, user等)正常显示
|
||||
builder.WriteString(fmt.Sprintf("[%s]: %s\n\n", role, content))
|
||||
}
|
||||
|
||||
return builder.String()
|
||||
}
|
||||
|
||||
// buildSimplePrompt 构建简化的prompt
|
||||
func (b *Builder) buildSimplePrompt(reactInput, modelOutput string) string {
|
||||
return fmt.Sprintf(`你是专业的安全测试分析师和攻击链构建专家。你的任务是根据对话记录和工具执行结果,构建一个逻辑清晰、有教育意义的攻击链图,完整展现渗透测试的思维过程和执行路径。
|
||||
return fmt.Sprintf(`你是专业的安全测试分析师和攻击链构建专家。你的任务是根据**当前任务轮次**的对话记录和工具执行结果,一次性输出攻击链 JSON(不要分多轮追问)。
|
||||
|
||||
## 输入范围(与「继续对话」续跑一致)
|
||||
- 下方「ReAct 轨迹」仅包含**最后一次用户提问之后**的消息与工具结果(last_react 当前任务轮次),不含更早的用户提问轮次。
|
||||
- 「助手结论」为同轮任务的最终输出摘要(last_react_output);节点须与轨迹中的实际工具执行一致,严禁编造。
|
||||
|
||||
## 核心目标
|
||||
|
||||
@@ -618,12 +632,9 @@ func (b *Builder) buildSimplePrompt(reactInput, modelOutput string) string {
|
||||
5. **漏洞确认**:如何确认漏洞存在?(action→vulnerability)
|
||||
6. **攻击路径**:完整的攻击路径是什么?(从target到vulnerability的路径)
|
||||
|
||||
## 最后一轮ReAct输入
|
||||
## 当前任务 ReAct 轨迹(含工具执行;助手结论见轨迹末尾 assistant)
|
||||
|
||||
%s
|
||||
|
||||
## 大模型输出
|
||||
|
||||
%s
|
||||
|
||||
## 输出格式
|
||||
@@ -752,7 +763,15 @@ func (b *Builder) buildSimplePrompt(reactInput, modelOutput string) string {
|
||||
9. **不要过度精简**:如果实际执行步骤较多,可以适当增加节点数量(最多20个),确保不遗漏关键步骤。
|
||||
10. **输出前验证**:在输出JSON前,必须验证所有边都满足source < target的条件,确保DAG结构正确。
|
||||
|
||||
现在开始分析并构建攻击链:`, reactInput, modelOutput)
|
||||
现在开始分析并构建攻击链:`, reactInput, assistantOutSection(modelOutput))
|
||||
}
|
||||
|
||||
func assistantOutSection(modelOutput string) string {
|
||||
modelOutput = strings.TrimSpace(modelOutput)
|
||||
if modelOutput == "" {
|
||||
return ""
|
||||
}
|
||||
return "\n## 助手结论(补充)\n\n" + modelOutput + "\n"
|
||||
}
|
||||
|
||||
// saveChain 保存攻击链到数据库
|
||||
@@ -812,7 +831,7 @@ func (b *Builder) callAIForChainGeneration(ctx context.Context, prompt string) (
|
||||
},
|
||||
},
|
||||
"temperature": 0.3,
|
||||
"max_completion_tokens": 80000,
|
||||
"max_completion_tokens": attackChainMaxCompletionTokens(b.maxTokens),
|
||||
}
|
||||
|
||||
var apiResponse struct {
|
||||
|
||||
@@ -0,0 +1,248 @@
|
||||
package attackchain
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const (
|
||||
attackChainTruncationMarker = "\n\n...[攻击链输入已截断 / attack chain input truncated]...\n\n"
|
||||
attackChainSystemReserve = 256
|
||||
attackChainSafetyReserve = 2048
|
||||
)
|
||||
|
||||
// attackChainMaxCompletionTokens 为攻击链 JSON 输出预留的 completion token 上限。
|
||||
func attackChainMaxCompletionTokens(maxTotal int) int {
|
||||
const capTokens = 16384
|
||||
if maxTotal <= 0 {
|
||||
return 8192
|
||||
}
|
||||
v := maxTotal / 8
|
||||
if v < 4096 {
|
||||
v = 4096
|
||||
}
|
||||
if v > capTokens {
|
||||
v = capTokens
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
func (b *Builder) modelName() string {
|
||||
if b.openAIConfig != nil && b.openAIConfig.Model != "" {
|
||||
return b.openAIConfig.Model
|
||||
}
|
||||
return "gpt-4"
|
||||
}
|
||||
|
||||
func (b *Builder) countTokens(text string) int {
|
||||
if text == "" {
|
||||
return 0
|
||||
}
|
||||
n, err := b.tokenCounter.Count(b.modelName(), text)
|
||||
if err != nil {
|
||||
return utf8.RuneCountInString(text) / 4
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
// attackChainPayloadTokenBudget 计算 reactInput + modelOutput 可用的 token 预算。
|
||||
func (b *Builder) attackChainPayloadTokenBudget() int {
|
||||
maxTotal := b.maxTokens
|
||||
if maxTotal <= 0 {
|
||||
maxTotal = 100000
|
||||
}
|
||||
templateTok := b.countTokens(b.buildSimplePrompt("", ""))
|
||||
completion := attackChainMaxCompletionTokens(maxTotal)
|
||||
reserve := templateTok + attackChainSystemReserve + completion + attackChainSafetyReserve
|
||||
budget := maxTotal - reserve
|
||||
minBudget := maxTotal * 35 / 100
|
||||
if budget < minBudget {
|
||||
budget = minBudget
|
||||
}
|
||||
if budget < 4096 {
|
||||
budget = 4096
|
||||
}
|
||||
return budget
|
||||
}
|
||||
|
||||
// fitAttackChainPayload 在构建最终 prompt 前压缩 ReAct 轨迹与模型输出,避免超出模型上下文。
|
||||
func (b *Builder) fitAttackChainPayload(reactInput, modelOutput string) (string, string, bool) {
|
||||
budget := b.attackChainPayloadTokenBudget()
|
||||
modelBudget := budget * 15 / 100
|
||||
if modelBudget < 512 {
|
||||
modelBudget = 512
|
||||
}
|
||||
reactBudget := budget - modelBudget
|
||||
|
||||
origReactTok := b.countTokens(reactInput)
|
||||
origModelTok := b.countTokens(modelOutput)
|
||||
truncated := false
|
||||
|
||||
outModel := modelOutput
|
||||
if origModelTok > modelBudget {
|
||||
outModel = truncateTextByTokens(b, modelOutput, modelBudget)
|
||||
truncated = true
|
||||
}
|
||||
|
||||
outReact := reactInput
|
||||
perToolLimits := []int{12000, 6000, 3000, 1500, 800}
|
||||
for _, lim := range perToolLimits {
|
||||
compact := compactFormattedToolBodies(outReact, lim)
|
||||
if compact != outReact {
|
||||
outReact = compact
|
||||
truncated = true
|
||||
}
|
||||
if b.countTokens(outReact) <= reactBudget {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if b.countTokens(outReact) > reactBudget {
|
||||
outReact = truncateTextByTokens(b, outReact, reactBudget)
|
||||
truncated = true
|
||||
}
|
||||
|
||||
if truncated {
|
||||
b.logger.Info("攻击链输入已按 token 预算截断",
|
||||
zap.Int("maxTotalTokens", b.maxTokens),
|
||||
zap.Int("payloadBudget", budget),
|
||||
zap.Int("reactBudget", reactBudget),
|
||||
zap.Int("modelBudget", modelBudget),
|
||||
zap.Int("reactInputTokensBefore", origReactTok),
|
||||
zap.Int("reactInputTokensAfter", b.countTokens(outReact)),
|
||||
zap.Int("modelOutputTokensBefore", origModelTok),
|
||||
zap.Int("modelOutputTokensAfter", b.countTokens(outModel)),
|
||||
zap.Int("maxCompletionTokens", attackChainMaxCompletionTokens(b.maxTokens)),
|
||||
)
|
||||
}
|
||||
|
||||
return outReact, outModel, truncated
|
||||
}
|
||||
|
||||
// compactFormattedToolBodies 缩短格式化 trace 中 [tool] 消息的正文,保留工具头与调用 ID。
|
||||
func compactFormattedToolBodies(s string, maxRunesPerBody int) string {
|
||||
if maxRunesPerBody <= 0 || s == "" {
|
||||
return s
|
||||
}
|
||||
const marker = "[tool]"
|
||||
var out strings.Builder
|
||||
remaining := s
|
||||
changed := false
|
||||
for {
|
||||
idx := strings.Index(remaining, marker)
|
||||
if idx < 0 {
|
||||
out.WriteString(remaining)
|
||||
break
|
||||
}
|
||||
out.WriteString(remaining[:idx])
|
||||
remaining = remaining[idx:]
|
||||
nl := strings.IndexByte(remaining, '\n')
|
||||
if nl < 0 {
|
||||
out.WriteString(remaining)
|
||||
break
|
||||
}
|
||||
header := remaining[:nl+1]
|
||||
remaining = remaining[nl+1:]
|
||||
bodyEnd := strings.Index(remaining, "\n\n[")
|
||||
var body, rest string
|
||||
if bodyEnd < 0 {
|
||||
body = remaining
|
||||
rest = ""
|
||||
} else {
|
||||
body = remaining[:bodyEnd]
|
||||
rest = remaining[bodyEnd:]
|
||||
}
|
||||
if runeLen(body) > maxRunesPerBody {
|
||||
body = truncateRunesWithNotice(body, maxRunesPerBody)
|
||||
changed = true
|
||||
}
|
||||
out.WriteString(header)
|
||||
out.WriteString(body)
|
||||
remaining = rest
|
||||
if rest == "" {
|
||||
break
|
||||
}
|
||||
}
|
||||
if !changed {
|
||||
return s
|
||||
}
|
||||
return out.String()
|
||||
}
|
||||
|
||||
func truncateTextByTokens(b *Builder, text string, maxTokens int) string {
|
||||
if maxTokens <= 0 || text == "" {
|
||||
return ""
|
||||
}
|
||||
if b.countTokens(text) <= maxTokens {
|
||||
return text
|
||||
}
|
||||
markerTok := b.countTokens(attackChainTruncationMarker)
|
||||
usable := maxTokens - markerTok
|
||||
if usable < 256 {
|
||||
usable = maxTokens / 2
|
||||
}
|
||||
headBudget := usable * 60 / 100
|
||||
tailBudget := usable - headBudget
|
||||
head := takeTokensFromStart(b, text, headBudget)
|
||||
tail := takeTokensFromEnd(b, text, tailBudget)
|
||||
return head + attackChainTruncationMarker + tail
|
||||
}
|
||||
|
||||
func takeTokensFromStart(b *Builder, text string, maxTokens int) string {
|
||||
rs := []rune(text)
|
||||
if len(rs) == 0 || maxTokens <= 0 {
|
||||
return ""
|
||||
}
|
||||
lo, hi := 0, len(rs)
|
||||
for lo < hi {
|
||||
mid := (lo + hi + 1) / 2
|
||||
if b.countTokens(string(rs[:mid])) <= maxTokens {
|
||||
lo = mid
|
||||
} else {
|
||||
hi = mid - 1
|
||||
}
|
||||
}
|
||||
return string(rs[:lo])
|
||||
}
|
||||
|
||||
func takeTokensFromEnd(b *Builder, text string, maxTokens int) string {
|
||||
rs := []rune(text)
|
||||
if len(rs) == 0 || maxTokens <= 0 {
|
||||
return ""
|
||||
}
|
||||
lo, hi := 0, len(rs)
|
||||
for lo < hi {
|
||||
mid := (lo + hi) / 2
|
||||
if b.countTokens(string(rs[mid:])) <= maxTokens {
|
||||
hi = mid
|
||||
} else {
|
||||
lo = mid + 1
|
||||
}
|
||||
}
|
||||
return string(rs[lo:])
|
||||
}
|
||||
|
||||
func truncateRunesWithNotice(s string, maxRunes int) string {
|
||||
rs := []rune(s)
|
||||
if len(rs) <= maxRunes {
|
||||
return s
|
||||
}
|
||||
const notice = "\n...[工具输出已截断 / tool output truncated]...\n"
|
||||
noticeRunes := []rune(notice)
|
||||
keep := maxRunes - len(noticeRunes)
|
||||
if keep < 200 {
|
||||
keep = maxRunes * 2 / 3
|
||||
}
|
||||
if keep < 1 {
|
||||
return notice
|
||||
}
|
||||
head := keep * 70 / 100
|
||||
tail := keep - head
|
||||
return string(rs[:head]) + notice + string(rs[len(rs)-tail:])
|
||||
}
|
||||
|
||||
func runeLen(s string) int {
|
||||
return len([]rune(s))
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
package attackchain
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/config"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func testBuilder(maxTotal int) *Builder {
|
||||
return &Builder{
|
||||
logger: zap.NewNop(),
|
||||
openAIConfig: &config.OpenAIConfig{Model: "gpt-4"},
|
||||
tokenCounter: agent.NewTikTokenCounter(),
|
||||
maxTokens: maxTotal,
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompactFormattedToolBodies(t *testing.T) {
|
||||
long := strings.Repeat("x", 20000)
|
||||
in := "[user]: hi\n\n[tool] (tool_call_id: abc):\n" + long + "\n\n[assistant]: done\n"
|
||||
out := compactFormattedToolBodies(in, 500)
|
||||
if strings.Contains(out, strings.Repeat("x", 10000)) {
|
||||
t.Fatal("expected tool body to be truncated")
|
||||
}
|
||||
if !strings.Contains(out, "[user]: hi") {
|
||||
t.Fatal("expected user header preserved")
|
||||
}
|
||||
if !strings.Contains(out, "[assistant]: done") {
|
||||
t.Fatal("expected assistant header preserved")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFitAttackChainPayloadWithinBudget(t *testing.T) {
|
||||
b := testBuilder(32000)
|
||||
react := strings.Repeat("scan ", 50000)
|
||||
model := strings.Repeat("result ", 10000)
|
||||
r, m, truncated := b.fitAttackChainPayload(react, model)
|
||||
if !truncated {
|
||||
t.Fatal("expected truncation for large payload")
|
||||
}
|
||||
prompt := b.buildSimplePrompt(r, m)
|
||||
total := b.countTokens(prompt) + attackChainMaxCompletionTokens(b.maxTokens) + attackChainSystemReserve
|
||||
if total > b.maxTokens+attackChainSafetyReserve {
|
||||
t.Fatalf("prompt still too large: estimated %d > max %d", total, b.maxTokens)
|
||||
}
|
||||
_ = m
|
||||
}
|
||||
|
||||
func TestAttackChainMaxCompletionTokens(t *testing.T) {
|
||||
if got := attackChainMaxCompletionTokens(120000); got != 15000 && got != 16384 {
|
||||
// 120000/8 = 15000
|
||||
if got < 4096 || got > 16384 {
|
||||
t.Fatalf("unexpected completion cap: %d", got)
|
||||
}
|
||||
}
|
||||
if got := attackChainMaxCompletionTokens(0); got != 8192 {
|
||||
t.Fatalf("expected default 8192, got %d", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,55 @@
|
||||
package audit
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// RegisterConversationCreateHook records platform audit rows for every new conversation.
|
||||
func RegisterConversationCreateHook(s *Service) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
database.SetConversationCreateHook(func(conv *database.Conversation, meta database.ConversationCreateMeta) {
|
||||
detail := map[string]interface{}{
|
||||
"title": conv.Title,
|
||||
"source": meta.Source,
|
||||
}
|
||||
if meta.WebShellConnectionID != "" {
|
||||
detail["webshell_connection_id"] = meta.WebShellConnectionID
|
||||
}
|
||||
s.Record(nil, Entry{
|
||||
Category: "conversation",
|
||||
Action: "create",
|
||||
Result: "success",
|
||||
Message: "创建对话",
|
||||
ResourceType: "conversation",
|
||||
ResourceID: conv.ID,
|
||||
Detail: detail,
|
||||
ClientIP: meta.ClientIP,
|
||||
SessionHint: meta.SessionHint,
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
// ConversationCreateMeta builds audit metadata for conversation creation.
|
||||
func ConversationCreateMeta(source string) database.ConversationCreateMeta {
|
||||
return database.ConversationCreateMeta{Source: strings.TrimSpace(source)}
|
||||
}
|
||||
|
||||
// ConversationCreateMetaFromGin includes client IP and session hint when available.
|
||||
func ConversationCreateMetaFromGin(c *gin.Context, source string) database.ConversationCreateMeta {
|
||||
m := ConversationCreateMeta(source)
|
||||
if c == nil {
|
||||
return m
|
||||
}
|
||||
m.ClientIP = c.ClientIP()
|
||||
if token := c.GetString(security.ContextAuthTokenKey); token != "" {
|
||||
m.SessionHint = sessionHint(token)
|
||||
}
|
||||
return m
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
package audit
|
||||
|
||||
// RetentionDays returns configured retention; 0 means keep forever.
|
||||
func (s *Service) RetentionDays() int {
|
||||
if s == nil || s.cfg == nil {
|
||||
return 0
|
||||
}
|
||||
return s.cfg.Audit.RetentionDaysEffective()
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
package audit
|
||||
|
||||
import "github.com/gin-gonic/gin"
|
||||
|
||||
// RecordAction writes a platform audit row with common defaults.
|
||||
func (s *Service) RecordAction(c *gin.Context, category, action, result, message, resourceType, resourceID string, detail map[string]interface{}) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.Record(c, Entry{
|
||||
Category: category,
|
||||
Action: action,
|
||||
Result: result,
|
||||
Message: message,
|
||||
ResourceType: resourceType,
|
||||
ResourceID: resourceID,
|
||||
Detail: detail,
|
||||
})
|
||||
}
|
||||
|
||||
// RecordOK is a shorthand for successful operations.
|
||||
func (s *Service) RecordOK(c *gin.Context, category, action, message, resourceType, resourceID string, detail map[string]interface{}) {
|
||||
s.RecordAction(c, category, action, "success", message, resourceType, resourceID, detail)
|
||||
}
|
||||
|
||||
// RecordFail is a shorthand for failed operations.
|
||||
func (s *Service) RecordFail(c *gin.Context, category, action, message string, detail map[string]interface{}) {
|
||||
s.RecordAction(c, category, action, "failure", message, "", "", detail)
|
||||
}
|
||||
@@ -0,0 +1,86 @@
|
||||
package audit
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
)
|
||||
|
||||
var auditActionsResourceRemoved = map[string]bool{
|
||||
"delete": true,
|
||||
"item_delete": true,
|
||||
"connection_delete": true,
|
||||
"listener_delete": true,
|
||||
"session_delete": true,
|
||||
"task_delete": true,
|
||||
"execution_delete": true,
|
||||
"execution_delete_batch": true,
|
||||
"delete_queue": true,
|
||||
"delete_batch_task": true,
|
||||
"markdown_delete": true,
|
||||
}
|
||||
|
||||
// ApplyResourceAvailability sets log.ResourceAvailable when the linked resource can be checked.
|
||||
func ApplyResourceAvailability(db *database.DB, log *database.AuditLog) {
|
||||
if log == nil || strings.TrimSpace(log.ResourceID) == "" {
|
||||
return
|
||||
}
|
||||
if auditActionsResourceRemoved[log.Action] {
|
||||
f := false
|
||||
log.ResourceAvailable = &f
|
||||
return
|
||||
}
|
||||
if db == nil {
|
||||
return
|
||||
}
|
||||
available, known := resourceStillExists(db, log.ResourceType, log.ResourceID)
|
||||
if known {
|
||||
log.ResourceAvailable = &available
|
||||
}
|
||||
}
|
||||
|
||||
func resourceStillExists(db *database.DB, resourceType, resourceID string) (bool, bool) {
|
||||
resourceID = strings.TrimSpace(resourceID)
|
||||
if resourceID == "" {
|
||||
return false, false
|
||||
}
|
||||
t := strings.TrimSpace(resourceType)
|
||||
if t == "" {
|
||||
if len(resourceID) > 8 && !strings.HasPrefix(resourceID, "c2_") {
|
||||
t = "conversation"
|
||||
} else {
|
||||
return false, false
|
||||
}
|
||||
}
|
||||
switch t {
|
||||
case "conversation":
|
||||
ok, err := db.ConversationExists(resourceID)
|
||||
return ok, err == nil
|
||||
case "vulnerability":
|
||||
_, err := db.GetVulnerability(resourceID)
|
||||
if err != nil {
|
||||
return false, strings.Contains(err.Error(), "不存在")
|
||||
}
|
||||
return true, true
|
||||
case "batch_queue":
|
||||
_, err := db.GetBatchQueue(resourceID)
|
||||
return err == nil, true
|
||||
case "c2_listener":
|
||||
_, err := db.GetC2Listener(resourceID)
|
||||
return err == nil, true
|
||||
case "c2_session":
|
||||
_, err := db.GetC2Session(resourceID)
|
||||
return err == nil, true
|
||||
case "c2_task":
|
||||
_, err := db.GetC2Task(resourceID)
|
||||
return err == nil, true
|
||||
case "webshell_connection":
|
||||
c, err := db.GetWebshellConnection(resourceID)
|
||||
return err == nil && c != nil, true
|
||||
case "tool_execution":
|
||||
_, err := db.GetToolExecution(resourceID)
|
||||
return err == nil, true
|
||||
default:
|
||||
return false, false
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
package audit
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// auditRetentionPurgeInterval is how often PurgeExpired runs while the process is up (startup also purges once).
|
||||
const auditRetentionPurgeInterval = time.Hour
|
||||
|
||||
// StartRetentionLoop periodically purges expired audit rows.
|
||||
func StartRetentionLoop(s *Service, logger *zap.Logger) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
go func() {
|
||||
ticker := time.NewTicker(auditRetentionPurgeInterval)
|
||||
defer ticker.Stop()
|
||||
for range ticker.C {
|
||||
s.PurgeExpired()
|
||||
if logger != nil {
|
||||
logger.Debug("audit retention tick completed")
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
package audit
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
)
|
||||
|
||||
var sensitiveKeySubstrings = []string{
|
||||
"password", "api_key", "apikey", "secret", "token", "authorization",
|
||||
"credential", "private_key", "access_key",
|
||||
}
|
||||
|
||||
// SanitizeDetail redacts sensitive keys and truncates serialized size.
|
||||
func SanitizeDetail(detail map[string]interface{}, maxBytes int) map[string]interface{} {
|
||||
if detail == nil {
|
||||
return nil
|
||||
}
|
||||
if maxBytes <= 0 {
|
||||
maxBytes = 8192
|
||||
}
|
||||
out := sanitizeValue("", detail)
|
||||
if m, ok := out.(map[string]interface{}); ok {
|
||||
b, _ := json.Marshal(m)
|
||||
if len(b) > maxBytes {
|
||||
return map[string]interface{}{
|
||||
"_truncated": true,
|
||||
"_preview": string(b[:maxBytes]),
|
||||
}
|
||||
}
|
||||
return m
|
||||
}
|
||||
return map[string]interface{}{"value": out}
|
||||
}
|
||||
|
||||
func sanitizeValue(key string, v interface{}) interface{} {
|
||||
kl := strings.ToLower(key)
|
||||
for _, sub := range sensitiveKeySubstrings {
|
||||
if strings.Contains(kl, sub) {
|
||||
return "***"
|
||||
}
|
||||
}
|
||||
switch t := v.(type) {
|
||||
case map[string]interface{}:
|
||||
m := make(map[string]interface{}, len(t))
|
||||
for k, val := range t {
|
||||
m[k] = sanitizeValue(k, val)
|
||||
}
|
||||
return m
|
||||
case []interface{}:
|
||||
arr := make([]interface{}, len(t))
|
||||
for i, val := range t {
|
||||
arr[i] = sanitizeValue(key, val)
|
||||
}
|
||||
return arr
|
||||
default:
|
||||
return v
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,172 @@
|
||||
package audit
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// Service persists platform audit logs.
|
||||
type Service struct {
|
||||
db *database.DB
|
||||
cfg *config.Config
|
||||
logger *zap.Logger
|
||||
failThrottle *failureThrottle
|
||||
}
|
||||
|
||||
// NewService creates an audit service.
|
||||
func NewService(db *database.DB, cfg *config.Config, logger *zap.Logger) *Service {
|
||||
return &Service{
|
||||
db: db,
|
||||
cfg: cfg,
|
||||
logger: logger,
|
||||
failThrottle: newFailureThrottle(),
|
||||
}
|
||||
}
|
||||
|
||||
// Enabled reports whether audit persistence is on.
|
||||
func (s *Service) Enabled() bool {
|
||||
if s == nil || s.cfg == nil {
|
||||
return false
|
||||
}
|
||||
return s.cfg.Audit.EnabledEffective()
|
||||
}
|
||||
|
||||
// Record writes one audit row from a Gin request context.
|
||||
func (s *Service) Record(c *gin.Context, e Entry) {
|
||||
if s == nil || !s.Enabled() || s.db == nil {
|
||||
return
|
||||
}
|
||||
if strings.TrimSpace(e.Category) == "" || strings.TrimSpace(e.Action) == "" {
|
||||
return
|
||||
}
|
||||
if e.Result == "failure" && !s.allowFailureAudit(c, e) {
|
||||
return
|
||||
}
|
||||
if strings.TrimSpace(e.Result) == "" {
|
||||
e.Result = "success"
|
||||
}
|
||||
if strings.TrimSpace(e.Level) == "" {
|
||||
if e.Result == "failure" {
|
||||
e.Level = "warn"
|
||||
} else {
|
||||
e.Level = "info"
|
||||
}
|
||||
}
|
||||
if strings.TrimSpace(e.Actor) == "" {
|
||||
e.Actor = "admin"
|
||||
}
|
||||
maxDetail := s.cfg.Audit.MaxDetailBytesEffective()
|
||||
detail := SanitizeDetail(e.Detail, maxDetail)
|
||||
|
||||
sessionHintVal := e.SessionHint
|
||||
if sessionHintVal == "" && c != nil {
|
||||
if token := c.GetString(security.ContextAuthTokenKey); token != "" {
|
||||
sessionHintVal = sessionHint(token)
|
||||
}
|
||||
}
|
||||
clientIPVal := e.ClientIP
|
||||
if clientIPVal == "" {
|
||||
clientIPVal = clientIP(c)
|
||||
}
|
||||
|
||||
row := &database.AuditLog{
|
||||
ID: "audit_" + strings.ReplaceAll(uuid.New().String(), "-", ""),
|
||||
CreatedAt: time.Now(),
|
||||
Level: e.Level,
|
||||
Category: e.Category,
|
||||
Action: e.Action,
|
||||
Result: e.Result,
|
||||
Actor: e.Actor,
|
||||
SessionHint: sessionHintVal,
|
||||
ClientIP: clientIPVal,
|
||||
UserAgent: userAgent(c),
|
||||
ResourceType: e.ResourceType,
|
||||
ResourceID: e.ResourceID,
|
||||
Message: e.Message,
|
||||
Detail: detail,
|
||||
}
|
||||
if err := s.db.AppendAuditLog(row); err != nil && s.logger != nil {
|
||||
s.logger.Warn("写入审计日志失败",
|
||||
zap.String("action", e.Action),
|
||||
zap.Error(err),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
// RecordSystem writes an audit row without HTTP context (e.g. retention cleanup).
|
||||
func (s *Service) RecordSystem(e Entry) {
|
||||
s.Record(nil, e)
|
||||
}
|
||||
|
||||
// PurgeExpired deletes rows older than retention_days when configured.
|
||||
func (s *Service) PurgeExpired() {
|
||||
if s == nil || s.db == nil || s.cfg == nil {
|
||||
return
|
||||
}
|
||||
days := s.cfg.Audit.RetentionDaysEffective()
|
||||
if days <= 0 {
|
||||
return
|
||||
}
|
||||
cutoff := time.Now().AddDate(0, 0, -days)
|
||||
n, err := s.db.DeleteAuditLogsBefore(cutoff)
|
||||
if err != nil {
|
||||
if s.logger != nil {
|
||||
s.logger.Warn("清理过期审计日志失败", zap.Error(err))
|
||||
}
|
||||
return
|
||||
}
|
||||
if n > 0 && s.logger != nil {
|
||||
s.logger.Info("已清理过期审计日志", zap.Int64("deleted", n))
|
||||
}
|
||||
}
|
||||
|
||||
// HintFromToken returns a short stable hash prefix for a session token.
|
||||
func HintFromToken(token string) string {
|
||||
return sessionHint(token)
|
||||
}
|
||||
|
||||
func sessionHint(token string) string {
|
||||
token = strings.TrimSpace(token)
|
||||
if token == "" {
|
||||
return ""
|
||||
}
|
||||
sum := sha256.Sum256([]byte(token))
|
||||
return hex.EncodeToString(sum[:4])
|
||||
}
|
||||
|
||||
func (s *Service) allowFailureAudit(c *gin.Context, e Entry) bool {
|
||||
if !isAuthFailureThrottled(e.Category, e.Action) {
|
||||
return true
|
||||
}
|
||||
cooldown := time.Duration(s.cfg.Audit.AuthFailureCooldownEffective()) * time.Second
|
||||
key := authFailureThrottleKey(e.Category, e.Action, clientIP(c))
|
||||
return s.failThrottle.allow(key, cooldown)
|
||||
}
|
||||
|
||||
func clientIP(c *gin.Context) string {
|
||||
if c == nil {
|
||||
return ""
|
||||
}
|
||||
return c.ClientIP()
|
||||
}
|
||||
|
||||
func userAgent(c *gin.Context) string {
|
||||
if c == nil {
|
||||
return ""
|
||||
}
|
||||
ua := c.GetHeader("User-Agent")
|
||||
if len(ua) > 512 {
|
||||
return ua[:512]
|
||||
}
|
||||
return ua
|
||||
}
|
||||
@@ -0,0 +1,55 @@
|
||||
package audit
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// failureThrottle deduplicates high-frequency failure audit rows (e.g. wrong password).
|
||||
type failureThrottle struct {
|
||||
mu sync.Mutex
|
||||
last map[string]time.Time
|
||||
}
|
||||
|
||||
func newFailureThrottle() *failureThrottle {
|
||||
return &failureThrottle{last: make(map[string]time.Time)}
|
||||
}
|
||||
|
||||
// allow reports whether a row with the given key may be written now.
|
||||
func (t *failureThrottle) allow(key string, cooldown time.Duration) bool {
|
||||
if t == nil || cooldown <= 0 || key == "" {
|
||||
return true
|
||||
}
|
||||
now := time.Now()
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
if prev, ok := t.last[key]; ok && now.Sub(prev) < cooldown {
|
||||
return false
|
||||
}
|
||||
t.last[key] = now
|
||||
if len(t.last) > 4096 {
|
||||
for k, ts := range t.last {
|
||||
if now.Sub(ts) > cooldown*2 {
|
||||
delete(t.last, k)
|
||||
}
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// authFailureThrottleKey builds a per-IP key for auth failure deduplication.
|
||||
func authFailureThrottleKey(category, action, clientIP string) string {
|
||||
return category + ":" + action + ":" + clientIP
|
||||
}
|
||||
|
||||
func isAuthFailureThrottled(category, action string) bool {
|
||||
if category != "auth" {
|
||||
return false
|
||||
}
|
||||
switch action {
|
||||
case "login", "change_password":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
package audit
|
||||
|
||||
// Entry describes one platform audit record (not chat/tool execution bodies).
|
||||
type Entry struct {
|
||||
Level string
|
||||
Category string
|
||||
Action string
|
||||
Result string // success | failure
|
||||
Actor string
|
||||
SessionHint string
|
||||
ResourceType string
|
||||
ResourceID string
|
||||
Message string
|
||||
Detail map[string]interface{}
|
||||
ClientIP string // optional when c is nil (robot, batch, DB hook)
|
||||
}
|
||||
@@ -239,13 +239,15 @@ func (m *Manager) StartListener(id string) (*database.C2Listener, error) {
|
||||
}
|
||||
cfg.ApplyDefaults()
|
||||
|
||||
// 通过工厂创建具体实现
|
||||
// 通过工厂创建具体实现。必须使用 rec 的副本:HTTP handler 在返回 JSON 前会清空
|
||||
// rec.ImplantToken / EncryptionKey 做脱敏,若 listener 实现持有同一指针会导致 beacon 鉴权永久失败。
|
||||
listenerRec := *rec
|
||||
factory := m.registry.Get(rec.Type)
|
||||
if factory == nil {
|
||||
return nil, ErrUnsupportedType
|
||||
}
|
||||
inst, err := factory(ListenerCreationCtx{
|
||||
Listener: rec,
|
||||
Listener: &listenerRec,
|
||||
Config: cfg,
|
||||
Manager: m,
|
||||
Logger: m.logger.With(zap.String("listener_id", rec.ID), zap.String("type", rec.Type)),
|
||||
|
||||
@@ -0,0 +1,74 @@
|
||||
package c2
|
||||
|
||||
import (
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// 回归:StartListener 返回的 rec 被 handler 脱敏清空 ImplantToken 后,运行中的 HTTP listener 仍能鉴权。
|
||||
func TestStartListener_ImplantTokenSurvivesHandlerRedaction(t *testing.T) {
|
||||
tmp := t.TempDir()
|
||||
db, err := database.NewDB(filepath.Join(tmp, "c2.sqlite"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
|
||||
lnPick, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
port := lnPick.Addr().(*net.TCPAddr).Port
|
||||
_ = lnPick.Close()
|
||||
|
||||
mgr := NewManager(db, zap.NewNop(), tmp)
|
||||
mgr.Registry().Register(string(ListenerTypeHTTPBeacon), NewHTTPBeaconListener)
|
||||
rec, err := mgr.CreateListener(CreateListenerInput{
|
||||
Name: "t",
|
||||
Type: string(ListenerTypeHTTPBeacon),
|
||||
BindHost: "127.0.0.1",
|
||||
BindPort: port,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
token := rec.ImplantToken
|
||||
|
||||
rec, err = mgr.StartListener(rec.ID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// 模拟 internal/handler/c2.go StartListener 在 JSON 响应前的脱敏
|
||||
rec.ImplantToken = ""
|
||||
rec.EncryptionKey = ""
|
||||
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
body := `{"hostname":"n","username":"u","os":"Linux","arch":"amd64","internal_ip":"10.0.0.1","pid":42}`
|
||||
req, _ := http.NewRequest(http.MethodPost, "http://127.0.0.1:"+strconv.Itoa(port)+"/check_in", strings.NewReader(body))
|
||||
req.Header.Set("X-Implant-Token", token)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
b, _ := io.ReadAll(resp.Body)
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
t.Fatalf("status=%d body=%s", resp.StatusCode, b)
|
||||
}
|
||||
if !strings.Contains(string(b), "session_id") {
|
||||
t.Fatalf("expected session_id in body: %s", b)
|
||||
}
|
||||
_ = mgr.StopListener(rec.ID)
|
||||
}
|
||||
+123
-11
@@ -26,6 +26,7 @@ type Config struct {
|
||||
Security SecurityConfig `yaml:"security"`
|
||||
Database DatabaseConfig `yaml:"database"`
|
||||
Auth AuthConfig `yaml:"auth"`
|
||||
Audit AuditConfig `yaml:"audit,omitempty" json:"audit,omitempty"`
|
||||
ExternalMCP ExternalMCPConfig `yaml:"external_mcp,omitempty"`
|
||||
Knowledge KnowledgeConfig `yaml:"knowledge,omitempty"`
|
||||
C2 C2Config `yaml:"c2,omitempty" json:"c2,omitempty"` // 内置 C2 总开关;未配置时默认启用
|
||||
@@ -35,13 +36,39 @@ type Config struct {
|
||||
SkillsDir string `yaml:"skills_dir,omitempty" json:"skills_dir,omitempty"` // Skills配置文件目录
|
||||
AgentsDir string `yaml:"agents_dir,omitempty" json:"agents_dir,omitempty"` // 多代理子 Agent Markdown 定义目录(*.md,YAML front matter)
|
||||
MultiAgent MultiAgentConfig `yaml:"multi_agent,omitempty" json:"multi_agent,omitempty"`
|
||||
Project ProjectConfig `yaml:"project,omitempty" json:"project,omitempty"`
|
||||
}
|
||||
|
||||
// ProjectConfig 项目黑板(跨对话共享事实)配置。
|
||||
type ProjectConfig struct {
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
DefaultProjectID string `yaml:"default_project_id,omitempty" json:"default_project_id,omitempty"` // 机器人/批量等无显式项目时绑定的默认项目
|
||||
FactIndexMaxRunes int `yaml:"fact_index_max_runes,omitempty" json:"fact_index_max_runes,omitempty"`
|
||||
FactSummaryMaxRunes int `yaml:"fact_summary_max_runes,omitempty" json:"fact_summary_max_runes,omitempty"`
|
||||
DefaultInjectDeprecated bool `yaml:"default_inject_deprecated,omitempty" json:"default_inject_deprecated,omitempty"`
|
||||
}
|
||||
|
||||
// FactIndexMaxRunesEffective 自动注入黑板索引的最大 rune 数。
|
||||
func (c ProjectConfig) FactIndexMaxRunesEffective() int {
|
||||
if c.FactIndexMaxRunes <= 0 {
|
||||
return 3500
|
||||
}
|
||||
return c.FactIndexMaxRunes
|
||||
}
|
||||
|
||||
// FactSummaryMaxRunesEffective upsert 时 summary 最大 rune 数(索引一行,宜含验证要点)。
|
||||
func (c ProjectConfig) FactSummaryMaxRunesEffective() int {
|
||||
if c.FactSummaryMaxRunes <= 0 {
|
||||
return 200
|
||||
}
|
||||
return c.FactSummaryMaxRunes
|
||||
}
|
||||
|
||||
// MultiAgentConfig 基于 CloudWeGo Eino adk/prebuilt 的多代理编排(deep | plan_execute | supervisor,与单 Agent /agent-loop 并存)。
|
||||
type MultiAgentConfig struct {
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
RobotUseMultiAgent bool `yaml:"robot_use_multi_agent" json:"robot_use_multi_agent"` // 为 true 时钉钉/飞书/企微机器人走 Eino 多代理
|
||||
BatchUseMultiAgent bool `yaml:"batch_use_multi_agent" json:"batch_use_multi_agent"` // 为 true 时批量任务队列中每子任务走 Eino 多代理
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
RobotDefaultAgentMode string `yaml:"robot_default_agent_mode,omitempty" json:"robot_default_agent_mode,omitempty"` // react | eino_single | deep | plan_execute | supervisor
|
||||
BatchUseMultiAgent bool `yaml:"batch_use_multi_agent" json:"batch_use_multi_agent"` // 为 true 时批量任务队列中每子任务走 Eino 多代理
|
||||
// Orchestration 已弃用:保留仅兼容旧版 config.yaml;编排由聊天/WebShell 请求体 orchestration 决定,未传时按 deep。
|
||||
Orchestration string `yaml:"orchestration,omitempty" json:"orchestration,omitempty"`
|
||||
MaxIteration int `yaml:"max_iteration" json:"max_iteration"` // 主代理 / 执行器最大推理轮次(Deep、Supervisor、plan_execute 的 Executor)
|
||||
@@ -227,6 +254,10 @@ type MultiAgentEinoMiddlewareConfig struct {
|
||||
DeepOutputKey string `yaml:"deep_output_key,omitempty" json:"deep_output_key,omitempty"`
|
||||
// DeepModelRetryMaxRetries > 0 enables deep.Config ModelRetryConfig (framework-level chat model retries).
|
||||
DeepModelRetryMaxRetries int `yaml:"deep_model_retry_max_retries,omitempty" json:"deep_model_retry_max_retries,omitempty"`
|
||||
// RunRetryMaxAttempts > 0:429/5xx/网络抖动时 handler 分段续跑次数;0=默认 10。
|
||||
RunRetryMaxAttempts int `yaml:"run_retry_max_attempts,omitempty" json:"run_retry_max_attempts,omitempty"`
|
||||
// RunRetryMaxBackoffSec 单次退避上限秒数;0=默认 30。
|
||||
RunRetryMaxBackoffSec int `yaml:"run_retry_max_backoff_sec,omitempty" json:"run_retry_max_backoff_sec,omitempty"`
|
||||
// TaskToolDescriptionPrefix when non-empty sets deep.Config TaskToolDescriptionGenerator (sub-agent names appended).
|
||||
TaskToolDescriptionPrefix string `yaml:"task_tool_description_prefix,omitempty" json:"task_tool_description_prefix,omitempty"`
|
||||
}
|
||||
@@ -362,9 +393,9 @@ type MultiAgentSubConfig struct {
|
||||
|
||||
// MultiAgentPublic 返回给前端的精简信息(不含子代理指令全文)。
|
||||
type MultiAgentPublic struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
RobotUseMultiAgent bool `json:"robot_use_multi_agent"`
|
||||
BatchUseMultiAgent bool `json:"batch_use_multi_agent"`
|
||||
Enabled bool `json:"enabled"`
|
||||
RobotDefaultAgentMode string `json:"robot_default_agent_mode,omitempty"`
|
||||
BatchUseMultiAgent bool `json:"batch_use_multi_agent"`
|
||||
SubAgentCount int `json:"sub_agent_count"`
|
||||
Orchestration string `json:"orchestration,omitempty"`
|
||||
PlanExecuteLoopMaxIterations int `json:"plan_execute_loop_max_iterations"`
|
||||
@@ -372,6 +403,18 @@ type MultiAgentPublic struct {
|
||||
ToolSearchAlwaysVisibleEffectiveTools []string `json:"tool_search_always_visible_effective_tools,omitempty"`
|
||||
}
|
||||
|
||||
// NormalizeRobotAgentMode 解析机器人默认对话模式(react | eino_single | deep | plan_execute | supervisor);空值视为 react。
|
||||
func NormalizeRobotAgentMode(ma MultiAgentConfig) string {
|
||||
s := strings.TrimSpace(strings.ToLower(ma.RobotDefaultAgentMode))
|
||||
if s == "" || s == "single" || s == "react" {
|
||||
return "react"
|
||||
}
|
||||
if s == "eino_single" {
|
||||
return "eino_single"
|
||||
}
|
||||
return NormalizeMultiAgentOrchestration(s)
|
||||
}
|
||||
|
||||
// NormalizeMultiAgentOrchestration 返回 deep、plan_execute 或 supervisor。
|
||||
func NormalizeMultiAgentOrchestration(s string) string {
|
||||
v := strings.TrimSpace(strings.ToLower(s))
|
||||
@@ -387,22 +430,35 @@ func NormalizeMultiAgentOrchestration(s string) string {
|
||||
|
||||
// MultiAgentAPIUpdate 设置页/API 仅更新多代理标量字段;写入 YAML 时不覆盖 sub_agents 等块。
|
||||
type MultiAgentAPIUpdate struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
RobotUseMultiAgent bool `json:"robot_use_multi_agent"`
|
||||
BatchUseMultiAgent bool `json:"batch_use_multi_agent"`
|
||||
Enabled bool `json:"enabled"`
|
||||
RobotDefaultAgentMode string `json:"robot_default_agent_mode,omitempty"`
|
||||
BatchUseMultiAgent bool `json:"batch_use_multi_agent"`
|
||||
PlanExecuteLoopMaxIterations *int `json:"plan_execute_loop_max_iterations,omitempty"`
|
||||
// 指针区分「JSON 未传该字段」与「传空数组要清空」;省略时不应覆盖 YAML 中的常驻工具白名单。
|
||||
ToolSearchAlwaysVisibleTools *[]string `json:"tool_search_always_visible_tools,omitempty"`
|
||||
}
|
||||
|
||||
// RobotsConfig 机器人配置(企业微信、钉钉、飞书等)
|
||||
// RobotsConfig 机器人配置(企业微信、钉钉、飞书、微信 iLink 等)
|
||||
type RobotsConfig struct {
|
||||
Session RobotSessionConfig `yaml:"session,omitempty" json:"session,omitempty"` // 机器人会话隔离策略
|
||||
Wechat RobotWechatConfig `yaml:"wechat,omitempty" json:"wechat,omitempty"` // 微信(iLink 扫码绑定)
|
||||
Wecom RobotWecomConfig `yaml:"wecom,omitempty" json:"wecom,omitempty"` // 企业微信
|
||||
Dingtalk RobotDingtalkConfig `yaml:"dingtalk,omitempty" json:"dingtalk,omitempty"` // 钉钉
|
||||
Lark RobotLarkConfig `yaml:"lark,omitempty" json:"lark,omitempty"` // 飞书
|
||||
}
|
||||
|
||||
// RobotWechatConfig 微信 iLink 机器人配置(个人微信 ClawBot / iLink 协议)
|
||||
type RobotWechatConfig struct {
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
BotToken string `yaml:"bot_token,omitempty" json:"bot_token,omitempty"`
|
||||
ILinkBotID string `yaml:"ilink_bot_id,omitempty" json:"ilink_bot_id,omitempty"`
|
||||
ILinkUserID string `yaml:"ilink_user_id,omitempty" json:"ilink_user_id,omitempty"`
|
||||
BaseURL string `yaml:"base_url,omitempty" json:"base_url,omitempty"` // 默认 https://ilinkai.weixin.qq.com
|
||||
BotType string `yaml:"bot_type,omitempty" json:"bot_type,omitempty"` // get_bot_qrcode 参数,默认 3
|
||||
BotAgent string `yaml:"bot_agent,omitempty" json:"bot_agent,omitempty"` // base_info.bot_agent
|
||||
GetUpdatesBuf string `yaml:"get_updates_buf,omitempty" json:"get_updates_buf,omitempty"` // 长轮询游标(运行时)
|
||||
}
|
||||
|
||||
// RobotSessionConfig 机器人会话隔离策略
|
||||
type RobotSessionConfig struct {
|
||||
StrictUserIdentity *bool `yaml:"strict_user_identity,omitempty" json:"strict_user_identity,omitempty"` // true 时只允许真实用户标识,不允许会话/群 ID 兜底
|
||||
@@ -484,7 +540,7 @@ type OpenAIConfig struct {
|
||||
type OpenAIReasoningConfig struct {
|
||||
// Mode: auto(默认)| on | off | default(与 auto 相同)。off 时不向模型附加推理扩展字段。
|
||||
Mode string `yaml:"mode,omitempty" json:"mode,omitempty"`
|
||||
// Effort: low | medium | high | max;空表示不单独指定强度(各 profile 行为见 internal/reasoning)。
|
||||
// Effort: low | medium | high | max | xhigh;max/xhigh 为不同网关最高档命名,原样下发、不互转。空表示不单独指定强度。
|
||||
Effort string `yaml:"effort,omitempty" json:"effort,omitempty"`
|
||||
// AllowClientReasoning 为 false 时忽略请求体 reasoning;nil 或未设置等同于 true。
|
||||
AllowClientReasoning *bool `yaml:"allow_client_reasoning,omitempty" json:"allow_client_reasoning,omitempty"`
|
||||
@@ -562,6 +618,51 @@ type AuthConfig struct {
|
||||
GeneratedPasswordPersistErr string `yaml:"-" json:"-"`
|
||||
}
|
||||
|
||||
// AuditConfig platform operation audit log settings (not chat/tool execution bodies).
|
||||
type AuditConfig struct {
|
||||
// Enabled nil or true enables persistence; explicit false disables.
|
||||
Enabled *bool `yaml:"enabled,omitempty" json:"enabled,omitempty"`
|
||||
RetentionDays int `yaml:"retention_days,omitempty" json:"retention_days,omitempty"`
|
||||
MaxDetailBytes int `yaml:"max_detail_bytes,omitempty" json:"max_detail_bytes,omitempty"`
|
||||
// AuthFailureCooldownSeconds: per-IP cooldown for auth login/change_password failure audit rows; -1 disables; 0 uses default 60.
|
||||
AuthFailureCooldownSeconds int `yaml:"auth_failure_cooldown_seconds,omitempty" json:"auth_failure_cooldown_seconds,omitempty"`
|
||||
}
|
||||
|
||||
// EnabledEffective returns true unless audit.enabled is explicitly false.
|
||||
func (a AuditConfig) EnabledEffective() bool {
|
||||
if a.Enabled == nil {
|
||||
return true
|
||||
}
|
||||
return *a.Enabled
|
||||
}
|
||||
|
||||
// RetentionDaysEffective returns retention; 0 means keep forever.
|
||||
func (a AuditConfig) RetentionDaysEffective() int {
|
||||
if a.RetentionDays < 0 {
|
||||
return 0
|
||||
}
|
||||
return a.RetentionDays
|
||||
}
|
||||
|
||||
// MaxDetailBytesEffective caps serialized detail JSON size.
|
||||
func (a AuditConfig) MaxDetailBytesEffective() int {
|
||||
if a.MaxDetailBytes <= 0 {
|
||||
return 8192
|
||||
}
|
||||
return a.MaxDetailBytes
|
||||
}
|
||||
|
||||
// AuthFailureCooldownEffective returns seconds between duplicate auth-failure audit rows per IP (default 60; -1 disables).
|
||||
func (a AuditConfig) AuthFailureCooldownEffective() int {
|
||||
if a.AuthFailureCooldownSeconds < 0 {
|
||||
return 0
|
||||
}
|
||||
if a.AuthFailureCooldownSeconds == 0 {
|
||||
return 60
|
||||
}
|
||||
return a.AuthFailureCooldownSeconds
|
||||
}
|
||||
|
||||
// ExternalMCPConfig 外部MCP配置
|
||||
type ExternalMCPConfig struct {
|
||||
Servers map[string]ExternalMCPServerConfig `yaml:"servers,omitempty" json:"servers,omitempty"`
|
||||
@@ -654,6 +755,9 @@ func Load(path string) (*Config, error) {
|
||||
if cfg.Auth.SessionDurationHours <= 0 {
|
||||
cfg.Auth.SessionDurationHours = 12
|
||||
}
|
||||
if cfg.Audit.MaxDetailBytes <= 0 {
|
||||
cfg.Audit.MaxDetailBytes = 8192
|
||||
}
|
||||
if strings.TrimSpace(cfg.Auth.Password) == "" {
|
||||
password, err := generateStrongPassword(24)
|
||||
if err != nil {
|
||||
@@ -1157,6 +1261,14 @@ func Default() *Config {
|
||||
Auth: AuthConfig{
|
||||
SessionDurationHours: 12,
|
||||
},
|
||||
Audit: func() AuditConfig {
|
||||
on := true
|
||||
return AuditConfig{
|
||||
RetentionDays: 90,
|
||||
MaxDetailBytes: 8192,
|
||||
Enabled: &on,
|
||||
}
|
||||
}(),
|
||||
Robots: RobotsConfig{
|
||||
Session: RobotSessionConfig{
|
||||
StrictUserIdentity: &strictRobotIdentity,
|
||||
|
||||
@@ -0,0 +1,210 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// AuditLog platform operation audit record.
|
||||
type AuditLog struct {
|
||||
ID string `json:"id"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
Level string `json:"level"`
|
||||
Category string `json:"category"`
|
||||
Action string `json:"action"`
|
||||
Result string `json:"result"`
|
||||
Actor string `json:"actor"`
|
||||
SessionHint string `json:"sessionHint,omitempty"`
|
||||
ClientIP string `json:"clientIp,omitempty"`
|
||||
UserAgent string `json:"userAgent,omitempty"`
|
||||
ResourceType string `json:"resourceType,omitempty"`
|
||||
ResourceID string `json:"resourceId,omitempty"`
|
||||
ResourceAvailable *bool `json:"resourceAvailable,omitempty"` // API-only: whether linked resource still exists
|
||||
Message string `json:"message"`
|
||||
Detail map[string]interface{} `json:"detail,omitempty"`
|
||||
}
|
||||
|
||||
// ListAuditLogsFilter query parameters.
|
||||
type ListAuditLogsFilter struct {
|
||||
Level string
|
||||
Category string
|
||||
Action string
|
||||
Result string
|
||||
Query string
|
||||
ResourceType string
|
||||
ResourceID string
|
||||
Since *time.Time
|
||||
Until *time.Time
|
||||
Limit int
|
||||
Offset int
|
||||
}
|
||||
|
||||
func buildAuditLogsWhere(filter ListAuditLogsFilter) (string, []interface{}) {
|
||||
conditions := []string{"1=1"}
|
||||
args := []interface{}{}
|
||||
if filter.Level != "" {
|
||||
conditions = append(conditions, "level = ?")
|
||||
args = append(args, filter.Level)
|
||||
}
|
||||
if filter.Category != "" {
|
||||
conditions = append(conditions, "category = ?")
|
||||
args = append(args, filter.Category)
|
||||
}
|
||||
if filter.Action != "" {
|
||||
conditions = append(conditions, "action = ?")
|
||||
args = append(args, filter.Action)
|
||||
}
|
||||
if filter.Result != "" {
|
||||
conditions = append(conditions, "result = ?")
|
||||
args = append(args, filter.Result)
|
||||
}
|
||||
if filter.ResourceType != "" {
|
||||
conditions = append(conditions, "resource_type = ?")
|
||||
args = append(args, filter.ResourceType)
|
||||
}
|
||||
if filter.ResourceID != "" {
|
||||
conditions = append(conditions, "resource_id = ?")
|
||||
args = append(args, filter.ResourceID)
|
||||
}
|
||||
if filter.Since != nil {
|
||||
conditions = append(conditions, "created_at >= ?")
|
||||
args = append(args, *filter.Since)
|
||||
}
|
||||
if filter.Until != nil {
|
||||
conditions = append(conditions, "created_at <= ?")
|
||||
args = append(args, *filter.Until)
|
||||
}
|
||||
if q := strings.TrimSpace(filter.Query); q != "" {
|
||||
like := "%" + q + "%"
|
||||
conditions = append(conditions, "(message LIKE ? OR resource_id LIKE ? OR action LIKE ? OR category LIKE ?)")
|
||||
args = append(args, like, like, like, like)
|
||||
}
|
||||
return strings.Join(conditions, " AND "), args
|
||||
}
|
||||
|
||||
// AppendAuditLog inserts one audit row.
|
||||
func (db *DB) AppendAuditLog(row *AuditLog) error {
|
||||
if row == nil {
|
||||
return errors.New("audit log is nil")
|
||||
}
|
||||
if strings.TrimSpace(row.ID) == "" {
|
||||
return errors.New("audit id is required")
|
||||
}
|
||||
if row.CreatedAt.IsZero() {
|
||||
row.CreatedAt = time.Now()
|
||||
}
|
||||
if strings.TrimSpace(row.Level) == "" {
|
||||
row.Level = "info"
|
||||
}
|
||||
detailJSON := ""
|
||||
if len(row.Detail) > 0 {
|
||||
if b, err := json.Marshal(row.Detail); err == nil {
|
||||
detailJSON = string(b)
|
||||
}
|
||||
}
|
||||
query := `
|
||||
INSERT INTO audit_logs (
|
||||
id, created_at, level, category, action, result, actor, session_hint,
|
||||
client_ip, user_agent, resource_type, resource_id, message, detail_json
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
`
|
||||
_, err := db.Exec(query,
|
||||
row.ID, row.CreatedAt, row.Level, row.Category, row.Action, row.Result,
|
||||
row.Actor, row.SessionHint, row.ClientIP, row.UserAgent,
|
||||
row.ResourceType, row.ResourceID, row.Message, detailJSON,
|
||||
)
|
||||
return err
|
||||
}
|
||||
|
||||
// GetAuditLogByID returns one row.
|
||||
func (db *DB) GetAuditLogByID(id string) (*AuditLog, error) {
|
||||
id = strings.TrimSpace(id)
|
||||
if id == "" {
|
||||
return nil, errors.New("id is required")
|
||||
}
|
||||
query := `
|
||||
SELECT id, created_at, level, category, action, result, actor,
|
||||
COALESCE(session_hint, ''), COALESCE(client_ip, ''), COALESCE(user_agent, ''),
|
||||
COALESCE(resource_type, ''), COALESCE(resource_id, ''), message, COALESCE(detail_json, '')
|
||||
FROM audit_logs WHERE id = ?
|
||||
`
|
||||
var row AuditLog
|
||||
var detailJSON string
|
||||
err := db.QueryRow(query, id).Scan(
|
||||
&row.ID, &row.CreatedAt, &row.Level, &row.Category, &row.Action, &row.Result, &row.Actor,
|
||||
&row.SessionHint, &row.ClientIP, &row.UserAgent,
|
||||
&row.ResourceType, &row.ResourceID, &row.Message, &detailJSON,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if detailJSON != "" {
|
||||
_ = json.Unmarshal([]byte(detailJSON), &row.Detail)
|
||||
}
|
||||
return &row, nil
|
||||
}
|
||||
|
||||
// CountAuditLogs counts rows matching filter.
|
||||
func (db *DB) CountAuditLogs(filter ListAuditLogsFilter) (int64, error) {
|
||||
where, args := buildAuditLogsWhere(filter)
|
||||
query := `SELECT COUNT(*) FROM audit_logs WHERE ` + where
|
||||
var n int64
|
||||
err := db.QueryRow(query, args...).Scan(&n)
|
||||
return n, err
|
||||
}
|
||||
|
||||
// ListAuditLogs lists audit rows newest first.
|
||||
func (db *DB) ListAuditLogs(filter ListAuditLogsFilter) ([]*AuditLog, error) {
|
||||
where, args := buildAuditLogsWhere(filter)
|
||||
limit := filter.Limit
|
||||
if limit <= 0 || limit > 500 {
|
||||
limit = 50
|
||||
}
|
||||
offset := filter.Offset
|
||||
if offset < 0 {
|
||||
offset = 0
|
||||
}
|
||||
query := `
|
||||
SELECT id, created_at, level, category, action, result, actor,
|
||||
COALESCE(session_hint, ''), COALESCE(client_ip, ''), COALESCE(user_agent, ''),
|
||||
COALESCE(resource_type, ''), COALESCE(resource_id, ''), message, COALESCE(detail_json, '')
|
||||
FROM audit_logs
|
||||
WHERE ` + where + `
|
||||
ORDER BY created_at DESC
|
||||
LIMIT ? OFFSET ?
|
||||
`
|
||||
args = append(args, limit, offset)
|
||||
rows, err := db.Query(query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var list []*AuditLog
|
||||
for rows.Next() {
|
||||
var row AuditLog
|
||||
var detailJSON string
|
||||
if err := rows.Scan(
|
||||
&row.ID, &row.CreatedAt, &row.Level, &row.Category, &row.Action, &row.Result, &row.Actor,
|
||||
&row.SessionHint, &row.ClientIP, &row.UserAgent,
|
||||
&row.ResourceType, &row.ResourceID, &row.Message, &detailJSON,
|
||||
); err != nil {
|
||||
continue
|
||||
}
|
||||
if detailJSON != "" {
|
||||
_ = json.Unmarshal([]byte(detailJSON), &row.Detail)
|
||||
}
|
||||
list = append(list, &row)
|
||||
}
|
||||
return list, rows.Err()
|
||||
}
|
||||
|
||||
// DeleteAuditLogsBefore removes rows older than cutoff.
|
||||
func (db *DB) DeleteAuditLogsBefore(cutoff time.Time) (int64, error) {
|
||||
res, err := db.Exec(`DELETE FROM audit_logs WHERE created_at < ?`, cutoff)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return res.RowsAffected()
|
||||
}
|
||||
@@ -22,6 +22,7 @@ type BatchTaskQueueRow struct {
|
||||
LastScheduleTriggerAt sql.NullTime
|
||||
LastScheduleError sql.NullString
|
||||
LastRunError sql.NullString
|
||||
ProjectID sql.NullString
|
||||
Status string
|
||||
CreatedAt time.Time
|
||||
StartedAt sql.NullTime
|
||||
@@ -51,6 +52,7 @@ func (db *DB) CreateBatchQueue(
|
||||
scheduleMode string,
|
||||
cronExpr string,
|
||||
nextRunAt *time.Time,
|
||||
projectID string,
|
||||
tasks []map[string]interface{},
|
||||
) error {
|
||||
tx, err := db.Begin()
|
||||
@@ -65,9 +67,13 @@ func (db *DB) CreateBatchQueue(
|
||||
nextRunAtValue = *nextRunAt
|
||||
}
|
||||
|
||||
var projectIDVal interface{}
|
||||
if strings.TrimSpace(projectID) != "" {
|
||||
projectIDVal = strings.TrimSpace(projectID)
|
||||
}
|
||||
_, err = tx.Exec(
|
||||
"INSERT INTO batch_task_queues (id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, status, created_at, current_index) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
|
||||
queueID, title, role, agentMode, scheduleMode, cronExpr, nextRunAtValue, 1, "pending", now, 0,
|
||||
"INSERT INTO batch_task_queues (id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, project_id, status, created_at, current_index) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
|
||||
queueID, title, role, agentMode, scheduleMode, cronExpr, nextRunAtValue, 1, projectIDVal, "pending", now, 0,
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("创建批量任务队列失败: %w", err)
|
||||
@@ -101,9 +107,9 @@ func (db *DB) GetBatchQueue(queueID string) (*BatchTaskQueueRow, error) {
|
||||
var row BatchTaskQueueRow
|
||||
var createdAt string
|
||||
err := db.QueryRow(
|
||||
"SELECT id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, last_schedule_trigger_at, last_schedule_error, last_run_error, status, created_at, started_at, completed_at, current_index FROM batch_task_queues WHERE id = ?",
|
||||
"SELECT id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, last_schedule_trigger_at, last_schedule_error, last_run_error, project_id, status, created_at, started_at, completed_at, current_index FROM batch_task_queues WHERE id = ?",
|
||||
queueID,
|
||||
).Scan(&row.ID, &row.Title, &row.Role, &row.AgentMode, &row.ScheduleMode, &row.CronExpr, &row.NextRunAt, &row.ScheduleEnabled, &row.LastScheduleTriggerAt, &row.LastScheduleError, &row.LastRunError, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex)
|
||||
).Scan(&row.ID, &row.Title, &row.Role, &row.AgentMode, &row.ScheduleMode, &row.CronExpr, &row.NextRunAt, &row.ScheduleEnabled, &row.LastScheduleTriggerAt, &row.LastScheduleError, &row.LastRunError, &row.ProjectID, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex)
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, nil
|
||||
}
|
||||
@@ -127,7 +133,7 @@ func (db *DB) GetBatchQueue(queueID string) (*BatchTaskQueueRow, error) {
|
||||
// GetAllBatchQueues 获取所有批量任务队列
|
||||
func (db *DB) GetAllBatchQueues() ([]*BatchTaskQueueRow, error) {
|
||||
rows, err := db.Query(
|
||||
"SELECT id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, last_schedule_trigger_at, last_schedule_error, last_run_error, status, created_at, started_at, completed_at, current_index FROM batch_task_queues ORDER BY created_at DESC",
|
||||
"SELECT id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, last_schedule_trigger_at, last_schedule_error, last_run_error, project_id, status, created_at, started_at, completed_at, current_index FROM batch_task_queues ORDER BY created_at DESC",
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询批量任务队列列表失败: %w", err)
|
||||
@@ -138,7 +144,7 @@ func (db *DB) GetAllBatchQueues() ([]*BatchTaskQueueRow, error) {
|
||||
for rows.Next() {
|
||||
var row BatchTaskQueueRow
|
||||
var createdAt string
|
||||
if err := rows.Scan(&row.ID, &row.Title, &row.Role, &row.AgentMode, &row.ScheduleMode, &row.CronExpr, &row.NextRunAt, &row.ScheduleEnabled, &row.LastScheduleTriggerAt, &row.LastScheduleError, &row.LastRunError, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex); err != nil {
|
||||
if err := rows.Scan(&row.ID, &row.Title, &row.Role, &row.AgentMode, &row.ScheduleMode, &row.CronExpr, &row.NextRunAt, &row.ScheduleEnabled, &row.LastScheduleTriggerAt, &row.LastScheduleError, &row.LastRunError, &row.ProjectID, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex); err != nil {
|
||||
return nil, fmt.Errorf("扫描批量任务队列失败: %w", err)
|
||||
}
|
||||
parsedTime, parseErr := time.Parse("2006-01-02 15:04:05", createdAt)
|
||||
@@ -158,7 +164,7 @@ func (db *DB) GetAllBatchQueues() ([]*BatchTaskQueueRow, error) {
|
||||
|
||||
// ListBatchQueues 列出批量任务队列(支持筛选和分页)
|
||||
func (db *DB) ListBatchQueues(limit, offset int, status, keyword string) ([]*BatchTaskQueueRow, error) {
|
||||
query := "SELECT id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, last_schedule_trigger_at, last_schedule_error, last_run_error, status, created_at, started_at, completed_at, current_index FROM batch_task_queues WHERE 1=1"
|
||||
query := "SELECT id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, last_schedule_trigger_at, last_schedule_error, last_run_error, project_id, status, created_at, started_at, completed_at, current_index FROM batch_task_queues WHERE 1=1"
|
||||
args := []interface{}{}
|
||||
|
||||
// 状态筛选
|
||||
@@ -186,7 +192,7 @@ func (db *DB) ListBatchQueues(limit, offset int, status, keyword string) ([]*Bat
|
||||
for rows.Next() {
|
||||
var row BatchTaskQueueRow
|
||||
var createdAt string
|
||||
if err := rows.Scan(&row.ID, &row.Title, &row.Role, &row.AgentMode, &row.ScheduleMode, &row.CronExpr, &row.NextRunAt, &row.ScheduleEnabled, &row.LastScheduleTriggerAt, &row.LastScheduleError, &row.LastRunError, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex); err != nil {
|
||||
if err := rows.Scan(&row.ID, &row.Title, &row.Role, &row.AgentMode, &row.ScheduleMode, &row.CronExpr, &row.NextRunAt, &row.ScheduleEnabled, &row.LastScheduleTriggerAt, &row.LastScheduleError, &row.LastRunError, &row.ProjectID, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex); err != nil {
|
||||
return nil, fmt.Errorf("扫描批量任务队列失败: %w", err)
|
||||
}
|
||||
parsedTime, parseErr := time.Parse("2006-01-02 15:04:05", createdAt)
|
||||
|
||||
@@ -17,6 +17,7 @@ import (
|
||||
type Conversation struct {
|
||||
ID string `json:"id"`
|
||||
Title string `json:"title"`
|
||||
ProjectID string `json:"projectId,omitempty"`
|
||||
Pinned bool `json:"pinned"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
@@ -37,22 +38,41 @@ type Message struct {
|
||||
}
|
||||
|
||||
// CreateConversation 创建新对话
|
||||
func (db *DB) CreateConversation(title string) (*Conversation, error) {
|
||||
return db.CreateConversationWithWebshell("", title)
|
||||
func (db *DB) CreateConversation(title string, meta ConversationCreateMeta) (*Conversation, error) {
|
||||
return db.CreateConversationWithWebshell("", title, meta)
|
||||
}
|
||||
|
||||
// CreateConversationWithWebshell 创建新对话,可选绑定 WebShell 连接 ID(为空则普通对话)
|
||||
func (db *DB) CreateConversationWithWebshell(webshellConnectionID, title string) (*Conversation, error) {
|
||||
func (db *DB) CreateConversationWithWebshell(webshellConnectionID, title string, meta ConversationCreateMeta) (*Conversation, error) {
|
||||
id := uuid.New().String()
|
||||
now := time.Now()
|
||||
|
||||
projectID := strings.TrimSpace(meta.ProjectID)
|
||||
if projectID != "" {
|
||||
if _, err := db.GetProject(projectID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
var err error
|
||||
if webshellConnectionID != "" {
|
||||
wsID := strings.TrimSpace(webshellConnectionID)
|
||||
switch {
|
||||
case wsID != "" && projectID != "":
|
||||
_, err = db.Exec(
|
||||
"INSERT INTO conversations (id, title, created_at, updated_at, webshell_connection_id, project_id) VALUES (?, ?, ?, ?, ?, ?)",
|
||||
id, title, now, now, wsID, projectID,
|
||||
)
|
||||
case wsID != "":
|
||||
_, err = db.Exec(
|
||||
"INSERT INTO conversations (id, title, created_at, updated_at, webshell_connection_id) VALUES (?, ?, ?, ?, ?)",
|
||||
id, title, now, now, webshellConnectionID,
|
||||
id, title, now, now, wsID,
|
||||
)
|
||||
} else {
|
||||
case projectID != "":
|
||||
_, err = db.Exec(
|
||||
"INSERT INTO conversations (id, title, created_at, updated_at, project_id) VALUES (?, ?, ?, ?, ?)",
|
||||
id, title, now, now, projectID,
|
||||
)
|
||||
default:
|
||||
_, err = db.Exec(
|
||||
"INSERT INTO conversations (id, title, created_at, updated_at) VALUES (?, ?, ?, ?)",
|
||||
id, title, now, now,
|
||||
@@ -62,12 +82,18 @@ func (db *DB) CreateConversationWithWebshell(webshellConnectionID, title string)
|
||||
return nil, fmt.Errorf("创建对话失败: %w", err)
|
||||
}
|
||||
|
||||
return &Conversation{
|
||||
conv := &Conversation{
|
||||
ID: id,
|
||||
Title: title,
|
||||
ProjectID: projectID,
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}, nil
|
||||
}
|
||||
if wsID != "" {
|
||||
meta.WebShellConnectionID = wsID
|
||||
}
|
||||
notifyConversationCreated(conv, meta)
|
||||
return conv, nil
|
||||
}
|
||||
|
||||
// GetConversationByWebshellConnectionID 根据 WebShell 连接 ID 获取该连接下最近一条对话(用于 AI 助手持久化)
|
||||
@@ -182,22 +208,43 @@ func (db *DB) ListConversationsByWebshellConnectionID(connectionID string) ([]We
|
||||
return list, rows.Err()
|
||||
}
|
||||
|
||||
// ConversationExists reports whether a conversation row exists (lightweight check for audit links).
|
||||
func (db *DB) ConversationExists(id string) (bool, error) {
|
||||
id = strings.TrimSpace(id)
|
||||
if id == "" {
|
||||
return false, nil
|
||||
}
|
||||
var one int
|
||||
err := db.QueryRow("SELECT 1 FROM conversations WHERE id = ? LIMIT 1", id).Scan(&one)
|
||||
if err == sql.ErrNoRows {
|
||||
return false, nil
|
||||
}
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
// GetConversation 获取对话
|
||||
func (db *DB) GetConversation(id string) (*Conversation, error) {
|
||||
var conv Conversation
|
||||
var createdAt, updatedAt string
|
||||
var pinned int
|
||||
|
||||
var projectID sql.NullString
|
||||
err := db.QueryRow(
|
||||
"SELECT id, title, pinned, created_at, updated_at FROM conversations WHERE id = ?",
|
||||
"SELECT id, title, pinned, created_at, updated_at, project_id FROM conversations WHERE id = ?",
|
||||
id,
|
||||
).Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt)
|
||||
).Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID)
|
||||
if err != nil {
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, fmt.Errorf("对话不存在")
|
||||
}
|
||||
return nil, fmt.Errorf("查询对话失败: %w", err)
|
||||
}
|
||||
if projectID.Valid {
|
||||
conv.ProjectID = strings.TrimSpace(projectID.String)
|
||||
}
|
||||
|
||||
// 尝试多种时间格式解析
|
||||
var err1, err2 error
|
||||
@@ -270,16 +317,20 @@ func (db *DB) GetConversationLite(id string) (*Conversation, error) {
|
||||
var createdAt, updatedAt string
|
||||
var pinned int
|
||||
|
||||
var projectID sql.NullString
|
||||
err := db.QueryRow(
|
||||
"SELECT id, title, pinned, created_at, updated_at FROM conversations WHERE id = ?",
|
||||
"SELECT id, title, pinned, created_at, updated_at, project_id FROM conversations WHERE id = ?",
|
||||
id,
|
||||
).Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt)
|
||||
).Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID)
|
||||
if err != nil {
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, fmt.Errorf("对话不存在")
|
||||
}
|
||||
return nil, fmt.Errorf("查询对话失败: %w", err)
|
||||
}
|
||||
if projectID.Valid {
|
||||
conv.ProjectID = strings.TrimSpace(projectID.String)
|
||||
}
|
||||
|
||||
// 尝试多种时间格式解析
|
||||
var err1, err2 error
|
||||
@@ -319,7 +370,7 @@ func (db *DB) ListConversations(limit, offset int, search string) ([]*Conversati
|
||||
// 使用 EXISTS 子查询代替 LEFT JOIN + DISTINCT,避免大表笛卡尔积
|
||||
searchPattern := "%" + search + "%"
|
||||
rows, err = db.Query(
|
||||
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at
|
||||
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id
|
||||
FROM conversations c
|
||||
WHERE c.title LIKE ?
|
||||
OR EXISTS (SELECT 1 FROM messages m WHERE m.conversation_id = c.id AND m.content LIKE ?)
|
||||
@@ -329,7 +380,7 @@ func (db *DB) ListConversations(limit, offset int, search string) ([]*Conversati
|
||||
)
|
||||
} else {
|
||||
rows, err = db.Query(
|
||||
"SELECT id, title, COALESCE(pinned, 0), created_at, updated_at FROM conversations ORDER BY updated_at DESC LIMIT ? OFFSET ?",
|
||||
"SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id FROM conversations ORDER BY updated_at DESC LIMIT ? OFFSET ?",
|
||||
limit, offset,
|
||||
)
|
||||
}
|
||||
@@ -344,10 +395,14 @@ func (db *DB) ListConversations(limit, offset int, search string) ([]*Conversati
|
||||
var conv Conversation
|
||||
var createdAt, updatedAt string
|
||||
var pinned int
|
||||
var projectID sql.NullString
|
||||
|
||||
if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt); err != nil {
|
||||
if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID); err != nil {
|
||||
return nil, fmt.Errorf("扫描对话失败: %w", err)
|
||||
}
|
||||
if projectID.Valid {
|
||||
conv.ProjectID = strings.TrimSpace(projectID.String)
|
||||
}
|
||||
|
||||
// 尝试多种时间格式解析
|
||||
var err1, err2 error
|
||||
|
||||
@@ -0,0 +1,30 @@
|
||||
package database
|
||||
|
||||
// ConversationCreateMeta describes how a conversation was created (for audit hooks).
|
||||
type ConversationCreateMeta struct {
|
||||
Source string
|
||||
WebShellConnectionID string
|
||||
ProjectID string
|
||||
ClientIP string
|
||||
SessionHint string
|
||||
}
|
||||
|
||||
// ConversationCreateHook is invoked after a conversation row is inserted.
|
||||
type ConversationCreateHook func(conv *Conversation, meta ConversationCreateMeta)
|
||||
|
||||
var conversationCreateHook ConversationCreateHook
|
||||
|
||||
// SetConversationCreateHook registers a global hook (e.g. platform audit).
|
||||
func SetConversationCreateHook(h ConversationCreateHook) {
|
||||
conversationCreateHook = h
|
||||
}
|
||||
|
||||
func notifyConversationCreated(conv *Conversation, meta ConversationCreateMeta) {
|
||||
if conversationCreateHook == nil || conv == nil {
|
||||
return
|
||||
}
|
||||
if meta.Source == "" {
|
||||
meta.Source = "unknown"
|
||||
}
|
||||
conversationCreateHook(conv, meta)
|
||||
}
|
||||
@@ -213,6 +213,40 @@ func (db *DB) initTables() error {
|
||||
FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE
|
||||
);`
|
||||
|
||||
// 创建项目表
|
||||
createProjectsTable := `
|
||||
CREATE TABLE IF NOT EXISTS projects (
|
||||
id TEXT PRIMARY KEY,
|
||||
name TEXT NOT NULL,
|
||||
description TEXT,
|
||||
scope_json TEXT,
|
||||
status TEXT NOT NULL DEFAULT 'active',
|
||||
pinned INTEGER NOT NULL DEFAULT 0,
|
||||
created_at DATETIME NOT NULL,
|
||||
updated_at DATETIME NOT NULL
|
||||
);`
|
||||
|
||||
// 创建项目事实表(黑板)
|
||||
createProjectFactsTable := `
|
||||
CREATE TABLE IF NOT EXISTS project_facts (
|
||||
id TEXT PRIMARY KEY,
|
||||
project_id TEXT NOT NULL,
|
||||
fact_key TEXT NOT NULL,
|
||||
category TEXT NOT NULL DEFAULT 'note',
|
||||
summary TEXT NOT NULL DEFAULT '',
|
||||
body TEXT,
|
||||
confidence TEXT NOT NULL DEFAULT 'tentative',
|
||||
source_conversation_id TEXT,
|
||||
source_message_id TEXT,
|
||||
pinned INTEGER NOT NULL DEFAULT 0,
|
||||
supersedes_fact_id TEXT,
|
||||
related_vulnerability_id TEXT,
|
||||
created_at DATETIME NOT NULL,
|
||||
updated_at DATETIME NOT NULL,
|
||||
FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE CASCADE,
|
||||
UNIQUE(project_id, fact_key)
|
||||
);`
|
||||
|
||||
// 创建漏洞表
|
||||
createVulnerabilitiesTable := `
|
||||
CREATE TABLE IF NOT EXISTS vulnerabilities (
|
||||
@@ -387,6 +421,24 @@ func (db *DB) initTables() error {
|
||||
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||||
);`
|
||||
|
||||
createAuditLogsTable := `
|
||||
CREATE TABLE IF NOT EXISTS audit_logs (
|
||||
id TEXT PRIMARY KEY,
|
||||
created_at DATETIME NOT NULL,
|
||||
level TEXT NOT NULL DEFAULT 'info',
|
||||
category TEXT NOT NULL,
|
||||
action TEXT NOT NULL,
|
||||
result TEXT NOT NULL,
|
||||
actor TEXT NOT NULL DEFAULT 'admin',
|
||||
session_hint TEXT,
|
||||
client_ip TEXT,
|
||||
user_agent TEXT,
|
||||
resource_type TEXT,
|
||||
resource_id TEXT,
|
||||
message TEXT NOT NULL,
|
||||
detail_json TEXT
|
||||
);`
|
||||
|
||||
createC2ProfilesTable := `
|
||||
CREATE TABLE IF NOT EXISTS c2_profiles (
|
||||
id TEXT PRIMARY KEY,
|
||||
@@ -427,6 +479,12 @@ func (db *DB) initTables() error {
|
||||
CREATE INDEX IF NOT EXISTS idx_vulnerabilities_severity ON vulnerabilities(severity);
|
||||
CREATE INDEX IF NOT EXISTS idx_vulnerabilities_status ON vulnerabilities(status);
|
||||
CREATE INDEX IF NOT EXISTS idx_vulnerabilities_created_at ON vulnerabilities(created_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_projects_status ON projects(status);
|
||||
CREATE INDEX IF NOT EXISTS idx_projects_updated_at ON projects(updated_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_project_facts_project_id ON project_facts(project_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_project_facts_confidence ON project_facts(confidence);
|
||||
CREATE INDEX IF NOT EXISTS idx_conversations_project_id ON conversations(project_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_vulnerabilities_project_id ON vulnerabilities(project_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_batch_tasks_queue_id ON batch_tasks(queue_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_batch_task_queues_created_at ON batch_task_queues(created_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_batch_task_queues_title ON batch_task_queues(title);
|
||||
@@ -445,6 +503,10 @@ func (db *DB) initTables() error {
|
||||
CREATE INDEX IF NOT EXISTS idx_c2_events_created_at ON c2_events(created_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_c2_events_category ON c2_events(category);
|
||||
CREATE INDEX IF NOT EXISTS idx_c2_events_session ON c2_events(session_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_audit_logs_created_at ON audit_logs(created_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_audit_logs_category ON audit_logs(category);
|
||||
CREATE INDEX IF NOT EXISTS idx_audit_logs_action ON audit_logs(action);
|
||||
CREATE INDEX IF NOT EXISTS idx_audit_logs_result ON audit_logs(result);
|
||||
`
|
||||
|
||||
if _, err := db.Exec(createConversationsTable); err != nil {
|
||||
@@ -494,6 +556,14 @@ func (db *DB) initTables() error {
|
||||
return fmt.Errorf("创建robot_user_sessions表失败: %w", err)
|
||||
}
|
||||
|
||||
if _, err := db.Exec(createProjectsTable); err != nil {
|
||||
return fmt.Errorf("创建projects表失败: %w", err)
|
||||
}
|
||||
|
||||
if _, err := db.Exec(createProjectFactsTable); err != nil {
|
||||
return fmt.Errorf("创建project_facts表失败: %w", err)
|
||||
}
|
||||
|
||||
if _, err := db.Exec(createVulnerabilitiesTable); err != nil {
|
||||
return fmt.Errorf("创建vulnerabilities表失败: %w", err)
|
||||
}
|
||||
@@ -514,6 +584,10 @@ func (db *DB) initTables() error {
|
||||
return fmt.Errorf("创建webshell_connection_states表失败: %w", err)
|
||||
}
|
||||
|
||||
if _, err := db.Exec(createAuditLogsTable); err != nil {
|
||||
return fmt.Errorf("创建audit_logs表失败: %w", err)
|
||||
}
|
||||
|
||||
for tableName, ddl := range map[string]string{
|
||||
"c2_listeners": createC2ListenersTable,
|
||||
"c2_sessions": createC2SessionsTable,
|
||||
@@ -557,6 +631,10 @@ func (db *DB) initTables() error {
|
||||
// 不返回错误,允许继续运行
|
||||
}
|
||||
|
||||
if err := db.migrateProjectsTable(); err != nil {
|
||||
db.logger.Warn("迁移projects相关表失败", zap.Error(err))
|
||||
}
|
||||
|
||||
if err := db.migrateWebshellConnectionsTable(); err != nil {
|
||||
db.logger.Warn("迁移webshell_connections表失败", zap.Error(err))
|
||||
// 不返回错误,允许继续运行
|
||||
@@ -904,6 +982,51 @@ func (db *DB) migrateBatchTaskQueuesTable() error {
|
||||
}
|
||||
}
|
||||
|
||||
var projectIDCount int
|
||||
err = db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('batch_task_queues') WHERE name='project_id'").Scan(&projectIDCount)
|
||||
if err != nil {
|
||||
if _, addErr := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN project_id TEXT"); addErr != nil {
|
||||
errMsg := strings.ToLower(addErr.Error())
|
||||
if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") {
|
||||
db.logger.Warn("添加batch_task_queues.project_id字段失败", zap.Error(addErr))
|
||||
}
|
||||
}
|
||||
} else if projectIDCount == 0 {
|
||||
if _, err := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN project_id TEXT"); err != nil {
|
||||
db.logger.Warn("添加batch_task_queues.project_id字段失败", zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// migrateProjectsTable 迁移 projects / conversations / vulnerabilities 的项目关联字段。
|
||||
func (db *DB) migrateProjectsTable() error {
|
||||
for _, col := range []struct {
|
||||
table string
|
||||
name string
|
||||
stmt string
|
||||
}{
|
||||
{"conversations", "project_id", "ALTER TABLE conversations ADD COLUMN project_id TEXT REFERENCES projects(id) ON DELETE SET NULL"},
|
||||
{"vulnerabilities", "project_id", "ALTER TABLE vulnerabilities ADD COLUMN project_id TEXT"},
|
||||
} {
|
||||
var count int
|
||||
err := db.QueryRow("SELECT COUNT(*) FROM pragma_table_info(?) WHERE name=?", col.table, col.name).Scan(&count)
|
||||
if err != nil {
|
||||
if _, addErr := db.Exec(col.stmt); addErr != nil {
|
||||
errMsg := strings.ToLower(addErr.Error())
|
||||
if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") {
|
||||
db.logger.Warn("添加字段失败", zap.String("table", col.table), zap.String("field", col.name), zap.Error(addErr))
|
||||
}
|
||||
}
|
||||
continue
|
||||
}
|
||||
if count == 0 {
|
||||
if _, addErr := db.Exec(col.stmt); addErr != nil {
|
||||
db.logger.Warn("添加字段失败", zap.String("table", col.table), zap.String("field", col.name), zap.Error(addErr))
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -915,6 +1038,7 @@ func (db *DB) migrateVulnerabilitiesTable() error {
|
||||
}{
|
||||
{name: "conversation_tag", stmt: "ALTER TABLE vulnerabilities ADD COLUMN conversation_tag TEXT"},
|
||||
{name: "task_tag", stmt: "ALTER TABLE vulnerabilities ADD COLUMN task_tag TEXT"},
|
||||
{name: "project_id", stmt: "ALTER TABLE vulnerabilities ADD COLUMN project_id TEXT"},
|
||||
}
|
||||
|
||||
for _, col := range columns {
|
||||
|
||||
@@ -0,0 +1,485 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
var factKeyPattern = regexp.MustCompile(`^[a-z0-9][a-z0-9._/-]*$`)
|
||||
|
||||
// ValidateFactKey 校验事实 key(项目内唯一标识)。
|
||||
func ValidateFactKey(key string) error {
|
||||
key = strings.TrimSpace(key)
|
||||
if key == "" {
|
||||
return fmt.Errorf("fact_key 不能为空")
|
||||
}
|
||||
if len(key) > 128 {
|
||||
return fmt.Errorf("fact_key 过长(最多 128 字符)")
|
||||
}
|
||||
if !factKeyPattern.MatchString(key) {
|
||||
return fmt.Errorf("fact_key 格式无效,仅允许小写字母、数字及 . _ / -,且须以小写字母或数字开头")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Project 渗透测试项目(跨对话共享黑板)。
|
||||
type Project struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description,omitempty"`
|
||||
ScopeJSON string `json:"scope_json,omitempty"`
|
||||
Status string `json:"status"` // active | archived
|
||||
Pinned bool `json:"pinned"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
// ProjectFact 项目事实(黑板条目)。
|
||||
type ProjectFact struct {
|
||||
ID string `json:"id"`
|
||||
ProjectID string `json:"project_id"`
|
||||
FactKey string `json:"fact_key"`
|
||||
Category string `json:"category"`
|
||||
Summary string `json:"summary"`
|
||||
Body string `json:"body"`
|
||||
Confidence string `json:"confidence"` // confirmed | tentative | deprecated
|
||||
SourceConversationID string `json:"source_conversation_id,omitempty"`
|
||||
SourceMessageID string `json:"source_message_id,omitempty"`
|
||||
Pinned bool `json:"pinned"`
|
||||
SupersedesFactID string `json:"supersedes_fact_id,omitempty"`
|
||||
RelatedVulnerabilityID string `json:"related_vulnerability_id,omitempty"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
// ProjectFactListFilter 事实列表筛选。
|
||||
type ProjectFactListFilter struct {
|
||||
Category string
|
||||
Confidence string
|
||||
Search string
|
||||
}
|
||||
|
||||
// CreateProject 创建项目。
|
||||
func (db *DB) CreateProject(p *Project) (*Project, error) {
|
||||
if p.ID == "" {
|
||||
p.ID = uuid.New().String()
|
||||
}
|
||||
if strings.TrimSpace(p.Status) == "" {
|
||||
p.Status = "active"
|
||||
}
|
||||
now := time.Now()
|
||||
if p.CreatedAt.IsZero() {
|
||||
p.CreatedAt = now
|
||||
}
|
||||
p.UpdatedAt = now
|
||||
|
||||
_, err := db.Exec(
|
||||
`INSERT INTO projects (id, name, description, scope_json, status, pinned, created_at, updated_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)`,
|
||||
p.ID, p.Name, p.Description, p.ScopeJSON, p.Status, boolToInt(p.Pinned), p.CreatedAt, p.UpdatedAt,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("创建项目失败: %w", err)
|
||||
}
|
||||
return p, nil
|
||||
}
|
||||
|
||||
// GetProject 获取项目。
|
||||
func (db *DB) GetProject(id string) (*Project, error) {
|
||||
var p Project
|
||||
var pinned int
|
||||
var createdAt, updatedAt string
|
||||
err := db.QueryRow(
|
||||
`SELECT id, name, COALESCE(description,''), COALESCE(scope_json,''), status, pinned, created_at, updated_at
|
||||
FROM projects WHERE id = ?`, id,
|
||||
).Scan(&p.ID, &p.Name, &p.Description, &p.ScopeJSON, &p.Status, &pinned, &createdAt, &updatedAt)
|
||||
if err != nil {
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, fmt.Errorf("项目不存在")
|
||||
}
|
||||
return nil, fmt.Errorf("获取项目失败: %w", err)
|
||||
}
|
||||
p.Pinned = pinned != 0
|
||||
p.CreatedAt = parseDBTime(createdAt)
|
||||
p.UpdatedAt = parseDBTime(updatedAt)
|
||||
return &p, nil
|
||||
}
|
||||
|
||||
// ListProjects 列出项目。
|
||||
func (db *DB) ListProjects(status string, limit, offset int) ([]*Project, error) {
|
||||
if limit <= 0 {
|
||||
limit = 200
|
||||
}
|
||||
query := `SELECT id, name, COALESCE(description,''), COALESCE(scope_json,''), status, pinned, created_at, updated_at
|
||||
FROM projects WHERE 1=1`
|
||||
args := []interface{}{}
|
||||
if s := strings.TrimSpace(status); s != "" {
|
||||
query += " AND status = ?"
|
||||
args = append(args, s)
|
||||
}
|
||||
query += " ORDER BY pinned DESC, updated_at DESC LIMIT ? OFFSET ?"
|
||||
args = append(args, limit, offset)
|
||||
|
||||
rows, err := db.Query(query, args...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("列出项目失败: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var out []*Project
|
||||
for rows.Next() {
|
||||
var p Project
|
||||
var pinned int
|
||||
var createdAt, updatedAt string
|
||||
if err := rows.Scan(&p.ID, &p.Name, &p.Description, &p.ScopeJSON, &p.Status, &pinned, &createdAt, &updatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
p.Pinned = pinned != 0
|
||||
p.CreatedAt = parseDBTime(createdAt)
|
||||
p.UpdatedAt = parseDBTime(updatedAt)
|
||||
out = append(out, &p)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// UpdateProject 更新项目。
|
||||
func (db *DB) UpdateProject(p *Project) error {
|
||||
p.UpdatedAt = time.Now()
|
||||
_, err := db.Exec(
|
||||
`UPDATE projects SET name = ?, description = ?, scope_json = ?, status = ?, pinned = ?, updated_at = ? WHERE id = ?`,
|
||||
p.Name, p.Description, p.ScopeJSON, p.Status, boolToInt(p.Pinned), p.UpdatedAt, p.ID,
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("更新项目失败: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteProject 删除项目(级联删除事实;对话 project_id 置空由 FK 处理)。
|
||||
func (db *DB) DeleteProject(id string) error {
|
||||
_, err := db.Exec(`DELETE FROM projects WHERE id = ?`, id)
|
||||
if err != nil {
|
||||
return fmt.Errorf("删除项目失败: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetConversationProjectID 返回对话绑定的项目 ID。
|
||||
func (db *DB) GetConversationProjectID(conversationID string) (string, error) {
|
||||
var pid sql.NullString
|
||||
err := db.QueryRow(`SELECT project_id FROM conversations WHERE id = ?`, conversationID).Scan(&pid)
|
||||
if err != nil {
|
||||
if err == sql.ErrNoRows {
|
||||
return "", fmt.Errorf("对话不存在")
|
||||
}
|
||||
return "", err
|
||||
}
|
||||
if pid.Valid {
|
||||
return strings.TrimSpace(pid.String), nil
|
||||
}
|
||||
return "", nil
|
||||
}
|
||||
|
||||
// SetConversationProjectID 设置对话所属项目(空字符串表示解除绑定)。
|
||||
func (db *DB) SetConversationProjectID(conversationID, projectID string) error {
|
||||
projectID = strings.TrimSpace(projectID)
|
||||
if projectID != "" {
|
||||
if _, err := db.GetProject(projectID); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
var val interface{}
|
||||
if projectID == "" {
|
||||
val = nil
|
||||
} else {
|
||||
val = projectID
|
||||
}
|
||||
_, err := db.Exec(`UPDATE conversations SET project_id = ?, updated_at = ? WHERE id = ?`, val, time.Now(), conversationID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("设置对话项目失败: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ListProjectFactsForIndex 列出用于黑板索引注入的事实(不含 deprecated,除非 includeDeprecated)。
|
||||
func (db *DB) ListProjectFactsForIndex(projectID string, includeDeprecated bool) ([]*ProjectFact, error) {
|
||||
query := `SELECT id, project_id, fact_key, category, summary, COALESCE(body,''), confidence,
|
||||
COALESCE(source_conversation_id,''), COALESCE(source_message_id,''), pinned,
|
||||
COALESCE(supersedes_fact_id,''), COALESCE(related_vulnerability_id,''), created_at, updated_at
|
||||
FROM project_facts WHERE project_id = ?`
|
||||
args := []interface{}{projectID}
|
||||
if !includeDeprecated {
|
||||
query += " AND confidence != 'deprecated'"
|
||||
}
|
||||
query += " ORDER BY pinned DESC, updated_at DESC"
|
||||
rows, err := db.Query(query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
return scanProjectFacts(rows)
|
||||
}
|
||||
|
||||
// ListProjectFacts 分页列出项目事实。
|
||||
func (db *DB) ListProjectFacts(projectID string, filter ProjectFactListFilter, limit, offset int) ([]*ProjectFact, error) {
|
||||
if limit <= 0 {
|
||||
limit = 100
|
||||
}
|
||||
query := `SELECT id, project_id, fact_key, category, summary, COALESCE(body,''), confidence,
|
||||
COALESCE(source_conversation_id,''), COALESCE(source_message_id,''), pinned,
|
||||
COALESCE(supersedes_fact_id,''), COALESCE(related_vulnerability_id,''), created_at, updated_at
|
||||
FROM project_facts WHERE project_id = ?`
|
||||
args := []interface{}{projectID}
|
||||
if c := strings.TrimSpace(filter.Category); c != "" {
|
||||
query += " AND category = ?"
|
||||
args = append(args, c)
|
||||
}
|
||||
if c := strings.TrimSpace(filter.Confidence); c != "" {
|
||||
query += " AND confidence = ?"
|
||||
args = append(args, c)
|
||||
}
|
||||
if s := strings.TrimSpace(filter.Search); s != "" {
|
||||
pat := "%" + s + "%"
|
||||
query += " AND (fact_key LIKE ? OR summary LIKE ? OR body LIKE ?)"
|
||||
args = append(args, pat, pat, pat)
|
||||
}
|
||||
query += " ORDER BY pinned DESC, updated_at DESC LIMIT ? OFFSET ?"
|
||||
args = append(args, limit, offset)
|
||||
|
||||
rows, err := db.Query(query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
return scanProjectFacts(rows)
|
||||
}
|
||||
|
||||
// GetProjectFactByKey 按 key 获取事实。
|
||||
func (db *DB) GetProjectFactByKey(projectID, factKey string) (*ProjectFact, error) {
|
||||
row := db.QueryRow(
|
||||
`SELECT id, project_id, fact_key, category, summary, COALESCE(body,''), confidence,
|
||||
COALESCE(source_conversation_id,''), COALESCE(source_message_id,''), pinned,
|
||||
COALESCE(supersedes_fact_id,''), COALESCE(related_vulnerability_id,''), created_at, updated_at
|
||||
FROM project_facts WHERE project_id = ? AND fact_key = ?`,
|
||||
projectID, factKey,
|
||||
)
|
||||
return scanProjectFactRow(row)
|
||||
}
|
||||
|
||||
// GetProjectFact 按 ID 获取事实。
|
||||
func (db *DB) GetProjectFact(id string) (*ProjectFact, error) {
|
||||
row := db.QueryRow(
|
||||
`SELECT id, project_id, fact_key, category, summary, COALESCE(body,''), confidence,
|
||||
COALESCE(source_conversation_id,''), COALESCE(source_message_id,''), pinned,
|
||||
COALESCE(supersedes_fact_id,''), COALESCE(related_vulnerability_id,''), created_at, updated_at
|
||||
FROM project_facts WHERE id = ?`, id,
|
||||
)
|
||||
return scanProjectFactRow(row)
|
||||
}
|
||||
|
||||
// mergeFactBodyOnUpdate 更新时若 incoming body 为空则保留已有内容,避免仅改 summary 时丢失攻击链。
|
||||
func mergeFactBodyOnUpdate(incoming, existing string) string {
|
||||
if strings.TrimSpace(incoming) == "" {
|
||||
return existing
|
||||
}
|
||||
return incoming
|
||||
}
|
||||
|
||||
// UpsertProjectFact 创建或更新事实(按 project_id + fact_key)。
|
||||
func (db *DB) UpsertProjectFact(f *ProjectFact) (*ProjectFact, error) {
|
||||
if err := ValidateFactKey(f.FactKey); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if strings.TrimSpace(f.Category) == "" {
|
||||
f.Category = "note"
|
||||
}
|
||||
if strings.TrimSpace(f.Confidence) == "" {
|
||||
f.Confidence = "tentative"
|
||||
}
|
||||
now := time.Now()
|
||||
|
||||
existing, err := db.GetProjectFactByKey(f.ProjectID, f.FactKey)
|
||||
if err == nil && existing != nil {
|
||||
f.ID = existing.ID
|
||||
f.CreatedAt = existing.CreatedAt
|
||||
f.UpdatedAt = now
|
||||
f.Body = mergeFactBodyOnUpdate(f.Body, existing.Body)
|
||||
_, err = db.Exec(
|
||||
`UPDATE project_facts SET category = ?, summary = ?, body = ?, confidence = ?,
|
||||
source_conversation_id = ?, source_message_id = ?, pinned = ?,
|
||||
supersedes_fact_id = ?, related_vulnerability_id = ?, updated_at = ?
|
||||
WHERE id = ?`,
|
||||
f.Category, f.Summary, f.Body, f.Confidence,
|
||||
nullIfEmpty(f.SourceConversationID), nullIfEmpty(f.SourceMessageID), boolToInt(f.Pinned),
|
||||
nullIfEmpty(f.SupersedesFactID), nullIfEmpty(f.RelatedVulnerabilityID), f.UpdatedAt, f.ID,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("更新事实失败: %w", err)
|
||||
}
|
||||
return f, nil
|
||||
}
|
||||
|
||||
if f.ID == "" {
|
||||
f.ID = uuid.New().String()
|
||||
}
|
||||
f.CreatedAt = now
|
||||
f.UpdatedAt = now
|
||||
_, err = db.Exec(
|
||||
`INSERT INTO project_facts (
|
||||
id, project_id, fact_key, category, summary, body, confidence,
|
||||
source_conversation_id, source_message_id, pinned, supersedes_fact_id, related_vulnerability_id,
|
||||
created_at, updated_at
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
|
||||
f.ID, f.ProjectID, f.FactKey, f.Category, f.Summary, f.Body, f.Confidence,
|
||||
nullIfEmpty(f.SourceConversationID), nullIfEmpty(f.SourceMessageID), boolToInt(f.Pinned),
|
||||
nullIfEmpty(f.SupersedesFactID), nullIfEmpty(f.RelatedVulnerabilityID),
|
||||
f.CreatedAt, f.UpdatedAt,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("创建事实失败: %w", err)
|
||||
}
|
||||
return f, nil
|
||||
}
|
||||
|
||||
// DeprecateProjectFact 将事实标记为 deprecated。
|
||||
func (db *DB) DeprecateProjectFact(projectID, factKey string) error {
|
||||
res, err := db.Exec(
|
||||
`UPDATE project_facts SET confidence = 'deprecated', updated_at = ? WHERE project_id = ? AND fact_key = ?`,
|
||||
time.Now(), projectID, factKey,
|
||||
)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
n, _ := res.RowsAffected()
|
||||
if n == 0 {
|
||||
return fmt.Errorf("事实不存在")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// RestoreProjectFact 将已废弃事实恢复为 tentative 或 confirmed(重新参与黑板索引)。
|
||||
func (db *DB) RestoreProjectFact(projectID, factKey, confidence string) error {
|
||||
confidence = strings.TrimSpace(strings.ToLower(confidence))
|
||||
if confidence == "" {
|
||||
confidence = "tentative"
|
||||
}
|
||||
if confidence != "confirmed" && confidence != "tentative" {
|
||||
return fmt.Errorf("confidence 须为 confirmed 或 tentative")
|
||||
}
|
||||
|
||||
existing, err := db.GetProjectFactByKey(projectID, factKey)
|
||||
if err != nil {
|
||||
return fmt.Errorf("事实不存在")
|
||||
}
|
||||
if strings.ToLower(strings.TrimSpace(existing.Confidence)) != "deprecated" {
|
||||
return fmt.Errorf("事实未处于废弃状态")
|
||||
}
|
||||
|
||||
_, err = db.Exec(
|
||||
`UPDATE project_facts SET confidence = ?, updated_at = ? WHERE project_id = ? AND fact_key = ?`,
|
||||
confidence, time.Now(), projectID, factKey,
|
||||
)
|
||||
return err
|
||||
}
|
||||
|
||||
// DeleteProjectFact 删除事实。
|
||||
func (db *DB) DeleteProjectFact(id string) error {
|
||||
_, err := db.Exec(`DELETE FROM project_facts WHERE id = ?`, id)
|
||||
return err
|
||||
}
|
||||
|
||||
func scanProjectFacts(rows *sql.Rows) ([]*ProjectFact, error) {
|
||||
var out []*ProjectFact
|
||||
for rows.Next() {
|
||||
f, err := scanProjectFactFromRows(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, f)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
func scanProjectFactRow(row *sql.Row) (*ProjectFact, error) {
|
||||
var f ProjectFact
|
||||
var pinned int
|
||||
var createdAt, updatedAt string
|
||||
err := row.Scan(
|
||||
&f.ID, &f.ProjectID, &f.FactKey, &f.Category, &f.Summary, &f.Body, &f.Confidence,
|
||||
&f.SourceConversationID, &f.SourceMessageID, &pinned,
|
||||
&f.SupersedesFactID, &f.RelatedVulnerabilityID, &createdAt, &updatedAt,
|
||||
)
|
||||
if err != nil {
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, fmt.Errorf("事实不存在")
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
f.Pinned = pinned != 0
|
||||
f.CreatedAt = parseDBTime(createdAt)
|
||||
f.UpdatedAt = parseDBTime(updatedAt)
|
||||
return &f, nil
|
||||
}
|
||||
|
||||
func scanProjectFactFromRows(rows *sql.Rows) (*ProjectFact, error) {
|
||||
var f ProjectFact
|
||||
var pinned int
|
||||
var createdAt, updatedAt string
|
||||
err := rows.Scan(
|
||||
&f.ID, &f.ProjectID, &f.FactKey, &f.Category, &f.Summary, &f.Body, &f.Confidence,
|
||||
&f.SourceConversationID, &f.SourceMessageID, &pinned,
|
||||
&f.SupersedesFactID, &f.RelatedVulnerabilityID, &createdAt, &updatedAt,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
f.Pinned = pinned != 0
|
||||
f.CreatedAt = parseDBTime(createdAt)
|
||||
f.UpdatedAt = parseDBTime(updatedAt)
|
||||
return &f, nil
|
||||
}
|
||||
|
||||
func boolToInt(b bool) int {
|
||||
if b {
|
||||
return 1
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func nullIfEmpty(s string) interface{} {
|
||||
if strings.TrimSpace(s) == "" {
|
||||
return nil
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func parseDBTime(s string) time.Time {
|
||||
s = strings.TrimSpace(s)
|
||||
if s == "" {
|
||||
return time.Time{}
|
||||
}
|
||||
// go-sqlite3 读 DATETIME 常返回 RFC3339(含 T),写入时可能是空格分隔格式,需兼容多种形态
|
||||
layouts := []string{
|
||||
time.RFC3339Nano,
|
||||
time.RFC3339,
|
||||
"2006-01-02 15:04:05.999999999-07:00",
|
||||
"2006-01-02 15:04:05-07:00",
|
||||
"2006-01-02T15:04:05.999999999-07:00",
|
||||
"2006-01-02T15:04:05-07:00",
|
||||
"2006-01-02 15:04:05.999999999",
|
||||
"2006-01-02 15:04:05",
|
||||
"2006-01-02T15:04:05.999999999",
|
||||
"2006-01-02T15:04:05",
|
||||
}
|
||||
for _, layout := range layouts {
|
||||
if t, e := time.Parse(layout, s); e == nil {
|
||||
return t
|
||||
}
|
||||
}
|
||||
return time.Time{}
|
||||
}
|
||||
@@ -0,0 +1,148 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestUpsertProjectFact_preservesBodyOnEmptyUpdate(t *testing.T) {
|
||||
dbPath := filepath.Join(t.TempDir(), "facts.db")
|
||||
db, err := NewDB(dbPath, zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
proj, err := db.CreateProject(&Project{Name: "test-facts"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
const body = "## 攻击链\n1. step\n```http\nGET / HTTP/1.1\n```\n"
|
||||
_, err = db.UpsertProjectFact(&ProjectFact{
|
||||
ProjectID: proj.ID,
|
||||
FactKey: "finding/sqli-login",
|
||||
Category: "finding",
|
||||
Summary: "SQLi on /login",
|
||||
Body: body,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
updated, err := db.UpsertProjectFact(&ProjectFact{
|
||||
ProjectID: proj.ID,
|
||||
FactKey: "finding/sqli-login",
|
||||
Summary: "SQLi on /login (confirmed)",
|
||||
Body: "",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if updated.Summary != "SQLi on /login (confirmed)" {
|
||||
t.Fatalf("summary=%q", updated.Summary)
|
||||
}
|
||||
if updated.Body != body {
|
||||
t.Fatalf("returned body=%q want preserved attack chain", updated.Body)
|
||||
}
|
||||
|
||||
fromDB, err := db.GetProjectFactByKey(proj.ID, "finding/sqli-login")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if fromDB.Body != body {
|
||||
t.Fatalf("stored body=%q want preserved", fromDB.Body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpsertProjectFact_replacesBodyWhenProvided(t *testing.T) {
|
||||
dbPath := filepath.Join(t.TempDir(), "facts.db")
|
||||
db, err := NewDB(dbPath, zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
proj, err := db.CreateProject(&Project{Name: "test-facts"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
_, err = db.UpsertProjectFact(&ProjectFact{
|
||||
ProjectID: proj.ID,
|
||||
FactKey: "target/primary",
|
||||
Summary: "v1",
|
||||
Body: "old body",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
const newBody = "new body with evidence"
|
||||
updated, err := db.UpsertProjectFact(&ProjectFact{
|
||||
ProjectID: proj.ID,
|
||||
FactKey: "target/primary",
|
||||
Summary: "v2",
|
||||
Body: newBody,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if updated.Body != newBody {
|
||||
t.Fatalf("body=%q want %q", updated.Body, newBody)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRestoreProjectFact(t *testing.T) {
|
||||
dbPath := filepath.Join(t.TempDir(), "facts.db")
|
||||
db, err := NewDB(dbPath, zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
proj, err := db.CreateProject(&Project{Name: "restore-test"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
key := "target/restore-me"
|
||||
_, err = db.UpsertProjectFact(&ProjectFact{
|
||||
ProjectID: proj.ID,
|
||||
FactKey: key,
|
||||
Summary: "s",
|
||||
Confidence: "confirmed",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.DeprecateProjectFact(proj.ID, key); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.RestoreProjectFact(proj.ID, key, "confirmed"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
f, err := db.GetProjectFactByKey(proj.ID, key)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if f.Confidence != "confirmed" {
|
||||
t.Fatalf("confidence=%q want confirmed", f.Confidence)
|
||||
}
|
||||
if err := db.RestoreProjectFact(proj.ID, key, ""); err == nil {
|
||||
t.Fatal("expected error when not deprecated")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMergeFactBodyOnUpdate(t *testing.T) {
|
||||
if got := mergeFactBodyOnUpdate("", "keep"); got != "keep" {
|
||||
t.Fatalf("empty incoming: got %q", got)
|
||||
}
|
||||
if got := mergeFactBodyOnUpdate(" ", "keep"); got != "keep" {
|
||||
t.Fatalf("whitespace incoming: got %q", got)
|
||||
}
|
||||
if got := mergeFactBodyOnUpdate("new", "old"); got != "new" {
|
||||
t.Fatalf("non-empty incoming: got %q", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,93 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestParseDBTime_projectFactFormats(t *testing.T) {
|
||||
cases := []string{
|
||||
"2026-05-26 11:13:07.442143+08:00",
|
||||
"2026-05-26 11:13:07",
|
||||
"2026-05-26T11:13:07.442143+08:00",
|
||||
}
|
||||
for _, s := range cases {
|
||||
got := parseDBTime(s)
|
||||
if got.IsZero() {
|
||||
t.Fatalf("parseDBTime(%q) returned zero", s)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestListProjectFacts_updatedAtJSON(t *testing.T) {
|
||||
root, err := os.Getwd()
|
||||
if err != nil {
|
||||
t.Skip(err)
|
||||
}
|
||||
dbPath := filepath.Join(root, "..", "..", "data", "conversations.db")
|
||||
if _, err := os.Stat(dbPath); err != nil {
|
||||
t.Skip("conversations.db not found")
|
||||
}
|
||||
db, err := NewDB(dbPath, zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
projects, err := db.ListProjects("", 1, 0)
|
||||
if err != nil || len(projects) == 0 {
|
||||
t.Skip("no projects")
|
||||
}
|
||||
pid := projects[0].ID
|
||||
|
||||
list, err := db.ListProjectFacts(pid, ProjectFactListFilter{}, 5, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(list) == 0 {
|
||||
t.Skip("no facts")
|
||||
}
|
||||
for _, f := range list {
|
||||
if f.UpdatedAt.IsZero() {
|
||||
t.Fatalf("fact %s UpdatedAt is zero after ListProjectFacts", f.FactKey)
|
||||
}
|
||||
b, err := json.Marshal(f)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var m map[string]interface{}
|
||||
if err := json.Unmarshal(b, &m); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
raw, ok := m["updated_at"].(string)
|
||||
if !ok || raw == "" || raw[:4] == "0001" {
|
||||
t.Fatalf("bad updated_at in JSON: %v", m["updated_at"])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseDBTime_zeroOnGarbage(t *testing.T) {
|
||||
if !parseDBTime("").IsZero() {
|
||||
t.Fatal("expected zero for empty")
|
||||
}
|
||||
}
|
||||
|
||||
// Ensure RFC3339 round-trip used by API is after year 2000.
|
||||
func TestParseDBTime_marshalRoundTrip(t *testing.T) {
|
||||
s := "2026-05-26 11:13:07.442143+08:00"
|
||||
tm := parseDBTime(s)
|
||||
b, err := json.Marshal(tm)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var back time.Time
|
||||
if err := json.Unmarshal(b, &back); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if back.IsZero() {
|
||||
t.Fatalf("unmarshal zero from %s", string(b))
|
||||
}
|
||||
}
|
||||
@@ -3,16 +3,94 @@ package database
|
||||
import (
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// VulnerabilityListFilter 列表/统计/导出共用的筛选条件
|
||||
type VulnerabilityListFilter struct {
|
||||
ID string
|
||||
Search string // 关键词模糊匹配(标题、描述、类型、目标等)
|
||||
ConversationID string
|
||||
ProjectID string
|
||||
Severity string
|
||||
Status string
|
||||
TaskID string
|
||||
ConversationTag string
|
||||
TaskTag string
|
||||
}
|
||||
|
||||
func escapeVulnerabilityLikePattern(s string) string {
|
||||
s = strings.ReplaceAll(s, `\`, `\\`)
|
||||
s = strings.ReplaceAll(s, `%`, `\%`)
|
||||
s = strings.ReplaceAll(s, `_`, `\_`)
|
||||
return "%" + s + "%"
|
||||
}
|
||||
|
||||
func (f VulnerabilityListFilter) appendWhere(query string, args []interface{}) (string, []interface{}) {
|
||||
if f.ID != "" {
|
||||
query += " AND id = ?"
|
||||
args = append(args, f.ID)
|
||||
}
|
||||
if f.ConversationID != "" {
|
||||
query += " AND conversation_id = ?"
|
||||
args = append(args, f.ConversationID)
|
||||
}
|
||||
if f.ProjectID != "" {
|
||||
query += " AND project_id = ?"
|
||||
args = append(args, f.ProjectID)
|
||||
}
|
||||
if f.TaskID != "" {
|
||||
query += " AND EXISTS (SELECT 1 FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id AND (bt.id = ? OR bt.queue_id = ?))"
|
||||
args = append(args, f.TaskID, f.TaskID)
|
||||
}
|
||||
if f.ConversationTag != "" {
|
||||
query += " AND conversation_tag = ?"
|
||||
args = append(args, f.ConversationTag)
|
||||
}
|
||||
if f.TaskTag != "" {
|
||||
query += " AND task_tag = ?"
|
||||
args = append(args, f.TaskTag)
|
||||
}
|
||||
if f.Severity != "" {
|
||||
query += " AND severity = ?"
|
||||
args = append(args, f.Severity)
|
||||
}
|
||||
if f.Status != "" {
|
||||
query += " AND status = ?"
|
||||
args = append(args, f.Status)
|
||||
}
|
||||
search := strings.TrimSpace(f.Search)
|
||||
if search != "" {
|
||||
pattern := escapeVulnerabilityLikePattern(search)
|
||||
query += ` AND (
|
||||
LOWER(id) LIKE LOWER(?) OR
|
||||
LOWER(title) LIKE LOWER(?) OR
|
||||
LOWER(COALESCE(description, '')) LIKE LOWER(?) OR
|
||||
LOWER(COALESCE(vulnerability_type, '')) LIKE LOWER(?) OR
|
||||
LOWER(COALESCE(target, '')) LIKE LOWER(?) OR
|
||||
LOWER(COALESCE(proof, '')) LIKE LOWER(?) OR
|
||||
LOWER(COALESCE(impact, '')) LIKE LOWER(?) OR
|
||||
LOWER(COALESCE(recommendation, '')) LIKE LOWER(?) OR
|
||||
LOWER(COALESCE(conversation_id, '')) LIKE LOWER(?) OR
|
||||
LOWER(COALESCE(conversation_tag, '')) LIKE LOWER(?) OR
|
||||
LOWER(COALESCE(task_tag, '')) LIKE LOWER(?)
|
||||
)`
|
||||
for i := 0; i < 11; i++ {
|
||||
args = append(args, pattern)
|
||||
}
|
||||
}
|
||||
return query, args
|
||||
}
|
||||
|
||||
// Vulnerability 漏洞
|
||||
type Vulnerability struct {
|
||||
ID string `json:"id"`
|
||||
ConversationID string `json:"conversation_id"`
|
||||
ProjectID string `json:"project_id,omitempty"`
|
||||
ConversationTag string `json:"conversation_tag,omitempty"`
|
||||
TaskTag string `json:"task_tag,omitempty"`
|
||||
TaskID string `json:"task_id,omitempty"`
|
||||
@@ -44,17 +122,23 @@ func (db *DB) CreateVulnerability(vuln *Vulnerability) (*Vulnerability, error) {
|
||||
}
|
||||
vuln.UpdatedAt = now
|
||||
|
||||
if strings.TrimSpace(vuln.ProjectID) == "" && vuln.ConversationID != "" {
|
||||
if pid, err := db.GetConversationProjectID(vuln.ConversationID); err == nil {
|
||||
vuln.ProjectID = pid
|
||||
}
|
||||
}
|
||||
|
||||
query := `
|
||||
INSERT INTO vulnerabilities (
|
||||
id, conversation_id, conversation_tag, task_tag, title, description, severity, status,
|
||||
id, conversation_id, project_id, conversation_tag, task_tag, title, description, severity, status,
|
||||
vulnerability_type, target, proof, impact, recommendation,
|
||||
created_at, updated_at
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
`
|
||||
|
||||
_, err := db.Exec(
|
||||
query,
|
||||
vuln.ID, vuln.ConversationID, vuln.ConversationTag, vuln.TaskTag, vuln.Title, vuln.Description,
|
||||
vuln.ID, vuln.ConversationID, nullIfEmpty(vuln.ProjectID), vuln.ConversationTag, vuln.TaskTag, vuln.Title, vuln.Description,
|
||||
vuln.Severity, vuln.Status, vuln.Type, vuln.Target,
|
||||
vuln.Proof, vuln.Impact, vuln.Recommendation,
|
||||
vuln.CreatedAt, vuln.UpdatedAt,
|
||||
@@ -70,7 +154,7 @@ func (db *DB) CreateVulnerability(vuln *Vulnerability) (*Vulnerability, error) {
|
||||
func (db *DB) GetVulnerability(id string) (*Vulnerability, error) {
|
||||
var vuln Vulnerability
|
||||
query := `
|
||||
SELECT id, conversation_id, title, description, severity, status,
|
||||
SELECT id, conversation_id, COALESCE(project_id,''), title, description, severity, status,
|
||||
conversation_tag, task_tag, vulnerability_type, target, proof, impact, recommendation,
|
||||
COALESCE((SELECT bt.id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_id,
|
||||
COALESCE((SELECT bt.queue_id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_queue_id,
|
||||
@@ -80,7 +164,7 @@ func (db *DB) GetVulnerability(id string) (*Vulnerability, error) {
|
||||
`
|
||||
|
||||
err := db.QueryRow(query, id).Scan(
|
||||
&vuln.ID, &vuln.ConversationID, &vuln.Title, &vuln.Description,
|
||||
&vuln.ID, &vuln.ConversationID, &vuln.ProjectID, &vuln.Title, &vuln.Description,
|
||||
&vuln.Severity, &vuln.Status, &vuln.ConversationTag, &vuln.TaskTag, &vuln.Type, &vuln.Target,
|
||||
&vuln.Proof, &vuln.Impact, &vuln.Recommendation,
|
||||
&vuln.TaskID, &vuln.TaskQueueID,
|
||||
@@ -97,9 +181,9 @@ func (db *DB) GetVulnerability(id string) (*Vulnerability, error) {
|
||||
}
|
||||
|
||||
// ListVulnerabilities 列出漏洞
|
||||
func (db *DB) ListVulnerabilities(limit, offset int, id, conversationID, severity, status, taskID, conversationTag, taskTag string) ([]*Vulnerability, error) {
|
||||
func (db *DB) ListVulnerabilities(limit, offset int, filter VulnerabilityListFilter) ([]*Vulnerability, error) {
|
||||
query := `
|
||||
SELECT id, conversation_id, title, description, severity, status, conversation_tag, task_tag,
|
||||
SELECT id, conversation_id, COALESCE(project_id,''), title, description, severity, status, conversation_tag, task_tag,
|
||||
vulnerability_type, target, proof, impact, recommendation,
|
||||
COALESCE((SELECT bt.id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_id,
|
||||
COALESCE((SELECT bt.queue_id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_queue_id,
|
||||
@@ -108,35 +192,7 @@ func (db *DB) ListVulnerabilities(limit, offset int, id, conversationID, severit
|
||||
WHERE 1=1
|
||||
`
|
||||
args := []interface{}{}
|
||||
|
||||
if id != "" {
|
||||
query += " AND id = ?"
|
||||
args = append(args, id)
|
||||
}
|
||||
if conversationID != "" {
|
||||
query += " AND conversation_id = ?"
|
||||
args = append(args, conversationID)
|
||||
}
|
||||
if taskID != "" {
|
||||
query += " AND EXISTS (SELECT 1 FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id AND (bt.id = ? OR bt.queue_id = ?))"
|
||||
args = append(args, taskID, taskID)
|
||||
}
|
||||
if conversationTag != "" {
|
||||
query += " AND conversation_tag = ?"
|
||||
args = append(args, conversationTag)
|
||||
}
|
||||
if taskTag != "" {
|
||||
query += " AND task_tag = ?"
|
||||
args = append(args, taskTag)
|
||||
}
|
||||
if severity != "" {
|
||||
query += " AND severity = ?"
|
||||
args = append(args, severity)
|
||||
}
|
||||
if status != "" {
|
||||
query += " AND status = ?"
|
||||
args = append(args, status)
|
||||
}
|
||||
query, args = filter.appendWhere(query, args)
|
||||
|
||||
query += " ORDER BY created_at DESC LIMIT ? OFFSET ?"
|
||||
args = append(args, limit, offset)
|
||||
@@ -151,7 +207,7 @@ func (db *DB) ListVulnerabilities(limit, offset int, id, conversationID, severit
|
||||
for rows.Next() {
|
||||
var vuln Vulnerability
|
||||
err := rows.Scan(
|
||||
&vuln.ID, &vuln.ConversationID, &vuln.Title, &vuln.Description,
|
||||
&vuln.ID, &vuln.ConversationID, &vuln.ProjectID, &vuln.Title, &vuln.Description,
|
||||
&vuln.Severity, &vuln.Status, &vuln.ConversationTag, &vuln.TaskTag, &vuln.Type, &vuln.Target,
|
||||
&vuln.Proof, &vuln.Impact, &vuln.Recommendation,
|
||||
&vuln.TaskID, &vuln.TaskQueueID,
|
||||
@@ -168,38 +224,10 @@ func (db *DB) ListVulnerabilities(limit, offset int, id, conversationID, severit
|
||||
}
|
||||
|
||||
// CountVulnerabilities 统计漏洞总数(支持筛选条件)
|
||||
func (db *DB) CountVulnerabilities(id, conversationID, severity, status, taskID, conversationTag, taskTag string) (int, error) {
|
||||
func (db *DB) CountVulnerabilities(filter VulnerabilityListFilter) (int, error) {
|
||||
query := "SELECT COUNT(*) FROM vulnerabilities WHERE 1=1"
|
||||
args := []interface{}{}
|
||||
|
||||
if id != "" {
|
||||
query += " AND id = ?"
|
||||
args = append(args, id)
|
||||
}
|
||||
if conversationID != "" {
|
||||
query += " AND conversation_id = ?"
|
||||
args = append(args, conversationID)
|
||||
}
|
||||
if taskID != "" {
|
||||
query += " AND EXISTS (SELECT 1 FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id AND (bt.id = ? OR bt.queue_id = ?))"
|
||||
args = append(args, taskID, taskID)
|
||||
}
|
||||
if conversationTag != "" {
|
||||
query += " AND conversation_tag = ?"
|
||||
args = append(args, conversationTag)
|
||||
}
|
||||
if taskTag != "" {
|
||||
query += " AND task_tag = ?"
|
||||
args = append(args, taskTag)
|
||||
}
|
||||
if severity != "" {
|
||||
query += " AND severity = ?"
|
||||
args = append(args, severity)
|
||||
}
|
||||
if status != "" {
|
||||
query += " AND status = ?"
|
||||
args = append(args, status)
|
||||
}
|
||||
query, args = filter.appendWhere(query, args)
|
||||
|
||||
var count int
|
||||
err := db.QueryRow(query, args...).Scan(&count)
|
||||
@@ -216,7 +244,7 @@ func (db *DB) UpdateVulnerability(id string, vuln *Vulnerability) error {
|
||||
|
||||
query := `
|
||||
UPDATE vulnerabilities
|
||||
SET conversation_tag = ?, task_tag = ?, title = ?, description = ?, severity = ?, status = ?,
|
||||
SET project_id = ?, conversation_tag = ?, task_tag = ?, title = ?, description = ?, severity = ?, status = ?,
|
||||
vulnerability_type = ?, target = ?, proof = ?, impact = ?,
|
||||
recommendation = ?, updated_at = ?
|
||||
WHERE id = ?
|
||||
@@ -224,7 +252,7 @@ func (db *DB) UpdateVulnerability(id string, vuln *Vulnerability) error {
|
||||
|
||||
_, err := db.Exec(
|
||||
query,
|
||||
vuln.ConversationTag, vuln.TaskTag, vuln.Title, vuln.Description, vuln.Severity, vuln.Status,
|
||||
nullIfEmpty(vuln.ProjectID), vuln.ConversationTag, vuln.TaskTag, vuln.Title, vuln.Description, vuln.Severity, vuln.Status,
|
||||
vuln.Type, vuln.Target, vuln.Proof, vuln.Impact,
|
||||
vuln.Recommendation, vuln.UpdatedAt, id,
|
||||
)
|
||||
@@ -245,19 +273,12 @@ func (db *DB) DeleteVulnerability(id string) error {
|
||||
}
|
||||
|
||||
// GetVulnerabilityStats 获取漏洞统计(筛选条件与 ListVulnerabilities / CountVulnerabilities 一致)
|
||||
func (db *DB) GetVulnerabilityStats(conversationID, taskID string) (map[string]interface{}, error) {
|
||||
func (db *DB) GetVulnerabilityStats(filter VulnerabilityListFilter) (map[string]interface{}, error) {
|
||||
stats := make(map[string]interface{})
|
||||
|
||||
where := "WHERE 1=1"
|
||||
args := []interface{}{}
|
||||
if conversationID != "" {
|
||||
where += " AND conversation_id = ?"
|
||||
args = append(args, conversationID)
|
||||
}
|
||||
if taskID != "" {
|
||||
where += " AND EXISTS (SELECT 1 FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id AND (bt.id = ? OR bt.queue_id = ?))"
|
||||
args = append(args, taskID, taskID)
|
||||
}
|
||||
where, args = filter.appendWhere(where, args)
|
||||
|
||||
// 总漏洞数
|
||||
var totalCount int
|
||||
@@ -357,10 +378,15 @@ func (db *DB) GetVulnerabilityFilterOptions() (map[string][]string, error) {
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询任务标签建议失败: %w", err)
|
||||
}
|
||||
projectIDs, err := collect(`SELECT DISTINCT project_id FROM vulnerabilities WHERE project_id IS NOT NULL AND project_id <> '' ORDER BY created_at DESC LIMIT 200`)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询项目ID建议失败: %w", err)
|
||||
}
|
||||
|
||||
return map[string][]string{
|
||||
"vulnerability_ids": vulnIDs,
|
||||
"conversation_ids": conversationIDs,
|
||||
"project_ids": projectIDs,
|
||||
"task_ids": taskIDs,
|
||||
"queue_ids": queueIDs,
|
||||
"conversation_tags": conversationTags,
|
||||
|
||||
+210
-55
@@ -17,12 +17,14 @@ import (
|
||||
"unicode/utf8"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/reasoning"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
"cyberstrike-ai/internal/multiagent"
|
||||
"cyberstrike-ai/internal/openai"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/robfig/cron/v3"
|
||||
@@ -130,6 +132,12 @@ type AgentHandler struct {
|
||||
batchRunning map[string]struct{}
|
||||
// hitlWhitelistSaver 侧栏「应用」HITL 时将会话增量白名单合并写入 config.yaml(可选)
|
||||
hitlWhitelistSaver HitlToolWhitelistSaver
|
||||
audit *audit.Service
|
||||
}
|
||||
|
||||
// SetAudit wires platform audit logging.
|
||||
func (h *AgentHandler) SetAudit(s *audit.Service) {
|
||||
h.audit = s
|
||||
}
|
||||
|
||||
// HitlToolWhitelistSaver 合并 HITL 免审批工具到全局配置并落盘
|
||||
@@ -206,7 +214,7 @@ type ChatAttachment struct {
|
||||
type ChatReasoningRequest struct {
|
||||
// Mode: default(跟随系统)| off | on | auto
|
||||
Mode string `json:"mode,omitempty"`
|
||||
// Effort: low | medium | high | max;空表示不指定(由系统默认与各 profile 决定)。
|
||||
// Effort: low | medium | high | max | xhigh(原样下发;不同网关最高档命名不同)。空表示不指定。
|
||||
Effort string `json:"effort,omitempty"`
|
||||
}
|
||||
|
||||
@@ -214,6 +222,7 @@ type ChatReasoningRequest struct {
|
||||
type ChatRequest struct {
|
||||
Message string `json:"message" binding:"required"`
|
||||
ConversationID string `json:"conversationId,omitempty"`
|
||||
ProjectID string `json:"projectId,omitempty"` // 新对话绑定的项目(可选;未指定时可用 config.project.default_project_id)
|
||||
Role string `json:"role,omitempty"` // 角色名称
|
||||
Attachments []ChatAttachment `json:"attachments,omitempty"`
|
||||
WebShellConnectionID string `json:"webshellConnectionId,omitempty"` // WebShell 管理 - AI 助手:当前选中的连接 ID,仅使用 webshell_* 工具
|
||||
@@ -552,7 +561,9 @@ func (h *AgentHandler) AgentLoop(c *gin.Context) {
|
||||
conversationID := req.ConversationID
|
||||
if conversationID == "" {
|
||||
title := safeTruncateString(req.Message, 50)
|
||||
conv, err := h.db.CreateConversation(title)
|
||||
meta := audit.ConversationCreateMetaFromGin(c, "agent_loop")
|
||||
meta.ProjectID = effectiveProjectID(h.config, req.ProjectID)
|
||||
conv, err := h.db.CreateConversation(title, meta)
|
||||
if err != nil {
|
||||
h.logger.Error("创建对话失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
@@ -627,6 +638,8 @@ func (h *AgentHandler) AgentLoop(c *gin.Context) {
|
||||
builtin.ToolWebshellFileRead,
|
||||
builtin.ToolWebshellFileWrite,
|
||||
builtin.ToolRecordVulnerability,
|
||||
builtin.ToolListVulnerabilities,
|
||||
builtin.ToolGetVulnerability,
|
||||
builtin.ToolListKnowledgeRiskTypes,
|
||||
builtin.ToolSearchKnowledgeBase,
|
||||
}
|
||||
@@ -674,7 +687,7 @@ func (h *AgentHandler) AgentLoop(c *gin.Context) {
|
||||
taskCtx = h.injectReactHITLInterceptor(taskCtx, cancelWithCause, conversationID, "", nil)
|
||||
|
||||
// 执行Agent Loop,传入历史消息和对话ID(使用包含角色提示词的finalMessage和角色工具列表)
|
||||
result, err := h.agent.AgentLoopWithProgress(taskCtx, finalMessage, agentHistoryMessages, conversationID, progressCallback, roleTools)
|
||||
result, err := h.agent.AgentLoopWithProgress(taskCtx, finalMessage, agentHistoryMessages, conversationID, progressCallback, roleTools, h.projectBlackboardBlock(conversationID))
|
||||
if err != nil {
|
||||
h.logger.Error("Agent Loop执行失败", zap.Error(err))
|
||||
|
||||
@@ -716,11 +729,45 @@ func (h *AgentHandler) AgentLoop(c *gin.Context) {
|
||||
})
|
||||
}
|
||||
|
||||
func (h *AgentHandler) finalizeRobotAgentError(ctx context.Context, assistantMessageID, conversationID string, resultMA *multiagent.RunResult, errMA error) (string, string, error) {
|
||||
if shouldPersistEinoAgentTraceAfterRunError(ctx) {
|
||||
h.persistEinoAgentTraceForResume(conversationID, resultMA)
|
||||
}
|
||||
errMsg := "执行失败: " + errMA.Error()
|
||||
if assistantMessageID != "" {
|
||||
_, _ = h.db.Exec("UPDATE messages SET content = ?, updated_at = ? WHERE id = ?", errMsg, time.Now(), assistantMessageID)
|
||||
_ = h.db.AddProcessDetail(assistantMessageID, conversationID, "error", errMsg, nil)
|
||||
}
|
||||
return "", conversationID, errMA
|
||||
}
|
||||
|
||||
func (h *AgentHandler) finalizeRobotAgentSuccess(assistantMessageID, conversationID string, resultMA *multiagent.RunResult) (string, string, error) {
|
||||
if assistantMessageID != "" {
|
||||
if errU := h.db.UpdateAssistantMessageFinalize(assistantMessageID, resultMA.Response, resultMA.MCPExecutionIDs, multiagent.AggregatedReasoningFromTraceJSON(resultMA.LastAgentTraceInput)); errU != nil {
|
||||
h.logger.Warn("机器人:更新助手消息失败", zap.Error(errU))
|
||||
}
|
||||
} else {
|
||||
if _, err := h.db.AddMessage(conversationID, "assistant", resultMA.Response, resultMA.MCPExecutionIDs); err != nil {
|
||||
h.logger.Warn("机器人:保存助手消息失败", zap.Error(err))
|
||||
}
|
||||
}
|
||||
if resultMA.LastAgentTraceInput != "" || resultMA.LastAgentTraceOutput != "" {
|
||||
_ = h.db.SaveAgentTrace(conversationID, resultMA.LastAgentTraceInput, resultMA.LastAgentTraceOutput)
|
||||
}
|
||||
return resultMA.Response, conversationID, nil
|
||||
}
|
||||
|
||||
// ProcessMessageForRobot 供机器人(企业微信/钉钉/飞书)调用:与 /api/agent-loop/stream 相同执行路径(含 progressCallback、过程详情),仅不发送 SSE,最后返回完整回复
|
||||
func (h *AgentHandler) ProcessMessageForRobot(ctx context.Context, conversationID, message, role string) (response string, convID string, err error) {
|
||||
func (h *AgentHandler) ProcessMessageForRobot(ctx context.Context, platform, conversationID, message, role string) (response string, convID string, err error) {
|
||||
if conversationID == "" {
|
||||
title := safeTruncateString(message, 50)
|
||||
conv, createErr := h.db.CreateConversation(title)
|
||||
src := "robot"
|
||||
if strings.TrimSpace(platform) != "" {
|
||||
src = "robot:" + strings.TrimSpace(platform)
|
||||
}
|
||||
meta := audit.ConversationCreateMeta(src)
|
||||
meta.ProjectID = effectiveProjectID(h.config, "")
|
||||
conv, createErr := h.db.CreateConversation(title, meta)
|
||||
if createErr != nil {
|
||||
return "", "", fmt.Errorf("创建对话失败: %w", createErr)
|
||||
}
|
||||
@@ -768,53 +815,92 @@ func (h *AgentHandler) ProcessMessageForRobot(ctx context.Context, conversationI
|
||||
if assistantMsg != nil {
|
||||
assistantMessageID = assistantMsg.ID
|
||||
}
|
||||
progressCallback := h.createProgressCallback(ctx, nil, conversationID, assistantMessageID, nil)
|
||||
|
||||
useRobotMulti := h.config != nil && h.config.MultiAgent.Enabled && h.config.MultiAgent.RobotUseMultiAgent
|
||||
if useRobotMulti {
|
||||
resultMA, errMA := multiagent.RunDeepAgent(
|
||||
ctx,
|
||||
h.config,
|
||||
&h.config.MultiAgent,
|
||||
h.agent,
|
||||
h.logger,
|
||||
conversationID,
|
||||
finalMessage,
|
||||
agentHistoryMessages,
|
||||
roleTools,
|
||||
progressCallback,
|
||||
h.agentsMarkdownDir,
|
||||
"deep",
|
||||
nil,
|
||||
)
|
||||
if errMA != nil {
|
||||
if shouldPersistEinoAgentTraceAfterRunError(ctx) {
|
||||
h.persistEinoAgentTraceForResume(conversationID, resultMA)
|
||||
}
|
||||
errMsg := "执行失败: " + errMA.Error()
|
||||
if assistantMessageID != "" {
|
||||
_, _ = h.db.Exec("UPDATE messages SET content = ?, updated_at = ? WHERE id = ?", errMsg, time.Now(), assistantMessageID)
|
||||
_ = h.db.AddProcessDetail(assistantMessageID, conversationID, "error", errMsg, nil)
|
||||
}
|
||||
return "", conversationID, errMA
|
||||
// 注册运行中任务并向 taskEventBus 镜像进度事件,供 Web 端 task-events 补流(与 agent-loop/stream 一致)。
|
||||
taskCtx, cancelWithCause := context.WithCancelCause(ctx)
|
||||
defer cancelWithCause(nil)
|
||||
taskStatus := "completed"
|
||||
defer func() {
|
||||
h.tasks.FinishTask(conversationID, taskStatus)
|
||||
}()
|
||||
if _, err := h.tasks.StartTask(conversationID, message, cancelWithCause); err != nil {
|
||||
if errors.Is(err, ErrTaskAlreadyRunning) {
|
||||
return "", conversationID, fmt.Errorf("当前会话已有任务正在执行中,请稍后再试")
|
||||
}
|
||||
if assistantMessageID != "" {
|
||||
if errU := h.db.UpdateAssistantMessageFinalize(assistantMessageID, resultMA.Response, resultMA.MCPExecutionIDs, multiagent.AggregatedReasoningFromTraceJSON(resultMA.LastAgentTraceInput)); errU != nil {
|
||||
h.logger.Warn("机器人:更新助手消息失败", zap.Error(errU))
|
||||
return "", conversationID, fmt.Errorf("无法启动任务: %w", err)
|
||||
}
|
||||
progressCallback := h.createProgressCallback(taskCtx, cancelWithCause, conversationID, assistantMessageID, nil)
|
||||
|
||||
robotMode := "react"
|
||||
if h.config != nil {
|
||||
robotMode = config.NormalizeRobotAgentMode(h.config.MultiAgent)
|
||||
}
|
||||
switch robotMode {
|
||||
case "eino_single":
|
||||
curHist := agentHistoryMessages
|
||||
curMsg := finalMessage
|
||||
segmentUserMessage := finalMessage
|
||||
var resultMA *multiagent.RunResult
|
||||
var errMA error
|
||||
var transientRunAttempts int
|
||||
for {
|
||||
resultMA, errMA = multiagent.RunEinoSingleChatModelAgent(
|
||||
taskCtx, h.config, &h.config.MultiAgent, h.agent, h.logger,
|
||||
conversationID, curMsg, curHist, roleTools, progressCallback, nil, h.projectBlackboardBlock(conversationID),
|
||||
)
|
||||
if errMA == nil {
|
||||
// 成功后重置 transient 重试窗口,下一次分段从第 1 次重试开始。
|
||||
transientRunAttempts = 0
|
||||
break
|
||||
}
|
||||
} else {
|
||||
if _, err = h.db.AddMessage(conversationID, "assistant", resultMA.Response, resultMA.MCPExecutionIDs); err != nil {
|
||||
h.logger.Warn("机器人:保存助手消息失败", zap.Error(err))
|
||||
if handled, _ := h.handleEinoTransientRetryContinue(
|
||||
taskCtx, conversationID, resultMA, errMA, &transientRunAttempts,
|
||||
&curHist, &curMsg, segmentUserMessage, progressCallback, nil,
|
||||
); handled {
|
||||
continue
|
||||
}
|
||||
taskStatus = "failed"
|
||||
return h.finalizeRobotAgentError(taskCtx, assistantMessageID, conversationID, resultMA, errMA)
|
||||
}
|
||||
if resultMA.LastAgentTraceInput != "" || resultMA.LastAgentTraceOutput != "" {
|
||||
_ = h.db.SaveAgentTrace(conversationID, resultMA.LastAgentTraceInput, resultMA.LastAgentTraceOutput)
|
||||
return h.finalizeRobotAgentSuccess(assistantMessageID, conversationID, resultMA)
|
||||
case "deep", "plan_execute", "supervisor":
|
||||
if h.config == nil || !h.config.MultiAgent.Enabled {
|
||||
h.logger.Warn("机器人配置为多代理模式但未启用 multi_agent,回退原生 ReAct",
|
||||
zap.String("robot_mode", robotMode))
|
||||
break
|
||||
}
|
||||
return resultMA.Response, conversationID, nil
|
||||
curHist := agentHistoryMessages
|
||||
curMsg := finalMessage
|
||||
segmentUserMessage := finalMessage
|
||||
var resultMA *multiagent.RunResult
|
||||
var errMA error
|
||||
var transientRunAttempts int
|
||||
for {
|
||||
resultMA, errMA = multiagent.RunDeepAgent(
|
||||
taskCtx, h.config, &h.config.MultiAgent, h.agent, h.logger,
|
||||
conversationID, curMsg, curHist, roleTools, progressCallback,
|
||||
h.agentsMarkdownDir, robotMode, nil, h.projectBlackboardBlock(conversationID),
|
||||
)
|
||||
if errMA == nil {
|
||||
// 成功后重置 transient 重试窗口,下一次分段从第 1 次重试开始。
|
||||
transientRunAttempts = 0
|
||||
break
|
||||
}
|
||||
if handled, _ := h.handleEinoTransientRetryContinue(
|
||||
taskCtx, conversationID, resultMA, errMA, &transientRunAttempts,
|
||||
&curHist, &curMsg, segmentUserMessage, progressCallback, nil,
|
||||
); handled {
|
||||
continue
|
||||
}
|
||||
taskStatus = "failed"
|
||||
return h.finalizeRobotAgentError(taskCtx, assistantMessageID, conversationID, resultMA, errMA)
|
||||
}
|
||||
return h.finalizeRobotAgentSuccess(assistantMessageID, conversationID, resultMA)
|
||||
}
|
||||
|
||||
result, err := h.agent.AgentLoopWithProgress(ctx, finalMessage, agentHistoryMessages, conversationID, progressCallback, roleTools)
|
||||
result, err := h.agent.AgentLoopWithProgress(taskCtx, finalMessage, agentHistoryMessages, conversationID, progressCallback, roleTools, h.projectBlackboardBlock(conversationID))
|
||||
if err != nil {
|
||||
taskStatus = "failed"
|
||||
errMsg := "执行失败: " + err.Error()
|
||||
if assistantMessageID != "" {
|
||||
_, _ = h.db.Exec("UPDATE messages SET content = ?, updated_at = ? WHERE id = ?", errMsg, time.Now(), assistantMessageID)
|
||||
@@ -846,6 +932,23 @@ type StreamEvent struct {
|
||||
Data interface{} `json:"data,omitempty"`
|
||||
}
|
||||
|
||||
// publishProgressToTaskEventBus 将进度事件镜像到 taskEventBus(机器人/无 HTTP SSE 客户端时供 Web task-events 订阅)。
|
||||
func (h *AgentHandler) publishProgressToTaskEventBus(conversationID, eventType, message string, data interface{}) {
|
||||
if h == nil || h.taskEventBus == nil || strings.TrimSpace(conversationID) == "" {
|
||||
return
|
||||
}
|
||||
event := StreamEvent{Type: eventType, Message: message, Data: data}
|
||||
eventJSON, err := json.Marshal(event)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
sseLine := make([]byte, 0, len(eventJSON)+8)
|
||||
sseLine = append(sseLine, []byte("data: ")...)
|
||||
sseLine = append(sseLine, eventJSON...)
|
||||
sseLine = append(sseLine, '\n', '\n')
|
||||
h.taskEventBus.Publish(conversationID, sseLine)
|
||||
}
|
||||
|
||||
// createProgressCallback 创建进度回调函数,用于保存processDetails
|
||||
// sendEventFunc: 可选的流式事件发送函数,如果为nil则不发送流式事件
|
||||
func (h *AgentHandler) createProgressCallback(runCtx context.Context, cancelRun context.CancelCauseFunc, conversationID, assistantMessageID string, sendEventFunc func(eventType, message string, data interface{})) agent.ProgressCallback {
|
||||
@@ -955,9 +1058,11 @@ func (h *AgentHandler) createProgressCallback(runCtx context.Context, cancelRun
|
||||
}
|
||||
|
||||
return func(eventType, message string, data interface{}) {
|
||||
// 如果提供了sendEventFunc,发送流式事件
|
||||
// 流式:写 HTTP SSE;非流式(机器人等):镜像到 taskEventBus 供 Web 订阅
|
||||
if sendEventFunc != nil {
|
||||
sendEventFunc(eventType, message, data)
|
||||
} else {
|
||||
h.publishProgressToTaskEventBus(conversationID, eventType, message, data)
|
||||
}
|
||||
|
||||
// 保存tool_call事件中的参数
|
||||
@@ -1158,7 +1263,16 @@ func (h *AgentHandler) createProgressCallback(runCtx context.Context, cancelRun
|
||||
return
|
||||
}
|
||||
if eventType == "response_delta" {
|
||||
respPlan.b.WriteString(message)
|
||||
if dataMap, ok := data.(map[string]interface{}); ok {
|
||||
if acc, okAcc := dataMap[openai.SSEAccumulatedKey].(string); okAcc {
|
||||
respPlan.b.Reset()
|
||||
respPlan.b.WriteString(acc)
|
||||
} else {
|
||||
respPlan.b.WriteString(message)
|
||||
}
|
||||
} else {
|
||||
respPlan.b.WriteString(message)
|
||||
}
|
||||
if dataMap, ok := data.(map[string]interface{}); ok && respPlan.meta == nil {
|
||||
respPlan.meta = make(map[string]interface{}, len(dataMap))
|
||||
for k, v := range dataMap {
|
||||
@@ -1213,8 +1327,12 @@ func (h *AgentHandler) createProgressCallback(runCtx context.Context, cancelRun
|
||||
} else if tb.persistAs == "" {
|
||||
tb.persistAs = persistAs
|
||||
}
|
||||
// delta 片段直接拼接
|
||||
tb.b.WriteString(message)
|
||||
if acc, okAcc := dataMap[openai.SSEAccumulatedKey].(string); okAcc {
|
||||
tb.b.Reset()
|
||||
tb.b.WriteString(acc)
|
||||
} else {
|
||||
tb.b.WriteString(message)
|
||||
}
|
||||
// 有时 delta 先到 start 未到,补充元信息
|
||||
for k, v := range dataMap {
|
||||
tb.meta[k] = v
|
||||
@@ -1406,10 +1524,13 @@ func (h *AgentHandler) AgentLoopStream(c *gin.Context) {
|
||||
title := safeTruncateString(req.Message, 50)
|
||||
var conv *database.Conversation
|
||||
var err error
|
||||
meta := audit.ConversationCreateMetaFromGin(c, "agent_loop_stream")
|
||||
meta.ProjectID = effectiveProjectID(h.config, req.ProjectID)
|
||||
if req.WebShellConnectionID != "" {
|
||||
conv, err = h.db.CreateConversationWithWebshell(strings.TrimSpace(req.WebShellConnectionID), title)
|
||||
meta.Source = "webshell_chat"
|
||||
conv, err = h.db.CreateConversationWithWebshell(strings.TrimSpace(req.WebShellConnectionID), title, meta)
|
||||
} else {
|
||||
conv, err = h.db.CreateConversation(title)
|
||||
conv, err = h.db.CreateConversation(title, meta)
|
||||
}
|
||||
if err != nil {
|
||||
h.logger.Error("创建对话失败", zap.Error(err))
|
||||
@@ -1482,6 +1603,8 @@ func (h *AgentHandler) AgentLoopStream(c *gin.Context) {
|
||||
builtin.ToolWebshellFileRead,
|
||||
builtin.ToolWebshellFileWrite,
|
||||
builtin.ToolRecordVulnerability,
|
||||
builtin.ToolListVulnerabilities,
|
||||
builtin.ToolGetVulnerability,
|
||||
builtin.ToolListKnowledgeRiskTypes,
|
||||
builtin.ToolSearchKnowledgeBase,
|
||||
}
|
||||
@@ -1612,7 +1735,7 @@ func (h *AgentHandler) AgentLoopStream(c *gin.Context) {
|
||||
go sseKeepalive(c, stopKeepalive, &sseWriteMu)
|
||||
defer close(stopKeepalive)
|
||||
|
||||
result, err := h.agent.AgentLoopWithProgress(taskCtx, finalMessage, agentHistoryMessages, conversationID, progressCallback, roleTools)
|
||||
result, err := h.agent.AgentLoopWithProgress(taskCtx, finalMessage, agentHistoryMessages, conversationID, progressCallback, roleTools, h.projectBlackboardBlock(conversationID))
|
||||
if err != nil {
|
||||
h.logger.Error("Agent Loop执行失败", zap.Error(err))
|
||||
cause := context.Cause(baseCtx)
|
||||
@@ -1924,6 +2047,7 @@ type BatchTaskRequest struct {
|
||||
ScheduleMode string `json:"scheduleMode,omitempty"` // manual | cron
|
||||
CronExpr string `json:"cronExpr,omitempty"` // scheduleMode=cron 时必填
|
||||
ExecuteNow bool `json:"executeNow,omitempty"` // 创建后是否立即执行(默认 false)
|
||||
ProjectID string `json:"projectId,omitempty"` // 队列内子对话绑定的项目(可选)
|
||||
}
|
||||
|
||||
func normalizeBatchQueueAgentMode(mode string) string {
|
||||
@@ -2004,7 +2128,7 @@ func (h *AgentHandler) CreateBatchQueue(c *gin.Context) {
|
||||
nextRunAt = &next
|
||||
}
|
||||
|
||||
queue, createErr := h.batchTaskManager.CreateBatchQueue(req.Title, req.Role, agentMode, scheduleMode, cronExpr, nextRunAt, validTasks)
|
||||
queue, createErr := h.batchTaskManager.CreateBatchQueue(req.Title, req.Role, agentMode, scheduleMode, cronExpr, req.ProjectID, nextRunAt, validTasks)
|
||||
if createErr != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": createErr.Error()})
|
||||
return
|
||||
@@ -2025,6 +2149,11 @@ func (h *AgentHandler) CreateBatchQueue(c *gin.Context) {
|
||||
queue = refreshed
|
||||
}
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "task", "create_queue", "创建批量任务队列", "batch_queue", queue.ID, map[string]interface{}{
|
||||
"task_count": len(validTasks), "started": started,
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"queueId": queue.ID,
|
||||
"queue": queue,
|
||||
@@ -2132,6 +2261,9 @@ func (h *AgentHandler) StartBatchQueue(c *gin.Context) {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "队列不存在"})
|
||||
return
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "task", "start_queue", "启动批量任务队列", "batch_queue", queueID, nil)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"message": "批量任务已开始执行", "queueId": queueID})
|
||||
}
|
||||
|
||||
@@ -2160,6 +2292,9 @@ func (h *AgentHandler) RerunBatchQueue(c *gin.Context) {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "启动失败"})
|
||||
return
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "task", "rerun_queue", "重跑批量任务队列", "batch_queue", queueID, nil)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"message": "批量任务已重新开始执行", "queueId": queueID})
|
||||
}
|
||||
|
||||
@@ -2171,6 +2306,9 @@ func (h *AgentHandler) PauseBatchQueue(c *gin.Context) {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "队列不存在或无法暂停"})
|
||||
return
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "task", "pause_queue", "暂停批量任务队列", "batch_queue", queueID, nil)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"message": "批量任务已暂停"})
|
||||
}
|
||||
|
||||
@@ -2266,6 +2404,16 @@ func (h *AgentHandler) DeleteBatchQueue(c *gin.Context) {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "队列不存在"})
|
||||
return
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.Record(c, audit.Entry{
|
||||
Category: "task",
|
||||
Action: "delete_queue",
|
||||
Result: "success",
|
||||
ResourceType: "batch_queue",
|
||||
ResourceID: queueID,
|
||||
Message: "删除批量任务队列",
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"message": "批量任务队列已删除"})
|
||||
}
|
||||
|
||||
@@ -2351,6 +2499,11 @@ func (h *AgentHandler) DeleteBatchTask(c *gin.Context) {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "队列不存在"})
|
||||
return
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "task", "delete_batch_task", "删除批量子任务", "batch_task", taskID, map[string]interface{}{
|
||||
"batch_queue_id": queueID,
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"message": "任务已删除", "queue": queue})
|
||||
}
|
||||
|
||||
@@ -2509,7 +2662,9 @@ func (h *AgentHandler) executeBatchQueue(queueID string) {
|
||||
|
||||
// 创建新对话
|
||||
title := safeTruncateString(task.Message, 50)
|
||||
conv, err := h.db.CreateConversation(title)
|
||||
batchMeta := audit.ConversationCreateMeta("batch_task")
|
||||
batchMeta.ProjectID = effectiveProjectID(h.config, queue.ProjectID)
|
||||
conv, err := h.db.CreateConversation(title, batchMeta)
|
||||
var conversationID string
|
||||
if err != nil {
|
||||
h.logger.Error("创建对话失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(err))
|
||||
@@ -2659,15 +2814,15 @@ func (h *AgentHandler) executeBatchQueue(queueID string) {
|
||||
var runErr error
|
||||
switch {
|
||||
case useBatchMulti:
|
||||
resultMA, runErr = multiagent.RunDeepAgent(taskCtx, h.config, &h.config.MultiAgent, h.agent, h.logger, conversationID, finalMessage, []agent.ChatMessage{}, roleTools, progressCallback, h.agentsMarkdownDir, batchOrch, nil)
|
||||
resultMA, runErr = multiagent.RunDeepAgent(taskCtx, h.config, &h.config.MultiAgent, h.agent, h.logger, conversationID, finalMessage, []agent.ChatMessage{}, roleTools, progressCallback, h.agentsMarkdownDir, batchOrch, nil, h.projectBlackboardBlock(conversationID))
|
||||
case useEinoSingle:
|
||||
if h.config == nil {
|
||||
runErr = fmt.Errorf("服务器配置未加载")
|
||||
} else {
|
||||
resultMA, runErr = multiagent.RunEinoSingleChatModelAgent(taskCtx, h.config, &h.config.MultiAgent, h.agent, h.logger, conversationID, finalMessage, []agent.ChatMessage{}, roleTools, progressCallback, nil)
|
||||
resultMA, runErr = multiagent.RunEinoSingleChatModelAgent(taskCtx, h.config, &h.config.MultiAgent, h.agent, h.logger, conversationID, finalMessage, []agent.ChatMessage{}, roleTools, progressCallback, nil, h.projectBlackboardBlock(conversationID))
|
||||
}
|
||||
default:
|
||||
result, runErr = h.agent.AgentLoopWithProgress(taskCtx, finalMessage, []agent.ChatMessage{}, conversationID, progressCallback, roleTools)
|
||||
result, runErr = h.agent.AgentLoopWithProgress(taskCtx, finalMessage, []agent.ChatMessage{}, conversationID, progressCallback, roleTools, h.projectBlackboardBlock(conversationID))
|
||||
}
|
||||
|
||||
if runErr != nil {
|
||||
|
||||
@@ -0,0 +1,147 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/database"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// AuditHandler serves platform audit log APIs.
|
||||
type AuditHandler struct {
|
||||
db *database.DB
|
||||
audit *audit.Service
|
||||
logger *zap.Logger
|
||||
}
|
||||
|
||||
// NewAuditHandler creates an audit log handler.
|
||||
func NewAuditHandler(db *database.DB, auditSvc *audit.Service, logger *zap.Logger) *AuditHandler {
|
||||
return &AuditHandler{db: db, audit: auditSvc, logger: logger}
|
||||
}
|
||||
|
||||
// Meta GET /api/audit/meta
|
||||
func (h *AuditHandler) Meta(c *gin.Context) {
|
||||
enabled := false
|
||||
retentionDays := 0
|
||||
if h.audit != nil {
|
||||
enabled = h.audit.Enabled()
|
||||
retentionDays = h.audit.RetentionDays()
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"enabled": enabled,
|
||||
"retention_days": retentionDays,
|
||||
"default_page_size": 20,
|
||||
"max_page_size": 100,
|
||||
"max_export": 5000,
|
||||
})
|
||||
}
|
||||
|
||||
// Summary GET /api/audit/summary
|
||||
func (h *AuditHandler) Summary(c *gin.Context) {
|
||||
if h.db == nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "database unavailable"})
|
||||
return
|
||||
}
|
||||
base := auditFilterFromQuery(c)
|
||||
total, err := h.db.CountAuditLogs(base)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
failFilter := base
|
||||
failFilter.Result = "failure"
|
||||
failures, err := h.db.CountAuditLogs(failFilter)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
since := time.Now().AddDate(0, 0, -7)
|
||||
recentFilter := base
|
||||
recentFilter.Since = &since
|
||||
recent7d, err := h.db.CountAuditLogs(recentFilter)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"total": total,
|
||||
"failures": failures,
|
||||
"recent_7d": recent7d,
|
||||
"has_filters": c.Query("category") != "" || c.Query("action") != "" || c.Query("result") != "" ||
|
||||
c.Query("q") != "" || c.Query("since") != "" || c.Query("until") != "",
|
||||
})
|
||||
}
|
||||
|
||||
// ListLogs GET /api/audit/logs
|
||||
func (h *AuditHandler) ListLogs(c *gin.Context) {
|
||||
if h.db == nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "database unavailable"})
|
||||
return
|
||||
}
|
||||
filter := auditFilterFromQuery(c)
|
||||
page, pageSize := auditPaginationFromQuery(c)
|
||||
filter.Limit = pageSize
|
||||
filter.Offset = (page - 1) * pageSize
|
||||
|
||||
logs, err := h.db.ListAuditLogs(filter)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
total, err := h.db.CountAuditLogs(filter)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"logs": logs,
|
||||
"total": total,
|
||||
"page": page,
|
||||
"page_size": pageSize,
|
||||
})
|
||||
}
|
||||
|
||||
// GetLog GET /api/audit/logs/:id
|
||||
func (h *AuditHandler) GetLog(c *gin.Context) {
|
||||
if h.db == nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "database unavailable"})
|
||||
return
|
||||
}
|
||||
row, err := h.db.GetAuditLogByID(c.Param("id"))
|
||||
if err != nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "审计记录不存在"})
|
||||
return
|
||||
}
|
||||
audit.ApplyResourceAvailability(h.db, row)
|
||||
c.JSON(http.StatusOK, gin.H{"log": row})
|
||||
}
|
||||
|
||||
// ExportLogs GET /api/audit/logs/export — JSON or CSV (?format=csv), max 5000 rows.
|
||||
func (h *AuditHandler) ExportLogs(c *gin.Context) {
|
||||
if h.db == nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "database unavailable"})
|
||||
return
|
||||
}
|
||||
filter := auditFilterFromQuery(c)
|
||||
filter.Limit = 5000
|
||||
filter.Offset = 0
|
||||
|
||||
logs, err := h.db.ListAuditLogs(filter)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if c.Query("format") == "csv" {
|
||||
writeAuditLogsCSV(c, logs)
|
||||
return
|
||||
}
|
||||
c.Header("Content-Disposition", `attachment; filename="audit-logs.json"`)
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"exported_at": time.Now().UTC().Format(time.RFC3339),
|
||||
"logs": logs,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"encoding/csv"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func writeAuditLogsCSV(c *gin.Context, logs []*database.AuditLog) {
|
||||
c.Header("Content-Type", "text/csv; charset=utf-8")
|
||||
c.Header("Content-Disposition", fmt.Sprintf(`attachment; filename="audit-logs-%s.csv"`, time.Now().Format("20060102")))
|
||||
|
||||
w := csv.NewWriter(c.Writer)
|
||||
_ = w.Write([]string{
|
||||
"id", "created_at", "level", "category", "action", "result", "actor",
|
||||
"session_hint", "client_ip", "resource_type", "resource_id", "message",
|
||||
})
|
||||
for _, row := range logs {
|
||||
if row == nil {
|
||||
continue
|
||||
}
|
||||
_ = w.Write([]string{
|
||||
row.ID,
|
||||
row.CreatedAt.UTC().Format(time.RFC3339),
|
||||
row.Level,
|
||||
row.Category,
|
||||
row.Action,
|
||||
row.Result,
|
||||
row.Actor,
|
||||
row.SessionHint,
|
||||
row.ClientIP,
|
||||
row.ResourceType,
|
||||
row.ResourceID,
|
||||
row.Message,
|
||||
})
|
||||
}
|
||||
w.Flush()
|
||||
}
|
||||
@@ -0,0 +1,48 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func auditFilterFromQuery(c *gin.Context) database.ListAuditLogsFilter {
|
||||
filter := database.ListAuditLogsFilter{
|
||||
Level: c.Query("level"),
|
||||
Category: c.Query("category"),
|
||||
Action: c.Query("action"),
|
||||
Result: c.Query("result"),
|
||||
Query: c.Query("q"),
|
||||
ResourceType: c.Query("resource_type"),
|
||||
ResourceID: c.Query("resource_id"),
|
||||
}
|
||||
if since := c.Query("since"); since != "" {
|
||||
if t, err := time.Parse(time.RFC3339, since); err == nil {
|
||||
filter.Since = &t
|
||||
}
|
||||
}
|
||||
if until := c.Query("until"); until != "" {
|
||||
if t, err := time.Parse(time.RFC3339, until); err == nil {
|
||||
filter.Until = &t
|
||||
}
|
||||
}
|
||||
return filter
|
||||
}
|
||||
|
||||
func auditPaginationFromQuery(c *gin.Context) (page, pageSize int) {
|
||||
page = 1
|
||||
pageSize = 20
|
||||
if p, err := strconv.Atoi(c.DefaultQuery("page", "1")); err == nil && p > 0 {
|
||||
page = p
|
||||
}
|
||||
if ps, err := strconv.Atoi(c.DefaultQuery("page_size", "20")); err == nil && ps > 0 {
|
||||
pageSize = ps
|
||||
if pageSize > 100 {
|
||||
pageSize = 100
|
||||
}
|
||||
}
|
||||
return page, pageSize
|
||||
}
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
@@ -18,6 +19,12 @@ type AuthHandler struct {
|
||||
config *config.Config
|
||||
configPath string
|
||||
logger *zap.Logger
|
||||
audit *audit.Service
|
||||
}
|
||||
|
||||
// SetAudit wires platform audit logging.
|
||||
func (h *AuthHandler) SetAudit(s *audit.Service) {
|
||||
h.audit = s
|
||||
}
|
||||
|
||||
// NewAuthHandler creates a new AuthHandler.
|
||||
@@ -49,10 +56,32 @@ func (h *AuthHandler) Login(c *gin.Context) {
|
||||
|
||||
token, expiresAt, err := h.manager.Authenticate(req.Password)
|
||||
if err != nil {
|
||||
if h.audit != nil {
|
||||
h.audit.Record(c, audit.Entry{
|
||||
Level: "warn",
|
||||
Category: "auth",
|
||||
Action: "login",
|
||||
Result: "failure",
|
||||
Message: "登录失败:密码错误",
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"error": "密码错误"})
|
||||
return
|
||||
}
|
||||
|
||||
if h.audit != nil {
|
||||
h.audit.Record(c, audit.Entry{
|
||||
Category: "auth",
|
||||
Action: "login",
|
||||
Result: "success",
|
||||
SessionHint: audit.HintFromToken(token),
|
||||
Message: "登录成功",
|
||||
Detail: map[string]interface{}{
|
||||
"expires_at": expiresAt.UTC().Format(time.RFC3339),
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"token": token,
|
||||
"expires_at": expiresAt.UTC().Format(time.RFC3339),
|
||||
@@ -73,6 +102,14 @@ func (h *AuthHandler) Logout(c *gin.Context) {
|
||||
}
|
||||
|
||||
h.manager.RevokeToken(token)
|
||||
if h.audit != nil {
|
||||
h.audit.Record(c, audit.Entry{
|
||||
Category: "auth",
|
||||
Action: "logout",
|
||||
Result: "success",
|
||||
Message: "退出登录",
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"message": "已退出登录"})
|
||||
}
|
||||
|
||||
@@ -103,6 +140,15 @@ func (h *AuthHandler) ChangePassword(c *gin.Context) {
|
||||
}
|
||||
|
||||
if !h.manager.CheckPassword(oldPassword) {
|
||||
if h.audit != nil {
|
||||
h.audit.Record(c, audit.Entry{
|
||||
Level: "warn",
|
||||
Category: "auth",
|
||||
Action: "change_password",
|
||||
Result: "failure",
|
||||
Message: "修改密码失败:当前密码不正确",
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "当前密码不正确"})
|
||||
return
|
||||
}
|
||||
@@ -132,6 +178,15 @@ func (h *AuthHandler) ChangePassword(c *gin.Context) {
|
||||
h.logger.Info("登录密码已更新,所有会话已失效")
|
||||
}
|
||||
|
||||
if h.audit != nil {
|
||||
h.audit.Record(c, audit.Entry{
|
||||
Category: "auth",
|
||||
Action: "change_password",
|
||||
Result: "success",
|
||||
Message: "登录密码已修改",
|
||||
})
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"message": "密码已更新,请使用新密码重新登录"})
|
||||
}
|
||||
|
||||
|
||||
@@ -65,6 +65,7 @@ type BatchTaskQueue struct {
|
||||
LastScheduleTriggerAt *time.Time `json:"lastScheduleTriggerAt,omitempty"`
|
||||
LastScheduleError string `json:"lastScheduleError,omitempty"`
|
||||
LastRunError string `json:"lastRunError,omitempty"`
|
||||
ProjectID string `json:"projectId,omitempty"`
|
||||
Tasks []*BatchTask `json:"tasks"`
|
||||
Status string `json:"status"` // pending, running, paused, completed, cancelled
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
@@ -103,7 +104,7 @@ func (m *BatchTaskManager) SetDB(db *database.DB) {
|
||||
|
||||
// CreateBatchQueue 创建批量任务队列
|
||||
func (m *BatchTaskManager) CreateBatchQueue(
|
||||
title, role, agentMode, scheduleMode, cronExpr string,
|
||||
title, role, agentMode, scheduleMode, cronExpr, projectID string,
|
||||
nextRunAt *time.Time,
|
||||
tasks []string,
|
||||
) (*BatchTaskQueue, error) {
|
||||
@@ -126,6 +127,7 @@ func (m *BatchTaskManager) CreateBatchQueue(
|
||||
ID: queueID,
|
||||
Title: title,
|
||||
Role: role,
|
||||
ProjectID: strings.TrimSpace(projectID),
|
||||
AgentMode: normalizeBatchQueueAgentMode(agentMode),
|
||||
ScheduleMode: normalizeBatchQueueScheduleMode(scheduleMode),
|
||||
CronExpr: strings.TrimSpace(cronExpr),
|
||||
@@ -171,6 +173,7 @@ func (m *BatchTaskManager) CreateBatchQueue(
|
||||
queue.ScheduleMode,
|
||||
queue.CronExpr,
|
||||
queue.NextRunAt,
|
||||
queue.ProjectID,
|
||||
dbTasks,
|
||||
); err != nil {
|
||||
m.logger.Warn("batch queue DB create failed", zap.String("queueId", queueID), zap.Error(err))
|
||||
@@ -263,6 +266,9 @@ func (m *BatchTaskManager) loadQueueFromDB(queueID string) *BatchTaskQueue {
|
||||
if queueRow.LastRunError.Valid {
|
||||
queue.LastRunError = strings.TrimSpace(queueRow.LastRunError.String)
|
||||
}
|
||||
if queueRow.ProjectID.Valid {
|
||||
queue.ProjectID = strings.TrimSpace(queueRow.ProjectID.String)
|
||||
}
|
||||
if queueRow.StartedAt.Valid {
|
||||
queue.StartedAt = &queueRow.StartedAt.Time
|
||||
}
|
||||
@@ -499,6 +505,9 @@ func (m *BatchTaskManager) LoadFromDB() error {
|
||||
if queueRow.LastRunError.Valid {
|
||||
queue.LastRunError = strings.TrimSpace(queueRow.LastRunError.String)
|
||||
}
|
||||
if queueRow.ProjectID.Valid {
|
||||
queue.ProjectID = strings.TrimSpace(queueRow.ProjectID.String)
|
||||
}
|
||||
if queueRow.StartedAt.Valid {
|
||||
queue.StartedAt = &queueRow.StartedAt.Time
|
||||
}
|
||||
|
||||
@@ -176,6 +176,10 @@ func RegisterBatchTaskMCPTools(mcpServer *mcp.Server, h *AgentHandler, logger *z
|
||||
"type": "boolean",
|
||||
"description": "创建后是否立即开始执行队列,默认 false(pending,需 batch_task_start)",
|
||||
},
|
||||
"project_id": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "队列内子对话绑定的项目 ID(可选,未指定时使用 config.project.default_project_id)",
|
||||
},
|
||||
},
|
||||
},
|
||||
}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
@@ -204,7 +208,8 @@ func RegisterBatchTaskMCPTools(mcpServer *mcp.Server, h *AgentHandler, logger *z
|
||||
if !ok {
|
||||
executeNow = false
|
||||
}
|
||||
queue, createErr := h.batchTaskManager.CreateBatchQueue(title, role, agentMode, scheduleMode, cronExpr, nextRunAt, tasks)
|
||||
projectID := strings.TrimSpace(mcpArgString(args, "project_id"))
|
||||
queue, createErr := h.batchTaskManager.CreateBatchQueue(title, role, agentMode, scheduleMode, cronExpr, projectID, nextRunAt, tasks)
|
||||
if createErr != nil {
|
||||
return batchMCPTextResult("创建队列失败: "+createErr.Error(), true), nil
|
||||
}
|
||||
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/c2"
|
||||
"cyberstrike-ai/internal/database"
|
||||
|
||||
@@ -25,6 +26,12 @@ import (
|
||||
type C2Handler struct {
|
||||
mgrPtr atomic.Pointer[c2.Manager]
|
||||
logger *zap.Logger
|
||||
audit *audit.Service
|
||||
}
|
||||
|
||||
// SetAudit wires platform audit logging.
|
||||
func (h *C2Handler) SetAudit(s *audit.Service) {
|
||||
h.audit = s
|
||||
}
|
||||
|
||||
// NewC2Handler 创建 C2 处理器;manager 可为 nil(功能关闭时)
|
||||
@@ -104,6 +111,11 @@ func (h *C2Handler) CreateListener(c *gin.Context) {
|
||||
implantToken := listener.ImplantToken
|
||||
listener.EncryptionKey = ""
|
||||
listener.ImplantToken = ""
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "c2", "listener_create", "创建 C2 监听器", "c2_listener", listener.ID, map[string]interface{}{
|
||||
"name": listener.Name, "bind": listener.BindHost, "port": listener.BindPort,
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"listener": listener, "implant_token": implantToken})
|
||||
}
|
||||
|
||||
@@ -205,6 +217,9 @@ func (h *C2Handler) DeleteListener(c *gin.Context) {
|
||||
c.JSON(code, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "c2", "listener_delete", "删除 C2 监听器", "c2_listener", id, nil)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"deleted": true})
|
||||
}
|
||||
|
||||
@@ -222,6 +237,9 @@ func (h *C2Handler) StartListener(c *gin.Context) {
|
||||
}
|
||||
listener.EncryptionKey = ""
|
||||
listener.ImplantToken = ""
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "c2", "listener_start", "启动 C2 监听器", "c2_listener", id, nil)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"listener": listener})
|
||||
}
|
||||
|
||||
@@ -236,6 +254,9 @@ func (h *C2Handler) StopListener(c *gin.Context) {
|
||||
c.JSON(code, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "c2", "listener_stop", "停止 C2 监听器", "c2_listener", id, nil)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"stopped": true})
|
||||
}
|
||||
|
||||
@@ -297,6 +318,9 @@ func (h *C2Handler) DeleteSession(c *gin.Context) {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "c2", "session_delete", "删除 C2 会话", "c2_session", id, nil)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"deleted": true})
|
||||
}
|
||||
|
||||
@@ -407,6 +431,11 @@ func (h *C2Handler) DeleteTasks(c *gin.Context) {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "c2", "task_delete", "批量删除 C2 任务", "c2_task", "", map[string]interface{}{
|
||||
"count": n, "ids": req.IDs,
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"deleted": n})
|
||||
}
|
||||
|
||||
@@ -457,6 +486,11 @@ func (h *C2Handler) CreateTask(c *gin.Context) {
|
||||
c.JSON(code, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "c2", "task_create", "创建 C2 任务", "c2_task", task.ID, map[string]interface{}{
|
||||
"session_id": req.SessionID, "task_type": req.TaskType,
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"task": task})
|
||||
}
|
||||
|
||||
@@ -471,6 +505,9 @@ func (h *C2Handler) CancelTask(c *gin.Context) {
|
||||
c.JSON(code, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "c2", "task_cancel", "取消 C2 任务", "c2_task", id, nil)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"cancelled": true})
|
||||
}
|
||||
|
||||
|
||||
@@ -12,6 +12,8 @@ import (
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"cyberstrike-ai/internal/audit"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
@@ -24,6 +26,12 @@ const (
|
||||
// ChatUploadsHandler 对话中上传附件(chat_uploads 目录)的管理 API
|
||||
type ChatUploadsHandler struct {
|
||||
logger *zap.Logger
|
||||
audit *audit.Service
|
||||
}
|
||||
|
||||
// SetAudit wires platform audit logging.
|
||||
func (h *ChatUploadsHandler) SetAudit(s *audit.Service) {
|
||||
h.audit = s
|
||||
}
|
||||
|
||||
// NewChatUploadsHandler 创建处理器
|
||||
@@ -230,6 +238,9 @@ func (h *ChatUploadsHandler) Delete(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "file", "delete", "删除对话附件", "chat_upload", body.Path, nil)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
}
|
||||
|
||||
@@ -503,6 +514,11 @@ func (h *ChatUploadsHandler) Upload(c *gin.Context) {
|
||||
}
|
||||
rel, _ := filepath.Rel(root, fullPath)
|
||||
absSaved, _ := filepath.Abs(fullPath)
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "file", "upload", "上传对话附件", "chat_upload", filepath.ToSlash(rel), map[string]interface{}{
|
||||
"name": unique,
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"ok": true,
|
||||
"relativePath": filepath.ToSlash(rel),
|
||||
|
||||
@@ -14,6 +14,7 @@ import (
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/agents"
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/knowledge"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
@@ -87,6 +88,7 @@ type ConfigHandler struct {
|
||||
knowledgeInitializer KnowledgeInitializer // 知识库初始化器(可选)
|
||||
appUpdater AppUpdater // App更新器(可选)
|
||||
robotRestarter RobotRestarter // 机器人连接重启器(可选),ApplyConfig 时重启钉钉/飞书
|
||||
audit *audit.Service
|
||||
logger *zap.Logger
|
||||
mu sync.RWMutex
|
||||
lastEmbeddingConfig *config.EmbeddingConfig // 上一次的嵌入模型配置(用于检测变更)
|
||||
@@ -206,6 +208,32 @@ func (h *ConfigHandler) SetRobotRestarter(restarter RobotRestarter) {
|
||||
h.robotRestarter = restarter
|
||||
}
|
||||
|
||||
// SetAudit wires platform audit logging.
|
||||
func (h *ConfigHandler) SetAudit(s *audit.Service) {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
h.audit = s
|
||||
}
|
||||
|
||||
// ApplyWechatRobotBinding 微信 iLink 扫码绑定成功后写入配置并重启机器人连接
|
||||
func (h *ConfigHandler) ApplyWechatRobotBinding(wc config.RobotWechatConfig) error {
|
||||
h.mu.Lock()
|
||||
wc.Enabled = true
|
||||
h.config.Robots.Wechat = wc
|
||||
h.mu.Unlock()
|
||||
if err := h.saveConfig(); err != nil {
|
||||
return err
|
||||
}
|
||||
if h.robotRestarter != nil {
|
||||
h.robotRestarter.RestartRobotConnections()
|
||||
}
|
||||
h.logger.Info("微信机器人绑定已保存",
|
||||
zap.String("ilink_bot_id", wc.ILinkBotID),
|
||||
zap.Bool("enabled", wc.Enabled),
|
||||
)
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetConfigResponse 获取配置响应
|
||||
type GetConfigResponse struct {
|
||||
OpenAI config.OpenAIConfig `json:"openai"`
|
||||
@@ -291,7 +319,7 @@ func (h *ConfigHandler) GetConfig(c *gin.Context) {
|
||||
}
|
||||
multiPub := config.MultiAgentPublic{
|
||||
Enabled: h.config.MultiAgent.Enabled,
|
||||
RobotUseMultiAgent: h.config.MultiAgent.RobotUseMultiAgent,
|
||||
RobotDefaultAgentMode: config.NormalizeRobotAgentMode(h.config.MultiAgent),
|
||||
BatchUseMultiAgent: h.config.MultiAgent.BatchUseMultiAgent,
|
||||
SubAgentCount: subAgentCount,
|
||||
Orchestration: config.NormalizeMultiAgentOrchestration(h.config.MultiAgent.Orchestration),
|
||||
@@ -735,6 +763,7 @@ func (h *ConfigHandler) UpdateConfig(c *gin.Context) {
|
||||
if req.Robots != nil {
|
||||
h.config.Robots = *req.Robots
|
||||
h.logger.Info("更新机器人配置",
|
||||
zap.Bool("wechat_enabled", h.config.Robots.Wechat.Enabled),
|
||||
zap.Bool("wecom_enabled", h.config.Robots.Wecom.Enabled),
|
||||
zap.Bool("dingtalk_enabled", h.config.Robots.Dingtalk.Enabled),
|
||||
zap.Bool("lark_enabled", h.config.Robots.Lark.Enabled),
|
||||
@@ -750,8 +779,12 @@ func (h *ConfigHandler) UpdateConfig(c *gin.Context) {
|
||||
// 多代理标量(sub_agents 等仍由 config.yaml 维护)
|
||||
if req.MultiAgent != nil {
|
||||
h.config.MultiAgent.Enabled = req.MultiAgent.Enabled
|
||||
h.config.MultiAgent.RobotUseMultiAgent = req.MultiAgent.RobotUseMultiAgent
|
||||
h.config.MultiAgent.BatchUseMultiAgent = req.MultiAgent.BatchUseMultiAgent
|
||||
if mode := strings.TrimSpace(req.MultiAgent.RobotDefaultAgentMode); mode != "" {
|
||||
h.config.MultiAgent.RobotDefaultAgentMode = mode
|
||||
} else {
|
||||
h.config.MultiAgent.RobotDefaultAgentMode = "react"
|
||||
}
|
||||
if req.MultiAgent.PlanExecuteLoopMaxIterations != nil {
|
||||
h.config.MultiAgent.PlanExecuteLoopMaxIterations = *req.MultiAgent.PlanExecuteLoopMaxIterations
|
||||
}
|
||||
@@ -760,7 +793,7 @@ func (h *ConfigHandler) UpdateConfig(c *gin.Context) {
|
||||
}
|
||||
h.logger.Info("更新多代理配置",
|
||||
zap.Bool("enabled", h.config.MultiAgent.Enabled),
|
||||
zap.Bool("robot_use_multi_agent", h.config.MultiAgent.RobotUseMultiAgent),
|
||||
zap.String("robot_default_agent_mode", config.NormalizeRobotAgentMode(h.config.MultiAgent)),
|
||||
zap.Bool("batch_use_multi_agent", h.config.MultiAgent.BatchUseMultiAgent),
|
||||
zap.Int("plan_execute_loop_max_iterations", h.config.MultiAgent.PlanExecuteLoopMaxIterations),
|
||||
zap.Int("tool_search_always_visible_tools", len(h.config.MultiAgent.EinoMiddleware.ToolSearchAlwaysVisibleTools)),
|
||||
@@ -883,6 +916,9 @@ func (h *ConfigHandler) UpdateConfig(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "config", "update", "更新内存配置", "config", "", nil)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"message": "配置已更新"})
|
||||
}
|
||||
|
||||
@@ -1013,6 +1049,9 @@ func (h *ConfigHandler) ApplyConfig(c *gin.Context) {
|
||||
h.logger.Info("检测到知识库从禁用变为启用,开始动态初始化知识库组件")
|
||||
if _, err := knowledgeInitializer(); err != nil {
|
||||
h.logger.Error("动态初始化知识库失败", zap.Error(err))
|
||||
if h.audit != nil {
|
||||
h.audit.RecordFail(c, "config", "apply", "应用配置失败:初始化知识库", map[string]interface{}{"error": err.Error()})
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "初始化知识库失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
@@ -1047,6 +1086,9 @@ func (h *ConfigHandler) ApplyConfig(c *gin.Context) {
|
||||
h.logger.Info("开始重新初始化知识库组件(嵌入模型配置已变更)")
|
||||
if _, err := reinitKnowledgeInitializer(); err != nil {
|
||||
h.logger.Error("重新初始化知识库失败", zap.Error(err))
|
||||
if h.audit != nil {
|
||||
h.audit.RecordFail(c, "config", "apply", "应用配置失败:重新初始化知识库", map[string]interface{}{"error": err.Error()})
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "重新初始化知识库失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
@@ -1060,6 +1102,9 @@ func (h *ConfigHandler) ApplyConfig(c *gin.Context) {
|
||||
if c2Rt != nil {
|
||||
if err := c2Rt.ReconcileC2AfterConfigApply(); err != nil {
|
||||
h.logger.Error("C2 配置应用失败", zap.Error(err))
|
||||
if h.audit != nil {
|
||||
h.audit.RecordFail(c, "config", "apply", "应用配置失败:C2", map[string]interface{}{"error": err.Error()})
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "C2 启动失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
@@ -1201,6 +1246,20 @@ func (h *ConfigHandler) ApplyConfig(c *gin.Context) {
|
||||
zap.Int("tools_count", len(h.config.Security.Tools)),
|
||||
)
|
||||
|
||||
if h.audit != nil {
|
||||
h.audit.Record(c, audit.Entry{
|
||||
Category: "config",
|
||||
Action: "apply",
|
||||
Result: "success",
|
||||
Message: "配置已应用",
|
||||
Detail: map[string]interface{}{
|
||||
"tools_count": len(h.config.Security.Tools),
|
||||
"knowledge_enabled": h.config.Knowledge.Enabled,
|
||||
"c2_enabled": h.config.C2.EnabledEffective(),
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"message": "配置已应用",
|
||||
"tools_count": len(h.config.Security.Tools),
|
||||
@@ -1481,6 +1540,15 @@ func updateRobotsConfig(doc *yaml.Node, cfg config.RobotsConfig) {
|
||||
setBoolInMap(sessionNode, "strict_user_identity", *cfg.Session.StrictUserIdentity)
|
||||
}
|
||||
|
||||
wechatNode := ensureMap(robotsNode, "wechat")
|
||||
setBoolInMap(wechatNode, "enabled", cfg.Wechat.Enabled)
|
||||
setStringInMap(wechatNode, "bot_token", cfg.Wechat.BotToken)
|
||||
setStringInMap(wechatNode, "ilink_bot_id", cfg.Wechat.ILinkBotID)
|
||||
setStringInMap(wechatNode, "ilink_user_id", cfg.Wechat.ILinkUserID)
|
||||
setStringInMap(wechatNode, "base_url", cfg.Wechat.BaseURL)
|
||||
setStringInMap(wechatNode, "bot_type", cfg.Wechat.BotType)
|
||||
setStringInMap(wechatNode, "bot_agent", cfg.Wechat.BotAgent)
|
||||
|
||||
wecomNode := ensureMap(robotsNode, "wecom")
|
||||
setBoolInMap(wecomNode, "enabled", cfg.Wecom.Enabled)
|
||||
setStringInMap(wecomNode, "token", cfg.Wecom.Token)
|
||||
@@ -1507,7 +1575,7 @@ func updateMultiAgentConfig(doc *yaml.Node, cfg config.MultiAgentConfig) {
|
||||
root := doc.Content[0]
|
||||
maNode := ensureMap(root, "multi_agent")
|
||||
setBoolInMap(maNode, "enabled", cfg.Enabled)
|
||||
setBoolInMap(maNode, "robot_use_multi_agent", cfg.RobotUseMultiAgent)
|
||||
setStringInMap(maNode, "robot_default_agent_mode", config.NormalizeRobotAgentMode(cfg))
|
||||
setBoolInMap(maNode, "batch_use_multi_agent", cfg.BatchUseMultiAgent)
|
||||
setIntInMap(maNode, "plan_execute_loop_max_iterations", cfg.PlanExecuteLoopMaxIterations)
|
||||
mwNode := ensureMap(maNode, "eino_middleware")
|
||||
|
||||
@@ -4,7 +4,9 @@ import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
@@ -14,6 +16,12 @@ import (
|
||||
type ConversationHandler struct {
|
||||
db *database.DB
|
||||
logger *zap.Logger
|
||||
audit *audit.Service
|
||||
}
|
||||
|
||||
// SetAudit wires platform audit logging.
|
||||
func (h *ConversationHandler) SetAudit(s *audit.Service) {
|
||||
h.audit = s
|
||||
}
|
||||
|
||||
// NewConversationHandler 创建新的对话处理器
|
||||
@@ -26,7 +34,13 @@ func NewConversationHandler(db *database.DB, logger *zap.Logger) *ConversationHa
|
||||
|
||||
// CreateConversationRequest 创建对话请求
|
||||
type CreateConversationRequest struct {
|
||||
Title string `json:"title"`
|
||||
Title string `json:"title"`
|
||||
ProjectID string `json:"projectId,omitempty"`
|
||||
}
|
||||
|
||||
// SetConversationProjectRequest 设置对话所属项目
|
||||
type SetConversationProjectRequest struct {
|
||||
ProjectID string `json:"projectId"` // 空字符串表示解除绑定
|
||||
}
|
||||
|
||||
// CreateConversation 创建新对话
|
||||
@@ -42,7 +56,9 @@ func (h *ConversationHandler) CreateConversation(c *gin.Context) {
|
||||
title = "新对话"
|
||||
}
|
||||
|
||||
conv, err := h.db.CreateConversation(title)
|
||||
meta := audit.ConversationCreateMetaFromGin(c, "api")
|
||||
meta.ProjectID = strings.TrimSpace(req.ProjectID)
|
||||
conv, err := h.db.CreateConversation(title, meta)
|
||||
if err != nil {
|
||||
h.logger.Error("创建对话失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
@@ -52,6 +68,25 @@ func (h *ConversationHandler) CreateConversation(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, conv)
|
||||
}
|
||||
|
||||
// SetConversationProject 设置或清除对话绑定的项目
|
||||
func (h *ConversationHandler) SetConversationProject(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
var req SetConversationProjectRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if _, err := h.db.GetConversation(id); err != nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "对话不存在"})
|
||||
return
|
||||
}
|
||||
if err := h.db.SetConversationProjectID(id, req.ProjectID); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"success": true, "projectId": strings.TrimSpace(req.ProjectID)})
|
||||
}
|
||||
|
||||
// ListConversations 列出对话
|
||||
func (h *ConversationHandler) ListConversations(c *gin.Context) {
|
||||
limitStr := c.DefaultQuery("limit", "50")
|
||||
@@ -189,6 +224,17 @@ func (h *ConversationHandler) DeleteConversation(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
if h.audit != nil {
|
||||
h.audit.Record(c, audit.Entry{
|
||||
Category: "conversation",
|
||||
Action: "delete",
|
||||
Result: "success",
|
||||
ResourceType: "conversation",
|
||||
ResourceID: id,
|
||||
Message: "删除对话",
|
||||
})
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"message": "删除成功"})
|
||||
}
|
||||
|
||||
@@ -227,6 +273,12 @@ func (h *ConversationHandler) DeleteConversationTurn(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "conversation", "delete_turn", "删除对话轮次", "conversation", conversationID, map[string]interface{}{
|
||||
"message_id": req.MessageID,
|
||||
"deleted": len(deletedIDs),
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"deletedMessageIds": deletedIDs,
|
||||
"message": "ok",
|
||||
|
||||
@@ -0,0 +1,122 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/multiagent"
|
||||
)
|
||||
|
||||
func (h *AgentHandler) einoRunRetryMaxAttempts() int {
|
||||
if h.config != nil {
|
||||
return multiagent.RunRetryMaxAttemptsFromConfig(&h.config.MultiAgent.EinoMiddleware)
|
||||
}
|
||||
return multiagent.RunRetryMaxAttemptsFromConfig(nil)
|
||||
}
|
||||
|
||||
func (h *AgentHandler) einoRunRetryMaxBackoffSec() int {
|
||||
if h.config != nil && h.config.MultiAgent.EinoMiddleware.RunRetryMaxBackoffSec > 0 {
|
||||
return h.config.MultiAgent.EinoMiddleware.RunRetryMaxBackoffSec
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// applyEinoTraceResumeSegment 中断并继续:persist last_react_* → loadHistory,可选替换下一段 user 文案。
|
||||
func (h *AgentHandler) applyEinoTraceResumeSegment(
|
||||
conversationID string,
|
||||
result *multiagent.RunResult,
|
||||
curHistory *[]agent.ChatMessage,
|
||||
curFinalMessage *string,
|
||||
segmentUserMessage string,
|
||||
) {
|
||||
if shouldPersistEinoAgentTraceAfterRunError(context.Background()) {
|
||||
h.persistEinoAgentTraceForResume(conversationID, result)
|
||||
}
|
||||
if hist, err := h.loadHistoryFromAgentTrace(conversationID); err == nil && len(hist) > 0 {
|
||||
*curHistory = hist
|
||||
}
|
||||
if segmentUserMessage != "" {
|
||||
*curFinalMessage = segmentUserMessage
|
||||
}
|
||||
}
|
||||
|
||||
// applyEinoTransientRetrySegment 临时错误重试:恢复轨迹并保留本请求原始 user 文案(不注入续跑说明)。
|
||||
// segmentUserMessage 为本轮 HTTP 请求开始时用户发送的内容,避免因清空 finalMessage 而丢失「你好」等短句。
|
||||
func (h *AgentHandler) applyEinoTransientRetrySegment(
|
||||
conversationID string,
|
||||
result *multiagent.RunResult,
|
||||
curHistory *[]agent.ChatMessage,
|
||||
curFinalMessage *string,
|
||||
segmentUserMessage string,
|
||||
) {
|
||||
if shouldPersistEinoAgentTraceAfterRunError(context.Background()) {
|
||||
h.persistEinoAgentTraceForResume(conversationID, result)
|
||||
}
|
||||
if hist, err := h.loadHistoryFromAgentTrace(conversationID); err == nil && len(hist) > 0 {
|
||||
*curHistory = hist
|
||||
}
|
||||
if s := strings.TrimSpace(segmentUserMessage); s != "" {
|
||||
*curFinalMessage = segmentUserMessage
|
||||
}
|
||||
}
|
||||
|
||||
// handleEinoTransientRetryContinue 在 SSE 任务循环内处理临时错误重试;返回 true 表示外层 for 应 continue。
|
||||
func (h *AgentHandler) handleEinoTransientRetryContinue(
|
||||
baseCtx context.Context,
|
||||
conversationID string,
|
||||
result *multiagent.RunResult,
|
||||
runErr error,
|
||||
transientAttempts *int,
|
||||
curHistory *[]agent.ChatMessage,
|
||||
curFinalMessage *string,
|
||||
segmentUserMessage string,
|
||||
progressCallback func(eventType, message string, data interface{}),
|
||||
sendProgress func(msg string, extra map[string]interface{}),
|
||||
) (handled bool, fatal error) {
|
||||
if !errors.Is(runErr, multiagent.ErrTransientRetryContinue) {
|
||||
return false, nil
|
||||
}
|
||||
maxAttempts := h.einoRunRetryMaxAttempts()
|
||||
*transientAttempts++
|
||||
if *transientAttempts > maxAttempts {
|
||||
if shouldPersistEinoAgentTraceAfterRunError(baseCtx) {
|
||||
h.persistEinoAgentTraceForResume(conversationID, result)
|
||||
}
|
||||
return false, errors.New("transient retry exhausted: " + runErr.Error())
|
||||
}
|
||||
attemptNo := *transientAttempts
|
||||
backoff := multiagent.TransientRetryBackoff(attemptNo-1, h.einoRunRetryMaxBackoffSec())
|
||||
if progressCallback != nil {
|
||||
progressCallback("eino_run_retry", fmt.Sprintf("遇到临时错误,%d 秒后第 %d/%d 次重试…", int(backoff.Seconds()), attemptNo, maxAttempts), map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
"attempt": attemptNo,
|
||||
"maxAttempts": maxAttempts,
|
||||
"backoffSec": int(backoff.Seconds()),
|
||||
})
|
||||
}
|
||||
select {
|
||||
case <-baseCtx.Done():
|
||||
return false, context.Cause(baseCtx)
|
||||
case <-time.After(backoff):
|
||||
}
|
||||
h.applyEinoTransientRetrySegment(conversationID, result, curHistory, curFinalMessage, segmentUserMessage)
|
||||
if progressCallback != nil {
|
||||
progressCallback("eino_run_retry", "已恢复上下文,正在重试…", map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
"attempt": attemptNo,
|
||||
})
|
||||
}
|
||||
if sendProgress != nil {
|
||||
sendProgress("正在重试…", map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"source": "transient_retry",
|
||||
})
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
@@ -90,7 +90,7 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
||||
zap.String("conversationId", req.ConversationID),
|
||||
)
|
||||
|
||||
prep, err := h.prepareMultiAgentSession(&req)
|
||||
prep, err := h.prepareMultiAgentSession(&req, c, "eino_agent_stream")
|
||||
if err != nil {
|
||||
sendEvent("error", err.Error(), nil)
|
||||
sendEvent("done", "", nil)
|
||||
@@ -119,6 +119,7 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
||||
|
||||
var cancelWithCause context.CancelCauseFunc
|
||||
curFinalMessage := prep.FinalMessage
|
||||
segmentUserMessage := prep.FinalMessage // 本请求原始用户句,临时重试时不得丢失
|
||||
curHistory := prep.History
|
||||
roleTools := prep.RoleTools
|
||||
|
||||
@@ -176,9 +177,41 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
||||
taskOwned = true
|
||||
|
||||
var cumulativeMCPExecutionIDs []string
|
||||
var transientRunAttempts int
|
||||
// 同一请求内分段续跑时,主代理 iteration 事件按偏移累计,避免 UI 出现「第3轮 → 第1轮」回跳。
|
||||
var mainIterationOffset int
|
||||
|
||||
for {
|
||||
progressCallback := h.createProgressCallback(taskCtx, cancelWithCause, conversationID, assistantMessageID, sendEvent)
|
||||
segmentMainIterationMax := 0
|
||||
rawProgressCallback := h.createProgressCallback(taskCtx, cancelWithCause, conversationID, assistantMessageID, sendEvent)
|
||||
progressCallback := func(eventType, message string, data interface{}) {
|
||||
if eventType == "iteration" {
|
||||
if m, ok := data.(map[string]interface{}); ok {
|
||||
if scope, _ := m["einoScope"].(string); scope == "main" {
|
||||
raw := 0
|
||||
switch v := m["iteration"].(type) {
|
||||
case int:
|
||||
raw = v
|
||||
case int32:
|
||||
raw = int(v)
|
||||
case int64:
|
||||
raw = int(v)
|
||||
case float64:
|
||||
raw = int(v)
|
||||
case float32:
|
||||
raw = int(v)
|
||||
}
|
||||
if raw > 0 {
|
||||
if raw > segmentMainIterationMax {
|
||||
segmentMainIterationMax = raw
|
||||
}
|
||||
m["iteration"] = raw + mainIterationOffset
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
rawProgressCallback(eventType, message, data)
|
||||
}
|
||||
taskCtxLoop := mcp.WithMCPConversationID(taskCtx, conversationID)
|
||||
taskCtxLoop = mcp.WithToolRunRegistry(taskCtxLoop, h.tasks)
|
||||
taskCtxLoop = multiagent.WithHITLToolInterceptor(taskCtxLoop, func(ctx context.Context, toolName, arguments string) (string, error) {
|
||||
@@ -197,17 +230,38 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
||||
roleTools,
|
||||
progressCallback,
|
||||
chatReasoningToClientIntent(req.Reasoning),
|
||||
h.projectBlackboardBlock(conversationID),
|
||||
)
|
||||
timeoutCancel()
|
||||
|
||||
if result != nil && len(result.MCPExecutionIDs) > 0 {
|
||||
cumulativeMCPExecutionIDs = mergeMCPExecutionIDLists(cumulativeMCPExecutionIDs, result.MCPExecutionIDs)
|
||||
}
|
||||
|
||||
if runErr == nil {
|
||||
// 任一段成功完成后,重置临时错误重试窗口(次数/退避从头开始)。
|
||||
transientRunAttempts = 0
|
||||
timeoutCancel()
|
||||
break
|
||||
}
|
||||
|
||||
handled, fatalErr := h.handleEinoTransientRetryContinue(
|
||||
baseCtx, conversationID, result, runErr, &transientRunAttempts,
|
||||
&curHistory, &curFinalMessage, segmentUserMessage, progressCallback,
|
||||
func(msg string, extra map[string]interface{}) { sendEvent("progress", msg, extra) },
|
||||
)
|
||||
if handled {
|
||||
mainIterationOffset += segmentMainIterationMax
|
||||
timeoutCancel()
|
||||
baseCtx, cancelWithCause = context.WithCancelCause(context.Background())
|
||||
h.tasks.BindTaskCancel(conversationID, cancelWithCause)
|
||||
taskCtx, timeoutCancel = context.WithTimeout(baseCtx, 600*time.Minute)
|
||||
h.tasks.UpdateTaskStatus(conversationID, "running")
|
||||
continue
|
||||
}
|
||||
if fatalErr != nil {
|
||||
runErr = fatalErr
|
||||
}
|
||||
|
||||
cause := context.Cause(baseCtx)
|
||||
if errors.Is(cause, multiagent.ErrInterruptContinue) {
|
||||
if shouldPersistEinoAgentTraceAfterRunError(baseCtx) {
|
||||
@@ -231,10 +285,14 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
||||
"conversationId": conversationID,
|
||||
"source": "interrupt_continue",
|
||||
})
|
||||
h.tasks.UpdateTaskStatus(conversationID, "running")
|
||||
mainIterationOffset += segmentMainIterationMax
|
||||
// 非临时错误分段续跑(用户中断并继续)时,清空 transient 计数,避免跨分段累加。
|
||||
transientRunAttempts = 0
|
||||
timeoutCancel()
|
||||
baseCtx, cancelWithCause = context.WithCancelCause(context.Background())
|
||||
h.tasks.BindTaskCancel(conversationID, cancelWithCause)
|
||||
taskCtx, timeoutCancel = context.WithTimeout(baseCtx, 600*time.Minute)
|
||||
h.tasks.UpdateTaskStatus(conversationID, "running")
|
||||
continue
|
||||
}
|
||||
|
||||
@@ -261,6 +319,7 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
||||
"messageId": assistantMessageID,
|
||||
})
|
||||
sendEvent("done", "", map[string]interface{}{"conversationId": conversationID})
|
||||
timeoutCancel()
|
||||
return
|
||||
}
|
||||
|
||||
@@ -278,6 +337,7 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
||||
"errorType": "timeout",
|
||||
})
|
||||
sendEvent("done", "", map[string]interface{}{"conversationId": conversationID})
|
||||
timeoutCancel()
|
||||
return
|
||||
}
|
||||
|
||||
@@ -294,9 +354,12 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
||||
"messageId": assistantMessageID,
|
||||
})
|
||||
sendEvent("done", "", map[string]interface{}{"conversationId": conversationID})
|
||||
timeoutCancel()
|
||||
return
|
||||
}
|
||||
|
||||
timeoutCancel()
|
||||
|
||||
if assistantMessageID != "" {
|
||||
_ = h.db.UpdateAssistantMessageFinalize(assistantMessageID, result.Response, cumulativeMCPExecutionIDs, multiagent.AggregatedReasoningFromTraceJSON(result.LastAgentTraceInput))
|
||||
}
|
||||
@@ -326,7 +389,7 @@ func (h *AgentHandler) EinoSingleAgentLoop(c *gin.Context) {
|
||||
|
||||
h.logger.Info("收到 Eino ADK 单代理非流式请求", zap.String("conversationId", req.ConversationID))
|
||||
|
||||
prep, err := h.prepareMultiAgentSession(&req)
|
||||
prep, err := h.prepareMultiAgentSession(&req, c, "eino_agent")
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
@@ -367,6 +430,7 @@ func (h *AgentHandler) EinoSingleAgentLoop(c *gin.Context) {
|
||||
prep.RoleTools,
|
||||
progressCallback,
|
||||
chatReasoningToClientIntent(req.Reasoning),
|
||||
h.projectBlackboardBlock(prep.ConversationID),
|
||||
)
|
||||
if runErr != nil {
|
||||
if shouldPersistEinoAgentTraceAfterRunError(baseCtx) {
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"os"
|
||||
"sync"
|
||||
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
|
||||
@@ -20,9 +21,15 @@ type ExternalMCPHandler struct {
|
||||
config *config.Config
|
||||
configPath string
|
||||
logger *zap.Logger
|
||||
audit *audit.Service
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
// SetAudit wires platform audit logging.
|
||||
func (h *ExternalMCPHandler) SetAudit(s *audit.Service) {
|
||||
h.audit = s
|
||||
}
|
||||
|
||||
// NewExternalMCPHandler 创建外部MCP处理器
|
||||
func NewExternalMCPHandler(manager *mcp.ExternalMCPManager, cfg *config.Config, configPath string, logger *zap.Logger) *ExternalMCPHandler {
|
||||
return &ExternalMCPHandler{
|
||||
@@ -180,6 +187,16 @@ func (h *ExternalMCPHandler) AddOrUpdateExternalMCP(c *gin.Context) {
|
||||
}
|
||||
|
||||
h.logger.Info("外部MCP配置已更新", zap.String("name", name))
|
||||
if h.audit != nil {
|
||||
h.audit.Record(c, audit.Entry{
|
||||
Category: "external_mcp",
|
||||
Action: "upsert",
|
||||
Result: "success",
|
||||
ResourceType: "external_mcp",
|
||||
ResourceID: name,
|
||||
Message: "更新外部 MCP 配置",
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"message": "配置已更新"})
|
||||
}
|
||||
|
||||
@@ -209,6 +226,16 @@ func (h *ExternalMCPHandler) DeleteExternalMCP(c *gin.Context) {
|
||||
}
|
||||
|
||||
h.logger.Info("外部MCP配置已删除", zap.String("name", name))
|
||||
if h.audit != nil {
|
||||
h.audit.Record(c, audit.Entry{
|
||||
Category: "external_mcp",
|
||||
Action: "delete",
|
||||
Result: "success",
|
||||
ResourceType: "external_mcp",
|
||||
ResourceID: name,
|
||||
Message: "删除外部 MCP 配置",
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"message": "配置已删除"})
|
||||
}
|
||||
|
||||
|
||||
@@ -616,6 +616,11 @@ func (h *AgentHandler) DecideHITLInterrupt(c *gin.Context) {
|
||||
c.JSON(http.StatusConflict, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "hitl", "decision", "HITL 审批决策", "hitl_interrupt", req.InterruptID, map[string]interface{}{
|
||||
"decision": req.Decision,
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
}
|
||||
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/knowledge"
|
||||
|
||||
@@ -20,6 +21,12 @@ type KnowledgeHandler struct {
|
||||
indexer *knowledge.Indexer
|
||||
db *database.DB
|
||||
logger *zap.Logger
|
||||
audit *audit.Service
|
||||
}
|
||||
|
||||
// SetAudit wires platform audit logging.
|
||||
func (h *KnowledgeHandler) SetAudit(s *audit.Service) {
|
||||
h.audit = s
|
||||
}
|
||||
|
||||
// NewKnowledgeHandler 创建新的知识库处理器
|
||||
@@ -303,6 +310,9 @@ func (h *KnowledgeHandler) DeleteItem(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "knowledge", "item_delete", "删除知识项", "knowledge_item", id, nil)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"message": "删除成功"})
|
||||
}
|
||||
|
||||
@@ -316,6 +326,9 @@ func (h *KnowledgeHandler) RebuildIndex(c *gin.Context) {
|
||||
}
|
||||
}()
|
||||
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "knowledge", "index_rebuild", "重建知识库索引", "knowledge", "", nil)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"message": "索引重建已开始,将在后台进行"})
|
||||
}
|
||||
|
||||
|
||||
@@ -9,6 +9,7 @@ import (
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/agents"
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/config"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
@@ -18,7 +19,8 @@ var markdownAgentFilenameRe = regexp.MustCompile(`^[a-zA-Z0-9][a-zA-Z0-9_.-]*\.m
|
||||
|
||||
// MarkdownAgentsHandler 管理 agents 目录下子代理 Markdown(增删改查)。
|
||||
type MarkdownAgentsHandler struct {
|
||||
dir string
|
||||
dir string
|
||||
audit *audit.Service
|
||||
}
|
||||
|
||||
// NewMarkdownAgentsHandler dir 须为已解析的绝对路径。
|
||||
@@ -26,6 +28,11 @@ func NewMarkdownAgentsHandler(dir string) *MarkdownAgentsHandler {
|
||||
return &MarkdownAgentsHandler{dir: strings.TrimSpace(dir)}
|
||||
}
|
||||
|
||||
// SetAudit wires platform audit logging.
|
||||
func (h *MarkdownAgentsHandler) SetAudit(s *audit.Service) {
|
||||
h.audit = s
|
||||
}
|
||||
|
||||
func (h *MarkdownAgentsHandler) safeJoin(filename string) (string, error) {
|
||||
filename = strings.TrimSpace(filename)
|
||||
if filename == "" || !markdownAgentFilenameRe.MatchString(filename) {
|
||||
@@ -227,6 +234,9 @@ func (h *MarkdownAgentsHandler) CreateMarkdownAgent(c *gin.Context) {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "agent", "markdown_create", "创建 Markdown 子代理", "markdown_agent", filepath.Base(path), nil)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"filename": filepath.Base(path), "message": "已创建"})
|
||||
}
|
||||
|
||||
@@ -294,6 +304,9 @@ func (h *MarkdownAgentsHandler) UpdateMarkdownAgent(c *gin.Context) {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "agent", "markdown_update", "更新 Markdown 子代理", "markdown_agent", filename, nil)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"message": "已保存"})
|
||||
}
|
||||
|
||||
@@ -313,5 +326,8 @@ func (h *MarkdownAgentsHandler) DeleteMarkdownAgent(c *gin.Context) {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "agent", "markdown_delete", "删除 Markdown 子代理", "markdown_agent", filename, nil)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"message": "已删除"})
|
||||
}
|
||||
|
||||
@@ -9,6 +9,7 @@ import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/security"
|
||||
@@ -23,6 +24,12 @@ type MonitorHandler struct {
|
||||
executor *security.Executor
|
||||
db *database.DB
|
||||
logger *zap.Logger
|
||||
audit *audit.Service
|
||||
}
|
||||
|
||||
// SetAudit wires platform audit logging.
|
||||
func (h *MonitorHandler) SetAudit(s *audit.Service) {
|
||||
h.audit = s
|
||||
}
|
||||
|
||||
// NewMonitorHandler 创建新的监控处理器
|
||||
@@ -365,6 +372,11 @@ func (h *MonitorHandler) DeleteExecution(c *gin.Context) {
|
||||
}
|
||||
|
||||
h.logger.Info("执行记录已从数据库删除", zap.String("executionId", id), zap.String("toolName", exec.ToolName))
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "tool", "execution_delete", "删除工具执行记录", "tool_execution", id, map[string]interface{}{
|
||||
"tool_name": exec.ToolName,
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"message": "执行记录已删除"})
|
||||
return
|
||||
}
|
||||
@@ -440,6 +452,11 @@ func (h *MonitorHandler) DeleteExecutions(c *gin.Context) {
|
||||
}
|
||||
|
||||
h.logger.Info("批量删除执行记录成功", zap.Int("count", len(request.IDs)))
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "tool", "execution_delete_batch", "批量删除工具执行记录", "tool_execution", "", map[string]interface{}{
|
||||
"count": len(request.IDs),
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"message": "成功删除执行记录", "deleted": len(executions)})
|
||||
return
|
||||
}
|
||||
|
||||
@@ -107,7 +107,7 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
||||
zap.String("conversationId", req.ConversationID),
|
||||
)
|
||||
|
||||
prep, err := h.prepareMultiAgentSession(&req)
|
||||
prep, err := h.prepareMultiAgentSession(&req, c, "multi_agent_stream")
|
||||
if err != nil {
|
||||
sendEvent("error", err.Error(), nil)
|
||||
sendEvent("done", "", nil)
|
||||
@@ -136,6 +136,7 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
||||
|
||||
var cancelWithCause context.CancelCauseFunc
|
||||
curFinalMessage := prep.FinalMessage
|
||||
segmentUserMessage := prep.FinalMessage // 本请求原始用户句,临时重试时不得丢失
|
||||
curHistory := prep.History
|
||||
roleTools := prep.RoleTools
|
||||
orch := strings.TrimSpace(req.Orchestration)
|
||||
@@ -186,9 +187,41 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
||||
|
||||
// 同一 HTTP 流内多段 Run(如中断并继续)合并 MCP execution id,供最终 response / 库表与工具芯片展示完整列表
|
||||
var cumulativeMCPExecutionIDs []string
|
||||
var transientRunAttempts int
|
||||
// 同一请求内分段续跑时,主代理 iteration 事件按偏移累计,避免 UI 出现「第3轮 → 第1轮」回跳。
|
||||
var mainIterationOffset int
|
||||
|
||||
for {
|
||||
progressCallback := h.createProgressCallback(taskCtx, cancelWithCause, conversationID, assistantMessageID, sendEvent)
|
||||
segmentMainIterationMax := 0
|
||||
rawProgressCallback := h.createProgressCallback(taskCtx, cancelWithCause, conversationID, assistantMessageID, sendEvent)
|
||||
progressCallback := func(eventType, message string, data interface{}) {
|
||||
if eventType == "iteration" {
|
||||
if m, ok := data.(map[string]interface{}); ok {
|
||||
if scope, _ := m["einoScope"].(string); scope == "main" {
|
||||
raw := 0
|
||||
switch v := m["iteration"].(type) {
|
||||
case int:
|
||||
raw = v
|
||||
case int32:
|
||||
raw = int(v)
|
||||
case int64:
|
||||
raw = int(v)
|
||||
case float64:
|
||||
raw = int(v)
|
||||
case float32:
|
||||
raw = int(v)
|
||||
}
|
||||
if raw > 0 {
|
||||
if raw > segmentMainIterationMax {
|
||||
segmentMainIterationMax = raw
|
||||
}
|
||||
m["iteration"] = raw + mainIterationOffset
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
rawProgressCallback(eventType, message, data)
|
||||
}
|
||||
taskCtxLoop := mcp.WithMCPConversationID(taskCtx, conversationID)
|
||||
taskCtxLoop = mcp.WithToolRunRegistry(taskCtxLoop, h.tasks)
|
||||
taskCtxLoop = multiagent.WithHITLToolInterceptor(taskCtxLoop, func(ctx context.Context, toolName, arguments string) (string, error) {
|
||||
@@ -209,17 +242,38 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
||||
h.agentsMarkdownDir,
|
||||
orch,
|
||||
chatReasoningToClientIntent(req.Reasoning),
|
||||
h.projectBlackboardBlock(conversationID),
|
||||
)
|
||||
timeoutCancel()
|
||||
|
||||
if result != nil && len(result.MCPExecutionIDs) > 0 {
|
||||
cumulativeMCPExecutionIDs = mergeMCPExecutionIDLists(cumulativeMCPExecutionIDs, result.MCPExecutionIDs)
|
||||
}
|
||||
|
||||
if runErr == nil {
|
||||
// 任一段成功完成后,重置临时错误重试窗口(次数/退避从头开始)。
|
||||
transientRunAttempts = 0
|
||||
timeoutCancel()
|
||||
break
|
||||
}
|
||||
|
||||
handled, fatalErr := h.handleEinoTransientRetryContinue(
|
||||
baseCtx, conversationID, result, runErr, &transientRunAttempts,
|
||||
&curHistory, &curFinalMessage, segmentUserMessage, progressCallback,
|
||||
func(msg string, extra map[string]interface{}) { sendEvent("progress", msg, extra) },
|
||||
)
|
||||
if handled {
|
||||
mainIterationOffset += segmentMainIterationMax
|
||||
timeoutCancel()
|
||||
baseCtx, cancelWithCause = context.WithCancelCause(context.Background())
|
||||
h.tasks.BindTaskCancel(conversationID, cancelWithCause)
|
||||
taskCtx, timeoutCancel = context.WithTimeout(baseCtx, 600*time.Minute)
|
||||
h.tasks.UpdateTaskStatus(conversationID, "running")
|
||||
continue
|
||||
}
|
||||
if fatalErr != nil {
|
||||
runErr = fatalErr
|
||||
}
|
||||
|
||||
cause := context.Cause(baseCtx)
|
||||
if errors.Is(cause, multiagent.ErrInterruptContinue) {
|
||||
if shouldPersistEinoAgentTraceAfterRunError(baseCtx) {
|
||||
@@ -243,10 +297,14 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
||||
"conversationId": conversationID,
|
||||
"source": "interrupt_continue",
|
||||
})
|
||||
h.tasks.UpdateTaskStatus(conversationID, "running")
|
||||
mainIterationOffset += segmentMainIterationMax
|
||||
// 非临时错误分段续跑(用户中断并继续)时,清空 transient 计数,避免跨分段累加。
|
||||
transientRunAttempts = 0
|
||||
timeoutCancel()
|
||||
baseCtx, cancelWithCause = context.WithCancelCause(context.Background())
|
||||
h.tasks.BindTaskCancel(conversationID, cancelWithCause)
|
||||
taskCtx, timeoutCancel = context.WithTimeout(baseCtx, 600*time.Minute)
|
||||
h.tasks.UpdateTaskStatus(conversationID, "running")
|
||||
continue
|
||||
}
|
||||
|
||||
@@ -273,6 +331,7 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
||||
"messageId": assistantMessageID,
|
||||
})
|
||||
sendEvent("done", "", map[string]interface{}{"conversationId": conversationID})
|
||||
timeoutCancel()
|
||||
return
|
||||
}
|
||||
|
||||
@@ -290,6 +349,7 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
||||
"errorType": "timeout",
|
||||
})
|
||||
sendEvent("done", "", map[string]interface{}{"conversationId": conversationID})
|
||||
timeoutCancel()
|
||||
return
|
||||
}
|
||||
|
||||
@@ -306,9 +366,12 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
||||
"messageId": assistantMessageID,
|
||||
})
|
||||
sendEvent("done", "", map[string]interface{}{"conversationId": conversationID})
|
||||
timeoutCancel()
|
||||
return
|
||||
}
|
||||
|
||||
timeoutCancel()
|
||||
|
||||
if assistantMessageID != "" {
|
||||
_ = h.db.UpdateAssistantMessageFinalize(assistantMessageID, result.Response, cumulativeMCPExecutionIDs, multiagent.AggregatedReasoningFromTraceJSON(result.LastAgentTraceInput))
|
||||
}
|
||||
@@ -347,7 +410,7 @@ func (h *AgentHandler) MultiAgentLoop(c *gin.Context) {
|
||||
|
||||
h.logger.Info("收到 Eino DeepAgent 非流式请求", zap.String("conversationId", req.ConversationID))
|
||||
|
||||
prep, err := h.prepareMultiAgentSession(&req)
|
||||
prep, err := h.prepareMultiAgentSession(&req, c, "multi_agent")
|
||||
if err != nil {
|
||||
status, msg := multiAgentHTTPErrorStatus(err)
|
||||
c.JSON(status, gin.H{"error": msg})
|
||||
@@ -381,6 +444,7 @@ func (h *AgentHandler) MultiAgentLoop(c *gin.Context) {
|
||||
h.agentsMarkdownDir,
|
||||
strings.TrimSpace(req.Orchestration),
|
||||
chatReasoningToClientIntent(req.Reasoning),
|
||||
h.projectBlackboardBlock(prep.ConversationID),
|
||||
)
|
||||
if runErr != nil {
|
||||
if shouldPersistEinoAgentTraceAfterRunError(baseCtx) {
|
||||
|
||||
@@ -5,9 +5,11 @@ import (
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
@@ -22,7 +24,7 @@ type multiAgentPrepared struct {
|
||||
UserMessageID string
|
||||
}
|
||||
|
||||
func (h *AgentHandler) prepareMultiAgentSession(req *ChatRequest) (*multiAgentPrepared, error) {
|
||||
func (h *AgentHandler) prepareMultiAgentSession(req *ChatRequest, c *gin.Context, source string) (*multiAgentPrepared, error) {
|
||||
if len(req.Attachments) > maxAttachments {
|
||||
return nil, fmt.Errorf("附件最多 %d 个", maxAttachments)
|
||||
}
|
||||
@@ -33,10 +35,14 @@ func (h *AgentHandler) prepareMultiAgentSession(req *ChatRequest) (*multiAgentPr
|
||||
title := safeTruncateString(req.Message, 50)
|
||||
var conv *database.Conversation
|
||||
var err error
|
||||
meta := audit.ConversationCreateMetaFromGin(c, source)
|
||||
meta.ProjectID = effectiveProjectID(h.config, req.ProjectID)
|
||||
if strings.TrimSpace(req.WebShellConnectionID) != "" {
|
||||
conv, err = h.db.CreateConversationWithWebshell(strings.TrimSpace(req.WebShellConnectionID), title)
|
||||
meta.Source = source + "_webshell"
|
||||
meta.WebShellConnectionID = strings.TrimSpace(req.WebShellConnectionID)
|
||||
conv, err = h.db.CreateConversationWithWebshell(meta.WebShellConnectionID, title, meta)
|
||||
} else {
|
||||
conv, err = h.db.CreateConversation(title)
|
||||
conv, err = h.db.CreateConversation(title, meta)
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("创建对话失败: %w", err)
|
||||
@@ -85,6 +91,14 @@ func (h *AgentHandler) prepareMultiAgentSession(req *ChatRequest) (*multiAgentPr
|
||||
builtin.ToolWebshellFileRead,
|
||||
builtin.ToolWebshellFileWrite,
|
||||
builtin.ToolRecordVulnerability,
|
||||
builtin.ToolListVulnerabilities,
|
||||
builtin.ToolGetVulnerability,
|
||||
builtin.ToolUpsertProjectFact,
|
||||
builtin.ToolGetProjectFact,
|
||||
builtin.ToolListProjectFacts,
|
||||
builtin.ToolSearchProjectFacts,
|
||||
builtin.ToolDeprecateProjectFact,
|
||||
builtin.ToolRestoreProjectFact,
|
||||
builtin.ToolListKnowledgeRiskTypes,
|
||||
builtin.ToolSearchKnowledgeBase,
|
||||
}
|
||||
|
||||
+139
-1
@@ -73,8 +73,22 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"description": "对话标题",
|
||||
"example": "Web应用安全测试",
|
||||
},
|
||||
"projectId": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "绑定的项目 ID(可选,共享事实黑板)",
|
||||
},
|
||||
},
|
||||
},
|
||||
"SetConversationProjectRequest": map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"projectId": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "项目 ID;空字符串表示解除绑定",
|
||||
},
|
||||
},
|
||||
"required": []string{"projectId"},
|
||||
},
|
||||
"Conversation": map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
@@ -98,6 +112,10 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"format": "date-time",
|
||||
"description": "更新时间",
|
||||
},
|
||||
"projectId": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "绑定的项目 ID(可选)",
|
||||
},
|
||||
},
|
||||
},
|
||||
"ConversationDetail": map[string]interface{}{
|
||||
@@ -1326,6 +1344,37 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
},
|
||||
},
|
||||
},
|
||||
"/api/conversations/{id}/project": map[string]interface{}{
|
||||
"put": map[string]interface{}{
|
||||
"tags": []string{"对话管理"},
|
||||
"summary": "设置对话所属项目",
|
||||
"description": "绑定或解除对话与项目的关联,用于共享事实黑板",
|
||||
"operationId": "setConversationProject",
|
||||
"parameters": []map[string]interface{}{
|
||||
{
|
||||
"name": "id", "in": "path", "required": true,
|
||||
"description": "对话ID",
|
||||
"schema": map[string]interface{}{"type": "string"},
|
||||
},
|
||||
},
|
||||
"requestBody": map[string]interface{}{
|
||||
"required": true,
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"$ref": "#/components/schemas/SetConversationProjectRequest",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
"responses": map[string]interface{}{
|
||||
"200": map[string]interface{}{"description": "设置成功"},
|
||||
"400": map[string]interface{}{"description": "项目不存在或参数错误"},
|
||||
"404": map[string]interface{}{"description": "对话不存在"},
|
||||
"401": map[string]interface{}{"description": "未授权"},
|
||||
},
|
||||
},
|
||||
},
|
||||
"/api/conversations/{id}/results": map[string]interface{}{
|
||||
"get": map[string]interface{}{
|
||||
"tags": []string{"对话管理"},
|
||||
@@ -2444,6 +2493,86 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
},
|
||||
},
|
||||
},
|
||||
"/api/projects": map[string]interface{}{
|
||||
"get": map[string]interface{}{
|
||||
"tags": []string{"项目管理"},
|
||||
"summary": "列出项目",
|
||||
"operationId": "listProjects",
|
||||
"parameters": []map[string]interface{}{
|
||||
{"name": "status", "in": "query", "schema": map[string]interface{}{"type": "string", "enum": []string{"active", "archived"}}},
|
||||
{"name": "limit", "in": "query", "schema": map[string]interface{}{"type": "integer", "default": 200}},
|
||||
},
|
||||
"responses": map[string]interface{}{
|
||||
"200": map[string]interface{}{"description": "项目列表"},
|
||||
"401": map[string]interface{}{"description": "未授权"},
|
||||
},
|
||||
},
|
||||
"post": map[string]interface{}{
|
||||
"tags": []string{"项目管理"},
|
||||
"summary": "创建项目",
|
||||
"operationId": "createProject",
|
||||
"requestBody": map[string]interface{}{
|
||||
"required": true,
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"name": map[string]interface{}{"type": "string"},
|
||||
"description": map[string]interface{}{"type": "string"},
|
||||
"scope_json": map[string]interface{}{"type": "string"},
|
||||
},
|
||||
"required": []string{"name"},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
"responses": map[string]interface{}{
|
||||
"200": map[string]interface{}{"description": "创建成功"},
|
||||
"401": map[string]interface{}{"description": "未授权"},
|
||||
},
|
||||
},
|
||||
},
|
||||
"/api/projects/{id}": map[string]interface{}{
|
||||
"get": map[string]interface{}{
|
||||
"tags": []string{"项目管理"}, "summary": "获取项目", "operationId": "getProject",
|
||||
"parameters": []map[string]interface{}{
|
||||
{"name": "id", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}},
|
||||
},
|
||||
"responses": map[string]interface{}{"200": map[string]interface{}{"description": "项目详情"}},
|
||||
},
|
||||
"put": map[string]interface{}{
|
||||
"tags": []string{"项目管理"}, "summary": "更新项目", "operationId": "updateProject",
|
||||
"parameters": []map[string]interface{}{
|
||||
{"name": "id", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}},
|
||||
},
|
||||
"responses": map[string]interface{}{"200": map[string]interface{}{"description": "更新成功"}},
|
||||
},
|
||||
"delete": map[string]interface{}{
|
||||
"tags": []string{"项目管理"}, "summary": "删除项目", "operationId": "deleteProject",
|
||||
"parameters": []map[string]interface{}{
|
||||
{"name": "id", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}},
|
||||
},
|
||||
"responses": map[string]interface{}{"200": map[string]interface{}{"description": "删除成功"}},
|
||||
},
|
||||
},
|
||||
"/api/projects/{id}/facts": map[string]interface{}{
|
||||
"get": map[string]interface{}{
|
||||
"tags": []string{"项目管理"}, "summary": "列出或按 key 获取事实", "operationId": "listProjectFacts",
|
||||
"parameters": []map[string]interface{}{
|
||||
{"name": "id", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}},
|
||||
{"name": "fact_key", "in": "query", "schema": map[string]interface{}{"type": "string"}},
|
||||
},
|
||||
"responses": map[string]interface{}{"200": map[string]interface{}{"description": "事实列表或单条"}},
|
||||
},
|
||||
"post": map[string]interface{}{
|
||||
"tags": []string{"项目管理"}, "summary": "创建/更新事实", "operationId": "upsertProjectFactREST",
|
||||
"parameters": []map[string]interface{}{
|
||||
{"name": "id", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}},
|
||||
},
|
||||
"responses": map[string]interface{}{"200": map[string]interface{}{"description": "成功"}},
|
||||
},
|
||||
},
|
||||
"/api/vulnerabilities": map[string]interface{}{
|
||||
"get": map[string]interface{}{
|
||||
"tags": []string{"漏洞管理"},
|
||||
@@ -2502,6 +2631,15 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"type": "string",
|
||||
},
|
||||
},
|
||||
{
|
||||
"name": "project_id",
|
||||
"in": "query",
|
||||
"required": false,
|
||||
"description": "项目ID",
|
||||
"schema": map[string]interface{}{
|
||||
"type": "string",
|
||||
},
|
||||
},
|
||||
{
|
||||
"name": "severity",
|
||||
"in": "query",
|
||||
@@ -6254,7 +6392,7 @@ func (h *OpenAPIHandler) GetConversationResults(c *gin.Context) {
|
||||
}
|
||||
|
||||
// 获取漏洞列表
|
||||
vulnList, err := h.db.ListVulnerabilities(1000, 0, "", conversationID, "", "", "", "", "")
|
||||
vulnList, err := h.db.ListVulnerabilities(1000, 0, database.VulnerabilityListFilter{ConversationID: conversationID})
|
||||
if err != nil {
|
||||
h.logger.Warn("获取漏洞列表失败", zap.Error(err))
|
||||
vulnList = []*database.Vulnerability{}
|
||||
|
||||
@@ -0,0 +1,283 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// ProjectHandler 项目管理处理器。
|
||||
type ProjectHandler struct {
|
||||
db *database.DB
|
||||
logger *zap.Logger
|
||||
}
|
||||
|
||||
// NewProjectHandler 创建项目管理处理器。
|
||||
func NewProjectHandler(db *database.DB, logger *zap.Logger) *ProjectHandler {
|
||||
return &ProjectHandler{db: db, logger: logger}
|
||||
}
|
||||
|
||||
type createProjectRequest struct {
|
||||
Name string `json:"name" binding:"required"`
|
||||
Description string `json:"description"`
|
||||
ScopeJSON string `json:"scope_json"`
|
||||
Status string `json:"status"`
|
||||
}
|
||||
|
||||
type updateProjectRequest struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
ScopeJSON string `json:"scope_json"`
|
||||
Status string `json:"status"`
|
||||
Pinned *bool `json:"pinned"`
|
||||
}
|
||||
|
||||
// CreateProject POST /api/projects
|
||||
func (h *ProjectHandler) CreateProject(c *gin.Context) {
|
||||
var req createProjectRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
p := &database.Project{
|
||||
Name: strings.TrimSpace(req.Name),
|
||||
Description: req.Description,
|
||||
ScopeJSON: req.ScopeJSON,
|
||||
Status: strings.TrimSpace(req.Status),
|
||||
}
|
||||
created, err := h.db.CreateProject(p)
|
||||
if err != nil {
|
||||
h.logger.Error("创建项目失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, created)
|
||||
}
|
||||
|
||||
// ListProjects GET /api/projects
|
||||
func (h *ProjectHandler) ListProjects(c *gin.Context) {
|
||||
status := c.Query("status")
|
||||
limit, _ := strconv.Atoi(c.DefaultQuery("limit", "200"))
|
||||
offset, _ := strconv.Atoi(c.Query("offset"))
|
||||
list, err := h.db.ListProjects(status, limit, offset)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if list == nil {
|
||||
list = []*database.Project{}
|
||||
}
|
||||
c.JSON(http.StatusOK, list)
|
||||
}
|
||||
|
||||
// GetProject GET /api/projects/:id
|
||||
func (h *ProjectHandler) GetProject(c *gin.Context) {
|
||||
p, err := h.db.GetProject(c.Param("id"))
|
||||
if err != nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "项目不存在"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, p)
|
||||
}
|
||||
|
||||
// UpdateProject PUT /api/projects/:id
|
||||
func (h *ProjectHandler) UpdateProject(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
p, err := h.db.GetProject(id)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "项目不存在"})
|
||||
return
|
||||
}
|
||||
var req updateProjectRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if s := strings.TrimSpace(req.Name); s != "" {
|
||||
p.Name = s
|
||||
}
|
||||
if req.Description != "" || c.Request.ContentLength > 0 {
|
||||
p.Description = req.Description
|
||||
}
|
||||
if req.ScopeJSON != "" || c.GetHeader("Content-Type") != "" {
|
||||
p.ScopeJSON = req.ScopeJSON
|
||||
}
|
||||
if s := strings.TrimSpace(req.Status); s != "" {
|
||||
p.Status = s
|
||||
}
|
||||
if req.Pinned != nil {
|
||||
p.Pinned = *req.Pinned
|
||||
}
|
||||
if err := h.db.UpdateProject(p); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, p)
|
||||
}
|
||||
|
||||
// DeleteProject DELETE /api/projects/:id
|
||||
func (h *ProjectHandler) DeleteProject(c *gin.Context) {
|
||||
if err := h.db.DeleteProject(c.Param("id")); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"success": true})
|
||||
}
|
||||
|
||||
type upsertFactRequest struct {
|
||||
FactKey string `json:"fact_key" binding:"required"`
|
||||
Category string `json:"category"`
|
||||
Summary string `json:"summary" binding:"required"`
|
||||
Body string `json:"body"`
|
||||
Confidence string `json:"confidence"`
|
||||
Pinned bool `json:"pinned"`
|
||||
RelatedVulnerabilityID string `json:"related_vulnerability_id"`
|
||||
}
|
||||
|
||||
// ListFacts GET /api/projects/:id/facts (fact_key 查询参数可获取单条详情)
|
||||
func (h *ProjectHandler) ListFacts(c *gin.Context) {
|
||||
projectID := c.Param("id")
|
||||
if key := strings.TrimSpace(c.Query("fact_key")); key != "" {
|
||||
f, err := h.db.GetProjectFactByKey(projectID, key)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, f)
|
||||
return
|
||||
}
|
||||
limit, _ := strconv.Atoi(c.DefaultQuery("limit", "100"))
|
||||
offset, _ := strconv.Atoi(c.Query("offset"))
|
||||
filter := database.ProjectFactListFilter{
|
||||
Category: c.Query("category"),
|
||||
Confidence: c.Query("confidence"),
|
||||
Search: c.Query("search"),
|
||||
}
|
||||
list, err := h.db.ListProjectFacts(projectID, filter, limit, offset)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if list == nil {
|
||||
list = []*database.ProjectFact{}
|
||||
}
|
||||
c.JSON(http.StatusOK, list)
|
||||
}
|
||||
|
||||
// CreateFact POST /api/projects/:id/facts
|
||||
func (h *ProjectHandler) CreateFact(c *gin.Context) {
|
||||
var req upsertFactRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
f := &database.ProjectFact{
|
||||
ProjectID: c.Param("id"),
|
||||
FactKey: req.FactKey,
|
||||
Category: req.Category,
|
||||
Summary: req.Summary,
|
||||
Body: req.Body,
|
||||
Confidence: req.Confidence,
|
||||
Pinned: req.Pinned,
|
||||
RelatedVulnerabilityID: req.RelatedVulnerabilityID,
|
||||
}
|
||||
created, err := h.db.UpsertProjectFact(f)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, created)
|
||||
}
|
||||
|
||||
// UpdateFact PUT /api/projects/:id/facts/:factId
|
||||
func (h *ProjectHandler) UpdateFact(c *gin.Context) {
|
||||
existing, err := h.db.GetProjectFact(c.Param("factId"))
|
||||
if err != nil || existing.ProjectID != c.Param("id") {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "事实不存在"})
|
||||
return
|
||||
}
|
||||
var req upsertFactRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if k := strings.TrimSpace(req.FactKey); k != "" {
|
||||
existing.FactKey = k
|
||||
}
|
||||
if req.Category != "" {
|
||||
existing.Category = req.Category
|
||||
}
|
||||
if req.Summary != "" {
|
||||
existing.Summary = req.Summary
|
||||
}
|
||||
if strings.TrimSpace(req.Body) != "" {
|
||||
existing.Body = req.Body
|
||||
}
|
||||
if req.Confidence != "" {
|
||||
existing.Confidence = req.Confidence
|
||||
}
|
||||
existing.Pinned = req.Pinned
|
||||
existing.RelatedVulnerabilityID = req.RelatedVulnerabilityID
|
||||
updated, err := h.db.UpsertProjectFact(existing)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, updated)
|
||||
}
|
||||
|
||||
// DeleteFact DELETE /api/projects/:id/facts/:factId
|
||||
func (h *ProjectHandler) DeleteFact(c *gin.Context) {
|
||||
existing, err := h.db.GetProjectFact(c.Param("factId"))
|
||||
if err != nil || existing.ProjectID != c.Param("id") {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "事实不存在"})
|
||||
return
|
||||
}
|
||||
if err := h.db.DeleteProjectFact(existing.ID); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"success": true})
|
||||
}
|
||||
|
||||
type deprecateFactRequest struct {
|
||||
FactKey string `json:"fact_key" binding:"required"`
|
||||
}
|
||||
|
||||
// DeprecateFact POST /api/projects/:id/facts/deprecate
|
||||
func (h *ProjectHandler) DeprecateFact(c *gin.Context) {
|
||||
var req deprecateFactRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if err := h.db.DeprecateProjectFact(c.Param("id"), req.FactKey); err != nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"success": true})
|
||||
}
|
||||
|
||||
type restoreFactRequest struct {
|
||||
FactKey string `json:"fact_key" binding:"required"`
|
||||
Confidence string `json:"confidence"` // 可选:confirmed | tentative,默认 tentative
|
||||
}
|
||||
|
||||
// RestoreFact POST /api/projects/:id/facts/restore
|
||||
func (h *ProjectHandler) RestoreFact(c *gin.Context) {
|
||||
var req restoreFactRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if err := h.db.RestoreProjectFact(c.Param("id"), req.FactKey, req.Confidence); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"success": true})
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/project"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// projectBlackboardBlock 根据对话 ID 构建项目事实索引块(用于注入 system prompt)。
|
||||
func (h *AgentHandler) projectBlackboardBlock(conversationID string) string {
|
||||
if h == nil || h.db == nil || h.config == nil {
|
||||
return ""
|
||||
}
|
||||
if !h.config.Project.Enabled {
|
||||
return ""
|
||||
}
|
||||
conversationID = strings.TrimSpace(conversationID)
|
||||
if conversationID == "" {
|
||||
return ""
|
||||
}
|
||||
projectID, err := h.db.GetConversationProjectID(conversationID)
|
||||
if err != nil || projectID == "" {
|
||||
return ""
|
||||
}
|
||||
block, err := project.BuildFactIndexBlock(h.db, projectID, h.config.Project)
|
||||
if err != nil {
|
||||
h.logger.Warn("构建项目黑板索引失败", zap.String("conversationId", conversationID), zap.Error(err))
|
||||
return ""
|
||||
}
|
||||
return strings.TrimSpace(block)
|
||||
}
|
||||
@@ -0,0 +1,18 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
)
|
||||
|
||||
// effectiveProjectID 请求/队列显式项目优先,否则使用 config.project.default_project_id。
|
||||
func effectiveProjectID(cfg *config.Config, explicit string) string {
|
||||
if pid := strings.TrimSpace(explicit); pid != "" {
|
||||
return pid
|
||||
}
|
||||
if cfg != nil {
|
||||
return strings.TrimSpace(cfg.Project.DefaultProjectID)
|
||||
}
|
||||
return ""
|
||||
}
|
||||
@@ -133,7 +133,9 @@ func (h *RobotHandler) getOrCreateConversation(platform, userID, title string) (
|
||||
} else {
|
||||
t = safeTruncateString(t, 50)
|
||||
}
|
||||
conv, err := h.db.CreateConversation(t)
|
||||
meta := database.ConversationCreateMeta{Source: "robot:" + platform}
|
||||
meta.ProjectID = effectiveProjectID(h.config, "")
|
||||
conv, err := h.db.CreateConversation(t, meta)
|
||||
if err != nil {
|
||||
h.logger.Warn("创建机器人会话失败", zap.Error(err))
|
||||
return "", false
|
||||
@@ -188,7 +190,9 @@ func (h *RobotHandler) setRole(platform, userID, roleName string) {
|
||||
// clearConversation 清空当前会话(切换到新对话)
|
||||
func (h *RobotHandler) clearConversation(platform, userID string) (newConvID string) {
|
||||
title := "新对话 " + time.Now().Format("01-02 15:04")
|
||||
conv, err := h.db.CreateConversation(title)
|
||||
meta := database.ConversationCreateMeta{Source: "robot:" + platform + ":new"}
|
||||
meta.ProjectID = effectiveProjectID(h.config, "")
|
||||
conv, err := h.db.CreateConversation(title, meta)
|
||||
if err != nil {
|
||||
h.logger.Warn("创建新对话失败", zap.Error(err))
|
||||
return ""
|
||||
@@ -242,7 +246,7 @@ func (h *RobotHandler) HandleMessage(platform, userID, text string) (reply strin
|
||||
h.cancelMu.Unlock()
|
||||
}()
|
||||
role := h.getRole(platform, userID)
|
||||
resp, newConvID, err := h.agentHandler.ProcessMessageForRobot(ctx, convID, text, role)
|
||||
resp, newConvID, err := h.agentHandler.ProcessMessageForRobot(ctx, platform, convID, text, role)
|
||||
if err != nil {
|
||||
h.logger.Warn("机器人 Agent 执行失败", zap.String("platform", platform), zap.String("userID", userID), zap.Error(err))
|
||||
if errors.Is(err, context.Canceled) {
|
||||
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
"regexp"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/config"
|
||||
|
||||
"gopkg.in/yaml.v3"
|
||||
@@ -21,6 +22,12 @@ type RoleHandler struct {
|
||||
config *config.Config
|
||||
configPath string
|
||||
logger *zap.Logger
|
||||
audit *audit.Service
|
||||
}
|
||||
|
||||
// SetAudit wires platform audit logging.
|
||||
func (h *RoleHandler) SetAudit(s *audit.Service) {
|
||||
h.audit = s
|
||||
}
|
||||
|
||||
// NewRoleHandler 创建新的角色处理器
|
||||
@@ -174,6 +181,9 @@ func (h *RoleHandler) UpdateRole(c *gin.Context) {
|
||||
}
|
||||
|
||||
h.logger.Info("更新角色", zap.String("oldKey", roleName), zap.String("newKey", finalKey), zap.String("name", req.Name))
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "role", "update", "更新角色", "role", finalKey, map[string]interface{}{"name": req.Name})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"message": "角色已更新",
|
||||
"role": req,
|
||||
@@ -219,6 +229,9 @@ func (h *RoleHandler) CreateRole(c *gin.Context) {
|
||||
}
|
||||
|
||||
h.logger.Info("创建角色", zap.String("roleName", req.Name))
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "role", "create", "创建角色", "role", req.Name, nil)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"message": "角色已创建",
|
||||
"role": req,
|
||||
@@ -287,6 +300,9 @@ func (h *RoleHandler) DeleteRole(c *gin.Context) {
|
||||
}
|
||||
|
||||
h.logger.Info("删除角色", zap.String("roleName", roleName))
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "role", "delete", "删除角色", "role", roleName, nil)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"message": "角色已删除",
|
||||
})
|
||||
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
"regexp"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/skillpackage"
|
||||
@@ -23,6 +24,12 @@ type SkillsHandler struct {
|
||||
configPath string
|
||||
logger *zap.Logger
|
||||
db *database.DB // 数据库连接(遗留统计;MCP list/read 已移除)
|
||||
audit *audit.Service
|
||||
}
|
||||
|
||||
// SetAudit wires platform audit logging.
|
||||
func (h *SkillsHandler) SetAudit(s *audit.Service) {
|
||||
h.audit = s
|
||||
}
|
||||
|
||||
// NewSkillsHandler 创建新的Skills处理器
|
||||
@@ -365,6 +372,9 @@ func (h *SkillsHandler) CreateSkill(c *gin.Context) {
|
||||
}
|
||||
|
||||
h.logger.Info("创建skill成功", zap.String("skill", req.Name))
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "skill", "create", "创建 Skill", "skill", req.Name, nil)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"message": "skill已创建",
|
||||
"skill": map[string]interface{}{
|
||||
@@ -425,6 +435,9 @@ func (h *SkillsHandler) UpdateSkill(c *gin.Context) {
|
||||
}
|
||||
|
||||
h.logger.Info("更新skill成功", zap.String("skill", skillName))
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "skill", "update", "更新 Skill", "skill", skillName, nil)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"message": "skill已更新",
|
||||
})
|
||||
@@ -459,6 +472,11 @@ func (h *SkillsHandler) DeleteSkill(c *gin.Context) {
|
||||
}
|
||||
|
||||
h.logger.Info("删除skill成功", zap.String("skill", skillName))
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "skill", "delete", "删除 Skill", "skill", skillName, map[string]interface{}{
|
||||
"affected_roles": affectedRoles,
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"message": responseMsg,
|
||||
"affected_roles": affectedRoles,
|
||||
|
||||
@@ -253,5 +253,5 @@ func (h *TerminalHandler) RunCommandStream(c *gin.Context) {
|
||||
flusher.Flush()
|
||||
}
|
||||
|
||||
runCommandStreamImpl(cmd, sendEvent, ctx)
|
||||
_ = runCommandStreamImpl(cmd, sendEvent, ctx)
|
||||
}
|
||||
|
||||
@@ -15,11 +15,11 @@ const ptyCols = 256
|
||||
const ptyRows = 40
|
||||
|
||||
// runCommandStreamImpl 在 Unix 下用 PTY 执行,使 ping 等命令按终端宽度排版(isatty 为真)
|
||||
func runCommandStreamImpl(cmd *exec.Cmd, sendEvent func(streamEvent), ctx context.Context) {
|
||||
func runCommandStreamImpl(cmd *exec.Cmd, sendEvent func(streamEvent), ctx context.Context) int {
|
||||
ptmx, err := pty.StartWithSize(cmd, &pty.Winsize{Cols: ptyCols, Rows: ptyRows})
|
||||
if err != nil {
|
||||
sendEvent(streamEvent{T: "exit", C: -1})
|
||||
return
|
||||
return -1
|
||||
}
|
||||
defer ptmx.Close()
|
||||
|
||||
@@ -43,4 +43,5 @@ func runCommandStreamImpl(cmd *exec.Cmd, sendEvent func(streamEvent), ctx contex
|
||||
exitCode = -1
|
||||
}
|
||||
sendEvent(streamEvent{T: "exit", C: exitCode})
|
||||
return exitCode
|
||||
}
|
||||
|
||||
@@ -11,20 +11,20 @@ import (
|
||||
)
|
||||
|
||||
// runCommandStreamImpl 在 Windows 下用 stdout/stderr 管道执行
|
||||
func runCommandStreamImpl(cmd *exec.Cmd, sendEvent func(streamEvent), ctx context.Context) {
|
||||
func runCommandStreamImpl(cmd *exec.Cmd, sendEvent func(streamEvent), ctx context.Context) int {
|
||||
stdoutPipe, err := cmd.StdoutPipe()
|
||||
if err != nil {
|
||||
sendEvent(streamEvent{T: "exit", C: -1})
|
||||
return
|
||||
return -1
|
||||
}
|
||||
stderrPipe, err := cmd.StderrPipe()
|
||||
if err != nil {
|
||||
sendEvent(streamEvent{T: "exit", C: -1})
|
||||
return
|
||||
return -1
|
||||
}
|
||||
if err := cmd.Start(); err != nil {
|
||||
sendEvent(streamEvent{T: "exit", C: -1})
|
||||
return
|
||||
return -1
|
||||
}
|
||||
|
||||
normalize := func(s string) string {
|
||||
@@ -62,4 +62,5 @@ func runCommandStreamImpl(cmd *exec.Cmd, sendEvent func(streamEvent), ctx contex
|
||||
exitCode = -1
|
||||
}
|
||||
sendEvent(streamEvent{T: "exit", C: exitCode})
|
||||
return exitCode
|
||||
}
|
||||
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
@@ -16,6 +17,12 @@ import (
|
||||
type VulnerabilityHandler struct {
|
||||
db *database.DB
|
||||
logger *zap.Logger
|
||||
audit *audit.Service
|
||||
}
|
||||
|
||||
// SetAudit wires platform audit logging.
|
||||
func (h *VulnerabilityHandler) SetAudit(s *audit.Service) {
|
||||
h.audit = s
|
||||
}
|
||||
|
||||
// NewVulnerabilityHandler 创建新的漏洞处理器
|
||||
@@ -29,6 +36,7 @@ func NewVulnerabilityHandler(db *database.DB, logger *zap.Logger) *Vulnerability
|
||||
// CreateVulnerabilityRequest 创建漏洞请求
|
||||
type CreateVulnerabilityRequest struct {
|
||||
ConversationID string `json:"conversation_id" binding:"required"`
|
||||
ProjectID string `json:"project_id"`
|
||||
ConversationTag string `json:"conversation_tag"`
|
||||
TaskTag string `json:"task_tag"`
|
||||
Title string `json:"title" binding:"required"`
|
||||
@@ -52,6 +60,7 @@ func (h *VulnerabilityHandler) CreateVulnerability(c *gin.Context) {
|
||||
|
||||
vuln := &database.Vulnerability{
|
||||
ConversationID: req.ConversationID,
|
||||
ProjectID: strings.TrimSpace(req.ProjectID),
|
||||
ConversationTag: req.ConversationTag,
|
||||
TaskTag: req.TaskTag,
|
||||
Title: req.Title,
|
||||
@@ -72,6 +81,11 @@ func (h *VulnerabilityHandler) CreateVulnerability(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "vulnerability", "create", "创建漏洞记录", "vulnerability", created.ID, map[string]interface{}{
|
||||
"severity": created.Severity, "title": created.Title,
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusOK, created)
|
||||
}
|
||||
|
||||
@@ -98,18 +112,30 @@ type ListVulnerabilitiesResponse struct {
|
||||
TotalPages int `json:"total_pages"`
|
||||
}
|
||||
|
||||
func parseVulnerabilityListFilter(c *gin.Context) database.VulnerabilityListFilter {
|
||||
q := strings.TrimSpace(c.Query("q"))
|
||||
if q == "" {
|
||||
q = strings.TrimSpace(c.Query("search"))
|
||||
}
|
||||
return database.VulnerabilityListFilter{
|
||||
ProjectID: c.Query("project_id"),
|
||||
ID: c.Query("id"),
|
||||
Search: q,
|
||||
ConversationID: c.Query("conversation_id"),
|
||||
Severity: c.Query("severity"),
|
||||
Status: c.Query("status"),
|
||||
TaskID: c.Query("task_id"),
|
||||
ConversationTag: c.Query("conversation_tag"),
|
||||
TaskTag: c.Query("task_tag"),
|
||||
}
|
||||
}
|
||||
|
||||
// ListVulnerabilities 列出漏洞
|
||||
func (h *VulnerabilityHandler) ListVulnerabilities(c *gin.Context) {
|
||||
limitStr := c.DefaultQuery("limit", "20")
|
||||
offsetStr := c.DefaultQuery("offset", "0")
|
||||
pageStr := c.Query("page")
|
||||
id := c.Query("id")
|
||||
conversationID := c.Query("conversation_id")
|
||||
severity := c.Query("severity")
|
||||
status := c.Query("status")
|
||||
taskID := c.Query("task_id")
|
||||
conversationTag := c.Query("conversation_tag")
|
||||
taskTag := c.Query("task_tag")
|
||||
filter := parseVulnerabilityListFilter(c)
|
||||
|
||||
limit, _ := strconv.Atoi(limitStr)
|
||||
offset, _ := strconv.Atoi(offsetStr)
|
||||
@@ -131,7 +157,7 @@ func (h *VulnerabilityHandler) ListVulnerabilities(c *gin.Context) {
|
||||
}
|
||||
|
||||
// 获取总数
|
||||
total, err := h.db.CountVulnerabilities(id, conversationID, severity, status, taskID, conversationTag, taskTag)
|
||||
total, err := h.db.CountVulnerabilities(filter)
|
||||
if err != nil {
|
||||
h.logger.Error("获取漏洞总数失败", zap.Error(err))
|
||||
// 继续执行,使用0作为总数
|
||||
@@ -139,7 +165,7 @@ func (h *VulnerabilityHandler) ListVulnerabilities(c *gin.Context) {
|
||||
}
|
||||
|
||||
// 获取漏洞列表
|
||||
vulnerabilities, err := h.db.ListVulnerabilities(limit, offset, id, conversationID, severity, status, taskID, conversationTag, taskTag)
|
||||
vulnerabilities, err := h.db.ListVulnerabilities(limit, offset, filter)
|
||||
if err != nil {
|
||||
h.logger.Error("获取漏洞列表失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
@@ -170,17 +196,18 @@ func (h *VulnerabilityHandler) ListVulnerabilities(c *gin.Context) {
|
||||
|
||||
// UpdateVulnerabilityRequest 更新漏洞请求
|
||||
type UpdateVulnerabilityRequest struct {
|
||||
ConversationTag string `json:"conversation_tag"`
|
||||
TaskTag string `json:"task_tag"`
|
||||
Title string `json:"title"`
|
||||
Description string `json:"description"`
|
||||
Severity string `json:"severity"`
|
||||
Status string `json:"status"`
|
||||
Type string `json:"type"`
|
||||
Target string `json:"target"`
|
||||
Proof string `json:"proof"`
|
||||
Impact string `json:"impact"`
|
||||
Recommendation string `json:"recommendation"`
|
||||
ProjectID *string `json:"project_id"`
|
||||
ConversationTag string `json:"conversation_tag"`
|
||||
TaskTag string `json:"task_tag"`
|
||||
Title string `json:"title"`
|
||||
Description string `json:"description"`
|
||||
Severity string `json:"severity"`
|
||||
Status string `json:"status"`
|
||||
Type string `json:"type"`
|
||||
Target string `json:"target"`
|
||||
Proof string `json:"proof"`
|
||||
Impact string `json:"impact"`
|
||||
Recommendation string `json:"recommendation"`
|
||||
}
|
||||
|
||||
// UpdateVulnerability 更新漏洞
|
||||
@@ -201,6 +228,9 @@ func (h *VulnerabilityHandler) UpdateVulnerability(c *gin.Context) {
|
||||
}
|
||||
|
||||
// 更新字段
|
||||
if req.ProjectID != nil {
|
||||
existing.ProjectID = strings.TrimSpace(*req.ProjectID)
|
||||
}
|
||||
if req.ConversationTag != "" {
|
||||
existing.ConversationTag = req.ConversationTag
|
||||
}
|
||||
@@ -249,6 +279,11 @@ func (h *VulnerabilityHandler) UpdateVulnerability(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "vulnerability", "update", "更新漏洞记录", "vulnerability", id, map[string]interface{}{
|
||||
"severity": updated.Severity, "status": updated.Status, "project_id": updated.ProjectID,
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusOK, updated)
|
||||
}
|
||||
|
||||
@@ -262,15 +297,25 @@ func (h *VulnerabilityHandler) DeleteVulnerability(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
if h.audit != nil {
|
||||
h.audit.Record(c, audit.Entry{
|
||||
Category: "vulnerability",
|
||||
Action: "delete",
|
||||
Result: "success",
|
||||
ResourceType: "vulnerability",
|
||||
ResourceID: id,
|
||||
Message: "删除漏洞记录",
|
||||
})
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"message": "删除成功"})
|
||||
}
|
||||
|
||||
// GetVulnerabilityStats 获取漏洞统计
|
||||
func (h *VulnerabilityHandler) GetVulnerabilityStats(c *gin.Context) {
|
||||
conversationID := c.Query("conversation_id")
|
||||
taskID := c.Query("task_id")
|
||||
filter := parseVulnerabilityListFilter(c)
|
||||
|
||||
stats, err := h.db.GetVulnerabilityStats(conversationID, taskID)
|
||||
stats, err := h.db.GetVulnerabilityStats(filter)
|
||||
if err != nil {
|
||||
h.logger.Error("获取漏洞统计失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
@@ -304,15 +349,9 @@ func (h *VulnerabilityHandler) ExportVulnerabilities(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
id := c.Query("id")
|
||||
conversationID := c.Query("conversation_id")
|
||||
severity := c.Query("severity")
|
||||
status := c.Query("status")
|
||||
taskID := c.Query("task_id")
|
||||
conversationTag := c.Query("conversation_tag")
|
||||
taskTag := c.Query("task_tag")
|
||||
filter := parseVulnerabilityListFilter(c)
|
||||
|
||||
total, err := h.db.CountVulnerabilities(id, conversationID, severity, status, taskID, conversationTag, taskTag)
|
||||
total, err := h.db.CountVulnerabilities(filter)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
@@ -322,7 +361,7 @@ func (h *VulnerabilityHandler) ExportVulnerabilities(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
items, err := h.db.ListVulnerabilities(total, 0, id, conversationID, severity, status, taskID, conversationTag, taskTag)
|
||||
items, err := h.db.ListVulnerabilities(total, 0, filter)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
|
||||
@@ -2,6 +2,7 @@ package handler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"crypto/tls"
|
||||
"database/sql"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
@@ -12,6 +13,7 @@ import (
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/database"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
@@ -304,6 +306,12 @@ type WebShellHandler struct {
|
||||
logger *zap.Logger
|
||||
client *http.Client
|
||||
db *database.DB
|
||||
audit *audit.Service
|
||||
}
|
||||
|
||||
// SetAudit wires platform audit logging.
|
||||
func (h *WebShellHandler) SetAudit(s *audit.Service) {
|
||||
h.audit = s
|
||||
}
|
||||
|
||||
// NewWebShellHandler 创建 WebShell 处理器,db 可为 nil(连接配置接口将不可用)
|
||||
@@ -311,8 +319,12 @@ func NewWebShellHandler(logger *zap.Logger, db *database.DB) *WebShellHandler {
|
||||
return &WebShellHandler{
|
||||
logger: logger,
|
||||
client: &http.Client{
|
||||
Timeout: 30 * time.Second,
|
||||
Transport: &http.Transport{DisableKeepAlives: false},
|
||||
Timeout: 30 * time.Second,
|
||||
Transport: &http.Transport{
|
||||
DisableKeepAlives: false,
|
||||
// WebShell 场景常见自签证书或 IP 访问(证书无 IP SAN);默认跳过校验,与蚁剑等客户端一致。
|
||||
TLSClientConfig: &tls.Config{InsecureSkipVerify: true}, //nolint:gosec // intentional for webshell proxy
|
||||
},
|
||||
},
|
||||
db: db,
|
||||
}
|
||||
@@ -403,6 +415,15 @@ func (h *WebShellHandler) CreateConnection(c *gin.Context) {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if h.audit != nil {
|
||||
host := req.URL
|
||||
if u, err := url.Parse(req.URL); err == nil {
|
||||
host = u.Host
|
||||
}
|
||||
h.audit.RecordOK(c, "webshell", "connection_create", "创建 WebShell 连接", "webshell_connection", conn.ID, map[string]interface{}{
|
||||
"host": host, "type": shellType,
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusOK, conn)
|
||||
}
|
||||
|
||||
@@ -485,6 +506,9 @@ func (h *WebShellHandler) DeleteConnection(c *gin.Context) {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "webshell", "connection_delete", "删除 WebShell 连接", "webshell_connection", id, nil)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
}
|
||||
|
||||
@@ -714,8 +738,9 @@ func (h *WebShellHandler) Exec(c *gin.Context) {
|
||||
output := decodeWebshellOutput(out, req.Encoding)
|
||||
httpCode := resp.StatusCode
|
||||
|
||||
ok := resp.StatusCode == http.StatusOK
|
||||
c.JSON(http.StatusOK, ExecResponse{
|
||||
OK: resp.StatusCode == http.StatusOK,
|
||||
OK: ok,
|
||||
Output: output,
|
||||
HTTPCode: httpCode,
|
||||
})
|
||||
|
||||
@@ -15,7 +15,7 @@ const WebshellSkillHintMultiAgent = "Skills 包请使用 Eino 多代理内置 `s
|
||||
|
||||
// webshellAssistantToolList AI 助手在 WebShell 上下文下允许使用的工具清单(展示给模型用)。
|
||||
// 注意:此处只是展示字符串,真正的权限限制是在调用方设置的 roleTools 切片里。
|
||||
const webshellAssistantToolList = "webshell_exec、webshell_file_list、webshell_file_read、webshell_file_write、record_vulnerability、list_knowledge_risk_types、search_knowledge_base"
|
||||
const webshellAssistantToolList = "webshell_exec、webshell_file_list、webshell_file_read、webshell_file_write、record_vulnerability、list_vulnerabilities、get_vulnerability、upsert_project_fact、get_project_fact、list_project_facts、search_project_facts、deprecate_project_fact、restore_project_fact、list_knowledge_risk_types、search_knowledge_base"
|
||||
|
||||
// BuildWebshellAssistantContext 根据连接信息与用户原始消息组装 AI 助手的上下文提示词。
|
||||
// 上下文包含:连接 ID、备注、目标系统(及对应命令集建议)、响应编码、可用工具清单、Skills 加载入口、
|
||||
|
||||
@@ -0,0 +1,293 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/robot/ilink"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const wechatLoginTTL = 5 * time.Minute
|
||||
|
||||
// WechatConfigSaver 绑定成功后写入配置并重启机器人连接
|
||||
type WechatConfigSaver interface {
|
||||
ApplyWechatRobotBinding(cfg config.RobotWechatConfig) error
|
||||
}
|
||||
|
||||
type wechatLoginSession struct {
|
||||
QRCode string
|
||||
QRCodeImgURL string
|
||||
PendingVerify string
|
||||
CurrentBaseURL string
|
||||
StartedAt time.Time
|
||||
}
|
||||
|
||||
// WechatRobotHandler 微信 iLink 机器人(扫码绑定 + 配置)
|
||||
type WechatRobotHandler struct {
|
||||
config *config.Config
|
||||
configSaver WechatConfigSaver
|
||||
logger *zap.Logger
|
||||
mu sync.Mutex
|
||||
logins map[string]*wechatLoginSession
|
||||
}
|
||||
|
||||
// NewWechatRobotHandler 创建微信机器人处理器
|
||||
func NewWechatRobotHandler(cfg *config.Config, saver WechatConfigSaver, logger *zap.Logger) *WechatRobotHandler {
|
||||
return &WechatRobotHandler{
|
||||
config: cfg,
|
||||
configSaver: saver,
|
||||
logger: logger,
|
||||
logins: make(map[string]*wechatLoginSession),
|
||||
}
|
||||
}
|
||||
|
||||
func (h *WechatRobotHandler) purgeExpiredLogins() {
|
||||
now := time.Now()
|
||||
for k, v := range h.logins {
|
||||
if now.Sub(v.StartedAt) > wechatLoginTTL {
|
||||
delete(h.logins, k)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (h *WechatRobotHandler) ilinkClient(baseURL string) *ilink.Client {
|
||||
ver := h.config.Version
|
||||
if ver == "" {
|
||||
ver = "1.0.0"
|
||||
}
|
||||
ver = strings.TrimPrefix(strings.TrimSpace(ver), "v")
|
||||
ver = strings.TrimPrefix(ver, "V")
|
||||
wc := h.config.Robots.Wechat
|
||||
return ilink.NewClient(baseURL, wc.BotToken, wc.BotAgent, ilink.BuildClientVersion(ver))
|
||||
}
|
||||
|
||||
// HandleWechatQRCode POST /api/robot/wechat/qrcode — 生成绑定二维码
|
||||
func (h *WechatRobotHandler) HandleWechatQRCode(c *gin.Context) {
|
||||
h.mu.Lock()
|
||||
h.purgeExpiredLogins()
|
||||
h.mu.Unlock()
|
||||
|
||||
var req struct {
|
||||
BotType string `json:"bot_type"`
|
||||
}
|
||||
_ = c.ShouldBindJSON(&req)
|
||||
|
||||
botType := req.BotType
|
||||
if botType == "" {
|
||||
botType = h.config.Robots.Wechat.BotType
|
||||
}
|
||||
if botType == "" {
|
||||
botType = ilink.DefaultBotType
|
||||
}
|
||||
baseURL := h.config.Robots.Wechat.BaseURL
|
||||
if baseURL == "" {
|
||||
baseURL = ilink.DefaultBaseURL
|
||||
}
|
||||
|
||||
var localTokens []string
|
||||
if t := h.config.Robots.Wechat.BotToken; t != "" {
|
||||
localTokens = []string{t}
|
||||
}
|
||||
|
||||
client := h.ilinkClient(baseURL)
|
||||
ctx, cancel := context.WithTimeout(c.Request.Context(), 30*time.Second)
|
||||
defer cancel()
|
||||
|
||||
qr, err := client.GetBotQRCode(ctx, botType, localTokens)
|
||||
if err != nil {
|
||||
h.logger.Warn("获取微信二维码失败", zap.Error(err))
|
||||
c.JSON(http.StatusBadGateway, gin.H{"error": "获取二维码失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
if qr.QRCode == "" || qr.QRCodeImgContent == "" {
|
||||
c.JSON(http.StatusBadGateway, gin.H{"error": "微信服务器未返回有效二维码"})
|
||||
return
|
||||
}
|
||||
|
||||
sessionKey := uuid.New().String()
|
||||
h.mu.Lock()
|
||||
h.logins[sessionKey] = &wechatLoginSession{
|
||||
QRCode: qr.QRCode,
|
||||
QRCodeImgURL: qr.QRCodeImgContent,
|
||||
CurrentBaseURL: baseURL,
|
||||
StartedAt: time.Now(),
|
||||
}
|
||||
h.mu.Unlock()
|
||||
|
||||
resp := gin.H{
|
||||
"session_key": sessionKey,
|
||||
"qrcode": qr.QRCode,
|
||||
"qrcode_open_url": qr.QRCodeImgContent,
|
||||
"message": "请使用微信扫描二维码并确认绑定",
|
||||
}
|
||||
if dataURL, err := ilink.QRCodeDataURL(qr.QRCodeImgContent, 256); err != nil {
|
||||
h.logger.Warn("生成二维码图片失败", zap.Error(err))
|
||||
} else {
|
||||
resp["qrcode_image_data_url"] = dataURL
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// HandleWechatQRCodeStatus GET /api/robot/wechat/qrcode/status — 轮询扫码状态
|
||||
func (h *WechatRobotHandler) HandleWechatQRCodeStatus(c *gin.Context) {
|
||||
sessionKey := c.Query("session_key")
|
||||
verifyCode := c.Query("verify_code")
|
||||
if sessionKey == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "缺少 session_key"})
|
||||
return
|
||||
}
|
||||
|
||||
h.mu.Lock()
|
||||
sess, ok := h.logins[sessionKey]
|
||||
h.mu.Unlock()
|
||||
if !ok {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "登录会话不存在或已过期,请重新生成二维码"})
|
||||
return
|
||||
}
|
||||
if time.Since(sess.StartedAt) > wechatLoginTTL {
|
||||
h.mu.Lock()
|
||||
delete(h.logins, sessionKey)
|
||||
h.mu.Unlock()
|
||||
c.JSON(http.StatusGone, gin.H{"error": "二维码已过期,请重新生成"})
|
||||
return
|
||||
}
|
||||
|
||||
baseURL := sess.CurrentBaseURL
|
||||
if baseURL == "" {
|
||||
baseURL = ilink.DefaultBaseURL
|
||||
}
|
||||
vc := verifyCode
|
||||
if vc == "" {
|
||||
vc = sess.PendingVerify
|
||||
}
|
||||
|
||||
client := h.ilinkClient(baseURL)
|
||||
ctx, cancel := context.WithTimeout(c.Request.Context(), 40*time.Second)
|
||||
defer cancel()
|
||||
|
||||
st, err := client.GetQRCodeStatus(ctx, sess.QRCode, vc)
|
||||
if err != nil {
|
||||
h.logger.Warn("轮询微信二维码状态失败", zap.Error(err))
|
||||
c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
switch st.Status {
|
||||
case "wait", "scaned":
|
||||
c.JSON(http.StatusOK, gin.H{"status": st.Status})
|
||||
return
|
||||
case "need_verifycode":
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"status": st.Status,
|
||||
"message": "请在手机微信查看配对数字,并在下方输入",
|
||||
})
|
||||
return
|
||||
case "scaned_but_redirect":
|
||||
if st.RedirectHost != "" {
|
||||
h.mu.Lock()
|
||||
if s, ok := h.logins[sessionKey]; ok {
|
||||
s.CurrentBaseURL = "https://" + st.RedirectHost
|
||||
}
|
||||
h.mu.Unlock()
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"status": st.Status})
|
||||
return
|
||||
case "binded_redirect":
|
||||
h.mu.Lock()
|
||||
delete(h.logins, sessionKey)
|
||||
h.mu.Unlock()
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"status": st.Status,
|
||||
"already_connected": true,
|
||||
"message": "该微信已绑定过,无需重复绑定",
|
||||
})
|
||||
return
|
||||
case "confirmed":
|
||||
if st.BotToken == "" || st.ILinkBotID == "" {
|
||||
c.JSON(http.StatusBadGateway, gin.H{"error": "绑定确认成功但缺少 bot_token"})
|
||||
return
|
||||
}
|
||||
saveBase := st.BaseURL
|
||||
if saveBase == "" {
|
||||
saveBase = baseURL
|
||||
}
|
||||
wc := h.config.Robots.Wechat
|
||||
wc.Enabled = true
|
||||
wc.BotToken = st.BotToken
|
||||
wc.ILinkBotID = st.ILinkBotID
|
||||
wc.ILinkUserID = st.ILinkUserID
|
||||
wc.BaseURL = saveBase
|
||||
if wc.BotType == "" {
|
||||
wc.BotType = ilink.DefaultBotType
|
||||
}
|
||||
if wc.BotAgent == "" {
|
||||
wc.BotAgent = ilink.DefaultBotAgent
|
||||
}
|
||||
if h.configSaver != nil {
|
||||
if err := h.configSaver.ApplyWechatRobotBinding(wc); err != nil {
|
||||
h.logger.Warn("保存微信机器人配置失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "保存配置失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
} else {
|
||||
h.config.Robots.Wechat = wc
|
||||
}
|
||||
h.mu.Lock()
|
||||
delete(h.logins, sessionKey)
|
||||
h.mu.Unlock()
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"status": "confirmed",
|
||||
"message": "绑定成功,微信机器人已启用",
|
||||
"ilink_bot_id": st.ILinkBotID,
|
||||
"ilink_user_id": st.ILinkUserID,
|
||||
})
|
||||
return
|
||||
default:
|
||||
c.JSON(http.StatusOK, gin.H{"status": st.Status})
|
||||
}
|
||||
}
|
||||
|
||||
// HandleWechatVerifyCode POST /api/robot/wechat/qrcode/verify — 提交手机配对数字
|
||||
func (h *WechatRobotHandler) HandleWechatVerifyCode(c *gin.Context) {
|
||||
var req struct {
|
||||
SessionKey string `json:"session_key"`
|
||||
VerifyCode string `json:"verify_code"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil || req.SessionKey == "" || req.VerifyCode == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "需要 session_key 与 verify_code"})
|
||||
return
|
||||
}
|
||||
h.mu.Lock()
|
||||
sess, ok := h.logins[req.SessionKey]
|
||||
if ok {
|
||||
sess.PendingVerify = req.VerifyCode
|
||||
}
|
||||
h.mu.Unlock()
|
||||
if !ok {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "登录会话不存在或已过期"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"message": "已提交配对码,请继续等待绑定"})
|
||||
}
|
||||
|
||||
// HandleWechatStatus GET /api/robot/wechat/status — 当前绑定状态(供前端展示)
|
||||
func (h *WechatRobotHandler) HandleWechatStatus(c *gin.Context) {
|
||||
wc := h.config.Robots.Wechat
|
||||
bound := wc.BotToken != "" && wc.ILinkBotID != ""
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"enabled": wc.Enabled,
|
||||
"bound": bound,
|
||||
"ilink_bot_id": wc.ILinkBotID,
|
||||
"ilink_user_id": wc.ILinkUserID,
|
||||
"base_url": wc.BaseURL,
|
||||
})
|
||||
}
|
||||
@@ -4,7 +4,17 @@ package builtin
|
||||
// 所有代码中使用内置工具名称的地方都应该使用这些常量,而不是硬编码字符串
|
||||
const (
|
||||
// 漏洞管理工具
|
||||
ToolRecordVulnerability = "record_vulnerability"
|
||||
ToolRecordVulnerability = "record_vulnerability"
|
||||
ToolListVulnerabilities = "list_vulnerabilities"
|
||||
ToolGetVulnerability = "get_vulnerability"
|
||||
|
||||
// 项目黑板(事实)工具
|
||||
ToolUpsertProjectFact = "upsert_project_fact"
|
||||
ToolGetProjectFact = "get_project_fact"
|
||||
ToolListProjectFacts = "list_project_facts"
|
||||
ToolSearchProjectFacts = "search_project_facts"
|
||||
ToolDeprecateProjectFact = "deprecate_project_fact"
|
||||
ToolRestoreProjectFact = "restore_project_fact"
|
||||
|
||||
// 知识库工具
|
||||
ToolListKnowledgeRiskTypes = "list_knowledge_risk_types"
|
||||
@@ -53,6 +63,14 @@ const (
|
||||
func IsBuiltinTool(toolName string) bool {
|
||||
switch toolName {
|
||||
case ToolRecordVulnerability,
|
||||
ToolListVulnerabilities,
|
||||
ToolGetVulnerability,
|
||||
ToolUpsertProjectFact,
|
||||
ToolGetProjectFact,
|
||||
ToolListProjectFacts,
|
||||
ToolSearchProjectFacts,
|
||||
ToolDeprecateProjectFact,
|
||||
ToolRestoreProjectFact,
|
||||
ToolListKnowledgeRiskTypes,
|
||||
ToolSearchKnowledgeBase,
|
||||
ToolWebshellExec,
|
||||
@@ -96,6 +114,14 @@ func IsBuiltinTool(toolName string) bool {
|
||||
func GetAllBuiltinTools() []string {
|
||||
return []string{
|
||||
ToolRecordVulnerability,
|
||||
ToolListVulnerabilities,
|
||||
ToolGetVulnerability,
|
||||
ToolUpsertProjectFact,
|
||||
ToolGetProjectFact,
|
||||
ToolListProjectFacts,
|
||||
ToolSearchProjectFacts,
|
||||
ToolDeprecateProjectFact,
|
||||
ToolRestoreProjectFact,
|
||||
ToolListKnowledgeRiskTypes,
|
||||
ToolSearchKnowledgeBase,
|
||||
ToolWebshellExec,
|
||||
|
||||
@@ -77,6 +77,9 @@ type einoADKRunLoopArgs struct {
|
||||
StreamsMainAssistant func(agent string) bool
|
||||
EinoRoleTag func(agent string) string
|
||||
CheckpointDir string
|
||||
// RunRetryMaxAttempts / RunRetryMaxBackoffSec:429、5xx、网络抖动时的指数退避续跑(0=默认 10 次 / 30s 上限)。
|
||||
RunRetryMaxAttempts int
|
||||
RunRetryMaxBackoffSec int
|
||||
|
||||
McpIDsMu *sync.Mutex
|
||||
McpIDs *[]string
|
||||
@@ -177,6 +180,8 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
||||
var einoMainRound int
|
||||
var einoLastAgent string
|
||||
subAgentToolStep := make(map[string]int)
|
||||
// mainAgentToolStep:主代理每次工具调用批次递增,供 UI 显示「第 N 轮」(单代理无子代理切换时原先会一直停在第 1 轮)。
|
||||
mainAgentToolStep := make(map[string]int)
|
||||
pendingByID := make(map[string]toolCallPendingInfo)
|
||||
pendingQueueByAgent := make(map[string][]string)
|
||||
markPending := func(tc toolCallPendingInfo) {
|
||||
@@ -435,6 +440,28 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
||||
return runErr
|
||||
}
|
||||
|
||||
// maybeRetryTransientRun:不在此层 runner.Run/Resume;由 handler 落库 + loadHistoryFromAgentTrace 分段续跑(同中断并继续)。
|
||||
maybeRetryTransientRun := func(runErr error) (retry bool, fatal error) {
|
||||
if runErr == nil || !isEinoTransientRunError(runErr) {
|
||||
return false, handleRunErr(runErr)
|
||||
}
|
||||
if logger != nil {
|
||||
logger.Warn("eino transient error, ending run segment for handler resume",
|
||||
zap.Error(runErr),
|
||||
zap.String("orchestration", orchMode))
|
||||
}
|
||||
if progress != nil {
|
||||
progress("eino_run_retry", "遇到临时错误(限流或网络波动),将保存上下文并重试…", map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
"orchestration": orchMode,
|
||||
"error": runErr.Error(),
|
||||
"resumeKind": "trace_segment",
|
||||
})
|
||||
}
|
||||
return false, ErrTransientRetryContinue
|
||||
}
|
||||
|
||||
takePartial := func(runErr error) (*RunResult, error) {
|
||||
if len(runAccumulatedMsgs) <= baseAccumulatedCount {
|
||||
return nil, runErr
|
||||
@@ -517,7 +544,7 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
||||
continue
|
||||
}
|
||||
if ev.Err != nil {
|
||||
if retErr := handleRunErr(ev.Err); retErr != nil {
|
||||
if _, retErr := maybeRetryTransientRun(ev.Err); retErr != nil {
|
||||
return takePartial(retErr)
|
||||
}
|
||||
}
|
||||
@@ -529,8 +556,10 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
||||
}
|
||||
}
|
||||
if streamsMainAssistant(ev.AgentName) {
|
||||
mainIterKey := einoMainIterationKey(iterEinoAgent, orchestratorName)
|
||||
if einoMainRound == 0 {
|
||||
einoMainRound = 1
|
||||
mainAgentToolStep[mainIterKey] = 1
|
||||
progress("iteration", "", map[string]interface{}{
|
||||
"iteration": 1,
|
||||
"einoScope": "main",
|
||||
@@ -540,17 +569,26 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
})
|
||||
} else if einoLastAgent != "" && !streamsMainAssistant(einoLastAgent) {
|
||||
einoMainRound++
|
||||
progress("iteration", "", map[string]interface{}{
|
||||
"iteration": einoMainRound,
|
||||
"einoScope": "main",
|
||||
"einoRole": "orchestrator",
|
||||
"einoAgent": iterEinoAgent,
|
||||
"orchestration": orchMode,
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
})
|
||||
} else if einoLastAgent != "" {
|
||||
needBump := false
|
||||
if !streamsMainAssistant(einoLastAgent) {
|
||||
needBump = true // 子代理 → 主代理
|
||||
} else if einoLastAgent != ev.AgentName {
|
||||
needBump = true // plan_execute:planner ↔ executor 等主代理切换
|
||||
}
|
||||
if needBump {
|
||||
einoMainRound++
|
||||
mainAgentToolStep[mainIterKey] = einoMainRound
|
||||
progress("iteration", "", map[string]interface{}{
|
||||
"iteration": einoMainRound,
|
||||
"einoScope": "main",
|
||||
"einoRole": "orchestrator",
|
||||
"einoAgent": iterEinoAgent,
|
||||
"orchestration": orchMode,
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
einoLastAgent = ev.AgentName
|
||||
@@ -644,9 +682,9 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
||||
"orchestration": orchMode,
|
||||
})
|
||||
}
|
||||
progress("reasoning_chain_stream_delta", displayDelta, map[string]interface{}{
|
||||
progress("reasoning_chain_stream_delta", displayDelta, openai.WithSSEAccumulated(map[string]interface{}{
|
||||
"streamId": reasoningStreamID,
|
||||
})
|
||||
}, fullDisplay))
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -676,13 +714,13 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
||||
})
|
||||
streamHeaderSent = true
|
||||
}
|
||||
progress("response_delta", contentDelta, map[string]interface{}{
|
||||
progress("response_delta", contentDelta, openai.WithSSEAccumulated(map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"mcpExecutionIds": snapshotMCPIDs(),
|
||||
"einoRole": "orchestrator",
|
||||
"einoAgent": ev.AgentName,
|
||||
"orchestration": orchMode,
|
||||
})
|
||||
}, mainAssistantBuf))
|
||||
mainAssistWireAccum, _ = normalizeStreamingDelta(mainAssistWireAccum, contentDelta)
|
||||
}
|
||||
}
|
||||
@@ -701,10 +739,10 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
||||
"source": "eino",
|
||||
})
|
||||
}
|
||||
progress("eino_agent_reply_stream_delta", subDelta, map[string]interface{}{
|
||||
progress("eino_agent_reply_stream_delta", subDelta, openai.WithSSEAccumulated(map[string]interface{}{
|
||||
"streamId": subReplyStreamID,
|
||||
"conversationId": conversationID,
|
||||
})
|
||||
}, subAssistantBuf))
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -743,13 +781,13 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
||||
"orchestration": orchMode,
|
||||
})
|
||||
}
|
||||
progress("response_delta", eofTail, map[string]interface{}{
|
||||
progress("response_delta", eofTail, openai.WithSSEAccumulated(map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"mcpExecutionIds": snapshotMCPIDs(),
|
||||
"einoRole": "orchestrator",
|
||||
"einoAgent": ev.AgentName,
|
||||
"orchestration": orchMode,
|
||||
})
|
||||
}, mainAssistantBuf))
|
||||
mainAssistWireAccum, _ = normalizeStreamingDelta(mainAssistWireAccum, eofTail)
|
||||
}
|
||||
}
|
||||
@@ -791,7 +829,7 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
||||
if merged := mergeStreamingToolCallFragments(toolStreamFragments); len(merged) > 0 {
|
||||
lastToolChunk = mergeMessageToolCalls(&schema.Message{ToolCalls: merged})
|
||||
}
|
||||
tryEmitToolCallsOnce(lastToolChunk, ev.AgentName, orchestratorName, conversationID, progress, toolEmitSeen, subAgentToolStep, markPending)
|
||||
tryEmitToolCallsOnce(lastToolChunk, ev.AgentName, orchestratorName, conversationID, orchMode, progress, toolEmitSeen, subAgentToolStep, mainAgentToolStep, markPending)
|
||||
// 流式路径此前只把 tool_calls 推给进度 UI,未写入 runAccumulatedMsgs;落库后 loadHistory→RepairOrphan 会删掉全部 tool 结果,表现为「续跑/下轮失忆」。
|
||||
if lastToolChunk != nil && len(lastToolChunk.ToolCalls) > 0 {
|
||||
runAccumulatedMsgs = append(runAccumulatedMsgs, schema.AssistantMessage("", lastToolChunk.ToolCalls))
|
||||
@@ -808,7 +846,7 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
||||
"einoRole": einoRoleTag(ev.AgentName),
|
||||
})
|
||||
}
|
||||
if retErr := handleRunErr(streamRecvErr); retErr != nil {
|
||||
if _, retErr := maybeRetryTransientRun(streamRecvErr); retErr != nil {
|
||||
return takePartial(retErr)
|
||||
}
|
||||
}
|
||||
@@ -820,7 +858,7 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
||||
continue
|
||||
}
|
||||
runAccumulatedMsgs = append(runAccumulatedMsgs, msg)
|
||||
tryEmitToolCallsOnce(mergeMessageToolCalls(msg), ev.AgentName, orchestratorName, conversationID, progress, toolEmitSeen, subAgentToolStep, markPending)
|
||||
tryEmitToolCallsOnce(mergeMessageToolCalls(msg), ev.AgentName, orchestratorName, conversationID, orchMode, progress, toolEmitSeen, subAgentToolStep, mainAgentToolStep, markPending)
|
||||
|
||||
if mv.Role == schema.Assistant {
|
||||
if progress != nil && strings.TrimSpace(msg.ReasoningContent) != "" {
|
||||
@@ -859,13 +897,13 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
||||
"einoAgent": ev.AgentName,
|
||||
"orchestration": orchMode,
|
||||
})
|
||||
progress("response_delta", body, map[string]interface{}{
|
||||
progress("response_delta", body, openai.WithSSEAccumulated(map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"mcpExecutionIds": snapshotMCPIDs(),
|
||||
"einoRole": "orchestrator",
|
||||
"einoAgent": ev.AgentName,
|
||||
"orchestration": orchMode,
|
||||
})
|
||||
}, body))
|
||||
}
|
||||
lastAssistant = body
|
||||
if orchMode == "plan_execute" && strings.EqualFold(strings.TrimSpace(ev.AgentName), "executor") {
|
||||
|
||||
@@ -13,12 +13,12 @@ import (
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/einomcp"
|
||||
"cyberstrike-ai/internal/openai"
|
||||
"cyberstrike-ai/internal/project"
|
||||
"cyberstrike-ai/internal/reasoning"
|
||||
|
||||
einoopenai "github.com/cloudwego/eino-ext/components/model/openai"
|
||||
"github.com/cloudwego/eino/adk"
|
||||
"github.com/cloudwego/eino/compose"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
@@ -39,6 +39,7 @@ func RunEinoSingleChatModelAgent(
|
||||
roleTools []string,
|
||||
progress func(eventType, message string, data interface{}),
|
||||
reasoningClient *reasoning.ClientIntent,
|
||||
systemPromptExtra string,
|
||||
) (*RunResult, error) {
|
||||
if appCfg == nil || ag == nil {
|
||||
return nil, fmt.Errorf("eino single: 配置或 Agent 为空")
|
||||
@@ -178,7 +179,8 @@ func RunEinoSingleChatModelAgent(
|
||||
},
|
||||
EmitInternalEvents: true,
|
||||
}
|
||||
ins := injectToolNamesOnlyInstruction(ctx, ag.EinoSingleAgentSystemInstruction(), mainTools, singleToolSearchActive)
|
||||
ins := project.AppendSystemPromptBlock(ag.EinoSingleAgentSystemInstruction(), systemPromptExtra)
|
||||
ins = injectToolNamesOnlyInstruction(ctx, ins, mainTools, singleToolSearchActive)
|
||||
if logger != nil {
|
||||
names := collectToolNames(ctx, mainTools)
|
||||
mountedNames := collectToolNames(ctx, mainToolsForCfg)
|
||||
@@ -213,7 +215,7 @@ func RunEinoSingleChatModelAgent(
|
||||
}
|
||||
|
||||
baseMsgs := historyToMessages(history, appCfg, &ma.EinoMiddleware)
|
||||
baseMsgs = append(baseMsgs, schema.UserMessage(userMessage))
|
||||
baseMsgs = appendUserMessageIfNeeded(baseMsgs, userMessage)
|
||||
|
||||
streamsMainAssistant := func(agent string) bool {
|
||||
return agent == "" || agent == einoSingleAgentName
|
||||
@@ -233,6 +235,8 @@ func RunEinoSingleChatModelAgent(
|
||||
StreamsMainAssistant: streamsMainAssistant,
|
||||
EinoRoleTag: einoRoleTag,
|
||||
CheckpointDir: ma.EinoMiddleware.CheckpointDir,
|
||||
RunRetryMaxAttempts: ma.EinoMiddleware.RunRetryMaxAttempts,
|
||||
RunRetryMaxBackoffSec: ma.EinoMiddleware.RunRetryMaxBackoffSec,
|
||||
McpIDsMu: &mcpIDsMu,
|
||||
McpIDs: &mcpIDs,
|
||||
FilesystemMonitorAgent: ag,
|
||||
|
||||
@@ -0,0 +1,173 @@
|
||||
package multiagent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
|
||||
"github.com/cloudwego/eino/adk"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultEinoRunRetryMaxAttempts = 10
|
||||
defaultEinoRunRetryMaxBackoff = 30 * time.Second
|
||||
)
|
||||
|
||||
// isEinoTransientRunError 判断 ADK 运行期错误是否适合指数退避续跑(429、5xx、网络抖动等)。
|
||||
// 用户取消、超时、迭代上限等由 run loop 单独处理,不在此列。
|
||||
func isEinoTransientRunError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
|
||||
return false
|
||||
}
|
||||
if isEinoIterationLimitError(err) {
|
||||
return false
|
||||
}
|
||||
msg := strings.ToLower(strings.TrimSpace(err.Error()))
|
||||
if msg == "" {
|
||||
return false
|
||||
}
|
||||
transientMarkers := []string{
|
||||
"406",
|
||||
"429",
|
||||
"too many requests",
|
||||
"rate limit",
|
||||
"rate_limit",
|
||||
"ratelimit",
|
||||
"quota exceeded",
|
||||
"overloaded",
|
||||
"capacity",
|
||||
"temporarily unavailable",
|
||||
"service unavailable",
|
||||
"bad gateway",
|
||||
"gateway timeout",
|
||||
"internal server error",
|
||||
"connection reset",
|
||||
"connection refused",
|
||||
"connection closed",
|
||||
"i/o timeout",
|
||||
"no such host",
|
||||
"network is unreachable",
|
||||
"broken pipe",
|
||||
"eof",
|
||||
"read tcp",
|
||||
"write tcp",
|
||||
"dial tcp",
|
||||
"tls handshake timeout",
|
||||
"stream error",
|
||||
"unexpected eof",
|
||||
"unexpected end of json",
|
||||
"status code: 406",
|
||||
"status code: 502",
|
||||
"502",
|
||||
"503",
|
||||
"504",
|
||||
"500",
|
||||
}
|
||||
for _, m := range transientMarkers {
|
||||
if strings.Contains(msg, m) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func einoRunRetryMaxAttempts(args *einoADKRunLoopArgs) int {
|
||||
if args != nil && args.RunRetryMaxAttempts > 0 {
|
||||
return args.RunRetryMaxAttempts
|
||||
}
|
||||
return defaultEinoRunRetryMaxAttempts
|
||||
}
|
||||
|
||||
// RunRetryMaxAttemptsFromConfig 供 handler 分段续跑计数(与 eino_middleware.run_retry_max_attempts 一致)。
|
||||
func RunRetryMaxAttemptsFromConfig(mw *config.MultiAgentEinoMiddlewareConfig) int {
|
||||
if mw != nil && mw.RunRetryMaxAttempts > 0 {
|
||||
return mw.RunRetryMaxAttempts
|
||||
}
|
||||
return defaultEinoRunRetryMaxAttempts
|
||||
}
|
||||
|
||||
// TransientRetryBackoff 供 handler 在分段续跑前退避。
|
||||
func TransientRetryBackoff(attempt int, maxBackoffSec int) time.Duration {
|
||||
max := defaultEinoRunRetryMaxBackoff
|
||||
if maxBackoffSec > 0 {
|
||||
max = time.Duration(maxBackoffSec) * time.Second
|
||||
}
|
||||
return einoTransientRetryBackoff(attempt, max)
|
||||
}
|
||||
|
||||
func einoRunRetryMaxBackoff(args *einoADKRunLoopArgs) time.Duration {
|
||||
if args != nil && args.RunRetryMaxBackoffSec > 0 {
|
||||
return time.Duration(args.RunRetryMaxBackoffSec) * time.Second
|
||||
}
|
||||
return defaultEinoRunRetryMaxBackoff
|
||||
}
|
||||
|
||||
// einoRunRestartContextSource 描述无 checkpoint Resume 时 Run 使用的消息来源(日志/SSE)。
|
||||
type einoRunRestartContextSource string
|
||||
|
||||
const (
|
||||
einoRestartContextInitial einoRunRestartContextSource = "initial"
|
||||
einoRestartContextAccumulated einoRunRestartContextSource = "accumulated"
|
||||
einoRestartContextModelTrace einoRunRestartContextSource = "model_trace"
|
||||
)
|
||||
|
||||
// einoMessagesForRunRestart 在退避后重新 Run 时选用最完整的上下文:
|
||||
// 1) ModelFacingTrace(与模型实际入参一致) 2) 事件流累积的 runAccumulatedMsgs 3) 初始 msgs。
|
||||
func einoMessagesForRunRestart(args *einoADKRunLoopArgs, baseMsgs, accumulated []adk.Message, baseCount int) ([]adk.Message, einoRunRestartContextSource) {
|
||||
if trace := persistTraceSource(args, nil); len(trace) > 0 {
|
||||
return append([]adk.Message(nil), trace...), einoRestartContextModelTrace
|
||||
}
|
||||
if len(accumulated) > baseCount {
|
||||
return append([]adk.Message(nil), accumulated...), einoRestartContextAccumulated
|
||||
}
|
||||
return append([]adk.Message(nil), baseMsgs...), einoRestartContextInitial
|
||||
}
|
||||
|
||||
// adkMessagesHasUserContent 从尾部向前查找,是否已有与 want 相同的 user 消息(避免重复 append)。
|
||||
func adkMessagesHasUserContent(msgs []adk.Message, want string) bool {
|
||||
want = strings.TrimSpace(want)
|
||||
if want == "" {
|
||||
return true
|
||||
}
|
||||
for i := len(msgs) - 1; i >= 0; i-- {
|
||||
m := msgs[i]
|
||||
if m == nil {
|
||||
continue
|
||||
}
|
||||
if m.Role == schema.User {
|
||||
return strings.TrimSpace(m.Content) == want
|
||||
}
|
||||
if m.Role == schema.Assistant || m.Role == schema.Tool {
|
||||
continue
|
||||
}
|
||||
break
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// appendUserMessageIfNeeded 在 history 轨迹之后追加本轮 user 消息(仅当轨迹中尚未包含该句)。
|
||||
func appendUserMessageIfNeeded(msgs []adk.Message, userMessage string) []adk.Message {
|
||||
if strings.TrimSpace(userMessage) == "" || adkMessagesHasUserContent(msgs, userMessage) {
|
||||
return msgs
|
||||
}
|
||||
return append(msgs, schema.UserMessage(userMessage))
|
||||
}
|
||||
|
||||
// einoTransientRetryBackoff 指数退避:2s, 4s, 8s… capped by maxBackoff。
|
||||
func einoTransientRetryBackoff(attempt int, maxBackoff time.Duration) time.Duration {
|
||||
if attempt < 0 {
|
||||
attempt = 0
|
||||
}
|
||||
backoff := time.Duration(1<<uint(attempt+1)) * time.Second
|
||||
if maxBackoff > 0 && backoff > maxBackoff {
|
||||
backoff = maxBackoff
|
||||
}
|
||||
return backoff
|
||||
}
|
||||
@@ -0,0 +1,104 @@
|
||||
package multiagent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/cloudwego/eino/adk"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
func TestIsEinoTransientRunError(t *testing.T) {
|
||||
t.Parallel()
|
||||
cases := []struct {
|
||||
name string
|
||||
err error
|
||||
want bool
|
||||
}{
|
||||
{"nil", nil, false},
|
||||
{"429", errors.New("HTTP 429 Too Many Requests"), true},
|
||||
{"rate limit", errors.New(`{"error":"rate limit exceeded"}`), true},
|
||||
{"connection reset", errors.New("read tcp: connection reset by peer"), true},
|
||||
{"503", errors.New("upstream returned 503"), true},
|
||||
{"iteration limit", errors.New("max iteration reached"), false},
|
||||
{"canceled", context.Canceled, false},
|
||||
{"deadline", context.DeadlineExceeded, false},
|
||||
{"auth", errors.New("invalid api key"), false},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
tc := tc
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
if got := isEinoTransientRunError(tc.err); got != tc.want {
|
||||
t.Fatalf("isEinoTransientRunError(%v) = %v, want %v", tc.err, got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestEinoTransientRetryBackoff(t *testing.T) {
|
||||
t.Parallel()
|
||||
max := 30 * time.Second
|
||||
if got := einoTransientRetryBackoff(0, max); got != 2*time.Second {
|
||||
t.Fatalf("attempt 0: got %v", got)
|
||||
}
|
||||
if got := einoTransientRetryBackoff(4, max); got != 30*time.Second {
|
||||
t.Fatalf("attempt 4 capped: got %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEinoMessagesForRunRestart(t *testing.T) {
|
||||
t.Parallel()
|
||||
base := []adk.Message{schema.UserMessage("hi")}
|
||||
acc := append([]adk.Message(nil), base...)
|
||||
acc = append(acc, schema.AssistantMessage("step1", nil))
|
||||
|
||||
got, src := einoMessagesForRunRestart(nil, base, acc, len(base))
|
||||
if src != einoRestartContextAccumulated || len(got) != 2 {
|
||||
t.Fatalf("accumulated: src=%s len=%d", src, len(got))
|
||||
}
|
||||
|
||||
holder := newModelFacingTraceHolder()
|
||||
holder.storeFromState(&adk.ChatModelAgentState{
|
||||
Messages: []adk.Message{schema.UserMessage("u"), schema.AssistantMessage("model-view", nil)},
|
||||
})
|
||||
got2, src2 := einoMessagesForRunRestart(&einoADKRunLoopArgs{ModelFacingTrace: holder}, base, acc, len(base))
|
||||
if src2 != einoRestartContextModelTrace || len(got2) != 2 {
|
||||
t.Fatalf("model trace: src=%s len=%d", src2, len(got2))
|
||||
}
|
||||
}
|
||||
|
||||
func TestEinoRunRetryMaxAttemptsFromArgs(t *testing.T) {
|
||||
t.Parallel()
|
||||
if einoRunRetryMaxAttempts(nil) != defaultEinoRunRetryMaxAttempts {
|
||||
t.Fatal("nil args should use default")
|
||||
}
|
||||
if einoRunRetryMaxAttempts(&einoADKRunLoopArgs{RunRetryMaxAttempts: 3}) != 3 {
|
||||
t.Fatal("custom max attempts")
|
||||
}
|
||||
if RunRetryMaxAttemptsFromConfig(nil) != defaultEinoRunRetryMaxAttempts {
|
||||
t.Fatal("config nil should use default")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAppendUserMessageIfNeeded(t *testing.T) {
|
||||
t.Parallel()
|
||||
msgs := []adk.Message{schema.UserMessage("old task")}
|
||||
out := appendUserMessageIfNeeded(msgs, "你好,你是谁")
|
||||
if len(out) != 2 || out[1].Content != "你好,你是谁" {
|
||||
t.Fatalf("should append user: len=%d", len(out))
|
||||
}
|
||||
dup := appendUserMessageIfNeeded(out, "你好,你是谁")
|
||||
if len(dup) != 2 {
|
||||
t.Fatalf("should not duplicate user message: len=%d", len(dup))
|
||||
}
|
||||
}
|
||||
|
||||
func TestErrTransientRetryContinue(t *testing.T) {
|
||||
t.Parallel()
|
||||
if !errors.Is(ErrTransientRetryContinue, ErrTransientRetryContinue) {
|
||||
t.Fatal("sentinel should match")
|
||||
}
|
||||
}
|
||||
@@ -5,3 +5,7 @@ import "errors"
|
||||
// ErrInterruptContinue 作为 context.CancelCause 使用:用户选择「中断并继续」且当前无进行中的 MCP 工具时,
|
||||
// 取消当前推理/流式输出,并在同一会话任务内携带用户补充说明自动续跑下一轮(类似 Hermes 式人机回合)。
|
||||
var ErrInterruptContinue = errors.New("agent interrupt: continue with user-supplied context")
|
||||
|
||||
// ErrTransientRetryContinue 表示 Run 因 429/网络等临时错误结束,应由 handler 落库轨迹后
|
||||
// loadHistoryFromAgentTrace 再开下一轮 Run(与 ErrInterruptContinue 同级的「分段续跑」语义)。
|
||||
var ErrTransientRetryContinue = errors.New("agent transient: retry after persisting trace")
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"cyberstrike-ai/internal/agents"
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
"cyberstrike-ai/internal/project"
|
||||
)
|
||||
|
||||
// DefaultPlanExecuteOrchestratorInstruction 当未配置 plan_execute 专用 Markdown / YAML 时的内置主代理(规划/重规划侧)提示。
|
||||
@@ -106,16 +107,18 @@ func DefaultPlanExecuteOrchestratorInstruction() string {
|
||||
|
||||
当工具返回错误时,错误信息会包含在工具响应中,请仔细阅读并做出合理的决策。
|
||||
|
||||
## 漏洞记录
|
||||
## 项目黑板(事实)与漏洞记录(分离)
|
||||
|
||||
发现有效漏洞时,必须使用 ` + builtin.ToolRecordVulnerability + ` 记录:标题、描述、严重程度、类型、目标、证明(POC)、影响、修复建议。
|
||||
绑定项目时会自动注入黑板索引(fact_key + 摘要)。**摘要不足必须 ` + builtin.ToolGetProjectFact + `(fact_key) 取 body,禁止臆造。** 环境认知用 ` + builtin.ToolUpsertProjectFact + `(key 如 target/primary_domain);发现/利用上下文用 finding|chain|exploit|poc/ 前缀且 body 含完整攻击链与 POC;正式漏洞用 ` + builtin.ToolRecordVulnerability + `(记前可先 ` + builtin.ToolListVulnerabilities + ` 防重复,详情用 ` + builtin.ToolGetVulnerability + `);二者可各记一次。误报用 ` + builtin.ToolDeprecateProjectFact + `。漏洞查询默认仅当前项目(未绑项目则仅当前会话)。
|
||||
|
||||
严重程度:critical / high / medium / low / info。证明须含足够证据(请求响应、截图、命令输出等)。记录后可在授权范围内继续测试。
|
||||
` + project.FactRecordingGuidanceBlock() + `
|
||||
|
||||
严重程度:critical / high / medium / low / info。证明须含足够证据。
|
||||
|
||||
## 技能库(Skills)与知识库
|
||||
|
||||
- 技能包位于服务器 skills/ 目录(各子目录 SKILL.md,遵循 agentskills.io);知识库用于向量检索片段,Skills 为可执行工作流指令。
|
||||
- plan_execute 执行器通过 MCP 使用知识库与漏洞记录等;Skills 的渐进式加载在「多代理 / Eino DeepAgent」等模式中由内置 skill 工具完成(需 multi_agent.eino_skills)。
|
||||
- plan_execute 执行器通过 MCP 使用知识库、项目事实与漏洞记录等;Skills 的渐进式加载在「多代理 / Eino DeepAgent」等模式中由内置 skill 工具完成(需 multi_agent.eino_skills)。
|
||||
- 若需要完整 Skill 工作流而当前会话无 skill 工具,请在计划或对用户说明中建议切换多代理或 Eino 编排会话。
|
||||
|
||||
## 执行器对用户输出(重要)
|
||||
@@ -206,7 +209,7 @@ func DefaultSupervisorOrchestratorInstruction() string {
|
||||
- **委派优先**:可独立封装、需要专项上下文的子目标(枚举、验证、归纳、报告素材)优先 transfer 给匹配子代理,并在委派说明中写清:子目标、约束、期望交付物结构、证据要求。
|
||||
- **亲自执行**:仅当无合适专家、需全局衔接或子代理结果不足时,由你直接调用工具。
|
||||
- **汇总**:子代理输出是证据来源;你要对齐矛盾、补全上下文,给出统一结论与可复现验证步骤,避免机械拼接。
|
||||
- **漏洞**:有效漏洞应通过 ` + builtin.ToolRecordVulnerability + ` 记录(含 POC 与严重性:critical / high / medium / low / info)。
|
||||
- **事实与漏洞**:环境认知用 ` + builtin.ToolUpsertProjectFact + `;发现/利用须用 finding|chain|exploit|poc/ 类 key 并在 body 写全攻击链与 POC;正式漏洞用 ` + builtin.ToolRecordVulnerability + `,查询用 ` + builtin.ToolListVulnerabilities + ` / ` + builtin.ToolGetVulnerability + `;索引摘要不足时必须 ` + builtin.ToolGetProjectFact + ` 取详情。
|
||||
|
||||
## transfer 交接与防重复劳动
|
||||
|
||||
|
||||
@@ -17,6 +17,7 @@ import (
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/einomcp"
|
||||
"cyberstrike-ai/internal/openai"
|
||||
"cyberstrike-ai/internal/project"
|
||||
"cyberstrike-ai/internal/reasoning"
|
||||
|
||||
einoopenai "github.com/cloudwego/eino-ext/components/model/openai"
|
||||
@@ -64,6 +65,7 @@ func RunDeepAgent(
|
||||
agentsMarkdownDir string,
|
||||
orchestrationOverride string,
|
||||
reasoningClient *reasoning.ClientIntent,
|
||||
systemPromptExtra string,
|
||||
) (*RunResult, error) {
|
||||
if appCfg == nil || ma == nil || ag == nil {
|
||||
return nil, fmt.Errorf("multiagent: 配置或 Agent 为空")
|
||||
@@ -339,6 +341,7 @@ func RunDeepAgent(
|
||||
return nil, err
|
||||
}
|
||||
|
||||
orchInstruction = project.AppendSystemPromptBlock(orchInstruction, systemPromptExtra)
|
||||
orchInstruction = injectToolNamesOnlyInstruction(ctx, orchInstruction, mainTools, mainToolSearchActive)
|
||||
if logger != nil {
|
||||
mainNames := collectToolNames(ctx, mainTools)
|
||||
@@ -387,7 +390,8 @@ func RunDeepAgent(
|
||||
|
||||
// noNestedTaskMiddleware 必须在最外层(最先拦截),防止 skill 或其他中间件内部触发 task 调用绕过检测。
|
||||
deepHandlers := []adk.ChatModelAgentMiddleware{newNoNestedTaskMiddleware()}
|
||||
if mw := newTaskContextEnrichMiddleware(userMessage, history, ma.SubAgentUserContextMaxRunes); mw != nil {
|
||||
taskEnrichExtra := systemPromptExtra
|
||||
if mw := newTaskContextEnrichMiddleware(userMessage, history, ma.SubAgentUserContextMaxRunes, taskEnrichExtra); mw != nil {
|
||||
deepHandlers = append(deepHandlers, mw)
|
||||
}
|
||||
if len(mainOrchestratorPre) > 0 {
|
||||
@@ -538,7 +542,7 @@ func RunDeepAgent(
|
||||
}
|
||||
|
||||
baseMsgs := historyToMessages(history, appCfg, &ma.EinoMiddleware)
|
||||
baseMsgs = append(baseMsgs, schema.UserMessage(userMessage))
|
||||
baseMsgs = appendUserMessageIfNeeded(baseMsgs, userMessage)
|
||||
|
||||
streamsMainAssistant := func(agent string) bool {
|
||||
if orchMode == "plan_execute" {
|
||||
@@ -566,6 +570,8 @@ func RunDeepAgent(
|
||||
StreamsMainAssistant: streamsMainAssistant,
|
||||
EinoRoleTag: einoRoleTag,
|
||||
CheckpointDir: ma.EinoMiddleware.CheckpointDir,
|
||||
RunRetryMaxAttempts: ma.EinoMiddleware.RunRetryMaxAttempts,
|
||||
RunRetryMaxBackoffSec: ma.EinoMiddleware.RunRetryMaxBackoffSec,
|
||||
McpIDsMu: &mcpIDsMu,
|
||||
McpIDs: &mcpIDs,
|
||||
FilesystemMonitorAgent: ag,
|
||||
@@ -595,6 +601,13 @@ func chatToolCallsToSchema(tcs []agent.ToolCall) []schema.ToolCall {
|
||||
argsStr = string(b)
|
||||
}
|
||||
}
|
||||
// Some OpenAI-compatible gateways require `function.arguments` to exist
|
||||
// on every assistant tool_call message. When args are empty, omitempty may
|
||||
// drop the field during serialization and cause "missing field arguments"
|
||||
// on the next turn history replay.
|
||||
if strings.TrimSpace(argsStr) == "" {
|
||||
argsStr = "{}"
|
||||
}
|
||||
typ := tc.Type
|
||||
if typ == "" {
|
||||
typ = "function"
|
||||
@@ -737,12 +750,23 @@ func toolCallsRichSignature(msg *schema.Message) string {
|
||||
return base + "|" + strings.Join(parts, ";")
|
||||
}
|
||||
|
||||
func einoMainIterationKey(agentName, orchestratorName string) string {
|
||||
key := strings.TrimSpace(agentName)
|
||||
if key == "" {
|
||||
key = strings.TrimSpace(orchestratorName)
|
||||
}
|
||||
if key == "" {
|
||||
return "_main"
|
||||
}
|
||||
return key
|
||||
}
|
||||
|
||||
func tryEmitToolCallsOnce(
|
||||
msg *schema.Message,
|
||||
agentName, orchestratorName, conversationID string,
|
||||
agentName, orchestratorName, conversationID, orchMode string,
|
||||
progress func(string, string, interface{}),
|
||||
seen map[string]struct{},
|
||||
subAgentToolStep map[string]int,
|
||||
subAgentToolStep, mainAgentToolStep map[string]int,
|
||||
markPending func(toolCallPendingInfo),
|
||||
) {
|
||||
if msg == nil || len(msg.ToolCalls) == 0 || progress == nil || seen == nil {
|
||||
@@ -756,14 +780,14 @@ func tryEmitToolCallsOnce(
|
||||
return
|
||||
}
|
||||
seen[sig] = struct{}{}
|
||||
emitToolCallsFromMessage(msg, agentName, orchestratorName, conversationID, progress, subAgentToolStep, markPending)
|
||||
emitToolCallsFromMessage(msg, agentName, orchestratorName, conversationID, orchMode, progress, subAgentToolStep, mainAgentToolStep, markPending)
|
||||
}
|
||||
|
||||
func emitToolCallsFromMessage(
|
||||
msg *schema.Message,
|
||||
agentName, orchestratorName, conversationID string,
|
||||
agentName, orchestratorName, conversationID, orchMode string,
|
||||
progress func(string, string, interface{}),
|
||||
subAgentToolStep map[string]int,
|
||||
subAgentToolStep, mainAgentToolStep map[string]int,
|
||||
markPending func(toolCallPendingInfo),
|
||||
) {
|
||||
if msg == nil || len(msg.ToolCalls) == 0 || progress == nil {
|
||||
@@ -784,6 +808,22 @@ func emitToolCallsFromMessage(
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
})
|
||||
} else if mainAgentToolStep != nil {
|
||||
key := einoMainIterationKey(agentName, orchestratorName)
|
||||
mainAgentToolStep[key]++
|
||||
n := mainAgentToolStep[key]
|
||||
// 第 1 轮已在主代理进入时发出;此后每次工具批次对应新一轮 ReAct(与子代理按工具计步一致)。
|
||||
if n > 1 {
|
||||
progress("iteration", "", map[string]interface{}{
|
||||
"iteration": n,
|
||||
"einoScope": "main",
|
||||
"einoRole": "orchestrator",
|
||||
"einoAgent": agentName,
|
||||
"orchestration": orchMode,
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
})
|
||||
}
|
||||
}
|
||||
role := "orchestrator"
|
||||
if isSubToolRound {
|
||||
|
||||
@@ -30,8 +30,15 @@ type taskContextEnrichMiddleware struct {
|
||||
// newTaskContextEnrichMiddleware returns a middleware that enriches task
|
||||
// descriptions with user conversation context. Returns nil if disabled
|
||||
// (maxRunes < 0) or no user messages exist.
|
||||
func newTaskContextEnrichMiddleware(userMessage string, history []agent.ChatMessage, maxRunes int) adk.ChatModelAgentMiddleware {
|
||||
func newTaskContextEnrichMiddleware(userMessage string, history []agent.ChatMessage, maxRunes int, projectBlackboard string) adk.ChatModelAgentMiddleware {
|
||||
supplement := buildUserContextSupplement(userMessage, history, maxRunes)
|
||||
if bb := strings.TrimSpace(projectBlackboard); bb != "" {
|
||||
if supplement != "" {
|
||||
supplement += "\n\n## 项目黑板索引\n" + bb
|
||||
} else {
|
||||
supplement = "\n\n## 项目黑板索引\n" + bb
|
||||
}
|
||||
}
|
||||
if supplement == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -105,6 +105,7 @@ func TestTaskContextEnrichMiddleware_EnrichesTaskDescription(t *testing.T) {
|
||||
"继续测试",
|
||||
[]agent.ChatMessage{{Role: "user", Content: "http://8.163.32.73:8081 pikachu靶场"}},
|
||||
0,
|
||||
"",
|
||||
)
|
||||
if mw == nil {
|
||||
t.Fatal("expected non-nil middleware")
|
||||
@@ -149,7 +150,7 @@ func TestTaskContextEnrichMiddleware_EnrichesTaskDescription(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestTaskContextEnrichMiddleware_IgnoresNonTaskTools(t *testing.T) {
|
||||
mw := newTaskContextEnrichMiddleware("test", nil, 0)
|
||||
mw := newTaskContextEnrichMiddleware("test", nil, 0, "")
|
||||
if mw == nil {
|
||||
t.Fatal("expected non-nil middleware")
|
||||
}
|
||||
@@ -175,7 +176,7 @@ func TestTaskContextEnrichMiddleware_IgnoresNonTaskTools(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestTaskContextEnrichMiddleware_NilWhenDisabled(t *testing.T) {
|
||||
mw := newTaskContextEnrichMiddleware("test", nil, -1)
|
||||
mw := newTaskContextEnrichMiddleware("test", nil, -1, "")
|
||||
if mw != nil {
|
||||
t.Error("middleware should be nil when disabled")
|
||||
}
|
||||
|
||||
@@ -0,0 +1,20 @@
|
||||
package openai
|
||||
|
||||
// SSEAccumulatedKey 为 SSE progress 事件 data 中的服务端权威流式全文快照字段。
|
||||
// 前端应优先用该字段更新 buffer,避免对 delta 二次 normalize 导致叠字。
|
||||
const SSEAccumulatedKey = "accumulated"
|
||||
|
||||
// WithSSEAccumulated 在 progress data 中附带当前流式累计全文(权威快照)。
|
||||
func WithSSEAccumulated(data map[string]interface{}, accumulated string) map[string]interface{} {
|
||||
if data == nil {
|
||||
data = make(map[string]interface{}, 1)
|
||||
}
|
||||
data[SSEAccumulatedKey] = accumulated
|
||||
return data
|
||||
}
|
||||
|
||||
// NormalizeStreamingDelta 将可能是“累计片段/重发片段”的内容归一化为“纯增量”。
|
||||
// 与 unexported normalizeStreamingDelta 相同,供 agent / multiagent 等包在发 SSE 前累计正文。
|
||||
func NormalizeStreamingDelta(current, incoming string) (next, delta string) {
|
||||
return normalizeStreamingDelta(current, incoming)
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
package project
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
)
|
||||
|
||||
// AppendSystemPromptBlock 将附加块追加到 system prompt。
|
||||
func AppendSystemPromptBlock(base, block string) string {
|
||||
base = strings.TrimSpace(base)
|
||||
block = strings.TrimSpace(block)
|
||||
if block == "" {
|
||||
return base
|
||||
}
|
||||
if base == "" {
|
||||
return block
|
||||
}
|
||||
return base + "\n\n" + block
|
||||
}
|
||||
|
||||
// BuildFactIndexBlock 为 Agent 系统提示生成项目黑板索引(仅 key + summary,不含 body)。
|
||||
func BuildFactIndexBlock(db *database.DB, projectID string, cfg config.ProjectConfig) (string, error) {
|
||||
if db == nil || !cfg.Enabled {
|
||||
return "", nil
|
||||
}
|
||||
projectID = strings.TrimSpace(projectID)
|
||||
if projectID == "" {
|
||||
return "", nil
|
||||
}
|
||||
|
||||
proj, err := db.GetProject(projectID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
facts, err := db.ListProjectFactsForIndex(projectID, cfg.DefaultInjectDeprecated)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if len(facts) == 0 {
|
||||
return fmt.Sprintf("## 项目黑板索引(project: %s, id: %s)\n(暂无事实)\n需要写入请使用 upsert_project_fact;需要详情请调用 get_project_fact(fact_key)。", proj.Name, proj.ID), nil
|
||||
}
|
||||
|
||||
sort.SliceStable(facts, func(i, j int) bool {
|
||||
if facts[i].Pinned != facts[j].Pinned {
|
||||
return facts[i].Pinned
|
||||
}
|
||||
return facts[i].UpdatedAt.After(facts[j].UpdatedAt)
|
||||
})
|
||||
|
||||
maxRunes := cfg.FactIndexMaxRunesEffective()
|
||||
var b strings.Builder
|
||||
b.WriteString(fmt.Sprintf("## 项目黑板索引(project: %s, id: %s)\n", proj.Name, proj.ID))
|
||||
used := len([]rune(b.String()))
|
||||
omitted := 0
|
||||
|
||||
for _, f := range facts {
|
||||
line := fmt.Sprintf("- [%s] %s — %s (%s)\n", f.FactKey, f.Category, strings.TrimSpace(f.Summary), f.Confidence)
|
||||
lineRunes := len([]rune(line))
|
||||
if used+lineRunes > maxRunes {
|
||||
omitted++
|
||||
continue
|
||||
}
|
||||
b.WriteString(line)
|
||||
used += lineRunes
|
||||
}
|
||||
|
||||
if omitted > 0 {
|
||||
b.WriteString(fmt.Sprintf("\n(另有 %d 条未列入索引,请使用 list_project_facts 或 search_project_facts 查询。)\n", omitted))
|
||||
}
|
||||
b.WriteString("需要完整内容(攻击链、POC、请求响应等)时必须调用 get_project_fact(fact_key),禁止凭摘要臆造细节。\n")
|
||||
b.WriteString("写入事实时:summary 写「什么+在哪+如何验证」;body 写可复现全流程(发现/利用类 fact_key 建议 finding|chain|exploit|poc/ 前缀)。\n")
|
||||
return b.String(), nil
|
||||
}
|
||||
@@ -0,0 +1,140 @@
|
||||
package project
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// 事实 category 常量(写入 upsert_project_fact 的 category 字段)。
|
||||
const (
|
||||
FactCategoryTarget = "target"
|
||||
FactCategoryAuth = "auth"
|
||||
FactCategoryInfra = "infra"
|
||||
FactCategoryBusiness = "business"
|
||||
FactCategoryFinding = "finding"
|
||||
FactCategoryChain = "chain"
|
||||
FactCategoryExploit = "exploit"
|
||||
FactCategoryPOC = "poc"
|
||||
FactCategoryNote = "note"
|
||||
)
|
||||
|
||||
// RequiresAttackChainBody 判断该事实是否应携带可复现的攻击链 / exploit 详情(写在 body,非仅 summary)。
|
||||
func RequiresAttackChainBody(category, factKey string) bool {
|
||||
c := strings.ToLower(strings.TrimSpace(category))
|
||||
switch c {
|
||||
case FactCategoryFinding, FactCategoryChain, FactCategoryExploit, FactCategoryPOC, "vuln":
|
||||
return true
|
||||
}
|
||||
key := strings.ToLower(strings.TrimSpace(factKey))
|
||||
for _, prefix := range []string{"finding/", "chain/", "exploit/", "poc/"} {
|
||||
if strings.HasPrefix(key, prefix) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// IsSparseFactBody 攻击链类事实 body 过短或缺少关键段落时返回 true(软校验,不阻断写入)。
|
||||
func IsSparseFactBody(category, factKey, body string) bool {
|
||||
if !RequiresAttackChainBody(category, factKey) {
|
||||
return false
|
||||
}
|
||||
body = strings.TrimSpace(body)
|
||||
if body == "" {
|
||||
return true
|
||||
}
|
||||
lower := strings.ToLower(body)
|
||||
// 至少应包含可复现线索:步骤/请求/命令/代码块 之一
|
||||
hasSteps := strings.Contains(lower, "攻击链") || strings.Contains(lower, "## 攻击") ||
|
||||
strings.Contains(lower, "## exploit") || strings.Contains(lower, "## poc")
|
||||
hasHTTP := strings.Contains(lower, "```http") || strings.Contains(lower, "```bash") ||
|
||||
strings.Contains(lower, "curl ") || strings.Contains(lower, "get ") || strings.Contains(lower, "post ")
|
||||
hasReq := strings.Contains(lower, "请求") || strings.Contains(lower, "响应") || strings.Contains(lower, "payload")
|
||||
// 无攻击链/POC/请求等结构线索,视为仅结论性描述(不论长短)
|
||||
return !(hasSteps || hasHTTP || hasReq)
|
||||
}
|
||||
|
||||
// FactBodyTemplate 按 category 返回建议的 body Markdown 骨架(供 Agent 填入真实内容)。
|
||||
func FactBodyTemplate(category, factKey string) string {
|
||||
if RequiresAttackChainBody(category, factKey) {
|
||||
return attackChainFactBodyTemplate
|
||||
}
|
||||
return envFactBodyTemplate
|
||||
}
|
||||
|
||||
const attackChainFactBodyTemplate = `## 结论(可验证,一句话)
|
||||
<勿仅写「存在漏洞」;写明类型 + 位置 + 触发条件>
|
||||
|
||||
## 目标与入口
|
||||
- 目标: <URL / IP:Port / 主机名>
|
||||
- 入口: <路径 / 接口 / 参数>
|
||||
- 前置条件: <匿名 / 角色 / Cookie / 其他依赖>
|
||||
|
||||
## 攻击链(逐步可复现)
|
||||
1. <侦察/发现>
|
||||
2. <利用/触发>
|
||||
3. <影响证明(读文件、RCE 回显、越权数据等)>
|
||||
|
||||
## Exploit / POC
|
||||
### 请求
|
||||
` + "```http\n<METHOD> <path> HTTP/1.1\nHost: ...\n...\n\n<body>\n```" + `
|
||||
|
||||
### 响应 / 现象
|
||||
<关键响应片段、状态码、差异点>
|
||||
|
||||
### 命令 / 脚本(如有)
|
||||
` + "```bash\n<command>\n```" + `
|
||||
|
||||
## 关键证据
|
||||
- <工具输出摘要 / 截图路径 / 会话或消息 ID>
|
||||
|
||||
## 关联
|
||||
- related_vulnerability_id: <可选,对应 record_vulnerability 的 id>
|
||||
- 依赖事实: <fact_key,如 auth/session_cookie>
|
||||
|
||||
## 备注与不确定性
|
||||
<待验证假设、环境差异、绕过尝试记录>`
|
||||
|
||||
const envFactBodyTemplate = `## 摘要
|
||||
<该事实的核心认知>
|
||||
|
||||
## 细节
|
||||
<端口/版本/路径/凭据特征/业务规则等>
|
||||
|
||||
## 来源与证据
|
||||
<命令输出、响应片段、发现时间>
|
||||
|
||||
## 关联
|
||||
- 相关 fact_key: <可选>`
|
||||
|
||||
// FactRecordingGuidanceBlock 写入系统提示:要求事实沉淀攻击链上下文而非仅结论。
|
||||
func FactRecordingGuidanceBlock() string {
|
||||
return `### 事实写入规范(审计复现 / 知识沉淀)
|
||||
|
||||
- **summary**:索引用一行,须含「什么 + 在哪 + 如何触发/验证」要点,禁止只写结论(如仅写「存在 SQLi」)。
|
||||
- **body**:完整可复现上下文,写入 ` + "`upsert_project_fact`" + ` 的 body 字段;索引不含 body,后续会话须靠 ` + "`get_project_fact`" + ` 取回。
|
||||
- **category / fact_key 建议**:
|
||||
- 环境认知:` + "`target/`" + `、` + "`auth/`" + `、` + "`infra/`" + `、` + "`business/`" + `(body 用环境模板即可)
|
||||
- 发现与利用:` + "`finding/`" + `、` + "`chain/`" + `、` + "`exploit/`" + `、` + "`poc/`" + `(**必须**用攻击链模板填满 body:入口、逐步攻击链、原始请求/响应或命令、证据、关联漏洞 ID)
|
||||
- **与漏洞记录分工**:` + "`record_vulnerability`" + ` 记可交付 findings;事实记**复现所需的全部上下文**(含失败尝试、绕过、依赖会话),二者可各记一次。
|
||||
- 更新同一发现时保持相同 ` + "`fact_key`" + ` 覆盖写入,勿散落多个 key 导致上下文丢失。`
|
||||
}
|
||||
|
||||
// SparseBodyWarning 攻击链类事实 body 不足时的工具返回提示(不阻断保存)。
|
||||
func SparseBodyWarning(category, factKey string) string {
|
||||
if !IsSparseFactBody(category, factKey, "") {
|
||||
return ""
|
||||
}
|
||||
return fmt.Sprintf(
|
||||
"\n\n⚠ 提示:category=%q / fact_key=%q 属于攻击链类事实,但 body 为空或过简。请补充完整攻击链与 POC(参考模板),便于后续审计复现。\n建议 body 骨架:\n%s",
|
||||
category, factKey, FactBodyTemplate(category, factKey),
|
||||
)
|
||||
}
|
||||
|
||||
// SparseBodyWarningIfNeeded 根据实际 body 判断是否追加警告。
|
||||
func SparseBodyWarningIfNeeded(category, factKey, body string) string {
|
||||
if !IsSparseFactBody(category, factKey, body) {
|
||||
return ""
|
||||
}
|
||||
return SparseBodyWarning(category, factKey)
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
package project
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestRequiresAttackChainBody(t *testing.T) {
|
||||
cases := []struct {
|
||||
cat, key string
|
||||
want bool
|
||||
}{
|
||||
{"finding", "note/misc", true},
|
||||
{"note", "finding/sqli-login", true},
|
||||
{"target", "target/primary_domain", false},
|
||||
{"auth", "auth/admin_cookie", false},
|
||||
{"chain", "x", true},
|
||||
{"", "exploit/rce-upload", true},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
if got := RequiresAttackChainBody(tc.cat, tc.key); got != tc.want {
|
||||
t.Errorf("RequiresAttackChainBody(%q,%q)=%v want %v", tc.cat, tc.key, got, tc.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsSparseFactBody(t *testing.T) {
|
||||
long := strings.Repeat("x", 150)
|
||||
if !IsSparseFactBody("finding", "finding/x", "") {
|
||||
t.Error("empty body should be sparse")
|
||||
}
|
||||
if !IsSparseFactBody("finding", "finding/x", long) {
|
||||
t.Error("body without repro clues should be sparse")
|
||||
}
|
||||
body := "## 攻击链\n1. step\n## Exploit\n```http\nGET / HTTP/1.1\n```\n"
|
||||
if IsSparseFactBody("finding", "finding/x", body) {
|
||||
t.Error("structured body should not be sparse")
|
||||
}
|
||||
if IsSparseFactBody("target", "target/x", "") {
|
||||
t.Error("env fact empty body is ok")
|
||||
}
|
||||
}
|
||||
@@ -149,13 +149,18 @@ func effectiveEffort(sr *config.OpenAIReasoningConfig, client *ClientIntent, all
|
||||
func normalizeEffort(s string) string {
|
||||
e := strings.ToLower(strings.TrimSpace(s))
|
||||
switch e {
|
||||
case "low", "medium", "high", "max":
|
||||
case "low", "medium", "high", "max", "xhigh":
|
||||
return e
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
// usesExtraFieldsReasoningEffort 为 Eino 无枚举的最高档 effort,经 ExtraFields 原样下发(max / xhigh 由网关自行识别,不做互转)。
|
||||
func usesExtraFieldsReasoningEffort(e string) bool {
|
||||
return e == "max" || e == "xhigh"
|
||||
}
|
||||
|
||||
func resolveWireProfile(oa *config.OpenAIConfig, sr *config.OpenAIReasoningConfig) wireProfile {
|
||||
if strings.EqualFold(strings.TrimSpace(oa.Provider), "claude") {
|
||||
return wireClaude
|
||||
@@ -210,11 +215,11 @@ func applyOpenAICompat(cfg *einoopenai.ChatModelConfig, mode, effort string) {
|
||||
if e == "" {
|
||||
return
|
||||
}
|
||||
if e == "max" {
|
||||
if usesExtraFieldsReasoningEffort(e) {
|
||||
if cfg.ExtraFields == nil {
|
||||
cfg.ExtraFields = make(map[string]any)
|
||||
}
|
||||
cfg.ExtraFields["reasoning_effort"] = "max"
|
||||
cfg.ExtraFields["reasoning_effort"] = effortStringForAPI(e)
|
||||
return
|
||||
}
|
||||
switch e {
|
||||
@@ -245,6 +250,6 @@ func applyOutputConfigEffort(cfg *einoopenai.ChatModelConfig, mode, effort strin
|
||||
}
|
||||
|
||||
func effortStringForAPI(e string) string {
|
||||
// Gateways expect lowercase strings; "max" kept as max.
|
||||
// 原样透传:OpenAI 官方多为 xhigh,部分兼容网关为 max,由配置/对话 effort 选择。
|
||||
return strings.ToLower(strings.TrimSpace(e))
|
||||
}
|
||||
|
||||
@@ -0,0 +1,66 @@
|
||||
package reasoning
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
|
||||
einoopenai "github.com/cloudwego/eino-ext/components/model/openai"
|
||||
)
|
||||
|
||||
func TestEffortStringForAPI_passthrough(t *testing.T) {
|
||||
cases := map[string]string{
|
||||
"max": "max",
|
||||
"xhigh": "xhigh",
|
||||
"HIGH": "high",
|
||||
"Medium": "medium",
|
||||
}
|
||||
for in, want := range cases {
|
||||
if got := effortStringForAPI(in); got != want {
|
||||
t.Fatalf("%q -> %q, want %q", in, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeEffort_maxAndXhigh(t *testing.T) {
|
||||
if normalizeEffort("xhigh") != "xhigh" {
|
||||
t.Fatal("xhigh not accepted")
|
||||
}
|
||||
if normalizeEffort("max") != "max" {
|
||||
t.Fatal("max not accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyOpenAICompat_xhighExtraField(t *testing.T) {
|
||||
cfg := &einoopenai.ChatModelConfig{}
|
||||
oa := &config.OpenAIConfig{
|
||||
Reasoning: config.OpenAIReasoningConfig{
|
||||
Profile: "openai_compat",
|
||||
Mode: "on",
|
||||
Effort: "xhigh",
|
||||
},
|
||||
}
|
||||
ApplyToEinoChatModelConfig(cfg, oa, nil)
|
||||
if cfg.ExtraFields == nil {
|
||||
t.Fatal("expected ExtraFields")
|
||||
}
|
||||
if got, _ := cfg.ExtraFields["reasoning_effort"].(string); got != "xhigh" {
|
||||
t.Fatalf("reasoning_effort=%q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyOpenAICompat_maxPassthrough(t *testing.T) {
|
||||
cfg := &einoopenai.ChatModelConfig{}
|
||||
oa := &config.OpenAIConfig{
|
||||
Reasoning: config.OpenAIReasoningConfig{
|
||||
Profile: "openai_compat",
|
||||
Mode: "on",
|
||||
Effort: "max",
|
||||
},
|
||||
}
|
||||
ApplyToEinoChatModelConfig(cfg, oa, nil)
|
||||
got, _ := cfg.ExtraFields["reasoning_effort"].(string)
|
||||
if got != "max" {
|
||||
t.Fatalf("max effort wire=%q, want max", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,316 @@
|
||||
package ilink
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
DefaultBaseURL = "https://ilinkai.weixin.qq.com"
|
||||
DefaultBotType = "3"
|
||||
DefaultBotAgent = "CyberStrikeAI/1.0"
|
||||
ILinkAppID = "bot"
|
||||
QRLongPollTimeout = 35 * time.Second
|
||||
APIDefaultTimeout = 15 * time.Second
|
||||
GetUpdatesTimeout = 35 * time.Second
|
||||
)
|
||||
|
||||
// Client 微信 iLink Bot HTTP 客户端(与 @tencent-weixin/openclaw-weixin 协议兼容)
|
||||
type Client struct {
|
||||
BaseURL string
|
||||
BotToken string
|
||||
BotAgent string
|
||||
ClientVersion uint32
|
||||
HTTP *http.Client
|
||||
}
|
||||
|
||||
func NewClient(baseURL, botToken, botAgent string, clientVersion uint32) *Client {
|
||||
base := strings.TrimSpace(baseURL)
|
||||
if base == "" {
|
||||
base = DefaultBaseURL
|
||||
}
|
||||
agent := strings.TrimSpace(botAgent)
|
||||
if agent == "" {
|
||||
agent = DefaultBotAgent
|
||||
}
|
||||
return &Client{
|
||||
BaseURL: strings.TrimRight(base, "/"),
|
||||
BotToken: strings.TrimSpace(botToken),
|
||||
BotAgent: sanitizeBotAgent(agent),
|
||||
ClientVersion: clientVersion,
|
||||
HTTP: &http.Client{Timeout: 0},
|
||||
}
|
||||
}
|
||||
|
||||
// BuildClientVersion 将 semver 编码为 iLink-App-ClientVersion(0x00MMNNPP)
|
||||
func BuildClientVersion(version string) uint32 {
|
||||
parts := strings.Split(version, ".")
|
||||
parse := func(i int) int {
|
||||
if i >= len(parts) {
|
||||
return 0
|
||||
}
|
||||
n, _ := strconv.Atoi(strings.TrimSpace(parts[i]))
|
||||
if n < 0 {
|
||||
return 0
|
||||
}
|
||||
return n
|
||||
}
|
||||
major := parse(0) & 0xff
|
||||
minor := parse(1) & 0xff
|
||||
patch := parse(2) & 0xff
|
||||
return uint32((major << 16) | (minor << 8) | patch)
|
||||
}
|
||||
|
||||
type baseInfo struct {
|
||||
ChannelVersion string `json:"channel_version"`
|
||||
BotAgent string `json:"bot_agent"`
|
||||
}
|
||||
|
||||
func (c *Client) buildBaseInfo() baseInfo {
|
||||
return baseInfo{
|
||||
ChannelVersion: "1.0.0",
|
||||
BotAgent: c.BotAgent,
|
||||
}
|
||||
}
|
||||
|
||||
func randomWechatUIN() string {
|
||||
var b [4]byte
|
||||
_, _ = rand.Read(b[:])
|
||||
u := uint32(b[0])<<24 | uint32(b[1])<<16 | uint32(b[2])<<8 | uint32(b[3])
|
||||
return base64.StdEncoding.EncodeToString([]byte(strconv.FormatUint(uint64(u), 10)))
|
||||
}
|
||||
|
||||
func (c *Client) commonHeaders() http.Header {
|
||||
h := http.Header{}
|
||||
h.Set("iLink-App-Id", ILinkAppID)
|
||||
h.Set("iLink-App-ClientVersion", strconv.FormatUint(uint64(c.ClientVersion), 10))
|
||||
return h
|
||||
}
|
||||
|
||||
func (c *Client) authHeaders() http.Header {
|
||||
h := c.commonHeaders()
|
||||
h.Set("Content-Type", "application/json")
|
||||
h.Set("AuthorizationType", "ilink_bot_token")
|
||||
h.Set("X-WECHAT-UIN", randomWechatUIN())
|
||||
if c.BotToken != "" {
|
||||
h.Set("Authorization", "Bearer "+c.BotToken)
|
||||
}
|
||||
return h
|
||||
}
|
||||
|
||||
func (c *Client) endpointURL(path string) (string, error) {
|
||||
u, err := url.Parse(c.BaseURL + "/")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
ref, err := url.Parse(path)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return u.ResolveReference(ref).String(), nil
|
||||
}
|
||||
|
||||
func (c *Client) doRequest(ctx context.Context, method, path string, body []byte, headers http.Header, timeout time.Duration) ([]byte, error) {
|
||||
reqURL, err := c.endpointURL(path)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var bodyReader io.Reader
|
||||
if len(body) > 0 {
|
||||
bodyReader = bytes.NewReader(body)
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, method, reqURL, bodyReader)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for k, vs := range headers {
|
||||
for _, v := range vs {
|
||||
req.Header.Add(k, v)
|
||||
}
|
||||
}
|
||||
client := c.HTTP
|
||||
if client == nil {
|
||||
client = http.DefaultClient
|
||||
}
|
||||
if timeout > 0 {
|
||||
ctx2, cancel := context.WithTimeout(ctx, timeout)
|
||||
defer cancel()
|
||||
req = req.WithContext(ctx2)
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
raw, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return nil, fmt.Errorf("ilink %s %s: %d %s", method, path, resp.StatusCode, string(raw))
|
||||
}
|
||||
return raw, nil
|
||||
}
|
||||
|
||||
// QRCodeResponse 获取二维码响应
|
||||
type QRCodeResponse struct {
|
||||
QRCode string `json:"qrcode"`
|
||||
QRCodeImgContent string `json:"qrcode_img_content"`
|
||||
}
|
||||
|
||||
// GetBotQRCode 获取绑定二维码
|
||||
func (c *Client) GetBotQRCode(ctx context.Context, botType string, localTokenList []string) (*QRCodeResponse, error) {
|
||||
if strings.TrimSpace(botType) == "" {
|
||||
botType = DefaultBotType
|
||||
}
|
||||
body, _ := json.Marshal(map[string]interface{}{
|
||||
"local_token_list": localTokenList,
|
||||
})
|
||||
path := "ilink/bot/get_bot_qrcode?bot_type=" + url.QueryEscape(botType)
|
||||
raw, err := c.doRequest(ctx, http.MethodPost, path, body, c.authHeaders(), APIDefaultTimeout)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var out QRCodeResponse
|
||||
if err := json.Unmarshal(raw, &out); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &out, nil
|
||||
}
|
||||
|
||||
// QRStatusResponse 二维码状态轮询响应
|
||||
type QRStatusResponse struct {
|
||||
Status string `json:"status"`
|
||||
BotToken string `json:"bot_token"`
|
||||
ILinkBotID string `json:"ilink_bot_id"`
|
||||
ILinkUserID string `json:"ilink_user_id"`
|
||||
BaseURL string `json:"baseurl"`
|
||||
RedirectHost string `json:"redirect_host"`
|
||||
}
|
||||
|
||||
// GetQRCodeStatus 长轮询二维码扫码状态
|
||||
func (c *Client) GetQRCodeStatus(ctx context.Context, qrcode, verifyCode string) (*QRStatusResponse, error) {
|
||||
path := "ilink/bot/get_qrcode_status?qrcode=" + url.QueryEscape(qrcode)
|
||||
if verifyCode != "" {
|
||||
path += "&verify_code=" + url.QueryEscape(verifyCode)
|
||||
}
|
||||
raw, err := c.doRequest(ctx, http.MethodGet, path, nil, c.commonHeaders(), QRLongPollTimeout)
|
||||
if err != nil {
|
||||
if ctx.Err() != nil {
|
||||
return &QRStatusResponse{Status: "wait"}, nil
|
||||
}
|
||||
return &QRStatusResponse{Status: "wait"}, nil
|
||||
}
|
||||
var out QRStatusResponse
|
||||
if err := json.Unmarshal(raw, &out); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &out, nil
|
||||
}
|
||||
|
||||
// MessageItem 消息内容项
|
||||
type MessageItem struct {
|
||||
Type int `json:"type"`
|
||||
TextItem *struct {
|
||||
Text string `json:"text"`
|
||||
} `json:"text_item,omitempty"`
|
||||
}
|
||||
|
||||
// WeixinMessage 入站消息
|
||||
type WeixinMessage struct {
|
||||
FromUserID string `json:"from_user_id"`
|
||||
MessageType int `json:"message_type"`
|
||||
MessageState int `json:"message_state"`
|
||||
ItemList []MessageItem `json:"item_list"`
|
||||
ContextToken string `json:"context_token"`
|
||||
}
|
||||
|
||||
// GetUpdatesResponse 长轮询消息响应
|
||||
type GetUpdatesResponse struct {
|
||||
Ret int `json:"ret"`
|
||||
ErrCode int `json:"errcode"`
|
||||
ErrMsg string `json:"errmsg"`
|
||||
Msgs []WeixinMessage `json:"msgs"`
|
||||
GetUpdatesBuf string `json:"get_updates_buf"`
|
||||
LongPollingTimeoutMs int `json:"longpolling_timeout_ms"`
|
||||
}
|
||||
|
||||
// GetUpdates 长轮询获取新消息
|
||||
func (c *Client) GetUpdates(ctx context.Context, getUpdatesBuf string) (*GetUpdatesResponse, error) {
|
||||
body, _ := json.Marshal(map[string]interface{}{
|
||||
"get_updates_buf": getUpdatesBuf,
|
||||
"base_info": c.buildBaseInfo(),
|
||||
})
|
||||
raw, err := c.doRequest(ctx, http.MethodPost, "ilink/bot/getupdates", body, c.authHeaders(), GetUpdatesTimeout)
|
||||
if err != nil {
|
||||
if ctx.Err() != nil {
|
||||
return &GetUpdatesResponse{Ret: 0, GetUpdatesBuf: getUpdatesBuf}, nil
|
||||
}
|
||||
return &GetUpdatesResponse{Ret: 0, GetUpdatesBuf: getUpdatesBuf}, nil
|
||||
}
|
||||
var out GetUpdatesResponse
|
||||
if err := json.Unmarshal(raw, &out); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &out, nil
|
||||
}
|
||||
|
||||
// SendTextMessage 发送文本回复
|
||||
func (c *Client) SendTextMessage(ctx context.Context, toUserID, contextToken, text, clientID string) error {
|
||||
if clientID == "" {
|
||||
clientID = randomClientID()
|
||||
}
|
||||
payload := map[string]interface{}{
|
||||
"msg": map[string]interface{}{
|
||||
"to_user_id": toUserID,
|
||||
"client_id": clientID,
|
||||
"message_type": 2,
|
||||
"message_state": 2,
|
||||
"context_token": contextToken,
|
||||
"item_list": []map[string]interface{}{
|
||||
{"type": 1, "text_item": map[string]string{"text": text}},
|
||||
},
|
||||
},
|
||||
"base_info": c.buildBaseInfo(),
|
||||
}
|
||||
body, _ := json.Marshal(payload)
|
||||
_, err := c.doRequest(ctx, http.MethodPost, "ilink/bot/sendmessage", body, c.authHeaders(), APIDefaultTimeout)
|
||||
return err
|
||||
}
|
||||
|
||||
func randomClientID() string {
|
||||
var b [8]byte
|
||||
_, _ = rand.Read(b[:])
|
||||
return fmt.Sprintf("%x", b)
|
||||
}
|
||||
|
||||
func sanitizeBotAgent(raw string) string {
|
||||
raw = strings.TrimSpace(raw)
|
||||
if raw == "" {
|
||||
return DefaultBotAgent
|
||||
}
|
||||
if len(raw) > 256 {
|
||||
return raw[:256]
|
||||
}
|
||||
return raw
|
||||
}
|
||||
|
||||
// ExtractText 从消息中提取首条文本
|
||||
func ExtractText(msg WeixinMessage) string {
|
||||
for _, item := range msg.ItemList {
|
||||
if item.Type == 1 && item.TextItem != nil {
|
||||
return strings.TrimSpace(item.TextItem.Text)
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
@@ -0,0 +1,26 @@
|
||||
package ilink
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/skip2/go-qrcode"
|
||||
)
|
||||
|
||||
// QRCodeDataURL 将扫码内容(一般为 liteapp 链接)编码为 PNG data URL,供 Web 端展示。
|
||||
// qrcode_img_content 不是图片直链,不能用作 <img src>。
|
||||
func QRCodeDataURL(content string, size int) (string, error) {
|
||||
content = strings.TrimSpace(content)
|
||||
if content == "" {
|
||||
return "", fmt.Errorf("empty qr content")
|
||||
}
|
||||
if size <= 0 {
|
||||
size = 256
|
||||
}
|
||||
png, err := qrcode.Encode(content, qrcode.Medium, size)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return "data:image/png;base64," + base64.StdEncoding.EncodeToString(png), nil
|
||||
}
|
||||
@@ -0,0 +1,96 @@
|
||||
package robot
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/robot/ilink"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const (
|
||||
wechatReconnectInitial = 5 * time.Second
|
||||
wechatReconnectMax = 60 * time.Second
|
||||
wechatPlatform = "wechat"
|
||||
)
|
||||
|
||||
// StartWechat 启动微信 iLink 长轮询(无需公网回调),收到消息后调用 handler 并回复。
|
||||
func StartWechat(ctx context.Context, robotsCfg config.RobotsConfig, h MessageHandler, appVersion string, logger *zap.Logger) {
|
||||
cfg := robotsCfg.Wechat
|
||||
if !cfg.Enabled || cfg.BotToken == "" {
|
||||
return
|
||||
}
|
||||
go runWechatLoop(ctx, cfg, h, appVersion, logger)
|
||||
}
|
||||
|
||||
func runWechatLoop(ctx context.Context, cfg config.RobotWechatConfig, h MessageHandler, appVersion string, logger *zap.Logger) {
|
||||
backoff := wechatReconnectInitial
|
||||
for {
|
||||
err := runWechatPoll(ctx, cfg, h, appVersion, logger)
|
||||
if ctx.Err() != nil {
|
||||
logger.Info("微信 iLink 长轮询已按配置关闭")
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
logger.Warn("微信 iLink 长轮询异常,将自动重连", zap.Error(err), zap.Duration("retry_after", backoff))
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-time.After(backoff):
|
||||
if backoff < wechatReconnectMax {
|
||||
backoff *= 2
|
||||
if backoff > wechatReconnectMax {
|
||||
backoff = wechatReconnectMax
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func runWechatPoll(ctx context.Context, cfg config.RobotWechatConfig, h MessageHandler, appVersion string, logger *zap.Logger) error {
|
||||
client := ilink.NewClient(cfg.BaseURL, cfg.BotToken, cfg.BotAgent, ilink.BuildClientVersion(appVersion))
|
||||
buf := cfg.GetUpdatesBuf
|
||||
logger.Info("微信 iLink 长轮询已启动", zap.String("ilink_bot_id", cfg.ILinkBotID))
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
default:
|
||||
}
|
||||
resp, err := client.GetUpdates(ctx, buf)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if resp.ErrCode != 0 && resp.Ret != 0 {
|
||||
logger.Warn("微信 getUpdates 返回错误", zap.Int("errcode", resp.ErrCode), zap.String("errmsg", resp.ErrMsg))
|
||||
}
|
||||
if resp.GetUpdatesBuf != "" {
|
||||
buf = resp.GetUpdatesBuf
|
||||
}
|
||||
for _, msg := range resp.Msgs {
|
||||
if msg.MessageType != 1 {
|
||||
continue
|
||||
}
|
||||
text := ilink.ExtractText(msg)
|
||||
if text == "" {
|
||||
continue
|
||||
}
|
||||
userID := strings.TrimSpace(msg.FromUserID)
|
||||
if userID == "" {
|
||||
continue
|
||||
}
|
||||
logger.Info("微信收到消息", zap.String("from", userID), zap.String("content", text))
|
||||
reply := h.HandleMessage(wechatPlatform, userID, text)
|
||||
if strings.TrimSpace(reply) == "" {
|
||||
continue
|
||||
}
|
||||
if err := client.SendTextMessage(ctx, userID, msg.ContextToken, reply, ""); err != nil {
|
||||
logger.Warn("微信发送回复失败", zap.String("to", userID), zap.Error(err))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -4,7 +4,7 @@
|
||||
|
||||
### What it does
|
||||
|
||||
- Configure **Host / Port / Password** and choose **Single-Agent** or **Multi-Agent**
|
||||
- Configure **Host / Port / HTTPS / Password** and choose an agent mode
|
||||
- Click **Validate** to login (`POST /api/auth/login`) and verify token (`GET /api/auth/validate`)
|
||||
- Right-click any HTTP message in Burp and send it to CyberStrikeAI for **streaming web pentest**
|
||||
- Keep a **test history sidebar** (searchable) so you can revisit previous runs
|
||||
@@ -63,6 +63,7 @@ If you already have Gradle available, you can still use `build.gradle` to build.
|
||||
|
||||
### Notes
|
||||
|
||||
- This extension connects to your CyberStrikeAI server (default is `http://127.0.0.1:8080`).
|
||||
- Default connection is `https://127.0.0.1:8080` (**HTTPS** checked). Self-signed / local certs are trusted automatically (no import).
|
||||
- Uncheck **HTTPS** only if your server runs plain HTTP.
|
||||
- It uses **Bearer Token** authentication obtained from the configured password.
|
||||
|
||||
|
||||
@@ -81,7 +81,8 @@ cd plugins/burp-suite/cyberstrikeai-burp-extension
|
||||
2) 填写:
|
||||
- **Host**:例如 `127.0.0.1`
|
||||
- **Port**:例如 `8080`
|
||||
- **Password**:你的 CyberStrikeAI 登录密码(对应服务端 `config.yaml` 的 `auth.password`)
|
||||
- **HTTPS**:默认勾选(对接 `config.yaml` 中 `tls_enabled` / 自签证书);插件会自动信任本地自签证书,无需导入
|
||||
- **Password**:你的 CyberStrikeAI 登录密码(对应服务端 `auth.password`)
|
||||
- **Agent mode**:选择 `Single Agent` 或 `Multi Agent`
|
||||
3) 点击 **Validate**
|
||||
- 成功:状态显示 `OK (token saved)`
|
||||
@@ -94,8 +95,9 @@ cd plugins/burp-suite/cyberstrikeai-burp-extension
|
||||
|
||||
- **Validate 失败 / 401**
|
||||
- 确认密码是否正确(服务端 `auth.password`)
|
||||
- 确认 IP/端口是否能访问(例如浏览器能打开 `http://IP:PORT/`)
|
||||
- 若服务器启用了反向代理/HTTPS,需要把插件里 baseUrl 改成对应协议与端口(当前插件默认使用 `http://`)
|
||||
- 确认 IP/端口是否能访问(例如浏览器能打开 `https://IP:PORT/`)
|
||||
- 服务端启用 TLS 时勾选 **HTTPS**(默认已勾选);自签证书无需手动导入
|
||||
- 若仍为纯 HTTP 部署,取消勾选 **HTTPS**
|
||||
|
||||
- **选择 Multi Agent 后提示“多代理未启用”**
|
||||
- 服务端需要开启:`config.yaml` 中 `multi_agent.enabled: true`
|
||||
|
||||
BIN
Binary file not shown.
+52
-11
@@ -73,15 +73,34 @@ public class BurpExtender implements IBurpExtender, IContextMenuFactory {
|
||||
public void onEvent(String type, String message, String rawJson) {
|
||||
if (type == null) type = "";
|
||||
switch (type) {
|
||||
case "response_start":
|
||||
tab.appendProgressToRun(runId, "\n\n[主回复]\n");
|
||||
break;
|
||||
case "response_delta":
|
||||
case "eino_agent_reply_stream_delta":
|
||||
tab.appendFinalToRun(runId, message);
|
||||
if (message != null && !message.isEmpty()) {
|
||||
tab.appendFinalToRun(runId, message);
|
||||
}
|
||||
break;
|
||||
case "response":
|
||||
tab.appendFinalToRun(runId, "\n\n--- Final Response ---\n");
|
||||
tab.appendFinalToRun(runId, message);
|
||||
tab.setFinalResponse(runId, message);
|
||||
break;
|
||||
case "eino_agent_reply_stream_start":
|
||||
tab.appendProgressToRun(runId, "\n\n[子代理回复]\n");
|
||||
break;
|
||||
case "eino_agent_reply_stream_delta":
|
||||
if (message != null && !message.isEmpty()) {
|
||||
tab.appendProgressToRun(runId, message);
|
||||
}
|
||||
break;
|
||||
case "eino_agent_reply_stream_end":
|
||||
tab.appendProgressToRun(runId, "\n");
|
||||
break;
|
||||
case "eino_agent_reply":
|
||||
if (message != null && !message.isEmpty()) {
|
||||
tab.appendProgressToRun(runId, "\n\n[子代理回复]\n" + message + "\n");
|
||||
}
|
||||
break;
|
||||
case "progress":
|
||||
tab.appendProgressToRun(runId, "\n[progress] " + message + "\n");
|
||||
tab.setRunStatus(runId, "running");
|
||||
@@ -94,21 +113,40 @@ public class BurpExtender implements IBurpExtender, IContextMenuFactory {
|
||||
tab.appendProgressToRun(runId, "\n[error] " + message + "\n");
|
||||
tab.setRunStatus(runId, "error");
|
||||
break;
|
||||
case "reasoning_chain_stream_start":
|
||||
tab.appendProgressToRun(runId, "\n\n[推理过程]\n");
|
||||
break;
|
||||
case "reasoning_chain_stream_delta":
|
||||
if (message != null && !message.isEmpty()) {
|
||||
tab.appendProgressToRun(runId, message);
|
||||
}
|
||||
break;
|
||||
case "reasoning_chain_stream_end":
|
||||
tab.appendProgressToRun(runId, "\n");
|
||||
break;
|
||||
case "reasoning_chain":
|
||||
if (message != null && !message.isEmpty()) {
|
||||
String streamId = rawJson != null ? SimpleJson.extractStringField(rawJson, "streamId") : "";
|
||||
if (streamId == null || streamId.isEmpty()) {
|
||||
tab.appendProgressToRun(runId, "\n\n[推理过程]\n" + message + "\n");
|
||||
}
|
||||
}
|
||||
break;
|
||||
case "thinking_stream_start":
|
||||
if (tab.isShowDebugEvents()) {
|
||||
tab.resetThinkingStream(runId);
|
||||
}
|
||||
break;
|
||||
case "thinking_stream_delta":
|
||||
if (tab.isShowDebugEvents() && message != null && !message.isEmpty()) {
|
||||
tab.appendProgressToRun(runId, message);
|
||||
}
|
||||
break;
|
||||
case "tool_call":
|
||||
case "tool_result":
|
||||
case "tool_result_delta":
|
||||
if (tab.isShowDebugEvents() && message != null && !message.isEmpty()) {
|
||||
if ("thinking_stream_delta".equals(type)) {
|
||||
tab.appendThinkingDelta(runId, message);
|
||||
} else {
|
||||
tab.appendProgressToRun(runId, "\n[" + type + "] " + message + "\n");
|
||||
}
|
||||
tab.appendProgressToRun(runId, "\n[" + type + "] " + message + "\n");
|
||||
}
|
||||
break;
|
||||
case "conversation":
|
||||
@@ -125,7 +163,9 @@ public class BurpExtender implements IBurpExtender, IContextMenuFactory {
|
||||
case "done":
|
||||
break;
|
||||
default:
|
||||
if (tab.isShowDebugEvents() && message != null && !message.isEmpty()) {
|
||||
if (tab.isShowDebugEvents() && message != null && !message.isEmpty()
|
||||
&& !type.endsWith("_stream_delta") && !type.endsWith("_stream_start")
|
||||
&& !type.endsWith("_stream_end")) {
|
||||
tab.appendProgressToRun(runId, "\n[" + type + "] " + message + "\n");
|
||||
}
|
||||
break;
|
||||
@@ -134,8 +174,9 @@ public class BurpExtender implements IBurpExtender, IContextMenuFactory {
|
||||
|
||||
@Override
|
||||
public void onError(String message, Exception e) {
|
||||
tab.appendProgressToRun(runId, "\n[error] " + message + "\n");
|
||||
tab.setRunStatus(runId, "error");
|
||||
boolean cancelled = message != null && message.toLowerCase().contains("cancel");
|
||||
tab.appendProgressToRun(runId, cancelled ? "\n[info] " + message + "\n" : "\n[error] " + message + "\n");
|
||||
tab.setRunStatus(runId, cancelled ? "cancelled" : "error");
|
||||
callbacks.printError("CyberStrikeAI stream error: " + message);
|
||||
if (e != null) {
|
||||
callbacks.printError(e.toString());
|
||||
|
||||
+127
-11
@@ -2,17 +2,29 @@ package burp;
|
||||
|
||||
import java.io.BufferedReader;
|
||||
import java.io.IOException;
|
||||
import java.io.InterruptedIOException;
|
||||
import java.io.InputStream;
|
||||
import java.io.InputStreamReader;
|
||||
import java.io.OutputStream;
|
||||
import java.net.HttpURLConnection;
|
||||
import java.net.SocketTimeoutException;
|
||||
import java.net.URL;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.util.HashMap;
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
|
||||
final class CyberStrikeAIClient {
|
||||
|
||||
private static final int AUTH_CONNECT_TIMEOUT_MS = 4_000;
|
||||
private static final int AUTH_READ_TIMEOUT_MS = 5_000;
|
||||
/** login + validate 整段上限,避免两次读超时叠加拖到半分钟 */
|
||||
private static final int AUTH_OVERALL_TIMEOUT_MS = 10_000;
|
||||
private static final int DEFAULT_READ_TIMEOUT_MS = 15_000;
|
||||
|
||||
private final AtomicReference<HttpURLConnection> activeConnection = new AtomicReference<>();
|
||||
private final AtomicReference<Thread> activeThread = new AtomicReference<>();
|
||||
|
||||
static final class Config {
|
||||
final String baseUrl; // e.g. http://127.0.0.1:8080
|
||||
final String password;
|
||||
@@ -49,15 +61,97 @@ final class CyberStrikeAIClient {
|
||||
void onDone();
|
||||
}
|
||||
|
||||
boolean hasActiveRequest() {
|
||||
return activeConnection.get() != null;
|
||||
}
|
||||
|
||||
void cancelActiveRequest() {
|
||||
HttpURLConnection conn = activeConnection.getAndSet(null);
|
||||
if (conn != null) {
|
||||
try {
|
||||
conn.disconnect();
|
||||
} catch (Exception ignored) {
|
||||
}
|
||||
}
|
||||
Thread t = activeThread.getAndSet(null);
|
||||
if (t != null) {
|
||||
t.interrupt();
|
||||
}
|
||||
}
|
||||
|
||||
String loginAndValidate(Config cfg) throws IOException {
|
||||
String token = login(cfg.baseUrl, cfg.password);
|
||||
validate(cfg.baseUrl, token);
|
||||
return token;
|
||||
Thread worker = Thread.currentThread();
|
||||
java.util.Timer deadline = new java.util.Timer("CyberStrikeAI-AuthDeadline", true);
|
||||
deadline.schedule(new java.util.TimerTask() {
|
||||
@Override
|
||||
public void run() {
|
||||
worker.interrupt();
|
||||
cancelActiveRequest();
|
||||
}
|
||||
}, AUTH_OVERALL_TIMEOUT_MS);
|
||||
try {
|
||||
String token = login(cfg.baseUrl, cfg.password);
|
||||
if (Thread.interrupted()) {
|
||||
throw timeoutIOException();
|
||||
}
|
||||
validate(cfg.baseUrl, token);
|
||||
if (Thread.interrupted()) {
|
||||
throw timeoutIOException();
|
||||
}
|
||||
return token;
|
||||
} catch (SocketTimeoutException e) {
|
||||
throw timeoutIOException();
|
||||
} finally {
|
||||
deadline.cancel();
|
||||
}
|
||||
}
|
||||
|
||||
private static IOException timeoutIOException() {
|
||||
return new IOException("Connection timed out (~" + (AUTH_OVERALL_TIMEOUT_MS / 1000)
|
||||
+ "s). Check host/port and HTTPS checkbox.");
|
||||
}
|
||||
|
||||
private void trackConnection(HttpURLConnection conn) {
|
||||
activeThread.set(Thread.currentThread());
|
||||
activeConnection.set(conn);
|
||||
}
|
||||
|
||||
private void releaseConnection(HttpURLConnection conn) {
|
||||
if (activeConnection.compareAndSet(conn, null)) {
|
||||
activeThread.set(null);
|
||||
}
|
||||
}
|
||||
|
||||
private static boolean isCancelled(Throwable e) {
|
||||
if (e == null) {
|
||||
return Thread.currentThread().isInterrupted();
|
||||
}
|
||||
if (Thread.currentThread().isInterrupted()) {
|
||||
return true;
|
||||
}
|
||||
if (e instanceof InterruptedIOException) {
|
||||
return true;
|
||||
}
|
||||
if (e instanceof SocketTimeoutException) {
|
||||
return false;
|
||||
}
|
||||
Throwable cause = e.getCause();
|
||||
if (cause != null && cause != e) {
|
||||
return isCancelled(cause);
|
||||
}
|
||||
String msg = e.getMessage();
|
||||
return msg != null && (
|
||||
msg.toLowerCase().contains("cancel")
|
||||
|| msg.toLowerCase().contains("abort")
|
||||
|| msg.toLowerCase().contains("closed")
|
||||
);
|
||||
}
|
||||
|
||||
private String login(String baseUrl, String password) throws IOException {
|
||||
URL url = new URL(baseUrl + "/api/auth/login");
|
||||
HttpURLConnection conn = (HttpURLConnection) url.openConnection();
|
||||
HttpURLConnection conn = SslTrustAll.open(url, AUTH_CONNECT_TIMEOUT_MS, AUTH_READ_TIMEOUT_MS);
|
||||
trackConnection(conn);
|
||||
try {
|
||||
conn.setRequestMethod("POST");
|
||||
conn.setDoOutput(true);
|
||||
conn.setRequestProperty("Content-Type", "application/json");
|
||||
@@ -92,11 +186,16 @@ final class CyberStrikeAIClient {
|
||||
throw new IOException("Login response missing token. Check backend address and credentials.");
|
||||
}
|
||||
return token;
|
||||
} finally {
|
||||
releaseConnection(conn);
|
||||
}
|
||||
}
|
||||
|
||||
private void validate(String baseUrl, String token) throws IOException {
|
||||
URL url = new URL(baseUrl + "/api/auth/validate");
|
||||
HttpURLConnection conn = (HttpURLConnection) url.openConnection();
|
||||
HttpURLConnection conn = SslTrustAll.open(url, AUTH_CONNECT_TIMEOUT_MS, AUTH_READ_TIMEOUT_MS);
|
||||
trackConnection(conn);
|
||||
try {
|
||||
conn.setRequestMethod("GET");
|
||||
conn.setRequestProperty("Authorization", "Bearer " + token);
|
||||
int code = conn.getResponseCode();
|
||||
@@ -104,6 +203,9 @@ final class CyberStrikeAIClient {
|
||||
if (code < 200 || code >= 300) {
|
||||
throw new IOException("Validate failed (" + code + "): " + resp);
|
||||
}
|
||||
} finally {
|
||||
releaseConnection(conn);
|
||||
}
|
||||
}
|
||||
|
||||
void streamTest(Config cfg, String token, String message, StreamListener listener) {
|
||||
@@ -117,11 +219,12 @@ final class CyberStrikeAIClient {
|
||||
payload.put("orchestration", cfg.agentMode.orchestration);
|
||||
}
|
||||
|
||||
new Thread(() -> {
|
||||
Thread worker = new Thread(() -> {
|
||||
HttpURLConnection conn = null;
|
||||
try {
|
||||
URL url = new URL(urlStr);
|
||||
conn = (HttpURLConnection) url.openConnection();
|
||||
conn = SslTrustAll.open(url, AUTH_CONNECT_TIMEOUT_MS, 0);
|
||||
trackConnection(conn);
|
||||
conn.setRequestMethod("POST");
|
||||
conn.setDoOutput(true);
|
||||
conn.setRequestProperty("Content-Type", "application/json");
|
||||
@@ -142,6 +245,9 @@ final class CyberStrikeAIClient {
|
||||
try (BufferedReader br = new BufferedReader(new InputStreamReader(is, StandardCharsets.UTF_8))) {
|
||||
String line;
|
||||
while ((line = br.readLine()) != null) {
|
||||
if (Thread.currentThread().isInterrupted()) {
|
||||
break;
|
||||
}
|
||||
// SSE format: "data: {json}"
|
||||
if (line.startsWith("data:")) {
|
||||
String json = line.substring("data:".length()).trim();
|
||||
@@ -156,15 +262,25 @@ final class CyberStrikeAIClient {
|
||||
}
|
||||
}
|
||||
}
|
||||
listener.onDone();
|
||||
if (Thread.currentThread().isInterrupted()) {
|
||||
listener.onError("Cancelled.", null);
|
||||
} else {
|
||||
listener.onDone();
|
||||
}
|
||||
} catch (Exception e) {
|
||||
listener.onError(e.getMessage(), e);
|
||||
if (isCancelled(e)) {
|
||||
listener.onError("Cancelled.", e);
|
||||
} else {
|
||||
listener.onError(e.getMessage(), e);
|
||||
}
|
||||
} finally {
|
||||
if (conn != null) {
|
||||
releaseConnection(conn);
|
||||
conn.disconnect();
|
||||
}
|
||||
}
|
||||
}, "CyberStrikeAI-Stream").start();
|
||||
}, "CyberStrikeAI-Stream");
|
||||
worker.start();
|
||||
}
|
||||
|
||||
void cancelByConversationId(String baseUrl, String token, String conversationId) throws IOException {
|
||||
@@ -172,7 +288,7 @@ final class CyberStrikeAIClient {
|
||||
throw new IOException("Missing conversationId.");
|
||||
}
|
||||
URL url = new URL(baseUrl + "/api/agent-loop/cancel");
|
||||
HttpURLConnection conn = (HttpURLConnection) url.openConnection();
|
||||
HttpURLConnection conn = SslTrustAll.open(url, AUTH_CONNECT_TIMEOUT_MS, AUTH_READ_TIMEOUT_MS);
|
||||
conn.setRequestMethod("POST");
|
||||
conn.setDoOutput(true);
|
||||
conn.setRequestProperty("Content-Type", "application/json");
|
||||
|
||||
+130
-34
@@ -14,6 +14,7 @@ final class CyberStrikeAITab implements ITab {
|
||||
|
||||
private final JTextField hostField = new JTextField("127.0.0.1");
|
||||
private final JTextField portField = new JTextField("8080");
|
||||
private final JCheckBox useHttpsBox = new JCheckBox("HTTPS", true);
|
||||
private final JPasswordField passwordField = new JPasswordField();
|
||||
private final JComboBox<String> agentModeBox = new JComboBox<>(new String[]{
|
||||
"Native ReAct", "Eino Single (ADK)", "Deep (DeepAgent)", "Plan-Execute", "Supervisor"
|
||||
@@ -29,6 +30,10 @@ final class CyberStrikeAITab implements ITab {
|
||||
|
||||
private final JTextArea progressArea = new JTextArea();
|
||||
private final JTextArea finalRawArea = new JTextArea(); // raw final stream / final response
|
||||
private JScrollPane progressScrollPane;
|
||||
private JScrollPane finalRawScrollPane;
|
||||
/** 距底部在此像素内视为「跟随滚动」,否则用户上拉阅读时不抢滚动条 */
|
||||
private static final int SCROLL_FOLLOW_THRESHOLD_PX = 48;
|
||||
private final JEditorPane markdownPane = new JEditorPane("text/html", "");
|
||||
private final CardLayout outputCardsLayout = new CardLayout();
|
||||
private final JPanel outputCards = new JPanel(outputCardsLayout);
|
||||
@@ -41,6 +46,7 @@ final class CyberStrikeAITab implements ITab {
|
||||
|
||||
private final CyberStrikeAIClient client = new CyberStrikeAIClient();
|
||||
private final AtomicReference<String> tokenRef = new AtomicReference<>("");
|
||||
private final AtomicReference<Thread> validateThreadRef = new AtomicReference<>();
|
||||
|
||||
private final DefaultListModel<TestRun> testListModel = new DefaultListModel<>();
|
||||
private final JList<TestRun> testList = new JList<>(testListModel);
|
||||
@@ -107,6 +113,8 @@ final class CyberStrikeAITab implements ITab {
|
||||
row1.add(hostField);
|
||||
row1.add(new JLabel("Port"));
|
||||
row1.add(portField);
|
||||
useHttpsBox.setToolTipText("Use https:// for CyberStrikeAI (self-signed certs are trusted automatically)");
|
||||
row1.add(useHttpsBox);
|
||||
row1.add(new JLabel("Password"));
|
||||
row1.add(passwordField);
|
||||
row1.add(validateButton);
|
||||
@@ -186,15 +194,22 @@ final class CyberStrikeAITab implements ITab {
|
||||
configureTextArea(requestArea, false);
|
||||
configureTextArea(responseArea, false);
|
||||
|
||||
outputCards.add(new JScrollPane(finalRawArea), "raw");
|
||||
finalRawScrollPane = new JScrollPane(finalRawArea);
|
||||
finalRawScrollPane.setHorizontalScrollBarPolicy(ScrollPaneConstants.HORIZONTAL_SCROLLBAR_NEVER);
|
||||
finalRawScrollPane.getVerticalScrollBar().setUnitIncrement(16);
|
||||
outputCards.add(finalRawScrollPane, "raw");
|
||||
outputCards.add(new JScrollPane(markdownPane), "md");
|
||||
|
||||
outputRoot.add(buildOutputHeader(), BorderLayout.NORTH);
|
||||
outputRoot.add(buildOutputBody(), BorderLayout.CENTER);
|
||||
|
||||
rightTabs.addTab("Output", outputRoot);
|
||||
rightTabs.addTab("Request", new JScrollPane(requestArea));
|
||||
rightTabs.addTab("Response", new JScrollPane(responseArea));
|
||||
JScrollPane requestScroll = new JScrollPane(requestArea);
|
||||
requestScroll.setHorizontalScrollBarPolicy(ScrollPaneConstants.HORIZONTAL_SCROLLBAR_NEVER);
|
||||
rightTabs.addTab("Request", requestScroll);
|
||||
JScrollPane responseScroll = new JScrollPane(responseArea);
|
||||
responseScroll.setHorizontalScrollBarPolicy(ScrollPaneConstants.HORIZONTAL_SCROLLBAR_NEVER);
|
||||
rightTabs.addTab("Response", responseScroll);
|
||||
return rightTabs;
|
||||
}
|
||||
|
||||
@@ -210,12 +225,13 @@ final class CyberStrikeAITab implements ITab {
|
||||
}
|
||||
|
||||
private JComponent buildOutputBody() {
|
||||
JScrollPane progressScroll = new JScrollPane(progressArea);
|
||||
progressScroll.setBorder(BorderFactory.createTitledBorder("Progress"));
|
||||
progressScroll.getVerticalScrollBar().setUnitIncrement(16);
|
||||
progressScrollPane = new JScrollPane(progressArea);
|
||||
progressScrollPane.setBorder(BorderFactory.createTitledBorder("Progress"));
|
||||
progressScrollPane.setHorizontalScrollBarPolicy(ScrollPaneConstants.HORIZONTAL_SCROLLBAR_NEVER);
|
||||
progressScrollPane.getVerticalScrollBar().setUnitIncrement(16);
|
||||
|
||||
JPanel empty = new JPanel();
|
||||
progressContainer.add(progressScroll, "show");
|
||||
progressContainer.add(progressScrollPane, "show");
|
||||
progressContainer.add(empty, "hide");
|
||||
((CardLayout) progressContainer.getLayout()).show(progressContainer, "show");
|
||||
|
||||
@@ -259,10 +275,27 @@ final class CyberStrikeAITab implements ITab {
|
||||
return split;
|
||||
}
|
||||
|
||||
private static boolean isScrollNearBottom(JScrollPane scrollPane) {
|
||||
if (scrollPane == null) {
|
||||
return true;
|
||||
}
|
||||
JScrollBar bar = scrollPane.getVerticalScrollBar();
|
||||
int max = Math.max(0, bar.getMaximum() - bar.getVisibleAmount());
|
||||
return bar.getValue() >= max - SCROLL_FOLLOW_THRESHOLD_PX;
|
||||
}
|
||||
|
||||
private static void scrollPaneToBottom(JScrollPane scrollPane) {
|
||||
if (scrollPane == null) {
|
||||
return;
|
||||
}
|
||||
JScrollBar bar = scrollPane.getVerticalScrollBar();
|
||||
bar.setValue(bar.getMaximum());
|
||||
}
|
||||
|
||||
private static void configureTextArea(JTextArea area, boolean monospaced) {
|
||||
area.setEditable(false);
|
||||
area.setLineWrap(false);
|
||||
area.setWrapStyleWord(false);
|
||||
area.setLineWrap(true);
|
||||
area.setWrapStyleWord(true);
|
||||
if (monospaced) {
|
||||
area.setFont(new Font(Font.MONOSPACED, Font.PLAIN, 12));
|
||||
} else {
|
||||
@@ -381,24 +414,44 @@ final class CyberStrikeAITab implements ITab {
|
||||
|
||||
private void wireActions() {
|
||||
validateButton.addActionListener(e -> {
|
||||
validateButton.setEnabled(false);
|
||||
if ("Cancel".equals(validateButton.getText())) {
|
||||
cancelValidateInProgress();
|
||||
return;
|
||||
}
|
||||
validateButton.setText("Cancel");
|
||||
validateButton.setEnabled(true);
|
||||
stopButton.setEnabled(true);
|
||||
statusLabel.setText("Validating...");
|
||||
log("Validating connection...");
|
||||
new Thread(() -> {
|
||||
log("Validating connection... (max ~10s; click Cancel or Stop to abort)");
|
||||
Thread worker = new Thread(() -> {
|
||||
try {
|
||||
CyberStrikeAIClient.Config cfg = currentConfig();
|
||||
String token = client.loginAndValidate(cfg);
|
||||
if (Thread.currentThread().isInterrupted()) {
|
||||
return;
|
||||
}
|
||||
tokenRef.set(token);
|
||||
SwingUtilities.invokeLater(() -> statusLabel.setText("OK (token saved)"));
|
||||
log("Validation OK.");
|
||||
} catch (Exception ex) {
|
||||
tokenRef.set("");
|
||||
SwingUtilities.invokeLater(() -> statusLabel.setText("Failed: " + ex.getMessage()));
|
||||
log("Validation failed: " + ex.getMessage());
|
||||
if (Thread.currentThread().isInterrupted()) {
|
||||
SwingUtilities.invokeLater(() -> statusLabel.setText("Cancelled"));
|
||||
log("Validation cancelled.");
|
||||
} else {
|
||||
SwingUtilities.invokeLater(() -> statusLabel.setText("Failed: " + ex.getMessage()));
|
||||
log("Validation failed: " + ex.getMessage());
|
||||
}
|
||||
} finally {
|
||||
SwingUtilities.invokeLater(() -> validateButton.setEnabled(true));
|
||||
validateThreadRef.set(null);
|
||||
SwingUtilities.invokeLater(() -> {
|
||||
validateButton.setText("Validate");
|
||||
validateButton.setEnabled(true);
|
||||
});
|
||||
}
|
||||
}, "CyberStrikeAI-Validate").start();
|
||||
}, "CyberStrikeAI-Validate");
|
||||
validateThreadRef.set(worker);
|
||||
worker.start();
|
||||
});
|
||||
|
||||
clearButton.addActionListener(e -> {
|
||||
@@ -435,10 +488,23 @@ final class CyberStrikeAITab implements ITab {
|
||||
});
|
||||
|
||||
stopButton.addActionListener(e -> {
|
||||
if ("Cancel".equals(validateButton.getText())) {
|
||||
cancelValidateInProgress();
|
||||
return;
|
||||
}
|
||||
|
||||
String runId = selectedRunId;
|
||||
if (runId != null && client.hasActiveRequest()) {
|
||||
client.cancelActiveRequest();
|
||||
appendProgressToRun(runId, "\n[info] Stream stopped.\n");
|
||||
setRunStatus(runId, "cancelled");
|
||||
return;
|
||||
}
|
||||
|
||||
if (runId == null) return;
|
||||
TestRun run = runs.get(runId);
|
||||
if (run == null) return;
|
||||
|
||||
String token = getToken();
|
||||
if (token == null || token.trim().isEmpty()) {
|
||||
appendProgressToRun(runId, "\n[error] Not validated.\n");
|
||||
@@ -483,7 +549,8 @@ final class CyberStrikeAITab implements ITab {
|
||||
String host = hostField.getText().trim();
|
||||
String port = portField.getText().trim();
|
||||
String password = new String(passwordField.getPassword());
|
||||
String baseUrl = "http://" + host + ":" + port;
|
||||
String scheme = useHttpsBox.isSelected() ? "https" : "http";
|
||||
String baseUrl = scheme + "://" + host + ":" + port;
|
||||
int idx = agentModeBox.getSelectedIndex();
|
||||
CyberStrikeAIClient.AgentMode mode = (idx >= 0 && idx < AGENT_MODES.length)
|
||||
? AGENT_MODES[idx]
|
||||
@@ -567,10 +634,31 @@ final class CyberStrikeAITab implements ITab {
|
||||
run.progressBuffer.append(s);
|
||||
}
|
||||
if (runId.equals(selectedRunId)) {
|
||||
SwingUtilities.invokeLater(() -> {
|
||||
progressArea.append(s);
|
||||
progressArea.setCaretPosition(progressArea.getDocument().getLength());
|
||||
});
|
||||
SwingUtilities.invokeLater(() -> appendProgressUi(s, false));
|
||||
}
|
||||
}
|
||||
|
||||
private void appendProgressUi(String s, boolean forceFollow) {
|
||||
JScrollBar bar = progressScrollPane != null ? progressScrollPane.getVerticalScrollBar() : null;
|
||||
int scrollBefore = bar != null ? bar.getValue() : 0;
|
||||
boolean follow = forceFollow || isScrollNearBottom(progressScrollPane);
|
||||
progressArea.append(s);
|
||||
if (follow) {
|
||||
scrollPaneToBottom(progressScrollPane);
|
||||
} else if (bar != null) {
|
||||
bar.setValue(scrollBefore);
|
||||
}
|
||||
}
|
||||
|
||||
private void appendFinalUi(String s, boolean forceFollow) {
|
||||
JScrollBar bar = finalRawScrollPane != null ? finalRawScrollPane.getVerticalScrollBar() : null;
|
||||
int scrollBefore = bar != null ? bar.getValue() : 0;
|
||||
boolean follow = forceFollow || isScrollNearBottom(finalRawScrollPane);
|
||||
finalRawArea.append(s);
|
||||
if (follow) {
|
||||
scrollPaneToBottom(finalRawScrollPane);
|
||||
} else if (bar != null) {
|
||||
bar.setValue(scrollBefore);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -620,10 +708,7 @@ final class CyberStrikeAITab implements ITab {
|
||||
run.finalBuffer.append(s);
|
||||
}
|
||||
if (runId.equals(selectedRunId)) {
|
||||
SwingUtilities.invokeLater(() -> {
|
||||
finalRawArea.append(s);
|
||||
finalRawArea.setCaretPosition(finalRawArea.getDocument().getLength());
|
||||
});
|
||||
SwingUtilities.invokeLater(() -> appendFinalUi(s, false));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -656,9 +741,9 @@ final class CyberStrikeAITab implements ITab {
|
||||
}
|
||||
SwingUtilities.invokeLater(() -> {
|
||||
progressArea.setText(progress);
|
||||
progressArea.setCaretPosition(progressArea.getDocument().getLength());
|
||||
scrollPaneToBottom(progressScrollPane);
|
||||
finalRawArea.setText(fin);
|
||||
finalRawArea.setCaretPosition(finalRawArea.getDocument().getLength());
|
||||
scrollPaneToBottom(finalRawScrollPane);
|
||||
requestArea.setText(run.requestRaw == null ? "" : run.requestRaw);
|
||||
responseArea.setText(run.responseRaw == null ? "" : run.responseRaw);
|
||||
refreshOutputView();
|
||||
@@ -682,25 +767,36 @@ final class CyberStrikeAITab implements ITab {
|
||||
|
||||
void clearAndShowStreamHeader(String title) {
|
||||
SwingUtilities.invokeLater(() -> {
|
||||
progressArea.setText("");
|
||||
finalRawArea.setText(title + "\n\n");
|
||||
progressArea.setText("[*] " + title + "\n\n");
|
||||
finalRawArea.setText("");
|
||||
markdownPane.setText("");
|
||||
});
|
||||
}
|
||||
|
||||
// Legacy helpers kept for Validate logging
|
||||
void appendStreamLine(String s) {
|
||||
if (s == null) return;
|
||||
SwingUtilities.invokeLater(() -> {
|
||||
progressArea.append(s);
|
||||
progressArea.append("\n");
|
||||
progressArea.setCaretPosition(progressArea.getDocument().getLength());
|
||||
});
|
||||
SwingUtilities.invokeLater(() -> appendProgressUi(s + "\n", false));
|
||||
}
|
||||
|
||||
private void log(String s) {
|
||||
appendStreamLine("[*] " + s);
|
||||
}
|
||||
|
||||
private void cancelValidateInProgress() {
|
||||
client.cancelActiveRequest();
|
||||
Thread t = validateThreadRef.getAndSet(null);
|
||||
if (t != null) {
|
||||
t.interrupt();
|
||||
}
|
||||
SwingUtilities.invokeLater(() -> {
|
||||
statusLabel.setText("Cancelled");
|
||||
validateButton.setText("Validate");
|
||||
validateButton.setEnabled(true);
|
||||
});
|
||||
log("Validation cancelled.");
|
||||
}
|
||||
|
||||
private void applyFilter() {
|
||||
String q = searchField.getText();
|
||||
if (q == null) q = "";
|
||||
|
||||
@@ -0,0 +1,149 @@
|
||||
package burp;
|
||||
|
||||
import javax.net.ssl.HostnameVerifier;
|
||||
import javax.net.ssl.HttpsURLConnection;
|
||||
import javax.net.ssl.SSLSocketFactory;
|
||||
import javax.net.ssl.SSLContext;
|
||||
import javax.net.ssl.TrustManager;
|
||||
import javax.net.ssl.X509TrustManager;
|
||||
import java.io.IOException;
|
||||
import java.net.HttpURLConnection;
|
||||
import java.net.InetSocketAddress;
|
||||
import java.net.Socket;
|
||||
import java.net.URL;
|
||||
import java.security.cert.X509Certificate;
|
||||
|
||||
/**
|
||||
* Opens HTTPS connections without validating server certificates (self-signed / local dev).
|
||||
* Applied per-connection only; does not change JVM-wide defaults for other Burp components.
|
||||
*/
|
||||
final class SslTrustAll {
|
||||
|
||||
private static volatile SSLSocketFactory socketFactory;
|
||||
private static final HostnameVerifier TRUST_ALL_HOSTS = (hostname, session) -> true;
|
||||
|
||||
private SslTrustAll() {
|
||||
}
|
||||
|
||||
static HttpURLConnection open(URL url) throws IOException {
|
||||
return open(url, 5_000, 30_000);
|
||||
}
|
||||
|
||||
static HttpURLConnection open(URL url, int connectTimeoutMs, int readTimeoutMs) throws IOException {
|
||||
HttpURLConnection conn = (HttpURLConnection) url.openConnection();
|
||||
conn.setConnectTimeout(connectTimeoutMs);
|
||||
conn.setReadTimeout(readTimeoutMs);
|
||||
if (conn instanceof HttpsURLConnection) {
|
||||
HttpsURLConnection https = (HttpsURLConnection) conn;
|
||||
https.setSSLSocketFactory(new TimeoutSslSocketFactory(socketFactory(), connectTimeoutMs, readTimeoutMs));
|
||||
https.setHostnameVerifier(TRUST_ALL_HOSTS);
|
||||
}
|
||||
return conn;
|
||||
}
|
||||
|
||||
private static SSLSocketFactory socketFactory() {
|
||||
SSLSocketFactory sf = socketFactory;
|
||||
if (sf != null) {
|
||||
return sf;
|
||||
}
|
||||
synchronized (SslTrustAll.class) {
|
||||
sf = socketFactory;
|
||||
if (sf != null) {
|
||||
return sf;
|
||||
}
|
||||
try {
|
||||
TrustManager[] trustAll = new TrustManager[]{
|
||||
new X509TrustManager() {
|
||||
@Override
|
||||
public X509Certificate[] getAcceptedIssuers() {
|
||||
return new X509Certificate[0];
|
||||
}
|
||||
|
||||
@Override
|
||||
public void checkClientTrusted(X509Certificate[] chain, String authType) {
|
||||
}
|
||||
|
||||
@Override
|
||||
public void checkServerTrusted(X509Certificate[] chain, String authType) {
|
||||
}
|
||||
}
|
||||
};
|
||||
SSLContext ctx = SSLContext.getInstance("TLS");
|
||||
ctx.init(null, trustAll, new java.security.SecureRandom());
|
||||
sf = ctx.getSocketFactory();
|
||||
socketFactory = sf;
|
||||
return sf;
|
||||
} catch (Exception e) {
|
||||
throw new RuntimeException("Failed to initialize trust-all TLS", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/** Ensures TCP connect + socket read respect timeouts (plain HttpURLConnection SSL can hang longer). */
|
||||
private static final class TimeoutSslSocketFactory extends SSLSocketFactory {
|
||||
private final SSLSocketFactory delegate;
|
||||
private final int connectTimeoutMs;
|
||||
private final int readTimeoutMs;
|
||||
|
||||
TimeoutSslSocketFactory(SSLSocketFactory delegate, int connectTimeoutMs, int readTimeoutMs) {
|
||||
this.delegate = delegate;
|
||||
this.connectTimeoutMs = connectTimeoutMs;
|
||||
this.readTimeoutMs = readTimeoutMs;
|
||||
}
|
||||
|
||||
@Override
|
||||
public String[] getDefaultCipherSuites() {
|
||||
return delegate.getDefaultCipherSuites();
|
||||
}
|
||||
|
||||
@Override
|
||||
public String[] getSupportedCipherSuites() {
|
||||
return delegate.getSupportedCipherSuites();
|
||||
}
|
||||
|
||||
@Override
|
||||
public Socket createSocket() throws IOException {
|
||||
return tune(delegate.createSocket());
|
||||
}
|
||||
|
||||
@Override
|
||||
public Socket createSocket(Socket s, String host, int port, boolean autoClose) throws IOException {
|
||||
return tune(delegate.createSocket(s, host, port, autoClose));
|
||||
}
|
||||
|
||||
@Override
|
||||
public Socket createSocket(String host, int port) throws IOException {
|
||||
Socket plain = new Socket();
|
||||
plain.connect(new InetSocketAddress(host, port), connectTimeoutMs);
|
||||
return tune(delegate.createSocket(plain, host, port, true));
|
||||
}
|
||||
|
||||
@Override
|
||||
public Socket createSocket(String host, int port, java.net.InetAddress localHost, int localPort) throws IOException {
|
||||
Socket plain = new Socket();
|
||||
plain.bind(new InetSocketAddress(localHost, localPort));
|
||||
plain.connect(new InetSocketAddress(host, port), connectTimeoutMs);
|
||||
return tune(delegate.createSocket(plain, host, port, true));
|
||||
}
|
||||
|
||||
@Override
|
||||
public Socket createSocket(java.net.InetAddress host, int port) throws IOException {
|
||||
Socket plain = new Socket();
|
||||
plain.connect(new InetSocketAddress(host, port), connectTimeoutMs);
|
||||
return tune(delegate.createSocket(plain, host.getHostName(), port, true));
|
||||
}
|
||||
|
||||
@Override
|
||||
public Socket createSocket(java.net.InetAddress address, int port, java.net.InetAddress localAddress, int localPort) throws IOException {
|
||||
Socket plain = new Socket();
|
||||
plain.bind(new InetSocketAddress(localAddress, localPort));
|
||||
plain.connect(new InetSocketAddress(address, port), connectTimeoutMs);
|
||||
return tune(delegate.createSocket(plain, address.getHostName(), port, true));
|
||||
}
|
||||
|
||||
private Socket tune(Socket socket) throws IOException {
|
||||
socket.setSoTimeout(readTimeoutMs);
|
||||
return socket;
|
||||
}
|
||||
}
|
||||
}
|
||||
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user