From 3aee7022c407382e382611fbcf9a8af955352fe3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=85=AC=E6=98=8E?= <83812544+Ed1s0nZ@users.noreply.github.com> Date: Sun, 11 Jan 2026 02:03:33 +0800 Subject: [PATCH] Add files via upload --- README.md | 50 +- README_CN.md | 50 +- internal/agent/agent.go | 43 +- internal/app/app.go | 15 +- internal/config/config.go | 135 ++++ internal/handler/agent.go | 75 +- internal/handler/config.go | 126 ++- internal/handler/role.go | 453 +++++++++++ internal/knowledge/retriever.go | 16 +- internal/knowledge/tool.go | 23 +- internal/mcp/builtin/constants.go | 33 + roles/API安全测试.yaml | 20 + roles/CTF.yaml | 33 + roles/Web应用扫描.yaml | 25 + roles/Web框架测试.yaml | 19 + roles/二进制分析.yaml | 31 + roles/云安全审计.yaml | 17 + roles/信息收集.yaml | 31 + roles/后渗透测试.yaml | 23 + roles/容器安全.yaml | 18 + roles/数字取证.yaml | 24 + roles/渗透测试.yaml | 33 + roles/综合漏洞扫描.yaml | 23 + roles/默认.yaml | 5 + web/static/css/style.css | 1005 ++++++++++++++++++++++- web/static/js/builtin-tools.js | 27 + web/static/js/chat.js | 63 +- web/static/js/roles.js | 1230 +++++++++++++++++++++++++++++ web/static/js/router.js | 25 +- web/templates/index.html | 174 +++- 30 files changed, 3759 insertions(+), 86 deletions(-) create mode 100644 internal/handler/role.go create mode 100644 internal/mcp/builtin/constants.go create mode 100644 roles/API安全测试.yaml create mode 100644 roles/CTF.yaml create mode 100644 roles/Web应用扫描.yaml create mode 100644 roles/Web框架测试.yaml create mode 100644 roles/二进制分析.yaml create mode 100644 roles/云安全审计.yaml create mode 100644 roles/信息收集.yaml create mode 100644 roles/后渗透测试.yaml create mode 100644 roles/容器安全.yaml create mode 100644 roles/数字取证.yaml create mode 100644 roles/渗透测试.yaml create mode 100644 roles/综合漏洞扫描.yaml create mode 100644 roles/默认.yaml create mode 100644 web/static/js/builtin-tools.js create mode 100644 web/static/js/roles.js diff --git a/README.md b/README.md index 5dd6b077..ad3ac528 100644 --- a/README.md +++ b/README.md @@ -42,6 +42,7 @@ CyberStrikeAI is an **AI-native security testing platform** built in Go. It inte - 📁 Conversation grouping with pinning, rename, and batch management - 🛡️ Vulnerability management with CRUD operations, severity tracking, status workflow, and statistics - 📋 Batch task management: create task queues, add multiple tasks, and execute them sequentially +- 🎭 Role-based testing: predefined security testing roles (Penetration Testing, CTF, Web App Scanning, etc.) with custom prompts and tool restrictions ## Tool Overview @@ -121,6 +122,7 @@ go build -o cyberstrike-ai cmd/server/main.go ### Core Workflows - **Conversation testing** – Natural-language prompts trigger toolchains with streaming SSE output. +- **Role-based testing** – Select from predefined security testing roles (Penetration Testing, CTF, Web App Scanning, API Security Testing, etc.) to customize AI behavior and tool availability. Each role applies custom system prompts and can restrict available tools for focused testing scenarios. - **Tool monitor** – Inspect running jobs, execution logs, and large-result attachments. - **History & audit** – Every conversation and tool invocation is stored in SQLite with replay. - **Conversation groups** – Organize conversations into groups, pin important groups, rename or delete groups via context menu. @@ -136,6 +138,28 @@ go build -o cyberstrike-ai cmd/server/main.go ## Advanced Usage +### Role-Based Testing +- **Predefined roles** – System includes 12+ predefined security testing roles (Penetration Testing, CTF, Web App Scanning, API Security Testing, Binary Analysis, Cloud Security Audit, etc.) in the `roles/` directory. +- **Custom prompts** – Each role can define a `user_prompt` that prepends to user messages, guiding the AI to adopt specialized testing methodologies and focus areas. +- **Tool restrictions** – Roles can specify a `tools` list to limit available tools, ensuring focused testing workflows (e.g., CTF role restricts to CTF-specific utilities). +- **Easy role creation** – Create custom roles by adding YAML files to the `roles/` directory. Each role defines `name`, `description`, `user_prompt`, `icon`, `tools`, and `enabled` fields. +- **Web UI integration** – Select roles from a dropdown in the chat interface. Role selection affects both AI behavior and available tool suggestions. + +**Creating a custom role (example):** +1. Create a YAML file in `roles/` (e.g., `roles/custom-role.yaml`): + ```yaml + name: Custom Role + description: Specialized testing scenario + user_prompt: You are a specialized security tester focusing on API security... + icon: "\U0001F4E1" + tools: + - api-fuzzer + - arjun + - graphql-scanner + enabled: true + ``` +2. Restart the server or reload configuration; the role appears in the role selector dropdown. + ### Tool Orchestration & Extensions - **YAML recipes** in `tools/*.yaml` describe commands, arguments, prompts, and metadata. - **Directory hot-reload** – pointing `security.tools_dir` to a folder is usually enough; inline definitions in `config.yaml` remain supported for quick experiments. @@ -292,7 +316,8 @@ A test SSE MCP server is available at `cmd/test-sse-mcp-server/` for validation ### Automation Hooks -- **REST APIs** – everything the UI uses (auth, conversations, tool runs, monitor, vulnerabilities) is available over JSON. +- **REST APIs** – everything the UI uses (auth, conversations, tool runs, monitor, vulnerabilities, roles) is available over JSON. +- **Role APIs** – manage security testing roles via `/api/roles` endpoints: `GET /api/roles` (list all roles), `GET /api/roles/:name` (get role), `POST /api/roles` (create role), `PUT /api/roles/:name` (update role), `DELETE /api/roles/:name` (delete role). Roles are stored as YAML files in the `roles/` directory and support hot-reload. - **Vulnerability APIs** – manage vulnerabilities via `/api/vulnerabilities` endpoints: `GET /api/vulnerabilities` (list with filters), `POST /api/vulnerabilities` (create), `GET /api/vulnerabilities/:id` (get), `PUT /api/vulnerabilities/:id` (update), `DELETE /api/vulnerabilities/:id` (delete), `GET /api/vulnerabilities/stats` (statistics). - **Batch Task APIs** – manage batch task queues via `/api/batch-tasks` endpoints: `POST /api/batch-tasks` (create queue), `GET /api/batch-tasks` (list queues), `GET /api/batch-tasks/:queueId` (get queue), `POST /api/batch-tasks/:queueId/start` (start execution), `POST /api/batch-tasks/:queueId/cancel` (cancel), `DELETE /api/batch-tasks/:queueId` (delete), `POST /api/batch-tasks/:queueId/tasks` (add task), `PUT /api/batch-tasks/:queueId/tasks/:taskId` (update task), `DELETE /api/batch-tasks/:queueId/tasks/:taskId` (delete task). Tasks execute sequentially, each creating a separate conversation with full status tracking. - **Task control** – pause/resume/stop long scans, re-run steps with new params, or stream transcripts. @@ -335,6 +360,7 @@ knowledge: top_k: 5 # Number of top results to return similarity_threshold: 0.7 # Minimum similarity score (0-1) hybrid_weight: 0.7 # Weight for vector search (1.0 = pure vector, 0.0 = pure keyword) +roles_dir: "roles" # Role configuration directory (relative to config file) ``` ### Tool Definition Example (`tools/nmap.yaml`) @@ -357,6 +383,26 @@ parameters: description: "Range, e.g. 1-1000" ``` +### Role Definition Example (`roles/penetration-testing.yaml`) + +```yaml +name: Penetration Testing +description: Professional penetration testing expert for comprehensive security testing +user_prompt: You are a professional cybersecurity penetration testing expert. Please use professional penetration testing methods and tools to conduct comprehensive security testing on targets, including but not limited to SQL injection, XSS, CSRF, file inclusion, command execution and other common vulnerabilities. +icon: "\U0001F3AF" +tools: + - nmap + - sqlmap + - nuclei + - burpsuite + - metasploit + - httpx + - record_vulnerability + - list_knowledge_risk_types + - search_knowledge_base +enabled: true +``` + ## Project Layout ``` @@ -365,6 +411,7 @@ CyberStrikeAI/ ├── internal/ # Agent, MCP core, handlers, security executor ├── web/ # Static SPA + templates ├── tools/ # YAML tool recipes (100+ examples provided) +├── roles/ # Role configurations (12+ predefined security testing roles) ├── img/ # Docs screenshots & diagrams ├── config.yaml # Runtime configuration ├── run.sh # Convenience launcher @@ -392,6 +439,7 @@ Build an attack chain for the latest engagement and export the node list with se ## Changelog (Recent) +- 2026-01-11 – Added role-based testing feature: predefined security testing roles with custom system prompts and tool restrictions. Users can select roles (Penetration Testing, CTF, Web App Scanning, etc.) from the chat interface to customize AI behavior and available tools. Roles are defined as YAML files in the `roles/` directory with support for hot-reload. - 2026-01-08 – Added SSE (Server-Sent Events) transport mode support for external MCP servers. External MCP federation now supports HTTP, stdio, and SSE modes. SSE mode enables real-time streaming communication for push-based scenarios. - 2026-01-01 – Added batch task management feature: create task queues with multiple tasks, add/edit/delete tasks before execution, and execute them sequentially. Each task runs as a separate conversation with status tracking (pending/running/completed/failed/cancelled). All queues and tasks are persisted in the database. - 2025-12-25 – Added vulnerability management feature: full CRUD operations for tracking vulnerabilities discovered during testing. Supports severity levels (critical/high/medium/low/info), status workflow (open/confirmed/fixed/false_positive), filtering by conversation/severity/status, and comprehensive statistics dashboard. diff --git a/README_CN.md b/README_CN.md index f4fa1a30..3560b073 100644 --- a/README_CN.md +++ b/README_CN.md @@ -41,6 +41,7 @@ CyberStrikeAI 是一款 **AI 原生安全测试平台**，基于 Go 构建，集 - 📁 对话分组管理：支持分组创建、置顶、重命名、删除等操作 - 🛡️ 漏洞管理功能：完整的漏洞 CRUD 操作，支持严重程度分级、状态流转、按对话/严重程度/状态过滤，以及统计看板 - 📋 批量任务管理：创建任务队列，批量添加任务，依次顺序执行，支持任务编辑与状态跟踪 +- 🎭 角色化测试：预设安全测试角色（渗透测试、CTF、Web 应用扫描等），支持自定义提示词和工具限制 ## 工具概览 @@ -120,6 +121,7 @@ go build -o cyberstrike-ai cmd/server/main.go ### 常用流程 - **对话测试**：自然语言触发多步工具编排，SSE 实时输出。 +- **角色化测试**：从预设的安全测试角色（渗透测试、CTF、Web 应用扫描、API 安全测试等）中选择，自定义 AI 行为和可用工具。每个角色可应用自定义系统提示词，并可限制可用工具列表，实现聚焦的测试场景。 - **工具监控**：查看任务队列、执行日志、大文件附件。 - **会话历史**：所有对话与工具调用保存在 SQLite，可随时重放。 - **对话分组**：将对话按项目或主题组织到不同分组，支持置顶、重命名、删除等操作，所有数据持久化存储。 @@ -135,6 +137,28 @@ go build -o cyberstrike-ai cmd/server/main.go ## 进阶使用 +### 角色化测试 +- **预设角色**：系统内置 12+ 个预设的安全测试角色（渗透测试、CTF、Web 应用扫描、API 安全测试、二进制分析、云安全审计等），位于 `roles/` 目录。 +- **自定义提示词**：每个角色可定义 `user_prompt`，会在用户消息前自动添加，引导 AI 采用特定的测试方法和关注重点。 +- **工具限制**：角色可指定 `tools` 列表，限制可用工具，实现聚焦的测试流程（如 CTF 角色限制为 CTF 专用工具）。 +- **轻松创建角色**：通过在 `roles/` 目录添加 YAML 文件即可创建自定义角色。每个角色定义 `name`、`description`、`user_prompt`、`icon`、`tools`、`enabled` 字段。 +- **Web 界面集成**：在聊天界面通过下拉菜单选择角色。角色选择会影响 AI 行为和可用工具建议。 + +**创建自定义角色示例：** +1. 在 `roles/` 目录创建 YAML 文件（如 `roles/custom-role.yaml`）： + ```yaml + name: 自定义角色 + description: 专用测试场景 + user_prompt: 你是一个专注于 API 安全的专业安全测试人员... + icon: "\U0001F4E1" + tools: + - api-fuzzer + - arjun + - graphql-scanner + enabled: true + ``` +2. 重启服务或重新加载配置，角色会出现在角色选择下拉菜单中。 + ### 工具编排与扩展 - `tools/*.yaml` 定义命令、参数、提示词与元数据，可热加载。 - `security.tools_dir` 指向目录即可批量启用；仍支持在主配置里内联定义。 @@ -291,7 +315,8 @@ CyberStrikeAI 支持通过三种传输模式连接外部 MCP 服务器： ### 自动化与安全 -- **REST API**：认证、会话、任务、监控、漏洞管理等接口全部开放，可与 CI/CD 集成。 +- **REST API**：认证、会话、任务、监控、漏洞管理、角色管理等接口全部开放，可与 CI/CD 集成。 +- **角色管理 API**：通过 `/api/roles` 端点管理安全测试角色：`GET /api/roles`（列表）、`GET /api/roles/:name`（获取角色）、`POST /api/roles`（创建角色）、`PUT /api/roles/:name`（更新角色）、`DELETE /api/roles/:name`（删除角色）。角色以 YAML 文件形式存储在 `roles/` 目录，支持热加载。 - **漏洞管理 API**：通过 `/api/vulnerabilities` 端点管理漏洞：`GET /api/vulnerabilities`（列表，支持过滤）、`POST /api/vulnerabilities`（创建）、`GET /api/vulnerabilities/:id`（获取）、`PUT /api/vulnerabilities/:id`（更新）、`DELETE /api/vulnerabilities/:id`（删除）、`GET /api/vulnerabilities/stats`（统计）。 - **批量任务 API**：通过 `/api/batch-tasks` 端点管理批量任务队列：`POST /api/batch-tasks`（创建队列）、`GET /api/batch-tasks`（列表）、`GET /api/batch-tasks/:queueId`（获取队列）、`POST /api/batch-tasks/:queueId/start`（开始执行）、`POST /api/batch-tasks/:queueId/cancel`（取消）、`DELETE /api/batch-tasks/:queueId`（删除队列）、`POST /api/batch-tasks/:queueId/tasks`（添加任务）、`PUT /api/batch-tasks/:queueId/tasks/:taskId`（更新任务）、`DELETE /api/batch-tasks/:queueId/tasks/:taskId`（删除任务）。任务依次顺序执行，每个任务创建独立对话，支持完整状态跟踪。 - **任务控制**：支持暂停/终止长任务、修改参数后重跑、流式获取日志。 @@ -334,6 +359,7 @@ knowledge: top_k: 5 # 检索返回的 Top-K 结果数量 similarity_threshold: 0.7 # 相似度阈值（0-1），低于此值的结果将被过滤 hybrid_weight: 0.7 # 混合检索权重（0-1），向量检索的权重，1.0 表示纯向量检索，0.0 表示纯关键词检索 +roles_dir: "roles" # 角色配置文件目录（相对于配置文件所在目录） ``` ### 工具模版示例（`tools/nmap.yaml`） @@ -356,6 +382,26 @@ parameters: description: "端口范围，如 1-1000" ``` +### 角色配置示例（`roles/渗透测试.yaml`） + +```yaml +name: 渗透测试 +description: 专业渗透测试专家，全面深入的漏洞检测 +user_prompt: 你是一个专业的网络安全渗透测试专家。请使用专业的渗透测试方法和工具，对目标进行全面的安全测试，包括但不限于SQL注入、XSS、CSRF、文件包含、命令执行等常见漏洞。 +icon: "\U0001F3AF" +tools: + - nmap + - sqlmap + - nuclei + - burpsuite + - metasploit + - httpx + - record_vulnerability + - list_knowledge_risk_types + - search_knowledge_base +enabled: true +``` + ## 项目结构 ``` @@ -364,6 +410,7 @@ CyberStrikeAI/ ├── internal/ # Agent、MCP 核心、路由与执行器 ├── web/ # 前端静态资源与模板 ├── tools/ # YAML 工具目录（含 100+ 示例） +├── roles/ # 角色配置文件目录（含 12+ 预设安全测试角色） ├── img/ # 文档配图 ├── config.yaml # 运行配置 ├── run.sh # 启动脚本 @@ -390,6 +437,7 @@ CyberStrikeAI/ ``` ## Changelog（近期） +- 2026-01-11 —— 新增角色化测试功能：预设安全测试角色，支持自定义系统提示词和工具限制。用户可在聊天界面选择角色（渗透测试、CTF、Web 应用扫描等），以自定义 AI 行为和可用工具。角色以 YAML 文件形式定义在 `roles/` 目录，支持热加载。 - 2026-01-08 —— 新增 SSE（Server-Sent Events）传输模式支持，外部 MCP 联邦现支持 HTTP、stdio 和 SSE 三种模式。SSE 模式支持实时流式通信，适用于基于推送的场景。 - 2026-01-01 —— 新增批量任务管理功能：支持创建任务队列，批量添加多个任务，执行前可编辑或删除任务，然后依次顺序执行。每个任务作为独立对话运行，支持状态跟踪（待执行/执行中/已完成/失败/已取消），所有队列和任务数据持久化存储到数据库。 - 2025-12-25 —— 新增漏洞管理功能：完整的漏洞 CRUD 操作，支持跟踪测试过程中发现的漏洞。支持严重程度分级（严重/高/中/低/信息）、状态流转（待确认/已确认/已修复/误报）、按对话/严重程度/状态过滤，以及统计看板。 diff --git a/internal/agent/agent.go b/internal/agent/agent.go index a319890c..b69c4368 100644 --- a/internal/agent/agent.go +++ b/internal/agent/agent.go @@ -12,6 +12,7 @@ import ( "cyberstrike-ai/internal/config" "cyberstrike-ai/internal/mcp" + "cyberstrike-ai/internal/mcp/builtin" "cyberstrike-ai/internal/openai" "cyberstrike-ai/internal/storage" @@ -302,16 +303,16 @@ type ProgressCallback func(eventType, message string, data interface{}) // AgentLoop 执行Agent循环 func (a *Agent) AgentLoop(ctx context.Context, userInput string, historyMessages []ChatMessage) (*AgentLoopResult, error) { - return a.AgentLoopWithProgress(ctx, userInput, historyMessages, "", nil) + return a.AgentLoopWithProgress(ctx, userInput, historyMessages, "", nil, nil) } // AgentLoopWithConversationID 执行Agent循环（带对话ID） func (a *Agent) AgentLoopWithConversationID(ctx context.Context, userInput string, historyMessages []ChatMessage, conversationID string) (*AgentLoopResult, error) { - return a.AgentLoopWithProgress(ctx, userInput, historyMessages, conversationID, nil) + return a.AgentLoopWithProgress(ctx, userInput, historyMessages, conversationID, nil, nil) } // AgentLoopWithProgress 执行Agent循环（带进度回调和对话ID） -func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, historyMessages []ChatMessage, conversationID string, callback ProgressCallback) (*AgentLoopResult, error) { +func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, historyMessages []ChatMessage, conversationID string, callback ProgressCallback, roleTools []string) (*AgentLoopResult, error) { // 设置当前对话ID a.mu.Lock() a.currentConversationID = conversationID @@ -401,8 +402,8 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his 当工具返回错误时，错误信息会包含在工具响应中，请仔细阅读并做出合理的决策。 漏洞记录要求： -- 当你发现有效漏洞时，必须使用 record_vulnerability 工具记录漏洞详情 -- 漏洞记录应包含：标题、描述、严重程度、类型、目标、证明（POC）、影响和修复建议 +- 当你发现有效漏洞时，必须使用 ` + builtin.ToolRecordVulnerability + ` 工具记录漏洞详情 +` + `- 漏洞记录应包含：标题、描述、严重程度、类型、目标、证明（POC）、影响和修复建议 - 严重程度评估标准： * critical（严重）：可导致系统完全被控制、数据泄露、服务中断等 * high（高）：可导致敏感信息泄露、权限提升、重要功能被绕过等 @@ -512,7 +513,7 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his } // 获取可用工具 - tools := a.getAvailableTools() + tools := a.getAvailableTools(roleTools) // 记录当前上下文的Token用量，展示压缩器运行状态 if a.memoryCompressor != nil { @@ -837,13 +838,29 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his // getAvailableTools 获取可用工具 // 从MCP服务器动态获取工具列表，使用简短描述以减少token消耗 -func (a *Agent) getAvailableTools() []Tool { +// roleTools: 角色配置的工具列表（toolKey格式），如果为空或nil，则使用所有工具（默认角色） +func (a *Agent) getAvailableTools(roleTools []string) []Tool { + // 构建角色工具集合（用于快速查找） + roleToolSet := make(map[string]bool) + if len(roleTools) > 0 { + for _, toolKey := range roleTools { + roleToolSet[toolKey] = true + } + } + // 从MCP服务器获取所有已注册的内部工具 mcpTools := a.mcpServer.GetAllTools() // 转换为OpenAI格式的工具定义 tools := make([]Tool, 0, len(mcpTools)) for _, mcpTool := range mcpTools { + // 如果指定了角色工具列表，只添加在列表中的工具 + if len(roleToolSet) > 0 { + toolKey := mcpTool.Name // 内置工具使用工具名称作为key + if !roleToolSet[toolKey] { + continue // 不在角色工具列表中，跳过 + } + } // 使用简短描述（如果存在），否则使用详细描述 description := mcpTool.ShortDescription if description == "" { @@ -883,6 +900,16 @@ func (a *Agent) getAvailableTools() []Tool { // 将外部MCP工具添加到工具列表（只添加启用的工具） for _, externalTool := range externalTools { + // 外部工具使用 "mcpName::toolName" 作为toolKey + externalToolKey := externalTool.Name + + // 如果指定了角色工具列表，只添加在列表中的工具 + if len(roleToolSet) > 0 { + if !roleToolSet[externalToolKey] { + continue // 不在角色工具列表中，跳过 + } + } + // 解析工具名称：mcpName::toolName var mcpName, actualToolName string if idx := strings.Index(externalTool.Name, "::"); idx > 0 { @@ -1136,7 +1163,7 @@ func (a *Agent) executeToolViaMCP(ctx context.Context, toolName string, args map ) // 如果是record_vulnerability工具，自动添加conversation_id - if toolName == "record_vulnerability" { + if toolName == builtin.ToolRecordVulnerability { a.mu.RLock() conversationID := a.currentConversationID a.mu.RUnlock() diff --git a/internal/app/app.go b/internal/app/app.go index a566e2f7..f283cd97 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -16,6 +16,7 @@ import ( "cyberstrike-ai/internal/knowledge" "cyberstrike-ai/internal/logger" "cyberstrike-ai/internal/mcp" + "cyberstrike-ai/internal/mcp/builtin" "cyberstrike-ai/internal/openai" "cyberstrike-ai/internal/security" "cyberstrike-ai/internal/storage" @@ -278,7 +279,7 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) { } // 创建处理器 - agentHandler := handler.NewAgentHandler(agent, db, log.Logger) + agentHandler := handler.NewAgentHandler(agent, db, cfg, log.Logger) // 如果知识库已启用，设置知识库管理器到AgentHandler以便记录检索日志 if knowledgeManager != nil { agentHandler.SetKnowledgeManager(knowledgeManager) @@ -292,6 +293,7 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) { vulnerabilityHandler := handler.NewVulnerabilityHandler(db, log.Logger) configHandler := handler.NewConfigHandler(configPath, cfg, mcpServer, executor, agent, attackChainHandler, externalMCPMgr, log.Logger) externalMCPHandler := handler.NewExternalMCPHandler(externalMCPMgr, cfg, configPath, log.Logger) + roleHandler := handler.NewRoleHandler(cfg, configPath, log.Logger) // 创建 App 实例（部分字段稍后填充） app := &App{ @@ -368,6 +370,7 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) { attackChainHandler, app, // 传递 App 实例以便动态获取 knowledgeHandler vulnerabilityHandler, + roleHandler, mcpServer, authManager, ) @@ -428,6 +431,7 @@ func setupRoutes( attackChainHandler *handler.AttackChainHandler, app *App, // 传递 App 实例以便动态获取 knowledgeHandler vulnerabilityHandler *handler.VulnerabilityHandler, + roleHandler *handler.RoleHandler, mcpServer *mcp.Server, authManager *security.AuthManager, ) { @@ -653,6 +657,13 @@ func setupRoutes( protected.PUT("/vulnerabilities/:id", vulnerabilityHandler.UpdateVulnerability) protected.DELETE("/vulnerabilities/:id", vulnerabilityHandler.DeleteVulnerability) + // 角色管理 + protected.GET("/roles", roleHandler.GetRoles) + protected.GET("/roles/:name", roleHandler.GetRole) + protected.POST("/roles", roleHandler.CreateRole) + protected.PUT("/roles/:name", roleHandler.UpdateRole) + protected.DELETE("/roles/:name", roleHandler.DeleteRole) + // MCP端点 protected.POST("/mcp", func(c *gin.Context) { mcpServer.HandleHTTP(c.Writer, c.Request) @@ -672,7 +683,7 @@ func setupRoutes( // registerVulnerabilityTool 注册漏洞记录工具到MCP服务器 func registerVulnerabilityTool(mcpServer *mcp.Server, db *database.DB, logger *zap.Logger) { tool := mcp.Tool{ - Name: "record_vulnerability", + Name: builtin.ToolRecordVulnerability, Description: "记录发现的漏洞详情到漏洞管理系统。当发现有效漏洞时，使用此工具记录漏洞信息，包括标题、描述、严重程度、类型、目标、证明、影响和建议等。", ShortDescription: "记录发现的漏洞详情到漏洞管理系统", InputSchema: map[string]interface{}{ diff --git a/internal/config/config.go b/internal/config/config.go index 85aaee3b..4c22f600 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -6,6 +6,7 @@ import ( "fmt" "os" "path/filepath" + "strconv" "strings" "gopkg.in/yaml.v3" @@ -22,6 +23,8 @@ type Config struct { Auth AuthConfig `yaml:"auth"` ExternalMCP ExternalMCPConfig `yaml:"external_mcp,omitempty"` Knowledge KnowledgeConfig `yaml:"knowledge,omitempty"` + RolesDir string `yaml:"roles_dir,omitempty" json:"roles_dir,omitempty"` // 角色配置文件目录（新方式） + Roles map[string]RoleConfig `yaml:"roles,omitempty" json:"roles,omitempty"` // 向后兼容：支持在主配置文件中定义角色 } type ServerConfig struct { @@ -207,6 +210,29 @@ func Load(path string) (*Config, error) { } } + // 从角色目录加载角色配置 + if cfg.RolesDir != "" { + configDir := filepath.Dir(path) + rolesDir := cfg.RolesDir + + // 如果是相对路径，相对于配置文件所在目录 + if !filepath.IsAbs(rolesDir) { + rolesDir = filepath.Join(configDir, rolesDir) + } + + roles, err := LoadRolesFromDir(rolesDir) + if err != nil { + return nil, fmt.Errorf("从角色目录加载角色配置失败: %w", err) + } + + cfg.Roles = roles + } else { + // 如果未配置 roles_dir，初始化为空 map + if cfg.Roles == nil { + cfg.Roles = make(map[string]RoleConfig) + } + } + return &cfg, nil } @@ -375,6 +401,98 @@ func LoadToolFromFile(path string) (*ToolConfig, error) { return &tool, nil } +// LoadRolesFromDir 从目录加载所有角色配置文件 +func LoadRolesFromDir(dir string) (map[string]RoleConfig, error) { + roles := make(map[string]RoleConfig) + + // 检查目录是否存在 + if _, err := os.Stat(dir); os.IsNotExist(err) { + return roles, nil // 目录不存在时返回空map，不报错 + } + + // 读取目录中的所有 .yaml 和 .yml 文件 + entries, err := os.ReadDir(dir) + if err != nil { + return nil, fmt.Errorf("读取角色目录失败: %w", err) + } + + for _, entry := range entries { + if entry.IsDir() { + continue + } + + name := entry.Name() + if !strings.HasSuffix(name, ".yaml") && !strings.HasSuffix(name, ".yml") { + continue + } + + filePath := filepath.Join(dir, name) + role, err := LoadRoleFromFile(filePath) + if err != nil { + // 记录错误但继续加载其他文件 + fmt.Printf("警告: 加载角色配置文件 %s 失败: %v\n", filePath, err) + continue + } + + // 使用角色名称作为key + roleName := role.Name + if roleName == "" { + // 如果角色名称为空，使用文件名（去掉扩展名）作为名称 + roleName = strings.TrimSuffix(strings.TrimSuffix(name, ".yaml"), ".yml") + role.Name = roleName + } + + roles[roleName] = *role + } + + return roles, nil +} + +// LoadRoleFromFile 从单个文件加载角色配置 +func LoadRoleFromFile(path string) (*RoleConfig, error) { + data, err := os.ReadFile(path) + if err != nil { + return nil, fmt.Errorf("读取文件失败: %w", err) + } + + var role RoleConfig + if err := yaml.Unmarshal(data, &role); err != nil { + return nil, fmt.Errorf("解析角色配置失败: %w", err) + } + + // 处理 icon 字段：如果包含 Unicode 转义格式（\U0001F3C6），转换为实际的 Unicode 字符 + // Go 的 yaml 库可能不会自动解析 \U 转义序列，需要手动转换 + if role.Icon != "" { + icon := role.Icon + // 去除可能的引号 + icon = strings.Trim(icon, `"`) + + // 检查是否是 Unicode 转义格式 \U0001F3C6（8位十六进制）或 \uXXXX（4位十六进制） + if len(icon) >= 3 && icon[0] == '\\' { + if icon[1] == 'U' && len(icon) >= 10 { + // \U0001F3C6 格式（8位十六进制） + if codePoint, err := strconv.ParseInt(icon[2:10], 16, 32); err == nil { + role.Icon = string(rune(codePoint)) + } + } else if icon[1] == 'u' && len(icon) >= 6 { + // \uXXXX 格式（4位十六进制） + if codePoint, err := strconv.ParseInt(icon[2:6], 16, 32); err == nil { + role.Icon = string(rune(codePoint)) + } + } + } + } + + // 验证必需字段 + if role.Name == "" { + // 如果名称为空，尝试从文件名获取 + baseName := filepath.Base(path) + role.Name = strings.TrimSuffix(strings.TrimSuffix(baseName, ".yaml"), ".yml") + } + + return &role, nil +} + func Default() *Config { return &Config{ Server: ServerConfig{ @@ -448,3 +566,20 @@ type RetrievalConfig struct { SimilarityThreshold float64 `yaml:"similarity_threshold" json:"similarity_threshold"` // 相似度阈值 HybridWeight float64 `yaml:"hybrid_weight" json:"hybrid_weight"` // 向量检索权重（0-1） } + +// RolesConfig 角色配置（已废弃，使用 map[string]RoleConfig 替代） +// 保留此类型以兼容旧代码，但建议直接使用 map[string]RoleConfig +type RolesConfig struct { + Roles map[string]RoleConfig `yaml:"roles,omitempty" json:"roles,omitempty"` +} + +// RoleConfig 单个角色配置 +type RoleConfig struct { + Name string `yaml:"name" json:"name"` // 角色名称 + Description string `yaml:"description" json:"description"` // 角色描述 + UserPrompt string `yaml:"user_prompt" json:"user_prompt"` // 用户提示词(追加到用户消息前) + Icon string `yaml:"icon,omitempty" json:"icon,omitempty"` // 角色图标（可选） + Tools []string `yaml:"tools,omitempty" json:"tools,omitempty"` // 关联的工具列表（toolKey格式，如 "toolName" 或 "mcpName::toolName"） + MCPs []string `yaml:"mcps,omitempty" json:"mcps,omitempty"` // 向后兼容：关联的MCP服务器列表（已废弃，使用tools替代） + Enabled bool `yaml:"enabled" json:"enabled"` // 是否启用 +} diff --git a/internal/handler/agent.go b/internal/handler/agent.go index 015cc5ab..8d80dbfe 100644 --- a/internal/handler/agent.go +++ b/internal/handler/agent.go @@ -12,7 +12,9 @@ import ( "unicode/utf8" "cyberstrike-ai/internal/agent" + "cyberstrike-ai/internal/config" "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/mcp/builtin" "github.com/gin-gonic/gin" "go.uber.org/zap" @@ -66,13 +68,14 @@ type AgentHandler struct { logger *zap.Logger tasks *AgentTaskManager batchTaskManager *BatchTaskManager - knowledgeManager interface { // 知识库管理器接口 + config *config.Config // 配置引用，用于获取角色信息 + knowledgeManager interface { // 知识库管理器接口 LogRetrieval(conversationID, messageID, query, riskType string, retrievedItems []string) error } } // NewAgentHandler 创建新的Agent处理器 -func NewAgentHandler(agent *agent.Agent, db *database.DB, logger *zap.Logger) *AgentHandler { +func NewAgentHandler(agent *agent.Agent, db *database.DB, cfg *config.Config, logger *zap.Logger) *AgentHandler { batchTaskManager := NewBatchTaskManager() batchTaskManager.SetDB(db) @@ -87,6 +90,7 @@ func NewAgentHandler(agent *agent.Agent, db *database.DB, logger *zap.Logger) *A logger: logger, tasks: NewAgentTaskManager(), batchTaskManager: batchTaskManager, + config: cfg, } } @@ -101,6 +105,7 @@ func (h *AgentHandler) SetKnowledgeManager(manager interface { type ChatRequest struct { Message string `json:"message" binding:"required"` ConversationID string `json:"conversationId,omitempty"` + Role string `json:"role,omitempty"` // 角色名称 } // ChatResponse 聊天响应 @@ -161,14 +166,34 @@ func (h *AgentHandler) AgentLoop(c *gin.Context) { h.logger.Info("从ReAct数据恢复历史上下文", zap.Int("count", len(agentHistoryMessages))) } - // 保存用户消息 + // 应用角色用户提示词和工具配置 + finalMessage := req.Message + var roleTools []string // 角色配置的工具列表 + if req.Role != "" && req.Role != "默认" { + if h.config.Roles != nil { + if role, exists := h.config.Roles[req.Role]; exists && role.Enabled { + // 应用用户提示词 + if role.UserPrompt != "" { + finalMessage = role.UserPrompt + "\n\n" + req.Message + h.logger.Info("应用角色用户提示词", zap.String("role", req.Role)) + } + // 获取角色配置的工具列表（优先使用tools字段，向后兼容mcps字段） + if len(role.Tools) > 0 { + roleTools = role.Tools + h.logger.Info("使用角色配置的工具列表", zap.String("role", req.Role), zap.Int("toolCount", len(roleTools))) + } + } + } + } + + // 保存用户消息（保存原始消息，不包含角色提示词） _, err = h.db.AddMessage(conversationID, "user", req.Message, nil) if err != nil { h.logger.Error("保存用户消息失败", zap.Error(err)) } - // 执行Agent Loop，传入历史消息和对话ID - result, err := h.agent.AgentLoopWithConversationID(c.Request.Context(), req.Message, agentHistoryMessages, conversationID) + // 执行Agent Loop，传入历史消息和对话ID（使用包含角色提示词的finalMessage和角色工具列表） + result, err := h.agent.AgentLoopWithProgress(c.Request.Context(), finalMessage, agentHistoryMessages, conversationID, nil, roleTools) if err != nil { h.logger.Error("Agent Loop执行失败", zap.Error(err)) @@ -231,7 +256,7 @@ func (h *AgentHandler) createProgressCallback(conversationID, assistantMessageID if eventType == "tool_call" { if dataMap, ok := data.(map[string]interface{}); ok { toolName, _ := dataMap["toolName"].(string) - if toolName == "search_knowledge_base" { + if toolName == builtin.ToolSearchKnowledgeBase { if toolCallId, ok := dataMap["toolCallId"].(string); ok && toolCallId != "" { if argumentsObj, ok := dataMap["argumentsObj"].(map[string]interface{}); ok { toolCallCache[toolCallId] = argumentsObj @@ -245,7 +270,7 @@ func (h *AgentHandler) createProgressCallback(conversationID, assistantMessageID if eventType == "tool_result" && h.knowledgeManager != nil { if dataMap, ok := data.(map[string]interface{}); ok { toolName, _ := dataMap["toolName"].(string) - if toolName == "search_knowledge_base" { + if toolName == builtin.ToolSearchKnowledgeBase { // 提取检索信息 query := "" riskType := "" @@ -470,7 +495,32 @@ func (h *AgentHandler) AgentLoopStream(c *gin.Context) { h.logger.Info("从ReAct数据恢复历史上下文", zap.Int("count", len(agentHistoryMessages))) } - // 保存用户消息 + // 应用角色用户提示词和工具配置 + finalMessage := req.Message + var roleTools []string // 角色配置的工具列表 + if req.Role != "" && req.Role != "默认" { + if h.config.Roles != nil { + if role, exists := h.config.Roles[req.Role]; exists && role.Enabled { + // 应用用户提示词 + if role.UserPrompt != "" { + finalMessage = role.UserPrompt + "\n\n" + req.Message + h.logger.Info("应用角色用户提示词", zap.String("role", req.Role)) + } + // 获取角色配置的工具列表（优先使用tools字段，向后兼容mcps字段） + if len(role.Tools) > 0 { + roleTools = role.Tools + h.logger.Info("使用角色配置的工具列表", zap.String("role", req.Role), zap.Int("toolCount", len(roleTools))) + } else if len(role.MCPs) > 0 { + // 向后兼容：如果只有mcps字段，暂时使用空列表（表示使用所有工具） + // 因为mcps是MCP服务器名称，不是工具列表 + h.logger.Info("角色配置使用旧的mcps字段，将使用所有工具", zap.String("role", req.Role)) + } + } + } + } + // 如果roleTools为空，表示使用所有工具（默认角色或未配置工具的角色） + + // 保存用户消息（保存原始消息，不包含角色提示词） _, err = h.db.AddMessage(conversationID, "user", req.Message, nil) if err != nil { h.logger.Error("保存用户消息失败", zap.Error(err)) @@ -547,9 +597,9 @@ func (h *AgentHandler) AgentLoopStream(c *gin.Context) { taskStatus := "completed" defer h.tasks.FinishTask(conversationID, taskStatus) - // 执行Agent Loop，传入独立的上下文，确保任务不会因客户端断开而中断 + // 执行Agent Loop，传入独立的上下文，确保任务不会因客户端断开而中断（使用包含角色提示词的finalMessage和角色工具列表） sendEvent("progress", "正在分析您的请求...", nil) - result, err := h.agent.AgentLoopWithProgress(taskCtx, req.Message, agentHistoryMessages, conversationID, progressCallback) + result, err := h.agent.AgentLoopWithProgress(taskCtx, finalMessage, agentHistoryMessages, conversationID, progressCallback, roleTools) if err != nil { h.logger.Error("Agent Loop执行失败", zap.Error(err)) cause := context.Cause(baseCtx) @@ -759,7 +809,7 @@ func (h *AgentHandler) ListCompletedTasks(c *gin.Context) { // BatchTaskRequest 批量任务请求 type BatchTaskRequest struct { - Title string `json:"title"` // 任务标题（可选） + Title string `json:"title"` // 任务标题（可选） Tasks []string `json:"tasks" binding:"required"` // 任务列表，每行一个任务 } @@ -1072,7 +1122,8 @@ func (h *AgentHandler) executeBatchQueue(queueID string) { ctx, cancel := context.WithTimeout(context.Background(), 30*time.Minute) // 存储取消函数，以便在取消队列时能够取消当前任务 h.batchTaskManager.SetTaskCancel(queueID, cancel) - result, err := h.agent.AgentLoopWithProgress(ctx, task.Message, []agent.ChatMessage{}, conversationID, progressCallback) + // 批量任务暂时不支持角色工具过滤，使用所有工具（传入nil） + result, err := h.agent.AgentLoopWithProgress(ctx, task.Message, []agent.ChatMessage{}, conversationID, progressCallback, nil) // 任务执行完成，清理取消函数 h.batchTaskManager.SetTaskCancel(queueID, nil) cancel() diff --git a/internal/handler/config.go b/internal/handler/config.go index a0965721..5dad5056 100644 --- a/internal/handler/config.go +++ b/internal/handler/config.go @@ -147,6 +147,7 @@ type ToolConfigInfo struct { Enabled bool `json:"enabled"` IsExternal bool `json:"is_external,omitempty"` // 是否为外部MCP工具 ExternalMCP string `json:"external_mcp,omitempty"` // 外部MCP名称（如果是外部工具） + RoleEnabled *bool `json:"role_enabled,omitempty"` // 该工具在当前角色中是否启用（nil表示未指定角色或使用所有工具） } // GetConfig 获取当前配置 @@ -272,11 +273,12 @@ func (h *ConfigHandler) GetConfig(c *gin.Context) { // GetToolsResponse 获取工具列表响应（分页） type GetToolsResponse struct { - Tools []ToolConfigInfo `json:"tools"` - Total int `json:"total"` - Page int `json:"page"` - PageSize int `json:"page_size"` - TotalPages int `json:"total_pages"` + Tools []ToolConfigInfo `json:"tools"` + Total int `json:"total"` + TotalEnabled int `json:"total_enabled"` // 已启用的工具总数 + Page int `json:"page"` + PageSize int `json:"page_size"` + TotalPages int `json:"total_pages"` } // GetTools 获取工具列表（支持分页和搜索） @@ -305,6 +307,23 @@ func (h *ConfigHandler) GetTools(c *gin.Context) { searchTermLower = strings.ToLower(searchTerm) } + // 解析角色参数，用于过滤工具并标注启用状态 + roleName := c.Query("role") + var roleToolsSet map[string]bool // 角色配置的工具集合 + var roleUsesAllTools bool = true // 角色是否使用所有工具（默认角色） + if roleName != "" && roleName != "默认" && h.config.Roles != nil { + if role, exists := h.config.Roles[roleName]; exists && role.Enabled { + if len(role.Tools) > 0 { + // 角色配置了工具列表，只使用这些工具 + roleToolsSet = make(map[string]bool) + for _, toolKey := range role.Tools { + roleToolsSet[toolKey] = true + } + roleUsesAllTools = false + } + } + } + // 获取所有内部工具并应用搜索过滤 configToolMap := make(map[string]bool) allTools := make([]ToolConfigInfo, 0, len(h.config.Security.Tools)) @@ -325,6 +344,31 @@ func (h *ConfigHandler) GetTools(c *gin.Context) { toolInfo.Description = desc } + // 根据角色配置标注工具状态 + if roleName != "" { + if roleUsesAllTools { + // 角色使用所有工具，标注启用的工具为role_enabled=true + if tool.Enabled { + roleEnabled := true + toolInfo.RoleEnabled = &roleEnabled + } else { + roleEnabled := false + toolInfo.RoleEnabled = &roleEnabled + } + } else { + // 角色配置了工具列表，检查工具是否在列表中 + // 内部工具使用工具名称作为key + if roleToolsSet[tool.Name] { + roleEnabled := tool.Enabled // 工具必须在角色列表中且本身启用 + toolInfo.RoleEnabled = &roleEnabled + } else { + // 不在角色列表中，标记为false + roleEnabled := false + toolInfo.RoleEnabled = &roleEnabled + } + } + } + // 如果有关键词，进行搜索过滤 if searchTermLower != "" { nameLower := strings.ToLower(toolInfo.Name) @@ -361,6 +405,26 @@ func (h *ConfigHandler) GetTools(c *gin.Context) { IsExternal: false, } + // 根据角色配置标注工具状态 + if roleName != "" { + if roleUsesAllTools { + // 角色使用所有工具，直接注册的工具默认启用 + roleEnabled := true + toolInfo.RoleEnabled = &roleEnabled + } else { + // 角色配置了工具列表，检查工具是否在列表中 + // 内部工具使用工具名称作为key + if roleToolsSet[mcpTool.Name] { + roleEnabled := true // 在角色列表中且工具本身启用 + toolInfo.RoleEnabled = &roleEnabled + } else { + // 不在角色列表中，标记为false + roleEnabled := false + toolInfo.RoleEnabled = &roleEnabled + } + } + } + // 如果有关键词，进行搜索过滤 if searchTermLower != "" { nameLower := strings.ToLower(toolInfo.Name) @@ -439,18 +503,55 @@ func (h *ConfigHandler) GetTools(c *gin.Context) { } } - allTools = append(allTools, ToolConfigInfo{ + toolInfo := ToolConfigInfo{ Name: actualToolName, // 显示实际工具名称，不带前缀 Description: description, Enabled: enabled, IsExternal: true, ExternalMCP: mcpName, - }) + } + + // 根据角色配置标注工具状态 + if roleName != "" { + if roleUsesAllTools { + // 角色使用所有工具，标注启用的工具为role_enabled=true + toolInfo.RoleEnabled = &enabled + } else { + // 角色配置了工具列表，检查工具是否在列表中 + // 外部工具使用 "mcpName::toolName" 格式作为key + externalToolKey := externalTool.Name // 这是 "mcpName::toolName" 格式 + if roleToolsSet[externalToolKey] { + roleEnabled := enabled // 工具必须在角色列表中且本身启用 + toolInfo.RoleEnabled = &roleEnabled + } else { + // 不在角色列表中，标记为false + roleEnabled := false + toolInfo.RoleEnabled = &roleEnabled + } + } + } + + allTools = append(allTools, toolInfo) } } } + // 如果角色配置了工具列表，过滤工具（只保留列表中的工具，但保留其他工具并标记为禁用） + // 注意：这里我们不直接过滤掉工具，而是保留所有工具，但通过 role_enabled 字段标注状态 + // 这样前端可以显示所有工具，并标注哪些工具在当前角色中可用 + total := len(allTools) + // 统计已启用的工具数（在角色中的启用工具数） + totalEnabled := 0 + for _, tool := range allTools { + if tool.RoleEnabled != nil && *tool.RoleEnabled { + totalEnabled++ + } else if tool.RoleEnabled == nil && tool.Enabled { + // 如果未指定角色，统计所有启用的工具 + totalEnabled++ + } + } + totalPages := (total + pageSize - 1) / pageSize if totalPages == 0 { totalPages = 1 @@ -471,11 +572,12 @@ func (h *ConfigHandler) GetTools(c *gin.Context) { } c.JSON(http.StatusOK, GetToolsResponse{ - Tools: tools, - Total: total, - Page: page, - PageSize: pageSize, - TotalPages: totalPages, + Tools: tools, + Total: total, + TotalEnabled: totalEnabled, + Page: page, + PageSize: pageSize, + TotalPages: totalPages, }) } diff --git a/internal/handler/role.go b/internal/handler/role.go new file mode 100644 index 00000000..85411b19 --- /dev/null +++ b/internal/handler/role.go @@ -0,0 +1,453 @@ +package handler + +import ( + "fmt" + "net/http" + "os" + "path/filepath" + "regexp" + "strings" + + "cyberstrike-ai/internal/config" + + "gopkg.in/yaml.v3" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +// RoleHandler 角色处理器 +type RoleHandler struct { + config *config.Config + configPath string + logger *zap.Logger +} + +// NewRoleHandler 创建新的角色处理器 +func NewRoleHandler(cfg *config.Config, configPath string, logger *zap.Logger) *RoleHandler { + return &RoleHandler{ + config: cfg, + configPath: configPath, + logger: logger, + } +} + +// GetRoles 获取所有角色 +func (h *RoleHandler) GetRoles(c *gin.Context) { + if h.config.Roles == nil { + h.config.Roles = make(map[string]config.RoleConfig) + } + + roles := make([]config.RoleConfig, 0, len(h.config.Roles)) + for key, role := range h.config.Roles { + // 确保角色的key与name一致 + if role.Name == "" { + role.Name = key + } + roles = append(roles, role) + } + + c.JSON(http.StatusOK, gin.H{ + "roles": roles, + }) +} + +// GetRole 获取单个角色 +func (h *RoleHandler) GetRole(c *gin.Context) { + roleName := c.Param("name") + if roleName == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "角色名称不能为空"}) + return + } + + if h.config.Roles == nil { + c.JSON(http.StatusNotFound, gin.H{"error": "角色不存在"}) + return + } + + role, exists := h.config.Roles[roleName] + if !exists { + c.JSON(http.StatusNotFound, gin.H{"error": "角色不存在"}) + return + } + + // 确保角色的name与key一致 + if role.Name == "" { + role.Name = roleName + } + + c.JSON(http.StatusOK, gin.H{ + "role": role, + }) +} + +// UpdateRole 更新角色 +func (h *RoleHandler) UpdateRole(c *gin.Context) { + roleName := c.Param("name") + if roleName == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "角色名称不能为空"}) + return + } + + var req config.RoleConfig + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()}) + return + } + + // 确保角色名称与请求中的name一致 + if req.Name == "" { + req.Name = roleName + } + + // 初始化Roles map + if h.config.Roles == nil { + h.config.Roles = make(map[string]config.RoleConfig) + } + + // 删除所有与角色name相同但key不同的旧角色（避免重复） + // 使用角色name作为key，确保唯一性 + finalKey := req.Name + keysToDelete := make([]string, 0) + for key := range h.config.Roles { + // 如果key与最终的key不同，但name相同，则标记为删除 + if key != finalKey { + role := h.config.Roles[key] + // 确保角色的name字段正确设置 + if role.Name == "" { + role.Name = key + } + if role.Name == req.Name { + keysToDelete = append(keysToDelete, key) + } + } + } + // 删除旧的角色 + for _, key := range keysToDelete { + delete(h.config.Roles, key) + h.logger.Info("删除重复的角色", zap.String("oldKey", key), zap.String("name", req.Name)) + } + + // 如果当前更新的key与最终key不同，也需要删除旧的 + if roleName != finalKey { + delete(h.config.Roles, roleName) + } + + // 如果角色名称改变，需要删除旧文件 + if roleName != finalKey { + configDir := filepath.Dir(h.configPath) + rolesDir := h.config.RolesDir + if rolesDir == "" { + rolesDir = "roles" // 默认目录 + } + + // 如果是相对路径，相对于配置文件所在目录 + if !filepath.IsAbs(rolesDir) { + rolesDir = filepath.Join(configDir, rolesDir) + } + + // 删除旧的角色文件 + oldSafeFileName := sanitizeFileName(roleName) + oldRoleFileYaml := filepath.Join(rolesDir, oldSafeFileName+".yaml") + oldRoleFileYml := filepath.Join(rolesDir, oldSafeFileName+".yml") + + if _, err := os.Stat(oldRoleFileYaml); err == nil { + if err := os.Remove(oldRoleFileYaml); err != nil { + h.logger.Warn("删除旧角色配置文件失败", zap.String("file", oldRoleFileYaml), zap.Error(err)) + } + } + if _, err := os.Stat(oldRoleFileYml); err == nil { + if err := os.Remove(oldRoleFileYml); err != nil { + h.logger.Warn("删除旧角色配置文件失败", zap.String("file", oldRoleFileYml), zap.Error(err)) + } + } + } + + // 使用角色name作为key来保存（确保唯一性） + h.config.Roles[finalKey] = req + + // 保存配置到文件 + if err := h.saveConfig(); err != nil { + h.logger.Error("保存配置失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": "保存配置失败: " + err.Error()}) + return + } + + h.logger.Info("更新角色", zap.String("oldKey", roleName), zap.String("newKey", finalKey), zap.String("name", req.Name)) + c.JSON(http.StatusOK, gin.H{ + "message": "角色已更新", + "role": req, + }) +} + +// CreateRole 创建新角色 +func (h *RoleHandler) CreateRole(c *gin.Context) { + var req config.RoleConfig + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()}) + return + } + + if req.Name == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "角色名称不能为空"}) + return + } + + // 初始化Roles map + if h.config.Roles == nil { + h.config.Roles = make(map[string]config.RoleConfig) + } + + // 检查角色是否已存在 + if _, exists := h.config.Roles[req.Name]; exists { + c.JSON(http.StatusBadRequest, gin.H{"error": "角色已存在"}) + return + } + + // 创建角色（默认启用） + if !req.Enabled { + req.Enabled = true + } + + h.config.Roles[req.Name] = req + + // 保存配置到文件 + if err := h.saveConfig(); err != nil { + h.logger.Error("保存配置失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": "保存配置失败: " + err.Error()}) + return + } + + h.logger.Info("创建角色", zap.String("roleName", req.Name)) + c.JSON(http.StatusOK, gin.H{ + "message": "角色已创建", + "role": req, + }) +} + +// DeleteRole 删除角色 +func (h *RoleHandler) DeleteRole(c *gin.Context) { + roleName := c.Param("name") + if roleName == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "角色名称不能为空"}) + return + } + + if h.config.Roles == nil { + c.JSON(http.StatusNotFound, gin.H{"error": "角色不存在"}) + return + } + + if _, exists := h.config.Roles[roleName]; !exists { + c.JSON(http.StatusNotFound, gin.H{"error": "角色不存在"}) + return + } + + // 不允许删除"默认"角色 + if roleName == "默认" { + c.JSON(http.StatusBadRequest, gin.H{"error": "不能删除默认角色"}) + return + } + + delete(h.config.Roles, roleName) + + // 删除对应的角色文件 + configDir := filepath.Dir(h.configPath) + rolesDir := h.config.RolesDir + if rolesDir == "" { + rolesDir = "roles" // 默认目录 + } + + // 如果是相对路径，相对于配置文件所在目录 + if !filepath.IsAbs(rolesDir) { + rolesDir = filepath.Join(configDir, rolesDir) + } + + // 尝试删除角色文件（.yaml 和 .yml） + safeFileName := sanitizeFileName(roleName) + roleFileYaml := filepath.Join(rolesDir, safeFileName+".yaml") + roleFileYml := filepath.Join(rolesDir, safeFileName+".yml") + + // 删除 .yaml 文件（如果存在） + if _, err := os.Stat(roleFileYaml); err == nil { + if err := os.Remove(roleFileYaml); err != nil { + h.logger.Warn("删除角色配置文件失败", zap.String("file", roleFileYaml), zap.Error(err)) + } else { + h.logger.Info("已删除角色配置文件", zap.String("file", roleFileYaml)) + } + } + + // 删除 .yml 文件（如果存在） + if _, err := os.Stat(roleFileYml); err == nil { + if err := os.Remove(roleFileYml); err != nil { + h.logger.Warn("删除角色配置文件失败", zap.String("file", roleFileYml), zap.Error(err)) + } else { + h.logger.Info("已删除角色配置文件", zap.String("file", roleFileYml)) + } + } + + h.logger.Info("删除角色", zap.String("roleName", roleName)) + c.JSON(http.StatusOK, gin.H{ + "message": "角色已删除", + }) +} + +// saveConfig 保存配置到目录中的文件 +func (h *RoleHandler) saveConfig() error { + configDir := filepath.Dir(h.configPath) + rolesDir := h.config.RolesDir + if rolesDir == "" { + rolesDir = "roles" // 默认目录 + } + + // 如果是相对路径，相对于配置文件所在目录 + if !filepath.IsAbs(rolesDir) { + rolesDir = filepath.Join(configDir, rolesDir) + } + + // 确保目录存在 + if err := os.MkdirAll(rolesDir, 0755); err != nil { + return fmt.Errorf("创建角色目录失败: %w", err) + } + + // 保存每个角色到独立的文件 + if h.config.Roles != nil { + for roleName, role := range h.config.Roles { + // 确保角色名称正确设置 + if role.Name == "" { + role.Name = roleName + } + + // 使用角色名称作为文件名（安全化文件名，避免特殊字符） + safeFileName := sanitizeFileName(role.Name) + roleFile := filepath.Join(rolesDir, safeFileName+".yaml") + + // 将角色配置序列化为YAML + roleData, err := yaml.Marshal(&role) + if err != nil { + h.logger.Error("序列化角色配置失败", zap.String("role", roleName), zap.Error(err)) + continue + } + + // 处理icon字段：确保包含\U的icon值被引号包围（YAML需要引号才能正确解析Unicode转义） + roleDataStr := string(roleData) + if role.Icon != "" && strings.HasPrefix(role.Icon, "\\U") { + // 匹配 icon: \UXXXXXXXX 格式（没有引号），排除已经有引号的情况 + // 使用负向前瞻确保后面没有引号，或者直接匹配没有引号的情况 + re := regexp.MustCompile(`(?m)^(icon:\s+)(\\U[0-9A-F]{8})(\s*)$`) + roleDataStr = re.ReplaceAllString(roleDataStr, `${1}"${2}"${3}`) + roleData = []byte(roleDataStr) + } + + // 写入文件 + if err := os.WriteFile(roleFile, roleData, 0644); err != nil { + h.logger.Error("保存角色配置文件失败", zap.String("role", roleName), zap.String("file", roleFile), zap.Error(err)) + continue + } + + h.logger.Info("角色配置已保存到文件", zap.String("role", roleName), zap.String("file", roleFile)) + } + } + + return nil +} + +// sanitizeFileName 将角色名称转换为安全的文件名 +func sanitizeFileName(name string) string { + // 替换可能不安全的字符 + replacer := map[rune]string{ + '/': "_", + '\\': "_", + ':': "_", + '*': "_", + '?': "_", + '"': "_", + '<': "_", + '>': "_", + '|': "_", + ' ': "_", + } + + var result []rune + for _, r := range name { + if replacement, ok := replacer[r]; ok { + result = append(result, []rune(replacement)...) + } else { + result = append(result, r) + } + } + + fileName := string(result) + // 如果文件名为空，使用默认名称 + if fileName == "" { + fileName = "role" + } + + return fileName +} + +// updateRolesConfig 更新角色配置 +func updateRolesConfig(doc *yaml.Node, cfg config.RolesConfig) { + root := doc.Content[0] + rolesNode := ensureMap(root, "roles") + + // 清空现有角色 + if rolesNode.Kind == yaml.MappingNode { + rolesNode.Content = nil + } + + // 添加新角色（使用name作为key，确保唯一性） + if cfg.Roles != nil { + // 先建立一个以name为key的map，去重（保留最后一个） + rolesByName := make(map[string]config.RoleConfig) + for roleKey, role := range cfg.Roles { + // 确保角色的name字段正确设置 + if role.Name == "" { + role.Name = roleKey + } + // 使用name作为最终key，如果有多个key对应相同的name，只保留最后一个 + rolesByName[role.Name] = role + } + + // 将去重后的角色写入YAML + for roleName, role := range rolesByName { + roleNode := ensureMap(rolesNode, roleName) + setStringInMap(roleNode, "name", role.Name) + setStringInMap(roleNode, "description", role.Description) + setStringInMap(roleNode, "user_prompt", role.UserPrompt) + if role.Icon != "" { + setStringInMap(roleNode, "icon", role.Icon) + } + setBoolInMap(roleNode, "enabled", role.Enabled) + + // 添加工具列表（优先使用tools字段） + if len(role.Tools) > 0 { + toolsNode := ensureArray(roleNode, "tools") + toolsNode.Content = nil + for _, toolKey := range role.Tools { + toolNode := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: toolKey} + toolsNode.Content = append(toolsNode.Content, toolNode) + } + } else if len(role.MCPs) > 0 { + // 向后兼容：如果没有tools但有mcps，保存mcps + mcpsNode := ensureArray(roleNode, "mcps") + mcpsNode.Content = nil + for _, mcpName := range role.MCPs { + mcpNode := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: mcpName} + mcpsNode.Content = append(mcpsNode.Content, mcpNode) + } + } + } + } +} + +// ensureArray 确保数组中存在指定key的数组节点 +func ensureArray(parent *yaml.Node, key string) *yaml.Node { + _, valueNode := ensureKeyValue(parent, key) + if valueNode.Kind != yaml.SequenceNode { + valueNode.Kind = yaml.SequenceNode + valueNode.Tag = "!!seq" + valueNode.Content = nil + } + return valueNode +} diff --git a/internal/knowledge/retriever.go b/internal/knowledge/retriever.go index 8ea77f2e..15e46a1f 100644 --- a/internal/knowledge/retriever.go +++ b/internal/knowledge/retriever.go @@ -161,14 +161,14 @@ func (r *Retriever) Search(ctx context.Context, req *SearchRequest) ([]*Retrieva // 查询所有向量（或按风险类型过滤） // 使用精确匹配（=）以提高性能和准确性 - // 由于系统提供了 list_knowledge_risk_types 工具，用户应该使用准确的category名称 - // 同时，向量嵌入中已包含category信息，即使SQL过滤不完全匹配，向量相似度也能帮助匹配 - var rows *sql.Rows - if req.RiskType != "" { - // 使用精确匹配（=），性能更好且更准确 - // 使用 COLLATE NOCASE 实现大小写不敏感匹配，提高容错性 - // 注意：如果用户输入的risk_type与category不完全一致，可能匹配不到 - // 建议用户先调用 list_knowledge_risk_types 获取准确的category名称 + // 由于系统提供了内置工具来获取风险类型列表，用户应该使用准确的category名称 + // 同时，向量嵌入中已包含category信息，即使SQL过滤不完全匹配，向量相似度也能帮助匹配 + var rows *sql.Rows + if req.RiskType != "" { + // 使用精确匹配（=），性能更好且更准确 + // 使用 COLLATE NOCASE 实现大小写不敏感匹配，提高容错性 + // 注意：如果用户输入的risk_type与category不完全一致，可能匹配不到 + // 建议用户先调用相应的内置工具获取准确的category名称 rows, err = r.db.Query(` SELECT e.id, e.item_id, e.chunk_index, e.chunk_text, e.embedding, i.category, i.title FROM knowledge_embeddings e diff --git a/internal/knowledge/tool.go b/internal/knowledge/tool.go index a1001bf6..31e52554 100644 --- a/internal/knowledge/tool.go +++ b/internal/knowledge/tool.go @@ -8,6 +8,7 @@ import ( "strings" "cyberstrike-ai/internal/mcp" + "cyberstrike-ai/internal/mcp/builtin" "go.uber.org/zap" ) @@ -21,7 +22,7 @@ func RegisterKnowledgeTool( ) { // 注册第一个工具：获取所有可用的风险类型列表 listRiskTypesTool := mcp.Tool{ - Name: "list_knowledge_risk_types", + Name: builtin.ToolListKnowledgeRiskTypes, Description: "获取知识库中所有可用的风险类型（risk_type）列表。在搜索知识库之前，可以先调用此工具获取可用的风险类型，然后使用正确的风险类型进行精确搜索，这样可以大幅减少检索时间并提高检索准确性。", ShortDescription: "获取知识库中所有可用的风险类型列表", InputSchema: map[string]interface{}{ @@ -62,7 +63,7 @@ func RegisterKnowledgeTool( for i, category := range categories { resultText.WriteString(fmt.Sprintf("%d. %s\n", i+1, category)) } - resultText.WriteString("\n提示：在调用 search_knowledge_base 工具时，可以使用上述风险类型之一作为 risk_type 参数，以缩小搜索范围并提高检索效率。") + resultText.WriteString("\n提示：在调用 " + builtin.ToolSearchKnowledgeBase + " 工具时，可以使用上述风险类型之一作为 risk_type 参数，以缩小搜索范围并提高检索效率。") return &mcp.ToolResult{ Content: []mcp.Content{ @@ -79,8 +80,8 @@ func RegisterKnowledgeTool( // 注册第二个工具：搜索知识库（保持原有功能） searchTool := mcp.Tool{ - Name: "search_knowledge_base", - Description: "在知识库中搜索相关的安全知识。当你需要了解特定漏洞类型、攻击技术、检测方法等安全知识时，可以使用此工具进行检索。工具使用向量检索和混合搜索技术，能够根据查询内容的语义相似度和关键词匹配，自动找到最相关的知识片段。建议：在搜索前可以先调用 list_knowledge_risk_types 工具获取可用的风险类型，然后使用正确的 risk_type 参数进行精确搜索，这样可以大幅减少检索时间。", + Name: builtin.ToolSearchKnowledgeBase, + Description: "在知识库中搜索相关的安全知识。当你需要了解特定漏洞类型、攻击技术、检测方法等安全知识时，可以使用此工具进行检索。工具使用向量检索和混合搜索技术，能够根据查询内容的语义相似度和关键词匹配，自动找到最相关的知识片段。建议：在搜索前可以先调用 " + builtin.ToolListKnowledgeRiskTypes + " 工具获取可用的风险类型，然后使用正确的 risk_type 参数进行精确搜索，这样可以大幅减少检索时间。", ShortDescription: "搜索知识库中的安全知识（支持向量检索和混合搜索）", InputSchema: map[string]interface{}{ "type": "object", @@ -91,7 +92,7 @@ func RegisterKnowledgeTool( }, "risk_type": map[string]interface{}{ "type": "string", - "description": "可选：指定风险类型（如：SQL注入、XSS、文件上传等）。建议先调用 list_knowledge_risk_types 工具获取可用的风险类型列表，然后使用正确的风险类型进行精确搜索，这样可以大幅减少检索时间。如果不指定则搜索所有类型。", + "description": "可选：指定风险类型（如：SQL注入、XSS、文件上传等）。建议先调用 " + builtin.ToolListKnowledgeRiskTypes + " 工具获取可用的风险类型列表，然后使用正确的风险类型进行精确搜索，这样可以大幅减少检索时间。如果不指定则搜索所有类型。", }, }, "required": []string{"query"}, @@ -165,9 +166,9 @@ func RegisterKnowledgeTool( // 按文档分组结果，以便更好地展示上下文 // 使用有序的slice来保持文档顺序（按最高混合分数） type itemGroup struct { - itemID string - results []*RetrievalResult - maxScore float64 // 该文档的最高混合分数 + itemID string + results []*RetrievalResult + maxScore float64 // 该文档的最高混合分数 } itemGroups := make([]*itemGroup, 0) itemMap := make(map[string]*itemGroup) @@ -177,8 +178,8 @@ func RegisterKnowledgeTool( group, exists := itemMap[itemID] if !exists { group = &itemGroup{ - itemID: itemID, - results: make([]*RetrievalResult, 0), + itemID: itemID, + results: make([]*RetrievalResult, 0), maxScore: result.Score, } itemMap[itemID] = group @@ -219,7 +220,7 @@ func RegisterKnowledgeTool( }) // 显示主结果（混合分数最高的，同时显示相似度和混合分数） - resultText.WriteString(fmt.Sprintf("--- 结果 %d (相似度: %.2f%%, 混合分数: %.2f%%) ---\n", + resultText.WriteString(fmt.Sprintf("--- 结果 %d (相似度: %.2f%%, 混合分数: %.2f%%) ---\n", resultIndex, mainResult.Similarity*100, mainResult.Score*100)) resultText.WriteString(fmt.Sprintf("来源: [%s] %s (ID: %s)\n", mainResult.Item.Category, mainResult.Item.Title, mainResult.Item.ID)) diff --git a/internal/mcp/builtin/constants.go b/internal/mcp/builtin/constants.go new file mode 100644 index 00000000..3ae41e00 --- /dev/null +++ b/internal/mcp/builtin/constants.go @@ -0,0 +1,33 @@ +package builtin + +// 内置工具名称常量 +// 所有代码中使用内置工具名称的地方都应该使用这些常量，而不是硬编码字符串 +const ( + // 漏洞管理工具 + ToolRecordVulnerability = "record_vulnerability" + + // 知识库工具 + ToolListKnowledgeRiskTypes = "list_knowledge_risk_types" + ToolSearchKnowledgeBase = "search_knowledge_base" +) + +// IsBuiltinTool 检查工具名称是否是内置工具 +func IsBuiltinTool(toolName string) bool { + switch toolName { + case ToolRecordVulnerability, + ToolListKnowledgeRiskTypes, + ToolSearchKnowledgeBase: + return true + default: + return false + } +} + +// GetAllBuiltinTools 返回所有内置工具名称列表 +func GetAllBuiltinTools() []string { + return []string{ + ToolRecordVulnerability, + ToolListKnowledgeRiskTypes, + ToolSearchKnowledgeBase, + } +} diff --git a/roles/API安全测试.yaml b/roles/API安全测试.yaml new file mode 100644 index 00000000..0b05d8bc --- /dev/null +++ b/roles/API安全测试.yaml @@ -0,0 +1,20 @@ +name: API安全测试 +description: API安全测试专家，专注于API接口安全检测 +user_prompt: 你是一个专业的API安全测试专家。请使用专业的API测试工具对目标API接口进行全面的安全检测，包括GraphQL安全、API参数fuzzing、JWT分析、API架构分析等工作。 +icon: "\U0001F4E1" +tools: + - api-fuzzer + - api-schema-analyzer + - graphql-scanner + - arjun + - jwt-analyzer + - http-intruder + - http-framework-test + - burpsuite + - httpx + - execute-python-script + - install-python-package + - record_vulnerability + - list_knowledge_risk_types + - search_knowledge_base +enabled: true diff --git a/roles/CTF.yaml b/roles/CTF.yaml new file mode 100644 index 00000000..1148351c --- /dev/null +++ b/roles/CTF.yaml @@ -0,0 +1,33 @@ +name: CTF +description: CTF竞赛专家，擅长解题和漏洞利用 +user_prompt: 你是一个CTF竞赛专家。请使用CTF解题思维和方法，快速定位和利用漏洞，解决各类CTF题目。 +icon: "\U0001F3C6" +tools: + - amass + - anew + - angr + - api-fuzzer + - api-schema-analyzer + - arjun + - arp-scan + - autorecon + - binwalk + - bloodhound + - burpsuite + - cat + - checkov + - checksec + - cloudmapper + - create-file + - cyberchef + - dalfox + - delete-file + - httpx + - http-framework-test + - exec + - execute-python-script + - install-python-package + - record_vulnerability + - list_knowledge_risk_types + - search_knowledge_base +enabled: true diff --git a/roles/Web应用扫描.yaml b/roles/Web应用扫描.yaml new file mode 100644 index 00000000..7462c71e --- /dev/null +++ b/roles/Web应用扫描.yaml @@ -0,0 +1,25 @@ +name: Web应用扫描 +description: Web应用漏洞扫描专家，全面的Web安全检测 +user_prompt: 你是一个专业的Web应用漏洞扫描专家。请使用各种Web扫描工具对目标Web应用进行全面的安全检测，包括目录枚举、文件扫描、漏洞识别等工作。 +icon: "\U0001F310" +tools: + - dirsearch + - dirb + - gobuster + - feroxbuster + - ffuf + - wfuzz + - sqlmap + - dalfox + - xsser + - nikto + - nuclei + - wpscan + - httpx + - http-framework-test + - execute-python-script + - install-python-package + - record_vulnerability + - list_knowledge_risk_types + - search_knowledge_base +enabled: true diff --git a/roles/Web框架测试.yaml b/roles/Web框架测试.yaml new file mode 100644 index 00000000..944bc87a --- /dev/null +++ b/roles/Web框架测试.yaml @@ -0,0 +1,19 @@ +name: Web框架测试 +description: Web框架安全测试专家，专注于Web应用框架漏洞检测 +user_prompt: 你是一个专业的Web框架安全测试专家。请使用专业的工具对Web应用框架进行安全测试，识别框架相关的安全漏洞和配置问题。 +icon: "\U0001F310" +tools: + - http-framework-test + - nikto + - nuclei + - wafw00f + - wpscan + - httpx + - burpsuite + - zap + - execute-python-script + - install-python-package + - record_vulnerability + - list_knowledge_risk_types + - search_knowledge_base +enabled: true diff --git a/roles/二进制分析.yaml b/roles/二进制分析.yaml new file mode 100644 index 00000000..56f467d9 --- /dev/null +++ b/roles/二进制分析.yaml @@ -0,0 +1,31 @@ +name: 二进制分析 +description: 二进制分析与利用专家，擅长逆向工程和密码破解 +user_prompt: 你是一个专业的二进制分析与利用专家。请使用逆向工程工具分析二进制文件，识别漏洞，进行利用开发。同时擅长密码破解、哈希分析等技术。 +icon: "\U0001F52C" +tools: + - dirsearch + - docker-bench-security + - exec + - execute-python-script + - install-python-package + - ghidra + - graphql-scanner + - hakrawler + - hash-identifier + - hashcat + - hashpump + - http-framework-test + - httpx + - gdb + - radare2 + - objdump + - strings + - binwalk + - ropper + - ropgadget + - john + - cyberchef + - record_vulnerability + - list_knowledge_risk_types + - search_knowledge_base +enabled: true diff --git a/roles/云安全审计.yaml b/roles/云安全审计.yaml new file mode 100644 index 00000000..1f4efc34 --- /dev/null +++ b/roles/云安全审计.yaml @@ -0,0 +1,17 @@ +name: 云安全审计 +description: 云安全审计专家，多云环境安全检测 +user_prompt: 你是一个专业的云安全审计专家。请使用专业的云安全工具对AWS、Azure、GCP等云环境进行全面的安全审计，包括配置检查、合规性评估、权限审计、安全最佳实践验证等工作。 +icon: "\U00002601" +tools: + - prowler + - scout-suite + - cloudmapper + - pacu + - terrascan + - checkov + - execute-python-script + - install-python-package + - record_vulnerability + - list_knowledge_risk_types + - search_knowledge_base +enabled: true diff --git a/roles/信息收集.yaml b/roles/信息收集.yaml new file mode 100644 index 00000000..98cc5fc9 --- /dev/null +++ b/roles/信息收集.yaml @@ -0,0 +1,31 @@ +name: 信息收集 +description: 资产发现与信息搜集专家 +user_prompt: 你是一个专业的信息收集专家。请使用各种信息收集技术和工具，对目标进行全面的资产发现、子域名枚举、端口扫描、服务识别等信息收集工作。 +icon: "\U0001F50D" +tools: + - amass + - subfinder + - dnsenum + - fierce + - fofa_search + - zoomeye_search + - nmap + - masscan + - rustscan + - arp-scan + - nbtscan + - httpx + - http-framework-test + - katana + - hakrawler + - waybackurls + - paramspider + - gau + - uro + - qsreplace + - execute-python-script + - install-python-package + - record_vulnerability + - list_knowledge_risk_types + - search_knowledge_base +enabled: true diff --git a/roles/后渗透测试.yaml b/roles/后渗透测试.yaml new file mode 100644 index 00000000..80593885 --- /dev/null +++ b/roles/后渗透测试.yaml @@ -0,0 +1,23 @@ +name: 后渗透测试 +description: 后渗透测试专家，权限维持与横向移动 +user_prompt: 你是一个专业的后渗透测试专家。请使用专业的后渗透工具在获得初始访问权限后进行权限提升、横向移动、权限维持、数据收集等后渗透测试工作。 +icon: "\U0001F575" +tools: + - linpeas + - winpeas + - mimikatz + - bloodhound + - impacket + - responder + - netexec + - rpcclient + - smbmap + - enum4linux + - enum4linux-ng + - exec + - execute-python-script + - install-python-package + - record_vulnerability + - list_knowledge_risk_types + - search_knowledge_base +enabled: true diff --git a/roles/容器安全.yaml b/roles/容器安全.yaml new file mode 100644 index 00000000..ce40edfe --- /dev/null +++ b/roles/容器安全.yaml @@ -0,0 +1,18 @@ +name: 容器安全 +description: 容器与Kubernetes安全专家，容器环境安全检测 +user_prompt: 你是一个专业的容器与Kubernetes安全专家。请使用专业的容器安全工具对Docker容器和Kubernetes集群进行全面的安全检测，包括镜像漏洞扫描、配置检查、运行时安全等工作。 +icon: "\U0001F6E1" +tools: + - trivy + - clair + - docker-bench-security + - kube-bench + - kube-hunter + - falco + - exec + - execute-python-script + - install-python-package + - record_vulnerability + - list_knowledge_risk_types + - search_knowledge_base +enabled: true diff --git a/roles/数字取证.yaml b/roles/数字取证.yaml new file mode 100644 index 00000000..7a1d7527 --- /dev/null +++ b/roles/数字取证.yaml @@ -0,0 +1,24 @@ +name: 数字取证 +description: 数字取证与隐写分析专家，文件与内存取证 +user_prompt: 你是一个专业的数字取证与隐写分析专家。请使用专业的取证工具对文件、磁盘镜像、内存转储进行分析，提取证据信息。同时擅长隐写分析、数据恢复、元数据提取等技术。 +icon: "\U0001F50E" +tools: + - volatility + - volatility3 + - foremost + - steghide + - stegsolve + - zsteg + - exiftool + - binwalk + - strings + - xxd + - fcrackzip + - pdfcrack + - exec + - execute-python-script + - install-python-package + - record_vulnerability + - list_knowledge_risk_types + - search_knowledge_base +enabled: true diff --git a/roles/渗透测试.yaml b/roles/渗透测试.yaml new file mode 100644 index 00000000..9c420417 --- /dev/null +++ b/roles/渗透测试.yaml @@ -0,0 +1,33 @@ +name: 渗透测试 +description: 专业渗透测试专家，全面深入的漏洞检测 +user_prompt: 你是一个专业的网络安全渗透测试专家。请使用专业的渗透测试方法和工具，对目标进行全面的安全测试，包括但不限于SQL注入、XSS、CSRF、文件包含、命令执行等常见漏洞。 +icon: "\U0001F3AF" +tools: + - http-framework-test + - httpx + - amass + - anew + - angr + - api-fuzzer + - api-schema-analyzer + - arjun + - arp-scan + - autorecon + - binwalk + - bloodhound + - burpsuite + - cat + - checkov + - checksec + - cloudmapper + - create-file + - cyberchef + - dalfox + - delete-file + - exec + - execute-python-script + - install-python-package + - record_vulnerability + - list_knowledge_risk_types + - search_knowledge_base +enabled: true diff --git a/roles/综合漏洞扫描.yaml b/roles/综合漏洞扫描.yaml new file mode 100644 index 00000000..51d08bac --- /dev/null +++ b/roles/综合漏洞扫描.yaml @@ -0,0 +1,23 @@ +name: 综合漏洞扫描 +description: 综合漏洞扫描专家，多类型漏洞检测 +user_prompt: 你是一个专业的综合漏洞扫描专家。请使用各种漏洞扫描工具对目标进行全面的安全检测，包括Web漏洞、网络服务漏洞、配置缺陷等多种类型的漏洞识别和分析。 +icon: "\U000026A0" +tools: + - nuclei + - nikto + - sqlmap + - nmap + - masscan + - rustscan + - wafw00f + - dalfox + - xsser + - jaeles + - httpx + - http-framework-test + - execute-python-script + - install-python-package + - record_vulnerability + - list_knowledge_risk_types + - search_knowledge_base +enabled: true diff --git a/roles/默认.yaml b/roles/默认.yaml new file mode 100644 index 00000000..62e596d7 --- /dev/null +++ b/roles/默认.yaml @@ -0,0 +1,5 @@ +name: 默认 +description: 默认角色，不额外携带用户提示词，使用默认MCP +user_prompt: "" +icon: "\U0001F535" +enabled: true diff --git a/web/static/css/style.css b/web/static/css/style.css index 91f637b3..9182bd3e 100644 --- a/web/static/css/style.css +++ b/web/static/css/style.css @@ -1476,15 +1476,23 @@ header { .chat-input-container { display: flex; - gap: 12px; - padding: 20px 24px; + flex-direction: row; + gap: 8px; + align-items: center; + padding: 12px 16px; background: var(--bg-primary); - border-top: 1px solid var(--border-color); + border-top: 1px solid rgba(0, 0, 0, 0.06); flex-shrink: 0; width: 100%; box-sizing: border-box; } +.chat-input-container > .chat-input-field { + flex: 1; + display: flex; + min-width: 0; +} + .chat-input-field { flex: 1; position: relative; @@ -1498,16 +1506,17 @@ header { .chat-input-container textarea { flex: 1; min-width: 0; - padding: 12px 16px; - border: 1px solid var(--border-color); - border-radius: 8px; + padding: 10px 14px; + border: 1px solid rgba(0, 0, 0, 0.1); + border-radius: 12px; font-size: 0.9375rem; outline: none; transition: all 0.2s; - background: var(--bg-primary); + background: #ffffff; color: var(--text-primary); resize: none; - height: 44px; + min-height: 40px; + max-height: 120px; font-family: inherit; line-height: 1.5; overflow-y: auto; @@ -1542,34 +1551,46 @@ header { .chat-input-container textarea:focus { border-color: var(--accent-color); - box-shadow: 0 0 0 3px rgba(0, 102, 255, 0.1); + box-shadow: 0 0 0 2px rgba(0, 102, 255, 0.08); + background: #ffffff; } .chat-input-container textarea::placeholder { color: var(--text-muted); } -.chat-input-container button { - padding: 12px 24px; +.chat-input-container .send-btn { + display: flex; + align-items: center; + gap: 6px; + padding: 10px 18px; + height: 40px; background: var(--accent-color); color: white; border: none; - border-radius: 8px; + border-radius: 12px; cursor: pointer; font-size: 0.9375rem; font-weight: 500; transition: all 0.2s; white-space: nowrap; + flex-shrink: 0; + box-sizing: border-box; } -.chat-input-container button:hover { +.chat-input-container .send-btn:hover { background: var(--accent-hover); transform: translateY(-1px); - box-shadow: var(--shadow-md); + box-shadow: 0 2px 8px rgba(0, 102, 255, 0.25); } -.chat-input-container button:active { +.chat-input-container .send-btn:active { transform: translateY(0); + box-shadow: 0 1px 4px rgba(0, 102, 255, 0.2); +} + +.chat-input-container .send-btn svg { + flex-shrink: 0; } .mention-suggestions { @@ -7639,3 +7660,957 @@ header { font-size: 0.8125rem !important; } } + +/* 角色选择器样式 */ +.role-selector-btn { + display: flex; + align-items: center; + gap: 6px; + padding: 8px 12px; + height: 40px; + background: #ffffff; + border: 1px solid rgba(0, 0, 0, 0.1); + border-radius: 12px; + color: var(--text-primary); + cursor: pointer; + transition: all 0.2s; + flex-shrink: 0; + font-size: 0.875rem; + font-weight: 500; + white-space: nowrap; + box-sizing: border-box; +} + +.role-selector-btn:hover { + background: #f8f9fa; + border-color: rgba(0, 102, 255, 0.3); + box-shadow: 0 0 0 2px rgba(0, 102, 255, 0.08); +} + +.role-selector-btn:active { + transform: scale(0.98); + background: #f0f2f5; +} + +.role-selector-btn.active { + background: #f0f2f5; + border-color: rgba(0, 102, 255, 0.4); + box-shadow: 0 0 0 2px rgba(0, 102, 255, 0.12); +} + +.role-selector-btn.active .role-selector-arrow { + transform: rotate(180deg); + color: rgba(0, 102, 255, 0.8); +} + +.role-selector-icon { + font-size: 1rem; + line-height: 1; + flex-shrink: 0; + display: flex; + align-items: center; + justify-content: center; +} + +.role-selector-text { + font-size: 0.875rem; + font-weight: 500; + color: var(--text-primary); + flex-shrink: 0; +} + +.role-selector-arrow { + width: 10px; + height: 10px; + color: rgba(0, 0, 0, 0.4); + flex-shrink: 0; + transition: transform 0.2s, color 0.2s; +} + +.role-selector-btn:hover .role-selector-arrow { + color: rgba(0, 102, 255, 0.6); +} + +/* 角色选择器包装器 */ +.role-selector-wrapper { + position: relative; + flex-shrink: 0; +} + +/* 主内容区域角色选择面板样式（下拉菜单形式） */ +.role-selection-panel { + position: absolute; + bottom: calc(100% + 8px); + left: 0; + width: 320px; + max-width: calc(100vw - 32px); + max-height: 60vh; + background: var(--bg-primary); + border: 1px solid var(--border-color); + border-radius: 8px; + padding: 8px; + box-shadow: 0 4px 16px rgba(0, 0, 0, 0.12), 0 2px 8px rgba(0, 0, 0, 0.08); + z-index: 1000; + overflow: hidden; + display: flex; + flex-direction: column; + animation: slideUp 0.3s ease-out; +} + +@keyframes slideUp { + from { + opacity: 0; + transform: translateY(10px); + } + to { + opacity: 1; + transform: translateY(0); + } +} + +.role-selection-panel-header { + display: flex; + justify-content: space-between; + align-items: center; + margin-bottom: 8px; + padding: 8px 4px; + border-bottom: 1px solid var(--border-color); + position: sticky; + top: 0; + background: var(--bg-primary); + z-index: 10; +} + +.role-selection-panel-title { + font-size: 0.875rem; + font-weight: 600; + color: var(--text-primary); + margin: 0; +} + +.role-selection-panel-close { + display: flex; + align-items: center; + justify-content: center; + width: 24px; + height: 24px; + border: none; + background: transparent; + border-radius: 4px; + cursor: pointer; + color: var(--text-secondary); + transition: all 0.2s; + padding: 0; +} + +.role-selection-panel-close:hover { + background: var(--bg-hover); + color: var(--text-primary); +} + +.role-selection-list-main { + display: flex; + flex-direction: column; + gap: 4px; + /* 限制显示8个角色：每个角色约70px高度 + gap，8个角色约580px */ + max-height: 580px; + overflow-y: auto; + padding-right: 4px; + flex: 1; +} + +.role-selection-list-main::-webkit-scrollbar { + width: 6px; +} + +.role-selection-list-main::-webkit-scrollbar-track { + background: transparent; +} + +.role-selection-list-main::-webkit-scrollbar-thumb { + background: rgba(0, 0, 0, 0.2); + border-radius: 3px; +} + +.role-selection-list-main::-webkit-scrollbar-thumb:hover { + background: rgba(0, 0, 0, 0.3); +} + +.role-selection-item-main { + display: flex; + flex-direction: row; + align-items: flex-start; + gap: 10px; + padding: 8px 10px; + border: 1px solid var(--border-color); + border-radius: 6px; + background: var(--bg-primary); + cursor: pointer; + transition: all 0.2s ease; + position: relative; +} + +.role-selection-item-main:hover { + background: var(--bg-secondary); + border-color: rgba(138, 43, 226, 0.3); +} + +.role-selection-item-main.selected { + background: rgba(138, 43, 226, 0.1); + border-color: #8a2be2; +} + +.role-selection-item-icon-main { + font-size: 1.125rem; + flex-shrink: 0; + width: 28px; + height: 28px; + display: flex; + align-items: center; + justify-content: center; + background: var(--bg-secondary); + border-radius: 6px; + transition: all 0.2s; +} + +.role-selection-item-main.selected .role-selection-item-icon-main { + background: rgba(138, 43, 226, 0.15); +} + +.role-selection-item-content-main { + display: flex; + flex-direction: column; + gap: 2px; + min-width: 0; + flex: 1; +} + +.role-selection-item-name-main { + font-weight: 500; + color: var(--text-primary); + font-size: 0.8125rem; + line-height: 1.3; + margin: 0; + transition: color 0.2s; +} + +.role-selection-item-main.selected .role-selection-item-name-main { + color: #8a2be2; + font-weight: 600; +} + +.role-selection-item-description-main { + font-size: 0.75rem; + color: var(--text-secondary); + line-height: 1.3; + margin: 0; + display: -webkit-box; + -webkit-line-clamp: 1; + line-clamp: 1; + -webkit-box-orient: vertical; + overflow: hidden; + text-overflow: ellipsis; +} + +.role-selection-checkmark-main { + position: absolute; + top: 6px; + right: 8px; + width: 16px; + height: 16px; + display: flex; + align-items: center; + justify-content: center; + background: #8a2be2; + color: white; + border-radius: 50%; + font-size: 0.625rem; + font-weight: 700; + flex-shrink: 0; +} + +@keyframes highlight-flash { + 0%, 100% { + background-color: transparent; + } + 50% { + background-color: rgba(138, 43, 226, 0.1); + } +} + +.role-select-item { + display: flex; + align-items: flex-start; + gap: 14px; + padding: 16px 18px; + border: 2px solid var(--border-color); + border-radius: 12px; + background: var(--bg-primary); + cursor: pointer; + transition: all 0.25s cubic-bezier(0.4, 0, 0.2, 1); + position: relative; + box-shadow: 0 1px 3px rgba(0, 0, 0, 0.05); +} + +.role-select-item::before { + content: ''; + position: absolute; + top: 0; + left: 0; + right: 0; + bottom: 0; + border-radius: 12px; + background: linear-gradient(135deg, rgba(138, 43, 226, 0.03) 0%, rgba(138, 43, 226, 0.01) 100%); + opacity: 0; + transition: opacity 0.25s; +} + +.role-select-item:hover { + background: linear-gradient(135deg, #ffffff 0%, #f8f9fa 100%); + border-color: rgba(138, 43, 226, 0.3); + box-shadow: 0 4px 12px rgba(138, 43, 226, 0.1), 0 2px 4px rgba(0, 0, 0, 0.08); + transform: translateY(-2px); +} + +.role-select-item:hover::before { + opacity: 1; +} + +.role-select-item.selected { + background: linear-gradient(135deg, rgba(138, 43, 226, 0.1) 0%, rgba(138, 43, 226, 0.05) 100%); + border-color: #8a2be2; + box-shadow: 0 4px 16px rgba(138, 43, 226, 0.2), 0 2px 8px rgba(138, 43, 226, 0.15); + transform: translateY(-2px); +} + +.role-select-item.selected::before { + opacity: 1; + background: linear-gradient(135deg, rgba(138, 43, 226, 0.08) 0%, rgba(138, 43, 226, 0.03) 100%); +} + +.role-select-item.selected::after { + content: '✓'; + position: absolute; + top: 12px; + right: 14px; + width: 24px; + height: 24px; + display: flex; + align-items: center; + justify-content: center; + background: linear-gradient(135deg, #8a2be2 0%, #6a1bb2 100%); + color: white; + border-radius: 50%; + font-size: 0.75rem; + font-weight: 700; + box-shadow: 0 2px 6px rgba(138, 43, 226, 0.4); +} + +.role-select-item-icon { + font-size: 1.5rem; + flex-shrink: 0; + width: 40px; + height: 40px; + display: flex; + align-items: center; + justify-content: center; + background: linear-gradient(135deg, #f8f9fa 0%, #e9ecef 100%); + border-radius: 10px; + transition: all 0.25s; + border: 1px solid rgba(0, 0, 0, 0.05); +} + +.role-select-item.selected .role-select-item-icon { + background: linear-gradient(135deg, rgba(138, 43, 226, 0.15) 0%, rgba(138, 43, 226, 0.08) 100%); + border-color: rgba(138, 43, 226, 0.2); + transform: scale(1.05); +} + +.role-select-item-content { + flex: 1; + display: flex; + flex-direction: column; + gap: 6px; + min-width: 0; +} + +.role-select-item-name { + font-weight: 600; + color: var(--text-primary); + font-size: 1rem; + line-height: 1.4; + margin: 0; + transition: color 0.2s; +} + +.role-select-item.selected .role-select-item-name { + color: #8a2be2; +} + +.role-select-item-description { + font-size: 0.8125rem; + color: var(--text-secondary); + line-height: 1.5; + margin: 0; + display: -webkit-box; + -webkit-line-clamp: 2; + line-clamp: 2; + -webkit-box-orient: vertical; + overflow: hidden; + text-overflow: ellipsis; +} + +/* 角色管理页面样式 */ +.roles-controls { + margin-bottom: 24px; +} + +.roles-stats-bar { + display: flex; + gap: 16px; + padding: 16px; + background: var(--bg-primary); + border: 1px solid var(--border-color); + border-radius: 8px; +} + +.role-stat-item { + display: flex; + flex-direction: column; + gap: 4px; +} + +.role-stat-label { + font-size: 0.8125rem; + color: var(--text-secondary); +} + +.role-stat-value { + font-size: 1.25rem; + font-weight: 600; + color: var(--text-primary); +} + +.roles-list { + display: flex; + flex-direction: column; + gap: 12px; +} + +.role-item { + display: flex; + justify-content: space-between; + align-items: flex-start; + gap: 16px; + padding: 16px; + background: var(--bg-primary); + border: 1px solid var(--border-color); + border-radius: 8px; + transition: all 0.2s; +} + +.role-item:hover { + box-shadow: var(--shadow-sm); + border-color: var(--accent-color); +} + +.role-item-content { + flex: 1; + display: flex; + flex-direction: column; + gap: 8px; +} + +.role-item-header { + display: flex; + align-items: center; + gap: 12px; +} + +.role-item-name { + font-size: 1rem; + font-weight: 600; + color: var(--text-primary); +} + +.role-item-badge { + padding: 4px 10px; + border-radius: 12px; + font-size: 0.75rem; + font-weight: 600; +} + +.role-item-badge.enabled { + background: rgba(40, 167, 69, 0.12); + color: var(--success-color); +} + +.role-item-badge.disabled { + background: rgba(220, 53, 69, 0.12); + color: var(--error-color); +} + +.role-item-description { + font-size: 0.875rem; + color: var(--text-secondary); + line-height: 1.5; +} + +.role-item-details { + display: flex; + flex-direction: column; + gap: 6px; + margin-top: 8px; +} + +.role-item-detail { + display: flex; + gap: 8px; + font-size: 0.8125rem; +} + +.role-item-detail-label { + color: var(--text-secondary); + font-weight: 500; + min-width: 80px; +} + +.role-item-detail-value { + color: var(--text-primary); + flex: 1; + word-break: break-word; +} + +.role-item-actions { + display: flex; + gap: 8px; + flex-shrink: 0; +} + +.empty-state { + padding: 48px 24px; + text-align: center; + color: var(--text-secondary); + font-size: 0.9375rem; +} + +/* 角色MCP选择列表样式 */ +.role-mcps-list { + margin-top: 8px; + border: 1px solid var(--border-color); + border-radius: 8px; + background: var(--bg-primary); + min-height: 120px; + max-height: 300px; + overflow-y: auto; + padding: 12px; +} + +.role-mcps-checkbox-list { + display: flex; + flex-direction: column; + gap: 8px; +} + +.role-mcp-checkbox-item { + display: flex; + align-items: center; + padding: 10px 12px; + border: 1px solid var(--border-color); + border-radius: 6px; + background: var(--bg-primary); + cursor: pointer; + transition: all 0.2s ease; + user-select: none; +} + +.role-mcp-checkbox-item:hover { + background: var(--bg-secondary); + border-color: var(--accent-color); + box-shadow: 0 2px 4px rgba(0, 102, 255, 0.1); +} + +.role-mcp-checkbox-item input[type="checkbox"] { + width: 18px; + height: 18px; + margin-right: 12px; + cursor: pointer; + accent-color: var(--accent-color); +} + +.role-mcp-checkbox-item:has(input[type="checkbox"]:checked) { + background: rgba(0, 102, 255, 0.08); + border-color: var(--accent-color); +} + +.role-mcp-checkbox-content { + flex: 1; + display: flex; + flex-direction: column; + gap: 4px; +} + +.role-mcp-name { + display: flex; + align-items: center; + gap: 8px; + font-weight: 500; + color: var(--text-primary); + font-size: 0.9375rem; +} + +.role-mcp-status { + padding: 2px 8px; + border-radius: 10px; + font-size: 0.75rem; + font-weight: 500; +} + +.role-mcp-status.status-connected { + background: rgba(40, 167, 69, 0.12); + color: var(--success-color); +} + +.role-mcp-status.status-disconnected { + background: rgba(108, 117, 125, 0.12); + color: var(--text-secondary); +} + +.role-mcp-status.status-connecting { + background: rgba(255, 193, 7, 0.12); + color: var(--warning-color); +} + +.role-mcp-status.status-error { + background: rgba(220, 53, 69, 0.12); + color: var(--error-color); +} + +.role-mcp-status.status-disabled { + background: rgba(108, 117, 125, 0.12); + color: var(--text-muted); +} + +.role-mcp-tool-count { + padding: 2px 6px; + background: rgba(0, 102, 255, 0.1); + color: var(--accent-color); + border-radius: 8px; + font-size: 0.75rem; + font-weight: 500; +} + +.mcp-loading, +.mcp-empty, +.mcp-error { + padding: 24px; + text-align: center; + color: var(--text-secondary); + font-size: 0.875rem; +} + +.mcp-error { + color: var(--error-color); +} + +/* 角色工具选择列表样式 */ +.role-tools-controls { + margin-bottom: 12px; +} + +.role-tools-actions { + display: flex; + gap: 8px; + align-items: center; + flex-wrap: wrap; + margin-bottom: 8px; +} + +.role-tools-search-box { + position: relative; + flex: 1; + min-width: 200px; + max-width: 400px; +} + +.role-tools-search-box input { + width: 100%; + padding: 8px 32px 8px 12px; + border: 1px solid var(--border-color); + border-radius: 6px; + font-size: 0.875rem; + background: var(--bg-primary); + color: var(--text-primary); + transition: all 0.2s; +} + +.role-tools-search-box input:focus { + outline: none; + border-color: var(--accent-color); + box-shadow: 0 0 0 2px rgba(0, 102, 255, 0.1); +} + +.role-tools-search-clear { + position: absolute; + right: 8px; + top: 50%; + transform: translateY(-50%); + width: 20px; + height: 20px; + padding: 0; + border: none; + background: transparent; + color: var(--text-muted); + cursor: pointer; + border-radius: 4px; + display: flex; + align-items: center; + justify-content: center; + transition: all 0.2s; +} + +.role-tools-search-clear:hover { + background: var(--bg-tertiary); + color: var(--text-primary); +} + +.role-tools-stats { + display: flex; + gap: 16px; + padding: 8px 12px; + background: var(--bg-secondary); + border: 1px solid var(--border-color); + border-radius: 6px; + font-size: 0.8125rem; + color: var(--text-secondary); +} + +.role-tools-stats span { + white-space: nowrap; +} + +.role-tools-stats strong { + color: var(--text-primary); + font-weight: 600; +} + +.role-tools-list { + margin-top: 8px; + border: 1px solid var(--border-color); + border-radius: 8px; + background: var(--bg-primary); + min-height: 200px; + max-height: 400px; + overflow-y: auto; + padding: 12px; +} + +.role-tools-list-items { + display: flex; + flex-direction: column; + gap: 8px; +} + +.role-tool-item { + display: flex; + align-items: flex-start; + gap: 12px; + padding: 10px 12px; + border: 1px solid var(--border-color); + border-radius: 6px; + background: var(--bg-primary); + transition: all 0.2s ease; +} + +.role-tool-item:hover { + background: var(--bg-secondary); + border-color: var(--accent-color); + box-shadow: 0 2px 4px rgba(0, 102, 255, 0.1); +} + +.role-tool-item input[type="checkbox"] { + width: 18px; + height: 18px; + margin-top: 2px; + cursor: pointer; + accent-color: var(--accent-color); + flex-shrink: 0; +} + +.role-tool-item-info { + flex: 1; + display: flex; + flex-direction: column; + gap: 4px; + min-width: 0; +} + +.role-tool-item-name { + display: flex; + align-items: center; + gap: 8px; + font-weight: 500; + color: var(--text-primary); + font-size: 0.9375rem; +} + +.role-tool-item-desc { + font-size: 0.8125rem; + color: var(--text-secondary); + line-height: 1.5; +} + +.external-tool-badge { + padding: 2px 6px; + background: rgba(0, 102, 255, 0.1); + color: var(--accent-color); + border-radius: 8px; + font-size: 0.75rem; + font-weight: 500; + white-space: nowrap; +} + +.role-tools-pagination { + display: flex; + justify-content: space-between; + align-items: center; + padding: 12px; + border-top: 1px solid var(--border-color); + background: var(--bg-secondary); + flex-wrap: wrap; + gap: 12px; +} + +.role-tools-pagination .pagination-info { + font-size: 0.8125rem; + color: var(--text-secondary); +} + +.role-tools-pagination .pagination-controls { + display: flex; + gap: 8px; + align-items: center; +} + +.role-tools-pagination .pagination-page { + font-size: 0.8125rem; + color: var(--text-secondary); + padding: 0 8px; +} + +.tools-loading, +.tools-empty, +.tools-error { + padding: 24px; + text-align: center; + color: var(--text-secondary); + font-size: 0.875rem; +} + +.tools-error { + color: var(--error-color); +} + +/* 默认角色提示信息样式 */ +.role-tools-default-hint { + margin-top: 8px; + padding: 16px; + border: 1px solid var(--border-color); + border-radius: 8px; + background: var(--bg-secondary); +} + +.role-tools-default-info { + display: flex; + align-items: flex-start; + gap: 12px; +} + +.role-tools-default-icon { + font-size: 1.25rem; + flex-shrink: 0; + margin-top: 2px; +} + +.role-tools-default-content { + flex: 1; + display: flex; + flex-direction: column; + gap: 4px; +} + +.role-tools-default-title { + font-weight: 600; + color: var(--text-primary); + font-size: 0.9375rem; +} + +.role-tools-default-desc { + color: var(--text-secondary); + font-size: 0.875rem; + line-height: 1.5; +} + +/* 角色选择面板响应式样式 */ +@media (max-width: 768px) { + .role-selection-panel { + width: calc(100vw - 16px); + max-width: calc(100vw - 16px); + left: -8px; + padding: 8px; + border-radius: 8px; + max-height: 60vh; + } + + .role-selection-panel-header { + margin-bottom: 8px; + padding: 6px 4px; + } + + .role-selection-panel-title { + font-size: 0.8125rem; + } + + .role-selection-list-main { + /* 在移动设备上限制显示8个角色，每个约60px，总计约480px */ + max-height: 480px; + gap: 4px; + } + + .role-selection-item-main { + padding: 8px 10px; + } + + .role-selection-item-icon-main { + width: 24px; + height: 24px; + font-size: 1rem; + } + + .role-selection-item-name-main { + font-size: 0.75rem; + } + + .role-selection-item-description-main { + font-size: 0.6875rem; + } +} + +@media (max-width: 480px) { + .role-selection-panel { + width: calc(100vw - 8px); + max-width: calc(100vw - 8px); + left: -4px; + padding: 6px; + max-height: 50vh; + } + + .role-selection-list-main { + /* 在小屏幕上限制显示8个角色，每个约55px，总计约450px */ + max-height: 450px; + } + + .role-selection-item-main { + padding: 6px 8px; + } +} diff --git a/web/static/js/builtin-tools.js b/web/static/js/builtin-tools.js new file mode 100644 index 00000000..942bfc5a --- /dev/null +++ b/web/static/js/builtin-tools.js @@ -0,0 +1,27 @@ +/** + * 内置工具名称常量 + * 所有前端代码中使用内置工具名称的地方都应该使用这些常量，而不是硬编码字符串 + * + * 注意：这些常量必须与后端的 internal/mcp/builtin/constants.go 中的常量保持一致 + */ + +// 内置工具名称常量 +const BuiltinTools = { + // 漏洞管理工具 + RECORD_VULNERABILITY: 'record_vulnerability', + + // 知识库工具 + LIST_KNOWLEDGE_RISK_TYPES: 'list_knowledge_risk_types', + SEARCH_KNOWLEDGE_BASE: 'search_knowledge_base' +}; + +// 检查是否是内置工具 +function isBuiltinTool(toolName) { + return Object.values(BuiltinTools).includes(toolName); +} + +// 获取所有内置工具名称列表 +function getAllBuiltinTools() { + return Object.values(BuiltinTools); +} + diff --git a/web/static/js/chat.js b/web/static/js/chat.js index 787df6ac..285a11f9 100644 --- a/web/static/js/chat.js +++ b/web/static/js/chat.js @@ -145,6 +145,9 @@ async function sendMessage() { let mcpExecutionIds = []; try { + // 获取当前选中的角色（从 roles.js 的函数获取） + const roleName = typeof getCurrentRole === 'function' ? getCurrentRole() : ''; + const response = await apiFetch('/api/agent-loop/stream', { method: 'POST', headers: { @@ -152,7 +155,8 @@ async function sendMessage() { }, body: JSON.stringify({ message: message, - conversationId: currentConversationId + conversationId: currentConversationId, + role: roleName || undefined }), }); @@ -252,6 +256,13 @@ if (typeof window !== 'undefined') { } function ensureMentionToolsLoaded() { + // 检查角色是否改变，如果改变则强制重新加载 + if (typeof window !== 'undefined' && window._mentionToolsRoleChanged) { + mentionToolsLoaded = false; + mentionTools = []; + delete window._mentionToolsRoleChanged; + } + if (mentionToolsLoaded) { return Promise.resolve(mentionTools); } @@ -282,6 +293,9 @@ async function fetchMentionTools() { const collected = []; try { + // 获取当前选中的角色（从 roles.js 的函数获取） + const roleName = typeof getCurrentRole === 'function' ? getCurrentRole() : ''; + // 同时获取外部MCP列表 try { const mcpResponse = await apiFetch('/api/external-mcp'); @@ -300,7 +314,13 @@ async function fetchMentionTools() { } while (page <= totalPages && page <= 20) { - const response = await apiFetch(`/api/config/tools?page=${page}&page_size=${pageSize}`); + // 构建API URL，如果指定了角色，添加role查询参数 + let url = `/api/config/tools?page=${page}&page_size=${pageSize}`; + if (roleName && roleName !== '默认') { + url += `&role=${encodeURIComponent(roleName)}`; + } + + const response = await apiFetch(url); if (!response.ok) { break; } @@ -316,10 +336,20 @@ async function fetchMentionTools() { return; } seen.add(toolKey); + + // 确定工具在当前角色中的启用状态 + // 如果有 role_enabled 字段，使用它（表示指定了角色） + // 否则使用 enabled 字段（表示未指定角色或使用所有工具） + let roleEnabled = tool.enabled !== false; + if (tool.role_enabled !== undefined && tool.role_enabled !== null) { + roleEnabled = tool.role_enabled; + } + collected.push({ name: tool.name, description: tool.description || '', - enabled: tool.enabled !== false, + enabled: tool.enabled !== false, // 工具本身的启用状态 + roleEnabled: roleEnabled, // 在当前角色中的启用状态 isExternal: !!tool.is_external, externalMcp: tool.external_mcp || '', toolKey: toolKey, // 保存唯一标识符 @@ -488,6 +518,15 @@ function updateMentionCandidates() { } filtered = filtered.slice().sort((a, b) => { + // 如果指定了角色，优先显示在当前角色中启用的工具 + if (a.roleEnabled !== undefined || b.roleEnabled !== undefined) { + const aRoleEnabled = a.roleEnabled !== undefined ? a.roleEnabled : a.enabled; + const bRoleEnabled = b.roleEnabled !== undefined ? b.roleEnabled : b.enabled; + if (aRoleEnabled !== bRoleEnabled) { + return aRoleEnabled ? -1 : 1; // 启用的工具排在前面 + } + } + if (normalizedQuery) { // 精确匹配MCP名称的工具优先显示 const aMcpExact = a.externalMcp && a.externalMcp.toLowerCase() === normalizedQuery; @@ -502,8 +541,11 @@ function updateMentionCandidates() { return aStarts ? -1 : 1; } } - if (a.enabled !== b.enabled) { - return a.enabled ? -1 : 1; + // 如果指定了角色，使用 roleEnabled；否则使用 enabled + const aEnabled = a.roleEnabled !== undefined ? a.roleEnabled : a.enabled; + const bEnabled = b.roleEnabled !== undefined ? b.roleEnabled : b.enabled; + if (aEnabled !== bEnabled) { + return aEnabled ? -1 : 1; } return a.name.localeCompare(b.name, 'zh-CN'); }); @@ -545,13 +587,16 @@ function renderMentionSuggestions({ showLoading = false } = {}) { const itemsHtml = mentionFilteredTools.map((tool, index) => { const activeClass = index === mentionState.selectedIndex ? 'active' : ''; - const disabledClass = tool.enabled ? '' : 'disabled'; + // 如果工具有 roleEnabled 字段（指定了角色），使用它；否则使用 enabled + const toolEnabled = tool.roleEnabled !== undefined ? tool.roleEnabled : tool.enabled; + const disabledClass = toolEnabled ? '' : 'disabled'; const badge = tool.isExternal ? '外部' : '内置'; const nameHtml = escapeHtml(tool.name); const description = tool.description && tool.description.length > 0 ? escapeHtml(tool.description) : '暂无描述'; const descHtml = `
${description}
`; - const statusLabel = tool.enabled ? '可用' : '已禁用'; - const statusClass = tool.enabled ? 'enabled' : 'disabled'; + // 根据工具在当前角色中的启用状态显示状态标签 + const statusLabel = toolEnabled ? '可用' : (tool.roleEnabled !== undefined ? '已禁用（当前角色）' : '已禁用'); + const statusClass = toolEnabled ? 'enabled' : 'disabled'; const originLabel = tool.isExternal ? (tool.externalMcp ? `来源：${escapeHtml(tool.externalMcp)}` : '来源：外部MCP') : '来源：内置工具'; @@ -1109,7 +1154,7 @@ function renderProcessDetails(messageId, processDetails) { itemTitle = `${statusIcon} 工具 ${escapeHtml(toolName)} 执行${success ? '完成' : '失败'}`; // 如果是知识检索工具，添加特殊标记 - if (toolName === 'search_knowledge_base' && success) { + if (toolName === BuiltinTools.SEARCH_KNOWLEDGE_BASE && success) { itemTitle = `📚 ${itemTitle} - 知识检索`; } } else if (eventType === 'knowledge_retrieval') { diff --git a/web/static/js/roles.js b/web/static/js/roles.js new file mode 100644 index 00000000..bd429e4d --- /dev/null +++ b/web/static/js/roles.js @@ -0,0 +1,1230 @@ +// 角色管理相关功能 +let currentRole = localStorage.getItem('currentRole') || ''; +let roles = []; +let allRoleTools = []; // 存储所有工具列表（用于角色工具选择） +let roleToolsPagination = { + page: 1, + pageSize: 20, + total: 0, + totalPages: 1 +}; +let roleToolsSearchKeyword = ''; // 工具搜索关键词 +let roleToolStateMap = new Map(); // 工具状态映射：toolKey -> { enabled: boolean, ... } +let roleUsesAllTools = false; // 标记角色是否使用所有工具（当没有配置tools时） +let totalEnabledToolsInMCP = 0; // 已启用的工具总数（从MCP管理中获取，从API响应中获取） +let roleConfiguredTools = new Set(); // 角色配置的工具列表（用于确定哪些工具应该被选中） + +// 对角色列表进行排序：默认角色排在第一个，其他按名称排序 +function sortRoles(rolesArray) { + const sortedRoles = [...rolesArray]; + // 将"默认"角色分离出来 + const defaultRole = sortedRoles.find(r => r.name === '默认'); + const otherRoles = sortedRoles.filter(r => r.name !== '默认'); + + // 其他角色按名称排序，保持固定顺序 + otherRoles.sort((a, b) => { + const nameA = a.name || ''; + const nameB = b.name || ''; + return nameA.localeCompare(nameB, 'zh-CN'); + }); + + // 将"默认"角色放在第一个，其他角色按排序后的顺序跟在后面 + const result = defaultRole ? [defaultRole, ...otherRoles] : otherRoles; + return result; +} + +// 加载所有角色 +async function loadRoles() { + try { + const response = await apiFetch('/api/roles'); + if (!response.ok) { + throw new Error('加载角色失败'); + } + const data = await response.json(); + roles = data.roles || []; + updateRoleSelectorDisplay(); + renderRoleSelectionSidebar(); // 渲染侧边栏角色列表 + return roles; + } catch (error) { + console.error('加载角色失败:', error); + showNotification('加载角色失败: ' + error.message, 'error'); + return []; + } +} + +// 处理角色变更 +function handleRoleChange(roleName) { + const oldRole = currentRole; + currentRole = roleName || ''; + localStorage.setItem('currentRole', currentRole); + updateRoleSelectorDisplay(); + renderRoleSelectionSidebar(); // 更新侧边栏选中状态 + + // 当角色切换时，如果工具列表已加载，标记为需要重新加载 + // 这样下次触发@工具建议时会使用新的角色重新加载工具列表 + if (oldRole !== currentRole && typeof window !== 'undefined') { + // 通过设置一个标记来通知chat.js需要重新加载工具列表 + window._mentionToolsRoleChanged = true; + } +} + +// 更新角色选择器显示 +function updateRoleSelectorDisplay() { + const roleSelectorBtn = document.getElementById('role-selector-btn'); + const roleSelectorIcon = document.getElementById('role-selector-icon'); + const roleSelectorText = document.getElementById('role-selector-text'); + + if (!roleSelectorBtn || !roleSelectorIcon || !roleSelectorText) return; + + let selectedRole; + if (currentRole && currentRole !== '默认') { + selectedRole = roles.find(r => r.name === currentRole); + } else { + selectedRole = roles.find(r => r.name === '默认'); + } + + if (selectedRole) { + // 使用配置中的图标，如果没有则使用默认图标 + let icon = selectedRole.icon || '🔵'; + // 如果 icon 是 Unicode 转义格式（\U0001F3C6），需要转换为 emoji + if (icon && typeof icon === 'string') { + const unicodeMatch = icon.match(/^"?\\U([0-9A-F]{8})"?$/i); + if (unicodeMatch) { + try { + const codePoint = parseInt(unicodeMatch[1], 16); + icon = String.fromCodePoint(codePoint); + } catch (e) { + // 如果转换失败，使用默认图标 + console.warn('转换 icon Unicode 转义失败:', icon, e); + icon = '🔵'; + } + } + } + roleSelectorIcon.textContent = icon; + roleSelectorText.textContent = selectedRole.name || '默认'; + } else { + // 默认角色 + roleSelectorIcon.textContent = '🔵'; + roleSelectorText.textContent = '默认'; + } +} + +// 渲染主内容区域角色选择列表 +function renderRoleSelectionSidebar() { + const roleList = document.getElementById('role-selection-list'); + if (!roleList) return; + + // 清空列表 + roleList.innerHTML = ''; + + // 根据角色配置获取图标，如果没有配置则使用默认图标 + function getRoleIcon(role) { + if (role.icon) { + // 如果 icon 是 Unicode 转义格式（\U0001F3C6），需要转换为 emoji + let icon = role.icon; + // 检查是否是 Unicode 转义格式（可能包含引号） + const unicodeMatch = icon.match(/^"?\\U([0-9A-F]{8})"?$/i); + if (unicodeMatch) { + try { + const codePoint = parseInt(unicodeMatch[1], 16); + icon = String.fromCodePoint(codePoint); + } catch (e) { + // 如果转换失败，使用原值 + console.warn('转换 icon Unicode 转义失败:', icon, e); + } + } + return icon; + } + // 如果没有配置图标，根据角色名称的首字符生成默认图标 + // 使用一些通用的默认图标 + return '👤'; + } + + // 对角色进行排序：默认角色第一个，其他按名称排序 + const sortedRoles = sortRoles(roles); + + // 只显示已启用的角色 + const enabledSortedRoles = sortedRoles.filter(r => r.enabled !== false); + + enabledSortedRoles.forEach(role => { + const isDefaultRole = role.name === '默认'; + const isSelected = isDefaultRole ? (currentRole === '' || currentRole === '默认') : (currentRole === role.name); + const roleItem = document.createElement('div'); + roleItem.className = 'role-selection-item-main' + (isSelected ? ' selected' : ''); + roleItem.onclick = () => { + selectRole(role.name); + closeRoleSelectionPanel(); // 选择后自动关闭面板 + }; + const icon = getRoleIcon(role); + + // 处理默认角色的描述 + let description = role.description || '暂无描述'; + if (isDefaultRole && !role.description) { + description = '默认角色，不额外携带用户提示词，使用默认MCP'; + } + + roleItem.innerHTML = ` +
${icon}
+
+
${escapeHtml(role.name)}
+
${escapeHtml(description)}
+
+ ${isSelected ? '
✓
' : ''} + `; + roleList.appendChild(roleItem); + }); +} + +// 选择角色 +function selectRole(roleName) { + // 将"默认"映射为空字符串（表示默认角色） + if (roleName === '默认') { + roleName = ''; + } + handleRoleChange(roleName); + renderRoleSelectionSidebar(); // 重新渲染以更新选中状态 +} + +// 切换角色选择面板显示/隐藏 +function toggleRoleSelectionPanel() { + const panel = document.getElementById('role-selection-panel'); + const roleSelectorBtn = document.getElementById('role-selector-btn'); + if (!panel) return; + + const isHidden = panel.style.display === 'none' || !panel.style.display; + + if (isHidden) { + panel.style.display = 'flex'; // 使用flex布局 + // 添加打开状态的视觉反馈 + if (roleSelectorBtn) { + roleSelectorBtn.classList.add('active'); + } + + // 确保面板渲染后再检查位置 + setTimeout(() => { + const wrapper = document.querySelector('.role-selector-wrapper'); + if (wrapper) { + const rect = wrapper.getBoundingClientRect(); + const panelHeight = panel.offsetHeight || 400; + const viewportHeight = window.innerHeight; + + // 如果面板顶部超出视窗，滚动到合适位置 + if (rect.top - panelHeight < 0) { + const scrollY = window.scrollY + rect.top - panelHeight - 20; + window.scrollTo({ top: Math.max(0, scrollY), behavior: 'smooth' }); + } + } + }, 10); + } else { + panel.style.display = 'none'; + // 移除打开状态的视觉反馈 + if (roleSelectorBtn) { + roleSelectorBtn.classList.remove('active'); + } + } +} + +// 关闭角色选择面板（选择角色后自动调用） +function closeRoleSelectionPanel() { + const panel = document.getElementById('role-selection-panel'); + const roleSelectorBtn = document.getElementById('role-selector-btn'); + if (panel) { + panel.style.display = 'none'; + } + if (roleSelectorBtn) { + roleSelectorBtn.classList.remove('active'); + } +} + +// 转义HTML +function escapeHtml(text) { + const div = document.createElement('div'); + div.textContent = text; + return div.innerHTML; +} + +// 刷新角色列表 +async function refreshRoles() { + await loadRoles(); + // 检查当前页面是否为角色管理页面 + const currentPage = typeof window.currentPage === 'function' ? window.currentPage() : (window.currentPage || 'chat'); + if (currentPage === 'roles-management') { + renderRolesList(); + } + // 始终更新侧边栏角色选择列表 + renderRoleSelectionSidebar(); +} + +// 渲染角色列表 +function renderRolesList() { + const rolesList = document.getElementById('roles-list'); + if (!rolesList) return; + + if (roles.length === 0) { + rolesList.innerHTML = '
暂无角色
'; + return; + } + + // 更新统计 + const stats = document.getElementById('roles-stats'); + if (stats) { + const totalRoles = roles.length; + const enabledRoles = roles.filter(r => r.enabled !== false).length; + stats.innerHTML = ` +
+ 总角色数 + ${totalRoles} +
+
+ 已启用 + ${enabledRoles} +
+ `; + } + + // 对角色进行排序：默认角色第一个，其他按名称排序 + const sortedRoles = sortRoles(roles); + + rolesList.innerHTML = sortedRoles.map(role => { + // 获取角色图标，如果是Unicode转义格式则转换为emoji + let roleIcon = role.icon || '👤'; + if (roleIcon && typeof roleIcon === 'string') { + // 检查是否是 Unicode 转义格式（可能包含引号） + const unicodeMatch = roleIcon.match(/^"?\\U([0-9A-F]{8})"?$/i); + if (unicodeMatch) { + try { + const codePoint = parseInt(unicodeMatch[1], 16); + roleIcon = String.fromCodePoint(codePoint); + } catch (e) { + // 如果转换失败，使用默认图标 + console.warn('转换 icon Unicode 转义失败:', roleIcon, e); + roleIcon = '👤'; + } + } + } + return ` +
+
+
+
+ ${roleIcon} + ${escapeHtml(role.name)} +
+
+ ${role.enabled !== false ? '已启用' : '已禁用'} +
+
+
${escapeHtml(role.description || '无描述')}
+
+
+ 用户提示词: + ${role.user_prompt ? (role.user_prompt.length > 100 ? escapeHtml(role.user_prompt.substring(0, 100)) + '...' : escapeHtml(role.user_prompt)) : '无'} +
+
+ 关联的工具: + ${ + role.name === '默认' + ? '使用所有工具' + : (role.tools && role.tools.length > 0 + ? `${role.tools.length} 个工具` + : (role.mcps && role.mcps.length > 0 + ? `${role.mcps.length} 个工具（兼容旧版）` + : '使用所有工具')) + } +
+
+
+
+ + ${role.name !== '默认' ? `` : ''} +
+
+ `; + }).join(''); +} + +// 生成工具唯一标识符（与settings.js中的getToolKey保持一致） +function getToolKey(tool) { + // 如果是外部工具，使用 external_mcp::tool.name 作为唯一标识符 + if (tool.is_external && tool.external_mcp) { + return `${tool.external_mcp}::${tool.name}`; + } + // 内置工具直接使用工具名称 + return tool.name; +} + +// 保存当前页的工具状态到全局映射 +function saveCurrentRolePageToolStates() { + document.querySelectorAll('#role-tools-list .role-tool-item').forEach(item => { + const toolKey = item.dataset.toolKey; + const checkbox = item.querySelector('input[type="checkbox"]'); + if (toolKey && checkbox) { + const toolName = item.dataset.toolName; + const isExternal = item.dataset.isExternal === 'true'; + const externalMcp = item.dataset.externalMcp || ''; + const existingState = roleToolStateMap.get(toolKey); + roleToolStateMap.set(toolKey, { + enabled: checkbox.checked, + is_external: isExternal, + external_mcp: externalMcp, + name: toolName, + mcpEnabled: existingState ? existingState.mcpEnabled : true // 保留MCP启用状态 + }); + } + }); +} + +// 加载所有工具列表（用于角色工具选择） +async function loadRoleTools(page = 1, searchKeyword = '') { + try { + // 在加载新页面之前，先保存当前页的状态到全局映射 + saveCurrentRolePageToolStates(); + + const pageSize = roleToolsPagination.pageSize; + let url = `/api/config/tools?page=${page}&page_size=${pageSize}`; + if (searchKeyword) { + url += `&search=${encodeURIComponent(searchKeyword)}`; + } + + const response = await apiFetch(url); + if (!response.ok) { + throw new Error('获取工具列表失败'); + } + + const result = await response.json(); + allRoleTools = result.tools || []; + roleToolsPagination = { + page: result.page || page, + pageSize: result.page_size || pageSize, + total: result.total || 0, + totalPages: result.total_pages || 1 + }; + + // 更新已启用的工具总数（从API响应中获取） + if (result.total_enabled !== undefined) { + totalEnabledToolsInMCP = result.total_enabled; + } + + // 初始化工具状态映射（如果工具不在映射中，使用服务器返回的状态） + // 但要注意：如果工具已经在映射中（比如编辑角色时预先设置的选中工具），则保留映射中的状态 + allRoleTools.forEach(tool => { + const toolKey = getToolKey(tool); + if (!roleToolStateMap.has(toolKey)) { + // 工具不在映射中 + let enabled = false; + if (roleUsesAllTools) { + // 如果使用所有工具，且工具在MCP管理中已启用，则标记为选中 + enabled = tool.enabled ? true : false; + } else { + // 如果不使用所有工具，只有工具在角色配置的工具列表中才标记为选中 + enabled = roleConfiguredTools.has(toolKey); + } + roleToolStateMap.set(toolKey, { + enabled: enabled, + is_external: tool.is_external || false, + external_mcp: tool.external_mcp || '', + name: tool.name, + mcpEnabled: tool.enabled // 保存MCP管理中的原始启用状态 + }); + } else { + // 工具已在映射中（可能是预先设置的选中工具或用户手动选择的），保留映射中的状态 + // 但如果使用所有工具，且工具在MCP管理中已启用，确保标记为选中 + const state = roleToolStateMap.get(toolKey); + if (roleUsesAllTools && tool.enabled) { + // 使用所有工具时，确保所有已启用的工具都被选中 + state.enabled = true; + } + // 如果不使用所有工具，保留映射中的状态（不要覆盖，因为状态已经在初始化时正确设置了） + state.is_external = tool.is_external || false; + state.external_mcp = tool.external_mcp || ''; + state.mcpEnabled = tool.enabled; // 更新MCP管理中的原始启用状态 + if (!state.name || state.name === toolKey.split('::').pop()) { + state.name = tool.name; // 更新工具名称 + } + } + }); + + renderRoleToolsList(); + renderRoleToolsPagination(); + updateRoleToolsStats(); + } catch (error) { + console.error('加载工具列表失败:', error); + const toolsList = document.getElementById('role-tools-list'); + if (toolsList) { + toolsList.innerHTML = `
加载工具列表失败: ${escapeHtml(error.message)}
`; + } + } +} + +// 渲染角色工具选择列表 +function renderRoleToolsList() { + const toolsList = document.getElementById('role-tools-list'); + if (!toolsList) return; + + // 清除加载提示和旧内容 + toolsList.innerHTML = ''; + + const listContainer = document.createElement('div'); + listContainer.className = 'role-tools-list-items'; + listContainer.innerHTML = ''; + + if (allRoleTools.length === 0) { + listContainer.innerHTML = '
暂无工具
'; + toolsList.appendChild(listContainer); + return; + } + + allRoleTools.forEach(tool => { + const toolKey = getToolKey(tool); + const toolItem = document.createElement('div'); + toolItem.className = 'role-tool-item'; + toolItem.dataset.toolKey = toolKey; + toolItem.dataset.toolName = tool.name; + toolItem.dataset.isExternal = tool.is_external ? 'true' : 'false'; + toolItem.dataset.externalMcp = tool.external_mcp || ''; + + // 从状态映射获取工具状态 + const toolState = roleToolStateMap.get(toolKey) || { + enabled: tool.enabled, + is_external: tool.is_external || false, + external_mcp: tool.external_mcp || '' + }; + + // 外部工具标签 + let externalBadge = ''; + if (toolState.is_external || tool.is_external) { + const externalMcpName = toolState.external_mcp || tool.external_mcp || ''; + const badgeText = externalMcpName ? `外部 (${escapeHtml(externalMcpName)})` : '外部'; + const badgeTitle = externalMcpName ? `外部MCP工具 - 来源：${escapeHtml(externalMcpName)}` : '外部MCP工具'; + externalBadge = `${badgeText}`; + } + + // 生成唯一的checkbox id + const checkboxId = `role-tool-${escapeHtml(toolKey).replace(/::/g, '--')}`; + + toolItem.innerHTML = ` + +
+
+ ${escapeHtml(tool.name)} + ${externalBadge} +
+
${escapeHtml(tool.description || '无描述')}
+
+ `; + listContainer.appendChild(toolItem); + }); + + toolsList.appendChild(listContainer); +} + +// 渲染工具列表分页控件 +function renderRoleToolsPagination() { + const toolsList = document.getElementById('role-tools-list'); + if (!toolsList) return; + + // 移除旧的分页控件 + const oldPagination = toolsList.querySelector('.role-tools-pagination'); + if (oldPagination) { + oldPagination.remove(); + } + + // 如果只有一页或没有数据，不显示分页 + if (roleToolsPagination.totalPages <= 1) { + return; + } + + const pagination = document.createElement('div'); + pagination.className = 'role-tools-pagination'; + + const { page, totalPages, total } = roleToolsPagination; + const startItem = (page - 1) * roleToolsPagination.pageSize + 1; + const endItem = Math.min(page * roleToolsPagination.pageSize, total); + + pagination.innerHTML = ` +
+ 显示 ${startItem}-${endItem} / 共 ${total} 个工具${roleToolsSearchKeyword ? ` (搜索: "${escapeHtml(roleToolsSearchKeyword)}")` : ''} +
+
+ + + 第 ${page} / ${totalPages} 页 + + +
+ `; + + toolsList.appendChild(pagination); +} + +// 处理工具checkbox状态变化 +function handleRoleToolCheckboxChange(toolKey, enabled) { + const toolItem = document.querySelector(`.role-tool-item[data-tool-key="${toolKey}"]`); + if (toolItem) { + const toolName = toolItem.dataset.toolName; + const isExternal = toolItem.dataset.isExternal === 'true'; + const externalMcp = toolItem.dataset.externalMcp || ''; + const existingState = roleToolStateMap.get(toolKey); + roleToolStateMap.set(toolKey, { + enabled: enabled, + is_external: isExternal, + external_mcp: externalMcp, + name: toolName, + mcpEnabled: existingState ? existingState.mcpEnabled : true // 保留MCP启用状态 + }); + } + updateRoleToolsStats(); +} + +// 全选工具 +function selectAllRoleTools() { + document.querySelectorAll('#role-tools-list input[type="checkbox"]').forEach(checkbox => { + const toolItem = checkbox.closest('.role-tool-item'); + if (toolItem) { + const toolKey = toolItem.dataset.toolKey; + const toolName = toolItem.dataset.toolName; + const isExternal = toolItem.dataset.isExternal === 'true'; + const externalMcp = toolItem.dataset.externalMcp || ''; + if (toolKey) { + const existingState = roleToolStateMap.get(toolKey); + // 只选中在MCP管理中已启用的工具 + const shouldEnable = existingState && existingState.mcpEnabled !== false; + checkbox.checked = shouldEnable; + roleToolStateMap.set(toolKey, { + enabled: shouldEnable, + is_external: isExternal, + external_mcp: externalMcp, + name: toolName, + mcpEnabled: existingState ? existingState.mcpEnabled : true + }); + } + } + }); + updateRoleToolsStats(); +} + +// 全不选工具 +function deselectAllRoleTools() { + document.querySelectorAll('#role-tools-list input[type="checkbox"]').forEach(checkbox => { + checkbox.checked = false; + const toolItem = checkbox.closest('.role-tool-item'); + if (toolItem) { + const toolKey = toolItem.dataset.toolKey; + const toolName = toolItem.dataset.toolName; + const isExternal = toolItem.dataset.isExternal === 'true'; + const externalMcp = toolItem.dataset.externalMcp || ''; + if (toolKey) { + const existingState = roleToolStateMap.get(toolKey); + roleToolStateMap.set(toolKey, { + enabled: false, + is_external: isExternal, + external_mcp: externalMcp, + name: toolName, + mcpEnabled: existingState ? existingState.mcpEnabled : true // 保留MCP启用状态 + }); + } + } + }); + updateRoleToolsStats(); +} + +// 搜索工具 +function searchRoleTools(keyword) { + roleToolsSearchKeyword = keyword; + const clearBtn = document.getElementById('role-tools-search-clear'); + if (clearBtn) { + clearBtn.style.display = keyword ? 'block' : 'none'; + } + loadRoleTools(1, keyword); +} + +// 清除搜索 +function clearRoleToolsSearch() { + document.getElementById('role-tools-search').value = ''; + searchRoleTools(''); +} + +// 更新工具统计信息 +function updateRoleToolsStats() { + const statsEl = document.getElementById('role-tools-stats'); + if (!statsEl) return; + + // 统计当前页已选中的工具数 + const currentPageEnabled = Array.from(document.querySelectorAll('#role-tools-list input[type="checkbox"]:checked')).length; + + // 统计当前页已启用的工具数（在MCP管理中已启用的工具） + // 优先从状态映射中获取，如果没有则从工具数据中获取 + let currentPageEnabledInMCP = 0; + allRoleTools.forEach(tool => { + const toolKey = getToolKey(tool); + const state = roleToolStateMap.get(toolKey); + // 如果工具在MCP管理中已启用（从状态映射或工具数据中获取），计入当前页已启用工具数 + const mcpEnabled = state ? (state.mcpEnabled !== false) : (tool.enabled !== false); + if (mcpEnabled) { + currentPageEnabledInMCP++; + } + }); + + // 如果使用所有工具，使用从API获取的已启用工具总数 + if (roleUsesAllTools) { + // 使用从API响应中获取的已启用工具总数 + const totalEnabled = totalEnabledToolsInMCP || 0; + // 当前页分母应该是当前页已启用的工具数，而不是所有工具数 + const currentPageDenominator = currentPageEnabledInMCP > 0 ? currentPageEnabledInMCP : document.querySelectorAll('#role-tools-list input[type="checkbox"]').length; + statsEl.innerHTML = ` + ✅ 当前页已选中: ${currentPageEnabled} / ${currentPageDenominator} + 📊 总计已选中: ${totalEnabled} / ${totalEnabled} (使用所有已启用工具) + `; + return; + } + + // 统计角色实际选中的工具数（只统计在MCP管理中已启用的工具） + let totalSelected = 0; + roleToolStateMap.forEach(state => { + // 只统计在MCP管理中已启用且被角色选中的工具 + if (state.enabled && state.mcpEnabled !== false) { + totalSelected++; + } + }); + + // 如果当前页有未保存的状态，需要合并计算 + document.querySelectorAll('#role-tools-list input[type="checkbox"]').forEach(checkbox => { + const toolItem = checkbox.closest('.role-tool-item'); + if (toolItem) { + const toolKey = toolItem.dataset.toolKey; + const savedState = roleToolStateMap.get(toolKey); + if (savedState && savedState.enabled !== checkbox.checked && savedState.mcpEnabled !== false) { + // 状态不一致，使用checkbox状态（但只统计MCP管理中已启用的工具） + if (checkbox.checked && !savedState.enabled) { + totalSelected++; + } else if (!checkbox.checked && savedState.enabled) { + totalSelected--; + } + } + } + }); + + // 角色可选择的所有已启用工具总数（应该基于MCP管理中的总数，而不是状态映射） + // 因为角色可以选择任意已启用的工具，所以总数应该是所有已启用工具的总数 + let totalEnabledForRole = totalEnabledToolsInMCP || 0; + + // 如果API返回的总数为0或未设置，尝试从状态映射中统计（作为备选方案） + if (totalEnabledForRole === 0) { + roleToolStateMap.forEach(state => { + // 只统计在MCP管理中已启用的工具 + if (state.mcpEnabled !== false) { // mcpEnabled 为 true 或 undefined（未设置时默认为启用） + totalEnabledForRole++; + } + }); + } + + // 当前页分母应该是当前页已启用的工具数，而不是所有工具数 + const currentPageDenominator = currentPageEnabledInMCP > 0 ? currentPageEnabledInMCP : document.querySelectorAll('#role-tools-list input[type="checkbox"]').length; + + statsEl.innerHTML = ` + ✅ 当前页已选中: ${currentPageEnabled} / ${currentPageDenominator} + 📊 总计已选中: ${totalSelected} / ${totalEnabledForRole} + `; +} + +// 获取选中的工具列表（返回toolKey数组） +async function getSelectedRoleTools() { + // 先保存当前页的状态 + saveCurrentRolePageToolStates(); + + // 如果没有搜索关键词，需要加载所有页面的工具来确保状态映射完整 + // 但为了性能，我们可以只从状态映射中获取已选中的工具 + // 问题是：如果用户只在某些页面选择了工具，其他页面的工具状态可能不在映射中 + + // 如果总工具数大于已加载的工具数，我们需要确保所有未加载页面的工具也被考虑 + // 但对于角色工具选择，我们只需要获取用户明确选择过的工具 + // 所以直接从状态映射获取已选中的工具即可 + + // 从状态映射获取所有选中的工具（只返回在MCP管理中已启用的工具） + const selectedTools = []; + roleToolStateMap.forEach((state, toolKey) => { + // 只返回在MCP管理中已启用且被角色选中的工具 + if (state.enabled && state.mcpEnabled !== false) { + selectedTools.push(toolKey); + } + }); + + // 如果用户可能在其他页面选择了工具，我们需要确保当前页的状态也被保存 + // 但状态映射应该已经包含了所有访问过的页面的状态 + + return selectedTools; +} + +// 设置选中的工具（用于编辑角色时） +function setSelectedRoleTools(selectedToolKeys) { + const selectedSet = new Set(selectedToolKeys || []); + + // 更新状态映射 + roleToolStateMap.forEach((state, toolKey) => { + state.enabled = selectedSet.has(toolKey); + }); + + // 更新当前页的checkbox状态 + document.querySelectorAll('#role-tools-list .role-tool-item').forEach(item => { + const toolKey = item.dataset.toolKey; + const checkbox = item.querySelector('input[type="checkbox"]'); + if (toolKey && checkbox) { + checkbox.checked = selectedSet.has(toolKey); + } + }); + + updateRoleToolsStats(); +} + +// 显示添加角色模态框 +async function showAddRoleModal() { + const modal = document.getElementById('role-modal'); + if (!modal) return; + + document.getElementById('role-modal-title').textContent = '添加角色'; + document.getElementById('role-name').value = ''; + document.getElementById('role-name').disabled = false; + document.getElementById('role-description').value = ''; + document.getElementById('role-icon').value = ''; + document.getElementById('role-user-prompt').value = ''; + document.getElementById('role-enabled').checked = true; + + // 添加角色时：显示工具选择界面，隐藏默认角色提示 + const toolsSection = document.getElementById('role-tools-section'); + const defaultHint = document.getElementById('role-tools-default-hint'); + const toolsControls = document.querySelector('.role-tools-controls'); + const toolsList = document.getElementById('role-tools-list'); + const formHint = toolsSection ? toolsSection.querySelector('.form-hint') : null; + + if (defaultHint) { + defaultHint.style.display = 'none'; + } + if (toolsControls) { + toolsControls.style.display = 'block'; + } + if (toolsList) { + toolsList.style.display = 'block'; + } + if (formHint) { + formHint.style.display = 'block'; + } + + // 重置工具状态 + roleToolStateMap.clear(); + roleConfiguredTools.clear(); // 清空角色配置的工具列表 + roleUsesAllTools = false; // 添加角色时默认不使用所有工具 + roleToolsSearchKeyword = ''; + const searchInput = document.getElementById('role-tools-search'); + if (searchInput) { + searchInput.value = ''; + } + const clearBtn = document.getElementById('role-tools-search-clear'); + if (clearBtn) { + clearBtn.style.display = 'none'; + } + + // 加载并渲染工具列表 + await loadRoleTools(1, ''); + + // 确保工具列表显示 + if (toolsList) { + toolsList.style.display = 'block'; + } + + modal.style.display = 'flex'; +} + +// 编辑角色 +async function editRole(roleName) { + const role = roles.find(r => r.name === roleName); + if (!role) { + showNotification('角色不存在', 'error'); + return; + } + + const modal = document.getElementById('role-modal'); + if (!modal) return; + + document.getElementById('role-modal-title').textContent = '编辑角色'; + document.getElementById('role-name').value = role.name; + document.getElementById('role-name').disabled = true; // 编辑时不允许修改名称 + document.getElementById('role-description').value = role.description || ''; + // 处理icon字段：如果是Unicode转义格式，转换为emoji；否则直接使用 + let iconValue = role.icon || ''; + if (iconValue && iconValue.startsWith('\\U')) { + // 转换Unicode转义格式（如 \U0001F3C6）为emoji + try { + const codePoint = parseInt(iconValue.substring(2), 16); + iconValue = String.fromCodePoint(codePoint); + } catch (e) { + // 如果转换失败，使用原值 + } + } + document.getElementById('role-icon').value = iconValue; + document.getElementById('role-user-prompt').value = role.user_prompt || ''; + document.getElementById('role-enabled').checked = role.enabled !== false; + + // 检查是否为默认角色 + const isDefaultRole = roleName === '默认'; + const toolsSection = document.getElementById('role-tools-section'); + const defaultHint = document.getElementById('role-tools-default-hint'); + const toolsControls = document.querySelector('.role-tools-controls'); + const toolsList = document.getElementById('role-tools-list'); + const formHint = toolsSection ? toolsSection.querySelector('.form-hint') : null; + + if (isDefaultRole) { + // 默认角色：隐藏工具选择界面，显示提示信息 + if (defaultHint) { + defaultHint.style.display = 'block'; + } + if (toolsControls) { + toolsControls.style.display = 'none'; + } + if (toolsList) { + toolsList.style.display = 'none'; + } + if (formHint) { + formHint.style.display = 'none'; + } + } else { + // 非默认角色：显示工具选择界面，隐藏提示信息 + if (defaultHint) { + defaultHint.style.display = 'none'; + } + if (toolsControls) { + toolsControls.style.display = 'block'; + } + if (toolsList) { + toolsList.style.display = 'block'; + } + if (formHint) { + formHint.style.display = 'block'; + } + + // 重置工具状态 + roleToolStateMap.clear(); + roleConfiguredTools.clear(); // 清空角色配置的工具列表 + roleToolsSearchKeyword = ''; + const searchInput = document.getElementById('role-tools-search'); + if (searchInput) { + searchInput.value = ''; + } + const clearBtn = document.getElementById('role-tools-search-clear'); + if (clearBtn) { + clearBtn.style.display = 'none'; + } + + // 优先使用tools字段，如果没有则使用mcps字段（向后兼容） + const selectedTools = role.tools || (role.mcps && role.mcps.length > 0 ? role.mcps : []); + + // 判断是否使用所有工具：如果没有配置tools（或tools为空数组），表示使用所有工具 + roleUsesAllTools = !role.tools || role.tools.length === 0; + + // 保存角色配置的工具列表 + if (selectedTools.length > 0) { + selectedTools.forEach(toolKey => { + roleConfiguredTools.add(toolKey); + }); + } + + // 如果有选中的工具，先初始化状态映射 + if (selectedTools.length > 0) { + roleUsesAllTools = false; // 有配置工具，不使用所有工具 + // 将选中的工具添加到状态映射（标记为选中） + selectedTools.forEach(toolKey => { + // 如果映射中还没有这个工具，先创建一个默认状态（enabled为true） + if (!roleToolStateMap.has(toolKey)) { + roleToolStateMap.set(toolKey, { + enabled: true, + is_external: false, + external_mcp: '', + name: toolKey.split('::').pop() || toolKey // 从toolKey中提取工具名称 + }); + } else { + // 如果已存在，更新为选中状态 + const state = roleToolStateMap.get(toolKey); + state.enabled = true; + } + }); + } + + // 加载工具列表（第一页） + await loadRoleTools(1, ''); + + // 如果使用所有工具，标记当前页所有已启用的工具为选中 + if (roleUsesAllTools) { + // 标记当前页所有在MCP管理中已启用的工具为选中 + document.querySelectorAll('#role-tools-list input[type="checkbox"]').forEach(checkbox => { + const toolItem = checkbox.closest('.role-tool-item'); + if (toolItem) { + const toolKey = toolItem.dataset.toolKey; + const toolName = toolItem.dataset.toolName; + const isExternal = toolItem.dataset.isExternal === 'true'; + const externalMcp = toolItem.dataset.externalMcp || ''; + if (toolKey) { + const state = roleToolStateMap.get(toolKey); + // 只选中在MCP管理中已启用的工具 + const shouldEnable = state && state.mcpEnabled !== false; + checkbox.checked = shouldEnable; + if (state) { + state.enabled = shouldEnable; + } else { + roleToolStateMap.set(toolKey, { + enabled: shouldEnable, + is_external: isExternal, + external_mcp: externalMcp, + name: toolName, + mcpEnabled: true // 假设已启用，实际值会在loadRoleTools中更新 + }); + } + } + } + }); + } else if (selectedTools.length > 0) { + // 加载完成后，再次设置选中状态（确保当前页的工具也被正确设置） + setSelectedRoleTools(selectedTools); + } + } + + modal.style.display = 'flex'; +} + +// 关闭角色模态框 +function closeRoleModal() { + const modal = document.getElementById('role-modal'); + if (modal) { + modal.style.display = 'none'; + } +} + +// 获取所有选中的工具（包括未在MCP管理中启用的工具） +function getAllSelectedRoleTools() { + // 先保存当前页的状态 + saveCurrentRolePageToolStates(); + + // 从状态映射获取所有选中的工具（不管是否在MCP管理中启用） + const selectedTools = []; + roleToolStateMap.forEach((state, toolKey) => { + if (state.enabled) { + selectedTools.push({ + key: toolKey, + name: state.name || toolKey.split('::').pop() || toolKey, + mcpEnabled: state.mcpEnabled !== false // mcpEnabled 为 false 时是未启用，其他情况视为已启用 + }); + } + }); + + return selectedTools; +} + +// 检查并获取未在MCP管理中启用的工具 +function getDisabledTools(selectedTools) { + return selectedTools.filter(tool => { + const state = roleToolStateMap.get(tool.key); + // 如果 mcpEnabled 明确为 false，则认为是未启用 + return state && state.mcpEnabled === false; + }); +} + +// 保存角色 +async function saveRole() { + const name = document.getElementById('role-name').value.trim(); + if (!name) { + showNotification('角色名称不能为空', 'error'); + return; + } + + const description = document.getElementById('role-description').value.trim(); + let icon = document.getElementById('role-icon').value.trim(); + // 将emoji转换为Unicode转义格式以匹配YAML格式（如 \U0001F3C6） + if (icon) { + // 获取第一个字符的Unicode代码点（处理emoji可能是多个字符的情况） + const codePoint = icon.codePointAt(0); + if (codePoint && codePoint > 0x7F) { + // 转换为8位十六进制格式（\U0001F3C6） + icon = '\\U' + codePoint.toString(16).toUpperCase().padStart(8, '0'); + } + } + const userPrompt = document.getElementById('role-user-prompt').value.trim(); + const enabled = document.getElementById('role-enabled').checked; + + // 检查是否为默认角色 + const isDefaultRole = name === '默认'; + + // 默认角色不保存tools字段（使用所有工具） + // 非默认角色：如果使用所有工具（roleUsesAllTools为true），也不保存tools字段 + let tools = []; + let disabledTools = []; // 存储未在MCP管理中启用的工具 + + if (!isDefaultRole && !roleUsesAllTools) { + // 保存当前页的状态 + saveCurrentRolePageToolStates(); + + // 收集所有选中的工具（包括未在MCP管理中启用的） + const allSelectedTools = getAllSelectedRoleTools(); + + // 检查哪些工具未在MCP管理中启用 + disabledTools = getDisabledTools(allSelectedTools); + + // 如果有未启用的工具，提示用户 + if (disabledTools.length > 0) { + const toolNames = disabledTools.map(t => t.name).join('、'); + const message = `以下 ${disabledTools.length} 个工具未在MCP管理中启用，无法在角色中配置：\n\n${toolNames}\n\n请先在"MCP管理"中启用这些工具，然后再在角色中配置。\n\n是否继续保存？（将只保存已启用的工具）`; + + if (!confirm(message)) { + return; // 用户取消保存 + } + } + + // 获取选中的工具列表（只包含在MCP管理中已启用的工具） + tools = await getSelectedRoleTools(); + } + + const roleData = { + name: name, + description: description, + icon: icon || undefined, // 如果为空字符串，则不发送该字段 + user_prompt: userPrompt, + tools: tools, // 默认角色为空数组，表示使用所有工具 + enabled: enabled + }; + + const isEdit = document.getElementById('role-name').disabled; + const url = isEdit ? `/api/roles/${encodeURIComponent(name)}` : '/api/roles'; + const method = isEdit ? 'PUT' : 'POST'; + + try { + const response = await apiFetch(url, { + method: method, + headers: { + 'Content-Type': 'application/json' + }, + body: JSON.stringify(roleData) + }); + + if (!response.ok) { + const error = await response.json(); + throw new Error(error.error || '保存角色失败'); + } + + // 如果有未启用的工具被过滤掉了，提示用户 + if (disabledTools.length > 0) { + let toolNames = disabledTools.map(t => t.name).join('、'); + // 如果工具名称列表太长，截断显示 + if (toolNames.length > 100) { + toolNames = toolNames.substring(0, 100) + '...'; + } + showNotification( + `${isEdit ? '角色已更新' : '角色已创建'}，但已过滤 ${disabledTools.length} 个未在MCP管理中启用的工具。请先在"MCP管理"中启用这些工具，然后再在角色中配置。`, + 'warning' + ); + } else { + showNotification(isEdit ? '角色已更新' : '角色已创建', 'success'); + } + + closeRoleModal(); + await refreshRoles(); + } catch (error) { + console.error('保存角色失败:', error); + showNotification('保存角色失败: ' + error.message, 'error'); + } +} + +// 删除角色 +async function deleteRole(roleName) { + if (roleName === '默认') { + showNotification('不能删除默认角色', 'error'); + return; + } + + if (!confirm(`确定要删除角色"${roleName}"吗？此操作不可撤销。`)) { + return; + } + + try { + const response = await apiFetch(`/api/roles/${encodeURIComponent(roleName)}`, { + method: 'DELETE' + }); + + if (!response.ok) { + const error = await response.json(); + throw new Error(error.error || '删除角色失败'); + } + + showNotification('角色已删除', 'success'); + + // 如果删除的是当前选中的角色,切换到默认角色 + if (currentRole === roleName) { + handleRoleChange(''); + } + + await refreshRoles(); + } catch (error) { + console.error('删除角色失败:', error); + showNotification('删除角色失败: ' + error.message, 'error'); + } +} + +// 在页面切换时初始化角色列表 +if (typeof switchPage === 'function') { + const originalSwitchPage = switchPage; + switchPage = function(page) { + originalSwitchPage(page); + if (page === 'roles-management') { + loadRoles().then(() => renderRolesList()); + } + }; +} + +// 点击模态框外部关闭 +document.addEventListener('click', (e) => { + const roleSelectModal = document.getElementById('role-select-modal'); + if (roleSelectModal && e.target === roleSelectModal) { + closeRoleSelectModal(); + } + + const roleModal = document.getElementById('role-modal'); + if (roleModal && e.target === roleModal) { + closeRoleModal(); + } + + // 点击角色选择面板外部关闭面板（但不包括角色选择按钮和面板本身） + const roleSelectionPanel = document.getElementById('role-selection-panel'); + const roleSelectorWrapper = document.querySelector('.role-selector-wrapper'); + if (roleSelectionPanel && roleSelectionPanel.style.display !== 'none' && roleSelectionPanel.style.display) { + // 检查点击是否在面板或包装器上 + if (!roleSelectorWrapper?.contains(e.target)) { + closeRoleSelectionPanel(); + } + } +}); + +// 页面加载时初始化 +document.addEventListener('DOMContentLoaded', () => { + loadRoles(); + updateRoleSelectorDisplay(); +}); + +// 获取当前选中的角色（供chat.js使用） +function getCurrentRole() { + return currentRole || ''; +} + +// 暴露函数到全局作用域 +if (typeof window !== 'undefined') { + window.getCurrentRole = getCurrentRole; + window.toggleRoleSelectionPanel = toggleRoleSelectionPanel; + window.closeRoleSelectionPanel = closeRoleSelectionPanel; + window.currentSelectedRole = getCurrentRole(); + + // 监听角色变化，更新全局变量 + const originalHandleRoleChange = handleRoleChange; + handleRoleChange = function(roleName) { + originalHandleRoleChange(roleName); + if (typeof window !== 'undefined') { + window.currentSelectedRole = getCurrentRole(); + } + }; +} + diff --git a/web/static/js/router.js b/web/static/js/router.js index fc810203..b14eba4f 100644 --- a/web/static/js/router.js +++ b/web/static/js/router.js @@ -8,7 +8,7 @@ function initRouter() { if (hash) { const hashParts = hash.split('?'); const pageId = hashParts[0]; - if (pageId && ['chat', 'vulnerabilities', 'mcp-monitor', 'mcp-management', 'knowledge-management', 'knowledge-retrieval-logs', 'settings', 'tasks'].includes(pageId)) { + if (pageId && ['chat', 'vulnerabilities', 'mcp-monitor', 'mcp-management', 'knowledge-management', 'knowledge-retrieval-logs', 'roles-management', 'settings', 'tasks'].includes(pageId)) { switchPage(pageId); // 如果是chat页面且带有conversation参数，加载对应对话 @@ -94,6 +94,19 @@ function updateNavState(pageId) { knowledgeItem.classList.add('expanded'); } + const submenuItem = document.querySelector(`.nav-submenu-item[data-page="${pageId}"]`); + if (submenuItem) { + submenuItem.classList.add('active'); + } + } else if (pageId === 'roles-management') { + // 角色子菜单项 + const rolesItem = document.querySelector('.nav-item[data-page="roles"]'); + if (rolesItem) { + rolesItem.classList.add('active'); + // 展开角色子菜单 + rolesItem.classList.add('expanded'); + } + const submenuItem = document.querySelector(`.nav-submenu-item[data-page="${pageId}"]`); if (submenuItem) { submenuItem.classList.add('active'); @@ -239,6 +252,16 @@ function initPage(pageId) { loadConfig(false); } break; + case 'roles-management': + // 初始化角色管理页面 + if (typeof loadRoles === 'function') { + loadRoles().then(() => { + if (typeof renderRolesList === 'function') { + renderRolesList(); + } + }); + } + break; } // 清理其他页面的定时器 diff --git a/web/templates/index.html b/web/templates/index.html index 4b9b8723..2db10660 100644 --- a/web/templates/index.html +++ b/web/templates/index.html @@ -136,6 +136,23 @@ +
+
+ + + + + 角色 + + + +
+
+
+ 角色管理 +
+
+
@@ -162,7 +179,7 @@
-
+
@@ -269,11 +286,37 @@
+
+ + + +
- +
@@ -605,6 +648,34 @@
+ +
+
+

角色管理

+
+ + +
+
+
+
+
+
+ 总角色数 + - +
+
+ 已启用 + - +
+
+
+
+ +
+
+
+
@@ -687,10 +758,6 @@
-
- - -
@@ -701,6 +768,10 @@ 留空则使用OpenAI配置的api_key
+
+ + +
检索配置
@@ -1329,6 +1400,96 @@
+ + + + + + + @@ -1337,6 +1498,7 @@ +