mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-07-31 16:17:35 +02:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
ca0bcc21d4 | ||
|
|
c33e5f2026 | ||
|
|
0251230654 | ||
|
|
b76e06ff92 | ||
|
|
aca97ffc94 | ||
|
|
24052717bd | ||
|
|
478b52b011 | ||
|
|
6f8e324a75 | ||
|
|
1abfd3d22a | ||
|
|
e87011b081 | ||
|
|
e020ffed49 | ||
|
|
823fb47a81 | ||
|
|
46a9b42fde | ||
|
|
2995847e0e | ||
|
|
1b64f5d8a0 | ||
|
|
a145687508 | ||
|
|
49d2175872 | ||
|
|
cd87a0b965 | ||
|
|
3bfa846db1 | ||
|
|
bc3246d157 | ||
|
|
542d1d2411 | ||
|
|
0ff8c58fbd | ||
|
|
25a76a8c97 | ||
|
|
a2744d6936 | ||
|
|
225d9ba0d9 | ||
|
|
9dfb9d6c68 | ||
|
|
04efc50161 | ||
|
|
46f891a114 | ||
|
|
bd80010cda | ||
|
|
a1201240d4 | ||
|
|
5266e2d95f | ||
|
|
1af455b762 | ||
|
|
8da02f94ac | ||
|
|
7f1bc1d229 | ||
|
|
0e61c17db3 | ||
|
|
2f1a95e8bf | ||
|
|
878cf158b0 | ||
|
|
e092f6c590 | ||
|
|
deafabc276 | ||
|
|
87c83c92fd | ||
|
|
898f69a388 | ||
|
|
0fa739d438 | ||
|
|
a7c3d28cbc | ||
|
|
1c576db75d | ||
|
|
56ee11a97a | ||
|
|
95da42effb | ||
|
|
64c3f28780 | ||
|
|
c8fece70ef | ||
|
|
46ecb2f13b | ||
|
|
6484950be3 | ||
|
|
f68d69d562 | ||
|
|
798a76ec3c | ||
|
|
37e553ba8a | ||
|
|
16a854d0f5 | ||
|
|
c57b681a19 | ||
|
|
f405e02d70 | ||
|
|
0531ee7292 | ||
|
|
69e8d7020c | ||
|
|
526e8626e9 | ||
|
|
cb52ff37f5 | ||
|
|
461682acaf | ||
|
|
7522bab98c |
+44
@@ -0,0 +1,44 @@
|
||||
# Runtime data
|
||||
data/
|
||||
*.db
|
||||
*.db-shm
|
||||
*.db-wal
|
||||
*.sqlite
|
||||
*.sqlite3
|
||||
|
||||
# Local configuration and secrets
|
||||
config.yaml
|
||||
config.local.yaml
|
||||
.env
|
||||
.env.*
|
||||
*.pem
|
||||
*.key
|
||||
*.crt
|
||||
|
||||
# Build outputs
|
||||
cyberstrike-ai
|
||||
bin/
|
||||
dist/
|
||||
build/
|
||||
coverage.out
|
||||
coverage.html
|
||||
|
||||
# Logs and temporary files
|
||||
*.log
|
||||
tmp/
|
||||
temp/
|
||||
|
||||
# Python
|
||||
venv/
|
||||
.venv/
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
.pytest_cache/
|
||||
|
||||
# Go
|
||||
vendor/
|
||||
|
||||
# macOS / editors
|
||||
.DS_Store
|
||||
.idea/
|
||||
.vscode/
|
||||
@@ -127,6 +127,7 @@ If CyberStrikeAI helps you, you can support the project via **WeChat Pay** or **
|
||||
- 🔗 Attack-chain intelligence with graph views, risk scoring, project facts, and step-by-step replay
|
||||
- 🧑⚖️ Human-in-the-loop governance with approval modes, allowlists, audit-agent review, and traceable decisions
|
||||
- 🔒 Password-protected web UI, audit logs, SQLite persistence, and operational evidence retention
|
||||
- 🔐 **Platform RBAC** with multi-user accounts, system/custom roles, per-permission scopes (`all` / `assigned` / `own`), ownership, and explicit assignments enforced across APIs, Agents, MCP, background jobs, and chatbots; see the [RBAC administration guide](docs/en-US/rbac.md)
|
||||
- 📚 Knowledge base (RAG): **Eino MultiQuery** query rewrite + multi-path vector retrieval + **HTTP rerank** (DashScope `gte-rerank` / Cohere-compatible) + post-processing (dedupe, budget); **Eino Compose** indexing pipeline
|
||||
- 📁 Conversation grouping with pinning, rename, and batch management
|
||||
- 📂 **Project management**: shared facts (blackboard) across sessions, `upsert_project_fact` + `links` to chain paths; attack-chain and project fact graph views
|
||||
@@ -149,6 +150,11 @@ CyberStrikeAI includes optional integrations under `plugins/`.
|
||||
- **Burp Suite extension**: `plugins/burp-suite/cyberstrikeai-burp-extension/`
|
||||
Build output: `plugins/burp-suite/cyberstrikeai-burp-extension/dist/cyberstrikeai-burp-extension.jar`
|
||||
Docs: `plugins/burp-suite/cyberstrikeai-burp-extension/README.md`
|
||||
- **Browser extension (Chrome / Edge)**: `plugins/browser-extension/cyberstrikeai-browser-extension/`
|
||||
Capture Network traffic in DevTools and send it to CyberStrikeAI for AI-assisted security testing—aligned with the Burp plugin.
|
||||
Install: `chrome://extensions/` → Load unpacked → F12 → **CyberStrikeAI** tab
|
||||
Package output: `plugins/browser-extension/cyberstrikeai-browser-extension/dist/cyberstrikeai-browser-extension.zip`
|
||||
Docs: `plugins/browser-extension/cyberstrikeai-browser-extension/README.md` / `README.zh-CN.md`
|
||||
|
||||
## Tool Overview
|
||||
|
||||
@@ -644,10 +650,11 @@ enabled: true
|
||||
- [API recipes](docs/en-US/api-recipes.md): examples for login, Agent, streaming, multi-agent, uploads, vulnerabilities, KB, and audit export.
|
||||
- [Configuration reference](docs/en-US/configuration.md): main `config.yaml` sections, recommended values, and update guidance.
|
||||
- [Security model](docs/en-US/security-model.md): authentication, tool execution, HITL, audit, C2/WebShell, and data safety boundaries.
|
||||
- [RBAC administration](docs/en-US/rbac.md): platform users, system/custom roles, permission catalog, per-permission scopes, resource assignments, Agent/MCP/robot boundaries, and API examples.
|
||||
- [API reference](docs/en-US/api-reference.md): OpenAPI, authentication, Agent, projects, knowledge base, C2, WebShell, and other API entry points.
|
||||
- [Multi-agent mode (Eino)](docs/en-US/MULTI_AGENT_EINO.md): **Deep**, **Plan-Execute**, **Supervisor**, `agents/*.md`, `eino_skills` / `eino_middleware`, APIs, and chat/stream behavior.
|
||||
- [Graph orchestration guide](docs/en-US/workflow-graph.md): visual workflow design, node configuration, `previous` / `outputs` variable passing, and role binding.
|
||||
- [Robot / Chatbot guide](docs/en-US/robot.md): Setup, commands, and troubleshooting for WeChat, WeCom, DingTalk, Lark, Telegram, Slack, Discord, and QQ Bot.
|
||||
- [Robot / Chatbot guide](docs/en-US/robot.md): Platform setup, RBAC user-binding/service-account modes, sender allowlists, commands, verification, and troubleshooting.
|
||||
- [HITL best practices](docs/en-US/hitl-best-practices.md): reviewer modes, allowlists, Audit Agent prompts, and separate small-model configuration.
|
||||
|
||||
## Project Layout
|
||||
|
||||
+8
-1
@@ -126,6 +126,7 @@ CyberStrikeAI 基于 Go 构建,为 AI 原生安全运营提供完整底座:1
|
||||
- 🔗 攻击链智能分析,支持图谱视图、风险打分、项目事实沉淀与步骤回放
|
||||
- 🧑⚖️ 人机协同治理,支持审批模式、免审批白名单、审计 Agent 复核与可追溯决策
|
||||
- 🔒 Web 登录保护、审计日志、SQLite 持久化与行动证据留存
|
||||
- 🔐 **平台 RBAC**:支持多用户、系统/自定义角色、逐权限 Scope(`all` / `assigned` / `own`)、资源归属与显式授权,并统一约束 API、Agent、MCP、后台任务和机器人;详见 [RBAC 权限管理](docs/zh-CN/rbac.md)
|
||||
- 📚 知识库(RAG):**Eino MultiQuery** 查询改写 + 多路向量检索 + **HTTP 精排**(DashScope `gte-rerank` / Cohere 兼容)+ 后处理(去重、预算);索引侧为 **Eino Compose** 流水线
|
||||
- 📁 对话分组管理:支持分组创建、置顶、重命名、删除等操作
|
||||
- 📂 **项目管理**:共享事实(黑板)跨会话沉淀认知,`upsert_project_fact` + `links` 串联攻击路径;聊天攻击链与项目事实图可视化
|
||||
@@ -148,6 +149,11 @@ CyberStrikeAI 基于 Go 构建,为 AI 原生安全运营提供完整底座:1
|
||||
- **Burp Suite 插件**:`plugins/burp-suite/cyberstrikeai-burp-extension/`
|
||||
构建产物:`plugins/burp-suite/cyberstrikeai-burp-extension/dist/cyberstrikeai-burp-extension.jar`
|
||||
说明文档:`plugins/burp-suite/cyberstrikeai-burp-extension/README.zh-CN.md`
|
||||
- **浏览器扩展(Chrome / Edge)**:`plugins/browser-extension/cyberstrikeai-browser-extension/`
|
||||
在 DevTools 中捕获 Network 流量并发送到 CyberStrikeAI 进行 AI 辅助安全测试,能力与 Burp 插件对齐。
|
||||
安装:`chrome://extensions/` → 加载已解压 → F12 → **CyberStrikeAI** 标签页
|
||||
打包产物:`plugins/browser-extension/cyberstrikeai-browser-extension/dist/cyberstrikeai-browser-extension.zip`
|
||||
说明文档:`plugins/browser-extension/cyberstrikeai-browser-extension/README.zh-CN.md`
|
||||
|
||||
## 工具概览
|
||||
|
||||
@@ -642,10 +648,11 @@ enabled: true
|
||||
- [API Recipes](docs/zh-CN/api-recipes.md):登录、Agent、流式、多代理、上传、漏洞、知识库和审计导出调用示例。
|
||||
- [配置参考](docs/zh-CN/configuration.md):`config.yaml` 各配置段、推荐值和修改建议。
|
||||
- [安全模型](docs/zh-CN/security-model.md):认证、工具执行、HITL、审计、C2/WebShell 和数据安全边界。
|
||||
- [RBAC 权限管理](docs/zh-CN/rbac.md):平台用户、系统/自定义角色、权限目录、逐权限 Scope、资源授权、Agent/MCP/机器人边界与 API 示例。
|
||||
- [API 参考](docs/zh-CN/api-reference.md):OpenAPI、认证、Agent、项目、知识库、C2、WebShell 等接口入口。
|
||||
- [多代理模式(Eino)](docs/zh-CN/MULTI_AGENT_EINO.md):**Deep**、**Plan-Execute**、**Supervisor**、`agents/*.md`、`eino_skills` / `eino_middleware`、接口与流式说明。
|
||||
- [图编排使用说明](docs/zh-CN/workflow-graph.md):可视化流程搭建、节点配置、`previous` / `outputs` 变量传参与角色绑定。
|
||||
- [机器人使用说明](docs/zh-CN/robot.md):个人微信、企业微信、钉钉、飞书、Telegram、Slack、Discord、QQ 机器人的配置、命令与排查。
|
||||
- [机器人使用说明](docs/zh-CN/robot.md):各平台接入、RBAC 逐用户绑定/服务账号模式、发送者白名单、命令、验证与排查。
|
||||
- [人机协同最佳实践](docs/zh-CN/hitl-best-practices.md):审批方模式、白名单、审计 Agent 提示词策略与独立小模型配置。
|
||||
|
||||
## 项目结构
|
||||
|
||||
@@ -35,11 +35,21 @@ func main() {
|
||||
fmt.Fprintf(os.Stderr, "无效的 -config 路径 %q。\n若同时需要 HTTPS,请写成: ./cyberstrike-ai --https -config config.yaml(-config 后必须是 yaml 文件路径)。\n", cp)
|
||||
os.Exit(2)
|
||||
}
|
||||
localConfig, err := config.EnsureLocalConfig(cp)
|
||||
if err != nil {
|
||||
fmt.Printf("加载配置失败: %v\n", err)
|
||||
return
|
||||
}
|
||||
|
||||
cfg, err := config.Load(cp)
|
||||
if err != nil {
|
||||
fmt.Printf("加载配置失败: %v\n", err)
|
||||
return
|
||||
}
|
||||
if localConfig.Created {
|
||||
cfg.Auth.GeneratedPassword = localConfig.GeneratedPassword
|
||||
cfg.Auth.GeneratedPasswordPersisted = true
|
||||
}
|
||||
|
||||
if *httpsBootstrap {
|
||||
config.ApplyDevHTTPSBootstrap(cfg)
|
||||
|
||||
@@ -10,7 +10,7 @@
|
||||
# ============================================
|
||||
|
||||
# 前端显示的版本号(可选,不填则显示默认版本)
|
||||
version: "v1.6.51"
|
||||
version: "v1.7.0"
|
||||
# 服务器配置
|
||||
server:
|
||||
host: 0.0.0.0 # 监听地址,0.0.0.0 表示监听所有网络接口
|
||||
@@ -203,14 +203,19 @@ multi_agent:
|
||||
tool_search_always_visible_tools: [read_file, glob, grep, analyze_image, 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, exec] # 后端内置常驻工具白名单(优先于 always_visible 数量策略)
|
||||
plantask_enable: true # P0:主代理挂载 TaskCreate/Get/Update/List 结构化任务板;需 eino_skills 可用且 skills_dir 存在
|
||||
plantask_rel_dir: .eino/plantask # 任务文件相对 skills_dir,按会话分子目录:skills/.eino/plantask/<conversationId>/
|
||||
reduction_enable: true # true:大工具输出截断/落盘以控上下文;依赖与 plantask 相同的 eino local 写盘后端,无后端时不挂载
|
||||
reduction_enable: true # true:大工具输出截断/落盘以控上下文;后端会独立创建,不依赖 eino_skills 是否启用
|
||||
reduction_max_length_for_trunc: 50000 # 单条工具结果超过该字符数(bytes)时截断并落盘(由 reduction 中间件处理)
|
||||
reduction_max_tokens_for_clear: 160000 # 历史工具结果清理阈值(tokens),超阈值时在模型调用前清理旧结果
|
||||
reduction_max_tokens_for_clear: 60000 # 历史工具结果清理阈值(tokens),应低于 max_total_tokens * summarization_trigger_ratio
|
||||
reduction_root_dir: "" # 非空:截断/清理内容落盘根路径;空:使用系统临时目录下按会话隔离的默认路径
|
||||
reduction_clear_exclude: [] # 不参与「清理阶段」的工具名额外列表(会与 task/transfer/exit 等内置排除项合并);需要时用 YAML 列表填写
|
||||
reduction_sub_agents: true # true:子代理也挂 reduction;false:仅编排主代理使用 reduction
|
||||
summarization_trigger_ratio: 0.8 # summarization 触发比例(max_total_tokens * ratio),建议 0.75~0.85
|
||||
summarization_emit_internal_events: true # true:发出 summarization 内部事件(便于诊断)
|
||||
summarization_user_intent_ledger_max_runes: 96000 # 压缩后注入模型上下文的「原始用户输入与约束账本」总字符上限;DB 原始消息不裁剪
|
||||
summarization_user_intent_ledger_entry_max_runes: 16000 # 账本中单条用户消息的字符上限;超出仅裁剪模型可见账本,不影响 DB 原文
|
||||
latest_user_message_max_runes: 48000 # 本轮最新 user 进入模型上下文的字符上限;超出时全文落盘,仅注入 head/tail 预览
|
||||
latest_user_message_head_runes: 24000 # 超长本轮 user 的头部预览字符数
|
||||
latest_user_message_tail_runes: 24000 # 超长本轮 user 的尾部预览字符数
|
||||
plan_execute_user_input_budget_ratio: 0.35 # plan_execute 中 userInput 预算比例(planner/replanner/executor 共用)
|
||||
plan_execute_executed_steps_budget_ratio: 0.2 # plan_execute 中 executed_steps 预算比例
|
||||
plan_execute_max_step_result_runes: 4000 # plan_execute 每步结果最大字符数(超出截断)
|
||||
@@ -267,10 +272,11 @@ security:
|
||||
# MCP (Model Context Protocol) 用于工具注册和调用
|
||||
mcp:
|
||||
enabled: false # 是否启用 MCP 服务器(http模式)
|
||||
host: 0.0.0.0 # MCP 服务器监听地址
|
||||
host: 127.0.0.1 # MCP 服务器监听地址;需要远程访问时再显式修改并配置网络层访问控制
|
||||
port: 8081 # MCP 服务器端口
|
||||
auth_header: "X-MCP-Token" # 鉴权:请求需携带该 header 且值与 auth_header_value 一致方可调用。留空表示不鉴权
|
||||
auth_header_value: "" # 鉴权密钥值(与 auth_header 配合使用,建议使用随机字符串)
|
||||
auth_header: "X-MCP-Token" # 可选的全局服务凭证 Header;普通调用请使用用户 Authorization: Bearer Token
|
||||
auth_header_value: "" # 全局服务凭证值,仅 allow_global_access=true 时生效
|
||||
allow_global_access: false # 高风险兼容模式:静态密钥映射为全局服务身份;默认请使用用户 Bearer Token
|
||||
# 外部 MCP 配置
|
||||
external_mcp:
|
||||
servers: {}
|
||||
@@ -339,6 +345,11 @@ knowledge:
|
||||
robots:
|
||||
wechat: # 微信 iLink(个人微信 ClawBot,扫码绑定)
|
||||
enabled: false
|
||||
# 鉴权默认 user_binding;专用机器人可改 service_account,并必须限制真实发送者
|
||||
auth:
|
||||
mode: user_binding # user_binding | service_account
|
||||
# service_user_id: "admin 或专用服务账号的 RBAC user ID"
|
||||
# allowed_external_users: ["t:tenant|u:sender"]
|
||||
bot_token: ""
|
||||
ilink_bot_id: ""
|
||||
ilink_user_id: ""
|
||||
@@ -347,6 +358,8 @@ robots:
|
||||
bot_agent: CyberStrikeAI/1.0
|
||||
wecom: # 企业微信
|
||||
enabled: false
|
||||
auth:
|
||||
mode: user_binding
|
||||
token: ""
|
||||
encoding_aes_key: ""
|
||||
corp_id: ""
|
||||
@@ -354,30 +367,42 @@ robots:
|
||||
agent_id: 0
|
||||
dingtalk: # 钉钉
|
||||
enabled: false
|
||||
auth:
|
||||
mode: user_binding
|
||||
client_id: ""
|
||||
client_secret: ""
|
||||
allow_conversation_id_fallback: false
|
||||
lark: # 飞书
|
||||
enabled: false
|
||||
auth:
|
||||
mode: user_binding
|
||||
app_id: ""
|
||||
app_secret: ""
|
||||
verify_token: ""
|
||||
allow_chat_id_fallback: false
|
||||
telegram: # Telegram
|
||||
enabled: false
|
||||
auth:
|
||||
mode: user_binding
|
||||
bot_token: ""
|
||||
bot_username: ""
|
||||
allow_group_messages: false
|
||||
slack: # Slack
|
||||
enabled: false
|
||||
auth:
|
||||
mode: user_binding
|
||||
bot_token: ""
|
||||
app_token: ""
|
||||
discord: # Discord
|
||||
enabled: false
|
||||
auth:
|
||||
mode: user_binding
|
||||
bot_token: ""
|
||||
allow_guild_messages: false
|
||||
qq: # QQ 机器人
|
||||
enabled: false
|
||||
auth:
|
||||
mode: user_binding
|
||||
app_id: ""
|
||||
client_secret: ""
|
||||
sandbox: true
|
||||
@@ -15,6 +15,7 @@ Documentation is split by language:
|
||||
- [贡献规范](zh-CN/contributing-guide.md)
|
||||
- [配置参考](zh-CN/configuration.md)
|
||||
- [安全模型](zh-CN/security-model.md)
|
||||
- [RBAC 权限管理](zh-CN/rbac.md)
|
||||
- [架构说明](zh-CN/architecture.md)
|
||||
- [API 参考](zh-CN/api-reference.md)
|
||||
- [排错指南](zh-CN/troubleshooting.md)
|
||||
@@ -45,6 +46,7 @@ Documentation is split by language:
|
||||
- [Contributing Guide](en-US/contributing-guide.md)
|
||||
- [Configuration Reference](en-US/configuration.md)
|
||||
- [Security Model](en-US/security-model.md)
|
||||
- [RBAC Administration](en-US/rbac.md)
|
||||
- [Architecture](en-US/architecture.md)
|
||||
- [API Reference](en-US/api-reference.md)
|
||||
- [Troubleshooting](en-US/troubleshooting.md)
|
||||
|
||||
@@ -8,6 +8,7 @@
|
||||
- [Contributing Guide](contributing-guide.md): checklists for APIs, config, tools, frontend, DB, high-risk features, and docs.
|
||||
- [Configuration Reference](configuration.md): `config.yaml` fields, hot-apply boundaries, recommended values, and source anchors.
|
||||
- [Security Model](security-model.md): trust boundaries, HITL, tool execution, C2/WebShell, and data safety.
|
||||
- [RBAC Administration](rbac.md): platform users, system/custom roles, permission catalog, per-permission scopes, resource assignments, Agent/MCP/robot boundaries, and API examples.
|
||||
- [Architecture](architecture.md): request flow, module relationships, complexity hotspots, and design trade-offs.
|
||||
- [API Reference](api-reference.md): authentication, OpenAPI, SSE, stability tiers, and common endpoints.
|
||||
- [Troubleshooting](troubleshooting.md): diagnostic order, minimal commands, common misdiagnoses, and issue template.
|
||||
|
||||
@@ -0,0 +1,374 @@
|
||||
# CyberStrikeAI RBAC Administration Guide
|
||||
|
||||
[中文](../zh-CN/rbac.md)
|
||||
|
||||
CyberStrikeAI can execute Agents, MCP tools, WebShell operations, C2 actions, and batch jobs. RBAC therefore applies beyond navigation visibility: it is enforced across HTTP APIs, resource queries, Agent contexts, built-in and external MCP tools, background jobs, and chatbot execution.
|
||||
|
||||
---
|
||||
|
||||
## 1. Two different kinds of roles
|
||||
|
||||
| Concept | Management location | Purpose |
|
||||
|---------|---------------------|---------|
|
||||
| **Platform role (RBAC Role)** | **Platform permissions** | Controls which features and resources a user may access |
|
||||
| **AI testing role (Agent Role)** | **Roles** / `roles/*.yaml` | Controls Agent prompts, methodology, and candidate tools |
|
||||
|
||||
An AI testing role is not an authorization boundary. Selecting a penetration-testing role does not grant platform permissions, and granting RBAC permissions does not change the Agent prompt.
|
||||
|
||||
---
|
||||
|
||||
## 2. Authorization model
|
||||
|
||||
An operation is allowed only when all relevant checks pass:
|
||||
|
||||
```text
|
||||
enabled account
|
||||
+ required permission for the route/tool
|
||||
+ scope attached to that permission
|
||||
+ resource owner / explicit assignment / parent inheritance
|
||||
+ additional rules for process-global operations
|
||||
```
|
||||
|
||||
Request flow:
|
||||
|
||||
1. Login issues a Bearer token whose session contains user, roles, permissions, and per-permission scopes.
|
||||
2. HTTP middleware maps the route to a permission, for example `GET /api/projects` → `project:read`.
|
||||
3. Resource-ID requests also check ownership, explicit assignments, or supported parent inheritance.
|
||||
4. Agent execution receives an immutable Principal through `context.Context`.
|
||||
5. Built-in MCP tools authorize both the tool and resource IDs in tool arguments. External MCP has separate restrictions.
|
||||
6. Denials are written to RBAC/audit logs.
|
||||
|
||||
Frontend button hiding is only a usability feature. The server is the security boundary.
|
||||
|
||||
---
|
||||
|
||||
## 3. Built-in platform roles
|
||||
|
||||
| Role | Scope | Default capability |
|
||||
|------|-------|--------------------|
|
||||
| **Administrator `admin`** | `all` | Every known permission, including RBAC, configuration, terminal, audit deletion, and global definition management |
|
||||
| **Operator `operator`** | `assigned` | Normal read/write/execute work; excludes RBAC, core configuration, terminal, audit management, external MCP execution, and several global definition writes |
|
||||
| **Auditor `auditor`** | `all` | Read permissions across modules plus `audit:read`; no writes |
|
||||
| **Viewer `viewer`** | `assigned` | Read-only access within authorized resources |
|
||||
|
||||
System roles cannot be edited or deleted. Their grants are rebuilt from the current permission catalog during upgrade, preventing stale grants from older versions. Create custom roles for different job functions.
|
||||
|
||||
An account without a role can still authenticate but has almost no business capability; do not treat “no role” as a complete job profile.
|
||||
|
||||
---
|
||||
|
||||
## 4. Permission catalog
|
||||
|
||||
Permissions use `module:action`. Common actions are `read`, `write`, `delete`, and `execute`. The authoritative catalog for the running build is available in Platform permissions or `GET /api/rbac/metadata`.
|
||||
|
||||
| Module | Permissions |
|
||||
|--------|-------------|
|
||||
| Account | `auth:self` |
|
||||
| Dashboard | `dashboard:read` |
|
||||
| Chat | `chat:read`, `chat:write`, `chat:delete` |
|
||||
| Agent | `agent:execute`, `agent:local-execute` |
|
||||
| HITL | `hitl:read`, `hitl:write` |
|
||||
| Tasks | `tasks:read`, `tasks:write`, `tasks:delete` |
|
||||
| Projects | `project:read`, `project:write`, `project:delete` |
|
||||
| Vulnerabilities | `vulnerability:read`, `vulnerability:write`, `vulnerability:delete` |
|
||||
| WebShell | `webshell:read`, `webshell:write`, `webshell:delete` |
|
||||
| C2 | `c2:read`, `c2:write`, `c2:delete` |
|
||||
| MCP | `mcp:read`, `mcp:execute`, `mcp:write`, `mcp:external:execute` |
|
||||
| Knowledge | `knowledge:read`, `knowledge:write`, `knowledge:delete` |
|
||||
| Skills | `skills:read`, `skills:write`, `skills:delete` |
|
||||
| Markdown Agents | `agents:read`, `agents:write`, `agents:delete` |
|
||||
| AI testing roles | `roles:read`, `roles:write`, `roles:delete` |
|
||||
| Workflows | `workflow:read`, `workflow:execute`, `workflow:write`, `workflow:delete` |
|
||||
| Configuration | `config:read`, `config:write` |
|
||||
| Terminal | `terminal:execute` |
|
||||
| Audit | `audit:read`, `audit:delete` |
|
||||
| RBAC | `rbac:read`, `rbac:write` |
|
||||
| Notifications | `notification:read`, `notification:write` |
|
||||
| Robots | `robot:read`, `robot:write` |
|
||||
| Files | `files:read`, `files:write`, `files:delete` |
|
||||
| Attack chain | `attackchain:read`, `attackchain:write` |
|
||||
| FOFA | `fofa:execute` |
|
||||
| OpenAPI | `openapi:read` |
|
||||
| Chat groups | `group:read`, `group:write`, `group:delete` |
|
||||
| Monitor | `monitor:read`, `monitor:write`, `monitor:delete` |
|
||||
|
||||
Important distinctions:
|
||||
|
||||
- `agent:execute` runs Agents but does not grant local filesystem, shell, or arbitrary configured command access.
|
||||
- `agent:local-execute` is the local execution fallback and should be limited to trusted operators.
|
||||
- `mcp:execute` protects the authenticated MCP HTTP entry point.
|
||||
- `mcp:external:execute` allows Agent calls to external MCP tools and currently also requires `all` scope.
|
||||
- `mcp:write` manages external MCP configuration; it is separate from external tool execution.
|
||||
- `robot:write` manages robot configuration and the test endpoint. Chatbot conversations use the bound user or configured service account's business permissions.
|
||||
|
||||
---
|
||||
|
||||
## 5. Resource scopes
|
||||
|
||||
Each role has one scope:
|
||||
|
||||
| Scope | Meaning | Typical use |
|
||||
|-------|---------|-------------|
|
||||
| `all` | All resources covered by the permission | Administrator, global auditor |
|
||||
| `assigned` | Explicitly assigned resources and supported parent-resource inheritance | Project member, assigned asset operator |
|
||||
| `own` | Primarily resources created by/owned by the user; some resource types also support explicit assignment or parent inheritance | Personal workspace, isolated robot identity |
|
||||
|
||||
Users may have multiple roles. Permissions are unioned, while scopes are merged **for the same permission only**:
|
||||
|
||||
```text
|
||||
all > assigned > own
|
||||
```
|
||||
|
||||
Example:
|
||||
|
||||
```text
|
||||
Global audit role: project:read + all
|
||||
Personal editor: project:write + own
|
||||
|
||||
Effective:
|
||||
project:read → all
|
||||
project:write → own
|
||||
```
|
||||
|
||||
A global read role does not widen an unrelated write permission. Authorization code must use `ScopeFor(permission)`, not the user's broadest display scope.
|
||||
|
||||
### Process-global restrictions
|
||||
|
||||
Some definitions have no owner. Their mutations require the corresponding permission with `all` scope even if the user has a `write` key:
|
||||
|
||||
- AI testing roles, Skills, and Markdown Agents.
|
||||
- External MCP configuration.
|
||||
- Robot configuration.
|
||||
- Workflow definitions.
|
||||
- Knowledge mutations other than search.
|
||||
- Global HITL allowlist, reviewer, and audit policy.
|
||||
- C2 Profile mutations.
|
||||
- Some global monitor statistics.
|
||||
|
||||
---
|
||||
|
||||
## 6. Ownership, assignments, and inheritance
|
||||
|
||||
Use Platform permissions → Member details → Resource assignments. Directly assignable resource types include:
|
||||
|
||||
- `project`
|
||||
- `conversation`
|
||||
- `vulnerability`
|
||||
- `webshell`
|
||||
- `batch_task`
|
||||
- `c2_listener`
|
||||
|
||||
A batch request accepts at most 100 resources. Duplicate grants are skipped.
|
||||
|
||||
Supported inheritance includes:
|
||||
|
||||
| Child resource | Parent access source |
|
||||
|----------------|----------------------|
|
||||
| Conversation | Project |
|
||||
| Vulnerability | Project or related conversation |
|
||||
| Message, process detail, attack chain | Conversation |
|
||||
| C2 Session | Listener |
|
||||
| C2 Task/file/event | Session, Task, or Listener chain |
|
||||
|
||||
Assigning a project therefore usually avoids assigning each conversation and vulnerability separately. The concrete route/tool server check remains authoritative.
|
||||
|
||||
---
|
||||
|
||||
## 7. Web administration workflow
|
||||
|
||||
### Create a user
|
||||
|
||||
1. Sign in as an administrator and open **Platform permissions**.
|
||||
2. Create a user with username, display name, an eight-character-or-longer password, and enabled status.
|
||||
3. Assign one or more platform roles.
|
||||
4. For `assigned` roles, configure resource assignments.
|
||||
5. Have the user sign in again and verify roles, permission count, and scope in the top-right user menu.
|
||||
|
||||
### Create a custom role
|
||||
|
||||
1. Give the role a job-oriented name and description.
|
||||
2. Select `all`, `assigned`, or `own`.
|
||||
3. Select only required permissions.
|
||||
4. Test list, detail, mutation, deletion, Agent, and tool behavior with a test account.
|
||||
5. Assign it to production users only after verification.
|
||||
|
||||
System roles are immutable; create a custom role instead of modifying them.
|
||||
|
||||
### When changes take effect
|
||||
|
||||
- Updating a user, password, enabled state, or role membership revokes that user's sessions; they must sign in again.
|
||||
- Updating or deleting a custom role revokes all sessions; all users must sign in again.
|
||||
- Robots resolve the bound user/service account on every message, so disablement and role changes affect the next message.
|
||||
- Background batch jobs resolve a Principal from the task owner rather than trusting frontend state.
|
||||
|
||||
---
|
||||
|
||||
## 8. Suggested role templates
|
||||
|
||||
### Read-only project member
|
||||
|
||||
```text
|
||||
Scope: assigned
|
||||
dashboard:read
|
||||
chat:read
|
||||
project:read
|
||||
vulnerability:read
|
||||
files:read
|
||||
attackchain:read
|
||||
```
|
||||
|
||||
### Daily security operator
|
||||
|
||||
```text
|
||||
Scope: assigned
|
||||
agent:execute
|
||||
chat:read / chat:write
|
||||
project:read / project:write
|
||||
vulnerability:read / vulnerability:write
|
||||
tasks:read / tasks:write
|
||||
files:read / files:write
|
||||
hitl:read / hitl:write
|
||||
```
|
||||
|
||||
Add `agent:local-execute` or `terminal:execute` only when local commands are required. Add individual `:delete` permissions only when deletion is part of the job.
|
||||
|
||||
### Robot service account
|
||||
|
||||
```text
|
||||
Scope: own (isolated workspace) or assigned (specific projects)
|
||||
agent:execute
|
||||
chat:read / chat:write
|
||||
optional project, vulnerability, and knowledge permissions
|
||||
```
|
||||
|
||||
`admin` can be used as a robot service account, but exact sender allowlisting still applies. Every allowlisted sender receives full permissions and shares admin-owned data. See the [Robot guide](robot.md).
|
||||
|
||||
---
|
||||
|
||||
## 9. Agent, MCP, and robot boundaries
|
||||
|
||||
### Agent
|
||||
|
||||
The HTTP user becomes an immutable Principal propagated to single-agent, multi-agent, workflow, and tool contexts. A long-running task may survive an SSE disconnect while retaining that identity.
|
||||
|
||||
### Built-in MCP
|
||||
|
||||
Every built-in tool requires an explicit authorization policy. WebShell tools check both `webshell:read/write/delete` and the target `connection_id`; project, vulnerability, task, and C2 tools validate resource arguments as well. An unregistered built-in policy fails closed. Other local/configured tools require `agent:local-execute`.
|
||||
|
||||
### External MCP
|
||||
|
||||
Agent calls to external MCP require `mcp:external:execute` with `all` scope because an external service's resource model is not protected by local ownership and assignments.
|
||||
|
||||
### Robots
|
||||
|
||||
- `user_binding`: each platform sender binds their own RBAC user.
|
||||
- `service_account`: exact allowlisted senders share one RBAC user.
|
||||
- Platform signature verification authenticates message origin, not business authorization.
|
||||
- Run `whoami` to inspect the effective Principal.
|
||||
|
||||
---
|
||||
|
||||
## 10. RBAC API
|
||||
|
||||
All requests use:
|
||||
|
||||
```http
|
||||
Authorization: Bearer <token>
|
||||
```
|
||||
|
||||
Management routes require `rbac:read` or `rbac:write`; the resource picker requires `rbac:write`.
|
||||
|
||||
| Method | Path | Purpose |
|
||||
|--------|------|---------|
|
||||
| GET | `/api/rbac/me` | Current user, roles, permissions, overall scope, per-permission scopes |
|
||||
| GET | `/api/rbac/metadata` | Permission catalog, roles, grants, and scopes |
|
||||
| GET/POST | `/api/rbac/users` | List/create users |
|
||||
| PUT/DELETE | `/api/rbac/users/:id` | Update/delete a user |
|
||||
| GET/POST | `/api/rbac/roles` | List/create roles |
|
||||
| PUT/DELETE | `/api/rbac/roles/:id` | Update/delete a custom role |
|
||||
| GET | `/api/rbac/resources?type=project&q=...` | Search assignable resources |
|
||||
| GET/POST | `/api/rbac/resource-assignments` | List/create assignments |
|
||||
| DELETE | `/api/rbac/resource-assignments/:id` | Revoke an assignment |
|
||||
|
||||
Create a user:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8080/api/rbac/users \
|
||||
-H "Authorization: Bearer $TOKEN" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"username": "operator01",
|
||||
"display_name": "Security Operator 01",
|
||||
"password": "change-me-123",
|
||||
"enabled": true,
|
||||
"roles": ["operator"]
|
||||
}'
|
||||
```
|
||||
|
||||
Create a custom role:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8080/api/rbac/roles \
|
||||
-H "Authorization: Bearer $TOKEN" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"name": "Project Auditor",
|
||||
"description": "Read assigned projects",
|
||||
"scope": "assigned",
|
||||
"permissions": ["chat:read", "project:read", "vulnerability:read"]
|
||||
}'
|
||||
```
|
||||
|
||||
Assign projects:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8080/api/rbac/resource-assignments \
|
||||
-H "Authorization: Bearer $TOKEN" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"user_id": "USER_ID",
|
||||
"resource_type": "project",
|
||||
"resource_ids": ["PROJECT_ID_1", "PROJECT_ID_2"]
|
||||
}'
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 11. Audit and operations recommendations
|
||||
|
||||
- Use individual administrator accounts instead of sharing one password.
|
||||
- Name custom roles by job function and document purpose/owner.
|
||||
- Review high-risk permissions separately: `terminal:execute`, `agent:local-execute`, `c2:write/delete`, `webshell:write/delete`, `rbac:write`, and `config:write`.
|
||||
- Periodically review `all` roles, service accounts, robot allowlists, and dormant users.
|
||||
- On offboarding, disable the account first, then revoke robot bindings, assignments, and sessions.
|
||||
- Monitor RBAC denials, user/role changes, resource assignments, and robot service-account execution in audit logs.
|
||||
- Pair RBAC with HITL for dangerous tools; permission to invoke does not bypass approval policy.
|
||||
|
||||
---
|
||||
|
||||
## 12. Troubleshooting
|
||||
|
||||
### A button is missing
|
||||
|
||||
The frontend hides actions based on `/api/rbac/me`. Verify the required permission. Direct API calls are still rejected server-side.
|
||||
|
||||
### Permission exists but the resource is denied
|
||||
|
||||
Inspect the scope for that specific permission, not only the overall display scope. Then check owner, explicit assignment, and parent assignment.
|
||||
|
||||
### Role changed but the user sees old access
|
||||
|
||||
Role changes revoke sessions. Sign in again. Robots resolve again on the next message.
|
||||
|
||||
### A global mutation is denied despite `write`
|
||||
|
||||
Process-global definitions require the corresponding permission with `all` scope. Create a dedicated global administration role instead of widening unrelated permissions.
|
||||
|
||||
### Agent chat works but commands fail
|
||||
|
||||
`agent:execute` and `agent:local-execute` are separate. Grant local execution only when necessary and combine it with HITL, tool allowlists, and audit.
|
||||
|
||||
### External MCP requires global scope
|
||||
|
||||
The user needs `mcp:external:execute`, and that permission's scope must be `all`.
|
||||
|
||||
+128
-22
@@ -2,7 +2,7 @@
|
||||
|
||||
[中文](../zh-CN/robot.md)
|
||||
|
||||
This document explains how to chat with CyberStrikeAI from **personal WeChat**, **DingTalk**, **Lark (Feishu)**, and **WeCom (Enterprise WeChat)** using long-lived connections or HTTP callbacks—no need to open a browser on the server. Following the steps below helps avoid common mistakes.
|
||||
This guide covers **Personal WeChat, WeCom, DingTalk, Lark, Telegram, Slack, Discord, and QQ Bot**, including platform connectivity, RBAC identity binding, service-account allowlists, commands, verification, and troubleshooting.
|
||||
|
||||
---
|
||||
|
||||
@@ -15,10 +15,18 @@ This document explains how to chat with CyberStrikeAI from **personal WeChat**,
|
||||
- **Personal WeChat**: Open **WeChat / iLink** → **Generate QR code and bind**, then scan with WeChat (see [Section 3.4](#34-personal-wechat-wechat--ilink))
|
||||
- **DingTalk**: Enable and fill in Client ID / Client Secret
|
||||
- **Lark**: Enable and fill in App ID / App Secret
|
||||
5. Click **Apply configuration** to save (WeChat binding saves and enables automatically on success—usually no extra click needed)
|
||||
6. **Restart the CyberStrikeAI process** (DingTalk/Lark: saving alone does not establish the connection; WeChat auto-restarts the iLink poll after binding—usually no manual restart needed)
|
||||
5. Click **Apply configuration** to save and automatically restart the corresponding bot connection. WeChat binding saves and enables automatically on success.
|
||||
|
||||
Settings are written to the `robots` section of `config.yaml`; you can also edit the file directly. **After changing DingTalk or Lark config, you must restart for the long-lived connection to take effect.** Personal WeChat binding automatically writes `robots.wechat` and restarts the iLink long poll.
|
||||
Settings are written to the `robots` section of `config.yaml`; you can also edit the file directly. Web-based **Apply configuration** restarts the corresponding connection automatically. Restart the CyberStrikeAI process only when editing YAML directly. Personal WeChat binding automatically writes `robots.wechat` and restarts the iLink long poll.
|
||||
|
||||
### Shortest path to first use
|
||||
|
||||
After the platform connection works, configure the business identity before sending normal prompts:
|
||||
|
||||
- **Multiple users**: choose User binding → each user generates a code from the top-right Web user menu → sends the bind command to the bot → runs `whoami` to verify.
|
||||
- **Only you**: run `whoami` first and copy the sender ID → choose Service account → set User ID to `admin` or another RBAC user → paste the exact sender allowlist → apply configuration → run `whoami` again.
|
||||
|
||||
Start normal AI chat only after the response shows an authorized status and the expected effective identity.
|
||||
|
||||
---
|
||||
|
||||
@@ -82,7 +90,7 @@ If you only have a **custom bot** Webhook URL (`oapi.dingtalk.com/robot/send?acc
|
||||
- In CyberStrikeAI: System settings → Robot settings → DingTalk.
|
||||
- Enable “Enable DingTalk robot”.
|
||||
- Paste the Client ID and Client Secret from step 3.
|
||||
- Click **Apply configuration**, then **restart CyberStrikeAI**.
|
||||
- Click **Apply configuration**; CyberStrikeAI restarts the DingTalk connection automatically.
|
||||
|
||||
---
|
||||
|
||||
@@ -105,7 +113,7 @@ If you only have a **custom bot** Webhook URL (`oapi.dingtalk.com/robot/send?acc
|
||||
| App Secret | From Lark open platform app credentials |
|
||||
| Verify Token | Optional; for event subscription |
|
||||
|
||||
**Lark setup in short**: Log in to [Lark Open Platform](https://open.feishu.cn) → Create an enterprise app → In “Credentials and basic info” get **App ID** and **App Secret** → In “Application capabilities” enable **Robot** and the right permissions → Add **event subscription** and **permissions** below → Publish the app → Enter App ID and App Secret in CyberStrikeAI robot settings → Save and **restart** the app.
|
||||
**Lark setup in short**: Log in to [Lark Open Platform](https://open.feishu.cn) → Create an enterprise app → In “Credentials and basic info” get **App ID** and **App Secret** → In “Application capabilities” enable **Robot** and the right permissions → Add **event subscription** and **permissions** below → Publish the app → Enter App ID and App Secret in CyberStrikeAI robot settings → **Apply configuration**.
|
||||
|
||||
**Event subscription**
|
||||
The long-lived connection only receives message events if you subscribe to them. In the app’s **Events and callbacks** (事件与回调) → **Event subscription** (事件订阅), add the event **Receive message** (**im.message.receive_v1**). Without it, the connection succeeds but no message events are delivered (no logs when users send messages).
|
||||
@@ -272,12 +280,79 @@ In **Permission management** (权限管理), enable the following (names and ide
|
||||
|
||||
---
|
||||
|
||||
## 4. Bot commands
|
||||
## 4. RBAC authorization and bot commands
|
||||
|
||||
Platform credentials and callback signatures authenticate the messaging platform. CyberStrikeAI RBAC determines what the sender can actually do. Each bot instance uses one authorization mode.
|
||||
|
||||
### 4.1 Choose an authorization mode
|
||||
|
||||
| Scenario | Recommended mode | Identity and data behavior |
|
||||
|----------|------------------|----------------------------|
|
||||
| Shared WeCom, Lark, DingTalk, or Slack bot | `user_binding` | Each sender binds their own Web user; permissions and resources remain isolated |
|
||||
| Personal WeChat, single-user bot, fixed automation entry | `service_account` | Allowlisted senders share the configured RBAC user's permissions and owned resources |
|
||||
|
||||
Both modes resolve user status, roles, per-permission scope, and resource assignments before every message. Basic AI chat requires:
|
||||
|
||||
```text
|
||||
agent:execute
|
||||
chat:read
|
||||
chat:write
|
||||
```
|
||||
|
||||
Grant project, role, local execution, WebShell, C2, or MCP permissions only when those features are required. Conversation deletion also requires `chat:delete`.
|
||||
|
||||
### 4.2 User-binding mode (default)
|
||||
|
||||
Administrator:
|
||||
|
||||
1. Open System settings → Robot settings → select a platform.
|
||||
2. Set Authorization policy to `user_binding` and apply the configuration.
|
||||
|
||||
Each user:
|
||||
|
||||
1. Sign in to the Web UI and open the top-right user menu → **Bind robot account**.
|
||||
2. Generate a binding code; a five-minute countdown starts.
|
||||
3. Send the full command to the target bot, for example `bind 7C6E-BD4C`.
|
||||
4. Send `whoami` and confirm the effective RBAC identity is their own Web user.
|
||||
|
||||
Codes are stored only as hashes and are single-use. When the countdown ends, the UI marks the code expired, disables copying, and refreshes the binding list; the server also rejects it. Generating a new code immediately invalidates the previous unused code. Users can send `unbind` or revoke a binding from the Web dialog.
|
||||
|
||||
### 4.3 Service-account mode
|
||||
|
||||
1. Connect the bot to its messaging platform.
|
||||
2. Have each intended sender run `whoami` and copy the exact sender ID. For Personal WeChat it usually resembles `xxxx@im.wechat`; never substitute `ilink_bot_id` or configured `ilink_user_id`.
|
||||
3. In Robot settings, select `service_account`.
|
||||
4. Enter the RBAC **User ID**, not its display name. `admin` is allowed; every allowlisted sender then receives full platform permissions and the UI shows a red warning.
|
||||
5. Add one exact sender ID per line. Matching is case-sensitive and `*` wildcards are rejected.
|
||||
6. Apply configuration and run `whoami` again to verify the effective user, roles, and scope.
|
||||
|
||||
Example:
|
||||
|
||||
```yaml
|
||||
robots:
|
||||
wechat:
|
||||
auth:
|
||||
mode: service_account
|
||||
service_user_id: admin
|
||||
allowed_external_users:
|
||||
- "o9cq806s32Sm2_kyOmkyaV7Rn1lU@im.wechat"
|
||||
```
|
||||
|
||||
Service-account mode rejects `bind` and `unbind`. All allowlisted senders share conversations, projects, and other resources owned by the service account. Use `user_binding` when that sharing is undesirable.
|
||||
|
||||
### 4.4 Inspect the effective identity
|
||||
|
||||
Send `whoami`. The response includes platform, exact sender ID, authorization mode and status, effective RBAC user and ID, roles, scope, and permission count. A non-allowlisted sender sees only the denial status and no service-account details.
|
||||
|
||||
### 4.5 Command list
|
||||
|
||||
Send these **text commands** to the bot on any connected platform (text only):
|
||||
|
||||
| Command | Description |
|
||||
|---------|-------------|
|
||||
| **绑定 \<code\>** or **bind \<code\>** | Bind the verified platform sender to the RBAC user that generated the code |
|
||||
| **解绑** or **unbind** | Remove the current platform identity binding |
|
||||
| **身份** or **whoami** | Show sender ID, authorization mode, binding status, and the effective RBAC user, roles, and scope |
|
||||
| **帮助** (help) | Show command help |
|
||||
| **列表** or **对话列表** (list) | List all conversation titles and IDs |
|
||||
| **切换 \<conversationID\>** or **继续 \<conversationID\>** | Continue in the given conversation |
|
||||
@@ -292,6 +367,8 @@ Send these **text commands** to the bot on any connected platform (text only):
|
||||
|
||||
Any other text is sent to the AI as a user message, same as in the web UI (e.g. penetration testing, security analysis).
|
||||
|
||||
Group messages are authorized as the actual sender, never as a group ID. In service-account mode, explicitly allowlisted senders intentionally share the configured account.
|
||||
|
||||
---
|
||||
|
||||
## 5. How to use (do I need to @ the bot?)
|
||||
@@ -310,14 +387,17 @@ Summary: **Personal WeChat and direct chat**—just send; **DingTalk/Lark in a g
|
||||
|
||||
1. CyberStrikeAI web UI → System settings → Robot settings → **WeChat / iLink** → **Generate QR code and bind**.
|
||||
2. Scan with WeChat and confirm (enter pairing code on the web page if prompted).
|
||||
3. After binding, send “帮助” in the WeChat private chat to test.
|
||||
3. Send `whoami` in the WeChat private chat and copy the sender ID.
|
||||
4. Choose `user_binding`, or configure `service_account` with the RBAC user and exact sender allowlist.
|
||||
5. Apply configuration, run `whoami` again, then send a normal message.
|
||||
|
||||
**DingTalk / Lark**
|
||||
|
||||
1. **In the open platform**: Complete app creation, copy credentials, enable the bot (DingTalk: **Stream mode**), set permissions, and publish (Section 3).
|
||||
2. **In CyberStrikeAI**: System settings → Robot settings → Enable the platform, paste Client ID/App ID and Client Secret/App Secret → **Apply configuration**.
|
||||
3. **Restart the CyberStrikeAI process** (otherwise the long-lived connection is not established).
|
||||
4. **On your phone**: Open DingTalk or Lark, find the bot (direct chat or @ in a group), send “帮助” or any message to test.
|
||||
3. **Choose authorization**: use `user_binding` for multiple users, or configure a service account and exact allowlist for a dedicated bot.
|
||||
4. **Apply configuration**; the Web UI restarts the corresponding connection automatically.
|
||||
5. **On your phone**: Open the bot, run `whoami` first, then send a normal message.
|
||||
|
||||
If the bot does not respond, see **Section 9 (troubleshooting)** and **Section 10 (common pitfalls)**.
|
||||
|
||||
@@ -331,6 +411,11 @@ Example `robots` section in `config.yaml`:
|
||||
robots:
|
||||
wechat: # Personal WeChat iLink (auto-filled after QR bind; usually no manual edit)
|
||||
enabled: true
|
||||
auth:
|
||||
mode: service_account
|
||||
service_user_id: admin
|
||||
allowed_external_users:
|
||||
- "exact sender ID copied from whoami"
|
||||
bot_token: "your_bot_token@im.bot:..."
|
||||
ilink_bot_id: "your_bot_id@im.bot"
|
||||
ilink_user_id: "your_user_id@im.wechat"
|
||||
@@ -339,10 +424,14 @@ robots:
|
||||
bot_agent: "CyberStrikeAI/1.0"
|
||||
dingtalk:
|
||||
enabled: true
|
||||
auth:
|
||||
mode: user_binding
|
||||
client_id: "your_dingtalk_app_key"
|
||||
client_secret: "your_dingtalk_app_secret"
|
||||
lark:
|
||||
enabled: true
|
||||
auth:
|
||||
mode: user_binding
|
||||
app_id: "your_lark_app_id"
|
||||
app_secret: "your_lark_app_secret"
|
||||
verify_token: ""
|
||||
@@ -372,7 +461,7 @@ robots:
|
||||
sandbox: true
|
||||
```
|
||||
|
||||
After changing DingTalk/Lark/WeCom/Telegram/Slack/Discord/QQ settings, **Apply configuration** restarts the corresponding connections. Personal WeChat QR binding saves and restarts automatically.
|
||||
Authorization is configured independently per platform; omitting `auth` defaults to `user_binding`. **Apply configuration** restarts the corresponding connections. Restart the process only after editing YAML directly. Personal WeChat QR binding saves and restarts automatically.
|
||||
|
||||
---
|
||||
|
||||
@@ -380,20 +469,24 @@ After changing DingTalk/Lark/WeCom/Telegram/Slack/Discord/QQ settings, **Apply c
|
||||
|
||||
You can verify bot logic with the **test API** (no DingTalk/Lark client needed):
|
||||
|
||||
1. Log in to the CyberStrikeAI web UI (so you have a session).
|
||||
2. Call the test endpoint with curl (include your session Cookie):
|
||||
1. Sign in with an account that has global `robot:write` permission and obtain a Bearer token.
|
||||
2. Call the test endpoint with curl:
|
||||
|
||||
```bash
|
||||
# Replace YOUR_COOKIE with the Cookie from your browser (F12 → Network → any request → Request headers → Cookie)
|
||||
# Adjust the URL, username, and password for your deployment
|
||||
TOKEN=$(curl -s -X POST "http://localhost:8080/api/auth/login" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"username":"admin","password":"YOUR_PASSWORD"}' | jq -r '.token')
|
||||
|
||||
curl -X POST "http://localhost:8080/api/robot/test" \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "Cookie: YOUR_COOKIE" \
|
||||
-H "Authorization: Bearer $TOKEN" \
|
||||
-d '{"platform":"dingtalk","user_id":"test_user","text":"帮助"}'
|
||||
```
|
||||
|
||||
If the JSON response contains `"reply":"【CyberStrikeAI 机器人命令】..."`, command handling works. You can also try `"text":"列表"` or `"text":"当前"`.
|
||||
If the JSON response contains `"reply":"【CyberStrikeAI 机器人命令】..."`, command handling works. `help`, `version`, and `whoami` work before binding. `list`, `current`, and normal AI messages enforce RBAC: the test `platform + user_id` must already be bound or exactly match the service-account allowlist.
|
||||
|
||||
API: `POST /api/robot/test` (requires login). Body: `{"platform":"optional","user_id":"optional","text":"required"}`. Response: `{"reply":"..."}`.
|
||||
API: `POST /api/robot/test` (requires global `robot:write`). Body: `{"platform":"optional","user_id":"optional","text":"required"}`. Response: `{"reply":"..."}`. This endpoint simulates bot business logic only; it does not validate a third-party callback signature or long-lived connection.
|
||||
|
||||
---
|
||||
|
||||
@@ -407,7 +500,7 @@ Check in this order:
|
||||
Robot settings should show “Connected” or a bound Bot ID; `robots.wechat.bot_token` in `config.yaml` must not be empty.
|
||||
|
||||
2. **Enabled?**
|
||||
Confirm “Enable WeChat robot” is checked; restart CyberStrikeAI if you just changed settings.
|
||||
Confirm “Enable WeChat robot” is checked and click **Apply configuration** if you just changed settings.
|
||||
|
||||
3. **Application logs**
|
||||
- On startup: `微信 iLink 长轮询已启动`;
|
||||
@@ -430,8 +523,8 @@ Check in this order:
|
||||
1. **Client ID / Client Secret match the open platform exactly**
|
||||
Copy from “Credentials and basic info”; avoid typing. Watch **0** vs **o** and **1** vs **l** (e.g. `ding9gf9tiozuc504aer` has **504**, not 5o4).
|
||||
|
||||
2. **Did you restart after saving?**
|
||||
The long-lived connection is created at **startup**. “Apply configuration” only updates the config file; you **must restart the CyberStrikeAI process** for the DingTalk connection to start.
|
||||
2. **Did you apply the configuration?**
|
||||
Web changes require **Apply configuration**, which restarts the corresponding connection automatically. Restart the process only after editing `config.yaml` directly.
|
||||
|
||||
3. **Application logs**
|
||||
- On startup you should see: `钉钉 Stream 正在连接…`, `钉钉 Stream 已启动(无需公网),等待收消息`.
|
||||
@@ -441,6 +534,15 @@ Check in this order:
|
||||
4. **Open platform**
|
||||
The app must be **published**. Under “Robot” you must enable **Stream** for receiving messages (HTTP callback only is not enough). Permission management must include robot receive/send message permissions.
|
||||
|
||||
### 9.3 Reply says unbound, sender denied, or permission missing
|
||||
|
||||
1. Run `whoami` and inspect the authorization mode and status.
|
||||
2. In `user_binding`, generate a code from the top-right Web user menu and send the complete bind command from the same platform identity. Regenerate expired or already-used codes.
|
||||
3. In `service_account`, copy the exact sender ID from `whoami` into that platform's allowlist. Preserve case, tenant prefixes, and suffixes such as `@im.wechat`.
|
||||
4. If an effective user is shown but permissions are missing, grant at least `agent:execute`, `chat:read`, and `chat:write` for normal AI chat.
|
||||
5. A missing or disabled service user is rejected when applying configuration.
|
||||
6. If `admin` is denied, the usual cause is an allowlist mismatch—not insufficient admin permissions.
|
||||
|
||||
---
|
||||
|
||||
## 10. Common pitfalls
|
||||
@@ -448,7 +550,11 @@ Check in this order:
|
||||
- **Personal WeChat vs WeCom**: Personal WeChat uses `robots.wechat` + web QR bind; WeCom uses `robots.wecom` + admin callback URL—they are completely different.
|
||||
- **WeChat QR expired**: QR codes last ~5 minutes; regenerate instead of reusing an old one.
|
||||
- **Wrong bot type**: The “Custom” bot added in a DingTalk **group** (Webhook + sign secret) **cannot** be used for two-way chat. Only the **enterprise internal app** bot from the open platform is supported.
|
||||
- **Saved but not restarted**: After changing DingTalk/Lark robot settings you **must restart** the app (WeChat QR bind restarts the connection automatically).
|
||||
- **Configuration not applied**: Click **Apply configuration** after Web changes; connections restart automatically. A process restart is needed only for direct YAML edits.
|
||||
- **Bot ID used as sender ID**: Copy the sender ID from `whoami`; do not use `ilink_bot_id`, configured `ilink_user_id`, a group ID, or a display name.
|
||||
- **Reusing an expired code**: Codes last five minutes and are single-use; generating a new code immediately invalidates the old one.
|
||||
- **Assuming service-account users are isolated**: All allowlisted senders share that account's conversations and owned resources. Use `user_binding` for isolation.
|
||||
- **Assuming admin removes the allowlist**: It does not. The sender must still match exactly, but every matching sender gets full permissions.
|
||||
- **Client ID typo**: If the platform shows `504`, use `504` (not `5o4`); prefer copy/paste.
|
||||
- **DingTalk: only HTTP callback, no Stream**: This app receives messages via **Stream**. In the open platform, message reception must be **Stream mode**.
|
||||
- **App not published**: After changing the bot or permissions in the open platform, **publish a new version** under “Version management and release”, or changes won’t apply.
|
||||
@@ -459,5 +565,5 @@ Check in this order:
|
||||
|
||||
- All platforms: **text messages only**; other types (e.g. image, voice) are not supported and may be ignored.
|
||||
- Personal WeChat: **private chat only**—group @-bot is not supported.
|
||||
- Conversations are shared with the web UI: conversations created from the bot appear in the web “Conversations” list and vice versa.
|
||||
- Bot data is shared with the web UI: under `user_binding` it belongs to the bound user; under `service_account` it belongs to the service account and is shared by allowlisted senders.
|
||||
- Bot execution uses the same **Eino single/multi-agent** path as the web UI (`ProcessMessageForRobot`, with progress callbacks and process details stored in the DB); only the final reply is sent back to personal WeChat/DingTalk/Lark/WeCom in one message (no SSE). Default: `robot_default_agent_mode: eino_single`.
|
||||
|
||||
@@ -8,6 +8,7 @@
|
||||
- [贡献规范](contributing-guide.md):新增 API、配置、工具、前端、数据库、高风险能力和文档的 checklist。
|
||||
- [配置参考](configuration.md):`config.yaml` 字段、热应用边界、参数建议和源码锚点。
|
||||
- [安全模型](security-model.md):信任边界、HITL、工具执行、C2/WebShell 与数据安全。
|
||||
- [RBAC 权限管理](rbac.md):平台用户、系统/自定义角色、权限目录、逐权限 Scope、资源授权、Agent/MCP/机器人边界与 API 示例。
|
||||
- [架构说明](architecture.md):请求路径、模块关系、复杂度热点和设计取舍。
|
||||
- [API 参考](api-reference.md):认证、OpenAPI、SSE、稳定性分层和常用接口。
|
||||
- [排错指南](troubleshooting.md):诊断顺序、最小命令、常见误判和故障模板。
|
||||
|
||||
@@ -0,0 +1,386 @@
|
||||
# CyberStrikeAI RBAC 使用与管理指南
|
||||
|
||||
[English](../en-US/rbac.md)
|
||||
|
||||
CyberStrikeAI 是可执行 Agent、MCP、WebShell、C2 和批量任务的安全自动化平台。RBAC 不仅控制页面是否可见,还会贯穿 HTTP API、资源查询、Agent 上下文、内置/外部 MCP 工具、后台任务和机器人执行链路。
|
||||
|
||||
---
|
||||
|
||||
## 一、先区分两种“角色”
|
||||
|
||||
| 概念 | 管理入口 | 作用 |
|
||||
|------|----------|------|
|
||||
| **平台角色(RBAC Role)** | 左侧 **平台权限** | 决定用户能调用哪些功能、能访问哪些资源 |
|
||||
| **AI 测试角色(Agent Role)** | 左侧 **角色** / `roles/*.yaml` | 决定 Agent 的提示词、测试方法和可选工具集合 |
|
||||
|
||||
AI 测试角色不是安全授权边界。即使选择了“渗透测试”角色,用户仍必须拥有对应的平台权限;反过来,RBAC 有权限也不会自动改变 Agent 提示词。
|
||||
|
||||
---
|
||||
|
||||
## 二、授权模型
|
||||
|
||||
一次访问同时满足以下条件才会放行:
|
||||
|
||||
```text
|
||||
有效账号
|
||||
+ 路由/工具所需 permission
|
||||
+ 该 permission 对应的 scope
|
||||
+ 目标资源 owner / 显式授权 / 父资源继承
|
||||
+ 全局操作的额外限制
|
||||
```
|
||||
|
||||
处理链路:
|
||||
|
||||
1. 登录后签发 Bearer Token,会话包含用户、角色、权限和逐权限 Scope。
|
||||
2. HTTP 中间件把路由映射为权限,例如 `GET /api/projects` → `project:read`。
|
||||
3. 对带资源 ID 的请求继续校验 owner、显式资源授权或父资源继承。
|
||||
4. Agent 启动时把不可变 Principal 写入 `context.Context`。
|
||||
5. 内置 MCP 工具根据工具和参数再次检查权限与资源;外部 MCP 也有独立限制。
|
||||
6. 拒绝事件写入 RBAC/审计日志。
|
||||
|
||||
前端隐藏按钮只用于改善体验,不构成安全边界;真正的拒绝发生在服务端。
|
||||
|
||||
---
|
||||
|
||||
## 三、内置平台角色
|
||||
|
||||
| 角色 | Scope | 默认能力 |
|
||||
|------|-------|----------|
|
||||
| **管理员 `admin`** | `all` | 所有已知权限,包括 RBAC、配置、终端、审计删除和全局定义管理 |
|
||||
| **操作员 `operator`** | `assigned` | 日常读写与执行能力;不含 RBAC、核心配置、终端、审计管理、外部 MCP 执行和部分全局定义写权限 |
|
||||
| **审计员 `auditor`** | `all` | 各模块只读权限和 `audit:read`,不执行写操作 |
|
||||
| **只读用户 `viewer`** | `assigned` | 各模块只读权限,仅查看被授权范围 |
|
||||
|
||||
系统角色不可修改或删除,升级时会按当前版本的权限目录重新构建授权,避免旧版本残留权限。需要不同组合时创建自定义角色。
|
||||
|
||||
没有分配任何角色的账号仍可登录,但基本没有业务权限;不要依赖“无角色”作为完整岗位配置。
|
||||
|
||||
---
|
||||
|
||||
## 四、权限命名与目录
|
||||
|
||||
权限使用 `模块:动作` 命名。常见动作:
|
||||
|
||||
- `read`:查看、列表、查询、导出。
|
||||
- `write`:创建、更新、执行或管理。
|
||||
- `delete`:删除。
|
||||
- `execute`:执行 Agent、终端、工作流或特定能力。
|
||||
|
||||
当前权限按模块分组如下;运行版本的权威目录以“平台权限”页面或 `GET /api/rbac/metadata` 为准。
|
||||
|
||||
| 模块 | 权限 |
|
||||
|------|------|
|
||||
| 账号 | `auth:self` |
|
||||
| 仪表盘 | `dashboard:read` |
|
||||
| 对话 | `chat:read`、`chat:write`、`chat:delete` |
|
||||
| Agent | `agent:execute`、`agent:local-execute` |
|
||||
| HITL | `hitl:read`、`hitl:write` |
|
||||
| 任务 | `tasks:read`、`tasks:write`、`tasks:delete` |
|
||||
| 项目 | `project:read`、`project:write`、`project:delete` |
|
||||
| 漏洞 | `vulnerability:read`、`vulnerability:write`、`vulnerability:delete` |
|
||||
| WebShell | `webshell:read`、`webshell:write`、`webshell:delete` |
|
||||
| C2 | `c2:read`、`c2:write`、`c2:delete` |
|
||||
| MCP | `mcp:read`、`mcp:execute`、`mcp:write`、`mcp:external:execute` |
|
||||
| 知识库 | `knowledge:read`、`knowledge:write`、`knowledge:delete` |
|
||||
| Skills | `skills:read`、`skills:write`、`skills:delete` |
|
||||
| Markdown Agents | `agents:read`、`agents:write`、`agents:delete` |
|
||||
| AI 测试角色 | `roles:read`、`roles:write`、`roles:delete` |
|
||||
| 工作流 | `workflow:read`、`workflow:execute`、`workflow:write`、`workflow:delete` |
|
||||
| 系统配置 | `config:read`、`config:write` |
|
||||
| 终端 | `terminal:execute` |
|
||||
| 审计 | `audit:read`、`audit:delete` |
|
||||
| RBAC | `rbac:read`、`rbac:write` |
|
||||
| 通知 | `notification:read`、`notification:write` |
|
||||
| 机器人 | `robot:read`、`robot:write` |
|
||||
| 文件 | `files:read`、`files:write`、`files:delete` |
|
||||
| 攻击链 | `attackchain:read`、`attackchain:write` |
|
||||
| FOFA | `fofa:execute` |
|
||||
| OpenAPI | `openapi:read` |
|
||||
| 对话分组 | `group:read`、`group:write`、`group:delete` |
|
||||
| 执行监控 | `monitor:read`、`monitor:write`、`monitor:delete` |
|
||||
|
||||
特殊权限说明:
|
||||
|
||||
- `agent:execute` 允许运行 Agent,但不自动允许本地文件系统、Shell 或任意配置命令。
|
||||
- `agent:local-execute` 是本地执行兜底权限,应仅授予可信操作员。
|
||||
- `mcp:execute` 用于访问认证后的 MCP HTTP 入口。
|
||||
- `mcp:external:execute` 用于 Agent 调用外部 MCP 工具,当前还要求该权限的 Scope 为 `all`。
|
||||
- 管理外部 MCP 配置使用 `mcp:write`,与执行外部工具是两项权限。
|
||||
- `robot:write` 管理机器人配置和测试入口;机器人聊天本身使用绑定用户或服务账号的业务权限。
|
||||
|
||||
---
|
||||
|
||||
## 五、资源 Scope
|
||||
|
||||
每个角色包含一个 Scope:
|
||||
|
||||
| Scope | 含义 | 适合场景 |
|
||||
|-------|------|----------|
|
||||
| `all` | 访问该权限覆盖的所有资源 | 管理员、全局审计员 |
|
||||
| `assigned` | 访问管理员指定的资源及系统支持的父资源继承范围 | 项目成员、指定资产操作员 |
|
||||
| `own` | 以本人创建/归属资源为主;部分资源仍可通过显式授权或父资源关系访问 | 个人工作区、机器人独立身份 |
|
||||
|
||||
权限和 Scope 是绑定在一起计算的。一个用户可拥有多个角色,权限取并集;**同一个权限**的 Scope 取最宽值:
|
||||
|
||||
```text
|
||||
all > assigned > own
|
||||
```
|
||||
|
||||
示例:
|
||||
|
||||
- “全局审计”角色:`project:read` + `all`
|
||||
- “个人项目编辑”角色:`project:write` + `own`
|
||||
|
||||
最终结果是:
|
||||
|
||||
```text
|
||||
project:read → all
|
||||
project:write → own
|
||||
```
|
||||
|
||||
全局读取不会把无关的写权限扩大为全局写入。服务端授权必须使用 `ScopeFor(permission)`,不能使用用户的最宽总 Scope。
|
||||
|
||||
### 全局对象限制
|
||||
|
||||
部分对象是进程级共享定义,没有 owner。即使用户拥有 `write`,若该权限 Scope 不是 `all`,服务端仍拒绝修改。例如:
|
||||
|
||||
- AI 测试角色、Skills、Markdown Agents。
|
||||
- 外部 MCP 配置。
|
||||
- 机器人配置。
|
||||
- 工作流定义。
|
||||
- 知识库写操作(搜索除外)。
|
||||
- HITL 全局白名单、默认审核方和审计策略。
|
||||
- C2 Profile 写操作。
|
||||
- 部分全局监控统计。
|
||||
|
||||
---
|
||||
|
||||
## 六、资源归属、显式授权与继承
|
||||
|
||||
可在“平台权限 → 成员详情 → 资源授权”中给用户分配资源。当前可直接选择的主要类型:
|
||||
|
||||
- 项目 `project`
|
||||
- 对话 `conversation`
|
||||
- 漏洞 `vulnerability`
|
||||
- WebShell `webshell`
|
||||
- 批量任务队列 `batch_task`
|
||||
- C2 Listener `c2_listener`
|
||||
|
||||
一次批量授权最多 100 个资源。重复授权会跳过,不会创建重复记录。
|
||||
|
||||
部分子资源会继承父资源访问能力:
|
||||
|
||||
| 子资源 | 可继承的父资源 |
|
||||
|--------|----------------|
|
||||
| 对话 | 所属项目 |
|
||||
| 漏洞 | 所属项目或关联对话 |
|
||||
| 消息、过程详情、攻击链 | 所属对话 |
|
||||
| C2 Session | Listener |
|
||||
| C2 Task / 文件 /事件 | Session、Task 或 Listener 链路 |
|
||||
|
||||
因此,给用户授权一个项目,通常不需要再逐个授权该项目中的每条对话和漏洞。仍应以具体页面/API 的服务端检查结果为准。
|
||||
|
||||
---
|
||||
|
||||
## 七、Web 管理流程
|
||||
|
||||
### 7.1 创建用户
|
||||
|
||||
1. 使用管理员进入左侧 **平台权限**。
|
||||
2. 创建平台用户,设置用户名、显示名称、至少 8 位密码和启用状态。
|
||||
3. 分配一个或多个平台角色。
|
||||
4. 若角色 Scope 为 `assigned`,继续配置资源授权。
|
||||
5. 让用户重新登录并在右上角用户菜单确认角色、权限数量和 Scope。
|
||||
|
||||
### 7.2 创建自定义角色
|
||||
|
||||
1. 新建平台角色并填写清晰的岗位名称与说明。
|
||||
2. 选择 `all`、`assigned` 或 `own`。
|
||||
3. 只勾选岗位实际需要的权限。
|
||||
4. 先用测试账号验证列表、详情、写操作、删除和 Agent 工具调用。
|
||||
5. 再批量分配给正式用户。
|
||||
|
||||
系统角色不可编辑;复制其思路创建自定义角色即可。
|
||||
|
||||
### 7.3 权限变更何时生效
|
||||
|
||||
- 更新用户、密码、启用状态或角色后,该用户现有会话会被撤销,需要重新登录。
|
||||
- 更新或删除自定义角色后,平台会撤销全部现有会话,所有用户需重新登录。
|
||||
- 机器人每条消息实时解析绑定用户/服务账号权限;用户禁用或角色调整会立即影响下一条消息。
|
||||
- 后台批量任务会根据任务 owner 重新解析 Principal,不应依赖创建任务时的前端状态。
|
||||
|
||||
---
|
||||
|
||||
## 八、推荐角色模板
|
||||
|
||||
以下是起点,不是固定策略。
|
||||
|
||||
### 只读项目成员
|
||||
|
||||
```text
|
||||
Scope: assigned
|
||||
dashboard:read
|
||||
chat:read
|
||||
project:read
|
||||
vulnerability:read
|
||||
files:read
|
||||
attackchain:read
|
||||
```
|
||||
|
||||
### 日常安全操作员
|
||||
|
||||
```text
|
||||
Scope: assigned
|
||||
agent:execute
|
||||
chat:read / chat:write
|
||||
project:read / project:write
|
||||
vulnerability:read / vulnerability:write
|
||||
tasks:read / tasks:write
|
||||
files:read / files:write
|
||||
hitl:read / hitl:write
|
||||
```
|
||||
|
||||
只有确实需要本机命令时才增加 `agent:local-execute` 或 `terminal:execute`;需要删除时再增加对应 `:delete`。
|
||||
|
||||
### 机器人专用账号
|
||||
|
||||
```text
|
||||
Scope: own(独立工作区)或 assigned(指定项目)
|
||||
agent:execute
|
||||
chat:read / chat:write
|
||||
按需增加 project、vulnerability、knowledge 等权限
|
||||
```
|
||||
|
||||
也可使用 `admin` 作为机器人服务账号,但发送者仍需精确白名单;白名单内每个人都会获得完整权限并共享 admin 数据。详见[机器人指南](robot.md)。
|
||||
|
||||
---
|
||||
|
||||
## 九、Agent、MCP 与机器人边界
|
||||
|
||||
### Agent
|
||||
|
||||
HTTP 登录用户会被转换为不可变 Principal,传入单 Agent、多 Agent、工作流和工具执行上下文。长任务脱离 SSE 连接后仍保留身份,但不会因为前端按钮可见而绕过服务端权限。
|
||||
|
||||
### 内置 MCP
|
||||
|
||||
每个内置工具必须有显式授权策略。例如 WebShell 工具会同时检查 `webshell:read/write/delete` 和 `connection_id` 的资源访问;漏洞、项目、任务与 C2 工具也会检查参数指向的资源。
|
||||
|
||||
未登记授权策略的内置工具默认拒绝。普通本地/配置工具需要 `agent:local-execute`。
|
||||
|
||||
### 外部 MCP
|
||||
|
||||
Agent 调用外部 MCP 工具需要 `mcp:external:execute`,且当前要求 Scope 为 `all`。这是因为外部服务的资源模型通常不受本地 owner/assignment 约束。
|
||||
|
||||
### 机器人
|
||||
|
||||
- `user_binding`:平台发送者绑定自己的 RBAC 用户。
|
||||
- `service_account`:精确白名单发送者统一使用一个 RBAC 用户。
|
||||
- 平台验签只做来源认证,不代替业务授权。
|
||||
- 发送 `身份` / `whoami` 可检查实际 Principal。
|
||||
|
||||
---
|
||||
|
||||
## 十、RBAC API
|
||||
|
||||
所有接口使用:
|
||||
|
||||
```http
|
||||
Authorization: Bearer <token>
|
||||
```
|
||||
|
||||
管理接口需要 `rbac:read` 或 `rbac:write`,资源选择器需要 `rbac:write`。
|
||||
|
||||
| 方法 | 路径 | 说明 |
|
||||
|------|------|------|
|
||||
| GET | `/api/rbac/me` | 当前用户、角色、权限、总 Scope 与逐权限 Scope |
|
||||
| GET | `/api/rbac/metadata` | 权限目录、角色、角色权限和 Scope 列表 |
|
||||
| GET/POST | `/api/rbac/users` | 列出/创建用户 |
|
||||
| PUT/DELETE | `/api/rbac/users/:id` | 更新/删除用户 |
|
||||
| GET/POST | `/api/rbac/roles` | 列出/创建角色 |
|
||||
| PUT/DELETE | `/api/rbac/roles/:id` | 更新/删除自定义角色 |
|
||||
| GET | `/api/rbac/resources?type=project&q=...` | 分页搜索可授权资源 |
|
||||
| GET/POST | `/api/rbac/resource-assignments` | 列出/创建资源授权 |
|
||||
| DELETE | `/api/rbac/resource-assignments/:id` | 撤销资源授权 |
|
||||
|
||||
创建用户示例:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8080/api/rbac/users \
|
||||
-H "Authorization: Bearer $TOKEN" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"username": "operator01",
|
||||
"display_name": "安全操作员 01",
|
||||
"password": "change-me-123",
|
||||
"enabled": true,
|
||||
"roles": ["operator"]
|
||||
}'
|
||||
```
|
||||
|
||||
创建自定义角色示例:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8080/api/rbac/roles \
|
||||
-H "Authorization: Bearer $TOKEN" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"name": "项目审计员",
|
||||
"description": "只读查看指定项目",
|
||||
"scope": "assigned",
|
||||
"permissions": ["chat:read", "project:read", "vulnerability:read"]
|
||||
}'
|
||||
```
|
||||
|
||||
批量授权项目示例:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8080/api/rbac/resource-assignments \
|
||||
-H "Authorization: Bearer $TOKEN" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"user_id": "USER_ID",
|
||||
"resource_type": "project",
|
||||
"resource_ids": ["PROJECT_ID_1", "PROJECT_ID_2"]
|
||||
}'
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 十一、审计与运维建议
|
||||
|
||||
- 使用个人账号管理平台,避免多人共享管理员密码。
|
||||
- 自定义角色按岗位命名,描述中写明用途和负责人。
|
||||
- 高风险权限单独审批:`terminal:execute`、`agent:local-execute`、`c2:write/delete`、`webshell:write/delete`、`rbac:write`、`config:write`。
|
||||
- 定期检查 `all` Scope 角色、服务账号、机器人白名单和长期未使用用户。
|
||||
- 用户离职时先禁用账号,再撤销机器人绑定、资源授权和会话。
|
||||
- 在日志审计中关注 `rbac/access_denied`、角色/用户变更、资源授权、机器人服务账号执行。
|
||||
- 配合 HITL 控制高风险工具;RBAC 允许调用不等于可以跳过审批。
|
||||
|
||||
---
|
||||
|
||||
## 十二、常见问题
|
||||
|
||||
### 页面按钮看不到
|
||||
|
||||
检查用户是否有对应权限;前端会根据 `/api/rbac/me` 隐藏无权操作。直接调用 API 仍会由服务端拒绝。
|
||||
|
||||
### 有权限但返回“无权访问该资源”
|
||||
|
||||
检查该权限的 Scope,而不是只看用户总 Scope;再检查资源 owner、显式授权和父资源授权。
|
||||
|
||||
### 角色改了但用户仍是旧权限
|
||||
|
||||
角色变更会撤销会话。让用户重新登录;机器人下一条消息会重新解析权限。
|
||||
|
||||
### `write` 权限存在但全局配置仍被拒绝
|
||||
|
||||
全局对象写操作要求对应权限的 Scope 为 `all`。创建一个 `all` Scope 的专用管理角色,而不是扩大无关权限。
|
||||
|
||||
### Agent 能对话但不能运行命令
|
||||
|
||||
`agent:execute` 与 `agent:local-execute` 分离。按需授予本地执行权限,并结合 HITL、工具白名单和审计。
|
||||
|
||||
### 外部 MCP 提示需要 global scope
|
||||
|
||||
除 `mcp:external:execute` 外,该权限的 Scope 还必须为 `all`。外部 MCP 的数据边界不由本地资源授权自动保护。
|
||||
|
||||
+134
-19
@@ -2,7 +2,7 @@
|
||||
|
||||
[English](../en-US/robot.md)
|
||||
|
||||
本文档说明如何通过**个人微信**、**钉钉**、**飞书**与 **企业微信** 与 CyberStrikeAI 对话(长连接 / 回调模式),在手机端即可使用,无需在服务器上打开网页。按下面步骤操作可避免常见弯路。
|
||||
本文档说明如何通过**个人微信、企业微信、钉钉、飞书、Telegram、Slack、Discord 和 QQ 机器人**使用 CyberStrikeAI,包括平台接入、RBAC 身份绑定、服务账号白名单、命令、验证与故障排查。
|
||||
|
||||
---
|
||||
|
||||
@@ -15,10 +15,18 @@
|
||||
- **个人微信**:点击「微信 / iLink」→「生成二维码并绑定」,用微信扫码确认(见 [3.4 个人微信](#34-个人微信-wechat--ilink))
|
||||
- **钉钉**:勾选并填写 Client ID / Client Secret
|
||||
- **飞书**:勾选并填写 App ID / App Secret
|
||||
5. 点击 **应用配置** 保存(微信扫码绑定成功后会**自动保存并启用**,一般无需再点)
|
||||
6. **重启 CyberStrikeAI 应用**(钉钉/飞书:只保存不重启,长连接不会建立;微信绑定成功后会自动重启连接,通常无需手动重启)
|
||||
5. 点击 **应用配置** 保存;程序会自动重启对应机器人连接。微信扫码绑定成功后会自动保存并启用,一般无需再点。
|
||||
|
||||
配置会写入 `config.yaml` 的 `robots` 段,也可在配置文件中直接编辑。**修改钉钉/飞书配置后必须重启,长连接才会生效。** 个人微信绑定成功后程序会自动写入 `robots.wechat` 并重启 iLink 长轮询。
|
||||
配置会写入 `config.yaml` 的 `robots` 段,也可在配置文件中直接编辑。通过 Web 点击“应用配置”会自动重启对应连接;直接手工修改 `config.yaml` 时,需要重启 CyberStrikeAI 进程。个人微信绑定成功后程序会自动写入 `robots.wechat` 并重启 iLink 长轮询。
|
||||
|
||||
### 最短使用路径
|
||||
|
||||
平台连接成功后,不要直接开始普通对话,先完成业务身份配置:
|
||||
|
||||
- **多人使用**:机器人设置选择“逐用户绑定” → 每位用户在 Web 右上角头像生成绑定码 → 在机器人中发送绑定命令 → 发送 `身份` 验证。
|
||||
- **只有自己使用**:先在机器人中发送 `身份` 复制发送者 ID → 机器人设置选择“专用服务账号” → User ID 填 `admin` 或其他 RBAC 用户 → 粘贴发送者白名单 → 应用配置 → 再次发送 `身份` 验证。
|
||||
|
||||
看到“鉴权状态:已授权”且“实际身份”正确后,即可直接发送普通文本与 AI 对话。
|
||||
|
||||
---
|
||||
|
||||
@@ -300,12 +308,88 @@
|
||||
|
||||
---
|
||||
|
||||
## 四、机器人命令
|
||||
## 四、RBAC 鉴权与机器人命令
|
||||
|
||||
平台 Token、签名或长连接凭证只负责证明“消息来自该平台”;真正能执行哪些操作,由 CyberStrikeAI 的 RBAC 决定。每个机器人实例都必须选择一种业务鉴权模式。
|
||||
|
||||
### 4.1 应该选择哪种模式
|
||||
|
||||
| 使用场景 | 推荐模式 | 身份与数据范围 |
|
||||
|----------|----------|----------------|
|
||||
| 企业微信、飞书、钉钉、Slack 等多人共享机器人 | `user_binding` | 每个发送者绑定自己的 Web 用户,权限和数据互相隔离 |
|
||||
| 个人微信、单人专属机器人、固定自动化入口 | `service_account` | 白名单发送者统一使用配置的 RBAC 用户,并共享该账号的数据 |
|
||||
|
||||
两种模式都会在**每条消息**执行前重新读取用户状态、角色、逐权限 Scope 和资源授权。用户被禁用或权限被收回后,下一条消息立即失效。
|
||||
|
||||
机器人执行普通 AI 对话至少需要以下权限:
|
||||
|
||||
```text
|
||||
agent:execute
|
||||
chat:read
|
||||
chat:write
|
||||
```
|
||||
|
||||
使用项目、角色、本地命令、WebShell、C2 或外部 MCP 时,还需按功能增加对应权限。删除对话需要 `chat:delete`。
|
||||
|
||||
### 4.2 逐用户绑定模式(默认)
|
||||
|
||||
管理员操作:
|
||||
|
||||
1. 系统设置 → 机器人设置 → 选择平台。
|
||||
2. 在“业务鉴权策略”中选择“逐用户绑定(`user_binding`)”。
|
||||
3. 点击“应用配置”。
|
||||
|
||||
每位使用者操作:
|
||||
|
||||
1. 登录 CyberStrikeAI Web,点击右上角头像 → **绑定机器人账号**。
|
||||
2. 点击 **生成绑定码**,页面开始 5 分钟倒计时。
|
||||
3. 在目标机器人中发送页面给出的完整命令,例如 `绑定 7C6E-BD4C`。
|
||||
4. 发送 `身份` 或 `whoami`,确认“鉴权状态:已授权”且“实际身份”是自己的 Web 用户。
|
||||
|
||||
绑定码仅保存哈希、只能使用一次。倒计时结束后前端会标记失效、禁用复制并刷新绑定列表;服务端也会拒绝过期码。重新生成会让此前尚未使用的旧码立即失效。用户可发送 `解绑`,或在 Web 绑定窗口中撤销绑定。
|
||||
|
||||
### 4.3 专用服务账号模式
|
||||
|
||||
1. 先让机器人正常连接平台。
|
||||
2. 目标使用者向机器人发送 `身份` / `whoami`,复制返回的完整“发送者 ID”。个人微信的 ID 通常形如 `xxxx@im.wechat`;必须以命令返回值为准,不能用 `ilink_bot_id` 或配置中的 `ilink_user_id` 代替。
|
||||
3. 系统设置 → 机器人设置 → 选择平台 → 业务鉴权策略选择“专用服务账号(`service_account`)”。
|
||||
4. 填写服务账号的 **RBAC User ID**,不是显示名称。可以填写 `admin`;此时白名单发送者拥有完整平台权限,界面会显示红色风险提示。
|
||||
5. 在“允许的平台发送者 ID”中每行填写一个完整 ID。必须精确匹配、区分大小写,不允许 `*` 通配符。
|
||||
6. 点击“应用配置”,再发送 `身份` 确认“实际身份”和角色正确。
|
||||
|
||||
示例:
|
||||
|
||||
```yaml
|
||||
robots:
|
||||
wechat:
|
||||
auth:
|
||||
mode: service_account
|
||||
service_user_id: admin
|
||||
allowed_external_users:
|
||||
- "o9cq806s32Sm2_kyOmkyaV7Rn1lU@im.wechat"
|
||||
```
|
||||
|
||||
服务账号模式不接受 `绑定` / `解绑` 命令。多个白名单发送者会共享服务账号创建的对话、项目和其他 `own` 范围资源;若不希望共享,请使用逐用户绑定。
|
||||
|
||||
### 4.4 如何检查当前身份
|
||||
|
||||
发送:
|
||||
|
||||
```text
|
||||
身份
|
||||
```
|
||||
|
||||
返回内容包含:平台、真实发送者 ID、鉴权模式、鉴权状态、实际 RBAC 用户、RBAC User ID、平台角色、资源范围和有效权限数量。不在服务账号白名单中的发送者只会看到拒绝状态,不会看到服务账号详情。
|
||||
|
||||
### 4.5 命令列表
|
||||
|
||||
在任一已接入平台(钉钉/飞书/微信/Telegram/Slack/Discord/QQ 等)向机器人发送以下**文本命令**(仅支持文本):
|
||||
|
||||
| 命令 | 说明 |
|
||||
|------|------|
|
||||
| **绑定 \<绑定码\>** | 将当前平台发送者绑定到生成绑定码的 RBAC 用户 |
|
||||
| **解绑** | 解除当前平台账号绑定;也可在 Web 端的绑定列表中撤销 |
|
||||
| **身份** 或 **whoami** | 显示平台发送者 ID、鉴权模式、绑定状态及当前实际 RBAC 用户、角色和资源范围 |
|
||||
| **帮助** | 显示命令帮助与说明 |
|
||||
| **列表** 或 **对话列表** | 列出所有对话的标题与对话 ID |
|
||||
| **切换 \<对话ID\>** 或 **继续 \<对话ID\>** | 指定对话 ID,后续消息在该对话中继续 |
|
||||
@@ -320,6 +404,8 @@
|
||||
|
||||
除以上命令外,**直接输入任意文字**会作为用户消息发给 AI,与 Web 端对话逻辑一致(渗透测试/安全分析等)。
|
||||
|
||||
群聊消息按实际发送者鉴权,不使用群 ID 作为业务身份。服务账号模式除外:白名单发送者会明确共享配置的服务账号权限和资源。
|
||||
|
||||
---
|
||||
|
||||
## 五、如何使用(要 @ 机器人吗?)
|
||||
@@ -338,14 +424,17 @@
|
||||
|
||||
1. CyberStrikeAI Web 端 → 系统设置 → 机器人设置 → **微信 / iLink** → **生成二维码并绑定**。
|
||||
2. 手机微信扫码确认(如需配对数字则在 Web 页填写)。
|
||||
3. 绑定成功后,在手机微信私聊中发「帮助」测试。
|
||||
3. 在手机微信私聊中发送 `身份`,复制发送者 ID。
|
||||
4. 回到机器人设置选择 `user_binding`,或选择 `service_account` 并填写服务账号与发送者白名单。
|
||||
5. 点击应用配置,在微信中再次发送 `身份`,确认实际 RBAC 身份后再发送普通消息。
|
||||
|
||||
**钉钉 / 飞书**
|
||||
|
||||
1. **在开放平台**:按第三节完成应用创建、凭证复制、机器人开通(钉钉务必选 **Stream 模式**)、权限与发布。
|
||||
2. **在 CyberStrikeAI**:系统设置 → 机器人设置 → 勾选对应平台,粘贴 Client ID/App ID、Client Secret/App Secret → 点击 **应用配置**。
|
||||
3. **重启 CyberStrikeAI 进程**(否则长连接不会建立)。
|
||||
4. **在手机钉钉/飞书**:找到该机器人(单聊直接发,群聊需 @机器人),发「帮助」或任意内容测试。
|
||||
3. **选择鉴权模式**:多人使用建议 `user_binding`;专用机器人配置服务账号与发送者白名单。
|
||||
4. **应用配置**:Web 会自动重启对应连接。
|
||||
5. **在手机钉钉/飞书**:找到机器人(单聊直接发,群聊需 @),先发 `身份` 检查鉴权,再发普通内容测试。
|
||||
|
||||
若发消息没反应,先看 **第九节排查** 和 **第十节常见弯路**。
|
||||
|
||||
@@ -359,6 +448,11 @@
|
||||
robots:
|
||||
wechat: # 个人微信 iLink(扫码绑定后自动写入,一般无需手填)
|
||||
enabled: true
|
||||
auth:
|
||||
mode: service_account
|
||||
service_user_id: admin
|
||||
allowed_external_users:
|
||||
- "从身份命令复制的完整发送者 ID"
|
||||
bot_token: "your_bot_token@im.bot:..."
|
||||
ilink_bot_id: "your_bot_id@im.bot"
|
||||
ilink_user_id: "your_user_id@im.wechat"
|
||||
@@ -367,10 +461,14 @@ robots:
|
||||
bot_agent: "CyberStrikeAI/1.0"
|
||||
dingtalk:
|
||||
enabled: true
|
||||
auth:
|
||||
mode: user_binding
|
||||
client_id: "your_dingtalk_app_key"
|
||||
client_secret: "your_dingtalk_app_secret"
|
||||
lark:
|
||||
enabled: true
|
||||
auth:
|
||||
mode: user_binding
|
||||
app_id: "your_lark_app_id"
|
||||
app_secret: "your_lark_app_secret"
|
||||
verify_token: ""
|
||||
@@ -400,7 +498,7 @@ robots:
|
||||
sandbox: true
|
||||
```
|
||||
|
||||
修改钉钉/飞书/企业微信/Telegram/Slack/Discord/QQ 配置后,点击 **应用配置** 会自动重启对应长连接。个人微信扫码绑定成功后会自动写入并重启 iLink 连接。
|
||||
每个平台的 `auth` 独立配置;省略时默认为 `user_binding`。修改配置后,在 Web 点击 **应用配置** 会自动重启对应连接;手工编辑 YAML 则需重启进程。个人微信扫码绑定成功后会自动写入并重启 iLink 连接。
|
||||
|
||||
---
|
||||
|
||||
@@ -408,20 +506,24 @@ robots:
|
||||
|
||||
在未安装钉钉或飞书时,可用**测试接口**验证机器人逻辑是否正常:
|
||||
|
||||
1. 先登录 CyberStrikeAI Web 端(保证有登录态)。
|
||||
2. 使用 curl 调用测试接口(需携带登录后的 Cookie):
|
||||
1. 使用具有全局 `robot:write` 权限的账号登录并获取 Bearer Token。
|
||||
2. 使用 curl 调用测试接口:
|
||||
|
||||
```bash
|
||||
# 将 YOUR_COOKIE 替换为登录后获得的 Cookie(浏览器 F12 → 网络 → 任意请求 → 请求头中的 Cookie)
|
||||
# 先登录;请按实际地址、用户名和密码修改
|
||||
TOKEN=$(curl -s -X POST "http://localhost:8080/api/auth/login" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"username":"admin","password":"YOUR_PASSWORD"}' | jq -r '.token')
|
||||
|
||||
curl -X POST "http://localhost:8080/api/robot/test" \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "Cookie: YOUR_COOKIE" \
|
||||
-H "Authorization: Bearer $TOKEN" \
|
||||
-d '{"platform":"dingtalk","user_id":"test_user","text":"帮助"}'
|
||||
```
|
||||
|
||||
若返回 JSON 中含有 `"reply":"【CyberStrikeAI 机器人命令】..."`,说明命令处理正常。可再试 `"text":"列表"`、`"text":"当前"` 等。
|
||||
若返回 JSON 中含有 `"reply":"【CyberStrikeAI 机器人命令】..."`,说明命令处理正常。`帮助`、`版本` 和 `身份` 可在未绑定时执行;`列表`、`当前` 和普通 AI 消息会走真实 RBAC,测试用 `platform + user_id` 必须已经绑定,或与服务账号模式的发送者白名单精确匹配。
|
||||
|
||||
接口说明:`POST /api/robot/test`(需登录),请求体 `{"platform":"可选","user_id":"可选","text":"必填"}`,响应 `{"reply":"回复内容"}`。
|
||||
接口说明:`POST /api/robot/test`(需全局 `robot:write`),请求体 `{"platform":"可选","user_id":"可选","text":"必填"}`,响应 `{"reply":"回复内容"}`。该接口仅模拟机器人业务逻辑,不验证第三方平台签名或长连接。
|
||||
|
||||
---
|
||||
|
||||
@@ -458,8 +560,8 @@ curl -X POST "http://localhost:8080/api/robot/test" \
|
||||
1. **Client ID / Client Secret 是否与开放平台完全一致**
|
||||
从「凭证与基础信息」里**复制粘贴**,不要手打。注意数字 **0** 与字母 **o**、数字 **1** 与字母 **l**(例如 `ding9gf9tiozuc504aer` 中间是 **504** 不是 5o4)。
|
||||
|
||||
2. **是否在保存配置后重启了应用**
|
||||
机器人长连接在**应用启动时**建立。在 Web 端点击「应用配置」只写入配置文件,**必须重启 CyberStrikeAI 进程**后钉钉连接才会生效。
|
||||
2. **配置是否已应用**
|
||||
在 Web 端修改后必须点击“应用配置”,程序会自动重启对应连接。若直接手工编辑 `config.yaml`,则需重启 CyberStrikeAI 进程。
|
||||
|
||||
3. **看程序日志**
|
||||
- 启动后应看到:`钉钉 Stream 正在连接…`、`钉钉 Stream 已启动(无需公网),等待收消息`。
|
||||
@@ -469,6 +571,15 @@ curl -X POST "http://localhost:8080/api/robot/test" \
|
||||
4. **开放平台侧**
|
||||
应用需已**发布**;在「机器人」能力中需开启**流式接入(Stream)** 用于接收消息(仅 HTTP 回调不够);权限管理里需有机器人接收、发送消息等权限。
|
||||
|
||||
### 9.3 收到回复但提示未绑定、白名单拒绝或权限不足
|
||||
|
||||
1. 先发送 `身份`,查看“鉴权模式”和“鉴权状态”。
|
||||
2. `user_binding` 显示未绑定:在 Web 右上角头像中生成绑定码,并在同一个平台账号中发送完整绑定命令。绑定码过期或已经使用时需重新生成。
|
||||
3. `service_account` 显示白名单拒绝:把 `身份` 返回的完整发送者 ID 原样加入当前平台的白名单,注意大小写、租户前缀和 `@im.wechat` 等后缀。
|
||||
4. 显示实际身份但提示缺少权限:在“平台权限”检查该 RBAC 用户的角色。普通 AI 对话至少需要 `agent:execute`、`chat:read`、`chat:write`。
|
||||
5. 服务账号不存在或被禁用:应用配置会拒绝保存;恢复用户或选择其他已启用 RBAC 用户。
|
||||
6. 使用 `admin` 时仍被拒绝:通常是发送者不在精确白名单中,而不是 admin 权限不足。
|
||||
|
||||
---
|
||||
|
||||
## 十、常见弯路(避免踩坑)
|
||||
@@ -476,7 +587,11 @@ curl -X POST "http://localhost:8080/api/robot/test" \
|
||||
- **个人微信与企业微信混淆**:个人微信走 `robots.wechat` + Web 扫码绑定;企业微信走 `robots.wecom` + 管理后台回调 URL,二者完全不同。
|
||||
- **个人微信二维码过期**:二维码约 5 分钟有效,过期需重新生成,不要一直扫旧码。
|
||||
- **用错了机器人类型**:在钉钉**群里**添加的「自定义」机器人(Webhook + 加签)**不能**用来做对话,本程序只支持**开放平台「企业内部应用」**里的机器人。
|
||||
- **只保存没重启**:钉钉/飞书改完配置后必须**重启应用**,否则长连接不会建立(个人微信扫码绑定会自动重启连接)。
|
||||
- **改完没有点应用配置**:Web 中修改机器人配置后要点击“应用配置”;程序会自动重启对应连接。只有手工编辑 YAML 时才需要重启进程。
|
||||
- **把 Bot ID 当成发送者 ID**:服务账号白名单必须填写 `身份` 命令返回的“发送者 ID”,不要填 `ilink_bot_id`、`ilink_user_id`、群 ID 或显示昵称。
|
||||
- **绑定码过期后继续使用**:绑定码 5 分钟有效且只能使用一次;新生成的码会让旧码立即失效。
|
||||
- **服务账号误以为数据隔离**:同一服务账号白名单中的发送者共享该账号的对话和 `own` 范围资源;需要隔离时应使用 `user_binding`。
|
||||
- **admin 配置后任意人都能用**:不会。即使服务账号是 `admin`,发送者仍必须与白名单精确匹配;但白名单中的人将拥有完整权限。
|
||||
- **Client ID 抄错**:开放平台是 `504` 就填 `504`,不要填成 `5o4`;尽量用复制粘贴。
|
||||
- **钉钉只开了 HTTP 回调没开 Stream**:本程序通过 **Stream 长连接**收消息,开放平台里机器人的消息接收方式必须选 **Stream 模式**。
|
||||
- **应用没发布**:开放平台里修改了机器人或权限后,要在「版本管理与发布」里**发布新版本**,否则不生效。
|
||||
@@ -487,5 +602,5 @@ curl -X POST "http://localhost:8080/api/robot/test" \
|
||||
|
||||
- 各平台均**仅处理文本消息**;其他类型(如图片、语音)会提示暂不支持或忽略。
|
||||
- 个人微信仅支持**私聊**,不支持群聊 @ 机器人。
|
||||
- 会话与 Web 端共用同一套对话数据:在机器人里创建的对话会在 Web 端「对话」列表中看到,反之亦然。
|
||||
- 会话与 Web 端共用同一套数据:`user_binding` 下归属于绑定用户;`service_account` 下归属于服务账号,并由白名单发送者共享。
|
||||
- 机器人执行与 **Eino 单/多代理** 相同逻辑(`ProcessMessageForRobot`,含进度回调与过程详情入库),仅不向客户端推送 SSE,最后一次性回复个人微信/钉钉/飞书/企业微信。默认 `robot_default_agent_mode: eino_single`。
|
||||
|
||||
+18
-15
@@ -91,15 +91,15 @@ func NewAgent(cfg *config.OpenAIConfig, agentCfg *config.AgentConfig, mcpServer
|
||||
llmClient := openai.NewClient(cfg, httpClient, logger)
|
||||
|
||||
return &Agent{
|
||||
openAIClient: llmClient,
|
||||
config: cfg,
|
||||
agentConfig: agentCfg,
|
||||
mcpServer: mcpServer,
|
||||
externalMCPMgr: externalMCPMgr,
|
||||
logger: logger,
|
||||
maxIterations: maxIterations,
|
||||
toolNameMapping: make(map[string]string), // 初始化工具名称映射
|
||||
toolDescriptionMode: "short",
|
||||
openAIClient: llmClient,
|
||||
config: cfg,
|
||||
agentConfig: agentCfg,
|
||||
mcpServer: mcpServer,
|
||||
externalMCPMgr: externalMCPMgr,
|
||||
logger: logger,
|
||||
maxIterations: maxIterations,
|
||||
toolNameMapping: make(map[string]string), // 初始化工具名称映射
|
||||
toolDescriptionMode: "short",
|
||||
}
|
||||
}
|
||||
|
||||
@@ -120,6 +120,9 @@ type ChatMessage struct {
|
||||
ToolName string `json:"tool_name,omitempty"`
|
||||
// ReasoningContent 对应 OpenAI/DeepSeek 的 reasoning_content;思考模式 + 工具调用后续跑须回传(见 DeepSeek 文档)。
|
||||
ReasoningContent string `json:"reasoning_content,omitempty"`
|
||||
// ModelFacingTrace is runtime-only metadata: true means Content was already the exact
|
||||
// payload seen at the model boundary and must be restored byte-for-byte.
|
||||
ModelFacingTrace bool `json:"-"`
|
||||
}
|
||||
|
||||
// MarshalJSON 自定义JSON序列化,将tool_calls中的arguments转换为JSON字符串
|
||||
@@ -780,25 +783,25 @@ func (a *Agent) ExecuteMCPToolForConversation(ctx context.Context, conversationI
|
||||
}
|
||||
|
||||
// BeginLocalToolExecution 在非 CallTool 路径工具开始时写入 running 状态,供 MCP 监控页展示「执行中」。
|
||||
func (a *Agent) BeginLocalToolExecution(toolName string, args map[string]interface{}) string {
|
||||
func (a *Agent) BeginLocalToolExecution(ctx context.Context, toolName string, args map[string]interface{}) string {
|
||||
if a == nil || a.mcpServer == nil {
|
||||
return ""
|
||||
}
|
||||
return a.mcpServer.BeginToolExecution(toolName, args)
|
||||
return a.mcpServer.BeginToolExecution(ctx, toolName, args)
|
||||
}
|
||||
|
||||
// FinishLocalToolExecution 完成 BeginLocalToolExecution 创建的记录;executionID 为空时一次性写入已完成记录。
|
||||
func (a *Agent) FinishLocalToolExecution(executionID, toolName string, args map[string]interface{}, resultText string, invokeErr error) string {
|
||||
func (a *Agent) FinishLocalToolExecution(ctx context.Context, executionID, toolName string, args map[string]interface{}, resultText string, invokeErr error) string {
|
||||
if a == nil || a.mcpServer == nil {
|
||||
return ""
|
||||
}
|
||||
return a.mcpServer.FinishToolExecution(executionID, toolName, args, resultText, invokeErr)
|
||||
return a.mcpServer.FinishToolExecution(ctx, executionID, toolName, args, resultText, invokeErr)
|
||||
}
|
||||
|
||||
// RecordLocalToolExecution 将非 CallTool 路径完成的工具调用写入 MCP 监控库(与 CallTool 落库一致),返回 executionId。
|
||||
// 用于 Eino filesystem execute 等场景,使助手气泡「渗透测试详情」与常规 MCP 一致可点进监控。
|
||||
func (a *Agent) RecordLocalToolExecution(toolName string, args map[string]interface{}, resultText string, invokeErr error) string {
|
||||
return a.FinishLocalToolExecution("", toolName, args, resultText, invokeErr)
|
||||
func (a *Agent) RecordLocalToolExecution(ctx context.Context, toolName string, args map[string]interface{}, resultText string, invokeErr error) string {
|
||||
return a.FinishLocalToolExecution(ctx, "", toolName, args, resultText, invokeErr)
|
||||
}
|
||||
|
||||
// UpdateMCPExecutionDisplayResult 将监控库中的工具结果更新为送入模型的展示正文(reduction 后)。
|
||||
|
||||
@@ -5,6 +5,31 @@ import (
|
||||
"strings"
|
||||
)
|
||||
|
||||
const ModelFacingTraceVersionKey = "cyberstrike_model_facing_trace_version"
|
||||
|
||||
// IsModelFacingTraceJSON reports whether a persisted trace was produced from the final
|
||||
// model-boundary state. Legacy traces have no version marker and require one-time migration.
|
||||
func IsModelFacingTraceJSON(traceInputJSON string) bool {
|
||||
var raw []map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(strings.TrimSpace(traceInputJSON)), &raw); err != nil {
|
||||
return false
|
||||
}
|
||||
for _, msg := range raw {
|
||||
extra, _ := msg["extra"].(map[string]interface{})
|
||||
v, ok := extra[ModelFacingTraceVersionKey]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
switch n := v.(type) {
|
||||
case float64:
|
||||
return n >= 1
|
||||
case int:
|
||||
return n >= 1
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// ParseTraceMessages 解析落库的 last_react_input(OpenAI 风格 messages JSON 数组)。
|
||||
func ParseTraceMessages(traceInputJSON string) ([]ChatMessage, error) {
|
||||
traceInputJSON = strings.TrimSpace(traceInputJSON)
|
||||
@@ -15,6 +40,7 @@ func ParseTraceMessages(traceInputJSON string) ([]ChatMessage, error) {
|
||||
if err := json.Unmarshal([]byte(traceInputJSON), &raw); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
modelFacing := IsModelFacingTraceJSON(traceInputJSON)
|
||||
out := make([]ChatMessage, 0, len(raw))
|
||||
for _, msgMap := range raw {
|
||||
msg := ChatMessage{}
|
||||
@@ -23,6 +49,7 @@ func ParseTraceMessages(traceInputJSON string) ([]ChatMessage, error) {
|
||||
continue
|
||||
}
|
||||
msg.Role = role
|
||||
msg.ModelFacingTrace = modelFacing
|
||||
if content, ok := msgMap["content"].(string); ok {
|
||||
msg.Content = content
|
||||
}
|
||||
|
||||
@@ -55,3 +55,22 @@ func TestMergeAssistantTraceOutput(t *testing.T) {
|
||||
t.Fatalf("expected merged output, got %q", out[len(out)-1].Content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseTraceMessagesMarksVersionedModelFacingTrace(t *testing.T) {
|
||||
raw := `[{"role":"system","content":"s","extra":{"cyberstrike_model_facing_trace_version":1}},{"role":"user","content":"u"},{"role":"tool","content":"exact","tool_call_id":"c1"}]`
|
||||
if !IsModelFacingTraceJSON(raw) {
|
||||
t.Fatal("versioned trace not detected")
|
||||
}
|
||||
msgs, err := ParseTraceMessages(raw)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for i, msg := range msgs {
|
||||
if !msg.ModelFacingTrace {
|
||||
t.Fatalf("message %d missing model-facing marker", i)
|
||||
}
|
||||
}
|
||||
if IsModelFacingTraceJSON(`[{"role":"user","content":"legacy"}]`) {
|
||||
t.Fatal("legacy trace incorrectly marked model-facing")
|
||||
}
|
||||
}
|
||||
|
||||
+103
-20
@@ -8,6 +8,7 @@ import (
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
@@ -16,6 +17,7 @@ import (
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/c2"
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
@@ -83,12 +85,6 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
||||
// CORS中间件
|
||||
router.Use(corsMiddleware())
|
||||
|
||||
// 认证管理器
|
||||
authManager, err := security.NewAuthManager(cfg.Auth.Password, cfg.Auth.SessionDurationHours)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("初始化认证失败: %w", err)
|
||||
}
|
||||
|
||||
// 初始化数据库
|
||||
dbPath := cfg.Database.Path
|
||||
if dbPath == "" {
|
||||
@@ -105,6 +101,21 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
||||
return nil, fmt.Errorf("初始化数据库失败: %w", err)
|
||||
}
|
||||
|
||||
// 认证管理器(数据库初始化后挂载 RBAC,以兼容旧的单密码配置)
|
||||
authManager, err := security.NewAuthManager(cfg.Auth.Password, cfg.Auth.SessionDurationHours)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("初始化认证失败: %w", err)
|
||||
}
|
||||
if err := authManager.AttachRBACStore(db); err != nil {
|
||||
return nil, fmt.Errorf("初始化RBAC失败: %w", err)
|
||||
}
|
||||
for platform, userID := range cfg.Robots.ServiceAccountUserIDs() {
|
||||
user, userErr := db.GetRBACUserByID(userID)
|
||||
if userErr != nil || !user.Enabled {
|
||||
return nil, fmt.Errorf("robots.%s.auth.service_user_id 必须指向已启用的 RBAC 用户", platform)
|
||||
}
|
||||
}
|
||||
|
||||
auditSvc := audit.NewService(db, cfg, log.Logger)
|
||||
audit.RegisterConversationCreateHook(auditSvc)
|
||||
auditSvc.PurgeExpired()
|
||||
@@ -120,6 +131,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
||||
|
||||
// 创建MCP服务器(带数据库持久化)
|
||||
mcpServer := mcp.NewServerWithStorage(log.Logger, db)
|
||||
mcpServer.SetToolAuthorizer(mcpToolAuthorizer(db))
|
||||
mcpServer.ConfigureHTTPToolCallTimeoutFromAgentMinutes(cfg.Agent.ToolTimeoutMinutes)
|
||||
|
||||
// 创建安全工具执行器
|
||||
@@ -143,6 +155,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
||||
|
||||
// 创建外部MCP管理器(使用与内部MCP服务器相同的存储)
|
||||
externalMCPMgr := mcp.NewExternalMCPManagerWithStorage(log.Logger, db)
|
||||
externalMCPMgr.SetToolAuthorizer(externalMCPToolAuthorizer())
|
||||
if cfg.ExternalMCP.Servers != nil {
|
||||
externalMCPMgr.LoadConfigs(&cfg.ExternalMCP)
|
||||
// 启动所有启用的外部MCP客户端
|
||||
@@ -360,17 +373,21 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
||||
attackChainHandler := handler.NewAttackChainHandler(db, &cfg.OpenAI, log.Logger)
|
||||
vulnerabilityHandler := handler.NewVulnerabilityHandler(db, log.Logger)
|
||||
projectHandler := handler.NewProjectHandler(db, log.Logger)
|
||||
rbacHandler := handler.NewRBACHandler(db, log.Logger)
|
||||
rbacHandler.SetAudit(auditSvc)
|
||||
rbacHandler.SetAuthManager(authManager)
|
||||
workflowHandler := handler.NewWorkflowHandler(db, log.Logger)
|
||||
workflowHandler.SetAudit(auditSvc)
|
||||
workflowHandler.SetRuntime(agent, cfg)
|
||||
vulnerabilityHandler.SetAudit(auditSvc)
|
||||
webshellHandler := handler.NewWebShellHandler(log.Logger, db)
|
||||
webshellHandler.SetAudit(auditSvc)
|
||||
chatUploadsHandler := handler.NewChatUploadsHandler(log.Logger)
|
||||
chatUploadsHandler := handler.NewChatUploadsHandler(log.Logger, db)
|
||||
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.SetDB(db)
|
||||
configHandler.SetAudit(auditSvc)
|
||||
agentHandler.SetHitlToolWhitelistSaver(configHandler)
|
||||
agentHandler.SetHitlAuditStrategySaver(configHandler)
|
||||
@@ -403,6 +420,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
||||
conversationHandler.SetTaskStopper(agentHandler)
|
||||
auditHandler := handler.NewAuditHandler(db, auditSvc, log.Logger)
|
||||
robotHandler := handler.NewRobotHandler(cfg, db, agentHandler, log.Logger)
|
||||
robotHandler.SetAudit(auditSvc)
|
||||
openAPIHandler := handler.NewOpenAPIHandler(db, log.Logger, conversationHandler, agentHandler)
|
||||
|
||||
// 创建 App 实例(部分字段稍后填充)
|
||||
@@ -535,6 +553,8 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
||||
terminalHandler,
|
||||
app.c2Handler,
|
||||
auditHandler,
|
||||
auditSvc,
|
||||
rbacHandler,
|
||||
mcpServer,
|
||||
authManager,
|
||||
openAPIHandler,
|
||||
@@ -547,17 +567,30 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error
|
||||
// mcpHandlerWithAuth 在鉴权通过后转发到 MCP 处理;若配置了 auth_header 则校验请求头,否则直接放行
|
||||
func (a *App) mcpHandlerWithAuth(w http.ResponseWriter, r *http.Request) {
|
||||
cfg := a.config.MCP
|
||||
if cfg.AuthHeader != "" {
|
||||
actual := []byte(r.Header.Get(cfg.AuthHeader))
|
||||
expected := []byte(cfg.AuthHeaderValue)
|
||||
if subtle.ConstantTimeCompare(actual, expected) != 1 {
|
||||
a.logger.Logger.Debug("MCP 鉴权失败:header 缺失或值不匹配", zap.String("header", cfg.AuthHeader))
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
w.Write([]byte(`{"error":"unauthorized"}`))
|
||||
if authHeader := strings.TrimSpace(r.Header.Get("Authorization")); len(authHeader) > 7 && strings.EqualFold(authHeader[:7], "Bearer ") {
|
||||
if session, ok := a.auth.ValidateToken(strings.TrimSpace(authHeader[7:])); ok && session.Permissions["mcp:execute"] {
|
||||
principal := authctx.NewPrincipalWithScopes(session.UserID, session.Username, session.Scope, session.Permissions, session.PermissionScopes)
|
||||
a.mcpServer.HandleHTTP(w, r.WithContext(authctx.WithPrincipal(r.Context(), principal)))
|
||||
return
|
||||
}
|
||||
}
|
||||
if !cfg.AllowGlobalAccess || strings.TrimSpace(cfg.AuthHeader) == "" || strings.TrimSpace(cfg.AuthHeaderValue) == "" {
|
||||
http.Error(w, "use an authorized user bearer token; global MCP service access is disabled", http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
if subtle.ConstantTimeCompare([]byte(r.Header.Get(cfg.AuthHeader)), []byte(cfg.AuthHeaderValue)) != 1 {
|
||||
a.logger.Logger.Debug("MCP 鉴权失败:header 缺失或值不匹配", zap.String("header", cfg.AuthHeader))
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
w.Write([]byte(`{"error":"unauthorized"}`))
|
||||
return
|
||||
}
|
||||
permissions := make(map[string]bool, len(security.PermissionCatalog))
|
||||
for permission := range security.PermissionCatalog {
|
||||
permissions[permission] = true
|
||||
}
|
||||
principal := authctx.NewPrincipal("service:mcp", "mcp-service", database.RBACScopeAll, permissions)
|
||||
r = r.WithContext(authctx.WithPrincipal(r.Context(), principal))
|
||||
a.mcpServer.HandleHTTP(w, r)
|
||||
}
|
||||
|
||||
@@ -818,6 +851,8 @@ func setupRoutes(
|
||||
terminalHandler *handler.TerminalHandler,
|
||||
c2Handler *handler.C2Handler,
|
||||
auditHandler *handler.AuditHandler,
|
||||
auditSvc *audit.Service,
|
||||
rbacHandler *handler.RBACHandler,
|
||||
mcpServer *mcp.Server,
|
||||
authManager *security.AuthManager,
|
||||
openAPIHandler *handler.OpenAPIHandler,
|
||||
@@ -827,11 +862,15 @@ func setupRoutes(
|
||||
|
||||
// 认证相关路由
|
||||
authRoutes := api.Group("/auth")
|
||||
loginRL := security.NewRateLimiter(10, 1*time.Minute)
|
||||
{
|
||||
authRoutes.POST("/login", authHandler.Login)
|
||||
authRoutes.POST("/login", security.RateLimitMiddleware(loginRL), authHandler.Login)
|
||||
authRoutes.POST("/logout", security.AuthMiddleware(authManager), authHandler.Logout)
|
||||
authRoutes.POST("/change-password", security.AuthMiddleware(authManager), authHandler.ChangePassword)
|
||||
authRoutes.POST("/change-password", security.AuthMiddleware(authManager), security.RequirePermission("auth:self"), authHandler.ChangePassword)
|
||||
authRoutes.GET("/validate", security.AuthMiddleware(authManager), authHandler.Validate)
|
||||
authRoutes.POST("/robot-binding-code", security.AuthMiddleware(authManager), security.RequirePermission("auth:self"), robotHandler.CreateRobotBindingCode)
|
||||
authRoutes.GET("/robot-bindings", security.AuthMiddleware(authManager), security.RequirePermission("auth:self"), robotHandler.ListMyRobotBindings)
|
||||
authRoutes.DELETE("/robot-bindings/:id", security.AuthMiddleware(authManager), security.RequirePermission("auth:self"), robotHandler.DeleteMyRobotBinding)
|
||||
}
|
||||
|
||||
// 机器人回调(无需登录,供企业微信/钉钉/飞书服务器调用)
|
||||
@@ -848,7 +887,31 @@ func setupRoutes(
|
||||
|
||||
protected := api.Group("")
|
||||
protected.Use(security.AuthMiddleware(authManager))
|
||||
protected.Use(security.RBACMiddlewareWithDenyHook(app.db, func(c *gin.Context, reason, permission string) {
|
||||
if auditSvc != nil {
|
||||
auditSvc.Record(c, audit.Entry{
|
||||
Level: "warn", Category: "rbac", Action: "access_denied", Result: "failure",
|
||||
Message: "RBAC 拒绝访问", ResourceType: "route", ResourceID: c.FullPath(),
|
||||
Detail: map[string]interface{}{"reason": reason, "permission": permission, "method": c.Request.Method},
|
||||
})
|
||||
}
|
||||
}))
|
||||
{
|
||||
protected.GET("/rbac/me", rbacHandler.Me)
|
||||
protected.GET("/rbac/metadata", rbacHandler.Metadata)
|
||||
protected.GET("/rbac/users", rbacHandler.ListUsers)
|
||||
protected.POST("/rbac/users", rbacHandler.CreateUser)
|
||||
protected.PUT("/rbac/users/:id", rbacHandler.UpdateUser)
|
||||
protected.DELETE("/rbac/users/:id", rbacHandler.DeleteUser)
|
||||
protected.GET("/rbac/roles", rbacHandler.ListRoles)
|
||||
protected.POST("/rbac/roles", rbacHandler.CreateRole)
|
||||
protected.PUT("/rbac/roles/:id", rbacHandler.UpdateRole)
|
||||
protected.DELETE("/rbac/roles/:id", rbacHandler.DeleteRole)
|
||||
protected.GET("/rbac/resource-assignments", rbacHandler.ListResourceAssignments)
|
||||
protected.GET("/rbac/resources", rbacHandler.ListAssignableResources)
|
||||
protected.POST("/rbac/resource-assignments", rbacHandler.AssignResource)
|
||||
protected.DELETE("/rbac/resource-assignments/:id", rbacHandler.DeleteResourceAssignment)
|
||||
|
||||
// 机器人测试(需登录):POST /api/robot/test,body: {"platform":"dingtalk","user_id":"test","text":"帮助"},用于验证机器人逻辑
|
||||
protected.POST("/robot/test", robotHandler.HandleRobotTest)
|
||||
|
||||
@@ -918,6 +981,7 @@ func setupRoutes(
|
||||
protected.GET("/conversations", conversationHandler.ListConversations)
|
||||
protected.GET("/conversations/:id", conversationHandler.GetConversation)
|
||||
protected.GET("/messages/:id/process-details", conversationHandler.GetMessageProcessDetails)
|
||||
protected.GET("/process-details/:id", conversationHandler.GetProcessDetail)
|
||||
protected.PUT("/conversations/:id", conversationHandler.UpdateConversation)
|
||||
protected.PUT("/conversations/:id/project", conversationHandler.SetConversationProject)
|
||||
protected.DELETE("/conversations/:id", conversationHandler.DeleteConversation)
|
||||
@@ -1465,7 +1529,13 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
|
||||
},
|
||||
}
|
||||
listHandler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
connections, err := db.ListWebshellConnections()
|
||||
connections := []database.WebShellConnection{}
|
||||
var err error
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||
connections, err = db.ListWebshellConnectionsForAccess(principal.UserID, principal.ScopeFor("webshell:read"))
|
||||
} else {
|
||||
return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "缺少认证身份"}}, IsError: true}, nil
|
||||
}
|
||||
if err != nil {
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{{Type: "text", Text: "获取连接列表失败: " + err.Error()}},
|
||||
@@ -1580,6 +1650,10 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
|
||||
IsError: true,
|
||||
}, nil
|
||||
}
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||
_ = db.SetResourceOwner("webshell", conn.ID, principal.UserID)
|
||||
_ = db.AssignResourceToUser(principal.UserID, "webshell", conn.ID)
|
||||
}
|
||||
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{{
|
||||
@@ -1975,8 +2049,17 @@ func initializeKnowledge(
|
||||
// corsMiddleware CORS中间件
|
||||
func corsMiddleware() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
c.Writer.Header().Set("Access-Control-Allow-Origin", "*")
|
||||
c.Writer.Header().Set("Access-Control-Allow-Credentials", "true")
|
||||
origin := strings.TrimSpace(c.GetHeader("Origin"))
|
||||
if origin != "" {
|
||||
parsed, err := url.Parse(origin)
|
||||
if err != nil || parsed.Host == "" || !strings.EqualFold(parsed.Host, c.Request.Host) {
|
||||
c.AbortWithStatusJSON(http.StatusForbidden, gin.H{"error": "cross-origin request denied"})
|
||||
return
|
||||
}
|
||||
c.Writer.Header().Set("Access-Control-Allow-Origin", origin)
|
||||
c.Writer.Header().Set("Access-Control-Allow-Credentials", "true")
|
||||
c.Writer.Header().Add("Vary", "Origin")
|
||||
}
|
||||
c.Writer.Header().Set("Access-Control-Allow-Headers", "Content-Type, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization, accept, origin, Cache-Control, X-Requested-With")
|
||||
c.Writer.Header().Set("Access-Control-Allow-Methods", "POST, OPTIONS, GET, PUT, DELETE")
|
||||
|
||||
|
||||
+29
-12
@@ -4,11 +4,13 @@ import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/c2"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
@@ -66,16 +68,16 @@ tcp_reverse 默认仅接受 CSB1 加密 Beacon(AES-GCM + ImplantToken)才登
|
||||
InputSchema: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"action": map[string]interface{}{"type": "string", "description": "操作: list/get/create/update/start/stop/delete", "enum": []string{"list", "get", "create", "update", "start", "stop", "delete"}},
|
||||
"listener_id": map[string]interface{}{"type": "string", "description": "监听器 ID(get/update/start/stop/delete 需要)"},
|
||||
"name": map[string]interface{}{"type": "string", "description": "监听器名称(create/update)"},
|
||||
"type": map[string]interface{}{"type": "string", "description": "监听器类型(create)", "enum": []string{"tcp_reverse", "http_beacon", "https_beacon", "websocket"}},
|
||||
"action": map[string]interface{}{"type": "string", "description": "操作: list/get/create/update/start/stop/delete", "enum": []string{"list", "get", "create", "update", "start", "stop", "delete"}},
|
||||
"listener_id": map[string]interface{}{"type": "string", "description": "监听器 ID(get/update/start/stop/delete 需要)"},
|
||||
"name": map[string]interface{}{"type": "string", "description": "监听器名称(create/update)"},
|
||||
"type": map[string]interface{}{"type": "string", "description": "监听器类型(create)", "enum": []string{"tcp_reverse", "http_beacon", "https_beacon", "websocket"}},
|
||||
"bind_host": map[string]interface{}{"type": "string", "description": "绑定地址,默认 127.0.0.1;外网监听常用 0.0.0.0"},
|
||||
"callback_host": map[string]interface{}{"type": "string", "description": "可选:植入端/Payload 回连主机名(公网 IP 或域名)。写入 config_json;生成 oneliner/beacon 时优先于 bind_host。update 时传入空字符串可清除"},
|
||||
"bind_port": map[string]interface{}{"type": "integer", "description": fmt.Sprintf("绑定端口(create 必填)。须 ≠ %d(当前本服务 Web/API 端口,配置 server.port)", webListenPort), "minimum": 1, "maximum": 65535},
|
||||
"profile_id": map[string]interface{}{"type": "string", "description": "Malleable Profile ID"},
|
||||
"remark": map[string]interface{}{"type": "string", "description": "备注"},
|
||||
"config": map[string]interface{}{"type": "object", "description": "高级配置(beacon 路径/TLS/OPSEC 等),create/update 可用。tcp_reverse 可选 allow_legacy_shell:true 允许未加密经典 shell(默认 false)"},
|
||||
"bind_port": map[string]interface{}{"type": "integer", "description": fmt.Sprintf("绑定端口(create 必填)。须 ≠ %d(当前本服务 Web/API 端口,配置 server.port)", webListenPort), "minimum": 1, "maximum": 65535},
|
||||
"profile_id": map[string]interface{}{"type": "string", "description": "Malleable Profile ID"},
|
||||
"remark": map[string]interface{}{"type": "string", "description": "备注"},
|
||||
"config": map[string]interface{}{"type": "object", "description": "高级配置(beacon 路径/TLS/OPSEC 等),create/update 可用。tcp_reverse 可选 allow_legacy_shell:true 允许未加密经典 shell(默认 false)"},
|
||||
},
|
||||
"required": []string{"action"},
|
||||
},
|
||||
@@ -85,7 +87,7 @@ tcp_reverse 默认仅接受 CSB1 加密 Beacon(AES-GCM + ImplantToken)才登
|
||||
|
||||
switch action {
|
||||
case "list":
|
||||
listeners, err := m.DB().ListC2Listeners()
|
||||
listeners, err := m.DB().ListC2ListenersForAccess(c2ToolAccess(ctx))
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
@@ -128,6 +130,10 @@ tcp_reverse 默认仅接受 CSB1 加密 Beacon(AES-GCM + ImplantToken)才登
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||
_ = m.DB().SetResourceOwner("c2_listener", listener.ID, principal.UserID)
|
||||
_ = m.DB().AssignResourceToUser(principal.UserID, "c2_listener", listener.ID)
|
||||
}
|
||||
implantToken := listener.ImplantToken
|
||||
listener.EncryptionKey = ""
|
||||
listener.ImplantToken = ""
|
||||
@@ -264,7 +270,7 @@ func registerC2SessionTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
|
||||
if v, ok := params["suspicious"].(bool); ok && v {
|
||||
filter.Suspicious = true
|
||||
}
|
||||
sessions, err := m.DB().ListC2Sessions(filter)
|
||||
sessions, err := m.DB().ListC2SessionsForAccess(filter, c2ToolAccess(ctx))
|
||||
return makeC2Result(map[string]interface{}{"sessions": sessions, "count": len(sessions)}, err)
|
||||
|
||||
case "get":
|
||||
@@ -494,7 +500,7 @@ func registerC2TaskManageTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
|
||||
if limit := int(getFloat64(params, "limit")); limit > 0 {
|
||||
filter.Limit = limit
|
||||
}
|
||||
tasks, err := m.DB().ListC2Tasks(filter)
|
||||
tasks, err := m.DB().ListC2TasksForAccess(filter, c2ToolAccess(ctx))
|
||||
return makeC2Result(map[string]interface{}{"tasks": tasks, "count": len(tasks)}, err)
|
||||
|
||||
case "cancel":
|
||||
@@ -602,6 +608,9 @@ func registerC2PayloadTool(s *mcp.Server, m *c2.Manager, l *zap.Logger, webListe
|
||||
if err != nil {
|
||||
return makeC2Result(nil, err)
|
||||
}
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||
_ = m.DB().RecordC2PayloadArtifact(filepath.Base(result.OutputPath), result.PayloadID, result.ListenerID, principal.UserID)
|
||||
}
|
||||
return makeC2Result(map[string]interface{}{
|
||||
"payload_id": result.PayloadID, "download_path": result.DownloadPath,
|
||||
"os": result.OS, "arch": result.Arch, "size_bytes": result.SizeBytes,
|
||||
@@ -648,11 +657,19 @@ func registerC2EventTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
|
||||
filter.Since = &t
|
||||
}
|
||||
}
|
||||
events, err := m.DB().ListC2Events(filter)
|
||||
events, err := m.DB().ListC2EventsForAccess(filter, c2ToolAccess(ctx))
|
||||
return makeC2Result(map[string]interface{}{"events": events, "count": len(events)}, err)
|
||||
})
|
||||
}
|
||||
|
||||
func c2ToolAccess(ctx context.Context) database.RBACListAccess {
|
||||
principal, ok := authctx.PrincipalFromContext(ctx)
|
||||
if !ok {
|
||||
return database.RBACListAccess{Scope: database.RBACScopeAssigned}
|
||||
}
|
||||
return database.RBACListAccess{UserID: principal.UserID, Scope: principal.ScopeFor("c2:read")}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// c2_profile — Malleable Profile 管理工具(新增)
|
||||
// ============================================================================
|
||||
|
||||
@@ -0,0 +1,34 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func TestCORSMiddlewareAllowsSameOriginAndRejectsForeignOrigin(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
router.Use(corsMiddleware())
|
||||
router.GET("/test", func(c *gin.Context) { c.Status(http.StatusNoContent) })
|
||||
|
||||
same := httptest.NewRequest(http.MethodGet, "http://app.example/test", nil)
|
||||
same.Host = "app.example"
|
||||
same.Header.Set("Origin", "http://app.example")
|
||||
sameW := httptest.NewRecorder()
|
||||
router.ServeHTTP(sameW, same)
|
||||
if sameW.Code != http.StatusNoContent || sameW.Header().Get("Access-Control-Allow-Origin") != "http://app.example" {
|
||||
t.Fatalf("same-origin response = %d, allow-origin=%q", sameW.Code, sameW.Header().Get("Access-Control-Allow-Origin"))
|
||||
}
|
||||
|
||||
foreign := httptest.NewRequest(http.MethodGet, "http://app.example/test", nil)
|
||||
foreign.Host = "app.example"
|
||||
foreign.Header.Set("Origin", "https://evil.example")
|
||||
foreignW := httptest.NewRecorder()
|
||||
router.ServeHTTP(foreignW, foreign)
|
||||
if foreignW.Code != http.StatusForbidden {
|
||||
t.Fatalf("foreign-origin response = %d, want %d", foreignW.Code, http.StatusForbidden)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,249 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
)
|
||||
|
||||
func mcpToolAuthorizer(db *database.DB) func(context.Context, string, map[string]interface{}) error {
|
||||
return func(ctx context.Context, toolName string, args map[string]interface{}) error {
|
||||
principal, ok := authctx.PrincipalFromContext(ctx)
|
||||
if !ok {
|
||||
return fmt.Errorf("missing authenticated principal")
|
||||
}
|
||||
require := func(permission string) error {
|
||||
if !principal.HasPermission(permission) {
|
||||
return fmt.Errorf("missing permission %s", permission)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
resource := func(permission, resourceType, argument string) error {
|
||||
if err := require(permission); err != nil {
|
||||
return err
|
||||
}
|
||||
id := mcpAuthorizationString(args, argument)
|
||||
if id == "" || db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), resourceType, id) {
|
||||
return fmt.Errorf("no access to %s %s", resourceType, id)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
switch toolName {
|
||||
case builtin.ToolWebshellExec, builtin.ToolWebshellFileWrite:
|
||||
return resource("webshell:write", "webshell", "connection_id")
|
||||
case builtin.ToolWebshellFileList, builtin.ToolWebshellFileRead:
|
||||
return resource("webshell:read", "webshell", "connection_id")
|
||||
case builtin.ToolManageWebshellList:
|
||||
return require("webshell:read")
|
||||
case builtin.ToolManageWebshellAdd:
|
||||
return require("webshell:write")
|
||||
case builtin.ToolManageWebshellUpdate, builtin.ToolManageWebshellTest:
|
||||
return resource("webshell:write", "webshell", "connection_id")
|
||||
case builtin.ToolManageWebshellDelete:
|
||||
return resource("webshell:delete", "webshell", "connection_id")
|
||||
case builtin.ToolRecordVulnerability:
|
||||
if err := require("vulnerability:write"); err != nil {
|
||||
return err
|
||||
}
|
||||
conversationID := mcpAuthorizationString(args, "conversation_id")
|
||||
if conversationID == "" {
|
||||
conversationID = mcpAuthorizationConversationID(ctx)
|
||||
}
|
||||
if conversationID == "" || db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor("vulnerability:write"), "conversation", conversationID) {
|
||||
return fmt.Errorf("no access to conversation %s", conversationID)
|
||||
}
|
||||
return nil
|
||||
case builtin.ToolListVulnerabilities:
|
||||
if err := require("vulnerability:read"); err != nil {
|
||||
return err
|
||||
}
|
||||
conversationID := mcpAuthorizationConversationID(ctx)
|
||||
if conversationID == "" || db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor("vulnerability:read"), "conversation", conversationID) {
|
||||
return fmt.Errorf("no access to conversation %s", conversationID)
|
||||
}
|
||||
return nil
|
||||
case builtin.ToolGetVulnerability:
|
||||
return resource("vulnerability:read", "vulnerability", "id")
|
||||
case builtin.ToolUpsertProjectFact, builtin.ToolDeprecateProjectFact, builtin.ToolRestoreProjectFact:
|
||||
return authorizeProjectTool(ctx, principal, db, "project:write")
|
||||
case builtin.ToolGetProjectFact, builtin.ToolListProjectFacts, builtin.ToolSearchProjectFacts:
|
||||
return authorizeProjectTool(ctx, principal, db, "project:read")
|
||||
case builtin.ToolListKnowledgeRiskTypes, builtin.ToolSearchKnowledgeBase:
|
||||
return require("knowledge:read")
|
||||
case builtin.ToolAnalyzeImage:
|
||||
return require("agent:execute")
|
||||
case builtin.ToolBatchTaskList:
|
||||
return require("tasks:read")
|
||||
case builtin.ToolBatchTaskGet:
|
||||
return resource("tasks:read", "batch_task", "queue_id")
|
||||
case builtin.ToolBatchTaskCreate:
|
||||
if err := require("tasks:write"); err != nil {
|
||||
return err
|
||||
}
|
||||
if projectID := mcpAuthorizationString(args, "project_id"); projectID != "" && (db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor("tasks:write"), "project", projectID)) {
|
||||
return fmt.Errorf("no access to project %s", projectID)
|
||||
}
|
||||
return nil
|
||||
case builtin.ToolBatchTaskDelete, builtin.ToolBatchTaskRemove:
|
||||
return resource("tasks:delete", "batch_task", "queue_id")
|
||||
case builtin.ToolBatchTaskStart, builtin.ToolBatchTaskRerun, builtin.ToolBatchTaskPause,
|
||||
builtin.ToolBatchTaskUpdateMetadata, builtin.ToolBatchTaskUpdateSchedule,
|
||||
builtin.ToolBatchTaskScheduleEnabled, builtin.ToolBatchTaskAdd, builtin.ToolBatchTaskUpdate:
|
||||
return resource("tasks:write", "batch_task", "queue_id")
|
||||
case builtin.ToolC2Listener:
|
||||
return authorizeC2Action(principal, db, args, "c2_listener", "listener_id")
|
||||
case builtin.ToolC2Session, builtin.ToolC2Task, builtin.ToolC2File:
|
||||
if toolName == builtin.ToolC2File && mcpAuthorizationString(args, "action") == "get_result" {
|
||||
return authorizeC2Action(principal, db, args, "c2_task", "task_id")
|
||||
}
|
||||
return authorizeC2Action(principal, db, args, "c2_session", "session_id")
|
||||
case builtin.ToolC2TaskManage:
|
||||
return authorizeC2Action(principal, db, args, "c2_task", "task_id")
|
||||
case builtin.ToolC2Payload:
|
||||
return resource("c2:write", "c2_listener", "listener_id")
|
||||
case builtin.ToolC2Event:
|
||||
if id := mcpAuthorizationString(args, "session_id"); id != "" {
|
||||
return resource("c2:read", "c2_session", "session_id")
|
||||
}
|
||||
if principal.ScopeFor("c2:read") != database.RBACScopeAll {
|
||||
return fmt.Errorf("unfiltered C2 event list requires global scope")
|
||||
}
|
||||
return require("c2:read")
|
||||
case builtin.ToolC2Profile:
|
||||
// Profiles are process-global and do not yet have an owner. Writes are
|
||||
// therefore reserved for global scope; reads require c2:read.
|
||||
if mcpAuthorizationString(args, "action") == "list" || mcpAuthorizationString(args, "action") == "get" {
|
||||
return require("c2:read")
|
||||
}
|
||||
permission := "c2:write"
|
||||
if mcpAuthorizationString(args, "action") == "delete" {
|
||||
permission = "c2:delete"
|
||||
}
|
||||
if principal.ScopeFor(permission) != database.RBACScopeAll {
|
||||
return fmt.Errorf("C2 profile mutation requires global scope")
|
||||
}
|
||||
if mcpAuthorizationString(args, "action") == "delete" {
|
||||
return require("c2:delete")
|
||||
}
|
||||
return require("c2:write")
|
||||
default:
|
||||
if builtin.IsBuiltinTool(toolName) {
|
||||
return fmt.Errorf("no authorization policy registered for builtin tool %s", toolName)
|
||||
}
|
||||
if principal.HasPermission("agent:local-execute") {
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("missing agent:local-execute")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func externalMCPToolAuthorizer() func(context.Context, string, map[string]interface{}) error {
|
||||
return func(ctx context.Context, toolName string, _ map[string]interface{}) error {
|
||||
principal, ok := authctx.PrincipalFromContext(ctx)
|
||||
if !ok {
|
||||
return fmt.Errorf("missing authenticated principal")
|
||||
}
|
||||
if !principal.HasPermission("mcp:external:execute") {
|
||||
return fmt.Errorf("missing permission mcp:external:execute")
|
||||
}
|
||||
if principal.ScopeFor("mcp:external:execute") != database.RBACScopeAll {
|
||||
return fmt.Errorf("external MCP invocation requires global scope")
|
||||
}
|
||||
if strings.TrimSpace(toolName) == "" {
|
||||
return fmt.Errorf("missing external tool name")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func authorizeC2Action(principal authctx.Principal, db *database.DB, args map[string]interface{}, resourceType, argument string) error {
|
||||
action := mcpAuthorizationString(args, "action")
|
||||
permission := "c2:write"
|
||||
if action == "list" || action == "get" || action == "get_result" || action == "wait" {
|
||||
permission = "c2:read"
|
||||
} else if action == "delete" || action == "delete_batch" {
|
||||
permission = "c2:delete"
|
||||
}
|
||||
if !principal.HasPermission(permission) {
|
||||
return fmt.Errorf("missing permission %s", permission)
|
||||
}
|
||||
id := mcpAuthorizationString(args, argument)
|
||||
if action == "delete_batch" {
|
||||
ids := mcpAuthorizationStrings(args, argument+"s")
|
||||
if len(ids) == 0 {
|
||||
return fmt.Errorf("missing resource identifiers %ss", argument)
|
||||
}
|
||||
for _, candidate := range ids {
|
||||
if db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), resourceType, candidate) {
|
||||
return fmt.Errorf("no access to %s %s", resourceType, candidate)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if id == "" {
|
||||
if action == "create" || action == "list" {
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("missing resource identifier %s", argument)
|
||||
}
|
||||
if db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), resourceType, id) {
|
||||
return fmt.Errorf("no access to %s %s", resourceType, id)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func mcpAuthorizationStrings(args map[string]interface{}, key string) []string {
|
||||
values := []string{}
|
||||
switch raw := args[key].(type) {
|
||||
case []string:
|
||||
for _, value := range raw {
|
||||
if value = strings.TrimSpace(value); value != "" {
|
||||
values = append(values, value)
|
||||
}
|
||||
}
|
||||
case []interface{}:
|
||||
for _, item := range raw {
|
||||
if value, ok := item.(string); ok {
|
||||
if value = strings.TrimSpace(value); value != "" {
|
||||
values = append(values, value)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return values
|
||||
}
|
||||
|
||||
func authorizeProjectTool(ctx context.Context, principal authctx.Principal, db *database.DB, permission string) error {
|
||||
if !principal.HasPermission(permission) {
|
||||
return fmt.Errorf("missing permission %s", permission)
|
||||
}
|
||||
conversationID := mcpAuthorizationConversationID(ctx)
|
||||
if conversationID == "" || db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), "conversation", conversationID) {
|
||||
return fmt.Errorf("no access to conversation %s", conversationID)
|
||||
}
|
||||
projectID, err := db.GetConversationProjectID(conversationID)
|
||||
if err != nil || strings.TrimSpace(projectID) == "" || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), "project", projectID) {
|
||||
return fmt.Errorf("no access to project %s", projectID)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func mcpAuthorizationConversationID(ctx context.Context) string {
|
||||
if id := strings.TrimSpace(agent.ConversationIDFromContext(ctx)); id != "" {
|
||||
return id
|
||||
}
|
||||
return strings.TrimSpace(mcp.MCPConversationIDFromContext(ctx))
|
||||
}
|
||||
|
||||
func mcpAuthorizationString(args map[string]interface{}, key string) string {
|
||||
value, _ := args[key].(string)
|
||||
return strings.TrimSpace(value)
|
||||
}
|
||||
@@ -0,0 +1,97 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"context"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestMCPToolAuthorizerEnforcesPermissionAndResource(t *testing.T) {
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-authz.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
user, err := db.CreateRBACUser("mcp-user", "MCP User", "hash", true, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, id := range []string{"ws_allowed", "ws_hidden"} {
|
||||
if err := db.CreateWebshellConnection(&database.WebShellConnection{ID: id, URL: "http://127.0.0.1/" + id, Type: "php", Method: "post", CmdParam: "cmd", CreatedAt: time.Now()}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if err := db.AssignResourceToUser(user.ID, "webshell", "ws_allowed"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
principal := authctx.NewPrincipal(user.ID, user.Username, database.RBACScopeAssigned, map[string]bool{"mcp:write": true, "webshell:write": true})
|
||||
ctx := authctx.WithPrincipal(context.Background(), principal)
|
||||
authorize := mcpToolAuthorizer(db)
|
||||
if err := authorize(ctx, builtin.ToolWebshellExec, map[string]interface{}{"connection_id": "ws_allowed"}); err != nil {
|
||||
t.Fatalf("allowed resource denied: %v", err)
|
||||
}
|
||||
if err := authorize(ctx, builtin.ToolWebshellExec, map[string]interface{}{"connection_id": "ws_hidden"}); err == nil {
|
||||
t.Fatal("foreign webshell resource was allowed")
|
||||
}
|
||||
if err := authorize(ctx, builtin.ToolManageWebshellDelete, map[string]interface{}{"connection_id": "ws_allowed"}); err == nil {
|
||||
t.Fatal("delete without webshell:delete was allowed")
|
||||
}
|
||||
}
|
||||
|
||||
func TestEveryBuiltinMCPToolHasExplicitAuthorizationPolicy(t *testing.T) {
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-policy-inventory.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
permissions := map[string]bool{}
|
||||
for permission := range security.PermissionCatalog {
|
||||
permissions[permission] = true
|
||||
}
|
||||
ctx := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("admin", "admin", database.RBACScopeAll, permissions))
|
||||
authorize := mcpToolAuthorizer(db)
|
||||
args := map[string]interface{}{
|
||||
"action": "get", "connection_id": "x", "queue_id": "x", "listener_id": "x",
|
||||
"session_id": "x", "task_id": "x", "id": "x", "conversation_id": "x",
|
||||
}
|
||||
for _, toolName := range builtin.GetAllBuiltinTools() {
|
||||
err := authorize(ctx, toolName, args)
|
||||
if err != nil && strings.Contains(err.Error(), "no authorization policy registered") {
|
||||
t.Errorf("builtin tool %s has no explicit policy", toolName)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestExternalMCPRequiresDedicatedPermission(t *testing.T) {
|
||||
authorize := externalMCPToolAuthorizer()
|
||||
ctx := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", database.RBACScopeAssigned, map[string]bool{"agent:execute": true}))
|
||||
if err := authorize(ctx, "server::tool", nil); err == nil {
|
||||
t.Fatal("agent:execute alone authorized an external MCP tool")
|
||||
}
|
||||
ctx = authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", database.RBACScopeAll, map[string]bool{"mcp:external:execute": true}))
|
||||
if err := authorize(ctx, "server::tool", nil); err != nil {
|
||||
t.Fatalf("dedicated external MCP permission rejected: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConfiguredCommandToolRequiresLocalExecutePermission(t *testing.T) {
|
||||
authorize := mcpToolAuthorizer(nil)
|
||||
agentOnly := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", database.RBACScopeAssigned, map[string]bool{"agent:execute": true}))
|
||||
if err := authorize(agentOnly, "nmap_scan", nil); err == nil {
|
||||
t.Fatal("agent:execute alone authorized a configured command tool")
|
||||
}
|
||||
local := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", database.RBACScopeAssigned, map[string]bool{"agent:local-execute": true}))
|
||||
if err := authorize(local, "nmap_scan", nil); err != nil {
|
||||
t.Fatalf("agent:local-execute rejected: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,55 @@
|
||||
package app
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestStandaloneMCPPrefersUserRBACAndDisablesGlobalTokenByDefault(t *testing.T) {
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-http-auth.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
auth, err := security.NewAuthManager("admin-secret", 12)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := auth.AttachRBACStore(db); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
token, _, err := auth.Authenticate("admin", "admin-secret")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
server := mcp.NewServer(zap.NewNop())
|
||||
server.SetToolAuthorizer(mcpToolAuthorizer(db))
|
||||
a := &App{config: &config.Config{MCP: config.MCPConfig{AuthHeader: "X-MCP-Token", AuthHeaderValue: "static-secret"}}, auth: auth, mcpServer: server}
|
||||
body := []byte(`{"jsonrpc":"2.0","id":1,"method":"tools/list"}`)
|
||||
|
||||
userReq := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader(body))
|
||||
userReq.Header.Set("Authorization", "Bearer "+token)
|
||||
userW := httptest.NewRecorder()
|
||||
a.mcpHandlerWithAuth(userW, userReq)
|
||||
if userW.Code != http.StatusOK {
|
||||
t.Fatalf("user bearer status = %d: %s", userW.Code, userW.Body.String())
|
||||
}
|
||||
|
||||
staticReq := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader(body))
|
||||
staticReq.Header.Set("X-MCP-Token", "static-secret")
|
||||
staticW := httptest.NewRecorder()
|
||||
a.mcpHandlerWithAuth(staticW, staticReq)
|
||||
if staticW.Code != http.StatusUnauthorized {
|
||||
t.Fatalf("global static token status = %d, want 401", staticW.Code)
|
||||
}
|
||||
}
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
@@ -313,6 +314,10 @@ func registerRecordVulnerabilityTool(mcpServer *mcp.Server, db *database.DB, log
|
||||
}
|
||||
return textResult(fmt.Sprintf("记录漏洞失败: %v", err), true), nil
|
||||
}
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||
_ = db.SetResourceOwner("vulnerability", created.ID, principal.UserID)
|
||||
_ = db.AssignResourceToUser(principal.UserID, "vulnerability", created.ID)
|
||||
}
|
||||
|
||||
if logger != nil {
|
||||
logger.Info("漏洞记录成功",
|
||||
|
||||
@@ -63,7 +63,12 @@ func (s *Service) Record(c *gin.Context, e Entry) {
|
||||
}
|
||||
}
|
||||
if strings.TrimSpace(e.Actor) == "" {
|
||||
e.Actor = "admin"
|
||||
if c != nil {
|
||||
e.Actor = strings.TrimSpace(c.GetString(security.ContextUsernameKey))
|
||||
}
|
||||
if e.Actor == "" {
|
||||
e.Actor = "admin"
|
||||
}
|
||||
}
|
||||
maxDetail := s.cfg.Audit.MaxDetailBytesEffective()
|
||||
detail := SanitizeDetail(e.Detail, maxDetail)
|
||||
|
||||
@@ -0,0 +1,72 @@
|
||||
package authctx
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Principal is the immutable authorization identity propagated beyond the
|
||||
// transport layer into Agent, MCP and background task contexts.
|
||||
type Principal struct {
|
||||
UserID string
|
||||
Username string
|
||||
Permissions map[string]bool
|
||||
PermissionScopes map[string]string
|
||||
Scope string
|
||||
}
|
||||
|
||||
type principalContextKey struct{}
|
||||
|
||||
func NewPrincipal(userID, username, scope string, permissions map[string]bool) Principal {
|
||||
return NewPrincipalWithScopes(userID, username, scope, permissions, nil)
|
||||
}
|
||||
|
||||
func NewPrincipalWithScopes(userID, username, scope string, permissions map[string]bool, permissionScopes map[string]string) Principal {
|
||||
permissionCopy := make(map[string]bool, len(permissions))
|
||||
scopeCopy := make(map[string]string, len(permissionScopes))
|
||||
for permission, allowed := range permissions {
|
||||
if allowed {
|
||||
permissionCopy[permission] = true
|
||||
if permissionScope := strings.TrimSpace(permissionScopes[permission]); permissionScope != "" {
|
||||
scopeCopy[permission] = permissionScope
|
||||
}
|
||||
}
|
||||
}
|
||||
return Principal{
|
||||
UserID: strings.TrimSpace(userID), Username: strings.TrimSpace(username),
|
||||
Scope: strings.TrimSpace(scope), Permissions: permissionCopy, PermissionScopes: scopeCopy,
|
||||
}
|
||||
}
|
||||
|
||||
func WithPrincipal(ctx context.Context, principal Principal) context.Context {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
if strings.TrimSpace(principal.UserID) == "" {
|
||||
return ctx
|
||||
}
|
||||
return context.WithValue(ctx, principalContextKey{}, principal)
|
||||
}
|
||||
|
||||
func PrincipalFromContext(ctx context.Context) (Principal, bool) {
|
||||
if ctx == nil {
|
||||
return Principal{}, false
|
||||
}
|
||||
principal, ok := ctx.Value(principalContextKey{}).(Principal)
|
||||
return principal, ok && strings.TrimSpace(principal.UserID) != ""
|
||||
}
|
||||
|
||||
func (p Principal) HasPermission(permission string) bool {
|
||||
return p.Permissions[strings.TrimSpace(permission)]
|
||||
}
|
||||
|
||||
// ScopeFor returns the scope attached to the permission that authorizes the
|
||||
// current action. Falling back to Scope keeps explicit service principals and
|
||||
// legacy callers compatible without reintroducing cross-role scope widening.
|
||||
func (p Principal) ScopeFor(permission string) string {
|
||||
permission = strings.TrimSpace(permission)
|
||||
if scope := strings.TrimSpace(p.PermissionScopes[permission]); scope != "" {
|
||||
return scope
|
||||
}
|
||||
return strings.TrimSpace(p.Scope)
|
||||
}
|
||||
@@ -5,9 +5,9 @@ import (
|
||||
"fmt"
|
||||
"net"
|
||||
"os"
|
||||
"strconv"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"text/template"
|
||||
|
||||
@@ -173,15 +173,16 @@ func (b *PayloadBuilder) BuildBeacon(in PayloadBuilderInput) (*BuildResult, erro
|
||||
}
|
||||
|
||||
// 交叉编译
|
||||
payloadID := "p_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:14]
|
||||
binName := strings.TrimSpace(in.OutputName)
|
||||
if binName == "" {
|
||||
binName = fmt.Sprintf("beacon_%s_%s", goos, goarch)
|
||||
binName = fmt.Sprintf("beacon_%s_%s_%s", goos, goarch, payloadID)
|
||||
}
|
||||
if goos == "windows" && !strings.HasSuffix(binName, ".exe") {
|
||||
binName += ".exe"
|
||||
}
|
||||
binPath := filepath.Join(b.outputDir, binName)
|
||||
|
||||
|
||||
if err := os.MkdirAll(b.outputDir, 0755); err != nil {
|
||||
return nil, fmt.Errorf("mkdir output: %w", err)
|
||||
}
|
||||
@@ -214,7 +215,6 @@ func (b *PayloadBuilder) BuildBeacon(in PayloadBuilderInput) (*BuildResult, erro
|
||||
return nil, fmt.Errorf("stat output: %w", err)
|
||||
}
|
||||
|
||||
payloadID := "p_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:14]
|
||||
return &BuildResult{
|
||||
PayloadID: payloadID,
|
||||
ListenerID: listener.ID,
|
||||
|
||||
+319
-61
@@ -6,6 +6,7 @@ import (
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io/fs"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
@@ -41,6 +42,20 @@ type Config struct {
|
||||
Vision VisionConfig `yaml:"vision,omitempty" json:"vision,omitempty"`
|
||||
}
|
||||
|
||||
type EnsureLocalConfigResult struct {
|
||||
Created bool
|
||||
GeneratedPassword string
|
||||
ExamplePath string
|
||||
}
|
||||
|
||||
const (
|
||||
DefaultSummarizationUserIntentLedgerMaxRunes = 96000
|
||||
DefaultSummarizationUserIntentLedgerEntryMaxRunes = 16000
|
||||
DefaultLatestUserMessageMaxRunes = 48000
|
||||
DefaultLatestUserMessageHeadRunes = 24000
|
||||
DefaultLatestUserMessageTailRunes = 24000
|
||||
)
|
||||
|
||||
// ProjectConfig 项目黑板(跨对话共享事实)配置。
|
||||
type ProjectConfig struct {
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
@@ -255,6 +270,16 @@ type MultiAgentEinoMiddlewareConfig struct {
|
||||
SummarizationTriggerRatio float64 `yaml:"summarization_trigger_ratio,omitempty" json:"summarization_trigger_ratio,omitempty"`
|
||||
// SummarizationEmitInternalEvents controls middleware internal event emission (default true).
|
||||
SummarizationEmitInternalEvents *bool `yaml:"summarization_emit_internal_events,omitempty" json:"summarization_emit_internal_events,omitempty"`
|
||||
// SummarizationUserIntentLedgerMaxRunes caps the DB-backed immutable user input ledger injected into model context.
|
||||
SummarizationUserIntentLedgerMaxRunes int `yaml:"summarization_user_intent_ledger_max_runes,omitempty" json:"summarization_user_intent_ledger_max_runes,omitempty"`
|
||||
// SummarizationUserIntentLedgerEntryMaxRunes caps each user message entry inside the immutable user input ledger.
|
||||
SummarizationUserIntentLedgerEntryMaxRunes int `yaml:"summarization_user_intent_ledger_entry_max_runes,omitempty" json:"summarization_user_intent_ledger_entry_max_runes,omitempty"`
|
||||
// LatestUserMessageMaxRunes caps the current user turn inserted into model context; full text is persisted as an artifact when capped.
|
||||
LatestUserMessageMaxRunes int `yaml:"latest_user_message_max_runes,omitempty" json:"latest_user_message_max_runes,omitempty"`
|
||||
// LatestUserMessageHeadRunes keeps the head preview for an oversized current user turn.
|
||||
LatestUserMessageHeadRunes int `yaml:"latest_user_message_head_runes,omitempty" json:"latest_user_message_head_runes,omitempty"`
|
||||
// LatestUserMessageTailRunes keeps the tail preview for an oversized current user turn.
|
||||
LatestUserMessageTailRunes int `yaml:"latest_user_message_tail_runes,omitempty" json:"latest_user_message_tail_runes,omitempty"`
|
||||
// SummarizationRetryMaxAttempts 已废弃:summarization 与 run loop 共用 run_retry_max_attempts 及 isEinoTransientRunError。
|
||||
SummarizationRetryMaxAttempts int `yaml:"summarization_retry_max_attempts,omitempty" json:"summarization_retry_max_attempts,omitempty"`
|
||||
// PlanExecuteUserInputBudgetRatio caps planner/replanner/executor userInput prompt budget ratio (default 0.35).
|
||||
@@ -302,6 +327,41 @@ func (c MultiAgentEinoMiddlewareConfig) SummarizationEmitInternalEventsEffective
|
||||
return true
|
||||
}
|
||||
|
||||
func (c MultiAgentEinoMiddlewareConfig) SummarizationUserIntentLedgerMaxRunesEffective() int {
|
||||
if c.SummarizationUserIntentLedgerMaxRunes > 0 {
|
||||
return c.SummarizationUserIntentLedgerMaxRunes
|
||||
}
|
||||
return DefaultSummarizationUserIntentLedgerMaxRunes
|
||||
}
|
||||
|
||||
func (c MultiAgentEinoMiddlewareConfig) SummarizationUserIntentLedgerEntryMaxRunesEffective() int {
|
||||
if c.SummarizationUserIntentLedgerEntryMaxRunes > 0 {
|
||||
return c.SummarizationUserIntentLedgerEntryMaxRunes
|
||||
}
|
||||
return DefaultSummarizationUserIntentLedgerEntryMaxRunes
|
||||
}
|
||||
|
||||
func (c MultiAgentEinoMiddlewareConfig) LatestUserMessageMaxRunesEffective() int {
|
||||
if c.LatestUserMessageMaxRunes > 0 {
|
||||
return c.LatestUserMessageMaxRunes
|
||||
}
|
||||
return DefaultLatestUserMessageMaxRunes
|
||||
}
|
||||
|
||||
func (c MultiAgentEinoMiddlewareConfig) LatestUserMessageHeadRunesEffective() int {
|
||||
if c.LatestUserMessageHeadRunes > 0 {
|
||||
return c.LatestUserMessageHeadRunes
|
||||
}
|
||||
return DefaultLatestUserMessageHeadRunes
|
||||
}
|
||||
|
||||
func (c MultiAgentEinoMiddlewareConfig) LatestUserMessageTailRunesEffective() int {
|
||||
if c.LatestUserMessageTailRunes > 0 {
|
||||
return c.LatestUserMessageTailRunes
|
||||
}
|
||||
return DefaultLatestUserMessageTailRunes
|
||||
}
|
||||
|
||||
func (c MultiAgentEinoMiddlewareConfig) PlanExecuteUserInputBudgetRatioEffective() float64 {
|
||||
v := c.PlanExecuteUserInputBudgetRatio
|
||||
if v <= 0 {
|
||||
@@ -398,14 +458,19 @@ type MultiAgentSubConfig struct {
|
||||
|
||||
// MultiAgentPublic 返回给前端的精简信息(不含子代理指令全文)。
|
||||
type MultiAgentPublic struct {
|
||||
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"`
|
||||
ToolSearchAlwaysVisibleTools []string `json:"tool_search_always_visible_tools,omitempty"`
|
||||
ToolSearchAlwaysVisibleEffectiveTools []string `json:"tool_search_always_visible_effective_tools,omitempty"`
|
||||
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"`
|
||||
SummarizationUserIntentLedgerMaxRunes int `json:"summarization_user_intent_ledger_max_runes"`
|
||||
SummarizationUserIntentLedgerEntryMaxRunes int `json:"summarization_user_intent_ledger_entry_max_runes"`
|
||||
LatestUserMessageMaxRunes int `json:"latest_user_message_max_runes"`
|
||||
LatestUserMessageHeadRunes int `json:"latest_user_message_head_runes"`
|
||||
LatestUserMessageTailRunes int `json:"latest_user_message_tail_runes"`
|
||||
ToolSearchAlwaysVisibleTools []string `json:"tool_search_always_visible_tools,omitempty"`
|
||||
ToolSearchAlwaysVisibleEffectiveTools []string `json:"tool_search_always_visible_effective_tools,omitempty"`
|
||||
}
|
||||
|
||||
// NormalizeAgentMode 解析代理模式(eino_single | deep | plan_execute | supervisor);空值默认 eino_single。
|
||||
@@ -445,10 +510,15 @@ func NormalizeMultiAgentOrchestration(s string) string {
|
||||
|
||||
// MultiAgentAPIUpdate 设置页/API 仅更新多代理标量字段;写入 YAML 时不覆盖 sub_agents 等块。
|
||||
type MultiAgentAPIUpdate struct {
|
||||
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"`
|
||||
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"`
|
||||
SummarizationUserIntentLedgerMaxRunes *int `json:"summarization_user_intent_ledger_max_runes,omitempty"`
|
||||
SummarizationUserIntentLedgerEntryMaxRunes *int `json:"summarization_user_intent_ledger_entry_max_runes,omitempty"`
|
||||
LatestUserMessageMaxRunes *int `json:"latest_user_message_max_runes,omitempty"`
|
||||
LatestUserMessageHeadRunes *int `json:"latest_user_message_head_runes,omitempty"`
|
||||
LatestUserMessageTailRunes *int `json:"latest_user_message_tail_runes,omitempty"`
|
||||
// 指针区分「JSON 未传该字段」与「传空数组要清空」;省略时不应覆盖 YAML 中的常驻工具白名单。
|
||||
ToolSearchAlwaysVisibleTools *[]string `json:"tool_search_always_visible_tools,omitempty"`
|
||||
}
|
||||
@@ -468,14 +538,50 @@ type RobotsConfig struct {
|
||||
|
||||
// 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"` // 长轮询游标(运行时)
|
||||
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"` // 长轮询游标(运行时)
|
||||
Auth RobotAuthorizationConfig `yaml:"auth,omitempty" json:"auth,omitempty"`
|
||||
}
|
||||
|
||||
const (
|
||||
RobotAuthModeUserBinding = "user_binding"
|
||||
RobotAuthModeServiceAccount = "service_account"
|
||||
)
|
||||
|
||||
// RobotAuthorizationConfig controls how a verified platform sender becomes
|
||||
// an RBAC principal. service_account is intentionally fail-closed unless an
|
||||
// explicit non-admin service user and sender allowlist are both configured.
|
||||
type RobotAuthorizationConfig struct {
|
||||
Mode string `yaml:"mode,omitempty" json:"mode,omitempty"`
|
||||
ServiceUserID string `yaml:"service_user_id,omitempty" json:"service_user_id,omitempty"`
|
||||
AllowedExternalUsers []string `yaml:"allowed_external_users,omitempty" json:"allowed_external_users,omitempty"`
|
||||
}
|
||||
|
||||
func (c RobotAuthorizationConfig) EffectiveMode() string {
|
||||
mode := strings.ToLower(strings.TrimSpace(c.Mode))
|
||||
if mode == "" {
|
||||
return RobotAuthModeUserBinding
|
||||
}
|
||||
return mode
|
||||
}
|
||||
|
||||
func (c RobotAuthorizationConfig) ExternalUserAllowed(externalUserID string) bool {
|
||||
externalUserID = strings.TrimSpace(externalUserID)
|
||||
if externalUserID == "" {
|
||||
return false
|
||||
}
|
||||
for _, allowed := range c.AllowedExternalUsers {
|
||||
if strings.TrimSpace(allowed) == externalUserID {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// RobotSessionConfig 机器人会话隔离策略
|
||||
@@ -493,12 +599,13 @@ func (c RobotSessionConfig) StrictUserIdentityEnabled() bool {
|
||||
|
||||
// RobotWecomConfig 企业微信机器人配置
|
||||
type RobotWecomConfig struct {
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
Token string `yaml:"token" json:"token"` // 回调 URL 校验 Token
|
||||
EncodingAESKey string `yaml:"encoding_aes_key" json:"encoding_aes_key"` // EncodingAESKey
|
||||
CorpID string `yaml:"corp_id" json:"corp_id"` // 企业 ID
|
||||
Secret string `yaml:"secret" json:"secret"` // 应用 Secret
|
||||
AgentID int64 `yaml:"agent_id" json:"agent_id"` // 应用 AgentId
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
Token string `yaml:"token" json:"token"` // 回调 URL 校验 Token
|
||||
EncodingAESKey string `yaml:"encoding_aes_key" json:"encoding_aes_key"` // EncodingAESKey
|
||||
CorpID string `yaml:"corp_id" json:"corp_id"` // 企业 ID
|
||||
Secret string `yaml:"secret" json:"secret"` // 应用 Secret
|
||||
AgentID int64 `yaml:"agent_id" json:"agent_id"` // 应用 AgentId
|
||||
Auth RobotAuthorizationConfig `yaml:"auth,omitempty" json:"auth,omitempty"`
|
||||
}
|
||||
|
||||
// ValidateWecomConfig 校验企业微信机器人配置;启用时必须配置 token,否则回调无法防伪造。
|
||||
@@ -514,50 +621,137 @@ func ValidateWecomConfig(w RobotWecomConfig) error {
|
||||
|
||||
// RobotDingtalkConfig 钉钉机器人配置
|
||||
type RobotDingtalkConfig struct {
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
ClientID string `yaml:"client_id" json:"client_id"` // 应用 Key (AppKey)
|
||||
ClientSecret string `yaml:"client_secret" json:"client_secret"` // 应用 Secret
|
||||
AllowConversationIDFallback bool `yaml:"allow_conversation_id_fallback" json:"allow_conversation_id_fallback"` // sender_id 缺失时是否允许回退到会话 ID
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
ClientID string `yaml:"client_id" json:"client_id"` // 应用 Key (AppKey)
|
||||
ClientSecret string `yaml:"client_secret" json:"client_secret"` // 应用 Secret
|
||||
AllowConversationIDFallback bool `yaml:"allow_conversation_id_fallback" json:"allow_conversation_id_fallback"` // sender_id 缺失时是否允许回退到会话 ID
|
||||
Auth RobotAuthorizationConfig `yaml:"auth,omitempty" json:"auth,omitempty"`
|
||||
}
|
||||
|
||||
// RobotLarkConfig 飞书机器人配置
|
||||
type RobotLarkConfig struct {
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
AppID string `yaml:"app_id" json:"app_id"` // 应用 App ID
|
||||
AppSecret string `yaml:"app_secret" json:"app_secret"` // 应用 App Secret
|
||||
VerifyToken string `yaml:"verify_token" json:"verify_token"` // 事件订阅 Verification Token(可选)
|
||||
AllowChatIDFallback bool `yaml:"allow_chat_id_fallback" json:"allow_chat_id_fallback"` // 用户 ID 缺失时是否允许回退到 chat_id
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
AppID string `yaml:"app_id" json:"app_id"` // 应用 App ID
|
||||
AppSecret string `yaml:"app_secret" json:"app_secret"` // 应用 App Secret
|
||||
VerifyToken string `yaml:"verify_token" json:"verify_token"` // 事件订阅 Verification Token(可选)
|
||||
AllowChatIDFallback bool `yaml:"allow_chat_id_fallback" json:"allow_chat_id_fallback"` // 用户 ID 缺失时是否允许回退到 chat_id
|
||||
Auth RobotAuthorizationConfig `yaml:"auth,omitempty" json:"auth,omitempty"`
|
||||
}
|
||||
|
||||
// RobotTelegramConfig Telegram 机器人配置(Bot API 长轮询)
|
||||
type RobotTelegramConfig struct {
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
BotToken string `yaml:"bot_token" json:"bot_token"`
|
||||
BotUsername string `yaml:"bot_username,omitempty" json:"bot_username,omitempty"` // 可选,用于群聊 @ 识别;留空则启动时 getMe
|
||||
AllowGroupMessages bool `yaml:"allow_group_messages" json:"allow_group_messages"` // 群聊中仅响应 @ 机器人
|
||||
UpdateOffset int64 `yaml:"update_offset,omitempty" json:"update_offset,omitempty"`
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
BotToken string `yaml:"bot_token" json:"bot_token"`
|
||||
BotUsername string `yaml:"bot_username,omitempty" json:"bot_username,omitempty"` // 可选,用于群聊 @ 识别;留空则启动时 getMe
|
||||
AllowGroupMessages bool `yaml:"allow_group_messages" json:"allow_group_messages"` // 群聊中仅响应 @ 机器人
|
||||
UpdateOffset int64 `yaml:"update_offset,omitempty" json:"update_offset,omitempty"`
|
||||
Auth RobotAuthorizationConfig `yaml:"auth,omitempty" json:"auth,omitempty"`
|
||||
}
|
||||
|
||||
// RobotSlackConfig Slack 机器人配置(Socket Mode,无需公网回调)
|
||||
type RobotSlackConfig struct {
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
BotToken string `yaml:"bot_token" json:"bot_token"` // xoxb-
|
||||
AppToken string `yaml:"app_token" json:"app_token"` // xapp-(connections:write)
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
BotToken string `yaml:"bot_token" json:"bot_token"` // xoxb-
|
||||
AppToken string `yaml:"app_token" json:"app_token"` // xapp-(connections:write)
|
||||
Auth RobotAuthorizationConfig `yaml:"auth,omitempty" json:"auth,omitempty"`
|
||||
}
|
||||
|
||||
// RobotDiscordConfig Discord 机器人配置(Gateway WebSocket)
|
||||
type RobotDiscordConfig struct {
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
BotToken string `yaml:"bot_token" json:"bot_token"`
|
||||
AllowGuildMessages bool `yaml:"allow_guild_messages" json:"allow_guild_messages"` // 服务器频道中仅响应 @ 机器人
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
BotToken string `yaml:"bot_token" json:"bot_token"`
|
||||
AllowGuildMessages bool `yaml:"allow_guild_messages" json:"allow_guild_messages"` // 服务器频道中仅响应 @ 机器人
|
||||
Auth RobotAuthorizationConfig `yaml:"auth,omitempty" json:"auth,omitempty"`
|
||||
}
|
||||
|
||||
// RobotQQConfig QQ 机器人配置(QQ 开放平台 WebSocket)
|
||||
type RobotQQConfig struct {
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
AppID string `yaml:"app_id" json:"app_id"`
|
||||
ClientSecret string `yaml:"client_secret" json:"client_secret"`
|
||||
Sandbox bool `yaml:"sandbox" json:"sandbox"` // 沙箱环境(上线前测试)
|
||||
Enabled bool `yaml:"enabled" json:"enabled"`
|
||||
AppID string `yaml:"app_id" json:"app_id"`
|
||||
ClientSecret string `yaml:"client_secret" json:"client_secret"`
|
||||
Sandbox bool `yaml:"sandbox" json:"sandbox"` // 沙箱环境(上线前测试)
|
||||
Auth RobotAuthorizationConfig `yaml:"auth,omitempty" json:"auth,omitempty"`
|
||||
}
|
||||
|
||||
func (c RobotsConfig) AuthorizationFor(platform string) RobotAuthorizationConfig {
|
||||
switch strings.ToLower(strings.TrimSpace(platform)) {
|
||||
case "wechat":
|
||||
return c.Wechat.Auth
|
||||
case "wecom":
|
||||
return c.Wecom.Auth
|
||||
case "dingtalk":
|
||||
return c.Dingtalk.Auth
|
||||
case "lark":
|
||||
return c.Lark.Auth
|
||||
case "telegram":
|
||||
return c.Telegram.Auth
|
||||
case "slack":
|
||||
return c.Slack.Auth
|
||||
case "discord":
|
||||
return c.Discord.Auth
|
||||
case "qq":
|
||||
return c.QQ.Auth
|
||||
default:
|
||||
return RobotAuthorizationConfig{}
|
||||
}
|
||||
}
|
||||
|
||||
func ValidateRobotAuthorization(c RobotAuthorizationConfig, path string) error {
|
||||
switch c.EffectiveMode() {
|
||||
case RobotAuthModeUserBinding:
|
||||
return nil
|
||||
case RobotAuthModeServiceAccount:
|
||||
serviceUserID := strings.TrimSpace(c.ServiceUserID)
|
||||
if serviceUserID == "" {
|
||||
return fmt.Errorf("%s.auth.service_user_id 不能为空", path)
|
||||
}
|
||||
if len(c.AllowedExternalUsers) == 0 {
|
||||
return fmt.Errorf("%s.auth.allowed_external_users 至少配置一个真实发送者", path)
|
||||
}
|
||||
seen := map[string]bool{}
|
||||
for _, userID := range c.AllowedExternalUsers {
|
||||
userID = strings.TrimSpace(userID)
|
||||
if userID == "" || userID == "*" {
|
||||
return fmt.Errorf("%s.auth.allowed_external_users 不允许空值或通配符", path)
|
||||
}
|
||||
if seen[userID] {
|
||||
return fmt.Errorf("%s.auth.allowed_external_users 包含重复用户", path)
|
||||
}
|
||||
seen[userID] = true
|
||||
}
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("%s.auth.mode 仅支持 user_binding 或 service_account", path)
|
||||
}
|
||||
}
|
||||
|
||||
func ValidateRobotsAuthorization(c RobotsConfig) error {
|
||||
items := []struct {
|
||||
path string
|
||||
auth RobotAuthorizationConfig
|
||||
}{
|
||||
{"robots.wechat", c.Wechat.Auth}, {"robots.wecom", c.Wecom.Auth},
|
||||
{"robots.dingtalk", c.Dingtalk.Auth}, {"robots.lark", c.Lark.Auth},
|
||||
{"robots.telegram", c.Telegram.Auth}, {"robots.slack", c.Slack.Auth},
|
||||
{"robots.discord", c.Discord.Auth}, {"robots.qq", c.QQ.Auth},
|
||||
}
|
||||
for _, item := range items {
|
||||
if err := ValidateRobotAuthorization(item.auth, item.path); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c RobotsConfig) ServiceAccountUserIDs() map[string]string {
|
||||
out := map[string]string{}
|
||||
for _, platform := range []string{"wechat", "wecom", "dingtalk", "lark", "telegram", "slack", "discord", "qq"} {
|
||||
auth := c.AuthorizationFor(platform)
|
||||
if auth.EffectiveMode() == RobotAuthModeServiceAccount {
|
||||
out[platform] = strings.TrimSpace(auth.ServiceUserID)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
type ServerConfig struct {
|
||||
@@ -580,11 +774,12 @@ type LogConfig struct {
|
||||
}
|
||||
|
||||
type MCPConfig struct {
|
||||
Enabled bool `yaml:"enabled"`
|
||||
Host string `yaml:"host"`
|
||||
Port int `yaml:"port"`
|
||||
AuthHeader string `yaml:"auth_header,omitempty"` // 鉴权 header 名,留空表示不鉴权
|
||||
AuthHeaderValue string `yaml:"auth_header_value,omitempty"` // 鉴权 header 值,需与请求中该 header 一致
|
||||
Enabled bool `yaml:"enabled"`
|
||||
Host string `yaml:"host"`
|
||||
Port int `yaml:"port"`
|
||||
AuthHeader string `yaml:"auth_header,omitempty"` // 可选的全局服务凭证 header;普通调用优先使用用户 Bearer Token
|
||||
AuthHeaderValue string `yaml:"auth_header_value,omitempty"` // 全局服务凭证,仅 allow_global_access=true 时接受
|
||||
AllowGlobalAccess bool `yaml:"allow_global_access,omitempty"` // 静态服务密钥是否映射为全局服务身份(默认关闭)
|
||||
}
|
||||
|
||||
type OpenAIConfig struct {
|
||||
@@ -1036,10 +1231,71 @@ func Load(path string) (*Config, error) {
|
||||
if err := ValidateWecomConfig(cfg.Robots.Wecom); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := ValidateRobotsAuthorization(cfg.Robots); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &cfg, nil
|
||||
}
|
||||
|
||||
func EnsureLocalConfig(path string) (EnsureLocalConfigResult, error) {
|
||||
path = strings.TrimSpace(path)
|
||||
if path == "" {
|
||||
path = "config.yaml"
|
||||
}
|
||||
|
||||
if _, err := os.Stat(path); err == nil {
|
||||
return EnsureLocalConfigResult{}, nil
|
||||
} else if err != nil && !os.IsNotExist(err) {
|
||||
return EnsureLocalConfigResult{}, fmt.Errorf("检查配置文件失败: %w", err)
|
||||
}
|
||||
|
||||
examplePath := filepath.Join(filepath.Dir(path), "config.example.yaml")
|
||||
if _, err := os.Stat(examplePath); err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
if alt := "config.example.yaml"; examplePath != alt {
|
||||
if _, altErr := os.Stat(alt); altErr == nil {
|
||||
examplePath = alt
|
||||
} else {
|
||||
return EnsureLocalConfigResult{}, fmt.Errorf("配置文件 %s 不存在,且未找到模板 %s", path, examplePath)
|
||||
}
|
||||
} else {
|
||||
return EnsureLocalConfigResult{}, fmt.Errorf("配置文件 %s 不存在,且未找到模板 %s", path, examplePath)
|
||||
}
|
||||
} else {
|
||||
return EnsureLocalConfigResult{}, fmt.Errorf("检查配置模板失败: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
data, err := os.ReadFile(examplePath)
|
||||
if err != nil {
|
||||
return EnsureLocalConfigResult{}, fmt.Errorf("读取配置模板失败: %w", err)
|
||||
}
|
||||
|
||||
if dir := filepath.Dir(path); dir != "." && dir != "" {
|
||||
if err := os.MkdirAll(dir, 0700); err != nil {
|
||||
return EnsureLocalConfigResult{}, fmt.Errorf("创建配置目录失败: %w", err)
|
||||
}
|
||||
}
|
||||
if err := os.WriteFile(path, data, fs.FileMode(0600)); err != nil {
|
||||
return EnsureLocalConfigResult{}, fmt.Errorf("创建配置文件失败: %w", err)
|
||||
}
|
||||
|
||||
password, err := generateStrongPassword(24)
|
||||
if err != nil {
|
||||
return EnsureLocalConfigResult{}, fmt.Errorf("生成默认密码失败: %w", err)
|
||||
}
|
||||
if err := PersistAuthPassword(path, password); err != nil {
|
||||
return EnsureLocalConfigResult{}, fmt.Errorf("写入默认密码失败: %w", err)
|
||||
}
|
||||
|
||||
return EnsureLocalConfigResult{
|
||||
Created: true,
|
||||
GeneratedPassword: password,
|
||||
ExamplePath: examplePath,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func generateStrongPassword(length int) (string, error) {
|
||||
if length <= 0 {
|
||||
length = 24
|
||||
@@ -1268,9 +1524,10 @@ func persistMCPAuth(path string, mcp *MCPConfig) error {
|
||||
return os.WriteFile(path, []byte(strings.Join(lines, "\n")), 0644)
|
||||
}
|
||||
|
||||
// EnsureMCPAuth 在 MCP 启用且 auth_header_value 为空时,自动生成随机密钥并写回配置
|
||||
// EnsureMCPAuth only provisions the privileged static service credential when
|
||||
// global service access was explicitly enabled.
|
||||
func EnsureMCPAuth(path string, cfg *Config) error {
|
||||
if !cfg.MCP.Enabled || strings.TrimSpace(cfg.MCP.AuthHeaderValue) != "" {
|
||||
if !cfg.MCP.Enabled || !cfg.MCP.AllowGlobalAccess || strings.TrimSpace(cfg.MCP.AuthHeaderValue) != "" {
|
||||
return nil
|
||||
}
|
||||
token, err := generateRandomToken()
|
||||
@@ -1294,8 +1551,9 @@ func PrintMCPConfigJSON(mcp MCPConfig) {
|
||||
hostForURL = "localhost"
|
||||
}
|
||||
url := fmt.Sprintf("http://%s:%d/mcp", hostForURL, mcp.Port)
|
||||
headers := map[string]string{}
|
||||
if mcp.AuthHeader != "" {
|
||||
headers := map[string]string{"Authorization": "Bearer <USER_SESSION_TOKEN>"}
|
||||
if mcp.AllowGlobalAccess && mcp.AuthHeader != "" {
|
||||
delete(headers, "Authorization")
|
||||
headers[mcp.AuthHeader] = mcp.AuthHeaderValue
|
||||
}
|
||||
serverEntry := map[string]interface{}{
|
||||
@@ -1555,8 +1813,8 @@ func Default() *Config {
|
||||
Output: "stdout",
|
||||
},
|
||||
MCP: MCPConfig{
|
||||
Enabled: true,
|
||||
Host: "0.0.0.0",
|
||||
Enabled: false,
|
||||
Host: "127.0.0.1",
|
||||
Port: 8081,
|
||||
},
|
||||
OpenAI: OpenAIConfig{
|
||||
|
||||
@@ -72,6 +72,56 @@ func TestPersistAuthPasswordDoesNotTreatQuotedHashAsComment(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestEnsureLocalConfigCreatesFromExampleWithGeneratedPassword(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
examplePath := filepath.Join(dir, "config.example.yaml")
|
||||
configPath := filepath.Join(dir, "config.yaml")
|
||||
|
||||
example := []byte(`auth:
|
||||
password: "change-me-use-a-long-random-password"
|
||||
session_duration_hours: 12
|
||||
server:
|
||||
host: 127.0.0.1
|
||||
port: 8080
|
||||
`)
|
||||
if err := os.WriteFile(examplePath, example, 0644); err != nil {
|
||||
t.Fatalf("write example: %v", err)
|
||||
}
|
||||
|
||||
result, err := EnsureLocalConfig(configPath)
|
||||
if err != nil {
|
||||
t.Fatalf("EnsureLocalConfig: %v", err)
|
||||
}
|
||||
if !result.Created {
|
||||
t.Fatal("Created = false, want true")
|
||||
}
|
||||
if result.GeneratedPassword == "" {
|
||||
t.Fatal("GeneratedPassword is empty")
|
||||
}
|
||||
if result.ExamplePath != examplePath {
|
||||
t.Fatalf("ExamplePath = %q, want %q", result.ExamplePath, examplePath)
|
||||
}
|
||||
|
||||
cfg, err := Load(configPath)
|
||||
if err != nil {
|
||||
t.Fatalf("Load generated config: %v", err)
|
||||
}
|
||||
if cfg.Auth.Password == "change-me-use-a-long-random-password" {
|
||||
t.Fatal("auth.password still contains the template placeholder")
|
||||
}
|
||||
if cfg.Auth.Password != result.GeneratedPassword {
|
||||
t.Fatalf("Auth.Password = %q, want generated password %q", cfg.Auth.Password, result.GeneratedPassword)
|
||||
}
|
||||
|
||||
second, err := EnsureLocalConfig(configPath)
|
||||
if err != nil {
|
||||
t.Fatalf("EnsureLocalConfig existing: %v", err)
|
||||
}
|
||||
if second.Created {
|
||||
t.Fatal("Created = true for existing config, want false")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHitlAuditModelEffectiveFallsBackToMainConfig(t *testing.T) {
|
||||
main := OpenAIConfig{
|
||||
Provider: "openai",
|
||||
@@ -91,3 +141,52 @@ func TestHitlAuditModelEffectiveFallsBackToMainConfig(t *testing.T) {
|
||||
t.Fatalf("expected audit model override, got %q", got.Model)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSummarizationUserIntentLedgerRunesEffective(t *testing.T) {
|
||||
var zero MultiAgentEinoMiddlewareConfig
|
||||
if got := zero.SummarizationUserIntentLedgerMaxRunesEffective(); got != DefaultSummarizationUserIntentLedgerMaxRunes {
|
||||
t.Fatalf("default ledger max runes = %d, want %d", got, DefaultSummarizationUserIntentLedgerMaxRunes)
|
||||
}
|
||||
if got := zero.SummarizationUserIntentLedgerEntryMaxRunesEffective(); got != DefaultSummarizationUserIntentLedgerEntryMaxRunes {
|
||||
t.Fatalf("default ledger entry max runes = %d, want %d", got, DefaultSummarizationUserIntentLedgerEntryMaxRunes)
|
||||
}
|
||||
|
||||
custom := MultiAgentEinoMiddlewareConfig{
|
||||
SummarizationUserIntentLedgerMaxRunes: 12345,
|
||||
SummarizationUserIntentLedgerEntryMaxRunes: 2345,
|
||||
}
|
||||
if got := custom.SummarizationUserIntentLedgerMaxRunesEffective(); got != 12345 {
|
||||
t.Fatalf("custom ledger max runes = %d", got)
|
||||
}
|
||||
if got := custom.SummarizationUserIntentLedgerEntryMaxRunesEffective(); got != 2345 {
|
||||
t.Fatalf("custom ledger entry max runes = %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLatestUserMessageRunesEffective(t *testing.T) {
|
||||
var zero MultiAgentEinoMiddlewareConfig
|
||||
if got := zero.LatestUserMessageMaxRunesEffective(); got != DefaultLatestUserMessageMaxRunes {
|
||||
t.Fatalf("default latest user max runes = %d, want %d", got, DefaultLatestUserMessageMaxRunes)
|
||||
}
|
||||
if got := zero.LatestUserMessageHeadRunesEffective(); got != DefaultLatestUserMessageHeadRunes {
|
||||
t.Fatalf("default latest user head runes = %d, want %d", got, DefaultLatestUserMessageHeadRunes)
|
||||
}
|
||||
if got := zero.LatestUserMessageTailRunesEffective(); got != DefaultLatestUserMessageTailRunes {
|
||||
t.Fatalf("default latest user tail runes = %d, want %d", got, DefaultLatestUserMessageTailRunes)
|
||||
}
|
||||
|
||||
custom := MultiAgentEinoMiddlewareConfig{
|
||||
LatestUserMessageMaxRunes: 100,
|
||||
LatestUserMessageHeadRunes: 40,
|
||||
LatestUserMessageTailRunes: 60,
|
||||
}
|
||||
if got := custom.LatestUserMessageMaxRunesEffective(); got != 100 {
|
||||
t.Fatalf("custom latest user max runes = %d", got)
|
||||
}
|
||||
if got := custom.LatestUserMessageHeadRunesEffective(); got != 40 {
|
||||
t.Fatalf("custom latest user head runes = %d", got)
|
||||
}
|
||||
if got := custom.LatestUserMessageTailRunesEffective(); got != 60 {
|
||||
t.Fatalf("custom latest user tail runes = %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -43,3 +43,27 @@ func TestValidateWecomConfig(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateRobotAuthorization(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
cfg RobotAuthorizationConfig
|
||||
wantErr bool
|
||||
}{
|
||||
{name: "default user binding", cfg: RobotAuthorizationConfig{}, wantErr: false},
|
||||
{name: "explicit user binding", cfg: RobotAuthorizationConfig{Mode: RobotAuthModeUserBinding}, wantErr: false},
|
||||
{name: "service account", cfg: RobotAuthorizationConfig{Mode: RobotAuthModeServiceAccount, ServiceUserID: "svc-1", AllowedExternalUsers: []string{"t:x|u:y"}}, wantErr: false},
|
||||
{name: "missing service user", cfg: RobotAuthorizationConfig{Mode: RobotAuthModeServiceAccount, AllowedExternalUsers: []string{"t:x|u:y"}}, wantErr: true},
|
||||
{name: "admin allowed with exact sender", cfg: RobotAuthorizationConfig{Mode: RobotAuthModeServiceAccount, ServiceUserID: "admin", AllowedExternalUsers: []string{"t:x|u:y"}}, wantErr: false},
|
||||
{name: "allowlist required", cfg: RobotAuthorizationConfig{Mode: RobotAuthModeServiceAccount, ServiceUserID: "svc-1"}, wantErr: true},
|
||||
{name: "wildcard forbidden", cfg: RobotAuthorizationConfig{Mode: RobotAuthModeServiceAccount, ServiceUserID: "svc-1", AllowedExternalUsers: []string{"*"}}, wantErr: true},
|
||||
{name: "unknown mode", cfg: RobotAuthorizationConfig{Mode: "open"}, wantErr: true},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
if err := ValidateRobotAuthorization(tt.cfg, "robots.lark"); (err != nil) != tt.wantErr {
|
||||
t.Fatalf("ValidateRobotAuthorization() error=%v wantErr=%v", err, tt.wantErr)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
+22
-17
@@ -9,25 +9,26 @@ import (
|
||||
|
||||
// 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"`
|
||||
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 {
|
||||
Actor string
|
||||
Level string
|
||||
Category string
|
||||
Action string
|
||||
@@ -44,6 +45,10 @@ type ListAuditLogsFilter struct {
|
||||
func buildAuditLogsWhere(filter ListAuditLogsFilter) (string, []interface{}) {
|
||||
conditions := []string{"1=1"}
|
||||
args := []interface{}{}
|
||||
if filter.Actor != "" {
|
||||
conditions = append(conditions, "actor = ?")
|
||||
args = append(args, filter.Actor)
|
||||
}
|
||||
if filter.Level != "" {
|
||||
conditions = append(conditions, "level = ?")
|
||||
args = append(args, filter.Level)
|
||||
@@ -78,8 +83,8 @@ func buildAuditLogsWhere(filter ListAuditLogsFilter) (string, []interface{}) {
|
||||
}
|
||||
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)
|
||||
conditions = append(conditions, "(message LIKE ? OR resource_id LIKE ? OR action LIKE ? OR category LIKE ? OR detail_json LIKE ?)")
|
||||
args = append(args, like, like, like, like, like)
|
||||
}
|
||||
return strings.Join(conditions, " AND "), args
|
||||
}
|
||||
|
||||
@@ -137,7 +137,7 @@ func (db *DB) GetBatchQueue(queueID string) (*BatchTaskQueueRow, error) {
|
||||
// GetAllBatchQueues 获取所有批量任务队列
|
||||
func (db *DB) GetAllBatchQueues() ([]*BatchTaskQueueRow, error) {
|
||||
rows, err := db.Query(
|
||||
"SELECT "+batchQueueSelectColumns+" FROM batch_task_queues ORDER BY created_at DESC",
|
||||
"SELECT " + batchQueueSelectColumns + " FROM batch_task_queues ORDER BY created_at DESC",
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询批量任务队列列表失败: %w", err)
|
||||
@@ -168,6 +168,10 @@ func (db *DB) GetAllBatchQueues() ([]*BatchTaskQueueRow, error) {
|
||||
|
||||
// ListBatchQueues 列出批量任务队列(支持筛选和分页)
|
||||
func (db *DB) ListBatchQueues(limit, offset int, status, keyword string) ([]*BatchTaskQueueRow, error) {
|
||||
return db.ListBatchQueuesForAccess(limit, offset, status, keyword, "", "")
|
||||
}
|
||||
|
||||
func (db *DB) ListBatchQueuesForAccess(limit, offset int, status, keyword, userID, scope string) ([]*BatchTaskQueueRow, error) {
|
||||
query := "SELECT " + batchQueueSelectColumns + " FROM batch_task_queues WHERE 1=1"
|
||||
args := []interface{}{}
|
||||
|
||||
@@ -182,6 +186,26 @@ func (db *DB) ListBatchQueues(limit, offset int, status, keyword string) ([]*Bat
|
||||
query += " AND (id LIKE ? OR title LIKE ?)"
|
||||
args = append(args, "%"+keyword+"%", "%"+keyword+"%")
|
||||
}
|
||||
userID = strings.TrimSpace(userID)
|
||||
if userID != "" && scope != RBACScopeAll {
|
||||
query += ` AND (
|
||||
owner_user_id = ?
|
||||
OR EXISTS (
|
||||
SELECT 1 FROM rbac_resource_assignments ra
|
||||
WHERE ra.user_id = ? AND ra.resource_type = 'batch_task' AND ra.resource_id = batch_task_queues.id
|
||||
)
|
||||
OR (
|
||||
project_id IS NOT NULL AND project_id <> '' AND (
|
||||
EXISTS (SELECT 1 FROM projects p WHERE p.id = batch_task_queues.project_id AND p.owner_user_id = ?)
|
||||
OR EXISTS (
|
||||
SELECT 1 FROM rbac_resource_assignments pra
|
||||
WHERE pra.user_id = ? AND pra.resource_type = 'project' AND pra.resource_id = batch_task_queues.project_id
|
||||
)
|
||||
)
|
||||
)
|
||||
)`
|
||||
args = append(args, userID, userID, userID, userID)
|
||||
}
|
||||
|
||||
query += " ORDER BY created_at DESC LIMIT ? OFFSET ?"
|
||||
args = append(args, limit, offset)
|
||||
@@ -216,6 +240,10 @@ func (db *DB) ListBatchQueues(limit, offset int, status, keyword string) ([]*Bat
|
||||
|
||||
// CountBatchQueues 统计批量任务队列总数(支持筛选条件)
|
||||
func (db *DB) CountBatchQueues(status, keyword string) (int, error) {
|
||||
return db.CountBatchQueuesForAccess(status, keyword, "", "")
|
||||
}
|
||||
|
||||
func (db *DB) CountBatchQueuesForAccess(status, keyword, userID, scope string) (int, error) {
|
||||
query := "SELECT COUNT(*) FROM batch_task_queues WHERE 1=1"
|
||||
args := []interface{}{}
|
||||
|
||||
@@ -230,6 +258,26 @@ func (db *DB) CountBatchQueues(status, keyword string) (int, error) {
|
||||
query += " AND (id LIKE ? OR title LIKE ?)"
|
||||
args = append(args, "%"+keyword+"%", "%"+keyword+"%")
|
||||
}
|
||||
userID = strings.TrimSpace(userID)
|
||||
if userID != "" && scope != RBACScopeAll {
|
||||
query += ` AND (
|
||||
owner_user_id = ?
|
||||
OR EXISTS (
|
||||
SELECT 1 FROM rbac_resource_assignments ra
|
||||
WHERE ra.user_id = ? AND ra.resource_type = 'batch_task' AND ra.resource_id = batch_task_queues.id
|
||||
)
|
||||
OR (
|
||||
project_id IS NOT NULL AND project_id <> '' AND (
|
||||
EXISTS (SELECT 1 FROM projects p WHERE p.id = batch_task_queues.project_id AND p.owner_user_id = ?)
|
||||
OR EXISTS (
|
||||
SELECT 1 FROM rbac_resource_assignments pra
|
||||
WHERE pra.user_id = ? AND pra.resource_type = 'project' AND pra.resource_id = batch_task_queues.project_id
|
||||
)
|
||||
)
|
||||
)
|
||||
)`
|
||||
args = append(args, userID, userID, userID, userID)
|
||||
}
|
||||
|
||||
var count int
|
||||
err := db.QueryRow(query, args...).Scan(&count)
|
||||
|
||||
+507
-32
@@ -45,20 +45,21 @@ func validC2TextIDForDelete(id string) bool {
|
||||
|
||||
// C2Listener 监听器实体
|
||||
type C2Listener struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Type string `json:"type"` // tcp_reverse|http_beacon|https_beacon|websocket|dns
|
||||
BindHost string `json:"bindHost"` // 默认 127.0.0.1
|
||||
BindPort int `json:"bindPort"` // 1-65535
|
||||
ProfileID string `json:"profileId"` // 可空:关联 c2_profiles.id
|
||||
EncryptionKey string `json:"-"` // base64(AES-256),前端不返回
|
||||
ImplantToken string `json:"-"` // beacon 携带的鉴权 token,前端不返回
|
||||
Status string `json:"status"` // stopped|running|error
|
||||
ConfigJSON string `json:"configJson"` // TLS 证书路径 / URI 模式 / 上限并发 等
|
||||
Remark string `json:"remark"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Type string `json:"type"` // tcp_reverse|http_beacon|https_beacon|websocket|dns
|
||||
BindHost string `json:"bindHost"` // 默认 127.0.0.1
|
||||
BindPort int `json:"bindPort"` // 1-65535
|
||||
ProfileID string `json:"profileId"` // 可空:关联 c2_profiles.id
|
||||
EncryptionKey string `json:"-"` // base64(AES-256),前端不返回
|
||||
ImplantToken string `json:"-"` // beacon 携带的鉴权 token,前端不返回
|
||||
Status string `json:"status"` // stopped|running|error
|
||||
ConfigJSON string `json:"configJson"` // TLS 证书路径 / URI 模式 / 上限并发 等
|
||||
Remark string `json:"remark"`
|
||||
OwnerUserID string `json:"ownerUserId,omitempty"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
StartedAt *time.Time `json:"startedAt,omitempty"`
|
||||
LastError string `json:"lastError,omitempty"`
|
||||
LastError string `json:"lastError,omitempty"`
|
||||
}
|
||||
|
||||
// C2Session 已上线会话
|
||||
@@ -132,17 +133,17 @@ type C2Event struct {
|
||||
|
||||
// C2Profile Malleable Profile
|
||||
type C2Profile struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
UserAgent string `json:"userAgent"`
|
||||
URIs []string `json:"uris"`
|
||||
RequestHeaders map[string]string `json:"requestHeaders,omitempty"`
|
||||
ResponseHeaders map[string]string `json:"responseHeaders,omitempty"`
|
||||
BodyTemplate string `json:"bodyTemplate"`
|
||||
JitterMinMS int `json:"jitterMinMs"`
|
||||
JitterMaxMS int `json:"jitterMaxMs"`
|
||||
Extra map[string]interface{} `json:"extra,omitempty"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
UserAgent string `json:"userAgent"`
|
||||
URIs []string `json:"uris"`
|
||||
RequestHeaders map[string]string `json:"requestHeaders,omitempty"`
|
||||
ResponseHeaders map[string]string `json:"responseHeaders,omitempty"`
|
||||
BodyTemplate string `json:"bodyTemplate"`
|
||||
JitterMinMS int `json:"jitterMinMs"`
|
||||
JitterMaxMS int `json:"jitterMaxMs"`
|
||||
Extra map[string]interface{} `json:"extra,omitempty"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
}
|
||||
|
||||
// ----------------------------------------------------------------------------
|
||||
@@ -165,12 +166,12 @@ func (db *DB) CreateC2Listener(l *C2Listener) error {
|
||||
}
|
||||
query := `
|
||||
INSERT INTO c2_listeners (id, name, type, bind_host, bind_port, profile_id, encryption_key,
|
||||
implant_token, status, config_json, remark, created_at, started_at, last_error)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
implant_token, status, config_json, remark, owner_user_id, created_at, started_at, last_error)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
`
|
||||
_, err := db.Exec(query,
|
||||
l.ID, l.Name, l.Type, l.BindHost, l.BindPort, l.ProfileID, l.EncryptionKey,
|
||||
l.ImplantToken, l.Status, l.ConfigJSON, l.Remark, l.CreatedAt, l.StartedAt, l.LastError,
|
||||
l.ImplantToken, l.Status, l.ConfigJSON, l.Remark, l.OwnerUserID, l.CreatedAt, l.StartedAt, l.LastError,
|
||||
)
|
||||
if err != nil {
|
||||
db.logger.Error("创建 C2 监听器失败", zap.Error(err), zap.String("id", l.ID))
|
||||
@@ -190,12 +191,12 @@ func (db *DB) UpdateC2Listener(l *C2Listener) error {
|
||||
query := `
|
||||
UPDATE c2_listeners SET
|
||||
name = ?, type = ?, bind_host = ?, bind_port = ?, profile_id = ?, encryption_key = ?,
|
||||
implant_token = ?, status = ?, config_json = ?, remark = ?, started_at = ?, last_error = ?
|
||||
implant_token = ?, status = ?, config_json = ?, remark = ?, owner_user_id = ?, started_at = ?, last_error = ?
|
||||
WHERE id = ?
|
||||
`
|
||||
res, err := db.Exec(query,
|
||||
l.Name, l.Type, l.BindHost, l.BindPort, l.ProfileID, l.EncryptionKey,
|
||||
l.ImplantToken, l.Status, l.ConfigJSON, l.Remark, l.StartedAt, l.LastError, l.ID,
|
||||
l.ImplantToken, l.Status, l.ConfigJSON, l.Remark, l.OwnerUserID, l.StartedAt, l.LastError, l.ID,
|
||||
)
|
||||
if err != nil {
|
||||
db.logger.Error("更新 C2 监听器失败", zap.Error(err), zap.String("id", l.ID))
|
||||
@@ -231,7 +232,7 @@ func (db *DB) GetC2Listener(id string) (*C2Listener, error) {
|
||||
SELECT id, name, type, bind_host, bind_port, COALESCE(profile_id, ''),
|
||||
COALESCE(encryption_key, ''), COALESCE(implant_token, ''), status,
|
||||
COALESCE(config_json, '{}'), COALESCE(remark, ''),
|
||||
created_at, started_at, COALESCE(last_error, '')
|
||||
COALESCE(owner_user_id, ''), created_at, started_at, COALESCE(last_error, '')
|
||||
FROM c2_listeners WHERE id = ?
|
||||
`
|
||||
var l C2Listener
|
||||
@@ -240,7 +241,7 @@ func (db *DB) GetC2Listener(id string) (*C2Listener, error) {
|
||||
&l.ID, &l.Name, &l.Type, &l.BindHost, &l.BindPort, &l.ProfileID,
|
||||
&l.EncryptionKey, &l.ImplantToken, &l.Status,
|
||||
&l.ConfigJSON, &l.Remark,
|
||||
&l.CreatedAt, &startedAt, &l.LastError,
|
||||
&l.OwnerUserID, &l.CreatedAt, &startedAt, &l.LastError,
|
||||
)
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, nil
|
||||
@@ -261,7 +262,7 @@ func (db *DB) ListC2Listeners() ([]*C2Listener, error) {
|
||||
SELECT id, name, type, bind_host, bind_port, COALESCE(profile_id, ''),
|
||||
COALESCE(encryption_key, ''), COALESCE(implant_token, ''), status,
|
||||
COALESCE(config_json, '{}'), COALESCE(remark, ''),
|
||||
created_at, started_at, COALESCE(last_error, '')
|
||||
COALESCE(owner_user_id, ''), created_at, started_at, COALESCE(last_error, '')
|
||||
FROM c2_listeners ORDER BY created_at DESC
|
||||
`
|
||||
rows, err := db.Query(query)
|
||||
@@ -277,6 +278,47 @@ func (db *DB) ListC2Listeners() ([]*C2Listener, error) {
|
||||
&l.ID, &l.Name, &l.Type, &l.BindHost, &l.BindPort, &l.ProfileID,
|
||||
&l.EncryptionKey, &l.ImplantToken, &l.Status,
|
||||
&l.ConfigJSON, &l.Remark,
|
||||
&l.OwnerUserID, &l.CreatedAt, &startedAt, &l.LastError,
|
||||
); err != nil {
|
||||
db.logger.Warn("扫描 c2_listeners 行失败", zap.Error(err))
|
||||
continue
|
||||
}
|
||||
if startedAt.Valid {
|
||||
t := startedAt.Time
|
||||
l.StartedAt = &t
|
||||
}
|
||||
list = append(list, &l)
|
||||
}
|
||||
return list, rows.Err()
|
||||
}
|
||||
|
||||
// ListC2ListenersForAccess lists listeners visible to the resolved RBAC scope.
|
||||
func (db *DB) ListC2ListenersForAccess(access RBACListAccess) ([]*C2Listener, error) {
|
||||
conditions := []string{"1=1"}
|
||||
args := []interface{}{}
|
||||
appendC2ListenerAccessFilter(&conditions, &args, access)
|
||||
query := `
|
||||
SELECT id, name, type, bind_host, bind_port, COALESCE(profile_id, ''),
|
||||
COALESCE(encryption_key, ''), COALESCE(implant_token, ''), status,
|
||||
COALESCE(config_json, '{}'), COALESCE(remark, ''),
|
||||
COALESCE(owner_user_id, ''), created_at, started_at, COALESCE(last_error, '')
|
||||
FROM c2_listeners
|
||||
WHERE ` + strings.Join(conditions, " AND ") + `
|
||||
ORDER BY created_at DESC
|
||||
`
|
||||
rows, err := db.Query(query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var list []*C2Listener
|
||||
for rows.Next() {
|
||||
var l C2Listener
|
||||
var startedAt sql.NullTime
|
||||
if err := rows.Scan(
|
||||
&l.ID, &l.Name, &l.Type, &l.BindHost, &l.BindPort, &l.ProfileID,
|
||||
&l.EncryptionKey, &l.ImplantToken, &l.Status,
|
||||
&l.ConfigJSON, &l.Remark, &l.OwnerUserID,
|
||||
&l.CreatedAt, &startedAt, &l.LastError,
|
||||
); err != nil {
|
||||
db.logger.Warn("扫描 c2_listeners 行失败", zap.Error(err))
|
||||
@@ -291,6 +333,26 @@ func (db *DB) ListC2Listeners() ([]*C2Listener, error) {
|
||||
return list, rows.Err()
|
||||
}
|
||||
|
||||
func appendC2ListenerAccessFilter(conditions *[]string, args *[]interface{}, access RBACListAccess) {
|
||||
if access.Scope == RBACScopeAll {
|
||||
return
|
||||
}
|
||||
if access.UserID == "" {
|
||||
*conditions = append(*conditions, "1=0")
|
||||
return
|
||||
}
|
||||
clauses := []string{"owner_user_id = ?"}
|
||||
*args = append(*args, access.UserID)
|
||||
if access.Scope == RBACScopeAssigned {
|
||||
clauses = append(clauses, `EXISTS (
|
||||
SELECT 1 FROM rbac_resource_assignments ra
|
||||
WHERE ra.user_id = ? AND ra.resource_type = 'c2_listener' AND ra.resource_id = c2_listeners.id
|
||||
)`)
|
||||
*args = append(*args, access.UserID)
|
||||
}
|
||||
*conditions = append(*conditions, "("+strings.Join(clauses, " OR ")+")")
|
||||
}
|
||||
|
||||
// DeleteC2Listener 级联删除(会话/任务/文件/事件随之消失)
|
||||
func (db *DB) DeleteC2Listener(id string) error {
|
||||
res, err := db.Exec(`DELETE FROM c2_listeners WHERE id = ?`, id)
|
||||
@@ -550,6 +612,109 @@ func (db *DB) ListC2Sessions(filter ListC2SessionsFilter) ([]*C2Session, error)
|
||||
return list, rows.Err()
|
||||
}
|
||||
|
||||
// ListC2SessionsForAccess lists sessions whose parent listener is visible.
|
||||
func (db *DB) ListC2SessionsForAccess(filter ListC2SessionsFilter, access RBACListAccess) ([]*C2Session, error) {
|
||||
conditions, args := buildC2SessionsWhere(filter)
|
||||
appendC2SessionAccessFilter(&conditions, &args, access)
|
||||
query := `
|
||||
SELECT id, listener_id, implant_uuid, COALESCE(hostname,''), COALESCE(username,''),
|
||||
COALESCE(os,''), COALESCE(arch,''), COALESCE(pid, 0), COALESCE(process_name,''),
|
||||
COALESCE(is_admin, 0), COALESCE(internal_ip,''), COALESCE(external_ip,''),
|
||||
COALESCE(user_agent,''), COALESCE(sleep_seconds, 5), COALESCE(jitter_percent, 0),
|
||||
status, first_seen_at, last_check_in, COALESCE(metadata_json, '{}'),
|
||||
COALESCE(note, '')
|
||||
FROM c2_sessions
|
||||
WHERE ` + strings.Join(conditions, " AND ") + `
|
||||
ORDER BY last_check_in DESC
|
||||
`
|
||||
if filter.Limit > 0 {
|
||||
query += fmt.Sprintf(" LIMIT %d", filter.Limit)
|
||||
}
|
||||
rows, err := db.Query(query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
return db.scanC2SessionRows(rows)
|
||||
}
|
||||
|
||||
func buildC2SessionsWhere(filter ListC2SessionsFilter) ([]string, []interface{}) {
|
||||
conditions := []string{"1=1"}
|
||||
args := []interface{}{}
|
||||
if filter.ListenerID != "" {
|
||||
conditions = append(conditions, "listener_id = ?")
|
||||
args = append(args, filter.ListenerID)
|
||||
}
|
||||
if filter.Status != "" {
|
||||
conditions = append(conditions, "status = ?")
|
||||
args = append(args, filter.Status)
|
||||
}
|
||||
if filter.OS != "" {
|
||||
conditions = append(conditions, "os = ?")
|
||||
args = append(args, filter.OS)
|
||||
}
|
||||
if filter.Search != "" {
|
||||
conditions = append(conditions, "(hostname LIKE ? OR username LIKE ? OR internal_ip LIKE ?)")
|
||||
kw := "%" + filter.Search + "%"
|
||||
args = append(args, kw, kw, kw)
|
||||
}
|
||||
if filter.Suspicious {
|
||||
conditions = append(conditions, `status = 'dead' AND (
|
||||
hostname LIKE 'tcp_%' OR LOWER(COALESCE(username,'')) = 'unknown' OR COALESCE(pid, 0) = 0
|
||||
)`)
|
||||
}
|
||||
return conditions, args
|
||||
}
|
||||
|
||||
func (db *DB) scanC2SessionRows(rows *sql.Rows) ([]*C2Session, error) {
|
||||
var list []*C2Session
|
||||
for rows.Next() {
|
||||
var s C2Session
|
||||
var isAdminInt int
|
||||
var metadataJSON string
|
||||
if err := rows.Scan(
|
||||
&s.ID, &s.ListenerID, &s.ImplantUUID, &s.Hostname, &s.Username,
|
||||
&s.OS, &s.Arch, &s.PID, &s.ProcessName,
|
||||
&isAdminInt, &s.InternalIP, &s.ExternalIP,
|
||||
&s.UserAgent, &s.SleepSeconds, &s.JitterPercent,
|
||||
&s.Status, &s.FirstSeenAt, &s.LastCheckIn, &metadataJSON,
|
||||
&s.Note,
|
||||
); err != nil {
|
||||
db.logger.Warn("扫描 c2_sessions 行失败", zap.Error(err))
|
||||
continue
|
||||
}
|
||||
s.IsAdmin = isAdminInt != 0
|
||||
if metadataJSON != "" && metadataJSON != "{}" {
|
||||
_ = json.Unmarshal([]byte(metadataJSON), &s.Metadata)
|
||||
}
|
||||
list = append(list, &s)
|
||||
}
|
||||
return list, rows.Err()
|
||||
}
|
||||
|
||||
func appendC2SessionAccessFilter(conditions *[]string, args *[]interface{}, access RBACListAccess) {
|
||||
if access.Scope == RBACScopeAll {
|
||||
return
|
||||
}
|
||||
if access.UserID == "" {
|
||||
*conditions = append(*conditions, "1=0")
|
||||
return
|
||||
}
|
||||
clauses := []string{`EXISTS (
|
||||
SELECT 1 FROM c2_listeners
|
||||
WHERE c2_listeners.id = c2_sessions.listener_id AND c2_listeners.owner_user_id = ?
|
||||
)`}
|
||||
*args = append(*args, access.UserID)
|
||||
if access.Scope == RBACScopeAssigned {
|
||||
clauses = append(clauses, `EXISTS (
|
||||
SELECT 1 FROM rbac_resource_assignments ra
|
||||
WHERE ra.user_id = ? AND ra.resource_type = 'c2_listener' AND ra.resource_id = c2_sessions.listener_id
|
||||
)`)
|
||||
*args = append(*args, access.UserID)
|
||||
}
|
||||
*conditions = append(*conditions, "("+strings.Join(clauses, " OR ")+")")
|
||||
}
|
||||
|
||||
// DeleteC2Session 级联删除其 tasks/files
|
||||
func (db *DB) DeleteC2Session(id string) error {
|
||||
res, err := db.Exec(`DELETE FROM c2_sessions WHERE id = ?`, id)
|
||||
@@ -601,6 +766,29 @@ func (db *DB) DeleteC2SessionsByIDs(ids []string) (int64, error) {
|
||||
return res.RowsAffected()
|
||||
}
|
||||
|
||||
func (db *DB) DeleteC2SessionsByIDsForAccess(ids []string, access RBACListAccess) (int64, error) {
|
||||
if access.Scope == RBACScopeAll {
|
||||
return db.DeleteC2SessionsByIDs(ids)
|
||||
}
|
||||
clean := cleanC2IDs(ids)
|
||||
if len(clean) == 0 {
|
||||
return 0, ErrNoValidC2SessionIDs
|
||||
}
|
||||
placeholders := strings.Repeat("?,", len(clean)-1) + "?"
|
||||
args := make([]interface{}, 0, len(clean)+2)
|
||||
for _, id := range clean {
|
||||
args = append(args, id)
|
||||
}
|
||||
conditions := []string{"id IN (" + placeholders + ")"}
|
||||
appendC2SessionAccessFilter(&conditions, &args, access)
|
||||
query := `DELETE FROM c2_sessions WHERE ` + strings.Join(conditions, " AND ")
|
||||
res, err := db.Exec(query, args...)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return res.RowsAffected()
|
||||
}
|
||||
|
||||
// ----------------------------------------------------------------------------
|
||||
// CRUD:C2 任务
|
||||
// ----------------------------------------------------------------------------
|
||||
@@ -778,6 +966,39 @@ func buildC2TasksWhere(filter ListC2TasksFilter) (where string, args []interface
|
||||
return strings.Join(conditions, " AND "), args
|
||||
}
|
||||
|
||||
func appendC2TaskAccessFilter(conditions *[]string, args *[]interface{}, access RBACListAccess) {
|
||||
if access.Scope == RBACScopeAll {
|
||||
return
|
||||
}
|
||||
if access.UserID == "" {
|
||||
*conditions = append(*conditions, "1=0")
|
||||
return
|
||||
}
|
||||
clauses := []string{`EXISTS (
|
||||
SELECT 1 FROM c2_sessions s
|
||||
JOIN c2_listeners l ON l.id = s.listener_id
|
||||
WHERE s.id = c2_tasks.session_id AND l.owner_user_id = ?
|
||||
)`}
|
||||
*args = append(*args, access.UserID)
|
||||
if access.Scope == RBACScopeAssigned {
|
||||
clauses = append(clauses, `EXISTS (
|
||||
SELECT 1 FROM c2_sessions s
|
||||
JOIN rbac_resource_assignments ra ON ra.resource_id = s.listener_id
|
||||
WHERE s.id = c2_tasks.session_id
|
||||
AND ra.user_id = ? AND ra.resource_type = 'c2_listener'
|
||||
)`)
|
||||
*args = append(*args, access.UserID)
|
||||
}
|
||||
*conditions = append(*conditions, "("+strings.Join(clauses, " OR ")+")")
|
||||
}
|
||||
|
||||
func buildC2TasksWhereForAccess(filter ListC2TasksFilter, access RBACListAccess) (string, []interface{}) {
|
||||
where, args := buildC2TasksWhere(filter)
|
||||
conditions := []string{where}
|
||||
appendC2TaskAccessFilter(&conditions, &args, access)
|
||||
return strings.Join(conditions, " AND "), args
|
||||
}
|
||||
|
||||
// CountC2Tasks 与 ListC2Tasks 相同过滤条件下的记录总数
|
||||
func (db *DB) CountC2Tasks(filter ListC2TasksFilter) (int64, error) {
|
||||
where, args := buildC2TasksWhere(filter)
|
||||
@@ -787,6 +1008,14 @@ func (db *DB) CountC2Tasks(filter ListC2TasksFilter) (int64, error) {
|
||||
return n, err
|
||||
}
|
||||
|
||||
func (db *DB) CountC2TasksForAccess(filter ListC2TasksFilter, access RBACListAccess) (int64, error) {
|
||||
where, args := buildC2TasksWhereForAccess(filter, access)
|
||||
query := `SELECT COUNT(*) FROM c2_tasks WHERE ` + where
|
||||
var n int64
|
||||
err := db.QueryRow(query, args...).Scan(&n)
|
||||
return n, err
|
||||
}
|
||||
|
||||
// CountC2TasksQueuedOrPending 统计 queued/pending 状态任务数(仪表盘「待审任务」)
|
||||
func (db *DB) CountC2TasksQueuedOrPending(sessionID string) (int64, error) {
|
||||
conditions := []string{"status IN ('queued', 'pending')"}
|
||||
@@ -801,6 +1030,15 @@ func (db *DB) CountC2TasksQueuedOrPending(sessionID string) (int64, error) {
|
||||
return n, err
|
||||
}
|
||||
|
||||
func (db *DB) CountC2TasksQueuedOrPendingForAccess(sessionID string, access RBACListAccess) (int64, error) {
|
||||
filter := ListC2TasksFilter{SessionID: sessionID}
|
||||
where, args := buildC2TasksWhereForAccess(filter, access)
|
||||
query := `SELECT COUNT(*) FROM c2_tasks WHERE status IN ('queued', 'pending') AND ` + where
|
||||
var n int64
|
||||
err := db.QueryRow(query, args...).Scan(&n)
|
||||
return n, err
|
||||
}
|
||||
|
||||
// ListC2Tasks 任务列表,按创建时间倒序
|
||||
func (db *DB) ListC2Tasks(filter ListC2TasksFilter) ([]*C2Task, error) {
|
||||
where, args := buildC2TasksWhere(filter)
|
||||
@@ -866,6 +1104,74 @@ func (db *DB) ListC2Tasks(filter ListC2TasksFilter) ([]*C2Task, error) {
|
||||
return list, rows.Err()
|
||||
}
|
||||
|
||||
func (db *DB) ListC2TasksForAccess(filter ListC2TasksFilter, access RBACListAccess) ([]*C2Task, error) {
|
||||
where, args := buildC2TasksWhereForAccess(filter, access)
|
||||
query := `
|
||||
SELECT id, session_id, task_type, COALESCE(payload_json, '{}'),
|
||||
status, COALESCE(result_text, ''), COALESCE(result_blob_path, ''),
|
||||
COALESCE(error, ''), COALESCE(source, 'manual'),
|
||||
COALESCE(conversation_id, ''), COALESCE(approval_status, ''),
|
||||
created_at, sent_at, started_at, completed_at, COALESCE(duration_ms, 0)
|
||||
FROM c2_tasks
|
||||
WHERE ` + where + `
|
||||
ORDER BY created_at DESC
|
||||
`
|
||||
limit := filter.Limit
|
||||
offset := filter.Offset
|
||||
if offset < 0 {
|
||||
offset = 0
|
||||
}
|
||||
if limit > 0 {
|
||||
if limit > 1000 {
|
||||
limit = 1000
|
||||
}
|
||||
query += ` LIMIT ? OFFSET ?`
|
||||
args = append(args, limit, offset)
|
||||
}
|
||||
rows, err := db.Query(query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
return db.scanC2TaskRows(rows)
|
||||
}
|
||||
|
||||
func (db *DB) scanC2TaskRows(rows *sql.Rows) ([]*C2Task, error) {
|
||||
var list []*C2Task
|
||||
for rows.Next() {
|
||||
var t C2Task
|
||||
var payloadJSON string
|
||||
var sentAt, startedAt, completedAt sql.NullTime
|
||||
if err := rows.Scan(
|
||||
&t.ID, &t.SessionID, &t.TaskType, &payloadJSON,
|
||||
&t.Status, &t.ResultText, &t.ResultBlobPath,
|
||||
&t.Error, &t.Source,
|
||||
&t.ConversationID, &t.ApprovalStatus,
|
||||
&t.CreatedAt, &sentAt, &startedAt, &completedAt, &t.DurationMS,
|
||||
); err != nil {
|
||||
db.logger.Warn("扫描 c2_tasks 行失败", zap.Error(err))
|
||||
continue
|
||||
}
|
||||
if payloadJSON != "" && payloadJSON != "{}" {
|
||||
_ = json.Unmarshal([]byte(payloadJSON), &t.Payload)
|
||||
}
|
||||
if sentAt.Valid {
|
||||
x := sentAt.Time
|
||||
t.SentAt = &x
|
||||
}
|
||||
if startedAt.Valid {
|
||||
x := startedAt.Time
|
||||
t.StartedAt = &x
|
||||
}
|
||||
if completedAt.Valid {
|
||||
x := completedAt.Time
|
||||
t.CompletedAt = &x
|
||||
}
|
||||
list = append(list, &t)
|
||||
}
|
||||
return list, rows.Err()
|
||||
}
|
||||
|
||||
// PopQueuedC2Tasks 取出某会话所有 queued/approved 任务(用于 beacon 拉取),原子置为 sent
|
||||
func (db *DB) PopQueuedC2Tasks(sessionID string, limit int) ([]*C2Task, error) {
|
||||
if limit <= 0 {
|
||||
@@ -978,6 +1284,29 @@ func (db *DB) DeleteC2TasksByIDs(ids []string) (int64, error) {
|
||||
return res.RowsAffected()
|
||||
}
|
||||
|
||||
func (db *DB) DeleteC2TasksByIDsForAccess(ids []string, access RBACListAccess) (int64, error) {
|
||||
if access.Scope == RBACScopeAll {
|
||||
return db.DeleteC2TasksByIDs(ids)
|
||||
}
|
||||
clean := cleanC2IDs(ids)
|
||||
if len(clean) == 0 {
|
||||
return 0, ErrNoValidC2TaskIDs
|
||||
}
|
||||
placeholders := strings.Repeat("?,", len(clean)-1) + "?"
|
||||
args := make([]interface{}, 0, len(clean)+2)
|
||||
for _, id := range clean {
|
||||
args = append(args, id)
|
||||
}
|
||||
conditions := []string{"id IN (" + placeholders + ")"}
|
||||
appendC2TaskAccessFilter(&conditions, &args, access)
|
||||
query := `DELETE FROM c2_tasks WHERE ` + strings.Join(conditions, " AND ")
|
||||
res, err := db.Exec(query, args...)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return res.RowsAffected()
|
||||
}
|
||||
|
||||
// ----------------------------------------------------------------------------
|
||||
// CRUD:C2 文件
|
||||
// ----------------------------------------------------------------------------
|
||||
@@ -1024,6 +1353,27 @@ func (db *DB) ListC2FilesBySession(sessionID string) ([]*C2File, error) {
|
||||
return list, rows.Err()
|
||||
}
|
||||
|
||||
func cleanC2IDs(ids []string) []string {
|
||||
const maxBatch = 500
|
||||
if len(ids) > maxBatch {
|
||||
ids = ids[:maxBatch]
|
||||
}
|
||||
clean := make([]string, 0, len(ids))
|
||||
seen := make(map[string]struct{}, len(ids))
|
||||
for _, id := range ids {
|
||||
id = strings.TrimSpace(id)
|
||||
if !validC2TextIDForDelete(id) {
|
||||
continue
|
||||
}
|
||||
if _, ok := seen[id]; ok {
|
||||
continue
|
||||
}
|
||||
seen[id] = struct{}{}
|
||||
clean = append(clean, id)
|
||||
}
|
||||
return clean
|
||||
}
|
||||
|
||||
// ----------------------------------------------------------------------------
|
||||
// CRUD:C2 事件审计
|
||||
// ----------------------------------------------------------------------------
|
||||
@@ -1093,6 +1443,56 @@ func buildC2EventsWhere(filter ListC2EventsFilter) (where string, args []interfa
|
||||
return strings.Join(conditions, " AND "), args
|
||||
}
|
||||
|
||||
func appendC2EventAccessFilter(conditions *[]string, args *[]interface{}, access RBACListAccess) {
|
||||
if access.Scope == RBACScopeAll {
|
||||
return
|
||||
}
|
||||
if access.UserID == "" {
|
||||
*conditions = append(*conditions, "1=0")
|
||||
return
|
||||
}
|
||||
clauses := []string{`EXISTS (
|
||||
SELECT 1 FROM c2_sessions s
|
||||
JOIN c2_listeners l ON l.id = s.listener_id
|
||||
WHERE s.id = c2_events.session_id AND l.owner_user_id = ?
|
||||
)`}
|
||||
*args = append(*args, access.UserID)
|
||||
if access.Scope == RBACScopeAssigned {
|
||||
clauses = append(clauses, `EXISTS (
|
||||
SELECT 1 FROM c2_sessions s
|
||||
JOIN rbac_resource_assignments ra ON ra.resource_id = s.listener_id
|
||||
WHERE s.id = c2_events.session_id
|
||||
AND ra.user_id = ? AND ra.resource_type = 'c2_listener'
|
||||
)`)
|
||||
*args = append(*args, access.UserID)
|
||||
}
|
||||
clauses = append(clauses, `EXISTS (
|
||||
SELECT 1 FROM c2_tasks t
|
||||
JOIN c2_sessions s ON s.id = t.session_id
|
||||
JOIN c2_listeners l ON l.id = s.listener_id
|
||||
WHERE t.id = c2_events.task_id AND l.owner_user_id = ?
|
||||
)`)
|
||||
*args = append(*args, access.UserID)
|
||||
if access.Scope == RBACScopeAssigned {
|
||||
clauses = append(clauses, `EXISTS (
|
||||
SELECT 1 FROM c2_tasks t
|
||||
JOIN c2_sessions s ON s.id = t.session_id
|
||||
JOIN rbac_resource_assignments ra ON ra.resource_id = s.listener_id
|
||||
WHERE t.id = c2_events.task_id
|
||||
AND ra.user_id = ? AND ra.resource_type = 'c2_listener'
|
||||
)`)
|
||||
*args = append(*args, access.UserID)
|
||||
}
|
||||
*conditions = append(*conditions, "("+strings.Join(clauses, " OR ")+")")
|
||||
}
|
||||
|
||||
func buildC2EventsWhereForAccess(filter ListC2EventsFilter, access RBACListAccess) (string, []interface{}) {
|
||||
where, args := buildC2EventsWhere(filter)
|
||||
conditions := []string{where}
|
||||
appendC2EventAccessFilter(&conditions, &args, access)
|
||||
return strings.Join(conditions, " AND "), args
|
||||
}
|
||||
|
||||
// CountC2Events 与 ListC2Events 相同过滤条件下的记录总数
|
||||
func (db *DB) CountC2Events(filter ListC2EventsFilter) (int64, error) {
|
||||
where, args := buildC2EventsWhere(filter)
|
||||
@@ -1102,6 +1502,14 @@ func (db *DB) CountC2Events(filter ListC2EventsFilter) (int64, error) {
|
||||
return n, err
|
||||
}
|
||||
|
||||
func (db *DB) CountC2EventsForAccess(filter ListC2EventsFilter, access RBACListAccess) (int64, error) {
|
||||
where, args := buildC2EventsWhereForAccess(filter, access)
|
||||
query := `SELECT COUNT(*) FROM c2_events WHERE ` + where
|
||||
var n int64
|
||||
err := db.QueryRow(query, args...).Scan(&n)
|
||||
return n, err
|
||||
}
|
||||
|
||||
// ListC2Events 事件查询,按创建时间倒序
|
||||
func (db *DB) ListC2Events(filter ListC2EventsFilter) ([]*C2Event, error) {
|
||||
where, args := buildC2EventsWhere(filter)
|
||||
@@ -1143,6 +1551,50 @@ func (db *DB) ListC2Events(filter ListC2EventsFilter) ([]*C2Event, error) {
|
||||
return list, rows.Err()
|
||||
}
|
||||
|
||||
func (db *DB) ListC2EventsForAccess(filter ListC2EventsFilter, access RBACListAccess) ([]*C2Event, error) {
|
||||
where, args := buildC2EventsWhereForAccess(filter, access)
|
||||
limit := filter.Limit
|
||||
if limit <= 0 || limit > 1000 {
|
||||
limit = 200
|
||||
}
|
||||
offset := filter.Offset
|
||||
if offset < 0 {
|
||||
offset = 0
|
||||
}
|
||||
query := `
|
||||
SELECT id, level, category, COALESCE(session_id, ''), COALESCE(task_id, ''),
|
||||
message, COALESCE(data_json, ''), created_at
|
||||
FROM c2_events
|
||||
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()
|
||||
return scanC2EventRows(rows)
|
||||
}
|
||||
|
||||
func scanC2EventRows(rows *sql.Rows) ([]*C2Event, error) {
|
||||
var list []*C2Event
|
||||
for rows.Next() {
|
||||
var e C2Event
|
||||
var dataJSON string
|
||||
if err := rows.Scan(&e.ID, &e.Level, &e.Category, &e.SessionID, &e.TaskID,
|
||||
&e.Message, &dataJSON, &e.CreatedAt); err != nil {
|
||||
continue
|
||||
}
|
||||
if dataJSON != "" {
|
||||
_ = json.Unmarshal([]byte(dataJSON), &e.Data)
|
||||
}
|
||||
list = append(list, &e)
|
||||
}
|
||||
return list, rows.Err()
|
||||
}
|
||||
|
||||
// DeleteC2EventsByIDs 按主键批量删除事件,返回实际删除行数
|
||||
func (db *DB) DeleteC2EventsByIDs(ids []string) (int64, error) {
|
||||
if len(ids) == 0 {
|
||||
@@ -1181,6 +1633,29 @@ func (db *DB) DeleteC2EventsByIDs(ids []string) (int64, error) {
|
||||
return res.RowsAffected()
|
||||
}
|
||||
|
||||
func (db *DB) DeleteC2EventsByIDsForAccess(ids []string, access RBACListAccess) (int64, error) {
|
||||
if access.Scope == RBACScopeAll {
|
||||
return db.DeleteC2EventsByIDs(ids)
|
||||
}
|
||||
clean := cleanC2IDs(ids)
|
||||
if len(clean) == 0 {
|
||||
return 0, ErrNoValidC2EventIDs
|
||||
}
|
||||
placeholders := strings.Repeat("?,", len(clean)-1) + "?"
|
||||
args := make([]interface{}, 0, len(clean)+4)
|
||||
for _, id := range clean {
|
||||
args = append(args, id)
|
||||
}
|
||||
conditions := []string{"id IN (" + placeholders + ")"}
|
||||
appendC2EventAccessFilter(&conditions, &args, access)
|
||||
query := `DELETE FROM c2_events WHERE ` + strings.Join(conditions, " AND ")
|
||||
res, err := db.Exec(query, args...)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return res.RowsAffected()
|
||||
}
|
||||
|
||||
// ----------------------------------------------------------------------------
|
||||
// CRUD:C2 Malleable Profile
|
||||
// ----------------------------------------------------------------------------
|
||||
|
||||
@@ -0,0 +1,30 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
func (db *DB) RecordC2PayloadArtifact(filename, payloadID, listenerID, ownerUserID string) error {
|
||||
filename = strings.TrimSpace(filename)
|
||||
if filename == "" || strings.TrimSpace(listenerID) == "" || strings.TrimSpace(ownerUserID) == "" {
|
||||
return nil
|
||||
}
|
||||
_, err := db.Exec(`
|
||||
INSERT INTO c2_payload_artifacts(filename, payload_id, listener_id, owner_user_id, created_at)
|
||||
VALUES (?, ?, ?, ?, ?)
|
||||
ON CONFLICT(filename) DO UPDATE SET payload_id=excluded.payload_id, listener_id=excluded.listener_id, owner_user_id=excluded.owner_user_id, created_at=excluded.created_at
|
||||
`, filename, payloadID, listenerID, ownerUserID, time.Now())
|
||||
return err
|
||||
}
|
||||
|
||||
func (db *DB) UserCanAccessC2Payload(userID, scope, filename string) bool {
|
||||
if scope == RBACScopeAll {
|
||||
return true
|
||||
}
|
||||
var listenerID, ownerUserID string
|
||||
if err := db.QueryRow(`SELECT listener_id, owner_user_id FROM c2_payload_artifacts WHERE filename = ?`, strings.TrimSpace(filename)).Scan(&listenerID, &ownerUserID); err != nil {
|
||||
return false
|
||||
}
|
||||
return ownerUserID == strings.TrimSpace(userID) || db.UserCanAccessResource(userID, scope, "c2_listener", listenerID)
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
func (db *DB) UpsertChatUploadArtifact(relativePath, conversationID, ownerUserID string) error {
|
||||
relativePath = strings.TrimSpace(relativePath)
|
||||
conversationID = strings.TrimSpace(conversationID)
|
||||
ownerUserID = strings.TrimSpace(ownerUserID)
|
||||
if relativePath == "" || conversationID == "" || ownerUserID == "" {
|
||||
return nil
|
||||
}
|
||||
_, err := db.Exec(`
|
||||
INSERT INTO chat_upload_artifacts(relative_path, conversation_id, owner_user_id, created_at)
|
||||
VALUES (?, ?, ?, ?)
|
||||
ON CONFLICT(relative_path) DO UPDATE SET conversation_id=excluded.conversation_id, owner_user_id=excluded.owner_user_id
|
||||
`, relativePath, conversationID, ownerUserID, time.Now())
|
||||
return err
|
||||
}
|
||||
|
||||
func (db *DB) GetChatUploadArtifact(relativePath string) (conversationID, ownerUserID string, ok bool) {
|
||||
err := db.QueryRow(`SELECT conversation_id, owner_user_id FROM chat_upload_artifacts WHERE relative_path = ?`, strings.TrimSpace(relativePath)).Scan(&conversationID, &ownerUserID)
|
||||
return conversationID, ownerUserID, err == nil
|
||||
}
|
||||
|
||||
func (db *DB) DeleteChatUploadArtifactPath(relativePath string) error {
|
||||
path := strings.Trim(strings.TrimSpace(relativePath), "/")
|
||||
if path == "" {
|
||||
return nil
|
||||
}
|
||||
_, err := db.Exec(`DELETE FROM chat_upload_artifacts WHERE relative_path = ? OR relative_path LIKE ? ESCAPE '\'`, path, escapeLikePrefix(path)+"/%")
|
||||
return err
|
||||
}
|
||||
|
||||
func (db *DB) RenameChatUploadArtifactPath(oldPath, newPath string) error {
|
||||
oldPath = strings.Trim(strings.TrimSpace(oldPath), "/")
|
||||
newPath = strings.Trim(strings.TrimSpace(newPath), "/")
|
||||
if oldPath == "" || newPath == "" {
|
||||
return nil
|
||||
}
|
||||
_, err := db.Exec(`
|
||||
UPDATE chat_upload_artifacts
|
||||
SET relative_path = CASE
|
||||
WHEN relative_path = ? THEN ?
|
||||
ELSE ? || substr(relative_path, length(?) + 1)
|
||||
END
|
||||
WHERE relative_path = ? OR relative_path LIKE ? ESCAPE '\'
|
||||
`, oldPath, newPath, newPath, oldPath, oldPath, escapeLikePrefix(oldPath)+"/%")
|
||||
return err
|
||||
}
|
||||
|
||||
func escapeLikePrefix(value string) string {
|
||||
return strings.NewReplacer(`\`, `\\`, `%`, `\%`, `_`, `\_`).Replace(value)
|
||||
}
|
||||
@@ -384,6 +384,31 @@ func appendConversationProjectFilter(where string, args []interface{}, projectID
|
||||
return where + fmt.Sprintf(" AND %s = ?", col), append(args, pid)
|
||||
}
|
||||
|
||||
func appendConversationAccessFilter(where string, args []interface{}, userID, scope, alias string) (string, []interface{}) {
|
||||
userID = strings.TrimSpace(userID)
|
||||
if userID == "" || scope == RBACScopeAll {
|
||||
return where, args
|
||||
}
|
||||
prefix := ""
|
||||
if alias != "" {
|
||||
prefix = alias + "."
|
||||
}
|
||||
where += fmt.Sprintf(` AND (%sowner_user_id = ? OR EXISTS (
|
||||
SELECT 1 FROM rbac_resource_assignments ra
|
||||
WHERE ra.user_id = ? AND ra.resource_type = 'conversation' AND ra.resource_id = %sid
|
||||
) OR EXISTS (
|
||||
SELECT 1 FROM projects p
|
||||
WHERE p.id = %sproject_id AND (
|
||||
p.owner_user_id = ? OR EXISTS (
|
||||
SELECT 1 FROM rbac_resource_assignments pra
|
||||
WHERE pra.user_id = ? AND pra.resource_type = 'project' AND pra.resource_id = p.id
|
||||
)
|
||||
)
|
||||
))`, prefix, prefix, prefix)
|
||||
args = append(args, userID, userID, userID, userID)
|
||||
return where, args
|
||||
}
|
||||
|
||||
// CountConversations 统计对话数量。
|
||||
func (db *DB) CountConversations(search, projectID string) (int, error) {
|
||||
var count int
|
||||
@@ -410,6 +435,33 @@ func (db *DB) CountConversations(search, projectID string) (int, error) {
|
||||
return count, nil
|
||||
}
|
||||
|
||||
func (db *DB) CountConversationsForAccess(search, projectID, userID, scope string) (int, error) {
|
||||
var count int
|
||||
var err error
|
||||
if search != "" {
|
||||
searchPattern := "%" + search + "%"
|
||||
where := ` WHERE (c.title LIKE ?
|
||||
OR EXISTS (SELECT 1 FROM messages m WHERE m.conversation_id = c.id AND m.content LIKE ?))`
|
||||
args := []interface{}{searchPattern, searchPattern}
|
||||
where, args = appendConversationProjectFilter(where, args, projectID, "c")
|
||||
where, args = appendConversationAccessFilter(where, args, userID, scope, "c")
|
||||
err = db.QueryRow(`SELECT COUNT(*) FROM conversations c`+where, args...).Scan(&count)
|
||||
} else {
|
||||
where := ""
|
||||
args := []interface{}{}
|
||||
where, args = appendConversationProjectFilter(where, args, projectID, "")
|
||||
where, args = appendConversationAccessFilter(where, args, userID, scope, "")
|
||||
if where != "" {
|
||||
where = " WHERE" + strings.TrimPrefix(where, " AND")
|
||||
}
|
||||
err = db.QueryRow(`SELECT COUNT(*) FROM conversations`+where, args...).Scan(&count)
|
||||
}
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("统计对话失败: %w", err)
|
||||
}
|
||||
return count, nil
|
||||
}
|
||||
|
||||
func conversationOrderClause(sortBy, tableAlias string) string {
|
||||
col := "updated_at"
|
||||
if strings.TrimSpace(strings.ToLower(sortBy)) == "created_at" {
|
||||
@@ -503,6 +555,81 @@ func (db *DB) ListConversations(limit, offset int, search, sortBy, projectID str
|
||||
return conversations, nil
|
||||
}
|
||||
|
||||
func (db *DB) ListConversationsForAccess(limit, offset int, search, sortBy, projectID, userID, scope string) ([]*Conversation, error) {
|
||||
if scope == RBACScopeAll || strings.TrimSpace(userID) == "" {
|
||||
return db.ListConversations(limit, offset, search, sortBy, projectID)
|
||||
}
|
||||
var rows *sql.Rows
|
||||
var err error
|
||||
if search != "" {
|
||||
searchPattern := "%" + search + "%"
|
||||
orderClause := conversationOrderClause(sortBy, "c")
|
||||
where := ` WHERE (c.title LIKE ?
|
||||
OR EXISTS (SELECT 1 FROM messages m WHERE m.conversation_id = c.id AND m.content LIKE ?))`
|
||||
args := []interface{}{searchPattern, searchPattern}
|
||||
where, args = appendConversationProjectFilter(where, args, projectID, "c")
|
||||
where, args = appendConversationAccessFilter(where, args, userID, scope, "c")
|
||||
args = append(args, limit, offset)
|
||||
rows, err = db.Query(
|
||||
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id
|
||||
FROM conversations c`+where+`
|
||||
`+orderClause+`
|
||||
LIMIT ? OFFSET ?`, args...)
|
||||
} else {
|
||||
orderClause := conversationOrderClause(sortBy, "")
|
||||
where := ""
|
||||
args := []interface{}{}
|
||||
where, args = appendConversationProjectFilter(where, args, projectID, "")
|
||||
where, args = appendConversationAccessFilter(where, args, userID, scope, "")
|
||||
if where != "" {
|
||||
where = " WHERE" + strings.TrimPrefix(where, " AND")
|
||||
}
|
||||
args = append(args, limit, offset)
|
||||
rows, err = db.Query(
|
||||
"SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id FROM conversations"+where+" "+orderClause+" LIMIT ? OFFSET ?",
|
||||
args...)
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询对话列表失败: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
return scanConversationRows(rows)
|
||||
}
|
||||
|
||||
func scanConversationRows(rows *sql.Rows) ([]*Conversation, error) {
|
||||
var conversations []*Conversation
|
||||
for rows.Next() {
|
||||
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, &projectID); err != nil {
|
||||
return nil, fmt.Errorf("扫描对话失败: %w", err)
|
||||
}
|
||||
if projectID.Valid {
|
||||
conv.ProjectID = strings.TrimSpace(projectID.String)
|
||||
}
|
||||
var err1, err2 error
|
||||
conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05.999999999-07:00", createdAt)
|
||||
if err1 != nil {
|
||||
conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05", createdAt)
|
||||
}
|
||||
if err1 != nil {
|
||||
conv.CreatedAt, _ = time.Parse(time.RFC3339, createdAt)
|
||||
}
|
||||
conv.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05.999999999-07:00", updatedAt)
|
||||
if err2 != nil {
|
||||
conv.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05", updatedAt)
|
||||
}
|
||||
if err2 != nil {
|
||||
conv.UpdatedAt, _ = time.Parse(time.RFC3339, updatedAt)
|
||||
}
|
||||
conv.Pinned = pinned != 0
|
||||
conversations = append(conversations, &conv)
|
||||
}
|
||||
return conversations, rows.Err()
|
||||
}
|
||||
|
||||
const ungroupedConversationsSQL = `
|
||||
FROM conversations c
|
||||
WHERE NOT EXISTS (
|
||||
@@ -521,6 +648,18 @@ func (db *DB) CountUngroupedConversations(projectID string) (int, error) {
|
||||
return count, nil
|
||||
}
|
||||
|
||||
func (db *DB) CountUngroupedConversationsForAccess(projectID, userID, scope string) (int, error) {
|
||||
where := ungroupedConversationsSQL
|
||||
args := []interface{}{}
|
||||
where, args = appendConversationProjectFilter(where, args, projectID, "c")
|
||||
where, args = appendConversationAccessFilter(where, args, userID, scope, "c")
|
||||
var count int
|
||||
if err := db.QueryRow(`SELECT COUNT(*) `+where, args...).Scan(&count); err != nil {
|
||||
return 0, fmt.Errorf("统计未分组对话失败: %w", err)
|
||||
}
|
||||
return count, nil
|
||||
}
|
||||
|
||||
// ListUngroupedConversations 列出不在任何分组中的对话(最近对话侧栏)。
|
||||
func (db *DB) ListUngroupedConversations(limit, offset int, sortBy, projectID string) ([]*Conversation, error) {
|
||||
orderClause := conversationOrderClause(sortBy, "c")
|
||||
@@ -578,6 +717,30 @@ func (db *DB) ListUngroupedConversations(limit, offset int, sortBy, projectID st
|
||||
return conversations, rows.Err()
|
||||
}
|
||||
|
||||
func (db *DB) ListUngroupedConversationsForAccess(limit, offset int, sortBy, projectID, userID, scope string) ([]*Conversation, error) {
|
||||
if scope == RBACScopeAll || strings.TrimSpace(userID) == "" {
|
||||
return db.ListUngroupedConversations(limit, offset, sortBy, projectID)
|
||||
}
|
||||
orderClause := conversationOrderClause(sortBy, "c")
|
||||
where := ungroupedConversationsSQL
|
||||
args := []interface{}{}
|
||||
where, args = appendConversationProjectFilter(where, args, projectID, "c")
|
||||
where, args = appendConversationAccessFilter(where, args, userID, scope, "c")
|
||||
args = append(args, limit, offset)
|
||||
rows, err := db.Query(
|
||||
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id `+
|
||||
where+`
|
||||
`+orderClause+`
|
||||
LIMIT ? OFFSET ?`,
|
||||
args...,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询未分组对话失败: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
return scanConversationRows(rows)
|
||||
}
|
||||
|
||||
// GetConversationTitle 获取对话标题(轻量查询,不加载消息)
|
||||
func (db *DB) GetConversationTitle(id string) (string, error) {
|
||||
var title string
|
||||
@@ -1144,6 +1307,12 @@ ORDER BY created_at ASC, rowid ASC`, assistantMessageID)
|
||||
|
||||
// AddProcessDetail 添加过程详情事件
|
||||
func (db *DB) AddProcessDetail(messageID, conversationID, eventType, message string, data interface{}) error {
|
||||
_, err := db.AddProcessDetailWithID(messageID, conversationID, eventType, message, data)
|
||||
return err
|
||||
}
|
||||
|
||||
// AddProcessDetailWithID 添加过程详情事件并返回记录 ID。
|
||||
func (db *DB) AddProcessDetailWithID(messageID, conversationID, eventType, message string, data interface{}) (string, error) {
|
||||
id := uuid.New().String()
|
||||
|
||||
var dataJSON string
|
||||
@@ -1161,10 +1330,10 @@ func (db *DB) AddProcessDetail(messageID, conversationID, eventType, message str
|
||||
id, messageID, conversationID, eventType, message, dataJSON, time.Now(),
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("添加过程详情失败: %w", err)
|
||||
return "", fmt.Errorf("添加过程详情失败: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
return id, nil
|
||||
}
|
||||
|
||||
// GetProcessDetails 获取消息的过程详情
|
||||
@@ -1203,11 +1372,45 @@ func (db *DB) GetProcessDetails(messageID string) ([]ProcessDetail, error) {
|
||||
return details, nil
|
||||
}
|
||||
|
||||
// GetProcessDetailByID 获取单条过程详情。
|
||||
func (db *DB) GetProcessDetailByID(id string) (*ProcessDetail, error) {
|
||||
var detail ProcessDetail
|
||||
var createdAt string
|
||||
err := db.QueryRow(
|
||||
"SELECT id, message_id, conversation_id, event_type, message, data, created_at FROM process_details WHERE id = ?",
|
||||
id,
|
||||
).Scan(&detail.ID, &detail.MessageID, &detail.ConversationID, &detail.EventType, &detail.Message, &detail.Data, &createdAt)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询过程详情失败: %w", err)
|
||||
}
|
||||
|
||||
var parseErr error
|
||||
detail.CreatedAt, parseErr = time.Parse("2006-01-02 15:04:05.999999999-07:00", createdAt)
|
||||
if parseErr != nil {
|
||||
detail.CreatedAt, parseErr = time.Parse("2006-01-02 15:04:05", createdAt)
|
||||
}
|
||||
if parseErr != nil {
|
||||
detail.CreatedAt, _ = time.Parse(time.RFC3339, createdAt)
|
||||
}
|
||||
return &detail, nil
|
||||
}
|
||||
|
||||
// ProcessDetailsSummary 过程详情摘要(用于折叠态展示,避免全量加载)。
|
||||
type ProcessDetailsSummary struct {
|
||||
Total int `json:"total"`
|
||||
IterationCount int `json:"iterationCount"`
|
||||
MaxIteration int `json:"maxIteration"`
|
||||
Total int `json:"total"`
|
||||
IterationCount int `json:"iterationCount"`
|
||||
MaxIteration int `json:"maxIteration"`
|
||||
ToolCount int `json:"toolCount"`
|
||||
ToolExecutions []ProcessDetailsToolExecution `json:"toolExecutions,omitempty"`
|
||||
MCPExecutionIDs []string `json:"mcpExecutionIds,omitempty"`
|
||||
}
|
||||
|
||||
type ProcessDetailsToolExecution struct {
|
||||
ProcessDetailID string `json:"processDetailId,omitempty"`
|
||||
ToolName string `json:"toolName,omitempty"`
|
||||
ToolCallID string `json:"toolCallId,omitempty"`
|
||||
ExecutionID string `json:"executionId,omitempty"`
|
||||
Status string `json:"status,omitempty"`
|
||||
}
|
||||
|
||||
// GetProcessDetailsSummary 统计消息的过程详情数量与迭代轮次。
|
||||
@@ -1225,6 +1428,144 @@ func (db *DB) GetProcessDetailsSummary(messageID string) (*ProcessDetailsSummary
|
||||
return summary, nil
|
||||
}
|
||||
|
||||
if err := db.QueryRow(
|
||||
"SELECT COUNT(*) FROM process_details WHERE message_id = ? AND event_type = 'tool_call'",
|
||||
messageID,
|
||||
).Scan(&summary.ToolCount); err != nil {
|
||||
return nil, fmt.Errorf("统计工具调用详情失败: %w", err)
|
||||
}
|
||||
|
||||
execRows, err := db.Query(
|
||||
"SELECT id, event_type, data FROM process_details WHERE message_id = ? AND event_type IN ('tool_call', 'tool_result') ORDER BY created_at ASC, rowid ASC",
|
||||
messageID,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询工具执行摘要失败: %w", err)
|
||||
}
|
||||
seenExecIDs := make(map[string]bool)
|
||||
// A provider may reuse a fallback toolCallId across streaming rounds. Keep a
|
||||
// FIFO per ID instead of a single index so every persisted call gets at most
|
||||
// one result. Results without an ID fall back to the oldest unmatched call.
|
||||
toolIndexesByCallID := make(map[string][]int)
|
||||
lastMatchedToolIndexByCallID := make(map[string]int)
|
||||
matchedToolIndexes := make([]bool, 0)
|
||||
nextUnmatchedToolIdx := 0
|
||||
for execRows.Next() {
|
||||
var detailID string
|
||||
var eventType string
|
||||
var dataJSON string
|
||||
if err := execRows.Scan(&detailID, &eventType, &dataJSON); err != nil {
|
||||
execRows.Close()
|
||||
return nil, fmt.Errorf("扫描工具执行摘要失败: %w", err)
|
||||
}
|
||||
if dataJSON == "" {
|
||||
continue
|
||||
}
|
||||
var payload map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(dataJSON), &payload); err != nil {
|
||||
continue
|
||||
}
|
||||
toolName, _ := payload["toolName"].(string)
|
||||
toolName = strings.TrimSpace(toolName)
|
||||
toolCallID, _ := payload["toolCallId"].(string)
|
||||
toolCallID = strings.TrimSpace(toolCallID)
|
||||
execID, _ := payload["executionId"].(string)
|
||||
execID = strings.TrimSpace(execID)
|
||||
status := ""
|
||||
if eventType == "tool_result" {
|
||||
if success, ok := payload["success"].(bool); ok {
|
||||
if success {
|
||||
status = "completed"
|
||||
} else {
|
||||
status = "failed"
|
||||
}
|
||||
} else if isErr, ok := payload["isError"].(bool); ok && isErr {
|
||||
status = "failed"
|
||||
}
|
||||
}
|
||||
if eventType == "tool_call" {
|
||||
summary.ToolExecutions = append(summary.ToolExecutions, ProcessDetailsToolExecution{
|
||||
ProcessDetailID: strings.TrimSpace(detailID),
|
||||
ToolName: toolName,
|
||||
ToolCallID: toolCallID,
|
||||
// This summary is reconstructed from persisted history, not live
|
||||
// execution state. Until a matching result is found the honest state
|
||||
// is "result_missing", never "running".
|
||||
Status: "result_missing",
|
||||
})
|
||||
matchedToolIndexes = append(matchedToolIndexes, false)
|
||||
if toolCallID != "" {
|
||||
toolIndexesByCallID[toolCallID] = append(toolIndexesByCallID[toolCallID], len(summary.ToolExecutions)-1)
|
||||
}
|
||||
}
|
||||
if eventType == "tool_result" {
|
||||
idx := -1
|
||||
if toolCallID != "" {
|
||||
queue := toolIndexesByCallID[toolCallID]
|
||||
for len(queue) > 0 {
|
||||
candidate := queue[0]
|
||||
queue = queue[1:]
|
||||
if candidate >= 0 && candidate < len(matchedToolIndexes) && !matchedToolIndexes[candidate] {
|
||||
idx = candidate
|
||||
break
|
||||
}
|
||||
}
|
||||
toolIndexesByCallID[toolCallID] = queue
|
||||
if idx < 0 {
|
||||
// Multiple persisted result events for one call (for example an
|
||||
// agent-facing reduced result replacing an earlier preview) update
|
||||
// that call instead of consuming an unrelated FIFO entry.
|
||||
if previous, ok := lastMatchedToolIndexByCallID[toolCallID]; ok {
|
||||
idx = previous
|
||||
}
|
||||
}
|
||||
}
|
||||
if idx < 0 {
|
||||
for nextUnmatchedToolIdx < len(matchedToolIndexes) && matchedToolIndexes[nextUnmatchedToolIdx] {
|
||||
nextUnmatchedToolIdx++
|
||||
}
|
||||
if nextUnmatchedToolIdx < len(matchedToolIndexes) {
|
||||
idx = nextUnmatchedToolIdx
|
||||
nextUnmatchedToolIdx++
|
||||
}
|
||||
}
|
||||
if idx >= 0 && idx < len(summary.ToolExecutions) {
|
||||
matchedToolIndexes[idx] = true
|
||||
if toolCallID != "" {
|
||||
lastMatchedToolIndexByCallID[toolCallID] = idx
|
||||
}
|
||||
if summary.ToolExecutions[idx].ToolName == "" {
|
||||
summary.ToolExecutions[idx].ToolName = toolName
|
||||
}
|
||||
if summary.ToolExecutions[idx].ToolCallID == "" {
|
||||
summary.ToolExecutions[idx].ToolCallID = toolCallID
|
||||
}
|
||||
summary.ToolExecutions[idx].ExecutionID = execID
|
||||
if status != "" {
|
||||
summary.ToolExecutions[idx].Status = status
|
||||
}
|
||||
} else {
|
||||
summary.ToolExecutions = append(summary.ToolExecutions, ProcessDetailsToolExecution{
|
||||
ProcessDetailID: strings.TrimSpace(detailID),
|
||||
ToolName: toolName,
|
||||
ToolCallID: toolCallID,
|
||||
ExecutionID: execID,
|
||||
Status: status,
|
||||
})
|
||||
matchedToolIndexes = append(matchedToolIndexes, true)
|
||||
}
|
||||
}
|
||||
if execID != "" && !seenExecIDs[execID] {
|
||||
seenExecIDs[execID] = true
|
||||
summary.MCPExecutionIDs = append(summary.MCPExecutionIDs, execID)
|
||||
}
|
||||
}
|
||||
if err := execRows.Err(); err != nil {
|
||||
execRows.Close()
|
||||
return nil, fmt.Errorf("遍历工具执行摘要失败: %w", err)
|
||||
}
|
||||
execRows.Close()
|
||||
|
||||
rows, err := db.Query(
|
||||
"SELECT data FROM process_details WHERE message_id = ? AND event_type = 'iteration' ORDER BY created_at ASC, rowid ASC",
|
||||
messageID,
|
||||
@@ -1304,6 +1645,36 @@ func (db *DB) GetProcessDetailsPage(messageID string, limit, offset int) ([]Proc
|
||||
return details, total, nil
|
||||
}
|
||||
|
||||
// GetProcessDetailOffset 返回某条过程详情在所属消息详情流中的零基 offset。
|
||||
func (db *DB) GetProcessDetailOffset(messageID, detailID string) (int, error) {
|
||||
messageID = strings.TrimSpace(messageID)
|
||||
detailID = strings.TrimSpace(detailID)
|
||||
if messageID == "" || detailID == "" {
|
||||
return 0, fmt.Errorf("messageID and detailID are required")
|
||||
}
|
||||
var createdAt string
|
||||
var rowID int64
|
||||
if err := db.QueryRow(
|
||||
"SELECT created_at, rowid FROM process_details WHERE message_id = ? AND id = ?",
|
||||
messageID, detailID,
|
||||
).Scan(&createdAt, &rowID); err != nil {
|
||||
if err == sql.ErrNoRows {
|
||||
return 0, fmt.Errorf("过程详情不存在")
|
||||
}
|
||||
return 0, fmt.Errorf("查询过程详情锚点失败: %w", err)
|
||||
}
|
||||
var offset int
|
||||
if err := db.QueryRow(
|
||||
`SELECT COUNT(*) FROM process_details
|
||||
WHERE message_id = ?
|
||||
AND (created_at < ? OR (created_at = ? AND rowid < ?))`,
|
||||
messageID, createdAt, createdAt, rowID,
|
||||
).Scan(&offset); err != nil {
|
||||
return 0, fmt.Errorf("计算过程详情锚点位置失败: %w", err)
|
||||
}
|
||||
return offset, nil
|
||||
}
|
||||
|
||||
// GetProcessDetailsByConversation 获取对话的所有过程详情(按消息分组)
|
||||
func (db *DB) GetProcessDetailsByConversation(conversationID string) (map[string][]ProcessDetail, error) {
|
||||
rows, err := db.Query(
|
||||
|
||||
@@ -225,6 +225,8 @@ func (db *DB) initTables() error {
|
||||
start_time DATETIME NOT NULL,
|
||||
end_time DATETIME,
|
||||
duration_ms INTEGER,
|
||||
owner_user_id TEXT,
|
||||
conversation_id TEXT,
|
||||
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||||
);`
|
||||
|
||||
@@ -300,6 +302,7 @@ func (db *DB) initTables() error {
|
||||
id TEXT PRIMARY KEY,
|
||||
name TEXT NOT NULL,
|
||||
icon TEXT,
|
||||
owner_user_id TEXT,
|
||||
created_at DATETIME NOT NULL,
|
||||
updated_at DATETIME NOT NULL
|
||||
);`
|
||||
@@ -478,6 +481,7 @@ func (db *DB) initTables() error {
|
||||
status TEXT NOT NULL DEFAULT 'stopped',
|
||||
config_json TEXT NOT NULL DEFAULT '{}',
|
||||
remark TEXT NOT NULL DEFAULT '',
|
||||
owner_user_id TEXT,
|
||||
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
started_at DATETIME,
|
||||
last_error TEXT
|
||||
@@ -783,6 +787,10 @@ func (db *DB) initTables() error {
|
||||
return fmt.Errorf("创建audit_logs表失败: %w", err)
|
||||
}
|
||||
|
||||
if err := db.initRBACTables(); err != nil {
|
||||
return fmt.Errorf("创建RBAC表失败: %w", err)
|
||||
}
|
||||
|
||||
for tableName, ddl := range map[string]string{
|
||||
"workflow_definitions": createWorkflowDefinitionsTable,
|
||||
"workflow_runs": createWorkflowRunsTable,
|
||||
@@ -853,6 +861,9 @@ func (db *DB) initTables() error {
|
||||
if err := db.migrateWorkflowRunsTable(); err != nil {
|
||||
db.logger.Warn("迁移workflow_runs表失败", zap.Error(err))
|
||||
}
|
||||
if err := db.migrateRBACOwnershipColumns(); err != nil {
|
||||
db.logger.Warn("迁移RBAC资源归属字段失败", zap.Error(err))
|
||||
}
|
||||
|
||||
if _, err := db.Exec(createIndexes); err != nil {
|
||||
return fmt.Errorf("创建索引失败: %w", err)
|
||||
|
||||
+60
-23
@@ -10,20 +10,28 @@ import (
|
||||
|
||||
// ConversationGroup 对话分组
|
||||
type ConversationGroup struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Icon string `json:"icon"`
|
||||
Pinned bool `json:"pinned"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Icon string `json:"icon"`
|
||||
Pinned bool `json:"pinned"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
OwnerUserID string `json:"-"`
|
||||
}
|
||||
|
||||
// GroupExistsByName 检查分组名称是否已存在
|
||||
func (db *DB) GroupExistsByName(name string, excludeID string) (bool, error) {
|
||||
return db.groupExistsByNameForOwner(name, excludeID, "")
|
||||
}
|
||||
|
||||
func (db *DB) groupExistsByNameForOwner(name, excludeID, ownerUserID string) (bool, error) {
|
||||
var count int
|
||||
var err error
|
||||
|
||||
if excludeID != "" {
|
||||
if ownerUserID != "" && excludeID != "" {
|
||||
err = db.QueryRow("SELECT COUNT(*) FROM conversation_groups WHERE name = ? AND owner_user_id = ? AND id != ?", name, ownerUserID, excludeID).Scan(&count)
|
||||
} else if ownerUserID != "" {
|
||||
err = db.QueryRow("SELECT COUNT(*) FROM conversation_groups WHERE name = ? AND owner_user_id = ?", name, ownerUserID).Scan(&count)
|
||||
} else if excludeID != "" {
|
||||
err = db.QueryRow(
|
||||
"SELECT COUNT(*) FROM conversation_groups WHERE name = ? AND id != ?",
|
||||
name, excludeID,
|
||||
@@ -43,9 +51,13 @@ func (db *DB) GroupExistsByName(name string, excludeID string) (bool, error) {
|
||||
}
|
||||
|
||||
// CreateGroup 创建分组
|
||||
func (db *DB) CreateGroup(name, icon string) (*ConversationGroup, error) {
|
||||
func (db *DB) CreateGroup(name, icon string, owners ...string) (*ConversationGroup, error) {
|
||||
ownerUserID := ""
|
||||
if len(owners) > 0 {
|
||||
ownerUserID = owners[0]
|
||||
}
|
||||
// 检查名称是否已存在
|
||||
exists, err := db.GroupExistsByName(name, "")
|
||||
exists, err := db.groupExistsByNameForOwner(name, "", ownerUserID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -61,27 +73,39 @@ func (db *DB) CreateGroup(name, icon string) (*ConversationGroup, error) {
|
||||
}
|
||||
|
||||
_, err = db.Exec(
|
||||
"INSERT INTO conversation_groups (id, name, icon, pinned, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?)",
|
||||
id, name, icon, 0, now, now,
|
||||
"INSERT INTO conversation_groups (id, name, icon, pinned, owner_user_id, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?)",
|
||||
id, name, icon, 0, ownerUserID, now, now,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("创建分组失败: %w", err)
|
||||
}
|
||||
|
||||
return &ConversationGroup{
|
||||
ID: id,
|
||||
Name: name,
|
||||
Icon: icon,
|
||||
Pinned: false,
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
ID: id,
|
||||
Name: name,
|
||||
Icon: icon,
|
||||
Pinned: false,
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
OwnerUserID: ownerUserID,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// ListGroups 列出所有分组
|
||||
func (db *DB) ListGroups() ([]*ConversationGroup, error) {
|
||||
return db.ListGroupsForAccess("", RBACScopeAll)
|
||||
}
|
||||
|
||||
func (db *DB) ListGroupsForAccess(userID, scope string) ([]*ConversationGroup, error) {
|
||||
query := "SELECT id, name, icon, COALESCE(pinned, 0), COALESCE(owner_user_id, ''), created_at, updated_at FROM conversation_groups"
|
||||
args := []interface{}{}
|
||||
if scope != RBACScopeAll {
|
||||
query += " WHERE owner_user_id = ?"
|
||||
args = append(args, userID)
|
||||
}
|
||||
query += " ORDER BY COALESCE(pinned, 0) DESC, created_at ASC"
|
||||
rows, err := db.Query(
|
||||
"SELECT id, name, icon, COALESCE(pinned, 0), created_at, updated_at FROM conversation_groups ORDER BY COALESCE(pinned, 0) DESC, created_at ASC",
|
||||
query, args...,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询分组列表失败: %w", err)
|
||||
@@ -94,7 +118,7 @@ func (db *DB) ListGroups() ([]*ConversationGroup, error) {
|
||||
var createdAt, updatedAt string
|
||||
var pinned int
|
||||
|
||||
if err := rows.Scan(&group.ID, &group.Name, &group.Icon, &pinned, &createdAt, &updatedAt); err != nil {
|
||||
if err := rows.Scan(&group.ID, &group.Name, &group.Icon, &pinned, &group.OwnerUserID, &createdAt, &updatedAt); err != nil {
|
||||
return nil, fmt.Errorf("扫描分组失败: %w", err)
|
||||
}
|
||||
|
||||
@@ -131,9 +155,9 @@ func (db *DB) GetGroup(id string) (*ConversationGroup, error) {
|
||||
var pinned int
|
||||
|
||||
err := db.QueryRow(
|
||||
"SELECT id, name, icon, COALESCE(pinned, 0), created_at, updated_at FROM conversation_groups WHERE id = ?",
|
||||
"SELECT id, name, icon, COALESCE(pinned, 0), COALESCE(owner_user_id, ''), created_at, updated_at FROM conversation_groups WHERE id = ?",
|
||||
id,
|
||||
).Scan(&group.ID, &group.Name, &group.Icon, &pinned, &createdAt, &updatedAt)
|
||||
).Scan(&group.ID, &group.Name, &group.Icon, &pinned, &group.OwnerUserID, &createdAt, &updatedAt)
|
||||
if err != nil {
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, fmt.Errorf("分组不存在")
|
||||
@@ -164,10 +188,23 @@ func (db *DB) GetGroup(id string) (*ConversationGroup, error) {
|
||||
return &group, nil
|
||||
}
|
||||
|
||||
func (db *DB) UserCanAccessGroup(userID, scope, groupID string) bool {
|
||||
if scope == RBACScopeAll {
|
||||
return true
|
||||
}
|
||||
var count int
|
||||
err := db.QueryRow(`SELECT COUNT(*) FROM conversation_groups WHERE id = ? AND owner_user_id = ?`, groupID, userID).Scan(&count)
|
||||
return err == nil && count > 0
|
||||
}
|
||||
|
||||
// UpdateGroup 更新分组
|
||||
func (db *DB) UpdateGroup(id, name, icon string) error {
|
||||
existing, err := db.GetGroup(id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// 检查名称是否已存在(排除当前分组)
|
||||
exists, err := db.GroupExistsByName(name, id)
|
||||
exists, err := db.groupExistsByNameForOwner(name, id, existing.OwnerUserID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -46,8 +46,8 @@ func (db *DB) SaveToolExecution(exec *mcp.ToolExecution) error {
|
||||
|
||||
query := `
|
||||
INSERT OR REPLACE INTO tool_executions
|
||||
(id, tool_name, arguments, status, result, error, start_time, end_time, duration_ms, created_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
(id, tool_name, arguments, status, result, error, start_time, end_time, duration_ms, owner_user_id, conversation_id, created_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
`
|
||||
|
||||
_, err = db.Exec(query,
|
||||
@@ -60,6 +60,8 @@ func (db *DB) SaveToolExecution(exec *mcp.ToolExecution) error {
|
||||
exec.StartTime,
|
||||
endTime,
|
||||
durationMs,
|
||||
strings.TrimSpace(exec.OwnerUserID),
|
||||
strings.TrimSpace(exec.ConversationID),
|
||||
time.Now(),
|
||||
)
|
||||
|
||||
@@ -90,6 +92,10 @@ func (db *DB) UpdateToolExecutionResult(id string, result *mcp.ToolResult) error
|
||||
|
||||
// CountToolExecutions 统计工具执行记录总数
|
||||
func (db *DB) CountToolExecutions(status, toolName string) (int, error) {
|
||||
return db.CountToolExecutionsForAccess(status, toolName, RBACListAccess{Scope: RBACScopeAll})
|
||||
}
|
||||
|
||||
func (db *DB) CountToolExecutionsForAccess(status, toolName string, access RBACListAccess) (int, error) {
|
||||
query := `SELECT COUNT(*) FROM tool_executions`
|
||||
args := []interface{}{}
|
||||
conditions := []string{}
|
||||
@@ -108,6 +114,7 @@ func (db *DB) CountToolExecutions(status, toolName string) (int, error) {
|
||||
query += ` AND ` + conditions[i]
|
||||
}
|
||||
}
|
||||
query, args = appendToolExecutionAccessSQL(query, args, access, len(conditions) > 0)
|
||||
var count int
|
||||
err := db.QueryRow(query, args...).Scan(&count)
|
||||
if err != nil {
|
||||
@@ -135,7 +142,7 @@ func (db *DB) LoadToolExecutionsWithPagination(offset, limit int, status, toolNa
|
||||
}
|
||||
|
||||
query := `
|
||||
SELECT id, tool_name, arguments, status, result, error, start_time, end_time, duration_ms
|
||||
SELECT id, tool_name, arguments, status, result, error, start_time, end_time, duration_ms, COALESCE(owner_user_id, ''), COALESCE(conversation_id, '')
|
||||
FROM tool_executions
|
||||
`
|
||||
args := []interface{}{}
|
||||
@@ -183,6 +190,8 @@ func (db *DB) LoadToolExecutionsWithPagination(offset, limit int, status, toolNa
|
||||
&exec.StartTime,
|
||||
&endTime,
|
||||
&durationMs,
|
||||
&exec.OwnerUserID,
|
||||
&exec.ConversationID,
|
||||
)
|
||||
if err != nil {
|
||||
db.logger.Warn("加载执行记录失败", zap.Error(err))
|
||||
@@ -335,8 +344,62 @@ func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (db *DB) LoadToolStatsSummaryForAccess(topN int, access RBACListAccess) (*ToolStatsSummaryResult, error) {
|
||||
if access.Scope == RBACScopeAll {
|
||||
return db.LoadToolStatsSummary(topN)
|
||||
}
|
||||
if topN <= 0 {
|
||||
topN = 6
|
||||
}
|
||||
if topN > 100 {
|
||||
topN = 100
|
||||
}
|
||||
result := &ToolStatsSummaryResult{TopTools: make([]*mcp.ToolStats, 0, topN)}
|
||||
fromSQL, args := appendToolExecutionAccessSQL(` FROM tool_executions`, nil, access, false)
|
||||
var lastCall sql.NullString
|
||||
err := db.QueryRow(`SELECT COUNT(*),
|
||||
COALESCE(SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END), 0),
|
||||
COALESCE(SUM(CASE WHEN status IN ('failed', 'cancelled') THEN 1 ELSE 0 END), 0),
|
||||
MAX(start_time), COUNT(DISTINCT tool_name)`+fromSQL, args...).Scan(
|
||||
&result.Summary.TotalCalls, &result.Summary.SuccessCalls, &result.Summary.FailedCalls,
|
||||
&lastCall, &result.Summary.ToolCount,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if lastCall.Valid {
|
||||
parsed := parseDBTime(lastCall.String)
|
||||
result.Summary.LastCallTime = &parsed
|
||||
}
|
||||
rows, err := db.Query(`SELECT tool_name, COUNT(*),
|
||||
SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END),
|
||||
SUM(CASE WHEN status IN ('failed', 'cancelled') THEN 1 ELSE 0 END), MAX(start_time)`+
|
||||
fromSQL+` GROUP BY tool_name ORDER BY COUNT(*) DESC, tool_name ASC LIMIT ?`, append(args, topN)...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
for rows.Next() {
|
||||
var stat mcp.ToolStats
|
||||
var last sql.NullString
|
||||
if err := rows.Scan(&stat.ToolName, &stat.TotalCalls, &stat.SuccessCalls, &stat.FailedCalls, &last); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if last.Valid {
|
||||
parsed := parseDBTime(last.String)
|
||||
stat.LastCallTime = &parsed
|
||||
}
|
||||
result.TopTools = append(result.TopTools, &stat)
|
||||
}
|
||||
return result, rows.Err()
|
||||
}
|
||||
|
||||
// LoadToolExecutionListPage 分页加载执行记录列表(不含 arguments/result,供监控列表使用)
|
||||
func (db *DB) LoadToolExecutionListPage(offset, limit int, status, toolName string) ([]*mcp.ToolExecution, error) {
|
||||
return db.LoadToolExecutionListPageForAccess(offset, limit, status, toolName, RBACListAccess{Scope: RBACScopeAll})
|
||||
}
|
||||
|
||||
func (db *DB) LoadToolExecutionListPageForAccess(offset, limit int, status, toolName string, access RBACListAccess) ([]*mcp.ToolExecution, error) {
|
||||
if limit <= 0 {
|
||||
limit = 20
|
||||
}
|
||||
@@ -345,11 +408,13 @@ func (db *DB) LoadToolExecutionListPage(offset, limit int, status, toolName stri
|
||||
}
|
||||
|
||||
query := `
|
||||
SELECT id, tool_name, status, start_time, end_time, duration_ms
|
||||
SELECT id, tool_name, status, start_time, end_time, duration_ms, COALESCE(owner_user_id, ''), COALESCE(conversation_id, '')
|
||||
FROM tool_executions
|
||||
`
|
||||
whereSQL, args := toolExecutionsFilterSQL(status, toolName)
|
||||
query += whereSQL + ` ORDER BY start_time DESC LIMIT ? OFFSET ?`
|
||||
query += whereSQL
|
||||
query, args = appendToolExecutionAccessSQL(query, args, access, whereSQL != "")
|
||||
query += ` ORDER BY start_time DESC LIMIT ? OFFSET ?`
|
||||
args = append(args, limit, offset)
|
||||
|
||||
rows, err := db.Query(query, args...)
|
||||
@@ -371,6 +436,8 @@ func (db *DB) LoadToolExecutionListPage(offset, limit int, status, toolName stri
|
||||
&exec.StartTime,
|
||||
&endTime,
|
||||
&durationMs,
|
||||
&exec.OwnerUserID,
|
||||
&exec.ConversationID,
|
||||
); err != nil {
|
||||
db.logger.Warn("加载执行记录列表失败", zap.Error(err))
|
||||
continue
|
||||
@@ -387,10 +454,35 @@ func (db *DB) LoadToolExecutionListPage(offset, limit int, status, toolName stri
|
||||
return executions, nil
|
||||
}
|
||||
|
||||
func appendToolExecutionAccessSQL(query string, args []interface{}, access RBACListAccess, hasWhere bool) (string, []interface{}) {
|
||||
if access.Scope == RBACScopeAll {
|
||||
return query, args
|
||||
}
|
||||
userID := strings.TrimSpace(access.UserID)
|
||||
joiner := " WHERE "
|
||||
if hasWhere {
|
||||
joiner = " AND "
|
||||
}
|
||||
if userID == "" {
|
||||
return query + joiner + "1=0", args
|
||||
}
|
||||
query += joiner + `(
|
||||
owner_user_id = ?
|
||||
OR (conversation_id IS NOT NULL AND conversation_id <> '' AND (
|
||||
EXISTS (SELECT 1 FROM conversations c WHERE c.id = tool_executions.conversation_id AND c.owner_user_id = ?)
|
||||
OR EXISTS (SELECT 1 FROM rbac_resource_assignments ra WHERE ra.user_id = ? AND ra.resource_type = 'conversation' AND ra.resource_id = tool_executions.conversation_id)
|
||||
OR EXISTS (SELECT 1 FROM conversations c JOIN projects p ON p.id = c.project_id WHERE c.id = tool_executions.conversation_id AND p.owner_user_id = ?)
|
||||
OR EXISTS (SELECT 1 FROM conversations c JOIN rbac_resource_assignments pra ON pra.resource_id = c.project_id WHERE c.id = tool_executions.conversation_id AND pra.user_id = ? AND pra.resource_type = 'project')
|
||||
))
|
||||
)`
|
||||
args = append(args, userID, userID, userID, userID, userID)
|
||||
return query, args
|
||||
}
|
||||
|
||||
// GetToolExecution 根据ID获取单条工具执行记录
|
||||
func (db *DB) GetToolExecution(id string) (*mcp.ToolExecution, error) {
|
||||
query := `
|
||||
SELECT id, tool_name, arguments, status, result, error, start_time, end_time, duration_ms
|
||||
SELECT id, tool_name, arguments, status, result, error, start_time, end_time, duration_ms, COALESCE(owner_user_id, ''), COALESCE(conversation_id, '')
|
||||
FROM tool_executions
|
||||
WHERE id = ?
|
||||
`
|
||||
@@ -414,6 +506,8 @@ func (db *DB) GetToolExecution(id string) (*mcp.ToolExecution, error) {
|
||||
&exec.StartTime,
|
||||
&endTime,
|
||||
&durationMs,
|
||||
&exec.OwnerUserID,
|
||||
&exec.ConversationID,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -448,6 +542,29 @@ func (db *DB) GetToolExecution(id string) (*mcp.ToolExecution, error) {
|
||||
return &exec, nil
|
||||
}
|
||||
|
||||
// UserCanAccessToolExecution enforces ownership for monitor detail and mutation
|
||||
// endpoints. Legacy records without an owner or conversation fail closed for
|
||||
// non-global users.
|
||||
func (db *DB) UserCanAccessToolExecution(userID, scope, executionID string) bool {
|
||||
userID = strings.TrimSpace(userID)
|
||||
executionID = strings.TrimSpace(executionID)
|
||||
if userID == "" || executionID == "" {
|
||||
return false
|
||||
}
|
||||
if scope == RBACScopeAll {
|
||||
return true
|
||||
}
|
||||
var ownerUserID, conversationID sql.NullString
|
||||
if err := db.QueryRow(`SELECT owner_user_id, conversation_id FROM tool_executions WHERE id = ?`, executionID).Scan(&ownerUserID, &conversationID); err != nil {
|
||||
return false
|
||||
}
|
||||
if strings.TrimSpace(ownerUserID.String) == userID {
|
||||
return true
|
||||
}
|
||||
conversation := strings.TrimSpace(conversationID.String)
|
||||
return conversation != "" && db.UserCanAccessResource(userID, scope, "conversation", conversation)
|
||||
}
|
||||
|
||||
// CancelOrphanedRunningToolExecutions 将仍为 running 的记录批量标记为 cancelled(如进程重启后无对应执行协程)。
|
||||
func (db *DB) CancelOrphanedRunningToolExecutions(endTime time.Time, errMsg string) (int64, error) {
|
||||
errMsg = strings.TrimSpace(errMsg)
|
||||
@@ -584,7 +701,7 @@ func (db *DB) GetToolExecutionsByIds(ids []string) ([]*mcp.ToolExecution, error)
|
||||
}
|
||||
|
||||
query := `
|
||||
SELECT id, tool_name, arguments, status, result, error, start_time, end_time, duration_ms
|
||||
SELECT id, tool_name, arguments, status, result, error, start_time, end_time, duration_ms, COALESCE(owner_user_id, ''), COALESCE(conversation_id, '')
|
||||
FROM tool_executions
|
||||
WHERE id IN (` + strings.Join(placeholders, ",") + `)
|
||||
`
|
||||
@@ -614,6 +731,8 @@ func (db *DB) GetToolExecutionsByIds(ids []string) ([]*mcp.ToolExecution, error)
|
||||
&exec.StartTime,
|
||||
&endTime,
|
||||
&durationMs,
|
||||
&exec.OwnerUserID,
|
||||
&exec.ConversationID,
|
||||
)
|
||||
if err != nil {
|
||||
db.logger.Warn("加载执行记录失败", zap.Error(err))
|
||||
|
||||
@@ -0,0 +1,108 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestProcessDetailsSummaryPairsMixedIdentifiedAndIDLessResults(t *testing.T) {
|
||||
db, conversationID, messageID := setupProcessDetailsSummaryTest(t)
|
||||
for _, id := range []string{"call-1", "call-2", "call-3", "call-4"} {
|
||||
if err := db.AddProcessDetail(messageID, conversationID, "tool_call", "call", map[string]interface{}{
|
||||
"toolName": "http-framework-test", "toolCallId": id,
|
||||
}); err != nil {
|
||||
t.Fatalf("AddProcessDetail(tool_call): %v", err)
|
||||
}
|
||||
}
|
||||
results := []map[string]interface{}{
|
||||
{"toolName": "http-framework-test", "toolCallId": "call-1", "success": true},
|
||||
{"toolName": "http-framework-test", "toolCallId": "call-2", "success": true},
|
||||
{"toolName": "http-framework-test", "success": true},
|
||||
{"toolName": "http-framework-test", "success": true},
|
||||
}
|
||||
for _, result := range results {
|
||||
if err := db.AddProcessDetail(messageID, conversationID, "tool_result", "result", result); err != nil {
|
||||
t.Fatalf("AddProcessDetail(tool_result): %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
summary, err := db.GetProcessDetailsSummary(messageID)
|
||||
if err != nil {
|
||||
t.Fatalf("GetProcessDetailsSummary: %v", err)
|
||||
}
|
||||
if len(summary.ToolExecutions) != 4 {
|
||||
t.Fatalf("tool executions = %d, want 4", len(summary.ToolExecutions))
|
||||
}
|
||||
for i, execution := range summary.ToolExecutions {
|
||||
if execution.Status != "completed" {
|
||||
t.Fatalf("execution %d status = %q, want completed", i, execution.Status)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestProcessDetailsSummaryPairsRepeatedToolCallIDsFIFO(t *testing.T) {
|
||||
db, conversationID, messageID := setupProcessDetailsSummaryTest(t)
|
||||
for i := 0; i < 2; i++ {
|
||||
if err := db.AddProcessDetail(messageID, conversationID, "tool_call", "call", map[string]interface{}{
|
||||
"toolName": "execute", "toolCallId": "legacy-reused-id",
|
||||
}); err != nil {
|
||||
t.Fatalf("AddProcessDetail(tool_call): %v", err)
|
||||
}
|
||||
}
|
||||
for i := 0; i < 2; i++ {
|
||||
if err := db.AddProcessDetail(messageID, conversationID, "tool_result", "result", map[string]interface{}{
|
||||
"toolName": "execute", "toolCallId": "legacy-reused-id", "success": true,
|
||||
}); err != nil {
|
||||
t.Fatalf("AddProcessDetail(tool_result): %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
summary, err := db.GetProcessDetailsSummary(messageID)
|
||||
if err != nil {
|
||||
t.Fatalf("GetProcessDetailsSummary: %v", err)
|
||||
}
|
||||
if len(summary.ToolExecutions) != 2 {
|
||||
t.Fatalf("tool executions = %d, want 2", len(summary.ToolExecutions))
|
||||
}
|
||||
for i, execution := range summary.ToolExecutions {
|
||||
if execution.Status != "completed" {
|
||||
t.Fatalf("execution %d status = %q, want completed", i, execution.Status)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestProcessDetailsSummaryDoesNotReportPersistedOrphanAsRunning(t *testing.T) {
|
||||
db, conversationID, messageID := setupProcessDetailsSummaryTest(t)
|
||||
if err := db.AddProcessDetail(messageID, conversationID, "tool_call", "call", map[string]interface{}{
|
||||
"toolName": "execute", "toolCallId": "orphan",
|
||||
}); err != nil {
|
||||
t.Fatalf("AddProcessDetail(tool_call): %v", err)
|
||||
}
|
||||
summary, err := db.GetProcessDetailsSummary(messageID)
|
||||
if err != nil {
|
||||
t.Fatalf("GetProcessDetailsSummary: %v", err)
|
||||
}
|
||||
if len(summary.ToolExecutions) != 1 || summary.ToolExecutions[0].Status != "result_missing" {
|
||||
t.Fatalf("tool executions = %#v, want result_missing", summary.ToolExecutions)
|
||||
}
|
||||
}
|
||||
|
||||
func setupProcessDetailsSummaryTest(t *testing.T) (*DB, string, string) {
|
||||
t.Helper()
|
||||
db, err := NewDB(filepath.Join(t.TempDir(), "process-details.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatalf("NewDB: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
conversation, err := db.CreateConversation("process details", ConversationCreateMeta{})
|
||||
if err != nil {
|
||||
t.Fatalf("CreateConversation: %v", err)
|
||||
}
|
||||
message, err := db.AddMessage(conversation.ID, "assistant", "done", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("AddMessage: %v", err)
|
||||
}
|
||||
return db, conversation.ID, message.ID
|
||||
}
|
||||
@@ -58,11 +58,11 @@ type ProjectFact struct {
|
||||
|
||||
// ProjectFactListFilter 事实列表筛选。
|
||||
type ProjectFactListFilter struct {
|
||||
Category string
|
||||
Confidence string
|
||||
Search string
|
||||
RelatedVulnerabilityID string
|
||||
ExcludeDeprecated bool // 为 true 时排除 confidence=deprecated
|
||||
Category string
|
||||
Confidence string
|
||||
Search string
|
||||
RelatedVulnerabilityID string
|
||||
ExcludeDeprecated bool // 为 true 时排除 confidence=deprecated
|
||||
}
|
||||
|
||||
// CreateProject 创建项目。
|
||||
@@ -143,6 +143,19 @@ func appendProjectListFilters(query string, args []interface{}, status, search s
|
||||
return query, args
|
||||
}
|
||||
|
||||
func appendProjectAccessFilter(query string, args []interface{}, userID, scope string) (string, []interface{}) {
|
||||
userID = strings.TrimSpace(userID)
|
||||
if userID == "" || scope == RBACScopeAll {
|
||||
return query, args
|
||||
}
|
||||
query += ` AND (owner_user_id = ? OR EXISTS (
|
||||
SELECT 1 FROM rbac_resource_assignments ra
|
||||
WHERE ra.user_id = ? AND ra.resource_type = 'project' AND ra.resource_id = projects.id
|
||||
))`
|
||||
args = append(args, userID, userID)
|
||||
return query, args
|
||||
}
|
||||
|
||||
// CountProjects 统计项目数量。
|
||||
func (db *DB) CountProjects(status, search string) (int, error) {
|
||||
query := `SELECT COUNT(*) FROM projects WHERE 1=1`
|
||||
@@ -155,6 +168,18 @@ func (db *DB) CountProjects(status, search string) (int, error) {
|
||||
return count, nil
|
||||
}
|
||||
|
||||
func (db *DB) CountProjectsForAccess(status, search, userID, scope string) (int, error) {
|
||||
query := `SELECT COUNT(*) FROM projects WHERE 1=1`
|
||||
args := []interface{}{}
|
||||
query, args = appendProjectListFilters(query, args, status, search)
|
||||
query, args = appendProjectAccessFilter(query, args, userID, scope)
|
||||
var count int
|
||||
if err := db.QueryRow(query, args...).Scan(&count); err != nil {
|
||||
return 0, fmt.Errorf("统计项目失败: %w", err)
|
||||
}
|
||||
return count, nil
|
||||
}
|
||||
|
||||
// ListProjects 列出项目。
|
||||
func (db *DB) ListProjects(status, search string, limit, offset int) ([]*Project, error) {
|
||||
if limit <= 0 {
|
||||
@@ -189,6 +214,42 @@ func (db *DB) ListProjects(status, search string, limit, offset int) ([]*Project
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
func (db *DB) ListProjectsForAccess(status, search string, limit, offset int, userID, scope string) ([]*Project, error) {
|
||||
if scope == RBACScopeAll || strings.TrimSpace(userID) == "" {
|
||||
return db.ListProjects(status, search, limit, offset)
|
||||
}
|
||||
if limit <= 0 {
|
||||
limit = 50
|
||||
}
|
||||
query := `SELECT id, name, COALESCE(description,''), COALESCE(scope_json,''), status, pinned, created_at, updated_at
|
||||
FROM projects WHERE 1=1`
|
||||
args := []interface{}{}
|
||||
query, args = appendProjectListFilters(query, args, status, search)
|
||||
query, args = appendProjectAccessFilter(query, args, userID, scope)
|
||||
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()
|
||||
|
||||
@@ -33,6 +33,10 @@ type ProjectDashboardSummary struct {
|
||||
|
||||
// GetProjectDashboardSummary 聚合跨项目近期事实(仅活跃项目、排除 deprecated)。
|
||||
func (db *DB) GetProjectDashboardSummary(factLimit int) (*ProjectDashboardSummary, error) {
|
||||
return db.GetProjectDashboardSummaryForAccess(factLimit, "", "")
|
||||
}
|
||||
|
||||
func (db *DB) GetProjectDashboardSummaryForAccess(factLimit int, userID, scope string) (*ProjectDashboardSummary, error) {
|
||||
if factLimit <= 0 {
|
||||
factLimit = 5
|
||||
}
|
||||
@@ -44,25 +48,42 @@ func (db *DB) GetProjectDashboardSummary(factLimit int) (*ProjectDashboardSummar
|
||||
RecentFacts: []ProjectDashboardFact{},
|
||||
}
|
||||
|
||||
if err := db.QueryRow(`SELECT COUNT(*) FROM projects WHERE status = 'active'`).Scan(&out.Totals.ActiveProjects); err != nil {
|
||||
projectAccess := ""
|
||||
args := []interface{}{}
|
||||
userID = strings.TrimSpace(userID)
|
||||
if userID != "" && scope != RBACScopeAll {
|
||||
projectAccess = ` AND (
|
||||
p.owner_user_id = ?
|
||||
OR EXISTS (
|
||||
SELECT 1 FROM rbac_resource_assignments ra
|
||||
WHERE ra.user_id = ? AND ra.resource_type = 'project' AND ra.resource_id = p.id
|
||||
)
|
||||
)`
|
||||
args = append(args, userID, userID)
|
||||
}
|
||||
|
||||
if err := db.QueryRow(`SELECT COUNT(*) FROM projects p WHERE p.status = 'active'`+projectAccess, args...).Scan(&out.Totals.ActiveProjects); err != nil {
|
||||
return nil, fmt.Errorf("统计活跃项目失败: %w", err)
|
||||
}
|
||||
if err := db.QueryRow(
|
||||
`SELECT COUNT(*) FROM project_facts f
|
||||
INNER JOIN projects p ON p.id = f.project_id
|
||||
WHERE f.confidence != 'deprecated' AND p.status = 'active'`,
|
||||
WHERE f.confidence != 'deprecated' AND p.status = 'active'`+projectAccess,
|
||||
args...,
|
||||
).Scan(&out.Totals.TotalFacts); err != nil {
|
||||
return nil, fmt.Errorf("统计事实失败: %w", err)
|
||||
}
|
||||
|
||||
queryArgs := append([]interface{}{}, args...)
|
||||
queryArgs = append(queryArgs, factLimit)
|
||||
rows, err := db.Query(
|
||||
`SELECT f.id, f.project_id, p.name, f.fact_key, f.category, f.summary, f.confidence, f.pinned, f.updated_at
|
||||
FROM project_facts f
|
||||
INNER JOIN projects p ON p.id = f.project_id
|
||||
WHERE f.confidence != 'deprecated' AND p.status = 'active'
|
||||
WHERE f.confidence != 'deprecated' AND p.status = 'active'`+projectAccess+`
|
||||
ORDER BY f.pinned DESC, f.updated_at DESC
|
||||
LIMIT ?`,
|
||||
factLimit,
|
||||
queryArgs...,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询近期事实失败: %w", err)
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,581 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func newRBACTestDB(t *testing.T) *DB {
|
||||
t.Helper()
|
||||
db, err := NewDB(filepath.Join(t.TempDir(), "rbac.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatalf("NewDB: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
return db
|
||||
}
|
||||
|
||||
func TestRBACToolExecutionOwnershipAccess(t *testing.T) {
|
||||
db := newRBACTestDB(t)
|
||||
for _, exec := range []*mcp.ToolExecution{
|
||||
{ID: "exec-u1", ToolName: "one", Status: "completed", StartTime: time.Now(), OwnerUserID: "u1"},
|
||||
{ID: "exec-u2", ToolName: "two", Status: "completed", StartTime: time.Now(), OwnerUserID: "u2"},
|
||||
{ID: "exec-legacy", ToolName: "legacy", Status: "completed", StartTime: time.Now()},
|
||||
} {
|
||||
if err := db.SaveToolExecution(exec); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
access := RBACListAccess{UserID: "u1", Scope: RBACScopeAssigned}
|
||||
rows, err := db.LoadToolExecutionListPageForAccess(0, 20, "", "", access)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(rows) != 1 || rows[0].ID != "exec-u1" {
|
||||
t.Fatalf("rows = %#v, want only exec-u1", rows)
|
||||
}
|
||||
summary, err := db.LoadToolStatsSummaryForAccess(10, access)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if summary.Summary.TotalCalls != 1 || summary.Summary.ToolCount != 1 || len(summary.TopTools) != 1 || summary.TopTools[0].ToolName != "one" {
|
||||
t.Fatalf("scoped summary = %#v", summary)
|
||||
}
|
||||
if !db.UserCanAccessToolExecution("u1", RBACScopeAssigned, "exec-u1") {
|
||||
t.Fatal("owner could not access execution")
|
||||
}
|
||||
if db.UserCanAccessToolExecution("u1", RBACScopeAssigned, "exec-u2") {
|
||||
t.Fatal("foreign execution was accessible")
|
||||
}
|
||||
if db.UserCanAccessToolExecution("u1", RBACScopeAssigned, "exec-legacy") {
|
||||
t.Fatal("ownerless legacy execution did not fail closed")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRBACGroupAndUploadOwnership(t *testing.T) {
|
||||
db := newRBACTestDB(t)
|
||||
group1, err := db.CreateGroup("u1 group", "", "u1")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
group2, err := db.CreateGroup("u2 group", "", "u2")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
groups, err := db.ListGroupsForAccess("u1", RBACScopeAssigned)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(groups) != 1 || groups[0].ID != group1.ID {
|
||||
t.Fatalf("groups = %#v, want only %s (not %s)", groups, group1.ID, group2.ID)
|
||||
}
|
||||
if db.UserCanAccessGroup("u1", RBACScopeAssigned, group2.ID) {
|
||||
t.Fatal("foreign group was accessible")
|
||||
}
|
||||
|
||||
conversation, err := db.CreateConversation("upload", ConversationCreateMeta{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.UpsertChatUploadArtifact("2026-07-10/"+conversation.ID+"/a.txt", conversation.ID, "u1"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if conv, owner, ok := db.GetChatUploadArtifact("2026-07-10/" + conversation.ID + "/a.txt"); !ok || conv != conversation.ID || owner != "u1" {
|
||||
t.Fatalf("artifact = conv=%q owner=%q ok=%v", conv, owner, ok)
|
||||
}
|
||||
if err := db.RenameChatUploadArtifactPath("2026-07-10/"+conversation.ID+"/a.txt", "2026-07-10/"+conversation.ID+"/b.txt"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, _, ok := db.GetChatUploadArtifact("2026-07-10/" + conversation.ID + "/b.txt"); !ok {
|
||||
t.Fatal("renamed artifact metadata missing")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSystemRoleBootstrapDoesNotLeakManagementReadPermissions(t *testing.T) {
|
||||
db := newRBACTestDB(t)
|
||||
catalog := map[string]string{
|
||||
"auth:self": "self", "project:read": "projects", "project:write": "project writes",
|
||||
"agent:local-execute": "local tools",
|
||||
"rbac:read": "rbac", "config:read": "config", "audit:read": "audit", "terminal:execute": "terminal",
|
||||
"mcp:execute": "invoke", "mcp:write": "manage", "mcp:external:execute": "external invoke",
|
||||
"workflow:execute": "run", "workflow:write": "manage definitions", "knowledge:write": "manage knowledge",
|
||||
}
|
||||
if err := db.BootstrapRBAC("hash", catalog); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
viewer, err := db.CreateRBACUser("viewer-policy", "Viewer", "hash", true, []string{RBACSystemRoleViewer})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
viewerAccess, err := db.ResolveRBACAccess(viewer.ID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !viewerAccess.Permissions["project:read"] || viewerAccess.Permissions["rbac:read"] || viewerAccess.Permissions["config:read"] || viewerAccess.Permissions["audit:read"] {
|
||||
t.Fatalf("unexpected viewer permissions: %#v", viewerAccess.Permissions)
|
||||
}
|
||||
auditor, err := db.CreateRBACUser("auditor-policy", "Auditor", "hash", true, []string{RBACSystemRoleAuditor})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
auditorAccess, err := db.ResolveRBACAccess(auditor.ID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !auditorAccess.Permissions["audit:read"] || auditorAccess.Permissions["config:read"] || auditorAccess.Permissions["rbac:read"] {
|
||||
t.Fatalf("unexpected auditor permissions: %#v", auditorAccess.Permissions)
|
||||
}
|
||||
operator, err := db.CreateRBACUser("operator-policy", "Operator", "hash", true, []string{RBACSystemRoleOperator})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
operatorAccess, err := db.ResolveRBACAccess(operator.ID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !operatorAccess.Permissions["mcp:execute"] || operatorAccess.Permissions["mcp:write"] || operatorAccess.Permissions["mcp:external:execute"] {
|
||||
t.Fatalf("unexpected operator MCP permissions: %#v", operatorAccess.Permissions)
|
||||
}
|
||||
if !operatorAccess.Permissions["workflow:execute"] || operatorAccess.Permissions["workflow:write"] || operatorAccess.Permissions["knowledge:write"] {
|
||||
t.Fatalf("operator received global definition mutation permissions: %#v", operatorAccess.Permissions)
|
||||
}
|
||||
if !operatorAccess.Permissions["agent:local-execute"] {
|
||||
t.Fatalf("operator is missing explicit local tool permission: %#v", operatorAccess.Permissions)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPermissionScopeDoesNotWidenAcrossUnrelatedRoles(t *testing.T) {
|
||||
db := newRBACTestDB(t)
|
||||
catalog := map[string]string{"auth:self": "self", "project:read": "read", "project:write": "write", "audit:read": "audit"}
|
||||
if err := db.BootstrapRBAC("hash", catalog); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ownWrite, err := db.UpsertRBACRole("", "own-writer", "", RBACScopeOwn, []string{"project:write"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
user, err := db.CreateRBACUser("mixed-scope", "Mixed", "hash", true, []string{RBACSystemRoleAuditor, ownWrite.ID})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
access, err := db.ResolveRBACAccess(user.ID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if access.Scope != RBACScopeAll {
|
||||
t.Fatalf("compatibility scope = %q, want all", access.Scope)
|
||||
}
|
||||
if got := access.PermissionScopes["project:read"]; got != RBACScopeAll {
|
||||
t.Fatalf("project:read scope = %q, want all", got)
|
||||
}
|
||||
if got := access.PermissionScopes["project:write"]; got != RBACScopeOwn {
|
||||
t.Fatalf("project:write scope widened to %q, want own", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRoleRejectsUnknownPermission(t *testing.T) {
|
||||
db := newRBACTestDB(t)
|
||||
if err := db.BootstrapRBAC("hash", map[string]string{"auth:self": "self"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := db.UpsertRBACRole("", "future-role", "", RBACScopeAssigned, []string{"future:permission"}); err == nil {
|
||||
t.Fatal("unknown permission was persisted")
|
||||
}
|
||||
if _, err := db.Exec(`INSERT INTO rbac_permissions (key, description, created_at) VALUES ('stale:permission', '', ?)`, time.Now()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.BootstrapRBAC("hash", map[string]string{"auth:self": "self"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var count int
|
||||
if err := db.QueryRow(`SELECT COUNT(*) FROM rbac_permissions WHERE key = 'stale:permission'`).Scan(&count); err != nil || count != 0 {
|
||||
t.Fatalf("stale permission survived bootstrap: count=%d err=%v", count, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRBACProjectAndConversationListAccess(t *testing.T) {
|
||||
db := newRBACTestDB(t)
|
||||
p1, _ := db.CreateProject(&Project{Name: "visible"})
|
||||
p2, _ := db.CreateProject(&Project{Name: "hidden"})
|
||||
if err := db.SetResourceOwner("project", p1.ID, "u1"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
c1, _ := db.CreateConversation("visible conv", ConversationCreateMeta{ProjectID: p1.ID})
|
||||
c2, _ := db.CreateConversation("hidden conv", ConversationCreateMeta{ProjectID: p2.ID})
|
||||
_ = db.SetResourceOwner("conversation", c1.ID, "u1")
|
||||
_ = db.SetResourceOwner("conversation", c2.ID, "u2")
|
||||
|
||||
projects, err := db.ListProjectsForAccess("", "", 50, 0, "u1", RBACScopeOwn)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(projects) != 1 || projects[0].ID != p1.ID {
|
||||
t.Fatalf("projects = %#v, want only %s", projects, p1.ID)
|
||||
}
|
||||
|
||||
convs, err := db.ListConversationsForAccess(50, 0, "", "", "", "u1", RBACScopeOwn)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(convs) != 1 || convs[0].ID != c1.ID {
|
||||
t.Fatalf("conversations = %#v, want only %s", convs, c1.ID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRBACVulnerabilityAccessInheritsProject(t *testing.T) {
|
||||
db := newRBACTestDB(t)
|
||||
user, err := db.CreateRBACUser("u1", "User 1", "hash", true, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
p1, _ := db.CreateProject(&Project{Name: "visible"})
|
||||
p2, _ := db.CreateProject(&Project{Name: "hidden"})
|
||||
if err := db.AssignResourceToUser(user.ID, "project", p1.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
v1, _ := db.CreateVulnerability(&Vulnerability{ProjectID: p1.ID, Title: "v1", Severity: "high"})
|
||||
v2, _ := db.CreateVulnerability(&Vulnerability{ProjectID: p2.ID, Title: "v2", Severity: "high"})
|
||||
|
||||
items, err := db.ListVulnerabilitiesForAccess(50, 0, VulnerabilityListFilter{}, RBACListAccess{UserID: user.ID, Scope: RBACScopeAssigned})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(items) != 1 || items[0].ID != v1.ID {
|
||||
t.Fatalf("vulnerabilities = %#v, want only %s; hidden %s", items, v1.ID, v2.ID)
|
||||
}
|
||||
if !db.UserCanAccessResource(user.ID, RBACScopeAssigned, "vulnerability", v1.ID) {
|
||||
t.Fatalf("expected project assignment to allow vulnerability detail")
|
||||
}
|
||||
if db.UserCanAccessResource(user.ID, RBACScopeAssigned, "vulnerability", v2.ID) {
|
||||
t.Fatalf("unexpected access to hidden vulnerability")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRBACConversationAccessInheritsProject(t *testing.T) {
|
||||
db := newRBACTestDB(t)
|
||||
user, err := db.CreateRBACUser("project-member", "Project Member", "hash", true, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
project, err := db.CreateProject(&Project{Name: "assigned project"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
conversation, err := db.CreateConversation("project conversation", ConversationCreateMeta{ProjectID: project.ID})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.AssignResourceToUser(user.ID, "project", project.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
rows, err := db.ListConversationsForAccess(50, 0, "", "", "", user.ID, RBACScopeAssigned)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(rows) != 1 || rows[0].ID != conversation.ID {
|
||||
t.Fatalf("conversations = %#v, want project conversation %s", rows, conversation.ID)
|
||||
}
|
||||
if !db.UserCanAccessResource(user.ID, RBACScopeAssigned, "conversation", conversation.ID) {
|
||||
t.Fatal("expected project assignment to allow conversation detail")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRBACBatchResourceAssignmentValidationAndAtomicity(t *testing.T) {
|
||||
db := newRBACTestDB(t)
|
||||
user, err := db.CreateRBACUser("batch-member", "Batch Member", "hash", true, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
p1, err := db.CreateProject(&Project{Name: "p1"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
p2, err := db.CreateProject(&Project{Name: "p2"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
p3, err := db.CreateProject(&Project{Name: "p3"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
options, err := db.ListAssignableRBACResources("project", "p1", 50)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(options) != 1 || options[0].ID != p1.ID || options[0].Label != "p1" {
|
||||
t.Fatalf("resource options = %#v, want p1", options)
|
||||
}
|
||||
firstPage, err := db.ListAssignableRBACResourcesPage("project", "", 2, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
secondPage, err := db.ListAssignableRBACResourcesPage("project", "", 2, 2)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(firstPage) != 2 || len(secondPage) != 1 {
|
||||
t.Fatalf("paged resource options = %d + %d, want 2 + 1", len(firstPage), len(secondPage))
|
||||
}
|
||||
seen := map[string]bool{}
|
||||
for _, option := range append(firstPage, secondPage...) {
|
||||
seen[option.ID] = true
|
||||
}
|
||||
if !seen[p1.ID] || !seen[p2.ID] || !seen[p3.ID] {
|
||||
t.Fatalf("paged resource options missed resources: %#v", seen)
|
||||
}
|
||||
if _, err := db.ListAssignableRBACResources("secret_table", "", 50); err == nil {
|
||||
t.Fatal("expected unsupported picker resource type to fail")
|
||||
}
|
||||
|
||||
if _, err := db.AssignResourcesToUser(user.ID, "unknown_type", []string{p1.ID}); err == nil {
|
||||
t.Fatal("expected unsupported resource type to fail")
|
||||
}
|
||||
if _, err := db.AssignResourcesToUser(user.ID, "project", []string{p1.ID, "missing-project"}); err == nil {
|
||||
t.Fatal("expected missing resource to fail the entire batch")
|
||||
}
|
||||
rows, err := db.ListRBACResourceAssignments(user.ID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(rows) != 0 {
|
||||
t.Fatalf("partial grants persisted after failed batch: %#v", rows)
|
||||
}
|
||||
|
||||
created, err := db.AssignResourcesToUser(user.ID, "project", []string{p1.ID, p1.ID, p2.ID})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if created != 2 {
|
||||
t.Fatalf("created = %d, want 2 unique grants", created)
|
||||
}
|
||||
created, err = db.AssignResourcesToUser(user.ID, "project", []string{p1.ID, p2.ID})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if created != 0 {
|
||||
t.Fatalf("idempotent retry created = %d, want 0", created)
|
||||
}
|
||||
rows, err = db.ListRBACResourceAssignments(user.ID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(rows) != 2 {
|
||||
t.Fatalf("assignment count = %d, want 2", len(rows))
|
||||
}
|
||||
}
|
||||
|
||||
func TestRBACWebshellAndBatchListAccess(t *testing.T) {
|
||||
db := newRBACTestDB(t)
|
||||
ws1 := WebShellConnection{ID: "ws_visible", URL: "http://a", Type: "php", Method: "post", CreatedAt: time.Now()}
|
||||
ws2 := WebShellConnection{ID: "ws_hidden", URL: "http://b", Type: "php", Method: "post", CreatedAt: time.Now()}
|
||||
if err := db.CreateWebshellConnection(&ws1); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.CreateWebshellConnection(&ws2); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_ = db.SetResourceOwner("webshell", ws1.ID, "u1")
|
||||
_ = db.SetResourceOwner("webshell", ws2.ID, "u2")
|
||||
webshells, err := db.ListWebshellConnectionsForAccess("u1", RBACScopeOwn)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(webshells) != 1 || webshells[0].ID != ws1.ID {
|
||||
t.Fatalf("webshells = %#v, want only %s", webshells, ws1.ID)
|
||||
}
|
||||
|
||||
if err := db.CreateBatchQueue("q_visible", "visible", "", "eino_single", "manual", "", nil, "", 1, []map[string]interface{}{{"id": "t1", "message": "a"}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.CreateBatchQueue("q_hidden", "hidden", "", "eino_single", "manual", "", nil, "", 1, []map[string]interface{}{{"id": "t2", "message": "b"}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_ = db.SetResourceOwner("batch_task", "q_visible", "u1")
|
||||
_ = db.SetResourceOwner("batch_task", "q_hidden", "u2")
|
||||
queues, err := db.ListBatchQueuesForAccess(50, 0, "all", "", "u1", RBACScopeOwn)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(queues) != 1 || queues[0].ID != "q_visible" {
|
||||
t.Fatalf("queues = %#v, want only q_visible", queues)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRBACC2AccessInheritsListener(t *testing.T) {
|
||||
db := newRBACTestDB(t)
|
||||
now := time.Now()
|
||||
l1 := &C2Listener{ID: "l_visible", Name: "visible", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9001, OwnerUserID: "u1", CreatedAt: now}
|
||||
l2 := &C2Listener{ID: "l_hidden", Name: "hidden", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9002, OwnerUserID: "u2", CreatedAt: now}
|
||||
if err := db.CreateC2Listener(l1); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.CreateC2Listener(l2); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.UpsertC2Session(&C2Session{ID: "s_visible", ListenerID: l1.ID, ImplantUUID: "implant-visible", Status: "active", FirstSeenAt: now, LastCheckIn: now}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.UpsertC2Session(&C2Session{ID: "s_hidden", ListenerID: l2.ID, ImplantUUID: "implant-hidden", Status: "active", FirstSeenAt: now, LastCheckIn: now}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.CreateC2Task(&C2Task{ID: "t_visible", SessionID: "s_visible", TaskType: "shell", Status: "queued", CreatedAt: now}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.CreateC2Task(&C2Task{ID: "t_hidden", SessionID: "s_hidden", TaskType: "shell", Status: "queued", CreatedAt: now}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.AppendC2Event(&C2Event{ID: "e_visible", Level: "info", Category: "task", SessionID: "s_visible", TaskID: "t_visible", Message: "visible", CreatedAt: now}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.AppendC2Event(&C2Event{ID: "e_hidden", Level: "info", Category: "task", SessionID: "s_hidden", TaskID: "t_hidden", Message: "hidden", CreatedAt: now}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
access := RBACListAccess{UserID: "u1", Scope: RBACScopeOwn}
|
||||
listeners, err := db.ListC2ListenersForAccess(access)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(listeners) != 1 || listeners[0].ID != l1.ID {
|
||||
t.Fatalf("listeners = %#v, want only %s", listeners, l1.ID)
|
||||
}
|
||||
sessions, err := db.ListC2SessionsForAccess(ListC2SessionsFilter{}, access)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(sessions) != 1 || sessions[0].ID != "s_visible" {
|
||||
t.Fatalf("sessions = %#v, want only s_visible", sessions)
|
||||
}
|
||||
tasks, err := db.ListC2TasksForAccess(ListC2TasksFilter{}, access)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(tasks) != 1 || tasks[0].ID != "t_visible" {
|
||||
t.Fatalf("tasks = %#v, want only t_visible", tasks)
|
||||
}
|
||||
events, err := db.ListC2EventsForAccess(ListC2EventsFilter{}, access)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(events) != 1 || events[0].ID != "e_visible" {
|
||||
t.Fatalf("events = %#v, want only e_visible", events)
|
||||
}
|
||||
if !db.UserCanAccessResource("u1", RBACScopeOwn, "c2_task", "t_visible") {
|
||||
t.Fatalf("expected listener ownership to allow task detail")
|
||||
}
|
||||
if db.UserCanAccessResource("u1", RBACScopeOwn, "c2_task", "t_hidden") {
|
||||
t.Fatalf("unexpected access to hidden task")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRBACC2AssignedDeleteIsScoped(t *testing.T) {
|
||||
db := newRBACTestDB(t)
|
||||
user, err := db.CreateRBACUser("u1", "User 1", "hash", true, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
now := time.Now()
|
||||
if err := db.CreateC2Listener(&C2Listener{ID: "l_assigned", Name: "assigned", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9001, CreatedAt: now}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.CreateC2Listener(&C2Listener{ID: "l_hidden", Name: "hidden", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9002, CreatedAt: now}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.AssignResourceToUser(user.ID, "c2_listener", "l_assigned"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, row := range []struct {
|
||||
sessionID string
|
||||
listener string
|
||||
taskID string
|
||||
eventID string
|
||||
}{
|
||||
{"s_assigned", "l_assigned", "t_assigned", "e_assigned"},
|
||||
{"s_hidden", "l_hidden", "t_hidden", "e_hidden"},
|
||||
} {
|
||||
if err := db.UpsertC2Session(&C2Session{ID: row.sessionID, ListenerID: row.listener, ImplantUUID: row.sessionID + "_uuid", Status: "active", FirstSeenAt: now, LastCheckIn: now}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.CreateC2Task(&C2Task{ID: row.taskID, SessionID: row.sessionID, TaskType: "shell", Status: "queued", CreatedAt: now}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.AppendC2Event(&C2Event{ID: row.eventID, Level: "info", Category: "task", SessionID: row.sessionID, TaskID: row.taskID, Message: row.eventID, CreatedAt: now}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
access := RBACListAccess{UserID: user.ID, Scope: RBACScopeAssigned}
|
||||
n, err := db.DeleteC2TasksByIDsForAccess([]string{"t_assigned", "t_hidden"}, access)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if n != 1 {
|
||||
t.Fatalf("deleted tasks = %d, want 1", n)
|
||||
}
|
||||
if task, _ := db.GetC2Task("t_hidden"); task == nil {
|
||||
t.Fatalf("hidden task was deleted")
|
||||
}
|
||||
n, err = db.DeleteC2EventsByIDsForAccess([]string{"e_assigned", "e_hidden"}, access)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if n != 1 {
|
||||
t.Fatalf("deleted events = %d, want 1", n)
|
||||
}
|
||||
hiddenEvents, err := db.ListC2Events(ListC2EventsFilter{TaskID: "t_hidden"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(hiddenEvents) != 1 {
|
||||
t.Fatalf("hidden event count = %d, want 1", len(hiddenEvents))
|
||||
}
|
||||
}
|
||||
|
||||
func TestRBACAssignmentLabelsAndWeakTitles(t *testing.T) {
|
||||
db := newRBACTestDB(t)
|
||||
user, err := db.CreateRBACUser("label-member", "Label Member", "hash", true, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
project, err := db.CreateProject(&Project{Name: "Alpha Project"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
conversation, err := db.CreateConversation("1", ConversationCreateMeta{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := db.AssignResourcesToUser(user.ID, "project", []string{project.ID}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
options, err := db.ListAssignableRBACResources("conversation", "", 10)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(options) == 0 {
|
||||
t.Fatal("expected conversation options")
|
||||
}
|
||||
for _, option := range options {
|
||||
if option.ID == conversation.ID && !strings.Contains(option.Label, "1 ·") {
|
||||
t.Fatalf("weak conversation label = %q, want suffix with short id", option.Label)
|
||||
}
|
||||
}
|
||||
|
||||
rows, err := db.ListRBACResourceAssignments(user.ID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(rows) != 1 {
|
||||
t.Fatalf("assignments = %#v, want 1", rows)
|
||||
}
|
||||
if rows[0].ResourceLabel != "Alpha Project" {
|
||||
t.Fatalf("assignment label = %q, want Alpha Project", rows[0].ResourceLabel)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,174 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// RobotUserBinding maps one tenant-scoped platform identity to one RBAC user.
|
||||
// external_user_id must be derived from the verified platform event, never
|
||||
// from user-controlled message content.
|
||||
type RobotUserBinding struct {
|
||||
ID string `json:"id"`
|
||||
Platform string `json:"platform"`
|
||||
ExternalUserID string `json:"externalUserId"`
|
||||
RBACUserID string `json:"rbacUserId"`
|
||||
Enabled bool `json:"enabled"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
}
|
||||
|
||||
func normalizeRobotIdentity(platform, externalUserID string) (string, string, error) {
|
||||
platform = strings.ToLower(strings.TrimSpace(platform))
|
||||
externalUserID = strings.TrimSpace(externalUserID)
|
||||
if platform == "" || externalUserID == "" {
|
||||
return "", "", fmt.Errorf("robot platform and external user identity are required")
|
||||
}
|
||||
return platform, externalUserID, nil
|
||||
}
|
||||
|
||||
func (db *DB) CreateRobotBindingCode(userID, codeHash string, expiresAt time.Time) error {
|
||||
userID = strings.TrimSpace(userID)
|
||||
codeHash = strings.TrimSpace(codeHash)
|
||||
if userID == "" || codeHash == "" || !expiresAt.After(time.Now()) {
|
||||
return fmt.Errorf("invalid robot binding code")
|
||||
}
|
||||
now := time.Now()
|
||||
tx, err := db.Begin()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
// Keep only the newest active code per user and remove expired/used secrets.
|
||||
if _, err = tx.Exec(`DELETE FROM robot_binding_codes WHERE rbac_user_id = ? OR expires_at <= ? OR used_at IS NOT NULL`, userID, now); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err = tx.Exec(`INSERT INTO robot_binding_codes (code_hash, rbac_user_id, expires_at, created_at) VALUES (?, ?, ?, ?)`, codeHash, userID, expiresAt, now); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
// ConsumeRobotBindingCode atomically consumes a single-use code and binds the
|
||||
// verified platform identity. Existing bindings are deliberately replaced so
|
||||
// users can recover from stale or incorrect associations with a fresh code.
|
||||
func (db *DB) ConsumeRobotBindingCode(platform, externalUserID, codeHash string) (*RBACUser, error) {
|
||||
platform, externalUserID, err := normalizeRobotIdentity(platform, externalUserID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
codeHash = strings.TrimSpace(codeHash)
|
||||
if codeHash == "" {
|
||||
return nil, fmt.Errorf("binding code is required")
|
||||
}
|
||||
tx, err := db.Begin()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
|
||||
var userID string
|
||||
now := time.Now()
|
||||
if err = tx.QueryRow(`
|
||||
SELECT c.rbac_user_id
|
||||
FROM robot_binding_codes c
|
||||
JOIN rbac_users u ON u.id = c.rbac_user_id
|
||||
WHERE c.code_hash = ? AND c.used_at IS NULL AND c.expires_at > ? AND u.enabled = 1
|
||||
`, codeHash, now).Scan(&userID); err != nil {
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, fmt.Errorf("binding code is invalid or expired")
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
result, err := tx.Exec(`UPDATE robot_binding_codes SET used_at = ? WHERE code_hash = ? AND used_at IS NULL`, now, codeHash)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if affected, _ := result.RowsAffected(); affected != 1 {
|
||||
return nil, fmt.Errorf("binding code has already been used")
|
||||
}
|
||||
if _, err = tx.Exec(`
|
||||
INSERT INTO robot_user_bindings (id, platform, external_user_id, rbac_user_id, enabled, created_at, updated_at)
|
||||
VALUES (?, ?, ?, ?, 1, ?, ?)
|
||||
ON CONFLICT(platform, external_user_id) DO UPDATE SET
|
||||
rbac_user_id = excluded.rbac_user_id,
|
||||
enabled = 1,
|
||||
updated_at = excluded.updated_at
|
||||
`, uuid.New().String(), platform, externalUserID, userID, now, now); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err = tx.Commit(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return db.GetRBACUserByID(userID)
|
||||
}
|
||||
|
||||
func (db *DB) ResolveRobotRBACAccess(platform, externalUserID string) (*RBACAccess, error) {
|
||||
platform, externalUserID, err := normalizeRobotIdentity(platform, externalUserID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var userID string
|
||||
err = db.QueryRow(`
|
||||
SELECT b.rbac_user_id
|
||||
FROM robot_user_bindings b
|
||||
JOIN rbac_users u ON u.id = b.rbac_user_id
|
||||
WHERE b.platform = ? AND b.external_user_id = ? AND b.enabled = 1 AND u.enabled = 1
|
||||
`, platform, externalUserID).Scan(&userID)
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, fmt.Errorf("robot identity is not bound")
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return db.ResolveRBACAccess(userID)
|
||||
}
|
||||
|
||||
func (db *DB) ListRobotUserBindings(userID string) ([]RobotUserBinding, error) {
|
||||
rows, err := db.Query(`
|
||||
SELECT id, platform, external_user_id, rbac_user_id, enabled, created_at, updated_at
|
||||
FROM robot_user_bindings WHERE rbac_user_id = ? ORDER BY updated_at DESC
|
||||
`, strings.TrimSpace(userID))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []RobotUserBinding
|
||||
for rows.Next() {
|
||||
var b RobotUserBinding
|
||||
var enabled int
|
||||
var createdAt, updatedAt string
|
||||
if err := rows.Scan(&b.ID, &b.Platform, &b.ExternalUserID, &b.RBACUserID, &enabled, &createdAt, &updatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
b.Enabled = enabled != 0
|
||||
b.CreatedAt = parseDBTime(createdAt)
|
||||
b.UpdatedAt = parseDBTime(updatedAt)
|
||||
out = append(out, b)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
func (db *DB) DeleteRobotUserBindingForUser(bindingID, userID string) error {
|
||||
result, err := db.Exec(`DELETE FROM robot_user_bindings WHERE id = ? AND rbac_user_id = ?`, strings.TrimSpace(bindingID), strings.TrimSpace(userID))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if affected, _ := result.RowsAffected(); affected != 1 {
|
||||
return sql.ErrNoRows
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (db *DB) DeleteRobotIdentityBinding(platform, externalUserID string) error {
|
||||
platform, externalUserID, err := normalizeRobotIdentity(platform, externalUserID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = db.Exec(`DELETE FROM robot_user_bindings WHERE platform = ? AND external_user_id = ?`, platform, externalUserID)
|
||||
return err
|
||||
}
|
||||
@@ -0,0 +1,89 @@
|
||||
package database_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestRobotBindingCodeIsSingleUseAndPermissionsAreResolvedLive(t *testing.T) {
|
||||
db, err := database.NewDB(t.TempDir()+"/robot-identity.db", zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
if err := db.BootstrapRBAC("hash", security.PermissionCatalog); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
user, err := db.CreateRBACUser("bound-user", "Bound User", "hash", true, []string{database.RBACSystemRoleOperator})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.CreateRobotBindingCode(user.ID, "code-hash", time.Now().Add(time.Minute)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
bound, err := db.ConsumeRobotBindingCode("LARK", "t:tenant|u:user", "code-hash")
|
||||
if err != nil || bound.ID != user.ID {
|
||||
t.Fatalf("consume binding code: user=%v err=%v", bound, err)
|
||||
}
|
||||
if _, err := db.ConsumeRobotBindingCode("lark", "t:tenant|u:other", "code-hash"); err == nil {
|
||||
t.Fatal("single-use binding code was accepted twice")
|
||||
}
|
||||
access, err := db.ResolveRobotRBACAccess("lark", "t:tenant|u:user")
|
||||
if err != nil || !access.Permissions["agent:execute"] {
|
||||
t.Fatalf("resolved access does not include live role permissions: %#v err=%v", access, err)
|
||||
}
|
||||
disabled := false
|
||||
if err := db.UpdateRBACUser(user.ID, user.DisplayName, &disabled, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := db.ResolveRobotRBACAccess("lark", "t:tenant|u:user"); err == nil {
|
||||
t.Fatal("disabled RBAC user retained robot access")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRobotBindingCodeExpiryAndOwnerScopedRevocation(t *testing.T) {
|
||||
db, err := database.NewDB(t.TempDir()+"/robot-revoke.db", zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
if err := db.BootstrapRBAC("hash", security.PermissionCatalog); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
u1, _ := db.CreateRBACUser("binding-owner", "Owner", "hash", true, nil)
|
||||
u2, _ := db.CreateRBACUser("binding-other", "Other", "hash", true, nil)
|
||||
now := time.Now()
|
||||
if _, err := db.Exec(`INSERT INTO robot_binding_codes (code_hash, rbac_user_id, expires_at, created_at) VALUES (?, ?, ?, ?)`, "expired-hash", u1.ID, now.Add(-time.Minute), now.Add(-2*time.Minute)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := db.ConsumeRobotBindingCode("wecom", "t:corp|u:expired", "expired-hash"); err == nil {
|
||||
t.Fatal("expired binding code was accepted")
|
||||
}
|
||||
if err := db.CreateRobotBindingCode(u1.ID, "valid-hash", time.Now().Add(time.Minute)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := db.ConsumeRobotBindingCode("wecom", "t:corp|u:one", "valid-hash"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
bindings, err := db.ListRobotUserBindings(u1.ID)
|
||||
if err != nil || len(bindings) != 1 {
|
||||
t.Fatalf("bindings=%v err=%v", bindings, err)
|
||||
}
|
||||
if err := db.DeleteRobotUserBindingForUser(bindings[0].ID, u2.ID); err == nil {
|
||||
t.Fatal("another user revoked a binding they do not own")
|
||||
}
|
||||
if _, err := db.ResolveRobotRBACAccess("wecom", "t:corp|u:one"); err != nil {
|
||||
t.Fatalf("unauthorized revocation changed binding: %v", err)
|
||||
}
|
||||
if err := db.DeleteRobotUserBindingForUser(bindings[0].ID, u1.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := db.ResolveRobotRBACAccess("wecom", "t:corp|u:one"); err == nil {
|
||||
t.Fatal("revoked binding still resolves")
|
||||
}
|
||||
}
|
||||
@@ -23,6 +23,11 @@ type VulnerabilityListFilter struct {
|
||||
TaskTag string
|
||||
}
|
||||
|
||||
type RBACListAccess struct {
|
||||
UserID string
|
||||
Scope string
|
||||
}
|
||||
|
||||
func escapeVulnerabilityLikePattern(s string) string {
|
||||
s = strings.ReplaceAll(s, `\`, `\\`)
|
||||
s = strings.ReplaceAll(s, `%`, `\%`)
|
||||
@@ -89,6 +94,40 @@ func (f VulnerabilityListFilter) appendWhere(query string, args []interface{}) (
|
||||
return query, args
|
||||
}
|
||||
|
||||
func appendVulnerabilityAccessFilter(query string, args []interface{}, access RBACListAccess) (string, []interface{}) {
|
||||
userID := strings.TrimSpace(access.UserID)
|
||||
if userID == "" || access.Scope == RBACScopeAll {
|
||||
return query, args
|
||||
}
|
||||
query += ` AND (
|
||||
owner_user_id = ?
|
||||
OR EXISTS (
|
||||
SELECT 1 FROM rbac_resource_assignments ra
|
||||
WHERE ra.user_id = ? AND ra.resource_type = 'vulnerability' AND ra.resource_id = vulnerabilities.id
|
||||
)
|
||||
OR (
|
||||
project_id IS NOT NULL AND project_id <> '' AND (
|
||||
EXISTS (SELECT 1 FROM projects p WHERE p.id = vulnerabilities.project_id AND p.owner_user_id = ?)
|
||||
OR EXISTS (
|
||||
SELECT 1 FROM rbac_resource_assignments pra
|
||||
WHERE pra.user_id = ? AND pra.resource_type = 'project' AND pra.resource_id = vulnerabilities.project_id
|
||||
)
|
||||
)
|
||||
)
|
||||
OR (
|
||||
conversation_id IS NOT NULL AND conversation_id <> '' AND (
|
||||
EXISTS (SELECT 1 FROM conversations c WHERE c.id = vulnerabilities.conversation_id AND c.owner_user_id = ?)
|
||||
OR EXISTS (
|
||||
SELECT 1 FROM rbac_resource_assignments cra
|
||||
WHERE cra.user_id = ? AND cra.resource_type = 'conversation' AND cra.resource_id = vulnerabilities.conversation_id
|
||||
)
|
||||
)
|
||||
)
|
||||
)`
|
||||
args = append(args, userID, userID, userID, userID, userID, userID)
|
||||
return query, args
|
||||
}
|
||||
|
||||
// Vulnerability 漏洞
|
||||
type Vulnerability struct {
|
||||
ID string `json:"id"`
|
||||
@@ -190,6 +229,10 @@ func (db *DB) GetVulnerability(id string) (*Vulnerability, error) {
|
||||
|
||||
// ListVulnerabilities 列出漏洞
|
||||
func (db *DB) ListVulnerabilities(limit, offset int, filter VulnerabilityListFilter) ([]*Vulnerability, error) {
|
||||
return db.ListVulnerabilitiesForAccess(limit, offset, filter, RBACListAccess{})
|
||||
}
|
||||
|
||||
func (db *DB) ListVulnerabilitiesForAccess(limit, offset int, filter VulnerabilityListFilter, access RBACListAccess) ([]*Vulnerability, error) {
|
||||
query := `
|
||||
SELECT id, COALESCE(conversation_id,''), COALESCE(project_id,''), title, description, severity, status, conversation_tag, task_tag,
|
||||
vulnerability_type, target,
|
||||
@@ -203,6 +246,7 @@ func (db *DB) ListVulnerabilities(limit, offset int, filter VulnerabilityListFil
|
||||
`
|
||||
args := []interface{}{}
|
||||
query, args = filter.appendWhere(query, args)
|
||||
query, args = appendVulnerabilityAccessFilter(query, args, access)
|
||||
|
||||
query += " ORDER BY created_at DESC LIMIT ? OFFSET ?"
|
||||
args = append(args, limit, offset)
|
||||
@@ -235,9 +279,14 @@ func (db *DB) ListVulnerabilities(limit, offset int, filter VulnerabilityListFil
|
||||
|
||||
// CountVulnerabilities 统计漏洞总数(支持筛选条件)
|
||||
func (db *DB) CountVulnerabilities(filter VulnerabilityListFilter) (int, error) {
|
||||
return db.CountVulnerabilitiesForAccess(filter, RBACListAccess{})
|
||||
}
|
||||
|
||||
func (db *DB) CountVulnerabilitiesForAccess(filter VulnerabilityListFilter, access RBACListAccess) (int, error) {
|
||||
query := "SELECT COUNT(*) FROM vulnerabilities WHERE 1=1"
|
||||
args := []interface{}{}
|
||||
query, args = filter.appendWhere(query, args)
|
||||
query, args = appendVulnerabilityAccessFilter(query, args, access)
|
||||
|
||||
var count int
|
||||
err := db.QueryRow(query, args...).Scan(&count)
|
||||
@@ -275,6 +324,10 @@ func (db *DB) UpdateVulnerability(id string, vuln *Vulnerability) error {
|
||||
|
||||
// DeleteVulnerabilitiesByFilter 按筛选条件批量删除漏洞,返回实际删除条数
|
||||
func (db *DB) DeleteVulnerabilitiesByFilter(filter VulnerabilityListFilter) (int64, error) {
|
||||
return db.DeleteVulnerabilitiesByFilterForAccess(filter, RBACListAccess{})
|
||||
}
|
||||
|
||||
func (db *DB) DeleteVulnerabilitiesByFilterForAccess(filter VulnerabilityListFilter, access RBACListAccess) (int64, error) {
|
||||
tx, err := db.Begin()
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("开启事务失败: %w", err)
|
||||
@@ -284,6 +337,7 @@ func (db *DB) DeleteVulnerabilitiesByFilter(filter VulnerabilityListFilter) (int
|
||||
where := "WHERE 1=1"
|
||||
args := []interface{}{}
|
||||
where, args = filter.appendWhere(where, args)
|
||||
where, args = appendVulnerabilityAccessFilter(where, args, access)
|
||||
|
||||
clearQuery := `UPDATE project_facts SET related_vulnerability_id = NULL
|
||||
WHERE related_vulnerability_id IN (SELECT id FROM vulnerabilities ` + where + `)`
|
||||
@@ -329,11 +383,16 @@ func (db *DB) DeleteVulnerability(id string) error {
|
||||
|
||||
// GetVulnerabilityStats 获取漏洞统计(筛选条件与 ListVulnerabilities / CountVulnerabilities 一致)
|
||||
func (db *DB) GetVulnerabilityStats(filter VulnerabilityListFilter) (map[string]interface{}, error) {
|
||||
return db.GetVulnerabilityStatsForAccess(filter, RBACListAccess{})
|
||||
}
|
||||
|
||||
func (db *DB) GetVulnerabilityStatsForAccess(filter VulnerabilityListFilter, access RBACListAccess) (map[string]interface{}, error) {
|
||||
stats := make(map[string]interface{})
|
||||
|
||||
where := "WHERE 1=1"
|
||||
args := []interface{}{}
|
||||
where, args = filter.appendWhere(where, args)
|
||||
where, args = appendVulnerabilityAccessFilter(where, args, access)
|
||||
|
||||
// 总漏洞数
|
||||
var totalCount int
|
||||
@@ -389,6 +448,10 @@ func (db *DB) GetVulnerabilityStats(filter VulnerabilityListFilter) (map[string]
|
||||
|
||||
// GetVulnerabilityFilterOptions 获取漏洞筛选建议项
|
||||
func (db *DB) GetVulnerabilityFilterOptions() (map[string][]string, error) {
|
||||
return db.GetVulnerabilityFilterOptionsForAccess(RBACListAccess{})
|
||||
}
|
||||
|
||||
func (db *DB) GetVulnerabilityFilterOptionsForAccess(access RBACListAccess) (map[string][]string, error) {
|
||||
collect := func(query string, args ...interface{}) ([]string, error) {
|
||||
rows, err := db.Query(query, args...)
|
||||
if err != nil {
|
||||
@@ -409,31 +472,35 @@ func (db *DB) GetVulnerabilityFilterOptions() (map[string][]string, error) {
|
||||
return items, nil
|
||||
}
|
||||
|
||||
vulnIDs, err := collect(`SELECT DISTINCT id FROM vulnerabilities ORDER BY created_at DESC LIMIT 500`)
|
||||
where := "WHERE 1=1"
|
||||
accessArgs := []interface{}{}
|
||||
where, accessArgs = appendVulnerabilityAccessFilter(where, accessArgs, access)
|
||||
|
||||
vulnIDs, err := collect(`SELECT DISTINCT id FROM vulnerabilities `+where+` ORDER BY created_at DESC LIMIT 500`, accessArgs...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询漏洞ID建议失败: %w", err)
|
||||
}
|
||||
conversationIDs, err := collect(`SELECT DISTINCT conversation_id FROM vulnerabilities WHERE conversation_id IS NOT NULL AND conversation_id <> '' ORDER BY created_at DESC LIMIT 500`)
|
||||
conversationIDs, err := collect(`SELECT DISTINCT conversation_id FROM vulnerabilities `+where+` AND conversation_id IS NOT NULL AND conversation_id <> '' ORDER BY created_at DESC LIMIT 500`, accessArgs...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询会话ID建议失败: %w", err)
|
||||
}
|
||||
taskIDs, err := collect(`SELECT DISTINCT id FROM batch_tasks WHERE id <> '' ORDER BY rowid DESC LIMIT 500`)
|
||||
taskIDs, err := collect(`SELECT DISTINCT bt.id FROM batch_tasks bt JOIN vulnerabilities ON bt.conversation_id = vulnerabilities.conversation_id `+where+` AND bt.id <> '' ORDER BY bt.rowid DESC LIMIT 500`, accessArgs...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询任务ID建议失败: %w", err)
|
||||
}
|
||||
queueIDs, err := collect(`SELECT DISTINCT queue_id FROM batch_tasks WHERE queue_id <> '' ORDER BY rowid DESC LIMIT 500`)
|
||||
queueIDs, err := collect(`SELECT DISTINCT bt.queue_id FROM batch_tasks bt JOIN vulnerabilities ON bt.conversation_id = vulnerabilities.conversation_id `+where+` AND bt.queue_id <> '' ORDER BY bt.rowid DESC LIMIT 500`, accessArgs...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询队列ID建议失败: %w", err)
|
||||
}
|
||||
conversationTags, err := collect(`SELECT DISTINCT conversation_tag FROM vulnerabilities WHERE conversation_tag IS NOT NULL AND conversation_tag <> '' ORDER BY conversation_tag LIMIT 500`)
|
||||
conversationTags, err := collect(`SELECT DISTINCT conversation_tag FROM vulnerabilities `+where+` AND conversation_tag IS NOT NULL AND conversation_tag <> '' ORDER BY conversation_tag LIMIT 500`, accessArgs...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询对话标签建议失败: %w", err)
|
||||
}
|
||||
taskTags, err := collect(`SELECT DISTINCT task_tag FROM vulnerabilities WHERE task_tag IS NOT NULL AND task_tag <> '' ORDER BY task_tag LIMIT 500`)
|
||||
taskTags, err := collect(`SELECT DISTINCT task_tag FROM vulnerabilities `+where+` AND task_tag IS NOT NULL AND task_tag <> '' ORDER BY task_tag LIMIT 500`, accessArgs...)
|
||||
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`)
|
||||
projectIDs, err := collect(`SELECT DISTINCT project_id FROM vulnerabilities `+where+` AND project_id IS NOT NULL AND project_id <> '' ORDER BY created_at DESC LIMIT 200`, accessArgs...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询项目ID建议失败: %w", err)
|
||||
}
|
||||
|
||||
@@ -2,6 +2,7 @@ package database
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
@@ -59,13 +60,30 @@ func (db *DB) UpsertWebshellConnectionState(connectionID, stateJSON string) erro
|
||||
|
||||
// ListWebshellConnections 列出所有 WebShell 连接,按创建时间倒序
|
||||
func (db *DB) ListWebshellConnections() ([]WebShellConnection, error) {
|
||||
return db.ListWebshellConnectionsForAccess("", "")
|
||||
}
|
||||
|
||||
func (db *DB) ListWebshellConnectionsForAccess(userID, scope string) ([]WebShellConnection, error) {
|
||||
query := `
|
||||
SELECT id, url, password, type, method, cmd_param, remark,
|
||||
COALESCE(encoding, '') AS encoding, COALESCE(os, '') AS os, created_at
|
||||
FROM webshell_connections
|
||||
ORDER BY created_at DESC
|
||||
WHERE 1=1
|
||||
`
|
||||
rows, err := db.Query(query)
|
||||
args := []interface{}{}
|
||||
userID = strings.TrimSpace(userID)
|
||||
if userID != "" && scope != RBACScopeAll {
|
||||
query += ` AND (
|
||||
owner_user_id = ?
|
||||
OR EXISTS (
|
||||
SELECT 1 FROM rbac_resource_assignments ra
|
||||
WHERE ra.user_id = ? AND ra.resource_type = 'webshell' AND ra.resource_id = webshell_connections.id
|
||||
)
|
||||
)`
|
||||
args = append(args, userID, userID)
|
||||
}
|
||||
query += ` ORDER BY created_at DESC`
|
||||
rows, err := db.Query(query, args...)
|
||||
if err != nil {
|
||||
db.logger.Error("查询 WebShell 连接列表失败", zap.Error(err))
|
||||
return nil, err
|
||||
|
||||
+138
-23
@@ -18,12 +18,14 @@ import (
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/reasoning"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
"cyberstrike-ai/internal/multiagent"
|
||||
"cyberstrike-ai/internal/openai"
|
||||
"cyberstrike-ai/internal/reasoning"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/robfig/cron/v3"
|
||||
@@ -188,8 +190,8 @@ type AgentHandler struct {
|
||||
hitlWhitelistSaver HitlToolWhitelistSaver
|
||||
hitlStrategySaver HitlAuditStrategySaver
|
||||
hitlDefaultReviewerSaver HitlDefaultReviewerSaver
|
||||
auditLLM *openai.Client
|
||||
audit *audit.Service
|
||||
auditLLM *openai.Client
|
||||
audit *audit.Service
|
||||
}
|
||||
|
||||
// SetAudit wires platform audit logging.
|
||||
@@ -332,13 +334,13 @@ type ChatReasoningRequest struct {
|
||||
|
||||
// ChatRequest 聊天请求
|
||||
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_* 工具
|
||||
Hitl *HITLRequest `json:"hitl,omitempty"`
|
||||
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_* 工具
|
||||
Hitl *HITLRequest `json:"hitl,omitempty"`
|
||||
Reasoning *ChatReasoningRequest `json:"reasoning,omitempty"`
|
||||
// Orchestration 仅对 /api/multi-agent、/api/multi-agent/stream:deep | plan_execute | supervisor;空则等同 deep。机器人/批量等无请求体时由服务端默认 deep。/api/eino-agent* 不使用此字段。
|
||||
Orchestration string `json:"orchestration,omitempty"`
|
||||
@@ -727,7 +729,15 @@ func (h *AgentHandler) runRobotMultiAgentWithRetry(
|
||||
}
|
||||
|
||||
// ProcessMessageForRobot 供机器人(企业微信/钉钉/飞书)调用:Eino 单/多代理执行路径(含 progressCallback、过程详情),仅不发送 SSE,最后返回完整回复
|
||||
func (h *AgentHandler) ProcessMessageForRobot(ctx context.Context, platform, conversationID, message, role string) (response string, convID string, err error) {
|
||||
func (h *AgentHandler) ProcessMessageForRobot(ctx context.Context, platform string, principal authctx.Principal, conversationID, message, role string) (response string, convID string, err error) {
|
||||
ownerUserID := strings.TrimSpace(principal.UserID)
|
||||
if ownerUserID == "" {
|
||||
return "", "", fmt.Errorf("authenticated robot principal is required")
|
||||
}
|
||||
if !principal.HasPermission("agent:execute") || !principal.HasPermission("chat:read") || !principal.HasPermission("chat:write") {
|
||||
return "", "", fmt.Errorf("机器人账号缺少 agent:execute、chat:read 或 chat:write 权限")
|
||||
}
|
||||
ctx = authctx.WithPrincipal(ctx, principal)
|
||||
if conversationID == "" {
|
||||
title := safeTruncateString(message, 50)
|
||||
src := "robot"
|
||||
@@ -736,13 +746,17 @@ func (h *AgentHandler) ProcessMessageForRobot(ctx context.Context, platform, con
|
||||
}
|
||||
meta := audit.ConversationCreateMeta(src)
|
||||
meta.ProjectID = effectiveProjectID(h.config, "")
|
||||
if meta.ProjectID != "" && (!principal.HasPermission("project:read") || !h.db.UserCanAccessResource(ownerUserID, principal.ScopeFor("project:read"), "project", meta.ProjectID)) {
|
||||
meta.ProjectID = ""
|
||||
}
|
||||
conv, createErr := h.db.CreateConversation(title, meta)
|
||||
if createErr != nil {
|
||||
return "", "", fmt.Errorf("创建对话失败: %w", createErr)
|
||||
}
|
||||
conversationID = conv.ID
|
||||
_ = h.db.SetResourceOwner("conversation", conversationID, ownerUserID)
|
||||
} else {
|
||||
if _, getErr := h.db.GetConversation(conversationID); getErr != nil {
|
||||
if _, getErr := h.db.GetConversation(conversationID); getErr != nil || !h.db.UserCanAccessResource(ownerUserID, principal.ScopeFor("chat:write"), "conversation", conversationID) {
|
||||
return "", "", fmt.Errorf("对话不存在")
|
||||
}
|
||||
}
|
||||
@@ -844,6 +858,36 @@ func (h *AgentHandler) publishProgressToTaskEventBus(conversationID, eventType,
|
||||
h.taskEventBus.Publish(conversationID, sseLine)
|
||||
}
|
||||
|
||||
// enrichProgressEventData 为 SSE / taskEventBus 事件补齐 conversationId、messageId,便于前端懒加载过程详情。
|
||||
func enrichProgressEventData(data interface{}, conversationID, assistantMessageID string) interface{} {
|
||||
if strings.TrimSpace(conversationID) == "" && strings.TrimSpace(assistantMessageID) == "" {
|
||||
return data
|
||||
}
|
||||
var m map[string]interface{}
|
||||
switch v := data.(type) {
|
||||
case map[string]interface{}:
|
||||
m = make(map[string]interface{}, len(v)+2)
|
||||
for k, val := range v {
|
||||
m[k] = val
|
||||
}
|
||||
case nil:
|
||||
m = make(map[string]interface{}, 2)
|
||||
default:
|
||||
m = map[string]interface{}{"payload": data}
|
||||
}
|
||||
if id := strings.TrimSpace(assistantMessageID); id != "" {
|
||||
if existing, ok := m["messageId"]; !ok || strings.TrimSpace(fmt.Sprint(existing)) == "" {
|
||||
m["messageId"] = id
|
||||
}
|
||||
}
|
||||
if id := strings.TrimSpace(conversationID); id != "" {
|
||||
if existing, ok := m["conversationId"]; !ok || strings.TrimSpace(fmt.Sprint(existing)) == "" {
|
||||
m["conversationId"] = id
|
||||
}
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
// 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 {
|
||||
@@ -988,11 +1032,21 @@ func (h *AgentHandler) createProgressCallback(runCtx context.Context, cancelRun
|
||||
}
|
||||
}
|
||||
|
||||
// 流式:写 HTTP SSE;非流式(机器人等):镜像到 taskEventBus 供 Web 订阅
|
||||
if sendEventFunc != nil {
|
||||
sendEventFunc(eventType, message, data)
|
||||
} else {
|
||||
h.publishProgressToTaskEventBus(conversationID, eventType, message, data)
|
||||
// 工具输出片段不在详情区实时展示;完整结果由 tool_result 落库后按需拉取。
|
||||
if eventType == "tool_result_delta" {
|
||||
return
|
||||
}
|
||||
|
||||
deferToolProgressSend := eventType == "tool_call" || eventType == "tool_result"
|
||||
// 流式:写 HTTP SSE;非流式(机器人等):镜像到 taskEventBus 供 Web 订阅。
|
||||
// 工具事件需先落库拿 processDetailId,再向前端发送摘要,避免大 payload 默认进入浏览器。
|
||||
if !deferToolProgressSend {
|
||||
clientData := enrichProgressEventData(data, conversationID, assistantMessageID)
|
||||
if sendEventFunc != nil {
|
||||
sendEventFunc(eventType, message, clientData)
|
||||
} else {
|
||||
h.publishProgressToTaskEventBus(conversationID, eventType, message, clientData)
|
||||
}
|
||||
}
|
||||
|
||||
// 保存tool_call事件中的参数
|
||||
@@ -1360,9 +1414,28 @@ func (h *AgentHandler) createProgressCallback(runCtx context.Context, cancelRun
|
||||
// 在关键过程事件落库前,先把「规划中」与聚合中的 thinking / reasoning_chain 流落库
|
||||
flushResponsePlan()
|
||||
flushThinkingStreams()
|
||||
if err := h.db.AddProcessDetail(assistantMessageID, conversationID, eventType, message, data); err != nil {
|
||||
processDetailID, err := h.db.AddProcessDetailWithID(assistantMessageID, conversationID, eventType, message, data)
|
||||
if err != nil {
|
||||
h.logger.Warn("保存过程详情失败", zap.Error(err), zap.String("eventType", eventType))
|
||||
}
|
||||
if deferToolProgressSend {
|
||||
clientData := enrichProgressEventData(summarizeProcessDetailData(eventType, data), conversationID, assistantMessageID)
|
||||
if m, ok := clientData.(map[string]interface{}); ok {
|
||||
m["processDetailId"] = processDetailID
|
||||
}
|
||||
if sendEventFunc != nil {
|
||||
sendEventFunc(eventType, message, clientData)
|
||||
} else {
|
||||
h.publishProgressToTaskEventBus(conversationID, eventType, message, clientData)
|
||||
}
|
||||
}
|
||||
} else if deferToolProgressSend {
|
||||
clientData := enrichProgressEventData(summarizeProcessDetailData(eventType, data), conversationID, assistantMessageID)
|
||||
if sendEventFunc != nil {
|
||||
sendEventFunc(eventType, message, clientData)
|
||||
} else {
|
||||
h.publishProgressToTaskEventBus(conversationID, eventType, message, clientData)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1429,6 +1502,10 @@ func (h *AgentHandler) CancelAgentLoop(c *gin.Context) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if !h.agentConversationAllowed(c, req.ConversationID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
|
||||
if req.ContinueAfter {
|
||||
if h.tasks.GetTask(req.ConversationID) == nil {
|
||||
@@ -1495,10 +1572,10 @@ func (h *AgentHandler) CancelAgentLoop(c *gin.Context) {
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"status": "cancelling",
|
||||
"conversationId": req.ConversationID,
|
||||
"message": msg,
|
||||
"continueAfter": false,
|
||||
"status": "cancelling",
|
||||
"conversationId": req.ConversationID,
|
||||
"message": msg,
|
||||
"continueAfter": false,
|
||||
"interruptWithNote": false,
|
||||
})
|
||||
}
|
||||
@@ -1510,6 +1587,10 @@ func (h *AgentHandler) SubscribeAgentTaskEvents(c *gin.Context) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "conversationId is required"})
|
||||
return
|
||||
}
|
||||
if !h.agentConversationAllowed(c, conversationID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
if h.tasks.GetTask(conversationID) == nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "no active task for this conversation"})
|
||||
return
|
||||
@@ -1588,6 +1669,9 @@ func (h *AgentHandler) enrichCompletedTasksWithConversationTitles(tasks []*Compl
|
||||
// ListAgentTasks 列出所有运行中的任务
|
||||
func (h *AgentHandler) ListAgentTasks(c *gin.Context) {
|
||||
tasks := h.tasks.GetActiveTasks()
|
||||
tasks = filterSlice(tasks, func(task *AgentTask) bool {
|
||||
return task != nil && h.agentConversationAllowed(c, task.ConversationID)
|
||||
})
|
||||
h.enrichAgentTasksWithConversationTitles(tasks)
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"tasks": tasks,
|
||||
@@ -1597,12 +1681,30 @@ func (h *AgentHandler) ListAgentTasks(c *gin.Context) {
|
||||
// ListCompletedTasks 列出最近完成的任务历史
|
||||
func (h *AgentHandler) ListCompletedTasks(c *gin.Context) {
|
||||
tasks := h.tasks.GetCompletedTasks()
|
||||
tasks = filterSlice(tasks, func(task *CompletedTask) bool {
|
||||
return task != nil && h.agentConversationAllowed(c, task.ConversationID)
|
||||
})
|
||||
h.enrichCompletedTasksWithConversationTitles(tasks)
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"tasks": tasks,
|
||||
})
|
||||
}
|
||||
|
||||
func (h *AgentHandler) agentConversationAllowed(c *gin.Context, conversationID string) bool {
|
||||
session, ok := security.CurrentSession(c)
|
||||
return ok && h.db != nil && h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", strings.TrimSpace(conversationID))
|
||||
}
|
||||
|
||||
func filterSlice[T any](items []T, keep func(T) bool) []T {
|
||||
out := make([]T, 0, len(items))
|
||||
for _, item := range items {
|
||||
if keep(item) {
|
||||
out = append(out, item)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// BatchTaskRequest 批量任务请求
|
||||
type BatchTaskRequest struct {
|
||||
Title string `json:"title"` // 任务标题(可选)
|
||||
@@ -1654,6 +1756,12 @@ func (h *AgentHandler) CreateBatchQueue(c *gin.Context) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "没有有效的任务"})
|
||||
return
|
||||
}
|
||||
if session, ok := security.CurrentSession(c); ok && h.db != nil && session.Scope != database.RBACScopeAll && strings.TrimSpace(req.ProjectID) != "" {
|
||||
if !h.db.UserCanAccessResource(session.UserID, session.Scope, "project", strings.TrimSpace(req.ProjectID)) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权在该项目下创建批量任务"})
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
agentMode := config.NormalizeAgentMode(req.AgentMode)
|
||||
scheduleMode := normalizeBatchQueueScheduleMode(req.ScheduleMode)
|
||||
@@ -1678,6 +1786,10 @@ func (h *AgentHandler) CreateBatchQueue(c *gin.Context) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": createErr.Error()})
|
||||
return
|
||||
}
|
||||
if session, ok := security.CurrentSession(c); ok && h.db != nil {
|
||||
_ = h.db.SetResourceOwner("batch_task", queue.ID, session.UserID)
|
||||
_ = h.db.AssignResourceToUser(session.UserID, "batch_task", queue.ID)
|
||||
}
|
||||
started := false
|
||||
if req.ExecuteNow {
|
||||
ok, err := h.startBatchQueueExecution(queue.ID, false)
|
||||
@@ -1765,7 +1877,8 @@ func (h *AgentHandler) ListBatchQueues(c *gin.Context) {
|
||||
}
|
||||
|
||||
// 获取队列列表和总数
|
||||
queues, total, err := h.batchTaskManager.ListQueues(limit, offset, status, keyword)
|
||||
session, _ := security.CurrentSession(c)
|
||||
queues, total, err := h.batchTaskManager.ListQueuesForAccess(limit, offset, status, keyword, session.UserID, session.Scope)
|
||||
if err != nil {
|
||||
h.logger.Error("获取批量任务队列列表失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
@@ -2233,6 +2346,7 @@ func (h *AgentHandler) loadHistoryFromAgentTrace(conversationID string) ([]agent
|
||||
}
|
||||
|
||||
messageCount := len(messagesArray)
|
||||
modelFacingTrace := agent.IsModelFacingTraceJSON(traceInputJSON)
|
||||
|
||||
h.logger.Info("使用保存的代理轨迹恢复历史上下文",
|
||||
zap.String("conversationId", conversationID),
|
||||
@@ -2247,6 +2361,7 @@ func (h *AgentHandler) loadHistoryFromAgentTrace(conversationID string) ([]agent
|
||||
agentMessages := make([]agent.ChatMessage, 0, len(messagesArray))
|
||||
for _, msgMap := range messagesArray {
|
||||
msg := agent.ChatMessage{}
|
||||
msg.ModelFacingTrace = modelFacingTrace
|
||||
|
||||
// 解析role
|
||||
if role, ok := msgMap["role"].(string); ok {
|
||||
|
||||
@@ -97,3 +97,26 @@ func TestCreateProgressCallback_FlushesReasoningOnDone(t *testing.T) {
|
||||
t.Fatalf("expected reasoning_chain persisted on done, got %+v", details)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEnrichProgressEventData(t *testing.T) {
|
||||
t.Run("fills ids", func(t *testing.T) {
|
||||
out := enrichProgressEventData(map[string]interface{}{"source": "eino"}, "conv-1", "msg-1")
|
||||
m, ok := out.(map[string]interface{})
|
||||
if !ok {
|
||||
t.Fatalf("expected map, got %T", out)
|
||||
}
|
||||
if m["conversationId"] != "conv-1" || m["messageId"] != "msg-1" {
|
||||
t.Fatalf("unexpected enrichment: %+v", m)
|
||||
}
|
||||
})
|
||||
t.Run("preserves existing ids", func(t *testing.T) {
|
||||
out := enrichProgressEventData(map[string]interface{}{
|
||||
"conversationId": "keep-conv",
|
||||
"messageId": "keep-msg",
|
||||
}, "conv-1", "msg-1")
|
||||
m := out.(map[string]interface{})
|
||||
if m["conversationId"] != "keep-conv" || m["messageId"] != "keep-msg" {
|
||||
t.Fatalf("should not overwrite existing ids: %+v", m)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
@@ -32,8 +33,8 @@ func (h *AuditHandler) Meta(c *gin.Context) {
|
||||
retentionDays = h.audit.RetentionDays()
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"enabled": enabled,
|
||||
"retention_days": retentionDays,
|
||||
"enabled": enabled,
|
||||
"retention_days": retentionDays,
|
||||
"default_page_size": 20,
|
||||
"max_page_size": 100,
|
||||
"max_export": 5000,
|
||||
@@ -46,7 +47,7 @@ func (h *AuditHandler) Summary(c *gin.Context) {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "database unavailable"})
|
||||
return
|
||||
}
|
||||
base := auditFilterFromQuery(c)
|
||||
base := auditFilterForAccess(c, auditFilterFromQuery(c))
|
||||
total, err := h.db.CountAuditLogs(base)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
@@ -68,9 +69,9 @@ func (h *AuditHandler) Summary(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"total": total,
|
||||
"failures": failures,
|
||||
"recent_7d": recent7d,
|
||||
"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") != "",
|
||||
})
|
||||
@@ -82,7 +83,7 @@ func (h *AuditHandler) ListLogs(c *gin.Context) {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "database unavailable"})
|
||||
return
|
||||
}
|
||||
filter := auditFilterFromQuery(c)
|
||||
filter := auditFilterForAccess(c, auditFilterFromQuery(c))
|
||||
page, pageSize := auditPaginationFromQuery(c)
|
||||
filter.Limit = pageSize
|
||||
filter.Offset = (page - 1) * pageSize
|
||||
@@ -116,6 +117,10 @@ func (h *AuditHandler) GetLog(c *gin.Context) {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "审计记录不存在"})
|
||||
return
|
||||
}
|
||||
if session, ok := security.CurrentSession(c); !ok || (session.Scope != database.RBACScopeAll && row.Actor != session.Username) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
audit.ApplyResourceAvailability(h.db, row)
|
||||
c.JSON(http.StatusOK, gin.H{"log": row})
|
||||
}
|
||||
@@ -126,7 +131,7 @@ func (h *AuditHandler) ExportLogs(c *gin.Context) {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "database unavailable"})
|
||||
return
|
||||
}
|
||||
filter := auditFilterFromQuery(c)
|
||||
filter := auditFilterForAccess(c, auditFilterFromQuery(c))
|
||||
filter.Limit = 5000
|
||||
filter.Offset = 0
|
||||
|
||||
@@ -145,3 +150,10 @@ func (h *AuditHandler) ExportLogs(c *gin.Context) {
|
||||
"logs": logs,
|
||||
})
|
||||
}
|
||||
|
||||
func auditFilterForAccess(c *gin.Context, filter database.ListAuditLogsFilter) database.ListAuditLogsFilter {
|
||||
if session, ok := security.CurrentSession(c); ok && session.Scope != database.RBACScopeAll {
|
||||
filter.Actor = session.Username
|
||||
}
|
||||
return filter
|
||||
}
|
||||
|
||||
+63
-19
@@ -38,6 +38,7 @@ func NewAuthHandler(manager *security.AuthManager, cfg *config.Config, configPat
|
||||
}
|
||||
|
||||
type loginRequest struct {
|
||||
Username string `json:"username"`
|
||||
Password string `json:"password" binding:"required"`
|
||||
}
|
||||
|
||||
@@ -54,7 +55,7 @@ func (h *AuthHandler) Login(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
token, expiresAt, err := h.manager.Authenticate(req.Password)
|
||||
token, expiresAt, err := h.manager.Authenticate(req.Username, req.Password)
|
||||
if err != nil {
|
||||
if h.audit != nil {
|
||||
h.audit.Record(c, audit.Entry{
|
||||
@@ -63,11 +64,13 @@ func (h *AuthHandler) Login(c *gin.Context) {
|
||||
Action: "login",
|
||||
Result: "failure",
|
||||
Message: "登录失败:密码错误",
|
||||
Actor: strings.TrimSpace(req.Username),
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"error": "密码错误"})
|
||||
return
|
||||
}
|
||||
session, _ := h.manager.ValidateToken(token)
|
||||
|
||||
if h.audit != nil {
|
||||
h.audit.Record(c, audit.Entry{
|
||||
@@ -76,6 +79,7 @@ func (h *AuthHandler) Login(c *gin.Context) {
|
||||
Result: "success",
|
||||
SessionHint: audit.HintFromToken(token),
|
||||
Message: "登录成功",
|
||||
Actor: session.Username,
|
||||
Detail: map[string]interface{}{
|
||||
"expires_at": expiresAt.UTC().Format(time.RFC3339),
|
||||
},
|
||||
@@ -86,6 +90,15 @@ func (h *AuthHandler) Login(c *gin.Context) {
|
||||
"token": token,
|
||||
"expires_at": expiresAt.UTC().Format(time.RFC3339),
|
||||
"session_duration_hr": h.manager.SessionDurationHours(),
|
||||
"user": gin.H{
|
||||
"id": session.UserID,
|
||||
"username": session.Username,
|
||||
"display_name": session.DisplayName,
|
||||
},
|
||||
"roles": session.Roles,
|
||||
"permissions": permissionKeys(session.Permissions),
|
||||
"permission_scopes": session.PermissionScopes,
|
||||
"scope": session.Scope,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -139,7 +152,11 @@ func (h *AuthHandler) ChangePassword(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
if !h.manager.CheckPassword(oldPassword) {
|
||||
session, _ := security.CurrentSession(c)
|
||||
if session.Username == "" {
|
||||
session.Username = "admin"
|
||||
}
|
||||
if !h.manager.CheckUserPassword(session.Username, oldPassword) {
|
||||
if h.audit != nil {
|
||||
h.audit.Record(c, audit.Entry{
|
||||
Level: "warn",
|
||||
@@ -153,27 +170,35 @@ func (h *AuthHandler) ChangePassword(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
if err := config.PersistAuthPassword(h.configPath, newPassword); err != nil {
|
||||
if h.logger != nil {
|
||||
h.logger.Error("保存新密码失败", zap.Error(err))
|
||||
if session.UserID == "" || session.UserID == "admin" {
|
||||
if err := config.PersistAuthPassword(h.configPath, newPassword); err != nil {
|
||||
if h.logger != nil {
|
||||
h.logger.Error("保存新密码失败", zap.Error(err))
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "保存新密码失败,请重试"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "保存新密码失败,请重试"})
|
||||
|
||||
if err := h.manager.UpdateConfig(newPassword, h.config.Auth.SessionDurationHours); err != nil {
|
||||
if h.logger != nil {
|
||||
h.logger.Error("更新认证配置失败", zap.Error(err))
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "更新认证配置失败"})
|
||||
return
|
||||
}
|
||||
|
||||
h.config.Auth.Password = newPassword
|
||||
h.config.Auth.GeneratedPassword = ""
|
||||
h.config.Auth.GeneratedPasswordPersisted = false
|
||||
h.config.Auth.GeneratedPasswordPersistErr = ""
|
||||
} else if err := h.manager.UpdateUserPassword(session.UserID, newPassword); err != nil {
|
||||
if h.logger != nil {
|
||||
h.logger.Error("更新用户密码失败", zap.Error(err))
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "更新用户密码失败"})
|
||||
return
|
||||
}
|
||||
|
||||
if err := h.manager.UpdateConfig(newPassword, h.config.Auth.SessionDurationHours); err != nil {
|
||||
if h.logger != nil {
|
||||
h.logger.Error("更新认证配置失败", zap.Error(err))
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "更新认证配置失败"})
|
||||
return
|
||||
}
|
||||
|
||||
h.config.Auth.Password = newPassword
|
||||
h.config.Auth.GeneratedPassword = ""
|
||||
h.config.Auth.GeneratedPasswordPersisted = false
|
||||
h.config.Auth.GeneratedPasswordPersistErr = ""
|
||||
|
||||
if h.logger != nil {
|
||||
h.logger.Info("登录密码已更新,所有会话已失效")
|
||||
}
|
||||
@@ -207,5 +232,24 @@ func (h *AuthHandler) Validate(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"token": session.Token,
|
||||
"expires_at": session.ExpiresAt.UTC().Format(time.RFC3339),
|
||||
"user": gin.H{
|
||||
"id": session.UserID,
|
||||
"username": session.Username,
|
||||
"display_name": session.DisplayName,
|
||||
},
|
||||
"roles": session.Roles,
|
||||
"permissions": permissionKeys(session.Permissions),
|
||||
"permission_scopes": session.PermissionScopes,
|
||||
"scope": session.Scope,
|
||||
})
|
||||
}
|
||||
|
||||
func permissionKeys(perms map[string]bool) []string {
|
||||
keys := make([]string, 0, len(perms))
|
||||
for key, ok := range perms {
|
||||
if ok {
|
||||
keys = append(keys, key)
|
||||
}
|
||||
}
|
||||
return keys
|
||||
}
|
||||
|
||||
@@ -11,6 +11,7 @@ import (
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/multiagent"
|
||||
@@ -109,6 +110,13 @@ func (h *AgentHandler) tryFinalizeBatchQueue(queueID string) {
|
||||
|
||||
// executeOneBatchSubTask 执行单条批量子任务(各自独立会话)。
|
||||
func (h *AgentHandler) executeOneBatchSubTask(queueID string, queue *BatchTaskQueue, task *BatchTask) {
|
||||
ownerUserID := h.db.GetResourceOwner("batch_task", queueID)
|
||||
access, accessErr := h.db.ResolveRBACAccess(ownerUserID)
|
||||
if accessErr != nil || access == nil || !access.User.Enabled {
|
||||
h.batchTaskManager.UpdateTaskStatus(queueID, task.ID, BatchTaskStatusFailed, "", "队列所有者不存在或已禁用")
|
||||
return
|
||||
}
|
||||
principal := authctx.NewPrincipalWithScopes(access.User.ID, access.User.Username, access.Scope, access.Permissions, access.PermissionScopes)
|
||||
title := safeTruncateString(task.Message, 50)
|
||||
batchMeta := audit.ConversationCreateMeta("batch_task")
|
||||
batchMeta.ProjectID = effectiveProjectID(h.config, queue.ProjectID)
|
||||
@@ -119,6 +127,8 @@ func (h *AgentHandler) executeOneBatchSubTask(queueID string, queue *BatchTaskQu
|
||||
return
|
||||
}
|
||||
conversationID := conv.ID
|
||||
_ = h.db.SetResourceOwner("conversation", conversationID, access.User.ID)
|
||||
_ = h.db.AssignResourceToUser(access.User.ID, "conversation", conversationID)
|
||||
|
||||
h.batchTaskManager.UpdateTaskStatusWithConversationID(queueID, task.ID, BatchTaskStatusRunning, "", "", conversationID)
|
||||
|
||||
@@ -156,7 +166,8 @@ func (h *AgentHandler) executeOneBatchSubTask(queueID string, queue *BatchTaskQu
|
||||
|
||||
h.logger.Info("执行批量任务", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("message", task.Message), zap.String("role", queue.Role), zap.String("conversationId", conversationID))
|
||||
|
||||
baseCtx, cancelWithCause := context.WithCancelCause(context.Background())
|
||||
principalCtx := authctx.WithPrincipal(context.Background(), principal)
|
||||
baseCtx, cancelWithCause := context.WithCancelCause(principalCtx)
|
||||
taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 6*time.Hour)
|
||||
|
||||
registered := false
|
||||
|
||||
@@ -98,8 +98,8 @@ type BatchTaskManager struct {
|
||||
logger *zap.Logger
|
||||
queues map[string]*BatchTaskQueue
|
||||
taskCancels map[string]map[string]context.CancelFunc // queueID -> taskID -> 取消函数
|
||||
singleRunTasks map[string]string // queueID -> taskID,单条执行完成后暂停队列
|
||||
queueExecutors map[string]struct{} // executeBatchQueue 协程活跃标记(与队列 status 解耦)
|
||||
singleRunTasks map[string]string // queueID -> taskID,单条执行完成后暂停队列
|
||||
queueExecutors map[string]struct{} // executeBatchQueue 协程活跃标记(与队列 status 解耦)
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
@@ -426,20 +426,24 @@ func (m *BatchTaskManager) GetAllQueues() []*BatchTaskQueue {
|
||||
|
||||
// ListQueues 列出队列(支持筛选和分页)
|
||||
func (m *BatchTaskManager) ListQueues(limit, offset int, status, keyword string) ([]*BatchTaskQueue, int, error) {
|
||||
return m.ListQueuesForAccess(limit, offset, status, keyword, "", "")
|
||||
}
|
||||
|
||||
func (m *BatchTaskManager) ListQueuesForAccess(limit, offset int, status, keyword, userID, scope string) ([]*BatchTaskQueue, int, error) {
|
||||
var queues []*BatchTaskQueue
|
||||
var total int
|
||||
|
||||
// 如果数据库可用,从数据库查询
|
||||
if m.db != nil {
|
||||
// 获取总数
|
||||
count, err := m.db.CountBatchQueues(status, keyword)
|
||||
count, err := m.db.CountBatchQueuesForAccess(status, keyword, userID, scope)
|
||||
if err != nil {
|
||||
return nil, 0, fmt.Errorf("统计队列总数失败: %w", err)
|
||||
}
|
||||
total = count
|
||||
|
||||
// 获取队列列表(只获取ID)
|
||||
queueRows, err := m.db.ListBatchQueues(limit, offset, status, keyword)
|
||||
queueRows, err := m.db.ListBatchQueuesForAccess(limit, offset, status, keyword, userID, scope)
|
||||
if err != nil {
|
||||
return nil, 0, fmt.Errorf("查询队列列表失败: %w", err)
|
||||
}
|
||||
|
||||
@@ -9,7 +9,9 @@ import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
|
||||
@@ -74,7 +76,14 @@ func RegisterBatchTaskMCPTools(mcpServer *mcp.Server, h *AgentHandler, logger *z
|
||||
if offset > 100000 {
|
||||
offset = 100000
|
||||
}
|
||||
queues, total, err := h.batchTaskManager.ListQueues(pageSize, offset, status, keyword)
|
||||
queues := []*BatchTaskQueue{}
|
||||
total := 0
|
||||
var err error
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||
queues, total, err = h.batchTaskManager.ListQueuesForAccess(pageSize, offset, status, keyword, principal.UserID, principal.ScopeFor("tasks:read"))
|
||||
} else {
|
||||
return batchMCPTextResult("缺少认证身份", true), nil
|
||||
}
|
||||
if err != nil {
|
||||
return batchMCPTextResult(fmt.Sprintf("列出队列失败: %v", err), true), nil
|
||||
}
|
||||
@@ -215,11 +224,20 @@ func RegisterBatchTaskMCPTools(mcpServer *mcp.Server, h *AgentHandler, logger *z
|
||||
executeNow = false
|
||||
}
|
||||
projectID := strings.TrimSpace(mcpArgString(args, "project_id"))
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok && projectID != "" && principal.ScopeFor("tasks:write") != database.RBACScopeAll {
|
||||
if h.db == nil || !h.db.UserCanAccessResource(principal.UserID, principal.ScopeFor("tasks:write"), "project", projectID) {
|
||||
return batchMCPTextResult("无权访问目标项目", true), nil
|
||||
}
|
||||
}
|
||||
concurrency := int(mcpArgFloat(args, "concurrency"))
|
||||
queue, createErr := h.batchTaskManager.CreateBatchQueue(title, role, agentMode, scheduleMode, cronExpr, projectID, nextRunAt, concurrency, tasks)
|
||||
if createErr != nil {
|
||||
return batchMCPTextResult("创建队列失败: "+createErr.Error(), true), nil
|
||||
}
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok && h.db != nil {
|
||||
_ = h.db.SetResourceOwner("batch_task", queue.ID, principal.UserID)
|
||||
_ = h.db.AssignResourceToUser(principal.UserID, "batch_task", queue.ID)
|
||||
}
|
||||
started := false
|
||||
if executeNow {
|
||||
ok, err := h.startBatchQueueExecution(queue.ID, false)
|
||||
|
||||
+104
-14
@@ -17,6 +17,7 @@ import (
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/c2"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
@@ -59,7 +60,7 @@ func (h *C2Handler) SetManager(m *c2.Manager) {
|
||||
|
||||
// ListListeners 获取监听器列表
|
||||
func (h *C2Handler) ListListeners(c *gin.Context) {
|
||||
listeners, err := h.mgr().DB().ListC2Listeners()
|
||||
listeners, err := h.mgr().DB().ListC2ListenersForAccess(c2AccessFromContext(c))
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
@@ -109,6 +110,11 @@ func (h *C2Handler) CreateListener(c *gin.Context) {
|
||||
c.JSON(code, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if session, ok := security.CurrentSession(c); ok {
|
||||
listener.OwnerUserID = session.UserID
|
||||
_ = h.mgr().DB().SetResourceOwner("c2_listener", listener.ID, session.UserID)
|
||||
_ = h.mgr().DB().AssignResourceToUser(session.UserID, "c2_listener", listener.ID)
|
||||
}
|
||||
implantToken := listener.ImplantToken
|
||||
listener.EncryptionKey = ""
|
||||
listener.ImplantToken = ""
|
||||
@@ -282,7 +288,7 @@ func (h *C2Handler) ListSessions(c *gin.Context) {
|
||||
filter.Suspicious = true
|
||||
}
|
||||
|
||||
sessions, err := h.mgr().DB().ListC2Sessions(filter)
|
||||
sessions, err := h.mgr().DB().ListC2SessionsForAccess(filter, c2AccessFromContext(c))
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
@@ -304,10 +310,10 @@ func (h *C2Handler) GetSession(c *gin.Context) {
|
||||
}
|
||||
|
||||
// 获取最近任务
|
||||
tasks, _ := h.mgr().DB().ListC2Tasks(database.ListC2TasksFilter{
|
||||
tasks, _ := h.mgr().DB().ListC2TasksForAccess(database.ListC2TasksFilter{
|
||||
SessionID: id,
|
||||
Limit: 20,
|
||||
})
|
||||
}, c2AccessFromContext(c))
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"session": session,
|
||||
@@ -341,7 +347,7 @@ func (h *C2Handler) DeleteSessions(c *gin.Context) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "ids is required"})
|
||||
return
|
||||
}
|
||||
n, err := h.mgr().DB().DeleteC2SessionsByIDs(req.IDs)
|
||||
n, err := h.mgr().DB().DeleteC2SessionsByIDsForAccess(req.IDs, c2AccessFromContext(c))
|
||||
if err != nil {
|
||||
if errors.Is(err, database.ErrNoValidC2SessionIDs) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
@@ -433,24 +439,25 @@ func (h *C2Handler) ListTasks(c *gin.Context) {
|
||||
}
|
||||
}
|
||||
|
||||
tasks, err := h.mgr().DB().ListC2Tasks(filter)
|
||||
access := c2AccessFromContext(c)
|
||||
tasks, err := h.mgr().DB().ListC2TasksForAccess(filter, access)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// 仪表盘「待审任务」为全局 queued/pending 数量,与列表 session 过滤无关
|
||||
pendingN, _ := h.mgr().DB().CountC2TasksQueuedOrPending("")
|
||||
pendingN, _ := h.mgr().DB().CountC2TasksQueuedOrPendingForAccess("", access)
|
||||
|
||||
if !paginated {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"tasks": tasks,
|
||||
"pending_queued_count": pendingN,
|
||||
"tasks": tasks,
|
||||
"pending_queued_count": pendingN,
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
total, err := h.mgr().DB().CountC2Tasks(filter)
|
||||
total, err := h.mgr().DB().CountC2TasksForAccess(filter, access)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
@@ -477,7 +484,7 @@ func (h *C2Handler) DeleteTasks(c *gin.Context) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "ids is required"})
|
||||
return
|
||||
}
|
||||
n, err := h.mgr().DB().DeleteC2TasksByIDs(req.IDs)
|
||||
n, err := h.mgr().DB().DeleteC2TasksByIDsForAccess(req.IDs, c2AccessFromContext(c))
|
||||
if err != nil {
|
||||
if errors.Is(err, database.ErrNoValidC2TaskIDs) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
@@ -522,6 +529,21 @@ func (h *C2Handler) CreateTask(c *gin.Context) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if strings.TrimSpace(req.SessionID) == "" {
|
||||
req.SessionID = strings.TrimSpace(c.Param("id"))
|
||||
}
|
||||
if !h.c2ResourceAllowed(c, "c2_session", req.SessionID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
if conversationID := strings.TrimSpace(req.ConversationID); conversationID != "" {
|
||||
session, ok := security.CurrentSession(c)
|
||||
if !ok || !h.mgr().DB().UserCanAccessResource(session.UserID, session.Scope, "conversation", conversationID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权关联目标对话"})
|
||||
return
|
||||
}
|
||||
req.ConversationID = conversationID
|
||||
}
|
||||
|
||||
input := c2.EnqueueTaskInput{
|
||||
SessionID: req.SessionID,
|
||||
@@ -621,6 +643,10 @@ func (h *C2Handler) PayloadOneliner(c *gin.Context) {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "listener not found"})
|
||||
return
|
||||
}
|
||||
if !h.c2ResourceAllowed(c, "c2_listener", req.ListenerID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
|
||||
host := c2.ResolveBeaconDialHost(listener, strings.TrimSpace(req.Host), h.logger, listener.ID)
|
||||
|
||||
@@ -684,6 +710,10 @@ func (h *C2Handler) PayloadBuild(c *gin.Context) {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "listener not found"})
|
||||
return
|
||||
}
|
||||
if !h.c2ResourceAllowed(c, "c2_listener", req.ListenerID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
|
||||
builder := c2.NewPayloadBuilder(h.mgr(), h.logger, "", "")
|
||||
input := c2.PayloadBuilderInput{
|
||||
@@ -700,6 +730,9 @@ func (h *C2Handler) PayloadBuild(c *gin.Context) {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if session, ok := security.CurrentSession(c); ok {
|
||||
_ = h.mgr().DB().RecordC2PayloadArtifact(filepath.Base(result.OutputPath), result.PayloadID, result.ListenerID, session.UserID)
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"payload": result,
|
||||
@@ -718,6 +751,11 @@ func (h *C2Handler) PayloadDownload(c *gin.Context) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid payload id"})
|
||||
return
|
||||
}
|
||||
session, ok := security.CurrentSession(c)
|
||||
if !ok || !h.mgr().DB().UserCanAccessC2Payload(session.UserID, session.Scope, filename) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
|
||||
builder := c2.NewPayloadBuilder(h.mgr(), h.logger, "", "")
|
||||
storageDir := builder.GetPayloadStoragePath()
|
||||
@@ -779,7 +817,8 @@ func (h *C2Handler) ListEvents(c *gin.Context) {
|
||||
}
|
||||
}
|
||||
|
||||
events, err := h.mgr().DB().ListC2Events(filter)
|
||||
access := c2AccessFromContext(c)
|
||||
events, err := h.mgr().DB().ListC2EventsForAccess(filter, access)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
@@ -788,7 +827,7 @@ func (h *C2Handler) ListEvents(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"events": events})
|
||||
return
|
||||
}
|
||||
total, err := h.mgr().DB().CountC2Events(filter)
|
||||
total, err := h.mgr().DB().CountC2EventsForAccess(filter, access)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
@@ -814,7 +853,7 @@ func (h *C2Handler) DeleteEvents(c *gin.Context) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "ids is required"})
|
||||
return
|
||||
}
|
||||
n, err := h.mgr().DB().DeleteC2EventsByIDs(req.IDs)
|
||||
n, err := h.mgr().DB().DeleteC2EventsByIDsForAccess(req.IDs, c2AccessFromContext(c))
|
||||
if err != nil {
|
||||
if errors.Is(err, database.ErrNoValidC2EventIDs) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
@@ -851,6 +890,9 @@ func (h *C2Handler) EventStream(c *gin.Context) {
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
if !h.c2EventAllowed(c, e) {
|
||||
return true
|
||||
}
|
||||
data, _ := json.Marshal(e)
|
||||
fmt.Fprintf(w, "data: %s\n\n", data)
|
||||
return true
|
||||
@@ -964,6 +1006,10 @@ func (h *C2Handler) UploadFileForImplant(c *gin.Context) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "session_id and remote_path required"})
|
||||
return
|
||||
}
|
||||
if !h.c2ResourceAllowed(c, "c2_session", sessionID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
|
||||
file, header, err := c.Request.FormFile("file")
|
||||
if err != nil {
|
||||
@@ -1018,6 +1064,10 @@ func (h *C2Handler) ListFiles(c *gin.Context) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "session_id required"})
|
||||
return
|
||||
}
|
||||
if !h.c2ResourceAllowed(c, "c2_session", sessionID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
files, err := h.mgr().DB().ListC2FilesBySession(sessionID)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
@@ -1038,6 +1088,10 @@ func (h *C2Handler) DownloadResultFile(c *gin.Context) {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "task not found"})
|
||||
return
|
||||
}
|
||||
if !h.c2ResourceAllowed(c, "c2_task", taskID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
if task.ResultBlobPath == "" {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "no result file for this task"})
|
||||
return
|
||||
@@ -1053,6 +1107,42 @@ func osCreate(path string) (*os.File, error) {
|
||||
return os.Create(path)
|
||||
}
|
||||
|
||||
func c2AccessFromContext(c *gin.Context) database.RBACListAccess {
|
||||
session, ok := security.CurrentSession(c)
|
||||
if !ok {
|
||||
return database.RBACListAccess{}
|
||||
}
|
||||
return database.RBACListAccess{UserID: session.UserID, Scope: session.Scope}
|
||||
}
|
||||
|
||||
func (h *C2Handler) c2ResourceAllowed(c *gin.Context, resourceType, resourceID string) bool {
|
||||
session, ok := security.CurrentSession(c)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
return h.mgr().DB().UserCanAccessResource(session.UserID, session.Scope, resourceType, resourceID)
|
||||
}
|
||||
|
||||
func (h *C2Handler) c2EventAllowed(c *gin.Context, e *c2.Event) bool {
|
||||
if e == nil {
|
||||
return false
|
||||
}
|
||||
session, ok := security.CurrentSession(c)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
if session.Scope == database.RBACScopeAll {
|
||||
return true
|
||||
}
|
||||
if strings.TrimSpace(e.SessionID) != "" {
|
||||
return h.mgr().DB().UserCanAccessResource(session.UserID, session.Scope, "c2_session", e.SessionID)
|
||||
}
|
||||
if strings.TrimSpace(e.TaskID) != "" {
|
||||
return h.mgr().DB().UserCanAccessResource(session.UserID, session.Scope, "c2_task", e.TaskID)
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// 辅助函数(firstNonEmpty 已在 vulnerability.go 中定义)
|
||||
// ============================================================================
|
||||
|
||||
@@ -13,6 +13,8 @@ import (
|
||||
"unicode/utf8"
|
||||
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
@@ -27,6 +29,7 @@ const (
|
||||
type ChatUploadsHandler struct {
|
||||
logger *zap.Logger
|
||||
audit *audit.Service
|
||||
db *database.DB
|
||||
}
|
||||
|
||||
// SetAudit wires platform audit logging.
|
||||
@@ -35,8 +38,32 @@ func (h *ChatUploadsHandler) SetAudit(s *audit.Service) {
|
||||
}
|
||||
|
||||
// NewChatUploadsHandler 创建处理器
|
||||
func NewChatUploadsHandler(logger *zap.Logger) *ChatUploadsHandler {
|
||||
return &ChatUploadsHandler{logger: logger}
|
||||
func NewChatUploadsHandler(logger *zap.Logger, databases ...*database.DB) *ChatUploadsHandler {
|
||||
h := &ChatUploadsHandler{logger: logger}
|
||||
if len(databases) > 0 {
|
||||
h.db = databases[0]
|
||||
}
|
||||
return h
|
||||
}
|
||||
|
||||
func (h *ChatUploadsHandler) pathAllowed(c *gin.Context, relativePath string) bool {
|
||||
session, ok := security.CurrentSession(c)
|
||||
if !ok || h.db == nil {
|
||||
return false
|
||||
}
|
||||
if session.Scope == database.RBACScopeAll {
|
||||
return true
|
||||
}
|
||||
rel := filepath.ToSlash(filepath.Clean(filepath.FromSlash(strings.TrimSpace(relativePath))))
|
||||
rel = strings.Trim(rel, "/")
|
||||
if conversationID, ownerUserID, found := h.db.GetChatUploadArtifact(rel); found {
|
||||
return strings.TrimSpace(ownerUserID) == session.UserID || h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", conversationID)
|
||||
}
|
||||
parts := strings.Split(strings.Trim(rel, "/"), "/")
|
||||
if len(parts) < 2 || parts[1] == "" || parts[1] == "_manual" {
|
||||
return false
|
||||
}
|
||||
return h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", parts[1])
|
||||
}
|
||||
|
||||
func (h *ChatUploadsHandler) absRoot() (string, error) {
|
||||
@@ -175,6 +202,21 @@ func (h *ChatUploadsHandler) List(c *gin.Context) {
|
||||
}
|
||||
folders = filteredFolders
|
||||
}
|
||||
files = filterSlice(files, func(file ChatUploadFileItem) bool {
|
||||
return h.pathAllowed(c, file.RelativePath)
|
||||
})
|
||||
folders = filterSlice(folders, func(folder string) bool {
|
||||
if h.pathAllowed(c, folder) {
|
||||
return true
|
||||
}
|
||||
prefix := strings.TrimSuffix(folder, "/") + "/"
|
||||
for _, file := range files {
|
||||
if strings.HasPrefix(file.RelativePath, prefix) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
})
|
||||
sort.Strings(folders)
|
||||
sort.Slice(files, func(i, j int) bool {
|
||||
return files[i].ModifiedUnix > files[j].ModifiedUnix
|
||||
@@ -185,6 +227,10 @@ func (h *ChatUploadsHandler) List(c *gin.Context) {
|
||||
// Download GET /api/chat-uploads/download?path=...
|
||||
func (h *ChatUploadsHandler) Download(c *gin.Context) {
|
||||
p := c.Query("path")
|
||||
if !h.pathAllowed(c, p) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
abs, err := h.resolveUnderChatUploads(p)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
@@ -209,6 +255,10 @@ func (h *ChatUploadsHandler) Delete(c *gin.Context) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid body"})
|
||||
return
|
||||
}
|
||||
if !h.pathAllowed(c, body.Path) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
abs, err := h.resolveUnderChatUploads(body.Path)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
@@ -238,6 +288,7 @@ func (h *ChatUploadsHandler) Delete(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
}
|
||||
_ = h.db.DeleteChatUploadArtifactPath(filepath.ToSlash(filepath.Clean(filepath.FromSlash(body.Path))))
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "file", "delete", "删除对话附件", "chat_upload", body.Path, nil)
|
||||
}
|
||||
@@ -272,6 +323,10 @@ func (h *ChatUploadsHandler) Mkdir(c *gin.Context) {
|
||||
if parent == "." {
|
||||
parent = ""
|
||||
}
|
||||
if !h.pathAllowed(c, parent) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
|
||||
root, err := h.absRoot()
|
||||
if err != nil {
|
||||
@@ -327,6 +382,10 @@ func (h *ChatUploadsHandler) Rename(c *gin.Context) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid body"})
|
||||
return
|
||||
}
|
||||
if !h.pathAllowed(c, body.Path) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
newName := strings.TrimSpace(body.NewName)
|
||||
if newName == "" || strings.ContainsAny(newName, `/\`) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid newName"})
|
||||
@@ -354,6 +413,8 @@ func (h *ChatUploadsHandler) Rename(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
newRel, _ := filepath.Rel(root, newAbs)
|
||||
oldRel := filepath.ToSlash(filepath.Clean(filepath.FromSlash(body.Path)))
|
||||
_ = h.db.RenameChatUploadArtifactPath(oldRel, filepath.ToSlash(newRel))
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "relativePath": filepath.ToSlash(newRel)})
|
||||
}
|
||||
|
||||
@@ -365,6 +426,10 @@ type chatUploadContentBody struct {
|
||||
// GetContent GET /api/chat-uploads/content?path=...
|
||||
func (h *ChatUploadsHandler) GetContent(c *gin.Context) {
|
||||
p := c.Query("path")
|
||||
if !h.pathAllowed(c, p) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
abs, err := h.resolveUnderChatUploads(p)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
@@ -398,6 +463,10 @@ func (h *ChatUploadsHandler) PutContent(c *gin.Context) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid body"})
|
||||
return
|
||||
}
|
||||
if !h.pathAllowed(c, body.Path) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
if !utf8.ValidString(body.Content) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "content must be valid UTF-8"})
|
||||
return
|
||||
@@ -444,6 +513,10 @@ func (h *ChatUploadsHandler) Upload(c *gin.Context) {
|
||||
var targetDir string
|
||||
targetRel := strings.TrimSpace(c.PostForm("relativeDir"))
|
||||
if targetRel != "" {
|
||||
if !h.pathAllowed(c, targetRel) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
absDir, err := h.resolveUnderChatUploads(targetRel)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
@@ -467,13 +540,17 @@ func (h *ChatUploadsHandler) Upload(c *gin.Context) {
|
||||
targetDir = absDir
|
||||
} else {
|
||||
convID := strings.TrimSpace(c.PostForm("conversationId"))
|
||||
dateStr := time.Now().Format("2006-01-02")
|
||||
if !h.pathAllowed(c, filepath.ToSlash(filepath.Join(dateStr, convID))) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
convDir := convID
|
||||
if convDir == "" {
|
||||
convDir = "_manual"
|
||||
} else {
|
||||
convDir = strings.ReplaceAll(convDir, string(filepath.Separator), "_")
|
||||
}
|
||||
dateStr := time.Now().Format("2006-01-02")
|
||||
targetDir = filepath.Join(root, dateStr, convDir)
|
||||
if err := os.MkdirAll(targetDir, 0755); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
@@ -514,6 +591,16 @@ func (h *ChatUploadsHandler) Upload(c *gin.Context) {
|
||||
}
|
||||
rel, _ := filepath.Rel(root, fullPath)
|
||||
absSaved, _ := filepath.Abs(fullPath)
|
||||
if session, ok := security.CurrentSession(c); ok {
|
||||
conversationID := strings.TrimSpace(c.PostForm("conversationId"))
|
||||
if conversationID == "" {
|
||||
parts := strings.Split(filepath.ToSlash(rel), "/")
|
||||
if len(parts) >= 2 {
|
||||
conversationID = parts[1]
|
||||
}
|
||||
}
|
||||
_ = h.db.UpsertChatUploadArtifact(filepath.ToSlash(rel), conversationID, session.UserID)
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "file", "upload", "上传对话附件", "chat_upload", filepath.ToSlash(rel), map[string]interface{}{
|
||||
"name": unique,
|
||||
|
||||
@@ -16,6 +16,7 @@ import (
|
||||
"cyberstrike-ai/internal/agents"
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/knowledge"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
@@ -89,11 +90,32 @@ type ConfigHandler struct {
|
||||
appUpdater AppUpdater // App更新器(可选)
|
||||
robotRestarter RobotRestarter // 机器人连接重启器(可选),ApplyConfig 时重启钉钉/飞书
|
||||
audit *audit.Service
|
||||
db *database.DB
|
||||
logger *zap.Logger
|
||||
mu sync.RWMutex
|
||||
lastEmbeddingConfig *config.EmbeddingConfig // 上一次的嵌入模型配置(用于检测变更)
|
||||
}
|
||||
|
||||
func (h *ConfigHandler) SetDB(db *database.DB) {
|
||||
h.db = db
|
||||
}
|
||||
|
||||
func (h *ConfigHandler) validateRobotServiceAccounts(robots config.RobotsConfig) error {
|
||||
if h.db == nil {
|
||||
return fmt.Errorf("RBAC 服务不可用,无法校验机器人服务账号")
|
||||
}
|
||||
for platform, userID := range robots.ServiceAccountUserIDs() {
|
||||
user, err := h.db.GetRBACUserByID(userID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("robots.%s.auth.service_user_id 对应用户不存在", platform)
|
||||
}
|
||||
if !user.Enabled {
|
||||
return fmt.Errorf("robots.%s.auth.service_user_id 对应用户已禁用", platform)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// AttackChainUpdater 攻击链处理器更新接口
|
||||
type AttackChainUpdater interface {
|
||||
UpdateConfig(cfg *config.OpenAIConfig)
|
||||
@@ -319,13 +341,18 @@ func (h *ConfigHandler) GetConfig(c *gin.Context) {
|
||||
subAgentCount = len(agents.MergeYAMLAndMarkdown(h.config.MultiAgent.SubAgents, load.SubAgents))
|
||||
}
|
||||
multiPub := config.MultiAgentPublic{
|
||||
Enabled: h.config.MultiAgent.Enabled,
|
||||
RobotDefaultAgentMode: config.NormalizeRobotAgentMode(h.config.MultiAgent),
|
||||
BatchUseMultiAgent: h.config.MultiAgent.BatchUseMultiAgent,
|
||||
SubAgentCount: subAgentCount,
|
||||
Orchestration: config.NormalizeMultiAgentOrchestration(h.config.MultiAgent.Orchestration),
|
||||
PlanExecuteLoopMaxIterations: h.config.MultiAgent.PlanExecuteLoopMaxIterations,
|
||||
ToolSearchAlwaysVisibleTools: append([]string(nil), h.config.MultiAgent.EinoMiddleware.ToolSearchAlwaysVisibleTools...),
|
||||
Enabled: h.config.MultiAgent.Enabled,
|
||||
RobotDefaultAgentMode: config.NormalizeRobotAgentMode(h.config.MultiAgent),
|
||||
BatchUseMultiAgent: h.config.MultiAgent.BatchUseMultiAgent,
|
||||
SubAgentCount: subAgentCount,
|
||||
Orchestration: config.NormalizeMultiAgentOrchestration(h.config.MultiAgent.Orchestration),
|
||||
PlanExecuteLoopMaxIterations: h.config.MultiAgent.PlanExecuteLoopMaxIterations,
|
||||
SummarizationUserIntentLedgerMaxRunes: h.config.MultiAgent.EinoMiddleware.SummarizationUserIntentLedgerMaxRunesEffective(),
|
||||
SummarizationUserIntentLedgerEntryMaxRunes: h.config.MultiAgent.EinoMiddleware.SummarizationUserIntentLedgerEntryMaxRunesEffective(),
|
||||
LatestUserMessageMaxRunes: h.config.MultiAgent.EinoMiddleware.LatestUserMessageMaxRunesEffective(),
|
||||
LatestUserMessageHeadRunes: h.config.MultiAgent.EinoMiddleware.LatestUserMessageHeadRunesEffective(),
|
||||
LatestUserMessageTailRunes: h.config.MultiAgent.EinoMiddleware.LatestUserMessageTailRunesEffective(),
|
||||
ToolSearchAlwaysVisibleTools: append([]string(nil), h.config.MultiAgent.EinoMiddleware.ToolSearchAlwaysVisibleTools...),
|
||||
ToolSearchAlwaysVisibleEffectiveTools: mergeToolNameLists(
|
||||
h.config.MultiAgent.EinoMiddleware.ToolSearchAlwaysVisibleTools,
|
||||
builtin.GetAllBuiltinTools(),
|
||||
@@ -822,6 +849,14 @@ func (h *ConfigHandler) UpdateConfig(c *gin.Context) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if err := config.ValidateRobotsAuthorization(*req.Robots); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if err := h.validateRobotServiceAccounts(*req.Robots); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
h.config.Robots = *req.Robots
|
||||
h.logger.Info("更新机器人配置",
|
||||
zap.Bool("wechat_enabled", h.config.Robots.Wechat.Enabled),
|
||||
@@ -853,6 +888,41 @@ func (h *ConfigHandler) UpdateConfig(c *gin.Context) {
|
||||
if req.MultiAgent.PlanExecuteLoopMaxIterations != nil {
|
||||
h.config.MultiAgent.PlanExecuteLoopMaxIterations = *req.MultiAgent.PlanExecuteLoopMaxIterations
|
||||
}
|
||||
if req.MultiAgent.SummarizationUserIntentLedgerMaxRunes != nil {
|
||||
v := *req.MultiAgent.SummarizationUserIntentLedgerMaxRunes
|
||||
if v < 0 {
|
||||
v = 0
|
||||
}
|
||||
h.config.MultiAgent.EinoMiddleware.SummarizationUserIntentLedgerMaxRunes = v
|
||||
}
|
||||
if req.MultiAgent.SummarizationUserIntentLedgerEntryMaxRunes != nil {
|
||||
v := *req.MultiAgent.SummarizationUserIntentLedgerEntryMaxRunes
|
||||
if v < 0 {
|
||||
v = 0
|
||||
}
|
||||
h.config.MultiAgent.EinoMiddleware.SummarizationUserIntentLedgerEntryMaxRunes = v
|
||||
}
|
||||
if req.MultiAgent.LatestUserMessageMaxRunes != nil {
|
||||
v := *req.MultiAgent.LatestUserMessageMaxRunes
|
||||
if v < 0 {
|
||||
v = 0
|
||||
}
|
||||
h.config.MultiAgent.EinoMiddleware.LatestUserMessageMaxRunes = v
|
||||
}
|
||||
if req.MultiAgent.LatestUserMessageHeadRunes != nil {
|
||||
v := *req.MultiAgent.LatestUserMessageHeadRunes
|
||||
if v < 0 {
|
||||
v = 0
|
||||
}
|
||||
h.config.MultiAgent.EinoMiddleware.LatestUserMessageHeadRunes = v
|
||||
}
|
||||
if req.MultiAgent.LatestUserMessageTailRunes != nil {
|
||||
v := *req.MultiAgent.LatestUserMessageTailRunes
|
||||
if v < 0 {
|
||||
v = 0
|
||||
}
|
||||
h.config.MultiAgent.EinoMiddleware.LatestUserMessageTailRunes = v
|
||||
}
|
||||
if req.MultiAgent.ToolSearchAlwaysVisibleTools != nil {
|
||||
h.config.MultiAgent.EinoMiddleware.ToolSearchAlwaysVisibleTools = dedupeToolNameList(*req.MultiAgent.ToolSearchAlwaysVisibleTools)
|
||||
}
|
||||
@@ -861,6 +931,11 @@ func (h *ConfigHandler) UpdateConfig(c *gin.Context) {
|
||||
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("summarization_user_intent_ledger_max_runes", h.config.MultiAgent.EinoMiddleware.SummarizationUserIntentLedgerMaxRunesEffective()),
|
||||
zap.Int("summarization_user_intent_ledger_entry_max_runes", h.config.MultiAgent.EinoMiddleware.SummarizationUserIntentLedgerEntryMaxRunesEffective()),
|
||||
zap.Int("latest_user_message_max_runes", h.config.MultiAgent.EinoMiddleware.LatestUserMessageMaxRunesEffective()),
|
||||
zap.Int("latest_user_message_head_runes", h.config.MultiAgent.EinoMiddleware.LatestUserMessageHeadRunesEffective()),
|
||||
zap.Int("latest_user_message_tail_runes", h.config.MultiAgent.EinoMiddleware.LatestUserMessageTailRunesEffective()),
|
||||
zap.Int("tool_search_always_visible_tools", len(h.config.MultiAgent.EinoMiddleware.ToolSearchAlwaysVisibleTools)),
|
||||
)
|
||||
}
|
||||
@@ -1932,6 +2007,11 @@ func updateMultiAgentConfig(doc *yaml.Node, cfg config.MultiAgentConfig) {
|
||||
setBoolInMap(maNode, "batch_use_multi_agent", cfg.BatchUseMultiAgent)
|
||||
setIntInMap(maNode, "plan_execute_loop_max_iterations", cfg.PlanExecuteLoopMaxIterations)
|
||||
mwNode := ensureMap(maNode, "eino_middleware")
|
||||
setIntInMap(mwNode, "summarization_user_intent_ledger_max_runes", cfg.EinoMiddleware.SummarizationUserIntentLedgerMaxRunesEffective())
|
||||
setIntInMap(mwNode, "summarization_user_intent_ledger_entry_max_runes", cfg.EinoMiddleware.SummarizationUserIntentLedgerEntryMaxRunesEffective())
|
||||
setIntInMap(mwNode, "latest_user_message_max_runes", cfg.EinoMiddleware.LatestUserMessageMaxRunesEffective())
|
||||
setIntInMap(mwNode, "latest_user_message_head_runes", cfg.EinoMiddleware.LatestUserMessageHeadRunesEffective())
|
||||
setIntInMap(mwNode, "latest_user_message_tail_runes", cfg.EinoMiddleware.LatestUserMessageTailRunesEffective())
|
||||
setFlowStringSliceInMap(mwNode, "tool_search_always_visible_tools", dedupeToolNameList(cfg.EinoMiddleware.ToolSearchAlwaysVisibleTools))
|
||||
}
|
||||
|
||||
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/security"
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
@@ -69,12 +70,23 @@ func (h *ConversationHandler) CreateConversation(c *gin.Context) {
|
||||
|
||||
meta := audit.ConversationCreateMetaFromGin(c, "api")
|
||||
meta.ProjectID = strings.TrimSpace(req.ProjectID)
|
||||
if !h.conversationProjectAllowed(c, meta.ProjectID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问目标项目"})
|
||||
return
|
||||
}
|
||||
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()})
|
||||
return
|
||||
}
|
||||
if session, ok := security.CurrentSession(c); ok {
|
||||
_ = h.db.SetResourceOwner("conversation", conv.ID, session.UserID)
|
||||
_ = h.db.AssignResourceToUser(session.UserID, "conversation", conv.ID)
|
||||
if conv.ProjectID != "" {
|
||||
_ = h.db.AssignResourceToUser(session.UserID, "project", conv.ProjectID)
|
||||
}
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, conv)
|
||||
}
|
||||
@@ -91,11 +103,28 @@ func (h *ConversationHandler) SetConversationProject(c *gin.Context) {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "对话不存在"})
|
||||
return
|
||||
}
|
||||
if err := h.db.SetConversationProjectID(id, req.ProjectID); err != nil {
|
||||
projectID := strings.TrimSpace(req.ProjectID)
|
||||
if !h.conversationProjectAllowed(c, projectID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问目标项目"})
|
||||
return
|
||||
}
|
||||
if err := h.db.SetConversationProjectID(id, 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)})
|
||||
c.JSON(http.StatusOK, gin.H{"success": true, "projectId": projectID})
|
||||
}
|
||||
|
||||
func (h *ConversationHandler) conversationProjectAllowed(c *gin.Context, projectID string) bool {
|
||||
projectID = strings.TrimSpace(projectID)
|
||||
if projectID == "" {
|
||||
return true
|
||||
}
|
||||
session, ok := security.CurrentSession(c)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
return h.db.UserCanAccessResource(session.UserID, session.Scope, "project", projectID)
|
||||
}
|
||||
|
||||
// ListConversations 列出对话
|
||||
@@ -118,19 +147,20 @@ func (h *ConversationHandler) ListConversations(c *gin.Context) {
|
||||
excludeGrouped := strings.TrimSpace(search) == "" && projectID == "" &&
|
||||
(c.Query("exclude_grouped") == "true" || c.Query("exclude_grouped") == "1")
|
||||
sortBy := strings.TrimSpace(c.Query("sort_by"))
|
||||
session, _ := security.CurrentSession(c)
|
||||
|
||||
var conversations []*database.Conversation
|
||||
var total int
|
||||
var err error
|
||||
if excludeGrouped {
|
||||
conversations, err = h.db.ListUngroupedConversations(limit, offset, sortBy, projectID)
|
||||
conversations, err = h.db.ListUngroupedConversationsForAccess(limit, offset, sortBy, projectID, session.UserID, session.Scope)
|
||||
if err == nil {
|
||||
total, err = h.db.CountUngroupedConversations(projectID)
|
||||
total, err = h.db.CountUngroupedConversationsForAccess(projectID, session.UserID, session.Scope)
|
||||
}
|
||||
} else {
|
||||
conversations, err = h.db.ListConversations(limit, offset, search, sortBy, projectID)
|
||||
conversations, err = h.db.ListConversationsForAccess(limit, offset, search, sortBy, projectID, session.UserID, session.Scope)
|
||||
if err == nil {
|
||||
total, err = h.db.CountConversations(search, projectID)
|
||||
total, err = h.db.CountConversationsForAccess(search, projectID, session.UserID, session.Scope)
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
@@ -176,10 +206,17 @@ func (h *ConversationHandler) GetConversation(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, conv)
|
||||
}
|
||||
|
||||
const (
|
||||
defaultProcessDetailsPageLimit = 50
|
||||
maxProcessDetailsPageLimit = 500
|
||||
)
|
||||
|
||||
// GetMessageProcessDetails 获取指定消息的过程详情(按需加载)
|
||||
// 查询参数:
|
||||
// - summary=1:仅返回摘要(total / iterationCount / maxIteration)
|
||||
// - limit + offset:分页返回 processDetails(未指定 limit 时保持全量兼容)
|
||||
// - limit + offset:分页返回 processDetails(未指定 limit 时默认 50 条)
|
||||
// - anchorId:返回包含该过程详情锚点的一页,适合从工具按钮精准定位
|
||||
// - full=1:显式返回全量 processDetails(用于导出/兼容旧集成,不建议 UI 展开时使用)
|
||||
func (h *ConversationHandler) GetMessageProcessDetails(c *gin.Context) {
|
||||
messageID := c.Param("id")
|
||||
if messageID == "" {
|
||||
@@ -199,52 +236,109 @@ func (h *ConversationHandler) GetMessageProcessDetails(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
limitStr := strings.TrimSpace(c.Query("limit"))
|
||||
if limitStr != "" {
|
||||
limit, err := strconv.Atoi(limitStr)
|
||||
if err != nil || limit <= 0 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid limit"})
|
||||
return
|
||||
}
|
||||
if limit > 500 {
|
||||
limit = 500
|
||||
}
|
||||
offset, _ := strconv.Atoi(strings.TrimSpace(c.Query("offset")))
|
||||
if offset < 0 {
|
||||
offset = 0
|
||||
}
|
||||
|
||||
details, total, err := h.db.GetProcessDetailsPage(messageID, limit, offset)
|
||||
fullStr := strings.TrimSpace(c.Query("full"))
|
||||
if fullStr == "1" || strings.EqualFold(fullStr, "true") || strings.EqualFold(fullStr, "yes") {
|
||||
details, err := h.db.GetProcessDetails(messageID)
|
||||
if err != nil {
|
||||
h.logger.Error("分页获取过程详情失败", zap.Error(err))
|
||||
h.logger.Error("获取过程详情失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
details = database.DedupeConsecutiveProcessDetails(details)
|
||||
out := processDetailsToJSON(h.logger, details)
|
||||
out := processDetailsToJSON(h.logger, details, true)
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"processDetails": out,
|
||||
"total": total,
|
||||
"offset": offset,
|
||||
"limit": limit,
|
||||
"hasMore": offset+len(out) < total,
|
||||
"total": len(out),
|
||||
"offset": 0,
|
||||
"limit": len(out),
|
||||
"hasMore": false,
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
details, err := h.db.GetProcessDetails(messageID)
|
||||
limitStr := strings.TrimSpace(c.Query("limit"))
|
||||
limit := defaultProcessDetailsPageLimit
|
||||
if limitStr != "" {
|
||||
parsedLimit, err := strconv.Atoi(limitStr)
|
||||
if err != nil || parsedLimit <= 0 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid limit"})
|
||||
return
|
||||
}
|
||||
limit = parsedLimit
|
||||
}
|
||||
if limit > maxProcessDetailsPageLimit {
|
||||
limit = maxProcessDetailsPageLimit
|
||||
}
|
||||
offset, _ := strconv.Atoi(strings.TrimSpace(c.Query("offset")))
|
||||
if offset < 0 {
|
||||
offset = 0
|
||||
}
|
||||
anchorID := strings.TrimSpace(c.Query("anchorId"))
|
||||
if anchorID != "" {
|
||||
anchorOffset, err := h.db.GetProcessDetailOffset(messageID, anchorID)
|
||||
if err != nil {
|
||||
h.logger.Warn("获取过程详情锚点位置失败", zap.Error(err), zap.String("messageID", messageID), zap.String("anchorID", anchorID))
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
offset = anchorOffset - limit/3
|
||||
if offset < 0 {
|
||||
offset = 0
|
||||
}
|
||||
}
|
||||
|
||||
details, total, err := h.db.GetProcessDetailsPage(messageID, limit, offset)
|
||||
if err != nil {
|
||||
h.logger.Error("获取过程详情失败", zap.Error(err))
|
||||
h.logger.Error("分页获取过程详情失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
details = database.DedupeConsecutiveProcessDetails(details)
|
||||
out := processDetailsToJSON(h.logger, details)
|
||||
c.JSON(http.StatusOK, gin.H{"processDetails": out, "total": len(out)})
|
||||
out := processDetailsToJSON(h.logger, details, false)
|
||||
// A page may end between tool_call and tool_result. Return the full-history
|
||||
// execution summary so the UI can render terminal status without pretending
|
||||
// that an unloaded result is still running.
|
||||
summary, summaryErr := h.db.GetProcessDetailsSummary(messageID)
|
||||
if summaryErr != nil {
|
||||
h.logger.Warn("获取分页工具执行状态失败", zap.Error(summaryErr), zap.String("messageID", messageID))
|
||||
}
|
||||
var toolExecutions []database.ProcessDetailsToolExecution
|
||||
if summary != nil {
|
||||
toolExecutions = summary.ToolExecutions
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"processDetails": out,
|
||||
"toolExecutions": toolExecutions,
|
||||
"total": total,
|
||||
"offset": offset,
|
||||
"limit": limit,
|
||||
"hasMore": offset+len(out) < total,
|
||||
})
|
||||
}
|
||||
|
||||
func processDetailsToJSON(logger *zap.Logger, details []database.ProcessDetail) []map[string]interface{} {
|
||||
// GetProcessDetail 获取单条完整过程详情。列表接口默认不给工具 payload,用户点开单条工具时再拉这里。
|
||||
func (h *ConversationHandler) GetProcessDetail(c *gin.Context) {
|
||||
id := strings.TrimSpace(c.Param("id"))
|
||||
if id == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "process detail id required"})
|
||||
return
|
||||
}
|
||||
detail, err := h.db.GetProcessDetailByID(id)
|
||||
if err != nil {
|
||||
h.logger.Error("获取过程详情失败", zap.Error(err))
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "过程详情不存在"})
|
||||
return
|
||||
}
|
||||
out := processDetailsToJSON(h.logger, []database.ProcessDetail{*detail}, true)
|
||||
if len(out) == 0 {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "过程详情不存在"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"processDetail": out[0]})
|
||||
}
|
||||
|
||||
func processDetailsToJSON(logger *zap.Logger, details []database.ProcessDetail, includeToolPayload bool) []map[string]interface{} {
|
||||
out := make([]map[string]interface{}, 0, len(details))
|
||||
for _, d := range details {
|
||||
var data interface{}
|
||||
@@ -253,6 +347,9 @@ func processDetailsToJSON(logger *zap.Logger, details []database.ProcessDetail)
|
||||
logger.Warn("解析过程详情数据失败", zap.Error(err))
|
||||
}
|
||||
}
|
||||
if !includeToolPayload {
|
||||
data = summarizeProcessDetailData(d.EventType, data)
|
||||
}
|
||||
out = append(out, map[string]interface{}{
|
||||
"id": d.ID,
|
||||
"messageId": d.MessageID,
|
||||
@@ -266,6 +363,27 @@ func processDetailsToJSON(logger *zap.Logger, details []database.ProcessDetail)
|
||||
return out
|
||||
}
|
||||
|
||||
func summarizeProcessDetailData(eventType string, data interface{}) interface{} {
|
||||
m, ok := data.(map[string]interface{})
|
||||
if !ok || (eventType != "tool_call" && eventType != "tool_result") {
|
||||
return data
|
||||
}
|
||||
allow := map[string]bool{
|
||||
"toolName": true, "toolCallId": true, "index": true, "total": true,
|
||||
"success": true, "isError": true, "executionId": true,
|
||||
"einoAgent": true, "einoRole": true, "einoScope": true, "orchestration": true,
|
||||
"agentFacing": true,
|
||||
}
|
||||
out := make(map[string]interface{}, len(allow)+1)
|
||||
for k, v := range m {
|
||||
if allow[k] {
|
||||
out[k] = v
|
||||
}
|
||||
}
|
||||
out["_payloadDeferred"] = true
|
||||
return out
|
||||
}
|
||||
|
||||
// UpdateConversationRequest 更新对话请求
|
||||
type UpdateConversationRequest struct {
|
||||
Title string `json:"title"`
|
||||
|
||||
@@ -0,0 +1,75 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http/httptest"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestProcessDetailsPageIncludesTerminalToolStatusAcrossPageBoundary(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "process-details-page.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatalf("NewDB: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
conversation, err := db.CreateConversation("page boundary", database.ConversationCreateMeta{})
|
||||
if err != nil {
|
||||
t.Fatalf("CreateConversation: %v", err)
|
||||
}
|
||||
message, err := db.AddMessage(conversation.ID, "assistant", "done", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("AddMessage: %v", err)
|
||||
}
|
||||
for i := 1; i <= 4; i++ {
|
||||
id := fmt.Sprintf("call-%d", i)
|
||||
if err := db.AddProcessDetail(message.ID, conversation.ID, "tool_call", "call", map[string]interface{}{
|
||||
"toolName": "http-framework-test", "toolCallId": id, "index": i, "total": 4,
|
||||
}); err != nil {
|
||||
t.Fatalf("AddProcessDetail(tool_call): %v", err)
|
||||
}
|
||||
}
|
||||
for i := 1; i <= 4; i++ {
|
||||
id := fmt.Sprintf("call-%d", i)
|
||||
if err := db.AddProcessDetail(message.ID, conversation.ID, "tool_result", "result", map[string]interface{}{
|
||||
"toolName": "http-framework-test", "toolCallId": id, "success": true,
|
||||
}); err != nil {
|
||||
t.Fatalf("AddProcessDetail(tool_result): %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest("GET", "/api/messages/"+message.ID+"/process-details?limit=6&offset=0", nil)
|
||||
c.Params = gin.Params{{Key: "id", Value: message.ID}}
|
||||
NewConversationHandler(db, zap.NewNop()).GetMessageProcessDetails(c)
|
||||
if w.Code != 200 {
|
||||
t.Fatalf("status = %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
var response struct {
|
||||
HasMore bool `json:"hasMore"`
|
||||
ProcessDetails []map[string]interface{} `json:"processDetails"`
|
||||
ToolExecutions []database.ProcessDetailsToolExecution `json:"toolExecutions"`
|
||||
}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil {
|
||||
t.Fatalf("decode response: %v", err)
|
||||
}
|
||||
if !response.HasMore || len(response.ProcessDetails) != 6 {
|
||||
t.Fatalf("page hasMore=%v details=%d, want true/6", response.HasMore, len(response.ProcessDetails))
|
||||
}
|
||||
if len(response.ToolExecutions) != 4 {
|
||||
t.Fatalf("tool executions = %d, want 4", len(response.ToolExecutions))
|
||||
}
|
||||
for i, execution := range response.ToolExecutions {
|
||||
if execution.Status != "completed" {
|
||||
t.Fatalf("execution %d status = %q, want completed", i, execution.Status)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,118 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestCreateConversationRequiresProjectAccess(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
db, user := setupConversationRBACTest(t)
|
||||
project, err := db.CreateProject(&database.Project{Name: "hidden"})
|
||||
if err != nil {
|
||||
t.Fatalf("CreateProject: %v", err)
|
||||
}
|
||||
handler := NewConversationHandler(db, zap.NewNop())
|
||||
|
||||
w := performConversationRequest(user, http.MethodPost, "/api/conversations", map[string]string{
|
||||
"title": "blocked",
|
||||
"projectId": project.ID,
|
||||
}, handler.CreateConversation)
|
||||
if w.Code != http.StatusForbidden {
|
||||
t.Fatalf("status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String())
|
||||
}
|
||||
|
||||
if err := db.AssignResourceToUser(user.ID, "project", project.ID); err != nil {
|
||||
t.Fatalf("AssignResourceToUser: %v", err)
|
||||
}
|
||||
w = performConversationRequest(user, http.MethodPost, "/api/conversations", map[string]string{
|
||||
"title": "allowed",
|
||||
"projectId": project.ID,
|
||||
}, handler.CreateConversation)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d: %s", w.Code, http.StatusOK, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetConversationProjectRequiresProjectAccess(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
db, user := setupConversationRBACTest(t)
|
||||
project, err := db.CreateProject(&database.Project{Name: "hidden"})
|
||||
if err != nil {
|
||||
t.Fatalf("CreateProject: %v", err)
|
||||
}
|
||||
conv, err := db.CreateConversation("owned", database.ConversationCreateMeta{})
|
||||
if err != nil {
|
||||
t.Fatalf("CreateConversation: %v", err)
|
||||
}
|
||||
if err := db.SetResourceOwner("conversation", conv.ID, user.ID); err != nil {
|
||||
t.Fatalf("SetResourceOwner: %v", err)
|
||||
}
|
||||
if err := db.AssignResourceToUser(user.ID, "conversation", conv.ID); err != nil {
|
||||
t.Fatalf("AssignResourceToUser conversation: %v", err)
|
||||
}
|
||||
handler := NewConversationHandler(db, zap.NewNop())
|
||||
|
||||
w := performConversationRequest(user, http.MethodPut, "/api/conversations/"+conv.ID+"/project", map[string]string{
|
||||
"projectId": project.ID,
|
||||
}, func(c *gin.Context) {
|
||||
c.Params = gin.Params{{Key: "id", Value: conv.ID}}
|
||||
handler.SetConversationProject(c)
|
||||
})
|
||||
if w.Code != http.StatusForbidden {
|
||||
t.Fatalf("status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String())
|
||||
}
|
||||
|
||||
if err := db.AssignResourceToUser(user.ID, "project", project.ID); err != nil {
|
||||
t.Fatalf("AssignResourceToUser project: %v", err)
|
||||
}
|
||||
w = performConversationRequest(user, http.MethodPut, "/api/conversations/"+conv.ID+"/project", map[string]string{
|
||||
"projectId": project.ID,
|
||||
}, func(c *gin.Context) {
|
||||
c.Params = gin.Params{{Key: "id", Value: conv.ID}}
|
||||
handler.SetConversationProject(c)
|
||||
})
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d: %s", w.Code, http.StatusOK, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func setupConversationRBACTest(t *testing.T) (*database.DB, *database.RBACUser) {
|
||||
t.Helper()
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "conversation-rbac.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatalf("NewDB: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
user, err := db.CreateRBACUser("operator1", "Operator One", "hash", true, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("CreateRBACUser: %v", err)
|
||||
}
|
||||
return db, user
|
||||
}
|
||||
|
||||
func performConversationRequest(user *database.RBACUser, method, path string, body map[string]string, handler gin.HandlerFunc) *httptest.ResponseRecorder {
|
||||
payload, _ := json.Marshal(body)
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(method, path, bytes.NewReader(payload))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
c.Set(security.ContextSessionKey, security.Session{
|
||||
UserID: user.ID,
|
||||
Username: user.Username,
|
||||
Permissions: map[string]bool{"chat:write": true},
|
||||
Scope: database.RBACScopeAssigned,
|
||||
})
|
||||
handler(c)
|
||||
return w
|
||||
}
|
||||
@@ -13,11 +13,11 @@ import (
|
||||
)
|
||||
|
||||
// rebindEinoRunningTask 中断并继续 / 空正文续跑:重建 cancel 链与超时 ctx,保持任务 running。
|
||||
func (h *AgentHandler) rebindEinoRunningTask(conversationID string, timeoutCancel context.CancelFunc) (context.Context, context.CancelCauseFunc, context.Context, context.CancelFunc) {
|
||||
func (h *AgentHandler) rebindEinoRunningTask(parent context.Context, conversationID string, timeoutCancel context.CancelFunc) (context.Context, context.CancelCauseFunc, context.Context, context.CancelFunc) {
|
||||
if timeoutCancel != nil {
|
||||
timeoutCancel()
|
||||
}
|
||||
baseCtx, cancelWithCause := context.WithCancelCause(context.Background())
|
||||
baseCtx, cancelWithCause := context.WithCancelCause(detachedAgentContext(parent))
|
||||
h.tasks.BindTaskCancel(conversationID, cancelWithCause)
|
||||
taskCtx, newTimeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute)
|
||||
h.tasks.UpdateTaskStatus(conversationID, "running")
|
||||
|
||||
@@ -116,7 +116,7 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
||||
"userMessageId": prep.UserMessageID,
|
||||
})
|
||||
}
|
||||
if h.runRoleWorkflowStreamIfBound(&req, prep, sendEvent) {
|
||||
if h.runRoleWorkflowStreamIfBound(c, &req, prep, sendEvent) {
|
||||
return
|
||||
}
|
||||
|
||||
@@ -154,7 +154,7 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
||||
var result *multiagent.RunResult
|
||||
var runErr error
|
||||
|
||||
baseCtx, cancelWithCause = context.WithCancelCause(context.Background())
|
||||
baseCtx, cancelWithCause = context.WithCancelCause(detachedAgentContext(c.Request.Context()))
|
||||
taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute)
|
||||
|
||||
if _, err := h.tasks.StartTask(conversationID, req.Message, cancelWithCause); err != nil {
|
||||
@@ -247,7 +247,7 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
||||
if h.tryContinueOnEinoEmptyResponse(taskCtx, mw, conversationID, result, &emptyResponseContinueAttempt, &curHistory, &curFinalMessage, progressCallback) {
|
||||
mainIterationOffset += segmentMainIterationMax
|
||||
timeoutCancel()
|
||||
baseCtx, cancelWithCause, taskCtx, timeoutCancel = h.rebindEinoRunningTask(conversationID, timeoutCancel)
|
||||
baseCtx, cancelWithCause, taskCtx, timeoutCancel = h.rebindEinoRunningTask(taskCtx, conversationID, timeoutCancel)
|
||||
continue
|
||||
}
|
||||
timeoutCancel()
|
||||
@@ -279,7 +279,7 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
||||
})
|
||||
mainIterationOffset += segmentMainIterationMax
|
||||
timeoutCancel()
|
||||
baseCtx, cancelWithCause = context.WithCancelCause(context.Background())
|
||||
baseCtx, cancelWithCause = context.WithCancelCause(detachedAgentContext(baseCtx))
|
||||
h.tasks.BindTaskCancel(conversationID, cancelWithCause)
|
||||
taskCtx, timeoutCancel = context.WithTimeout(baseCtx, 600*time.Minute)
|
||||
h.tasks.UpdateTaskStatus(conversationID, "running")
|
||||
@@ -381,7 +381,8 @@ func (h *AgentHandler) EinoSingleAgentLoop(c *gin.Context) {
|
||||
|
||||
prep, err := h.prepareMultiAgentSession(&req, c, "eino_agent")
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
status, msg := multiAgentHTTPErrorStatus(err)
|
||||
c.JSON(status, gin.H{"error": msg})
|
||||
return
|
||||
}
|
||||
h.activateHITLForConversation(prep.ConversationID, req.Hitl)
|
||||
|
||||
@@ -9,6 +9,7 @@ import (
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
@@ -67,7 +68,7 @@ func (h *ExternalMCPHandler) GetExternalMCPs(c *gin.Context) {
|
||||
errorMsg := externalMCPStatusError(h.manager, name, status)
|
||||
|
||||
result[name] = ExternalMCPResponse{
|
||||
Config: cfg,
|
||||
Config: externalMCPConfigForResponse(c, cfg),
|
||||
Status: status,
|
||||
ToolCount: toolCount,
|
||||
Error: errorMsg,
|
||||
@@ -113,13 +114,33 @@ func (h *ExternalMCPHandler) GetExternalMCP(c *gin.Context) {
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, ExternalMCPResponse{
|
||||
Config: cfg,
|
||||
Config: externalMCPConfigForResponse(c, cfg),
|
||||
Status: status,
|
||||
ToolCount: toolCount,
|
||||
Error: externalMCPStatusError(h.manager, name, status),
|
||||
})
|
||||
}
|
||||
|
||||
func externalMCPConfigForResponse(c *gin.Context, cfg config.ExternalMCPServerConfig) config.ExternalMCPServerConfig {
|
||||
if security.SessionHasPermission(c, "mcp:write") {
|
||||
return cfg
|
||||
}
|
||||
copyCfg := cfg
|
||||
if len(cfg.Env) > 0 {
|
||||
copyCfg.Env = make(map[string]string, len(cfg.Env))
|
||||
for key := range cfg.Env {
|
||||
copyCfg.Env[key] = "***"
|
||||
}
|
||||
}
|
||||
if len(cfg.Headers) > 0 {
|
||||
copyCfg.Headers = make(map[string]string, len(cfg.Headers))
|
||||
for key := range cfg.Headers {
|
||||
copyCfg.Headers[key] = "***"
|
||||
}
|
||||
}
|
||||
return copyCfg
|
||||
}
|
||||
|
||||
// externalMCPStatusError 在 error/disconnected 状态下返回最近错误(含断连原因)。
|
||||
func externalMCPStatusError(manager *mcp.ExternalMCPManager, name, status string) string {
|
||||
if status != "error" && status != "disconnected" {
|
||||
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
@@ -43,7 +44,8 @@ func (h *GroupHandler) CreateGroup(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
group, err := h.db.CreateGroup(req.Name, req.Icon)
|
||||
session, _ := security.CurrentSession(c)
|
||||
group, err := h.db.CreateGroup(req.Name, req.Icon, session.UserID)
|
||||
if err != nil {
|
||||
h.logger.Error("创建分组失败", zap.Error(err))
|
||||
// 如果是名称重复错误,返回400状态码
|
||||
@@ -60,7 +62,8 @@ func (h *GroupHandler) CreateGroup(c *gin.Context) {
|
||||
|
||||
// ListGroups 列出所有分组
|
||||
func (h *GroupHandler) ListGroups(c *gin.Context) {
|
||||
groups, err := h.db.ListGroups()
|
||||
session, _ := security.CurrentSession(c)
|
||||
groups, err := h.db.ListGroupsForAccess(session.UserID, session.Scope)
|
||||
if err != nil {
|
||||
h.logger.Error("获取分组列表失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
@@ -73,6 +76,10 @@ func (h *GroupHandler) ListGroups(c *gin.Context) {
|
||||
// GetGroup 获取分组
|
||||
func (h *GroupHandler) GetGroup(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
if !h.groupAllowed(c, id) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
|
||||
group, err := h.db.GetGroup(id)
|
||||
if err != nil {
|
||||
@@ -93,6 +100,10 @@ type UpdateGroupRequest struct {
|
||||
// UpdateGroup 更新分组
|
||||
func (h *GroupHandler) UpdateGroup(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
if !h.groupAllowed(c, id) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
|
||||
var req UpdateGroupRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
@@ -129,6 +140,10 @@ func (h *GroupHandler) UpdateGroup(c *gin.Context) {
|
||||
// DeleteGroup 删除分组
|
||||
func (h *GroupHandler) DeleteGroup(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
if !h.groupAllowed(c, id) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
|
||||
if err := h.db.DeleteGroup(id); err != nil {
|
||||
h.logger.Error("删除分组失败", zap.Error(err))
|
||||
@@ -152,6 +167,14 @@ func (h *GroupHandler) AddConversationToGroup(c *gin.Context) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if !h.groupConversationAllowed(c, req.ConversationID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
if !h.groupAllowed(c, req.GroupID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该分组"})
|
||||
return
|
||||
}
|
||||
|
||||
if err := h.db.AddConversationToGroup(req.ConversationID, req.GroupID); err != nil {
|
||||
h.logger.Error("添加对话到分组失败", zap.Error(err))
|
||||
@@ -166,6 +189,14 @@ func (h *GroupHandler) AddConversationToGroup(c *gin.Context) {
|
||||
func (h *GroupHandler) RemoveConversationFromGroup(c *gin.Context) {
|
||||
conversationID := c.Param("conversationId")
|
||||
groupID := c.Param("id")
|
||||
if !h.groupAllowed(c, groupID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该分组"})
|
||||
return
|
||||
}
|
||||
if !h.groupConversationAllowed(c, conversationID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
|
||||
if err := h.db.RemoveConversationFromGroup(conversationID, groupID); err != nil {
|
||||
h.logger.Error("从分组中移除对话失败", zap.Error(err))
|
||||
@@ -189,6 +220,10 @@ type GroupConversation struct {
|
||||
// GetGroupConversations 获取分组中的所有对话
|
||||
func (h *GroupHandler) GetGroupConversations(c *gin.Context) {
|
||||
groupID := c.Param("id")
|
||||
if !h.groupAllowed(c, groupID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该分组"})
|
||||
return
|
||||
}
|
||||
searchQuery := c.Query("search") // 获取搜索参数
|
||||
|
||||
var conversations []*database.Conversation
|
||||
@@ -210,6 +245,9 @@ func (h *GroupHandler) GetGroupConversations(c *gin.Context) {
|
||||
// 获取每个对话在分组中的置顶状态
|
||||
groupConvs := make([]GroupConversation, 0, len(conversations))
|
||||
for _, conv := range conversations {
|
||||
if conv == nil || !h.groupConversationAllowed(c, conv.ID) {
|
||||
continue
|
||||
}
|
||||
// 查询分组内置顶状态
|
||||
var groupPinned int
|
||||
err := h.db.QueryRow(
|
||||
@@ -242,8 +280,14 @@ func (h *GroupHandler) GetAllMappings(c *gin.Context) {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
filtered := mappings[:0]
|
||||
for _, mapping := range mappings {
|
||||
if h.groupConversationAllowed(c, mapping.ConversationID) && h.groupAllowed(c, mapping.GroupID) {
|
||||
filtered = append(filtered, mapping)
|
||||
}
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, mappings)
|
||||
c.JSON(http.StatusOK, filtered)
|
||||
}
|
||||
|
||||
// UpdateConversationPinnedRequest 更新对话置顶状态请求
|
||||
@@ -254,6 +298,10 @@ type UpdateConversationPinnedRequest struct {
|
||||
// UpdateConversationPinned 更新对话置顶状态
|
||||
func (h *GroupHandler) UpdateConversationPinned(c *gin.Context) {
|
||||
conversationID := c.Param("id")
|
||||
if !h.groupConversationAllowed(c, conversationID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
|
||||
var req UpdateConversationPinnedRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
@@ -278,6 +326,10 @@ type UpdateGroupPinnedRequest struct {
|
||||
// UpdateGroupPinned 更新分组置顶状态
|
||||
func (h *GroupHandler) UpdateGroupPinned(c *gin.Context) {
|
||||
groupID := c.Param("id")
|
||||
if !h.groupAllowed(c, groupID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该分组"})
|
||||
return
|
||||
}
|
||||
|
||||
var req UpdateGroupPinnedRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
@@ -303,6 +355,14 @@ type UpdateConversationPinnedInGroupRequest struct {
|
||||
func (h *GroupHandler) UpdateConversationPinnedInGroup(c *gin.Context) {
|
||||
groupID := c.Param("id")
|
||||
conversationID := c.Param("conversationId")
|
||||
if !h.groupAllowed(c, groupID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该分组"})
|
||||
return
|
||||
}
|
||||
if !h.groupConversationAllowed(c, conversationID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
|
||||
var req UpdateConversationPinnedInGroupRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
@@ -318,3 +378,16 @@ func (h *GroupHandler) UpdateConversationPinnedInGroup(c *gin.Context) {
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"message": "更新成功"})
|
||||
}
|
||||
|
||||
func (h *GroupHandler) groupConversationAllowed(c *gin.Context, conversationID string) bool {
|
||||
session, ok := security.CurrentSession(c)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
return h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", conversationID)
|
||||
}
|
||||
|
||||
func (h *GroupHandler) groupAllowed(c *gin.Context, groupID string) bool {
|
||||
session, ok := security.CurrentSession(c)
|
||||
return ok && h.db.UserCanAccessGroup(session.UserID, session.Scope, groupID)
|
||||
}
|
||||
|
||||
@@ -611,6 +611,7 @@ func (h *AgentHandler) ListHITLPending(c *gin.Context) {
|
||||
offset := (page - 1) * pageSize
|
||||
q, args := h.buildHitlListQuery(false)
|
||||
q, args = h.appendHitlListFilters(q, args, c)
|
||||
q, args = appendConversationAccessSQL(q, args, "conversation_id", notificationAccessFromContext(c))
|
||||
total, err := h.countHitlQuery(q, args)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
@@ -649,6 +650,10 @@ func (h *AgentHandler) DecideHITLInterrupt(c *gin.Context) {
|
||||
c.JSON(500, gin.H{"error": "hitl manager unavailable"})
|
||||
return
|
||||
}
|
||||
if !h.hitlInterruptAllowed(c, req.InterruptID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
if err := h.hitlManager.ResolveInterrupt(req.InterruptID, req.Decision, req.Comment, req.EditedArguments); err != nil {
|
||||
c.JSON(http.StatusConflict, gin.H{"error": err.Error()})
|
||||
return
|
||||
@@ -673,6 +678,10 @@ func (h *AgentHandler) DismissHITLInterrupt(c *gin.Context) {
|
||||
c.JSON(500, gin.H{"error": "hitl manager unavailable"})
|
||||
return
|
||||
}
|
||||
if !h.hitlInterruptAllowed(c, req.InterruptID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
res, err := h.db.Exec(`UPDATE hitl_interrupts SET status='cancelled', decision='reject',
|
||||
decision_comment='dismissed by user', decided_at=CURRENT_TIMESTAMP, decided_by='human'
|
||||
WHERE id=? AND status='pending'`, req.InterruptID)
|
||||
@@ -728,7 +737,6 @@ func (h *AgentHandler) interceptHITLForEinoTool(runCtx context.Context, cancelRu
|
||||
return arguments, nil
|
||||
}
|
||||
|
||||
|
||||
type hitlConfigReq struct {
|
||||
ConversationID string `json:"conversationId" binding:"required"`
|
||||
HITLRequest
|
||||
@@ -740,6 +748,10 @@ func (h *AgentHandler) GetHITLConversationConfig(c *gin.Context) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "conversationId is required"})
|
||||
return
|
||||
}
|
||||
if !h.hitlConversationAllowed(c, conversationID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
cfg, err := h.loadHITLConversationConfig(conversationID)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
@@ -770,6 +782,10 @@ func (h *AgentHandler) UpsertHITLConversationConfig(c *gin.Context) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if !h.hitlConversationAllowed(c, req.ConversationID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
req.Mode = normalizeHitlMode(req.Mode)
|
||||
req.Reviewer = normalizeHitlReviewer(req.Reviewer)
|
||||
if strings.TrimSpace(req.Reviewer) == "" {
|
||||
@@ -868,8 +884,8 @@ func (h *AgentHandler) SetHITLGlobalToolWhitelist(c *gin.Context) {
|
||||
h.audit.RecordOK(c, "hitl", "tool_whitelist_update", "HITL 全局白名单更新", "hitl_config", "tool_whitelist", nil)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"ok": true,
|
||||
"toolWhitelist": h.hitlConfigGlobalToolWhitelist(),
|
||||
"ok": true,
|
||||
"toolWhitelist": h.hitlConfigGlobalToolWhitelist(),
|
||||
"hitlGlobalToolWhitelist": h.hitlConfigGlobalToolWhitelist(),
|
||||
"hitlGlobalWhitelistMerged": false,
|
||||
})
|
||||
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
@@ -163,6 +164,7 @@ func (h *AgentHandler) ListHITLLogs(c *gin.Context) {
|
||||
|
||||
q, args := h.buildHitlListQuery(true)
|
||||
q, args = h.appendHitlListFilters(q, args, c)
|
||||
q, args = appendConversationAccessSQL(q, args, "conversation_id", notificationAccessFromContext(c))
|
||||
total, err := h.countHitlQuery(q, args)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
@@ -207,6 +209,7 @@ func (h *AgentHandler) DeleteHITLLogs(c *gin.Context) {
|
||||
if request.All {
|
||||
where, args := h.buildHitlLogsWhere(true)
|
||||
where, args = h.appendHitlListFilters(where, args, c)
|
||||
where, args = appendConversationAccessSQL(where, args, "conversation_id", notificationAccessFromContext(c))
|
||||
deleted, err = h.db.DeleteHitlInterruptLogsMatching(where, args)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
@@ -222,7 +225,12 @@ func (h *AgentHandler) DeleteHITLLogs(c *gin.Context) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "审计日志 ID 列表不能为空"})
|
||||
return
|
||||
}
|
||||
deleted, err = h.db.DeleteHitlInterruptLogsByIDs(request.IDs)
|
||||
ids, filterErr := h.filterAllowedHitlInterruptIDs(c, request.IDs)
|
||||
if filterErr != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": filterErr.Error()})
|
||||
return
|
||||
}
|
||||
deleted, err = h.db.DeleteHitlInterruptLogsByIDs(ids)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
@@ -259,5 +267,65 @@ func (h *AgentHandler) GetHITLLog(c *gin.Context) {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if !h.hitlConversationAllowed(c, cid) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, hitlInterruptRowToMap(rowID, cid, mode, toolName, toolCallID, payload, rowStatus, decidedBy, messageID, decision, comment, createdAt, decidedAt))
|
||||
}
|
||||
|
||||
func (h *AgentHandler) filterAllowedHitlInterruptIDs(c *gin.Context, ids []string) ([]string, error) {
|
||||
clean := make([]string, 0, len(ids))
|
||||
seen := map[string]struct{}{}
|
||||
for _, id := range ids {
|
||||
id = strings.TrimSpace(id)
|
||||
if id == "" {
|
||||
continue
|
||||
}
|
||||
if _, ok := seen[id]; ok {
|
||||
continue
|
||||
}
|
||||
seen[id] = struct{}{}
|
||||
clean = append(clean, id)
|
||||
}
|
||||
if len(clean) == 0 {
|
||||
return clean, nil
|
||||
}
|
||||
query := `SELECT id, conversation_id FROM hitl_interrupts WHERE id IN (` + buildPlaceholders(len(clean)) + `)`
|
||||
args := make([]interface{}, 0, len(clean))
|
||||
for _, id := range clean {
|
||||
args = append(args, id)
|
||||
}
|
||||
rows, err := h.db.Query(query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
allowed := make([]string, 0, len(clean))
|
||||
for rows.Next() {
|
||||
var id, conversationID string
|
||||
if err := rows.Scan(&id, &conversationID); err != nil {
|
||||
continue
|
||||
}
|
||||
if h.hitlConversationAllowed(c, conversationID) {
|
||||
allowed = append(allowed, id)
|
||||
}
|
||||
}
|
||||
return allowed, rows.Err()
|
||||
}
|
||||
|
||||
func (h *AgentHandler) hitlInterruptAllowed(c *gin.Context, interruptID string) bool {
|
||||
var conversationID string
|
||||
if err := h.db.QueryRow(`SELECT conversation_id FROM hitl_interrupts WHERE id = ?`, strings.TrimSpace(interruptID)).Scan(&conversationID); err != nil {
|
||||
return false
|
||||
}
|
||||
return h.hitlConversationAllowed(c, conversationID)
|
||||
}
|
||||
|
||||
func (h *AgentHandler) hitlConversationAllowed(c *gin.Context, conversationID string) bool {
|
||||
session, ok := security.CurrentSession(c)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
return h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", conversationID)
|
||||
}
|
||||
|
||||
+122
-14
@@ -120,9 +120,22 @@ func (h *MonitorHandler) Monitor(c *gin.Context) {
|
||||
// 解析工具筛选参数(兼容 mcp__tool 与内部 mcp::tool)
|
||||
toolName := normalizeToolNameFilter(c.Query("tool"))
|
||||
|
||||
executions, total := h.loadExecutionListWithPagination(page, pageSize, status, toolName)
|
||||
access := notificationAccessFromContext(c)
|
||||
executions, total := h.loadExecutionListWithPagination(page, pageSize, status, toolName, access)
|
||||
h.enrichExecutionsConversationID(executions)
|
||||
summary, topTools := h.loadStatsSummary(monitorPageTopTools)
|
||||
var summary *MonitorStatsSummary
|
||||
var topTools []*mcp.ToolStats
|
||||
if access.Scope == database.RBACScopeAll {
|
||||
summary, topTools = h.loadStatsSummary(monitorPageTopTools)
|
||||
} else if h.db != nil {
|
||||
if scoped, err := h.db.LoadToolStatsSummaryForAccess(monitorPageTopTools, access); err == nil {
|
||||
summary, topTools = dbStatsSummaryToMonitor(scoped), scoped.TopTools
|
||||
} else {
|
||||
summary, topTools = summarizeAccessibleExecutionPage(executions, monitorPageTopTools)
|
||||
}
|
||||
} else {
|
||||
summary, topTools = summarizeAccessibleExecutionPage(executions, monitorPageTopTools)
|
||||
}
|
||||
|
||||
totalPages := (total + pageSize - 1) / pageSize
|
||||
if totalPages == 0 {
|
||||
@@ -142,6 +155,31 @@ func (h *MonitorHandler) Monitor(c *gin.Context) {
|
||||
})
|
||||
}
|
||||
|
||||
func summarizeAccessibleExecutionPage(executions []*mcp.ToolExecution, topN int) (*MonitorStatsSummary, []*mcp.ToolStats) {
|
||||
stats := map[string]*mcp.ToolStats{}
|
||||
for _, exec := range executions {
|
||||
if exec == nil {
|
||||
continue
|
||||
}
|
||||
stat := stats[exec.ToolName]
|
||||
if stat == nil {
|
||||
stat = &mcp.ToolStats{ToolName: exec.ToolName}
|
||||
stats[exec.ToolName] = stat
|
||||
}
|
||||
stat.TotalCalls++
|
||||
if exec.Status == "failed" || exec.Status == "cancelled" {
|
||||
stat.FailedCalls++
|
||||
} else if exec.Status == "completed" {
|
||||
stat.SuccessCalls++
|
||||
}
|
||||
started := exec.StartTime
|
||||
if stat.LastCallTime == nil || started.After(*stat.LastCallTime) {
|
||||
stat.LastCallTime = &started
|
||||
}
|
||||
}
|
||||
return summarizeToolStats(stats, topN)
|
||||
}
|
||||
|
||||
func (h *MonitorHandler) monitorRetentionDays() int {
|
||||
if h.monitorRetention != nil {
|
||||
return h.monitorRetention.RetentionDays()
|
||||
@@ -154,9 +192,9 @@ func (h *MonitorHandler) loadExecutions() []*mcp.ToolExecution {
|
||||
return executions
|
||||
}
|
||||
|
||||
func (h *MonitorHandler) loadExecutionListWithPagination(page, pageSize int, status, toolName string) ([]*mcp.ToolExecution, int) {
|
||||
func (h *MonitorHandler) loadExecutionListWithPagination(page, pageSize int, status, toolName string, access database.RBACListAccess) ([]*mcp.ToolExecution, int) {
|
||||
if h.db == nil {
|
||||
allExecutions := h.mcpServer.GetAllExecutions()
|
||||
allExecutions := filterToolExecutionsForAccess(h.mcpServer.GetAllExecutions(), access, h.db)
|
||||
if status != "" || toolName != "" {
|
||||
filtered := make([]*mcp.ToolExecution, 0)
|
||||
for _, exec := range allExecutions {
|
||||
@@ -189,13 +227,13 @@ func (h *MonitorHandler) loadExecutionListWithPagination(page, pageSize int, sta
|
||||
}
|
||||
|
||||
offset := (page - 1) * pageSize
|
||||
executions, err := h.db.LoadToolExecutionListPage(offset, pageSize, status, toolName)
|
||||
executions, err := h.db.LoadToolExecutionListPageForAccess(offset, pageSize, status, toolName, access)
|
||||
if err != nil {
|
||||
h.logger.Warn("从数据库加载执行记录列表失败,回退到内存数据", zap.Error(err))
|
||||
return h.loadExecutionListWithPaginationFromMemory(page, pageSize, status, toolName)
|
||||
return h.loadExecutionListWithPaginationFromMemory(page, pageSize, status, toolName, access)
|
||||
}
|
||||
|
||||
total, err := h.db.CountToolExecutions(status, toolName)
|
||||
total, err := h.db.CountToolExecutionsForAccess(status, toolName, access)
|
||||
if err != nil {
|
||||
h.logger.Warn("获取执行记录总数失败", zap.Error(err))
|
||||
total = offset + len(executions)
|
||||
@@ -207,8 +245,8 @@ func (h *MonitorHandler) loadExecutionListWithPagination(page, pageSize int, sta
|
||||
return executions, total
|
||||
}
|
||||
|
||||
func (h *MonitorHandler) loadExecutionListWithPaginationFromMemory(page, pageSize int, status, toolName string) ([]*mcp.ToolExecution, int) {
|
||||
allExecutions := h.mcpServer.GetAllExecutions()
|
||||
func (h *MonitorHandler) loadExecutionListWithPaginationFromMemory(page, pageSize int, status, toolName string, access database.RBACListAccess) ([]*mcp.ToolExecution, int) {
|
||||
allExecutions := filterToolExecutionsForAccess(h.mcpServer.GetAllExecutions(), access, h.db)
|
||||
if status != "" || toolName != "" {
|
||||
filtered := make([]*mcp.ToolExecution, 0)
|
||||
for _, exec := range allExecutions {
|
||||
@@ -260,6 +298,50 @@ func slimToolExecution(exec *mcp.ToolExecution) *mcp.ToolExecution {
|
||||
return slim
|
||||
}
|
||||
|
||||
func filterToolExecutionsForAccess(executions []*mcp.ToolExecution, access database.RBACListAccess, db *database.DB) []*mcp.ToolExecution {
|
||||
if access.Scope == database.RBACScopeAll {
|
||||
return executions
|
||||
}
|
||||
out := make([]*mcp.ToolExecution, 0, len(executions))
|
||||
for _, exec := range executions {
|
||||
if toolExecutionVisible(exec, access, db) {
|
||||
out = append(out, exec)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func toolExecutionVisible(exec *mcp.ToolExecution, access database.RBACListAccess, db *database.DB) bool {
|
||||
if exec == nil || strings.TrimSpace(access.UserID) == "" {
|
||||
return false
|
||||
}
|
||||
if access.Scope == database.RBACScopeAll || strings.TrimSpace(exec.OwnerUserID) == strings.TrimSpace(access.UserID) {
|
||||
return true
|
||||
}
|
||||
conversationID := strings.TrimSpace(exec.ConversationID)
|
||||
return conversationID != "" && db != nil && db.UserCanAccessResource(access.UserID, access.Scope, "conversation", conversationID)
|
||||
}
|
||||
|
||||
func (h *MonitorHandler) monitorExecutionAllowed(c *gin.Context, id string) bool {
|
||||
access := notificationAccessFromContext(c)
|
||||
if access.Scope == database.RBACScopeAll {
|
||||
return true
|
||||
}
|
||||
id = strings.TrimSpace(id)
|
||||
if id == "" {
|
||||
return false
|
||||
}
|
||||
if exec, ok := h.mcpServer.GetExecution(id); ok {
|
||||
return toolExecutionVisible(exec, access, h.db)
|
||||
}
|
||||
if h.externalMCPMgr != nil {
|
||||
if exec, ok := h.externalMCPMgr.GetExecution(id); ok {
|
||||
return toolExecutionVisible(exec, access, h.db)
|
||||
}
|
||||
}
|
||||
return h.db != nil && h.db.UserCanAccessToolExecution(access.UserID, access.Scope, id)
|
||||
}
|
||||
|
||||
func (h *MonitorHandler) loadExecutionsWithPagination(page, pageSize int, status, toolName string) ([]*mcp.ToolExecution, int) {
|
||||
if h.db == nil {
|
||||
allExecutions := h.mcpServer.GetAllExecutions()
|
||||
@@ -453,6 +535,10 @@ func (h *MonitorHandler) loadStatsMap() map[string]*mcp.ToolStats {
|
||||
// GetExecution 获取特定执行记录
|
||||
func (h *MonitorHandler) GetExecution(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
if !h.monitorExecutionAllowed(c, id) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
|
||||
// 先从内部MCP服务器查找
|
||||
exec, exists := h.mcpServer.GetExecution(id)
|
||||
@@ -493,6 +579,10 @@ func (h *MonitorHandler) CancelExecution(c *gin.Context) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "执行记录ID不能为空"})
|
||||
return
|
||||
}
|
||||
if !h.monitorExecutionAllowed(c, id) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
note := ""
|
||||
dec := json.NewDecoder(c.Request.Body)
|
||||
var body struct {
|
||||
@@ -575,7 +665,7 @@ func (h *MonitorHandler) lookupExecution(id string) *mcp.ToolExecution {
|
||||
return nil
|
||||
}
|
||||
|
||||
// BatchGetToolNames 批量获取工具执行的工具名称(消除前端 N+1 请求)
|
||||
// BatchGetToolNames 批量获取工具执行摘要(消除前端 N+1 请求)
|
||||
func (h *MonitorHandler) BatchGetToolNames(c *gin.Context) {
|
||||
var req struct {
|
||||
IDs []string `json:"ids"`
|
||||
@@ -585,24 +675,32 @@ func (h *MonitorHandler) BatchGetToolNames(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
result := make(map[string]string, len(req.IDs))
|
||||
type executionSummary struct {
|
||||
ToolName string `json:"toolName"`
|
||||
Status string `json:"status"`
|
||||
}
|
||||
|
||||
result := make(map[string]executionSummary, len(req.IDs))
|
||||
for _, id := range req.IDs {
|
||||
if !h.monitorExecutionAllowed(c, id) {
|
||||
continue
|
||||
}
|
||||
// 先从内部MCP服务器查找
|
||||
if exec, exists := h.mcpServer.GetExecution(id); exists {
|
||||
result[id] = exec.ToolName
|
||||
result[id] = executionSummary{ToolName: exec.ToolName, Status: exec.Status}
|
||||
continue
|
||||
}
|
||||
// 再从外部MCP管理器查找
|
||||
if h.externalMCPMgr != nil {
|
||||
if exec, exists := h.externalMCPMgr.GetExecution(id); exists {
|
||||
result[id] = exec.ToolName
|
||||
result[id] = executionSummary{ToolName: exec.ToolName, Status: exec.Status}
|
||||
continue
|
||||
}
|
||||
}
|
||||
// 最后从数据库查找
|
||||
if h.db != nil {
|
||||
if exec, err := h.db.GetToolExecution(id); err == nil && exec != nil {
|
||||
result[id] = exec.ToolName
|
||||
result[id] = executionSummary{ToolName: exec.ToolName, Status: exec.Status}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -750,6 +848,10 @@ func (h *MonitorHandler) DeleteExecution(c *gin.Context) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "执行记录ID不能为空"})
|
||||
return
|
||||
}
|
||||
if !h.monitorExecutionAllowed(c, id) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
|
||||
// 如果使用数据库,先获取执行记录信息,然后删除并更新统计
|
||||
if h.db != nil {
|
||||
@@ -818,6 +920,12 @@ func (h *MonitorHandler) DeleteExecutions(c *gin.Context) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "执行记录ID列表不能为空"})
|
||||
return
|
||||
}
|
||||
for _, id := range request.IDs {
|
||||
if !h.monitorExecutionAllowed(c, id) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问一个或多个执行记录"})
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// 如果使用数据库,先获取执行记录信息,然后删除并更新统计
|
||||
if h.db != nil {
|
||||
|
||||
@@ -133,7 +133,7 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
||||
"userMessageId": prep.UserMessageID,
|
||||
})
|
||||
}
|
||||
if h.runRoleWorkflowStreamIfBound(&req, prep, sendEvent) {
|
||||
if h.runRoleWorkflowStreamIfBound(c, &req, prep, sendEvent) {
|
||||
return
|
||||
}
|
||||
|
||||
@@ -163,7 +163,7 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
||||
var result *multiagent.RunResult
|
||||
var runErr error
|
||||
|
||||
baseCtx, cancelWithCause = context.WithCancelCause(context.Background())
|
||||
baseCtx, cancelWithCause = context.WithCancelCause(detachedAgentContext(c.Request.Context()))
|
||||
taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute)
|
||||
|
||||
if _, err := h.tasks.StartTask(conversationID, req.Message, cancelWithCause); err != nil {
|
||||
@@ -259,7 +259,7 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
||||
if h.tryContinueOnEinoEmptyResponse(taskCtx, mw, conversationID, result, &emptyResponseContinueAttempt, &curHistory, &curFinalMessage, progressCallback) {
|
||||
mainIterationOffset += segmentMainIterationMax
|
||||
timeoutCancel()
|
||||
baseCtx, cancelWithCause, taskCtx, timeoutCancel = h.rebindEinoRunningTask(conversationID, timeoutCancel)
|
||||
baseCtx, cancelWithCause, taskCtx, timeoutCancel = h.rebindEinoRunningTask(taskCtx, conversationID, timeoutCancel)
|
||||
continue
|
||||
}
|
||||
timeoutCancel()
|
||||
@@ -291,7 +291,7 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
||||
})
|
||||
mainIterationOffset += segmentMainIterationMax
|
||||
timeoutCancel()
|
||||
baseCtx, cancelWithCause = context.WithCancelCause(context.Background())
|
||||
baseCtx, cancelWithCause = context.WithCancelCause(detachedAgentContext(baseCtx))
|
||||
h.tasks.BindTaskCancel(conversationID, cancelWithCause)
|
||||
taskCtx, timeoutCancel = context.WithTimeout(baseCtx, 600*time.Minute)
|
||||
h.tasks.UpdateTaskStatus(conversationID, "running")
|
||||
@@ -541,6 +541,8 @@ func formatInterruptContinueUserMessage(note string) string {
|
||||
func multiAgentHTTPErrorStatus(err error) (int, string) {
|
||||
msg := err.Error()
|
||||
switch {
|
||||
case strings.Contains(msg, "无权访问"):
|
||||
return http.StatusForbidden, msg
|
||||
case strings.Contains(msg, "对话不存在"):
|
||||
return http.StatusNotFound, msg
|
||||
case strings.Contains(msg, "未找到该 WebShell"):
|
||||
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
@@ -30,16 +31,34 @@ func (h *AgentHandler) prepareMultiAgentSession(req *ChatRequest, c *gin.Context
|
||||
}
|
||||
|
||||
conversationID := strings.TrimSpace(req.ConversationID)
|
||||
projectID := strings.TrimSpace(effectiveProjectID(h.config, req.ProjectID))
|
||||
webshellID := strings.TrimSpace(req.WebShellConnectionID)
|
||||
session, hasSession := security.CurrentSession(c)
|
||||
if !hasSession || !session.Permissions["chat:write"] {
|
||||
return nil, fmt.Errorf("无权写入对话")
|
||||
}
|
||||
canAccess := func(resourceType, resourceID string) bool {
|
||||
if !hasSession || h.db == nil || strings.TrimSpace(resourceID) == "" {
|
||||
return false
|
||||
}
|
||||
return h.db.UserCanAccessResource(session.UserID, session.Scope, resourceType, resourceID)
|
||||
}
|
||||
if projectID != "" && (!session.Permissions["project:read"] || !canAccess("project", projectID)) {
|
||||
return nil, fmt.Errorf("无权访问目标项目")
|
||||
}
|
||||
if webshellID != "" && (!session.Permissions["webshell:write"] || !canAccess("webshell", webshellID)) {
|
||||
return nil, fmt.Errorf("无权访问该 WebShell 连接")
|
||||
}
|
||||
createdNew := false
|
||||
if conversationID == "" {
|
||||
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) != "" {
|
||||
meta.ProjectID = projectID
|
||||
if webshellID != "" {
|
||||
meta.Source = source + "_webshell"
|
||||
meta.WebShellConnectionID = strings.TrimSpace(req.WebShellConnectionID)
|
||||
meta.WebShellConnectionID = webshellID
|
||||
conv, err = h.db.CreateConversationWithWebshell(meta.WebShellConnectionID, title, meta)
|
||||
} else {
|
||||
conv, err = h.db.CreateConversation(title, meta)
|
||||
@@ -49,10 +68,17 @@ func (h *AgentHandler) prepareMultiAgentSession(req *ChatRequest, c *gin.Context
|
||||
}
|
||||
conversationID = conv.ID
|
||||
createdNew = true
|
||||
if hasSession {
|
||||
_ = h.db.SetResourceOwner("conversation", conversationID, session.UserID)
|
||||
_ = h.db.AssignResourceToUser(session.UserID, "conversation", conversationID)
|
||||
}
|
||||
} else {
|
||||
if _, err := h.db.GetConversation(conversationID); err != nil {
|
||||
return nil, fmt.Errorf("对话不存在")
|
||||
}
|
||||
if !canAccess("conversation", conversationID) {
|
||||
return nil, fmt.Errorf("无权访问该对话")
|
||||
}
|
||||
}
|
||||
|
||||
agentHistoryMessages, err := h.loadHistoryFromAgentTrace(conversationID)
|
||||
@@ -67,8 +93,8 @@ func (h *AgentHandler) prepareMultiAgentSession(req *ChatRequest, c *gin.Context
|
||||
|
||||
finalMessage := req.Message
|
||||
var roleTools []string
|
||||
if req.WebShellConnectionID != "" {
|
||||
conn, errConn := h.db.GetWebshellConnection(strings.TrimSpace(req.WebShellConnectionID))
|
||||
if webshellID != "" {
|
||||
conn, errConn := h.db.GetWebshellConnection(webshellID)
|
||||
if errConn != nil || conn == nil {
|
||||
h.logger.Warn("WebShell AI 助手:未找到连接", zap.String("id", req.WebShellConnectionID), zap.Error(errConn))
|
||||
return nil, fmt.Errorf("未找到该 WebShell 连接")
|
||||
|
||||
@@ -9,6 +9,8 @@ import (
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
@@ -89,6 +91,10 @@ func normalizedSinceSec(sinceMs int64) int64 {
|
||||
return 0
|
||||
}
|
||||
|
||||
func ptrTime(t time.Time) *time.Time {
|
||||
return &t
|
||||
}
|
||||
|
||||
func normalizeSinceMs(raw int64) int64 {
|
||||
if raw > 0 {
|
||||
return raw
|
||||
@@ -126,8 +132,82 @@ func i18nText(english bool, zh string, en string) string {
|
||||
return zh
|
||||
}
|
||||
|
||||
func (h *NotificationHandler) loadPendingHITLItems(limit int, english bool) ([]NotificationSummaryItem, error) {
|
||||
rows, err := h.db.Query(`
|
||||
func notificationAccessFromContext(c *gin.Context) database.RBACListAccess {
|
||||
session, ok := security.CurrentSession(c)
|
||||
if !ok {
|
||||
return database.RBACListAccess{}
|
||||
}
|
||||
return database.RBACListAccess{UserID: session.UserID, Scope: session.Scope}
|
||||
}
|
||||
|
||||
func appendConversationAccessSQL(query string, args []interface{}, column string, access database.RBACListAccess) (string, []interface{}) {
|
||||
userID := strings.TrimSpace(access.UserID)
|
||||
if access.Scope == database.RBACScopeAll {
|
||||
return query, args
|
||||
}
|
||||
if userID == "" {
|
||||
return query + ` AND 1=0`, args
|
||||
}
|
||||
query += ` AND ` + column + ` IS NOT NULL AND ` + column + ` <> '' AND (
|
||||
EXISTS (SELECT 1 FROM conversations c WHERE c.id = ` + column + ` AND c.owner_user_id = ?)
|
||||
OR EXISTS (
|
||||
SELECT 1 FROM rbac_resource_assignments ra
|
||||
WHERE ra.user_id = ? AND ra.resource_type = 'conversation' AND ra.resource_id = ` + column + `
|
||||
)
|
||||
OR EXISTS (
|
||||
SELECT 1 FROM conversations c
|
||||
JOIN projects p ON p.id = c.project_id
|
||||
WHERE c.id = ` + column + ` AND p.owner_user_id = ?
|
||||
)
|
||||
OR EXISTS (
|
||||
SELECT 1 FROM conversations c
|
||||
JOIN rbac_resource_assignments pra ON pra.resource_id = c.project_id
|
||||
WHERE c.id = ` + column + ` AND pra.user_id = ? AND pra.resource_type = 'project'
|
||||
)
|
||||
)`
|
||||
args = append(args, userID, userID, userID, userID)
|
||||
return query, args
|
||||
}
|
||||
|
||||
func appendVulnerabilityNotificationAccessSQL(query string, args []interface{}, access database.RBACListAccess) (string, []interface{}) {
|
||||
userID := strings.TrimSpace(access.UserID)
|
||||
if access.Scope == database.RBACScopeAll {
|
||||
return query, args
|
||||
}
|
||||
if userID == "" {
|
||||
return query + ` AND 1=0`, args
|
||||
}
|
||||
query += ` AND (
|
||||
owner_user_id = ?
|
||||
OR EXISTS (
|
||||
SELECT 1 FROM rbac_resource_assignments ra
|
||||
WHERE ra.user_id = ? AND ra.resource_type = 'vulnerability' AND ra.resource_id = vulnerabilities.id
|
||||
)
|
||||
OR (
|
||||
project_id IS NOT NULL AND project_id <> '' AND (
|
||||
EXISTS (SELECT 1 FROM projects p WHERE p.id = vulnerabilities.project_id AND p.owner_user_id = ?)
|
||||
OR EXISTS (
|
||||
SELECT 1 FROM rbac_resource_assignments pra
|
||||
WHERE pra.user_id = ? AND pra.resource_type = 'project' AND pra.resource_id = vulnerabilities.project_id
|
||||
)
|
||||
)
|
||||
)
|
||||
OR (
|
||||
conversation_id IS NOT NULL AND conversation_id <> '' AND (
|
||||
EXISTS (SELECT 1 FROM conversations c WHERE c.id = vulnerabilities.conversation_id AND c.owner_user_id = ?)
|
||||
OR EXISTS (
|
||||
SELECT 1 FROM rbac_resource_assignments cra
|
||||
WHERE cra.user_id = ? AND cra.resource_type = 'conversation' AND cra.resource_id = vulnerabilities.conversation_id
|
||||
)
|
||||
)
|
||||
)
|
||||
)`
|
||||
args = append(args, userID, userID, userID, userID, userID, userID)
|
||||
return query, args
|
||||
}
|
||||
|
||||
func (h *NotificationHandler) loadPendingHITLItems(limit int, english bool, access database.RBACListAccess) ([]NotificationSummaryItem, error) {
|
||||
query := `
|
||||
SELECT
|
||||
id,
|
||||
conversation_id,
|
||||
@@ -135,9 +215,14 @@ func (h *NotificationHandler) loadPendingHITLItems(limit int, english bool) ([]N
|
||||
COALESCE(CAST(strftime('%s', created_at) AS INTEGER), 0)
|
||||
FROM hitl_interrupts
|
||||
WHERE status = 'pending'
|
||||
ORDER BY created_at DESC
|
||||
`
|
||||
args := []interface{}{}
|
||||
query, args = appendConversationAccessSQL(query, args, "conversation_id", access)
|
||||
query += ` ORDER BY created_at DESC
|
||||
LIMIT ?
|
||||
`, limit)
|
||||
`
|
||||
args = append(args, limit)
|
||||
rows, err := h.db.Query(query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -170,9 +255,9 @@ func (h *NotificationHandler) loadPendingHITLItems(limit int, english bool) ([]N
|
||||
return items, nil
|
||||
}
|
||||
|
||||
func (h *NotificationHandler) loadVulnerabilityItems(sinceMs int64, limit int, english bool) ([]NotificationSummaryItem, map[string]int, error) {
|
||||
func (h *NotificationHandler) loadVulnerabilityItems(sinceMs int64, limit int, english bool, access database.RBACListAccess) ([]NotificationSummaryItem, map[string]int, error) {
|
||||
sinceSec := normalizedSinceSec(sinceMs)
|
||||
rows, err := h.db.Query(`
|
||||
query := `
|
||||
SELECT
|
||||
id,
|
||||
title,
|
||||
@@ -181,9 +266,15 @@ func (h *NotificationHandler) loadVulnerabilityItems(sinceMs int64, limit int, e
|
||||
COALESCE(CAST(strftime('%s', created_at) AS INTEGER), 0)
|
||||
FROM vulnerabilities
|
||||
WHERE CAST(strftime('%s', created_at) AS INTEGER) > ?
|
||||
`
|
||||
args := []interface{}{sinceSec}
|
||||
query, args = appendVulnerabilityNotificationAccessSQL(query, args, access)
|
||||
query += `
|
||||
ORDER BY created_at DESC
|
||||
LIMIT ?
|
||||
`, sinceSec, limit)
|
||||
`
|
||||
args = append(args, limit)
|
||||
rows, err := h.db.Query(query, args...)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
@@ -241,29 +332,23 @@ func (h *NotificationHandler) loadVulnerabilityItems(sinceMs int64, limit int, e
|
||||
}
|
||||
|
||||
// loadC2SessionOnlineEvents 新会话上线(c2_events:session + critical,与 Manager.IngestCheckIn 一致)
|
||||
func (h *NotificationHandler) loadC2SessionOnlineEvents(sinceMs int64, limit int, english bool) ([]NotificationSummaryItem, int, error) {
|
||||
func (h *NotificationHandler) loadC2SessionOnlineEvents(sinceMs int64, limit int, english bool, access database.RBACListAccess) ([]NotificationSummaryItem, int, error) {
|
||||
sinceSec := normalizedSinceSec(sinceMs)
|
||||
rows, err := h.db.Query(`
|
||||
SELECT id, message, COALESCE(session_id, ''),
|
||||
COALESCE(CAST(strftime('%s', created_at) AS INTEGER), 0)
|
||||
FROM c2_events
|
||||
WHERE category = 'session' AND level = 'critical'
|
||||
AND CAST(strftime('%s', created_at) AS INTEGER) > ?
|
||||
ORDER BY created_at DESC
|
||||
LIMIT ?
|
||||
`, sinceSec, limit)
|
||||
events, err := h.db.ListC2EventsForAccess(database.ListC2EventsFilter{
|
||||
Category: "session",
|
||||
Level: "critical",
|
||||
Since: ptrTime(time.Unix(sinceSec, 0)),
|
||||
Limit: limit,
|
||||
}, access)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := make([]NotificationSummaryItem, 0, limit)
|
||||
for rows.Next() {
|
||||
var id, message, sessionID string
|
||||
var createdSec int64
|
||||
if err := rows.Scan(&id, &message, &sessionID, &createdSec); err != nil {
|
||||
for _, e := range events {
|
||||
if e == nil {
|
||||
continue
|
||||
}
|
||||
desc := strings.TrimSpace(message)
|
||||
desc := strings.TrimSpace(e.Message)
|
||||
if len(desc) > 220 {
|
||||
desc = desc[:200] + "…"
|
||||
}
|
||||
@@ -271,19 +356,19 @@ func (h *NotificationHandler) loadC2SessionOnlineEvents(sinceMs int64, limit int
|
||||
desc = i18nText(english, "新会话已建立", "A new session was created")
|
||||
}
|
||||
items = append(items, NotificationSummaryItem{
|
||||
ID: "c2evt:" + id,
|
||||
ID: "c2evt:" + e.ID,
|
||||
Level: "p0",
|
||||
Type: "c2_session_online",
|
||||
Title: i18nText(english, "C2 新会话上线", "C2 new session online"),
|
||||
Desc: desc,
|
||||
Ts: unixSecToRFC3339(createdSec),
|
||||
Ts: e.CreatedAt.UTC().Format(time.RFC3339),
|
||||
Count: 1,
|
||||
Actionable: false,
|
||||
Read: false,
|
||||
SessionID: sessionID,
|
||||
SessionID: e.SessionID,
|
||||
})
|
||||
}
|
||||
return items, len(items), rows.Err()
|
||||
return items, len(items), nil
|
||||
}
|
||||
|
||||
func (h *NotificationHandler) loadFailedExecutionItems(sinceMs int64, limit int, english bool) ([]NotificationSummaryItem, int, error) {
|
||||
@@ -331,7 +416,7 @@ func (h *NotificationHandler) loadFailedExecutionItems(sinceMs int64, limit int,
|
||||
return items, count, nil
|
||||
}
|
||||
|
||||
func (h *NotificationHandler) summarizeLongRunningTasks(threshold time.Duration, english bool) ([]NotificationSummaryItem, int) {
|
||||
func (h *NotificationHandler) summarizeLongRunningTasks(threshold time.Duration, english bool, access database.RBACListAccess) ([]NotificationSummaryItem, int) {
|
||||
if h.agentHandler == nil || h.agentHandler.tasks == nil {
|
||||
return nil, 0
|
||||
}
|
||||
@@ -342,6 +427,9 @@ func (h *NotificationHandler) summarizeLongRunningTasks(threshold time.Duration,
|
||||
if t == nil {
|
||||
continue
|
||||
}
|
||||
if !h.notificationConversationAllowed(access, t.ConversationID) {
|
||||
continue
|
||||
}
|
||||
if now.Sub(t.StartedAt) >= threshold {
|
||||
items = append(items, NotificationSummaryItem{
|
||||
ID: "task_long:" + t.ConversationID,
|
||||
@@ -360,7 +448,7 @@ func (h *NotificationHandler) summarizeLongRunningTasks(threshold time.Duration,
|
||||
return items, len(items)
|
||||
}
|
||||
|
||||
func (h *NotificationHandler) summarizeCompletedTasksSince(sinceMs int64, limit int, english bool) ([]NotificationSummaryItem, int) {
|
||||
func (h *NotificationHandler) summarizeCompletedTasksSince(sinceMs int64, limit int, english bool, access database.RBACListAccess) ([]NotificationSummaryItem, int) {
|
||||
if h.agentHandler == nil || h.agentHandler.tasks == nil {
|
||||
return nil, 0
|
||||
}
|
||||
@@ -371,6 +459,9 @@ func (h *NotificationHandler) summarizeCompletedTasksSince(sinceMs int64, limit
|
||||
if t == nil {
|
||||
continue
|
||||
}
|
||||
if !h.notificationConversationAllowed(access, t.ConversationID) {
|
||||
continue
|
||||
}
|
||||
if t.CompletedAt.After(since) {
|
||||
items = append(items, NotificationSummaryItem{
|
||||
ID: "task_completed:" + t.ConversationID + ":" + strconv.FormatInt(t.CompletedAt.Unix(), 10),
|
||||
@@ -403,14 +494,16 @@ func buildPlaceholders(n int) string {
|
||||
return strings.Join(out, ",")
|
||||
}
|
||||
|
||||
func (h *NotificationHandler) readStatesByIDs(ids []string) (map[string]bool, error) {
|
||||
func (h *NotificationHandler) readStatesByIDs(userID string, ids []string) (map[string]bool, error) {
|
||||
result := make(map[string]bool, len(ids))
|
||||
if len(ids) == 0 {
|
||||
userID = strings.TrimSpace(userID)
|
||||
if len(ids) == 0 || userID == "" {
|
||||
return result, nil
|
||||
}
|
||||
holders := buildPlaceholders(len(ids))
|
||||
query := "SELECT event_id FROM notification_reads WHERE event_id IN (" + holders + ")"
|
||||
args := make([]interface{}, 0, len(ids))
|
||||
query := "SELECT event_id FROM notification_reads_by_user WHERE user_id = ? AND event_id IN (" + holders + ")"
|
||||
args := make([]interface{}, 0, len(ids)+1)
|
||||
args = append(args, userID)
|
||||
for _, id := range ids {
|
||||
args = append(args, id)
|
||||
}
|
||||
@@ -429,7 +522,7 @@ func (h *NotificationHandler) readStatesByIDs(ids []string) (map[string]bool, er
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (h *NotificationHandler) applyReadStates(items []NotificationSummaryItem) ([]NotificationSummaryItem, error) {
|
||||
func (h *NotificationHandler) applyReadStates(userID string, items []NotificationSummaryItem) ([]NotificationSummaryItem, error) {
|
||||
markableIDs := make([]string, 0, len(items))
|
||||
for _, item := range items {
|
||||
if item.Actionable {
|
||||
@@ -437,7 +530,7 @@ func (h *NotificationHandler) applyReadStates(items []NotificationSummaryItem) (
|
||||
}
|
||||
markableIDs = append(markableIDs, item.ID)
|
||||
}
|
||||
readMap, err := h.readStatesByIDs(markableIDs)
|
||||
readMap, err := h.readStatesByIDs(userID, markableIDs)
|
||||
if err != nil {
|
||||
return items, err
|
||||
}
|
||||
@@ -494,34 +587,38 @@ func createNotificationReadTableIfNeeded(db *database.DB) error {
|
||||
return fmt.Errorf("db is nil")
|
||||
}
|
||||
_, err := db.Exec(`
|
||||
CREATE TABLE IF NOT EXISTS notification_reads (
|
||||
event_id TEXT PRIMARY KEY,
|
||||
read_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||||
CREATE TABLE IF NOT EXISTS notification_reads_by_user (
|
||||
user_id TEXT NOT NULL,
|
||||
event_id TEXT NOT NULL,
|
||||
read_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
PRIMARY KEY(user_id, event_id)
|
||||
);
|
||||
`)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, idxErr := db.Exec(`CREATE INDEX IF NOT EXISTS idx_notification_reads_read_at ON notification_reads(read_at DESC);`)
|
||||
_, idxErr := db.Exec(`CREATE INDEX IF NOT EXISTS idx_notification_reads_user_read_at ON notification_reads_by_user(user_id, read_at DESC);`)
|
||||
return idxErr
|
||||
}
|
||||
|
||||
func pruneNotificationReads(db *database.DB, maxRows int) error {
|
||||
func pruneNotificationReads(db *database.DB, userID string, maxRows int) error {
|
||||
if db == nil {
|
||||
return fmt.Errorf("db is nil")
|
||||
}
|
||||
if maxRows <= 0 {
|
||||
userID = strings.TrimSpace(userID)
|
||||
if maxRows <= 0 || userID == "" {
|
||||
return nil
|
||||
}
|
||||
_, err := db.Exec(`
|
||||
DELETE FROM notification_reads
|
||||
WHERE event_id NOT IN (
|
||||
DELETE FROM notification_reads_by_user
|
||||
WHERE user_id = ? AND event_id NOT IN (
|
||||
SELECT event_id
|
||||
FROM notification_reads
|
||||
FROM notification_reads_by_user
|
||||
WHERE user_id = ?
|
||||
ORDER BY read_at DESC, rowid DESC
|
||||
LIMIT ?
|
||||
)
|
||||
`, maxRows)
|
||||
`, userID, userID, maxRows)
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -564,6 +661,11 @@ func (h *NotificationHandler) MarkRead(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "marked": 0})
|
||||
return
|
||||
}
|
||||
session, ok := security.CurrentSession(c)
|
||||
if !ok || strings.TrimSpace(session.UserID) == "" {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"error": "missing authenticated user"})
|
||||
return
|
||||
}
|
||||
tx, err := h.db.Begin()
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to begin transaction"})
|
||||
@@ -573,9 +675,9 @@ func (h *NotificationHandler) MarkRead(c *gin.Context) {
|
||||
_ = tx.Rollback()
|
||||
}()
|
||||
stmt, err := tx.Prepare(`
|
||||
INSERT INTO notification_reads(event_id, read_at)
|
||||
VALUES(?, CURRENT_TIMESTAMP)
|
||||
ON CONFLICT(event_id) DO UPDATE SET read_at = CURRENT_TIMESTAMP
|
||||
INSERT INTO notification_reads_by_user(user_id, event_id, read_at)
|
||||
VALUES(?, ?, CURRENT_TIMESTAMP)
|
||||
ON CONFLICT(user_id, event_id) DO UPDATE SET read_at = CURRENT_TIMESTAMP
|
||||
`)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to prepare statement"})
|
||||
@@ -588,7 +690,7 @@ func (h *NotificationHandler) MarkRead(c *gin.Context) {
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if _, err := stmt.Exec(id); err != nil {
|
||||
if _, err := stmt.Exec(session.UserID, id); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to mark read"})
|
||||
return
|
||||
}
|
||||
@@ -598,7 +700,7 @@ func (h *NotificationHandler) MarkRead(c *gin.Context) {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to commit read marks"})
|
||||
return
|
||||
}
|
||||
if err := pruneNotificationReads(h.db, notificationReadMaxRows); err != nil {
|
||||
if err := pruneNotificationReads(h.db, session.UserID, notificationReadMaxRows); err != nil {
|
||||
h.logger.Warn("裁剪通知已读记录失败", zap.Error(err))
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true, "marked": marked})
|
||||
@@ -626,30 +728,57 @@ func (h *NotificationHandler) GetSummary(c *gin.Context) {
|
||||
if limit > 200 {
|
||||
limit = 200
|
||||
}
|
||||
access := notificationAccessFromContext(c)
|
||||
|
||||
hitlItems, err := h.loadPendingHITLItems(limit, english)
|
||||
if err != nil {
|
||||
h.logger.Warn("加载 HITL 通知失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to summarize hitl notifications"})
|
||||
return
|
||||
hitlItems := []NotificationSummaryItem{}
|
||||
if security.SessionHasPermission(c, "hitl:read") {
|
||||
var err error
|
||||
hitlItems, err = h.loadPendingHITLItems(limit, english, access)
|
||||
if err != nil {
|
||||
h.logger.Warn("加载 HITL 通知失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to summarize hitl notifications"})
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
vulnItems, vulnCounts, err := h.loadVulnerabilityItems(sinceMs, limit, english)
|
||||
if err != nil {
|
||||
h.logger.Warn("加载漏洞通知失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to summarize vulnerabilities"})
|
||||
return
|
||||
vulnItems := []NotificationSummaryItem{}
|
||||
vulnCounts := map[string]int{
|
||||
"newCriticalVulns": 0,
|
||||
"newHighVulns": 0,
|
||||
"newMediumVulns": 0,
|
||||
"newLowVulns": 0,
|
||||
"newInfoVulns": 0,
|
||||
}
|
||||
if security.SessionHasPermission(c, "vulnerability:read") {
|
||||
var err error
|
||||
vulnItems, vulnCounts, err = h.loadVulnerabilityItems(sinceMs, limit, english, access)
|
||||
if err != nil {
|
||||
h.logger.Warn("加载漏洞通知失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to summarize vulnerabilities"})
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
c2OnlineItems, c2OnlineCount, err := h.loadC2SessionOnlineEvents(sinceMs, limit, english)
|
||||
if err != nil {
|
||||
h.logger.Warn("加载 C2 会话上线通知失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to summarize c2 session events"})
|
||||
return
|
||||
c2OnlineItems := []NotificationSummaryItem{}
|
||||
c2OnlineCount := 0
|
||||
if security.SessionHasPermission(c, "c2:read") {
|
||||
var err error
|
||||
c2OnlineItems, c2OnlineCount, err = h.loadC2SessionOnlineEvents(sinceMs, limit, english, access)
|
||||
if err != nil {
|
||||
h.logger.Warn("加载 C2 会话上线通知失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to summarize c2 session events"})
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
longRunningItems, longRunningCount := h.summarizeLongRunningTasks(15*time.Minute, english)
|
||||
completedItems, completedCount := h.summarizeCompletedTasksSince(sinceMs, limit, english)
|
||||
longRunningItems := []NotificationSummaryItem{}
|
||||
completedItems := []NotificationSummaryItem{}
|
||||
longRunningCount := 0
|
||||
completedCount := 0
|
||||
if security.SessionHasPermission(c, "tasks:read") || security.SessionHasPermission(c, "chat:read") {
|
||||
longRunningItems, longRunningCount = h.summarizeLongRunningTasks(15*time.Minute, english, access)
|
||||
completedItems, completedCount = h.summarizeCompletedTasksSince(sinceMs, limit, english, access)
|
||||
}
|
||||
|
||||
items := make([]NotificationSummaryItem, 0, len(hitlItems)+len(vulnItems)+len(c2OnlineItems)+len(longRunningItems)+len(completedItems))
|
||||
items = append(items, hitlItems...)
|
||||
@@ -658,7 +787,8 @@ func (h *NotificationHandler) GetSummary(c *gin.Context) {
|
||||
items = append(items, longRunningItems...)
|
||||
items = append(items, completedItems...)
|
||||
|
||||
items, err = h.applyReadStates(items)
|
||||
session, _ := security.CurrentSession(c)
|
||||
items, err := h.applyReadStates(session.UserID, items)
|
||||
if err != nil {
|
||||
h.logger.Warn("加载通知已读状态失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to load notification read states"})
|
||||
@@ -697,3 +827,11 @@ func (h *NotificationHandler) GetSummary(c *gin.Context) {
|
||||
Items: items,
|
||||
})
|
||||
}
|
||||
|
||||
func (h *NotificationHandler) notificationConversationAllowed(access database.RBACListAccess, conversationID string) bool {
|
||||
conversationID = strings.TrimSpace(conversationID)
|
||||
if conversationID == "" {
|
||||
return access.Scope == database.RBACScopeAll
|
||||
}
|
||||
return h.db.UserCanAccessResource(access.UserID, access.Scope, "conversation", conversationID)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestNotificationReadStateIsPerUser(t *testing.T) {
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "notification-rbac.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
h := NewNotificationHandler(db, nil, zap.NewNop())
|
||||
router := gin.New()
|
||||
router.Use(func(c *gin.Context) {
|
||||
c.Set(security.ContextSessionKey, security.Session{UserID: "u1", Scope: database.RBACScopeAssigned})
|
||||
c.Next()
|
||||
})
|
||||
router.POST("/notifications/read", h.MarkRead)
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/notifications/read", bytes.NewBufferString(`{"eventIds":["vuln:v1"]}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
router.ServeHTTP(w, req)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("mark read status = %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
u1, err := h.readStatesByIDs("u1", []string{"vuln:v1"})
|
||||
if err != nil || !u1["vuln:v1"] {
|
||||
t.Fatalf("u1 read state = %#v, err=%v", u1, err)
|
||||
}
|
||||
u2, err := h.readStatesByIDs("u2", []string{"vuln:v1"})
|
||||
if err != nil || u2["vuln:v1"] {
|
||||
t.Fatalf("u2 inherited u1 read state = %#v, err=%v", u2, err)
|
||||
}
|
||||
}
|
||||
@@ -4735,7 +4735,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"get": map[string]interface{}{
|
||||
"tags": []string{"对话交互"},
|
||||
"summary": "获取消息过程详情",
|
||||
"description": "按需加载指定消息的执行过程详情,包括工具调用、思考过程等事件。",
|
||||
"description": "按需分页加载指定消息的执行过程详情,包括工具调用、思考过程等事件。默认返回 50 条;导出或旧集成需要全量时可显式传 full=1。",
|
||||
"operationId": "getMessageProcessDetails",
|
||||
"parameters": []map[string]interface{}{
|
||||
{
|
||||
@@ -4745,6 +4745,34 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"description": "消息ID",
|
||||
"schema": map[string]interface{}{"type": "string"},
|
||||
},
|
||||
{
|
||||
"name": "summary",
|
||||
"in": "query",
|
||||
"required": false,
|
||||
"description": "仅返回过程详情摘要(total / iterationCount / maxIteration)",
|
||||
"schema": map[string]interface{}{"type": "boolean"},
|
||||
},
|
||||
{
|
||||
"name": "limit",
|
||||
"in": "query",
|
||||
"required": false,
|
||||
"description": "分页大小,默认 50,最大 500",
|
||||
"schema": map[string]interface{}{"type": "integer", "default": 50, "maximum": 500},
|
||||
},
|
||||
{
|
||||
"name": "offset",
|
||||
"in": "query",
|
||||
"required": false,
|
||||
"description": "分页偏移量,默认 0",
|
||||
"schema": map[string]interface{}{"type": "integer", "default": 0},
|
||||
},
|
||||
{
|
||||
"name": "full",
|
||||
"in": "query",
|
||||
"required": false,
|
||||
"description": "显式返回全量过程详情;仅建议导出/兼容旧集成使用",
|
||||
"schema": map[string]interface{}{"type": "boolean"},
|
||||
},
|
||||
},
|
||||
"responses": map[string]interface{}{
|
||||
"200": map[string]interface{}{
|
||||
@@ -4769,6 +4797,14 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
},
|
||||
},
|
||||
},
|
||||
"total": map[string]interface{}{"type": "integer", "description": "过程详情总数"},
|
||||
"offset": map[string]interface{}{"type": "integer", "description": "当前分页偏移量"},
|
||||
"limit": map[string]interface{}{"type": "integer", "description": "当前分页大小"},
|
||||
"hasMore": map[string]interface{}{"type": "boolean", "description": "是否还有更多过程详情"},
|
||||
"summary": map[string]interface{}{
|
||||
"type": "object",
|
||||
"description": "summary=1 时返回的摘要",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
|
||||
@@ -9,6 +9,7 @@ import (
|
||||
"cyberstrike-ai/internal/attackchain"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/project"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
@@ -70,6 +71,10 @@ func (h *ProjectHandler) CreateProject(c *gin.Context) {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if session, ok := security.CurrentSession(c); ok {
|
||||
_ = h.db.SetResourceOwner("project", created.ID, session.UserID)
|
||||
_ = h.db.AssignResourceToUser(session.UserID, "project", created.ID)
|
||||
}
|
||||
c.JSON(http.StatusOK, created)
|
||||
}
|
||||
|
||||
@@ -82,7 +87,8 @@ func (h *ProjectHandler) GetDashboardSummary(c *gin.Context) {
|
||||
if limit > 50 {
|
||||
limit = 50
|
||||
}
|
||||
summary, err := h.db.GetProjectDashboardSummary(limit)
|
||||
session, _ := security.CurrentSession(c)
|
||||
summary, err := h.db.GetProjectDashboardSummaryForAccess(limit, session.UserID, session.Scope)
|
||||
if err != nil {
|
||||
h.logger.Error("获取项目仪表盘摘要失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
@@ -106,7 +112,8 @@ func (h *ProjectHandler) ListProjects(c *gin.Context) {
|
||||
if limit > 500 {
|
||||
limit = 500
|
||||
}
|
||||
list, err := h.db.ListProjects(status, search, limit, offset)
|
||||
session, _ := security.CurrentSession(c)
|
||||
list, err := h.db.ListProjectsForAccess(status, search, limit, offset, session.UserID, session.Scope)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
@@ -114,7 +121,7 @@ func (h *ProjectHandler) ListProjects(c *gin.Context) {
|
||||
if list == nil {
|
||||
list = []*database.Project{}
|
||||
}
|
||||
total, err := h.db.CountProjects(status, search)
|
||||
total, err := h.db.CountProjectsForAccess(status, search, session.UserID, session.Scope)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
@@ -635,6 +642,12 @@ func (h *ProjectHandler) DeleteFactEdge(c *gin.Context) {
|
||||
func (h *ProjectHandler) PromoteAttackChain(c *gin.Context) {
|
||||
projectID := c.Param("id")
|
||||
conversationID := c.Param("conversationId")
|
||||
session, ok := security.CurrentSession(c)
|
||||
if !ok || !h.db.UserCanAccessResource(session.UserID, session.Scope, "project", projectID) ||
|
||||
!h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", conversationID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问目标项目或来源对话"})
|
||||
return
|
||||
}
|
||||
result, err := attackchain.PromoteToProject(h.db, projectID, conversationID)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
|
||||
@@ -0,0 +1,406 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
type RBACHandler struct {
|
||||
db *database.DB
|
||||
logger *zap.Logger
|
||||
audit *audit.Service
|
||||
auth *security.AuthManager
|
||||
}
|
||||
|
||||
func NewRBACHandler(db *database.DB, logger *zap.Logger) *RBACHandler {
|
||||
return &RBACHandler{db: db, logger: logger}
|
||||
}
|
||||
|
||||
func (h *RBACHandler) SetAudit(s *audit.Service) {
|
||||
h.audit = s
|
||||
}
|
||||
|
||||
func (h *RBACHandler) SetAuthManager(m *security.AuthManager) {
|
||||
h.auth = m
|
||||
}
|
||||
|
||||
func (h *RBACHandler) Me(c *gin.Context) {
|
||||
session, _ := security.CurrentSession(c)
|
||||
resolvedScope := session.Scope
|
||||
permissionScopes := session.PermissionScopes
|
||||
if principal, ok := authctx.PrincipalFromContext(c.Request.Context()); ok {
|
||||
resolvedScope = principal.Scope
|
||||
permissionScopes = principal.PermissionScopes
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"user": gin.H{
|
||||
"id": session.UserID,
|
||||
"username": session.Username,
|
||||
"display_name": session.DisplayName,
|
||||
},
|
||||
"roles": session.Roles,
|
||||
"permissions": permissionKeys(session.Permissions),
|
||||
"scope": resolvedScope,
|
||||
"permission_scopes": permissionScopes,
|
||||
})
|
||||
}
|
||||
|
||||
func (h *RBACHandler) Metadata(c *gin.Context) {
|
||||
roles, err := h.db.ListRBACRoles()
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
rolePermissions := map[string][]string{}
|
||||
for _, role := range roles {
|
||||
keys, _ := h.db.ListRBACRolePermissionKeys(role.ID)
|
||||
rolePermissions[role.ID] = keys
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"permissions": security.PermissionCatalog,
|
||||
"roles": roles,
|
||||
"role_permissions": rolePermissions,
|
||||
"scopes": []string{database.RBACScopeAll, database.RBACScopeAssigned, database.RBACScopeOwn},
|
||||
})
|
||||
}
|
||||
|
||||
func (h *RBACHandler) ListRoles(c *gin.Context) {
|
||||
roles, err := h.db.ListRBACRoles()
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
out := make([]gin.H, 0, len(roles))
|
||||
for _, role := range roles {
|
||||
keys, _ := h.db.ListRBACRolePermissionKeys(role.ID)
|
||||
out = append(out, gin.H{
|
||||
"id": role.ID,
|
||||
"name": role.Name,
|
||||
"description": role.Description,
|
||||
"scope": role.Scope,
|
||||
"is_system": role.IsSystem,
|
||||
"permissions": keys,
|
||||
"created_at": role.CreatedAt,
|
||||
"updated_at": role.UpdatedAt,
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"roles": out})
|
||||
}
|
||||
|
||||
type upsertRBACRoleRequest struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name" binding:"required"`
|
||||
Description string `json:"description"`
|
||||
Scope string `json:"scope"`
|
||||
Permissions []string `json:"permissions"`
|
||||
}
|
||||
|
||||
func validateRBACPermissionKeys(keys []string) error {
|
||||
for _, key := range keys {
|
||||
key = strings.TrimSpace(key)
|
||||
if key == "" {
|
||||
continue
|
||||
}
|
||||
if _, ok := security.PermissionCatalog[key]; !ok {
|
||||
return fmt.Errorf("未知权限: %s", key)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (h *RBACHandler) CreateRole(c *gin.Context) {
|
||||
var req upsertRBACRoleRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if err := validateRBACPermissionKeys(req.Permissions); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
role, err := h.db.UpsertRBACRole("", req.Name, req.Description, req.Scope, req.Permissions)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "rbac", "create_role", "创建平台角色", "role", role.ID, nil)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"role": role})
|
||||
}
|
||||
|
||||
func (h *RBACHandler) UpdateRole(c *gin.Context) {
|
||||
id := strings.TrimSpace(c.Param("id"))
|
||||
existing, err := h.db.GetRBACRoleByID(id)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "角色不存在"})
|
||||
return
|
||||
}
|
||||
if existing.IsSystem {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "系统内置角色不可修改,请创建自定义角色"})
|
||||
return
|
||||
}
|
||||
var req upsertRBACRoleRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if err := validateRBACPermissionKeys(req.Permissions); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
role, err := h.db.UpsertRBACRole(id, req.Name, req.Description, req.Scope, req.Permissions)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if existing.IsSystem {
|
||||
role.IsSystem = true
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "rbac", "update_role", "更新平台角色", "role", id, nil)
|
||||
}
|
||||
if h.auth != nil {
|
||||
h.auth.RevokeAllSessions()
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"role": role})
|
||||
}
|
||||
|
||||
func (h *RBACHandler) DeleteRole(c *gin.Context) {
|
||||
id := strings.TrimSpace(c.Param("id"))
|
||||
if err := h.db.DeleteRBACRole(id); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "rbac", "delete_role", "删除平台角色", "role", id, nil)
|
||||
}
|
||||
if h.auth != nil {
|
||||
h.auth.RevokeAllSessions()
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"success": true})
|
||||
}
|
||||
|
||||
func (h *RBACHandler) ListUsers(c *gin.Context) {
|
||||
users, err := h.db.ListRBACUsers()
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
out := make([]gin.H, 0, len(users))
|
||||
for _, user := range users {
|
||||
roleIDs, _ := h.db.ListRBACUserRoleIDs(user.ID)
|
||||
out = append(out, gin.H{
|
||||
"id": user.ID,
|
||||
"username": user.Username,
|
||||
"display_name": user.DisplayName,
|
||||
"enabled": user.Enabled,
|
||||
"is_builtin": user.IsBuiltin,
|
||||
"roles": roleIDs,
|
||||
"created_at": user.CreatedAt,
|
||||
"updated_at": user.UpdatedAt,
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"users": out})
|
||||
}
|
||||
|
||||
type createRBACUserRequest struct {
|
||||
Username string `json:"username" binding:"required"`
|
||||
DisplayName string `json:"display_name"`
|
||||
Password string `json:"password" binding:"required"`
|
||||
Enabled *bool `json:"enabled"`
|
||||
Roles []string `json:"roles"`
|
||||
}
|
||||
|
||||
func (h *RBACHandler) CreateUser(c *gin.Context) {
|
||||
var req createRBACUserRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if len(strings.TrimSpace(req.Password)) < 8 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "密码长度至少需要 8 位"})
|
||||
return
|
||||
}
|
||||
hash, err := security.HashPassword(req.Password)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
enabled := true
|
||||
if req.Enabled != nil {
|
||||
enabled = *req.Enabled
|
||||
}
|
||||
user, err := h.db.CreateRBACUser(req.Username, req.DisplayName, hash, enabled, req.Roles)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "rbac", "create_user", "创建平台用户", "user", user.ID, map[string]interface{}{"username": user.Username})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"user": user})
|
||||
}
|
||||
|
||||
type updateRBACUserRequest struct {
|
||||
DisplayName *string `json:"display_name"`
|
||||
Password *string `json:"password"`
|
||||
Enabled *bool `json:"enabled"`
|
||||
Roles *[]string `json:"roles"`
|
||||
}
|
||||
|
||||
func (h *RBACHandler) UpdateUser(c *gin.Context) {
|
||||
id := strings.TrimSpace(c.Param("id"))
|
||||
user, err := h.db.GetRBACUserByID(id)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "用户不存在"})
|
||||
return
|
||||
}
|
||||
var req updateRBACUserRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
displayName := user.DisplayName
|
||||
if req.DisplayName != nil {
|
||||
displayName = *req.DisplayName
|
||||
}
|
||||
if err := h.db.UpdateRBACUser(id, displayName, req.Enabled, req.Roles); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if req.Password != nil && strings.TrimSpace(*req.Password) != "" {
|
||||
if len(strings.TrimSpace(*req.Password)) < 8 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "密码长度至少需要 8 位"})
|
||||
return
|
||||
}
|
||||
hash, err := security.HashPassword(*req.Password)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if err := h.db.UpdateRBACUserPassword(id, hash); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "rbac", "update_user", "更新平台用户", "user", id, nil)
|
||||
}
|
||||
if h.auth != nil {
|
||||
h.auth.RevokeUserSessions(id)
|
||||
}
|
||||
updated, _ := h.db.GetRBACUserByID(id)
|
||||
c.JSON(http.StatusOK, gin.H{"user": updated})
|
||||
}
|
||||
|
||||
func (h *RBACHandler) DeleteUser(c *gin.Context) {
|
||||
id := strings.TrimSpace(c.Param("id"))
|
||||
if err := h.db.DeleteRBACUser(id); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "rbac", "delete_user", "删除平台用户", "user", id, nil)
|
||||
}
|
||||
if h.auth != nil {
|
||||
h.auth.RevokeUserSessions(id)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"success": true})
|
||||
}
|
||||
|
||||
type assignResourceRequest struct {
|
||||
UserID string `json:"user_id" binding:"required"`
|
||||
ResourceType string `json:"resource_type" binding:"required"`
|
||||
ResourceID string `json:"resource_id"`
|
||||
ResourceIDs []string `json:"resource_ids"`
|
||||
}
|
||||
|
||||
func (h *RBACHandler) AssignResource(c *gin.Context) {
|
||||
var req assignResourceRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
resourceIDs := append([]string(nil), req.ResourceIDs...)
|
||||
if strings.TrimSpace(req.ResourceID) != "" {
|
||||
resourceIDs = append(resourceIDs, req.ResourceID)
|
||||
}
|
||||
if len(resourceIDs) == 0 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "至少需要一个资源 ID"})
|
||||
return
|
||||
}
|
||||
created, err := h.db.AssignResourcesToUser(req.UserID, req.ResourceType, resourceIDs)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if h.audit != nil {
|
||||
for _, resourceID := range resourceIDs {
|
||||
h.audit.RecordOK(c, "rbac", "assign_resource", "授权资源访问", req.ResourceType, strings.TrimSpace(resourceID), map[string]interface{}{"user_id": req.UserID})
|
||||
}
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"requested": len(resourceIDs),
|
||||
"created": created,
|
||||
"skipped": int64(len(resourceIDs)) - created,
|
||||
})
|
||||
}
|
||||
|
||||
func (h *RBACHandler) ListResourceAssignments(c *gin.Context) {
|
||||
rows, err := h.db.ListRBACResourceAssignments(c.Query("user_id"))
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"assignments": rows})
|
||||
}
|
||||
|
||||
func (h *RBACHandler) ListAssignableResources(c *gin.Context) {
|
||||
limit, _ := strconv.Atoi(c.DefaultQuery("limit", "50"))
|
||||
if limit <= 0 || limit > 50 {
|
||||
limit = 50
|
||||
}
|
||||
offset, _ := strconv.Atoi(c.DefaultQuery("offset", "0"))
|
||||
if offset < 0 {
|
||||
offset = 0
|
||||
}
|
||||
resources, err := h.db.ListAssignableRBACResourcesPage(c.Query("type"), c.Query("q"), limit+1, offset)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
hasMore := len(resources) > limit
|
||||
if hasMore {
|
||||
resources = resources[:limit]
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"resources": resources,
|
||||
"has_more": hasMore,
|
||||
"limit": limit,
|
||||
"offset": offset,
|
||||
})
|
||||
}
|
||||
|
||||
func (h *RBACHandler) DeleteResourceAssignment(c *gin.Context) {
|
||||
id := strings.TrimSpace(c.Param("id"))
|
||||
if err := h.db.DeleteRBACResourceAssignment(id); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "rbac", "delete_resource_assignment", "撤销资源授权", "resource_assignment", id, nil)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"success": true})
|
||||
}
|
||||
@@ -0,0 +1,206 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestDetachedAgentContextRetainsPrincipalWithoutParentCancellation(t *testing.T) {
|
||||
parent, cancel := context.WithCancel(context.Background())
|
||||
parent = authctx.WithPrincipal(parent, authctx.NewPrincipal("u1", "user", database.RBACScopeAssigned, map[string]bool{"agent:execute": true}))
|
||||
detached := detachedAgentContext(parent)
|
||||
cancel()
|
||||
if err := detached.Err(); err != nil {
|
||||
t.Fatalf("detached context inherited cancellation: %v", err)
|
||||
}
|
||||
principal, ok := authctx.PrincipalFromContext(detached)
|
||||
if !ok || principal.UserID != "u1" || !principal.HasPermission("agent:execute") {
|
||||
t.Fatalf("detached context lost principal: %#v, ok=%v", principal, ok)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPromoteAttackChainRequiresSourceConversationAccess(t *testing.T) {
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "promote-rbac.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
project, _ := db.CreateProject(&database.Project{Name: "owned"})
|
||||
conversation, _ := db.CreateConversation("foreign", database.ConversationCreateMeta{})
|
||||
_ = db.SetResourceOwner("project", project.ID, "u1")
|
||||
_ = db.SetResourceOwner("conversation", conversation.ID, "u2")
|
||||
h := NewProjectHandler(db, zap.NewNop())
|
||||
router := gin.New()
|
||||
router.Use(func(c *gin.Context) {
|
||||
c.Set(security.ContextSessionKey, security.Session{UserID: "u1", Scope: database.RBACScopeOwn})
|
||||
c.Next()
|
||||
})
|
||||
router.POST("/api/projects/:id/promote-attack-chain/:conversationId", h.PromoteAttackChain)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, httptest.NewRequest(http.MethodPost, "/api/projects/"+project.ID+"/promote-attack-chain/"+conversation.ID, nil))
|
||||
if w.Code != http.StatusForbidden {
|
||||
t.Fatalf("status = %d, want 403: %s", w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestVulnerabilityCannotBeReparentedToForeignProject(t *testing.T) {
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "vuln-reparent-rbac.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
owned, _ := db.CreateProject(&database.Project{Name: "owned"})
|
||||
foreign, _ := db.CreateProject(&database.Project{Name: "foreign"})
|
||||
_ = db.SetResourceOwner("project", owned.ID, "u1")
|
||||
_ = db.SetResourceOwner("project", foreign.ID, "u2")
|
||||
vulnerability, err := db.CreateVulnerability(&database.Vulnerability{Title: "v", Severity: "high", ProjectID: owned.ID})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_ = db.SetResourceOwner("vulnerability", vulnerability.ID, "u1")
|
||||
h := NewVulnerabilityHandler(db, zap.NewNop())
|
||||
router := gin.New()
|
||||
router.Use(func(c *gin.Context) {
|
||||
c.Set(security.ContextSessionKey, security.Session{UserID: "u1", Scope: database.RBACScopeOwn})
|
||||
c.Next()
|
||||
})
|
||||
router.PUT("/api/vulnerabilities/:id", h.UpdateVulnerability)
|
||||
body, _ := json.Marshal(map[string]interface{}{"project_id": foreign.ID})
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, httptest.NewRequest(http.MethodPut, "/api/vulnerabilities/"+vulnerability.ID, bytes.NewReader(body)))
|
||||
if w.Code != http.StatusForbidden {
|
||||
t.Fatalf("status = %d, want 403: %s", w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentTaskEndpointsFilterAndRejectForeignConversations(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
db, user := setupConversationRBACTest(t)
|
||||
allowed, _ := db.CreateConversation("allowed", database.ConversationCreateMeta{})
|
||||
hidden, _ := db.CreateConversation("hidden", database.ConversationCreateMeta{})
|
||||
if err := db.AssignResourceToUser(user.ID, "conversation", allowed.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
tasks := NewAgentTaskManager()
|
||||
if _, err := tasks.StartTask(allowed.ID, "visible", func(error) {}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := tasks.StartTask(hidden.ID, "secret", func(error) {}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
h := &AgentHandler{db: db, tasks: tasks, logger: zap.NewNop()}
|
||||
|
||||
w := performAssignedHandler(user, http.MethodGet, "/api/agent-loop/tasks", nil, h.ListAgentTasks)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("list status = %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
var response struct {
|
||||
Tasks []*AgentTask `json:"tasks"`
|
||||
}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(response.Tasks) != 1 || response.Tasks[0].ConversationID != allowed.ID {
|
||||
t.Fatalf("tasks = %#v, want only %s", response.Tasks, allowed.ID)
|
||||
}
|
||||
|
||||
w = performAssignedHandler(user, http.MethodPost, "/api/agent-loop/cancel", map[string]string{"conversationId": hidden.ID}, h.CancelAgentLoop)
|
||||
if w.Code != http.StatusForbidden {
|
||||
t.Fatalf("cancel status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatUploadPathAuthorizationFollowsConversationAccess(t *testing.T) {
|
||||
db, user := setupConversationRBACTest(t)
|
||||
allowed, _ := db.CreateConversation("allowed", database.ConversationCreateMeta{})
|
||||
hidden, _ := db.CreateConversation("hidden", database.ConversationCreateMeta{})
|
||||
if err := db.AssignResourceToUser(user.ID, "conversation", allowed.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
h := NewChatUploadsHandler(zap.NewNop(), db)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Set(security.ContextSessionKey, security.Session{UserID: user.ID, Scope: database.RBACScopeAssigned, Permissions: map[string]bool{"chat:write": true}})
|
||||
|
||||
if !h.pathAllowed(c, filepath.ToSlash(filepath.Join("2026-07-10", allowed.ID, "a.txt"))) {
|
||||
t.Fatal("assigned conversation attachment should be accessible")
|
||||
}
|
||||
if h.pathAllowed(c, filepath.ToSlash(filepath.Join("2026-07-10", hidden.ID, "secret.txt"))) {
|
||||
t.Fatal("foreign conversation attachment should be denied")
|
||||
}
|
||||
if h.pathAllowed(c, "2026-07-10/_manual/secret.txt") {
|
||||
t.Fatal("unowned manual attachment should fail closed")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPrepareMultiAgentSessionRejectsForeignConversation(t *testing.T) {
|
||||
db, user := setupConversationRBACTest(t)
|
||||
hidden, _ := db.CreateConversation("hidden", database.ConversationCreateMeta{})
|
||||
h := &AgentHandler{db: db, logger: zap.NewNop()}
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Set(security.ContextSessionKey, security.Session{UserID: user.ID, Scope: database.RBACScopeAssigned, Permissions: map[string]bool{"chat:write": true}})
|
||||
|
||||
_, err := h.prepareMultiAgentSession(&ChatRequest{ConversationID: hidden.ID, Message: "write"}, c, "test")
|
||||
if err == nil || err.Error() != "无权访问该对话" {
|
||||
t.Fatalf("err = %v, want unauthorized conversation", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMonitorExecutionDetailRejectsForeignOwner(t *testing.T) {
|
||||
db, user := setupConversationRBACTest(t)
|
||||
for _, exec := range []*mcp.ToolExecution{
|
||||
{ID: "exec-allowed", ToolName: "allowed", Status: "completed", StartTime: time.Now(), OwnerUserID: user.ID},
|
||||
{ID: "exec-hidden", ToolName: "hidden", Status: "completed", StartTime: time.Now(), OwnerUserID: "another-user"},
|
||||
} {
|
||||
if err := db.SaveToolExecution(exec); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
h := NewMonitorHandler(mcp.NewServerWithStorage(zap.NewNop(), db), nil, db, zap.NewNop())
|
||||
|
||||
request := func(id string) *httptest.ResponseRecorder {
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/api/monitor/execution/"+id, nil)
|
||||
c.Params = gin.Params{{Key: "id", Value: id}}
|
||||
c.Set(security.ContextSessionKey, security.Session{UserID: user.ID, Scope: database.RBACScopeAssigned})
|
||||
h.GetExecution(c)
|
||||
return w
|
||||
}
|
||||
if w := request("exec-hidden"); w.Code != http.StatusForbidden {
|
||||
t.Fatalf("hidden status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String())
|
||||
}
|
||||
if w := request("exec-allowed"); w.Code != http.StatusOK {
|
||||
t.Fatalf("allowed status = %d, want %d: %s", w.Code, http.StatusOK, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func performAssignedHandler(user *database.RBACUser, method, path string, body interface{}, handler gin.HandlerFunc) *httptest.ResponseRecorder {
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
var req *http.Request
|
||||
if body == nil {
|
||||
req = httptest.NewRequest(method, path, nil)
|
||||
} else {
|
||||
payload, _ := json.Marshal(body)
|
||||
req = httptest.NewRequest(method, path, bytes.NewReader(payload))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
}
|
||||
c.Request = req
|
||||
c.Set(security.ContextSessionKey, security.Session{UserID: user.ID, Scope: database.RBACScopeAssigned})
|
||||
handler(c)
|
||||
return w
|
||||
}
|
||||
@@ -0,0 +1,12 @@
|
||||
package handler
|
||||
|
||||
import "context"
|
||||
|
||||
// detachedAgentContext lets a long-running Agent survive an SSE disconnect
|
||||
// while retaining immutable request values such as the authenticated Principal.
|
||||
func detachedAgentContext(parent context.Context) context.Context {
|
||||
if parent == nil {
|
||||
parent = context.Background()
|
||||
}
|
||||
return context.WithoutCancel(parent)
|
||||
}
|
||||
@@ -0,0 +1,118 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestRBACAssignResourceBatchIsAtomicAndLegacyCompatible(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "rbac-handler.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
user, err := db.CreateRBACUser("api-member", "API Member", "hash", true, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
p1, _ := db.CreateProject(&database.Project{Name: "p1"})
|
||||
p2, _ := db.CreateProject(&database.Project{Name: "p2"})
|
||||
p3, _ := db.CreateProject(&database.Project{Name: "p3"})
|
||||
|
||||
h := NewRBACHandler(db, zap.NewNop())
|
||||
router := gin.New()
|
||||
router.POST("/api/rbac/resource-assignments", h.AssignResource)
|
||||
|
||||
batch := performRBACJSONRequest(t, router, map[string]interface{}{
|
||||
"user_id": user.ID, "resource_type": "project", "resource_ids": []string{p1.ID, p2.ID},
|
||||
})
|
||||
if batch.Code != http.StatusOK {
|
||||
t.Fatalf("batch status = %d, body = %s", batch.Code, batch.Body.String())
|
||||
}
|
||||
var batchBody map[string]interface{}
|
||||
if err := json.Unmarshal(batch.Body.Bytes(), &batchBody); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if batchBody["created"] != float64(2) {
|
||||
t.Fatalf("batch response = %#v, want created=2", batchBody)
|
||||
}
|
||||
|
||||
invalid := performRBACJSONRequest(t, router, map[string]interface{}{
|
||||
"user_id": user.ID, "resource_type": "project", "resource_ids": []string{p3.ID, "missing"},
|
||||
})
|
||||
if invalid.Code != http.StatusBadRequest {
|
||||
t.Fatalf("invalid status = %d, body = %s", invalid.Code, invalid.Body.String())
|
||||
}
|
||||
rows, err := db.ListRBACResourceAssignments(user.ID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(rows) != 2 {
|
||||
t.Fatalf("failed batch persisted partial data: %#v", rows)
|
||||
}
|
||||
|
||||
legacy := performRBACJSONRequest(t, router, map[string]interface{}{
|
||||
"user_id": user.ID, "resource_type": "project", "resource_id": p3.ID,
|
||||
})
|
||||
if legacy.Code != http.StatusOK {
|
||||
t.Fatalf("legacy status = %d, body = %s", legacy.Code, legacy.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestRBACAssignableResourcesArePaged(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "rbac-picker.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
for _, name := range []string{"p1", "p2", "p3"} {
|
||||
if _, err := db.CreateProject(&database.Project{Name: name}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
h := NewRBACHandler(db, zap.NewNop())
|
||||
router := gin.New()
|
||||
router.GET("/api/rbac/resources", h.ListAssignableResources)
|
||||
|
||||
request := httptest.NewRequest(http.MethodGet, "/api/rbac/resources?type=project&limit=2&offset=0", nil)
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, request)
|
||||
if recorder.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, body = %s", recorder.Code, recorder.Body.String())
|
||||
}
|
||||
var body struct {
|
||||
Resources []database.RBACResourceOption `json:"resources"`
|
||||
HasMore bool `json:"has_more"`
|
||||
}
|
||||
if err := json.Unmarshal(recorder.Body.Bytes(), &body); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(body.Resources) != 2 || !body.HasMore {
|
||||
t.Fatalf("page = %#v, has_more = %v; want two rows and another page", body.Resources, body.HasMore)
|
||||
}
|
||||
}
|
||||
|
||||
func performRBACJSONRequest(t *testing.T, router http.Handler, payload map[string]interface{}) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
body, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/rbac/resource-assignments", bytes.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, req)
|
||||
return recorder
|
||||
}
|
||||
+450
-44
@@ -7,6 +7,8 @@ import (
|
||||
"crypto/cipher"
|
||||
"crypto/rand"
|
||||
"crypto/sha1"
|
||||
"crypto/sha256"
|
||||
"encoding/base32"
|
||||
"encoding/base64"
|
||||
"encoding/binary"
|
||||
"encoding/json"
|
||||
@@ -16,37 +18,45 @@ import (
|
||||
"io"
|
||||
"net/http"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const (
|
||||
robotCmdHelp = "帮助"
|
||||
robotCmdList = "列表"
|
||||
robotCmdListAlt = "对话列表"
|
||||
robotCmdSwitch = "切换"
|
||||
robotCmdContinue = "继续"
|
||||
robotCmdNew = "新对话"
|
||||
robotCmdClear = "清空"
|
||||
robotCmdCurrent = "当前"
|
||||
robotCmdStop = "停止"
|
||||
robotCmdRoles = "角色"
|
||||
robotCmdRolesList = "角色列表"
|
||||
robotCmdSwitchRole = "切换角色"
|
||||
robotCmdDelete = "删除"
|
||||
robotCmdVersion = "版本"
|
||||
robotCmdProjects = "项目"
|
||||
robotCmdProjectsList = "项目列表"
|
||||
robotCmdBindProject = "绑定项目"
|
||||
robotCmdNewProject = "新建项目"
|
||||
robotCmdHelp = "帮助"
|
||||
robotCmdList = "列表"
|
||||
robotCmdListAlt = "对话列表"
|
||||
robotCmdSwitch = "切换"
|
||||
robotCmdContinue = "继续"
|
||||
robotCmdNew = "新对话"
|
||||
robotCmdClear = "清空"
|
||||
robotCmdCurrent = "当前"
|
||||
robotCmdStop = "停止"
|
||||
robotCmdRoles = "角色"
|
||||
robotCmdRolesList = "角色列表"
|
||||
robotCmdSwitchRole = "切换角色"
|
||||
robotCmdDelete = "删除"
|
||||
robotCmdVersion = "版本"
|
||||
robotCmdProjects = "项目"
|
||||
robotCmdProjectsList = "项目列表"
|
||||
robotCmdBindProject = "绑定项目"
|
||||
robotCmdNewProject = "新建项目"
|
||||
robotCmdUnbindProject = "解除项目"
|
||||
robotCmdBindUser = "绑定"
|
||||
robotCmdUnbindUser = "解绑"
|
||||
robotCmdIdentity = "身份"
|
||||
robotBindingCodeTTL = 5 * time.Minute
|
||||
)
|
||||
|
||||
// RobotHandler 企业微信/钉钉/飞书等机器人回调处理
|
||||
@@ -60,6 +70,8 @@ type RobotHandler struct {
|
||||
sessionRoles map[string]string // key: "platform_userID", value: roleName(默认"默认")
|
||||
cancelMu sync.Mutex // 保护 runningCancels
|
||||
runningCancels map[string]context.CancelFunc // key: "platform_userID", 用于停止命令中断任务
|
||||
wecomReplay map[string]time.Time
|
||||
audit *audit.Service
|
||||
}
|
||||
|
||||
// NewRobotHandler 创建机器人处理器
|
||||
@@ -72,14 +84,97 @@ func NewRobotHandler(cfg *config.Config, db *database.DB, agentHandler *AgentHan
|
||||
sessions: make(map[string]string),
|
||||
sessionRoles: make(map[string]string),
|
||||
runningCancels: make(map[string]context.CancelFunc),
|
||||
wecomReplay: make(map[string]time.Time),
|
||||
}
|
||||
}
|
||||
|
||||
func (h *RobotHandler) SetAudit(s *audit.Service) {
|
||||
h.audit = s
|
||||
}
|
||||
|
||||
func (h *RobotHandler) acceptFreshWecomRequest(timestamp, nonce, signature string) bool {
|
||||
unixSeconds, err := strconv.ParseInt(strings.TrimSpace(timestamp), 10, 64)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
now := time.Now()
|
||||
requestTime := time.Unix(unixSeconds, 0)
|
||||
if requestTime.Before(now.Add(-5*time.Minute)) || requestTime.After(now.Add(5*time.Minute)) {
|
||||
return false
|
||||
}
|
||||
key := strings.TrimSpace(timestamp) + "\x00" + strings.TrimSpace(nonce) + "\x00" + strings.TrimSpace(signature)
|
||||
if strings.TrimSpace(nonce) == "" || strings.TrimSpace(signature) == "" {
|
||||
return false
|
||||
}
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
for replayKey, seenAt := range h.wecomReplay {
|
||||
if now.Sub(seenAt) > 10*time.Minute {
|
||||
delete(h.wecomReplay, replayKey)
|
||||
}
|
||||
}
|
||||
if _, exists := h.wecomReplay[key]; exists {
|
||||
return false
|
||||
}
|
||||
h.wecomReplay[key] = now
|
||||
return true
|
||||
}
|
||||
|
||||
// sessionKey 生成会话 key
|
||||
func (h *RobotHandler) sessionKey(platform, userID string) string {
|
||||
return platform + "_" + userID
|
||||
}
|
||||
|
||||
func normalizeRobotBindingCode(code string) string {
|
||||
return strings.ToUpper(strings.ReplaceAll(strings.TrimSpace(code), "-", ""))
|
||||
}
|
||||
|
||||
func hashRobotBindingCode(code string) string {
|
||||
sum := sha256.Sum256([]byte(normalizeRobotBindingCode(code)))
|
||||
return fmt.Sprintf("%x", sum[:])
|
||||
}
|
||||
|
||||
func (h *RobotHandler) resolveRobotAccess(platform, userID string) (*database.RBACAccess, error) {
|
||||
if h.db == nil {
|
||||
return nil, fmt.Errorf("机器人鉴权服务不可用")
|
||||
}
|
||||
authorization := h.config.Robots.AuthorizationFor(platform)
|
||||
var access *database.RBACAccess
|
||||
var err error
|
||||
switch authorization.EffectiveMode() {
|
||||
case config.RobotAuthModeUserBinding:
|
||||
access, err = h.db.ResolveRobotRBACAccess(platform, userID)
|
||||
case config.RobotAuthModeServiceAccount:
|
||||
if !authorization.ExternalUserAllowed(userID) {
|
||||
return nil, fmt.Errorf("机器人发送者不在服务账号白名单中")
|
||||
}
|
||||
access, err = h.db.ResolveRBACAccess(strings.TrimSpace(authorization.ServiceUserID))
|
||||
default:
|
||||
return nil, fmt.Errorf("机器人鉴权模式无效")
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !access.User.Enabled {
|
||||
return nil, fmt.Errorf("绑定的平台账号已被禁用")
|
||||
}
|
||||
return access, nil
|
||||
}
|
||||
|
||||
func robotPrincipal(access *database.RBACAccess) authctx.Principal {
|
||||
if access == nil {
|
||||
return authctx.Principal{}
|
||||
}
|
||||
return authctx.NewPrincipalWithScopes(access.User.ID, access.User.Username, access.Scope, access.Permissions, access.PermissionScopes)
|
||||
}
|
||||
|
||||
func (h *RobotHandler) robotAccessDeniedMessage(platform string) string {
|
||||
if h.config.Robots.AuthorizationFor(platform).EffectiveMode() == config.RobotAuthModeServiceAccount {
|
||||
return "当前平台账号不在该机器人的服务账号白名单中,或服务账号不可用。"
|
||||
}
|
||||
return "当前平台账号尚未绑定 CyberStrikeAI 用户。请先在网页端生成绑定码,然后发送:绑定 XXXX-XXXX"
|
||||
}
|
||||
|
||||
func (h *RobotHandler) loadSessionBinding(sk string) (convID, role string) {
|
||||
if h.db == nil || strings.TrimSpace(sk) == "" {
|
||||
return "", ""
|
||||
@@ -114,23 +209,29 @@ func (h *RobotHandler) deleteSessionBinding(sk string) {
|
||||
}
|
||||
|
||||
// getOrCreateConversation 获取或创建当前会话,title 用于新对话的标题(取用户首条消息前50字)
|
||||
func (h *RobotHandler) getOrCreateConversation(platform, userID, title string) (convID string, isNew bool) {
|
||||
func (h *RobotHandler) getOrCreateConversation(platform, userID, title string, access *database.RBACAccess) (convID string, isNew bool) {
|
||||
sk := h.sessionKey(platform, userID)
|
||||
h.mu.RLock()
|
||||
convID = h.sessions[sk]
|
||||
h.mu.RUnlock()
|
||||
if convID != "" {
|
||||
ownerID := access.User.ID
|
||||
readScope := robotPrincipal(access).ScopeFor("chat:read")
|
||||
if convID != "" && access.Permissions["chat:read"] && h.db.UserCanAccessResource(ownerID, readScope, "conversation", convID) {
|
||||
return convID, false
|
||||
}
|
||||
if persistedConvID, persistedRole := h.loadSessionBinding(sk); strings.TrimSpace(persistedConvID) != "" {
|
||||
// 会话绑定持久化:服务重启后也可恢复当前对话和角色。
|
||||
h.mu.Lock()
|
||||
h.sessions[sk] = persistedConvID
|
||||
if strings.TrimSpace(persistedRole) != "" {
|
||||
h.sessionRoles[sk] = persistedRole
|
||||
if !access.Permissions["chat:read"] || !h.db.UserCanAccessResource(ownerID, readScope, "conversation", persistedConvID) {
|
||||
h.deleteSessionBinding(sk)
|
||||
} else {
|
||||
// 会话绑定持久化:服务重启后也可恢复当前对话和角色。
|
||||
h.mu.Lock()
|
||||
h.sessions[sk] = persistedConvID
|
||||
if strings.TrimSpace(persistedRole) != "" {
|
||||
h.sessionRoles[sk] = persistedRole
|
||||
}
|
||||
h.mu.Unlock()
|
||||
return persistedConvID, false
|
||||
}
|
||||
h.mu.Unlock()
|
||||
return persistedConvID, false
|
||||
}
|
||||
t := strings.TrimSpace(title)
|
||||
if t == "" {
|
||||
@@ -139,13 +240,20 @@ func (h *RobotHandler) getOrCreateConversation(platform, userID, title string) (
|
||||
t = safeTruncateString(t, 50)
|
||||
}
|
||||
meta := database.ConversationCreateMeta{Source: "robot:" + platform}
|
||||
if !access.Permissions["chat:write"] {
|
||||
return "", false
|
||||
}
|
||||
meta.ProjectID = effectiveProjectID(h.config, "")
|
||||
if meta.ProjectID != "" && (!access.Permissions["project:read"] || !h.db.UserCanAccessResource(ownerID, robotPrincipal(access).ScopeFor("project:read"), "project", meta.ProjectID)) {
|
||||
meta.ProjectID = ""
|
||||
}
|
||||
conv, err := h.db.CreateConversation(t, meta)
|
||||
if err != nil {
|
||||
h.logger.Warn("创建机器人会话失败", zap.Error(err))
|
||||
return "", false
|
||||
}
|
||||
convID = conv.ID
|
||||
_ = h.db.SetResourceOwner("conversation", convID, ownerID)
|
||||
h.mu.Lock()
|
||||
role := h.sessionRoles[sk]
|
||||
h.sessions[sk] = convID
|
||||
@@ -193,15 +301,20 @@ func (h *RobotHandler) setRole(platform, userID, roleName string) {
|
||||
}
|
||||
|
||||
// clearConversation 清空当前会话(切换到新对话)
|
||||
func (h *RobotHandler) clearConversation(platform, userID string) (newConvID string) {
|
||||
func (h *RobotHandler) clearConversation(platform, userID string, access *database.RBACAccess) (newConvID string) {
|
||||
title := "新对话 " + time.Now().Format("01-02 15:04")
|
||||
meta := database.ConversationCreateMeta{Source: "robot:" + platform + ":new"}
|
||||
meta.ProjectID = effectiveProjectID(h.config, "")
|
||||
ownerID := access.User.ID
|
||||
if meta.ProjectID != "" && (!access.Permissions["project:read"] || !h.db.UserCanAccessResource(ownerID, robotPrincipal(access).ScopeFor("project:read"), "project", meta.ProjectID)) {
|
||||
meta.ProjectID = ""
|
||||
}
|
||||
conv, err := h.db.CreateConversation(title, meta)
|
||||
if err != nil {
|
||||
h.logger.Warn("创建新对话失败", zap.Error(err))
|
||||
return ""
|
||||
}
|
||||
_ = h.db.SetResourceOwner("conversation", conv.ID, ownerID)
|
||||
h.setConversation(platform, userID, conv.ID)
|
||||
return conv.ID
|
||||
}
|
||||
@@ -226,9 +339,24 @@ func (h *RobotHandler) HandleMessage(platform, userID, text string) (reply strin
|
||||
if cmdReply, ok := h.handleRobotCommand(platform, userID, text); ok {
|
||||
return cmdReply
|
||||
}
|
||||
access, err := h.resolveRobotAccess(platform, userID)
|
||||
if err != nil {
|
||||
return h.robotAccessDeniedMessage(platform)
|
||||
}
|
||||
if !access.Permissions["agent:execute"] || !access.Permissions["chat:read"] || !access.Permissions["chat:write"] {
|
||||
return "权限不足:机器人对话需要 agent:execute、chat:read 和 chat:write 权限。"
|
||||
}
|
||||
if h.audit != nil && h.config.Robots.AuthorizationFor(platform).EffectiveMode() == config.RobotAuthModeServiceAccount {
|
||||
hint := sha256.Sum256([]byte(userID))
|
||||
h.audit.RecordSystem(audit.Entry{
|
||||
Category: "robot", Action: "service_account_execute", Result: "success", Actor: access.User.Username,
|
||||
ResourceType: "robot_sender", ResourceID: platform + ":" + fmt.Sprintf("%x", hint[:4]),
|
||||
Message: "白名单平台发送者使用机器人服务账号执行 Agent",
|
||||
})
|
||||
}
|
||||
|
||||
// 普通消息:走 Agent
|
||||
convID, _ := h.getOrCreateConversation(platform, userID, text)
|
||||
convID, _ := h.getOrCreateConversation(platform, userID, text, access)
|
||||
if convID == "" {
|
||||
return "无法创建或获取对话,请稍后再试。"
|
||||
}
|
||||
@@ -251,7 +379,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, platform, convID, text, role)
|
||||
resp, newConvID, err := h.agentHandler.ProcessMessageForRobot(ctx, platform, robotPrincipal(access), 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) {
|
||||
@@ -279,6 +407,9 @@ func (h *RobotHandler) cmdHelp() string {
|
||||
b.WriteString("【通用 General】\n")
|
||||
b.WriteString("· 帮助 / help — 显示本帮助\n")
|
||||
b.WriteString("· 版本 / version — 显示当前版本号\n")
|
||||
b.WriteString("· 绑定 <绑定码> / bind <code> — 绑定网页端 RBAC 用户\n")
|
||||
b.WriteString("· 解绑 / unbind — 解除当前平台账号绑定\n")
|
||||
b.WriteString("· 身份 / whoami — 显示平台发送者、鉴权模式及当前实际 RBAC 身份\n")
|
||||
b.WriteString("\n【对话 Conversation】\n")
|
||||
b.WriteString("· 列表 / list — 列出所有对话标题与 ID\n")
|
||||
b.WriteString("· 切换 <ID> / switch <ID> — 指定对话继续\n")
|
||||
@@ -306,15 +437,20 @@ func (h *RobotHandler) projectsEnabled() bool {
|
||||
return h.config != nil && h.config.Project.Enabled
|
||||
}
|
||||
|
||||
func (h *RobotHandler) resolveProjectByIDOrName(idOrName string) (*database.Project, string) {
|
||||
func (h *RobotHandler) resolveProjectByIDOrName(access *database.RBACAccess, idOrName string) (*database.Project, string) {
|
||||
idOrName = strings.TrimSpace(idOrName)
|
||||
if idOrName == "" {
|
||||
return nil, "请指定项目 ID 或名称,例如:绑定项目 xxx-xxx"
|
||||
}
|
||||
ownerID := access.User.ID
|
||||
scope := robotPrincipal(access).ScopeFor("project:read")
|
||||
if p, err := h.db.GetProject(idOrName); err == nil {
|
||||
return p, ""
|
||||
if h.db.UserCanAccessResource(ownerID, scope, "project", p.ID) {
|
||||
return p, ""
|
||||
}
|
||||
return nil, "项目不存在或无权访问。"
|
||||
}
|
||||
list, err := h.db.ListProjects("", "", 200, 0)
|
||||
list, err := h.db.ListProjectsForAccess("", "", 200, 0, ownerID, scope)
|
||||
if err != nil {
|
||||
return nil, "查询项目失败: " + err.Error()
|
||||
}
|
||||
@@ -349,11 +485,15 @@ func (h *RobotHandler) formatProjectLabel(projectID string) string {
|
||||
return projectID
|
||||
}
|
||||
|
||||
func (h *RobotHandler) cmdProjects() string {
|
||||
func (h *RobotHandler) cmdProjects(platform, userID string) string {
|
||||
if !h.projectsEnabled() {
|
||||
return "项目功能未启用(config.project.enabled)。"
|
||||
}
|
||||
list, err := h.db.ListProjects("", "", 50, 0)
|
||||
access, err := h.resolveRobotAccess(platform, userID)
|
||||
if err != nil {
|
||||
return "当前平台账号尚未绑定。"
|
||||
}
|
||||
list, err := h.db.ListProjectsForAccess("", "", 50, 0, access.User.ID, robotPrincipal(access).ScopeFor("project:read"))
|
||||
if err != nil {
|
||||
return "获取项目列表失败: " + err.Error()
|
||||
}
|
||||
@@ -380,11 +520,15 @@ func (h *RobotHandler) cmdBindProject(platform, userID, idOrName string) string
|
||||
if !h.projectsEnabled() {
|
||||
return "项目功能未启用(config.project.enabled)。"
|
||||
}
|
||||
p, errMsg := h.resolveProjectByIDOrName(idOrName)
|
||||
access, err := h.resolveRobotAccess(platform, userID)
|
||||
if err != nil {
|
||||
return "当前平台账号尚未绑定。"
|
||||
}
|
||||
p, errMsg := h.resolveProjectByIDOrName(access, idOrName)
|
||||
if p == nil {
|
||||
return errMsg
|
||||
}
|
||||
convID, _ := h.getOrCreateConversation(platform, userID, "")
|
||||
convID, _ := h.getOrCreateConversation(platform, userID, "", access)
|
||||
if convID == "" {
|
||||
return "无法获取当前对话,请稍后再试。"
|
||||
}
|
||||
@@ -402,12 +546,17 @@ func (h *RobotHandler) cmdNewProject(platform, userID, name string) string {
|
||||
if name == "" {
|
||||
return "请指定项目名称,例如:新建项目 某目标渗透"
|
||||
}
|
||||
access, accessErr := h.resolveRobotAccess(platform, userID)
|
||||
if accessErr != nil {
|
||||
return "当前平台账号尚未绑定。"
|
||||
}
|
||||
p := &database.Project{Name: name, Status: "active"}
|
||||
created, err := h.db.CreateProject(p)
|
||||
if err != nil {
|
||||
return "创建项目失败: " + err.Error()
|
||||
}
|
||||
convID, _ := h.getOrCreateConversation(platform, userID, name)
|
||||
_ = h.db.SetResourceOwner("project", created.ID, access.User.ID)
|
||||
convID, _ := h.getOrCreateConversation(platform, userID, name, access)
|
||||
if convID == "" {
|
||||
return fmt.Sprintf("项目已创建:「%s」\nID: %s\n(绑定当前对话失败,请手动发送「绑定项目 %s」)", created.Name, created.ID, created.ID)
|
||||
}
|
||||
@@ -430,6 +579,13 @@ func (h *RobotHandler) cmdUnbindProject(platform, userID string) string {
|
||||
convID = persistedConvID
|
||||
}
|
||||
}
|
||||
access, err := h.resolveRobotAccess(platform, userID)
|
||||
if err != nil {
|
||||
return "当前平台账号尚未绑定。"
|
||||
}
|
||||
if !h.db.UserCanAccessResource(access.User.ID, robotPrincipal(access).ScopeFor("chat:write"), "conversation", convID) {
|
||||
return "当前对话不存在或无权访问。"
|
||||
}
|
||||
if convID == "" {
|
||||
return "当前没有进行中的对话,无需解除绑定。"
|
||||
}
|
||||
@@ -446,8 +602,12 @@ func (h *RobotHandler) cmdUnbindProject(platform, userID string) string {
|
||||
return "已解除当前对话的项目绑定。"
|
||||
}
|
||||
|
||||
func (h *RobotHandler) cmdList() string {
|
||||
convs, err := h.db.ListConversations(50, 0, "", "", "")
|
||||
func (h *RobotHandler) cmdList(platform, userID string) string {
|
||||
access, err := h.resolveRobotAccess(platform, userID)
|
||||
if err != nil {
|
||||
return "当前平台账号尚未绑定。"
|
||||
}
|
||||
convs, err := h.db.ListConversationsForAccess(50, 0, "", "", "", access.User.ID, robotPrincipal(access).ScopeFor("chat:read"))
|
||||
if err != nil {
|
||||
return "获取对话列表失败: " + err.Error()
|
||||
}
|
||||
@@ -470,8 +630,12 @@ func (h *RobotHandler) cmdSwitch(platform, userID, convID string) string {
|
||||
if convID == "" {
|
||||
return "请指定对话 ID,例如:切换 xxx-xxx-xxx"
|
||||
}
|
||||
access, accessErr := h.resolveRobotAccess(platform, userID)
|
||||
if accessErr != nil {
|
||||
return "当前平台账号尚未绑定。"
|
||||
}
|
||||
conv, err := h.db.GetConversation(convID)
|
||||
if err != nil {
|
||||
if err != nil || !h.db.UserCanAccessResource(access.User.ID, robotPrincipal(access).ScopeFor("chat:read"), "conversation", convID) {
|
||||
return "对话不存在或 ID 错误。"
|
||||
}
|
||||
h.setConversation(platform, userID, conv.ID)
|
||||
@@ -479,7 +643,11 @@ func (h *RobotHandler) cmdSwitch(platform, userID, convID string) string {
|
||||
}
|
||||
|
||||
func (h *RobotHandler) cmdNew(platform, userID string) string {
|
||||
newID := h.clearConversation(platform, userID)
|
||||
access, err := h.resolveRobotAccess(platform, userID)
|
||||
if err != nil {
|
||||
return "当前平台账号尚未绑定。"
|
||||
}
|
||||
newID := h.clearConversation(platform, userID, access)
|
||||
if newID == "" {
|
||||
return "创建新对话失败,请重试。"
|
||||
}
|
||||
@@ -512,6 +680,13 @@ func (h *RobotHandler) cmdCurrent(platform, userID string) string {
|
||||
if convID == "" {
|
||||
return "当前没有进行中的对话。发送任意内容将创建新对话。"
|
||||
}
|
||||
access, err := h.resolveRobotAccess(platform, userID)
|
||||
if err != nil {
|
||||
return "当前平台账号尚未绑定。"
|
||||
}
|
||||
if !h.db.UserCanAccessResource(access.User.ID, robotPrincipal(access).ScopeFor("chat:read"), "conversation", convID) {
|
||||
return "当前对话不存在或无权访问。"
|
||||
}
|
||||
conv, err := h.db.GetConversation(convID)
|
||||
if err != nil {
|
||||
return "当前对话 ID: " + convID + "(获取标题失败)"
|
||||
@@ -582,6 +757,13 @@ func (h *RobotHandler) cmdDelete(platform, userID, convID string) string {
|
||||
if convID == "" {
|
||||
return "请指定对话 ID,例如:删除 xxx-xxx-xxx"
|
||||
}
|
||||
access, err := h.resolveRobotAccess(platform, userID)
|
||||
if err != nil {
|
||||
return "当前平台账号尚未绑定。"
|
||||
}
|
||||
if !h.db.UserCanAccessResource(access.User.ID, robotPrincipal(access).ScopeFor("chat:delete"), "conversation", convID) {
|
||||
return "对话不存在或无权访问。"
|
||||
}
|
||||
sk := h.sessionKey(platform, userID)
|
||||
h.mu.RLock()
|
||||
currentConvID := h.sessions[sk]
|
||||
@@ -611,13 +793,164 @@ func (h *RobotHandler) cmdVersion() string {
|
||||
return "CyberStrikeAI " + v
|
||||
}
|
||||
|
||||
func (h *RobotHandler) cmdIdentity(platform, userID string) string {
|
||||
authorization := h.config.Robots.AuthorizationFor(platform)
|
||||
mode := authorization.EffectiveMode()
|
||||
modeLabel := "逐用户绑定(user_binding)"
|
||||
if mode == config.RobotAuthModeServiceAccount {
|
||||
modeLabel = "专用服务账号(service_account)"
|
||||
}
|
||||
var b strings.Builder
|
||||
b.WriteString("【机器人身份】\n")
|
||||
b.WriteString("平台:" + platform + "\n")
|
||||
b.WriteString("发送者 ID:" + userID + "\n")
|
||||
b.WriteString("鉴权模式:" + modeLabel + "\n")
|
||||
|
||||
access, err := h.resolveRobotAccess(platform, userID)
|
||||
if err != nil {
|
||||
if mode == config.RobotAuthModeServiceAccount {
|
||||
b.WriteString("鉴权状态:拒绝(发送者不在白名单中,或服务账号不可用)")
|
||||
} else {
|
||||
b.WriteString("鉴权状态:未绑定\n")
|
||||
b.WriteString("操作提示:请在 Web 端生成绑定码,然后发送“绑定 XXXX-XXXX”")
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
name := strings.TrimSpace(access.User.DisplayName)
|
||||
if name == "" {
|
||||
name = access.User.Username
|
||||
}
|
||||
roleNames := make([]string, 0, len(access.Roles))
|
||||
for _, role := range access.Roles {
|
||||
roleNames = append(roleNames, role.Name)
|
||||
}
|
||||
if len(roleNames) == 0 {
|
||||
roleNames = append(roleNames, "未分配角色")
|
||||
}
|
||||
b.WriteString("鉴权状态:已授权\n")
|
||||
b.WriteString("实际身份:" + name + " (" + access.User.Username + ")\n")
|
||||
b.WriteString("RBAC User ID:" + access.User.ID + "\n")
|
||||
b.WriteString("平台角色:" + strings.Join(roleNames, "、") + "\n")
|
||||
b.WriteString("资源范围:" + access.Scope + "\n")
|
||||
b.WriteString(fmt.Sprintf("有效权限:%d 项", len(access.Permissions)))
|
||||
return b.String()
|
||||
}
|
||||
|
||||
func robotCommandPermission(text string) (string, bool) {
|
||||
switch {
|
||||
case text == robotCmdHelp || text == "help" || text == "?" || text == "?", text == robotCmdVersion || text == "version", text == robotCmdIdentity || text == "whoami":
|
||||
return "", true
|
||||
case text == robotCmdList || text == robotCmdListAlt || text == "list",
|
||||
strings.HasPrefix(text, robotCmdSwitch+" "), strings.HasPrefix(text, robotCmdContinue+" "),
|
||||
strings.HasPrefix(text, "switch "), strings.HasPrefix(text, "continue "),
|
||||
text == robotCmdCurrent || text == "current":
|
||||
return "chat:read", true
|
||||
case text == robotCmdNew || text == "new", text == robotCmdClear || text == "clear":
|
||||
return "chat:write", true
|
||||
case strings.HasPrefix(text, robotCmdDelete+" "), strings.HasPrefix(text, "delete "):
|
||||
return "chat:delete", true
|
||||
case text == robotCmdStop || text == "stop":
|
||||
return "agent:execute", true
|
||||
case text == robotCmdRoles || text == robotCmdRolesList || text == "roles",
|
||||
strings.HasPrefix(text, robotCmdRoles+" "), strings.HasPrefix(text, robotCmdSwitchRole+" "), strings.HasPrefix(text, "role "):
|
||||
return "roles:read", true
|
||||
case text == robotCmdProjects || text == robotCmdProjectsList || text == "projects":
|
||||
return "project:read", true
|
||||
case text == robotCmdUnbindProject || text == "unbind project",
|
||||
strings.HasPrefix(text, robotCmdNewProject+" "), strings.HasPrefix(text, "new project "),
|
||||
strings.HasPrefix(text, robotCmdBindProject+" "), strings.HasPrefix(text, "bind project "):
|
||||
return "project:write", true
|
||||
default:
|
||||
return "", false
|
||||
}
|
||||
}
|
||||
|
||||
func (h *RobotHandler) cmdBindUser(platform, userID, code string) string {
|
||||
if h.config.Robots.AuthorizationFor(platform).EffectiveMode() != config.RobotAuthModeUserBinding {
|
||||
return "该机器人使用受控服务账号模式,不接受用户绑定。"
|
||||
}
|
||||
code = normalizeRobotBindingCode(code)
|
||||
if code == "" {
|
||||
return "请提供绑定码,例如:绑定 ABCD-1234"
|
||||
}
|
||||
user, err := h.db.ConsumeRobotBindingCode(platform, userID, hashRobotBindingCode(code))
|
||||
if err != nil {
|
||||
return "绑定失败:绑定码无效、已使用或已过期。请在网页端重新生成。"
|
||||
}
|
||||
// Never carry an old synthetic-owner conversation into the RBAC identity.
|
||||
sk := h.sessionKey(platform, userID)
|
||||
h.mu.Lock()
|
||||
delete(h.sessions, sk)
|
||||
delete(h.sessionRoles, sk)
|
||||
h.mu.Unlock()
|
||||
h.deleteSessionBinding(sk)
|
||||
name := strings.TrimSpace(user.DisplayName)
|
||||
if name == "" {
|
||||
name = user.Username
|
||||
}
|
||||
if h.audit != nil {
|
||||
hint := sha256.Sum256([]byte(userID))
|
||||
h.audit.RecordSystem(audit.Entry{
|
||||
Category: "auth", Action: "robot_bind", Result: "success", Actor: user.Username,
|
||||
ResourceType: "robot_binding", ResourceID: platform + ":" + fmt.Sprintf("%x", hint[:4]), Message: "机器人平台账号绑定成功",
|
||||
})
|
||||
}
|
||||
return fmt.Sprintf("绑定成功,当前身份:%s。后续操作将实时使用该用户的 RBAC 权限。", name)
|
||||
}
|
||||
|
||||
func (h *RobotHandler) cmdUnbindUser(platform, userID string) string {
|
||||
if h.config.Robots.AuthorizationFor(platform).EffectiveMode() != config.RobotAuthModeUserBinding {
|
||||
return "该机器人使用受控服务账号模式,无需用户解绑。"
|
||||
}
|
||||
access, accessErr := h.resolveRobotAccess(platform, userID)
|
||||
if accessErr != nil {
|
||||
return "当前平台账号尚未绑定。"
|
||||
}
|
||||
if err := h.db.DeleteRobotIdentityBinding(platform, userID); err != nil {
|
||||
return "解绑失败,请稍后重试。"
|
||||
}
|
||||
sk := h.sessionKey(platform, userID)
|
||||
h.mu.Lock()
|
||||
delete(h.sessions, sk)
|
||||
delete(h.sessionRoles, sk)
|
||||
h.mu.Unlock()
|
||||
h.deleteSessionBinding(sk)
|
||||
if h.audit != nil {
|
||||
hint := sha256.Sum256([]byte(userID))
|
||||
h.audit.RecordSystem(audit.Entry{
|
||||
Category: "auth", Action: "robot_unbind", Result: "success", Actor: access.User.Username,
|
||||
ResourceType: "robot_binding", ResourceID: platform + ":" + fmt.Sprintf("%x", hint[:4]), Message: "机器人平台账号解绑成功",
|
||||
})
|
||||
}
|
||||
return "已解除当前平台账号与 CyberStrikeAI 用户的绑定。"
|
||||
}
|
||||
|
||||
// handleRobotCommand 处理机器人内置命令;若匹配到命令返回 (回复内容, true),否则返回 ("", false)
|
||||
func (h *RobotHandler) handleRobotCommand(platform, userID, text string) (string, bool) {
|
||||
if (strings.HasPrefix(text, robotCmdBindUser+" ") || strings.HasPrefix(text, "bind ")) && !strings.HasPrefix(text, "bind project ") {
|
||||
parts := strings.SplitN(text, " ", 2)
|
||||
return h.cmdBindUser(platform, userID, strings.TrimSpace(parts[1])), true
|
||||
}
|
||||
if text == robotCmdUnbindUser || text == "unbind" {
|
||||
return h.cmdUnbindUser(platform, userID), true
|
||||
}
|
||||
if permission, recognized := robotCommandPermission(text); recognized && permission != "" {
|
||||
access, err := h.resolveRobotAccess(platform, userID)
|
||||
if err != nil {
|
||||
return h.robotAccessDeniedMessage(platform), true
|
||||
}
|
||||
if !access.Permissions[permission] {
|
||||
return fmt.Sprintf("权限不足:缺少 %s 权限。", permission), true
|
||||
}
|
||||
}
|
||||
switch {
|
||||
case text == robotCmdHelp || text == "help" || text == "?" || text == "?":
|
||||
return h.cmdHelp(), true
|
||||
case text == robotCmdIdentity || text == "whoami":
|
||||
return h.cmdIdentity(platform, userID), true
|
||||
case text == robotCmdList || text == robotCmdListAlt || text == "list":
|
||||
return h.cmdList(), true
|
||||
return h.cmdList(platform, userID), true
|
||||
case strings.HasPrefix(text, robotCmdSwitch+" ") || strings.HasPrefix(text, robotCmdContinue+" ") || strings.HasPrefix(text, "switch ") || strings.HasPrefix(text, "continue "):
|
||||
var id string
|
||||
switch {
|
||||
@@ -663,7 +996,7 @@ func (h *RobotHandler) handleRobotCommand(platform, userID, text string) (string
|
||||
case text == robotCmdVersion || text == "version":
|
||||
return h.cmdVersion(), true
|
||||
case text == robotCmdProjects || text == robotCmdProjectsList || text == "projects":
|
||||
return h.cmdProjects(), true
|
||||
return h.cmdProjects(platform, userID), true
|
||||
case text == robotCmdUnbindProject || text == "unbind project":
|
||||
return h.cmdUnbindProject(platform, userID), true
|
||||
case strings.HasPrefix(text, robotCmdNewProject+" ") || strings.HasPrefix(text, "new project "):
|
||||
@@ -903,6 +1236,11 @@ func (h *RobotHandler) HandleWecomPOST(c *gin.Context) {
|
||||
c.String(http.StatusOK, "")
|
||||
return
|
||||
}
|
||||
if !h.acceptFreshWecomRequest(timestamp, nonce, msgSignature) {
|
||||
h.logger.Warn("企业微信 POST 时间戳过期或请求重放,已拒绝")
|
||||
c.String(http.StatusOK, "")
|
||||
return
|
||||
}
|
||||
|
||||
var body wecomXML
|
||||
if err := xml.Unmarshal(bodyRaw, &body); err != nil {
|
||||
@@ -1079,6 +1417,74 @@ func (h *RobotHandler) sendWecomReply(c *gin.Context, toUser, fromUser, content,
|
||||
|
||||
// —————— 测试接口(需登录,用于验证机器人逻辑,无需钉钉/飞书客户端) ——————
|
||||
|
||||
// CreateRobotBindingCode creates a short-lived, single-use secret for the
|
||||
// currently authenticated RBAC user. Only its hash is persisted.
|
||||
func (h *RobotHandler) CreateRobotBindingCode(c *gin.Context) {
|
||||
session, ok := security.CurrentSession(c)
|
||||
if !ok || strings.TrimSpace(session.UserID) == "" {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"error": "未授权访问"})
|
||||
return
|
||||
}
|
||||
random := make([]byte, 5)
|
||||
if _, err := rand.Read(random); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "生成绑定码失败"})
|
||||
return
|
||||
}
|
||||
raw := base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(random)
|
||||
code := raw[:4] + "-" + raw[4:]
|
||||
expiresAt := time.Now().Add(robotBindingCodeTTL)
|
||||
if err := h.db.CreateRobotBindingCode(session.UserID, hashRobotBindingCode(code), expiresAt); err != nil {
|
||||
h.logger.Warn("创建机器人绑定码失败", zap.String("user_id", session.UserID), zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "生成绑定码失败"})
|
||||
return
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.Record(c, audit.Entry{Category: "auth", Action: "robot_binding_code_create", Result: "success", ResourceType: "user", ResourceID: session.UserID, Message: "生成机器人一次性绑定码"})
|
||||
}
|
||||
c.Header("Cache-Control", "no-store")
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"code": code, "expires_at": expiresAt.UTC().Format(time.RFC3339), "expires_in_seconds": int(robotBindingCodeTTL.Seconds()),
|
||||
})
|
||||
}
|
||||
|
||||
func (h *RobotHandler) ListMyRobotBindings(c *gin.Context) {
|
||||
session, ok := security.CurrentSession(c)
|
||||
if !ok {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"error": "未授权访问"})
|
||||
return
|
||||
}
|
||||
bindings, err := h.db.ListRobotUserBindings(session.UserID)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "获取机器人绑定失败"})
|
||||
return
|
||||
}
|
||||
items := make([]gin.H, 0, len(bindings))
|
||||
for _, binding := range bindings {
|
||||
sum := sha256.Sum256([]byte(binding.ExternalUserID))
|
||||
items = append(items, gin.H{
|
||||
"id": binding.ID, "platform": binding.Platform, "external_user_hint": fmt.Sprintf("%x", sum[:4]),
|
||||
"enabled": binding.Enabled, "created_at": binding.CreatedAt, "updated_at": binding.UpdatedAt,
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"bindings": items})
|
||||
}
|
||||
|
||||
func (h *RobotHandler) DeleteMyRobotBinding(c *gin.Context) {
|
||||
session, ok := security.CurrentSession(c)
|
||||
if !ok {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"error": "未授权访问"})
|
||||
return
|
||||
}
|
||||
if err := h.db.DeleteRobotUserBindingForUser(c.Param("id"), session.UserID); err != nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "绑定不存在"})
|
||||
return
|
||||
}
|
||||
if h.audit != nil {
|
||||
h.audit.Record(c, audit.Entry{Category: "auth", Action: "robot_binding_revoke", Result: "success", ResourceType: "robot_binding", ResourceID: c.Param("id"), Message: "撤销机器人平台账号绑定"})
|
||||
}
|
||||
c.Status(http.StatusNoContent)
|
||||
}
|
||||
|
||||
// RobotTestRequest 模拟机器人消息请求
|
||||
type RobotTestRequest struct {
|
||||
Platform string `json:"platform"` // 如 "dingtalk"、"lark"、"wecom"
|
||||
|
||||
@@ -0,0 +1,138 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestRobotUsersAreResourceIsolated(t *testing.T) {
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "robot-rbac.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
cfg := &config.Config{}
|
||||
cfg.Project.Enabled = true
|
||||
h := NewRobotHandler(cfg, db, nil, zap.NewNop())
|
||||
if err := db.BootstrapRBAC("hash", security.PermissionCatalog); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
alice, err := db.CreateRBACUser("robot-alice", "Robot Alice", "hash", true, []string{database.RBACSystemRoleOperator})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
bob, err := db.CreateRBACUser("robot-bob", "Robot Bob", "hash", true, []string{database.RBACSystemRoleOperator})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for externalID, user := range map[string]*database.RBACUser{"alice": alice, "bob": bob} {
|
||||
code := "TEST-" + strings.ToUpper(externalID)
|
||||
if err := db.CreateRobotBindingCode(user.ID, hashRobotBindingCode(code), time.Now().Add(time.Minute)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := db.ConsumeRobotBindingCode("wecom", externalID, hashRobotBindingCode(code)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
aliceAccess, err := db.ResolveRobotRBACAccess("wecom", "alice")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := h.HandleMessage("wecom", "alice", "身份"); !strings.Contains(got, "Robot Alice") || !strings.Contains(got, alice.ID) || !strings.Contains(got, "user_binding") {
|
||||
t.Fatalf("bound identity output is incomplete: %s", got)
|
||||
}
|
||||
|
||||
conversationID, _ := h.getOrCreateConversation("wecom", "alice", "alice conversation", aliceAccess)
|
||||
if conversationID == "" {
|
||||
t.Fatal("alice conversation was not created")
|
||||
}
|
||||
if got := h.cmdList("wecom", "bob"); strings.Contains(got, conversationID) {
|
||||
t.Fatalf("bob listed alice conversation: %s", got)
|
||||
}
|
||||
if got := h.cmdSwitch("wecom", "bob", conversationID); !strings.Contains(got, "不存在") && !strings.Contains(got, "无权访问") {
|
||||
t.Fatalf("bob switched to alice conversation: %s", got)
|
||||
}
|
||||
if got := h.cmdDelete("wecom", "bob", conversationID); !strings.Contains(got, "无权访问") {
|
||||
t.Fatalf("bob deleted alice conversation: %s", got)
|
||||
}
|
||||
if _, err := db.GetConversation(conversationID); err != nil {
|
||||
t.Fatalf("alice conversation was deleted: %v", err)
|
||||
}
|
||||
|
||||
createReply := h.cmdNewProject("wecom", "alice", "alice project")
|
||||
if !strings.Contains(createReply, "已创建项目") {
|
||||
t.Fatalf("create project reply: %s", createReply)
|
||||
}
|
||||
if got := h.cmdProjects("wecom", "bob"); strings.Contains(got, "alice project") {
|
||||
t.Fatalf("bob listed alice project: %s", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWecomReplayGuardRequiresFreshUniqueRequest(t *testing.T) {
|
||||
h := NewRobotHandler(&config.Config{}, nil, nil, zap.NewNop())
|
||||
timestamp := time.Now().Unix()
|
||||
if !h.acceptFreshWecomRequest(fmt.Sprintf("%d", timestamp), "nonce", "signature") {
|
||||
t.Fatal("fresh request was rejected")
|
||||
}
|
||||
if h.acceptFreshWecomRequest(fmt.Sprintf("%d", timestamp), "nonce", "signature") {
|
||||
t.Fatal("duplicate request was accepted")
|
||||
}
|
||||
if h.acceptFreshWecomRequest(fmt.Sprintf("%d", timestamp-600), "old", "signature") {
|
||||
t.Fatal("stale request was accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRobotServiceAccountRequiresExactSenderAllowlist(t *testing.T) {
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "robot-service.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
if err := db.BootstrapRBAC("hash", security.PermissionCatalog); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
serviceUser, err := db.CreateRBACUser("robot-service-user", "Robot Service", "hash", true, []string{database.RBACSystemRoleOperator})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg := &config.Config{}
|
||||
cfg.Robots.Lark.Auth = config.RobotAuthorizationConfig{
|
||||
Mode: config.RobotAuthModeServiceAccount, ServiceUserID: serviceUser.ID,
|
||||
AllowedExternalUsers: []string{"t:tenant|u:allowed"},
|
||||
}
|
||||
h := NewRobotHandler(cfg, db, nil, zap.NewNop())
|
||||
if got := h.HandleMessage("lark", "t:tenant|u:allowed", "列表"); strings.Contains(got, "白名单") || strings.Contains(got, "尚未绑定") {
|
||||
t.Fatalf("allowed service account sender was denied: %s", got)
|
||||
}
|
||||
if got := h.HandleMessage("lark", "t:tenant|u:denied", "列表"); !strings.Contains(got, "白名单") {
|
||||
t.Fatalf("non-allowlisted sender was not denied: %s", got)
|
||||
}
|
||||
if got := h.HandleMessage("lark", "t:tenant|u:allowed", "绑定 ABCD-1234"); !strings.Contains(got, "服务账号模式") {
|
||||
t.Fatalf("service-account robot accepted user binding: %s", got)
|
||||
}
|
||||
if got := h.HandleMessage("lark", "t:tenant|u:allowed", "whoami"); !strings.Contains(got, "Robot Service") || !strings.Contains(got, serviceUser.ID) || !strings.Contains(got, "service_account") {
|
||||
t.Fatalf("service-account identity output is incomplete: %s", got)
|
||||
}
|
||||
if got := h.HandleMessage("lark", "t:tenant|u:denied", "whoami"); !strings.Contains(got, "鉴权状态:拒绝") || strings.Contains(got, serviceUser.ID) {
|
||||
t.Fatalf("denied identity output leaked or omitted status: %s", got)
|
||||
}
|
||||
|
||||
adminCfg := &config.Config{}
|
||||
adminCfg.Robots.Lark.Auth = config.RobotAuthorizationConfig{
|
||||
Mode: config.RobotAuthModeServiceAccount, ServiceUserID: "admin",
|
||||
AllowedExternalUsers: []string{"t:tenant|u:owner"},
|
||||
}
|
||||
adminHandler := NewRobotHandler(adminCfg, db, nil, zap.NewNop())
|
||||
if got := adminHandler.HandleMessage("lark", "t:tenant|u:owner", "身份"); !strings.Contains(got, "admin") || !strings.Contains(got, "鉴权状态:已授权") {
|
||||
t.Fatalf("allowlisted admin service account was denied: %s", got)
|
||||
}
|
||||
}
|
||||
@@ -9,6 +9,7 @@ import (
|
||||
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/security"
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
@@ -60,6 +61,16 @@ func (h *VulnerabilityHandler) CreateVulnerability(c *gin.Context) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if session, ok := security.CurrentSession(c); ok && session.Scope != database.RBACScopeAll {
|
||||
if strings.TrimSpace(req.ConversationID) != "" && !h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", strings.TrimSpace(req.ConversationID)) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权在该对话下创建漏洞"})
|
||||
return
|
||||
}
|
||||
if strings.TrimSpace(req.ProjectID) != "" && !h.db.UserCanAccessResource(session.UserID, session.Scope, "project", strings.TrimSpace(req.ProjectID)) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权在该项目下创建漏洞"})
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
vuln := &database.Vulnerability{
|
||||
ConversationID: req.ConversationID,
|
||||
@@ -86,6 +97,10 @@ func (h *VulnerabilityHandler) CreateVulnerability(c *gin.Context) {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if session, ok := security.CurrentSession(c); ok {
|
||||
_ = h.db.SetResourceOwner("vulnerability", created.ID, session.UserID)
|
||||
_ = h.db.AssignResourceToUser(session.UserID, "vulnerability", created.ID)
|
||||
}
|
||||
|
||||
if h.audit != nil {
|
||||
h.audit.RecordOK(c, "vulnerability", "create", "创建漏洞记录", "vulnerability", created.ID, map[string]interface{}{
|
||||
@@ -142,6 +157,7 @@ func (h *VulnerabilityHandler) ListVulnerabilities(c *gin.Context) {
|
||||
offsetStr := c.DefaultQuery("offset", "0")
|
||||
pageStr := c.Query("page")
|
||||
filter := parseVulnerabilityListFilter(c)
|
||||
access := vulnerabilityAccessFromContext(c)
|
||||
|
||||
limit, _ := strconv.Atoi(limitStr)
|
||||
offset, _ := strconv.Atoi(offsetStr)
|
||||
@@ -163,7 +179,7 @@ func (h *VulnerabilityHandler) ListVulnerabilities(c *gin.Context) {
|
||||
}
|
||||
|
||||
// 获取总数
|
||||
total, err := h.db.CountVulnerabilities(filter)
|
||||
total, err := h.db.CountVulnerabilitiesForAccess(filter, access)
|
||||
if err != nil {
|
||||
h.logger.Error("获取漏洞总数失败", zap.Error(err))
|
||||
// 继续执行,使用0作为总数
|
||||
@@ -171,7 +187,7 @@ func (h *VulnerabilityHandler) ListVulnerabilities(c *gin.Context) {
|
||||
}
|
||||
|
||||
// 获取漏洞列表
|
||||
vulnerabilities, err := h.db.ListVulnerabilities(limit, offset, filter)
|
||||
vulnerabilities, err := h.db.ListVulnerabilitiesForAccess(limit, offset, filter, access)
|
||||
if err != nil {
|
||||
h.logger.Error("获取漏洞列表失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
@@ -238,7 +254,15 @@ func (h *VulnerabilityHandler) UpdateVulnerability(c *gin.Context) {
|
||||
|
||||
// 更新字段
|
||||
if req.ProjectID != nil {
|
||||
existing.ProjectID = strings.TrimSpace(*req.ProjectID)
|
||||
targetProjectID := strings.TrimSpace(*req.ProjectID)
|
||||
if targetProjectID != "" {
|
||||
session, ok := security.CurrentSession(c)
|
||||
if !ok || !h.db.UserCanAccessResource(session.UserID, session.Scope, "project", targetProjectID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权关联目标项目"})
|
||||
return
|
||||
}
|
||||
}
|
||||
existing.ProjectID = targetProjectID
|
||||
}
|
||||
if req.ConversationTag != nil {
|
||||
existing.ConversationTag = *req.ConversationTag
|
||||
@@ -332,8 +356,9 @@ func (h *VulnerabilityHandler) DeleteVulnerability(c *gin.Context) {
|
||||
// BatchDeleteVulnerabilities 按当前筛选条件批量删除漏洞
|
||||
func (h *VulnerabilityHandler) BatchDeleteVulnerabilities(c *gin.Context) {
|
||||
filter := parseVulnerabilityListFilter(c)
|
||||
access := vulnerabilityAccessFromContext(c)
|
||||
|
||||
total, err := h.db.CountVulnerabilities(filter)
|
||||
total, err := h.db.CountVulnerabilitiesForAccess(filter, access)
|
||||
if err != nil {
|
||||
h.logger.Error("统计待删除漏洞失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
@@ -344,7 +369,7 @@ func (h *VulnerabilityHandler) BatchDeleteVulnerabilities(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
deleted, err := h.db.DeleteVulnerabilitiesByFilter(filter)
|
||||
deleted, err := h.db.DeleteVulnerabilitiesByFilterForAccess(filter, access)
|
||||
if err != nil {
|
||||
h.logger.Error("批量删除漏洞失败", zap.Error(err), zap.Int("count", total))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
@@ -364,8 +389,9 @@ func (h *VulnerabilityHandler) BatchDeleteVulnerabilities(c *gin.Context) {
|
||||
// GetVulnerabilityStats 获取漏洞统计
|
||||
func (h *VulnerabilityHandler) GetVulnerabilityStats(c *gin.Context) {
|
||||
filter := parseVulnerabilityListFilter(c)
|
||||
access := vulnerabilityAccessFromContext(c)
|
||||
|
||||
stats, err := h.db.GetVulnerabilityStats(filter)
|
||||
stats, err := h.db.GetVulnerabilityStatsForAccess(filter, access)
|
||||
if err != nil {
|
||||
h.logger.Error("获取漏洞统计失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
@@ -377,7 +403,7 @@ func (h *VulnerabilityHandler) GetVulnerabilityStats(c *gin.Context) {
|
||||
|
||||
// GetVulnerabilityFilterOptions 获取漏洞筛选建议项
|
||||
func (h *VulnerabilityHandler) GetVulnerabilityFilterOptions(c *gin.Context) {
|
||||
options, err := h.db.GetVulnerabilityFilterOptions()
|
||||
options, err := h.db.GetVulnerabilityFilterOptionsForAccess(vulnerabilityAccessFromContext(c))
|
||||
if err != nil {
|
||||
h.logger.Error("获取漏洞筛选建议失败", zap.Error(err))
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
@@ -400,8 +426,9 @@ func (h *VulnerabilityHandler) ExportVulnerabilities(c *gin.Context) {
|
||||
}
|
||||
|
||||
filter := parseVulnerabilityListFilter(c)
|
||||
access := vulnerabilityAccessFromContext(c)
|
||||
|
||||
total, err := h.db.CountVulnerabilities(filter)
|
||||
total, err := h.db.CountVulnerabilitiesForAccess(filter, access)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
@@ -411,7 +438,7 @@ func (h *VulnerabilityHandler) ExportVulnerabilities(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
items, err := h.db.ListVulnerabilities(total, 0, filter)
|
||||
items, err := h.db.ListVulnerabilitiesForAccess(total, 0, filter, access)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
@@ -477,6 +504,14 @@ func (h *VulnerabilityHandler) ExportVulnerabilities(c *gin.Context) {
|
||||
})
|
||||
}
|
||||
|
||||
func vulnerabilityAccessFromContext(c *gin.Context) database.RBACListAccess {
|
||||
session, ok := security.CurrentSession(c)
|
||||
if !ok {
|
||||
return database.RBACListAccess{}
|
||||
}
|
||||
return database.RBACListAccess{UserID: session.UserID, Scope: session.Scope}
|
||||
}
|
||||
|
||||
// appendVulnerabilityMarkdown 单条漏洞的 Markdown 片段(与单文件下载字段对齐,缺省字段不写)
|
||||
func appendVulnerabilityMarkdown(b *strings.Builder, v *database.Vulnerability, titleHeading string) {
|
||||
b.WriteString(fmt.Sprintf("%s %s\n\n", titleHeading, v.Title))
|
||||
|
||||
@@ -15,6 +15,7 @@ import (
|
||||
|
||||
"cyberstrike-ai/internal/audit"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
@@ -144,13 +145,13 @@ func normalizeWindowsCmdPath(p string) string {
|
||||
return strings.ReplaceAll(s, "/", "\\")
|
||||
}
|
||||
|
||||
// quotePsSingle 把字符串按 PowerShell 单引号字符串规则转义(内部 ' → '')。
|
||||
// quotePsSingle 把字符串按 PowerShell 单引号字符串规则转义(内部 ' → ”)。
|
||||
// 供 PowerShell 脚本参数使用,全脚本只用单引号,外层 cmd 再用双引号包裹即可安全传递。
|
||||
func quotePsSingle(s string) string {
|
||||
return "'" + strings.ReplaceAll(s, "'", "''") + "'"
|
||||
}
|
||||
|
||||
// quoteShellSinglePosix 把路径按 POSIX sh 单引号规则转义(内部 ' → '\'')
|
||||
// quoteShellSinglePosix 把路径按 POSIX sh 单引号规则转义(内部 ' → '\”)
|
||||
func quoteShellSinglePosix(p string) string {
|
||||
if p == "" {
|
||||
return "."
|
||||
@@ -379,7 +380,8 @@ func (h *WebShellHandler) ListConnections(c *gin.Context) {
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"error": "database not available"})
|
||||
return
|
||||
}
|
||||
list, err := h.db.ListWebshellConnections()
|
||||
session, _ := security.CurrentSession(c)
|
||||
list, err := h.db.ListWebshellConnectionsForAccess(session.UserID, session.Scope)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
@@ -434,6 +436,10 @@ func (h *WebShellHandler) CreateConnection(c *gin.Context) {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if session, ok := security.CurrentSession(c); ok {
|
||||
_ = h.db.SetResourceOwner("webshell", conn.ID, session.UserID)
|
||||
_ = h.db.AssignResourceToUser(session.UserID, "webshell", conn.ID)
|
||||
}
|
||||
if h.audit != nil {
|
||||
host := req.URL
|
||||
if u, err := url.Parse(req.URL); err == nil {
|
||||
@@ -659,14 +665,15 @@ func (h *WebShellHandler) ListAIConversations(c *gin.Context) {
|
||||
|
||||
// ExecRequest 执行命令请求(前端传入连接信息 + 命令)
|
||||
type ExecRequest struct {
|
||||
URL string `json:"url" binding:"required"`
|
||||
Password string `json:"password"`
|
||||
Type string `json:"type"` // php, asp, aspx, jsp, custom
|
||||
Method string `json:"method"` // GET 或 POST,空则默认 POST
|
||||
CmdParam string `json:"cmd_param"` // 命令参数名,如 cmd/xxx,空则默认 cmd
|
||||
Encoding string `json:"encoding"` // 响应编码:auto / utf-8 / gbk / gb18030,空则 auto
|
||||
OS string `json:"os"` // 目标操作系统:auto / linux / windows,当前 exec 不用它,保留字段便于未来扩展
|
||||
Command string `json:"command" binding:"required"`
|
||||
URL string `json:"url" binding:"required"`
|
||||
Password string `json:"password"`
|
||||
Type string `json:"type"` // php, asp, aspx, jsp, custom
|
||||
Method string `json:"method"` // GET 或 POST,空则默认 POST
|
||||
CmdParam string `json:"cmd_param"` // 命令参数名,如 cmd/xxx,空则默认 cmd
|
||||
Encoding string `json:"encoding"` // 响应编码:auto / utf-8 / gbk / gb18030,空则 auto
|
||||
OS string `json:"os"` // 目标操作系统:auto / linux / windows,当前 exec 不用它,保留字段便于未来扩展
|
||||
ConnectionID string `json:"connection_id,omitempty"`
|
||||
Command string `json:"command" binding:"required"`
|
||||
}
|
||||
|
||||
// ExecResponse 执行命令响应
|
||||
@@ -714,6 +721,15 @@ func (h *WebShellHandler) Exec(c *gin.Context) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "url and command are required"})
|
||||
return
|
||||
}
|
||||
conn, allowed := h.authorizedWebshellConnection(c, req.ConnectionID, req.URL)
|
||||
if !allowed {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
// The database record is authoritative. Never let a caller pair an
|
||||
// authorized ID with attacker-controlled transport credentials or a URL.
|
||||
req.URL, req.Password, req.Type = conn.URL, conn.Password, conn.Type
|
||||
req.Method, req.CmdParam, req.Encoding = conn.Method, conn.CmdParam, conn.Encoding
|
||||
|
||||
parsed, err := url.Parse(req.URL)
|
||||
if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") {
|
||||
@@ -807,6 +823,13 @@ func (h *WebShellHandler) FileOp(c *gin.Context) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "url and action are required"})
|
||||
return
|
||||
}
|
||||
conn, allowed := h.authorizedWebshellConnection(c, req.ConnectionID, req.URL)
|
||||
if !allowed {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
req.URL, req.Password, req.Type = conn.URL, conn.Password, conn.Type
|
||||
req.Method, req.CmdParam, req.Encoding, req.OS = conn.Method, conn.CmdParam, conn.Encoding, conn.OS
|
||||
|
||||
parsed, err := url.Parse(req.URL)
|
||||
if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") {
|
||||
@@ -883,6 +906,28 @@ func (h *WebShellHandler) FileOp(c *gin.Context) {
|
||||
})
|
||||
}
|
||||
|
||||
func (h *WebShellHandler) authorizedWebshellConnection(c *gin.Context, connectionID, requestURL string) (*database.WebShellConnection, bool) {
|
||||
connectionID = strings.TrimSpace(connectionID)
|
||||
if connectionID == "" {
|
||||
return nil, false
|
||||
}
|
||||
if h.db == nil {
|
||||
return nil, false
|
||||
}
|
||||
session, ok := security.CurrentSession(c)
|
||||
if !ok || !h.db.UserCanAccessResource(session.UserID, session.Scope, "webshell", connectionID) {
|
||||
return nil, false
|
||||
}
|
||||
conn, err := h.db.GetWebshellConnection(connectionID)
|
||||
if err != nil || conn == nil {
|
||||
return nil, false
|
||||
}
|
||||
if requestURL = strings.TrimSpace(requestURL); requestURL != "" && strings.TrimSpace(conn.URL) != requestURL {
|
||||
return nil, false
|
||||
}
|
||||
return conn, true
|
||||
}
|
||||
|
||||
// ExecWithConnection 在指定 WebShell 连接上执行命令(供 MCP/Agent 等非 HTTP 调用)
|
||||
func (h *WebShellHandler) ExecWithConnection(conn *database.WebShellConnection, command string) (output string, ok bool, errMsg string) {
|
||||
if conn == nil {
|
||||
|
||||
@@ -0,0 +1,124 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestWebshellExecRequiresConnectionAccessWhenConnectionIDProvided(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
db, user, allowed, hidden := setupWebshellRBACTest(t)
|
||||
handler := NewWebShellHandler(zap.NewNop(), db)
|
||||
|
||||
w := performWebshellJSON(user, http.MethodPost, "/api/webshell/exec", map[string]interface{}{
|
||||
"url": hidden.URL,
|
||||
"connection_id": hidden.ID,
|
||||
"command": "id",
|
||||
}, handler.Exec)
|
||||
if w.Code != http.StatusForbidden {
|
||||
t.Fatalf("hidden connection status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String())
|
||||
}
|
||||
|
||||
w = performWebshellJSON(user, http.MethodPost, "/api/webshell/exec", map[string]interface{}{
|
||||
"url": hidden.URL,
|
||||
"connection_id": allowed.ID,
|
||||
"command": "id",
|
||||
}, handler.Exec)
|
||||
if w.Code != http.StatusForbidden {
|
||||
t.Fatalf("mismatched URL status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebshellExecRejectsAdHocURLWithoutConnectionID(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
_, user, _, _ := setupWebshellRBACTest(t)
|
||||
handler := NewWebShellHandler(zap.NewNop(), nil)
|
||||
w := performWebshellJSON(user, http.MethodPost, "/api/webshell/exec", map[string]interface{}{
|
||||
"url": "http://127.0.0.1/admin", "command": "id",
|
||||
}, handler.Exec)
|
||||
if w.Code != http.StatusForbidden {
|
||||
t.Fatalf("ad-hoc URL status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebshellFileOpRequiresConnectionAccessWhenConnectionIDProvided(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
db, user, _, hidden := setupWebshellRBACTest(t)
|
||||
handler := NewWebShellHandler(zap.NewNop(), db)
|
||||
|
||||
w := performWebshellJSON(user, http.MethodPost, "/api/webshell/file", map[string]interface{}{
|
||||
"url": hidden.URL,
|
||||
"connection_id": hidden.ID,
|
||||
"action": "list",
|
||||
"path": ".",
|
||||
}, handler.FileOp)
|
||||
if w.Code != http.StatusForbidden {
|
||||
t.Fatalf("status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func setupWebshellRBACTest(t *testing.T) (*database.DB, *database.RBACUser, *database.WebShellConnection, *database.WebShellConnection) {
|
||||
t.Helper()
|
||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "webshell-rbac.db"), zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatalf("NewDB: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
user, err := db.CreateRBACUser("operator1", "Operator One", "hash", true, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("CreateRBACUser: %v", err)
|
||||
}
|
||||
allowed := &database.WebShellConnection{
|
||||
ID: "ws_allowed",
|
||||
URL: "http://127.0.0.1/allowed.php",
|
||||
Type: "php",
|
||||
Method: "post",
|
||||
CmdParam: "cmd",
|
||||
CreatedAt: time.Now(),
|
||||
}
|
||||
hidden := &database.WebShellConnection{
|
||||
ID: "ws_hidden",
|
||||
URL: "http://127.0.0.1/hidden.php",
|
||||
Type: "php",
|
||||
Method: "post",
|
||||
CmdParam: "cmd",
|
||||
CreatedAt: time.Now(),
|
||||
}
|
||||
if err := db.CreateWebshellConnection(allowed); err != nil {
|
||||
t.Fatalf("CreateWebshellConnection allowed: %v", err)
|
||||
}
|
||||
if err := db.CreateWebshellConnection(hidden); err != nil {
|
||||
t.Fatalf("CreateWebshellConnection hidden: %v", err)
|
||||
}
|
||||
if err := db.AssignResourceToUser(user.ID, "webshell", allowed.ID); err != nil {
|
||||
t.Fatalf("AssignResourceToUser: %v", err)
|
||||
}
|
||||
return db, user, allowed, hidden
|
||||
}
|
||||
|
||||
func performWebshellJSON(user *database.RBACUser, method, path string, body map[string]interface{}, handler gin.HandlerFunc) *httptest.ResponseRecorder {
|
||||
payload, _ := json.Marshal(body)
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(method, path, bytes.NewReader(payload))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
c.Set(security.ContextSessionKey, security.Session{
|
||||
UserID: user.ID,
|
||||
Username: user.Username,
|
||||
Permissions: map[string]bool{"webshell:write": true},
|
||||
Scope: database.RBACScopeAssigned,
|
||||
})
|
||||
handler(c)
|
||||
return w
|
||||
}
|
||||
@@ -36,6 +36,7 @@ func (h *AgentHandler) roleForWorkflow(req *ChatRequest) (config.RoleConfig, boo
|
||||
}
|
||||
|
||||
func (h *AgentHandler) runRoleWorkflowStreamIfBound(
|
||||
c *gin.Context,
|
||||
req *ChatRequest,
|
||||
prep *multiAgentPrepared,
|
||||
sendEvent func(eventType, message string, data interface{}),
|
||||
@@ -60,7 +61,10 @@ func (h *AgentHandler) runRoleWorkflowStreamIfBound(
|
||||
}
|
||||
}()
|
||||
|
||||
baseCtx, cancelWithCause := context.WithCancelCause(context.Background())
|
||||
if c == nil || c.Request == nil {
|
||||
return false
|
||||
}
|
||||
baseCtx, cancelWithCause := context.WithCancelCause(detachedAgentContext(c.Request.Context()))
|
||||
defer cancelWithCause(nil)
|
||||
taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute)
|
||||
defer timeoutCancel()
|
||||
|
||||
@@ -8,6 +8,8 @@ import (
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
"cyberstrike-ai/internal/security"
|
||||
workflowrunner "cyberstrike-ai/internal/workflow"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
@@ -20,6 +22,10 @@ func (h *WorkflowHandler) SetRuntime(agent *agent.Agent, cfg *config.Config) {
|
||||
|
||||
func (h *WorkflowHandler) GetRun(c *gin.Context) {
|
||||
runID := strings.TrimSpace(c.Param("runId"))
|
||||
if !h.workflowRunAllowed(c, runID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
run, err := h.db.GetWorkflowRun(runID)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
@@ -39,6 +45,10 @@ func (h *WorkflowHandler) GetRun(c *gin.Context) {
|
||||
|
||||
func (h *WorkflowHandler) ReplayRun(c *gin.Context) {
|
||||
runID := strings.TrimSpace(c.Param("runId"))
|
||||
if !h.workflowRunAllowed(c, runID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
nodeRuns, err := h.db.ListWorkflowNodeRuns(runID)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
@@ -67,11 +77,18 @@ func (h *WorkflowHandler) ReplayRun(c *gin.Context) {
|
||||
|
||||
func (h *WorkflowHandler) ListPendingRuns(c *gin.Context) {
|
||||
conversationID := strings.TrimSpace(c.Query("conversationId"))
|
||||
if conversationID != "" && !h.workflowConversationAllowed(c, conversationID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
runs, err := h.db.ListWorkflowRunsAwaitingHITLFiltered(conversationID, 50)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
runs = filterSlice(runs, func(run *database.WorkflowRun) bool {
|
||||
return run != nil && h.workflowConversationAllowed(c, run.ConversationID)
|
||||
})
|
||||
c.JSON(http.StatusOK, gin.H{"runs": runs})
|
||||
}
|
||||
|
||||
@@ -86,6 +103,10 @@ func (h *WorkflowHandler) ResumeRun(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
runID := strings.TrimSpace(c.Param("runId"))
|
||||
if !h.workflowRunAllowed(c, runID) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||
return
|
||||
}
|
||||
var req workflowResumeRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()})
|
||||
@@ -160,3 +181,20 @@ func (h *WorkflowHandler) ResumeRun(c *gin.Context) {
|
||||
"awaitingHitl": result.AwaitingHITL,
|
||||
})
|
||||
}
|
||||
|
||||
func (h *WorkflowHandler) workflowConversationAllowed(c *gin.Context, conversationID string) bool {
|
||||
session, ok := security.CurrentSession(c)
|
||||
return ok && h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", strings.TrimSpace(conversationID))
|
||||
}
|
||||
|
||||
func (h *WorkflowHandler) workflowRunAllowed(c *gin.Context, runID string) bool {
|
||||
session, ok := security.CurrentSession(c)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
if session.Scope == database.RBACScopeAll {
|
||||
return true
|
||||
}
|
||||
run, err := h.db.GetWorkflowRun(strings.TrimSpace(runID))
|
||||
return err == nil && run != nil && h.workflowConversationAllowed(c, run.ConversationID)
|
||||
}
|
||||
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/config"
|
||||
|
||||
"github.com/google/uuid"
|
||||
@@ -37,29 +38,30 @@ type listToolsInflight struct {
|
||||
|
||||
// ExternalMCPManager 外部MCP管理器
|
||||
type ExternalMCPManager struct {
|
||||
clients map[string]ExternalMCPClient
|
||||
configs map[string]config.ExternalMCPServerConfig
|
||||
logger *zap.Logger
|
||||
storage MonitorStorage // 可选的持久化存储
|
||||
executions map[string]*ToolExecution // 执行记录
|
||||
stats map[string]*ToolStats // 工具统计信息
|
||||
errors map[string]string // 错误信息
|
||||
toolCounts map[string]int // 工具数量缓存
|
||||
toolCountsMu sync.RWMutex // 工具数量缓存的锁
|
||||
toolCache map[string]toolListCacheEntry // 工具列表缓存:MCP名称 -> 工具列表
|
||||
toolCacheMu sync.RWMutex // 工具列表缓存的锁
|
||||
listToolsMu sync.Mutex
|
||||
clients map[string]ExternalMCPClient
|
||||
configs map[string]config.ExternalMCPServerConfig
|
||||
logger *zap.Logger
|
||||
storage MonitorStorage // 可选的持久化存储
|
||||
executions map[string]*ToolExecution // 执行记录
|
||||
stats map[string]*ToolStats // 工具统计信息
|
||||
errors map[string]string // 错误信息
|
||||
toolCounts map[string]int // 工具数量缓存
|
||||
toolCountsMu sync.RWMutex // 工具数量缓存的锁
|
||||
toolCache map[string]toolListCacheEntry // 工具列表缓存:MCP名称 -> 工具列表
|
||||
toolCacheMu sync.RWMutex // 工具列表缓存的锁
|
||||
listToolsMu sync.Mutex
|
||||
listToolsInflight map[string]*listToolsInflight
|
||||
stopRefresh chan struct{} // 停止后台刷新的信号
|
||||
refreshWg sync.WaitGroup // 等待后台刷新goroutine完成
|
||||
refreshing atomic.Bool // 防止 refreshToolCounts 并发堆积
|
||||
mu sync.RWMutex
|
||||
stopRefresh chan struct{} // 停止后台刷新的信号
|
||||
refreshWg sync.WaitGroup // 等待后台刷新goroutine完成
|
||||
refreshing atomic.Bool // 防止 refreshToolCounts 并发堆积
|
||||
mu sync.RWMutex
|
||||
runningCancels map[string]context.CancelFunc
|
||||
abortUserNotes map[string]string
|
||||
reconnectMu sync.Mutex
|
||||
reconnecting map[string]bool
|
||||
reconnectLastTry map[string]time.Time
|
||||
reconnectAttempts map[string]int
|
||||
toolAuthorizer func(context.Context, string, map[string]interface{}) error
|
||||
}
|
||||
|
||||
// NewExternalMCPManager 创建外部MCP管理器
|
||||
@@ -67,16 +69,24 @@ func NewExternalMCPManager(logger *zap.Logger) *ExternalMCPManager {
|
||||
return NewExternalMCPManagerWithStorage(logger, nil)
|
||||
}
|
||||
|
||||
// SetToolAuthorizer installs the policy decision point for all external MCP
|
||||
// invocations. App wiring configures this before any Agent can call a tool.
|
||||
func (m *ExternalMCPManager) SetToolAuthorizer(authorizer func(context.Context, string, map[string]interface{}) error) {
|
||||
m.mu.Lock()
|
||||
m.toolAuthorizer = authorizer
|
||||
m.mu.Unlock()
|
||||
}
|
||||
|
||||
// NewExternalMCPManagerWithStorage 创建外部MCP管理器(带持久化存储)
|
||||
func NewExternalMCPManagerWithStorage(logger *zap.Logger, storage MonitorStorage) *ExternalMCPManager {
|
||||
manager := &ExternalMCPManager{
|
||||
clients: make(map[string]ExternalMCPClient),
|
||||
configs: make(map[string]config.ExternalMCPServerConfig),
|
||||
logger: logger,
|
||||
storage: storage,
|
||||
executions: make(map[string]*ToolExecution),
|
||||
stats: make(map[string]*ToolStats),
|
||||
errors: make(map[string]string),
|
||||
clients: make(map[string]ExternalMCPClient),
|
||||
configs: make(map[string]config.ExternalMCPServerConfig),
|
||||
logger: logger,
|
||||
storage: storage,
|
||||
executions: make(map[string]*ToolExecution),
|
||||
stats: make(map[string]*ToolStats),
|
||||
errors: make(map[string]string),
|
||||
toolCounts: make(map[string]int),
|
||||
toolCache: make(map[string]toolListCacheEntry),
|
||||
listToolsInflight: make(map[string]*listToolsInflight),
|
||||
@@ -554,6 +564,17 @@ func (m *ExternalMCPManager) updateToolCache(name string, tools []Tool) {
|
||||
|
||||
// CallTool 调用外部MCP工具(返回执行ID)
|
||||
func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args map[string]interface{}) (*ToolResult, string, error) {
|
||||
_, authenticated := authctx.PrincipalFromContext(ctx)
|
||||
m.mu.RLock()
|
||||
authorizer := m.toolAuthorizer
|
||||
m.mu.RUnlock()
|
||||
if authorizer != nil {
|
||||
if err := authorizer(ctx, toolName, args); err != nil {
|
||||
return nil, "", fmt.Errorf("external tool authorization denied: %w", err)
|
||||
}
|
||||
} else if authenticated {
|
||||
return nil, "", fmt.Errorf("external tool authorization policy is not configured")
|
||||
}
|
||||
// 解析工具名称:name::toolName
|
||||
var mcpName, actualToolName string
|
||||
if idx := findSubstring(toolName, "::"); idx > 0 {
|
||||
@@ -591,6 +612,10 @@ func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args
|
||||
Status: "running",
|
||||
StartTime: time.Now(),
|
||||
}
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||
execution.OwnerUserID = principal.UserID
|
||||
}
|
||||
execution.ConversationID = MCPConversationIDFromContext(ctx)
|
||||
|
||||
m.mu.Lock()
|
||||
m.executions[executionID] = execution
|
||||
|
||||
@@ -2,14 +2,30 @@ package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
"cyberstrike-ai/internal/config"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestExternalManagerEnforcesConfiguredAuthorizer(t *testing.T) {
|
||||
manager := NewExternalMCPManager(zap.NewNop())
|
||||
t.Cleanup(manager.StopAll)
|
||||
manager.SetToolAuthorizer(func(context.Context, string, map[string]interface{}) error {
|
||||
return errors.New("denied by policy")
|
||||
})
|
||||
ctx := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", "assigned", map[string]bool{"agent:execute": true}))
|
||||
_, _, err := manager.CallTool(ctx, "server::tool", map[string]interface{}{})
|
||||
if err == nil || !strings.Contains(err.Error(), "authorization denied") {
|
||||
t.Fatalf("external call bypassed authorizer: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExternalMCPManager_AddOrUpdateConfig(t *testing.T) {
|
||||
logger := zap.NewNop()
|
||||
manager := NewExternalMCPManager(logger)
|
||||
|
||||
+74
-17
@@ -14,6 +14,8 @@ import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
@@ -49,6 +51,18 @@ type Server struct {
|
||||
// nil 表示未配置,沿用默认 30 分钟;指向 0 表示不限制;>0 为分钟数。
|
||||
httpToolTimeoutMinutes *int
|
||||
httpToolTimeoutMu sync.RWMutex
|
||||
toolAuthorizer func(context.Context, string, map[string]interface{}) error
|
||||
}
|
||||
|
||||
// SetToolAuthorizer installs the common policy decision point for every
|
||||
// user-attributed tool call, whether it originates from HTTP or an Agent.
|
||||
func (s *Server) SetToolAuthorizer(authorizer func(context.Context, string, map[string]interface{}) error) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.toolAuthorizer = authorizer
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
type sseClient struct {
|
||||
@@ -111,21 +125,24 @@ func (s *Server) ConfigureHTTPToolCallTimeoutFromAgentMinutes(minutes int) {
|
||||
s.httpToolTimeoutMinutes = &v
|
||||
}
|
||||
|
||||
func (s *Server) effectiveHTTPToolCallDeadline() (context.Context, context.CancelFunc) {
|
||||
func (s *Server) effectiveHTTPToolCallDeadline(parent context.Context) (context.Context, context.CancelFunc) {
|
||||
const defaultDur = 30 * time.Minute
|
||||
if parent == nil {
|
||||
parent = context.Background()
|
||||
}
|
||||
if s == nil {
|
||||
return context.WithTimeout(context.Background(), defaultDur)
|
||||
return context.WithTimeout(parent, defaultDur)
|
||||
}
|
||||
s.httpToolTimeoutMu.RLock()
|
||||
mPtr := s.httpToolTimeoutMinutes
|
||||
s.httpToolTimeoutMu.RUnlock()
|
||||
if mPtr == nil {
|
||||
return context.WithTimeout(context.Background(), defaultDur)
|
||||
return context.WithTimeout(parent, defaultDur)
|
||||
}
|
||||
if *mPtr <= 0 {
|
||||
return context.WithCancel(context.Background())
|
||||
return context.WithCancel(parent)
|
||||
}
|
||||
return context.WithTimeout(context.Background(), time.Duration(*mPtr)*time.Minute)
|
||||
return context.WithTimeout(parent, time.Duration(*mPtr)*time.Minute)
|
||||
}
|
||||
|
||||
// RegisterTool 注册工具
|
||||
@@ -196,7 +213,7 @@ func (s *Server) HandleHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
|
||||
response := s.handleMessage(&msg)
|
||||
response := s.handleMessage(r.Context(), &msg)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
json.NewEncoder(w).Encode(response)
|
||||
}
|
||||
@@ -223,7 +240,7 @@ func (s *Server) serveSSESessionMessage(w http.ResponseWriter, r *http.Request,
|
||||
return
|
||||
}
|
||||
|
||||
response := s.handleMessage(&msg)
|
||||
response := s.handleMessage(r.Context(), &msg)
|
||||
if response == nil {
|
||||
w.WriteHeader(http.StatusAccepted)
|
||||
return
|
||||
@@ -317,7 +334,7 @@ func (s *Server) removeSSEClient(id string) {
|
||||
}
|
||||
|
||||
// handleMessage 处理MCP消息
|
||||
func (s *Server) handleMessage(msg *Message) *Message {
|
||||
func (s *Server) handleMessage(ctx context.Context, msg *Message) *Message {
|
||||
// 检查是否是通知(notification)- 通知没有id字段,不需要响应
|
||||
isNotification := msg.ID.Value() == nil || msg.ID.String() == ""
|
||||
|
||||
@@ -332,7 +349,7 @@ func (s *Server) handleMessage(msg *Message) *Message {
|
||||
case "tools/list":
|
||||
return s.handleListTools(msg)
|
||||
case "tools/call":
|
||||
return s.handleCallTool(msg)
|
||||
return s.handleCallTool(ctx, msg)
|
||||
case "prompts/list":
|
||||
return s.handleListPrompts(msg)
|
||||
case "prompts/get":
|
||||
@@ -433,7 +450,7 @@ func (s *Server) handleListTools(msg *Message) *Message {
|
||||
}
|
||||
|
||||
// handleCallTool 处理工具调用请求
|
||||
func (s *Server) handleCallTool(msg *Message) *Message {
|
||||
func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Message {
|
||||
var req CallToolRequest
|
||||
if err := json.Unmarshal(msg.Params, &req); err != nil {
|
||||
return &Message{
|
||||
@@ -443,6 +460,17 @@ func (s *Server) handleCallTool(msg *Message) *Message {
|
||||
Error: &Error{Code: -32602, Message: "Invalid params"},
|
||||
}
|
||||
}
|
||||
_, authenticated := authctx.PrincipalFromContext(requestCtx)
|
||||
s.mu.RLock()
|
||||
authorizer := s.toolAuthorizer
|
||||
s.mu.RUnlock()
|
||||
if authorizer != nil {
|
||||
if err := authorizer(requestCtx, req.Name, req.Arguments); err != nil {
|
||||
return &Message{ID: msg.ID, Type: MessageTypeError, Version: "2.0", Error: &Error{Code: -32003, Message: "Forbidden", Data: err.Error()}}
|
||||
}
|
||||
} else if authenticated {
|
||||
return &Message{ID: msg.ID, Type: MessageTypeError, Version: "2.0", Error: &Error{Code: -32003, Message: "Tool authorization policy is not configured"}}
|
||||
}
|
||||
|
||||
executionID := uuid.New().String()
|
||||
execution := &ToolExecution{
|
||||
@@ -452,6 +480,10 @@ func (s *Server) handleCallTool(msg *Message) *Message {
|
||||
Status: "running",
|
||||
StartTime: time.Now(),
|
||||
}
|
||||
if principal, ok := authctx.PrincipalFromContext(requestCtx); ok {
|
||||
execution.OwnerUserID = principal.UserID
|
||||
}
|
||||
execution.ConversationID = MCPConversationIDFromContext(requestCtx)
|
||||
|
||||
s.mu.Lock()
|
||||
s.executions[executionID] = execution
|
||||
@@ -495,7 +527,7 @@ func (s *Server) handleCallTool(msg *Message) *Message {
|
||||
}
|
||||
}
|
||||
|
||||
baseCtx, timeoutCancel := s.effectiveHTTPToolCallDeadline()
|
||||
baseCtx, timeoutCancel := s.effectiveHTTPToolCallDeadline(requestCtx)
|
||||
defer timeoutCancel()
|
||||
execCtx, runCancel := context.WithCancel(baseCtx)
|
||||
s.registerRunningCancel(executionID, runCancel)
|
||||
@@ -809,6 +841,17 @@ func (s *Server) GetAllTools() []Tool {
|
||||
|
||||
// CallTool 直接调用工具(用于内部调用)
|
||||
func (s *Server) CallTool(ctx context.Context, toolName string, args map[string]interface{}) (*ToolResult, string, error) {
|
||||
_, authenticated := authctx.PrincipalFromContext(ctx)
|
||||
s.mu.RLock()
|
||||
authorizer := s.toolAuthorizer
|
||||
s.mu.RUnlock()
|
||||
if authorizer != nil {
|
||||
if err := authorizer(ctx, toolName, args); err != nil {
|
||||
return nil, "", fmt.Errorf("tool authorization denied: %w", err)
|
||||
}
|
||||
} else if authenticated {
|
||||
return nil, "", errors.New("tool authorization policy is not configured")
|
||||
}
|
||||
s.mu.RLock()
|
||||
handler, exists := s.tools[toolName]
|
||||
s.mu.RUnlock()
|
||||
@@ -826,6 +869,10 @@ func (s *Server) CallTool(ctx context.Context, toolName string, args map[string]
|
||||
Status: "running",
|
||||
StartTime: time.Now(),
|
||||
}
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||
execution.OwnerUserID = principal.UserID
|
||||
}
|
||||
execution.ConversationID = MCPConversationIDFromContext(ctx)
|
||||
|
||||
s.mu.Lock()
|
||||
s.executions[executionID] = execution
|
||||
@@ -922,7 +969,7 @@ func (s *Server) CallTool(ctx context.Context, toolName string, args map[string]
|
||||
}
|
||||
|
||||
// BeginToolExecution 创建 running 状态的执行记录,供 Eino 等非 CallTool 路径在工具开始时落库。
|
||||
func (s *Server) BeginToolExecution(toolName string, args map[string]interface{}) string {
|
||||
func (s *Server) BeginToolExecution(ctx context.Context, toolName string, args map[string]interface{}) string {
|
||||
if s == nil {
|
||||
return ""
|
||||
}
|
||||
@@ -937,6 +984,10 @@ func (s *Server) BeginToolExecution(toolName string, args map[string]interface{}
|
||||
Status: "running",
|
||||
StartTime: time.Now(),
|
||||
}
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||
execution.OwnerUserID = principal.UserID
|
||||
}
|
||||
execution.ConversationID = MCPConversationIDFromContext(ctx)
|
||||
|
||||
s.mu.Lock()
|
||||
s.executions[executionID] = execution
|
||||
@@ -952,7 +1003,7 @@ func (s *Server) BeginToolExecution(toolName string, args map[string]interface{}
|
||||
}
|
||||
|
||||
// FinishToolExecution 完成先前 BeginToolExecution 创建的记录;executionID 为空时等同 RecordCompletedToolInvocation。
|
||||
func (s *Server) FinishToolExecution(executionID, toolName string, args map[string]interface{}, resultText string, invokeErr error) string {
|
||||
func (s *Server) FinishToolExecution(ctx context.Context, executionID, toolName string, args map[string]interface{}, resultText string, invokeErr error) string {
|
||||
if s == nil {
|
||||
return ""
|
||||
}
|
||||
@@ -961,7 +1012,7 @@ func (s *Server) FinishToolExecution(executionID, toolName string, args map[stri
|
||||
}
|
||||
id := strings.TrimSpace(executionID)
|
||||
if id == "" {
|
||||
return s.RecordCompletedToolInvocation(toolName, args, resultText, invokeErr)
|
||||
id = uuid.New().String()
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
@@ -984,6 +1035,12 @@ func (s *Server) FinishToolExecution(executionID, toolName string, args map[stri
|
||||
if len(args) > 0 {
|
||||
exec.Arguments = args
|
||||
}
|
||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||
exec.OwnerUserID = principal.UserID
|
||||
}
|
||||
if conversationID := MCPConversationIDFromContext(ctx); conversationID != "" {
|
||||
exec.ConversationID = conversationID
|
||||
}
|
||||
exec.EndTime = &now
|
||||
if exec.StartTime.IsZero() {
|
||||
exec.StartTime = now
|
||||
@@ -1027,8 +1084,8 @@ func (s *Server) FinishToolExecution(executionID, toolName string, args map[stri
|
||||
|
||||
// RecordCompletedToolInvocation 将已在其它路径完成的工具调用写入监控存储(格式与 CallTool 结束后一致),
|
||||
// 用于 Eino ADK filesystem execute 等未经过 CallTool 的场景;返回 executionId 供助手消息 mcpExecutionIds 关联。
|
||||
func (s *Server) RecordCompletedToolInvocation(toolName string, args map[string]interface{}, resultText string, invokeErr error) string {
|
||||
return s.FinishToolExecution("", toolName, args, resultText, invokeErr)
|
||||
func (s *Server) RecordCompletedToolInvocation(ctx context.Context, toolName string, args map[string]interface{}, resultText string, invokeErr error) string {
|
||||
return s.FinishToolExecution(ctx, "", toolName, args, resultText, invokeErr)
|
||||
}
|
||||
|
||||
// UpdateToolExecutionResult 将监控库中的工具结果更新为送入模型的展示正文(如 reduction 后的 persisted-output)。
|
||||
@@ -1519,7 +1576,7 @@ func (s *Server) HandleStdio() error {
|
||||
}
|
||||
|
||||
// 处理消息
|
||||
response := s.handleMessage(&msg)
|
||||
response := s.handleMessage(context.Background(), &msg)
|
||||
|
||||
// 如果是通知(response 为 nil),不需要发送响应
|
||||
if response == nil {
|
||||
|
||||
@@ -0,0 +1,36 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"cyberstrike-ai/internal/authctx"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func TestToolAuthorizerIsUniversalAndExecutionKeepsOwner(t *testing.T) {
|
||||
server := NewServer(zap.NewNop())
|
||||
server.RegisterTool(Tool{Name: "echo", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
|
||||
return &ToolResult{Content: []Content{{Type: "text", Text: "ok"}}}, nil
|
||||
})
|
||||
server.SetToolAuthorizer(func(ctx context.Context, toolName string, args map[string]interface{}) error {
|
||||
if _, ok := authctx.PrincipalFromContext(ctx); !ok {
|
||||
return errors.New("principal required")
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if _, _, err := server.CallTool(context.Background(), "echo", nil); err == nil {
|
||||
t.Fatal("tool call without principal was allowed")
|
||||
}
|
||||
ctx := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", "assigned", map[string]bool{"mcp:execute": true}))
|
||||
_, executionID, err := server.CallTool(ctx, "echo", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
execution, ok := server.GetExecution(executionID)
|
||||
if !ok || execution.OwnerUserID != "u1" {
|
||||
t.Fatalf("execution owner = %#v, want u1", execution)
|
||||
}
|
||||
}
|
||||
@@ -201,6 +201,7 @@ type ToolExecution struct {
|
||||
Duration time.Duration `json:"duration,omitempty"`
|
||||
// ConversationID 仅 API 展示用(进行中的 Agent 任务),不写入 tool_executions 表。
|
||||
ConversationID string `json:"conversationId,omitempty"`
|
||||
OwnerUserID string `json:"-"`
|
||||
}
|
||||
|
||||
// ToolStats 工具统计信息
|
||||
|
||||
@@ -79,7 +79,7 @@ type einoADKRunLoopArgs struct {
|
||||
EinoRoleTag func(agent string) string
|
||||
CheckpointDir string
|
||||
// RunRetryMaxAttempts / RunRetryMaxBackoffSec:429、5xx、网络抖动时的指数退避续跑(0=默认 10 次 / 30s 上限)。
|
||||
RunRetryMaxAttempts int
|
||||
RunRetryMaxAttempts int
|
||||
RunRetryMaxBackoffSec int
|
||||
|
||||
McpIDsMu *sync.Mutex
|
||||
@@ -200,6 +200,7 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
||||
markPendingWithMonitor := func(tc toolCallPendingInfo) {
|
||||
markPending(tc)
|
||||
beginEinoADKFilesystemToolMonitor(
|
||||
ctx,
|
||||
args.FilesystemMonitorAgent,
|
||||
args.FilesystemMonitorRecord,
|
||||
args.MCPExecutionBinder,
|
||||
@@ -301,9 +302,6 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
||||
tryEmitToolResultProgress := func(toolName, content, toolCallID string, isErr bool, agentName string) {
|
||||
// 仅由 ADK schema.Tool 事件调用;MCP/execute 桥在 reduction 前的 ToolInvokeNotify 不得推送 tool_result,
|
||||
// 否则全量输出会先占位并触发 toolResultSent 去重,导致 UI/监控展示与 agent 实际收到的截断正文不一致。
|
||||
if progress == nil {
|
||||
return
|
||||
}
|
||||
toolName = strings.TrimSpace(toolName)
|
||||
if toolName == "" {
|
||||
toolName = "unknown"
|
||||
@@ -345,13 +343,15 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
||||
toolCallID = tid
|
||||
}
|
||||
recordPendingExecuteStdoutDup(toolName, content, isErr)
|
||||
recordEinoADKFilesystemToolMonitor(args.FilesystemMonitorAgent, args.FilesystemMonitorRecord, args.MCPExecutionBinder, toolName, toolCallID, runAccumulatedMsgs, content, isErr)
|
||||
recordEinoADKFilesystemToolMonitor(ctx, args.FilesystemMonitorAgent, args.FilesystemMonitorRecord, args.MCPExecutionBinder, toolName, toolCallID, runAccumulatedMsgs, content, isErr)
|
||||
if args.FilesystemMonitorAgent != nil && args.MCPExecutionBinder != nil {
|
||||
if execID := args.MCPExecutionBinder.ExecutionID(toolCallID); execID != "" {
|
||||
args.FilesystemMonitorAgent.UpdateMCPExecutionDisplayResult(execID, content)
|
||||
}
|
||||
}
|
||||
progress("tool_result", fmt.Sprintf("工具结果 (%s)", toolName), data)
|
||||
if progress != nil {
|
||||
progress("tool_result", fmt.Sprintf("工具结果 (%s)", toolName), data)
|
||||
}
|
||||
}
|
||||
|
||||
if args.EinoCallbacks != nil {
|
||||
@@ -560,7 +560,7 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
||||
}
|
||||
ids := snapshotMCPIDs()
|
||||
return buildEinoRunResultFromAccumulated(
|
||||
orchMode, runAccumulatedMsgs, persistTraceSource(args, runAccumulatedMsgs),
|
||||
orchMode, runAccumulatedMsgs, modelFacingTraceSnapshot(args),
|
||||
lastAssistant, lastPlanExecuteExecutor, emptyHint, ids, true,
|
||||
), runErr
|
||||
}
|
||||
@@ -840,7 +840,7 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
||||
"einoRole": "orchestrator",
|
||||
"einoAgent": ev.AgentName,
|
||||
"orchestration": orchMode,
|
||||
"iteration": einoMainRound,
|
||||
"iteration": einoMainRound,
|
||||
"streamId": mainStreamID,
|
||||
}, mainAssistantBuf))
|
||||
mainAssistWireAccum, _ = normalizeStreamingDelta(mainAssistWireAccum, contentDelta)
|
||||
@@ -1089,19 +1089,22 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
||||
mcpIDsMu.Unlock()
|
||||
|
||||
out := buildEinoRunResultFromAccumulated(
|
||||
orchMode, runAccumulatedMsgs, persistTraceSource(args, runAccumulatedMsgs),
|
||||
orchMode, runAccumulatedMsgs, modelFacingTraceSnapshot(args),
|
||||
lastAssistant, lastPlanExecuteExecutor, emptyHint, ids, false,
|
||||
)
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func persistTraceSource(args *einoADKRunLoopArgs, fallback []adk.Message) []adk.Message {
|
||||
// modelFacingTraceSnapshot returns only the state that actually reached the model boundary.
|
||||
// Never fall back to event-stream accumulation here: it can contain pre-reduction tool output
|
||||
// that the model never received (for example when summarization failed before the first call).
|
||||
func modelFacingTraceSnapshot(args *einoADKRunLoopArgs) []adk.Message {
|
||||
if args != nil && args.ModelFacingTrace != nil {
|
||||
if snap := args.ModelFacingTrace.Snapshot(); len(snap) > 0 {
|
||||
return snap
|
||||
}
|
||||
}
|
||||
return fallback
|
||||
return nil
|
||||
}
|
||||
|
||||
func einoPartialRunLastOutputHint() string {
|
||||
@@ -1233,10 +1236,13 @@ func buildEinoRunResultFromAccumulated(
|
||||
partial bool,
|
||||
) *RunResult {
|
||||
traceForJSON := persistMsgs
|
||||
if len(traceForJSON) == 0 {
|
||||
traceForJSON = runAccumulatedMsgs
|
||||
traceJSON := ""
|
||||
if len(traceForJSON) > 0 {
|
||||
traceForJSON = markModelFacingTraceForPersistence(traceForJSON)
|
||||
if histJSON, err := json.Marshal(traceForJSON); err == nil {
|
||||
traceJSON = string(histJSON)
|
||||
}
|
||||
}
|
||||
histJSON, _ := json.Marshal(traceForJSON)
|
||||
cleaned := strings.TrimSpace(lastAssistant)
|
||||
if orchMode == "plan_execute" {
|
||||
if e := strings.TrimSpace(lastPlanExecuteExecutor); e != "" {
|
||||
@@ -1266,7 +1272,7 @@ func buildEinoRunResultFromAccumulated(
|
||||
out := &RunResult{
|
||||
Response: resp,
|
||||
MCPExecutionIDs: mcpIDs,
|
||||
LastAgentTraceInput: string(histJSON),
|
||||
LastAgentTraceInput: traceJSON,
|
||||
LastAgentTraceOutput: lastOut,
|
||||
}
|
||||
if !partial && out.Response == "" {
|
||||
@@ -1276,6 +1282,18 @@ func buildEinoRunResultFromAccumulated(
|
||||
return out
|
||||
}
|
||||
|
||||
func markModelFacingTraceForPersistence(msgs []adk.Message) []adk.Message {
|
||||
out := cloneADKMessagesForTrace(msgs)
|
||||
if len(out) == 0 || out[0] == nil {
|
||||
return out
|
||||
}
|
||||
if out[0].Extra == nil {
|
||||
out[0].Extra = make(map[string]any, 1)
|
||||
}
|
||||
out[0].Extra[agent.ModelFacingTraceVersionKey] = 1
|
||||
return out
|
||||
}
|
||||
|
||||
// einoExtractFallbackAssistantFromMsgs 在「主通道未产出助手正文」时,从 Eino ADK 轨迹中回填用户可见回复。
|
||||
// 典型场景:监督者仅调用 exit(final_result 落在 Tool 消息中),或工具结果已写入历史但 lastAssistant 未更新。
|
||||
//
|
||||
|
||||
@@ -11,15 +11,19 @@ import (
|
||||
// Order (best practice):
|
||||
// 1. system merge — accurate token count for summarization
|
||||
// 2. continuation user dedup — drop stale session-resume injections
|
||||
// 3. summarization
|
||||
// 4. orphan tool prune
|
||||
// 5. telemetry
|
||||
// 6. model-facing trace snapshot
|
||||
// 3. pre-summarization tool-call/result reconciliation
|
||||
// 4. summarization
|
||||
// 5. total model-input hard budget
|
||||
// 6. final tool-call/result reconciliation
|
||||
// 7. orphan tool prune (defense in depth)
|
||||
// 8. telemetry
|
||||
// 9. model-facing trace snapshot
|
||||
type einoChatModelTailConfig struct {
|
||||
logger *zap.Logger
|
||||
phase string
|
||||
summarization adk.ChatModelAgentMiddleware
|
||||
modelName string
|
||||
maxTotalTokens int
|
||||
conversationID string
|
||||
trace *modelFacingTraceHolder
|
||||
skipOrphanPruner bool
|
||||
@@ -31,8 +35,13 @@ func appendEinoChatModelTailMiddlewares(handlers []adk.ChatModelAgentMiddleware,
|
||||
handlers = append(handlers, newSystemMessageNormalizerMiddleware(cfg.logger, cfg.phase))
|
||||
handlers = append(handlers, newContinuationUserDedupMiddleware(cfg.logger, cfg.phase))
|
||||
if cfg.summarization != nil {
|
||||
// Summarization invokes the model internally, so its input needs the same
|
||||
// structural guarantee as the agent's final model call.
|
||||
handlers = append(handlers, newToolPairReconcilerMiddleware(cfg.logger, cfg.phase+"_pre_summarization"))
|
||||
handlers = append(handlers, cfg.summarization)
|
||||
}
|
||||
handlers = append(handlers, newModelInputBudgetMiddleware(cfg.maxTotalTokens, cfg.modelName, cfg.logger, cfg.phase))
|
||||
handlers = append(handlers, newToolPairReconcilerMiddleware(cfg.logger, cfg.phase))
|
||||
if !cfg.skipOrphanPruner {
|
||||
handlers = append(handlers, newOrphanToolPrunerMiddleware(cfg.logger, cfg.phase))
|
||||
}
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package multiagent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
@@ -9,7 +10,7 @@ import (
|
||||
|
||||
// newEinoExecuteMonitorCallbacks 在 Eino filesystem execute 开始/结束时写入 MCP 监控库并 recorder(executionId),
|
||||
// 与 CallTool 路径一致,使监控页能展示「执行中」状态。
|
||||
func newEinoExecuteMonitorCallbacks(ag *agent.Agent, recorder einomcp.ExecutionRecorder) (
|
||||
func newEinoExecuteMonitorCallbacks(ctx context.Context, ag *agent.Agent, recorder einomcp.ExecutionRecorder) (
|
||||
begin func(toolCallID, command string) string,
|
||||
finish func(executionID, toolCallID, command, stdout string, success bool, invokeErr error),
|
||||
) {
|
||||
@@ -18,7 +19,7 @@ func newEinoExecuteMonitorCallbacks(ag *agent.Agent, recorder einomcp.ExecutionR
|
||||
return ""
|
||||
}
|
||||
args := map[string]interface{}{"command": command}
|
||||
id := ag.BeginLocalToolExecution("execute", args)
|
||||
id := ag.BeginLocalToolExecution(ctx, "execute", args)
|
||||
if id != "" && recorder != nil {
|
||||
recorder(id, toolCallID)
|
||||
}
|
||||
@@ -37,7 +38,7 @@ func newEinoExecuteMonitorCallbacks(ag *agent.Agent, recorder einomcp.ExecutionR
|
||||
}
|
||||
}
|
||||
args := map[string]interface{}{"command": command}
|
||||
id := ag.FinishLocalToolExecution(executionID, "execute", args, stdout, err)
|
||||
id := ag.FinishLocalToolExecution(ctx, executionID, "execute", args, stdout, err)
|
||||
if id != "" && recorder != nil && executionID == "" {
|
||||
recorder(id, toolCallID)
|
||||
}
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package multiagent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"strings"
|
||||
@@ -65,6 +66,7 @@ func toolCallArgsFromAccumulated(msgs []adk.Message, toolCallID, expectToolName
|
||||
|
||||
// beginEinoADKFilesystemToolMonitor 在 Eino ADK filesystem 工具开始调用时写入 running 状态。
|
||||
func beginEinoADKFilesystemToolMonitor(
|
||||
ctx context.Context,
|
||||
ag *agent.Agent,
|
||||
rec einomcp.ExecutionRecorder,
|
||||
binder *MCPExecutionBinder,
|
||||
@@ -85,7 +87,7 @@ func beginEinoADKFilesystemToolMonitor(
|
||||
return
|
||||
}
|
||||
storedName := "eino_fs::" + strings.ToLower(name)
|
||||
id := ag.BeginLocalToolExecution(storedName, map[string]interface{}{})
|
||||
id := ag.BeginLocalToolExecution(ctx, storedName, map[string]interface{}{})
|
||||
if id == "" {
|
||||
return
|
||||
}
|
||||
@@ -97,6 +99,7 @@ func beginEinoADKFilesystemToolMonitor(
|
||||
|
||||
// recordEinoADKFilesystemToolMonitor 将 Eino ADK filesystem 中间件工具结果写入 MCP 监控(与 execute / MCP 桥芯片一致)。
|
||||
func recordEinoADKFilesystemToolMonitor(
|
||||
ctx context.Context,
|
||||
ag *agent.Agent,
|
||||
rec einomcp.ExecutionRecorder,
|
||||
binder *MCPExecutionBinder,
|
||||
@@ -131,7 +134,7 @@ func recordEinoADKFilesystemToolMonitor(
|
||||
if binder != nil {
|
||||
execID = binder.ExecutionID(toolCallID)
|
||||
}
|
||||
id := ag.FinishLocalToolExecution(execID, storedName, args, resultText, invErr)
|
||||
id := ag.FinishLocalToolExecution(ctx, execID, storedName, args, resultText, invErr)
|
||||
if id != "" && execID == "" {
|
||||
rec(id, toolCallID)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,34 @@
|
||||
package multiagent
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
func TestBuildEinoRunResultNeverPersistsRawAccumulationWithoutModelFacingTrace(t *testing.T) {
|
||||
raw := []schema.Message{*schema.ToolMessage(strings.Repeat("raw-tool-output", 1000), "call-1")}
|
||||
rawMsgs := make([]*schema.Message, len(raw))
|
||||
for i := range raw {
|
||||
rawMsgs[i] = &raw[i]
|
||||
}
|
||||
result := buildEinoRunResultFromAccumulated("deep", rawMsgs, nil, "", "", "empty", nil, true)
|
||||
if result.LastAgentTraceInput != "" {
|
||||
t.Fatalf("pre-model raw accumulation must not be persisted: %d bytes", len(result.LastAgentTraceInput))
|
||||
}
|
||||
|
||||
modelFacing := []*schema.Message{schema.UserMessage("bounded-model-view")}
|
||||
result = buildEinoRunResultFromAccumulated("deep", rawMsgs, modelFacing, "ok", "", "empty", nil, false)
|
||||
if !strings.Contains(result.LastAgentTraceInput, "bounded-model-view") {
|
||||
t.Fatalf("model-facing trace missing: %s", result.LastAgentTraceInput)
|
||||
}
|
||||
if strings.Contains(result.LastAgentTraceInput, "raw-tool-output") {
|
||||
t.Fatal("raw accumulation leaked into persisted model-facing trace")
|
||||
}
|
||||
if !agent.IsModelFacingTraceJSON(result.LastAgentTraceInput) {
|
||||
t.Fatal("persisted model-facing trace is missing its version marker")
|
||||
}
|
||||
}
|
||||
@@ -131,6 +131,7 @@ func buildPlanExecuteExecutorHandlers(ctx context.Context, a *PlanExecuteRootArg
|
||||
phase: "plan_execute_executor",
|
||||
summarization: sumMw,
|
||||
modelName: a.ModelName,
|
||||
maxTotalTokens: a.AppCfg.OpenAI.MaxTotalTokens,
|
||||
conversationID: a.ConversationID,
|
||||
trace: a.ModelFacingTrace,
|
||||
})
|
||||
|
||||
@@ -50,6 +50,7 @@ func RunEinoSingleChatModelAgent(
|
||||
if ma == nil {
|
||||
return nil, fmt.Errorf("eino single: multi_agent 配置为空")
|
||||
}
|
||||
runtimeUserMessage := prepareLatestUserMessageForModel(userMessage, appCfg, &ma.EinoMiddleware, conversationID, logger)
|
||||
|
||||
einoLoc, einoSkillMW, einoFSTools, skillsRoot, einoErr := prepareEinoSkills(ctx, appCfg.SkillsDir, ma, logger)
|
||||
if einoErr != nil {
|
||||
@@ -81,7 +82,7 @@ func RunEinoSingleChatModelAgent(
|
||||
}
|
||||
|
||||
toolInvokeNotify := einomcp.NewToolInvokeNotifyHolder()
|
||||
einoExecBegin, einoExecFinish := newEinoExecuteMonitorCallbacks(ag, recorder)
|
||||
einoExecBegin, einoExecFinish := newEinoExecuteMonitorCallbacks(ctx, ag, recorder)
|
||||
mainDefs := ag.ToolsForRole(roleTools)
|
||||
mainTools, err := einomcp.ToolsFromDefinitions(ag, holder, mainDefs, recorder, nil, toolInvokeNotify, einoSingleAgentName)
|
||||
if err != nil {
|
||||
@@ -149,6 +150,7 @@ func RunEinoSingleChatModelAgent(
|
||||
phase: "eino_single",
|
||||
summarization: mainSumMw,
|
||||
modelName: appCfg.OpenAI.Model,
|
||||
maxTotalTokens: appCfg.OpenAI.MaxTotalTokens,
|
||||
conversationID: conversationID,
|
||||
trace: modelFacingTrace,
|
||||
})
|
||||
@@ -160,6 +162,7 @@ func RunEinoSingleChatModelAgent(
|
||||
Tools: mainToolsForCfg,
|
||||
UnknownToolsHandler: einomcp.UnknownToolReminderHandler(),
|
||||
ToolCallMiddlewares: []compose.ToolMiddleware{
|
||||
localToolRBACMiddleware(),
|
||||
hitlToolCallMiddleware(),
|
||||
softRecoveryToolMiddleware(),
|
||||
},
|
||||
@@ -201,7 +204,7 @@ func RunEinoSingleChatModelAgent(
|
||||
}
|
||||
|
||||
baseMsgs := historyToMessages(history, appCfg, &ma.EinoMiddleware)
|
||||
baseMsgs = appendUserMessageIfNeeded(baseMsgs, userMessage)
|
||||
baseMsgs = appendUserMessageIfNeeded(baseMsgs, runtimeUserMessage)
|
||||
|
||||
streamsMainAssistant := func(agent string) bool {
|
||||
return agent == "" || agent == einoSingleAgentName
|
||||
|
||||
@@ -18,8 +18,9 @@ import (
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// prepareEinoSkills builds Eino official skill backend + middleware, and a shared local disk backend
|
||||
// for skill discovery and (optionally) filesystem/execute tools. Returns nils when disabled or dir missing.
|
||||
// prepareEinoSkills builds Eino official skill backend + middleware, and a shared local disk backend.
|
||||
// The local backend is also required by reduction, so reduction must not silently disappear merely
|
||||
// because Skills are disabled or skills_dir is unavailable.
|
||||
// skillsRoot is the absolute skills directory (empty when skills are not active).
|
||||
func prepareEinoSkills(
|
||||
ctx context.Context,
|
||||
@@ -27,15 +28,34 @@ func prepareEinoSkills(
|
||||
ma *config.MultiAgentConfig,
|
||||
logger *zap.Logger,
|
||||
) (loc *localbk.Local, skillMW adk.ChatModelAgentMiddleware, fsTools bool, skillsRoot string, err error) {
|
||||
if ma == nil || ma.EinoSkills.Disable {
|
||||
if ma == nil {
|
||||
return nil, nil, false, "", nil
|
||||
}
|
||||
needLocalBackend := ma.EinoMiddleware.ReductionEnable
|
||||
newLocalBackend := func() (*localbk.Local, error) {
|
||||
backend, backendErr := localbk.NewBackend(ctx, &localbk.Config{})
|
||||
if backendErr != nil {
|
||||
return nil, fmt.Errorf("eino local backend: %w", backendErr)
|
||||
}
|
||||
return backend, nil
|
||||
}
|
||||
if ma.EinoSkills.Disable {
|
||||
if !needLocalBackend {
|
||||
return nil, nil, false, "", nil
|
||||
}
|
||||
loc, err = newLocalBackend()
|
||||
return loc, nil, false, "", err
|
||||
}
|
||||
root := strings.TrimSpace(skillsDir)
|
||||
if root == "" {
|
||||
if logger != nil {
|
||||
logger.Warn("eino skills: skills_dir empty, skip")
|
||||
}
|
||||
return nil, nil, false, "", nil
|
||||
if !needLocalBackend {
|
||||
return nil, nil, false, "", nil
|
||||
}
|
||||
loc, err = newLocalBackend()
|
||||
return loc, nil, false, "", err
|
||||
}
|
||||
abs, err := filepath.Abs(root)
|
||||
if err != nil {
|
||||
@@ -45,12 +65,16 @@ func prepareEinoSkills(
|
||||
if logger != nil {
|
||||
logger.Warn("eino skills: directory missing, skip", zap.String("dir", abs), zap.Error(err))
|
||||
}
|
||||
return nil, nil, false, "", nil
|
||||
if !needLocalBackend {
|
||||
return nil, nil, false, "", nil
|
||||
}
|
||||
loc, err = newLocalBackend()
|
||||
return loc, nil, false, "", err
|
||||
}
|
||||
|
||||
loc, err = localbk.NewBackend(ctx, &localbk.Config{})
|
||||
loc, err = newLocalBackend()
|
||||
if err != nil {
|
||||
return nil, nil, false, "", fmt.Errorf("eino local backend: %w", err)
|
||||
return nil, nil, false, "", err
|
||||
}
|
||||
|
||||
skillBE, err := skill.NewBackendFromFilesystem(ctx, &skill.BackendFromFilesystemConfig{
|
||||
|
||||
@@ -0,0 +1,27 @@
|
||||
package multiagent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
)
|
||||
|
||||
func TestPrepareEinoSkillsStillCreatesReductionBackendWhenSkillsDisabled(t *testing.T) {
|
||||
ma := &config.MultiAgentConfig{
|
||||
EinoSkills: config.MultiAgentEinoSkillsConfig{Disable: true},
|
||||
EinoMiddleware: config.MultiAgentEinoMiddlewareConfig{
|
||||
ReductionEnable: true,
|
||||
},
|
||||
}
|
||||
loc, skillMW, fsTools, skillsRoot, err := prepareEinoSkills(context.Background(), "", ma, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if loc == nil {
|
||||
t.Fatal("reduction backend must exist even when Skills are disabled")
|
||||
}
|
||||
if skillMW != nil || fsTools || skillsRoot != "" {
|
||||
t.Fatalf("Skills unexpectedly enabled: mw=%v fs=%v root=%q", skillMW, fsTools, skillsRoot)
|
||||
}
|
||||
}
|
||||
@@ -14,11 +14,11 @@ import (
|
||||
"cyberstrike-ai/internal/project"
|
||||
|
||||
"github.com/bytedance/sonic"
|
||||
einoopenai "github.com/cloudwego/eino-ext/components/model/openai"
|
||||
"github.com/cloudwego/eino/adk"
|
||||
"github.com/cloudwego/eino/adk/middlewares/summarization"
|
||||
"github.com/cloudwego/eino/components/model"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
einoopenai "github.com/cloudwego/eino-ext/components/model/openai"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
@@ -98,10 +98,20 @@ func newEinoSummarizationMiddleware(
|
||||
}
|
||||
triggerRatio := 0.8
|
||||
emitInternalEvents := true
|
||||
userLedgerMaxRunes := config.DefaultSummarizationUserIntentLedgerMaxRunes
|
||||
userLedgerEntryMaxRunes := config.DefaultSummarizationUserIntentLedgerEntryMaxRunes
|
||||
if mwCfg != nil {
|
||||
triggerRatio = mwCfg.SummarizationTriggerRatioEffective()
|
||||
emitInternalEvents = mwCfg.SummarizationEmitInternalEventsEffective()
|
||||
userLedgerMaxRunes = mwCfg.SummarizationUserIntentLedgerMaxRunesEffective()
|
||||
userLedgerEntryMaxRunes = mwCfg.SummarizationUserIntentLedgerEntryMaxRunesEffective()
|
||||
}
|
||||
// The ledger is merged into the leading system message and cannot be removed as
|
||||
// an ordinary conversation round. Bound it relative to the configured window so
|
||||
// it cannot crowd out the summary/latest turn or trip the final fail-closed guard.
|
||||
ledgerWindowCap := modelFacingRuneBudget(maxTotal, 0.20)
|
||||
userLedgerMaxRunes = minPositiveInt(userLedgerMaxRunes, ledgerWindowCap)
|
||||
userLedgerEntryMaxRunes = minPositiveInt(userLedgerEntryMaxRunes, userLedgerMaxRunes)
|
||||
// Keep enough safety margin for tokenizer/model-side accounting mismatch.
|
||||
trigger := int(float64(maxTotal) * triggerRatio)
|
||||
if trigger < 4096 {
|
||||
@@ -127,6 +137,14 @@ func newEinoSummarizationMiddleware(
|
||||
if recentTrailMax > trigger/2 {
|
||||
recentTrailMax = trigger / 2
|
||||
}
|
||||
// The summarization request itself needs output headroom. A trigger is not a hard
|
||||
// request limit: one large turn can jump far beyond it. Bound the actual summary
|
||||
// model input to 60% of the configured context window and keep complete recent
|
||||
// rounds so tool_call/tool_result pairs are never split.
|
||||
summaryInputMax := int(float64(maxTotal) * 0.6)
|
||||
if summaryInputMax < 4096 {
|
||||
summaryInputMax = 4096
|
||||
}
|
||||
transcriptPath := ""
|
||||
if conv := strings.TrimSpace(conversationID); conv != "" {
|
||||
baseRoot := filepath.Join(os.TempDir(), "cyberstrike-summarization")
|
||||
@@ -164,6 +182,30 @@ func newEinoSummarizationMiddleware(
|
||||
mw, err := summarization.New(ctx, &summarization.Config{
|
||||
Model: summaryModel,
|
||||
ModelOptions: summaryModelOpts,
|
||||
GenModelInput: func(ctx context.Context, sysInstruction, userInstruction adk.Message, originalMsgs []adk.Message) ([]adk.Message, error) {
|
||||
if transcriptPath != "" && len(originalMsgs) > 0 {
|
||||
if werr := writeSummarizationTranscript(transcriptPath, originalMsgs); werr != nil && logger != nil {
|
||||
logger.Warn("eino summarization transcript preflight 写入失败",
|
||||
zap.String("path", transcriptPath), zap.Error(werr))
|
||||
}
|
||||
}
|
||||
input, dropped, berr := buildBudgetedSummarizationModelInput(
|
||||
ctx, sysInstruction, userInstruction, originalMsgs, tokenCounter, summaryInputMax,
|
||||
)
|
||||
if logger != nil && (berr != nil || dropped > 0) {
|
||||
fields := []zap.Field{
|
||||
zap.Int("max_input_tokens", summaryInputMax),
|
||||
zap.Int("dropped_rounds", dropped),
|
||||
}
|
||||
if berr != nil {
|
||||
fields = append(fields, zap.Error(berr))
|
||||
logger.Warn("eino summarization input budget failed", fields...)
|
||||
} else {
|
||||
logger.Info("eino summarization input bounded", fields...)
|
||||
}
|
||||
}
|
||||
return input, berr
|
||||
},
|
||||
Trigger: &summarization.TriggerCondition{
|
||||
ContextTokens: trigger,
|
||||
},
|
||||
@@ -190,10 +232,13 @@ func newEinoSummarizationMiddleware(
|
||||
},
|
||||
Finalize: func(ctx context.Context, originalMessages []adk.Message, summary adk.Message) ([]adk.Message, error) {
|
||||
summary = stripAnalysisFromSummarizationMessage(summary)
|
||||
out, ferr := summarizeFinalizeWithRecentAssistantToolTrail(ctx, originalMessages, summary, tokenCounter, recentTrailMax)
|
||||
userLedger := buildOriginalUserIntentLedgerMessage(originalMessages, userLedgerMaxRunes, userLedgerEntryMaxRunes)
|
||||
compactionMessages := stripOriginalUserIntentLedgerFromMessages(originalMessages)
|
||||
out, ferr := summarizeFinalizeWithRecentAssistantToolTrail(ctx, compactionMessages, summary, tokenCounter, recentTrailMax)
|
||||
if ferr != nil {
|
||||
return nil, ferr
|
||||
}
|
||||
out = mergeMessageIntoLeadingSystem(out, userLedger)
|
||||
if appCfg != nil {
|
||||
out = refreshFactIndexInMessages(out, db, projectID, appCfg.Project, logger)
|
||||
}
|
||||
@@ -230,6 +275,71 @@ func newEinoSummarizationMiddleware(
|
||||
return mw, nil
|
||||
}
|
||||
|
||||
// buildBudgetedSummarizationModelInput builds the exact payload sent to the summary model.
|
||||
// It retains the newest complete conversation rounds within budget and emits an explicit
|
||||
// marker when older rounds are omitted. The full pre-compaction transcript is persisted
|
||||
// separately; omitted raw messages never become model-facing history again.
|
||||
func buildBudgetedSummarizationModelInput(
|
||||
ctx context.Context,
|
||||
sysInstruction adk.Message,
|
||||
userInstruction adk.Message,
|
||||
originalMsgs []adk.Message,
|
||||
tokenCounter summarization.TokenCounterFunc,
|
||||
maxTokens int,
|
||||
) ([]adk.Message, int, error) {
|
||||
base := []adk.Message{sysInstruction, userInstruction}
|
||||
baseTokens, err := tokenCounter(ctx, &summarization.TokenCounterInput{Messages: base})
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
remaining := maxTokens - baseTokens
|
||||
markerTemplate := schema.UserMessage("[Context budget guard omitted older conversation rounds; summarize the retained recent rounds and preserve the omission marker.]")
|
||||
markerTokens, err := tokenCounter(ctx, &summarization.TokenCounterInput{Messages: []adk.Message{markerTemplate}})
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
remaining -= markerTokens
|
||||
if remaining < 0 {
|
||||
remaining = 0
|
||||
}
|
||||
|
||||
contextMsgs := make([]adk.Message, 0, len(originalMsgs))
|
||||
for _, msg := range originalMsgs {
|
||||
if msg != nil && msg.Role != schema.System {
|
||||
contextMsgs = append(contextMsgs, msg)
|
||||
}
|
||||
}
|
||||
rounds := splitMessagesIntoRounds(contextMsgs)
|
||||
selectedReverse := make([]messageRound, 0, len(rounds))
|
||||
used := 0
|
||||
for i := len(rounds) - 1; i >= 0; i-- {
|
||||
n, countErr := tokenCounter(ctx, &summarization.TokenCounterInput{Messages: rounds[i].messages})
|
||||
if countErr != nil {
|
||||
return nil, 0, countErr
|
||||
}
|
||||
if used+n > remaining {
|
||||
break
|
||||
}
|
||||
used += n
|
||||
selectedReverse = append(selectedReverse, rounds[i])
|
||||
}
|
||||
|
||||
dropped := len(rounds) - len(selectedReverse)
|
||||
input := make([]adk.Message, 0, 3+len(contextMsgs))
|
||||
input = append(input, sysInstruction)
|
||||
if dropped > 0 {
|
||||
input = append(input, schema.UserMessage(fmt.Sprintf(
|
||||
"[Context budget guard omitted %d older conversation round(s); summarize the retained recent rounds and preserve the omission marker.]",
|
||||
dropped,
|
||||
)))
|
||||
}
|
||||
for i := len(selectedReverse) - 1; i >= 0; i-- {
|
||||
input = append(input, selectedReverse[i].messages...)
|
||||
}
|
||||
input = append(input, userInstruction)
|
||||
return input, dropped, nil
|
||||
}
|
||||
|
||||
// refreshFactIndexInMessages 在 summarization 压缩后,用 DB 最新索引替换 system 中已有的项目黑板索引段。
|
||||
func refreshFactIndexInMessages(msgs []adk.Message, db *database.DB, projectID string, cfg config.ProjectConfig, logger *zap.Logger) []adk.Message {
|
||||
if db == nil || !cfg.Enabled {
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package multiagent
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strings"
|
||||
|
||||
@@ -11,6 +12,13 @@ import (
|
||||
var (
|
||||
summarizationAnalysisBlockRegex = regexp.MustCompile(`(?is)<analysis>\s*.*?\s*</analysis>`)
|
||||
summarizationSummaryBlockRegex = regexp.MustCompile(`(?is)<summary>\s*(.*?)\s*</summary>`)
|
||||
userIntentLedgerBlockRegex = regexp.MustCompile(`(?is)<original_user_intent_ledger>\s*(.*?)\s*</original_user_intent_ledger>`)
|
||||
userIntentLedgerSectionRegex = regexp.MustCompile(`(?is)\s*## 原始用户输入与约束账本(系统保真)\s*<original_user_intent_ledger>\s*.*?\s*</original_user_intent_ledger>\s*`)
|
||||
)
|
||||
|
||||
const (
|
||||
userIntentLedgerStartMarker = "<original_user_intent_ledger>"
|
||||
userIntentLedgerEndMarker = "</original_user_intent_ledger>"
|
||||
)
|
||||
|
||||
// stripAnalysisFromSummarizationMessage removes the <analysis> block from a post-processed
|
||||
@@ -71,3 +79,212 @@ func extractSummarizationSummaryBody(text string) (string, bool) {
|
||||
}
|
||||
return body, true
|
||||
}
|
||||
|
||||
// buildOriginalUserIntentLedgerMessage returns a deterministic, host-generated
|
||||
// context anchor for raw user inputs. Keep it separate from the model-generated
|
||||
// working summary so compaction can rewrite the summary without rewriting user intent.
|
||||
func buildOriginalUserIntentLedgerMessage(originalMessages []adk.Message, maxRunes, entryMaxRunes int) adk.Message {
|
||||
ledger := buildOriginalUserIntentLedger(originalMessages, maxRunes, entryMaxRunes)
|
||||
if strings.TrimSpace(ledger) == "" {
|
||||
return nil
|
||||
}
|
||||
return schema.SystemMessage(wrapUserIntentLedger(ledger))
|
||||
}
|
||||
|
||||
func mergeMessageIntoLeadingSystem(msgs []adk.Message, msg adk.Message) []adk.Message {
|
||||
if msg == nil {
|
||||
return msgs
|
||||
}
|
||||
for i, existing := range msgs {
|
||||
if existing == nil || existing.Role != schema.System {
|
||||
continue
|
||||
}
|
||||
cloned := *existing
|
||||
cloned.Content = strings.TrimSpace(strings.TrimSpace(cloned.Content) + "\n\n" + strings.TrimSpace(msg.Content))
|
||||
out := make([]adk.Message, len(msgs))
|
||||
copy(out, msgs)
|
||||
out[i] = &cloned
|
||||
return out
|
||||
}
|
||||
out := make([]adk.Message, 0, len(msgs)+1)
|
||||
out = append(out, msg)
|
||||
out = append(out, msgs...)
|
||||
return out
|
||||
}
|
||||
|
||||
func stripOriginalUserIntentLedgerFromMessages(msgs []adk.Message) []adk.Message {
|
||||
if len(msgs) == 0 {
|
||||
return msgs
|
||||
}
|
||||
out := make([]adk.Message, 0, len(msgs))
|
||||
for _, msg := range msgs {
|
||||
if msg == nil {
|
||||
continue
|
||||
}
|
||||
cloned := *msg
|
||||
cloned.Content = stripOriginalUserIntentLedgerFromText(cloned.Content)
|
||||
if len(cloned.UserInputMultiContent) > 0 {
|
||||
parts := make([]schema.MessageInputPart, len(cloned.UserInputMultiContent))
|
||||
copy(parts, cloned.UserInputMultiContent)
|
||||
for i := range parts {
|
||||
if parts[i].Type == schema.ChatMessagePartTypeText {
|
||||
parts[i].Text = stripOriginalUserIntentLedgerFromText(parts[i].Text)
|
||||
}
|
||||
}
|
||||
cloned.UserInputMultiContent = parts
|
||||
}
|
||||
if strings.TrimSpace(cloned.Content) == "" && len(cloned.UserInputMultiContent) == 0 && len(cloned.ToolCalls) == 0 && cloned.ReasoningContent == "" {
|
||||
continue
|
||||
}
|
||||
out = append(out, &cloned)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func buildOriginalUserIntentLedger(msgs []adk.Message, maxRunes, entryMaxRunes int) string {
|
||||
entries := collectOriginalUserIntentEntries(msgs)
|
||||
if len(entries) == 0 {
|
||||
return ""
|
||||
}
|
||||
var sb strings.Builder
|
||||
for i, entry := range entries {
|
||||
line := fmt.Sprintf("- [U%03d] %s\n", i+1, sanitizeUserIntentLedgerEntry(entry, entryMaxRunes))
|
||||
if maxRunes > 0 && utf8RuneLen(sb.String())+utf8RuneLen(line) > maxRunes {
|
||||
sb.WriteString("- [...truncated] 用户原始输入账本超过预算;完整压缩前记录见 summarization transcript。\n")
|
||||
break
|
||||
}
|
||||
sb.WriteString(line)
|
||||
}
|
||||
return strings.TrimSpace(sb.String())
|
||||
}
|
||||
|
||||
func collectOriginalUserIntentEntries(msgs []adk.Message) []string {
|
||||
seen := make(map[string]struct{})
|
||||
entries := make([]string, 0, 16)
|
||||
add := func(s string) {
|
||||
s = strings.TrimSpace(s)
|
||||
if s == "" || isSyntheticContinuationUserText(s) {
|
||||
return
|
||||
}
|
||||
key := normalizeUserIntentLedgerText(s)
|
||||
if key == "" {
|
||||
return
|
||||
}
|
||||
if _, ok := seen[key]; ok {
|
||||
return
|
||||
}
|
||||
seen[key] = struct{}{}
|
||||
entries = append(entries, s)
|
||||
}
|
||||
for _, msg := range msgs {
|
||||
if msg == nil {
|
||||
continue
|
||||
}
|
||||
for _, entry := range extractExistingUserIntentLedgerEntries(messageTextForLedger(msg)) {
|
||||
add(entry)
|
||||
}
|
||||
if msg.Role == schema.User {
|
||||
add(adkUserMessageText(msg))
|
||||
}
|
||||
}
|
||||
return entries
|
||||
}
|
||||
|
||||
func messageTextForLedger(msg adk.Message) string {
|
||||
if msg == nil {
|
||||
return ""
|
||||
}
|
||||
var b strings.Builder
|
||||
if strings.TrimSpace(msg.Content) != "" {
|
||||
b.WriteString(msg.Content)
|
||||
}
|
||||
for _, part := range msg.UserInputMultiContent {
|
||||
if part.Type != schema.ChatMessagePartTypeText || strings.TrimSpace(part.Text) == "" {
|
||||
continue
|
||||
}
|
||||
if b.Len() > 0 {
|
||||
b.WriteByte('\n')
|
||||
}
|
||||
b.WriteString(part.Text)
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
func extractExistingUserIntentLedgerEntries(text string) []string {
|
||||
matches := userIntentLedgerBlockRegex.FindAllStringSubmatch(text, -1)
|
||||
if len(matches) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]string, 0, len(matches))
|
||||
for _, match := range matches {
|
||||
if len(match) < 2 {
|
||||
continue
|
||||
}
|
||||
for _, line := range strings.Split(match[1], "\n") {
|
||||
line = strings.TrimSpace(line)
|
||||
if line == "" || strings.HasPrefix(line, "- [...truncated]") {
|
||||
continue
|
||||
}
|
||||
if strings.HasPrefix(line, "- [U") {
|
||||
if idx := strings.Index(line, "]"); idx >= 0 && idx+1 < len(line) {
|
||||
line = strings.TrimSpace(line[idx+1:])
|
||||
}
|
||||
}
|
||||
if line != "" {
|
||||
out = append(out, line)
|
||||
}
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func stripOriginalUserIntentLedgerFromText(text string) string {
|
||||
text = strings.TrimSpace(text)
|
||||
if text == "" {
|
||||
return ""
|
||||
}
|
||||
text = userIntentLedgerSectionRegex.ReplaceAllString(text, "\n")
|
||||
text = userIntentLedgerBlockRegex.ReplaceAllString(text, "\n")
|
||||
return strings.TrimSpace(text)
|
||||
}
|
||||
|
||||
func wrapUserIntentLedger(ledger string) string {
|
||||
return strings.TrimSpace("## 原始用户输入与约束账本(系统保真)\n" +
|
||||
userIntentLedgerStartMarker + "\n" +
|
||||
strings.TrimSpace(ledger) + "\n" +
|
||||
userIntentLedgerEndMarker)
|
||||
}
|
||||
|
||||
func sanitizeUserIntentLedgerEntry(s string, maxRunes int) string {
|
||||
s = strings.ReplaceAll(strings.TrimSpace(s), userIntentLedgerStartMarker, "[ledger-start]")
|
||||
s = strings.ReplaceAll(s, userIntentLedgerEndMarker, "[ledger-end]")
|
||||
s = strings.Join(strings.Fields(s), " ")
|
||||
if maxRunes > 0 && utf8RuneLen(s) > maxRunes {
|
||||
return truncateUserIntentLedgerRunes(s, maxRunes) + "..."
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func normalizeUserIntentLedgerText(s string) string {
|
||||
return strings.Join(strings.Fields(strings.TrimSpace(s)), " ")
|
||||
}
|
||||
|
||||
func isSyntheticContinuationUserText(s string) bool {
|
||||
return strings.Contains(s, continuationSessionMarker) ||
|
||||
strings.Contains(s, "【系统自动续跑 / Auto resume】")
|
||||
}
|
||||
|
||||
func utf8RuneLen(s string) int {
|
||||
return len([]rune(s))
|
||||
}
|
||||
|
||||
func truncateUserIntentLedgerRunes(s string, n int) string {
|
||||
if n <= 0 {
|
||||
return ""
|
||||
}
|
||||
runes := []rune(s)
|
||||
if len(runes) <= n {
|
||||
return s
|
||||
}
|
||||
return string(runes[:n])
|
||||
}
|
||||
|
||||
@@ -4,6 +4,7 @@ import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/cloudwego/eino/adk"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
@@ -65,3 +66,110 @@ func TestStripAnalysisFromSummarizationText_NoAnalysisUnchanged(t *testing.T) {
|
||||
t.Fatalf("expected unchanged text, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildOriginalUserIntentLedgerMessage_AppendsRawUserMessages(t *testing.T) {
|
||||
original := []adk.Message{
|
||||
schema.UserMessage("第一轮:只测 staging,不要碰 prod。"),
|
||||
schema.AssistantMessage("ok", nil),
|
||||
schema.UserMessage("第二轮:优先验证 /api/login 的 SQL 注入。"),
|
||||
schema.UserMessage(FormatEmptyResponseContinueUserMessage()),
|
||||
}
|
||||
|
||||
out := buildOriginalUserIntentLedgerMessage(original, 96000, 16000)
|
||||
if out == nil {
|
||||
t.Fatal("ledger message should be non-nil")
|
||||
}
|
||||
if out.Role != schema.System {
|
||||
t.Fatalf("ledger should be a system anchor, got %s", out.Role)
|
||||
}
|
||||
body := out.Content
|
||||
if !strings.Contains(body, userIntentLedgerStartMarker) || !strings.Contains(body, userIntentLedgerEndMarker) {
|
||||
t.Fatalf("ledger markers missing: %q", body)
|
||||
}
|
||||
if !strings.Contains(body, "只测 staging,不要碰 prod") {
|
||||
t.Fatalf("first user constraint missing: %q", body)
|
||||
}
|
||||
if !strings.Contains(body, "优先验证 /api/login") {
|
||||
t.Fatalf("second user request missing: %q", body)
|
||||
}
|
||||
if strings.Contains(body, "系统自动续跑") {
|
||||
t.Fatalf("synthetic auto-resume user message should be skipped: %q", body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildOriginalUserIntentLedgerMessage_CarriesPreviousLedgerAndDedups(t *testing.T) {
|
||||
prevSummary := schema.AssistantMessage(wrapUserIntentLedger("- [U001] 原始目标:example.com\n- [U002] 禁止高危破坏性操作"), nil)
|
||||
original := []adk.Message{
|
||||
prevSummary,
|
||||
schema.UserMessage("禁止高危破坏性操作"),
|
||||
schema.UserMessage("新增约束:只输出中文报告"),
|
||||
}
|
||||
|
||||
out := buildOriginalUserIntentLedgerMessage(original, 96000, 16000)
|
||||
body := out.Content
|
||||
if !strings.Contains(body, "原始目标:example.com") || !strings.Contains(body, "新增约束:只输出中文报告") {
|
||||
t.Fatalf("ledger did not carry old and new entries: %q", body)
|
||||
}
|
||||
if strings.Count(body, "禁止高危破坏性操作") != 1 {
|
||||
t.Fatalf("duplicate ledger entry was not deduped: %q", body)
|
||||
}
|
||||
if strings.Count(body, userIntentLedgerStartMarker) != 1 {
|
||||
t.Fatalf("expected one ledger block: %q", body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStripOriginalUserIntentLedgerFromMessages_RemovesOldLedgerBeforeRebuild(t *testing.T) {
|
||||
msgs := []adk.Message{
|
||||
schema.SystemMessage("sys\n\n" + wrapUserIntentLedger("- [U001] old goal")),
|
||||
schema.AssistantMessage("summary\n\n"+wrapUserIntentLedger("- [U001] old goal"), nil),
|
||||
}
|
||||
|
||||
out := stripOriginalUserIntentLedgerFromMessages(msgs)
|
||||
if len(out) != 2 {
|
||||
t.Fatalf("expected 2 messages, got %d", len(out))
|
||||
}
|
||||
for _, msg := range out {
|
||||
if strings.Contains(msg.Content, userIntentLedgerStartMarker) || strings.Contains(msg.Content, "old goal") {
|
||||
t.Fatalf("old ledger leaked after strip: %q", msg.Content)
|
||||
}
|
||||
}
|
||||
if out[0].Content != "sys" {
|
||||
t.Fatalf("non-ledger system content should remain: %q", out[0].Content)
|
||||
}
|
||||
if out[1].Content != "summary" {
|
||||
t.Fatalf("non-ledger summary content should remain: %q", out[1].Content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMergeMessageIntoLeadingSystem_KeepsLedgerSeparateFromSummary(t *testing.T) {
|
||||
sys := schema.SystemMessage("sys")
|
||||
ledger := schema.SystemMessage(wrapUserIntentLedger("- [U001] goal"))
|
||||
summary := schema.AssistantMessage("<summary>work state</summary>", nil)
|
||||
|
||||
out := mergeMessageIntoLeadingSystem([]adk.Message{sys, summary}, ledger)
|
||||
if len(out) != 2 {
|
||||
t.Fatalf("expected 2 messages, got %d", len(out))
|
||||
}
|
||||
if out[0].Role != schema.System || !strings.Contains(out[0].Content, "sys") || !strings.Contains(out[0].Content, userIntentLedgerStartMarker) {
|
||||
t.Fatalf("ledger should be merged into leading system: %+v", out[0])
|
||||
}
|
||||
if out[1] != summary {
|
||||
t.Fatalf("summary should remain second message")
|
||||
}
|
||||
if strings.Contains(out[1].Content, userIntentLedgerStartMarker) {
|
||||
t.Fatalf("summary should not carry ledger: %q", out[1].Content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMergeMessageIntoLeadingSystem_NoSystemPrependsLedger(t *testing.T) {
|
||||
ledger := schema.SystemMessage(wrapUserIntentLedger("- [U001] goal"))
|
||||
summary := schema.AssistantMessage("<summary>work state</summary>", nil)
|
||||
|
||||
out := mergeMessageIntoLeadingSystem([]adk.Message{summary}, ledger)
|
||||
if len(out) != 2 {
|
||||
t.Fatalf("expected 2 messages, got %d", len(out))
|
||||
}
|
||||
if out[0] != ledger || out[1] != summary {
|
||||
t.Fatalf("ledger should be prepended when no system exists: %+v", out)
|
||||
}
|
||||
}
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user