diff --git a/internal/handler/agent.go b/internal/handler/agent.go new file mode 100644 index 00000000..350aee0a --- /dev/null +++ b/internal/handler/agent.go @@ -0,0 +1,2611 @@ +package handler + +import ( + "context" + "crypto/rand" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "net/http" + "os" + "path/filepath" + "strconv" + "strings" + "sync" + "time" + "unicode/utf8" + + "cyberstrike-ai/internal/agent" + "cyberstrike-ai/internal/audit" + "cyberstrike-ai/internal/authctx" + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/mcp/builtin" + "cyberstrike-ai/internal/multiagent" + "cyberstrike-ai/internal/openai" + "cyberstrike-ai/internal/reasoning" + "cyberstrike-ai/internal/security" + + "github.com/gin-gonic/gin" + "github.com/robfig/cron/v3" + "go.uber.org/zap" +) + +// safeTruncateString 安全截断字符串，避免在 UTF-8 字符中间截断 +func safeTruncateString(s string, maxLen int) string { + if maxLen <= 0 { + return "" + } + if utf8.RuneCountInString(s) <= maxLen { + return s + } + + // 将字符串转换为 rune 切片以正确计算字符数 + runes := []rune(s) + if len(runes) <= maxLen { + return s + } + + // 截断到最大长度 + truncated := string(runes[:maxLen]) + + // 尝试在标点符号或空格处截断，使截断更自然 + // 在截断点往前查找合适的断点（不超过20%的长度） + searchRange := maxLen / 5 + if searchRange > maxLen { + searchRange = maxLen + } + breakChars := []rune("，。、 ,.;:!?！？/\\-_") + bestBreakPos := len(runes[:maxLen]) + + for i := bestBreakPos - 1; i >= bestBreakPos-searchRange && i >= 0; i-- { + for _, breakChar := range breakChars { + if runes[i] == breakChar { + bestBreakPos = i + 1 // 在标点符号后断开 + goto found + } + } + } + +found: + truncated = string(runes[:bestBreakPos]) + return truncated + "..." +} + +// responsePlanAgg buffers main-assistant response_stream chunks for one "planning" process_detail row. +type responsePlanAgg struct { + meta map[string]interface{} + b strings.Builder + detailID string + lastPersistAt time.Time + lastPersistSize int +} + +// thinkingBuf aggregates thinking_stream_* / reasoning_chain_stream_* before flush to process_details. +type thinkingBuf struct { + b strings.Builder + meta map[string]interface{} + persistAs string // "thinking" | "reasoning_chain" +} + +func normalizeProcessDetailText(s string) string { + s = strings.ReplaceAll(s, "\r\n", "\n") + s = strings.ReplaceAll(s, "\r", "\n") + return strings.TrimSpace(s) +} + +// discardPlanningIfEchoesToolResult drops buffered planning text when it only repeats the +// upcoming tool_result body. Streaming models often echo tool stdout in chunk.Content; flushing +// that into "planning" before persisting tool_result duplicates the output after page refresh. +// sameResponseStreamMeta 判断是否为同一段主通道流（Eino ADK 可能对同一 MessageStream 重复发 response_start）。 +func sameResponseStreamMeta(a, b map[string]interface{}) bool { + if a == nil || b == nil { + return false + } + agentA, _ := a["einoAgent"].(string) + agentB, _ := b["einoAgent"].(string) + agentA = strings.TrimSpace(agentA) + agentB = strings.TrimSpace(agentB) + if agentA == "" || !strings.EqualFold(agentA, agentB) { + return false + } + orchA, _ := a["orchestration"].(string) + orchB, _ := b["orchestration"].(string) + if strings.TrimSpace(orchA) != strings.TrimSpace(orchB) { + return false + } + iterA := responseStreamIterationFromMeta(a) + iterB := responseStreamIterationFromMeta(b) + if iterA != 0 && iterB != 0 && iterA != iterB { + return false + } + streamA, _ := a["streamId"].(string) + streamB, _ := b["streamId"].(string) + streamA = strings.TrimSpace(streamA) + streamB = strings.TrimSpace(streamB) + if streamA != "" && streamB != "" && streamA != streamB { + return false + } + return true +} + +func responseStreamIterationFromMeta(m map[string]interface{}) int { + if m == nil { + return 0 + } + switch v := m["iteration"].(type) { + case int: + return v + case int32: + return int(v) + case int64: + return int(v) + case float64: + return int(v) + default: + return 0 + } +} + +func discardPlanningIfEchoesToolResult(respPlan *responsePlanAgg, toolData interface{}) string { + if respPlan == nil { + return "" + } + plan := normalizeProcessDetailText(respPlan.b.String()) + if plan == "" { + return "" + } + dataMap, ok := toolData.(map[string]interface{}) + if !ok { + return "" + } + res, ok := dataMap["result"].(string) + if !ok { + return "" + } + r := normalizeProcessDetailText(res) + if r == "" { + return "" + } + if plan == r || strings.HasSuffix(plan, r) { + detailID := respPlan.detailID + respPlan.meta = nil + respPlan.b.Reset() + respPlan.detailID = "" + respPlan.lastPersistAt = time.Time{} + respPlan.lastPersistSize = 0 + return detailID + } + return "" +} + +// AgentHandler Agent处理器 +type AgentHandler struct { + agent *agent.Agent + db *database.DB + logger *zap.Logger + tasks *AgentTaskManager + taskEventBus *TaskEventBus // 镜像 SSE 事件，供刷新后订阅同一运行中任务 + batchTaskManager *BatchTaskManager + hitlManager *HITLManager + config *config.Config // 配置引用，用于获取角色信息 + knowledgeManager interface { // 知识库管理器接口 + LogRetrieval(conversationID, messageID, query, riskType string, retrievedItems []string) error + } + agentsMarkdownDir string // 多代理：Markdown 子 Agent 目录（绝对路径，空则不从磁盘合并） + batchCronParser cron.Parser + // hitlWhitelistSaver 侧栏「应用」HITL 时将会话增量白名单合并写入 config.yaml（可选） + hitlWhitelistSaver HitlToolWhitelistSaver + hitlStrategySaver HitlAuditStrategySaver + hitlDefaultReviewerSaver HitlDefaultReviewerSaver + auditLLM *openai.Client + audit *audit.Service +} + +// SetAudit wires platform audit logging. +func (h *AgentHandler) SetAudit(s *audit.Service) { + h.audit = s +} + +// TaskManager 返回 Agent 任务管理器（供 MCP 监控页终止 Eino execute 等）。 +func (h *AgentHandler) TaskManager() *AgentTaskManager { + if h == nil { + return nil + } + return h.tasks +} + +// CancelRunningTaskForConversation stops any in-flight agent work for the conversation (idempotent). +func (h *AgentHandler) CancelRunningTaskForConversation(conversationID string) { + if h == nil || conversationID == "" || h.tasks == nil { + return + } + h.cancelRunningMCPToolsForConversation(conversationID) + h.tasks.AbortActiveEinoExecute(conversationID, "") + if ok, err := h.tasks.CancelTask(conversationID, ErrTaskCancelled); ok { + h.logger.Info("已取消会话运行中任务", zap.String("conversationId", conversationID)) + } else if err != nil { + h.logger.Warn("取消会话运行中任务失败", zap.String("conversationId", conversationID), zap.Error(err)) + } +} + +// ConversationTaskRuntimeState exposes the authoritative live state and start +// time used to scope persisted TaskCreate files to the current run. A task +// already entering cancellation must stop driving progress UI immediately. +func (h *AgentHandler) ConversationTaskRuntimeState(conversationID string) (bool, time.Time) { + if h == nil || h.tasks == nil || strings.TrimSpace(conversationID) == "" { + return false, time.Time{} + } + task := h.tasks.GetTaskSnapshot(strings.TrimSpace(conversationID)) + if task == nil || !strings.EqualFold(strings.TrimSpace(task.Status), "running") { + return false, time.Time{} + } + return true, task.StartedAt +} + +func (h *AgentHandler) cancelRunningMCPToolsForConversation(conversationID string) { + if h == nil || h.agent == nil { + return + } + n := h.agent.CancelRunningMCPToolsForConversation(conversationID, "会话已结束，自动终止仍在运行的工具") + if n > 0 && h.logger != nil { + h.logger.Info("已终止会话仍在运行的 MCP 工具", zap.String("conversationId", conversationID), zap.Int("count", n)) + } +} + +// HitlToolWhitelistSaver 合并/设置 HITL 免审批工具到全局配置并落盘 +type HitlToolWhitelistSaver interface { + MergeHitlToolWhitelistIntoConfig(add []string) error + SetHitlToolWhitelist(tools []string) error +} + +// NewAgentHandler 创建新的Agent处理器 +func NewAgentHandler(agent *agent.Agent, db *database.DB, cfg *config.Config, logger *zap.Logger) *AgentHandler { + batchTaskManager := NewBatchTaskManager(logger) + batchTaskManager.SetDB(db) + + // 从数据库加载所有批量任务队列 + if err := batchTaskManager.LoadFromDB(); err != nil { + logger.Warn("从数据库加载批量任务队列失败", zap.Error(err)) + } + + bus := NewTaskEventBus() + tm := NewAgentTaskManager() + tm.SetTaskEventBus(bus) + llmHTTP := &http.Client{Timeout: 2 * time.Minute} + var llmCfg *config.OpenAIConfig + if cfg != nil { + llmCfg = &cfg.OpenAI + } + handler := &AgentHandler{ + agent: agent, + db: db, + logger: logger, + tasks: tm, + taskEventBus: bus, + batchTaskManager: batchTaskManager, + config: cfg, + hitlManager: NewHITLManager(db, logger), + batchCronParser: cron.NewParser(cron.Minute | cron.Hour | cron.Dom | cron.Month | cron.Dow | cron.Descriptor), + auditLLM: openai.NewClient(llmCfg, llmHTTP, logger), + } + tm.SetToolCanceler(handler.cancelRunningMCPToolsForConversation) + if err := handler.hitlManager.EnsureSchema(); err != nil { + logger.Warn("初始化 HITL 表失败", zap.Error(err)) + } + go handler.batchQueueSchedulerLoop() + return handler +} + +// SetKnowledgeManager 设置知识库管理器（用于记录检索日志） +func (h *AgentHandler) SetKnowledgeManager(manager interface { + LogRetrieval(conversationID, messageID, query, riskType string, retrievedItems []string) error +}) { + h.knowledgeManager = manager +} + +// SetAgentsMarkdownDir 设置 agents/*.md 子代理目录（绝对路径）；空表示仅使用 config.yaml 中的 sub_agents。 +func (h *AgentHandler) SetAgentsMarkdownDir(absDir string) { + h.agentsMarkdownDir = strings.TrimSpace(absDir) +} + +// SetHitlToolWhitelistSaver 设置 HITL 白名单落盘（与 ConfigHandler 配合，避免循环引用用接口） +func (h *AgentHandler) SetHitlToolWhitelistSaver(s HitlToolWhitelistSaver) { + h.hitlWhitelistSaver = s +} + +// HitlDefaultReviewerSaver 持久化全局默认审批方到 config.yaml。 +type HitlDefaultReviewerSaver interface { + UpdateHitlDefaultReviewer(reviewer string) error +} + +// SetHitlDefaultReviewerSaver 设置 HITL 默认审批方落盘。 +func (h *AgentHandler) SetHitlDefaultReviewerSaver(s HitlDefaultReviewerSaver) { + h.hitlDefaultReviewerSaver = s +} + +func (h *AgentHandler) hitlEffectiveDefaultReviewer() string { + if h != nil && h.config != nil { + return normalizeHitlReviewer(h.config.Hitl.EffectiveDefaultReviewer()) + } + return "human" +} + +// HITLNeedsToolApproval 供 C2 危险任务门控：与会话侧人机协同及免审批白名单判定一致。 +func (h *AgentHandler) HITLNeedsToolApproval(conversationID, toolName string) bool { + if h == nil || h.hitlManager == nil { + return false + } + return h.hitlManager.NeedsToolApproval(conversationID, toolName) +} + +// ChatAttachment 聊天附件（用户上传的文件） +type ChatAttachment struct { + FileName string `json:"fileName"` // 展示用文件名 + Content string `json:"content,omitempty"` // 文本或 base64；若已预先上传到服务器可留空 + MimeType string `json:"mimeType,omitempty"` + ServerPath string `json:"serverPath,omitempty"` // 已保存在 chat_uploads 下的绝对路径（由 POST /api/chat-uploads 返回） +} + +// ChatReasoningRequest 对话页「模型推理」意图（Eino 单/多代理路径消费）。 +type ChatReasoningRequest struct { + // Mode: default（跟随系统）| off | on | auto + Mode string `json:"mode,omitempty"` + // Effort: low | medium | high | max | xhigh（原样下发；不同网关最高档命名不同）。空表示不指定。 + Effort string `json:"effort,omitempty"` +} + +// ChatFinalizationRequest is a caller-provided delivery policy. The server does +// not infer execution intent from natural-language user text. +type ChatFinalizationRequest struct { + RequireExecutionEvidence *bool `json:"requireExecutionEvidence,omitempty"` +} + +// ChatRequest 聊天请求 +type ChatRequest struct { + Message string `json:"message" binding:"required"` + ConversationID string `json:"conversationId,omitempty"` + ProjectID string `json:"projectId,omitempty"` // 新对话绑定的项目（可选；未指定时可用 config.project.default_project_id） + Role string `json:"role,omitempty"` // 角色名称 + Attachments []ChatAttachment `json:"attachments,omitempty"` + WebShellConnectionID string `json:"webshellConnectionId,omitempty"` // WebShell 管理 - AI 助手：当前选中的连接 ID，仅使用 webshell_* 工具 + AIChannelID string `json:"aiChannelId,omitempty"` // 会话级 AI 通道；空则使用 ai.default_channel + Hitl *HITLRequest `json:"hitl,omitempty"` + Reasoning *ChatReasoningRequest `json:"reasoning,omitempty"` + Finalization ChatFinalizationRequest `json:"finalization,omitempty"` + // Orchestration 仅对 /api/multi-agent、/api/multi-agent/stream：deep | plan_execute | supervisor；空则等同 deep。机器人/批量等无请求体时由服务端默认 deep。/api/eino-agent* 不使用此字段。 + Orchestration string `json:"orchestration,omitempty"` +} + +func (h *AgentHandler) configForAIChannel(channelID string) (*config.Config, string, error) { + if h == nil || h.config == nil { + return nil, "", fmt.Errorf("服务器配置未加载") + } + oa, resolvedID, ok := h.config.ResolveAIChannel(channelID) + if !ok { + return nil, resolvedID, fmt.Errorf("AI 通道不存在: %s", resolvedID) + } + cfgCopy := *h.config + cfgCopy.OpenAI = oa + return &cfgCopy, resolvedID, nil +} + +func chatReasoningToClientIntent(r *ChatReasoningRequest) *reasoning.ClientIntent { + if r == nil { + return nil + } + return &reasoning.ClientIntent{Mode: r.Mode, Effort: r.Effort} +} + +type HITLRequest struct { + Enabled bool `json:"enabled"` + Mode string `json:"mode,omitempty"` + Reviewer string `json:"reviewer,omitempty"` // human | audit_agent + SensitiveTools []string `json:"sensitiveTools,omitempty"` + TimeoutSeconds int `json:"timeoutSeconds,omitempty"` +} + +const ( + maxAttachments = 10 + chatUploadsDirName = "chat_uploads" // 对话附件保存的根目录（相对当前工作目录） +) + +// validateChatAttachmentServerPath 校验绝对路径落在工作目录 chat_uploads 下且为普通文件（防路径穿越） +func validateChatAttachmentServerPath(abs string) (string, error) { + p := strings.TrimSpace(abs) + if p == "" { + return "", fmt.Errorf("empty path") + } + cwd, err := os.Getwd() + if err != nil { + return "", fmt.Errorf("获取当前工作目录失败: %w", err) + } + root := filepath.Join(cwd, chatUploadsDirName) + rootAbs, err := filepath.Abs(filepath.Clean(root)) + if err != nil { + return "", err + } + pathAbs, err := filepath.Abs(filepath.Clean(p)) + if err != nil { + return "", err + } + sep := string(filepath.Separator) + if pathAbs != rootAbs && !strings.HasPrefix(pathAbs, rootAbs+sep) { + return "", fmt.Errorf("path outside chat_uploads") + } + st, err := os.Stat(pathAbs) + if err != nil { + return "", err + } + if st.IsDir() { + return "", fmt.Errorf("not a regular file") + } + return pathAbs, nil +} + +// avoidChatUploadDestCollision 若 path 已存在则生成带时间戳+随机后缀的新文件名（与上传接口命名风格一致） +func avoidChatUploadDestCollision(path string) string { + if _, err := os.Stat(path); os.IsNotExist(err) { + return path + } + dir := filepath.Dir(path) + base := filepath.Base(path) + ext := filepath.Ext(base) + nameNoExt := strings.TrimSuffix(base, ext) + suffix := fmt.Sprintf("_%s_%s", time.Now().Format("150405"), shortRand(6)) + var unique string + if ext != "" { + unique = nameNoExt + suffix + ext + } else { + unique = base + suffix + } + return filepath.Join(dir, unique) +} + +// relocateManualOrNewUploadToConversation 无会话 ID 时前端会上传到 …/日期/_manual；首条消息创建会话后，将文件移入 …/日期/{conversationId}/ 以便按对话隔离。 +func relocateManualOrNewUploadToConversation(absPath, conversationID string, logger *zap.Logger) (string, error) { + conv := strings.TrimSpace(conversationID) + if conv == "" { + return absPath, nil + } + convSan := strings.ReplaceAll(conv, string(filepath.Separator), "_") + if convSan == "" || convSan == "_manual" || convSan == "_new" { + return absPath, nil + } + cwd, err := os.Getwd() + if err != nil { + return absPath, err + } + rootAbs, err := filepath.Abs(filepath.Join(cwd, chatUploadsDirName)) + if err != nil { + return absPath, err + } + rel, err := filepath.Rel(rootAbs, absPath) + if err != nil { + return absPath, nil + } + rel = filepath.ToSlash(filepath.Clean(rel)) + var segs []string + for _, p := range strings.Split(rel, "/") { + if p != "" && p != "." { + segs = append(segs, p) + } + } + // 仅处理扁平结构：日期/_manual|_new/文件名 + if len(segs) != 3 { + return absPath, nil + } + datePart, placeFolder, baseName := segs[0], segs[1], segs[2] + if placeFolder != "_manual" && placeFolder != "_new" { + return absPath, nil + } + targetDir := filepath.Join(rootAbs, datePart, convSan) + if err := os.MkdirAll(targetDir, 0755); err != nil { + return "", fmt.Errorf("创建会话附件目录失败: %w", err) + } + dest := filepath.Join(targetDir, baseName) + dest = avoidChatUploadDestCollision(dest) + if err := os.Rename(absPath, dest); err != nil { + return "", fmt.Errorf("将附件移入会话目录失败: %w", err) + } + out, _ := filepath.Abs(dest) + if logger != nil { + logger.Info("对话附件已从占位目录移入会话目录", + zap.String("from", absPath), + zap.String("to", out), + zap.String("conversationId", conv)) + } + return out, nil +} + +// saveAttachmentsToDateAndConversationDir 处理附件：若带 serverPath 则仅校验已存在文件；否则将 content 写入 chat_uploads/YYYY-MM-DD/{conversationID}/。 +// conversationID 为空时使用 "_new" 作为目录名（新对话尚未有 ID） +func saveAttachmentsToDateAndConversationDir(attachments []ChatAttachment, conversationID string, logger *zap.Logger) (savedPaths []string, err error) { + if len(attachments) == 0 { + return nil, nil + } + cwd, err := os.Getwd() + if err != nil { + return nil, fmt.Errorf("获取当前工作目录失败: %w", err) + } + dateDir := filepath.Join(cwd, chatUploadsDirName, time.Now().Format("2006-01-02")) + convDirName := strings.TrimSpace(conversationID) + if convDirName == "" { + convDirName = "_new" + } else { + convDirName = strings.ReplaceAll(convDirName, string(filepath.Separator), "_") + } + targetDir := filepath.Join(dateDir, convDirName) + if err = os.MkdirAll(targetDir, 0755); err != nil { + return nil, fmt.Errorf("创建上传目录失败: %w", err) + } + savedPaths = make([]string, 0, len(attachments)) + for i, a := range attachments { + if sp := strings.TrimSpace(a.ServerPath); sp != "" { + valid, verr := validateChatAttachmentServerPath(sp) + if verr != nil { + return nil, fmt.Errorf("附件 %s: %w", a.FileName, verr) + } + finalPath, rerr := relocateManualOrNewUploadToConversation(valid, conversationID, logger) + if rerr != nil { + return nil, fmt.Errorf("附件 %s: %w", a.FileName, rerr) + } + savedPaths = append(savedPaths, finalPath) + if logger != nil { + logger.Debug("对话附件使用已上传路径", zap.Int("index", i+1), zap.String("fileName", a.FileName), zap.String("path", finalPath)) + } + continue + } + if strings.TrimSpace(a.Content) == "" { + return nil, fmt.Errorf("附件 %s 缺少内容或未提供 serverPath", a.FileName) + } + raw, decErr := attachmentContentToBytes(a) + if decErr != nil { + return nil, fmt.Errorf("附件 %s 解码失败: %w", a.FileName, decErr) + } + baseName := filepath.Base(a.FileName) + if baseName == "" || baseName == "." { + baseName = "file" + } + baseName = strings.ReplaceAll(baseName, string(filepath.Separator), "_") + ext := filepath.Ext(baseName) + nameNoExt := strings.TrimSuffix(baseName, ext) + suffix := fmt.Sprintf("_%s_%s", time.Now().Format("150405"), shortRand(6)) + var unique string + if ext != "" { + unique = nameNoExt + suffix + ext + } else { + unique = baseName + suffix + } + fullPath := filepath.Join(targetDir, unique) + if err = os.WriteFile(fullPath, raw, 0644); err != nil { + return nil, fmt.Errorf("写入文件 %s 失败: %w", a.FileName, err) + } + absPath, _ := filepath.Abs(fullPath) + savedPaths = append(savedPaths, absPath) + if logger != nil { + logger.Debug("对话附件已保存", zap.Int("index", i+1), zap.String("fileName", a.FileName), zap.String("path", absPath)) + } + } + return savedPaths, nil +} + +func shortRand(n int) string { + const letters = "0123456789abcdef" + b := make([]byte, n) + _, _ = rand.Read(b) + for i := range b { + b[i] = letters[int(b[i])%len(letters)] + } + return string(b) +} + +func attachmentContentToBytes(a ChatAttachment) ([]byte, error) { + content := a.Content + if decoded, err := base64.StdEncoding.DecodeString(content); err == nil && len(decoded) > 0 { + return decoded, nil + } + return []byte(content), nil +} + +// userMessageContentForStorage 返回要存入数据库的用户消息内容：有附件时在正文后追加附件名（及路径），刷新后仍能显示，继续对话时大模型也能从历史中拿到路径 +func userMessageContentForStorage(message string, attachments []ChatAttachment, savedPaths []string) string { + if len(attachments) == 0 { + return message + } + var b strings.Builder + b.WriteString(message) + for i, a := range attachments { + b.WriteString("\n📎 ") + b.WriteString(a.FileName) + if i < len(savedPaths) && savedPaths[i] != "" { + b.WriteString(": ") + b.WriteString(savedPaths[i]) + } + } + return b.String() +} + +// appendAttachmentsToMessage 仅将附件的保存路径追加到用户消息末尾，不再内联附件内容，避免上下文过长 +func appendAttachmentsToMessage(msg string, attachments []ChatAttachment, savedPaths []string) string { + if len(attachments) == 0 { + return msg + } + var b strings.Builder + b.WriteString(msg) + b.WriteString("\n\n[用户上传的文件]\n") + for i, a := range attachments { + if i < len(savedPaths) && savedPaths[i] != "" { + b.WriteString(fmt.Sprintf("- %s: %s\n", a.FileName, savedPaths[i])) + } else { + b.WriteString(fmt.Sprintf("- %s: （路径未知，可能保存失败）\n", a.FileName)) + } + } + return b.String() +} + +// appendAssistantMessageNotice 在助手消息末尾追加提示，避免覆盖已生成内容。 +// 若消息为空则直接写入提示；若已包含相同提示则保持不变。 +func (h *AgentHandler) appendAssistantMessageNotice(messageID, notice string) error { + trimmedNotice := strings.TrimSpace(notice) + if strings.TrimSpace(messageID) == "" || trimmedNotice == "" { + return nil + } + _, err := h.db.Exec( + `UPDATE messages + SET content = CASE + WHEN content IS NULL OR TRIM(content) = '' THEN ? + WHEN INSTR(content, ?) > 0 THEN content + ELSE content || '\n\n' || ? + END, + updated_at = ? + WHERE id = ?`, + trimmedNotice, + trimmedNotice, + trimmedNotice, + time.Now(), + messageID, + ) + return err +} + +// mergeAssistantMessagePartialOnCancel 将取消前已生成的部分回复尽量合并进消息： +// - content 为空或仅占位（处理中...）时，直接替换为 partial； +// - 已有正文时，仅在尚未包含 partial 时追加，避免丢失与重复。 +func (h *AgentHandler) mergeAssistantMessagePartialOnCancel(messageID, partial string) error { + trimmedPartial := strings.TrimSpace(partial) + if strings.TrimSpace(messageID) == "" || trimmedPartial == "" { + return nil + } + _, err := h.db.Exec( + `UPDATE messages + SET content = CASE + WHEN content IS NULL OR TRIM(content) = '' OR TRIM(content) = '处理中...' THEN ? + WHEN INSTR(content, ?) > 0 THEN content + ELSE content || '\n\n' || ? + END, + updated_at = ? + WHERE id = ?`, + trimmedPartial, + trimmedPartial, + trimmedPartial, + time.Now(), + messageID, + ) + return err +} + +// ChatResponse 聊天响应 +type ChatResponse struct { + Response string `json:"response"` + MCPExecutionIDs []string `json:"mcpExecutionIds,omitempty"` // 本次对话中执行的MCP调用ID列表 + ConversationID string `json:"conversationId"` // 对话ID + Time time.Time `json:"time"` + Finalizable bool `json:"finalizable"` + Finalized bool `json:"finalized"` + Status string `json:"status,omitempty"` + CompletionReason string `json:"completionReason,omitempty"` + EvidenceVerified bool `json:"evidenceVerified"` + EvidenceRefs []string `json:"evidenceRefs,omitempty"` + PendingExecutionIDs []string `json:"pendingExecutionIds,omitempty"` + MissingChecks []string `json:"missingChecks,omitempty"` +} + +func (h *AgentHandler) finalizeRobotAgentError(ctx context.Context, assistantMessageID, conversationID string, resultMA *multiagent.RunResult, errMA error) (string, string, error) { + if shouldPersistEinoAgentTraceAfterRunError(ctx) { + h.persistEinoAgentTraceForResume(conversationID, resultMA) + } + errMsg := "执行失败: " + errMA.Error() + if assistantMessageID != "" { + _, _ = h.db.Exec("UPDATE messages SET content = ?, updated_at = ? WHERE id = ?", errMsg, time.Now(), assistantMessageID) + _ = h.db.AddProcessDetail(assistantMessageID, conversationID, "error", errMsg, nil) + } + return "", conversationID, errMA +} + +func (h *AgentHandler) finalizeRobotAgentSuccess(assistantMessageID, conversationID string, resultMA *multiagent.RunResult) (string, string, error) { + decision := h.finalizeAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, "robot", resultMA, resultMA.MCPExecutionIDs, multiagent.AggregatedReasoningFromTraceJSON(resultMA.LastAgentTraceInput), true) + responseText := decision.FinalText + if !decision.Finalizable { + responseText = finalizationBlockedMessage(decision) + } + if assistantMessageID == "" { + if _, err := h.db.AddMessage(conversationID, "assistant", responseText, resultMA.MCPExecutionIDs); err != nil { + h.logger.Warn("机器人：保存助手消息失败", zap.Error(err)) + } + } + if resultMA.LastAgentTraceInput != "" || resultMA.LastAgentTraceOutput != "" { + _ = h.db.SaveAgentTrace(conversationID, resultMA.LastAgentTraceInput, resultMA.LastAgentTraceOutput) + } + return responseText, conversationID, nil +} + +func (h *AgentHandler) runRobotEinoSingleWithRetry( + taskCtx context.Context, + conversationID, finalMessage string, + history []agent.ChatMessage, + roleTools []string, + progressCallback agent.ProgressCallback, + assistantMessageID string, + taskStatus *string, +) (string, string, error) { + resultMA, errMA := multiagent.RunEinoSingleChatModelAgent( + taskCtx, h.config, &h.config.MultiAgent, h.agent, h.db, h.logger, + conversationID, h.conversationProjectID(conversationID), finalMessage, history, roleTools, progressCallback, nil, h.agentSessionContextBlock(conversationID), + ) + if errMA != nil { + *taskStatus = "failed" + return h.finalizeRobotAgentError(taskCtx, assistantMessageID, conversationID, resultMA, errMA) + } + return h.finalizeRobotAgentSuccess(assistantMessageID, conversationID, resultMA) +} + +func (h *AgentHandler) runRobotMultiAgentWithRetry( + taskCtx context.Context, + conversationID, finalMessage, orchestration string, + history []agent.ChatMessage, + roleTools []string, + progressCallback agent.ProgressCallback, + assistantMessageID string, + taskStatus *string, +) (string, string, error) { + resultMA, errMA := multiagent.RunDeepAgent( + taskCtx, h.config, &h.config.MultiAgent, h.agent, h.db, h.logger, + conversationID, h.conversationProjectID(conversationID), finalMessage, history, roleTools, progressCallback, + h.agentsMarkdownDir, orchestration, nil, h.agentSessionContextBlock(conversationID), + ) + if errMA != nil { + *taskStatus = "failed" + return h.finalizeRobotAgentError(taskCtx, assistantMessageID, conversationID, resultMA, errMA) + } + return h.finalizeRobotAgentSuccess(assistantMessageID, conversationID, resultMA) +} + +// ProcessMessageForRobot 供机器人（企业微信/钉钉/飞书）调用：Eino 单/多代理执行路径（含 progressCallback、过程详情），仅不发送 SSE，最后返回完整回复 +func (h *AgentHandler) ProcessMessageForRobot(ctx context.Context, platform string, principal authctx.Principal, conversationID, message, role, agentMode string) (response string, convID string, err error) { + ownerUserID := strings.TrimSpace(principal.UserID) + if ownerUserID == "" { + return "", "", fmt.Errorf("authenticated robot principal is required") + } + if !principal.HasPermission("agent:execute") || !principal.HasPermission("chat:read") || !principal.HasPermission("chat:write") { + return "", "", fmt.Errorf("机器人账号缺少 agent:execute、chat:read 或 chat:write 权限") + } + ctx = authctx.WithPrincipal(ctx, principal) + if conversationID == "" { + title := safeTruncateString(message, 50) + src := "robot" + if strings.TrimSpace(platform) != "" { + src = "robot:" + strings.TrimSpace(platform) + } + meta := audit.ConversationCreateMeta(src) + meta.ProjectID = effectiveProjectID(h.config, "") + if meta.ProjectID != "" && (!principal.HasPermission("project:read") || !h.db.UserCanAccessResource(ownerUserID, principal.ScopeFor("project:read"), "project", meta.ProjectID)) { + meta.ProjectID = "" + } + conv, createErr := h.db.CreateConversation(title, meta) + if createErr != nil { + return "", "", fmt.Errorf("创建对话失败: %w", createErr) + } + conversationID = conv.ID + _ = h.db.SetResourceOwner("conversation", conversationID, ownerUserID) + } else { + if _, getErr := h.db.GetConversation(conversationID); getErr != nil || !h.db.UserCanAccessResource(ownerUserID, principal.ScopeFor("chat:write"), "conversation", conversationID) { + return "", "", fmt.Errorf("对话不存在") + } + } + + agentHistoryMessages, err := h.loadHistoryFromAgentTrace(conversationID) + if err != nil { + historyMessages, getErr := h.db.GetMessages(conversationID) + if getErr != nil { + agentHistoryMessages = []agent.ChatMessage{} + } else { + agentHistoryMessages = make([]agent.ChatMessage, 0, len(historyMessages)) + for _, msg := range historyMessages { + agentHistoryMessages = append(agentHistoryMessages, agent.ChatMessage{Role: msg.Role, Content: msg.Content}) + } + } + } + + finalMessage := message + var roleTools []string + if role != "" && role != "默认" && h.config.Roles != nil { + if r, exists := h.config.Roles[role]; exists && r.Enabled { + if r.UserPrompt != "" { + finalMessage = r.UserPrompt + "\n\n" + message + } + roleTools = r.Tools + } + } + + if _, err = h.db.AddMessage(conversationID, "user", message, nil); err != nil { + return "", "", fmt.Errorf("保存用户消息失败: %w", err) + } + + // 与 Eino 流式对话一致：先创建助手消息占位，用 progressCallback 写过程详情（不发送 SSE） + assistantMsg, err := h.db.AddMessage(conversationID, "assistant", "处理中...", nil) + if err != nil { + h.logger.Warn("机器人：创建助手消息占位失败", zap.Error(err)) + } + var assistantMessageID string + if assistantMsg != nil { + assistantMessageID = assistantMsg.ID + } + + // 注册运行中任务并向 taskEventBus 镜像进度事件，供 Web 端 task-events 补流。 + taskCtx, cancelWithCause := context.WithCancelCause(ctx) + defer cancelWithCause(nil) + taskStatus := "completed" + defer func() { + h.tasks.FinishTask(conversationID, taskStatus) + }() + if _, err := h.tasks.StartTask(conversationID, message, cancelWithCause); err != nil { + if errors.Is(err, ErrTaskAlreadyRunning) { + return "", conversationID, fmt.Errorf("当前会话已有任务正在执行中，请稍后再试") + } + return "", conversationID, fmt.Errorf("无法启动任务: %w", err) + } + progressCallback := h.createProgressCallback(taskCtx, cancelWithCause, conversationID, assistantMessageID, nil) + + robotMode := config.NormalizeAgentMode(agentMode) + if err := h.db.SetConversationAgentMode(conversationID, robotMode); err != nil { + h.logger.Warn("机器人：更新对话模式失败", zap.String("conversationId", conversationID), zap.String("agentMode", robotMode), zap.Error(err)) + } + switch robotMode { + case "eino_single": + return h.runRobotEinoSingleWithRetry(taskCtx, conversationID, finalMessage, agentHistoryMessages, roleTools, progressCallback, assistantMessageID, &taskStatus) + case "deep", "plan_execute", "supervisor": + if h.config == nil || !h.config.MultiAgent.Enabled { + taskStatus = "failed" + return "", conversationID, fmt.Errorf("机器人对话模式 %s 需要启用 Eino 多代理", robotMode) + } + return h.runRobotMultiAgentWithRetry(taskCtx, conversationID, finalMessage, robotMode, agentHistoryMessages, roleTools, progressCallback, assistantMessageID, &taskStatus) + } + + taskStatus = "failed" + return "", conversationID, fmt.Errorf("不支持的机器人代理模式: %s", robotMode) +} + +// StreamEvent 流式事件 +type StreamEvent struct { + Type string `json:"type"` // conversation, progress, tool_call, tool_result, response, error, cancelled, done + Message string `json:"message"` // 显示消息 + Data interface{} `json:"data,omitempty"` +} + +// publishProgressToTaskEventBus 将进度事件镜像到 taskEventBus（机器人/无 HTTP SSE 客户端时供 Web task-events 订阅）。 +func (h *AgentHandler) publishProgressToTaskEventBus(conversationID, eventType, message string, data interface{}) { + if h == nil || h.taskEventBus == nil || strings.TrimSpace(conversationID) == "" { + return + } + event := StreamEvent{Type: eventType, Message: message, Data: data} + eventJSON, err := json.Marshal(event) + if err != nil { + return + } + sseLine := make([]byte, 0, len(eventJSON)+8) + sseLine = append(sseLine, []byte("data: ")...) + sseLine = append(sseLine, eventJSON...) + sseLine = append(sseLine, '\n', '\n') + h.taskEventBus.Publish(conversationID, sseLine) +} + +func isInternalEinoDiagnosticProgress(eventType, message string, data interface{}) bool { + switch eventType { + case "model_output_rejected": + return true + case "progress": + msg := strings.TrimSpace(message) + if msg == "Eino TurnLoop 常驻多轮 runtime 已接管本轮会话。" || + msg == "Eino TurnLoop 已在安全点切换到用户补充后的下一轮。" || + msg == "已将用户补充推入 Eino TurnLoop，正在等待安全点切换…" { + return true + } + m, ok := data.(map[string]interface{}) + if !ok { + return false + } + switch strings.TrimSpace(fmt.Sprint(m["kind"])) { + case "turn_loop_takeover", "turn_loop_preempted": + return true + default: + return false + } + default: + return false + } +} + +// enrichProgressEventData 为 SSE / taskEventBus 事件补齐 conversationId、messageId，便于前端懒加载过程详情。 +func enrichProgressEventData(data interface{}, conversationID, assistantMessageID string) interface{} { + if strings.TrimSpace(conversationID) == "" && strings.TrimSpace(assistantMessageID) == "" { + return data + } + var m map[string]interface{} + switch v := data.(type) { + case map[string]interface{}: + m = make(map[string]interface{}, len(v)+2) + for k, val := range v { + m[k] = val + } + case nil: + m = make(map[string]interface{}, 2) + default: + m = map[string]interface{}{"payload": data} + } + if id := strings.TrimSpace(assistantMessageID); id != "" { + if existing, ok := m["messageId"]; !ok || strings.TrimSpace(fmt.Sprint(existing)) == "" { + m["messageId"] = id + } + } + if id := strings.TrimSpace(conversationID); id != "" { + if existing, ok := m["conversationId"]; !ok || strings.TrimSpace(fmt.Sprint(existing)) == "" { + m["conversationId"] = id + } + } + return m +} + +// createProgressCallback 创建进度回调函数，用于保存processDetails +// sendEventFunc: 可选的流式事件发送函数，如果为nil则不发送流式事件 +func (h *AgentHandler) createProgressCallback(runCtx context.Context, cancelRun context.CancelCauseFunc, conversationID, assistantMessageID string, sendEventFunc func(eventType, message string, data interface{})) agent.ProgressCallback { + // 用于保存tool_call事件中的参数，以便在tool_result时使用 + toolCallCache := make(map[string]map[string]interface{}) // toolCallId -> arguments + skillCallCache := make(map[string]string) // toolCallId -> skillName + skillToolName := "skill" + if h.config != nil { + if customName := strings.TrimSpace(h.config.MultiAgent.EinoSkills.SkillToolName); customName != "" { + skillToolName = customName + } + } + + extractSkillName := func(args map[string]interface{}) string { + if len(args) == 0 { + return "" + } + for _, key := range []string{"skill_name", "skillName", "name", "skill", "id", "skill_id", "skillId"} { + if v, ok := args[key]; ok { + switch vv := v.(type) { + case string: + if s := strings.TrimSpace(vv); s != "" { + return s + } + case map[string]interface{}: + for _, nestedKey := range []string{"name", "id", "skill_name", "skillId"} { + if nestedV, nestedOK := vv[nestedKey].(string); nestedOK { + if s := strings.TrimSpace(nestedV); s != "" { + return s + } + } + } + } + } + } + return "" + } + + // thinking_stream_*（ReAct 等助手正文流）与 reasoning_chain_stream_*（Eino ReasoningContent）： + // 不逐条落库，按 streamId 聚合，flush 时分别落 thinking / reasoning_chain。 + thinkingStreams := make(map[string]*thinkingBuf) // streamId -> buf + flushedThinking := make(map[string]bool) // streamId -> flushed + seenToolCallSigs := make(map[string]string) // toolCallId -> payload signature + seenToolResultSigs := make(map[string]string) // toolCallId -> payload signature + + // progressMu 保护闭包内 map 与聚合状态。Eino parallelRunToolCall 会在多 goroutine 中并发回调 + // progress（ToolInvokeNotifyHolder.Fire → createProgressCallback），未加锁的 map 会触发 fatal panic。 + var progressMu sync.Mutex + + // response_start + response_delta：前端时间线显示为「📝 规划中」（monitor.js），不落逐条 delta； + // 聚合为一条 planning 写入 process_details，刷新后与线上一致。 + var respPlan responsePlanAgg + if assistantMessageID != "" { + h.tasks.SetHitlAssistantMessageID(conversationID, assistantMessageID) + } + syncHitlCognition := func() { + h.syncHitlCognitionFromProgress(conversationID, assistantMessageID, thinkingStreams, &respPlan) + } + persistResponsePlan := func(reset bool) { + if assistantMessageID == "" { + return + } + content := strings.TrimSpace(respPlan.b.String()) + if content == "" { + if reset { + respPlan = responsePlanAgg{} + } + return + } + data := map[string]interface{}{ + "source": "response_stream", + } + for k, v := range respPlan.meta { + data[k] = v + } + var err error + if respPlan.detailID == "" { + respPlan.detailID, err = h.db.AddProcessDetailWithID( + assistantMessageID, conversationID, "planning", content, data, + ) + } else { + err = h.db.UpdateProcessDetailContent(respPlan.detailID, content, data) + } + if err != nil { + h.logger.Warn("保存过程详情失败", zap.Error(err), zap.String("eventType", "planning")) + } else { + respPlan.lastPersistAt = time.Now() + respPlan.lastPersistSize = respPlan.b.Len() + } + syncHitlCognition() + if reset { + respPlan = responsePlanAgg{} + } + } + flushResponsePlan := func() { persistResponsePlan(true) } + + flushThinkingStreams := func() { + if assistantMessageID == "" { + return + } + for sid, tb := range thinkingStreams { + if sid == "" || flushedThinking[sid] || tb == nil { + continue + } + content := strings.TrimSpace(tb.b.String()) + if content == "" { + flushedThinking[sid] = true + continue + } + data := map[string]interface{}{ + "streamId": sid, + } + for k, v := range tb.meta { + // 避免覆盖 streamId + if k == "streamId" { + continue + } + data[k] = v + } + persist := tb.persistAs + if persist != "reasoning_chain" { + persist = "thinking" + } + if err := h.db.AddProcessDetail(assistantMessageID, conversationID, persist, content, data); err != nil { + h.logger.Warn("保存过程详情失败", zap.Error(err), zap.String("eventType", persist)) + } + flushedThinking[sid] = true + } + syncHitlCognition() + } + + return func(eventType, message string, data interface{}) { + progressMu.Lock() + defer progressMu.Unlock() + + if isInternalEinoDiagnosticProgress(eventType, message, data) { + return + } + + // 上游在重试/补偿时可能重复回调相同 tool_call/tool_result。 + // 这里做幂等过滤，保证前端展示和 process_details 都以唯一事件为准。 + if (eventType == "tool_call" || eventType == "tool_result") && data != nil { + if dataMap, ok := data.(map[string]interface{}); ok { + toolCallID := strings.TrimSpace(fmt.Sprint(dataMap["toolCallId"])) + if toolCallID != "" && toolCallID != "" { + payloadJSON, _ := json.Marshal(dataMap) + sig := eventType + "|" + message + "|" + string(payloadJSON) + seen := seenToolCallSigs + if eventType == "tool_result" { + seen = seenToolResultSigs + } + if prev, exists := seen[toolCallID]; exists && prev == sig { + h.logger.Debug("跳过重复工具进度事件", + zap.String("eventType", eventType), + zap.String("toolCallId", toolCallID)) + return + } + seen[toolCallID] = sig + } + } + } + + // 工具输出片段不在详情区实时展示；完整结果由 tool_result 落库后按需拉取。 + if eventType == "tool_result_delta" { + return + } + + deferToolProgressSend := eventType == "tool_call" || eventType == "tool_result" + // 主 HTTP SSE 与 taskEventBus 必须同时写入：页面刷新会切断原连接，刷新后的 + // GET task-events 订阅依赖 eventBus 才能继续收到后续迭代。机器人等无主 SSE + // 的来源同样只写 eventBus。工具事件需先落库拿 processDetailId，再发送摘要。 + if !deferToolProgressSend { + clientData := enrichProgressEventData(data, conversationID, assistantMessageID) + if sendEventFunc != nil { + sendEventFunc(eventType, message, clientData) + } + h.publishProgressToTaskEventBus(conversationID, eventType, message, clientData) + } + + // 保存tool_call事件中的参数 + if eventType == "tool_call" { + if dataMap, ok := data.(map[string]interface{}); ok { + toolName, _ := dataMap["toolName"].(string) + if toolName == builtin.ToolSearchKnowledgeBase { + if toolCallId, ok := dataMap["toolCallId"].(string); ok && toolCallId != "" { + if argumentsObj, ok := dataMap["argumentsObj"].(map[string]interface{}); ok { + toolCallCache[toolCallId] = argumentsObj + } + } + } + if strings.EqualFold(strings.TrimSpace(toolName), skillToolName) { + toolCallID, _ := dataMap["toolCallId"].(string) + if toolCallID != "" { + if argumentsObj, ok := dataMap["argumentsObj"].(map[string]interface{}); ok { + if skillName := extractSkillName(argumentsObj); skillName != "" { + skillCallCache[toolCallID] = skillName + } + } + } + } + } + } + + if eventType == "tool_result" { + if dataMap, ok := data.(map[string]interface{}); ok { + toolName, _ := dataMap["toolName"].(string) + toolCallID, _ := dataMap["toolCallId"].(string) + success := true + if v, ok := dataMap["success"].(bool); ok { + success = v + } + resultText := "" + if r, ok := dataMap["result"].(string); ok { + resultText = r + } + if strings.TrimSpace(resultText) == "" { + resultText = message + } + h.recordHitlToolExecutionResult(conversationID, toolCallID, toolName, success, resultText) + } + } + + // 处理知识检索日志记录 + if eventType == "tool_result" && h.knowledgeManager != nil { + if dataMap, ok := data.(map[string]interface{}); ok { + toolName, _ := dataMap["toolName"].(string) + if toolName == builtin.ToolSearchKnowledgeBase { + // 提取检索信息 + query := "" + riskType := "" + var retrievedItems []string + + // 首先尝试从tool_call缓存中获取参数 + if toolCallId, ok := dataMap["toolCallId"].(string); ok && toolCallId != "" { + if cachedArgs, exists := toolCallCache[toolCallId]; exists { + if q, ok := cachedArgs["query"].(string); ok && q != "" { + query = q + } + if rt, ok := cachedArgs["risk_type"].(string); ok && rt != "" { + riskType = rt + } + // 使用后清理缓存 + delete(toolCallCache, toolCallId) + } + } + + // 如果缓存中没有，尝试从argumentsObj中提取 + if query == "" { + if arguments, ok := dataMap["argumentsObj"].(map[string]interface{}); ok { + if q, ok := arguments["query"].(string); ok && q != "" { + query = q + } + if rt, ok := arguments["risk_type"].(string); ok && rt != "" { + riskType = rt + } + } + } + + // 如果query仍然为空，尝试从result中提取（从结果文本的第一行） + if query == "" { + if result, ok := dataMap["result"].(string); ok && result != "" { + // 尝试从结果中提取查询内容（如果结果包含"未找到与查询 'xxx' 相关的知识"） + if strings.Contains(result, "未找到与查询 '") { + start := strings.Index(result, "未找到与查询 '") + len("未找到与查询 '") + end := strings.Index(result[start:], "'") + if end > 0 { + query = result[start : start+end] + } + } + } + // 如果还是为空，使用默认值 + if query == "" { + query = "未知查询" + } + } + + // 从工具结果中提取检索到的知识项ID + // 结果格式："找到 X 条相关知识：\n\n--- 结果 1 (相似度: XX.XX%) ---\n来源: [分类] 标题\n...\n" + if result, ok := dataMap["result"].(string); ok && result != "" { + // 尝试从元数据中提取知识项ID + metadataMatch := strings.Index(result, "") + if metadataEnd > 0 { + metadataJSON := result[metadataStart : metadataStart+metadataEnd] + var metadata map[string]interface{} + if err := json.Unmarshal([]byte(metadataJSON), &metadata); err == nil { + if meta, ok := metadata["_metadata"].(map[string]interface{}); ok { + if ids, ok := meta["retrievedItemIDs"].([]interface{}); ok { + retrievedItems = make([]string, 0, len(ids)) + for _, id := range ids { + if idStr, ok := id.(string); ok { + retrievedItems = append(retrievedItems, idStr) + } + } + } + } + } + } + } + + // 如果没有从元数据中提取到，但结果包含"找到 X 条"，至少标记为有结果 + if len(retrievedItems) == 0 && strings.Contains(result, "找到") && !strings.Contains(result, "未找到") { + // 有结果，但无法准确提取ID，使用特殊标记 + retrievedItems = []string{"_has_results"} + } + } + + // 记录检索日志（异步，不阻塞） + go func() { + if err := h.knowledgeManager.LogRetrieval(conversationID, assistantMessageID, query, riskType, retrievedItems); err != nil { + h.logger.Warn("记录知识检索日志失败", zap.Error(err)) + } + }() + + // 添加知识检索事件到processDetails + if assistantMessageID != "" { + retrievalData := map[string]interface{}{ + "query": query, + "riskType": riskType, + "toolName": toolName, + } + if err := h.db.AddProcessDetail(assistantMessageID, conversationID, "knowledge_retrieval", fmt.Sprintf("检索知识: %s", query), retrievalData); err != nil { + h.logger.Warn("保存知识检索详情失败", zap.Error(err)) + } + } + } + } + } + + // 记录 skills 调用统计（tool_call + tool_result 关联） + if eventType == "tool_result" && h.db != nil { + if dataMap, ok := data.(map[string]interface{}); ok { + toolName, _ := dataMap["toolName"].(string) + if strings.EqualFold(strings.TrimSpace(toolName), skillToolName) { + toolCallID, _ := dataMap["toolCallId"].(string) + skillName := "" + if toolCallID != "" { + skillName = strings.TrimSpace(skillCallCache[toolCallID]) + delete(skillCallCache, toolCallID) + } + if skillName == "" { + if argumentsObj, ok := dataMap["argumentsObj"].(map[string]interface{}); ok { + skillName = strings.TrimSpace(extractSkillName(argumentsObj)) + } + } + if skillName != "" { + success, ok := dataMap["success"].(bool) + if !ok { + if isError, okErr := dataMap["isError"].(bool); okErr { + success = !isError + } + } + successCalls := 0 + failedCalls := 0 + if success { + successCalls = 1 + } else { + failedCalls = 1 + } + now := time.Now() + if err := h.db.UpdateSkillStats(skillName, 1, successCalls, failedCalls, &now); err != nil { + h.logger.Warn("更新Skills调用统计失败", zap.Error(err), zap.String("skill", skillName)) + } + } + } + } + } + + // 子代理回复流式增量不落库；结束时合并为一条 eino_agent_reply + if assistantMessageID != "" && eventType == "eino_agent_reply_stream_end" { + flushResponsePlan() + // 确保思考流在子代理回复前能持久化（刷新后可读） + flushThinkingStreams() + if err := h.db.AddProcessDetail(assistantMessageID, conversationID, "eino_agent_reply", message, data); err != nil { + h.logger.Warn("保存过程详情失败", zap.Error(err), zap.String("eventType", eventType)) + } + return + } + + // 多代理主代理「规划中」：response_start / response_delta 仅用于 SSE，聚合落一条 planning + if eventType == "response_start" { + if dataMap, ok := data.(map[string]interface{}); ok { + if sameResponseStreamMeta(respPlan.meta, dataMap) { + if respPlan.meta == nil { + respPlan.meta = make(map[string]interface{}, len(dataMap)) + } + for k, v := range dataMap { + respPlan.meta[k] = v + } + return + } + } + flushResponsePlan() + // 助手正文开始前，推理流通常已结束；落库以便刷新后「渗透测试详情」可回放 + flushThinkingStreams() + respPlan.meta = nil + if dataMap, ok := data.(map[string]interface{}); ok { + respPlan.meta = make(map[string]interface{}, len(dataMap)) + for k, v := range dataMap { + respPlan.meta[k] = v + } + } + respPlan.b.Reset() + return + } + if eventType == "response_delta" { + if dataMap, ok := data.(map[string]interface{}); ok { + if acc, okAcc := dataMap[openai.SSEAccumulatedKey].(string); okAcc { + respPlan.b.Reset() + respPlan.b.WriteString(acc) + } else { + respPlan.b.WriteString(message) + } + } else { + respPlan.b.WriteString(message) + } + if dataMap, ok := data.(map[string]interface{}); ok && respPlan.meta == nil { + respPlan.meta = make(map[string]interface{}, len(dataMap)) + for k, v := range dataMap { + respPlan.meta[k] = v + } + } else if dataMap, ok := data.(map[string]interface{}); ok { + for k, v := range dataMap { + respPlan.meta[k] = v + } + } + // 运行中的主回复不能只保存在内存：刷新会销毁旧页面，新的 task-events + // 订阅只能收到未来增量。按时间或增量大小节流更新同一条 planning 记录， + // 这样刷新时能从数据库恢复刷新前已经展示的全部文本。 + if respPlan.lastPersistAt.IsZero() || + time.Since(respPlan.lastPersistAt) >= 300*time.Millisecond || + respPlan.b.Len()-respPlan.lastPersistSize >= 1024 { + persistResponsePlan(false) + } + syncHitlCognition() + return + } + if eventType == "response" { + flushResponsePlan() + flushThinkingStreams() + return + } + if eventType == "done" { + flushResponsePlan() + flushThinkingStreams() + return + } + + // 流式思考/推理结束：聚合落库（与 eino_agent_reply_stream_end 同理） + if eventType == "thinking_stream_end" || eventType == "reasoning_chain_stream_end" { + flushResponsePlan() + flushThinkingStreams() + return + } + + // 聚合 thinking_stream_* / reasoning_chain_stream_*，不逐条落库 + if eventType == "thinking_stream_start" || eventType == "reasoning_chain_stream_start" { + persistAs := "thinking" + if eventType == "reasoning_chain_stream_start" { + persistAs = "reasoning_chain" + } + if dataMap, ok := data.(map[string]interface{}); ok { + if sid, ok2 := dataMap["streamId"].(string); ok2 && sid != "" { + tb := thinkingStreams[sid] + if tb == nil { + tb = &thinkingBuf{meta: map[string]interface{}{}, persistAs: persistAs} + thinkingStreams[sid] = tb + } else { + tb.persistAs = persistAs + } + // 记录元信息（source/einoAgent/einoRole/iteration 等） + for k, v := range dataMap { + tb.meta[k] = v + } + } + } + return + } + if eventType == "thinking_stream_delta" || eventType == "reasoning_chain_stream_delta" { + persistAs := "thinking" + if eventType == "reasoning_chain_stream_delta" { + persistAs = "reasoning_chain" + } + if dataMap, ok := data.(map[string]interface{}); ok { + if sid, ok2 := dataMap["streamId"].(string); ok2 && sid != "" { + tb := thinkingStreams[sid] + if tb == nil { + tb = &thinkingBuf{meta: map[string]interface{}{}, persistAs: persistAs} + thinkingStreams[sid] = tb + } else if tb.persistAs == "" { + tb.persistAs = persistAs + } + if acc, okAcc := dataMap[openai.SSEAccumulatedKey].(string); okAcc { + tb.b.Reset() + tb.b.WriteString(acc) + } else { + tb.b.WriteString(message) + } + // 有时 delta 先到 start 未到，补充元信息 + for k, v := range dataMap { + tb.meta[k] = v + } + } + } + syncHitlCognition() + return + } + + // 当 Agent 同时发送 *_stream_* 与同名 streamId 的 thinking/reasoning_chain 时， + // 流式聚合已会在 flushThinkingStreams() 落库；此处跳过逐条重复。 + if eventType == "thinking" || eventType == "reasoning_chain" { + if dataMap, ok := data.(map[string]interface{}); ok { + if sid, ok2 := dataMap["streamId"].(string); ok2 && sid != "" { + if tb, exists := thinkingStreams[sid]; exists && tb != nil { + if strings.TrimSpace(tb.b.String()) != "" { + return + } + } + if flushedThinking[sid] { + return + } + } + } + } + + // 保存过程详情到数据库（排除 response/done；response 正文已在 messages 表） + // response_start/response_delta 已聚合为 planning，不落逐条。 + // [Eino] agent 心跳 progress 仅用于实时进度标题，不落库以免时间线刷屏。 + skipEinoAgentHeartbeat := eventType == "progress" && strings.HasPrefix(strings.TrimSpace(message), "[Eino] ") + if assistantMessageID != "" && + !skipEinoAgentHeartbeat && + eventType != "response" && + eventType != "done" && + eventType != "response_start" && + eventType != "response_delta" && + eventType != "tool_result_delta" && + eventType != "eino_trace_run" && + eventType != "eino_trace_start" && + eventType != "eino_trace_end" && + eventType != "eino_trace_error" && + eventType != "eino_agent_reply_stream_start" && + eventType != "eino_agent_reply_stream_delta" && + eventType != "eino_agent_reply_stream_end" { + if eventType == "tool_result" { + if detailID := discardPlanningIfEchoesToolResult(&respPlan, data); detailID != "" { + if err := h.db.DeleteProcessDetail(detailID); err != nil { + h.logger.Warn("删除工具结果回显规划失败", zap.Error(err), zap.String("processDetailId", detailID)) + } + } + } + // 在关键过程事件落库前，先把「规划中」与聚合中的 thinking / reasoning_chain 流落库 + flushResponsePlan() + flushThinkingStreams() + processDetailID, err := h.db.AddProcessDetailWithID(assistantMessageID, conversationID, eventType, message, data) + if err != nil { + h.logger.Warn("保存过程详情失败", zap.Error(err), zap.String("eventType", eventType)) + } + if deferToolProgressSend { + clientData := enrichProgressEventData(summarizeProcessDetailData(eventType, data), conversationID, assistantMessageID) + if m, ok := clientData.(map[string]interface{}); ok { + m["processDetailId"] = processDetailID + } + if sendEventFunc != nil { + sendEventFunc(eventType, message, clientData) + } + h.publishProgressToTaskEventBus(conversationID, eventType, message, clientData) + } + } else if deferToolProgressSend { + clientData := enrichProgressEventData(summarizeProcessDetailData(eventType, data), conversationID, assistantMessageID) + if sendEventFunc != nil { + sendEventFunc(eventType, message, clientData) + } + h.publishProgressToTaskEventBus(conversationID, eventType, message, clientData) + } + } +} + +// cancelToolContinueAfter 仅终止当前工具调用，不停止整条 Agent 任务（对话「中断并继续」与 MCP 监控终止共用）。 +func (h *AgentHandler) cancelToolContinueAfter(conversationID, preferredExecID, note string) (bool, gin.H) { + conversationID = strings.TrimSpace(conversationID) + if conversationID == "" || h.tasks.GetTask(conversationID) == nil { + return false, nil + } + note = strings.TrimSpace(note) + execID := strings.TrimSpace(preferredExecID) + if execID == "" { + execID = h.tasks.ActiveMCPExecutionID(conversationID) + } + if execID != "" { + if h.agent.CancelMCPToolExecutionWithNote(execID, note) { + return true, gin.H{ + "status": "tool_abort_requested", + "conversationId": conversationID, + "executionId": execID, + "message": "已请求终止当前工具调用；工具返回后本轮推理将继续（与 MCP 监控页终止一致）。", + "continueAfter": true, + "interruptWithNote": note != "", + "continueWithoutTool": false, + } + } + if h.tasks.AbortActiveEinoExecute(conversationID, note) { + return true, gin.H{ + "status": "tool_abort_requested", + "conversationId": conversationID, + "executionId": execID, + "message": "已请求终止当前 execute 命令；命令返回后本轮推理将继续。", + "continueAfter": true, + "interruptWithNote": note != "", + "continueWithoutTool": false, + } + } + return false, nil + } + if h.tasks.AbortActiveEinoExecute(conversationID, note) { + return true, gin.H{ + "status": "tool_abort_requested", + "conversationId": conversationID, + "message": "已请求终止当前 execute 命令；命令返回后本轮推理将继续。", + "continueAfter": true, + "interruptWithNote": note != "", + "continueWithoutTool": false, + } + } + return false, nil +} + +// CancelAgentLoop 取消正在执行的任务 +func (h *AgentHandler) CancelAgentLoop(c *gin.Context) { + var req struct { + ConversationID string `json:"conversationId" binding:"required"` + ExecutionID string `json:"executionId,omitempty"` + Reason string `json:"reason,omitempty"` + ContinueAfter bool `json:"continueAfter,omitempty"` + } + + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if !h.agentConversationAllowed(c, req.ConversationID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + + if req.ContinueAfter { + if h.tasks.GetTask(req.ConversationID) == nil { + c.JSON(http.StatusNotFound, gin.H{"error": "未找到正在执行的任务"}) + return + } + note := strings.TrimSpace(req.Reason) + activeExec := strings.TrimSpace(h.tasks.ActiveMCPExecutionID(req.ConversationID)) + if ok, payload := h.cancelToolContinueAfter(req.ConversationID, strings.TrimSpace(req.ExecutionID), note); ok { + execID, _ := payload["executionId"].(string) + h.logger.Info("对话页仅终止当前工具", + zap.String("conversationId", req.ConversationID), + zap.String("executionId", execID), + zap.Bool("hasNote", note != ""), + ) + c.JSON(http.StatusOK, payload) + return + } + if activeExec != "" { + c.JSON(http.StatusNotFound, gin.H{"error": "未找到进行中的工具执行或该调用已结束"}) + return + } + // 无进行中的 MCP 工具（模型纯推理/流式输出阶段）：取消当前上下文并由 Eino 流式处理器合并用户补充后自动续跑。 + h.tasks.SetInterruptContinueNote(req.ConversationID, note) + ok, err := h.tasks.CancelTask(req.ConversationID, multiagent.ErrInterruptContinue) + if err != nil { + h.logger.Error("中断并继续（无工具）失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if !ok { + c.JSON(http.StatusNotFound, gin.H{"error": "未找到正在执行的任务"}) + return + } + h.logger.Info("对话页中断并继续（无 MCP 工具，将自动续跑）", + zap.String("conversationId", req.ConversationID), + zap.Bool("hasNote", note != ""), + ) + c.JSON(http.StatusOK, gin.H{ + "status": "interrupt_continue_scheduled", + "conversationId": req.ConversationID, + "message": "已请求暂停当前推理；用户补充将合并到上下文并自动继续执行（无需整轮停止）。", + "continueAfter": true, + "interruptWithNote": note != "", + "continueWithoutTool": true, + }) + return + } + + var cause error = ErrTaskCancelled + msg := "已提交取消请求，任务将在当前步骤完成后停止。" + h.cancelRunningMCPToolsForConversation(req.ConversationID) + h.tasks.AbortActiveEinoExecute(req.ConversationID, "") + ok, err := h.tasks.CancelTask(req.ConversationID, cause) + if err != nil { + h.logger.Error("取消任务失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + if !ok { + c.JSON(http.StatusNotFound, gin.H{"error": "未找到正在执行的任务"}) + return + } + + c.JSON(http.StatusOK, gin.H{ + "status": "cancelling", + "conversationId": req.ConversationID, + "message": msg, + "continueAfter": false, + "interruptWithNote": false, + }) +} + +// SubscribeAgentTaskEvents GET SSE：订阅指定会话当前运行中任务的事件镜像（帧格式与 POST .../stream 一致），用于刷新页面或断线后接续 UI。 +func (h *AgentHandler) SubscribeAgentTaskEvents(c *gin.Context) { + conversationID := strings.TrimSpace(c.Query("conversationId")) + if conversationID == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "conversationId is required"}) + return + } + if !h.agentConversationAllowed(c, conversationID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + if h.tasks.GetTask(conversationID) == nil { + c.JSON(http.StatusNotFound, gin.H{"error": "no active task for this conversation"}) + return + } + if h.taskEventBus == nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "task event bus unavailable"}) + return + } + + c.Header("Content-Type", "text/event-stream; charset=utf-8") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + c.Header("X-Accel-Buffering", "no") + + sub, ch := h.taskEventBus.Subscribe(conversationID) + defer h.taskEventBus.Unsubscribe(conversationID, sub) + + flusher, _ := c.Writer.(http.Flusher) + ctx := c.Request.Context() + var writeMu sync.Mutex + stopKeepalive := runSSEKeepalive(c, &writeMu) + defer stopKeepalive() + + for { + select { + case <-ctx.Done(): + return + case chunk, ok := <-ch: + if !ok { + return + } + writeMu.Lock() + if _, err := c.Writer.Write(chunk); err != nil { + writeMu.Unlock() + return + } + if flusher != nil { + flusher.Flush() + } + writeMu.Unlock() + } + } +} + +// enrichAgentTasksWithConversationTitles 为任务列表附加当前会话标题（供顶栏/任务页展示，重命名后自动同步） +func (h *AgentHandler) enrichAgentTasksWithConversationTitles(tasks []*AgentTask) { + if h == nil || h.db == nil { + return + } + for _, task := range tasks { + if task == nil || strings.TrimSpace(task.ConversationID) == "" { + continue + } + if title, err := h.db.GetConversationTitle(task.ConversationID); err == nil { + task.Title = strings.TrimSpace(title) + } + } +} + +// enrichCompletedTasksWithConversationTitles 为已完成任务附加当前会话标题 +func (h *AgentHandler) enrichCompletedTasksWithConversationTitles(tasks []*CompletedTask) { + if h == nil || h.db == nil { + return + } + for _, task := range tasks { + if task == nil || strings.TrimSpace(task.ConversationID) == "" { + continue + } + if title, err := h.db.GetConversationTitle(task.ConversationID); err == nil { + task.Title = strings.TrimSpace(title) + } + } +} + +// ListAgentTasks 列出所有运行中的任务 +func (h *AgentHandler) ListAgentTasks(c *gin.Context) { + tasks := h.tasks.GetActiveTasks() + tasks = filterSlice(tasks, func(task *AgentTask) bool { + return task != nil && h.agentConversationAllowed(c, task.ConversationID) + }) + h.enrichAgentTasksWithConversationTitles(tasks) + c.JSON(http.StatusOK, gin.H{ + "tasks": tasks, + }) +} + +// ListCompletedTasks 列出最近完成的任务历史 +func (h *AgentHandler) ListCompletedTasks(c *gin.Context) { + tasks := h.tasks.GetCompletedTasks() + tasks = filterSlice(tasks, func(task *CompletedTask) bool { + return task != nil && h.agentConversationAllowed(c, task.ConversationID) + }) + h.enrichCompletedTasksWithConversationTitles(tasks) + c.JSON(http.StatusOK, gin.H{ + "tasks": tasks, + }) +} + +func (h *AgentHandler) agentConversationAllowed(c *gin.Context, conversationID string) bool { + session, ok := security.CurrentSession(c) + return ok && h.db != nil && h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", strings.TrimSpace(conversationID)) +} + +func filterSlice[T any](items []T, keep func(T) bool) []T { + out := make([]T, 0, len(items)) + for _, item := range items { + if keep(item) { + out = append(out, item) + } + } + return out +} + +// BatchTaskRequest 批量任务请求 +type BatchTaskRequest struct { + Title string `json:"title"` // 任务标题（可选） + Tasks []string `json:"tasks" binding:"required"` // 任务列表，每行一个任务 + Role string `json:"role,omitempty"` // 角色名称（可选，空字符串表示默认角色） + AgentMode string `json:"agentMode,omitempty"` // eino_single | deep | plan_execute | supervisor + ScheduleMode string `json:"scheduleMode,omitempty"` // manual | cron + CronExpr string `json:"cronExpr,omitempty"` // scheduleMode=cron 时必填 + ExecuteNow bool `json:"executeNow,omitempty"` // 创建后是否立即执行（默认 false） + ProjectID string `json:"projectId,omitempty"` // 队列内子对话绑定的项目（可选） + Concurrency int `json:"concurrency,omitempty"` // 同时执行的子任务数，默认 1，最大 8 +} + +// batchQueueWantsEino 队列是否配置为走 Eino 多代理。 +func batchQueueWantsEino(agentMode string) bool { + m := strings.TrimSpace(strings.ToLower(agentMode)) + return m == "deep" || m == "plan_execute" || m == "supervisor" +} + +func normalizeBatchQueueScheduleMode(mode string) string { + if strings.TrimSpace(mode) == "cron" { + return "cron" + } + return "manual" +} + +// CreateBatchQueue 创建批量任务队列 +func (h *AgentHandler) CreateBatchQueue(c *gin.Context) { + var req BatchTaskRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + if len(req.Tasks) == 0 { + c.JSON(http.StatusBadRequest, gin.H{"error": "任务列表不能为空"}) + return + } + + // 过滤空任务 + validTasks := make([]string, 0, len(req.Tasks)) + for _, task := range req.Tasks { + if task != "" { + validTasks = append(validTasks, task) + } + } + + if len(validTasks) == 0 { + c.JSON(http.StatusBadRequest, gin.H{"error": "没有有效的任务"}) + return + } + if session, ok := security.CurrentSession(c); ok && h.db != nil && session.Scope != database.RBACScopeAll && strings.TrimSpace(req.ProjectID) != "" { + if !h.db.UserCanAccessResource(session.UserID, session.Scope, "project", strings.TrimSpace(req.ProjectID)) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权在该项目下创建批量任务"}) + return + } + } + + agentMode := config.NormalizeAgentMode(req.AgentMode) + scheduleMode := normalizeBatchQueueScheduleMode(req.ScheduleMode) + cronExpr := strings.TrimSpace(req.CronExpr) + var nextRunAt *time.Time + if scheduleMode == "cron" { + if cronExpr == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "启用 Cron 调度时，调度表达式不能为空"}) + return + } + schedule, err := h.batchCronParser.Parse(cronExpr) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "无效的 Cron 表达式: " + err.Error()}) + return + } + next := schedule.Next(time.Now()) + nextRunAt = &next + } + + queue, createErr := h.batchTaskManager.CreateBatchQueue(req.Title, req.Role, agentMode, scheduleMode, cronExpr, req.ProjectID, nextRunAt, req.Concurrency, validTasks) + if createErr != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": createErr.Error()}) + return + } + if session, ok := security.CurrentSession(c); ok && h.db != nil { + _ = h.db.SetResourceOwner("batch_task", queue.ID, session.UserID) + _ = h.db.AssignResourceToUser(session.UserID, "batch_task", queue.ID) + } + started := false + if req.ExecuteNow { + ok, err := h.startBatchQueueExecution(queue.ID, false) + if !ok { + c.JSON(http.StatusNotFound, gin.H{"error": "队列不存在"}) + return + } + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error(), "queueId": queue.ID}) + return + } + started = true + if refreshed, exists := h.batchTaskManager.GetBatchQueue(queue.ID); exists { + queue = refreshed + } + } + if h.audit != nil { + h.audit.RecordOK(c, "task", "create_queue", "创建批量任务队列", "batch_queue", queue.ID, map[string]interface{}{ + "task_count": len(validTasks), "started": started, + }) + } + c.JSON(http.StatusOK, gin.H{ + "queueId": queue.ID, + "queue": queue, + "started": started, + }) +} + +// GetBatchQueue 获取批量任务队列 +func (h *AgentHandler) GetBatchQueue(c *gin.Context) { + queueID := c.Param("queueId") + queue, exists := h.batchTaskManager.GetBatchQueue(queueID) + if !exists { + c.JSON(http.StatusNotFound, gin.H{"error": "队列不存在"}) + return + } + c.JSON(http.StatusOK, gin.H{"queue": queue}) +} + +// ListBatchQueuesResponse 批量任务队列列表响应 +type ListBatchQueuesResponse struct { + Queues []*BatchTaskQueue `json:"queues"` + Total int `json:"total"` + Page int `json:"page"` + PageSize int `json:"page_size"` + TotalPages int `json:"total_pages"` +} + +// ListBatchQueues 列出所有批量任务队列（支持筛选和分页） +func (h *AgentHandler) ListBatchQueues(c *gin.Context) { + limitStr := c.DefaultQuery("limit", "10") + offsetStr := c.DefaultQuery("offset", "0") + pageStr := c.Query("page") + status := c.Query("status") + keyword := c.Query("keyword") + + limit, _ := strconv.Atoi(limitStr) + offset, _ := strconv.Atoi(offsetStr) + page := 1 + + // 如果提供了page参数，优先使用page计算offset + if pageStr != "" { + if p, err := strconv.Atoi(pageStr); err == nil && p > 0 { + page = p + offset = (page - 1) * limit + } + } + + // 限制pageSize范围 + if limit <= 0 || limit > 100 { + limit = 10 + } + if offset < 0 { + offset = 0 + } + // 防止恶意大 offset 导致 DB 性能问题 + const maxOffset = 100000 + if offset > maxOffset { + offset = maxOffset + } + + // 默认status为"all" + if status == "" { + status = "all" + } + + // 获取队列列表和总数 + session, _ := security.CurrentSession(c) + queues, total, err := h.batchTaskManager.ListQueuesForAccess(limit, offset, status, keyword, session.UserID, session.Scope) + if err != nil { + h.logger.Error("获取批量任务队列列表失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + // 计算总页数 + totalPages := (total + limit - 1) / limit + if totalPages == 0 { + totalPages = 1 + } + + // 如果使用offset计算page，需要重新计算 + if pageStr == "" { + page = (offset / limit) + 1 + } + + response := ListBatchQueuesResponse{ + Queues: queues, + Total: total, + Page: page, + PageSize: limit, + TotalPages: totalPages, + } + + c.JSON(http.StatusOK, response) +} + +// StartBatchQueue 开始执行批量任务队列 +func (h *AgentHandler) StartBatchQueue(c *gin.Context) { + queueID := c.Param("queueId") + h.batchTaskManager.ClearSingleRunTask(queueID) + ok, err := h.startBatchQueueExecution(queueID, false) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if !ok { + c.JSON(http.StatusNotFound, gin.H{"error": "队列不存在"}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "task", "start_queue", "启动批量任务队列", "batch_queue", queueID, nil) + } + c.JSON(http.StatusOK, gin.H{"message": "批量任务已开始执行", "queueId": queueID}) +} + +// RerunBatchQueue 重跑批量任务队列（重置所有子任务后重新执行） +func (h *AgentHandler) RerunBatchQueue(c *gin.Context) { + queueID := c.Param("queueId") + queue, exists := h.batchTaskManager.GetBatchQueue(queueID) + if !exists { + c.JSON(http.StatusNotFound, gin.H{"error": "队列不存在"}) + return + } + if queue.Status != "completed" && queue.Status != "cancelled" { + c.JSON(http.StatusBadRequest, gin.H{"error": "仅已完成或已取消的队列可以重跑"}) + return + } + if !h.batchTaskManager.ResetQueueForRerun(queueID) { + c.JSON(http.StatusInternalServerError, gin.H{"error": "重置队列失败"}) + return + } + h.batchTaskManager.ClearSingleRunTask(queueID) + ok, err := h.startBatchQueueExecution(queueID, false) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if !ok { + c.JSON(http.StatusInternalServerError, gin.H{"error": "启动失败"}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "task", "rerun_queue", "重跑批量任务队列", "batch_queue", queueID, nil) + } + c.JSON(http.StatusOK, gin.H{"message": "批量任务已重新开始执行", "queueId": queueID}) +} + +// PauseBatchQueue 暂停批量任务队列 +func (h *AgentHandler) PauseBatchQueue(c *gin.Context) { + queueID := c.Param("queueId") + success := h.batchTaskManager.PauseQueue(queueID) + if !success { + c.JSON(http.StatusNotFound, gin.H{"error": "队列不存在或无法暂停"}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "task", "pause_queue", "暂停批量任务队列", "batch_queue", queueID, nil) + } + c.JSON(http.StatusOK, gin.H{"message": "批量任务已暂停"}) +} + +// UpdateBatchQueueMetadata 修改批量任务队列的标题、角色和代理模式 +func (h *AgentHandler) UpdateBatchQueueMetadata(c *gin.Context) { + queueID := c.Param("queueId") + var req struct { + Title string `json:"title"` + Role string `json:"role"` + AgentMode string `json:"agentMode"` + Concurrency *int `json:"concurrency"` + } + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if err := h.batchTaskManager.UpdateQueueMetadata(queueID, req.Title, req.Role, req.AgentMode, req.Concurrency); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + updated, _ := h.batchTaskManager.GetBatchQueue(queueID) + c.JSON(http.StatusOK, gin.H{"queue": updated}) +} + +// UpdateBatchQueueSchedule 修改批量任务队列的调度配置（scheduleMode / cronExpr） +func (h *AgentHandler) UpdateBatchQueueSchedule(c *gin.Context) { + queueID := c.Param("queueId") + queue, exists := h.batchTaskManager.GetBatchQueue(queueID) + if !exists { + c.JSON(http.StatusNotFound, gin.H{"error": "队列不存在"}) + return + } + // 仅在非 running 状态下允许修改调度 + if queue.Status == "running" { + c.JSON(http.StatusBadRequest, gin.H{"error": "队列正在运行中，无法修改调度配置"}) + return + } + var req struct { + ScheduleMode string `json:"scheduleMode"` + CronExpr string `json:"cronExpr"` + } + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + scheduleMode := normalizeBatchQueueScheduleMode(req.ScheduleMode) + cronExpr := strings.TrimSpace(req.CronExpr) + var nextRunAt *time.Time + if scheduleMode == "cron" { + if cronExpr == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "启用 Cron 调度时，调度表达式不能为空"}) + return + } + schedule, err := h.batchCronParser.Parse(cronExpr) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "无效的 Cron 表达式: " + err.Error()}) + return + } + next := schedule.Next(time.Now()) + nextRunAt = &next + } + h.batchTaskManager.UpdateQueueSchedule(queueID, scheduleMode, cronExpr, nextRunAt) + updated, _ := h.batchTaskManager.GetBatchQueue(queueID) + c.JSON(http.StatusOK, gin.H{"queue": updated}) +} + +// SetBatchQueueScheduleEnabled 开启/关闭 Cron 自动调度（手工执行不受影响） +func (h *AgentHandler) SetBatchQueueScheduleEnabled(c *gin.Context) { + queueID := c.Param("queueId") + if _, exists := h.batchTaskManager.GetBatchQueue(queueID); !exists { + c.JSON(http.StatusNotFound, gin.H{"error": "队列不存在"}) + return + } + var req struct { + ScheduleEnabled bool `json:"scheduleEnabled"` + } + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if !h.batchTaskManager.SetScheduleEnabled(queueID, req.ScheduleEnabled) { + c.JSON(http.StatusNotFound, gin.H{"error": "队列不存在"}) + return + } + queue, _ := h.batchTaskManager.GetBatchQueue(queueID) + c.JSON(http.StatusOK, gin.H{"queue": queue}) +} + +// DeleteBatchQueue 删除批量任务队列 +func (h *AgentHandler) DeleteBatchQueue(c *gin.Context) { + queueID := c.Param("queueId") + if err := h.batchTaskManager.DeleteQueue(queueID); err != nil { + switch { + case errors.Is(err, ErrBatchQueueNotFound): + c.JSON(http.StatusNotFound, gin.H{"error": "队列不存在"}) + case errors.Is(err, ErrBatchQueueExecutorActive): + c.JSON(http.StatusConflict, gin.H{"error": "队列执行器仍在运行，请稍后再删除"}) + case errors.Is(err, ErrBatchQueueStillRunning): + c.JSON(http.StatusConflict, gin.H{"error": "队列正在运行中，无法删除"}) + default: + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + } + return + } + if h.audit != nil { + h.audit.Record(c, audit.Entry{ + Category: "task", + Action: "delete_queue", + Result: "success", + ResourceType: "batch_queue", + ResourceID: queueID, + Message: "删除批量任务队列", + }) + } + c.JSON(http.StatusOK, gin.H{"message": "批量任务队列已删除"}) +} + +// UpdateBatchTask 更新批量任务消息 +func (h *AgentHandler) UpdateBatchTask(c *gin.Context) { + queueID := c.Param("queueId") + taskID := c.Param("taskId") + + var req struct { + Message string `json:"message" binding:"required"` + } + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()}) + return + } + + if req.Message == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "任务消息不能为空"}) + return + } + + err := h.batchTaskManager.UpdateTaskMessage(queueID, taskID, req.Message) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + // 返回更新后的队列信息 + queue, exists := h.batchTaskManager.GetBatchQueue(queueID) + if !exists { + c.JSON(http.StatusNotFound, gin.H{"error": "队列不存在"}) + return + } + c.JSON(http.StatusOK, gin.H{"message": "任务已更新", "queue": queue}) +} + +// AddBatchTask 添加任务到批量任务队列 +func (h *AgentHandler) AddBatchTask(c *gin.Context) { + queueID := c.Param("queueId") + + var req struct { + Message string `json:"message" binding:"required"` + } + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()}) + return + } + + if req.Message == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "任务消息不能为空"}) + return + } + + task, err := h.batchTaskManager.AddTaskToQueue(queueID, req.Message) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + // 返回更新后的队列信息 + queue, exists := h.batchTaskManager.GetBatchQueue(queueID) + if !exists { + c.JSON(http.StatusNotFound, gin.H{"error": "队列不存在"}) + return + } + c.JSON(http.StatusOK, gin.H{"message": "任务已添加", "task": task, "queue": queue}) +} + +// RunSingleBatchTask 单条执行指定子任务（可覆盖已成功项），完成后暂停队列 +func (h *AgentHandler) RunSingleBatchTask(c *gin.Context) { + queueID := c.Param("queueId") + taskID := c.Param("taskId") + + if err := h.batchTaskManager.PrepareSingleTaskRun(queueID, taskID); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + h.batchTaskManager.SetSingleRunTask(queueID, taskID) + + // 暂停态单条执行：旧批量协程可能仍占用执行槽，先回收以便重新启动 + if queue, ok := h.batchTaskManager.GetBatchQueue(queueID); ok && queue.Status == BatchQueueStatusPaused { + h.batchTaskManager.ForceUnmarkQueueExecutor(queueID) + } + + autoStarted := true + autoStartMsg := "已开始单条执行" + ok, startErr := h.startBatchQueueExecution(queueID, false) + if startErr != nil { + h.batchTaskManager.ClearSingleRunTask(queueID) + autoStarted = false + autoStartMsg = "任务已准备就绪，但自动启动失败: " + startErr.Error() + } else if !ok { + h.batchTaskManager.ClearSingleRunTask(queueID) + autoStarted = false + autoStartMsg = "任务已准备就绪，但队列不存在" + } + + queue, exists := h.batchTaskManager.GetBatchQueue(queueID) + if !exists { + c.JSON(http.StatusNotFound, gin.H{"error": "队列不存在"}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "task", "run_single_batch_task", "单条执行批量子任务", "batch_task", taskID, map[string]interface{}{ + "batch_queue_id": queueID, + "auto_started": autoStarted, + }) + } + c.JSON(http.StatusOK, gin.H{ + "message": autoStartMsg, + "queue": queue, + "autoStarted": autoStarted, + }) +} + +// DeleteBatchTask 删除批量任务 +func (h *AgentHandler) DeleteBatchTask(c *gin.Context) { + queueID := c.Param("queueId") + taskID := c.Param("taskId") + + err := h.batchTaskManager.DeleteTask(queueID, taskID) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + // 返回更新后的队列信息 + queue, exists := h.batchTaskManager.GetBatchQueue(queueID) + if !exists { + c.JSON(http.StatusNotFound, gin.H{"error": "队列不存在"}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "task", "delete_batch_task", "删除批量子任务", "batch_task", taskID, map[string]interface{}{ + "batch_queue_id": queueID, + }) + } + c.JSON(http.StatusOK, gin.H{"message": "任务已删除", "queue": queue}) +} + +func (h *AgentHandler) nextBatchQueueRunAt(cronExpr string, from time.Time) (*time.Time, error) { + expr := strings.TrimSpace(cronExpr) + if expr == "" { + return nil, nil + } + schedule, err := h.batchCronParser.Parse(expr) + if err != nil { + return nil, err + } + next := schedule.Next(from) + return &next, nil +} + +func (h *AgentHandler) startBatchQueueExecution(queueID string, scheduled bool) (bool, error) { + // 先获取执行互斥门，再读取队列状态，避免基于过时快照做判断 + if !h.batchTaskManager.TryMarkQueueExecutor(queueID) { + return true, nil + } + + queue, exists := h.batchTaskManager.GetBatchQueue(queueID) + if !exists { + h.batchTaskManager.UnmarkQueueExecutor(queueID) + return false, nil + } + + if scheduled { + if queue.ScheduleMode != "cron" { + h.batchTaskManager.UnmarkQueueExecutor(queueID) + err := fmt.Errorf("队列未启用 cron 调度") + h.batchTaskManager.SetLastScheduleError(queueID, err.Error()) + return true, err + } + if queue.Status == "running" || queue.Status == "paused" || queue.Status == "cancelled" { + h.batchTaskManager.UnmarkQueueExecutor(queueID) + err := fmt.Errorf("当前队列状态不允许被调度执行") + h.batchTaskManager.SetLastScheduleError(queueID, err.Error()) + return true, err + } + if !h.batchTaskManager.ResetQueueForRerun(queueID) { + h.batchTaskManager.UnmarkQueueExecutor(queueID) + err := fmt.Errorf("重置队列失败") + h.batchTaskManager.SetLastScheduleError(queueID, err.Error()) + return true, err + } + queue, _ = h.batchTaskManager.GetBatchQueue(queueID) + } else if queue.Status != "pending" && queue.Status != "paused" { + h.batchTaskManager.UnmarkQueueExecutor(queueID) + return true, fmt.Errorf("队列状态不允许启动") + } + + if queue != nil && batchQueueWantsEino(queue.AgentMode) && (h.config == nil || !h.config.MultiAgent.Enabled) { + h.batchTaskManager.UnmarkQueueExecutor(queueID) + err := fmt.Errorf("当前队列配置为 Eino 多代理，但系统未启用多代理") + if scheduled { + h.batchTaskManager.SetLastScheduleError(queueID, err.Error()) + } + return true, err + } + + if scheduled { + h.batchTaskManager.RecordScheduledRunStart(queueID) + } + h.batchTaskManager.UpdateQueueStatus(queueID, "running") + if queue != nil && queue.ScheduleMode == "cron" { + nextRunAt, err := h.nextBatchQueueRunAt(queue.CronExpr, time.Now()) + if err == nil { + h.batchTaskManager.UpdateQueueSchedule(queueID, "cron", queue.CronExpr, nextRunAt) + } + } + + go h.executeBatchQueue(queueID) + return true, nil +} + +func (h *AgentHandler) batchQueueSchedulerLoop() { + ticker := time.NewTicker(20 * time.Second) + defer ticker.Stop() + for range ticker.C { + queues := h.batchTaskManager.GetLoadedQueues() + now := time.Now() + for _, queue := range queues { + if queue == nil || queue.ScheduleMode != "cron" || !queue.ScheduleEnabled || queue.Status == "cancelled" || queue.Status == "running" || queue.Status == "paused" { + continue + } + nextRunAt := queue.NextRunAt + if nextRunAt == nil { + next, err := h.nextBatchQueueRunAt(queue.CronExpr, now) + if err != nil { + h.logger.Warn("批量任务 cron 表达式无效，跳过调度", zap.String("queueId", queue.ID), zap.String("cronExpr", queue.CronExpr), zap.Error(err)) + continue + } + h.batchTaskManager.UpdateQueueSchedule(queue.ID, "cron", queue.CronExpr, next) + nextRunAt = next + } + if nextRunAt != nil && (nextRunAt.Before(now) || nextRunAt.Equal(now)) { + if _, err := h.startBatchQueueExecution(queue.ID, true); err != nil { + h.logger.Warn("自动调度批量任务失败", zap.String("queueId", queue.ID), zap.Error(err)) + } + } + } + } +} + +// loadHistoryFromAgentTrace 从库中保存的代理消息轨迹恢复历史（列 last_react_*；含单代理与 Eino）。 +// 逻辑与攻击链一致：优先用已保存的 JSON 消息带 + 最后一轮助手摘要，否则回退消息表。 +func (h *AgentHandler) loadHistoryFromAgentTrace(conversationID string) ([]agent.ChatMessage, error) { + traceInputJSON, assistantOut, err := h.db.GetAgentTrace(conversationID) + if err != nil { + return nil, fmt.Errorf("获取代理轨迹失败: %w", err) + } + + if traceInputJSON == "" { + return nil, fmt.Errorf("代理轨迹为空，将使用消息表") + } + + dataSource := "database_last_agent_trace" + + var messagesArray []map[string]interface{} + if err := json.Unmarshal([]byte(traceInputJSON), &messagesArray); err != nil { + return nil, fmt.Errorf("解析代理轨迹 JSON 失败: %w", err) + } + + messageCount := len(messagesArray) + modelFacingTrace := agent.IsModelFacingTraceJSON(traceInputJSON) + + h.logger.Info("使用保存的代理轨迹恢复历史上下文", + zap.String("conversationId", conversationID), + zap.String("dataSource", dataSource), + zap.Int("traceInputSize", len(traceInputJSON)), + zap.Int("messageCount", messageCount), + zap.Int("assistantOutSize", len(assistantOut)), + ) + // fmt.Println("messagesArray:", messagesArray)//debug + + // 转换为Agent消息格式 + agentMessages := make([]agent.ChatMessage, 0, len(messagesArray)) + for _, msgMap := range messagesArray { + msg := agent.ChatMessage{} + msg.ModelFacingTrace = modelFacingTrace + + // 解析role + if role, ok := msgMap["role"].(string); ok { + msg.Role = role + } else { + continue // 跳过无效消息 + } + + // 跳过 system 消息（由 Eino Instruction 提供） + if msg.Role == "system" { + continue + } + + // 解析content + if content, ok := msgMap["content"].(string); ok { + msg.Content = content + } + // DeepSeek 思考模式：含工具调用的 assistant 须在后续请求中回传 reasoning_content + if rc, ok := msgMap["reasoning_content"].(string); ok && strings.TrimSpace(rc) != "" { + msg.ReasoningContent = rc + } + + // 解析tool_calls（如果存在） + if toolCallsRaw, ok := msgMap["tool_calls"]; ok && toolCallsRaw != nil { + if toolCallsArray, ok := toolCallsRaw.([]interface{}); ok { + msg.ToolCalls = make([]agent.ToolCall, 0, len(toolCallsArray)) + for _, tcRaw := range toolCallsArray { + if tcMap, ok := tcRaw.(map[string]interface{}); ok { + toolCall := agent.ToolCall{} + + // 解析ID + if id, ok := tcMap["id"].(string); ok { + toolCall.ID = id + } + + // 解析Type + if toolType, ok := tcMap["type"].(string); ok { + toolCall.Type = toolType + } + + // 解析Function + if funcMap, ok := tcMap["function"].(map[string]interface{}); ok { + toolCall.Function = agent.FunctionCall{} + + // 解析函数名 + if name, ok := funcMap["name"].(string); ok { + toolCall.Function.Name = name + } + + // 解析arguments（可能是字符串或对象） + if argsRaw, ok := funcMap["arguments"]; ok { + if argsStr, ok := argsRaw.(string); ok { + // 如果是字符串，解析为JSON + var argsMap map[string]interface{} + if err := json.Unmarshal([]byte(argsStr), &argsMap); err == nil { + toolCall.Function.Arguments = argsMap + } + } else if argsMap, ok := argsRaw.(map[string]interface{}); ok { + // 如果已经是对象，直接使用 + toolCall.Function.Arguments = argsMap + } + } + } + + if toolCall.ID != "" { + msg.ToolCalls = append(msg.ToolCalls, toolCall) + } + } + } + } + } + + // 解析tool_call_id（tool角色消息） + if toolCallID, ok := msgMap["tool_call_id"].(string); ok { + msg.ToolCallID = toolCallID + } + if tn, ok := msgMap["tool_name"].(string); ok && strings.TrimSpace(tn) != "" { + msg.ToolName = strings.TrimSpace(tn) + } else if tn, ok := msgMap["name"].(string); ok && strings.TrimSpace(tn) != "" && strings.EqualFold(msg.Role, "tool") { + msg.ToolName = strings.TrimSpace(tn) + } + + agentMessages = append(agentMessages, msg) + } + + // 若存在 last_react_output（助手摘要），合并为最后一条 assistant（与保存格式一致） + if assistantOut != "" { + if len(agentMessages) > 0 { + lastMsg := &agentMessages[len(agentMessages)-1] + if strings.EqualFold(lastMsg.Role, "assistant") && len(lastMsg.ToolCalls) == 0 { + lastMsg.Content = assistantOut + } else { + agentMessages = append(agentMessages, agent.ChatMessage{ + Role: "assistant", + Content: assistantOut, + }) + } + } else { + agentMessages = append(agentMessages, agent.ChatMessage{ + Role: "assistant", + Content: assistantOut, + }) + } + } + + if len(agentMessages) == 0 { + return nil, fmt.Errorf("从代理轨迹解析的消息为空") + } + + if h.agent != nil { + if fixed := h.agent.RepairOrphanToolMessages(&agentMessages); fixed { + h.logger.Info("修复了从代理轨迹恢复的历史消息中的失配 tool 消息", + zap.String("conversationId", conversationID), + ) + } + } + + h.logger.Info("从代理轨迹恢复历史消息完成", + zap.String("conversationId", conversationID), + zap.String("dataSource", dataSource), + zap.Int("originalMessageCount", messageCount), + zap.Int("finalMessageCount", len(agentMessages)), + zap.Bool("hasAssistantOut", assistantOut != ""), + ) + return agentMessages, nil +} + +// dbMessagesToAgentChatMessages maps DB rows to agent ChatMessage for history fallback +// (includes reasoning_content for DeepSeek thinking + tool replay). +func dbMessagesToAgentChatMessages(msgs []database.Message) []agent.ChatMessage { + out := make([]agent.ChatMessage, 0, len(msgs)) + for i := range msgs { + m := msgs[i] + out = append(out, agent.ChatMessage{ + Role: m.Role, + Content: m.Content, + ReasoningContent: m.ReasoningContent, + }) + } + return out +} diff --git a/internal/handler/agent_progress_callback_test.go b/internal/handler/agent_progress_callback_test.go new file mode 100644 index 00000000..8c2b0da5 --- /dev/null +++ b/internal/handler/agent_progress_callback_test.go @@ -0,0 +1,242 @@ +package handler + +import ( + "context" + "fmt" + "os" + "path/filepath" + "strings" + "sync" + "testing" + "time" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/openai" + + "go.uber.org/zap" +) + +// TestCreateProgressCallback_ConcurrentToolEvents 回归 issue #142：并行 tool 回调不得 concurrent map panic。 +func TestCreateProgressCallback_ConcurrentToolEvents(t *testing.T) { + logger := zap.NewNop() + h := &AgentHandler{ + logger: logger, + config: &config.Config{}, + } + cb := h.createProgressCallback(context.Background(), nil, "conv-race-test", "", nil) + + const workers = 64 + var wg sync.WaitGroup + wg.Add(workers * 2) + for i := 0; i < workers; i++ { + i := i + go func() { + defer wg.Done() + toolCallID := fmt.Sprintf("tc-%d", i) + cb("tool_call", "calling skill", map[string]interface{}{ + "toolCallId": toolCallID, + "toolName": "skill", + "argumentsObj": map[string]interface{}{"skill_name": "demo-skill"}, + }) + }() + go func() { + defer wg.Done() + toolCallID := fmt.Sprintf("tc-%d", i) + cb("tool_result", "skill done", map[string]interface{}{ + "toolCallId": toolCallID, + "toolName": "skill", + "success": true, + }) + }() + } + wg.Wait() +} + +// TestCreateProgressCallback_MirrorsWebStreamEvents 页面刷新后 task-events 订阅必须 +// 继续收到原 Web SSE 任务的后续事件，不能只等数据库最终结果。 +func TestCreateProgressCallback_MirrorsWebStreamEvents(t *testing.T) { + bus := NewTaskEventBus() + h := &AgentHandler{logger: zap.NewNop(), config: &config.Config{}, taskEventBus: bus} + _, events := bus.Subscribe("conv-refresh-stream") + primaryCalls := 0 + cb := h.createProgressCallback( + context.Background(), nil, "conv-refresh-stream", "", + func(eventType, message string, data interface{}) { primaryCalls++ }, + ) + + cb("progress", "第 3 轮", map[string]interface{}{"iteration": 3}) + if primaryCalls != 1 { + t.Fatalf("expected primary SSE callback once, got %d", primaryCalls) + } + select { + case payload := <-events: + body := string(payload) + if !strings.Contains(body, `"type":"progress"`) || !strings.Contains(body, `"conversationId":"conv-refresh-stream"`) { + t.Fatalf("unexpected mirrored event: %s", body) + } + case <-time.After(time.Second): + t.Fatal("expected progress event mirrored to task event bus") + } +} + +func TestCreateProgressCallback_HidesInternalEinoDiagnostics(t *testing.T) { + tmp := t.TempDir() + db, err := database.NewDB(filepath.Join(tmp, "test.sqlite"), zap.NewNop()) + if err != nil { + t.Fatalf("NewDB: %v", err) + } + conv, err := db.CreateConversation("diag-hidden", database.ConversationCreateMeta{}) + if err != nil { + t.Fatalf("CreateConversation: %v", err) + } + asst, err := db.AddMessage(conv.ID, "assistant", "处理中...", nil) + if err != nil { + t.Fatalf("AddMessage: %v", err) + } + bus := NewTaskEventBus() + h := &AgentHandler{logger: zap.NewNop(), db: db, taskEventBus: bus} + _, events := bus.Subscribe(conv.ID) + primaryCalls := 0 + cb := h.createProgressCallback( + context.Background(), nil, conv.ID, asst.ID, + func(string, string, interface{}) { primaryCalls++ }, + ) + + cb("model_output_rejected", "模型工具调用不完整或参数不安全，已阻止执行并要求重写。", map[string]interface{}{ + "reason": "invalid_tool_arguments_json", + }) + cb("progress", "Eino TurnLoop 常驻多轮 runtime 已接管本轮会话。", map[string]interface{}{ + "kind": "turn_loop_takeover", + }) + + if primaryCalls != 0 { + t.Fatalf("primary SSE calls = %d, want hidden diagnostics", primaryCalls) + } + select { + case payload := <-events: + t.Fatalf("unexpected mirrored diagnostic event: %s", string(payload)) + default: + } + details, err := db.GetProcessDetails(asst.ID) + if err != nil { + t.Fatalf("GetProcessDetails: %v", err) + } + if len(details) != 0 { + t.Fatalf("process details = %+v, want no diagnostics persisted", details) + } +} + +func TestCreateProgressCallback_PersistsRunningResponseBeforeDone(t *testing.T) { + tmp := t.TempDir() + db, err := database.NewDB(filepath.Join(tmp, "test.sqlite"), zap.NewNop()) + if err != nil { + t.Fatalf("NewDB: %v", err) + } + conv, err := db.CreateConversation("refresh-running", database.ConversationCreateMeta{}) + if err != nil { + t.Fatalf("CreateConversation: %v", err) + } + asst, err := db.AddMessage(conv.ID, "assistant", "处理中...", nil) + if err != nil { + t.Fatalf("AddMessage: %v", err) + } + + h := &AgentHandler{logger: zap.NewNop(), db: db} + cb := h.createProgressCallback(context.Background(), nil, conv.ID, asst.ID, nil) + meta := map[string]interface{}{ + "streamId": "response-refresh-1", + "einoAgent": "cyberstrike-eino-single", + "orchestration": "eino_single", + } + cb("response_start", "", meta) + cb("response_delta", "刷新前已生成的第一部分", openai.WithSSEAccumulated(meta, "刷新前已生成的第一部分")) + + details, err := db.GetProcessDetails(asst.ID) + if err != nil { + t.Fatalf("GetProcessDetails: %v", err) + } + if len(details) != 1 || details[0].EventType != "planning" || details[0].Message != "刷新前已生成的第一部分" { + t.Fatalf("expected one running planning snapshot, got %+v", details) + } + + longer := "刷新前已生成的第一部分" + strings.Repeat("继续迭代", 300) + cb("response_delta", "继续迭代", openai.WithSSEAccumulated(meta, longer)) + details, err = db.GetProcessDetails(asst.ID) + if err != nil { + t.Fatalf("GetProcessDetails after update: %v", err) + } + if len(details) != 1 || details[0].Message != longer { + t.Fatalf("running snapshot should update in-place, rows=%d len=%d", len(details), len(details[0].Message)) + } +} + +// TestCreateProgressCallback_FlushesReasoningOnDone 流式推理聚合须在 done/response 时落库，刷新后可回放。 +func TestCreateProgressCallback_FlushesReasoningOnDone(t *testing.T) { + tmp := t.TempDir() + db, err := database.NewDB(filepath.Join(tmp, "test.sqlite"), zap.NewNop()) + if err != nil { + t.Fatalf("NewDB: %v", err) + } + defer os.RemoveAll(tmp) + + conv, err := db.CreateConversation("test", database.ConversationCreateMeta{}) + if err != nil { + t.Fatalf("CreateConversation: %v", err) + } + asst, err := db.AddMessage(conv.ID, "assistant", "处理中...", nil) + if err != nil { + t.Fatalf("AddMessage: %v", err) + } + + h := &AgentHandler{logger: zap.NewNop(), db: db} + cb := h.createProgressCallback(context.Background(), nil, conv.ID, asst.ID, nil) + + streamID := "eino-reasoning-test-1" + cb("reasoning_chain_stream_start", " ", map[string]interface{}{ + "streamId": streamID, + "source": "eino", + }) + cb("reasoning_chain_stream_delta", "step one", openai.WithSSEAccumulated(map[string]interface{}{ + "streamId": streamID, + }, "step one")) + cb("done", "", map[string]interface{}{"conversationId": conv.ID}) + + details, err := db.GetProcessDetails(asst.ID) + if err != nil { + t.Fatalf("GetProcessDetails: %v", err) + } + found := false + for _, d := range details { + if d.EventType == "reasoning_chain" && d.Message == "step one" { + found = true + break + } + } + if !found { + t.Fatalf("expected reasoning_chain persisted on done, got %+v", details) + } +} + +func TestEnrichProgressEventData(t *testing.T) { + t.Run("fills ids", func(t *testing.T) { + out := enrichProgressEventData(map[string]interface{}{"source": "eino"}, "conv-1", "msg-1") + m, ok := out.(map[string]interface{}) + if !ok { + t.Fatalf("expected map, got %T", out) + } + if m["conversationId"] != "conv-1" || m["messageId"] != "msg-1" { + t.Fatalf("unexpected enrichment: %+v", m) + } + }) + t.Run("preserves existing ids", func(t *testing.T) { + out := enrichProgressEventData(map[string]interface{}{ + "conversationId": "keep-conv", + "messageId": "keep-msg", + }, "conv-1", "msg-1") + m := out.(map[string]interface{}) + if m["conversationId"] != "keep-conv" || m["messageId"] != "keep-msg" { + t.Fatalf("should not overwrite existing ids: %+v", m) + } + }) +} diff --git a/internal/handler/asset.go b/internal/handler/asset.go new file mode 100644 index 00000000..0eb878db --- /dev/null +++ b/internal/handler/asset.go @@ -0,0 +1,540 @@ +package handler + +import ( + "errors" + "net/http" + "strconv" + "strings" + "time" + + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +type AssetHandler struct { + db *database.DB + logger *zap.Logger +} + +const ( + maxAssetImportBatch = 100000 + maxAssetOperationBatch = 10000 +) + +func NewAssetHandler(db *database.DB, logger *zap.Logger) *AssetHandler { + return &AssetHandler{db: db, logger: logger} +} + +func assetAccess(c *gin.Context) database.RBACListAccess { + if session, ok := security.CurrentSession(c); ok { + return database.RBACListAccess{UserID: session.UserID, Scope: session.Scope} + } + return database.RBACListAccess{} +} + +func assetAccessForPermission(c *gin.Context, permission string) database.RBACListAccess { + if session, ok := security.CurrentSession(c); ok { + return database.RBACListAccess{UserID: session.UserID, Scope: session.ScopeFor(permission)} + } + return database.RBACListAccess{} +} + +type importAssetsRequest struct { + Assets []*database.Asset `json:"assets" binding:"required"` + Source string `json:"source"` + SourceQuery string `json:"source_query"` +} + +type assetScanLink struct { + AssetID string `json:"asset_id" binding:"required"` + ConversationID string `json:"conversation_id"` + QueueID string `json:"queue_id"` + TaskID string `json:"task_id"` +} + +type recordAssetScansRequest struct { + Scans []assetScanLink `json:"scans" binding:"required"` +} + +type updateAssetsProjectRequest struct { + AssetIDs []string `json:"asset_ids" binding:"required"` + ProjectID string `json:"project_id"` +} + +type bulkUpdateAssetsRequest struct { + AssetIDs []string `json:"asset_ids" binding:"required"` + Status *string `json:"status"` + ResponsiblePerson *string `json:"responsible_person"` + Department *string `json:"department"` + BusinessSystem *string `json:"business_system"` + Environment *string `json:"environment"` + Criticality *string `json:"criticality"` + AddTags []string `json:"add_tags"` + RemoveTags []string `json:"remove_tags"` +} + +type assetIDsRequest struct { + AssetIDs []string `json:"asset_ids" binding:"required"` +} + +type mergeAssetsRequest struct { + AssetIDs []string `json:"asset_ids" binding:"required"` + PrimaryID string `json:"primary_id"` +} + +func (h *AssetHandler) Import(c *gin.Context) { + var req importAssetsRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if len(req.Assets) == 0 || len(req.Assets) > maxAssetImportBatch { + c.JSON(http.StatusBadRequest, gin.H{"error": "assets 数量必须在 1-100000 之间"}) + return + } + owner := "" + allowGlobal := false + if session, ok := security.CurrentSession(c); ok { + owner = session.UserID + allowGlobal = session.Scope == database.RBACScopeAll + } + for _, asset := range req.Assets { + if asset == nil { + continue + } + if strings.TrimSpace(asset.ProjectID) != "" { + if session, ok := security.CurrentSession(c); ok && !h.db.UserCanAccessResource(session.UserID, session.Scope, "project", strings.TrimSpace(asset.ProjectID)) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权绑定该项目"}) + return + } + } + if strings.TrimSpace(asset.Source) == "" { + asset.Source = strings.TrimSpace(req.Source) + } + if strings.TrimSpace(asset.SourceQuery) == "" { + asset.SourceQuery = strings.TrimSpace(req.SourceQuery) + } + } + result, err := h.db.UpsertAssets(req.Assets, owner, allowGlobal) + if err != nil { + var validationErr *database.AssetValidationError + if errors.As(err, &validationErr) { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + h.logger.Error("导入资产失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, result) +} + +func (h *AssetHandler) List(c *gin.Context) { + page, _ := strconv.Atoi(c.DefaultQuery("page", "1")) + pageSize, _ := strconv.Atoi(c.DefaultQuery("page_size", "20")) + if page < 1 { + page = 1 + } + if pageSize < 1 || pageSize > 100 { + pageSize = 20 + } + filter, err := assetListFilterFromQuery(c) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + assets, total, err := h.db.ListAssets(pageSize, (page-1)*pageSize, filter, assetAccess(c)) + if err != nil { + h.logger.Error("加载资产失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + totalPages := (total + pageSize - 1) / pageSize + if totalPages < 1 { + totalPages = 1 + } + c.JSON(http.StatusOK, gin.H{"assets": assets, "total": total, "page": page, "page_size": pageSize, "total_pages": totalPages}) +} + +func assetListFilterFromQuery(c *gin.Context) (database.AssetListFilter, error) { + filter := database.AssetListFilter{ + Search: strings.TrimSpace(c.Query("q")), Status: strings.ToLower(strings.TrimSpace(c.Query("status"))), + Protocol: strings.ToLower(strings.TrimSpace(c.Query("protocol"))), ProjectID: strings.TrimSpace(c.Query("project_id")), + Source: strings.TrimSpace(c.Query("source")), Tag: strings.TrimSpace(c.Query("tag")), Host: strings.TrimSpace(c.Query("host")), + IP: strings.TrimSpace(c.Query("ip")), Domain: strings.TrimSpace(c.Query("domain")), ScanState: strings.ToLower(strings.TrimSpace(c.Query("scan_state"))), + SortBy: strings.ToLower(strings.TrimSpace(c.Query("sort_by"))), SortOrder: strings.ToLower(strings.TrimSpace(c.Query("sort_order"))), + RiskLevel: strings.ToLower(strings.TrimSpace(c.Query("risk_level"))), + Country: strings.TrimSpace(c.Query("country")), Province: strings.TrimSpace(c.Query("province")), City: strings.TrimSpace(c.Query("city")), + ResponsiblePerson: strings.TrimSpace(c.Query("responsible_person")), Department: strings.TrimSpace(c.Query("department")), + BusinessSystem: strings.TrimSpace(c.Query("business_system")), Environment: strings.ToLower(strings.TrimSpace(c.Query("environment"))), + Criticality: strings.ToLower(strings.TrimSpace(c.Query("criticality"))), + } + if raw := strings.TrimSpace(c.Query("port")); raw != "" { + port, err := strconv.Atoi(raw) + if err != nil || port < 0 || port > 65535 { + return filter, &assetQueryError{field: "port", value: raw} + } + filter.Port = &port + } + for field, target := range map[string]**int{ + "min_vulnerabilities": &filter.MinVulnerabilities, + "max_vulnerabilities": &filter.MaxVulnerabilities, + "scan_overdue_days": &filter.ScanOverdueDays, + } { + raw := strings.TrimSpace(c.Query(field)) + if raw == "" { + continue + } + value, err := strconv.Atoi(raw) + if err != nil || value < 0 || (field == "scan_overdue_days" && value == 0) { + return filter, &assetQueryError{field: field, value: raw} + } + *target = &value + } + var err error + if filter.LastScanBefore, err = parseAssetQueryTime("last_scan_before", c.Query("last_scan_before")); err != nil { + return filter, err + } + if filter.LastScanAfter, err = parseAssetQueryTime("last_scan_after", c.Query("last_scan_after")); err != nil { + return filter, err + } + if filter.FirstSeenBefore, err = parseAssetQueryTime("first_seen_before", c.Query("first_seen_before")); err != nil { + return filter, err + } + if filter.FirstSeenAfter, err = parseAssetQueryTime("first_seen_after", c.Query("first_seen_after")); err != nil { + return filter, err + } + if filter.LastSeenBefore, err = parseAssetQueryTime("last_seen_before", c.Query("last_seen_before")); err != nil { + return filter, err + } + if filter.LastSeenAfter, err = parseAssetQueryTime("last_seen_after", c.Query("last_seen_after")); err != nil { + return filter, err + } + return filter, nil +} + +// Selection resolves all assets matching the current filter for cross-page actions. +func (h *AssetHandler) Selection(c *gin.Context) { + filter, err := assetListFilterFromQuery(c) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + assets, total, err := h.db.ListAssetsForOperation(maxAssetOperationBatch, filter, assetAccess(c)) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error(), "total": total}) + return + } + c.JSON(http.StatusOK, gin.H{"assets": assets, "total": total}) +} + +type assetQueryError struct{ field, value string } + +func (e *assetQueryError) Error() string { + return e.field + " 参数无效: " + e.value +} + +func parseAssetQueryTime(field, value string) (*time.Time, error) { + value = strings.TrimSpace(value) + if value == "" { + return nil, nil + } + for _, layout := range []string{time.RFC3339, "2006-01-02"} { + if parsed, err := time.Parse(layout, value); err == nil { + return &parsed, nil + } + } + return nil, &assetQueryError{field: field, value: value} +} + +func (h *AssetHandler) Stats(c *gin.Context) { + days := 30 + if raw := strings.TrimSpace(c.Query("days")); raw != "" { + parsed, err := strconv.Atoi(raw) + if err != nil || (parsed != 7 && parsed != 30 && parsed != 90) { + c.JSON(http.StatusBadRequest, gin.H{"error": "days 仅支持 7、30 或 90"}) + return + } + days = parsed + } + stats, err := h.db.GetAssetStats(assetAccess(c), days) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, stats) +} + +// RecordScans stores the execution link created by the asset-library scan action. +func (h *AssetHandler) RecordScans(c *gin.Context) { + var req recordAssetScansRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if len(req.Scans) == 0 || len(req.Scans) > maxAssetOperationBatch { + c.JSON(http.StatusBadRequest, gin.H{"error": "scans 数量必须在 1-10000 之间"}) + return + } + access := assetAccess(c) + for _, scan := range req.Scans { + conversationID := strings.TrimSpace(scan.ConversationID) + queueID := strings.TrimSpace(scan.QueueID) + taskID := strings.TrimSpace(scan.TaskID) + if conversationID == "" && taskID == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "conversation_id 或 task_id 至少需要一个"}) + return + } + if taskID != "" && (queueID == "" || !h.db.BatchTaskBelongsToQueue(taskID, queueID)) { + c.JSON(http.StatusBadRequest, gin.H{"error": "任务不属于指定队列"}) + return + } + if _, err := h.db.GetAsset(strings.TrimSpace(scan.AssetID), access); err != nil { + c.JSON(http.StatusForbidden, gin.H{"error": "资产不存在或无权扫描"}) + return + } + if session, ok := security.CurrentSession(c); ok { + if id := conversationID; id != "" && !h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", id) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权关联该对话"}) + return + } + if id := queueID; id != "" && !h.db.UserCanAccessResource(session.UserID, session.Scope, "batch_task", id) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权关联该任务队列"}) + return + } + } + } + for _, scan := range req.Scans { + if err := h.db.MarkAssetScanned(scan.AssetID, scan.ConversationID, scan.QueueID, scan.TaskID, access); err != nil { + h.logger.Error("记录资产扫描失败", zap.String("asset_id", scan.AssetID), zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + } + c.JSON(http.StatusOK, gin.H{"updated": len(req.Scans)}) +} + +func (h *AssetHandler) Update(c *gin.Context) { + var asset database.Asset + if err := c.ShouldBindJSON(&asset); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if asset.ProjectID != "" { + if session, ok := security.CurrentSession(c); ok && !h.db.UserCanAccessResource(session.UserID, session.Scope, "project", asset.ProjectID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权绑定该项目"}) + return + } + } + if err := h.db.UpdateAsset(c.Param("id"), &asset, assetAccess(c)); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + updated, err := h.db.GetAsset(c.Param("id"), assetAccess(c)) + if err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": "资产不存在"}) + return + } + c.JSON(http.StatusOK, updated) +} + +// UpdateProjectBinding replaces the project binding for a selected asset set. +func (h *AssetHandler) UpdateProjectBinding(c *gin.Context) { + var req updateAssetsProjectRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if len(req.AssetIDs) == 0 || len(req.AssetIDs) > maxAssetOperationBatch { + c.JSON(http.StatusBadRequest, gin.H{"error": "asset_ids 数量必须在 1-10000 之间"}) + return + } + req.ProjectID = strings.TrimSpace(req.ProjectID) + if req.ProjectID != "" { + if _, err := h.db.GetProject(req.ProjectID); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "项目不存在"}) + return + } + if session, ok := security.CurrentSession(c); ok && !h.db.UserCanAccessResource(session.UserID, session.Scope, "project", req.ProjectID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权绑定该项目"}) + return + } + } + updated, err := h.db.UpdateAssetsProject(req.AssetIDs, req.ProjectID, assetAccess(c)) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"updated": updated, "project_id": req.ProjectID}) +} + +func (h *AssetHandler) BulkUpdate(c *gin.Context) { + var req bulkUpdateAssetsRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if len(req.AssetIDs) == 0 || len(req.AssetIDs) > maxAssetOperationBatch { + c.JSON(http.StatusBadRequest, gin.H{"error": "asset_ids 数量必须在 1-10000 之间"}) + return + } + updated, err := h.db.UpdateAssetsBulk(req.AssetIDs, database.AssetBulkPatch{ + Status: req.Status, ResponsiblePerson: req.ResponsiblePerson, Department: req.Department, + BusinessSystem: req.BusinessSystem, Environment: req.Environment, Criticality: req.Criticality, + AddTags: req.AddTags, RemoveTags: req.RemoveTags, + }, assetAccess(c)) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"updated": updated}) +} + +func (h *AssetHandler) BatchDelete(c *gin.Context) { + var req assetIDsRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if len(req.AssetIDs) == 0 || len(req.AssetIDs) > maxAssetOperationBatch { + c.JSON(http.StatusBadRequest, gin.H{"error": "asset_ids 数量必须在 1-10000 之间"}) + return + } + deleted, err := h.db.DeleteAssets(req.AssetIDs, assetAccess(c)) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"deleted": deleted}) +} + +func assetIdentityKeys(asset *database.Asset) map[string]struct{} { + keys := map[string]struct{}{} + if value := strings.ToLower(strings.TrimSpace(asset.Domain)); value != "" { + keys["domain:"+value] = struct{}{} + } + if value := strings.ToLower(strings.Trim(strings.TrimSpace(asset.IP), "[]")); value != "" { + keys["ip:"+value] = struct{}{} + } + if value := strings.ToLower(strings.TrimSpace(asset.Host)); value != "" { + keys["host:"+value] = struct{}{} + } + return keys +} + +func shareAssetIdentity(left, right *database.Asset) bool { + for key := range assetIdentityKeys(left) { + if _, ok := assetIdentityKeys(right)[key]; ok { + return true + } + } + return false +} + +// Merge keeps the selected primary asset and safely combines compatible duplicate metadata. +func (h *AssetHandler) Merge(c *gin.Context) { + var req mergeAssetsRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if len(req.AssetIDs) < 2 || len(req.AssetIDs) > 100 { + c.JSON(http.StatusBadRequest, gin.H{"error": "合并资产数量必须在 2-100 之间"}) + return + } + writeAccess := assetAccessForPermission(c, "asset:write") + deleteAccess := assetAccessForPermission(c, "asset:delete") + primaryID := strings.TrimSpace(req.PrimaryID) + if primaryID == "" { + primaryID = strings.TrimSpace(req.AssetIDs[0]) + } + primary, err := h.db.GetAsset(primaryID, writeAccess) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "主资产不存在或无权访问"}) + return + } + others := make([]*database.Asset, 0, len(req.AssetIDs)-1) + seen := map[string]struct{}{primaryID: {}} + for _, id := range req.AssetIDs { + id = strings.TrimSpace(id) + if id == "" || id == primaryID { + continue + } + if _, ok := seen[id]; ok { + continue + } + seen[id] = struct{}{} + item, err := h.db.GetAsset(id, writeAccess) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "部分资产不存在或无权访问"}) + return + } + if !shareAssetIdentity(primary, item) { + c.JSON(http.StatusBadRequest, gin.H{"error": "所选资产没有共同域名、IP 或 Host，不能判定为重复资产"}) + return + } + others = append(others, item) + } + if len(others) == 0 { + c.JSON(http.StatusBadRequest, gin.H{"error": "至少需要两个不同资产"}) + return + } + mergeText := func(dst *string, src string) { + if strings.TrimSpace(*dst) == "" && strings.TrimSpace(src) != "" { + *dst = src + } + } + tagSet := map[string]struct{}{} + for _, tag := range primary.Tags { + tagSet[tag] = struct{}{} + } + for _, item := range others { + mergeText(&primary.ProjectID, item.ProjectID) + mergeText(&primary.Host, item.Host) + mergeText(&primary.IP, item.IP) + mergeText(&primary.Domain, item.Domain) + mergeText(&primary.Protocol, item.Protocol) + mergeText(&primary.Title, item.Title) + mergeText(&primary.Server, item.Server) + mergeText(&primary.Country, item.Country) + mergeText(&primary.Province, item.Province) + mergeText(&primary.City, item.City) + mergeText(&primary.ResponsiblePerson, item.ResponsiblePerson) + mergeText(&primary.Department, item.Department) + mergeText(&primary.BusinessSystem, item.BusinessSystem) + mergeText(&primary.Environment, item.Environment) + mergeText(&primary.Criticality, item.Criticality) + for _, tag := range item.Tags { + tagSet[tag] = struct{}{} + } + } + primary.Tags = primary.Tags[:0] + for tag := range tagSet { + primary.Tags = append(primary.Tags, tag) + } + if len(primary.Tags) > 30 { + c.JSON(http.StatusBadRequest, gin.H{"error": "合并后标签超过 30 个"}) + return + } + ids := make([]string, 0, len(others)) + for _, item := range others { + ids = append(ids, item.ID) + } + merged, err := h.db.MergeAssets(primary, ids, writeAccess, deleteAccess) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + updated, _ := h.db.GetAsset(primary.ID, writeAccess) + c.JSON(http.StatusOK, gin.H{"merged": merged, "asset": updated}) +} + +func (h *AssetHandler) Delete(c *gin.Context) { + if err := h.db.DeleteAsset(c.Param("id"), assetAccess(c)); err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": "资产不存在或无权删除"}) + return + } + c.JSON(http.StatusOK, gin.H{"success": true}) +} diff --git a/internal/handler/asset_test.go b/internal/handler/asset_test.go new file mode 100644 index 00000000..b45cebb3 --- /dev/null +++ b/internal/handler/asset_test.go @@ -0,0 +1,83 @@ +package handler + +import ( + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "path/filepath" + "testing" + + "cyberstrike-ai/internal/database" + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +func TestAssetListPaginatesWithinProject(t *testing.T) { + gin.SetMode(gin.TestMode) + db, err := database.NewDB(filepath.Join(t.TempDir(), "asset-list-pagination.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + defer db.Close() + + project, err := db.CreateProject(&database.Project{Name: "Paged Project", Status: "active"}) + if err != nil { + t.Fatal(err) + } + otherProject, err := db.CreateProject(&database.Project{Name: "Other Project", Status: "active"}) + if err != nil { + t.Fatal(err) + } + + assets := make([]*database.Asset, 0, 8) + for i := 1; i <= 7; i++ { + assets = append(assets, &database.Asset{ + ProjectID: project.ID, + IP: fmt.Sprintf("192.0.2.%d", i), + Port: 80, + Protocol: "http", + }) + } + assets = append(assets, &database.Asset{ + ProjectID: otherProject.ID, + IP: "198.51.100.1", + Port: 443, + Protocol: "https", + }) + if _, err := db.UpsertAssets(assets, "", true); err != nil { + t.Fatal(err) + } + + router := gin.New() + router.GET("/api/assets", NewAssetHandler(db, zap.NewNop()).List) + request := httptest.NewRequest(http.MethodGet, "/api/assets?project_id="+project.ID+"&page=2&page_size=3", nil) + response := httptest.NewRecorder() + router.ServeHTTP(response, request) + + if response.Code != http.StatusOK { + t.Fatalf("unexpected status %d: %s", response.Code, response.Body.String()) + } + var payload struct { + Assets []*database.Asset `json:"assets"` + Total int `json:"total"` + Page int `json:"page"` + PageSize int `json:"page_size"` + TotalPages int `json:"total_pages"` + } + if err := json.Unmarshal(response.Body.Bytes(), &payload); err != nil { + t.Fatal(err) + } + if payload.Total != 7 || payload.Page != 2 || payload.PageSize != 3 || payload.TotalPages != 3 { + t.Fatalf("unexpected pagination: total=%d page=%d page_size=%d total_pages=%d", + payload.Total, payload.Page, payload.PageSize, payload.TotalPages) + } + if len(payload.Assets) != 3 { + t.Fatalf("expected 3 assets on page 2, got %d", len(payload.Assets)) + } + for _, asset := range payload.Assets { + if asset.ProjectID != project.ID { + t.Fatalf("asset from another project leaked into page: %#v", asset) + } + } +} diff --git a/internal/handler/attackchain.go b/internal/handler/attackchain.go new file mode 100644 index 00000000..837516e8 --- /dev/null +++ b/internal/handler/attackchain.go @@ -0,0 +1,172 @@ +package handler + +import ( + "context" + "net/http" + "sync" + "time" + + "cyberstrike-ai/internal/attackchain" + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +// AttackChainHandler 攻击链处理器 +type AttackChainHandler struct { + db *database.DB + logger *zap.Logger + openAIConfig *config.OpenAIConfig + mu sync.RWMutex // 保护 openAIConfig 的并发访问 + // 用于防止同一对话的并发生成 + generatingLocks sync.Map // map[string]*sync.Mutex +} + +// NewAttackChainHandler 创建新的攻击链处理器 +func NewAttackChainHandler(db *database.DB, openAIConfig *config.OpenAIConfig, logger *zap.Logger) *AttackChainHandler { + return &AttackChainHandler{ + db: db, + logger: logger, + openAIConfig: openAIConfig, + } +} + +// UpdateConfig 更新OpenAI配置 +func (h *AttackChainHandler) UpdateConfig(cfg *config.OpenAIConfig) { + h.mu.Lock() + defer h.mu.Unlock() + h.openAIConfig = cfg + h.logger.Info("AttackChainHandler配置已更新", + zap.String("base_url", cfg.BaseURL), + zap.String("model", cfg.Model), + ) +} + +// getOpenAIConfig 获取OpenAI配置（线程安全） +func (h *AttackChainHandler) getOpenAIConfig() *config.OpenAIConfig { + h.mu.RLock() + defer h.mu.RUnlock() + return h.openAIConfig +} + +// GetAttackChain 获取攻击链（按需生成） +// GET /api/attack-chain/:conversationId +func (h *AttackChainHandler) GetAttackChain(c *gin.Context) { + conversationID := c.Param("conversationId") + if conversationID == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "conversationId is required"}) + return + } + + // 检查对话是否存在 + _, err := h.db.GetConversation(conversationID) + if err != nil { + h.logger.Warn("对话不存在", zap.String("conversationId", conversationID), zap.Error(err)) + c.JSON(http.StatusNotFound, gin.H{"error": "对话不存在"}) + return + } + + // 先尝试从数据库加载（如果已生成过） + openAIConfig := h.getOpenAIConfig() + builder := attackchain.NewBuilder(h.db, openAIConfig, h.logger) + chain, err := builder.LoadChainFromDatabase(conversationID) + if err == nil && len(chain.Nodes) > 0 { + // 如果已存在，直接返回 + h.logger.Info("返回已存在的攻击链", zap.String("conversationId", conversationID)) + c.JSON(http.StatusOK, chain) + return + } + + // 如果不存在，则生成新的攻击链（按需生成） + // 使用锁机制防止同一对话的并发生成 + lockInterface, _ := h.generatingLocks.LoadOrStore(conversationID, &sync.Mutex{}) + lock := lockInterface.(*sync.Mutex) + + // 尝试获取锁，如果正在生成则返回错误 + acquired := lock.TryLock() + if !acquired { + h.logger.Info("攻击链正在生成中，请稍后再试", zap.String("conversationId", conversationID)) + c.JSON(http.StatusConflict, gin.H{"error": "攻击链正在生成中，请稍后再试"}) + return + } + defer lock.Unlock() + + // 再次检查是否已生成（可能在等待锁的过程中已经生成完成） + chain, err = builder.LoadChainFromDatabase(conversationID) + if err == nil && len(chain.Nodes) > 0 { + h.logger.Info("返回已存在的攻击链（在锁等待期间已生成）", zap.String("conversationId", conversationID)) + c.JSON(http.StatusOK, chain) + return + } + + h.logger.Info("开始生成攻击链", zap.String("conversationId", conversationID)) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute) + defer cancel() + + chain, err = builder.BuildChainFromConversation(ctx, conversationID) + if err != nil { + h.logger.Error("生成攻击链失败", zap.String("conversationId", conversationID), zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": "生成攻击链失败: " + err.Error()}) + return + } + + // 生成完成后，从锁映射中删除（可选，保留也可以用于防止短时间内重复生成） + // h.generatingLocks.Delete(conversationID) + + c.JSON(http.StatusOK, chain) +} + +// RegenerateAttackChain 重新生成攻击链 +// POST /api/attack-chain/:conversationId/regenerate +func (h *AttackChainHandler) RegenerateAttackChain(c *gin.Context) { + conversationID := c.Param("conversationId") + if conversationID == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "conversationId is required"}) + return + } + + // 检查对话是否存在 + _, err := h.db.GetConversation(conversationID) + if err != nil { + h.logger.Warn("对话不存在", zap.String("conversationId", conversationID), zap.Error(err)) + c.JSON(http.StatusNotFound, gin.H{"error": "对话不存在"}) + return + } + + // 删除旧的攻击链 + if err := h.db.DeleteAttackChain(conversationID); err != nil { + h.logger.Warn("删除旧攻击链失败", zap.Error(err)) + } + + // 使用锁机制防止并发生成 + lockInterface, _ := h.generatingLocks.LoadOrStore(conversationID, &sync.Mutex{}) + lock := lockInterface.(*sync.Mutex) + + acquired := lock.TryLock() + if !acquired { + h.logger.Info("攻击链正在生成中，请稍后再试", zap.String("conversationId", conversationID)) + c.JSON(http.StatusConflict, gin.H{"error": "攻击链正在生成中，请稍后再试"}) + return + } + defer lock.Unlock() + + // 生成新的攻击链 + h.logger.Info("重新生成攻击链", zap.String("conversationId", conversationID)) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute) + defer cancel() + + openAIConfig := h.getOpenAIConfig() + builder := attackchain.NewBuilder(h.db, openAIConfig, h.logger) + chain, err := builder.BuildChainFromConversation(ctx, conversationID) + if err != nil { + h.logger.Error("生成攻击链失败", zap.String("conversationId", conversationID), zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": "生成攻击链失败: " + err.Error()}) + return + } + + c.JSON(http.StatusOK, chain) +} diff --git a/internal/handler/audit.go b/internal/handler/audit.go new file mode 100644 index 00000000..6baba312 --- /dev/null +++ b/internal/handler/audit.go @@ -0,0 +1,159 @@ +package handler + +import ( + "net/http" + "time" + + "cyberstrike-ai/internal/audit" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +// AuditHandler serves platform audit log APIs. +type AuditHandler struct { + db *database.DB + audit *audit.Service + logger *zap.Logger +} + +// NewAuditHandler creates an audit log handler. +func NewAuditHandler(db *database.DB, auditSvc *audit.Service, logger *zap.Logger) *AuditHandler { + return &AuditHandler{db: db, audit: auditSvc, logger: logger} +} + +// Meta GET /api/audit/meta +func (h *AuditHandler) Meta(c *gin.Context) { + enabled := false + retentionDays := 0 + if h.audit != nil { + enabled = h.audit.Enabled() + retentionDays = h.audit.RetentionDays() + } + c.JSON(http.StatusOK, gin.H{ + "enabled": enabled, + "retention_days": retentionDays, + "default_page_size": 20, + "max_page_size": 100, + "max_export": 5000, + }) +} + +// Summary GET /api/audit/summary +func (h *AuditHandler) Summary(c *gin.Context) { + if h.db == nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "database unavailable"}) + return + } + base := auditFilterForAccess(c, auditFilterFromQuery(c)) + total, err := h.db.CountAuditLogs(base) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + failFilter := base + failFilter.Result = "failure" + failures, err := h.db.CountAuditLogs(failFilter) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + since := time.Now().AddDate(0, 0, -7) + recentFilter := base + recentFilter.Since = &since + recent7d, err := h.db.CountAuditLogs(recentFilter) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{ + "total": total, + "failures": failures, + "recent_7d": recent7d, + "has_filters": c.Query("category") != "" || c.Query("action") != "" || c.Query("result") != "" || + c.Query("q") != "" || c.Query("since") != "" || c.Query("until") != "", + }) +} + +// ListLogs GET /api/audit/logs +func (h *AuditHandler) ListLogs(c *gin.Context) { + if h.db == nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "database unavailable"}) + return + } + filter := auditFilterForAccess(c, auditFilterFromQuery(c)) + page, pageSize := auditPaginationFromQuery(c) + filter.Limit = pageSize + filter.Offset = (page - 1) * pageSize + + logs, err := h.db.ListAuditLogs(filter) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + total, err := h.db.CountAuditLogs(filter) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{ + "logs": logs, + "total": total, + "page": page, + "page_size": pageSize, + }) +} + +// GetLog GET /api/audit/logs/:id +func (h *AuditHandler) GetLog(c *gin.Context) { + if h.db == nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "database unavailable"}) + return + } + row, err := h.db.GetAuditLogByID(c.Param("id")) + if err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": "审计记录不存在"}) + return + } + if session, ok := security.CurrentSession(c); !ok || (session.Scope != database.RBACScopeAll && row.Actor != session.Username) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + audit.ApplyResourceAvailability(h.db, row) + c.JSON(http.StatusOK, gin.H{"log": row}) +} + +// ExportLogs GET /api/audit/logs/export — JSON or CSV (?format=csv), max 5000 rows. +func (h *AuditHandler) ExportLogs(c *gin.Context) { + if h.db == nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "database unavailable"}) + return + } + filter := auditFilterForAccess(c, auditFilterFromQuery(c)) + filter.Limit = 5000 + filter.Offset = 0 + + logs, err := h.db.ListAuditLogs(filter) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if c.Query("format") == "csv" { + writeAuditLogsCSV(c, logs) + return + } + c.Header("Content-Disposition", `attachment; filename="audit-logs.json"`) + c.JSON(http.StatusOK, gin.H{ + "exported_at": time.Now().UTC().Format(time.RFC3339), + "logs": logs, + }) +} + +func auditFilterForAccess(c *gin.Context, filter database.ListAuditLogsFilter) database.ListAuditLogsFilter { + if session, ok := security.CurrentSession(c); ok && session.Scope != database.RBACScopeAll { + filter.Actor = session.Username + } + return filter +} diff --git a/internal/handler/audit_export_csv.go b/internal/handler/audit_export_csv.go new file mode 100644 index 00000000..debf10c9 --- /dev/null +++ b/internal/handler/audit_export_csv.go @@ -0,0 +1,42 @@ +package handler + +import ( + "encoding/csv" + "fmt" + "time" + + "cyberstrike-ai/internal/database" + + "github.com/gin-gonic/gin" +) + +func writeAuditLogsCSV(c *gin.Context, logs []*database.AuditLog) { + c.Header("Content-Type", "text/csv; charset=utf-8") + c.Header("Content-Disposition", fmt.Sprintf(`attachment; filename="audit-logs-%s.csv"`, time.Now().Format("20060102"))) + + w := csv.NewWriter(c.Writer) + _ = w.Write([]string{ + "id", "created_at", "level", "category", "action", "result", "actor", + "session_hint", "client_ip", "resource_type", "resource_id", "message", + }) + for _, row := range logs { + if row == nil { + continue + } + _ = w.Write([]string{ + row.ID, + row.CreatedAt.UTC().Format(time.RFC3339), + row.Level, + row.Category, + row.Action, + row.Result, + row.Actor, + row.SessionHint, + row.ClientIP, + row.ResourceType, + row.ResourceID, + row.Message, + }) + } + w.Flush() +} diff --git a/internal/handler/audit_query.go b/internal/handler/audit_query.go new file mode 100644 index 00000000..a355fe2e --- /dev/null +++ b/internal/handler/audit_query.go @@ -0,0 +1,49 @@ +package handler + +import ( + "strconv" + + "cyberstrike-ai/internal/database" + + "github.com/gin-gonic/gin" +) + +func auditFilterFromQuery(c *gin.Context) database.ListAuditLogsFilter { + filter := database.ListAuditLogsFilter{ + Actor: c.Query("actor"), + Level: c.Query("level"), + Category: c.Query("category"), + Action: c.Query("action"), + Result: c.Query("result"), + Query: c.Query("q"), + ResourceType: c.Query("resource_type"), + ResourceID: c.Query("resource_id"), + RelatedUserID: c.Query("related_user_id"), + } + if since := c.Query("since"); since != "" { + if t, err := database.ParseRFC3339Time(since); err == nil { + filter.Since = &t + } + } + if until := c.Query("until"); until != "" { + if t, err := database.ParseRFC3339Time(until); err == nil { + filter.Until = &t + } + } + return filter +} + +func auditPaginationFromQuery(c *gin.Context) (page, pageSize int) { + page = 1 + pageSize = 20 + if p, err := strconv.Atoi(c.DefaultQuery("page", "1")); err == nil && p > 0 { + page = p + } + if ps, err := strconv.Atoi(c.DefaultQuery("page_size", "20")); err == nil && ps > 0 { + pageSize = ps + if pageSize > 100 { + pageSize = 100 + } + } + return page, pageSize +} diff --git a/internal/handler/audit_query_test.go b/internal/handler/audit_query_test.go new file mode 100644 index 00000000..0e85513a --- /dev/null +++ b/internal/handler/audit_query_test.go @@ -0,0 +1,19 @@ +package handler + +import ( + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" +) + +func TestAuditFilterFromQueryIncludesActorAndMemberFilters(t *testing.T) { + gin.SetMode(gin.TestMode) + context, _ := gin.CreateTestContext(httptest.NewRecorder()) + context.Request = httptest.NewRequest("GET", "/api/audit/logs?actor=operator-user&action=assign_resource&resource_type=conversation&related_user_id=user-1", nil) + + filter := auditFilterFromQuery(context) + if filter.Actor != "operator-user" || filter.Action != "assign_resource" || filter.ResourceType != "conversation" || filter.RelatedUserID != "user-1" { + t.Fatalf("filter = %#v", filter) + } +} diff --git a/internal/handler/auth.go b/internal/handler/auth.go new file mode 100644 index 00000000..4157115e --- /dev/null +++ b/internal/handler/auth.go @@ -0,0 +1,237 @@ +package handler + +import ( + "net/http" + "strings" + "time" + + "cyberstrike-ai/internal/audit" + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/security" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +// AuthHandler handles authentication-related endpoints. +type AuthHandler struct { + manager *security.AuthManager + config *config.Config + configPath string + logger *zap.Logger + audit *audit.Service +} + +// SetAudit wires platform audit logging. +func (h *AuthHandler) SetAudit(s *audit.Service) { + h.audit = s +} + +// NewAuthHandler creates a new AuthHandler. +func NewAuthHandler(manager *security.AuthManager, cfg *config.Config, configPath string, logger *zap.Logger) *AuthHandler { + return &AuthHandler{ + manager: manager, + config: cfg, + configPath: configPath, + logger: logger, + } +} + +type loginRequest struct { + Username string `json:"username"` + Password string `json:"password" binding:"required"` +} + +type changePasswordRequest struct { + OldPassword string `json:"oldPassword"` + NewPassword string `json:"newPassword"` +} + +// Login verifies password and returns a session token. +func (h *AuthHandler) Login(c *gin.Context) { + var req loginRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "密码不能为空"}) + return + } + + token, expiresAt, err := h.manager.Authenticate(req.Username, req.Password) + if err != nil { + if h.audit != nil { + h.audit.Record(c, audit.Entry{ + Level: "warn", + Category: "auth", + Action: "login", + Result: "failure", + Message: "登录失败：密码错误", + Actor: strings.TrimSpace(req.Username), + }) + } + c.JSON(http.StatusUnauthorized, gin.H{"error": "密码错误"}) + return + } + session, _ := h.manager.ValidateToken(token) + + if h.audit != nil { + h.audit.Record(c, audit.Entry{ + Category: "auth", + Action: "login", + Result: "success", + SessionHint: audit.HintFromToken(token), + Message: "登录成功", + Actor: session.Username, + Detail: map[string]interface{}{ + "expires_at": expiresAt.UTC().Format(time.RFC3339), + }, + }) + } + + c.JSON(http.StatusOK, gin.H{ + "token": token, + "expires_at": expiresAt.UTC().Format(time.RFC3339), + "session_duration_hr": h.manager.SessionDurationHours(), + "user": gin.H{ + "id": session.UserID, + "username": session.Username, + "display_name": session.DisplayName, + }, + "roles": session.Roles, + "permissions": permissionKeys(session.Permissions), + "permission_scopes": session.PermissionScopes, + "scope": session.Scope, + }) +} + +// Logout revokes the current session token. +func (h *AuthHandler) Logout(c *gin.Context) { + token := c.GetString(security.ContextAuthTokenKey) + if token == "" { + authHeader := c.GetHeader("Authorization") + if len(authHeader) > 7 && strings.EqualFold(authHeader[:7], "Bearer ") { + token = strings.TrimSpace(authHeader[7:]) + } else { + token = strings.TrimSpace(authHeader) + } + } + + h.manager.RevokeToken(token) + if h.audit != nil { + h.audit.Record(c, audit.Entry{ + Category: "auth", + Action: "logout", + Result: "success", + Message: "退出登录", + }) + } + c.JSON(http.StatusOK, gin.H{"message": "已退出登录"}) +} + +// ChangePassword updates the login password. +func (h *AuthHandler) ChangePassword(c *gin.Context) { + var req changePasswordRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "参数无效"}) + return + } + + oldPassword := strings.TrimSpace(req.OldPassword) + newPassword := strings.TrimSpace(req.NewPassword) + + if oldPassword == "" || newPassword == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "当前密码和新密码均不能为空"}) + return + } + + if len(newPassword) < 8 { + c.JSON(http.StatusBadRequest, gin.H{"error": "新密码长度至少需要 8 位"}) + return + } + + if oldPassword == newPassword { + c.JSON(http.StatusBadRequest, gin.H{"error": "新密码不能与旧密码相同"}) + return + } + + session, _ := security.CurrentSession(c) + if session.Username == "" { + session.Username = "admin" + } + if !h.manager.CheckUserPassword(session.Username, oldPassword) { + if h.audit != nil { + h.audit.Record(c, audit.Entry{ + Level: "warn", + Category: "auth", + Action: "change_password", + Result: "failure", + Message: "修改密码失败：当前密码不正确", + }) + } + c.JSON(http.StatusBadRequest, gin.H{"error": "当前密码不正确"}) + return + } + + if session.UserID == "" { + session.UserID = "admin" + } + if err := h.manager.UpdateUserPassword(session.UserID, newPassword); err != nil { + if h.logger != nil { + h.logger.Error("更新用户密码失败", zap.Error(err)) + } + c.JSON(http.StatusInternalServerError, gin.H{"error": "更新用户密码失败"}) + return + } + + if h.logger != nil { + h.logger.Info("登录密码已更新，所有会话已失效") + } + + if h.audit != nil { + h.audit.Record(c, audit.Entry{ + Category: "auth", + Action: "change_password", + Result: "success", + Message: "登录密码已修改", + }) + } + + c.JSON(http.StatusOK, gin.H{"message": "密码已更新，请使用新密码重新登录"}) +} + +// Validate returns the current session status. +func (h *AuthHandler) Validate(c *gin.Context) { + token := c.GetString(security.ContextAuthTokenKey) + if token == "" { + c.JSON(http.StatusUnauthorized, gin.H{"error": "会话无效"}) + return + } + + session, ok := h.manager.ValidateToken(token) + if !ok { + c.JSON(http.StatusUnauthorized, gin.H{"error": "会话已过期"}) + return + } + + c.JSON(http.StatusOK, gin.H{ + "token": session.Token, + "expires_at": session.ExpiresAt.UTC().Format(time.RFC3339), + "user": gin.H{ + "id": session.UserID, + "username": session.Username, + "display_name": session.DisplayName, + }, + "roles": session.Roles, + "permissions": permissionKeys(session.Permissions), + "permission_scopes": session.PermissionScopes, + "scope": session.Scope, + }) +} + +func permissionKeys(perms map[string]bool) []string { + keys := make([]string, 0, len(perms)) + for key, ok := range perms { + if ok { + keys = append(keys, key) + } + } + return keys +} diff --git a/internal/handler/batch_queue_executor.go b/internal/handler/batch_queue_executor.go new file mode 100644 index 00000000..14456126 --- /dev/null +++ b/internal/handler/batch_queue_executor.go @@ -0,0 +1,402 @@ +package handler + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "strings" + "sync" + "time" + + "cyberstrike-ai/internal/agent" + "cyberstrike-ai/internal/audit" + "cyberstrike-ai/internal/authctx" + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/mcp" + "cyberstrike-ai/internal/multiagent" + + "go.uber.org/zap" +) + +const batchQueueWorkerIdlePoll = 200 * time.Millisecond + +// executeBatchQueue 使用并发 worker 池执行批量任务队列。 +func (h *AgentHandler) executeBatchQueue(queueID string) { + defer h.batchTaskManager.UnmarkQueueExecutor(queueID) + + queue, exists := h.batchTaskManager.GetBatchQueue(queueID) + if !exists { + return + } + concurrency := normalizeBatchQueueConcurrency(queue.Concurrency) + h.logger.Info("开始执行批量任务队列", zap.String("queueId", queueID), zap.Int("concurrency", concurrency)) + + var wg sync.WaitGroup + for i := 0; i < concurrency; i++ { + wg.Add(1) + go func() { + defer wg.Done() + h.runBatchQueueWorker(queueID) + }() + } + wg.Wait() + + h.tryFinalizeBatchQueue(queueID) +} + +func (h *AgentHandler) runBatchQueueWorker(queueID string) { + for { + queue, exists := h.batchTaskManager.GetBatchQueue(queueID) + if batchQueueExecutionShouldStop(queue, exists) { + return + } + + task, ok := h.batchTaskManager.ClaimNextPendingTask(queueID) + if !ok { + if !h.batchTaskManager.HasRunningTasks(queueID) { + return + } + time.Sleep(batchQueueWorkerIdlePoll) + continue + } + + queue, _ = h.batchTaskManager.GetBatchQueue(queueID) + if queue == nil { + return + } + + h.batchTaskManager.UpdateTaskStatus(queueID, task.ID, BatchTaskStatusRunning, "", "") + h.executeOneBatchSubTask(queueID, queue, task) + + if h.batchTaskManager.TakeSingleRunTaskIfMatch(queueID, task.ID) { + h.batchTaskManager.UpdateQueueStatus(queueID, BatchQueueStatusPaused) + h.logger.Info("单条执行完成，队列已暂停", zap.String("queueId", queueID), zap.String("taskId", task.ID)) + return + } + + queue, exists = h.batchTaskManager.GetBatchQueue(queueID) + if batchQueueExecutionShouldStop(queue, exists) { + if !exists { + h.logger.Warn("批量队列在执行收尾时已不存在，安全退出", zap.String("queueId", queueID)) + } + return + } + } +} + +func (h *AgentHandler) tryFinalizeBatchQueue(queueID string) { + queue, exists := h.batchTaskManager.GetBatchQueue(queueID) + if !exists || queue == nil { + return + } + if queue.Status != BatchQueueStatusRunning { + return + } + if h.batchTaskManager.HasPendingOrRunningTasks(queueID) { + return + } + + lastRunErr := "" + for _, t := range queue.Tasks { + if t != nil && t.Status == BatchTaskStatusFailed && t.Error != "" { + lastRunErr = t.Error + } + } + h.batchTaskManager.SetLastRunError(queueID, lastRunErr) + h.batchTaskManager.UpdateQueueStatus(queueID, BatchQueueStatusCompleted) + h.logger.Info("批量任务队列执行完成", zap.String("queueId", queueID)) +} + +// executeOneBatchSubTask 执行单条批量子任务（各自独立会话）。 +func (h *AgentHandler) executeOneBatchSubTask(queueID string, queue *BatchTaskQueue, task *BatchTask) { + ownerUserID := h.db.GetResourceOwner("batch_task", queueID) + access, accessErr := h.db.ResolveRBACAccess(ownerUserID) + if accessErr != nil || access == nil || !access.User.Enabled { + h.batchTaskManager.UpdateTaskStatus(queueID, task.ID, BatchTaskStatusFailed, "", "队列所有者不存在或已禁用") + return + } + principal := authctx.NewPrincipalWithScopes(access.User.ID, access.User.Username, access.Scope, access.Permissions, access.PermissionScopes) + title := safeTruncateString(task.Message, 50) + batchMeta := batchSubTaskConversationMeta(h.config, queue) + conv, err := h.db.CreateConversation(title, batchMeta) + if err != nil { + h.logger.Error("创建对话失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(err)) + h.batchTaskManager.UpdateTaskStatus(queueID, task.ID, BatchTaskStatusFailed, "", "创建对话失败: "+err.Error()) + return + } + conversationID := conv.ID + _ = h.db.SetResourceOwner("conversation", conversationID, access.User.ID) + _ = h.db.AssignResourceToUser(access.User.ID, "conversation", conversationID) + + h.batchTaskManager.UpdateTaskStatusWithConversationID(queueID, task.ID, BatchTaskStatusRunning, "", "", conversationID) + + finalMessage := task.Message + var roleTools []string + if queue.Role != "" && queue.Role != "默认" { + if h.config.Roles != nil { + if role, exists := h.config.Roles[queue.Role]; exists && role.Enabled { + if role.UserPrompt != "" { + finalMessage = role.UserPrompt + "\n\n" + task.Message + h.logger.Info("应用角色用户提示词", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("role", queue.Role)) + } + if len(role.Tools) > 0 { + roleTools = role.Tools + h.logger.Info("使用角色配置的工具列表", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("role", queue.Role), zap.Int("toolCount", len(roleTools))) + } + } + } + } + + if _, err = h.db.AddMessage(conversationID, "user", task.Message, nil); err != nil { + h.logger.Error("保存用户消息失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID), zap.Error(err)) + } + + assistantMsg, err := h.db.AddMessage(conversationID, "assistant", "处理中...", nil) + if err != nil { + h.logger.Error("创建助手消息失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID), zap.Error(err)) + assistantMsg = nil + } + + var assistantMessageID string + if assistantMsg != nil { + assistantMessageID = assistantMsg.ID + } + + h.logger.Info("执行批量任务", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("message", task.Message), zap.String("role", queue.Role), zap.String("conversationId", conversationID)) + + principalCtx := authctx.WithPrincipal(context.Background(), principal) + baseCtx, cancelWithCause := context.WithCancelCause(principalCtx) + taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 6*time.Hour) + + registered := false + finishStatus := "completed" + + defer func() { + h.batchTaskManager.SetTaskCancel(queueID, task.ID, nil) + timeoutCancel() + if registered { + if h.taskEventBus != nil { + ev := StreamEvent{Type: "done", Message: "", Data: map[string]interface{}{"conversationId": conversationID}} + if b, err := json.Marshal(ev); err == nil { + h.taskEventBus.Publish(conversationID, append(append([]byte("data: "), b...), '\n', '\n')) + } + } + h.tasks.FinishTask(conversationID, finishStatus) + } + cancelWithCause(nil) + }() + + sendEvent := func(eventType, message string, data interface{}) { + if h.taskEventBus == nil { + return + } + ev := StreamEvent{Type: eventType, Message: message, Data: data} + b, err := json.Marshal(ev) + if err != nil { + b = []byte(`{"type":"error","message":"marshal failed"}`) + } + line := make([]byte, 0, len(b)+8) + line = append(line, []byte("data: ")...) + line = append(line, b...) + line = append(line, '\n', '\n') + h.taskEventBus.Publish(conversationID, line) + } + + if _, err := h.tasks.StartTask(conversationID, task.Message, cancelWithCause); err != nil { + h.logger.Warn("批量队列子任务注册会话运行状态失败", + zap.String("queueId", queueID), + zap.String("taskId", task.ID), + zap.String("conversationId", conversationID), + zap.Error(err)) + failMsg := err.Error() + if errors.Is(err, ErrTaskAlreadyRunning) { + failMsg = "会话已有任务正在执行，无法在该会话上并行启动批量子任务" + } + h.batchTaskManager.UpdateTaskStatus(queueID, task.ID, BatchTaskStatusFailed, "", failMsg) + return + } + registered = true + h.batchTaskManager.SetTaskCancel(queueID, task.ID, timeoutCancel) + + progressCallback := h.createProgressCallback(taskCtx, cancelWithCause, conversationID, assistantMessageID, sendEvent) + taskCtx = mcp.WithMCPConversationID(taskCtx, conversationID) + taskCtx = mcp.WithToolRunRegistry(taskCtx, h.tasks) + taskCtx = mcp.WithEinoExecuteRunRegistry(taskCtx, h.tasks) + + useBatchMulti := false + batchOrch := "deep" + am := strings.TrimSpace(strings.ToLower(queue.AgentMode)) + if am == "multi" { + am = "deep" + } + if batchQueueWantsEino(queue.AgentMode) && h.config != nil && h.config.MultiAgent.Enabled { + useBatchMulti = true + batchOrch = config.NormalizeMultiAgentOrchestration(am) + } else if queue.AgentMode == "" && h.config != nil && h.config.MultiAgent.Enabled && h.config.MultiAgent.BatchUseMultiAgent { + useBatchMulti = true + batchOrch = "deep" + } + if useBatchMulti { + _ = h.db.SetConversationAgentMode(conversationID, batchOrch) + } else { + _ = h.db.SetConversationAgentMode(conversationID, "eino_single") + } + + var resultMA *multiagent.RunResult + var runErr error + switch { + case useBatchMulti: + resultMA, runErr = multiagent.RunDeepAgent(taskCtx, h.config, &h.config.MultiAgent, h.agent, h.db, h.logger, conversationID, h.conversationProjectID(conversationID), finalMessage, []agent.ChatMessage{}, roleTools, progressCallback, h.agentsMarkdownDir, batchOrch, nil, h.agentSessionContextBlock(conversationID)) + default: + if h.config == nil { + runErr = fmt.Errorf("服务器配置未加载") + } else { + resultMA, runErr = multiagent.RunEinoSingleChatModelAgent(taskCtx, h.config, &h.config.MultiAgent, h.agent, h.db, h.logger, conversationID, h.conversationProjectID(conversationID), finalMessage, []agent.ChatMessage{}, roleTools, progressCallback, nil, h.agentSessionContextBlock(conversationID)) + } + } + + if runErr != nil { + h.handleBatchSubTaskRunError(queueID, task, conversationID, assistantMessageID, baseCtx, taskCtx, resultMA, runErr, &finishStatus) + return + } + + if resultMA == nil { + h.logger.Error("批量任务执行成功但无结果对象", + zap.String("queueId", queueID), + zap.String("taskId", task.ID), + zap.String("conversationId", conversationID)) + h.batchTaskManager.UpdateTaskStatus(queueID, task.ID, BatchTaskStatusFailed, "", "内部错误：无执行结果") + return + } + + h.logger.Info("批量任务执行成功", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID)) + + mcpIDs := resultMA.MCPExecutionIDs + lastIn := resultMA.LastAgentTraceInput + lastOut := resultMA.LastAgentTraceOutput + reasoningContent := multiagent.AggregatedReasoningFromTraceJSON(lastIn) + agentMode := "batch_eino_single" + if useBatchMulti { + agentMode = "batch_eino_" + batchOrch + } + decision := h.finalizeAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, agentMode, resultMA, mcpIDs, reasoningContent, true) + resText := decision.FinalText + if !decision.Finalizable { + resText = finalizationBlockedMessage(decision) + finishStatus = decision.Status + sendEvent("finalization_check", resText, decision) + } + sendEvent("response", resText, finalizationResponsePayload(decision, map[string]interface{}{ + "conversationId": conversationID, + "messageId": assistantMessageID, + "agentMode": agentMode, + "mcpExecutionIds": mcpIDs, + "batchQueueId": queueID, + "batchTaskId": task.ID, + "batchTaskStatus": map[bool]string{true: string(BatchTaskStatusCompleted), false: string(BatchTaskStatusFailed)}[decision.Finalizable], + "candidatePreview": safeTruncateString(resultMA.Response, 500), + })) + + if assistantMessageID == "" { + _, err = h.db.AddMessage(conversationID, "assistant", resText, mcpIDs) + } else if !decision.Finalizable { + err = nil + } + if err != nil { + h.logger.Error("保存助手消息失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID), zap.Error(err)) + } + + if lastIn != "" || lastOut != "" { + if err := h.db.SaveAgentTrace(conversationID, lastIn, lastOut); err != nil { + h.logger.Warn("保存代理轨迹失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(err)) + } + } + + if !decision.Finalizable { + h.batchTaskManager.UpdateTaskStatusWithConversationID(queueID, task.ID, BatchTaskStatusFailed, resText, finalizationCheckMessage(decision), conversationID) + return + } + h.batchTaskManager.UpdateTaskStatusWithConversationID(queueID, task.ID, BatchTaskStatusCompleted, resText, "", conversationID) +} + +func batchSubTaskConversationMeta(cfg *config.Config, queue *BatchTaskQueue) database.ConversationCreateMeta { + meta := audit.ConversationCreateMeta("batch_task") + if queue == nil { + meta.ProjectID = effectiveProjectID(cfg, "") + return meta + } + meta.ProjectID = effectiveProjectID(cfg, queue.ProjectID) + meta.RoleName = strings.TrimSpace(queue.Role) + return meta +} + +func (h *AgentHandler) handleBatchSubTaskRunError( + queueID string, + task *BatchTask, + conversationID, assistantMessageID string, + baseCtx, taskCtx context.Context, + resultMA *multiagent.RunResult, + runErr error, + finishStatus *string, +) { + if shouldPersistEinoAgentTraceAfterRunError(baseCtx) { + h.persistEinoAgentTraceForResume(conversationID, resultMA) + } + errStr := runErr.Error() + partialResp := "" + if resultMA != nil { + partialResp = resultMA.Response + } + isCancelled := errors.Is(context.Cause(baseCtx), ErrTaskCancelled) || + errors.Is(runErr, context.Canceled) || + strings.Contains(strings.ToLower(errStr), "context canceled") || + strings.Contains(strings.ToLower(errStr), "context cancelled") || + (partialResp != "" && (strings.Contains(partialResp, "任务已被取消") || strings.Contains(partialResp, "任务执行中断"))) + isTimeout := errors.Is(runErr, context.DeadlineExceeded) || errors.Is(context.Cause(taskCtx), context.DeadlineExceeded) + + if isTimeout { + *finishStatus = "timeout" + } else if isCancelled { + *finishStatus = "cancelled" + } else { + *finishStatus = "failed" + } + + if isCancelled { + h.logger.Info("批量任务被取消", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID)) + cancelMsg := "任务已被用户取消，后续操作已停止。" + if partialResp != "" && (strings.Contains(partialResp, "任务已被取消") || strings.Contains(partialResp, "任务执行中断")) { + cancelMsg = partialResp + } + if assistantMessageID != "" { + if updateErr := h.appendAssistantMessageNotice(assistantMessageID, cancelMsg); updateErr != nil { + h.logger.Warn("更新取消后的助手消息失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(updateErr)) + } + if err := h.db.AddProcessDetail(assistantMessageID, conversationID, "cancelled", cancelMsg, nil); err != nil { + h.logger.Warn("保存取消详情失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(err)) + } + } else if _, errMsg := h.db.AddMessage(conversationID, "assistant", cancelMsg, nil); errMsg != nil { + h.logger.Warn("保存取消消息失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(errMsg)) + } + h.batchTaskManager.UpdateTaskStatusWithConversationID(queueID, task.ID, BatchTaskStatusCancelled, cancelMsg, "", conversationID) + return + } + + h.logger.Error("批量任务执行失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID), zap.Error(runErr)) + errorMsg := "执行失败: " + runErr.Error() + if assistantMessageID != "" { + if _, updateErr := h.db.Exec( + "UPDATE messages SET content = ?, updated_at = ? WHERE id = ?", + errorMsg, + time.Now(), assistantMessageID, + ); updateErr != nil { + h.logger.Warn("更新失败后的助手消息失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(updateErr)) + } + if err := h.db.AddProcessDetail(assistantMessageID, conversationID, "error", errorMsg, nil); err != nil { + h.logger.Warn("保存错误详情失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(err)) + } + } + h.batchTaskManager.UpdateTaskStatus(queueID, task.ID, BatchTaskStatusFailed, "", runErr.Error()) +} diff --git a/internal/handler/batch_task_manager.go b/internal/handler/batch_task_manager.go new file mode 100644 index 00000000..9db9f266 --- /dev/null +++ b/internal/handler/batch_task_manager.go @@ -0,0 +1,1457 @@ +package handler + +import ( + "context" + "crypto/rand" + "encoding/hex" + "errors" + "fmt" + "sort" + "strings" + "sync" + "time" + "unicode/utf8" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + + "go.uber.org/zap" +) + +var ( + // ErrBatchQueueNotFound 队列不存在或已从内存卸载。 + ErrBatchQueueNotFound = errors.New("batch queue not found") + // ErrBatchQueueExecutorActive executeBatchQueue 协程仍在收尾，禁止删除。 + ErrBatchQueueExecutorActive = errors.New("batch queue executor is still active") + // ErrBatchQueueStillRunning 队列状态仍为 running（无活跃执行器时的兜底保护）。 + ErrBatchQueueStillRunning = errors.New("batch queue is still running") +) + +// 批量任务状态常量 +const ( + BatchQueueStatusPending = "pending" + BatchQueueStatusRunning = "running" + BatchQueueStatusPaused = "paused" + BatchQueueStatusCompleted = "completed" + BatchQueueStatusCancelled = "cancelled" + + BatchTaskStatusPending = "pending" + BatchTaskStatusRunning = "running" + BatchTaskStatusCompleted = "completed" + BatchTaskStatusFailed = "failed" + BatchTaskStatusCancelled = "cancelled" + + // MaxBatchTasksPerQueue 单个队列最大任务数 + MaxBatchTasksPerQueue = 10000 + + // MaxBatchQueueTitleLen 队列标题最大长度 + MaxBatchQueueTitleLen = 200 + + // MaxBatchQueueRoleLen 角色名最大长度 + MaxBatchQueueRoleLen = 100 + + // DefaultBatchQueueConcurrency 批量队列默认并发数（串行） + DefaultBatchQueueConcurrency = 1 + + // MaxBatchQueueConcurrency 批量队列最大并发数 + MaxBatchQueueConcurrency = 8 +) + +// BatchTask 批量任务项 +type BatchTask struct { + ID string `json:"id"` + Message string `json:"message"` + ConversationID string `json:"conversationId,omitempty"` + Status string `json:"status"` // pending, running, completed, failed, cancelled + StartedAt *time.Time `json:"startedAt,omitempty"` + CompletedAt *time.Time `json:"completedAt,omitempty"` + Error string `json:"error,omitempty"` + Result string `json:"result,omitempty"` +} + +// BatchTaskQueue 批量任务队列 +type BatchTaskQueue struct { + ID string `json:"id"` + Title string `json:"title,omitempty"` + Role string `json:"role,omitempty"` // 角色名称（空字符串表示默认角色） + AgentMode string `json:"agentMode"` // single | eino_single | deep | plan_execute | supervisor + ScheduleMode string `json:"scheduleMode"` // manual | cron + CronExpr string `json:"cronExpr,omitempty"` + NextRunAt *time.Time `json:"nextRunAt,omitempty"` + ScheduleEnabled bool `json:"scheduleEnabled"` + LastScheduleTriggerAt *time.Time `json:"lastScheduleTriggerAt,omitempty"` + LastScheduleError string `json:"lastScheduleError,omitempty"` + LastRunError string `json:"lastRunError,omitempty"` + ProjectID string `json:"projectId,omitempty"` + Concurrency int `json:"concurrency"` // 同时执行的子任务数，默认 1 + Tasks []*BatchTask `json:"tasks"` + Status string `json:"status"` // pending, running, paused, completed, cancelled + CreatedAt time.Time `json:"createdAt"` + StartedAt *time.Time `json:"startedAt,omitempty"` + CompletedAt *time.Time `json:"completedAt,omitempty"` + CurrentIndex int `json:"currentIndex"` +} + +// BatchTaskManager 批量任务管理器 +type BatchTaskManager struct { + db *database.DB + logger *zap.Logger + queues map[string]*BatchTaskQueue + taskCancels map[string]map[string]context.CancelFunc // queueID -> taskID -> 取消函数 + singleRunTasks map[string]string // queueID -> taskID，单条执行完成后暂停队列 + queueExecutors map[string]struct{} // executeBatchQueue 协程活跃标记（与队列 status 解耦） + mu sync.RWMutex +} + +// NewBatchTaskManager 创建批量任务管理器 +func NewBatchTaskManager(logger *zap.Logger) *BatchTaskManager { + if logger == nil { + logger = zap.NewNop() + } + return &BatchTaskManager{ + logger: logger, + queues: make(map[string]*BatchTaskQueue), + taskCancels: make(map[string]map[string]context.CancelFunc), + singleRunTasks: make(map[string]string), + queueExecutors: make(map[string]struct{}), + } +} + +// batchQueueExecutionShouldStop 判断 executeBatchQueue 主循环是否应退出。 +func batchQueueExecutionShouldStop(queue *BatchTaskQueue, exists bool) bool { + if !exists || queue == nil { + return true + } + switch queue.Status { + case BatchQueueStatusCancelled, BatchQueueStatusCompleted, BatchQueueStatusPaused: + return true + default: + return false + } +} + +// TryMarkQueueExecutor 标记队列执行协程已启动；若已有执行协程则返回 false。 +func (m *BatchTaskManager) TryMarkQueueExecutor(queueID string) bool { + m.mu.Lock() + defer m.mu.Unlock() + if _, exists := m.queueExecutors[queueID]; exists { + return false + } + m.queueExecutors[queueID] = struct{}{} + return true +} + +// UnmarkQueueExecutor 清除队列执行协程标记（executeBatchQueue defer 调用）。 +func (m *BatchTaskManager) UnmarkQueueExecutor(queueID string) { + m.mu.Lock() + defer m.mu.Unlock() + delete(m.queueExecutors, queueID) +} + +// ForceUnmarkQueueExecutor 强制清除执行协程标记（暂停态单条重跑等场景回收陈旧槽位）。 +func (m *BatchTaskManager) ForceUnmarkQueueExecutor(queueID string) { + m.UnmarkQueueExecutor(queueID) +} + +// IsQueueExecutorActive 队列 executeBatchQueue 协程是否仍在运行。 +func (m *BatchTaskManager) IsQueueExecutorActive(queueID string) bool { + m.mu.RLock() + defer m.mu.RUnlock() + _, ok := m.queueExecutors[queueID] + return ok +} + +// SetDB 设置数据库连接 +func (m *BatchTaskManager) SetDB(db *database.DB) { + m.mu.Lock() + defer m.mu.Unlock() + m.db = db +} + +// normalizeBatchQueueConcurrency 规范化队列并发数。 +func normalizeBatchQueueConcurrency(n int) int { + if n < 1 { + return DefaultBatchQueueConcurrency + } + if n > MaxBatchQueueConcurrency { + return MaxBatchQueueConcurrency + } + return n +} + +// CreateBatchQueue 创建批量任务队列 +func (m *BatchTaskManager) CreateBatchQueue( + title, role, agentMode, scheduleMode, cronExpr, projectID string, + nextRunAt *time.Time, + concurrency int, + tasks []string, +) (*BatchTaskQueue, error) { + // 输入校验 + if utf8.RuneCountInString(title) > MaxBatchQueueTitleLen { + return nil, fmt.Errorf("标题不能超过 %d 个字符", MaxBatchQueueTitleLen) + } + if utf8.RuneCountInString(role) > MaxBatchQueueRoleLen { + return nil, fmt.Errorf("角色名不能超过 %d 个字符", MaxBatchQueueRoleLen) + } + if len(tasks) > MaxBatchTasksPerQueue { + return nil, fmt.Errorf("单个队列最多 %d 条任务", MaxBatchTasksPerQueue) + } + + m.mu.Lock() + defer m.mu.Unlock() + + queueID := time.Now().Format("20060102150405") + "-" + generateShortID() + queue := &BatchTaskQueue{ + ID: queueID, + Title: title, + Role: role, + ProjectID: strings.TrimSpace(projectID), + AgentMode: config.NormalizeAgentMode(agentMode), + ScheduleMode: normalizeBatchQueueScheduleMode(scheduleMode), + CronExpr: strings.TrimSpace(cronExpr), + NextRunAt: nextRunAt, + ScheduleEnabled: true, + Concurrency: normalizeBatchQueueConcurrency(concurrency), + Tasks: make([]*BatchTask, 0, len(tasks)), + Status: BatchQueueStatusPending, + CreatedAt: time.Now(), + CurrentIndex: 0, + } + if queue.ScheduleMode != "cron" { + queue.CronExpr = "" + queue.NextRunAt = nil + } + + // 准备数据库保存的任务数据 + dbTasks := make([]map[string]interface{}, 0, len(tasks)) + + for _, message := range tasks { + if message == "" { + continue // 跳过空行 + } + taskID := generateShortID() + task := &BatchTask{ + ID: taskID, + Message: message, + Status: BatchTaskStatusPending, + } + queue.Tasks = append(queue.Tasks, task) + dbTasks = append(dbTasks, map[string]interface{}{ + "id": taskID, + "message": message, + }) + } + + // 保存到数据库 + if m.db != nil { + if err := m.db.CreateBatchQueue( + queueID, + title, + role, + queue.AgentMode, + queue.ScheduleMode, + queue.CronExpr, + queue.NextRunAt, + queue.ProjectID, + queue.Concurrency, + dbTasks, + ); err != nil { + m.logger.Warn("batch queue DB create failed", zap.String("queueId", queueID), zap.Error(err)) + } + } + + m.queues[queueID] = queue + return queue, nil +} + +// GetBatchQueue 获取批量任务队列 +func (m *BatchTaskManager) GetBatchQueue(queueID string) (*BatchTaskQueue, bool) { + m.mu.RLock() + queue, exists := m.queues[queueID] + m.mu.RUnlock() + + if exists { + return queue, true + } + + // 如果内存中不存在，尝试从数据库加载 + if m.db != nil { + if queue := m.loadQueueFromDB(queueID); queue != nil { + m.mu.Lock() + m.queues[queueID] = queue + m.mu.Unlock() + return queue, true + } + } + + return nil, false +} + +// loadQueueFromDB 从数据库加载单个队列 +func (m *BatchTaskManager) loadQueueFromDB(queueID string) *BatchTaskQueue { + if m.db == nil { + return nil + } + + queueRow, err := m.db.GetBatchQueue(queueID) + if err != nil || queueRow == nil { + return nil + } + + taskRows, err := m.db.GetBatchTasks(queueID) + if err != nil { + return nil + } + + queue := &BatchTaskQueue{ + ID: queueRow.ID, + AgentMode: "eino_single", + ScheduleMode: "manual", + Status: queueRow.Status, + CreatedAt: queueRow.CreatedAt, + CurrentIndex: queueRow.CurrentIndex, + Tasks: make([]*BatchTask, 0, len(taskRows)), + } + + if queueRow.Title.Valid { + queue.Title = queueRow.Title.String + } + if queueRow.Role.Valid { + queue.Role = queueRow.Role.String + } + if queueRow.AgentMode.Valid { + queue.AgentMode = config.NormalizeAgentMode(queueRow.AgentMode.String) + } + if queueRow.ScheduleMode.Valid { + queue.ScheduleMode = normalizeBatchQueueScheduleMode(queueRow.ScheduleMode.String) + } + if queueRow.CronExpr.Valid && queue.ScheduleMode == "cron" { + queue.CronExpr = strings.TrimSpace(queueRow.CronExpr.String) + } + if queueRow.NextRunAt.Valid && queue.ScheduleMode == "cron" { + t := queueRow.NextRunAt.Time + queue.NextRunAt = &t + } + queue.ScheduleEnabled = true + if queueRow.ScheduleEnabled.Valid && queueRow.ScheduleEnabled.Int64 == 0 { + queue.ScheduleEnabled = false + } + if queueRow.LastScheduleTriggerAt.Valid { + t := queueRow.LastScheduleTriggerAt.Time + queue.LastScheduleTriggerAt = &t + } + if queueRow.LastScheduleError.Valid { + queue.LastScheduleError = strings.TrimSpace(queueRow.LastScheduleError.String) + } + if queueRow.LastRunError.Valid { + queue.LastRunError = strings.TrimSpace(queueRow.LastRunError.String) + } + if queueRow.ProjectID.Valid { + queue.ProjectID = strings.TrimSpace(queueRow.ProjectID.String) + } + queue.Concurrency = batchQueueConcurrencyFromRow(queueRow) + if queueRow.StartedAt.Valid { + queue.StartedAt = &queueRow.StartedAt.Time + } + if queueRow.CompletedAt.Valid { + queue.CompletedAt = &queueRow.CompletedAt.Time + } + + for _, taskRow := range taskRows { + task := &BatchTask{ + ID: taskRow.ID, + Message: taskRow.Message, + Status: taskRow.Status, + } + if taskRow.ConversationID.Valid { + task.ConversationID = taskRow.ConversationID.String + } + if taskRow.StartedAt.Valid { + task.StartedAt = &taskRow.StartedAt.Time + } + if taskRow.CompletedAt.Valid { + task.CompletedAt = &taskRow.CompletedAt.Time + } + if taskRow.Error.Valid { + task.Error = taskRow.Error.String + } + if taskRow.Result.Valid { + task.Result = taskRow.Result.String + } + queue.Tasks = append(queue.Tasks, task) + } + + return queue +} + +// GetLoadedQueues 获取内存中已加载的队列（不触发 DB 加载，仅用 RLock） +func (m *BatchTaskManager) GetLoadedQueues() []*BatchTaskQueue { + m.mu.RLock() + result := make([]*BatchTaskQueue, 0, len(m.queues)) + for _, queue := range m.queues { + result = append(result, queue) + } + m.mu.RUnlock() + return result +} + +// GetAllQueues 获取所有队列 +func (m *BatchTaskManager) GetAllQueues() []*BatchTaskQueue { + m.mu.RLock() + result := make([]*BatchTaskQueue, 0, len(m.queues)) + for _, queue := range m.queues { + result = append(result, queue) + } + m.mu.RUnlock() + + // 如果数据库可用，确保所有数据库中的队列都已加载到内存 + if m.db != nil { + dbQueues, err := m.db.GetAllBatchQueues() + if err == nil { + m.mu.Lock() + for _, queueRow := range dbQueues { + if _, exists := m.queues[queueRow.ID]; !exists { + if queue := m.loadQueueFromDB(queueRow.ID); queue != nil { + m.queues[queueRow.ID] = queue + result = append(result, queue) + } + } + } + m.mu.Unlock() + } + } + + return result +} + +// ListQueues 列出队列（支持筛选和分页） +func (m *BatchTaskManager) ListQueues(limit, offset int, status, keyword string) ([]*BatchTaskQueue, int, error) { + return m.ListQueuesForAccess(limit, offset, status, keyword, "", "") +} + +func (m *BatchTaskManager) ListQueuesForAccess(limit, offset int, status, keyword, userID, scope string) ([]*BatchTaskQueue, int, error) { + var queues []*BatchTaskQueue + var total int + + // 如果数据库可用，从数据库查询 + if m.db != nil { + // 获取总数 + count, err := m.db.CountBatchQueuesForAccess(status, keyword, userID, scope) + if err != nil { + return nil, 0, fmt.Errorf("统计队列总数失败: %w", err) + } + total = count + + // 获取队列列表（只获取ID） + queueRows, err := m.db.ListBatchQueuesForAccess(limit, offset, status, keyword, userID, scope) + if err != nil { + return nil, 0, fmt.Errorf("查询队列列表失败: %w", err) + } + + // 加载完整的队列信息（从内存或数据库） + m.mu.Lock() + for _, queueRow := range queueRows { + var queue *BatchTaskQueue + // 先从内存查找 + if cached, exists := m.queues[queueRow.ID]; exists { + queue = cached + } else { + // 从数据库加载 + queue = m.loadQueueFromDB(queueRow.ID) + if queue != nil { + m.queues[queueRow.ID] = queue + } + } + if queue != nil { + queues = append(queues, queue) + } + } + m.mu.Unlock() + } else { + // 没有数据库，从内存中筛选和分页 + m.mu.RLock() + allQueues := make([]*BatchTaskQueue, 0, len(m.queues)) + for _, queue := range m.queues { + allQueues = append(allQueues, queue) + } + m.mu.RUnlock() + + // 筛选 + filtered := make([]*BatchTaskQueue, 0) + for _, queue := range allQueues { + // 状态筛选 + if status != "" && status != "all" && queue.Status != status { + continue + } + // 关键字搜索（搜索队列ID和标题） + if keyword != "" { + keywordLower := strings.ToLower(keyword) + queueIDLower := strings.ToLower(queue.ID) + queueTitleLower := strings.ToLower(queue.Title) + if !strings.Contains(queueIDLower, keywordLower) && !strings.Contains(queueTitleLower, keywordLower) { + // 也可以搜索创建时间 + createdAtStr := queue.CreatedAt.Format("2006-01-02 15:04:05") + if !strings.Contains(createdAtStr, keyword) { + continue + } + } + } + filtered = append(filtered, queue) + } + + // 按创建时间倒序排序 + sort.Slice(filtered, func(i, j int) bool { + return filtered[i].CreatedAt.After(filtered[j].CreatedAt) + }) + + total = len(filtered) + + // 分页 + start := offset + if start > len(filtered) { + start = len(filtered) + } + end := start + limit + if end > len(filtered) { + end = len(filtered) + } + if start < len(filtered) { + queues = filtered[start:end] + } + } + + return queues, total, nil +} + +// LoadFromDB 从数据库加载所有队列 +func (m *BatchTaskManager) LoadFromDB() error { + if m.db == nil { + return nil + } + + queueRows, err := m.db.GetAllBatchQueues() + if err != nil { + return err + } + + m.mu.Lock() + defer m.mu.Unlock() + + for _, queueRow := range queueRows { + if _, exists := m.queues[queueRow.ID]; exists { + continue // 已存在，跳过 + } + + taskRows, err := m.db.GetBatchTasks(queueRow.ID) + if err != nil { + continue // 跳过加载失败的任务 + } + + queue := &BatchTaskQueue{ + ID: queueRow.ID, + AgentMode: "eino_single", + ScheduleMode: "manual", + Status: queueRow.Status, + CreatedAt: queueRow.CreatedAt, + CurrentIndex: queueRow.CurrentIndex, + Tasks: make([]*BatchTask, 0, len(taskRows)), + } + + if queueRow.Title.Valid { + queue.Title = queueRow.Title.String + } + if queueRow.Role.Valid { + queue.Role = queueRow.Role.String + } + if queueRow.AgentMode.Valid { + queue.AgentMode = config.NormalizeAgentMode(queueRow.AgentMode.String) + } + if queueRow.ScheduleMode.Valid { + queue.ScheduleMode = normalizeBatchQueueScheduleMode(queueRow.ScheduleMode.String) + } + if queueRow.CronExpr.Valid && queue.ScheduleMode == "cron" { + queue.CronExpr = strings.TrimSpace(queueRow.CronExpr.String) + } + if queueRow.NextRunAt.Valid && queue.ScheduleMode == "cron" { + t := queueRow.NextRunAt.Time + queue.NextRunAt = &t + } + queue.ScheduleEnabled = true + if queueRow.ScheduleEnabled.Valid && queueRow.ScheduleEnabled.Int64 == 0 { + queue.ScheduleEnabled = false + } + if queueRow.LastScheduleTriggerAt.Valid { + t := queueRow.LastScheduleTriggerAt.Time + queue.LastScheduleTriggerAt = &t + } + if queueRow.LastScheduleError.Valid { + queue.LastScheduleError = strings.TrimSpace(queueRow.LastScheduleError.String) + } + if queueRow.LastRunError.Valid { + queue.LastRunError = strings.TrimSpace(queueRow.LastRunError.String) + } + if queueRow.ProjectID.Valid { + queue.ProjectID = strings.TrimSpace(queueRow.ProjectID.String) + } + queue.Concurrency = batchQueueConcurrencyFromRow(queueRow) + if queueRow.StartedAt.Valid { + queue.StartedAt = &queueRow.StartedAt.Time + } + if queueRow.CompletedAt.Valid { + queue.CompletedAt = &queueRow.CompletedAt.Time + } + + for _, taskRow := range taskRows { + task := &BatchTask{ + ID: taskRow.ID, + Message: taskRow.Message, + Status: taskRow.Status, + } + if taskRow.ConversationID.Valid { + task.ConversationID = taskRow.ConversationID.String + } + if taskRow.StartedAt.Valid { + task.StartedAt = &taskRow.StartedAt.Time + } + if taskRow.CompletedAt.Valid { + task.CompletedAt = &taskRow.CompletedAt.Time + } + if taskRow.Error.Valid { + task.Error = taskRow.Error.String + } + if taskRow.Result.Valid { + task.Result = taskRow.Result.String + } + queue.Tasks = append(queue.Tasks, task) + } + + m.queues[queueRow.ID] = queue + } + + return nil +} + +// UpdateTaskStatus 更新任务状态 +func (m *BatchTaskManager) UpdateTaskStatus(queueID, taskID, status string, result, errorMsg string) { + m.UpdateTaskStatusWithConversationID(queueID, taskID, status, result, errorMsg, "") +} + +// UpdateTaskStatusWithConversationID 更新任务状态（包含conversationId） +func (m *BatchTaskManager) UpdateTaskStatusWithConversationID(queueID, taskID, status string, result, errorMsg, conversationID string) { + m.mu.Lock() + defer m.mu.Unlock() + + queue, exists := m.queues[queueID] + if !exists { + return + } + + // DB 优先：先持久化，成功后再更新内存，避免重启后状态不一致 + if m.db != nil { + if err := m.db.UpdateBatchTaskStatus(queueID, taskID, status, conversationID, result, errorMsg); err != nil { + m.logger.Warn("batch task DB status update failed, skipping memory update", + zap.String("queueId", queueID), zap.String("taskId", taskID), zap.Error(err)) + return + } + } + + for _, task := range queue.Tasks { + if task.ID == taskID { + task.Status = status + if result != "" { + task.Result = result + } + if errorMsg != "" { + task.Error = errorMsg + } + if conversationID != "" { + task.ConversationID = conversationID + } + now := time.Now() + if status == BatchTaskStatusRunning && task.StartedAt == nil { + task.StartedAt = &now + } + if status == BatchTaskStatusCompleted || status == BatchTaskStatusFailed || status == BatchTaskStatusCancelled { + task.CompletedAt = &now + } + break + } + } +} + +// UpdateQueueStatus 更新队列状态 +func (m *BatchTaskManager) UpdateQueueStatus(queueID, status string) { + m.mu.Lock() + defer m.mu.Unlock() + + queue, exists := m.queues[queueID] + if !exists { + return + } + + // DB 优先：先持久化，成功后再更新内存 + if m.db != nil { + if err := m.db.UpdateBatchQueueStatus(queueID, status); err != nil { + m.logger.Warn("batch queue DB status update failed, skipping memory update", + zap.String("queueId", queueID), zap.Error(err)) + return + } + } + + queue.Status = status + now := time.Now() + if status == BatchQueueStatusRunning && queue.StartedAt == nil { + queue.StartedAt = &now + } + if status == BatchQueueStatusCompleted || status == BatchQueueStatusCancelled { + queue.CompletedAt = &now + } +} + +// UpdateQueueSchedule 更新队列调度配置 +func (m *BatchTaskManager) UpdateQueueSchedule(queueID, scheduleMode, cronExpr string, nextRunAt *time.Time) { + m.mu.Lock() + defer m.mu.Unlock() + + queue, exists := m.queues[queueID] + if !exists { + return + } + + queue.ScheduleMode = normalizeBatchQueueScheduleMode(scheduleMode) + if queue.ScheduleMode == "cron" { + queue.CronExpr = strings.TrimSpace(cronExpr) + queue.NextRunAt = nextRunAt + } else { + queue.CronExpr = "" + queue.NextRunAt = nil + } + + if m.db != nil { + if err := m.db.UpdateBatchQueueSchedule(queueID, queue.ScheduleMode, queue.CronExpr, queue.NextRunAt); err != nil { + m.logger.Warn("batch queue DB schedule update failed", zap.String("queueId", queueID), zap.Error(err)) + } + } +} + +// batchQueueConcurrencyFromRow 从数据库行读取并发数（缺省为 1）。 +func batchQueueConcurrencyFromRow(row *database.BatchTaskQueueRow) int { + if row == nil || !row.Concurrency.Valid { + return DefaultBatchQueueConcurrency + } + return normalizeBatchQueueConcurrency(int(row.Concurrency.Int64)) +} + +// UpdateQueueMetadata 更新队列标题、角色、代理模式和并发数（非 running 时可用） +func (m *BatchTaskManager) UpdateQueueMetadata(queueID, title, role, agentMode string, concurrency *int) error { + if utf8.RuneCountInString(title) > MaxBatchQueueTitleLen { + return fmt.Errorf("标题不能超过 %d 个字符", MaxBatchQueueTitleLen) + } + if utf8.RuneCountInString(role) > MaxBatchQueueRoleLen { + return fmt.Errorf("角色名不能超过 %d 个字符", MaxBatchQueueRoleLen) + } + m.mu.Lock() + defer m.mu.Unlock() + + queue, exists := m.queues[queueID] + if !exists { + return fmt.Errorf("队列不存在") + } + if queue.Status == BatchQueueStatusRunning { + return fmt.Errorf("队列正在运行中，无法修改") + } + + // 如果未传 agentMode，保留原值 + if strings.TrimSpace(agentMode) != "" { + agentMode = config.NormalizeAgentMode(agentMode) + } else { + agentMode = queue.AgentMode + } + + queue.Title = title + queue.Role = role + queue.AgentMode = agentMode + if concurrency != nil { + queue.Concurrency = normalizeBatchQueueConcurrency(*concurrency) + } + + if m.db != nil { + if err := m.db.UpdateBatchQueueMetadata(queueID, title, role, agentMode, queue.Concurrency); err != nil { + m.logger.Warn("batch queue DB metadata update failed", zap.String("queueId", queueID), zap.Error(err)) + } + } + return nil +} + +// SetScheduleEnabled 暂停/恢复 Cron 自动调度（不影响手工执行） +func (m *BatchTaskManager) SetScheduleEnabled(queueID string, enabled bool) bool { + m.mu.Lock() + defer m.mu.Unlock() + + queue, exists := m.queues[queueID] + if !exists { + return false + } + queue.ScheduleEnabled = enabled + if m.db != nil { + _ = m.db.UpdateBatchQueueScheduleEnabled(queueID, enabled) + } + return true +} + +// RecordScheduledRunStart Cron 触发成功、即将执行子任务时调用 +func (m *BatchTaskManager) RecordScheduledRunStart(queueID string) { + now := time.Now() + m.mu.Lock() + defer m.mu.Unlock() + + queue, exists := m.queues[queueID] + if !exists { + return + } + queue.LastScheduleTriggerAt = &now + queue.LastScheduleError = "" + if m.db != nil { + _ = m.db.RecordBatchQueueScheduledTriggerStart(queueID, now) + } +} + +// SetLastScheduleError 调度层失败（未成功开始执行） +func (m *BatchTaskManager) SetLastScheduleError(queueID, msg string) { + m.mu.Lock() + defer m.mu.Unlock() + + queue, exists := m.queues[queueID] + if !exists { + return + } + queue.LastScheduleError = strings.TrimSpace(msg) + if m.db != nil { + _ = m.db.SetBatchQueueLastScheduleError(queueID, queue.LastScheduleError) + } +} + +// SetLastRunError 最近一轮批量执行中的失败摘要 +func (m *BatchTaskManager) SetLastRunError(queueID, msg string) { + msg = strings.TrimSpace(msg) + m.mu.Lock() + defer m.mu.Unlock() + + queue, exists := m.queues[queueID] + if !exists { + return + } + queue.LastRunError = msg + if m.db != nil { + _ = m.db.SetBatchQueueLastRunError(queueID, msg) + } +} + +// ResetQueueForRerun 重置队列与子任务状态，供 cron 下一轮执行 +func (m *BatchTaskManager) ResetQueueForRerun(queueID string) bool { + m.mu.Lock() + defer m.mu.Unlock() + + queue, exists := m.queues[queueID] + if !exists { + return false + } + + // DB 优先：先持久化重置，成功后再更新内存，避免 DB 失败导致内存脏状态 + if m.db != nil { + if err := m.db.ResetBatchQueueForRerun(queueID); err != nil { + m.logger.Warn("batch queue DB reset for rerun failed, skipping memory update", + zap.String("queueId", queueID), zap.Error(err)) + return false + } + } + + queue.Status = BatchQueueStatusPending + queue.CurrentIndex = 0 + queue.StartedAt = nil + queue.CompletedAt = nil + queue.NextRunAt = nil + queue.LastRunError = "" + queue.LastScheduleError = "" + for _, task := range queue.Tasks { + task.Status = BatchTaskStatusPending + task.ConversationID = "" + task.StartedAt = nil + task.CompletedAt = nil + task.Error = "" + task.Result = "" + } + return true +} + +// UpdateTaskMessage 更新任务消息（队列空闲时可改；任务需非 running） +func (m *BatchTaskManager) UpdateTaskMessage(queueID, taskID, message string) error { + m.mu.Lock() + defer m.mu.Unlock() + + queue, exists := m.queues[queueID] + if !exists { + return fmt.Errorf("队列不存在") + } + + if !queueAllowsTaskListMutationLocked(queue) { + return fmt.Errorf("队列正在执行或未就绪，无法编辑任务") + } + + // 查找并更新任务 + for _, task := range queue.Tasks { + if task.ID == taskID { + if task.Status == BatchTaskStatusRunning { + return fmt.Errorf("执行中的任务不能编辑") + } + task.Message = message + + // 同步到数据库 + if m.db != nil { + if err := m.db.UpdateBatchTaskMessage(queueID, taskID, message); err != nil { + return fmt.Errorf("更新任务消息失败: %w", err) + } + } + return nil + } + } + + return fmt.Errorf("任务不存在") +} + +// AddTaskToQueue 添加任务到队列（队列空闲时可添加：含 cron 本轮 completed、手动暂停后等） +func (m *BatchTaskManager) AddTaskToQueue(queueID, message string) (*BatchTask, error) { + m.mu.Lock() + defer m.mu.Unlock() + + queue, exists := m.queues[queueID] + if !exists { + return nil, fmt.Errorf("队列不存在") + } + + if !queueAllowsTaskListMutationLocked(queue) { + return nil, fmt.Errorf("队列正在执行或未就绪，无法添加任务") + } + + if message == "" { + return nil, fmt.Errorf("任务消息不能为空") + } + + // 生成任务ID + taskID := generateShortID() + task := &BatchTask{ + ID: taskID, + Message: message, + Status: BatchTaskStatusPending, + } + + // 添加到内存队列 + queue.Tasks = append(queue.Tasks, task) + + // 同步到数据库 + if m.db != nil { + if err := m.db.AddBatchTask(queueID, taskID, message); err != nil { + // 如果数据库保存失败，从内存中移除 + queue.Tasks = queue.Tasks[:len(queue.Tasks)-1] + return nil, fmt.Errorf("添加任务失败: %w", err) + } + } + + return task, nil +} + +// PrepareSingleTaskRun 准备单条执行：重置目标任务（若已有结果）并定位队列索引 +func (m *BatchTaskManager) PrepareSingleTaskRun(queueID, taskID string) error { + var siblingRunningIDs []string + + m.mu.Lock() + queue, exists := m.queues[queueID] + if !exists { + m.mu.Unlock() + return fmt.Errorf("队列不存在") + } + + var task *BatchTask + taskIndex := -1 + for i, t := range queue.Tasks { + if t.ID == taskID { + taskIndex = i + task = t + break + } + } + if task == nil { + m.mu.Unlock() + return fmt.Errorf("任务不存在") + } + + if !queueAllowsSingleTaskRunLocked(queue, task) { + m.mu.Unlock() + return fmt.Errorf("队列正在执行或未就绪，无法单条执行") + } + + // 暂停态：中止在途子任务并收口仍标记 running 的其它子任务，以便单条执行非冲突项 + var cancelFuncs []context.CancelFunc + if queue.Status == BatchQueueStatusPaused { + cancelFuncs = m.drainTaskCancelsLocked(queueID) + for _, t := range queue.Tasks { + if t != nil && t.ID != taskID && t.Status == BatchTaskStatusRunning { + siblingRunningIDs = append(siblingRunningIDs, t.ID) + } + } + } + + needsReset := task.Status != BatchTaskStatusPending + resumeQueue := queue.Status == BatchQueueStatusCompleted || queue.Status == BatchQueueStatusCancelled + m.mu.Unlock() + + for _, c := range cancelFuncs { + if c != nil { + c() + } + } + const staleRunMsg = "为单条执行其它任务，已中止" + for _, sid := range siblingRunningIDs { + m.UpdateTaskStatus(queueID, sid, BatchTaskStatusCancelled, "", staleRunMsg) + } + + m.mu.Lock() + defer m.mu.Unlock() + + queue, exists = m.queues[queueID] + if !exists { + return fmt.Errorf("队列不存在") + } + + task = nil + taskIndex = -1 + for i, t := range queue.Tasks { + if t.ID == taskID { + taskIndex = i + task = t + break + } + } + if task == nil { + return fmt.Errorf("任务不存在") + } + + if m.db != nil { + if err := m.db.PrepareBatchSingleTaskRun(queueID, taskID, taskIndex, needsReset, resumeQueue); err != nil { + return fmt.Errorf("准备单条执行失败: %w", err) + } + } + + if needsReset { + task.Status = BatchTaskStatusPending + task.ConversationID = "" + task.StartedAt = nil + task.CompletedAt = nil + task.Error = "" + task.Result = "" + } + queue.CurrentIndex = taskIndex + queue.LastRunError = "" + if resumeQueue { + queue.Status = BatchQueueStatusPaused + queue.CompletedAt = nil + } + + return nil +} + +// SetSingleRunTask 标记队列仅执行指定子任务，完成后自动暂停 +func (m *BatchTaskManager) SetSingleRunTask(queueID, taskID string) { + m.mu.Lock() + defer m.mu.Unlock() + if m.singleRunTasks == nil { + m.singleRunTasks = make(map[string]string) + } + m.singleRunTasks[queueID] = taskID +} + +// ClearSingleRunTask 清除单条执行标记 +func (m *BatchTaskManager) ClearSingleRunTask(queueID string) { + m.mu.Lock() + defer m.mu.Unlock() + delete(m.singleRunTasks, queueID) +} + +// TakeSingleRunTaskIfMatch 若刚完成的子任务为单条执行目标，则清除标记并返回 true +func (m *BatchTaskManager) TakeSingleRunTaskIfMatch(queueID, taskID string) bool { + m.mu.Lock() + defer m.mu.Unlock() + if m.singleRunTasks == nil { + return false + } + if m.singleRunTasks[queueID] != taskID { + return false + } + delete(m.singleRunTasks, queueID) + return true +} + +// DeleteTask 删除任务（队列空闲时可删；执行中任务不可删） +func (m *BatchTaskManager) DeleteTask(queueID, taskID string) error { + m.mu.Lock() + defer m.mu.Unlock() + + queue, exists := m.queues[queueID] + if !exists { + return fmt.Errorf("队列不存在") + } + + if !queueAllowsTaskListMutationLocked(queue) { + return fmt.Errorf("队列正在执行或未就绪，无法删除任务") + } + + // 查找任务 + taskIndex := -1 + for i, task := range queue.Tasks { + if task.ID == taskID { + if task.Status == BatchTaskStatusRunning { + return fmt.Errorf("执行中的任务不能删除") + } + taskIndex = i + break + } + } + + if taskIndex == -1 { + return fmt.Errorf("任务不存在") + } + + // DB 优先：先从数据库删除，成功后再从内存移除 + if m.db != nil { + if err := m.db.DeleteBatchTask(queueID, taskID); err != nil { + return fmt.Errorf("删除任务失败: %w", err) + } + } + + queue.Tasks = append(queue.Tasks[:taskIndex], queue.Tasks[taskIndex+1:]...) + return nil +} + +func queueHasRunningTaskLocked(queue *BatchTaskQueue) bool { + if queue == nil { + return false + } + for _, t := range queue.Tasks { + if t != nil && t.Status == BatchTaskStatusRunning { + return true + } + } + return false +} + +// queueAllowsTaskListMutationLocked 是否允许增删改子任务文案/列表（必须在持有 BatchTaskManager.mu 下调用） +func queueAllowsTaskListMutationLocked(queue *BatchTaskQueue) bool { + if queue == nil { + return false + } + if queue.Status == BatchQueueStatusRunning { + return false + } + if queueHasRunningTaskLocked(queue) { + return false + } + switch queue.Status { + case BatchQueueStatusPending, BatchQueueStatusPaused, BatchQueueStatusCompleted, BatchQueueStatusCancelled: + return true + default: + return false + } +} + +// queueAllowsSingleTaskRunLocked 是否允许对指定子任务发起单条执行（必须在持有 BatchTaskManager.mu 下调用） +func queueAllowsSingleTaskRunLocked(queue *BatchTaskQueue, task *BatchTask) bool { + if queue == nil || task == nil { + return false + } + if task.Status == BatchTaskStatusRunning { + return false + } + if queue.Status == BatchQueueStatusRunning { + return false + } + switch queue.Status { + case BatchQueueStatusPending, BatchQueueStatusPaused, BatchQueueStatusCompleted, BatchQueueStatusCancelled: + return true + default: + return false + } +} + +// ClaimNextPendingTask 原子领取下一个待执行子任务（并发 worker 安全）。 +func (m *BatchTaskManager) ClaimNextPendingTask(queueID string) (*BatchTask, bool) { + m.mu.Lock() + defer m.mu.Unlock() + + queue, exists := m.queues[queueID] + if !exists || queue == nil { + return nil, false + } + if queue.Status == BatchQueueStatusCancelled || queue.Status == BatchQueueStatusCompleted || queue.Status == BatchQueueStatusPaused { + return nil, false + } + + onlyTaskID := "" + if m.singleRunTasks != nil { + onlyTaskID = m.singleRunTasks[queueID] + } + + for i, task := range queue.Tasks { + if task == nil || task.Status != BatchTaskStatusPending { + continue + } + if onlyTaskID != "" && task.ID != onlyTaskID { + continue + } + task.Status = BatchTaskStatusRunning + queue.CurrentIndex = i + return task, true + } + return nil, false +} + +// HasRunningTasks 队列是否仍有 running 状态的子任务。 +func (m *BatchTaskManager) HasRunningTasks(queueID string) bool { + m.mu.RLock() + defer m.mu.RUnlock() + queue, exists := m.queues[queueID] + if !exists || queue == nil { + return false + } + for _, task := range queue.Tasks { + if task != nil && task.Status == BatchTaskStatusRunning { + return true + } + } + return false +} + +// HasPendingOrRunningTasks 队列是否仍有未完成的子任务。 +func (m *BatchTaskManager) HasPendingOrRunningTasks(queueID string) bool { + m.mu.RLock() + defer m.mu.RUnlock() + queue, exists := m.queues[queueID] + if !exists || queue == nil { + return false + } + for _, task := range queue.Tasks { + if task == nil { + continue + } + if task.Status == BatchTaskStatusPending || task.Status == BatchTaskStatusRunning { + return true + } + } + return false +} + +// drainTaskCancelsLocked 取出并清空队列下所有子任务取消函数（调用方须已持 m.mu）。 +func (m *BatchTaskManager) drainTaskCancelsLocked(queueID string) []context.CancelFunc { + taskMap, ok := m.taskCancels[queueID] + if !ok || len(taskMap) == 0 { + return nil + } + cancels := make([]context.CancelFunc, 0, len(taskMap)) + for _, c := range taskMap { + if c != nil { + cancels = append(cancels, c) + } + } + delete(m.taskCancels, queueID) + return cancels +} + +// GetNextTask 获取下一个待执行的任务（串行兼容，优先使用 ClaimNextPendingTask） +func (m *BatchTaskManager) GetNextTask(queueID string) (*BatchTask, bool) { + m.mu.Lock() + defer m.mu.Unlock() + + queue, exists := m.queues[queueID] + if !exists { + return nil, false + } + + for i := queue.CurrentIndex; i < len(queue.Tasks); i++ { + task := queue.Tasks[i] + if task.Status == BatchTaskStatusPending { + queue.CurrentIndex = i + return task, true + } + } + + return nil, false +} + +// MoveToNextTask 移动到下一个任务 +func (m *BatchTaskManager) MoveToNextTask(queueID string) { + m.mu.Lock() + defer m.mu.Unlock() + + queue, exists := m.queues[queueID] + if !exists { + return + } + + queue.CurrentIndex++ + + // 同步到数据库 + if m.db != nil { + if err := m.db.UpdateBatchQueueCurrentIndex(queueID, queue.CurrentIndex); err != nil { + m.logger.Warn("batch queue DB index update failed", zap.String("queueId", queueID), zap.Error(err)) + } + } +} + +// SetTaskCancel 设置子任务的取消函数 +func (m *BatchTaskManager) SetTaskCancel(queueID, taskID string, cancel context.CancelFunc) { + m.mu.Lock() + defer m.mu.Unlock() + if cancel == nil { + if taskMap, ok := m.taskCancels[queueID]; ok { + delete(taskMap, taskID) + if len(taskMap) == 0 { + delete(m.taskCancels, queueID) + } + } + return + } + if m.taskCancels[queueID] == nil { + m.taskCancels[queueID] = make(map[string]context.CancelFunc) + } + m.taskCancels[queueID][taskID] = cancel +} + +// PauseQueue 暂停队列 +func (m *BatchTaskManager) PauseQueue(queueID string) bool { + var cancelFuncs []context.CancelFunc + + m.mu.Lock() + queue, exists := m.queues[queueID] + if !exists { + m.mu.Unlock() + return false + } + + if queue.Status != BatchQueueStatusRunning { + m.mu.Unlock() + return false + } + + // DB 优先：先持久化，成功后再更新内存 + if m.db != nil { + if err := m.db.UpdateBatchQueueStatus(queueID, BatchQueueStatusPaused); err != nil { + m.logger.Warn("batch queue DB pause update failed, skipping memory update", + zap.String("queueId", queueID), zap.Error(err)) + m.mu.Unlock() + return false + } + } + + queue.Status = BatchQueueStatusPaused + cancelFuncs = m.drainTaskCancelsLocked(queueID) + m.mu.Unlock() + + for _, c := range cancelFuncs { + c() + } + + return true +} + +// CancelQueue 取消队列（保留此方法以保持向后兼容，但建议使用PauseQueue） +func (m *BatchTaskManager) CancelQueue(queueID string) bool { + now := time.Now() + var cancelFuncs []context.CancelFunc + + m.mu.Lock() + queue, exists := m.queues[queueID] + if !exists { + m.mu.Unlock() + return false + } + + if queue.Status == BatchQueueStatusCompleted || queue.Status == BatchQueueStatusCancelled { + m.mu.Unlock() + return false + } + + // DB 优先：先持久化，成功后再更新内存 + if m.db != nil { + if err := m.db.CancelPendingBatchTasks(queueID, now); err != nil { + m.logger.Warn("batch task DB batch cancel failed, skipping memory update", + zap.String("queueId", queueID), zap.Error(err)) + m.mu.Unlock() + return false + } + if err := m.db.UpdateBatchQueueStatus(queueID, BatchQueueStatusCancelled); err != nil { + m.logger.Warn("batch queue DB cancel update failed, skipping memory update", + zap.String("queueId", queueID), zap.Error(err)) + m.mu.Unlock() + return false + } + } + + queue.Status = BatchQueueStatusCancelled + queue.CompletedAt = &now + + // 内存中批量标记所有 pending 任务为 cancelled + for _, task := range queue.Tasks { + if task.Status == BatchTaskStatusPending { + task.Status = BatchTaskStatusCancelled + task.CompletedAt = &now + } + } + + cancelFuncs = m.drainTaskCancelsLocked(queueID) + m.mu.Unlock() + + for _, c := range cancelFuncs { + c() + } + + return true +} + +// DeleteQueue 删除队列。执行协程活跃或 status 为 running 时拒绝删除，避免 executeBatchQueue 空指针 panic。 +func (m *BatchTaskManager) DeleteQueue(queueID string) error { + m.mu.Lock() + defer m.mu.Unlock() + + queue, exists := m.queues[queueID] + if !exists { + return ErrBatchQueueNotFound + } + + if _, exec := m.queueExecutors[queueID]; exec { + return ErrBatchQueueExecutorActive + } + + // 运行中的队列不允许删除，防止孤儿协程和数据丢失 + if queue.Status == BatchQueueStatusRunning { + return ErrBatchQueueStillRunning + } + + // 清理取消函数 + delete(m.taskCancels, queueID) + + // 从数据库删除 + if m.db != nil { + if err := m.db.DeleteBatchQueue(queueID); err != nil { + m.logger.Warn("batch queue DB delete failed", zap.String("queueId", queueID), zap.Error(err)) + } + } + + delete(m.queues, queueID) + return nil +} + +// generateShortID 生成短ID +func generateShortID() string { + b := make([]byte, 4) + rand.Read(b) + return time.Now().Format("150405") + "-" + hex.EncodeToString(b) +} diff --git a/internal/handler/batch_task_manager_test.go b/internal/handler/batch_task_manager_test.go new file mode 100644 index 00000000..1c41c333 --- /dev/null +++ b/internal/handler/batch_task_manager_test.go @@ -0,0 +1,133 @@ +package handler + +import ( + "errors" + "testing" + + "go.uber.org/zap" +) + +func TestNormalizeBatchQueueConcurrency(t *testing.T) { + if got := normalizeBatchQueueConcurrency(0); got != DefaultBatchQueueConcurrency { + t.Fatalf("expected default %d, got %d", DefaultBatchQueueConcurrency, got) + } + if got := normalizeBatchQueueConcurrency(99); got != MaxBatchQueueConcurrency { + t.Fatalf("expected max %d, got %d", MaxBatchQueueConcurrency, got) + } +} + +func TestClaimNextPendingTaskParallel(t *testing.T) { + m := NewBatchTaskManager(zap.NewNop()) + queue, err := m.CreateBatchQueue("test", "", "eino_single", "manual", "", "", nil, 3, []string{"a", "b", "c"}) + if err != nil { + t.Fatalf("CreateBatchQueue: %v", err) + } + m.UpdateQueueStatus(queue.ID, BatchQueueStatusRunning) + + t1, ok1 := m.ClaimNextPendingTask(queue.ID) + t2, ok2 := m.ClaimNextPendingTask(queue.ID) + if !ok1 || !ok2 || t1.ID == t2.ID { + t.Fatalf("expected two distinct claims, got ok1=%v ok2=%v t1=%v t2=%v", ok1, ok2, t1, t2) + } + if t1.Status != BatchTaskStatusRunning || t2.Status != BatchTaskStatusRunning { + t.Fatalf("claimed tasks should be running") + } + t3, ok3 := m.ClaimNextPendingTask(queue.ID) + if !ok3 { + t.Fatal("expected third claim") + } + _, ok4 := m.ClaimNextPendingTask(queue.ID) + if ok4 { + t.Fatal("expected no fourth pending task") + } + _ = t3 +} + +func TestBatchQueueExecutionShouldStop(t *testing.T) { + t.Parallel() + if !batchQueueExecutionShouldStop(nil, false) { + t.Fatal("expected stop when queue missing") + } + if !batchQueueExecutionShouldStop(nil, true) { + t.Fatal("expected stop when queue is nil but exists=true") + } + q := &BatchTaskQueue{Status: BatchQueueStatusRunning} + if batchQueueExecutionShouldStop(q, true) { + t.Fatal("expected continue when running") + } + q.Status = BatchQueueStatusCancelled + if !batchQueueExecutionShouldStop(q, true) { + t.Fatal("expected stop when cancelled") + } +} + +func TestBatchSubTaskConversationMetaKeepsQueueRole(t *testing.T) { + t.Parallel() + + meta := batchSubTaskConversationMeta(nil, &BatchTaskQueue{Role: " 渗透测试 "}) + if meta.Source != "batch_task" { + t.Fatalf("expected batch_task source, got %q", meta.Source) + } + if meta.RoleName != "渗透测试" { + t.Fatalf("expected queue role to be stored on child conversation, got %q", meta.RoleName) + } +} + +func TestDeleteQueueBlockedWhileExecutorActive(t *testing.T) { + t.Parallel() + m := NewBatchTaskManager(zap.NewNop()) + queue, err := m.CreateBatchQueue("test", "", "eino_single", "manual", "", "", nil, 1, []string{"hello"}) + if err != nil { + t.Fatalf("CreateBatchQueue: %v", err) + } + if !m.TryMarkQueueExecutor(queue.ID) { + t.Fatal("expected to mark executor") + } + m.UpdateQueueStatus(queue.ID, BatchQueueStatusCancelled) + + err = m.DeleteQueue(queue.ID) + if !errors.Is(err, ErrBatchQueueExecutorActive) { + t.Fatalf("expected ErrBatchQueueExecutorActive, got %v", err) + } + if _, ok := m.GetBatchQueue(queue.ID); !ok { + t.Fatal("queue should still exist while executor active") + } + + m.UnmarkQueueExecutor(queue.ID) + if err := m.DeleteQueue(queue.ID); err != nil { + t.Fatalf("expected delete after executor unmarked, got %v", err) + } + if _, ok := m.GetBatchQueue(queue.ID); ok { + t.Fatal("queue should be deleted") + } +} + +func TestDeleteQueueBlockedWhileRunning(t *testing.T) { + t.Parallel() + m := NewBatchTaskManager(zap.NewNop()) + queue, err := m.CreateBatchQueue("test", "", "eino_single", "manual", "", "", nil, 1, []string{"hello"}) + if err != nil { + t.Fatalf("CreateBatchQueue: %v", err) + } + m.UpdateQueueStatus(queue.ID, BatchQueueStatusRunning) + + err = m.DeleteQueue(queue.ID) + if !errors.Is(err, ErrBatchQueueStillRunning) { + t.Fatalf("expected ErrBatchQueueStillRunning, got %v", err) + } +} + +func TestTryMarkQueueExecutorDedupes(t *testing.T) { + t.Parallel() + m := NewBatchTaskManager(zap.NewNop()) + if !m.TryMarkQueueExecutor("q-1") { + t.Fatal("first mark should succeed") + } + if m.TryMarkQueueExecutor("q-1") { + t.Fatal("second mark should fail") + } + m.UnmarkQueueExecutor("q-1") + if !m.TryMarkQueueExecutor("q-1") { + t.Fatal("mark after unmark should succeed") + } +} diff --git a/internal/handler/batch_task_mcp.go b/internal/handler/batch_task_mcp.go new file mode 100644 index 00000000..fbcce38b --- /dev/null +++ b/internal/handler/batch_task_mcp.go @@ -0,0 +1,875 @@ +package handler + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "strconv" + "strings" + "time" + + "cyberstrike-ai/internal/authctx" + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/mcp" + "cyberstrike-ai/internal/mcp/builtin" + + "go.uber.org/zap" +) + +// RegisterBatchTaskMCPTools 注册批量任务队列相关 MCP 工具（需传入已初始化 DB 的 AgentHandler） +func RegisterBatchTaskMCPTools(mcpServer *mcp.Server, h *AgentHandler, logger *zap.Logger) { + if mcpServer == nil || h == nil || logger == nil { + return + } + + reg := func(tool mcp.Tool, fn func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error)) { + mcpServer.RegisterTool(tool, fn) + } + + // --- list --- + reg(mcp.Tool{ + Name: builtin.ToolBatchTaskList, + Description: "列出批量任务队列（精简摘要，省上下文）。含队列元数据、子任务 id/status/截断后的 message、各状态计数。完整子任务（含 result/error/conversationId/时间等）请用 batch_task_get(queue_id)。\n\n⚠️ 调用约束：本工具属于「任务管理」模块，仅当用户明确提及查看/管理批量任务、任务队列时才可调用。不要在用户未要求时自行调用。", + ShortDescription: "列出批量任务队列", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "status": map[string]interface{}{ + "type": "string", + "description": "筛选状态：all（默认）、pending、running、paused、completed、cancelled", + "enum": []string{"all", "pending", "running", "paused", "completed", "cancelled"}, + }, + "keyword": map[string]interface{}{ + "type": "string", + "description": "按队列 ID 或标题模糊搜索", + }, + "page": map[string]interface{}{ + "type": "integer", + "description": "页码，从 1 开始，默认 1", + }, + "page_size": map[string]interface{}{ + "type": "integer", + "description": "每页条数，默认 20，最大 100", + }, + }, + }, + }, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + status := mcpArgString(args, "status") + if status == "" { + status = "all" + } + keyword := mcpArgString(args, "keyword") + page := int(mcpArgFloat(args, "page")) + if page <= 0 { + page = 1 + } + pageSize := int(mcpArgFloat(args, "page_size")) + if pageSize <= 0 { + pageSize = 20 + } + if pageSize > 100 { + pageSize = 100 + } + offset := (page - 1) * pageSize + if offset > 100000 { + offset = 100000 + } + queues := []*BatchTaskQueue{} + total := 0 + var err error + if principal, ok := authctx.PrincipalFromContext(ctx); ok { + queues, total, err = h.batchTaskManager.ListQueuesForAccess(pageSize, offset, status, keyword, principal.UserID, principal.ScopeFor("tasks:read")) + } else { + return batchMCPTextResult("缺少认证身份", true), nil + } + if err != nil { + return batchMCPTextResult(fmt.Sprintf("列出队列失败: %v", err), true), nil + } + totalPages := (total + pageSize - 1) / pageSize + if totalPages == 0 { + totalPages = 1 + } + slim := make([]batchTaskQueueMCPListItem, 0, len(queues)) + for _, q := range queues { + if q == nil { + continue + } + slim = append(slim, toBatchTaskQueueMCPListItem(q)) + } + payload := map[string]interface{}{ + "queues": slim, + "total": total, + "page": page, + "page_size": pageSize, + "total_pages": totalPages, + } + logger.Info("MCP batch_task_list", zap.String("status", status), zap.Int("total", total)) + return batchMCPJSONResult(payload) + }) + + // --- get --- + reg(mcp.Tool{ + Name: builtin.ToolBatchTaskGet, + Description: "根据 queue_id 获取单个批量任务队列详情（含子任务列表、Cron、调度开关与最近错误信息）。\n\n⚠️ 调用约束：本工具属于「任务管理」模块，仅当用户明确提及查看/管理批量任务、任务队列时才可调用。不要在用户未要求时自行调用。", + ShortDescription: "获取批量任务队列详情", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "queue_id": map[string]interface{}{ + "type": "string", + "description": "队列 ID", + }, + }, + "required": []string{"queue_id"}, + }, + }, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + qid := mcpArgString(args, "queue_id") + if qid == "" { + return batchMCPTextResult("queue_id 不能为空", true), nil + } + queue, ok := h.batchTaskManager.GetBatchQueue(qid) + if !ok { + return batchMCPTextResult("队列不存在: "+qid, true), nil + } + return batchMCPJSONResult(queue) + }) + + // --- create --- + reg(mcp.Tool{ + Name: builtin.ToolBatchTaskCreate, + Description: `⚠️ 调用约束：本工具属于「任务管理」模块，仅当用户明确要求创建批量任务、任务队列时才可调用。禁止在用户未提及”批量任务””任务队列””定时任务”等关键词时自行调用。如果用户只是让你做某件事，请在当前对话中直接完成，不要自作主张创建任务队列。 + +【用途】应用内「任务管理 / 批量任务队列」：把多条彼此独立的用户指令登记成一条队列，便于在界面里查看进度、暂停/继续、定时重跑等。这是队列数据与调度入口，不是再开一个”子代理会话”替你探索当前问题。 + +【何时用】用户明确要批量排队执行、Cron 周期跑同一批指令、或需要与任务管理页面对齐时调用。需要即时追问、强依赖当前对话上下文的分析/编码，应在本对话内直接完成，不要为了”委派”而创建队列。 + +【参数】tasks（字符串数组）或 tasks_text（多行，每行一条）二选一；每项是一条将来由系统按队列顺序执行的指令文案。agent_mode：eino_single（Eino ADK 单代理，默认）、deep / plan_execute / supervisor（需系统启用多代理）。非”把主对话拆给子代理”。schedule_mode：manual（默认）或 cron；cron 须填 cron_expr（5 段，如 “0 */6 * * *”）。 + +【执行】默认创建后为 pending，不自动跑。execute_now=true 可创建后立即跑；否则之后调用 batch_task_start。Cron 自动下一轮需 schedule_enabled 为 true（可用 batch_task_schedule_enabled）。`, + ShortDescription: "任务管理：创建批量任务队列（登记多条指令，可选立即或 Cron）", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "title": map[string]interface{}{ + "type": "string", + "description": "可选队列标题，便于在任务管理中识别", + }, + "role": map[string]interface{}{ + "type": "string", + "description": "队列使用的角色名，空表示默认", + }, + "tasks": map[string]interface{}{ + "type": "array", + "description": "队列中的子任务指令，每项一条独立待执行文案（与 tasks_text 二选一）", + "items": map[string]interface{}{"type": "string"}, + }, + "tasks_text": map[string]interface{}{ + "type": "string", + "description": "多行文本，每行一条子任务指令（与 tasks 二选一）", + }, + "agent_mode": map[string]interface{}{ + "type": "string", + "description": "执行模式：eino_single（Eino ADK，默认）、deep/plan_execute/supervisor（Eino 编排，需启用多代理）", + "enum": []string{"eino_single", "deep", "plan_execute", "supervisor"}, + }, + "schedule_mode": map[string]interface{}{ + "type": "string", + "description": "manual（仅手工/启动后跑）或 cron（按表达式触发）", + "enum": []string{"manual", "cron"}, + }, + "cron_expr": map[string]interface{}{ + "type": "string", + "description": "schedule_mode 为 cron 时必填。标准 5 段：分钟 小时 日 月 星期，例如 \"0 */6 * * *\"、\"30 2 * * 1-5\"", + }, + "execute_now": map[string]interface{}{ + "type": "boolean", + "description": "创建后是否立即开始执行队列，默认 false（pending，需 batch_task_start）", + }, + "project_id": map[string]interface{}{ + "type": "string", + "description": "队列内子对话绑定的项目 ID（可选，未指定时使用 config.project.default_project_id）", + }, + "concurrency": map[string]interface{}{ + "type": "integer", + "description": "同时执行的子任务数，默认 1（串行），最大 8。含扫描类工具时建议 1-2。", + }, + }, + }, + }, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + tasks, errMsg := batchMCPTasksFromArgs(args) + if errMsg != "" { + return batchMCPTextResult(errMsg, true), nil + } + title := mcpArgString(args, "title") + role := mcpArgString(args, "role") + agentMode := config.NormalizeAgentMode(mcpArgString(args, "agent_mode")) + scheduleMode := normalizeBatchQueueScheduleMode(mcpArgString(args, "schedule_mode")) + cronExpr := strings.TrimSpace(mcpArgString(args, "cron_expr")) + var nextRunAt *time.Time + if scheduleMode == "cron" { + if cronExpr == "" { + return batchMCPTextResult("Cron 调度模式下 cron_expr 不能为空", true), nil + } + sch, err := h.batchCronParser.Parse(cronExpr) + if err != nil { + return batchMCPTextResult("无效的 Cron 表达式: "+err.Error(), true), nil + } + n := sch.Next(time.Now()) + nextRunAt = &n + } + executeNow, ok := mcpArgBool(args, "execute_now") + if !ok { + executeNow = false + } + projectID := strings.TrimSpace(mcpArgString(args, "project_id")) + if principal, ok := authctx.PrincipalFromContext(ctx); ok && projectID != "" && principal.ScopeFor("tasks:write") != database.RBACScopeAll { + if h.db == nil || !h.db.UserCanAccessResource(principal.UserID, principal.ScopeFor("tasks:write"), "project", projectID) { + return batchMCPTextResult("无权访问目标项目", true), nil + } + } + concurrency := int(mcpArgFloat(args, "concurrency")) + queue, createErr := h.batchTaskManager.CreateBatchQueue(title, role, agentMode, scheduleMode, cronExpr, projectID, nextRunAt, concurrency, tasks) + if createErr != nil { + return batchMCPTextResult("创建队列失败: "+createErr.Error(), true), nil + } + if principal, ok := authctx.PrincipalFromContext(ctx); ok && h.db != nil { + _ = h.db.SetResourceOwner("batch_task", queue.ID, principal.UserID) + _ = h.db.AssignResourceToUser(principal.UserID, "batch_task", queue.ID) + } + started := false + if executeNow { + ok, err := h.startBatchQueueExecution(queue.ID, false) + if !ok { + return batchMCPTextResult("队列不存在: "+queue.ID, true), nil + } + if err != nil { + return batchMCPTextResult("创建成功但启动失败: "+err.Error(), true), nil + } + started = true + if refreshed, exists := h.batchTaskManager.GetBatchQueue(queue.ID); exists { + queue = refreshed + } + } + logger.Info("MCP batch_task_create", zap.String("queueId", queue.ID), zap.Int("taskCount", len(tasks))) + return batchMCPJSONResult(map[string]interface{}{ + "queue_id": queue.ID, + "queue": queue, + "started": started, + "execute_now": executeNow, + "reminder": func() string { + if started { + return "队列已创建并立即启动。" + } + return "队列已创建，当前为 pending。需要开始执行时请调用 MCP 工具 batch_task_start（queue_id 同上）。Cron 自动调度需 schedule_enabled 为 true，可用 batch_task_schedule_enabled。" + }(), + }) + }) + + // --- start --- + reg(mcp.Tool{ + Name: builtin.ToolBatchTaskStart, + Description: `启动或继续执行批量任务队列（pending / paused）。 +与 batch_task_create 配合使用：仅创建队列不会自动执行，需调用本工具才会开始跑子任务。 + +⚠️ 调用约束：本工具属于「任务管理」模块，仅当用户明确要求启动/继续批量任务时才可调用。不要在用户未要求时自行调用。`, + ShortDescription: "启动/继续批量任务队列（创建后需调用才会执行）", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "queue_id": map[string]interface{}{ + "type": "string", + "description": "队列 ID", + }, + }, + "required": []string{"queue_id"}, + }, + }, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + qid := mcpArgString(args, "queue_id") + if qid == "" { + return batchMCPTextResult("queue_id 不能为空", true), nil + } + ok, err := h.startBatchQueueExecution(qid, false) + if !ok { + return batchMCPTextResult("队列不存在: "+qid, true), nil + } + if err != nil { + return batchMCPTextResult("启动失败: "+err.Error(), true), nil + } + logger.Info("MCP batch_task_start", zap.String("queueId", qid)) + return batchMCPTextResult("已提交启动，队列将开始执行。", false), nil + }) + + // --- rerun (reset + start for completed/cancelled queues) --- + reg(mcp.Tool{ + Name: builtin.ToolBatchTaskRerun, + Description: "重跑已完成或已取消的批量任务队列。会重置所有子任务状态后重新执行一轮。\n\n⚠️ 调用约束：本工具属于「任务管理」模块，仅当用户明确要求重跑批量任务时才可调用。不要在用户未要求时自行调用。", + ShortDescription: "重跑批量任务队列", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "queue_id": map[string]interface{}{ + "type": "string", + "description": "队列 ID", + }, + }, + "required": []string{"queue_id"}, + }, + }, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + qid := mcpArgString(args, "queue_id") + if qid == "" { + return batchMCPTextResult("queue_id 不能为空", true), nil + } + queue, exists := h.batchTaskManager.GetBatchQueue(qid) + if !exists { + return batchMCPTextResult("队列不存在: "+qid, true), nil + } + if queue.Status != "completed" && queue.Status != "cancelled" { + return batchMCPTextResult("仅已完成或已取消的队列可以重跑，当前状态: "+queue.Status, true), nil + } + if !h.batchTaskManager.ResetQueueForRerun(qid) { + return batchMCPTextResult("重置队列失败", true), nil + } + ok, err := h.startBatchQueueExecution(qid, false) + if !ok { + return batchMCPTextResult("启动失败", true), nil + } + if err != nil { + return batchMCPTextResult("启动失败: "+err.Error(), true), nil + } + logger.Info("MCP batch_task_rerun", zap.String("queueId", qid)) + return batchMCPTextResult("已重置并重新启动队列。", false), nil + }) + + // --- pause --- + reg(mcp.Tool{ + Name: builtin.ToolBatchTaskPause, + Description: "暂停正在运行的批量任务队列（当前子任务会被取消）。\n\n⚠️ 调用约束：本工具属于「任务管理」模块，仅当用户明确要求暂停批量任务时才可调用。不要在用户未要求时自行调用。", + ShortDescription: "暂停批量任务队列", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "queue_id": map[string]interface{}{ + "type": "string", + "description": "队列 ID", + }, + }, + "required": []string{"queue_id"}, + }, + }, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + qid := mcpArgString(args, "queue_id") + if qid == "" { + return batchMCPTextResult("queue_id 不能为空", true), nil + } + if !h.batchTaskManager.PauseQueue(qid) { + return batchMCPTextResult("无法暂停：队列不存在或当前非 running 状态", true), nil + } + logger.Info("MCP batch_task_pause", zap.String("queueId", qid)) + return batchMCPTextResult("队列已暂停。", false), nil + }) + + // --- delete queue --- + reg(mcp.Tool{ + Name: builtin.ToolBatchTaskDelete, + Description: "删除批量任务队列及其子任务记录。\n\n⚠️ 调用约束：本工具属于「任务管理」模块，仅当用户明确要求删除批量任务队列时才可调用。不要在用户未要求时自行调用。", + ShortDescription: "删除批量任务队列", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "queue_id": map[string]interface{}{ + "type": "string", + "description": "队列 ID", + }, + }, + "required": []string{"queue_id"}, + }, + }, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + qid := mcpArgString(args, "queue_id") + if qid == "" { + return batchMCPTextResult("queue_id 不能为空", true), nil + } + if err := h.batchTaskManager.DeleteQueue(qid); err != nil { + switch { + case errors.Is(err, ErrBatchQueueNotFound): + return batchMCPTextResult("删除失败：队列不存在", true), nil + case errors.Is(err, ErrBatchQueueExecutorActive): + return batchMCPTextResult("删除失败：队列执行器仍在运行，请稍后再试", true), nil + case errors.Is(err, ErrBatchQueueStillRunning): + return batchMCPTextResult("删除失败：队列正在运行中", true), nil + default: + return batchMCPTextResult("删除失败："+err.Error(), true), nil + } + } + logger.Info("MCP batch_task_delete", zap.String("queueId", qid)) + return batchMCPTextResult("队列已删除。", false), nil + }) + + // --- update metadata (title/role/agentMode) --- + reg(mcp.Tool{ + Name: builtin.ToolBatchTaskUpdateMetadata, + Description: "修改批量任务队列的标题、角色和代理模式。仅在队列非 running 状态下可修改。\n\n⚠️ 调用约束：本工具属于「任务管理」模块，仅当用户明确要求修改批量任务队列属性时才可调用。不要在用户未要求时自行调用。", + ShortDescription: "修改批量任务队列标题/角色/代理模式", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "queue_id": map[string]interface{}{ + "type": "string", + "description": "队列 ID", + }, + "title": map[string]interface{}{ + "type": "string", + "description": "新标题（空字符串清除标题）", + }, + "role": map[string]interface{}{ + "type": "string", + "description": "新角色名（空字符串使用默认角色）", + }, + "agent_mode": map[string]interface{}{ + "type": "string", + "description": "代理模式：eino_single、deep、plan_execute、supervisor", + "enum": []string{"eino_single", "deep", "plan_execute", "supervisor"}, + }, + "concurrency": map[string]interface{}{ + "type": "integer", + "description": "同时执行的子任务数，默认 1，最大 8", + }, + }, + "required": []string{"queue_id"}, + }, + }, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + qid := mcpArgString(args, "queue_id") + if qid == "" { + return batchMCPTextResult("queue_id 不能为空", true), nil + } + title := mcpArgString(args, "title") + role := mcpArgString(args, "role") + agentMode := mcpArgString(args, "agent_mode") + var concurrency *int + if raw, ok := args["concurrency"]; ok && raw != nil { + v := int(mcpArgFloat(args, "concurrency")) + concurrency = &v + } + if err := h.batchTaskManager.UpdateQueueMetadata(qid, title, role, agentMode, concurrency); err != nil { + return batchMCPTextResult(err.Error(), true), nil + } + updated, _ := h.batchTaskManager.GetBatchQueue(qid) + logger.Info("MCP batch_task_update_metadata", zap.String("queueId", qid)) + return batchMCPJSONResult(updated) + }) + + // --- update schedule --- + reg(mcp.Tool{ + Name: builtin.ToolBatchTaskUpdateSchedule, + Description: `修改批量任务队列的调度方式和 Cron 表达式。仅在队列非 running 状态下可修改。 +schedule_mode 为 cron 时必须提供有效 cron_expr；为 manual 时会清除 Cron 配置。 + +⚠️ 调用约束：本工具属于「任务管理」模块，仅当用户明确要求修改批量任务调度配置时才可调用。不要在用户未要求时自行调用。`, + ShortDescription: "修改批量任务调度配置（Cron 表达式）", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "queue_id": map[string]interface{}{ + "type": "string", + "description": "队列 ID", + }, + "schedule_mode": map[string]interface{}{ + "type": "string", + "description": "manual 或 cron", + "enum": []string{"manual", "cron"}, + }, + "cron_expr": map[string]interface{}{ + "type": "string", + "description": "Cron 表达式（schedule_mode 为 cron 时必填）。标准 5 段格式：分钟 小时 日 月 星期，如 \"0 */6 * * *\"（每6小时）、\"30 2 * * 1-5\"（工作日凌晨2:30）", + }, + }, + "required": []string{"queue_id", "schedule_mode"}, + }, + }, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + qid := mcpArgString(args, "queue_id") + if qid == "" { + return batchMCPTextResult("queue_id 不能为空", true), nil + } + queue, exists := h.batchTaskManager.GetBatchQueue(qid) + if !exists { + return batchMCPTextResult("队列不存在: "+qid, true), nil + } + if queue.Status == "running" { + return batchMCPTextResult("队列正在运行中，无法修改调度配置", true), nil + } + scheduleMode := normalizeBatchQueueScheduleMode(mcpArgString(args, "schedule_mode")) + cronExpr := strings.TrimSpace(mcpArgString(args, "cron_expr")) + var nextRunAt *time.Time + if scheduleMode == "cron" { + if cronExpr == "" { + return batchMCPTextResult("Cron 调度模式下 cron_expr 不能为空", true), nil + } + sch, err := h.batchCronParser.Parse(cronExpr) + if err != nil { + return batchMCPTextResult("无效的 Cron 表达式: "+err.Error(), true), nil + } + n := sch.Next(time.Now()) + nextRunAt = &n + } + h.batchTaskManager.UpdateQueueSchedule(qid, scheduleMode, cronExpr, nextRunAt) + updated, _ := h.batchTaskManager.GetBatchQueue(qid) + logger.Info("MCP batch_task_update_schedule", zap.String("queueId", qid), zap.String("scheduleMode", scheduleMode), zap.String("cronExpr", cronExpr)) + return batchMCPJSONResult(updated) + }) + + // --- schedule enabled --- + reg(mcp.Tool{ + Name: builtin.ToolBatchTaskScheduleEnabled, + Description: `设置是否允许 Cron 自动触发该队列。关闭后仍保留 Cron 表达式，仅停止定时自动跑；可用手工「启动」执行。 +仅对 schedule_mode 为 cron 的队列有意义。 + +⚠️ 调用约束：本工具属于「任务管理」模块，仅当用户明确要求开关批量任务自动调度时才可调用。不要在用户未要求时自行调用。`, + ShortDescription: "开关批量任务 Cron 自动调度", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "queue_id": map[string]interface{}{ + "type": "string", + "description": "队列 ID", + }, + "schedule_enabled": map[string]interface{}{ + "type": "boolean", + "description": "true 允许定时触发，false 仅手工执行", + }, + }, + "required": []string{"queue_id", "schedule_enabled"}, + }, + }, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + qid := mcpArgString(args, "queue_id") + if qid == "" { + return batchMCPTextResult("queue_id 不能为空", true), nil + } + en, ok := mcpArgBool(args, "schedule_enabled") + if !ok { + return batchMCPTextResult("schedule_enabled 必须为布尔值", true), nil + } + if _, exists := h.batchTaskManager.GetBatchQueue(qid); !exists { + return batchMCPTextResult("队列不存在", true), nil + } + if !h.batchTaskManager.SetScheduleEnabled(qid, en) { + return batchMCPTextResult("更新失败", true), nil + } + queue, _ := h.batchTaskManager.GetBatchQueue(qid) + logger.Info("MCP batch_task_schedule_enabled", zap.String("queueId", qid), zap.Bool("enabled", en)) + return batchMCPJSONResult(queue) + }) + + // --- add task --- + reg(mcp.Tool{ + Name: builtin.ToolBatchTaskAdd, + Description: "向处于 pending 状态的队列追加一条子任务。\n\n⚠️ 调用约束：本工具属于「任务管理」模块，仅当用户明确要求向批量任务队列添加子任务时才可调用。不要在用户未要求时自行调用。", + ShortDescription: "批量队列添加子任务", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "queue_id": map[string]interface{}{ + "type": "string", + "description": "队列 ID", + }, + "message": map[string]interface{}{ + "type": "string", + "description": "任务指令内容", + }, + }, + "required": []string{"queue_id", "message"}, + }, + }, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + qid := mcpArgString(args, "queue_id") + msg := strings.TrimSpace(mcpArgString(args, "message")) + if qid == "" || msg == "" { + return batchMCPTextResult("queue_id 与 message 均不能为空", true), nil + } + task, err := h.batchTaskManager.AddTaskToQueue(qid, msg) + if err != nil { + return batchMCPTextResult(err.Error(), true), nil + } + queue, _ := h.batchTaskManager.GetBatchQueue(qid) + logger.Info("MCP batch_task_add_task", zap.String("queueId", qid), zap.String("taskId", task.ID)) + return batchMCPJSONResult(map[string]interface{}{"task": task, "queue": queue}) + }) + + // --- update task --- + reg(mcp.Tool{ + Name: builtin.ToolBatchTaskUpdate, + Description: "修改 pending 队列中仍为 pending 的子任务文案。\n\n⚠️ 调用约束：本工具属于「任务管理」模块，仅当用户明确要求修改批量子任务内容时才可调用。不要在用户未要求时自行调用。", + ShortDescription: "更新批量子任务内容", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "queue_id": map[string]interface{}{ + "type": "string", + "description": "队列 ID", + }, + "task_id": map[string]interface{}{ + "type": "string", + "description": "子任务 ID", + }, + "message": map[string]interface{}{ + "type": "string", + "description": "新的任务指令", + }, + }, + "required": []string{"queue_id", "task_id", "message"}, + }, + }, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + qid := mcpArgString(args, "queue_id") + tid := mcpArgString(args, "task_id") + msg := strings.TrimSpace(mcpArgString(args, "message")) + if qid == "" || tid == "" || msg == "" { + return batchMCPTextResult("queue_id、task_id、message 均不能为空", true), nil + } + if err := h.batchTaskManager.UpdateTaskMessage(qid, tid, msg); err != nil { + return batchMCPTextResult(err.Error(), true), nil + } + queue, _ := h.batchTaskManager.GetBatchQueue(qid) + logger.Info("MCP batch_task_update_task", zap.String("queueId", qid), zap.String("taskId", tid)) + return batchMCPJSONResult(queue) + }) + + // --- remove task --- + reg(mcp.Tool{ + Name: builtin.ToolBatchTaskRemove, + Description: "从 pending 队列中删除仍为 pending 的子任务。\n\n⚠️ 调用约束：本工具属于「任务管理」模块，仅当用户明确要求删除批量子任务时才可调用。不要在用户未要求时自行调用。", + ShortDescription: "删除批量子任务", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "queue_id": map[string]interface{}{ + "type": "string", + "description": "队列 ID", + }, + "task_id": map[string]interface{}{ + "type": "string", + "description": "子任务 ID", + }, + }, + "required": []string{"queue_id", "task_id"}, + }, + }, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + qid := mcpArgString(args, "queue_id") + tid := mcpArgString(args, "task_id") + if qid == "" || tid == "" { + return batchMCPTextResult("queue_id 与 task_id 均不能为空", true), nil + } + if err := h.batchTaskManager.DeleteTask(qid, tid); err != nil { + return batchMCPTextResult(err.Error(), true), nil + } + queue, _ := h.batchTaskManager.GetBatchQueue(qid) + logger.Info("MCP batch_task_remove_task", zap.String("queueId", qid), zap.String("taskId", tid)) + return batchMCPJSONResult(queue) + }) + + logger.Debug("批量任务 MCP 工具已注册", zap.Int("count", 12)) +} + +// --- batch_task_list 精简结构（避免把每条子任务的 result 等大段文本塞进列表上下文） --- + +const mcpBatchListTaskMessageMaxRunes = 160 + +// batchTaskMCPListSummary 列表中的子任务摘要（完整字段用 batch_task_get） +type batchTaskMCPListSummary struct { + ID string `json:"id"` + Status string `json:"status"` + Message string `json:"message,omitempty"` +} + +// batchTaskQueueMCPListItem 列表中的队列摘要 +type batchTaskQueueMCPListItem struct { + ID string `json:"id"` + Title string `json:"title,omitempty"` + Role string `json:"role,omitempty"` + AgentMode string `json:"agentMode"` + ScheduleMode string `json:"scheduleMode"` + CronExpr string `json:"cronExpr,omitempty"` + NextRunAt *time.Time `json:"nextRunAt,omitempty"` + ScheduleEnabled bool `json:"scheduleEnabled"` + LastScheduleTriggerAt *time.Time `json:"lastScheduleTriggerAt,omitempty"` + Status string `json:"status"` + CreatedAt time.Time `json:"createdAt"` + StartedAt *time.Time `json:"startedAt,omitempty"` + CompletedAt *time.Time `json:"completedAt,omitempty"` + CurrentIndex int `json:"currentIndex"` + Concurrency int `json:"concurrency"` + TaskTotal int `json:"task_total"` + TaskCounts map[string]int `json:"task_counts"` + Tasks []batchTaskMCPListSummary `json:"tasks"` +} + +func truncateStringRunes(s string, maxRunes int) string { + if maxRunes <= 0 { + return "" + } + n := 0 + for i := range s { + if n == maxRunes { + out := strings.TrimSpace(s[:i]) + if out == "" { + return "…" + } + return out + "…" + } + n++ + } + return s +} + +const mcpBatchListMaxTasksPerQueue = 200 // 列表中每个队列最多返回的子任务摘要数 + +func toBatchTaskQueueMCPListItem(q *BatchTaskQueue) batchTaskQueueMCPListItem { + counts := map[string]int{ + "pending": 0, + "running": 0, + "completed": 0, + "failed": 0, + "cancelled": 0, + } + tasks := make([]batchTaskMCPListSummary, 0, len(q.Tasks)) + for _, t := range q.Tasks { + if t == nil { + continue + } + counts[t.Status]++ + // 列表视图限制子任务摘要数量，完整列表通过 batch_task_get 查看 + if len(tasks) < mcpBatchListMaxTasksPerQueue { + tasks = append(tasks, batchTaskMCPListSummary{ + ID: t.ID, + Status: t.Status, + Message: truncateStringRunes(t.Message, mcpBatchListTaskMessageMaxRunes), + }) + } + } + return batchTaskQueueMCPListItem{ + ID: q.ID, + Title: q.Title, + Role: q.Role, + AgentMode: q.AgentMode, + ScheduleMode: q.ScheduleMode, + CronExpr: q.CronExpr, + NextRunAt: q.NextRunAt, + ScheduleEnabled: q.ScheduleEnabled, + LastScheduleTriggerAt: q.LastScheduleTriggerAt, + Status: q.Status, + CreatedAt: q.CreatedAt, + StartedAt: q.StartedAt, + CompletedAt: q.CompletedAt, + CurrentIndex: q.CurrentIndex, + Concurrency: q.Concurrency, + TaskTotal: len(tasks), + TaskCounts: counts, + Tasks: tasks, + } +} + +func batchMCPTextResult(text string, isErr bool) *mcp.ToolResult { + return &mcp.ToolResult{ + Content: []mcp.Content{{Type: "text", Text: text}}, + IsError: isErr, + } +} + +func batchMCPJSONResult(v interface{}) (*mcp.ToolResult, error) { + b, err := json.MarshalIndent(v, "", " ") + if err != nil { + return batchMCPTextResult(fmt.Sprintf("JSON 编码失败: %v", err), true), nil + } + return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: string(b)}}}, nil +} + +func batchMCPTasksFromArgs(args map[string]interface{}) ([]string, string) { + if raw, ok := args["tasks"]; ok && raw != nil { + switch t := raw.(type) { + case []interface{}: + out := make([]string, 0, len(t)) + for _, x := range t { + if s, ok := x.(string); ok { + if tr := strings.TrimSpace(s); tr != "" { + out = append(out, tr) + } + } + } + if len(out) > 0 { + return out, "" + } + } + } + if txt := mcpArgString(args, "tasks_text"); txt != "" { + lines := strings.Split(txt, "\n") + out := make([]string, 0, len(lines)) + for _, line := range lines { + if tr := strings.TrimSpace(line); tr != "" { + out = append(out, tr) + } + } + if len(out) > 0 { + return out, "" + } + } + return nil, "需要提供 tasks（字符串数组）或 tasks_text（多行文本，每行一条任务）" +} + +func mcpArgString(args map[string]interface{}, key string) string { + v, ok := args[key] + if !ok || v == nil { + return "" + } + switch t := v.(type) { + case string: + return strings.TrimSpace(t) + case float64: + return strings.TrimSpace(strconv.FormatFloat(t, 'f', -1, 64)) + case json.Number: + return strings.TrimSpace(t.String()) + default: + return strings.TrimSpace(fmt.Sprint(t)) + } +} + +func mcpArgFloat(args map[string]interface{}, key string) float64 { + v, ok := args[key] + if !ok || v == nil { + return 0 + } + switch t := v.(type) { + case float64: + return t + case int: + return float64(t) + case int64: + return float64(t) + case json.Number: + f, _ := t.Float64() + return f + case string: + f, _ := strconv.ParseFloat(strings.TrimSpace(t), 64) + return f + default: + return 0 + } +} + +func mcpArgBool(args map[string]interface{}, key string) (val bool, ok bool) { + v, exists := args[key] + if !exists { + return false, false + } + switch t := v.(type) { + case bool: + return t, true + case string: + s := strings.ToLower(strings.TrimSpace(t)) + if s == "true" || s == "1" || s == "yes" { + return true, true + } + if s == "false" || s == "0" || s == "no" { + return false, true + } + case float64: + return t != 0, true + } + return false, false +} diff --git a/internal/handler/c2.go b/internal/handler/c2.go new file mode 100644 index 00000000..aa7d7fb9 --- /dev/null +++ b/internal/handler/c2.go @@ -0,0 +1,1237 @@ +package handler + +import ( + "database/sql" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "os" + "path/filepath" + "strconv" + "strings" + "sync/atomic" + "time" + + "cyberstrike-ai/internal/audit" + "cyberstrike-ai/internal/c2" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + + "github.com/gin-gonic/gin" + "github.com/google/uuid" + "go.uber.org/zap" +) + +// C2Handler 处理 C2 相关的 REST API（manager 可在运行时置 nil 以关闭 C2） +type C2Handler struct { + mgrPtr atomic.Pointer[c2.Manager] + logger *zap.Logger + audit *audit.Service +} + +// SetAudit wires platform audit logging. +func (h *C2Handler) SetAudit(s *audit.Service) { + h.audit = s +} + +// NewC2Handler 创建 C2 处理器；manager 可为 nil（功能关闭时） +func NewC2Handler(manager *c2.Manager, logger *zap.Logger) *C2Handler { + h := &C2Handler{logger: logger} + if manager != nil { + h.mgrPtr.Store(manager) + } + return h +} + +func (h *C2Handler) mgr() *c2.Manager { + return h.mgrPtr.Load() +} + +// SetManager 运行时切换或清空 C2 Manager（与 App 启停同步） +func (h *C2Handler) SetManager(m *c2.Manager) { + h.mgrPtr.Store(m) +} + +// ============================================================================ +// 监听器 API +// ============================================================================ + +// ListListeners 获取监听器列表 +func (h *C2Handler) ListListeners(c *gin.Context) { + listeners, err := h.mgr().DB().ListC2ListenersForAccess(c2AccessFromContext(c), c.Query("project_id")) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + // 移除敏感字段 + for _, l := range listeners { + l.EncryptionKey = "" + l.ImplantToken = "" + } + c.JSON(http.StatusOK, gin.H{"listeners": listeners}) +} + +// CreateListener 创建监听器 +func (h *C2Handler) CreateListener(c *gin.Context) { + var req struct { + Name string `json:"name"` + ProjectID string `json:"project_id,omitempty"` + Type string `json:"type"` + BindHost string `json:"bind_host"` + BindPort int `json:"bind_port"` + ProfileID string `json:"profile_id,omitempty"` + Remark string `json:"remark,omitempty"` + CallbackHost string `json:"callback_host,omitempty"` + Config *c2.ListenerConfig `json:"config,omitempty"` + } + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + input := c2.CreateListenerInput{ + Name: req.Name, + ProjectID: req.ProjectID, + Type: req.Type, + BindHost: req.BindHost, + BindPort: req.BindPort, + ProfileID: req.ProfileID, + Remark: req.Remark, + Config: req.Config, + CallbackHost: strings.TrimSpace(req.CallbackHost), + } + if !h.canAccessProject(c, input.ProjectID) { + c.JSON(http.StatusForbidden, gin.H{"error": "project access denied"}) + return + } + + listener, err := h.mgr().CreateListener(input) + if err != nil { + code := http.StatusInternalServerError + if e, ok := err.(*c2.CommonError); ok { + code = e.HTTP + } + c.JSON(code, gin.H{"error": err.Error()}) + return + } + if session, ok := security.CurrentSession(c); ok { + listener.OwnerUserID = session.UserID + _ = h.mgr().DB().SetResourceOwner("c2_listener", listener.ID, session.UserID) + _ = h.mgr().DB().AssignResourceToUser(session.UserID, "c2_listener", listener.ID) + } + implantToken := listener.ImplantToken + listener.EncryptionKey = "" + listener.ImplantToken = "" + if h.audit != nil { + h.audit.RecordOK(c, "c2", "listener_create", "创建 C2 监听器", "c2_listener", listener.ID, map[string]interface{}{ + "name": listener.Name, "bind": listener.BindHost, "port": listener.BindPort, + }) + } + c.JSON(http.StatusOK, gin.H{"listener": listener, "implant_token": implantToken}) +} + +// GetListener 获取单个监听器 +func (h *C2Handler) GetListener(c *gin.Context) { + id := c.Param("id") + listener, err := h.mgr().DB().GetC2Listener(id) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if listener == nil { + c.JSON(http.StatusNotFound, gin.H{"error": "listener not found"}) + return + } + listener.EncryptionKey = "" + listener.ImplantToken = "" + c.JSON(http.StatusOK, gin.H{"listener": listener}) +} + +// UpdateListener 更新监听器 +func (h *C2Handler) UpdateListener(c *gin.Context) { + id := c.Param("id") + listener, err := h.mgr().DB().GetC2Listener(id) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if listener == nil { + c.JSON(http.StatusNotFound, gin.H{"error": "listener not found"}) + return + } + + var req struct { + Name string `json:"name"` + ProjectID string `json:"project_id"` + BindHost string `json:"bind_host"` + BindPort int `json:"bind_port"` + ProfileID string `json:"profile_id"` + Remark string `json:"remark"` + CallbackHost *string `json:"callback_host"` + Config *c2.ListenerConfig `json:"config,omitempty"` + } + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + // 若监听器在运行，不能修改关键字段 + if h.mgr().IsListenerRunning(id) { + if req.BindHost != listener.BindHost || req.BindPort != listener.BindPort { + c.JSON(http.StatusConflict, gin.H{"error": "cannot modify bind address while listener is running"}) + return + } + } + + listener.Name = req.Name + listener.ProjectID = strings.TrimSpace(req.ProjectID) + listener.BindHost = req.BindHost + listener.BindPort = req.BindPort + listener.ProfileID = req.ProfileID + listener.Remark = req.Remark + if req.Config != nil { + cfgJSON, _ := json.Marshal(req.Config) + listener.ConfigJSON = string(cfgJSON) + } + if !h.canAccessProject(c, listener.ProjectID) { + c.JSON(http.StatusForbidden, gin.H{"error": "project access denied"}) + return + } + if req.CallbackHost != nil { + cfg := &c2.ListenerConfig{} + raw := strings.TrimSpace(listener.ConfigJSON) + if raw == "" { + raw = "{}" + } + _ = json.Unmarshal([]byte(raw), cfg) + cfg.CallbackHost = strings.TrimSpace(*req.CallbackHost) + cfg.ApplyDefaults() + cfgJSON, err := json.Marshal(cfg) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + listener.ConfigJSON = string(cfgJSON) + } + + if err := h.mgr().DB().UpdateC2Listener(listener); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + listener.EncryptionKey = "" + listener.ImplantToken = "" + c.JSON(http.StatusOK, gin.H{"listener": listener}) +} + +// DeleteListener 删除监听器 +func (h *C2Handler) DeleteListener(c *gin.Context) { + id := c.Param("id") + if err := h.mgr().DeleteListener(id); err != nil { + code := http.StatusInternalServerError + if e, ok := err.(*c2.CommonError); ok { + code = e.HTTP + } + c.JSON(code, gin.H{"error": err.Error()}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "c2", "listener_delete", "删除 C2 监听器", "c2_listener", id, nil) + } + c.JSON(http.StatusOK, gin.H{"deleted": true}) +} + +// StartListener 启动监听器 +func (h *C2Handler) StartListener(c *gin.Context) { + id := c.Param("id") + listener, err := h.mgr().StartListener(id) + if err != nil { + code := http.StatusInternalServerError + if e, ok := err.(*c2.CommonError); ok { + code = e.HTTP + } + c.JSON(code, gin.H{"error": err.Error()}) + return + } + listener.EncryptionKey = "" + listener.ImplantToken = "" + if h.audit != nil { + h.audit.RecordOK(c, "c2", "listener_start", "启动 C2 监听器", "c2_listener", id, nil) + } + c.JSON(http.StatusOK, gin.H{"listener": listener}) +} + +// StopListener 停止监听器 +func (h *C2Handler) StopListener(c *gin.Context) { + id := c.Param("id") + if err := h.mgr().StopListener(id); err != nil { + code := http.StatusInternalServerError + if e, ok := err.(*c2.CommonError); ok { + code = e.HTTP + } + c.JSON(code, gin.H{"error": err.Error()}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "c2", "listener_stop", "停止 C2 监听器", "c2_listener", id, nil) + } + c.JSON(http.StatusOK, gin.H{"stopped": true}) +} + +// ============================================================================ +// 会话 API +// ============================================================================ + +// ListSessions 获取会话列表 +func (h *C2Handler) ListSessions(c *gin.Context) { + filter := database.ListC2SessionsFilter{ + ListenerID: c.Query("listener_id"), + ProjectID: c.Query("project_id"), + Status: c.Query("status"), + OS: c.Query("os"), + Search: c.Query("search"), + } + if limit := c.Query("limit"); limit != "" { + if n, err := strconv.Atoi(limit); err == nil && n > 0 { + filter.Limit = n + } + } + if c.Query("suspicious") == "1" || strings.EqualFold(c.Query("suspicious"), "true") { + filter.Suspicious = true + } + + sessions, err := h.mgr().DB().ListC2SessionsForAccess(filter, c2AccessFromContext(c)) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"sessions": sessions}) +} + +// GetSession 获取单个会话 +func (h *C2Handler) GetSession(c *gin.Context) { + id := c.Param("id") + session, err := h.mgr().DB().GetC2Session(id) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if session == nil { + c.JSON(http.StatusNotFound, gin.H{"error": "session not found"}) + return + } + + // 获取最近任务 + tasks, _ := h.mgr().DB().ListC2TasksForAccess(database.ListC2TasksFilter{ + SessionID: id, + Limit: 20, + }, c2AccessFromContext(c)) + + c.JSON(http.StatusOK, gin.H{ + "session": session, + "tasks": tasks, + }) +} + +// DeleteSession 删除会话 +func (h *C2Handler) DeleteSession(c *gin.Context) { + id := c.Param("id") + if err := h.mgr().DB().DeleteC2Session(id); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "c2", "session_delete", "删除 C2 会话", "c2_session", id, nil) + } + c.JSON(http.StatusOK, gin.H{"deleted": true}) +} + +// DeleteSessions 批量删除会话（请求体 JSON: {"ids":["s_xxx",...]}） +func (h *C2Handler) DeleteSessions(c *gin.Context) { + var req struct { + IDs []string `json:"ids"` + } + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid json: " + err.Error()}) + return + } + if len(req.IDs) == 0 { + c.JSON(http.StatusBadRequest, gin.H{"error": "ids is required"}) + return + } + n, err := h.mgr().DB().DeleteC2SessionsByIDsForAccess(req.IDs, c2AccessFromContext(c)) + if err != nil { + if errors.Is(err, database.ErrNoValidC2SessionIDs) { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "c2", "session_delete", "批量删除 C2 会话", "c2_session", "", map[string]interface{}{ + "count": n, "ids": req.IDs, + }) + } + c.JSON(http.StatusOK, gin.H{"deleted": n}) +} + +// SetSessionSleep 设置会话的 sleep/jitter，并下发 sleep 任务到植入体 +func (h *C2Handler) SetSessionSleep(c *gin.Context) { + id := c.Param("id") + var req struct { + SleepSeconds int `json:"sleep_seconds"` + JitterPercent int `json:"jitter_percent"` + } + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if req.SleepSeconds < 1 { + c.JSON(http.StatusBadRequest, gin.H{"error": "sleep_seconds must be >= 1"}) + return + } + if req.JitterPercent < 0 || req.JitterPercent > 100 { + c.JSON(http.StatusBadRequest, gin.H{"error": "jitter_percent must be 0-100"}) + return + } + + task, err := h.mgr().SetSessionSleep(id, req.SleepSeconds, req.JitterPercent) + if err != nil { + if errors.Is(err, sql.ErrNoRows) { + c.JSON(http.StatusNotFound, gin.H{"error": "session not found"}) + return + } + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + out := gin.H{ + "updated": true, + "sleep_seconds": req.SleepSeconds, + "jitter_percent": req.JitterPercent, + } + if task != nil { + out["task_id"] = task.ID + } + c.JSON(http.StatusOK, out) +} + +// SetSessionNote 更新会话备注（仅服务端元数据，不下发植入体） +func (h *C2Handler) SetSessionNote(c *gin.Context) { + id := c.Param("id") + var req struct { + Note string `json:"note"` + } + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + note := strings.TrimSpace(req.Note) + if len(note) > 2000 { + c.JSON(http.StatusBadRequest, gin.H{"error": "note too long (max 2000 characters)"}) + return + } + + session, err := h.mgr().DB().GetC2Session(id) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if session == nil { + c.JSON(http.StatusNotFound, gin.H{"error": "session not found"}) + return + } + + if err := h.mgr().DB().SetC2SessionNote(id, note); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "c2", "session_note", "更新 C2 会话备注", "c2_session", id, map[string]interface{}{ + "note_len": len(note), + }) + } + c.JSON(http.StatusOK, gin.H{ + "updated": true, + "note": note, + }) +} + +// ============================================================================ +// 任务 API +// ============================================================================ + +// ListTasks 获取任务列表 +func (h *C2Handler) ListTasks(c *gin.Context) { + filter := database.ListC2TasksFilter{ + SessionID: c.Query("session_id"), + ProjectID: c.Query("project_id"), + Status: c.Query("status"), + TaskType: c.Query("task_type"), + } + if since := c.Query("since"); since != "" { + if t, err := database.ParseRFC3339Time(since); err == nil { + filter.Since = &t + } + } + + paginated := false + page := 1 + pageSize := 10 + if c.Query("page") != "" || c.Query("page_size") != "" { + paginated = true + if p, err := strconv.Atoi(c.DefaultQuery("page", "1")); err == nil && p > 0 { + page = p + } + if ps, err := strconv.Atoi(c.DefaultQuery("page_size", "10")); err == nil && ps > 0 { + pageSize = ps + if pageSize > 100 { + pageSize = 100 + } + } + filter.Limit = pageSize + filter.Offset = (page - 1) * pageSize + } else { + if limit := c.Query("limit"); limit != "" { + if n, err := strconv.Atoi(limit); err == nil && n > 0 { + filter.Limit = n + } + } + } + + access := c2AccessFromContext(c) + tasks, err := h.mgr().DB().ListC2TasksForAccess(filter, access) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + // 仪表盘「待审任务」为全局 queued/pending 数量，与列表 session 过滤无关 + pendingN, _ := h.mgr().DB().CountC2TasksQueuedOrPendingForAccess("", filter.ProjectID, access) + + if !paginated { + c.JSON(http.StatusOK, gin.H{ + "tasks": tasks, + "pending_queued_count": pendingN, + }) + return + } + + total, err := h.mgr().DB().CountC2TasksForAccess(filter, access) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + statusCounts, err := h.mgr().DB().CountC2TasksByStatusForAccess(filter, access) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{ + "tasks": tasks, + "total": total, + "status_counts": statusCounts, + "page": page, + "page_size": pageSize, + "pending_queued_count": pendingN, + }) +} + +// DeleteTasks 批量删除任务（请求体 JSON: {"ids":["t_xxx",...]}） +func (h *C2Handler) DeleteTasks(c *gin.Context) { + var req struct { + IDs []string `json:"ids"` + } + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid json: " + err.Error()}) + return + } + if len(req.IDs) == 0 { + c.JSON(http.StatusBadRequest, gin.H{"error": "ids is required"}) + return + } + n, err := h.mgr().DB().DeleteC2TasksByIDsForAccess(req.IDs, c2AccessFromContext(c)) + if err != nil { + if errors.Is(err, database.ErrNoValidC2TaskIDs) { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "c2", "task_delete", "批量删除 C2 任务", "c2_task", "", map[string]interface{}{ + "count": n, "ids": req.IDs, + }) + } + c.JSON(http.StatusOK, gin.H{"deleted": n}) +} + +// GetTask 获取单个任务 +func (h *C2Handler) GetTask(c *gin.Context) { + id := c.Param("id") + task, err := h.mgr().DB().GetC2Task(id) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if task == nil { + c.JSON(http.StatusNotFound, gin.H{"error": "task not found"}) + return + } + c.JSON(http.StatusOK, gin.H{"task": task}) +} + +// CreateTask 创建任务 +func (h *C2Handler) CreateTask(c *gin.Context) { + var req struct { + SessionID string `json:"session_id"` + TaskType string `json:"task_type"` + Payload map[string]interface{} `json:"payload"` + Source string `json:"source"` + ConversationID string `json:"conversation_id"` + } + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if strings.TrimSpace(req.SessionID) == "" { + req.SessionID = strings.TrimSpace(c.Param("id")) + } + if !h.c2ResourceAllowed(c, "c2_session", req.SessionID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + if conversationID := strings.TrimSpace(req.ConversationID); conversationID != "" { + session, ok := security.CurrentSession(c) + if !ok || !h.mgr().DB().UserCanAccessResource(session.UserID, session.Scope, "conversation", conversationID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权关联目标对话"}) + return + } + req.ConversationID = conversationID + } + + input := c2.EnqueueTaskInput{ + SessionID: req.SessionID, + TaskType: c2.TaskType(req.TaskType), + Payload: req.Payload, + Source: firstNonEmpty(req.Source, "manual"), + ConversationID: req.ConversationID, + UserCtx: c.Request.Context(), + } + + task, err := h.mgr().EnqueueTask(input) + if err != nil { + code := http.StatusInternalServerError + if e, ok := err.(*c2.CommonError); ok { + code = e.HTTP + } + c.JSON(code, gin.H{"error": err.Error()}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "c2", "task_create", "创建 C2 任务", "c2_task", task.ID, map[string]interface{}{ + "session_id": req.SessionID, "task_type": req.TaskType, + }) + } + c.JSON(http.StatusOK, gin.H{"task": task}) +} + +// CancelTask 取消任务 +func (h *C2Handler) CancelTask(c *gin.Context) { + id := c.Param("id") + if err := h.mgr().CancelTask(id); err != nil { + code := http.StatusInternalServerError + if e, ok := err.(*c2.CommonError); ok { + code = e.HTTP + } + c.JSON(code, gin.H{"error": err.Error()}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "c2", "task_cancel", "取消 C2 任务", "c2_task", id, nil) + } + c.JSON(http.StatusOK, gin.H{"cancelled": true}) +} + +// WaitTask 等待任务完成 +func (h *C2Handler) WaitTask(c *gin.Context) { + id := c.Param("id") + timeout := 60 * time.Second + if t := c.Query("timeout"); t != "" { + if n, err := strconv.Atoi(t); err == nil && n > 0 { + timeout = time.Duration(n) * time.Second + } + } + + deadline := time.Now().Add(timeout) + for time.Now().Before(deadline) { + task, err := h.mgr().DB().GetC2Task(id) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if task == nil { + c.JSON(http.StatusNotFound, gin.H{"error": "task not found"}) + return + } + if task.Status == "success" || task.Status == "failed" || task.Status == "cancelled" { + c.JSON(http.StatusOK, gin.H{"task": task}) + return + } + time.Sleep(500 * time.Millisecond) + } + c.JSON(http.StatusRequestTimeout, gin.H{"error": "timeout waiting for task completion"}) +} + +// ============================================================================ +// Payload API +// ============================================================================ + +// PayloadOneliner 生成单行 payload +func (h *C2Handler) PayloadOneliner(c *gin.Context) { + var req struct { + ListenerID string `json:"listener_id"` + Kind string `json:"kind"` // bash, python, powershell, curl_beacon + Host string `json:"host"` + } + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + listener, err := h.mgr().DB().GetC2Listener(req.ListenerID) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if listener == nil { + c.JSON(http.StatusNotFound, gin.H{"error": "listener not found"}) + return + } + if !h.c2ResourceAllowed(c, "c2_listener", req.ListenerID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + + host := c2.ResolveBeaconDialHost(listener, strings.TrimSpace(req.Host), h.logger, listener.ID) + + kind := c2.OnelinerKind(req.Kind) + if !c2.IsOnelinerCompatible(listener.Type, kind) { + compatible := c2.OnelinerKindsForListener(listener.Type) + names := make([]string, len(compatible)) + for i, k := range compatible { + names[i] = string(k) + } + c.JSON(http.StatusBadRequest, gin.H{ + "error": fmt.Sprintf("监听器类型 %s 不支持 %s 类型的 oneliner，请选择兼容的类型", listener.Type, req.Kind), + "compatible_kinds": names, + }) + return + } + + input := c2.OnelinerInput{ + Kind: kind, + Host: host, + Port: listener.BindPort, + HTTPBaseURL: fmt.Sprintf("http://%s:%d", host, listener.BindPort), + ImplantToken: listener.ImplantToken, + } + + oneliner, err := c2.GenerateOneliner(input) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + c.JSON(http.StatusOK, gin.H{ + "oneliner": oneliner, + "kind": req.Kind, + "host": host, + "port": listener.BindPort, + }) +} + +// PayloadBuild 构建 beacon 二进制 +func (h *C2Handler) PayloadBuild(c *gin.Context) { + var req struct { + ListenerID string `json:"listener_id"` + OS string `json:"os"` + Arch string `json:"arch"` + SleepSeconds int `json:"sleep_seconds"` + JitterPercent int `json:"jitter_percent"` + Host string `json:"host"` // 可选：编译进 Beacon 的回连地址，覆盖监听器 bind_host + } + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + listener, err := h.mgr().DB().GetC2Listener(req.ListenerID) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if listener == nil { + c.JSON(http.StatusNotFound, gin.H{"error": "listener not found"}) + return + } + if !h.c2ResourceAllowed(c, "c2_listener", req.ListenerID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + + builder := c2.NewPayloadBuilder(h.mgr(), h.logger, "", "") + input := c2.PayloadBuilderInput{ + ListenerID: req.ListenerID, + OS: req.OS, + Arch: req.Arch, + SleepSeconds: req.SleepSeconds, + JitterPercent: req.JitterPercent, + Host: strings.TrimSpace(req.Host), + } + + result, err := builder.BuildBeacon(input) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if session, ok := security.CurrentSession(c); ok { + _ = h.mgr().DB().RecordC2PayloadArtifact(filepath.Base(result.OutputPath), result.PayloadID, result.ListenerID, session.UserID) + } + + c.JSON(http.StatusOK, gin.H{ + "payload": result, + }) +} + +// PayloadDownload 下载 payload +func (h *C2Handler) PayloadDownload(c *gin.Context) { + id := c.Param("id") + filename := id + if !strings.HasPrefix(filename, "beacon_") { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid payload id"}) + return + } + if strings.Contains(filename, "/") || strings.Contains(filename, "\\") || strings.Contains(filename, "..") { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid payload id"}) + return + } + session, ok := security.CurrentSession(c) + if !ok || !h.mgr().DB().UserCanAccessC2Payload(session.UserID, session.Scope, filename) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + + builder := c2.NewPayloadBuilder(h.mgr(), h.logger, "", "") + storageDir := builder.GetPayloadStoragePath() + targetPath := filepath.Join(storageDir, filename) + + absTarget, err := filepath.Abs(targetPath) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid path"}) + return + } + absDir, err := filepath.Abs(storageDir) + if err != nil || !strings.HasPrefix(absTarget, absDir+string(filepath.Separator)) { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid payload id"}) + return + } + + c.FileAttachment(absTarget, filepath.Base(absTarget)) +} + +// ============================================================================ +// 事件 API +// ============================================================================ + +// ListEvents 获取事件列表 +func (h *C2Handler) ListEvents(c *gin.Context) { + filter := database.ListC2EventsFilter{ + Level: c.Query("level"), + Category: c.Query("category"), + ProjectID: c.Query("project_id"), + SessionID: c.Query("session_id"), + TaskID: c.Query("task_id"), + } + if since := c.Query("since"); since != "" { + if t, err := database.ParseRFC3339Time(since); err == nil { + filter.Since = &t + } + } + + paginated := false + page := 1 + pageSize := 10 + if c.Query("page") != "" || c.Query("page_size") != "" { + paginated = true + if p, err := strconv.Atoi(c.DefaultQuery("page", "1")); err == nil && p > 0 { + page = p + } + if ps, err := strconv.Atoi(c.DefaultQuery("page_size", "10")); err == nil && ps > 0 { + pageSize = ps + if pageSize > 100 { + pageSize = 100 + } + } + filter.Limit = pageSize + filter.Offset = (page - 1) * pageSize + } else { + if limit := c.Query("limit"); limit != "" { + if n, err := strconv.Atoi(limit); err == nil && n > 0 { + filter.Limit = n + } + } + } + + access := c2AccessFromContext(c) + events, err := h.mgr().DB().ListC2EventsForAccess(filter, access) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if !paginated { + c.JSON(http.StatusOK, gin.H{"events": events}) + return + } + total, err := h.mgr().DB().CountC2EventsForAccess(filter, access) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + levelCounts, err := h.mgr().DB().CountC2EventsByLevelForAccess(filter, access) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{ + "events": events, + "total": total, + "level_counts": levelCounts, + "page": page, + "page_size": pageSize, + }) +} + +// DeleteEvents 批量删除事件（请求体 JSON: {"ids":["e_xxx",...]}） +func (h *C2Handler) DeleteEvents(c *gin.Context) { + var req struct { + IDs []string `json:"ids"` + } + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid json: " + err.Error()}) + return + } + if len(req.IDs) == 0 { + c.JSON(http.StatusBadRequest, gin.H{"error": "ids is required"}) + return + } + n, err := h.mgr().DB().DeleteC2EventsByIDsForAccess(req.IDs, c2AccessFromContext(c)) + if err != nil { + if errors.Is(err, database.ErrNoValidC2EventIDs) { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"deleted": n}) +} + +// EventStream SSE 实时事件流 +func (h *C2Handler) EventStream(c *gin.Context) { + c.Header("Content-Type", "text/event-stream") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + + sessionFilter := c.Query("session_id") + categoryFilter := c.Query("category") + levels := c.QueryArray("level") + + sub := h.mgr().EventBus().Subscribe( + "sse-"+uuid.New().String(), + 128, + sessionFilter, + categoryFilter, + levels, + ) + defer h.mgr().EventBus().Unsubscribe(sub.ID) + + c.Stream(func(w io.Writer) bool { + select { + case e, ok := <-sub.Ch: + if !ok { + return false + } + if !h.c2EventAllowed(c, e) { + return true + } + data, _ := json.Marshal(e) + fmt.Fprintf(w, "data: %s\n\n", data) + return true + case <-c.Request.Context().Done(): + return false + } + }) +} + +// ============================================================================ +// Profile API +// ============================================================================ + +// ListProfiles 获取 Malleable Profile 列表 +func (h *C2Handler) ListProfiles(c *gin.Context) { + profiles, err := h.mgr().DB().ListC2Profiles() + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"profiles": profiles}) +} + +// GetProfile 获取单个 Profile +func (h *C2Handler) GetProfile(c *gin.Context) { + id := c.Param("id") + profile, err := h.mgr().DB().GetC2Profile(id) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if profile == nil { + c.JSON(http.StatusNotFound, gin.H{"error": "profile not found"}) + return + } + c.JSON(http.StatusOK, gin.H{"profile": profile}) +} + +// CreateProfile 创建 Profile +func (h *C2Handler) CreateProfile(c *gin.Context) { + var req database.C2Profile + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + req.ID = "p_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:14] + req.CreatedAt = time.Now() + + if err := h.mgr().DB().CreateC2Profile(&req); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"profile": req}) +} + +// UpdateProfile 更新 Profile +func (h *C2Handler) UpdateProfile(c *gin.Context) { + id := c.Param("id") + profile, err := h.mgr().DB().GetC2Profile(id) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if profile == nil { + c.JSON(http.StatusNotFound, gin.H{"error": "profile not found"}) + return + } + + var req database.C2Profile + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + profile.Name = req.Name + profile.UserAgent = req.UserAgent + profile.URIs = req.URIs + profile.RequestHeaders = req.RequestHeaders + profile.ResponseHeaders = req.ResponseHeaders + profile.BodyTemplate = req.BodyTemplate + profile.JitterMinMS = req.JitterMinMS + profile.JitterMaxMS = req.JitterMaxMS + + if err := h.mgr().DB().UpdateC2Profile(profile); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"profile": profile}) +} + +// DeleteProfile 删除 Profile +func (h *C2Handler) DeleteProfile(c *gin.Context) { + id := c.Param("id") + if err := h.mgr().DB().DeleteC2Profile(id); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"deleted": true}) +} + +// ============================================================================ +// 文件管理 API（C2 Upload 任务需要先通过此 API 上传文件到 downstream 目录） +// ============================================================================ + +// UploadFileForImplant 操作员上传文件，供 upload 任务推送给 implant +func (h *C2Handler) UploadFileForImplant(c *gin.Context) { + sessionID := strings.TrimSpace(c.PostForm("session_id")) + remotePath := strings.TrimSpace(c.PostForm("remote_path")) + if sessionID == "" || remotePath == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "session_id and remote_path required"}) + return + } + if !h.c2ResourceAllowed(c, "c2_session", sessionID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + + file, header, err := c.Request.FormFile("file") + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "file field required: " + err.Error()}) + return + } + defer file.Close() + + fileID := "f_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:14] + dir := filepath.Join(h.mgr().StorageDir(), "downstream") + if err := osMkdirAll(dir); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + dstPath := filepath.Join(dir, fileID+".bin") + dst, err := osCreate(dstPath) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + n, err := io.Copy(dst, file) + dst.Close() + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + // Record in DB + dbFile := &database.C2File{ + ID: fileID, + SessionID: sessionID, + Direction: "upload", + RemotePath: remotePath, + LocalPath: dstPath, + SizeBytes: n, + CreatedAt: time.Now(), + } + _ = h.mgr().DB().CreateC2File(dbFile) + + c.JSON(http.StatusOK, gin.H{ + "file_id": fileID, + "size": n, + "filename": header.Filename, + "remote_path": remotePath, + }) +} + +// ListFiles 列出某会话的文件记录 +func (h *C2Handler) ListFiles(c *gin.Context) { + sessionID := c.Query("session_id") + if sessionID == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "session_id required"}) + return + } + if !h.c2ResourceAllowed(c, "c2_session", sessionID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + files, err := h.mgr().DB().ListC2FilesBySession(sessionID) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"files": files}) +} + +// DownloadResultFile 下载任务结果文件（截图等 blob 结果） +func (h *C2Handler) DownloadResultFile(c *gin.Context) { + taskID := c.Param("id") + task, err := h.mgr().DB().GetC2Task(taskID) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if task == nil { + c.JSON(http.StatusNotFound, gin.H{"error": "task not found"}) + return + } + if !h.c2ResourceAllowed(c, "c2_task", taskID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + if task.ResultBlobPath == "" { + c.JSON(http.StatusNotFound, gin.H{"error": "no result file for this task"}) + return + } + c.FileAttachment(task.ResultBlobPath, filepath.Base(task.ResultBlobPath)) +} + +func osMkdirAll(path string) error { + return os.MkdirAll(path, 0o755) +} + +func osCreate(path string) (*os.File, error) { + return os.Create(path) +} + +func c2AccessFromContext(c *gin.Context) database.RBACListAccess { + session, ok := security.CurrentSession(c) + if !ok { + return database.RBACListAccess{} + } + return database.RBACListAccess{UserID: session.UserID, Scope: session.Scope} +} + +func (h *C2Handler) canAccessProject(c *gin.Context, projectID string) bool { + projectID = strings.TrimSpace(projectID) + if projectID == "" { + return true + } + session, ok := security.CurrentSession(c) + if !ok { + return false + } + if session.Scope == database.RBACScopeAll { + return true + } + return h.mgr().DB().UserCanAccessResource(session.UserID, session.Scope, "project", projectID) +} + +func (h *C2Handler) c2ResourceAllowed(c *gin.Context, resourceType, resourceID string) bool { + session, ok := security.CurrentSession(c) + if !ok { + return false + } + return h.mgr().DB().UserCanAccessResource(session.UserID, session.Scope, resourceType, resourceID) +} + +func (h *C2Handler) c2EventAllowed(c *gin.Context, e *c2.Event) bool { + if e == nil { + return false + } + session, ok := security.CurrentSession(c) + if !ok { + return false + } + if session.Scope == database.RBACScopeAll { + return true + } + if strings.TrimSpace(e.SessionID) != "" { + return h.mgr().DB().UserCanAccessResource(session.UserID, session.Scope, "c2_session", e.SessionID) + } + if strings.TrimSpace(e.TaskID) != "" { + return h.mgr().DB().UserCanAccessResource(session.UserID, session.Scope, "c2_task", e.TaskID) + } + return false +} + +// ============================================================================ +// 辅助函数（firstNonEmpty 已在 vulnerability.go 中定义） +// ============================================================================ diff --git a/internal/handler/chat_uploads.go b/internal/handler/chat_uploads.go new file mode 100644 index 00000000..6edabcf1 --- /dev/null +++ b/internal/handler/chat_uploads.go @@ -0,0 +1,1516 @@ +package handler + +import ( + "archive/zip" + "crypto/rand" + "encoding/json" + "fmt" + "io" + "mime" + "net/http" + "os" + "path/filepath" + "sort" + "strconv" + "strings" + "time" + "unicode/utf8" + + "cyberstrike-ai/internal/audit" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +const ( + chatUploadsRootDirName = "chat_uploads" + reductionRootDirName = "tmp/reduction" + workspaceRootDirName = "tmp/workspace" + artifactsRootDirName = "data/conversation_artifacts" + reductionVirtualPrefix = "__reduction__/" + workspaceVirtualPrefix = "__workspace__/" + artifactVirtualPrefix = "__conversation_artifact__/" + chatUploadSourceUpload = "upload" + chatUploadSourceReduction = "reduction" + chatUploadSourceWorkspace = "workspace" + chatUploadSourceConversation = "conversation_artifact" + maxChatUploadEditBytes = 2 * 1024 * 1024 // 文本编辑上限 +) + +// ChatUploadsHandler 对话中上传附件（chat_uploads 目录）的管理 API +type ChatUploadsHandler struct { + logger *zap.Logger + audit *audit.Service + db *database.DB +} + +// SetAudit wires platform audit logging. +func (h *ChatUploadsHandler) SetAudit(s *audit.Service) { + h.audit = s +} + +// NewChatUploadsHandler 创建处理器 +func NewChatUploadsHandler(logger *zap.Logger, databases ...*database.DB) *ChatUploadsHandler { + h := &ChatUploadsHandler{logger: logger} + if len(databases) > 0 { + h.db = databases[0] + } + return h +} + +func (h *ChatUploadsHandler) pathAllowed(c *gin.Context, relativePath string) bool { + session, ok := security.CurrentSession(c) + if !ok || h.db == nil { + return false + } + if session.Scope == database.RBACScopeAll { + return true + } + rel := filepath.ToSlash(filepath.Clean(filepath.FromSlash(strings.TrimSpace(relativePath)))) + rel = strings.Trim(rel, "/") + if conversationID, ownerUserID, found := h.db.GetChatUploadArtifact(rel); found { + return strings.TrimSpace(ownerUserID) == session.UserID || h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", conversationID) + } + parts := strings.Split(strings.Trim(rel, "/"), "/") + if len(parts) < 2 || parts[1] == "" || parts[1] == "_manual" { + return false + } + return h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", parts[1]) +} + +func (h *ChatUploadsHandler) reductionPathAllowed(c *gin.Context, scope, id string) bool { + session, ok := security.CurrentSession(c) + if !ok || h.db == nil { + return false + } + if session.Scope == database.RBACScopeAll { + return true + } + id = strings.TrimSpace(id) + switch scope { + case "conversations": + if id == "" || id == "default" { + return false + } + return h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", id) + case "projects": + if id == "" { + return false + } + return h.db.UserCanAccessResource(session.UserID, session.Scope, "project", id) + default: + return false + } +} + +func (h *ChatUploadsHandler) reductionVirtualPathAllowed(c *gin.Context, relativePath string) bool { + rel := strings.TrimPrefix(strings.TrimSpace(relativePath), reductionVirtualPrefix) + rel = strings.Trim(filepath.ToSlash(filepath.Clean(filepath.FromSlash(rel))), "/") + if rel == "projects" || rel == "conversations" { + _, ok := security.CurrentSession(c) + return ok + } + parts := strings.Split(filepath.ToSlash(rel), "/") + if len(parts) < 2 { + return false + } + return h.reductionPathAllowed(c, parts[0], parts[1]) +} + +func (h *ChatUploadsHandler) workspacePathAllowed(c *gin.Context, scope, id string) bool { + return h.reductionPathAllowed(c, scope, id) +} + +func (h *ChatUploadsHandler) workspaceVirtualPathAllowed(c *gin.Context, relativePath string) bool { + rel := strings.TrimPrefix(strings.TrimSpace(relativePath), workspaceVirtualPrefix) + rel = strings.Trim(filepath.ToSlash(filepath.Clean(filepath.FromSlash(rel))), "/") + if rel == "projects" || rel == "conversations" { + _, ok := security.CurrentSession(c) + return ok + } + parts := strings.Split(filepath.ToSlash(rel), "/") + if len(parts) < 2 { + return false + } + return h.workspacePathAllowed(c, parts[0], parts[1]) +} + +func (h *ChatUploadsHandler) conversationArtifactPathAllowed(c *gin.Context, conversationID string) bool { + session, ok := security.CurrentSession(c) + if !ok || h.db == nil { + return false + } + if session.Scope == database.RBACScopeAll { + return true + } + conversationID = strings.TrimSpace(conversationID) + if conversationID == "" || conversationID == "default" { + return false + } + return h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", conversationID) +} + +func (h *ChatUploadsHandler) conversationArtifactVirtualPathAllowed(c *gin.Context, relativePath string) bool { + rel := strings.TrimPrefix(strings.TrimSpace(relativePath), artifactVirtualPrefix) + parts := strings.Split(filepath.ToSlash(rel), "/") + if len(parts) < 1 { + return false + } + return h.conversationArtifactPathAllowed(c, parts[0]) +} + +func (h *ChatUploadsHandler) absRoot() (string, error) { + cwd, err := os.Getwd() + if err != nil { + return "", err + } + return filepath.Abs(filepath.Join(cwd, chatUploadsRootDirName)) +} + +func (h *ChatUploadsHandler) absReductionRoot() (string, error) { + if h.db != nil { + if base := strings.TrimSpace(h.db.EinoReductionBaseDir()); base != "" { + if filepath.IsAbs(base) { + return filepath.Abs(base) + } + cwd, err := os.Getwd() + if err != nil { + return "", err + } + return filepath.Abs(filepath.Join(cwd, base)) + } + } + cwd, err := os.Getwd() + if err != nil { + return "", err + } + return filepath.Abs(filepath.Join(cwd, reductionRootDirName)) +} + +func (h *ChatUploadsHandler) absWorkspaceRoot() (string, error) { + if h.db != nil { + if base := strings.TrimSpace(h.db.EinoWorkspaceBaseDir()); base != "" { + if filepath.IsAbs(base) { + return filepath.Abs(base) + } + cwd, err := os.Getwd() + if err != nil { + return "", err + } + return filepath.Abs(filepath.Join(cwd, base)) + } + } + cwd, err := os.Getwd() + if err != nil { + return "", err + } + return filepath.Abs(filepath.Join(cwd, workspaceRootDirName)) +} + +func (h *ChatUploadsHandler) absConversationArtifactsRoot() (string, error) { + if h.db != nil { + if base := strings.TrimSpace(h.db.ConversationArtifactsBaseDir()); base != "" { + if filepath.IsAbs(base) { + return filepath.Abs(base) + } + cwd, err := os.Getwd() + if err != nil { + return "", err + } + return filepath.Abs(filepath.Join(cwd, base)) + } + } + cwd, err := os.Getwd() + if err != nil { + return "", err + } + return filepath.Abs(filepath.Join(cwd, artifactsRootDirName)) +} + +// resolveUnderChatUploads 校验 relativePath（使用 / 分隔）对应文件必须在 chat_uploads 根下 +func (h *ChatUploadsHandler) resolveUnderChatUploads(relativePath string) (abs string, err error) { + root, err := h.absRoot() + if err != nil { + return "", err + } + rel := strings.TrimSpace(relativePath) + if rel == "" { + return "", fmt.Errorf("empty path") + } + rel = filepath.Clean(filepath.FromSlash(rel)) + if rel == "." || strings.HasPrefix(rel, "..") { + return "", fmt.Errorf("invalid path") + } + full := filepath.Join(root, rel) + full, err = filepath.Abs(full) + if err != nil { + return "", err + } + rootAbs, _ := filepath.Abs(root) + if full != rootAbs && !strings.HasPrefix(full, rootAbs+string(filepath.Separator)) { + return "", fmt.Errorf("path escapes chat_uploads root") + } + return full, nil +} + +// ChatUploadFileItem 列表项 +type ChatUploadFileItem struct { + RelativePath string `json:"relativePath"` + AbsolutePath string `json:"absolutePath"` // 服务器上的绝对路径，便于在对话中引用（与附件落盘路径一致） + Name string `json:"name"` + Source string `json:"source,omitempty"` + Size int64 `json:"size"` + ModifiedUnix int64 `json:"modifiedUnix"` + Date string `json:"date"` + ConversationID string `json:"conversationId"` + ConversationTitle string `json:"conversationTitle,omitempty"` + ProjectID string `json:"projectId,omitempty"` + ProjectName string `json:"projectName,omitempty"` + // SubPath 为日期、会话目录之下的子路径（不含文件名），如 date/conv/a/b/file 则为 "a/b"；无嵌套则为 ""。 + SubPath string `json:"subPath"` +} + +func (h *ChatUploadsHandler) conversationProjectID(conversationID string, cache map[string]string) string { + conversationID = strings.TrimSpace(conversationID) + if conversationID == "" || conversationID == "_manual" || conversationID == "_new" || h.db == nil { + return "" + } + if v, ok := cache[conversationID]; ok { + return v + } + projectID, err := h.db.GetConversationProjectID(conversationID) + if err != nil { + projectID = "" + } + cache[conversationID] = projectID + return projectID +} + +func (h *ChatUploadsHandler) conversationTitle(conversationID string, cache map[string]string) string { + conversationID = strings.TrimSpace(conversationID) + if conversationID == "" || conversationID == "_manual" || conversationID == "_new" || h.db == nil { + return "" + } + if v, ok := cache[conversationID]; ok { + return v + } + title, err := h.db.GetConversationTitle(conversationID) + if err != nil { + title = "" + } + cache[conversationID] = title + return title +} + +func (h *ChatUploadsHandler) projectName(projectID string, cache map[string]string) string { + projectID = strings.TrimSpace(projectID) + if projectID == "" || h.db == nil { + return "" + } + if v, ok := cache[projectID]; ok { + return v + } + name, err := h.db.GetProjectName(projectID) + if err != nil { + name = "" + } + cache[projectID] = name + return name +} + +func (h *ChatUploadsHandler) collectFiles(c *gin.Context, conversationFilter, projectFilter string) ([]ChatUploadFileItem, []string, error) { + root, err := h.absRoot() + if err != nil { + return nil, nil, err + } + // 保证根目录存在，否则「按文件夹」浏览时无法 mkdir，且首次列表为空时界面无路径工具栏 + if err := os.MkdirAll(root, 0755); err != nil { + return nil, nil, err + } + var files []ChatUploadFileItem + var folders []string + projectCache := make(map[string]string) + projectNameCache := make(map[string]string) + conversationTitleCache := make(map[string]string) + err = filepath.WalkDir(root, func(path string, d os.DirEntry, walkErr error) error { + if walkErr != nil { + return walkErr + } + rel, err := filepath.Rel(root, path) + if err != nil { + return err + } + if rel == "." { + return nil + } + relSlash := filepath.ToSlash(rel) + if d.IsDir() { + folders = append(folders, relSlash) + return nil + } + info, err := d.Info() + if err != nil { + return err + } + parts := strings.Split(relSlash, "/") + var dateStr, convID string + if len(parts) >= 2 { + dateStr = parts[0] + } + if len(parts) >= 3 { + convID = parts[1] + } + projectID := h.conversationProjectID(convID, projectCache) + var subPath string + if len(parts) >= 4 { + subPath = strings.Join(parts[2:len(parts)-1], "/") + } + if conversationFilter != "" && convID != conversationFilter { + return nil + } + if projectFilter != "" && projectID != projectFilter { + return nil + } + absPath, _ := filepath.Abs(path) + files = append(files, ChatUploadFileItem{ + RelativePath: relSlash, + AbsolutePath: absPath, + Name: d.Name(), + Source: chatUploadSourceUpload, + Size: info.Size(), + ModifiedUnix: info.ModTime().Unix(), + Date: dateStr, + ConversationID: convID, + ConversationTitle: h.conversationTitle(convID, conversationTitleCache), + ProjectID: projectID, + ProjectName: h.projectName(projectID, projectNameCache), + SubPath: subPath, + }) + return nil + }) + if err != nil { + return nil, nil, err + } + if conversationFilter != "" || projectFilter != "" { + filteredFolders := make([]string, 0, len(folders)) + for _, rel := range folders { + parts := strings.Split(rel, "/") + if len(parts) >= 2 && (conversationFilter == "" || parts[1] == conversationFilter) && (projectFilter == "" || h.conversationProjectID(parts[1], projectCache) == projectFilter) { + filteredFolders = append(filteredFolders, rel) + continue + } + if len(parts) == 1 { + prefix := rel + "/" + for _, f := range files { + if strings.HasPrefix(f.RelativePath, prefix) { + filteredFolders = append(filteredFolders, rel) + break + } + } + } + } + folders = filteredFolders + } + files = filterSlice(files, func(file ChatUploadFileItem) bool { + return h.pathAllowed(c, file.RelativePath) + }) + folders = filterSlice(folders, func(folder string) bool { + if h.pathAllowed(c, folder) { + return true + } + prefix := strings.TrimSuffix(folder, "/") + "/" + for _, file := range files { + if strings.HasPrefix(file.RelativePath, prefix) { + return true + } + } + return false + }) + sort.Strings(folders) + sort.Slice(files, func(i, j int) bool { + return files[i].ModifiedUnix > files[j].ModifiedUnix + }) + reductionFiles, err := h.collectReductionFiles(c, conversationFilter, projectFilter) + if err != nil { + h.logger.Warn("列举 reduction 产物失败", zap.Error(err)) + } else if len(reductionFiles) > 0 { + files = append(files, reductionFiles...) + } + workspaceFiles, err := h.collectWorkspaceFiles(c, conversationFilter, projectFilter) + if err != nil { + h.logger.Warn("列举 workspace 产物失败", zap.Error(err)) + } else if len(workspaceFiles) > 0 { + files = append(files, workspaceFiles...) + } + artifactFiles, err := h.collectConversationArtifactFiles(c, conversationFilter, projectFilter) + if err != nil { + h.logger.Warn("列举 conversation_artifacts 产物失败", zap.Error(err)) + } else if len(artifactFiles) > 0 { + files = append(files, artifactFiles...) + } + sort.Slice(files, func(i, j int) bool { + return files[i].ModifiedUnix > files[j].ModifiedUnix + }) + return files, folders, nil +} + +func (h *ChatUploadsHandler) collectReductionFiles(c *gin.Context, conversationFilter, projectFilter string) ([]ChatUploadFileItem, error) { + root, err := h.absReductionRoot() + if err != nil { + return nil, err + } + if st, err := os.Stat(root); err != nil || !st.IsDir() { + return nil, nil + } + projectCache := make(map[string]string) + projectNameCache := make(map[string]string) + conversationTitleCache := make(map[string]string) + files := make([]ChatUploadFileItem, 0) + for _, scope := range []string{"conversations", "projects"} { + scopeRoot := filepath.Join(root, scope) + _ = filepath.WalkDir(scopeRoot, func(path string, d os.DirEntry, walkErr error) error { + if walkErr != nil || d == nil || d.IsDir() { + return nil + } + rel, err := filepath.Rel(root, path) + if err != nil { + return nil + } + relSlash := filepath.ToSlash(rel) + parts := strings.Split(relSlash, "/") + if len(parts) < 3 { + return nil + } + ownerID := parts[1] + if !h.reductionPathAllowed(c, scope, ownerID) { + return nil + } + var convID, projectID string + if scope == "conversations" { + convID = ownerID + projectID = h.conversationProjectID(convID, projectCache) + } else { + projectID = ownerID + } + if conversationFilter != "" && convID != conversationFilter { + if scope != "projects" || h.conversationProjectID(conversationFilter, projectCache) != projectID { + return nil + } + convID = conversationFilter + } + if projectFilter != "" && projectID != projectFilter { + return nil + } + info, err := d.Info() + if err != nil { + return nil + } + name := d.Name() + if filepath.Ext(name) == "" { + name += ".txt" + } + abs, _ := filepath.Abs(path) + files = append(files, ChatUploadFileItem{ + RelativePath: reductionVirtualPrefix + relSlash, + AbsolutePath: abs, + Name: name, + Source: chatUploadSourceReduction, + Size: info.Size(), + ModifiedUnix: info.ModTime().Unix(), + Date: info.ModTime().Format("2006-01-02"), + ConversationID: convID, + ConversationTitle: h.conversationTitle(convID, conversationTitleCache), + ProjectID: projectID, + ProjectName: h.projectName(projectID, projectNameCache), + SubPath: strings.Join(parts[2:len(parts)-1], "/"), + }) + return nil + }) + } + return files, nil +} + +func (h *ChatUploadsHandler) collectWorkspaceFiles(c *gin.Context, conversationFilter, projectFilter string) ([]ChatUploadFileItem, error) { + root, err := h.absWorkspaceRoot() + if err != nil { + return nil, err + } + if st, err := os.Stat(root); err != nil || !st.IsDir() { + return nil, nil + } + projectCache := make(map[string]string) + projectNameCache := make(map[string]string) + conversationTitleCache := make(map[string]string) + files := make([]ChatUploadFileItem, 0) + for _, scope := range []string{"conversations", "projects"} { + scopeRoot := filepath.Join(root, scope) + _ = filepath.WalkDir(scopeRoot, func(path string, d os.DirEntry, walkErr error) error { + if walkErr != nil || d == nil || d.IsDir() { + return nil + } + rel, err := filepath.Rel(root, path) + if err != nil { + return nil + } + relSlash := filepath.ToSlash(rel) + parts := strings.Split(relSlash, "/") + if len(parts) < 3 { + return nil + } + ownerID := parts[1] + if !h.workspacePathAllowed(c, scope, ownerID) { + return nil + } + var convID, projectID string + if scope == "conversations" { + convID = ownerID + projectID = h.conversationProjectID(convID, projectCache) + } else { + projectID = ownerID + } + if conversationFilter != "" && convID != conversationFilter { + if scope != "projects" || h.conversationProjectID(conversationFilter, projectCache) != projectID { + return nil + } + convID = conversationFilter + } + if projectFilter != "" && projectID != projectFilter { + return nil + } + info, err := d.Info() + if err != nil { + return nil + } + abs, _ := filepath.Abs(path) + files = append(files, ChatUploadFileItem{ + RelativePath: workspaceVirtualPrefix + relSlash, + AbsolutePath: abs, + Name: d.Name(), + Source: chatUploadSourceWorkspace, + Size: info.Size(), + ModifiedUnix: info.ModTime().Unix(), + Date: info.ModTime().Format("2006-01-02"), + ConversationID: convID, + ConversationTitle: h.conversationTitle(convID, conversationTitleCache), + ProjectID: projectID, + ProjectName: h.projectName(projectID, projectNameCache), + SubPath: strings.Join(parts[2:len(parts)-1], "/"), + }) + return nil + }) + } + return files, nil +} + +func (h *ChatUploadsHandler) collectConversationArtifactFiles(c *gin.Context, conversationFilter, projectFilter string) ([]ChatUploadFileItem, error) { + root, err := h.absConversationArtifactsRoot() + if err != nil { + return nil, err + } + if st, err := os.Stat(root); err != nil || !st.IsDir() { + return nil, nil + } + projectCache := make(map[string]string) + projectNameCache := make(map[string]string) + conversationTitleCache := make(map[string]string) + files := make([]ChatUploadFileItem, 0) + _ = filepath.WalkDir(root, func(path string, d os.DirEntry, walkErr error) error { + if walkErr != nil || d == nil || d.IsDir() { + return nil + } + rel, err := filepath.Rel(root, path) + if err != nil { + return nil + } + relSlash := filepath.ToSlash(rel) + parts := strings.Split(relSlash, "/") + if len(parts) < 2 { + return nil + } + convID := parts[0] + if !h.conversationArtifactPathAllowed(c, convID) { + return nil + } + projectID := h.conversationProjectID(convID, projectCache) + if conversationFilter != "" && convID != conversationFilter { + return nil + } + if projectFilter != "" && projectID != projectFilter { + return nil + } + info, err := d.Info() + if err != nil { + return nil + } + name := d.Name() + if filepath.Ext(name) == "" { + name += ".txt" + } + abs, _ := filepath.Abs(path) + files = append(files, ChatUploadFileItem{ + RelativePath: artifactVirtualPrefix + relSlash, + AbsolutePath: abs, + Name: name, + Source: chatUploadSourceConversation, + Size: info.Size(), + ModifiedUnix: info.ModTime().Unix(), + Date: info.ModTime().Format("2006-01-02"), + ConversationID: convID, + ConversationTitle: h.conversationTitle(convID, conversationTitleCache), + ProjectID: projectID, + ProjectName: h.projectName(projectID, projectNameCache), + SubPath: strings.Join(parts[1:len(parts)-1], "/"), + }) + return nil + }) + return files, nil +} + +func (h *ChatUploadsHandler) resolveReductionVirtualPath(relativePath string) (string, error) { + rel := strings.TrimPrefix(strings.TrimSpace(relativePath), reductionVirtualPrefix) + rel = filepath.Clean(filepath.FromSlash(rel)) + if rel == "." || strings.HasPrefix(rel, "..") { + return "", fmt.Errorf("invalid path") + } + root, err := h.absReductionRoot() + if err != nil { + return "", err + } + full, err := filepath.Abs(filepath.Join(root, rel)) + if err != nil { + return "", err + } + rootAbs, _ := filepath.Abs(root) + if full != rootAbs && !strings.HasPrefix(full, rootAbs+string(filepath.Separator)) { + return "", fmt.Errorf("path escapes reduction root") + } + return full, nil +} + +func (h *ChatUploadsHandler) resolveWorkspaceVirtualPath(relativePath string) (string, error) { + rel := strings.TrimPrefix(strings.TrimSpace(relativePath), workspaceVirtualPrefix) + rel = filepath.Clean(filepath.FromSlash(rel)) + if rel == "." || strings.HasPrefix(rel, "..") { + return "", fmt.Errorf("invalid path") + } + root, err := h.absWorkspaceRoot() + if err != nil { + return "", err + } + full, err := filepath.Abs(filepath.Join(root, rel)) + if err != nil { + return "", err + } + rootAbs, _ := filepath.Abs(root) + if full != rootAbs && !strings.HasPrefix(full, rootAbs+string(filepath.Separator)) { + return "", fmt.Errorf("path escapes workspace root") + } + return full, nil +} + +func (h *ChatUploadsHandler) resolveConversationArtifactVirtualPath(relativePath string) (string, error) { + rel := strings.TrimPrefix(strings.TrimSpace(relativePath), artifactVirtualPrefix) + rel = filepath.Clean(filepath.FromSlash(rel)) + if rel == "." || strings.HasPrefix(rel, "..") { + return "", fmt.Errorf("invalid path") + } + root, err := h.absConversationArtifactsRoot() + if err != nil { + return "", err + } + full, err := filepath.Abs(filepath.Join(root, rel)) + if err != nil { + return "", err + } + rootAbs, _ := filepath.Abs(root) + if full != rootAbs && !strings.HasPrefix(full, rootAbs+string(filepath.Separator)) { + return "", fmt.Errorf("path escapes conversation artifacts root") + } + return full, nil +} + +func chatUploadItemIsInternal(item ChatUploadFileItem) bool { + return item.Source == chatUploadSourceReduction || + item.Source == chatUploadSourceWorkspace || + item.Source == chatUploadSourceConversation || + strings.HasPrefix(item.RelativePath, reductionVirtualPrefix) || + strings.HasPrefix(item.RelativePath, workspaceVirtualPrefix) || + strings.HasPrefix(item.RelativePath, artifactVirtualPrefix) +} + +func (h *ChatUploadsHandler) resolveListedFilePath(item ChatUploadFileItem) (string, error) { + switch { + case item.Source == chatUploadSourceReduction || strings.HasPrefix(item.RelativePath, reductionVirtualPrefix): + return h.resolveReductionVirtualPath(item.RelativePath) + case item.Source == chatUploadSourceWorkspace || strings.HasPrefix(item.RelativePath, workspaceVirtualPrefix): + return h.resolveWorkspaceVirtualPath(item.RelativePath) + case item.Source == chatUploadSourceConversation || strings.HasPrefix(item.RelativePath, artifactVirtualPrefix): + return h.resolveConversationArtifactVirtualPath(item.RelativePath) + default: + return h.resolveUnderChatUploads(item.RelativePath) + } +} + +func chatUploadSourceMatches(item ChatUploadFileItem, sourceFilter string) bool { + sourceFilter = strings.TrimSpace(sourceFilter) + if sourceFilter == "" || sourceFilter == "all" { + return true + } + source := strings.TrimSpace(item.Source) + if source == "" { + source = chatUploadSourceUpload + } + return source == sourceFilter +} + +func chatUploadSearchMatches(item ChatUploadFileItem, search string) bool { + search = strings.ToLower(strings.TrimSpace(search)) + if search == "" { + return true + } + values := []string{ + item.RelativePath, + item.Name, + item.Source, + item.Date, + item.ConversationID, + item.ProjectID, + item.SubPath, + } + for _, value := range values { + if strings.Contains(strings.ToLower(value), search) { + return true + } + } + return false +} + +func filterChatUploadItems(files []ChatUploadFileItem, sourceFilter, search string) []ChatUploadFileItem { + if strings.TrimSpace(sourceFilter) == "" && strings.TrimSpace(search) == "" { + return files + } + out := make([]ChatUploadFileItem, 0, len(files)) + for _, item := range files { + if chatUploadSourceMatches(item, sourceFilter) && chatUploadSearchMatches(item, search) { + out = append(out, item) + } + } + return out +} + +func parsePositiveIntQuery(c *gin.Context, key string, def, max int) int { + raw := strings.TrimSpace(c.Query(key)) + if raw == "" { + return def + } + if strings.EqualFold(raw, "all") { + return 0 + } + n, err := strconv.Atoi(raw) + if err != nil || n <= 0 { + return def + } + if max > 0 && n > max { + return max + } + return n +} + +func paginateChatUploadItems(files []ChatUploadFileItem, page, pageSize int) []ChatUploadFileItem { + if pageSize <= 0 { + return files + } + if page <= 0 { + page = 1 + } + start := (page - 1) * pageSize + if start >= len(files) { + return []ChatUploadFileItem{} + } + end := start + pageSize + if end > len(files) { + end = len(files) + } + return files[start:end] +} + +// List GET /api/chat-uploads +func (h *ChatUploadsHandler) List(c *gin.Context) { + conversationFilter := strings.TrimSpace(c.Query("conversation")) + projectFilter := strings.TrimSpace(c.Query("project")) + sourceFilter := strings.TrimSpace(c.Query("source")) + search := strings.TrimSpace(c.Query("search")) + page := parsePositiveIntQuery(c, "page", 1, 0) + pageSize := parsePositiveIntQuery(c, "pageSize", 20, 200) + files, folders, err := h.collectFiles(c, conversationFilter, projectFilter) + if err != nil { + h.logger.Warn("列举对话附件失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + files = filterChatUploadItems(files, sourceFilter, search) + total := len(files) + paged := paginateChatUploadItems(files, page, pageSize) + c.JSON(http.StatusOK, gin.H{ + "files": paged, + "folders": folders, + "total": total, + "page": page, + "pageSize": pageSize, + }) +} + +// Export GET /api/chat-uploads/export?conversation=...&project=... +func (h *ChatUploadsHandler) Export(c *gin.Context) { + conversationFilter := strings.TrimSpace(c.Query("conversation")) + projectFilter := strings.TrimSpace(c.Query("project")) + sourceFilter := strings.TrimSpace(c.Query("source")) + search := strings.TrimSpace(c.Query("search")) + files, _, err := h.collectFiles(c, conversationFilter, projectFilter) + if err != nil { + h.logger.Warn("导出对话附件失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + files = filterChatUploadItems(files, sourceFilter, search) + if len(files) == 0 { + c.JSON(http.StatusNotFound, gin.H{"error": "no files to export"}) + return + } + nameParts := []string{"chat-files"} + if projectFilter != "" { + nameParts = append(nameParts, "project-"+projectFilter) + } + if conversationFilter != "" { + nameParts = append(nameParts, "conversation-"+conversationFilter) + } + nameParts = append(nameParts, time.Now().Format("20060102-150405")) + filename := strings.Join(nameParts, "-") + ".zip" + c.Header("Content-Type", "application/zip") + c.Header("Content-Disposition", mime.FormatMediaType("attachment", map[string]string{"filename": filename})) + + zw := zip.NewWriter(c.Writer) + defer zw.Close() + + manifest := gin.H{ + "exportedAt": time.Now().UTC().Format(time.RFC3339), + "conversationId": conversationFilter, + "projectId": projectFilter, + "source": sourceFilter, + "search": search, + "fileCount": len(files), + "files": files, + "layout": "chat uploads are stored under conversations//; internal outputs are stored under internal//", + "sourceDirectory": []string{chatUploadsRootDirName, reductionRootDirName, workspaceRootDirName, artifactsRootDirName}, + } + manifestBytes, _ := json.MarshalIndent(manifest, "", " ") + mw, err := zw.Create("manifest.json") + if err != nil { + h.logger.Warn("写入附件导出清单失败", zap.Error(err)) + return + } + _, _ = mw.Write(manifestBytes) + + used := make(map[string]int) + for _, item := range files { + abs, err := h.resolveListedFilePath(item) + if err != nil { + continue + } + st, err := os.Stat(abs) + if err != nil || st.IsDir() { + continue + } + var zipName string + if item.Source == chatUploadSourceReduction || strings.HasPrefix(item.RelativePath, reductionVirtualPrefix) { + rel := strings.TrimPrefix(item.RelativePath, reductionVirtualPrefix) + zipName = filepath.ToSlash(filepath.Join("internal", "reduction", rel)) + if filepath.Ext(zipName) == "" { + zipName += ".txt" + } + } else if item.Source == chatUploadSourceWorkspace || strings.HasPrefix(item.RelativePath, workspaceVirtualPrefix) { + rel := strings.TrimPrefix(item.RelativePath, workspaceVirtualPrefix) + zipName = filepath.ToSlash(filepath.Join("internal", "workspace", rel)) + } else if item.Source == chatUploadSourceConversation || strings.HasPrefix(item.RelativePath, artifactVirtualPrefix) { + rel := strings.TrimPrefix(item.RelativePath, artifactVirtualPrefix) + zipName = filepath.ToSlash(filepath.Join("internal", "conversation_artifacts", rel)) + if filepath.Ext(zipName) == "" { + zipName += ".txt" + } + } else { + conv := strings.TrimSpace(item.ConversationID) + if conv == "" || conv == "_manual" || conv == "_new" { + conv = "manual" + } + zipName = filepath.ToSlash(filepath.Join("conversations", conv, strings.TrimSpace(item.SubPath), item.Name)) + } + if used[zipName] > 0 { + ext := filepath.Ext(zipName) + zipName = strings.TrimSuffix(zipName, ext) + fmt.Sprintf("-%d", used[zipName]+1) + ext + } + used[zipName]++ + fw, err := zw.Create(zipName) + if err != nil { + h.logger.Warn("创建附件导出项失败", zap.String("path", item.RelativePath), zap.Error(err)) + continue + } + src, err := os.Open(abs) + if err != nil { + continue + } + _, copyErr := io.Copy(fw, src) + _ = src.Close() + if copyErr != nil { + h.logger.Warn("复制附件导出项失败", zap.String("path", item.RelativePath), zap.Error(copyErr)) + return + } + } + if h.audit != nil { + h.audit.RecordOK(c, "file", "export", "导出对话附件", "chat_upload", filename, map[string]interface{}{ + "conversation_id": conversationFilter, + "project_id": projectFilter, + "file_count": len(files), + }) + } +} + +// Download GET /api/chat-uploads/download?path=... +func (h *ChatUploadsHandler) Download(c *gin.Context) { + p := c.Query("path") + if strings.HasPrefix(strings.TrimSpace(p), reductionVirtualPrefix) { + if !h.reductionVirtualPathAllowed(c, p) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + abs, err := h.resolveReductionVirtualPath(p) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + st, err := os.Stat(abs) + if err != nil || st.IsDir() { + c.JSON(http.StatusNotFound, gin.H{"error": "file not found"}) + return + } + name := filepath.Base(abs) + if filepath.Ext(name) == "" { + name += ".txt" + } + c.FileAttachment(abs, name) + return + } + if strings.HasPrefix(strings.TrimSpace(p), workspaceVirtualPrefix) { + if !h.workspaceVirtualPathAllowed(c, p) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + abs, err := h.resolveWorkspaceVirtualPath(p) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + st, err := os.Stat(abs) + if err != nil || st.IsDir() { + c.JSON(http.StatusNotFound, gin.H{"error": "file not found"}) + return + } + c.FileAttachment(abs, filepath.Base(abs)) + return + } + if strings.HasPrefix(strings.TrimSpace(p), artifactVirtualPrefix) { + if !h.conversationArtifactVirtualPathAllowed(c, p) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + abs, err := h.resolveConversationArtifactVirtualPath(p) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + st, err := os.Stat(abs) + if err != nil || st.IsDir() { + c.JSON(http.StatusNotFound, gin.H{"error": "file not found"}) + return + } + name := filepath.Base(abs) + if filepath.Ext(name) == "" { + name += ".txt" + } + c.FileAttachment(abs, name) + return + } + if !h.pathAllowed(c, p) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + abs, err := h.resolveUnderChatUploads(p) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + st, err := os.Stat(abs) + if err != nil || st.IsDir() { + c.JSON(http.StatusNotFound, gin.H{"error": "file not found"}) + return + } + c.FileAttachment(abs, filepath.Base(abs)) +} + +// ResolvePath GET /api/chat-uploads/path?path=...&kind=file|directory +func (h *ChatUploadsHandler) ResolvePath(c *gin.Context) { + p := strings.TrimSpace(c.Query("path")) + kind := strings.TrimSpace(c.Query("kind")) + if kind == "" { + kind = "file" + } + var abs string + var err error + switch { + case strings.HasPrefix(p, reductionVirtualPrefix): + if strings.Trim(strings.TrimPrefix(p, reductionVirtualPrefix), "/") == "" { + if _, ok := security.CurrentSession(c); !ok { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + abs, err = h.absReductionRoot() + break + } + if !h.reductionVirtualPathAllowed(c, p) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + abs, err = h.resolveReductionVirtualPath(p) + case strings.HasPrefix(p, workspaceVirtualPrefix): + if strings.Trim(strings.TrimPrefix(p, workspaceVirtualPrefix), "/") == "" { + if _, ok := security.CurrentSession(c); !ok { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + abs, err = h.absWorkspaceRoot() + break + } + if !h.workspaceVirtualPathAllowed(c, p) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + abs, err = h.resolveWorkspaceVirtualPath(p) + case strings.HasPrefix(p, artifactVirtualPrefix): + if strings.Trim(strings.TrimPrefix(p, artifactVirtualPrefix), "/") == "" { + if _, ok := security.CurrentSession(c); !ok { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + abs, err = h.absConversationArtifactsRoot() + break + } + if !h.conversationArtifactVirtualPathAllowed(c, p) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + abs, err = h.resolveConversationArtifactVirtualPath(p) + default: + if p == "" || p == "." { + if _, ok := security.CurrentSession(c); !ok { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + abs, err = h.absRoot() + break + } + if !h.pathAllowed(c, p) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + abs, err = h.resolveUnderChatUploads(p) + } + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + st, err := os.Stat(abs) + if err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": "path not found"}) + return + } + if kind == "directory" && !st.IsDir() { + c.JSON(http.StatusBadRequest, gin.H{"error": "not a directory"}) + return + } + if kind != "directory" && st.IsDir() { + c.JSON(http.StatusBadRequest, gin.H{"error": "not a file"}) + return + } + c.JSON(http.StatusOK, gin.H{"absolutePath": abs, "isDir": st.IsDir()}) +} + +type chatUploadPathBody struct { + Path string `json:"path"` +} + +// Delete DELETE /api/chat-uploads +func (h *ChatUploadsHandler) Delete(c *gin.Context) { + var body chatUploadPathBody + if err := c.ShouldBindJSON(&body); err != nil || strings.TrimSpace(body.Path) == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid body"}) + return + } + if !h.pathAllowed(c, body.Path) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + abs, err := h.resolveUnderChatUploads(body.Path) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + st, err := os.Stat(abs) + if err != nil { + if os.IsNotExist(err) { + c.JSON(http.StatusNotFound, gin.H{"error": "file not found"}) + return + } + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if st.IsDir() { + if err := os.RemoveAll(abs); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + } else { + if err := os.Remove(abs); err != nil { + if os.IsNotExist(err) { + c.JSON(http.StatusNotFound, gin.H{"error": "file not found"}) + return + } + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + } + _ = h.db.DeleteChatUploadArtifactPath(filepath.ToSlash(filepath.Clean(filepath.FromSlash(body.Path)))) + if h.audit != nil { + h.audit.RecordOK(c, "file", "delete", "删除对话附件", "chat_upload", body.Path, nil) + } + c.JSON(http.StatusOK, gin.H{"ok": true}) +} + +type chatUploadMkdirBody struct { + Parent string `json:"parent"` + Name string `json:"name"` +} + +// Mkdir POST /api/chat-uploads/mkdir — 在 parent 目录下新建子目录（parent 为 chat_uploads 下相对路径，空表示根目录；name 为单段目录名） +func (h *ChatUploadsHandler) Mkdir(c *gin.Context) { + var body chatUploadMkdirBody + if err := c.ShouldBindJSON(&body); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid body"}) + return + } + name := strings.TrimSpace(body.Name) + if name == "" || strings.ContainsAny(name, `/\`) || name == "." || name == ".." { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid name"}) + return + } + if utf8.RuneCountInString(name) > 200 { + c.JSON(http.StatusBadRequest, gin.H{"error": "name too long"}) + return + } + + parent := strings.TrimSpace(body.Parent) + parent = filepath.ToSlash(filepath.Clean(filepath.FromSlash(parent))) + parent = strings.Trim(parent, "/") + if parent == "." { + parent = "" + } + if !h.pathAllowed(c, parent) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + + root, err := h.absRoot() + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + if parent != "" { + absParent, err := h.resolveUnderChatUploads(parent) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + st, err := os.Stat(absParent) + if err != nil || !st.IsDir() { + c.JSON(http.StatusBadRequest, gin.H{"error": "parent not found"}) + return + } + } + + var rel string + if parent == "" { + rel = name + } else { + rel = parent + "/" + name + } + absNew, err := h.resolveUnderChatUploads(rel) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if _, err := os.Stat(absNew); err == nil { + c.JSON(http.StatusConflict, gin.H{"error": "already exists"}) + return + } + if err := os.Mkdir(absNew, 0755); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + relOut, _ := filepath.Rel(root, absNew) + c.JSON(http.StatusOK, gin.H{"ok": true, "relativePath": filepath.ToSlash(relOut)}) +} + +type chatUploadRenameBody struct { + Path string `json:"path"` + NewName string `json:"newName"` +} + +// Rename PUT /api/chat-uploads/rename +func (h *ChatUploadsHandler) Rename(c *gin.Context) { + var body chatUploadRenameBody + if err := c.ShouldBindJSON(&body); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid body"}) + return + } + if !h.pathAllowed(c, body.Path) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + newName := strings.TrimSpace(body.NewName) + if newName == "" || strings.ContainsAny(newName, `/\`) { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid newName"}) + return + } + abs, err := h.resolveUnderChatUploads(body.Path) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + dir := filepath.Dir(abs) + newAbs := filepath.Join(dir, filepath.Base(newName)) + root, _ := h.absRoot() + newAbs, _ = filepath.Abs(newAbs) + if newAbs != root && !strings.HasPrefix(newAbs, root+string(filepath.Separator)) { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid target path"}) + return + } + if err := os.Rename(abs, newAbs); err != nil { + if os.IsNotExist(err) { + c.JSON(http.StatusNotFound, gin.H{"error": "file not found"}) + return + } + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + newRel, _ := filepath.Rel(root, newAbs) + oldRel := filepath.ToSlash(filepath.Clean(filepath.FromSlash(body.Path))) + _ = h.db.RenameChatUploadArtifactPath(oldRel, filepath.ToSlash(newRel)) + c.JSON(http.StatusOK, gin.H{"ok": true, "relativePath": filepath.ToSlash(newRel)}) +} + +type chatUploadContentBody struct { + Path string `json:"path"` + Content string `json:"content"` +} + +// GetContent GET /api/chat-uploads/content?path=... +func (h *ChatUploadsHandler) GetContent(c *gin.Context) { + p := c.Query("path") + if !h.pathAllowed(c, p) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + abs, err := h.resolveUnderChatUploads(p) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + st, err := os.Stat(abs) + if err != nil || st.IsDir() { + c.JSON(http.StatusNotFound, gin.H{"error": "file not found"}) + return + } + if st.Size() > maxChatUploadEditBytes { + c.JSON(http.StatusRequestEntityTooLarge, gin.H{"error": "file too large for editor"}) + return + } + b, err := os.ReadFile(abs) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if !utf8.Valid(b) { + c.JSON(http.StatusBadRequest, gin.H{"error": "binary file not editable in UI"}) + return + } + c.JSON(http.StatusOK, gin.H{"content": string(b)}) +} + +// PutContent PUT /api/chat-uploads/content +func (h *ChatUploadsHandler) PutContent(c *gin.Context) { + var body chatUploadContentBody + if err := c.ShouldBindJSON(&body); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid body"}) + return + } + if !h.pathAllowed(c, body.Path) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + if !utf8.ValidString(body.Content) { + c.JSON(http.StatusBadRequest, gin.H{"error": "content must be valid UTF-8"}) + return + } + if len(body.Content) > maxChatUploadEditBytes { + c.JSON(http.StatusRequestEntityTooLarge, gin.H{"error": "content too large"}) + return + } + abs, err := h.resolveUnderChatUploads(body.Path) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if err := os.WriteFile(abs, []byte(body.Content), 0644); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"ok": true}) +} + +func chatUploadShortRand(n int) string { + const letters = "0123456789abcdef" + b := make([]byte, n) + _, _ = rand.Read(b) + for i := range b { + b[i] = letters[int(b[i])%len(letters)] + } + return string(b) +} + +// Upload POST /api/chat-uploads multipart: file；conversationId 可选；relativeDir 可选（chat_uploads 下目录的相对路径，将文件直接上传至该目录） +func (h *ChatUploadsHandler) Upload(c *gin.Context) { + fh, err := c.FormFile("file") + if err != nil || fh == nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "missing file"}) + return + } + root, err := h.absRoot() + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + var targetDir string + targetRel := strings.TrimSpace(c.PostForm("relativeDir")) + if targetRel != "" { + if !h.pathAllowed(c, targetRel) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + absDir, err := h.resolveUnderChatUploads(targetRel) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + st, err := os.Stat(absDir) + if err != nil { + if os.IsNotExist(err) { + if err := os.MkdirAll(absDir, 0755); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + } else { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + } else if !st.IsDir() { + c.JSON(http.StatusBadRequest, gin.H{"error": "relativeDir is not a directory"}) + return + } + targetDir = absDir + } else { + convID := strings.TrimSpace(c.PostForm("conversationId")) + dateStr := time.Now().Format("2006-01-02") + if !h.pathAllowed(c, filepath.ToSlash(filepath.Join(dateStr, convID))) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + convDir := convID + if convDir == "" { + convDir = "_manual" + } else { + convDir = strings.ReplaceAll(convDir, string(filepath.Separator), "_") + } + targetDir = filepath.Join(root, dateStr, convDir) + if err := os.MkdirAll(targetDir, 0755); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + } + baseName := filepath.Base(fh.Filename) + if baseName == "" || baseName == "." { + baseName = "file" + } + baseName = strings.ReplaceAll(baseName, string(filepath.Separator), "_") + ext := filepath.Ext(baseName) + nameNoExt := strings.TrimSuffix(baseName, ext) + suffix := fmt.Sprintf("_%s_%s", time.Now().Format("150405"), chatUploadShortRand(6)) + var unique string + if ext != "" { + unique = nameNoExt + suffix + ext + } else { + unique = baseName + suffix + } + fullPath := filepath.Join(targetDir, unique) + src, err := fh.Open() + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + defer src.Close() + dst, err := os.Create(fullPath) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + defer dst.Close() + if _, err := io.Copy(dst, src); err != nil { + _ = os.Remove(fullPath) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + rel, _ := filepath.Rel(root, fullPath) + absSaved, _ := filepath.Abs(fullPath) + if session, ok := security.CurrentSession(c); ok { + conversationID := strings.TrimSpace(c.PostForm("conversationId")) + if conversationID == "" { + parts := strings.Split(filepath.ToSlash(rel), "/") + if len(parts) >= 2 { + conversationID = parts[1] + } + } + _ = h.db.UpsertChatUploadArtifact(filepath.ToSlash(rel), conversationID, session.UserID) + } + if h.audit != nil { + h.audit.RecordOK(c, "file", "upload", "上传对话附件", "chat_upload", filepath.ToSlash(rel), map[string]interface{}{ + "name": unique, + }) + } + c.JSON(http.StatusOK, gin.H{ + "ok": true, + "relativePath": filepath.ToSlash(rel), + "absolutePath": absSaved, + "name": unique, + }) +} diff --git a/internal/handler/config.go b/internal/handler/config.go new file mode 100644 index 00000000..cde56cdd --- /dev/null +++ b/internal/handler/config.go @@ -0,0 +1,2656 @@ +package handler + +import ( + "bytes" + "context" + "fmt" + "net/http" + "os" + "path/filepath" + "sort" + "strconv" + "strings" + "sync" + "time" + + "cyberstrike-ai/internal/agents" + "cyberstrike-ai/internal/audit" + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/knowledge" + "cyberstrike-ai/internal/mcp" + "cyberstrike-ai/internal/mcp/builtin" + "cyberstrike-ai/internal/openai" + "cyberstrike-ai/internal/security" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" + "gopkg.in/yaml.v3" +) + +// KnowledgeToolRegistrar 知识库工具注册器接口 +type KnowledgeToolRegistrar func() error + +// VulnerabilityToolRegistrar 漏洞工具注册器接口 +type VulnerabilityToolRegistrar func() error + +// WebshellToolRegistrar WebShell 工具注册器接口（ApplyConfig 时重新注册） +type WebshellToolRegistrar func() error + +// SkillsToolRegistrar Skills工具注册器接口 +type SkillsToolRegistrar func() error + +// BatchTaskToolRegistrar 批量任务 MCP 工具注册器（ApplyConfig 时重新注册） +type BatchTaskToolRegistrar func() error + +// C2ToolRegistrar C2 MCP 工具注册器（ApplyConfig 时 ClearTools 之后调用） +type C2ToolRegistrar func() error + +// C2Runtime ApplyConfig 时按配置启停 C2 子系统（由 internal/app.App 实现） +type C2Runtime interface { + ReconcileC2AfterConfigApply() error +} + +// RetrieverUpdater 检索器更新接口 +type RetrieverUpdater interface { + UpdateConfig(config *knowledge.RetrievalConfig) +} + +// KnowledgeInitializer 知识库初始化器接口 +type KnowledgeInitializer func() (*KnowledgeHandler, error) + +// AppUpdater App更新接口（用于更新App中的知识库组件） +type AppUpdater interface { + UpdateKnowledgeComponents(handler *KnowledgeHandler, manager interface{}, retriever interface{}, indexer interface{}) +} + +// RobotRestarter 机器人连接重启器（用于配置应用后重启钉钉/飞书长连接） +type RobotRestarter interface { + RestartRobotConnections() +} + +// ConfigHandler 配置处理器 +type ConfigHandler struct { + configPath string + config *config.Config + mcpServer *mcp.Server + executor *security.Executor + agent AgentUpdater // Agent接口，用于更新Agent配置 + attackChainHandler AttackChainUpdater // 攻击链处理器接口，用于更新配置 + externalMCPMgr *mcp.ExternalMCPManager // 外部MCP管理器 + knowledgeToolRegistrar KnowledgeToolRegistrar // 知识库工具注册器（可选） + vulnerabilityToolRegistrar VulnerabilityToolRegistrar // 漏洞工具注册器（可选） + webshellToolRegistrar WebshellToolRegistrar // WebShell 工具注册器（可选） + skillsToolRegistrar SkillsToolRegistrar // Skills工具注册器（可选） + batchTaskToolRegistrar BatchTaskToolRegistrar // 批量任务 MCP 工具（可选） + c2ToolRegistrar C2ToolRegistrar // C2 MCP 工具（可选） + c2Runtime C2Runtime // C2 启停（可选） + retrieverUpdater RetrieverUpdater // 检索器更新器（可选） + knowledgeInitializer KnowledgeInitializer // 知识库初始化器（可选） + appUpdater AppUpdater // App更新器（可选） + robotRestarter RobotRestarter // 机器人连接重启器（可选），ApplyConfig 时重启钉钉/飞书 + audit *audit.Service + db *database.DB + logger *zap.Logger + mu sync.RWMutex + lastEmbeddingConfig *config.EmbeddingConfig // 上一次的嵌入模型配置（用于检测变更） +} + +func (h *ConfigHandler) SetDB(db *database.DB) { + h.db = db +} + +func (h *ConfigHandler) validateRobotServiceAccounts(robots config.RobotsConfig) error { + if h.db == nil { + return fmt.Errorf("RBAC 服务不可用，无法校验机器人服务账号") + } + for platform, userID := range robots.ServiceAccountUserIDs() { + user, err := h.db.GetRBACUserByID(userID) + if err != nil { + return fmt.Errorf("robots.%s.auth.service_user_id 对应用户不存在", platform) + } + if !user.Enabled { + return fmt.Errorf("robots.%s.auth.service_user_id 对应用户已禁用", platform) + } + } + return nil +} + +// AttackChainUpdater 攻击链处理器更新接口 +type AttackChainUpdater interface { + UpdateConfig(cfg *config.OpenAIConfig) +} + +// AgentUpdater Agent更新接口 +type AgentUpdater interface { + UpdateConfig(cfg *config.OpenAIConfig) + UpdateMaxIterations(maxIterations int) + UpdateToolDescriptionMode(mode string) +} + +// NewConfigHandler 创建新的配置处理器 +func NewConfigHandler(configPath string, cfg *config.Config, mcpServer *mcp.Server, executor *security.Executor, agent AgentUpdater, attackChainHandler AttackChainUpdater, externalMCPMgr *mcp.ExternalMCPManager, logger *zap.Logger) *ConfigHandler { + // 保存初始的嵌入模型配置（如果知识库已启用） + var lastEmbeddingConfig *config.EmbeddingConfig + if cfg.Knowledge.Enabled { + lastEmbeddingConfig = &config.EmbeddingConfig{ + Provider: cfg.Knowledge.Embedding.Provider, + Model: cfg.Knowledge.Embedding.Model, + BaseURL: cfg.Knowledge.Embedding.BaseURL, + APIKey: cfg.Knowledge.Embedding.APIKey, + } + } + return &ConfigHandler{ + configPath: configPath, + config: cfg, + mcpServer: mcpServer, + executor: executor, + agent: agent, + attackChainHandler: attackChainHandler, + externalMCPMgr: externalMCPMgr, + logger: logger, + lastEmbeddingConfig: lastEmbeddingConfig, + } +} + +// SetKnowledgeToolRegistrar 设置知识库工具注册器 +func (h *ConfigHandler) SetKnowledgeToolRegistrar(registrar KnowledgeToolRegistrar) { + h.mu.Lock() + defer h.mu.Unlock() + h.knowledgeToolRegistrar = registrar +} + +// SetVulnerabilityToolRegistrar 设置漏洞工具注册器 +func (h *ConfigHandler) SetVulnerabilityToolRegistrar(registrar VulnerabilityToolRegistrar) { + h.mu.Lock() + defer h.mu.Unlock() + h.vulnerabilityToolRegistrar = registrar +} + +// SetWebshellToolRegistrar 设置 WebShell 工具注册器 +func (h *ConfigHandler) SetWebshellToolRegistrar(registrar WebshellToolRegistrar) { + h.mu.Lock() + defer h.mu.Unlock() + h.webshellToolRegistrar = registrar +} + +// SetSkillsToolRegistrar 设置Skills工具注册器 +func (h *ConfigHandler) SetSkillsToolRegistrar(registrar SkillsToolRegistrar) { + h.mu.Lock() + defer h.mu.Unlock() + h.skillsToolRegistrar = registrar +} + +// SetBatchTaskToolRegistrar 设置批量任务 MCP 工具注册器 +func (h *ConfigHandler) SetBatchTaskToolRegistrar(registrar BatchTaskToolRegistrar) { + h.mu.Lock() + defer h.mu.Unlock() + h.batchTaskToolRegistrar = registrar +} + +// SetC2ToolRegistrar 设置 C2 MCP 工具注册器 +func (h *ConfigHandler) SetC2ToolRegistrar(registrar C2ToolRegistrar) { + h.mu.Lock() + defer h.mu.Unlock() + h.c2ToolRegistrar = registrar +} + +// SetC2Runtime 设置 C2 运行时（Apply 时启停） +func (h *ConfigHandler) SetC2Runtime(rt C2Runtime) { + h.mu.Lock() + defer h.mu.Unlock() + h.c2Runtime = rt +} + +// SetRetrieverUpdater 设置检索器更新器 +func (h *ConfigHandler) SetRetrieverUpdater(updater RetrieverUpdater) { + h.mu.Lock() + defer h.mu.Unlock() + h.retrieverUpdater = updater +} + +// SetKnowledgeInitializer 设置知识库初始化器 +func (h *ConfigHandler) SetKnowledgeInitializer(initializer KnowledgeInitializer) { + h.mu.Lock() + defer h.mu.Unlock() + h.knowledgeInitializer = initializer +} + +// SetAppUpdater 设置App更新器 +func (h *ConfigHandler) SetAppUpdater(updater AppUpdater) { + h.mu.Lock() + defer h.mu.Unlock() + h.appUpdater = updater +} + +// SetRobotRestarter 设置机器人连接重启器（ApplyConfig 时用于重启钉钉/飞书长连接） +func (h *ConfigHandler) SetRobotRestarter(restarter RobotRestarter) { + h.mu.Lock() + defer h.mu.Unlock() + h.robotRestarter = restarter +} + +// SetAudit wires platform audit logging. +func (h *ConfigHandler) SetAudit(s *audit.Service) { + h.mu.Lock() + defer h.mu.Unlock() + h.audit = s +} + +// ApplyWechatRobotBinding 微信 iLink 扫码绑定成功后写入配置并重启机器人连接 +func (h *ConfigHandler) ApplyWechatRobotBinding(wc config.RobotWechatConfig) error { + h.mu.Lock() + wc.Enabled = true + h.config.Robots.Wechat = wc + h.mu.Unlock() + if err := h.saveConfig(); err != nil { + return err + } + if h.robotRestarter != nil { + h.robotRestarter.RestartRobotConnections() + } + h.logger.Info("微信机器人绑定已保存", + zap.String("ilink_bot_id", wc.ILinkBotID), + zap.Bool("enabled", wc.Enabled), + ) + return nil +} + +// GetConfigResponse 获取配置响应 +type GetConfigResponse struct { + AI config.AIConfig `json:"ai"` + OpenAI config.OpenAIConfig `json:"openai"` + Vision config.VisionConfig `json:"vision"` + FOFA config.FofaConfig `json:"fofa"` + ZoomEye config.SpaceSearchConfig `json:"zoomeye"` + Quake config.SpaceSearchConfig `json:"quake"` + Shodan config.SpaceSearchConfig `json:"shodan"` + MCP config.MCPConfig `json:"mcp"` + Tools []ToolConfigInfo `json:"tools"` + Agent config.AgentConfig `json:"agent"` + Hitl config.HitlConfig `json:"hitl,omitempty"` + Knowledge config.KnowledgeConfig `json:"knowledge"` + Robots config.RobotsConfig `json:"robots,omitempty"` + MultiAgent config.MultiAgentPublic `json:"multi_agent,omitempty"` + C2 config.C2Public `json:"c2"` +} + +// ToolConfigInfo 工具配置信息 +type ToolConfigInfo struct { + Name string `json:"name"` + Description string `json:"description"` + 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表示未指定角色或使用所有工具） + InputSchema map[string]interface{} `json:"input_schema,omitempty"` // 工具参数 JSON Schema（用于前端展示详情） +} + +// GetConfig 获取当前配置 +func (h *ConfigHandler) GetConfig(c *gin.Context) { + h.mu.RLock() + defer h.mu.RUnlock() + + // 获取工具列表（包含内部和外部工具） + // 首先从配置文件获取工具 + configToolMap := make(map[string]bool) + tools := make([]ToolConfigInfo, 0, len(h.config.Security.Tools)) + + for _, tool := range h.config.Security.Tools { + configToolMap[tool.Name] = true + info := ToolConfigInfo{ + Name: tool.Name, + Description: h.pickToolDescription(tool.ShortDescription, tool.Description), + Enabled: tool.Enabled, + IsExternal: false, + } + tools = append(tools, info) + } + + // 从MCP服务器获取所有已注册的工具（包括直接注册的工具，如知识检索工具） + if h.mcpServer != nil { + mcpTools := h.mcpServer.GetAllTools() + for _, mcpTool := range mcpTools { + if configToolMap[mcpTool.Name] { + continue + } + description := h.pickToolDescription(mcpTool.ShortDescription, mcpTool.Description) + tools = append(tools, ToolConfigInfo{ + Name: mcpTool.Name, + Description: description, + Enabled: true, + IsExternal: false, + }) + } + } + + // 获取外部MCP工具（走缓存，持锁期间通常不阻塞） + if h.externalMCPMgr != nil { + ctx := context.Background() + externalTools := h.getExternalMCPTools(ctx) + for _, toolInfo := range externalTools { + tools = append(tools, toolInfo) + } + } + + subAgentCount := len(h.config.MultiAgent.SubAgents) + agentsDir := strings.TrimSpace(h.config.AgentsDir) + if agentsDir == "" { + agentsDir = "agents" + } + if !filepath.IsAbs(agentsDir) { + agentsDir = filepath.Join(filepath.Dir(h.configPath), agentsDir) + } + if load, err := agents.LoadMarkdownAgentsDir(agentsDir); err == nil { + subAgentCount = len(agents.MergeYAMLAndMarkdown(h.config.MultiAgent.SubAgents, load.SubAgents)) + } + multiPub := config.MultiAgentPublic{ + Enabled: h.config.MultiAgent.Enabled, + RobotDefaultAgentMode: config.NormalizeRobotAgentMode(h.config.MultiAgent), + BatchUseMultiAgent: h.config.MultiAgent.BatchUseMultiAgent, + SubAgentCount: subAgentCount, + Orchestration: config.NormalizeMultiAgentOrchestration(h.config.MultiAgent.Orchestration), + PlanExecuteLoopMaxIterations: h.config.MultiAgent.PlanExecuteLoopMaxIterations, + SummarizationUserIntentLedgerMaxRunes: h.config.MultiAgent.EinoMiddleware.SummarizationUserIntentLedgerMaxRunesEffective(), + SummarizationUserIntentLedgerEntryMaxRunes: h.config.MultiAgent.EinoMiddleware.SummarizationUserIntentLedgerEntryMaxRunesEffective(), + LatestUserMessageMaxRunes: h.config.MultiAgent.EinoMiddleware.LatestUserMessageMaxRunesEffective(), + LatestUserMessageHeadRunes: h.config.MultiAgent.EinoMiddleware.LatestUserMessageHeadRunesEffective(), + LatestUserMessageTailRunes: h.config.MultiAgent.EinoMiddleware.LatestUserMessageTailRunesEffective(), + ModelRetryMaxRetries: h.config.MultiAgent.EinoMiddleware.ModelRetryMaxRetries, + ModelRetryMaxBackoffSec: h.config.MultiAgent.EinoMiddleware.ModelRetryMaxBackoffSec, + ModelFailoverChannels: append([]string(nil), h.config.MultiAgent.EinoMiddleware.ModelFailoverChannels...), + ModelFailoverMaxRetries: h.config.MultiAgent.EinoMiddleware.ModelFailoverMaxRetries, + ToolSearchAlwaysVisibleTools: append([]string(nil), h.config.MultiAgent.EinoMiddleware.ToolSearchAlwaysVisibleTools...), + ToolSearchAlwaysVisibleEffectiveTools: mergeToolNameLists( + h.config.MultiAgent.EinoMiddleware.ToolSearchAlwaysVisibleTools, + builtin.GetAllBuiltinTools(), + ), + } + + c.JSON(http.StatusOK, GetConfigResponse{ + AI: h.config.AI, + OpenAI: h.config.OpenAI, + Vision: h.config.Vision, + FOFA: h.config.FOFA, + ZoomEye: h.config.ZoomEye, + Quake: h.config.Quake, + Shodan: h.config.Shodan, + MCP: h.config.MCP, + Tools: tools, + Agent: h.config.Agent, + Hitl: h.config.Hitl, + Knowledge: h.config.Knowledge, + C2: h.config.C2.Public(), + Robots: h.config.Robots, + MultiAgent: multiPub, + }) +} + +// GetToolsResponse 获取工具列表响应（分页） +type GetToolsResponse struct { + 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 获取工具列表（支持分页和搜索） +func (h *ConfigHandler) GetTools(c *gin.Context) { + c.Header("Cache-Control", "no-store, no-cache, must-revalidate") + + // 解析分页参数 + page := 1 + pageSize := 20 + if pageStr := c.Query("page"); pageStr != "" { + if p, err := strconv.Atoi(pageStr); err == nil && p > 0 { + page = p + } + } + if pageSizeStr := c.Query("page_size"); pageSizeStr != "" { + if ps, err := strconv.Atoi(pageSizeStr); err == nil && ps > 0 && ps <= 100 { + pageSize = ps + } + } + + // 解析搜索参数 + searchTerm := c.Query("search") + searchTermLower := "" + if searchTerm != "" { + searchTermLower = strings.ToLower(searchTerm) + } + + // 解析状态筛选: tool_filter=on|off（角色弹窗等优先，避免与网关/代理对 enabled 的特殊处理冲突） + // 兼容旧参数 enabled=true|false + var filterEnabled *bool + toolFilter := strings.TrimSpace(strings.ToLower(c.Query("tool_filter"))) + switch toolFilter { + case "on", "1", "true", "enabled": + v := true + filterEnabled = &v + case "off", "0", "false", "disabled": + v := false + filterEnabled = &v + default: + enabledFilter := strings.TrimSpace(c.Query("enabled")) + if enabledFilter == "true" { + v := true + filterEnabled = &v + } else if enabledFilter == "false" { + v := false + filterEnabled = &v + } + } + + includeExternal := true + if v := strings.TrimSpace(strings.ToLower(c.Query("include_external"))); v == "0" || v == "false" || v == "no" { + includeExternal = false + } + refreshExternal := false + if v := strings.TrimSpace(strings.ToLower(c.Query("refresh_external"))); v == "1" || v == "true" || v == "yes" { + refreshExternal = true + } + + // 按外部 MCP 名称筛选（MCP 管理页左侧卡片 → 右侧工具列表联动） + externalMCPFilter := strings.TrimSpace(c.Query("external_mcp")) + + // 快照配置后立即释放锁，避免外部 MCP 网络 IO 阻塞整个配置子系统 + h.mu.RLock() + securityTools := append([]config.ToolConfig(nil), h.config.Security.Tools...) + roles := h.config.Roles + toolDescriptionMode := h.config.Security.ToolDescriptionMode + mcpServer := h.mcpServer + externalMCPMgr := h.externalMCPMgr + h.mu.RUnlock() + + pickDesc := func(shortDesc, fullDesc string) string { + return pickToolDescriptionWithMode(toolDescriptionMode, shortDesc, fullDesc) + } + + // 解析角色参数，用于过滤工具并标注启用状态 + roleName := c.Query("role") + var roleToolsSet map[string]bool // 角色配置的工具集合 + var roleUsesAllTools bool = true // 角色是否使用所有工具（默认角色） + if roleName != "" && roleName != "默认" && roles != nil { + if role, exists := 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(securityTools)) + for _, tool := range securityTools { + configToolMap[tool.Name] = true + toolInfo := ToolConfigInfo{ + Name: tool.Name, + Description: pickDesc(tool.ShortDescription, tool.Description), + Enabled: tool.Enabled, + IsExternal: false, + } + + // 根据角色配置标注工具状态 + 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) + descLower := strings.ToLower(toolInfo.Description) + if !strings.Contains(nameLower, searchTermLower) && !strings.Contains(descLower, searchTermLower) { + continue // 不匹配，跳过 + } + } + + // 状态筛选 + if filterEnabled != nil && toolInfo.Enabled != *filterEnabled { + continue + } + + allTools = append(allTools, toolInfo) + } + + // 从MCP服务器获取所有已注册的工具（包括直接注册的工具，如知识检索工具） + if mcpServer != nil { + mcpTools := mcpServer.GetAllTools() + for _, mcpTool := range mcpTools { + // 跳过已经在配置文件中的工具（避免重复） + if configToolMap[mcpTool.Name] { + continue + } + + description := pickDesc(mcpTool.ShortDescription, mcpTool.Description) + + toolInfo := ToolConfigInfo{ + Name: mcpTool.Name, + Description: description, + Enabled: true, + 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) + descLower := strings.ToLower(toolInfo.Description) + if !strings.Contains(nameLower, searchTermLower) && !strings.Contains(descLower, searchTermLower) { + continue // 不匹配，跳过 + } + } + + // 状态筛选 + if filterEnabled != nil && toolInfo.Enabled != *filterEnabled { + continue + } + + allTools = append(allTools, toolInfo) + } + } + + // 获取外部MCP工具（可走缓存，不持有 config 锁） + if includeExternal && externalMCPMgr != nil { + if refreshExternal { + externalMCPMgr.InvalidateAllToolCaches() + } + ctx := context.Background() + externalTools := h.getExternalMCPToolsWithManager(ctx, externalMCPMgr, pickDesc) + + // 应用搜索过滤和角色配置 + for _, toolInfo := range externalTools { + // 搜索过滤 + if searchTermLower != "" { + nameLower := strings.ToLower(toolInfo.Name) + descLower := strings.ToLower(toolInfo.Description) + if !strings.Contains(nameLower, searchTermLower) && !strings.Contains(descLower, searchTermLower) { + continue // 不匹配，跳过 + } + } + + // 根据角色配置标注工具状态 + if roleName != "" { + if roleUsesAllTools { + // 角色使用所有工具，标注启用的工具为role_enabled=true + roleEnabled := toolInfo.Enabled + toolInfo.RoleEnabled = &roleEnabled + } else { + // 角色配置了工具列表，检查工具是否在列表中 + // 外部工具使用 "mcpName::toolName" 格式作为key + externalToolKey := fmt.Sprintf("%s::%s", toolInfo.ExternalMCP, toolInfo.Name) + if roleToolsSet[externalToolKey] { + roleEnabled := toolInfo.Enabled // 工具必须在角色列表中且本身启用 + toolInfo.RoleEnabled = &roleEnabled + } else { + // 不在角色列表中，标记为false + roleEnabled := false + toolInfo.RoleEnabled = &roleEnabled + } + } + } + + // 状态筛选 + if filterEnabled != nil && toolInfo.Enabled != *filterEnabled { + continue + } + + allTools = append(allTools, toolInfo) + } + } + + // 如果角色配置了工具列表，过滤工具（只保留列表中的工具，但保留其他工具并标记为禁用） + // 注意：这里我们不直接过滤掉工具，而是保留所有工具，但通过 role_enabled 字段标注状态 + // 这样前端可以显示所有工具，并标注哪些工具在当前角色中可用 + + if externalMCPFilter != "" { + filtered := make([]ToolConfigInfo, 0) + for _, tool := range allTools { + if tool.IsExternal && tool.ExternalMCP == externalMCPFilter { + filtered = append(filtered, tool) + } + } + allTools = filtered + } + + // 统一按名称排序后再分页，避免配置文件中顺序导致「全部」与「仅已启用」前几页看起来完全一致 + sort.SliceStable(allTools, func(i, j int) bool { + key := func(t ToolConfigInfo) string { + if t.IsExternal && t.ExternalMCP != "" { + return strings.ToLower(t.ExternalMCP + "::" + t.Name) + } + return strings.ToLower(t.Name) + } + return key(allTools[i]) < key(allTools[j]) + }) + + 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 + } + + // 计算分页范围 + offset := (page - 1) * pageSize + end := offset + pageSize + if end > total { + end = total + } + + var tools []ToolConfigInfo + if offset < total { + tools = allTools[offset:end] + } else { + tools = []ToolConfigInfo{} + } + + c.JSON(http.StatusOK, GetToolsResponse{ + Tools: tools, + Total: total, + TotalEnabled: totalEnabled, + Page: page, + PageSize: pageSize, + TotalPages: totalPages, + }) +} + +// UpdateConfigRequest 更新配置请求 +type UpdateConfigRequest struct { + AI *config.AIConfig `json:"ai,omitempty"` + OpenAI *config.OpenAIConfig `json:"openai,omitempty"` + Vision *config.VisionConfig `json:"vision,omitempty"` + FOFA *config.FofaConfig `json:"fofa,omitempty"` + ZoomEye *config.SpaceSearchConfig `json:"zoomeye,omitempty"` + Quake *config.SpaceSearchConfig `json:"quake,omitempty"` + Shodan *config.SpaceSearchConfig `json:"shodan,omitempty"` + MCP *config.MCPConfig `json:"mcp,omitempty"` + Tools []ToolEnableStatus `json:"tools,omitempty"` + Agent *AgentConfigUpdate `json:"agent,omitempty"` + Hitl *config.HitlConfig `json:"hitl,omitempty"` + Knowledge *config.KnowledgeConfig `json:"knowledge,omitempty"` + Robots *config.RobotsConfig `json:"robots,omitempty"` + MultiAgent *config.MultiAgentAPIUpdate `json:"multi_agent,omitempty"` + C2 *config.C2APIUpdate `json:"c2,omitempty"` +} + +// AgentConfigUpdate 用于 PATCH /api/config 的 agent 段：仅 JSON 中出现的字段（指针非 nil）覆盖内存配置。 +// 避免旧版「整包替换 *AgentConfig」时，未传的整型字段被反序列化为 0 误覆盖（例如 tool_timeout_minutes 变成 0）。 +type AgentConfigUpdate struct { + MaxIterations *int `json:"max_iterations,omitempty"` + ToolTimeoutMinutes *int `json:"tool_timeout_minutes,omitempty"` + ToolWaitTimeoutSeconds *int `json:"tool_wait_timeout_seconds,omitempty"` + ExternalMCPMaxConcurrentPerServer *int `json:"external_mcp_max_concurrent_per_server,omitempty"` + ExternalMCPMaxConcurrentTotal *int `json:"external_mcp_max_concurrent_total,omitempty"` + ExternalMCPCircuitFailureThreshold *int `json:"external_mcp_circuit_failure_threshold,omitempty"` + ExternalMCPCircuitCooldownSeconds *int `json:"external_mcp_circuit_cooldown_seconds,omitempty"` + SystemPromptPath *string `json:"system_prompt_path,omitempty"` +} + +func applyAgentConfigUpdate(dst *config.AgentConfig, src *AgentConfigUpdate) { + if dst == nil || src == nil { + return + } + if src.MaxIterations != nil { + dst.MaxIterations = *src.MaxIterations + } + if src.ToolTimeoutMinutes != nil { + dst.ToolTimeoutMinutes = *src.ToolTimeoutMinutes + } + if src.ToolWaitTimeoutSeconds != nil { + dst.ToolWaitTimeoutSeconds = *src.ToolWaitTimeoutSeconds + } + if src.ExternalMCPMaxConcurrentPerServer != nil { + dst.ExternalMCPMaxConcurrentPerServer = *src.ExternalMCPMaxConcurrentPerServer + } + if src.ExternalMCPMaxConcurrentTotal != nil { + dst.ExternalMCPMaxConcurrentTotal = *src.ExternalMCPMaxConcurrentTotal + } + if src.ExternalMCPCircuitFailureThreshold != nil { + dst.ExternalMCPCircuitFailureThreshold = *src.ExternalMCPCircuitFailureThreshold + } + if src.ExternalMCPCircuitCooldownSeconds != nil { + dst.ExternalMCPCircuitCooldownSeconds = *src.ExternalMCPCircuitCooldownSeconds + } + if src.SystemPromptPath != nil { + dst.SystemPromptPath = *src.SystemPromptPath + } +} + +// ToolEnableStatus 工具启用状态 +type ToolEnableStatus struct { + Name string `json:"name"` + Enabled bool `json:"enabled"` + IsExternal bool `json:"is_external,omitempty"` // 是否为外部MCP工具 + ExternalMCP string `json:"external_mcp,omitempty"` // 外部MCP名称（如果是外部工具） +} + +// UpdateConfig 更新配置 +func (h *ConfigHandler) UpdateConfig(c *gin.Context) { + var req UpdateConfigRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()}) + return + } + + h.mu.Lock() + defer h.mu.Unlock() + + // 更新OpenAI配置 + if req.AI != nil { + h.config.AI = *req.AI + h.config.ApplyDefaultAIChannel() + h.logger.Info("更新 AI 通道配置", + zap.String("default_channel", h.config.AI.DefaultChannel), + zap.Int("channels", len(h.config.AI.Channels)), + ) + } + if req.OpenAI != nil { + h.config.OpenAI = *req.OpenAI + h.config.AI.EnsureDefaultFromOpenAI(h.config.OpenAI) + if def := config.NormalizeAIChannelID(h.config.AI.DefaultChannel); def != "" { + h.config.AI.Channels[def] = config.AIChannelFromOpenAI(def, "Default", h.config.OpenAI) + } + h.logger.Info("更新OpenAI配置", + zap.String("base_url", h.config.OpenAI.BaseURL), + zap.String("model", h.config.OpenAI.Model), + ) + } + + if req.Vision != nil { + h.config.Vision = *req.Vision + h.logger.Info("更新 Vision 配置", + zap.Bool("enabled", h.config.Vision.Enabled), + zap.String("model", h.config.Vision.Model), + ) + } + + // 更新FOFA配置 + if req.FOFA != nil { + h.config.FOFA = *req.FOFA + h.logger.Info("更新FOFA配置", zap.String("base_url", h.config.FOFA.BaseURL)) + } + if req.ZoomEye != nil { + h.config.ZoomEye = *req.ZoomEye + h.logger.Info("更新ZoomEye配置", zap.String("base_url", h.config.ZoomEye.BaseURL)) + } + if req.Quake != nil { + h.config.Quake = *req.Quake + h.logger.Info("更新Quake配置", zap.String("base_url", h.config.Quake.BaseURL)) + } + if req.Shodan != nil { + h.config.Shodan = *req.Shodan + h.logger.Info("更新Shodan配置", zap.String("base_url", h.config.Shodan.BaseURL)) + } + + // 更新MCP配置 + if req.MCP != nil { + h.config.MCP = *req.MCP + h.logger.Info("更新MCP配置", + zap.Bool("enabled", h.config.MCP.Enabled), + zap.String("host", h.config.MCP.Host), + zap.Int("port", h.config.MCP.Port), + ) + } + + // 更新Agent配置（按字段合并，避免部分 JSON 把未出现的字段写成 0） + if req.Agent != nil { + applyAgentConfigUpdate(&h.config.Agent, req.Agent) + h.logger.Info("更新Agent配置", + zap.Int("max_iterations", h.config.Agent.MaxIterations), + zap.Int("tool_timeout_minutes", h.config.Agent.ToolTimeoutMinutes), + zap.Int("tool_wait_timeout_seconds", h.config.Agent.ToolWaitTimeoutSeconds), + zap.Int("external_mcp_max_concurrent_per_server", h.config.Agent.ExternalMCPMaxConcurrentPerServer), + zap.Int("external_mcp_max_concurrent_total", h.config.Agent.ExternalMCPMaxConcurrentTotal), + zap.Int("external_mcp_circuit_failure_threshold", h.config.Agent.ExternalMCPCircuitFailureThreshold), + zap.Int("external_mcp_circuit_cooldown_seconds", h.config.Agent.ExternalMCPCircuitCooldownSeconds), + ) + if h.agent != nil && req.Agent.MaxIterations != nil { + h.agent.UpdateMaxIterations(h.config.Agent.MaxIterations) + } + if h.executor != nil { + h.executor.SetToolOutputMaxBytes(h.config.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective()) + h.executor.SetToolOutputSpillRoot(h.config.MultiAgent.EinoMiddleware.ReductionRootDir) + } + if h.mcpServer != nil { + h.mcpServer.ConfigureHTTPToolCallTimeoutFromAgentMinutes(h.config.Agent.ToolTimeoutMinutes) + h.mcpServer.ConfigureToolWaitTimeoutSeconds(h.config.Agent.ToolWaitTimeoutSeconds) + h.mcpServer.ConfigureToolResultMaxBytes(h.config.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective()) + h.mcpServer.ConfigureToolResultSpillRoot(h.config.MultiAgent.EinoMiddleware.ReductionRootDir) + } + if h.externalMCPMgr != nil { + h.externalMCPMgr.ConfigureToolWaitTimeoutSeconds(h.config.Agent.ToolWaitTimeoutSeconds) + h.externalMCPMgr.ConfigureToolResultMaxBytes(h.config.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective()) + h.externalMCPMgr.ConfigureToolResultSpillRoot(h.config.MultiAgent.EinoMiddleware.ReductionRootDir) + h.externalMCPMgr.ConfigureResilience(mcp.ExternalMCPResilienceConfig{ + MaxConcurrentPerServer: h.config.Agent.ExternalMCPMaxConcurrentPerServer, + MaxConcurrentTotal: h.config.Agent.ExternalMCPMaxConcurrentTotal, + CircuitFailureThreshold: h.config.Agent.ExternalMCPCircuitFailureThreshold, + CircuitCooldown: time.Duration(h.config.Agent.ExternalMCPCircuitCooldownSeconds) * time.Second, + }) + } + } + + if req.Hitl != nil { + h.config.Hitl.AuditModel = req.Hitl.AuditModel + h.config.Hitl.ToolWhitelist = mergeHitlToolWhitelistSlice(nil, req.Hitl.ToolWhitelist) + h.config.Hitl.DefaultReviewer = req.Hitl.EffectiveDefaultReviewer() + h.config.Hitl.AuditAgentPrompt = strings.TrimSpace(req.Hitl.AuditAgentPrompt) + h.config.Hitl.AuditAgentPromptReviewEdit = strings.TrimSpace(req.Hitl.AuditAgentPromptReviewEdit) + if req.Hitl.RetentionDays != nil { + v := *req.Hitl.RetentionDays + if v < 0 { + v = 0 + } + h.config.Hitl.RetentionDays = &v + } + h.logger.Info("更新HITL配置", + zap.String("default_reviewer", h.config.Hitl.DefaultReviewer), + zap.Int("tool_whitelist", len(h.config.Hitl.ToolWhitelist)), + ) + } + + // 更新Knowledge配置 + if req.Knowledge != nil { + // 保存旧的嵌入模型配置（用于检测变更） + if h.config.Knowledge.Enabled { + h.lastEmbeddingConfig = &config.EmbeddingConfig{ + Provider: h.config.Knowledge.Embedding.Provider, + Model: h.config.Knowledge.Embedding.Model, + BaseURL: h.config.Knowledge.Embedding.BaseURL, + APIKey: h.config.Knowledge.Embedding.APIKey, + } + } + h.config.Knowledge = *req.Knowledge + h.logger.Info("更新Knowledge配置", + zap.Bool("enabled", h.config.Knowledge.Enabled), + zap.String("base_path", h.config.Knowledge.BasePath), + zap.String("embedding_model", h.config.Knowledge.Embedding.Model), + zap.Int("retrieval_top_k", h.config.Knowledge.Retrieval.TopK), + zap.Float64("similarity_threshold", h.config.Knowledge.Retrieval.SimilarityThreshold), + ) + } + + // 更新机器人配置 + if req.Robots != nil { + if err := config.ValidateWecomConfig(req.Robots.Wecom); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if err := config.ValidateRobotsAuthorization(*req.Robots); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if err := h.validateRobotServiceAccounts(*req.Robots); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + h.config.Robots = *req.Robots + h.logger.Info("更新机器人配置", + zap.Bool("wechat_enabled", h.config.Robots.Wechat.Enabled), + zap.Bool("wecom_enabled", h.config.Robots.Wecom.Enabled), + zap.Bool("dingtalk_enabled", h.config.Robots.Dingtalk.Enabled), + zap.Bool("lark_enabled", h.config.Robots.Lark.Enabled), + zap.Bool("telegram_enabled", h.config.Robots.Telegram.Enabled), + zap.Bool("slack_enabled", h.config.Robots.Slack.Enabled), + zap.Bool("discord_enabled", h.config.Robots.Discord.Enabled), + zap.Bool("qq_enabled", h.config.Robots.QQ.Enabled), + ) + } + + if req.C2 != nil { + v := req.C2.Enabled + h.config.C2.Enabled = &v + h.logger.Info("更新C2配置", zap.Bool("enabled", v)) + } + + // 多代理标量（sub_agents 等仍由 config.yaml 维护） + if req.MultiAgent != nil { + h.config.MultiAgent.Enabled = req.MultiAgent.Enabled + h.config.MultiAgent.BatchUseMultiAgent = req.MultiAgent.BatchUseMultiAgent + if mode := strings.TrimSpace(req.MultiAgent.RobotDefaultAgentMode); mode != "" { + h.config.MultiAgent.RobotDefaultAgentMode = mode + } else { + h.config.MultiAgent.RobotDefaultAgentMode = "eino_single" + } + if req.MultiAgent.PlanExecuteLoopMaxIterations != nil { + h.config.MultiAgent.PlanExecuteLoopMaxIterations = *req.MultiAgent.PlanExecuteLoopMaxIterations + } + if req.MultiAgent.SummarizationUserIntentLedgerMaxRunes != nil { + v := *req.MultiAgent.SummarizationUserIntentLedgerMaxRunes + if v < 0 { + v = 0 + } + h.config.MultiAgent.EinoMiddleware.SummarizationUserIntentLedgerMaxRunes = v + } + if req.MultiAgent.SummarizationUserIntentLedgerEntryMaxRunes != nil { + v := *req.MultiAgent.SummarizationUserIntentLedgerEntryMaxRunes + if v < 0 { + v = 0 + } + h.config.MultiAgent.EinoMiddleware.SummarizationUserIntentLedgerEntryMaxRunes = v + } + if req.MultiAgent.LatestUserMessageMaxRunes != nil { + v := *req.MultiAgent.LatestUserMessageMaxRunes + if v < 0 { + v = 0 + } + h.config.MultiAgent.EinoMiddleware.LatestUserMessageMaxRunes = v + } + if req.MultiAgent.LatestUserMessageHeadRunes != nil { + v := *req.MultiAgent.LatestUserMessageHeadRunes + if v < 0 { + v = 0 + } + h.config.MultiAgent.EinoMiddleware.LatestUserMessageHeadRunes = v + } + if req.MultiAgent.LatestUserMessageTailRunes != nil { + v := *req.MultiAgent.LatestUserMessageTailRunes + if v < 0 { + v = 0 + } + h.config.MultiAgent.EinoMiddleware.LatestUserMessageTailRunes = v + } + if req.MultiAgent.ModelRetryMaxRetries != nil { + v := *req.MultiAgent.ModelRetryMaxRetries + if v < 0 { + v = 0 + } + h.config.MultiAgent.EinoMiddleware.ModelRetryMaxRetries = v + } + if req.MultiAgent.ModelRetryMaxBackoffSec != nil { + v := *req.MultiAgent.ModelRetryMaxBackoffSec + if v < 0 { + v = 0 + } + h.config.MultiAgent.EinoMiddleware.ModelRetryMaxBackoffSec = v + } + if req.MultiAgent.ModelFailoverChannels != nil { + h.config.MultiAgent.EinoMiddleware.ModelFailoverChannels = dedupeTrimmedStringList(*req.MultiAgent.ModelFailoverChannels) + } + if req.MultiAgent.ModelFailoverMaxRetries != nil { + v := *req.MultiAgent.ModelFailoverMaxRetries + if v < 0 { + v = 0 + } + h.config.MultiAgent.EinoMiddleware.ModelFailoverMaxRetries = v + } + if req.MultiAgent.ToolSearchAlwaysVisibleTools != nil { + h.config.MultiAgent.EinoMiddleware.ToolSearchAlwaysVisibleTools = dedupeToolNameList(*req.MultiAgent.ToolSearchAlwaysVisibleTools) + } + h.logger.Info("更新多代理配置", + zap.Bool("enabled", h.config.MultiAgent.Enabled), + zap.String("robot_default_agent_mode", config.NormalizeRobotAgentMode(h.config.MultiAgent)), + zap.Bool("batch_use_multi_agent", h.config.MultiAgent.BatchUseMultiAgent), + zap.Int("plan_execute_loop_max_iterations", h.config.MultiAgent.PlanExecuteLoopMaxIterations), + zap.Int("summarization_user_intent_ledger_max_runes", h.config.MultiAgent.EinoMiddleware.SummarizationUserIntentLedgerMaxRunesEffective()), + zap.Int("summarization_user_intent_ledger_entry_max_runes", h.config.MultiAgent.EinoMiddleware.SummarizationUserIntentLedgerEntryMaxRunesEffective()), + zap.Int("latest_user_message_max_runes", h.config.MultiAgent.EinoMiddleware.LatestUserMessageMaxRunesEffective()), + zap.Int("latest_user_message_head_runes", h.config.MultiAgent.EinoMiddleware.LatestUserMessageHeadRunesEffective()), + zap.Int("latest_user_message_tail_runes", h.config.MultiAgent.EinoMiddleware.LatestUserMessageTailRunesEffective()), + zap.Int("model_retry_max_retries", h.config.MultiAgent.EinoMiddleware.ModelRetryMaxRetries), + zap.Int("model_retry_max_backoff_sec", h.config.MultiAgent.EinoMiddleware.ModelRetryMaxBackoffSec), + zap.Int("model_failover_channels", len(h.config.MultiAgent.EinoMiddleware.ModelFailoverChannels)), + zap.Int("model_failover_max_retries", h.config.MultiAgent.EinoMiddleware.ModelFailoverMaxRetries), + zap.Int("tool_search_always_visible_tools", len(h.config.MultiAgent.EinoMiddleware.ToolSearchAlwaysVisibleTools)), + ) + } + + // 更新工具启用状态 + if req.Tools != nil { + // 分离内部工具和外部工具 + internalToolMap := make(map[string]bool) + // 外部工具状态：MCP名称 -> 工具名称 -> 启用状态 + externalMCPToolMap := make(map[string]map[string]bool) + + for _, toolStatus := range req.Tools { + if toolStatus.IsExternal && toolStatus.ExternalMCP != "" { + // 外部工具：保存每个工具的独立状态 + mcpName := toolStatus.ExternalMCP + if externalMCPToolMap[mcpName] == nil { + externalMCPToolMap[mcpName] = make(map[string]bool) + } + externalMCPToolMap[mcpName][toolStatus.Name] = toolStatus.Enabled + } else { + // 内部工具 + internalToolMap[toolStatus.Name] = toolStatus.Enabled + } + } + + // 更新内部工具状态 + for i := range h.config.Security.Tools { + if enabled, ok := internalToolMap[h.config.Security.Tools[i].Name]; ok { + h.config.Security.Tools[i].Enabled = enabled + h.logger.Info("更新工具启用状态", + zap.String("tool", h.config.Security.Tools[i].Name), + zap.Bool("enabled", enabled), + ) + } + } + + // 更新外部MCP工具状态 + if h.externalMCPMgr != nil { + for mcpName, toolStates := range externalMCPToolMap { + // 更新配置中的工具启用状态 + if h.config.ExternalMCP.Servers == nil { + h.config.ExternalMCP.Servers = make(map[string]config.ExternalMCPServerConfig) + } + cfg, exists := h.config.ExternalMCP.Servers[mcpName] + if !exists { + h.logger.Warn("外部MCP配置不存在", zap.String("mcp", mcpName)) + continue + } + + // 初始化ToolEnabled map + if cfg.ToolEnabled == nil { + cfg.ToolEnabled = make(map[string]bool) + } + + // 更新每个工具的启用状态 + for toolName, enabled := range toolStates { + cfg.ToolEnabled[toolName] = enabled + h.logger.Info("更新外部工具启用状态", + zap.String("mcp", mcpName), + zap.String("tool", toolName), + zap.Bool("enabled", enabled), + ) + } + + // 检查是否有任何工具启用，如果有则启用MCP + hasEnabledTool := false + for _, enabled := range cfg.ToolEnabled { + if enabled { + hasEnabledTool = true + break + } + } + + // 如果MCP之前未启用，但现在有工具启用，则启用MCP + // 如果MCP之前已启用，保持启用状态（允许部分工具禁用） + if !cfg.ExternalMCPEnable && hasEnabledTool { + cfg.ExternalMCPEnable = true + h.logger.Info("自动启用外部MCP（因为有工具启用）", zap.String("mcp", mcpName)) + } + + h.config.ExternalMCP.Servers[mcpName] = cfg + } + + // 同步更新 externalMCPMgr 中的配置，确保 GetConfigs() 返回最新配置 + // 在循环外部统一更新，避免重复调用 + h.externalMCPMgr.LoadConfigs(&h.config.ExternalMCP) + + // 处理MCP连接状态（异步启动，避免阻塞） + for mcpName := range externalMCPToolMap { + cfg := h.config.ExternalMCP.Servers[mcpName] + // 如果MCP需要启用，确保客户端已启动 + if cfg.ExternalMCPEnable { + // 启动外部MCP（如果未启动）- 异步执行，避免阻塞 + client, exists := h.externalMCPMgr.GetClient(mcpName) + if !exists || !client.IsConnected() { + go func(name string) { + if err := h.externalMCPMgr.StartClient(name); err != nil { + h.logger.Warn("启动外部MCP失败", + zap.String("mcp", name), + zap.Error(err), + ) + } else { + h.logger.Info("启动外部MCP", + zap.String("mcp", name), + ) + } + }(mcpName) + } + } + } + } + } + + // 保存配置到文件 + if err := h.saveConfig(); err != nil { + h.logger.Error("保存配置失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": "保存配置失败: " + err.Error()}) + return + } + + if h.audit != nil { + h.audit.RecordOK(c, "config", "update", "更新内存配置", "config", "", nil) + } + c.JSON(http.StatusOK, gin.H{"message": "配置已更新"}) +} + +// TestOpenAIRequest 测试OpenAI连接请求 +type TestOpenAIRequest struct { + Provider string `json:"provider"` + BaseURL string `json:"base_url"` + APIKey string `json:"api_key"` + Model string `json:"model"` +} + +// TestOpenAI 测试OpenAI API连接是否可用 +func (h *ConfigHandler) TestOpenAI(c *gin.Context) { + var req TestOpenAIRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()}) + return + } + + if strings.TrimSpace(req.APIKey) == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "API Key 不能为空"}) + return + } + if strings.TrimSpace(req.Model) == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "模型不能为空"}) + return + } + + baseURL := strings.TrimSuffix(strings.TrimSpace(req.BaseURL), "/") + if baseURL == "" { + if strings.EqualFold(strings.TrimSpace(req.Provider), "claude") { + baseURL = "https://api.anthropic.com" + } else { + baseURL = "https://api.openai.com/v1" + } + } + + // 构造一个最小的 chat completion 请求 + payload := map[string]interface{}{ + "model": req.Model, + "messages": []map[string]string{ + {"role": "user", "content": "Hi"}, + }, + "max_completion_tokens": 5, + } + + // 使用内部 openai Client 进行测试，若 provider 为 claude 会自动走桥接层 + tmpCfg := &config.OpenAIConfig{ + Provider: req.Provider, + BaseURL: baseURL, + APIKey: strings.TrimSpace(req.APIKey), + Model: req.Model, + } + client := openai.NewClient(tmpCfg, nil, h.logger) + + ctx, cancel := context.WithTimeout(c.Request.Context(), 30*time.Second) + defer cancel() + + start := time.Now() + var chatResp struct { + ID string `json:"id"` + Object string `json:"object"` + Model string `json:"model"` + Choices []struct { + Message struct { + Role string `json:"role"` + Content string `json:"content"` + } `json:"message"` + } `json:"choices"` + } + err := client.ChatCompletion(ctx, payload, &chatResp) + latency := time.Since(start) + + if err != nil { + if apiErr, ok := err.(*openai.APIError); ok { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "error": fmt.Sprintf("API 返回错误 (HTTP %d): %s", apiErr.StatusCode, apiErr.Body), + "status_code": apiErr.StatusCode, + }) + return + } + c.JSON(http.StatusOK, gin.H{ + "success": false, + "error": "连接失败: " + err.Error(), + }) + return + } + + // 严格校验：必须包含 choices 且有 assistant 回复 + if len(chatResp.Choices) == 0 { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "error": "API 响应缺少 choices 字段，请检查 Base URL 路径是否正确", + }) + return + } + if chatResp.ID == "" && chatResp.Model == "" { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "error": "API 响应格式不符合预期，请检查 Base URL 是否正确", + }) + return + } + + c.JSON(http.StatusOK, gin.H{ + "success": true, + "model": chatResp.Model, + "latency_ms": latency.Milliseconds(), + }) +} + +// ListModelsRequest 获取模型列表请求（OpenAI 兼容 GET /models）。 +type ListModelsRequest struct { + Provider string `json:"provider"` + BaseURL string `json:"base_url"` + APIKey string `json:"api_key"` +} + +// ListModels 代理调用上游 GET /models，返回可用模型 id 列表。 +func (h *ConfigHandler) ListModels(c *gin.Context) { + var req ListModelsRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()}) + return + } + + provider := strings.TrimSpace(req.Provider) + if provider == "" { + provider = "openai" + } + if strings.EqualFold(provider, "claude") { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "supported": false, + "error": "Claude (Anthropic Messages API) 不支持自动获取模型列表，请手动填写", + }) + return + } + + if strings.TrimSpace(req.APIKey) == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "API Key 不能为空"}) + return + } + + baseURL := strings.TrimSuffix(strings.TrimSpace(req.BaseURL), "/") + if baseURL == "" { + baseURL = "https://api.openai.com/v1" + } + + tmpCfg := &config.OpenAIConfig{ + Provider: provider, + BaseURL: baseURL, + APIKey: strings.TrimSpace(req.APIKey), + } + client := openai.NewClient(tmpCfg, nil, h.logger) + + ctx, cancel := context.WithTimeout(c.Request.Context(), 30*time.Second) + defer cancel() + + models, err := client.ListModels(ctx) + if err != nil { + if apiErr, ok := err.(*openai.APIError); ok { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "supported": true, + "error": fmt.Sprintf("API 返回错误 (HTTP %d): %s", apiErr.StatusCode, apiErr.Body), + }) + return + } + c.JSON(http.StatusOK, gin.H{ + "success": false, + "supported": true, + "error": err.Error(), + }) + return + } + + c.JSON(http.StatusOK, gin.H{ + "success": true, + "supported": true, + "models": models, + "count": len(models), + }) +} + +// TestVisionRequest 测试 Vision 模型连接；vision.api_key/base_url 留空时可传 openai 段作回退。 +type TestVisionRequest struct { + Vision config.VisionConfig `json:"vision"` + OpenAI config.OpenAIConfig `json:"openai,omitempty"` +} + +// TestVision 测试视觉模型 API 连接（最小 chat completion）。 +func (h *ConfigHandler) TestVision(c *gin.Context) { + var req TestVisionRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()}) + return + } + oa := req.Vision.OpenAICfgEffective(req.OpenAI) + if strings.TrimSpace(oa.APIKey) == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "API Key 不能为空（可填写 vision.api_key 或 openai.api_key）"}) + return + } + if strings.TrimSpace(oa.Model) == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "视觉模型不能为空"}) + return + } + + baseURL := strings.TrimSuffix(strings.TrimSpace(oa.BaseURL), "/") + if baseURL == "" { + if strings.EqualFold(strings.TrimSpace(oa.Provider), "claude") { + baseURL = "https://api.anthropic.com" + } else { + baseURL = "https://api.openai.com/v1" + } + } + + payload := map[string]interface{}{ + "model": oa.Model, + "messages": []map[string]string{ + {"role": "user", "content": "Hi"}, + }, + "max_completion_tokens": 5, + } + + tmpCfg := &config.OpenAIConfig{ + Provider: oa.Provider, + BaseURL: baseURL, + APIKey: strings.TrimSpace(oa.APIKey), + Model: oa.Model, + } + client := openai.NewClient(tmpCfg, nil, h.logger) + + ctx, cancel := context.WithTimeout(c.Request.Context(), 30*time.Second) + defer cancel() + + start := time.Now() + var chatResp struct { + Model string `json:"model"` + Choices []struct { + Message struct { + Content string `json:"content"` + } `json:"message"` + } `json:"choices"` + } + err := client.ChatCompletion(ctx, payload, &chatResp) + latency := time.Since(start) + + if err != nil { + if apiErr, ok := err.(*openai.APIError); ok { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "error": fmt.Sprintf("API 返回错误 (HTTP %d): %s", apiErr.StatusCode, apiErr.Body), + "status_code": apiErr.StatusCode, + }) + return + } + c.JSON(http.StatusOK, gin.H{ + "success": false, + "error": "连接失败: " + err.Error(), + }) + return + } + if len(chatResp.Choices) == 0 { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "error": "API 响应缺少 choices 字段，请检查 Base URL 与视觉模型名称", + }) + return + } + + c.JSON(http.StatusOK, gin.H{ + "success": true, + "model": chatResp.Model, + "latency_ms": latency.Milliseconds(), + }) +} + +// ApplyConfig 应用配置（重新加载并重启相关服务） +func (h *ConfigHandler) ApplyConfig(c *gin.Context) { + // 先检查是否需要动态初始化知识库（在锁外执行，避免阻塞其他请求） + var needInitKnowledge bool + var knowledgeInitializer KnowledgeInitializer + + h.mu.RLock() + needInitKnowledge = h.config.Knowledge.Enabled && h.knowledgeToolRegistrar == nil && h.knowledgeInitializer != nil + if needInitKnowledge { + knowledgeInitializer = h.knowledgeInitializer + } + h.mu.RUnlock() + + // 如果需要动态初始化知识库，在锁外执行（这是耗时操作） + if needInitKnowledge { + h.logger.Info("检测到知识库从禁用变为启用，开始动态初始化知识库组件") + if _, err := knowledgeInitializer(); err != nil { + h.logger.Error("动态初始化知识库失败", zap.Error(err)) + if h.audit != nil { + h.audit.RecordFail(c, "config", "apply", "应用配置失败：初始化知识库", map[string]interface{}{"error": err.Error()}) + } + c.JSON(http.StatusInternalServerError, gin.H{"error": "初始化知识库失败: " + err.Error()}) + return + } + h.logger.Debug("知识库动态初始化完成，工具已注册") + } + + // 检查嵌入模型配置是否变更（需要在锁外执行，避免阻塞） + var needReinitKnowledge bool + var reinitKnowledgeInitializer KnowledgeInitializer + h.mu.RLock() + if h.config.Knowledge.Enabled && h.knowledgeInitializer != nil && h.lastEmbeddingConfig != nil { + // 检查嵌入模型配置是否变更 + currentEmbedding := h.config.Knowledge.Embedding + if currentEmbedding.Provider != h.lastEmbeddingConfig.Provider || + currentEmbedding.Model != h.lastEmbeddingConfig.Model || + currentEmbedding.BaseURL != h.lastEmbeddingConfig.BaseURL || + currentEmbedding.APIKey != h.lastEmbeddingConfig.APIKey { + needReinitKnowledge = true + reinitKnowledgeInitializer = h.knowledgeInitializer + h.logger.Info("检测到嵌入模型配置变更，需要重新初始化知识库组件", + zap.String("old_model", h.lastEmbeddingConfig.Model), + zap.String("new_model", currentEmbedding.Model), + zap.String("old_base_url", h.lastEmbeddingConfig.BaseURL), + zap.String("new_base_url", currentEmbedding.BaseURL), + ) + } + } + h.mu.RUnlock() + + // 如果需要重新初始化知识库（嵌入模型配置变更），在锁外执行 + if needReinitKnowledge { + h.logger.Info("开始重新初始化知识库组件（嵌入模型配置已变更）") + if _, err := reinitKnowledgeInitializer(); err != nil { + h.logger.Error("重新初始化知识库失败", zap.Error(err)) + if h.audit != nil { + h.audit.RecordFail(c, "config", "apply", "应用配置失败：重新初始化知识库", map[string]interface{}{"error": err.Error()}) + } + c.JSON(http.StatusInternalServerError, gin.H{"error": "重新初始化知识库失败: " + err.Error()}) + return + } + h.logger.Info("知识库组件重新初始化完成") + } + + // C2：在 ClearTools 之前按配置启停（随后由 c2ToolRegistrar 注册 MCP 工具） + h.mu.RLock() + c2Rt := h.c2Runtime + h.mu.RUnlock() + if c2Rt != nil { + if err := c2Rt.ReconcileC2AfterConfigApply(); err != nil { + h.logger.Error("C2 配置应用失败", zap.Error(err)) + if h.audit != nil { + h.audit.RecordFail(c, "config", "apply", "应用配置失败：C2", map[string]interface{}{"error": err.Error()}) + } + c.JSON(http.StatusInternalServerError, gin.H{"error": "C2 启动失败: " + err.Error()}) + return + } + } + + // 现在获取写锁，执行快速的操作 + h.mu.Lock() + defer h.mu.Unlock() + + // 如果重新初始化了知识库，更新嵌入模型配置记录 + if needReinitKnowledge && h.config.Knowledge.Enabled { + h.lastEmbeddingConfig = &config.EmbeddingConfig{ + Provider: h.config.Knowledge.Embedding.Provider, + Model: h.config.Knowledge.Embedding.Model, + BaseURL: h.config.Knowledge.Embedding.BaseURL, + APIKey: h.config.Knowledge.Embedding.APIKey, + } + h.logger.Info("已更新嵌入模型配置记录") + } + + // 从 tools 目录重新加载工具配置（新增/修改/删除 yaml 后无需重启） + if err := config.ReloadSecurityToolsFromDir(h.config, h.configPath); err != nil { + h.logger.Error("重新加载工具配置失败", zap.Error(err)) + if h.audit != nil { + h.audit.RecordFail(c, "config", "apply", "应用配置失败：重新加载工具", map[string]interface{}{"error": err.Error()}) + } + c.JSON(http.StatusInternalServerError, gin.H{"error": "重新加载工具配置失败: " + err.Error()}) + return + } + h.logger.Debug("已从 tools 目录重新加载工具配置", zap.Int("tools_count", len(h.config.Security.Tools))) + + // 重新注册工具（根据新的启用状态） + h.logger.Debug("重新注册工具") + + // 清空MCP服务器中的工具 + h.mcpServer.ClearTools() + + // 重新注册安全工具 + h.executor.SetToolOutputMaxBytes(h.config.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective()) + h.executor.SetToolOutputSpillRoot(h.config.MultiAgent.EinoMiddleware.ReductionRootDir) + h.executor.RegisterTools(h.mcpServer) + mcp.RegisterExecutionControlTools(h.mcpServer, h.externalMCPMgr) + + // 重新注册漏洞记录工具（内置工具，必须注册） + if h.vulnerabilityToolRegistrar != nil { + h.logger.Info("重新注册漏洞记录工具") + if err := h.vulnerabilityToolRegistrar(); err != nil { + h.logger.Error("重新注册漏洞记录工具失败", zap.Error(err)) + } else { + h.logger.Info("漏洞记录工具已重新注册") + } + } + + // 重新注册 WebShell 工具（内置工具，必须注册） + if h.webshellToolRegistrar != nil { + h.logger.Info("重新注册 WebShell 工具") + if err := h.webshellToolRegistrar(); err != nil { + h.logger.Error("重新注册 WebShell 工具失败", zap.Error(err)) + } else { + h.logger.Info("WebShell 工具已重新注册") + } + } + + // 重新注册Skills工具（内置工具，必须注册） + if h.skillsToolRegistrar != nil { + h.logger.Info("重新注册Skills工具") + if err := h.skillsToolRegistrar(); err != nil { + h.logger.Error("重新注册Skills工具失败", zap.Error(err)) + } else { + h.logger.Info("Skills工具已重新注册") + } + } + + // 重新注册批量任务 MCP 工具 + if h.batchTaskToolRegistrar != nil { + h.logger.Info("重新注册批量任务 MCP 工具") + if err := h.batchTaskToolRegistrar(); err != nil { + h.logger.Error("重新注册批量任务 MCP 工具失败", zap.Error(err)) + } else { + h.logger.Info("批量任务 MCP 工具已重新注册") + } + } + + // 重新注册 C2 MCP 工具（仅当 C2 已启动） + if h.c2ToolRegistrar != nil { + h.logger.Info("重新注册 C2 MCP 工具") + if err := h.c2ToolRegistrar(); err != nil { + h.logger.Error("重新注册 C2 MCP 工具失败", zap.Error(err)) + } else { + h.logger.Info("C2 MCP 工具已处理") + } + } + + // 如果知识库启用，重新注册知识库工具 + if h.config.Knowledge.Enabled && h.knowledgeToolRegistrar != nil { + h.logger.Info("重新注册知识库工具") + if err := h.knowledgeToolRegistrar(); err != nil { + h.logger.Error("重新注册知识库工具失败", zap.Error(err)) + } else { + h.logger.Info("知识库工具已重新注册") + } + } + + // 更新Agent的OpenAI配置 + if h.agent != nil { + h.agent.UpdateConfig(&h.config.OpenAI) + h.agent.UpdateMaxIterations(h.config.Agent.MaxIterations) + h.agent.UpdateToolDescriptionMode(h.config.Security.ToolDescriptionMode) + h.logger.Info("Agent配置已更新") + } + if h.mcpServer != nil { + h.mcpServer.ConfigureHTTPToolCallTimeoutFromAgentMinutes(h.config.Agent.ToolTimeoutMinutes) + h.mcpServer.ConfigureToolWaitTimeoutSeconds(h.config.Agent.ToolWaitTimeoutSeconds) + h.mcpServer.ConfigureToolResultMaxBytes(h.config.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective()) + h.mcpServer.ConfigureToolResultSpillRoot(h.config.MultiAgent.EinoMiddleware.ReductionRootDir) + } + if h.executor != nil { + h.executor.SetToolOutputMaxBytes(h.config.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective()) + h.executor.SetToolOutputSpillRoot(h.config.MultiAgent.EinoMiddleware.ReductionRootDir) + } + if h.externalMCPMgr != nil { + h.externalMCPMgr.ConfigureToolWaitTimeoutSeconds(h.config.Agent.ToolWaitTimeoutSeconds) + h.externalMCPMgr.ConfigureToolResultMaxBytes(h.config.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective()) + h.externalMCPMgr.ConfigureToolResultSpillRoot(h.config.MultiAgent.EinoMiddleware.ReductionRootDir) + h.externalMCPMgr.ConfigureResilience(mcp.ExternalMCPResilienceConfig{ + MaxConcurrentPerServer: h.config.Agent.ExternalMCPMaxConcurrentPerServer, + MaxConcurrentTotal: h.config.Agent.ExternalMCPMaxConcurrentTotal, + CircuitFailureThreshold: h.config.Agent.ExternalMCPCircuitFailureThreshold, + CircuitCooldown: time.Duration(h.config.Agent.ExternalMCPCircuitCooldownSeconds) * time.Second, + }) + } + + // 更新AttackChainHandler的OpenAI配置 + if h.attackChainHandler != nil { + h.attackChainHandler.UpdateConfig(&h.config.OpenAI) + h.logger.Info("AttackChainHandler配置已更新") + } + + // 更新检索器配置（如果知识库启用） + if h.config.Knowledge.Enabled && h.retrieverUpdater != nil { + retrievalConfig := knowledge.RetrievalConfigFromYAML(h.config.Knowledge.Retrieval) + h.retrieverUpdater.UpdateConfig(retrievalConfig) + h.logger.Info("检索器配置已更新", + zap.Int("top_k", retrievalConfig.TopK), + zap.Float64("similarity_threshold", retrievalConfig.SimilarityThreshold), + ) + } + + // 更新嵌入模型配置记录（如果知识库启用） + if h.config.Knowledge.Enabled { + h.lastEmbeddingConfig = &config.EmbeddingConfig{ + Provider: h.config.Knowledge.Embedding.Provider, + Model: h.config.Knowledge.Embedding.Model, + BaseURL: h.config.Knowledge.Embedding.BaseURL, + APIKey: h.config.Knowledge.Embedding.APIKey, + } + } + + // 重启钉钉/飞书长连接，使前端修改的机器人配置立即生效（无需重启服务） + if h.robotRestarter != nil { + h.robotRestarter.RestartRobotConnections() + h.logger.Info("已触发机器人连接重启（钉钉/飞书）") + } + + h.logger.Info("配置已应用", + zap.Int("tools_count", len(h.config.Security.Tools)), + ) + + if h.audit != nil { + h.audit.Record(c, audit.Entry{ + Category: "config", + Action: "apply", + Result: "success", + Message: "配置已应用", + Detail: map[string]interface{}{ + "tools_count": len(h.config.Security.Tools), + "knowledge_enabled": h.config.Knowledge.Enabled, + "c2_enabled": h.config.C2.EnabledEffective(), + }, + }) + } + + c.JSON(http.StatusOK, gin.H{ + "message": "配置已应用", + "tools_count": len(h.config.Security.Tools), + }) +} + +// saveConfig 保存配置到文件 +func (h *ConfigHandler) saveConfig() error { + // 读取现有配置文件并创建备份 + data, err := os.ReadFile(h.configPath) + if err != nil { + return fmt.Errorf("读取配置文件失败: %w", err) + } + + if err := os.WriteFile(h.configPath+".backup", data, 0644); err != nil { + h.logger.Warn("创建配置备份失败", zap.Error(err)) + } + + root, err := loadYAMLDocument(h.configPath) + if err != nil { + return fmt.Errorf("解析配置文件失败: %w", err) + } + + updateAgentConfig(root, h.config.Agent) + updateMCPConfig(root, h.config.MCP) + updateAIConfig(root, h.config.AI) + removeKeyFromMap(root.Content[0], "openai") + updateVisionConfig(root, h.config.Vision) + updateFOFAConfig(root, h.config.FOFA) + updateSpaceSearchConfig(root, "zoomeye", h.config.ZoomEye) + updateSpaceSearchConfig(root, "quake", h.config.Quake) + updateSpaceSearchConfig(root, "shodan", h.config.Shodan) + updateKnowledgeConfig(root, h.config.Knowledge) + updateC2Config(root, h.config.C2) + updateRobotsConfig(root, h.config.Robots) + updateHitlConfig(root, h.config.Hitl) + updateMultiAgentConfig(root, h.config.MultiAgent) + // 更新外部MCP配置（使用external_mcp.go中的函数，同一包中可直接调用） + updateExternalMCPConfig(root, h.config.ExternalMCP) + + if err := writeYAMLDocument(h.configPath, root); err != nil { + return fmt.Errorf("保存配置文件失败: %w", err) + } + + // 更新工具配置文件中的enabled状态 + if h.config.Security.ToolsDir != "" { + configDir := filepath.Dir(h.configPath) + toolsDir := h.config.Security.ToolsDir + if !filepath.IsAbs(toolsDir) { + toolsDir = filepath.Join(configDir, toolsDir) + } + + for _, tool := range h.config.Security.Tools { + toolFile := filepath.Join(toolsDir, tool.Name+".yaml") + // 检查文件是否存在 + if _, err := os.Stat(toolFile); os.IsNotExist(err) { + // 尝试.yml扩展名 + toolFile = filepath.Join(toolsDir, tool.Name+".yml") + if _, err := os.Stat(toolFile); os.IsNotExist(err) { + h.logger.Warn("工具配置文件不存在", zap.String("tool", tool.Name)) + continue + } + } + + toolDoc, err := loadYAMLDocument(toolFile) + if err != nil { + h.logger.Warn("解析工具配置失败", zap.String("tool", tool.Name), zap.Error(err)) + continue + } + + setBoolInMap(toolDoc.Content[0], "enabled", tool.Enabled) + + if err := writeYAMLDocument(toolFile, toolDoc); err != nil { + h.logger.Warn("保存工具配置文件失败", zap.String("tool", tool.Name), zap.Error(err)) + continue + } + + h.logger.Info("更新工具配置", zap.String("tool", tool.Name), zap.Bool("enabled", tool.Enabled)) + } + } + + h.logger.Info("配置已保存", zap.String("path", h.configPath)) + return nil +} + +func loadYAMLDocument(path string) (*yaml.Node, error) { + data, err := os.ReadFile(path) + if err != nil { + return nil, err + } + + if len(bytes.TrimSpace(data)) == 0 { + return newEmptyYAMLDocument(), nil + } + + var doc yaml.Node + if err := yaml.Unmarshal(data, &doc); err != nil { + return nil, err + } + + if doc.Kind != yaml.DocumentNode || len(doc.Content) == 0 { + return newEmptyYAMLDocument(), nil + } + + if doc.Content[0].Kind != yaml.MappingNode { + root := &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"} + doc.Content = []*yaml.Node{root} + } + + return &doc, nil +} + +func newEmptyYAMLDocument() *yaml.Node { + root := &yaml.Node{ + Kind: yaml.DocumentNode, + Content: []*yaml.Node{{Kind: yaml.MappingNode, Tag: "!!map"}}, + } + return root +} + +func writeYAMLDocument(path string, doc *yaml.Node) error { + var buf bytes.Buffer + encoder := yaml.NewEncoder(&buf) + encoder.SetIndent(2) + if err := encoder.Encode(doc); err != nil { + return err + } + if err := encoder.Close(); err != nil { + return err + } + return os.WriteFile(path, buf.Bytes(), 0644) +} + +func updateAgentConfig(doc *yaml.Node, agent config.AgentConfig) { + root := doc.Content[0] + agentNode := ensureMap(root, "agent") + setIntInMap(agentNode, "max_iterations", agent.MaxIterations) + setIntInMap(agentNode, "tool_timeout_minutes", agent.ToolTimeoutMinutes) + setIntInMap(agentNode, "tool_wait_timeout_seconds", agent.ToolWaitTimeoutSeconds) + setIntInMap(agentNode, "external_mcp_max_concurrent_per_server", agent.ExternalMCPMaxConcurrentPerServer) + setIntInMap(agentNode, "external_mcp_max_concurrent_total", agent.ExternalMCPMaxConcurrentTotal) + setIntInMap(agentNode, "external_mcp_circuit_failure_threshold", agent.ExternalMCPCircuitFailureThreshold) + setIntInMap(agentNode, "external_mcp_circuit_cooldown_seconds", agent.ExternalMCPCircuitCooldownSeconds) + setStringInMap(agentNode, "system_prompt_path", agent.SystemPromptPath) +} + +func updateMCPConfig(doc *yaml.Node, cfg config.MCPConfig) { + root := doc.Content[0] + mcpNode := ensureMap(root, "mcp") + setBoolInMap(mcpNode, "enabled", cfg.Enabled) + setStringInMap(mcpNode, "host", cfg.Host) + setIntInMap(mcpNode, "port", cfg.Port) +} + +func updateVisionConfig(doc *yaml.Node, cfg config.VisionConfig) { + root := doc.Content[0] + visionNode := ensureMap(root, "vision") + setBoolInMap(visionNode, "enabled", cfg.Enabled) + if strings.TrimSpace(cfg.APIKey) != "" { + setStringInMap(visionNode, "api_key", cfg.APIKey) + } else { + setStringInMap(visionNode, "api_key", "") + } + if strings.TrimSpace(cfg.BaseURL) != "" { + setStringInMap(visionNode, "base_url", cfg.BaseURL) + } else { + setStringInMap(visionNode, "base_url", "") + } + setStringInMap(visionNode, "model", cfg.Model) + if strings.TrimSpace(cfg.Provider) != "" { + setStringInMap(visionNode, "provider", cfg.Provider) + } + if cfg.TimeoutSeconds > 0 { + setIntInMap(visionNode, "timeout_seconds", cfg.TimeoutSeconds) + } + if cfg.MaxImageBytes > 0 { + setIntInMap(visionNode, "max_image_bytes", int(cfg.MaxImageBytes)) + } + if cfg.MaxDimension > 0 { + setIntInMap(visionNode, "max_dimension", cfg.MaxDimension) + } + if cfg.JPEGQuality > 0 { + setIntInMap(visionNode, "jpeg_quality", cfg.JPEGQuality) + } + if cfg.MaxPayloadBytes > 0 { + setIntInMap(visionNode, "max_payload_bytes", int(cfg.MaxPayloadBytes)) + } + setIntInMap(visionNode, "skip_preprocess_below_bytes", int(cfg.SkipPreprocessBelowBytes)) + if strings.TrimSpace(cfg.Detail) != "" { + setStringInMap(visionNode, "detail", cfg.Detail) + } +} + +func updateOpenAIConfig(doc *yaml.Node, cfg config.OpenAIConfig) { + root := doc.Content[0] + openaiNode := ensureMap(root, "openai") + if cfg.Provider != "" { + setStringInMap(openaiNode, "provider", cfg.Provider) + } + setStringInMap(openaiNode, "api_key", cfg.APIKey) + setStringInMap(openaiNode, "base_url", cfg.BaseURL) + setStringInMap(openaiNode, "model", cfg.Model) + if cfg.MaxTotalTokens > 0 { + setIntInMap(openaiNode, "max_total_tokens", cfg.MaxTotalTokens) + } + rn := ensureMap(openaiNode, "reasoning") + if strings.TrimSpace(cfg.Reasoning.Mode) != "" { + setStringInMap(rn, "mode", cfg.Reasoning.Mode) + } + if strings.TrimSpace(cfg.Reasoning.Effort) != "" { + setStringInMap(rn, "effort", cfg.Reasoning.Effort) + } + if cfg.Reasoning.AllowClientReasoning != nil { + setBoolInMap(rn, "allow_client_reasoning", *cfg.Reasoning.AllowClientReasoning) + } + if strings.TrimSpace(cfg.Reasoning.Profile) != "" { + setStringInMap(rn, "profile", cfg.Reasoning.Profile) + } +} + +func updateAIConfig(doc *yaml.Node, cfg config.AIConfig) { + root := doc.Content[0] + aiNode := ensureMap(root, "ai") + if strings.TrimSpace(cfg.DefaultChannel) != "" { + setStringInMap(aiNode, "default_channel", config.NormalizeAIChannelID(cfg.DefaultChannel)) + } + channelsNode := ensureMap(aiNode, "channels") + channelsNode.Content = nil + normalized := make(map[string]config.AIChannelConfig, len(cfg.Channels)) + ids := make([]string, 0, len(cfg.Channels)) + for id, ch := range cfg.Channels { + nid := config.NormalizeAIChannelID(id) + if nid == "" { + continue + } + if _, exists := normalized[nid]; !exists { + ids = append(ids, nid) + } + normalized[nid] = ch + } + sort.Strings(ids) + seen := make(map[string]bool, len(ids)) + for _, id := range ids { + if seen[id] { + continue + } + seen[id] = true + ch := normalized[id] + keyNode := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: id} + channelNode := &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"} + channelsNode.Content = append(channelsNode.Content, keyNode, channelNode) + setStringInMap(channelNode, "name", ch.Name) + if strings.TrimSpace(ch.Provider) != "" { + setStringInMap(channelNode, "provider", ch.Provider) + } + setStringInMap(channelNode, "api_key", ch.APIKey) + setStringInMap(channelNode, "base_url", ch.BaseURL) + setStringInMap(channelNode, "model", ch.Model) + if ch.MaxTotalTokens > 0 { + setIntInMap(channelNode, "max_total_tokens", ch.MaxTotalTokens) + } + if ch.MaxCompletionTokens > 0 { + setIntInMap(channelNode, "max_completion_tokens", ch.MaxCompletionTokens) + } + rn := ensureMap(channelNode, "reasoning") + if strings.TrimSpace(ch.Reasoning.Mode) != "" { + setStringInMap(rn, "mode", ch.Reasoning.Mode) + } + if strings.TrimSpace(ch.Reasoning.Effort) != "" { + setStringInMap(rn, "effort", ch.Reasoning.Effort) + } + if ch.Reasoning.AllowClientReasoning != nil { + setBoolInMap(rn, "allow_client_reasoning", *ch.Reasoning.AllowClientReasoning) + } + if strings.TrimSpace(ch.Reasoning.Profile) != "" { + setStringInMap(rn, "profile", ch.Reasoning.Profile) + } + if len(rn.Content) == 0 { + removeKeyFromMap(channelNode, "reasoning") + } + } +} + +func updateFOFAConfig(doc *yaml.Node, cfg config.FofaConfig) { + root := doc.Content[0] + fofaNode := ensureMap(root, "fofa") + setStringInMap(fofaNode, "base_url", cfg.BaseURL) + removeKeyFromMap(fofaNode, "email") + setStringInMap(fofaNode, "api_key", cfg.APIKey) +} + +func updateSpaceSearchConfig(doc *yaml.Node, key string, cfg config.SpaceSearchConfig) { + root := doc.Content[0] + node := ensureMap(root, key) + setStringInMap(node, "base_url", cfg.BaseURL) + setStringInMap(node, "api_key", cfg.APIKey) +} + +func updateKnowledgeConfig(doc *yaml.Node, cfg config.KnowledgeConfig) { + root := doc.Content[0] + knowledgeNode := ensureMap(root, "knowledge") + setBoolInMap(knowledgeNode, "enabled", cfg.Enabled) + setStringInMap(knowledgeNode, "base_path", cfg.BasePath) + + // 更新嵌入配置 + embeddingNode := ensureMap(knowledgeNode, "embedding") + setStringInMap(embeddingNode, "provider", cfg.Embedding.Provider) + setStringInMap(embeddingNode, "model", cfg.Embedding.Model) + if cfg.Embedding.BaseURL != "" { + setStringInMap(embeddingNode, "base_url", cfg.Embedding.BaseURL) + } + if cfg.Embedding.APIKey != "" { + setStringInMap(embeddingNode, "api_key", cfg.Embedding.APIKey) + } + + // 更新检索配置 + retrievalNode := ensureMap(knowledgeNode, "retrieval") + setIntInMap(retrievalNode, "top_k", cfg.Retrieval.TopK) + setFloatInMap(retrievalNode, "similarity_threshold", cfg.Retrieval.SimilarityThreshold) + setStringInMap(retrievalNode, "sub_index_filter", cfg.Retrieval.SubIndexFilter) + mqNode := ensureMap(retrievalNode, "multi_query") + setIntInMap(mqNode, "max_queries", cfg.Retrieval.MultiQuery.MaxQueries) + rerankNode := ensureMap(retrievalNode, "rerank") + setStringInMap(rerankNode, "provider", cfg.Retrieval.Rerank.Provider) + setStringInMap(rerankNode, "model", cfg.Retrieval.Rerank.Model) + setStringInMap(rerankNode, "base_url", cfg.Retrieval.Rerank.BaseURL) + setStringInMap(rerankNode, "api_key", cfg.Retrieval.Rerank.APIKey) + postNode := ensureMap(retrievalNode, "post_retrieve") + setIntInMap(postNode, "prefetch_top_k", cfg.Retrieval.PostRetrieve.PrefetchTopK) + setIntInMap(postNode, "max_context_chars", cfg.Retrieval.PostRetrieve.MaxContextChars) + setIntInMap(postNode, "max_context_tokens", cfg.Retrieval.PostRetrieve.MaxContextTokens) + + // 更新索引配置 + indexingNode := ensureMap(knowledgeNode, "indexing") + setStringInMap(indexingNode, "chunk_strategy", cfg.Indexing.ChunkStrategy) + setIntInMap(indexingNode, "request_timeout_seconds", cfg.Indexing.RequestTimeoutSeconds) + setIntInMap(indexingNode, "chunk_size", cfg.Indexing.ChunkSize) + setIntInMap(indexingNode, "chunk_overlap", cfg.Indexing.ChunkOverlap) + setIntInMap(indexingNode, "max_chunks_per_item", cfg.Indexing.MaxChunksPerItem) + setBoolInMap(indexingNode, "prefer_source_file", cfg.Indexing.PreferSourceFile) + setIntInMap(indexingNode, "batch_size", cfg.Indexing.BatchSize) + setStringSliceInMap(indexingNode, "sub_indexes", cfg.Indexing.SubIndexes) + setIntInMap(indexingNode, "max_rpm", cfg.Indexing.MaxRPM) + setIntInMap(indexingNode, "rate_limit_delay_ms", cfg.Indexing.RateLimitDelayMs) + setIntInMap(indexingNode, "max_retries", cfg.Indexing.MaxRetries) + setIntInMap(indexingNode, "retry_delay_ms", cfg.Indexing.RetryDelayMs) +} + +func updateC2Config(doc *yaml.Node, cfg config.C2Config) { + root := doc.Content[0] + c2Node := ensureMap(root, "c2") + setBoolInMap(c2Node, "enabled", cfg.EnabledEffective()) +} + +func mergeHitlToolWhitelistSlice(existing, add []string) []string { + seen := make(map[string]struct{}) + out := make([]string, 0, len(existing)+len(add)) + for _, list := range [][]string{existing, add} { + for _, t := range list { + n := strings.ToLower(strings.TrimSpace(t)) + if n == "" { + continue + } + if _, ok := seen[n]; ok { + continue + } + seen[n] = struct{}{} + out = append(out, strings.TrimSpace(t)) + } + } + return out +} + +// SetHitlToolWhitelist 将全局免审批工具白名单整表写入 config.yaml（替换，非合并）。 +func (h *ConfigHandler) SetHitlToolWhitelist(tools []string) error { + h.mu.Lock() + defer h.mu.Unlock() + h.config.Hitl.ToolWhitelist = mergeHitlToolWhitelistSlice(nil, tools) + if err := h.saveConfig(); err != nil { + return err + } + h.logger.Info("HITL 全局工具白名单已写入配置文件", + zap.Int("count", len(h.config.Hitl.ToolWhitelist)), + ) + return nil +} + +// MergeHitlToolWhitelistIntoConfig 将会话侧栏提交的免审批工具名合并进内存配置并写入 config.yaml（与全局白名单去重规则一致：小写键、保留首次出现的原始大小写）。 +func (h *ConfigHandler) MergeHitlToolWhitelistIntoConfig(add []string) error { + h.mu.Lock() + defer h.mu.Unlock() + merged := mergeHitlToolWhitelistSlice(h.config.Hitl.ToolWhitelist, add) + h.config.Hitl.ToolWhitelist = merged + if err := h.saveConfig(); err != nil { + return err + } + h.logger.Info("HITL 全局工具白名单已合并写入配置文件", + zap.Int("count", len(merged)), + ) + return nil +} + +func updateHitlConfig(doc *yaml.Node, cfg config.HitlConfig) { + root := doc.Content[0] + hitlNode := ensureMap(root, "hitl") + auditModelNode := ensureMap(hitlNode, "audit_model") + setStringInMap(auditModelNode, "provider", cfg.AuditModel.Provider) + setStringInMap(auditModelNode, "base_url", cfg.AuditModel.BaseURL) + setStringInMap(auditModelNode, "api_key", cfg.AuditModel.APIKey) + setStringInMap(auditModelNode, "model", cfg.AuditModel.Model) + // flow 样式 [a, b, c] 单行展示，工具多时比块序列省行数 + setFlowStringSliceInMap(hitlNode, "tool_whitelist", cfg.ToolWhitelist) + setStringInMap(hitlNode, "default_reviewer", cfg.EffectiveDefaultReviewer()) + setIntInMap(hitlNode, "retention_days", cfg.RetentionDaysEffective()) + setStringInMap(hitlNode, "audit_agent_prompt", cfg.AuditAgentPrompt) + setStringInMap(hitlNode, "audit_agent_prompt_review_edit", cfg.AuditAgentPromptReviewEdit) +} + +// UpdateHitlDefaultReviewer 更新全局默认审批方并写入 config.yaml。 +func (h *ConfigHandler) UpdateHitlDefaultReviewer(reviewer string) error { + h.mu.Lock() + defer h.mu.Unlock() + h.config.Hitl.DefaultReviewer = config.HitlConfig{DefaultReviewer: reviewer}.EffectiveDefaultReviewer() + if err := h.saveConfig(); err != nil { + return err + } + h.logger.Info("HITL 全局默认审批方已写入配置文件", zap.String("default_reviewer", h.config.Hitl.DefaultReviewer)) + return nil +} + +// UpdateHitlAuditAgentStrategy 更新审批/审查编辑两套审计 Agent 提示词并写入 config.yaml。 +func (h *ConfigHandler) UpdateHitlAuditAgentStrategy(approvalPrompt, reviewEditPrompt string) error { + h.mu.Lock() + defer h.mu.Unlock() + h.config.Hitl.AuditAgentPrompt = strings.TrimSpace(approvalPrompt) + h.config.Hitl.AuditAgentPromptReviewEdit = strings.TrimSpace(reviewEditPrompt) + if err := h.saveConfig(); err != nil { + return err + } + h.logger.Info("HITL 审计 Agent 提示词已写入配置文件") + return nil +} + +func updateRobotsConfig(doc *yaml.Node, cfg config.RobotsConfig) { + root := doc.Content[0] + robotsNode := ensureMap(root, "robots") + + if cfg.Session.StrictUserIdentity != nil { + sessionNode := ensureMap(robotsNode, "session") + setBoolInMap(sessionNode, "strict_user_identity", *cfg.Session.StrictUserIdentity) + } + + wechatNode := ensureMap(robotsNode, "wechat") + setBoolInMap(wechatNode, "enabled", cfg.Wechat.Enabled) + setStringInMap(wechatNode, "bot_token", cfg.Wechat.BotToken) + setStringInMap(wechatNode, "ilink_bot_id", cfg.Wechat.ILinkBotID) + setStringInMap(wechatNode, "ilink_user_id", cfg.Wechat.ILinkUserID) + setStringInMap(wechatNode, "base_url", cfg.Wechat.BaseURL) + setStringInMap(wechatNode, "bot_type", cfg.Wechat.BotType) + setStringInMap(wechatNode, "bot_agent", cfg.Wechat.BotAgent) + + wecomNode := ensureMap(robotsNode, "wecom") + setBoolInMap(wecomNode, "enabled", cfg.Wecom.Enabled) + setStringInMap(wecomNode, "token", cfg.Wecom.Token) + setStringInMap(wecomNode, "encoding_aes_key", cfg.Wecom.EncodingAESKey) + setStringInMap(wecomNode, "corp_id", cfg.Wecom.CorpID) + setStringInMap(wecomNode, "secret", cfg.Wecom.Secret) + setIntInMap(wecomNode, "agent_id", int(cfg.Wecom.AgentID)) + + dingtalkNode := ensureMap(robotsNode, "dingtalk") + setBoolInMap(dingtalkNode, "enabled", cfg.Dingtalk.Enabled) + setStringInMap(dingtalkNode, "client_id", cfg.Dingtalk.ClientID) + setStringInMap(dingtalkNode, "client_secret", cfg.Dingtalk.ClientSecret) + setBoolInMap(dingtalkNode, "allow_conversation_id_fallback", cfg.Dingtalk.AllowConversationIDFallback) + + larkNode := ensureMap(robotsNode, "lark") + setBoolInMap(larkNode, "enabled", cfg.Lark.Enabled) + setStringInMap(larkNode, "app_id", cfg.Lark.AppID) + setStringInMap(larkNode, "app_secret", cfg.Lark.AppSecret) + setStringInMap(larkNode, "verify_token", cfg.Lark.VerifyToken) + setBoolInMap(larkNode, "allow_chat_id_fallback", cfg.Lark.AllowChatIDFallback) + + telegramNode := ensureMap(robotsNode, "telegram") + setBoolInMap(telegramNode, "enabled", cfg.Telegram.Enabled) + setStringInMap(telegramNode, "bot_token", cfg.Telegram.BotToken) + setStringInMap(telegramNode, "bot_username", cfg.Telegram.BotUsername) + setBoolInMap(telegramNode, "allow_group_messages", cfg.Telegram.AllowGroupMessages) + + slackNode := ensureMap(robotsNode, "slack") + setBoolInMap(slackNode, "enabled", cfg.Slack.Enabled) + setStringInMap(slackNode, "bot_token", cfg.Slack.BotToken) + setStringInMap(slackNode, "app_token", cfg.Slack.AppToken) + + discordNode := ensureMap(robotsNode, "discord") + setBoolInMap(discordNode, "enabled", cfg.Discord.Enabled) + setStringInMap(discordNode, "bot_token", cfg.Discord.BotToken) + setBoolInMap(discordNode, "allow_guild_messages", cfg.Discord.AllowGuildMessages) + + qqNode := ensureMap(robotsNode, "qq") + setBoolInMap(qqNode, "enabled", cfg.QQ.Enabled) + setStringInMap(qqNode, "app_id", cfg.QQ.AppID) + setStringInMap(qqNode, "client_secret", cfg.QQ.ClientSecret) + setBoolInMap(qqNode, "sandbox", cfg.QQ.Sandbox) +} + +func updateMultiAgentConfig(doc *yaml.Node, cfg config.MultiAgentConfig) { + root := doc.Content[0] + maNode := ensureMap(root, "multi_agent") + setBoolInMap(maNode, "enabled", cfg.Enabled) + setStringInMap(maNode, "robot_default_agent_mode", config.NormalizeRobotAgentMode(cfg)) + setBoolInMap(maNode, "batch_use_multi_agent", cfg.BatchUseMultiAgent) + setIntInMap(maNode, "plan_execute_loop_max_iterations", cfg.PlanExecuteLoopMaxIterations) + mwNode := ensureMap(maNode, "eino_middleware") + setIntInMap(mwNode, "summarization_user_intent_ledger_max_runes", cfg.EinoMiddleware.SummarizationUserIntentLedgerMaxRunesEffective()) + setIntInMap(mwNode, "summarization_user_intent_ledger_entry_max_runes", cfg.EinoMiddleware.SummarizationUserIntentLedgerEntryMaxRunesEffective()) + setIntInMap(mwNode, "latest_user_message_max_runes", cfg.EinoMiddleware.LatestUserMessageMaxRunesEffective()) + setIntInMap(mwNode, "latest_user_message_head_runes", cfg.EinoMiddleware.LatestUserMessageHeadRunesEffective()) + setIntInMap(mwNode, "latest_user_message_tail_runes", cfg.EinoMiddleware.LatestUserMessageTailRunesEffective()) + setIntInMap(mwNode, "model_retry_max_retries", cfg.EinoMiddleware.ModelRetryMaxRetries) + setIntInMap(mwNode, "model_retry_max_backoff_sec", cfg.EinoMiddleware.ModelRetryMaxBackoffSec) + setFlowStringSliceInMap(mwNode, "model_failover_channels", dedupeTrimmedStringList(cfg.EinoMiddleware.ModelFailoverChannels)) + setIntInMap(mwNode, "model_failover_max_retries", cfg.EinoMiddleware.ModelFailoverMaxRetries) + setFlowStringSliceInMap(mwNode, "tool_search_always_visible_tools", dedupeToolNameList(cfg.EinoMiddleware.ToolSearchAlwaysVisibleTools)) +} + +func dedupeToolNameList(in []string) []string { + return dedupeTrimmedStringList(in) +} + +func dedupeTrimmedStringList(in []string) []string { + if len(in) == 0 { + return []string{} + } + seen := make(map[string]struct{}, len(in)) + out := make([]string, 0, len(in)) + for _, name := range in { + n := strings.TrimSpace(name) + if n == "" { + continue + } + key := strings.ToLower(n) + if _, ok := seen[key]; ok { + continue + } + seen[key] = struct{}{} + out = append(out, n) + } + return out +} + +func mergeToolNameLists(a, b []string) []string { + return dedupeToolNameList(append(append([]string{}, a...), b...)) +} + +func ensureMap(parent *yaml.Node, path ...string) *yaml.Node { + current := parent + for _, key := range path { + value := findMapValue(current, key) + if value == nil { + keyNode := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: key} + mapNode := &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"} + current.Content = append(current.Content, keyNode, mapNode) + value = mapNode + } + + if value.Kind != yaml.MappingNode { + value.Kind = yaml.MappingNode + value.Tag = "!!map" + value.Style = 0 + value.Content = nil + } + + current = value + } + + return current +} + +func findMapValue(mapNode *yaml.Node, key string) *yaml.Node { + if mapNode == nil || mapNode.Kind != yaml.MappingNode { + return nil + } + + for i := 0; i < len(mapNode.Content); i += 2 { + if mapNode.Content[i].Value == key { + return mapNode.Content[i+1] + } + } + return nil +} + +func ensureKeyValue(mapNode *yaml.Node, key string) (*yaml.Node, *yaml.Node) { + if mapNode == nil || mapNode.Kind != yaml.MappingNode { + return nil, nil + } + + for i := 0; i < len(mapNode.Content); i += 2 { + if mapNode.Content[i].Value == key { + return mapNode.Content[i], mapNode.Content[i+1] + } + } + + keyNode := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: key} + valueNode := &yaml.Node{} + mapNode.Content = append(mapNode.Content, keyNode, valueNode) + return keyNode, valueNode +} + +func setStringInMap(mapNode *yaml.Node, key, value string) { + _, valueNode := ensureKeyValue(mapNode, key) + valueNode.Kind = yaml.ScalarNode + valueNode.Tag = "!!str" + valueNode.Style = 0 + valueNode.Value = value +} + +func removeKeyFromMap(mapNode *yaml.Node, key string) { + if mapNode == nil || mapNode.Kind != yaml.MappingNode { + return + } + for i := 0; i+1 < len(mapNode.Content); i += 2 { + if mapNode.Content[i].Value == key { + mapNode.Content = append(mapNode.Content[:i], mapNode.Content[i+2:]...) + return + } + } +} + +func setStringSliceInMap(mapNode *yaml.Node, key string, values []string) { + _, valueNode := ensureKeyValue(mapNode, key) + valueNode.Kind = yaml.SequenceNode + valueNode.Tag = "!!seq" + valueNode.Style = 0 + valueNode.Content = nil + for _, v := range values { + valueNode.Content = append(valueNode.Content, &yaml.Node{ + Kind: yaml.ScalarNode, + Tag: "!!str", + Value: v, + }) + } +} + +func setFlowStringSliceInMap(mapNode *yaml.Node, key string, values []string) { + _, valueNode := ensureKeyValue(mapNode, key) + valueNode.Kind = yaml.SequenceNode + valueNode.Tag = "!!seq" + valueNode.Style = yaml.FlowStyle + valueNode.Content = nil + for _, v := range values { + valueNode.Content = append(valueNode.Content, &yaml.Node{ + Kind: yaml.ScalarNode, + Tag: "!!str", + Value: v, + }) + } +} + +func setIntInMap(mapNode *yaml.Node, key string, value int) { + _, valueNode := ensureKeyValue(mapNode, key) + valueNode.Kind = yaml.ScalarNode + valueNode.Tag = "!!int" + valueNode.Style = 0 + valueNode.Value = fmt.Sprintf("%d", value) +} + +func findBoolInMap(mapNode *yaml.Node, key string) *bool { + if mapNode == nil || mapNode.Kind != yaml.MappingNode { + return nil + } + + for i := 0; i < len(mapNode.Content); i += 2 { + if i+1 >= len(mapNode.Content) { + break + } + keyNode := mapNode.Content[i] + valueNode := mapNode.Content[i+1] + + if keyNode.Kind == yaml.ScalarNode && keyNode.Value == key { + if valueNode.Kind == yaml.ScalarNode { + if valueNode.Value == "true" { + result := true + return &result + } else if valueNode.Value == "false" { + result := false + return &result + } + } + return nil + } + } + return nil +} + +func setBoolInMap(mapNode *yaml.Node, key string, value bool) { + _, valueNode := ensureKeyValue(mapNode, key) + valueNode.Kind = yaml.ScalarNode + valueNode.Tag = "!!bool" + valueNode.Style = 0 + if value { + valueNode.Value = "true" + } else { + valueNode.Value = "false" + } +} + +func setFloatInMap(mapNode *yaml.Node, key string, value float64) { + _, valueNode := ensureKeyValue(mapNode, key) + valueNode.Kind = yaml.ScalarNode + valueNode.Tag = "!!float" + valueNode.Style = 0 + // 对于0.0到1.0之间的值（如 similarity_threshold），使用%.1f确保0.0被明确序列化为"0.0" + // 对于其他值，使用%g自动选择最合适的格式 + if value >= 0.0 && value <= 1.0 { + valueNode.Value = fmt.Sprintf("%.1f", value) + } else { + valueNode.Value = fmt.Sprintf("%g", value) + } +} + +// getExternalMCPTools 获取外部MCP工具列表（公共方法） +func (h *ConfigHandler) getExternalMCPTools(ctx context.Context) []ToolConfigInfo { + if h.externalMCPMgr == nil { + return nil + } + return h.getExternalMCPToolsWithManager(ctx, h.externalMCPMgr, h.pickToolDescription) +} + +// getExternalMCPToolsWithManager 获取外部 MCP 工具（不持有 config 锁，供 GetTools 等热路径使用） +func (h *ConfigHandler) getExternalMCPToolsWithManager( + ctx context.Context, + mgr *mcp.ExternalMCPManager, + pickDesc func(shortDesc, fullDesc string) string, +) []ToolConfigInfo { + var result []ToolConfigInfo + if mgr == nil { + return result + } + + timeoutCtx, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + + externalTools, err := mgr.GetAllTools(timeoutCtx) + if err != nil { + h.logger.Warn("获取外部MCP工具失败（可能连接断开），尝试返回缓存的工具", + zap.Error(err), + zap.String("hint", "如果外部MCP工具未显示，请检查连接状态或点击刷新按钮"), + ) + } + + if len(externalTools) == 0 { + return result + } + + externalMCPConfigs := mgr.GetConfigs() + + for _, externalTool := range externalTools { + mcpName, actualToolName := h.parseExternalToolName(externalTool.Name) + if mcpName == "" || actualToolName == "" { + continue + } + + enabled := h.calculateExternalToolEnabledWithManager(mcpName, actualToolName, externalMCPConfigs, mgr) + + result = append(result, ToolConfigInfo{ + Name: actualToolName, + Description: pickDesc(externalTool.ShortDescription, externalTool.Description), + Enabled: enabled, + IsExternal: true, + ExternalMCP: mcpName, + }) + } + + return result +} + +// parseExternalToolName 解析外部工具名称（格式：mcpName::toolName） +func (h *ConfigHandler) parseExternalToolName(fullName string) (mcpName, toolName string) { + idx := strings.Index(fullName, "::") + if idx > 0 { + return fullName[:idx], fullName[idx+2:] + } + return "", "" +} + +// calculateExternalToolEnabled 计算外部工具的启用状态 +func (h *ConfigHandler) calculateExternalToolEnabled(mcpName, toolName string, configs map[string]config.ExternalMCPServerConfig) bool { + return h.calculateExternalToolEnabledWithManager(mcpName, toolName, configs, h.externalMCPMgr) +} + +func (h *ConfigHandler) calculateExternalToolEnabledWithManager( + mcpName, toolName string, + configs map[string]config.ExternalMCPServerConfig, + mgr *mcp.ExternalMCPManager, +) bool { + cfg, exists := configs[mcpName] + if !exists { + return false + } + + if !cfg.ExternalMCPEnable { + return false + } + + if cfg.ToolEnabled != nil { + if toolEnabled, exists := cfg.ToolEnabled[toolName]; exists && !toolEnabled { + return false + } + } + + if mgr == nil { + return false + } + client, exists := mgr.GetClient(mcpName) + if !exists || !client.IsConnected() { + return false + } + + return true +} + +// pickToolDescription 根据 security.tool_description_mode 选择 short 或 full 描述并限制长度。 +// 调用方若已持有 h.mu 读锁，须直接读 mode 并调用 pickToolDescriptionWithMode，避免嵌套 RLock 死锁。 +func (h *ConfigHandler) pickToolDescription(shortDesc, fullDesc string) string { + return pickToolDescriptionWithMode(h.config.Security.ToolDescriptionMode, shortDesc, fullDesc) +} + +func pickToolDescriptionWithMode(mode, shortDesc, fullDesc string) string { + useFull := strings.TrimSpace(strings.ToLower(mode)) == "full" + description := shortDesc + if useFull { + description = fullDesc + } else if description == "" { + description = fullDesc + } + if len(description) > 10000 { + description = description[:10000] + "..." + } + return description +} + +// GetToolSchema 获取单个工具的 inputSchema（按需加载，避免列表接口返回大量 schema 数据） +func (h *ConfigHandler) GetToolSchema(c *gin.Context) { + toolName := c.Param("name") + if toolName == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "工具名称不能为空"}) + return + } + + externalMCP := c.Query("external_mcp") + if externalMCP != "" { + h.mu.RLock() + externalMCPMgr := h.externalMCPMgr + h.mu.RUnlock() + + if externalMCPMgr != nil { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + externalTools, _ := externalMCPMgr.GetAllTools(ctx) + fullName := externalMCP + "::" + toolName + for _, t := range externalTools { + if t.Name == fullName { + c.JSON(http.StatusOK, gin.H{"input_schema": t.InputSchema}) + return + } + } + } + c.JSON(http.StatusNotFound, gin.H{"error": "外部工具未找到"}) + return + } + + h.mu.RLock() + securityTools := append([]config.ToolConfig(nil), h.config.Security.Tools...) + mcpServer := h.mcpServer + h.mu.RUnlock() + + for _, tool := range securityTools { + if tool.Name == toolName { + c.JSON(http.StatusOK, gin.H{"input_schema": buildInputSchemaFromParams(tool.Parameters)}) + return + } + } + + // MCP 注册工具（如知识检索） + if mcpServer != nil { + for _, mt := range mcpServer.GetAllTools() { + if mt.Name == toolName { + c.JSON(http.StatusOK, gin.H{"input_schema": mt.InputSchema}) + return + } + } + } + + c.JSON(http.StatusNotFound, gin.H{"error": "工具未找到"}) +} + +// buildInputSchemaFromParams 从 YAML 工具的 ParameterConfig 构建 JSON Schema（用于前端展示）。 +// 不依赖 MCP 服务器注册状态，所有工具（包括未启用的）都能返回参数定义。 +func buildInputSchemaFromParams(params []config.ParameterConfig) map[string]interface{} { + if len(params) == 0 { + return nil + } + + properties := make(map[string]interface{}) + required := make([]string, 0) + + for _, p := range params { + name := strings.TrimSpace(p.Name) + if name == "" { + continue + } + prop := map[string]interface{}{ + "type": convertParamType(p.Type), + "description": p.Description, + } + if p.Default != nil { + prop["default"] = p.Default + } + if len(p.Options) > 0 { + prop["enum"] = p.Options + } + properties[name] = prop + if p.Required { + required = append(required, name) + } + } + + schema := map[string]interface{}{ + "type": "object", + "properties": properties, + } + if len(required) > 0 { + schema["required"] = required + } + return schema +} + +func convertParamType(t string) string { + switch strings.TrimSpace(strings.ToLower(t)) { + case "int", "integer", "number": + return "number" + case "bool", "boolean": + return "boolean" + case "array", "list": + return "array" + default: + return "string" + } +} diff --git a/internal/handler/config_eino_resilience_test.go b/internal/handler/config_eino_resilience_test.go new file mode 100644 index 00000000..079d6acc --- /dev/null +++ b/internal/handler/config_eino_resilience_test.go @@ -0,0 +1,64 @@ +package handler + +import ( + "testing" + + "cyberstrike-ai/internal/config" + "gopkg.in/yaml.v3" +) + +func TestUpdateMultiAgentConfigWritesEinoModelResilience(t *testing.T) { + doc := &yaml.Node{ + Kind: yaml.DocumentNode, + Content: []*yaml.Node{{ + Kind: yaml.MappingNode, + Tag: "!!map", + }}, + } + + updateMultiAgentConfig(doc, config.MultiAgentConfig{ + Enabled: true, + RobotDefaultAgentMode: "deep", + PlanExecuteLoopMaxIterations: 3, + EinoMiddleware: config.MultiAgentEinoMiddlewareConfig{ + ModelRetryMaxRetries: 5, + ModelRetryMaxBackoffSec: 45, + ModelFailoverChannels: []string{"backup-openai", "backup-claude", "backup-openai"}, + ModelFailoverMaxRetries: 2, + }, + }) + + var got struct { + MultiAgent struct { + EinoMiddleware struct { + ModelRetryMaxRetries int `yaml:"model_retry_max_retries"` + ModelRetryMaxBackoffSec int `yaml:"model_retry_max_backoff_sec"` + ModelFailoverChannels []string `yaml:"model_failover_channels"` + ModelFailoverMaxRetries int `yaml:"model_failover_max_retries"` + } `yaml:"eino_middleware"` + } `yaml:"multi_agent"` + } + if err := doc.Decode(&got); err != nil { + t.Fatalf("decode config yaml: %v", err) + } + + mw := got.MultiAgent.EinoMiddleware + if mw.ModelRetryMaxRetries != 5 { + t.Fatalf("model_retry_max_retries = %d, want 5", mw.ModelRetryMaxRetries) + } + if mw.ModelRetryMaxBackoffSec != 45 { + t.Fatalf("model_retry_max_backoff_sec = %d, want 45", mw.ModelRetryMaxBackoffSec) + } + if mw.ModelFailoverMaxRetries != 2 { + t.Fatalf("model_failover_max_retries = %d, want 2", mw.ModelFailoverMaxRetries) + } + wantChannels := []string{"backup-openai", "backup-claude"} + if len(mw.ModelFailoverChannels) != len(wantChannels) { + t.Fatalf("model_failover_channels = %#v, want %#v", mw.ModelFailoverChannels, wantChannels) + } + for i, want := range wantChannels { + if mw.ModelFailoverChannels[i] != want { + t.Fatalf("model_failover_channels[%d] = %q, want %q", i, mw.ModelFailoverChannels[i], want) + } + } +} diff --git a/internal/handler/conversation.go b/internal/handler/conversation.go new file mode 100644 index 00000000..3cc5e4fb --- /dev/null +++ b/internal/handler/conversation.go @@ -0,0 +1,625 @@ +package handler + +import ( + "encoding/json" + "fmt" + "net/http" + "strconv" + "strings" + "time" + + "cyberstrike-ai/internal/audit" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +// ConversationTaskStopper cancels in-flight agent work when a conversation is removed. +type ConversationTaskStopper interface { + CancelRunningTaskForConversation(conversationID string) +} + +// ConversationTaskStateProvider reports whether the in-memory agent task for +// a conversation is still genuinely running. Plan files may survive a service +// restart or cancellation, so their status alone is not authoritative. +type ConversationTaskStateProvider interface { + ConversationTaskRuntimeState(conversationID string) (running bool, startedAt time.Time) +} + +// ConversationHandler 对话处理器 +type ConversationHandler struct { + db *database.DB + logger *zap.Logger + audit *audit.Service + taskStopper ConversationTaskStopper + taskState ConversationTaskStateProvider +} + +// SetAudit wires platform audit logging. +func (h *ConversationHandler) SetAudit(s *audit.Service) { + h.audit = s +} + +// SetTaskStopper wires cancellation of in-flight agent tasks on conversation delete. +func (h *ConversationHandler) SetTaskStopper(stopper ConversationTaskStopper) { + h.taskStopper = stopper +} + +// SetTaskStateProvider wires the live agent task registry used by supplemental +// conversation UI such as the agent-maintained plan list. +func (h *ConversationHandler) SetTaskStateProvider(provider ConversationTaskStateProvider) { + h.taskState = provider +} + +// NewConversationHandler 创建新的对话处理器 +func NewConversationHandler(db *database.DB, logger *zap.Logger) *ConversationHandler { + return &ConversationHandler{ + db: db, + logger: logger, + } +} + +// CreateConversationRequest 创建对话请求 +type CreateConversationRequest struct { + Title string `json:"title"` + ProjectID string `json:"projectId,omitempty"` +} + +// SetConversationProjectRequest 设置对话所属项目 +type SetConversationProjectRequest struct { + ProjectID string `json:"projectId"` // 空字符串表示解除绑定 +} + +// CreateConversation 创建新对话 +func (h *ConversationHandler) CreateConversation(c *gin.Context) { + var req CreateConversationRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + title := req.Title + if title == "" { + title = "新对话" + } + + meta := audit.ConversationCreateMetaFromGin(c, "api") + meta.ProjectID = strings.TrimSpace(req.ProjectID) + if !h.conversationProjectAllowed(c, meta.ProjectID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问目标项目"}) + return + } + conv, err := h.db.CreateConversation(title, meta) + if err != nil { + h.logger.Error("创建对话失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if session, ok := security.CurrentSession(c); ok { + _ = h.db.SetResourceOwner("conversation", conv.ID, session.UserID) + _ = h.db.AssignResourceToUser(session.UserID, "conversation", conv.ID) + if conv.ProjectID != "" { + _ = h.db.AssignResourceToUser(session.UserID, "project", conv.ProjectID) + } + } + + c.JSON(http.StatusOK, conv) +} + +// SetConversationProject 设置或清除对话绑定的项目 +func (h *ConversationHandler) SetConversationProject(c *gin.Context) { + id := c.Param("id") + var req SetConversationProjectRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if _, err := h.db.GetConversation(id); err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": "对话不存在"}) + return + } + projectID := strings.TrimSpace(req.ProjectID) + if !h.conversationProjectAllowed(c, projectID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问目标项目"}) + return + } + if err := h.db.SetConversationProjectID(id, projectID); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"success": true, "projectId": projectID}) +} + +func (h *ConversationHandler) conversationProjectAllowed(c *gin.Context, projectID string) bool { + projectID = strings.TrimSpace(projectID) + if projectID == "" { + return true + } + session, ok := security.CurrentSession(c) + if !ok { + return false + } + return h.db.UserCanAccessResource(session.UserID, session.Scope, "project", projectID) +} + +// ListConversations 列出对话 +func (h *ConversationHandler) ListConversations(c *gin.Context) { + limitStr := c.DefaultQuery("limit", "50") + offsetStr := c.DefaultQuery("offset", "0") + search := c.Query("search") // 获取搜索参数 + projectID := strings.TrimSpace(c.Query("project_id")) + + limit, _ := strconv.Atoi(limitStr) + offset, _ := strconv.Atoi(offsetStr) + + if limit <= 0 { + limit = 50 + } + if limit > 1000 { + limit = 1000 + } + + excludeGrouped := strings.TrimSpace(search) == "" && projectID == "" && + (c.Query("exclude_grouped") == "true" || c.Query("exclude_grouped") == "1") + sortBy := strings.TrimSpace(c.Query("sort_by")) + session, _ := security.CurrentSession(c) + + var conversations []*database.Conversation + var total int + var err error + if excludeGrouped { + conversations, err = h.db.ListUngroupedConversationsForAccess(limit, offset, sortBy, projectID, session.UserID, session.Scope) + if err == nil { + total, err = h.db.CountUngroupedConversationsForAccess(projectID, session.UserID, session.Scope) + } + } else { + conversations, err = h.db.ListConversationsForAccess(limit, offset, search, sortBy, projectID, session.UserID, session.Scope) + if err == nil { + total, err = h.db.CountConversationsForAccess(search, projectID, session.UserID, session.Scope) + } + } + if err != nil { + h.logger.Error("获取对话列表失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if conversations == nil { + conversations = []*database.Conversation{} + } + c.JSON(http.StatusOK, gin.H{ + "conversations": conversations, + "total": total, + "limit": limit, + "offset": offset, + }) +} + +// GetConversation 获取对话 +func (h *ConversationHandler) GetConversation(c *gin.Context) { + id := c.Param("id") + + // 默认轻量加载，只有用户需要展开详情时再按需拉取 + // include_process_details=1/true 时返回全量 processDetails（兼容旧行为） + includeStr := c.DefaultQuery("include_process_details", "0") + include := includeStr == "1" || includeStr == "true" || includeStr == "yes" + + var ( + conv *database.Conversation + err error + ) + if include { + conv, err = h.db.GetConversation(id) + } else { + conv, err = h.db.GetConversationLite(id) + } + if err != nil { + h.logger.Error("获取对话失败", zap.Error(err)) + c.JSON(http.StatusNotFound, gin.H{"error": "对话不存在"}) + return + } + + c.JSON(http.StatusOK, conv) +} + +// GetConversationPlanTasks returns the task list maintained by the agent's +// TaskCreate/TaskUpdate tools for this conversation. +func (h *ConversationHandler) GetConversationPlanTasks(c *gin.Context) { + id := strings.TrimSpace(c.Param("id")) + session, ok := security.CurrentSession(c) + if !ok || !h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", id) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该对话"}) + return + } + if _, err := h.db.GetConversationLite(id); err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": "对话不存在"}) + return + } + running := false + startedAt := time.Time{} + if h.taskState != nil { + running, startedAt = h.taskState.ConversationTaskRuntimeState(id) + } + if !running { + c.JSON(http.StatusOK, gin.H{ + "tasks": []database.ConversationPlanTask{}, "total": 0, + "completed": 0, "activeStep": 0, "running": false, + }) + return + } + tasks, err := h.db.ListConversationPlanTasksSince(id, startedAt) + if err != nil { + h.logger.Error("获取对话任务列表失败", zap.String("conversationId", id), zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": "获取任务列表失败"}) + return + } + + completed := 0 + activeStep := 0 + for i, task := range tasks { + status := strings.ToLower(strings.TrimSpace(task.Status)) + if status == "completed" { + completed++ + } + if activeStep == 0 && status == "in_progress" { + activeStep = i + 1 + } + } + if activeStep == 0 { + for i, task := range tasks { + if strings.ToLower(strings.TrimSpace(task.Status)) != "completed" { + activeStep = i + 1 + break + } + } + } + if activeStep == 0 && len(tasks) > 0 { + activeStep = len(tasks) + } + + c.JSON(http.StatusOK, gin.H{ + "tasks": tasks, + "total": len(tasks), + "completed": completed, + "activeStep": activeStep, + "running": true, + }) +} + +const ( + defaultProcessDetailsPageLimit = 50 + maxProcessDetailsPageLimit = 500 +) + +// GetMessageProcessDetails 获取指定消息的过程详情（按需加载） +// 查询参数： +// - summary=1：仅返回摘要（total / iterationCount / maxIteration） +// - limit + offset：分页返回 processDetails（未指定 limit 时默认 50 条） +// - anchorId：返回包含该过程详情锚点的一页，适合从工具按钮精准定位 +// - full=1：显式返回全量 processDetails（用于导出/兼容旧集成，不建议 UI 展开时使用） +func (h *ConversationHandler) GetMessageProcessDetails(c *gin.Context) { + messageID := c.Param("id") + if messageID == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "message id required"}) + return + } + + summaryStr := strings.TrimSpace(c.Query("summary")) + if summaryStr == "1" || strings.EqualFold(summaryStr, "true") || strings.EqualFold(summaryStr, "yes") { + summary, err := h.db.GetProcessDetailsSummary(messageID) + if err != nil { + h.logger.Error("获取过程详情摘要失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"summary": summary}) + return + } + + fullStr := strings.TrimSpace(c.Query("full")) + if fullStr == "1" || strings.EqualFold(fullStr, "true") || strings.EqualFold(fullStr, "yes") { + details, err := h.db.GetProcessDetails(messageID) + if err != nil { + h.logger.Error("获取过程详情失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + details = database.DedupeConsecutiveProcessDetails(details) + out := processDetailsToJSON(h.logger, h.db, details, true) + c.JSON(http.StatusOK, gin.H{ + "processDetails": out, + "total": len(out), + "offset": 0, + "limit": len(out), + "hasMore": false, + }) + return + } + + limitStr := strings.TrimSpace(c.Query("limit")) + limit := defaultProcessDetailsPageLimit + if limitStr != "" { + parsedLimit, err := strconv.Atoi(limitStr) + if err != nil || parsedLimit <= 0 { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid limit"}) + return + } + limit = parsedLimit + } + if limit > maxProcessDetailsPageLimit { + limit = maxProcessDetailsPageLimit + } + offset, _ := strconv.Atoi(strings.TrimSpace(c.Query("offset"))) + if offset < 0 { + offset = 0 + } + anchorID := strings.TrimSpace(c.Query("anchorId")) + if anchorID != "" { + anchorOffset, err := h.db.GetProcessDetailOffset(messageID, anchorID) + if err != nil { + h.logger.Warn("获取过程详情锚点位置失败", zap.Error(err), zap.String("messageID", messageID), zap.String("anchorID", anchorID)) + c.JSON(http.StatusNotFound, gin.H{"error": err.Error()}) + return + } + offset = anchorOffset - limit/3 + if offset < 0 { + offset = 0 + } + } + + details, total, err := h.db.GetProcessDetailsPage(messageID, limit, offset) + if err != nil { + h.logger.Error("分页获取过程详情失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + details = database.DedupeConsecutiveProcessDetails(details) + out := processDetailsToJSON(h.logger, h.db, details, false) + // A page may end between tool_call and tool_result. Return the full-history + // execution summary so the UI can render terminal status without pretending + // that an unloaded result is still running. + summary, summaryErr := h.db.GetProcessDetailsSummary(messageID) + if summaryErr != nil { + h.logger.Warn("获取分页工具执行状态失败", zap.Error(summaryErr), zap.String("messageID", messageID)) + } + var toolExecutions []database.ProcessDetailsToolExecution + if summary != nil { + toolExecutions = summary.ToolExecutions + } + c.JSON(http.StatusOK, gin.H{ + "processDetails": out, + "toolExecutions": toolExecutions, + "total": total, + "offset": offset, + "limit": limit, + "hasMore": offset+len(out) < total, + }) +} + +// GetProcessDetail 获取单条完整过程详情。列表接口默认不给工具 payload，用户点开单条工具时再拉这里。 +func (h *ConversationHandler) GetProcessDetail(c *gin.Context) { + id := strings.TrimSpace(c.Param("id")) + if id == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "process detail id required"}) + return + } + detail, err := h.db.GetProcessDetailByID(id) + if err != nil { + h.logger.Error("获取过程详情失败", zap.Error(err)) + c.JSON(http.StatusNotFound, gin.H{"error": "过程详情不存在"}) + return + } + out := processDetailsToJSON(h.logger, h.db, []database.ProcessDetail{*detail}, true) + if len(out) == 0 { + c.JSON(http.StatusNotFound, gin.H{"error": "过程详情不存在"}) + return + } + c.JSON(http.StatusOK, gin.H{"processDetail": out[0]}) +} + +func processDetailsToJSON(logger *zap.Logger, db *database.DB, details []database.ProcessDetail, includeToolPayload bool) []map[string]interface{} { + out := make([]map[string]interface{}, 0, len(details)) + for _, d := range details { + var data interface{} + if d.Data != "" { + if err := json.Unmarshal([]byte(d.Data), &data); err != nil { + logger.Warn("解析过程详情数据失败", zap.Error(err)) + } + } + if m, ok := data.(map[string]interface{}); ok { + enrichEmptyToolCallArgumentsFromExecution(logger, db, d, m) + } + if !includeToolPayload { + data = summarizeProcessDetailData(d.EventType, data) + } + out = append(out, map[string]interface{}{ + "id": d.ID, + "messageId": d.MessageID, + "conversationId": d.ConversationID, + "eventType": d.EventType, + "message": d.Message, + "data": data, + "createdAt": d.CreatedAt, + }) + } + return out +} + +func enrichEmptyToolCallArgumentsFromExecution(logger *zap.Logger, db *database.DB, detail database.ProcessDetail, data map[string]interface{}) { + if db == nil || detail.EventType != "tool_call" || !toolCallArgumentsEmpty(data) { + return + } + toolName := strings.TrimSpace(fmt.Sprint(data["toolName"])) + if toolName == "" || detail.ConversationID == "" || detail.CreatedAt.IsZero() { + return + } + execID, args, err := db.FindNearestToolExecutionArguments(detail.ConversationID, toolName, detail.CreatedAt, 5*time.Second) + if err != nil { + if logger != nil { + logger.Debug("未能从工具执行记录补全过程详情参数", + zap.Error(err), + zap.String("processDetailId", detail.ID), + zap.String("toolName", toolName)) + } + return + } + if len(args) == 0 { + return + } + data["argumentsObj"] = args + if b, err := json.Marshal(args); err == nil { + data["arguments"] = string(b) + } + if strings.TrimSpace(execID) != "" { + data["executionId"] = strings.TrimSpace(execID) + } +} + +func toolCallArgumentsEmpty(data map[string]interface{}) bool { + if data == nil { + return true + } + if args, ok := data["argumentsObj"].(map[string]interface{}); ok && len(args) > 0 { + return false + } + if raw, ok := data["arguments"]; ok { + s := strings.TrimSpace(fmt.Sprint(raw)) + return s == "" || s == "{}" || s == "null" + } + return true +} + +func summarizeProcessDetailData(eventType string, data interface{}) interface{} { + m, ok := data.(map[string]interface{}) + if !ok || (eventType != "tool_call" && eventType != "tool_result") { + return data + } + allow := map[string]bool{ + "toolName": true, "toolCallId": true, "index": true, "total": true, + "success": true, "isError": true, "executionId": true, + "einoAgent": true, "einoRole": true, "einoScope": true, "orchestration": true, + "agentFacing": true, + "status": true, "modelFacingIsError": true, "resultPreview": true, + } + out := make(map[string]interface{}, len(allow)+1) + for k, v := range m { + if allow[k] { + out[k] = v + } + } + out["_payloadDeferred"] = true + return out +} + +// UpdateConversationRequest 更新对话请求 +type UpdateConversationRequest struct { + Title string `json:"title"` +} + +// UpdateConversation 更新对话 +func (h *ConversationHandler) UpdateConversation(c *gin.Context) { + id := c.Param("id") + + var req UpdateConversationRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + if req.Title == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "标题不能为空"}) + return + } + + if err := h.db.UpdateConversationTitle(id, req.Title); err != nil { + h.logger.Error("更新对话失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + // 返回更新后的对话 + conv, err := h.db.GetConversation(id) + if err != nil { + h.logger.Error("获取更新后的对话失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + c.JSON(http.StatusOK, conv) +} + +// DeleteConversation 删除对话 +func (h *ConversationHandler) DeleteConversation(c *gin.Context) { + id := c.Param("id") + + if h.taskStopper != nil { + h.taskStopper.CancelRunningTaskForConversation(id) + } + + if err := h.db.DeleteConversation(id); err != nil { + h.logger.Error("删除对话失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + if h.audit != nil { + h.audit.Record(c, audit.Entry{ + Category: "conversation", + Action: "delete", + Result: "success", + ResourceType: "conversation", + ResourceID: id, + Message: "删除对话", + }) + } + + c.JSON(http.StatusOK, gin.H{"message": "删除成功"}) +} + +// DeleteTurnRequest 删除一轮对话（POST /api/conversations/:id/delete-turn） +type DeleteTurnRequest struct { + MessageID string `json:"messageId"` +} + +// DeleteConversationTurn 删除锚点消息所在轮次（从该轮 user 到下一轮 user 之前），并清空 last_react_*。 +func (h *ConversationHandler) DeleteConversationTurn(c *gin.Context) { + conversationID := c.Param("id") + if conversationID == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "conversation id required"}) + return + } + + var req DeleteTurnRequest + if err := c.ShouldBindJSON(&req); err != nil || req.MessageID == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "messageId required"}) + return + } + + if _, err := h.db.GetConversation(conversationID); err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": "对话不存在"}) + return + } + + deletedIDs, err := h.db.DeleteConversationTurn(conversationID, req.MessageID) + if err != nil { + h.logger.Warn("删除对话轮次失败", + zap.String("conversationId", conversationID), + zap.String("messageId", req.MessageID), + zap.Error(err), + ) + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + if h.audit != nil { + h.audit.RecordOK(c, "conversation", "delete_turn", "删除对话轮次", "conversation", conversationID, map[string]interface{}{ + "message_id": req.MessageID, + "deleted": len(deletedIDs), + }) + } + c.JSON(http.StatusOK, gin.H{ + "deletedMessageIds": deletedIDs, + "message": "ok", + }) +} diff --git a/internal/handler/conversation_delete_task_test.go b/internal/handler/conversation_delete_task_test.go new file mode 100644 index 00000000..39ef06c0 --- /dev/null +++ b/internal/handler/conversation_delete_task_test.go @@ -0,0 +1,30 @@ +package handler + +import ( + "context" + "testing" + "time" + + "go.uber.org/zap" +) + +func TestConversationHandlerDeleteConversationCancelsRunningTask(t *testing.T) { + tm := NewAgentTaskManager() + ctx, cancel := context.WithCancelCause(context.Background()) + _, err := tm.StartTask("conv-1", "hello", cancel) + if err != nil { + t.Fatalf("StartTask: %v", err) + } + + h := &AgentHandler{tasks: tm, logger: zap.NewNop()} + h.CancelRunningTaskForConversation("conv-1") + + select { + case <-ctx.Done(): + case <-time.After(2 * time.Second): + t.Fatal("task context was not cancelled") + } + if cause := context.Cause(ctx); cause != ErrTaskCancelled { + t.Fatalf("expected ErrTaskCancelled, got %v", cause) + } +} diff --git a/internal/handler/conversation_plantask_test.go b/internal/handler/conversation_plantask_test.go new file mode 100644 index 00000000..0162be34 --- /dev/null +++ b/internal/handler/conversation_plantask_test.go @@ -0,0 +1,151 @@ +package handler + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "testing" + "time" + + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +type staticConversationTaskState struct { + running bool + startedAt time.Time +} + +func (s staticConversationTaskState) ConversationTaskRuntimeState(string) (bool, time.Time) { + return s.running, s.startedAt +} + +func TestGetConversationPlanTasksRequiresAccessAndReportsProgress(t *testing.T) { + gin.SetMode(gin.TestMode) + tmp := t.TempDir() + db, err := database.NewDB(filepath.Join(tmp, "conversation-plantask.db"), zap.NewNop()) + if err != nil { + t.Fatalf("NewDB: %v", err) + } + t.Cleanup(func() { _ = db.Close() }) + conversation, err := db.CreateConversation("plan", database.ConversationCreateMeta{}) + if err != nil { + t.Fatalf("CreateConversation: %v", err) + } + user, err := db.CreateRBACUser("plan-user", "Plan User", "hash", true, nil) + if err != nil { + t.Fatalf("CreateRBACUser: %v", err) + } + base := filepath.Join(tmp, "plantask") + db.SetEinoConversationDirs(base, "", "", "") + dir := filepath.Join(base, conversation.ID) + if err := os.MkdirAll(dir, 0o755); err != nil { + t.Fatalf("MkdirAll: %v", err) + } + for name, content := range map[string]string{ + "1.json": `{"id":"1","subject":"完成项","status":"completed"}`, + "2.json": `{"id":"2","subject":"当前项","status":"in_progress"}`, + "3.json": `{"id":"3","subject":"等待项","status":"pending"}`, + } { + if err := os.WriteFile(filepath.Join(dir, name), []byte(content), 0o644); err != nil { + t.Fatalf("WriteFile: %v", err) + } + } + + handler := NewConversationHandler(db, zap.NewNop()) + handler.SetTaskStateProvider(staticConversationTaskState{running: true}) + request := func() *httptest.ResponseRecorder { + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodGet, "/api/conversations/"+conversation.ID+"/plan-tasks", nil) + c.Params = gin.Params{{Key: "id", Value: conversation.ID}} + c.Set(security.ContextSessionKey, security.Session{ + UserID: user.ID, + Scope: database.RBACScopeAssigned, + }) + handler.GetConversationPlanTasks(c) + return w + } + + w := request() + if w.Code != http.StatusForbidden { + t.Fatalf("unassigned status = %d, want %d", w.Code, http.StatusForbidden) + } + if err := db.AssignResourceToUser(user.ID, "conversation", conversation.ID); err != nil { + t.Fatalf("AssignResourceToUser: %v", err) + } + w = request() + if w.Code != http.StatusOK { + t.Fatalf("assigned status = %d: %s", w.Code, w.Body.String()) + } + var response struct { + Total int `json:"total"` + Completed int `json:"completed"` + ActiveStep int `json:"activeStep"` + Tasks []database.ConversationPlanTask `json:"tasks"` + Running bool `json:"running"` + } + if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil { + t.Fatalf("decode response: %v", err) + } + if response.Total != 3 || response.Completed != 1 || response.ActiveStep != 2 || !response.Running { + t.Fatalf("progress = %#v", response) + } +} + +func TestGetConversationPlanTasksReportsStoppedLiveTask(t *testing.T) { + gin.SetMode(gin.TestMode) + tmp := t.TempDir() + db, err := database.NewDB(filepath.Join(tmp, "conversation-plantask-stopped.db"), zap.NewNop()) + if err != nil { + t.Fatalf("NewDB: %v", err) + } + t.Cleanup(func() { _ = db.Close() }) + conversation, err := db.CreateConversation("stopped plan", database.ConversationCreateMeta{}) + if err != nil { + t.Fatalf("CreateConversation: %v", err) + } + user, err := db.CreateRBACUser("stopped-plan-user", "Stopped Plan User", "hash", true, nil) + if err != nil { + t.Fatalf("CreateRBACUser: %v", err) + } + if err := db.AssignResourceToUser(user.ID, "conversation", conversation.ID); err != nil { + t.Fatalf("AssignResourceToUser: %v", err) + } + base := filepath.Join(tmp, "plantask") + db.SetEinoConversationDirs(base, "", "", "") + dir := filepath.Join(base, conversation.ID) + if err := os.MkdirAll(dir, 0o755); err != nil { + t.Fatalf("MkdirAll: %v", err) + } + if err := os.WriteFile(filepath.Join(dir, "1.json"), []byte(`{"id":"1","subject":"残留项","status":"in_progress"}`), 0o644); err != nil { + t.Fatalf("WriteFile: %v", err) + } + + handler := NewConversationHandler(db, zap.NewNop()) + handler.SetTaskStateProvider(staticConversationTaskState{running: false}) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodGet, "/api/conversations/"+conversation.ID+"/plan-tasks", nil) + c.Params = gin.Params{{Key: "id", Value: conversation.ID}} + c.Set(security.ContextSessionKey, security.Session{UserID: user.ID, Scope: database.RBACScopeAssigned}) + handler.GetConversationPlanTasks(c) + if w.Code != http.StatusOK { + t.Fatalf("status = %d: %s", w.Code, w.Body.String()) + } + var response struct { + Running bool `json:"running"` + Total int `json:"total"` + } + if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil { + t.Fatalf("decode response: %v", err) + } + if response.Running || response.Total != 0 { + t.Fatalf("response = %#v", response) + } +} diff --git a/internal/handler/conversation_process_details_test.go b/internal/handler/conversation_process_details_test.go new file mode 100644 index 00000000..ff53898e --- /dev/null +++ b/internal/handler/conversation_process_details_test.go @@ -0,0 +1,141 @@ +package handler + +import ( + "encoding/json" + "fmt" + "net/http/httptest" + "path/filepath" + "testing" + "time" + + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/mcp" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +func TestProcessDetailsPageIncludesTerminalToolStatusAcrossPageBoundary(t *testing.T) { + gin.SetMode(gin.TestMode) + db, err := database.NewDB(filepath.Join(t.TempDir(), "process-details-page.db"), zap.NewNop()) + if err != nil { + t.Fatalf("NewDB: %v", err) + } + t.Cleanup(func() { _ = db.Close() }) + conversation, err := db.CreateConversation("page boundary", database.ConversationCreateMeta{}) + if err != nil { + t.Fatalf("CreateConversation: %v", err) + } + message, err := db.AddMessage(conversation.ID, "assistant", "done", nil) + if err != nil { + t.Fatalf("AddMessage: %v", err) + } + for i := 1; i <= 4; i++ { + id := fmt.Sprintf("call-%d", i) + if err := db.AddProcessDetail(message.ID, conversation.ID, "tool_call", "call", map[string]interface{}{ + "toolName": "http-framework-test", "toolCallId": id, "index": i, "total": 4, + }); err != nil { + t.Fatalf("AddProcessDetail(tool_call): %v", err) + } + } + for i := 1; i <= 4; i++ { + id := fmt.Sprintf("call-%d", i) + if err := db.AddProcessDetail(message.ID, conversation.ID, "tool_result", "result", map[string]interface{}{ + "toolName": "http-framework-test", "toolCallId": id, "success": true, + }); err != nil { + t.Fatalf("AddProcessDetail(tool_result): %v", err) + } + } + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("GET", "/api/messages/"+message.ID+"/process-details?limit=6&offset=0", nil) + c.Params = gin.Params{{Key: "id", Value: message.ID}} + NewConversationHandler(db, zap.NewNop()).GetMessageProcessDetails(c) + if w.Code != 200 { + t.Fatalf("status = %d: %s", w.Code, w.Body.String()) + } + var response struct { + HasMore bool `json:"hasMore"` + ProcessDetails []map[string]interface{} `json:"processDetails"` + ToolExecutions []database.ProcessDetailsToolExecution `json:"toolExecutions"` + } + if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil { + t.Fatalf("decode response: %v", err) + } + if !response.HasMore || len(response.ProcessDetails) != 6 { + t.Fatalf("page hasMore=%v details=%d, want true/6", response.HasMore, len(response.ProcessDetails)) + } + if len(response.ToolExecutions) != 4 { + t.Fatalf("tool executions = %d, want 4", len(response.ToolExecutions)) + } + for i, execution := range response.ToolExecutions { + if execution.Status != "completed" { + t.Fatalf("execution %d status = %q, want completed", i, execution.Status) + } + } +} + +func TestProcessDetailsFullBackfillsEmptyToolCallArgumentsFromExecution(t *testing.T) { + gin.SetMode(gin.TestMode) + db, err := database.NewDB(filepath.Join(t.TempDir(), "process-details-args.db"), zap.NewNop()) + if err != nil { + t.Fatalf("NewDB: %v", err) + } + t.Cleanup(func() { _ = db.Close() }) + conversation, err := db.CreateConversation("empty args", database.ConversationCreateMeta{}) + if err != nil { + t.Fatalf("CreateConversation: %v", err) + } + message, err := db.AddMessage(conversation.ID, "assistant", "done", nil) + if err != nil { + t.Fatalf("AddMessage: %v", err) + } + if err := db.AddProcessDetail(message.ID, conversation.ID, "tool_call", "calling exec", map[string]interface{}{ + "toolName": "exec", "toolCallId": "call-empty", "arguments": "", "argumentsObj": nil, + }); err != nil { + t.Fatalf("AddProcessDetail(tool_call): %v", err) + } + if err := db.SaveToolExecution(&mcp.ToolExecution{ + ID: "exec-whoami", + ToolName: "exec", + Arguments: map[string]interface{}{"command": "whoami"}, + Status: "completed", + StartTime: time.Now(), + ConversationID: conversation.ID, + }); err != nil { + t.Fatalf("SaveToolExecution: %v", err) + } + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("GET", "/api/messages/"+message.ID+"/process-details?full=1", nil) + c.Params = gin.Params{{Key: "id", Value: message.ID}} + NewConversationHandler(db, zap.NewNop()).GetMessageProcessDetails(c) + if w.Code != 200 { + t.Fatalf("status = %d: %s", w.Code, w.Body.String()) + } + var response struct { + ProcessDetails []map[string]interface{} `json:"processDetails"` + } + if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil { + t.Fatalf("decode response: %v", err) + } + if len(response.ProcessDetails) != 1 { + t.Fatalf("process details = %d, want 1", len(response.ProcessDetails)) + } + data, ok := response.ProcessDetails[0]["data"].(map[string]interface{}) + if !ok { + t.Fatalf("data = %#v", response.ProcessDetails[0]["data"]) + } + args, ok := data["argumentsObj"].(map[string]interface{}) + if !ok { + t.Fatalf("argumentsObj = %#v", data["argumentsObj"]) + } + if args["command"] != "whoami" { + t.Fatalf("command = %#v, want whoami", args["command"]) + } + if data["executionId"] != "exec-whoami" { + t.Fatalf("executionId = %#v, want exec-whoami", data["executionId"]) + } +} diff --git a/internal/handler/conversation_rbac_test.go b/internal/handler/conversation_rbac_test.go new file mode 100644 index 00000000..ea560e21 --- /dev/null +++ b/internal/handler/conversation_rbac_test.go @@ -0,0 +1,118 @@ +package handler + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "path/filepath" + "testing" + + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +func TestCreateConversationRequiresProjectAccess(t *testing.T) { + gin.SetMode(gin.TestMode) + db, user := setupConversationRBACTest(t) + project, err := db.CreateProject(&database.Project{Name: "hidden"}) + if err != nil { + t.Fatalf("CreateProject: %v", err) + } + handler := NewConversationHandler(db, zap.NewNop()) + + w := performConversationRequest(user, http.MethodPost, "/api/conversations", map[string]string{ + "title": "blocked", + "projectId": project.ID, + }, handler.CreateConversation) + if w.Code != http.StatusForbidden { + t.Fatalf("status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String()) + } + + if err := db.AssignResourceToUser(user.ID, "project", project.ID); err != nil { + t.Fatalf("AssignResourceToUser: %v", err) + } + w = performConversationRequest(user, http.MethodPost, "/api/conversations", map[string]string{ + "title": "allowed", + "projectId": project.ID, + }, handler.CreateConversation) + if w.Code != http.StatusOK { + t.Fatalf("status = %d, want %d: %s", w.Code, http.StatusOK, w.Body.String()) + } +} + +func TestSetConversationProjectRequiresProjectAccess(t *testing.T) { + gin.SetMode(gin.TestMode) + db, user := setupConversationRBACTest(t) + project, err := db.CreateProject(&database.Project{Name: "hidden"}) + if err != nil { + t.Fatalf("CreateProject: %v", err) + } + conv, err := db.CreateConversation("owned", database.ConversationCreateMeta{}) + if err != nil { + t.Fatalf("CreateConversation: %v", err) + } + if err := db.SetResourceOwner("conversation", conv.ID, user.ID); err != nil { + t.Fatalf("SetResourceOwner: %v", err) + } + if err := db.AssignResourceToUser(user.ID, "conversation", conv.ID); err != nil { + t.Fatalf("AssignResourceToUser conversation: %v", err) + } + handler := NewConversationHandler(db, zap.NewNop()) + + w := performConversationRequest(user, http.MethodPut, "/api/conversations/"+conv.ID+"/project", map[string]string{ + "projectId": project.ID, + }, func(c *gin.Context) { + c.Params = gin.Params{{Key: "id", Value: conv.ID}} + handler.SetConversationProject(c) + }) + if w.Code != http.StatusForbidden { + t.Fatalf("status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String()) + } + + if err := db.AssignResourceToUser(user.ID, "project", project.ID); err != nil { + t.Fatalf("AssignResourceToUser project: %v", err) + } + w = performConversationRequest(user, http.MethodPut, "/api/conversations/"+conv.ID+"/project", map[string]string{ + "projectId": project.ID, + }, func(c *gin.Context) { + c.Params = gin.Params{{Key: "id", Value: conv.ID}} + handler.SetConversationProject(c) + }) + if w.Code != http.StatusOK { + t.Fatalf("status = %d, want %d: %s", w.Code, http.StatusOK, w.Body.String()) + } +} + +func setupConversationRBACTest(t *testing.T) (*database.DB, *database.RBACUser) { + t.Helper() + db, err := database.NewDB(filepath.Join(t.TempDir(), "conversation-rbac.db"), zap.NewNop()) + if err != nil { + t.Fatalf("NewDB: %v", err) + } + t.Cleanup(func() { _ = db.Close() }) + user, err := db.CreateRBACUser("operator1", "Operator One", "hash", true, nil) + if err != nil { + t.Fatalf("CreateRBACUser: %v", err) + } + return db, user +} + +func performConversationRequest(user *database.RBACUser, method, path string, body map[string]string, handler gin.HandlerFunc) *httptest.ResponseRecorder { + payload, _ := json.Marshal(body) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(method, path, bytes.NewReader(payload)) + c.Request.Header.Set("Content-Type", "application/json") + c.Set(security.ContextSessionKey, security.Session{ + UserID: user.ID, + Username: user.Username, + Permissions: map[string]bool{"chat:write": true}, + Scope: database.RBACScopeAssigned, + }) + handler(c) + return w +} diff --git a/internal/handler/eino_empty_response_continue.go b/internal/handler/eino_empty_response_continue.go new file mode 100644 index 00000000..2c57df86 --- /dev/null +++ b/internal/handler/eino_empty_response_continue.go @@ -0,0 +1,83 @@ +package handler + +import ( + "context" + "fmt" + "time" + + "cyberstrike-ai/internal/agent" + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/multiagent" + + "go.uber.org/zap" +) + +// rebindEinoRunningTask 中断并继续 / 空正文续跑：重建 cancel 链与超时 ctx，保持任务 running。 +func (h *AgentHandler) rebindEinoRunningTask(parent context.Context, conversationID string, timeoutCancel context.CancelFunc) (context.Context, context.CancelCauseFunc, context.Context, context.CancelFunc) { + if timeoutCancel != nil { + timeoutCancel() + } + baseCtx, cancelWithCause := context.WithCancelCause(detachedAgentContext(parent)) + h.tasks.BindTaskCancel(conversationID, cancelWithCause) + taskCtx, newTimeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute) + h.tasks.UpdateTaskStatus(conversationID, "running") + return baseCtx, cancelWithCause, taskCtx, newTimeoutCancel +} + +// tryContinueOnEinoEmptyResponse Run 成功但 Response 为 emptyHint 时退避续跑；true 表示已准备下一段 Run。 +func (h *AgentHandler) tryContinueOnEinoEmptyResponse( + taskCtx context.Context, + mw *config.MultiAgentEinoMiddlewareConfig, + conversationID string, + result *multiagent.RunResult, + attempt *int, + curHistory *[]agent.ChatMessage, + curFinalMessage *string, + progressCallback func(eventType, message string, data interface{}), +) bool { + if result == nil || !multiagent.IsEinoEmptyResponseResult(result) || !multiagent.HasEinoResumeTrace(result) { + return false + } + maxAttempts := multiagent.EmptyResponseContinueMaxAttemptsFromConfig(mw) + if *attempt >= maxAttempts { + if h.logger != nil { + h.logger.Warn("eino empty response continue exhausted", + zap.String("conversationId", conversationID), + zap.Int("maxAttempts", maxAttempts)) + } + return false + } + *attempt++ + h.persistEinoAgentTraceForResume(conversationID, result) + + backoff := multiagent.EmptyResponseContinueBackoff(*attempt-1, mw) + waitMsg := fmt.Sprintf("会话已结束但未捕获到助手正文，%d 秒后第 %d/%d 次自动续跑…", + int(backoff.Seconds()), *attempt, maxAttempts) + if progressCallback != nil { + progressCallback("eino_empty_response_continue", waitMsg, map[string]interface{}{ + "conversationId": conversationID, + "source": "eino", + "attempt": *attempt, + "maxAttempts": maxAttempts, + "backoffSec": int(backoff.Seconds()), + }) + } + select { + case <-taskCtx.Done(): + return false + case <-time.After(backoff): + } + + h.applyEinoTraceResumeSegment(conversationID, result, curHistory, curFinalMessage, "") + if progressCallback != nil { + progressCallback("eino_empty_response_continue", "已恢复上下文，正在续跑…", map[string]interface{}{ + "conversationId": conversationID, + "source": "eino", + "attempt": *attempt, + "maxAttempts": maxAttempts, + "contextSource": "empty_response_continue", + "contextInjection": false, + }) + } + return true +} diff --git a/internal/handler/eino_resume_segment.go b/internal/handler/eino_resume_segment.go new file mode 100644 index 00000000..811162a1 --- /dev/null +++ b/internal/handler/eino_resume_segment.go @@ -0,0 +1,27 @@ +package handler + +import ( + "context" + + "cyberstrike-ai/internal/agent" + "cyberstrike-ai/internal/multiagent" +) + +// applyEinoTraceResumeSegment 中断并继续：persist last_react_* → loadHistory，可选替换下一段 user 文案。 +func (h *AgentHandler) applyEinoTraceResumeSegment( + conversationID string, + result *multiagent.RunResult, + curHistory *[]agent.ChatMessage, + curFinalMessage *string, + segmentUserMessage string, +) { + if shouldPersistEinoAgentTraceAfterRunError(context.Background()) { + h.persistEinoAgentTraceForResume(conversationID, result) + } + if hist, err := h.loadHistoryFromAgentTrace(conversationID); err == nil && len(hist) > 0 { + *curHistory = hist + } + if segmentUserMessage != "" { + *curFinalMessage = segmentUserMessage + } +} diff --git a/internal/handler/eino_single_agent.go b/internal/handler/eino_single_agent.go new file mode 100644 index 00000000..8ec130f0 --- /dev/null +++ b/internal/handler/eino_single_agent.go @@ -0,0 +1,526 @@ +package handler + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "strings" + "sync" + "time" + + "cyberstrike-ai/internal/agentfinalizer" + "cyberstrike-ai/internal/mcp" + "cyberstrike-ai/internal/multiagent" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +// EinoSingleAgentLoopStream Eino ADK 单代理（ChatModelAgent + Runner）流式对话；不依赖 multi_agent.enabled。 +func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) { + c.Header("Content-Type", "text/event-stream; charset=utf-8") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + + var req ChatRequest + if err := c.ShouldBindJSON(&req); err != nil { + ev := StreamEvent{Type: "error", Message: "请求参数错误: " + err.Error()} + b, _ := json.Marshal(ev) + fmt.Fprintf(c.Writer, "data: %s\n\n", b) + done := StreamEvent{Type: "done", Message: ""} + db, _ := json.Marshal(done) + fmt.Fprintf(c.Writer, "data: %s\n\n", db) + if flusher, ok := c.Writer.(http.Flusher); ok { + flusher.Flush() + } + return + } + + c.Header("X-Accel-Buffering", "no") + + var baseCtx context.Context + clientDisconnected := false + var sseWriteMu sync.Mutex + var ssePublishConversationID string + sendEvent := func(eventType, message string, data interface{}) { + if eventType == "error" && baseCtx != nil { + cause := context.Cause(baseCtx) + if errors.Is(cause, ErrTaskCancelled) || errors.Is(cause, multiagent.ErrInterruptContinue) { + return + } + } + ev := StreamEvent{Type: eventType, Message: message, Data: data} + b, errMarshal := json.Marshal(ev) + if errMarshal != nil { + b = []byte(`{"type":"error","message":"marshal failed"}`) + } + sseLine := make([]byte, 0, len(b)+8) + sseLine = append(sseLine, []byte("data: ")...) + sseLine = append(sseLine, b...) + sseLine = append(sseLine, '\n', '\n') + if ssePublishConversationID != "" && h.taskEventBus != nil { + h.taskEventBus.Publish(ssePublishConversationID, sseLine) + } + if clientDisconnected { + return + } + select { + case <-c.Request.Context().Done(): + clientDisconnected = true + return + default: + } + sseWriteMu.Lock() + _, err := c.Writer.Write(sseLine) + if err != nil { + sseWriteMu.Unlock() + clientDisconnected = true + return + } + if flusher, ok := c.Writer.(http.Flusher); ok { + flusher.Flush() + } else { + c.Writer.Flush() + } + sseWriteMu.Unlock() + } + + h.logger.Info("收到 Eino ADK 单代理流式请求", + zap.String("conversationId", req.ConversationID), + ) + + prep, err := h.prepareMultiAgentSession(&req, c, "eino_agent_stream") + if err != nil { + sendEvent("error", err.Error(), nil) + sendEvent("done", "", nil) + return + } + ssePublishConversationID = prep.ConversationID + if prep.CreatedNew { + sendEvent("conversation", "会话已创建", map[string]interface{}{ + "conversationId": prep.ConversationID, + }) + } + + conversationID := prep.ConversationID + assistantMessageID := prep.AssistantMessageID + h.activateHITLForConversation(conversationID, req.Hitl) + if h.hitlManager != nil { + defer h.hitlManager.DeactivateConversation(conversationID) + } + + if prep.UserMessageID != "" { + sendEvent("message_saved", "", map[string]interface{}{ + "conversationId": conversationID, + "userMessageId": prep.UserMessageID, + }) + } + if h.runRoleWorkflowStreamIfBound(c, &req, prep, sendEvent) { + return + } + + var cancelWithCause context.CancelCauseFunc + curFinalMessage := prep.FinalMessage + curHistory := prep.History + roleTools := prep.RoleTools + + taskStatus := "completed" + // 仅在成功 StartTask 后再 FinishTask。若 StartTask 因 ErrTaskAlreadyRunning 失败仍 defer FinishTask， + // 会误删其他连接上正在运行的同会话任务，导致「第一次拦截、第二次却放行」。 + taskOwned := false + defer func() { + if taskOwned { + h.tasks.FinishTask(conversationID, taskStatus) + } + }() + + sendEvent("progress", "正在启动 Eino ADK 单代理（ChatModelAgent）...", map[string]interface{}{ + "conversationId": conversationID, + }) + + stopKeepalive := runSSEKeepalive(c, &sseWriteMu) + defer stopKeepalive() + + if h.config == nil { + taskStatus = "failed" + h.tasks.UpdateTaskStatus(conversationID, taskStatus) + sendEvent("error", "服务器配置未加载", nil) + sendEvent("done", "", map[string]interface{}{"conversationId": conversationID}) + return + } + runCfg, resolvedAIChannelID, err := h.configForAIChannel(req.AIChannelID) + if err != nil { + taskStatus = "failed" + h.tasks.UpdateTaskStatus(conversationID, taskStatus) + sendEvent("error", err.Error(), nil) + sendEvent("done", "", map[string]interface{}{"conversationId": conversationID}) + return + } + + var result *multiagent.RunResult + var runErr error + + baseCtx, cancelWithCause = context.WithCancelCause(detachedAgentContext(c.Request.Context())) + taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute) + + if _, err := h.tasks.StartTask(conversationID, req.Message, cancelWithCause); err != nil { + var errorMsg string + if errors.Is(err, ErrTaskAlreadyRunning) { + errorMsg = "⚠️ 当前会话已有任务正在执行中，请等待当前任务完成或点击「停止任务」后再尝试。" + sendEvent("error", errorMsg, map[string]interface{}{ + "conversationId": conversationID, + "errorType": "task_already_running", + }) + } else { + errorMsg = "❌ 无法启动任务: " + err.Error() + sendEvent("error", errorMsg, nil) + } + if assistantMessageID != "" { + _, _ = h.db.Exec("UPDATE messages SET content = ?, updated_at = ? WHERE id = ?", errorMsg, time.Now(), assistantMessageID) + } + sendEvent("done", "", map[string]interface{}{"conversationId": conversationID}) + timeoutCancel() + return + } + taskOwned = true + + var cumulativeMCPExecutionIDs []string + // 同一请求内分段续跑时，主代理 iteration 事件按偏移累计，避免 UI 出现「第3轮 → 第1轮」回跳。 + var mainIterationOffset int + var emptyResponseContinueAttempt int + var finalizationAutoContinueAttempt int + var decision agentfinalizer.Decision + + for { + segmentMainIterationMax := 0 + rawProgressCallback := h.createProgressCallback(taskCtx, cancelWithCause, conversationID, assistantMessageID, sendEvent) + progressCallback := func(eventType, message string, data interface{}) { + if eventType == "iteration" { + if m, ok := data.(map[string]interface{}); ok { + if scope, _ := m["einoScope"].(string); scope == "main" { + raw := 0 + switch v := m["iteration"].(type) { + case int: + raw = v + case int32: + raw = int(v) + case int64: + raw = int(v) + case float64: + raw = int(v) + case float32: + raw = int(v) + } + if raw > 0 { + if raw > segmentMainIterationMax { + segmentMainIterationMax = raw + } + m["iteration"] = raw + mainIterationOffset + } + } + } + } + rawProgressCallback(eventType, message, data) + } + taskCtxLoop := mcp.WithMCPConversationID(taskCtx, conversationID) + taskCtxLoop = mcp.WithToolRunRegistry(taskCtxLoop, h.tasks) + taskCtxLoop = mcp.WithEinoExecuteRunRegistry(taskCtxLoop, h.tasks) + taskCtxLoop = multiagent.WithAgentRuntimeCancelRegistrar(taskCtxLoop, func(cancel func(error) bool) func() { + return h.tasks.BindAgentRuntimeCancel(conversationID, cancel) + }) + taskCtxLoop = multiagent.WithAgentTurnLoopInterruptRegistrar(taskCtxLoop, func(push func(string) bool) func() { + return h.tasks.BindAgentTurnLoopInterrupt(conversationID, push) + }) + taskCtxLoop = multiagent.WithHITLToolInterceptor(taskCtxLoop, func(ctx context.Context, toolName, arguments string) (string, error) { + return h.interceptHITLForEinoTool(ctx, cancelWithCause, conversationID, assistantMessageID, sendEvent, toolName, arguments) + }) + + result, runErr = multiagent.RunEinoSingleChatModelAgent( + taskCtxLoop, + runCfg, + &runCfg.MultiAgent, + h.agent, + h.db, + h.logger, + conversationID, + h.conversationProjectID(conversationID), + curFinalMessage, + curHistory, + roleTools, + progressCallback, + chatReasoningToClientIntent(req.Reasoning), + h.agentSessionContextBlock(conversationID), + ) + _ = resolvedAIChannelID + + if result != nil && len(result.MCPExecutionIDs) > 0 { + cumulativeMCPExecutionIDs = mergeMCPExecutionIDLists(cumulativeMCPExecutionIDs, result.MCPExecutionIDs) + } + + if runErr == nil { + mw := &h.config.MultiAgent.EinoMiddleware + if h.tryContinueOnEinoEmptyResponse(taskCtx, mw, conversationID, result, &emptyResponseContinueAttempt, &curHistory, &curFinalMessage, progressCallback) { + mainIterationOffset += segmentMainIterationMax + timeoutCancel() + baseCtx, cancelWithCause, taskCtx, timeoutCancel = h.rebindEinoRunningTask(taskCtx, conversationID, timeoutCancel) + continue + } + decision = h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, "eino_single", result, cumulativeMCPExecutionIDs, requestRequiresExecutionEvidence(&req)) + if h.tryAutoContinueAfterFinalization(taskCtx, conversationID, result, decision, &finalizationAutoContinueAttempt, &curHistory, &curFinalMessage, progressCallback) { + mainIterationOffset += segmentMainIterationMax + timeoutCancel() + baseCtx, cancelWithCause, taskCtx, timeoutCancel = h.rebindEinoRunningTask(taskCtx, conversationID, timeoutCancel) + continue + } + timeoutCancel() + break + } + + cause := context.Cause(baseCtx) + if cause == nil { + switch { + case errors.Is(runErr, multiagent.ErrInterruptContinue): + cause = multiagent.ErrInterruptContinue + case errors.Is(runErr, ErrTaskCancelled): + cause = ErrTaskCancelled + } + } + if errors.Is(cause, multiagent.ErrInterruptContinue) { + if shouldPersistEinoAgentTraceAfterRunError(baseCtx) { + h.persistEinoAgentTraceForResume(conversationID, result) + } + note := h.tasks.TakeInterruptContinueNote(conversationID) + icSummary := interruptContinueTimelineSummary(note) + progressCallback("user_interrupt_continue", icSummary, map[string]interface{}{ + "conversationId": conversationID, + "rawReason": strings.TrimSpace(note), + "emptyReason": strings.TrimSpace(note) == "", + "kind": "no_active_mcp_tool", + }) + inject := formatInterruptContinueUserMessage(note) + // 不写入 messages 表为 user 气泡：避免主对话流出现大段模板；说明已由 user_interrupt_continue 记入助手 process_details（迭代详情）。 + if hist, err := h.loadHistoryFromAgentTrace(conversationID); err == nil && len(hist) > 0 { + curHistory = hist + } + curFinalMessage = inject + sendEvent("progress", "已合并用户补充与最新轨迹，正在继续推理…", map[string]interface{}{ + "conversationId": conversationID, + "source": "interrupt_continue", + }) + mainIterationOffset += segmentMainIterationMax + timeoutCancel() + baseCtx, cancelWithCause = context.WithCancelCause(detachedAgentContext(baseCtx)) + h.tasks.BindTaskCancel(conversationID, cancelWithCause) + taskCtx, timeoutCancel = context.WithTimeout(baseCtx, 600*time.Minute) + h.tasks.UpdateTaskStatus(conversationID, "running") + continue + } + + if shouldPersistEinoAgentTraceAfterRunError(baseCtx) { + h.persistEinoAgentTraceForResume(conversationID, result) + } + if errors.Is(cause, ErrTaskCancelled) { + taskStatus = "cancelled" + h.tasks.UpdateTaskStatus(conversationID, taskStatus) + cancelMsg := "任务已被用户取消，后续操作已停止。" + if assistantMessageID != "" { + if result != nil { + if err := h.mergeAssistantMessagePartialOnCancel(assistantMessageID, result.Response); err != nil { + h.logger.Warn("合并取消前的部分回复失败", zap.Error(err)) + } + } + if err := h.appendAssistantMessageNotice(assistantMessageID, cancelMsg); err != nil { + h.logger.Warn("更新取消后的助手消息失败", zap.Error(err)) + } + _ = h.db.AddProcessDetail(assistantMessageID, conversationID, "cancelled", cancelMsg, nil) + } + sendEvent("cancelled", cancelMsg, map[string]interface{}{ + "conversationId": conversationID, + "messageId": assistantMessageID, + }) + sendEvent("done", "", map[string]interface{}{"conversationId": conversationID}) + timeoutCancel() + return + } + + if errors.Is(runErr, context.DeadlineExceeded) || errors.Is(context.Cause(taskCtx), context.DeadlineExceeded) { + taskStatus = "timeout" + h.tasks.UpdateTaskStatus(conversationID, taskStatus) + timeoutMsg := "任务执行超时，已自动终止。" + if assistantMessageID != "" { + _, _ = h.db.Exec("UPDATE messages SET content = ?, updated_at = ? WHERE id = ?", timeoutMsg, time.Now(), assistantMessageID) + _ = h.db.AddProcessDetail(assistantMessageID, conversationID, "timeout", timeoutMsg, nil) + } + sendEvent("error", timeoutMsg, map[string]interface{}{ + "conversationId": conversationID, + "messageId": assistantMessageID, + "errorType": "timeout", + }) + sendEvent("done", "", map[string]interface{}{"conversationId": conversationID}) + timeoutCancel() + return + } + + h.logger.Error("Eino ADK 单代理执行失败", zap.Error(runErr)) + taskStatus = "failed" + h.tasks.UpdateTaskStatus(conversationID, taskStatus) + errMsg := "执行失败: " + runErr.Error() + if assistantMessageID != "" { + _, _ = h.db.Exec("UPDATE messages SET content = ?, updated_at = ? WHERE id = ?", errMsg, time.Now(), assistantMessageID) + _ = h.db.AddProcessDetail(assistantMessageID, conversationID, "error", errMsg, nil) + } + sendEvent("error", errMsg, map[string]interface{}{ + "conversationId": conversationID, + "messageId": assistantMessageID, + }) + sendEvent("done", "", map[string]interface{}{"conversationId": conversationID}) + timeoutCancel() + return + } + + timeoutCancel() + + if decision.CompletionReason == "" { + decision = h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, "eino_single", result, cumulativeMCPExecutionIDs, requestRequiresExecutionEvidence(&req)) + } + h.persistFinalizationDecision(conversationID, assistantMessageID, "eino_single", cumulativeMCPExecutionIDs, multiagent.AggregatedReasoningFromTraceJSON(result.LastAgentTraceInput), decision) + + if result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "" { + if err := h.db.SaveAgentTrace(conversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); err != nil { + h.logger.Warn("保存代理轨迹失败", zap.Error(err)) + } + } + + responseText := decision.FinalText + if !decision.Finalizable { + responseText = finalizationBlockedMessage(decision) + sendEvent("finalization_check", responseText, decision) + taskStatus = decision.Status + h.tasks.UpdateTaskStatus(conversationID, taskStatus) + } + sendEvent("response", responseText, finalizationResponsePayload(decision, map[string]interface{}{ + "mcpExecutionIds": cumulativeMCPExecutionIDs, + "conversationId": conversationID, + "messageId": assistantMessageID, + "agentMode": "eino_single", + })) + sendEvent("done", "", map[string]interface{}{"conversationId": conversationID}) +} + +// EinoSingleAgentLoop Eino ADK 单代理非流式对话。 +func (h *AgentHandler) EinoSingleAgentLoop(c *gin.Context) { + var req ChatRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + h.logger.Info("收到 Eino ADK 单代理非流式请求", zap.String("conversationId", req.ConversationID)) + + prep, err := h.prepareMultiAgentSession(&req, c, "eino_agent") + if err != nil { + status, msg := multiAgentHTTPErrorStatus(err) + c.JSON(status, gin.H{"error": msg}) + return + } + h.activateHITLForConversation(prep.ConversationID, req.Hitl) + if h.hitlManager != nil { + defer h.hitlManager.DeactivateConversation(prep.ConversationID) + } + if h.runRoleWorkflowJSONIfBound(c, &req, prep) { + return + } + + var progressBuf strings.Builder + progressCallbackRaw := func(eventType, message string, data interface{}) { + progressBuf.WriteString(eventType) + progressBuf.WriteByte('\n') + } + baseCtx, cancelWithCause := context.WithCancelCause(c.Request.Context()) + defer cancelWithCause(nil) + taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute) + defer timeoutCancel() + progressCallback := h.createProgressCallback(taskCtx, cancelWithCause, prep.ConversationID, prep.AssistantMessageID, progressCallbackRaw) + taskCtx = multiagent.WithHITLToolInterceptor(taskCtx, func(ctx context.Context, toolName, arguments string) (string, error) { + return h.interceptHITLForEinoTool(ctx, cancelWithCause, prep.ConversationID, prep.AssistantMessageID, nil, toolName, arguments) + }) + + if h.config == nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "服务器配置未加载"}) + return + } + runCfg, _, err := h.configForAIChannel(req.AIChannelID) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + curHist := prep.History + curMsg := prep.FinalMessage + var result *multiagent.RunResult + var runErr error + var emptyResponseContinueAttempt int + var finalizationAutoContinueAttempt int + var decision agentfinalizer.Decision + for { + result, runErr = multiagent.RunEinoSingleChatModelAgent( + taskCtx, + runCfg, + &runCfg.MultiAgent, + h.agent, + h.db, + h.logger, + prep.ConversationID, + h.conversationProjectID(prep.ConversationID), + curMsg, + curHist, + prep.RoleTools, + progressCallback, + chatReasoningToClientIntent(req.Reasoning), + h.agentSessionContextBlock(prep.ConversationID), + ) + if runErr != nil { + if shouldPersistEinoAgentTraceAfterRunError(baseCtx) { + h.persistEinoAgentTraceForResume(prep.ConversationID, result) + } + c.JSON(http.StatusInternalServerError, gin.H{"error": runErr.Error()}) + return + } + mw := &h.config.MultiAgent.EinoMiddleware + if h.tryContinueOnEinoEmptyResponse(taskCtx, mw, prep.ConversationID, result, &emptyResponseContinueAttempt, &curHist, &curMsg, progressCallback) { + continue + } + decision = h.decideAgentRunForDeliveryWithPolicy(prep.ConversationID, prep.AssistantMessageID, "eino_single", result, result.MCPExecutionIDs, requestRequiresExecutionEvidence(&req)) + if h.tryAutoContinueAfterFinalization(taskCtx, prep.ConversationID, result, decision, &finalizationAutoContinueAttempt, &curHist, &curMsg, progressCallback) { + continue + } + break + } + + h.persistFinalizationDecision(prep.ConversationID, prep.AssistantMessageID, "eino_single", result.MCPExecutionIDs, multiagent.AggregatedReasoningFromTraceJSON(result.LastAgentTraceInput), decision) + if result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "" { + _ = h.db.SaveAgentTrace(prep.ConversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput) + } + + responseText := decision.FinalText + if !decision.Finalizable { + responseText = finalizationBlockedMessage(decision) + } + c.JSON(http.StatusOK, gin.H{ + "response": responseText, + "conversationId": prep.ConversationID, + "mcpExecutionIds": result.MCPExecutionIDs, + "assistantMessageId": prep.AssistantMessageID, + "agentMode": "eino_single", + "finalized": decision.Finalized, + "finalizable": decision.Finalizable, + "status": decision.Status, + "completionReason": decision.CompletionReason, + "evidenceVerified": decision.EvidenceVerified, + "evidenceRefs": decision.EvidenceRefs, + "pendingExecutionIds": decision.PendingExecutionIDs, + "missingChecks": decision.MissingChecks, + }) +} diff --git a/internal/handler/external_mcp.go b/internal/handler/external_mcp.go new file mode 100644 index 00000000..fff5784d --- /dev/null +++ b/internal/handler/external_mcp.go @@ -0,0 +1,506 @@ +package handler + +import ( + "fmt" + "net/http" + "os" + "sync" + + "cyberstrike-ai/internal/audit" + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/mcp" + "cyberstrike-ai/internal/security" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" + "gopkg.in/yaml.v3" +) + +// ExternalMCPHandler 外部MCP处理器 +type ExternalMCPHandler struct { + manager *mcp.ExternalMCPManager + config *config.Config + configPath string + logger *zap.Logger + audit *audit.Service + mu sync.RWMutex +} + +// SetAudit wires platform audit logging. +func (h *ExternalMCPHandler) SetAudit(s *audit.Service) { + h.audit = s +} + +// NewExternalMCPHandler 创建外部MCP处理器 +func NewExternalMCPHandler(manager *mcp.ExternalMCPManager, cfg *config.Config, configPath string, logger *zap.Logger) *ExternalMCPHandler { + return &ExternalMCPHandler{ + manager: manager, + config: cfg, + configPath: configPath, + logger: logger, + } +} + +// GetExternalMCPs 获取所有外部MCP配置 +func (h *ExternalMCPHandler) GetExternalMCPs(c *gin.Context) { + h.mu.RLock() + defer h.mu.RUnlock() + + configs := h.manager.GetConfigs() + + // 获取所有外部MCP的工具数量 + toolCounts := h.manager.GetToolCounts() + + // 转换为响应格式 + result := make(map[string]ExternalMCPResponse) + for name, cfg := range configs { + client, exists := h.manager.GetClient(name) + status := "disconnected" + if exists { + status = client.GetStatus() + } else if h.isEnabled(cfg) { + status = "disconnected" + } else { + status = "disabled" + } + + toolCount := toolCounts[name] + errorMsg := externalMCPStatusError(h.manager, name, status) + + result[name] = ExternalMCPResponse{ + Config: externalMCPConfigForResponse(c, cfg), + Status: status, + ToolCount: toolCount, + Error: errorMsg, + } + } + + c.JSON(http.StatusOK, gin.H{ + "servers": result, + "stats": h.manager.GetStats(), + }) +} + +// GetExternalMCP 获取单个外部MCP配置 +func (h *ExternalMCPHandler) GetExternalMCP(c *gin.Context) { + name := c.Param("name") + + h.mu.RLock() + defer h.mu.RUnlock() + + configs := h.manager.GetConfigs() + cfg, exists := configs[name] + if !exists { + c.JSON(http.StatusNotFound, gin.H{"error": "外部MCP配置不存在"}) + return + } + + client, clientExists := h.manager.GetClient(name) + status := "disconnected" + if clientExists { + status = client.GetStatus() + } else if h.isEnabled(cfg) { + status = "disconnected" + } else { + status = "disabled" + } + + // 获取工具数量 + toolCount := 0 + if clientExists && client.IsConnected() { + if count, err := h.manager.GetToolCount(name); err == nil { + toolCount = count + } + } + + c.JSON(http.StatusOK, ExternalMCPResponse{ + Config: externalMCPConfigForResponse(c, cfg), + Status: status, + ToolCount: toolCount, + Error: externalMCPStatusError(h.manager, name, status), + }) +} + +func externalMCPConfigForResponse(c *gin.Context, cfg config.ExternalMCPServerConfig) config.ExternalMCPServerConfig { + if security.SessionHasPermission(c, "mcp:write") { + return cfg + } + copyCfg := cfg + if len(cfg.Env) > 0 { + copyCfg.Env = make(map[string]string, len(cfg.Env)) + for key := range cfg.Env { + copyCfg.Env[key] = "***" + } + } + if len(cfg.Headers) > 0 { + copyCfg.Headers = make(map[string]string, len(cfg.Headers)) + for key := range cfg.Headers { + copyCfg.Headers[key] = "***" + } + } + return copyCfg +} + +// externalMCPStatusError 在 error/disconnected 状态下返回最近错误（含断连原因）。 +func externalMCPStatusError(manager *mcp.ExternalMCPManager, name, status string) string { + if status != "error" && status != "disconnected" { + return "" + } + return manager.GetError(name) +} + +// AddOrUpdateExternalMCP 添加或更新外部MCP配置 +func (h *ExternalMCPHandler) AddOrUpdateExternalMCP(c *gin.Context) { + var req AddOrUpdateExternalMCPRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()}) + return + } + + name := c.Param("name") + if name == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "名称不能为空"}) + return + } + + // 验证配置 + if err := h.validateConfig(req.Config); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + h.mu.Lock() + defer h.mu.Unlock() + + // 添加或更新配置 + if err := h.manager.AddOrUpdateConfig(name, req.Config); err != nil { + h.logger.Error("添加或更新外部MCP配置失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": "添加或更新配置失败: " + err.Error()}) + return + } + + // 更新内存中的配置 + if h.config.ExternalMCP.Servers == nil { + h.config.ExternalMCP.Servers = make(map[string]config.ExternalMCPServerConfig) + } + + cfg := req.Config + + // 官方 disabled 字段 → ExternalMCPEnable 取反 + if cfg.Disabled { + cfg.ExternalMCPEnable = false + } else if !cfg.ExternalMCPEnable { + // 用户未显式设置 external_mcp_enable，官方配置默认就是启用的 + cfg.ExternalMCPEnable = true + } + + // 展开 ${VAR} 环境变量 + config.ExpandConfigEnv(&cfg) + + h.config.ExternalMCP.Servers[name] = cfg + + // 保存到配置文件 + 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("外部MCP配置已更新", zap.String("name", name)) + if h.audit != nil { + h.audit.Record(c, audit.Entry{ + Category: "external_mcp", + Action: "upsert", + Result: "success", + ResourceType: "external_mcp", + ResourceID: name, + Message: "更新外部 MCP 配置", + }) + } + c.JSON(http.StatusOK, gin.H{"message": "配置已更新"}) +} + +// DeleteExternalMCP 删除外部MCP配置 +func (h *ExternalMCPHandler) DeleteExternalMCP(c *gin.Context) { + name := c.Param("name") + + h.mu.Lock() + defer h.mu.Unlock() + + // 移除配置 + if err := h.manager.RemoveConfig(name); err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": "配置不存在"}) + return + } + + // 从内存配置中删除 + if h.config.ExternalMCP.Servers != nil { + delete(h.config.ExternalMCP.Servers, name) + } + + // 保存到配置文件 + 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("外部MCP配置已删除", zap.String("name", name)) + if h.audit != nil { + h.audit.Record(c, audit.Entry{ + Category: "external_mcp", + Action: "delete", + Result: "success", + ResourceType: "external_mcp", + ResourceID: name, + Message: "删除外部 MCP 配置", + }) + } + c.JSON(http.StatusOK, gin.H{"message": "配置已删除"}) +} + +// StartExternalMCP 启动外部MCP +func (h *ExternalMCPHandler) StartExternalMCP(c *gin.Context) { + name := c.Param("name") + + h.mu.Lock() + defer h.mu.Unlock() + + // 更新配置为启用 + if h.config.ExternalMCP.Servers == nil { + h.config.ExternalMCP.Servers = make(map[string]config.ExternalMCPServerConfig) + } + cfg := h.config.ExternalMCP.Servers[name] + cfg.ExternalMCPEnable = true + h.config.ExternalMCP.Servers[name] = cfg + + // 保存到配置文件 + if err := h.saveConfig(); err != nil { + h.logger.Error("保存配置失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": "保存配置失败: " + err.Error()}) + return + } + + // 启动客户端（立即创建客户端并设置状态为connecting，实际连接在后台进行） + h.logger.Info("开始启动外部MCP", zap.String("name", name)) + if err := h.manager.StartClient(name); err != nil { + h.logger.Error("启动外部MCP失败", zap.String("name", name), zap.Error(err)) + c.JSON(http.StatusBadRequest, gin.H{ + "error": err.Error(), + "status": "error", + }) + return + } + + // 获取客户端状态（应该是connecting） + client, exists := h.manager.GetClient(name) + status := "connecting" + if exists { + status = client.GetStatus() + } + + // 立即返回，不等待连接完成 + // 客户端会在后台异步连接，用户可以通过状态查询接口查看连接状态 + c.JSON(http.StatusOK, gin.H{ + "message": "外部MCP启动请求已提交，正在后台连接中", + "status": status, + }) +} + +// StopExternalMCP 停止外部MCP +func (h *ExternalMCPHandler) StopExternalMCP(c *gin.Context) { + name := c.Param("name") + + h.mu.Lock() + defer h.mu.Unlock() + + // 停止客户端 + if err := h.manager.StopClient(name); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + // 更新配置 + if h.config.ExternalMCP.Servers == nil { + h.config.ExternalMCP.Servers = make(map[string]config.ExternalMCPServerConfig) + } + cfg := h.config.ExternalMCP.Servers[name] + cfg.ExternalMCPEnable = false + h.config.ExternalMCP.Servers[name] = cfg + + // 保存到配置文件 + 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("外部MCP已停止", zap.String("name", name)) + c.JSON(http.StatusOK, gin.H{"message": "外部MCP已停止"}) +} + +// GetExternalMCPStats 获取统计信息 +func (h *ExternalMCPHandler) GetExternalMCPStats(c *gin.Context) { + stats := h.manager.GetStats() + c.JSON(http.StatusOK, stats) +} + +// validateConfig 验证配置（同时支持官方 type 字段和旧版 transport 字段） +func (h *ExternalMCPHandler) validateConfig(cfg config.ExternalMCPServerConfig) error { + transport := cfg.GetTransportType() + if transport == "" { + return fmt.Errorf("需要指定 command（stdio模式）或 url + type（http/sse模式）") + } + + switch transport { + case "http": + if cfg.URL == "" { + return fmt.Errorf("HTTP模式需要 url") + } + case "stdio": + if cfg.Command == "" { + return fmt.Errorf("stdio模式需要 command") + } + case "sse": + if cfg.URL == "" { + return fmt.Errorf("SSE模式需要 url") + } + default: + return fmt.Errorf("不支持的传输模式: %s，支持的模式: http, stdio, sse", transport) + } + + return nil +} + +// isEnabled 检查是否启用 +func (h *ExternalMCPHandler) isEnabled(cfg config.ExternalMCPServerConfig) bool { + return cfg.ExternalMCPEnable +} + +// saveConfig 保存配置到文件 +func (h *ExternalMCPHandler) saveConfig() error { + data, err := os.ReadFile(h.configPath) + if err != nil { + return fmt.Errorf("读取配置文件失败: %w", err) + } + + if err := os.WriteFile(h.configPath+".backup", data, 0644); err != nil { + h.logger.Warn("创建配置备份失败", zap.Error(err)) + } + + root, err := loadYAMLDocument(h.configPath) + if err != nil { + return fmt.Errorf("解析配置文件失败: %w", err) + } + + updateExternalMCPConfig(root, h.config.ExternalMCP) + + if err := writeYAMLDocument(h.configPath, root); err != nil { + return fmt.Errorf("保存配置文件失败: %w", err) + } + + h.logger.Info("配置已保存", zap.String("path", h.configPath)) + return nil +} + +// updateExternalMCPConfig 更新外部MCP配置 +func updateExternalMCPConfig(doc *yaml.Node, cfg config.ExternalMCPConfig) { + root := doc.Content[0] + externalMCPNode := ensureMap(root, "external_mcp") + serversNode := ensureMap(externalMCPNode, "servers") + + // 清空现有服务器配置 + serversNode.Content = nil + + // 添加新的服务器配置 + for name, serverCfg := range cfg.Servers { + nameNode := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: name} + serverNode := &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"} + serversNode.Content = append(serversNode.Content, nameNode, serverNode) + + // type（官方 MCP 传输类型） + effectiveType := serverCfg.GetTransportType() + if effectiveType != "" && effectiveType != "stdio" { + // stdio 可省略（有 command 时自动推断） + setStringInMap(serverNode, "type", effectiveType) + } + if serverCfg.Command != "" { + setStringInMap(serverNode, "command", serverCfg.Command) + } + if len(serverCfg.Args) > 0 { + setStringArrayInMap(serverNode, "args", serverCfg.Args) + } + if serverCfg.Env != nil && len(serverCfg.Env) > 0 { + envNode := ensureMap(serverNode, "env") + for envKey, envValue := range serverCfg.Env { + setStringInMap(envNode, envKey, envValue) + } + } + if serverCfg.URL != "" { + setStringInMap(serverNode, "url", serverCfg.URL) + } + if serverCfg.Headers != nil && len(serverCfg.Headers) > 0 { + headersNode := ensureMap(serverNode, "headers") + for k, v := range serverCfg.Headers { + setStringInMap(headersNode, k, v) + } + } + if serverCfg.Description != "" { + setStringInMap(serverNode, "description", serverCfg.Description) + } + if serverCfg.Timeout > 0 { + setIntInMap(serverNode, "timeout", serverCfg.Timeout) + } + // 官方标准字段 + if serverCfg.Disabled { + setBoolInMap(serverNode, "disabled", true) + } + if len(serverCfg.AutoApprove) > 0 { + setStringArrayInMap(serverNode, "autoApprove", serverCfg.AutoApprove) + } + + // SDK 高级配置 + if serverCfg.MaxRetries > 0 { + setIntInMap(serverNode, "max_retries", serverCfg.MaxRetries) + } + if serverCfg.TerminateDuration > 0 { + setIntInMap(serverNode, "terminate_duration", serverCfg.TerminateDuration) + } + if serverCfg.KeepAlive > 0 { + setIntInMap(serverNode, "keep_alive", serverCfg.KeepAlive) + } + + setBoolInMap(serverNode, "external_mcp_enable", serverCfg.ExternalMCPEnable) + if serverCfg.ToolEnabled != nil && len(serverCfg.ToolEnabled) > 0 { + toolEnabledNode := ensureMap(serverNode, "tool_enabled") + for toolName, enabled := range serverCfg.ToolEnabled { + setBoolInMap(toolEnabledNode, toolName, enabled) + } + } + } +} + +// setStringArrayInMap 设置字符串数组 +func setStringArrayInMap(mapNode *yaml.Node, key string, values []string) { + _, valueNode := ensureKeyValue(mapNode, key) + valueNode.Kind = yaml.SequenceNode + valueNode.Tag = "!!seq" + valueNode.Content = nil + for _, v := range values { + itemNode := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: v} + valueNode.Content = append(valueNode.Content, itemNode) + } +} + +// AddOrUpdateExternalMCPRequest 添加或更新外部MCP请求 +type AddOrUpdateExternalMCPRequest struct { + Config config.ExternalMCPServerConfig `json:"config"` +} + +// ExternalMCPResponse 外部MCP响应 +type ExternalMCPResponse struct { + Config config.ExternalMCPServerConfig `json:"config"` + Status string `json:"status"` // "connected", "disconnected", "disabled", "error", "connecting" + ToolCount int `json:"tool_count"` // 工具数量 + Error string `json:"error,omitempty"` // 错误信息（仅在status为error时存在） +} diff --git a/internal/handler/external_mcp_test.go b/internal/handler/external_mcp_test.go new file mode 100644 index 00000000..e4cf3c1f --- /dev/null +++ b/internal/handler/external_mcp_test.go @@ -0,0 +1,518 @@ +package handler + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/mcp" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +func setupTestRouter() (*gin.Engine, *ExternalMCPHandler, string) { + gin.SetMode(gin.TestMode) + router := gin.New() + + // 创建临时配置文件 + tmpFile, err := os.CreateTemp("", "test-config-*.yaml") + if err != nil { + panic(err) + } + tmpFile.WriteString("server:\n host: 0.0.0.0\n port: 8080\n") + tmpFile.Close() + configPath := tmpFile.Name() + + logger := zap.NewNop() + manager := mcp.NewExternalMCPManager(logger) + cfg := &config.Config{ + ExternalMCP: config.ExternalMCPConfig{ + Servers: make(map[string]config.ExternalMCPServerConfig), + }, + } + + handler := NewExternalMCPHandler(manager, cfg, configPath, logger) + + api := router.Group("/api") + api.GET("/external-mcp", handler.GetExternalMCPs) + api.GET("/external-mcp/stats", handler.GetExternalMCPStats) + api.GET("/external-mcp/:name", handler.GetExternalMCP) + api.PUT("/external-mcp/:name", handler.AddOrUpdateExternalMCP) + api.DELETE("/external-mcp/:name", handler.DeleteExternalMCP) + api.POST("/external-mcp/:name/start", handler.StartExternalMCP) + api.POST("/external-mcp/:name/stop", handler.StopExternalMCP) + + return router, handler, configPath +} + +func cleanupTestConfig(configPath string) { + os.Remove(configPath) + os.Remove(configPath + ".backup") +} + +func TestExternalMCPHandler_AddOrUpdateExternalMCP_Stdio(t *testing.T) { + router, _, configPath := setupTestRouter() + defer cleanupTestConfig(configPath) + + // 测试添加stdio模式的配置（官方格式：有 command 时 type 可省略） + configJSON := `{ + "command": "python3", + "args": ["/path/to/script.py", "--server", "http://example.com"], + "description": "Test stdio MCP", + "timeout": 300, + "external_mcp_enable": true + }` + + var configObj config.ExternalMCPServerConfig + if err := json.Unmarshal([]byte(configJSON), &configObj); err != nil { + t.Fatalf("解析配置JSON失败: %v", err) + } + + reqBody := AddOrUpdateExternalMCPRequest{ + Config: configObj, + } + + body, _ := json.Marshal(reqBody) + req := httptest.NewRequest("PUT", "/api/external-mcp/test-stdio", bytes.NewBuffer(body)) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + router.ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("期望状态码200，实际%d: %s", w.Code, w.Body.String()) + } + + // 验证配置已添加 + req2 := httptest.NewRequest("GET", "/api/external-mcp/test-stdio", nil) + w2 := httptest.NewRecorder() + router.ServeHTTP(w2, req2) + + if w2.Code != http.StatusOK { + t.Fatalf("期望状态码200，实际%d: %s", w2.Code, w2.Body.String()) + } + + var response ExternalMCPResponse + if err := json.Unmarshal(w2.Body.Bytes(), &response); err != nil { + t.Fatalf("解析响应失败: %v", err) + } + + if response.Config.Command != "python3" { + t.Errorf("期望command为python3，实际%s", response.Config.Command) + } + if len(response.Config.Args) != 3 { + t.Errorf("期望args长度为3，实际%d", len(response.Config.Args)) + } + if response.Config.Description != "Test stdio MCP" { + t.Errorf("期望description为'Test stdio MCP'，实际%s", response.Config.Description) + } + if response.Config.Timeout != 300 { + t.Errorf("期望timeout为300，实际%d", response.Config.Timeout) + } +} + +func TestExternalMCPHandler_AddOrUpdateExternalMCP_HTTP(t *testing.T) { + router, _, configPath := setupTestRouter() + defer cleanupTestConfig(configPath) + + // 测试添加HTTP模式的配置（使用官方 type 字段） + configJSON := `{ + "type": "http", + "url": "http://127.0.0.1:8081/mcp", + "external_mcp_enable": true + }` + + var configObj config.ExternalMCPServerConfig + if err := json.Unmarshal([]byte(configJSON), &configObj); err != nil { + t.Fatalf("解析配置JSON失败: %v", err) + } + + reqBody := AddOrUpdateExternalMCPRequest{ + Config: configObj, + } + + body, _ := json.Marshal(reqBody) + req := httptest.NewRequest("PUT", "/api/external-mcp/test-http", bytes.NewBuffer(body)) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + router.ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("期望状态码200，实际%d: %s", w.Code, w.Body.String()) + } + + // 验证配置已添加 + req2 := httptest.NewRequest("GET", "/api/external-mcp/test-http", nil) + w2 := httptest.NewRecorder() + router.ServeHTTP(w2, req2) + + if w2.Code != http.StatusOK { + t.Fatalf("期望状态码200，实际%d: %s", w2.Code, w2.Body.String()) + } + + var response ExternalMCPResponse + if err := json.Unmarshal(w2.Body.Bytes(), &response); err != nil { + t.Fatalf("解析响应失败: %v", err) + } + + if response.Config.Type != "http" { + t.Errorf("期望type为http，实际%s", response.Config.Type) + } + if response.Config.URL != "http://127.0.0.1:8081/mcp" { + t.Errorf("期望url为'http://127.0.0.1:8081/mcp'，实际%s", response.Config.URL) + } +} + +func TestExternalMCPHandler_AddOrUpdateExternalMCP_InvalidConfig(t *testing.T) { + router, _, configPath := setupTestRouter() + defer cleanupTestConfig(configPath) + + testCases := []struct { + name string + configJSON string + expectedErr string + }{ + { + name: "缺少command和url", + configJSON: `{"external_mcp_enable": true}`, + expectedErr: "需要指定 command（stdio模式）或 url + type（http/sse模式）", + }, + { + name: "stdio模式缺少command", + configJSON: `{"args": ["test"], "external_mcp_enable": true}`, + expectedErr: "stdio模式需要command", + }, + { + name: "http模式缺少url", + configJSON: `{"type": "http", "external_mcp_enable": true}`, + expectedErr: "HTTP模式需要 url", + }, + { + name: "无效的type", + configJSON: `{"type": "invalid", "external_mcp_enable": true}`, + expectedErr: "不支持的传输模式", + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + var configObj config.ExternalMCPServerConfig + if err := json.Unmarshal([]byte(tc.configJSON), &configObj); err != nil { + t.Fatalf("解析配置JSON失败: %v", err) + } + + reqBody := AddOrUpdateExternalMCPRequest{ + Config: configObj, + } + + body, _ := json.Marshal(reqBody) + req := httptest.NewRequest("PUT", "/api/external-mcp/test-invalid", bytes.NewBuffer(body)) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + router.ServeHTTP(w, req) + + if w.Code != http.StatusBadRequest { + t.Errorf("期望状态码400，实际%d: %s", w.Code, w.Body.String()) + } + + var response map[string]interface{} + if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil { + t.Fatalf("解析响应失败: %v", err) + } + + errorMsg := response["error"].(string) + // 对于stdio模式缺少command的情况，错误信息可能略有不同 + if tc.name == "stdio模式缺少command" { + if !strings.Contains(errorMsg, "stdio") && !strings.Contains(errorMsg, "command") { + t.Errorf("期望错误信息包含'stdio'或'command'，实际'%s'", errorMsg) + } + } else if !strings.Contains(errorMsg, tc.expectedErr) { + t.Errorf("期望错误信息包含'%s'，实际'%s'", tc.expectedErr, errorMsg) + } + }) + } +} + +func TestExternalMCPHandler_DeleteExternalMCP(t *testing.T) { + router, handler, configPath := setupTestRouter() + defer cleanupTestConfig(configPath) + + // 先添加一个配置 + configObj := config.ExternalMCPServerConfig{ + Command: "python3", + ExternalMCPEnable: true, + } + handler.manager.AddOrUpdateConfig("test-delete", configObj) + + // 删除配置 + req := httptest.NewRequest("DELETE", "/api/external-mcp/test-delete", nil) + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("期望状态码200，实际%d: %s", w.Code, w.Body.String()) + } + + // 验证配置已删除 + req2 := httptest.NewRequest("GET", "/api/external-mcp/test-delete", nil) + w2 := httptest.NewRecorder() + router.ServeHTTP(w2, req2) + + if w2.Code != http.StatusNotFound { + t.Errorf("期望状态码404，实际%d: %s", w2.Code, w2.Body.String()) + } +} + +func TestExternalMCPStatusError(t *testing.T) { + manager := mcp.NewExternalMCPManager(zap.NewNop()) + if got := externalMCPStatusError(manager, "x", "connected"); got != "" { + t.Fatalf("connected status should not return error, got %q", got) + } + if got := externalMCPStatusError(manager, "x", "connecting"); got != "" { + t.Fatalf("connecting status should not return error, got %q", got) + } +} + +func TestExternalMCPHandler_GetExternalMCPs(t *testing.T) { + router, handler, _ := setupTestRouter() + + // 添加多个配置 + handler.manager.AddOrUpdateConfig("test1", config.ExternalMCPServerConfig{ + Command: "python3", + ExternalMCPEnable: true, + }) + handler.manager.AddOrUpdateConfig("test2", config.ExternalMCPServerConfig{ + URL: "http://127.0.0.1:8081/mcp", + ExternalMCPEnable: false, + }) + + req := httptest.NewRequest("GET", "/api/external-mcp", nil) + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("期望状态码200，实际%d: %s", w.Code, w.Body.String()) + } + + var response map[string]interface{} + if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil { + t.Fatalf("解析响应失败: %v", err) + } + + servers := response["servers"].(map[string]interface{}) + if len(servers) != 2 { + t.Errorf("期望2个服务器，实际%d", len(servers)) + } + if _, ok := servers["test1"]; !ok { + t.Error("期望包含test1") + } + if _, ok := servers["test2"]; !ok { + t.Error("期望包含test2") + } + + stats := response["stats"].(map[string]interface{}) + if int(stats["total"].(float64)) != 2 { + t.Errorf("期望总数为2，实际%d", int(stats["total"].(float64))) + } +} + +func TestExternalMCPHandler_GetExternalMCPStats(t *testing.T) { + router, handler, _ := setupTestRouter() + + // 添加配置 + handler.manager.AddOrUpdateConfig("enabled1", config.ExternalMCPServerConfig{ + Command: "python3", + ExternalMCPEnable: true, + }) + handler.manager.AddOrUpdateConfig("enabled2", config.ExternalMCPServerConfig{ + URL: "http://127.0.0.1:8081/mcp", + ExternalMCPEnable: true, + }) + handler.manager.AddOrUpdateConfig("disabled1", config.ExternalMCPServerConfig{ + Command: "python3", + }) + + req := httptest.NewRequest("GET", "/api/external-mcp/stats", nil) + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("期望状态码200，实际%d: %s", w.Code, w.Body.String()) + } + + var stats map[string]interface{} + if err := json.Unmarshal(w.Body.Bytes(), &stats); err != nil { + t.Fatalf("解析响应失败: %v", err) + } + + if int(stats["total"].(float64)) != 3 { + t.Errorf("期望总数为3，实际%d", int(stats["total"].(float64))) + } + if int(stats["enabled"].(float64)) != 2 { + t.Errorf("期望启用数为2，实际%d", int(stats["enabled"].(float64))) + } + if int(stats["disabled"].(float64)) != 1 { + t.Errorf("期望停用数为1，实际%d", int(stats["disabled"].(float64))) + } +} + +func TestExternalMCPHandler_StartStopExternalMCP(t *testing.T) { + router, handler, configPath := setupTestRouter() + defer cleanupTestConfig(configPath) + + // 添加一个禁用的配置 + handler.manager.AddOrUpdateConfig("test-start-stop", config.ExternalMCPServerConfig{ + Command: "python3", + }) + + // 测试启动（可能会失败，因为没有真实的服务器） + req := httptest.NewRequest("POST", "/api/external-mcp/test-start-stop/start", nil) + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + + // 启动可能会失败，但应该返回合理的状态码 + if w.Code != http.StatusOK { + // 如果启动失败，应该是400或500 + if w.Code != http.StatusBadRequest && w.Code != http.StatusInternalServerError { + t.Errorf("期望状态码200/400/500，实际%d: %s", w.Code, w.Body.String()) + } + } + + // 测试停止 + req2 := httptest.NewRequest("POST", "/api/external-mcp/test-start-stop/stop", nil) + w2 := httptest.NewRecorder() + router.ServeHTTP(w2, req2) + + if w2.Code != http.StatusOK { + t.Errorf("期望状态码200，实际%d: %s", w2.Code, w2.Body.String()) + } +} + +func TestExternalMCPHandler_GetExternalMCP_NotFound(t *testing.T) { + router, _, _ := setupTestRouter() + + req := httptest.NewRequest("GET", "/api/external-mcp/nonexistent", nil) + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + + if w.Code != http.StatusNotFound { + t.Errorf("期望状态码404，实际%d: %s", w.Code, w.Body.String()) + } +} + +func TestExternalMCPHandler_DeleteExternalMCP_NotFound(t *testing.T) { + router, _, configPath := setupTestRouter() + defer cleanupTestConfig(configPath) + + req := httptest.NewRequest("DELETE", "/api/external-mcp/nonexistent", nil) + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + + // 删除不存在的配置可能返回200（幂等操作）或404，都是合理的 + if w.Code != http.StatusNotFound && w.Code != http.StatusOK { + t.Errorf("期望状态码404或200，实际%d: %s", w.Code, w.Body.String()) + } +} + +func TestExternalMCPHandler_AddOrUpdateExternalMCP_EmptyName(t *testing.T) { + router, _, _ := setupTestRouter() + + configObj := config.ExternalMCPServerConfig{ + Command: "python3", + ExternalMCPEnable: true, + } + + reqBody := AddOrUpdateExternalMCPRequest{ + Config: configObj, + } + + body, _ := json.Marshal(reqBody) + req := httptest.NewRequest("PUT", "/api/external-mcp/", bytes.NewBuffer(body)) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + router.ServeHTTP(w, req) + + // 空名称应该返回404或400 + if w.Code != http.StatusNotFound && w.Code != http.StatusBadRequest { + t.Errorf("期望状态码404或400，实际%d: %s", w.Code, w.Body.String()) + } +} + +func TestExternalMCPHandler_AddOrUpdateExternalMCP_InvalidJSON(t *testing.T) { + router, _, _ := setupTestRouter() + + // 发送无效的JSON + body := []byte(`{"config": invalid json}`) + req := httptest.NewRequest("PUT", "/api/external-mcp/test", bytes.NewBuffer(body)) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + router.ServeHTTP(w, req) + + if w.Code != http.StatusBadRequest { + t.Errorf("期望状态码400，实际%d: %s", w.Code, w.Body.String()) + } +} + +func TestExternalMCPHandler_UpdateExistingConfig(t *testing.T) { + router, handler, configPath := setupTestRouter() + defer cleanupTestConfig(configPath) + + // 先添加配置 + config1 := config.ExternalMCPServerConfig{ + Command: "python3", + ExternalMCPEnable: true, + } + handler.manager.AddOrUpdateConfig("test-update", config1) + + // 更新配置 + config2 := config.ExternalMCPServerConfig{ + URL: "http://127.0.0.1:8081/mcp", + ExternalMCPEnable: true, + } + + reqBody := AddOrUpdateExternalMCPRequest{ + Config: config2, + } + + body, _ := json.Marshal(reqBody) + req := httptest.NewRequest("PUT", "/api/external-mcp/test-update", bytes.NewBuffer(body)) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + router.ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("期望状态码200，实际%d: %s", w.Code, w.Body.String()) + } + + // 验证配置已更新 + req2 := httptest.NewRequest("GET", "/api/external-mcp/test-update", nil) + w2 := httptest.NewRecorder() + router.ServeHTTP(w2, req2) + + if w2.Code != http.StatusOK { + t.Fatalf("期望状态码200，实际%d: %s", w2.Code, w2.Body.String()) + } + + var response ExternalMCPResponse + if err := json.Unmarshal(w2.Body.Bytes(), &response); err != nil { + t.Fatalf("解析响应失败: %v", err) + } + + if response.Config.URL != "http://127.0.0.1:8081/mcp" { + t.Errorf("期望url为'http://127.0.0.1:8081/mcp'，实际%s", response.Config.URL) + } + if response.Config.Command != "" { + t.Errorf("期望command为空，实际%s", response.Config.Command) + } +} diff --git a/internal/handler/finalization_auto_continue.go b/internal/handler/finalization_auto_continue.go new file mode 100644 index 00000000..2aca373c --- /dev/null +++ b/internal/handler/finalization_auto_continue.go @@ -0,0 +1,77 @@ +package handler + +import ( + "context" + "time" + + "cyberstrike-ai/internal/agent" + "cyberstrike-ai/internal/agentfinalizer" + "cyberstrike-ai/internal/multiagent" + + "go.uber.org/zap" +) + +const finalizationAutoContinueMaxAttempts = 2 + +func shouldAutoContinueAfterFinalization(d agentfinalizer.Decision, attempt int) bool { + if d.Finalizable || d.Finalized { + return false + } + if attempt >= finalizationAutoContinueMaxAttempts { + return false + } + return d.CompletionReason == agentfinalizer.ReasonMissingEvidence +} + +func (h *AgentHandler) tryAutoContinueAfterFinalization( + taskCtx context.Context, + conversationID string, + result *multiagent.RunResult, + decision agentfinalizer.Decision, + attempt *int, + curHistory *[]agent.ChatMessage, + curFinalMessage *string, + progressCallback func(eventType, message string, data interface{}), +) bool { + if !shouldAutoContinueAfterFinalization(decision, *attempt) || result == nil || !multiagent.HasEinoResumeTrace(result) { + return false + } + *attempt++ + h.persistEinoAgentTraceForResume(conversationID, result) + if hist, err := h.loadHistoryFromAgentTrace(conversationID); err == nil && len(hist) > 0 { + *curHistory = hist + } else if h.logger != nil { + h.logger.Warn("finalization auto-continue could not restore trace", + zap.String("conversationId", conversationID), + zap.Error(err)) + return false + } + // Agent 无感续跑：不追加新的 user/system 文案，只使用上一段模型可见轨迹继续 Runner。 + *curFinalMessage = "" + if progressCallback != nil { + progressCallback("finalization_auto_continue", "最终回复检查尚未收敛，正在基于已有轨迹继续执行…", map[string]interface{}{ + "conversationId": conversationID, + "source": "finalizer", + "attempt": *attempt, + "maxAttempts": finalizationAutoContinueMaxAttempts, + "status": decision.Status, + "completionReason": decision.CompletionReason, + "missingChecks": decision.MissingChecks, + "pendingExecutionIds": decision.PendingExecutionIDs, + "contextInjection": false, + }) + } + select { + case <-taskCtx.Done(): + return false + case <-time.After(finalizationAutoContinueBackoff(*attempt)): + return true + } +} + +func finalizationAutoContinueBackoff(attempt int) time.Duration { + if attempt <= 1 { + return 500 * time.Millisecond + } + return time.Duration(attempt) * time.Second +} diff --git a/internal/handler/finalization_auto_continue_test.go b/internal/handler/finalization_auto_continue_test.go new file mode 100644 index 00000000..1ce68d4c --- /dev/null +++ b/internal/handler/finalization_auto_continue_test.go @@ -0,0 +1,59 @@ +package handler + +import ( + "testing" + + "cyberstrike-ai/internal/agentfinalizer" +) + +func TestShouldAutoContinueAfterFinalization(t *testing.T) { + missingEvidence := agentfinalizer.Decision{ + Status: agentfinalizer.StatusBlocked, + CompletionReason: agentfinalizer.ReasonMissingEvidence, + } + if !shouldAutoContinueAfterFinalization(missingEvidence, 0) { + t.Fatal("missing execution evidence should trigger auto-continue") + } + if shouldAutoContinueAfterFinalization(missingEvidence, finalizationAutoContinueMaxAttempts) { + t.Fatal("auto-continue should stop at max attempts") + } + + finalized := agentfinalizer.Decision{ + Status: agentfinalizer.StatusCompleted, + CompletionReason: agentfinalizer.ReasonVerified, + Finalizable: true, + Finalized: true, + } + if shouldAutoContinueAfterFinalization(finalized, 0) { + t.Fatal("finalized decision should not auto-continue") + } + + awaitingHITL := agentfinalizer.Decision{ + Status: agentfinalizer.StatusAwaitingHITL, + CompletionReason: agentfinalizer.ReasonAwaitingHITL, + } + if shouldAutoContinueAfterFinalization(awaitingHITL, 0) { + t.Fatal("awaiting HITL should not auto-continue without approval") + } +} + +func TestRequestRequiresExecutionEvidenceUsesExplicitPolicyOnly(t *testing.T) { + if requestRequiresExecutionEvidence(nil) { + t.Fatal("nil request should not require execution evidence") + } + if requestRequiresExecutionEvidence(&ChatRequest{}) { + t.Fatal("missing finalization policy should not require execution evidence") + } + require := true + if !requestRequiresExecutionEvidence(&ChatRequest{ + Finalization: ChatFinalizationRequest{RequireExecutionEvidence: &require}, + }) { + t.Fatal("explicit true policy should require execution evidence") + } + require = false + if requestRequiresExecutionEvidence(&ChatRequest{ + Finalization: ChatFinalizationRequest{RequireExecutionEvidence: &require}, + }) { + t.Fatal("explicit false policy should not require execution evidence") + } +} diff --git a/internal/handler/finalization_helpers.go b/internal/handler/finalization_helpers.go new file mode 100644 index 00000000..f9fdb8ba --- /dev/null +++ b/internal/handler/finalization_helpers.go @@ -0,0 +1,171 @@ +package handler + +import ( + "fmt" + "strings" + "time" + + "cyberstrike-ai/internal/agentfinalizer" + "cyberstrike-ai/internal/multiagent" + + "go.uber.org/zap" +) + +func (h *AgentHandler) finalizeAgentRunForDelivery( + conversationID string, + assistantMessageID string, + agentMode string, + result *multiagent.RunResult, + mcpExecutionIDs []string, + reasoningContent string, +) agentfinalizer.Decision { + return h.finalizeAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, agentMode, result, mcpExecutionIDs, reasoningContent, false) +} + +func (h *AgentHandler) finalizeAgentRunForDeliveryWithPolicy( + conversationID string, + assistantMessageID string, + agentMode string, + result *multiagent.RunResult, + mcpExecutionIDs []string, + reasoningContent string, + requireExecutionEvidence bool, +) agentfinalizer.Decision { + decision := agentfinalizer.FromRunResult(h.db, result, agentfinalizer.Input{ + ConversationID: conversationID, + AssistantMessageID: assistantMessageID, + AgentMode: agentMode, + MCPExecutionIDs: mcpExecutionIDs, + RequireExecutionEvidence: requireExecutionEvidence, + }) + h.persistFinalizationDecision(conversationID, assistantMessageID, agentMode, mcpExecutionIDs, reasoningContent, decision) + return decision +} + +func (h *AgentHandler) decideAgentRunForDeliveryWithPolicy( + conversationID string, + assistantMessageID string, + agentMode string, + result *multiagent.RunResult, + mcpExecutionIDs []string, + requireExecutionEvidence bool, +) agentfinalizer.Decision { + return agentfinalizer.FromRunResult(h.db, result, agentfinalizer.Input{ + ConversationID: conversationID, + AssistantMessageID: assistantMessageID, + AgentMode: agentMode, + MCPExecutionIDs: mcpExecutionIDs, + RequireExecutionEvidence: requireExecutionEvidence, + }) +} + +func (h *AgentHandler) decideAgentRunForDelivery( + conversationID string, + assistantMessageID string, + agentMode string, + result *multiagent.RunResult, + mcpExecutionIDs []string, +) agentfinalizer.Decision { + return agentfinalizer.FromRunResult(h.db, result, agentfinalizer.Input{ + ConversationID: conversationID, + AssistantMessageID: assistantMessageID, + AgentMode: agentMode, + MCPExecutionIDs: mcpExecutionIDs, + RequireExecutionEvidence: false, + }) +} + +func (h *AgentHandler) persistFinalizationDecision( + conversationID string, + assistantMessageID string, + agentMode string, + mcpExecutionIDs []string, + reasoningContent string, + decision agentfinalizer.Decision, +) { + if assistantMessageID == "" || h.db == nil { + return + } + _ = h.db.AddProcessDetail(assistantMessageID, conversationID, "finalization_check", finalizationCheckMessage(decision), decision) + if decision.Finalizable { + if err := h.db.UpdateAssistantMessageFinalize(assistantMessageID, decision.FinalText, mcpExecutionIDs, reasoningContent); err != nil && h.logger != nil { + h.logger.Warn("更新最终助手消息失败", zap.Error(err), zap.String("conversationId", conversationID), zap.String("agentMode", agentMode)) + } + return + } + _, _ = h.db.Exec("UPDATE messages SET content = ?, updated_at = ? WHERE id = ?", finalizationBlockedMessage(decision), time.Now(), assistantMessageID) +} + +func (h *AgentHandler) finalizeCandidateForDelivery( + conversationID string, + assistantMessageID string, + agentMode string, + response string, + mcpExecutionIDs []string, + awaitingHITL bool, + reasoningContent string, +) agentfinalizer.Decision { + return h.finalizeCandidateForDeliveryWithPolicy(conversationID, assistantMessageID, agentMode, response, mcpExecutionIDs, awaitingHITL, reasoningContent, false) +} + +func (h *AgentHandler) finalizeCandidateForDeliveryWithPolicy( + conversationID string, + assistantMessageID string, + agentMode string, + response string, + mcpExecutionIDs []string, + awaitingHITL bool, + reasoningContent string, + requireExecutionEvidence bool, +) agentfinalizer.Decision { + decision := agentfinalizer.Decide(h.db, agentfinalizer.Input{ + Response: response, + ConversationID: conversationID, + AssistantMessageID: assistantMessageID, + AgentMode: agentMode, + MCPExecutionIDs: mcpExecutionIDs, + AwaitingHITL: awaitingHITL, + RequireExecutionEvidence: requireExecutionEvidence, + }) + if assistantMessageID == "" || h.db == nil { + return decision + } + _ = h.db.AddProcessDetail(assistantMessageID, conversationID, "finalization_check", finalizationCheckMessage(decision), decision) + if decision.Finalizable { + if err := h.db.UpdateAssistantMessageFinalize(assistantMessageID, decision.FinalText, mcpExecutionIDs, reasoningContent); err != nil && h.logger != nil { + h.logger.Warn("更新最终助手消息失败", zap.Error(err), zap.String("conversationId", conversationID), zap.String("agentMode", agentMode)) + } + return decision + } + _, _ = h.db.Exec("UPDATE messages SET content = ?, updated_at = ? WHERE id = ?", finalizationBlockedMessage(decision), time.Now(), assistantMessageID) + return decision +} + +func finalizationCheckMessage(d agentfinalizer.Decision) string { + if d.Finalizable { + return "最终回复检查通过。" + } + return finalizationBlockedMessage(d) +} + +func finalizationBlockedMessage(d agentfinalizer.Decision) string { + parts := []string{"任务尚未达到最终回复条件，暂不生成成功结论。"} + if d.CompletionReason != "" { + parts = append(parts, "原因: "+d.CompletionReason) + } + if len(d.PendingExecutionIDs) > 0 { + parts = append(parts, fmt.Sprintf("仍有 %d 个工具执行未结束: %s", len(d.PendingExecutionIDs), strings.Join(d.PendingExecutionIDs, ", "))) + } + if len(d.MissingChecks) > 0 { + parts = append(parts, "缺失检查: "+strings.Join(d.MissingChecks, "; ")) + } + return strings.Join(parts, "\n") +} + +func finalizationResponsePayload(d agentfinalizer.Decision, extra map[string]interface{}) map[string]interface{} { + return agentfinalizer.ResponsePayload(d, extra) +} + +func requestRequiresExecutionEvidence(req *ChatRequest) bool { + return req != nil && req.Finalization.RequireExecutionEvidence != nil && *req.Finalization.RequireExecutionEvidence +} diff --git a/internal/handler/fofa.go b/internal/handler/fofa.go new file mode 100644 index 00000000..4bfa6156 --- /dev/null +++ b/internal/handler/fofa.go @@ -0,0 +1,1046 @@ +package handler + +import ( + "context" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "net" + "net/http" + "net/url" + "os" + "strings" + "time" + + "cyberstrike-ai/internal/config" + openaiClient "cyberstrike-ai/internal/openai" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +type FofaHandler struct { + cfg *config.Config + logger *zap.Logger + client *http.Client + openAIClient *openaiClient.Client +} + +func NewFofaHandler(cfg *config.Config, logger *zap.Logger) *FofaHandler { + // LLM 请求通常比 FOFA 查询更慢一点，单独给一个更宽松的超时。 + llmHTTPClient := &http.Client{Timeout: 2 * time.Minute} + var llmCfg *config.OpenAIConfig + if cfg != nil { + llmCfg = &cfg.OpenAI + } + return &FofaHandler{ + cfg: cfg, + logger: logger, + client: &http.Client{Timeout: 60 * time.Second}, + openAIClient: openaiClient.NewClient(llmCfg, llmHTTPClient, logger), + } +} + +type fofaSearchRequest struct { + Provider string `json:"provider,omitempty"` + Query string `json:"query" binding:"required"` + Size int `json:"size,omitempty"` + Page int `json:"page,omitempty"` + Fields string `json:"fields,omitempty"` + Full bool `json:"full,omitempty"` +} + +type fofaParseRequest struct { + Provider string `json:"provider,omitempty"` + Text string `json:"text" binding:"required"` +} + +type fofaParseResponse struct { + Query string `json:"query"` + Explanation string `json:"explanation,omitempty"` + Warnings []string `json:"warnings,omitempty"` +} + +type fofaAPIResponse struct { + Error bool `json:"error"` + ErrMsg string `json:"errmsg"` + Size int `json:"size"` + Page int `json:"page"` + Total int `json:"total"` + Mode string `json:"mode"` + Query string `json:"query"` + Results [][]interface{} `json:"results"` +} + +type fofaSearchResponse struct { + Provider string `json:"provider,omitempty"` + Query string `json:"query"` + Size int `json:"size"` + Page int `json:"page"` + Total int `json:"total"` + Fields []string `json:"fields"` + ResultsCount int `json:"results_count"` + ExpectedCount int `json:"expected_count,omitempty"` + Shortfall int `json:"shortfall,omitempty"` + Warning string `json:"warning,omitempty"` + Results []map[string]interface{} `json:"results"` +} + +func normalizeSpaceSearchProvider(provider string) string { + switch strings.ToLower(strings.TrimSpace(provider)) { + case "", "fofa": + return "fofa" + case "zoomeye", "zoom-eye": + return "zoomeye" + case "quake": + return "quake" + case "shodan": + return "shodan" + default: + return "" + } +} + +func providerDisplayName(provider string) string { + switch normalizeSpaceSearchProvider(provider) { + case "zoomeye": + return "ZoomEye" + case "quake": + return "Quake" + case "shodan": + return "Shodan" + default: + return "FOFA" + } +} + +func defaultFieldsForProvider(provider string) string { + switch normalizeSpaceSearchProvider(provider) { + case "zoomeye": + return "ip,port,domain,hostname,title,service,app,country,city" + case "quake": + return "ip,port,domain,service.name,service.http.title,location.country_cn,location.province_cn,location.city_cn" + case "shodan": + return "ip_str,port,hostnames,domains,org,isp,location.country_name,location.city,product,transport" + default: + return "host,ip,port,domain,title,protocol,country,province,city,server" + } +} + +func (h *FofaHandler) resolveAPIKey(provider string) string { + // 优先环境变量（便于容器部署），其次配置文件。 + provider = normalizeSpaceSearchProvider(provider) + envKey := map[string]string{ + "fofa": "FOFA_API_KEY", + "zoomeye": "ZOOMEYE_API_KEY", + "quake": "QUAKE_API_KEY", + "shodan": "SHODAN_API_KEY", + }[provider] + if apiKey := strings.TrimSpace(os.Getenv(envKey)); apiKey != "" { + return apiKey + } + if h.cfg != nil { + switch provider { + case "zoomeye": + return strings.TrimSpace(h.cfg.ZoomEye.APIKey) + case "quake": + return strings.TrimSpace(h.cfg.Quake.APIKey) + case "shodan": + return strings.TrimSpace(h.cfg.Shodan.APIKey) + default: + return strings.TrimSpace(h.cfg.FOFA.APIKey) + } + } + return "" +} + +func (h *FofaHandler) resolveBaseURL(provider string) string { + if h.cfg != nil { + switch normalizeSpaceSearchProvider(provider) { + case "zoomeye": + if v := strings.TrimSpace(h.cfg.ZoomEye.BaseURL); v != "" { + return v + } + case "quake": + if v := strings.TrimSpace(h.cfg.Quake.BaseURL); v != "" { + return v + } + case "shodan": + if v := strings.TrimSpace(h.cfg.Shodan.BaseURL); v != "" { + return v + } + default: + if v := strings.TrimSpace(h.cfg.FOFA.BaseURL); v != "" { + return v + } + } + } + switch normalizeSpaceSearchProvider(provider) { + case "zoomeye": + return "https://api.zoomeye.org/v2/search" + case "quake": + return "https://quake.360.cn/api/v3/search/quake_service" + case "shodan": + return "https://api.shodan.io" + default: + return "https://fofa.info/api/v1/search/all" + } +} + +// ParseNaturalLanguage 将自然语言解析为 FOFA 查询语法（仅生成，不执行查询） +func (h *FofaHandler) ParseNaturalLanguage(c *gin.Context) { + var req fofaParseRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()}) + return + } + req.Text = strings.TrimSpace(req.Text) + provider := normalizeSpaceSearchProvider(req.Provider) + if provider == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "provider 不支持，可选：fofa、zoomeye、quake、shodan"}) + return + } + if req.Text == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "text 不能为空"}) + return + } + + if h.cfg == nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "系统配置未初始化"}) + return + } + if strings.TrimSpace(h.cfg.OpenAI.APIKey) == "" || strings.TrimSpace(h.cfg.OpenAI.Model) == "" { + c.JSON(http.StatusBadRequest, gin.H{ + "error": "未配置 AI 模型：请在系统设置中填写 openai.api_key 与 openai.model（支持 OpenAI 兼容 API，如 DeepSeek）", + "need": []string{"openai.api_key", "openai.model"}, + }) + return + } + if h.openAIClient == nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "AI 客户端未初始化"}) + return + } + + engineName := providerDisplayName(provider) + syntaxNotes := map[string]string{ + "fofa": ` +FOFA 官方查询语法参考： +- 基本格式：field="value"，字符串值使用英文双引号；多个条件用 &&（与）、||（或）、!（非）连接。 +- 组合优先级：复杂表达式必须使用 () 明确优先级，例如：(app="Apache" || app="nginx") && country="CN"。 +- 常用字段：app、title、body、header、host、domain、ip、port、protocol、country、province、city、server、icp、cert、icon_hash、fid。 +- 字段示例： + - app="Apache" + - title="后台管理" + - body="Powered by" + - header="JSESSIONID" + - domain="example.com" + - host="https://example.com" + - ip="1.1.1.1" + - port="443" + - country="CN" + - city="Hangzhou" + - cert="example.com" + - icon_hash="-247388890" +- 组合示例： + - app="Apache" && country="CN" + - title="login" || title="登录" + - (app="Apache" || app="nginx") && port="443" + - domain="example.com" && !title="404" + - cert="example.com" && port="443" + - header="JSESSIONID" && country="CN" +- 生成注意： + - 用户说“排除/不要/非”时优先使用 !field="value"。 + - 用户说“标题包含/页面标题”映射为 title；说“正文包含/页面包含”映射为 body；说“响应头/cookie/header”映射为 header。 + - 端口在 FOFA 中通常写成 port="443"。 +`, + "zoomeye": ` +ZoomEye 查询语法参考： +- 基本格式：field="value" 或 field=value；字符串/短语建议使用英文双引号。 +- 逻辑连接：可使用 && / || / !，也可使用 AND / OR / NOT；复杂表达式使用 () 明确优先级。 +- 常用字段：app、service、title、domain、hostname、ip、port、country、city、org、isp、asn、cidr、ssl、ssl.cert.fingerprint、iconhash。 +- 字段示例： + - app="Apache" + - service="ssh" + - title="登录" + - domain="example.com" + - hostname="example.com" + - ip="1.1.1.1" + - cidr="1.1.1.0/24" + - port=443 + - country="CN" + - city="Beijing" + - org="Tencent" + - ssl="example.com" + - ssl.cert.fingerprint="F3C98F223D82CC41CF83D94671CCC6C69873FABF" + - iconhash="-247388890" +- 组合示例： + - app="nginx" && country="CN" + - service="http" && (title="login" || title="登录") + - domain="example.com" && !app="cloudflare" + - port=443 && country="US" + - app="Elasticsearch" && port=9200 +- 生成注意： + - 用户说“服务/协议是 SSH、HTTP、RDP”优先映射为 service。 + - 用户说“站点/网站标题”映射为 title；说“域名/主域”优先映射为 domain 或 hostname。 + - 端口可以不加引号，例如 port=443；如果用户原文已给出冒号风格表达式且接近 ZoomEye 语法，可原样保留。 +`, + "quake": ` +Quake 查询语法参考： +- 基本格式：field:"value" 或 field:value；字符串/中文/短语使用英文双引号。 +- 逻辑连接：使用 AND、OR、NOT；复杂表达式必须使用 () 明确优先级。 +- 常用字段：domain、ip、port、service.name、service.http.title、service.http.server、service.http.response.header、service.http.favicon.hash、country_cn、province_cn、city_cn、location.country_cn、location.province_cn、location.city_cn、asn、org。 +- 字段示例： + - domain:"example.com" + - ip:"1.1.1.1" + - port:443 + - service.name:"http" + - service.name:"ssh" + - service.http.title:"登录" + - service.http.server:"nginx" + - service.http.response.header:"JSESSIONID" + - service.http.favicon.hash:"-247388890" + - country_cn:"中国" + - province_cn:"浙江" + - city_cn:"杭州" +- 组合示例： + - service.name:"http" AND country_cn:"中国" + - (service.name:"http" OR service.name:"https") AND port:443 + - domain:"example.com" AND NOT service.http.title:"404" + - service.http.title:"login" AND port:443 + - service.name:"ssh" AND country_cn:"中国" +- 生成注意： + - 用户说“中国/浙江/杭州”等中文地理位置时，Quake 优先使用 country_cn/province_cn/city_cn 并保留中文值。 + - 用户说“标题”映射为 service.http.title；说“Server/服务端软件”映射为 service.http.server；说“favicon/hash/icon”映射为 service.http.favicon.hash。 + - Quake 不使用 && / || 作为首选输出；优先输出 AND / OR / NOT。 +`, + "shodan": ` +Shodan 官方查询语法参考： +- 默认裸关键词只搜索 banner 的 data 内容；精确条件使用 filter:value。 +- filter 与 value 中间不能有空格；值包含空格时用英文双引号，例如 org:"Amazon Web Services"。 +- 多个过滤器并列表示同时满足（AND）；Shodan 查询不要使用 &&、||，除非用户明确给出并要求保留。 +- 常用过滤器：product、port、country、city、org、asn、hostname、net、ssl、ssl.cert.subject.cn、http.title、has_screenshot、vuln。 +- 字段示例： + - product:nginx + - port:443 + - country:CN + - city:Shanghai + - org:"Amazon" + - asn:AS15169 + - hostname:example.com + - ssl.cert.subject.cn:example.com + - http.title:"Dashboard" + - has_screenshot:true + - vuln:CVE-2021-41773 +- 组合示例： + - product:nginx country:CN + - apache country:DE + - org:"Amazon" port:443 + - ssl.cert.subject.cn:example.com port:443 + - http.title:"login" country:CN + - ssl:true port:443 hostname:example.com +- 生成注意： + - 用户说“产品/组件/服务软件”优先映射为 product；说“组织/公司/云厂商”映射为 org；说“证书 CN/SAN/域名证书”优先映射为 ssl.cert.subject.cn。 + - 国家用两位国家代码；如果用户给出中文国家名且无法确定代码，把推断写入 explanation 或 warnings。 + - Shodan 没有通用 NOT 排除语法；遇到“排除/不要”时应在 warnings 说明可能需要人工调整，不要强行编造过滤器。 +`, + }[provider] + + systemPrompt := strings.TrimSpace(fmt.Sprintf(` +你是“%s 查询语法生成器”。任务：把用户输入的自然语言搜索意图，转换成 %s 查询语法。 + +输出要求（非常重要）： +1) 只输出 JSON（不要 markdown、不要代码块、不要额外解释文本） +2) JSON 结构必须是： +{ + "query": "string，%s 查询语法（可直接粘贴到 %s 或本系统查询框）", + "explanation": "string，可选，解释你如何映射字段/逻辑", + "warnings": ["string"...] 可选，列出歧义/风险/需要人工确认的点 +} +3) 如果用户输入本身已经是 %s 查询语法（或非常接近该语法的表达式），应当“原样返回”为 query： + - 不要擅自改写字段名、操作符、括号结构 + - 不要改写任何字符串值（尤其是地理位置类值），不要做缩写/同义词替换/翻译/音译 + +当前搜索引擎语法速查： +%s + +通用生成约束： +- 严格遵守“当前搜索引擎语法速查”里的字段名、操作符和示例风格；不同数据源语法不同，不要混用。 +- 字符串值保持用户原意：不要无依据缩写、翻译、音译、替换同义词或改写大小写。 +- 地理位置、组织名、产品名、域名、证书名、CVE 编号等实体值必须尽量保留原文；确需推断（如“中国”到 CN）时在 explanation 或 warnings 中说明。 +- 不要捏造字段。不确定字段是否支持时，选择更通用且确定的字段，或把不确定点写进 warnings。 +- 当用户描述里有多个与/或条件，必须使用该数据源支持的括号和逻辑操作符明确优先级。 +- 如果用户输入已经是当前数据源查询语法或非常接近，应原样返回；只在明显有语法错误且能确定修复方式时轻微修正，并在 explanation 说明。 +- 如果需求范围过大、关键目标缺失或语义矛盾，允许 query 为空字符串，并在 warnings 中明确需要补充的信息。 +- 只生成资产测绘/信息收集查询语法，不生成扫描、利用、爆破、绕过、命令执行或攻击步骤。 +`, engineName, engineName, engineName, engineName, engineName, syntaxNotes)) + + userPrompt := fmt.Sprintf("自然语言意图：%s", req.Text) + + requestBody := map[string]interface{}{ + "model": h.cfg.OpenAI.Model, + "messages": []map[string]interface{}{ + {"role": "system", "content": systemPrompt}, + {"role": "user", "content": userPrompt}, + }, + "temperature": 0.1, + "max_completion_tokens": 12000, + } + + // OpenAI 返回结构：只需要 choices[0].message.content + var apiResponse struct { + Choices []struct { + Message struct { + Content string `json:"content"` + } `json:"message"` + } `json:"choices"` + } + + ctx, cancel := context.WithTimeout(c.Request.Context(), 90*time.Second) + defer cancel() + + if err := h.openAIClient.ChatCompletion(ctx, requestBody, &apiResponse); err != nil { + var apiErr *openaiClient.APIError + if errors.As(err, &apiErr) { + h.logger.Warn("FOFA自然语言解析：LLM返回错误", zap.Int("status", apiErr.StatusCode)) + c.JSON(http.StatusBadGateway, gin.H{"error": "AI 解析失败（上游返回非 200），请检查模型配置或稍后重试"}) + return + } + c.JSON(http.StatusBadGateway, gin.H{"error": "AI 解析失败: " + err.Error()}) + return + } + if len(apiResponse.Choices) == 0 { + c.JSON(http.StatusBadGateway, gin.H{"error": "AI 未返回有效结果"}) + return + } + + content := strings.TrimSpace(apiResponse.Choices[0].Message.Content) + jsonContent, extractErr := extractInfoCollectJSONObject(content) + if extractErr != nil { + snippet := trimSnippet(content, 1200) + c.JSON(http.StatusBadGateway, gin.H{ + "error": "AI 返回内容无法解析为 JSON，请稍后重试或换个描述方式", + "snippet": snippet, + }) + return + } + + var parsed fofaParseResponse + if err := json.Unmarshal([]byte(jsonContent), &parsed); err != nil { + // 直接回传一部分原文，方便排查，但避免太大 + snippet := trimSnippet(content, 1200) + c.JSON(http.StatusBadGateway, gin.H{ + "error": "AI 返回内容无法解析为 JSON，请稍后重试或换个描述方式", + "snippet": snippet, + }) + return + } + parsed.Query = strings.TrimSpace(parsed.Query) + if parsed.Query == "" { + // query 允许为空（表示需求不明确），但前端需要明确提示 + if len(parsed.Warnings) == 0 { + parsed.Warnings = []string{"需求信息不足，未能生成可用的 " + engineName + " 查询语法，请补充关键条件（如国家/端口/产品/域名等）。"} + } + } + + c.JSON(http.StatusOK, parsed) +} + +func extractInfoCollectJSONObject(content string) (string, error) { + content = strings.TrimSpace(content) + if content == "" { + return "", errors.New("empty content") + } + candidates := []string{content} + if fenced := extractFencedJSON(content); fenced != "" { + candidates = append([]string{fenced}, candidates...) + } + for _, candidate := range candidates { + candidate = strings.TrimSpace(candidate) + if candidate == "" { + continue + } + if json.Valid([]byte(candidate)) { + return candidate, nil + } + if obj := scanBalancedJSONObject(candidate); obj != "" && json.Valid([]byte(obj)) { + return obj, nil + } + } + return "", errors.New("json object not found") +} + +func extractFencedJSON(content string) string { + start := strings.Index(content, "```") + if start < 0 { + return "" + } + rest := content[start+3:] + if nl := strings.Index(rest, "\n"); nl >= 0 { + lang := strings.ToLower(strings.TrimSpace(rest[:nl])) + if lang == "" || lang == "json" || strings.HasPrefix(lang, "json ") { + rest = rest[nl+1:] + } + } + end := strings.Index(rest, "```") + if end < 0 { + return "" + } + return strings.TrimSpace(rest[:end]) +} + +func scanBalancedJSONObject(content string) string { + start := strings.Index(content, "{") + if start < 0 { + return "" + } + depth := 0 + inString := false + escaped := false + for i := start; i < len(content); i++ { + ch := content[i] + if inString { + if escaped { + escaped = false + continue + } + switch ch { + case '\\': + escaped = true + case '"': + inString = false + } + continue + } + switch ch { + case '"': + inString = true + case '{': + depth++ + case '}': + depth-- + if depth == 0 { + return strings.TrimSpace(content[start : i+1]) + } + } + } + return "" +} + +func trimSnippet(s string, maxRunes int) string { + s = strings.TrimSpace(s) + if maxRunes <= 0 { + return "" + } + runes := []rune(s) + if len(runes) <= maxRunes { + return s + } + return string(runes[:maxRunes]) +} + +// Search FOFA 查询（后端代理，避免前端暴露 key） +func (h *FofaHandler) Search(c *gin.Context) { + var req fofaSearchRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()}) + return + } + provider := normalizeSpaceSearchProvider(req.Provider) + if provider == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "provider 不支持，可选：fofa、zoomeye、quake、shodan"}) + return + } + + req.Query = strings.TrimSpace(req.Query) + if req.Query == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "query 不能为空"}) + return + } + if req.Size <= 0 { + req.Size = 100 + } + if req.Page <= 0 { + req.Page = 1 + } + // FOFA 接口 size 上限和账户权限相关，这里只做一个合理的保护 + if req.Size > 10000 { + req.Size = 10000 + } + if req.Fields == "" { + req.Fields = defaultFieldsForProvider(provider) + } + + if provider != "fofa" { + h.searchExternalProvider(c, provider, req) + return + } + + apiKey := h.resolveAPIKey(provider) + if apiKey == "" { + c.JSON(http.StatusBadRequest, gin.H{ + "error": "FOFA 未配置：请在系统设置的资产管理中填写 FOFA API Key，或设置环境变量 FOFA_API_KEY", + "need": []string{"fofa.api_key"}, + "env_key": []string{"FOFA_API_KEY"}, + }) + return + } + + baseURL := h.resolveBaseURL(provider) + qb64 := base64.StdEncoding.EncodeToString([]byte(req.Query)) + + u, err := url.Parse(baseURL) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "FOFA base_url 无效: " + err.Error()}) + return + } + + params := u.Query() + params.Set("key", apiKey) + params.Set("qbase64", qb64) + params.Set("size", fmt.Sprintf("%d", req.Size)) + params.Set("page", fmt.Sprintf("%d", req.Page)) + params.Set("fields", strings.TrimSpace(req.Fields)) + if req.Full { + params.Set("full", "true") + } else { + // 明确传 false，便于排查 + params.Set("full", "false") + } + u.RawQuery = params.Encode() + + httpReq, err := http.NewRequestWithContext(c.Request.Context(), http.MethodGet, u.String(), nil) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "创建请求失败: " + err.Error()}) + return + } + httpReq.Header.Set("User-Agent", "CyberStrikeAI/1.7.4") + httpReq.Header.Set("Accept", "application/json") + + resp, err := h.client.Do(httpReq) + if err != nil { + status, message, timeout := safeFofaRequestError(err) + h.logger.Warn("请求 FOFA 失败", + zap.String("host", u.Host), + zap.Bool("timeout", timeout), + zap.String("error_type", fmt.Sprintf("%T", err)), + ) + c.JSON(status, gin.H{"error": message}) + return + } + defer resp.Body.Close() + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + c.JSON(http.StatusBadGateway, gin.H{"error": fmt.Sprintf("FOFA 返回非 2xx: %d", resp.StatusCode)}) + return + } + + var apiResp fofaAPIResponse + if err := json.NewDecoder(resp.Body).Decode(&apiResp); err != nil { + c.JSON(http.StatusBadGateway, gin.H{"error": "解析 FOFA 响应失败: " + err.Error()}) + return + } + if apiResp.Error { + msg := strings.TrimSpace(apiResp.ErrMsg) + if msg == "" { + msg = "FOFA 返回错误" + } + c.JSON(http.StatusBadGateway, gin.H{"error": msg}) + return + } + + fields := splitAndCleanCSV(req.Fields) + results := make([]map[string]interface{}, 0, len(apiResp.Results)) + for _, row := range apiResp.Results { + item := make(map[string]interface{}, len(fields)) + for i, f := range fields { + if i < len(row) { + item[f] = row[i] + } else { + item[f] = nil + } + } + results = append(results, item) + } + + c.JSON(http.StatusOK, fofaSearchResponse{ + Provider: provider, + Query: req.Query, + Size: apiResp.Size, + Page: apiResp.Page, + Total: apiResp.Total, + Fields: fields, + ResultsCount: len(results), + Results: results, + }) +} + +func (h *FofaHandler) searchExternalProvider(c *gin.Context, provider string, req fofaSearchRequest) { + apiKey := h.resolveAPIKey(provider) + if apiKey == "" { + envKey := map[string]string{ + "zoomeye": "ZOOMEYE_API_KEY", + "quake": "QUAKE_API_KEY", + "shodan": "SHODAN_API_KEY", + }[provider] + c.JSON(http.StatusBadRequest, gin.H{ + "error": providerDisplayName(provider) + " 未配置：请在 config.yaml 中填写 api_key，或设置环境变量 " + envKey, + "need": []string{provider + ".api_key"}, + "env_key": []string{envKey}, + }) + return + } + + switch provider { + case "zoomeye": + h.searchZoomEye(c, req, apiKey) + case "quake": + h.searchQuake(c, req, apiKey) + case "shodan": + h.searchShodan(c, req, apiKey) + } +} + +func (h *FofaHandler) searchZoomEye(c *gin.Context, req fofaSearchRequest, apiKey string) { + baseURL := h.resolveBaseURL("zoomeye") + u, err := url.Parse(baseURL) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "ZoomEye base_url 无效: " + err.Error()}) + return + } + body := map[string]interface{}{ + "qbase64": base64.StdEncoding.EncodeToString([]byte(req.Query)), + "page": req.Page, + "pagesize": req.Size, + } + if fields := strings.TrimSpace(req.Fields); fields != "" { + body["fields"] = fields + } + var apiResp struct { + Code int `json:"code"` + Message string `json:"message"` + Query string `json:"query"` + Total int `json:"total"` + Page int `json:"page"` + PageSize int `json:"pagesize"` + Data []map[string]interface{} `json:"data"` + } + if !h.doJSONRequest(c, http.MethodPost, u.String(), apiKey, "API-KEY", body, &apiResp, "ZoomEye") { + return + } + if apiResp.Code != 60000 { + msg := strings.TrimSpace(apiResp.Message) + if msg == "" { + msg = "ZoomEye 返回错误" + } + c.JSON(http.StatusBadGateway, gin.H{"error": msg}) + return + } + fields := splitAndCleanCSV(req.Fields) + c.JSON(http.StatusOK, fofaSearchResponse{ + Provider: "zoomeye", + Query: firstNonEmptySpaceSearchValue(apiResp.Query, req.Query), + Size: firstPositive(apiResp.PageSize, req.Size), + Page: firstPositive(apiResp.Page, req.Page), + Total: apiResp.Total, + Fields: fields, + ResultsCount: len(apiResp.Data), + Results: projectRows(apiResp.Data, fields), + }) +} + +func (h *FofaHandler) searchQuake(c *gin.Context, req fofaSearchRequest, apiKey string) { + baseURL := h.resolveBaseURL("quake") + u, err := url.Parse(baseURL) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "Quake base_url 无效: " + err.Error()}) + return + } + fields := splitAndCleanCSV(req.Fields) + body := map[string]interface{}{ + "query": req.Query, + "size": req.Size, + "start": (req.Page - 1) * req.Size, + "latest": req.Full, + } + if len(fields) > 0 { + body["include"] = fields + } + var apiResp struct { + Code interface{} `json:"code"` + Message string `json:"message"` + TotalCount int `json:"total_count"` + Data []map[string]interface{} `json:"data"` + Meta struct { + Pagination struct { + Total int `json:"total"` + } `json:"pagination"` + } `json:"meta"` + } + if !h.doJSONRequest(c, http.MethodPost, u.String(), apiKey, "X-QuakeToken", body, &apiResp, "Quake") { + return + } + if !isZeroSpaceSearchCode(apiResp.Code) { + msg := strings.TrimSpace(apiResp.Message) + if msg == "" { + msg = "Quake 返回错误" + } + c.JSON(http.StatusBadGateway, gin.H{"error": msg}) + return + } + total := firstPositive(apiResp.TotalCount, apiResp.Meta.Pagination.Total) + c.JSON(http.StatusOK, fofaSearchResponse{ + Provider: "quake", + Query: req.Query, + Size: req.Size, + Page: req.Page, + Total: total, + Fields: fields, + ResultsCount: len(apiResp.Data), + Results: projectRows(apiResp.Data, fields), + }) +} + +func isZeroSpaceSearchCode(code interface{}) bool { + switch v := code.(type) { + case nil: + return false + case int: + return v == 0 + case int64: + return v == 0 + case float64: + return v == 0 + case string: + return strings.TrimSpace(v) == "0" + default: + return false + } +} + +func (h *FofaHandler) searchShodan(c *gin.Context, req fofaSearchRequest, apiKey string) { + baseURL := strings.TrimRight(h.resolveBaseURL("shodan"), "/") + "/shodan/host/search" + u, err := url.Parse(baseURL) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "Shodan base_url 无效: " + err.Error()}) + return + } + + var apiResp struct { + Total int `json:"total"` + Matches []map[string]interface{} `json:"matches"` + Error string `json:"error"` + } + targetSize := req.Size + if targetSize <= 0 { + targetSize = 100 + } + if targetSize > 1000 { + targetSize = 1000 + } + page := req.Page + matches := make([]map[string]interface{}, 0, targetSize) + pagesNeeded := (targetSize + 99) / 100 + for i := 0; i < pagesNeeded; i++ { + pageURL := *u + params := pageURL.Query() + params.Set("key", apiKey) + params.Set("query", req.Query) + params.Set("page", fmt.Sprintf("%d", page+i)) + params.Set("minify", "false") + if fields := strings.TrimSpace(req.Fields); fields != "" { + params.Set("fields", fields) + } + pageURL.RawQuery = params.Encode() + apiResp.Matches = nil + apiResp.Error = "" + if !h.doJSONRequest(c, http.MethodGet, pageURL.String(), "", "", nil, &apiResp, "Shodan") { + return + } + if strings.TrimSpace(apiResp.Error) != "" { + c.JSON(http.StatusBadGateway, gin.H{"error": apiResp.Error}) + return + } + if len(apiResp.Matches) == 0 { + break + } + matches = append(matches, apiResp.Matches...) + if len(matches) >= targetSize { + matches = matches[:targetSize] + break + } + } + fields := splitAndCleanCSV(req.Fields) + expectedCount := shodanExpectedResultCount(apiResp.Total, page, targetSize) + shortfall := expectedCount - len(matches) + warning := "" + if shortfall > 0 { + warning = fmt.Sprintf("Shodan 统计总数为 %d，但本次分页实际只返回 %d/%d 条明细", apiResp.Total, len(matches), expectedCount) + } + c.JSON(http.StatusOK, fofaSearchResponse{ + Provider: "shodan", + Query: req.Query, + Size: targetSize, + Page: page, + Total: apiResp.Total, + Fields: fields, + ResultsCount: len(matches), + ExpectedCount: expectedCount, + Shortfall: max(0, shortfall), + Warning: warning, + Results: projectRows(matches, fields), + }) +} + +func shodanExpectedResultCount(total, page, size int) int { + if total <= 0 || size <= 0 { + return 0 + } + if page <= 0 { + page = 1 + } + startOffset := (page - 1) * 100 + remaining := total - startOffset + if remaining <= 0 { + return 0 + } + if remaining < size { + return remaining + } + return size +} + +func (h *FofaHandler) doJSONRequest(c *gin.Context, method, endpoint, apiKey, headerName string, body interface{}, out interface{}, label string) bool { + var reqBody *strings.Reader + if body != nil { + b, err := json.Marshal(body) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "创建请求失败: " + err.Error()}) + return false + } + reqBody = strings.NewReader(string(b)) + } else { + reqBody = strings.NewReader("") + } + httpReq, err := http.NewRequestWithContext(c.Request.Context(), method, endpoint, reqBody) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "创建请求失败: " + err.Error()}) + return false + } + httpReq.Header.Set("User-Agent", "CyberStrikeAI/1.7.4") + httpReq.Header.Set("Accept", "application/json") + if body != nil { + httpReq.Header.Set("Content-Type", "application/json") + } + if headerName != "" && apiKey != "" { + httpReq.Header.Set(headerName, apiKey) + } + resp, err := h.client.Do(httpReq) + if err != nil { + status, message, timeout := safeFofaRequestError(err) + h.logger.Warn("请求空间测绘搜索失败", + zap.String("provider", label), + zap.Bool("timeout", timeout), + zap.String("error_type", fmt.Sprintf("%T", err)), + ) + c.JSON(status, gin.H{"error": strings.Replace(message, "FOFA", label, 1)}) + return false + } + defer resp.Body.Close() + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + c.JSON(http.StatusBadGateway, gin.H{"error": fmt.Sprintf("%s 返回非 2xx: %d", label, resp.StatusCode)}) + return false + } + if err := json.NewDecoder(resp.Body).Decode(out); err != nil { + c.JSON(http.StatusBadGateway, gin.H{"error": "解析 " + label + " 响应失败: " + err.Error()}) + return false + } + return true +} + +func safeFofaRequestError(err error) (status int, message string, timeout bool) { + var netErr net.Error + timeout = errors.Is(err, context.DeadlineExceeded) || + (errors.As(err, &netErr) && netErr.Timeout()) + if timeout { + return http.StatusGatewayTimeout, + "FOFA 请求超时（60 秒）：请稍后重试，或减少返回数量和返回字段", + true + } + return http.StatusBadGateway, + "无法连接 FOFA 服务，请检查服务器网络或代理配置", + false +} + +func splitAndCleanCSV(s string) []string { + parts := strings.Split(s, ",") + out := make([]string, 0, len(parts)) + seen := make(map[string]struct{}, len(parts)) + for _, p := range parts { + v := strings.TrimSpace(p) + if v == "" { + continue + } + if _, ok := seen[v]; ok { + continue + } + seen[v] = struct{}{} + out = append(out, v) + } + return out +} + +func projectRows(rows []map[string]interface{}, fields []string) []map[string]interface{} { + if len(fields) == 0 { + return rows + } + out := make([]map[string]interface{}, 0, len(rows)) + for _, row := range rows { + item := make(map[string]interface{}, len(fields)) + for _, field := range fields { + item[field] = valueByPath(row, field) + } + out = append(out, item) + } + return out +} + +func valueByPath(row map[string]interface{}, path string) interface{} { + if row == nil { + return nil + } + if v, ok := row[path]; ok { + return v + } + parts := strings.Split(path, ".") + var current interface{} = row + for _, part := range parts { + m, ok := current.(map[string]interface{}) + if !ok { + return nil + } + current, ok = m[part] + if !ok { + return nil + } + } + return current +} + +func firstPositive(values ...int) int { + for _, v := range values { + if v > 0 { + return v + } + } + return 0 +} + +func firstNonEmptySpaceSearchValue(values ...string) string { + for _, v := range values { + if strings.TrimSpace(v) != "" { + return v + } + } + return "" +} diff --git a/internal/handler/fofa_test.go b/internal/handler/fofa_test.go new file mode 100644 index 00000000..4fa629c0 --- /dev/null +++ b/internal/handler/fofa_test.go @@ -0,0 +1,239 @@ +package handler + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" + + "cyberstrike-ai/internal/config" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +func TestFofaSearchUsesAPIKeyWithoutEmail(t *testing.T) { + gin.SetMode(gin.TestMode) + t.Setenv("FOFA_API_KEY", "") + t.Setenv("FOFA_EMAIL", "legacy@example.com") + + var receivedEmail string + var receivedKey string + fofaServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + receivedEmail = r.URL.Query().Get("email") + receivedKey = r.URL.Query().Get("key") + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"error":false,"size":1,"page":1,"results":[["https://example.com"]]}`)) + })) + defer fofaServer.Close() + + h := NewFofaHandler(&config.Config{ + FOFA: config.FofaConfig{ + BaseURL: fofaServer.URL, + APIKey: "test-api-key", + }, + }, zap.NewNop()) + + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + body := `{"query":"domain=\"example.com\"","fields":"host"}` + ctx.Request = httptest.NewRequest(http.MethodPost, "/api/fofa/search", strings.NewReader(body)) + ctx.Request.Header.Set("Content-Type", "application/json") + + h.Search(ctx) + + if recorder.Code != http.StatusOK { + t.Fatalf("Search() status = %d, body = %s", recorder.Code, recorder.Body.String()) + } + if receivedEmail != "" { + t.Fatalf("FOFA request unexpectedly included email = %q", receivedEmail) + } + if receivedKey != "test-api-key" { + t.Fatalf("FOFA request key = %q, want %q", receivedKey, "test-api-key") + } + + var response fofaSearchResponse + if err := json.Unmarshal(recorder.Body.Bytes(), &response); err != nil { + t.Fatalf("decode response: %v", err) + } + if response.ResultsCount != 1 { + t.Fatalf("results_count = %d, want 1", response.ResultsCount) + } +} + +func TestSafeFofaRequestErrorDoesNotExposeURLOrAPIKey(t *testing.T) { + const secretURL = "https://fofa.info/api/v1/search/all?key=secret-api-key" + err := &url.Error{ + Op: http.MethodGet, + URL: secretURL, + Err: context.DeadlineExceeded, + } + + status, message, timeout := safeFofaRequestError(err) + + if status != http.StatusGatewayTimeout { + t.Fatalf("status = %d, want %d", status, http.StatusGatewayTimeout) + } + if !timeout { + t.Fatal("timeout = false, want true") + } + if strings.Contains(message, "secret-api-key") || strings.Contains(message, secretURL) { + t.Fatalf("safe error exposed request URL or API key: %q", message) + } +} + +func TestShodanSearchReportsShortfallWhenTotalExceedsMatches(t *testing.T) { + gin.SetMode(gin.TestMode) + t.Setenv("SHODAN_API_KEY", "") + + shodanServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/shodan/host/search" { + t.Fatalf("unexpected path: %s", r.URL.Path) + } + if got := r.URL.Query().Get("key"); got != "test-shodan-key" { + t.Fatalf("Shodan key = %q, want test-shodan-key", got) + } + page := r.URL.Query().Get("page") + count := 0 + switch page { + case "1": + count = 100 + case "2": + count = 3 + default: + count = 0 + } + matches := make([]map[string]interface{}, 0, count) + for i := 0; i < count; i++ { + matches = append(matches, map[string]interface{}{ + "ip_str": fmt.Sprintf("192.0.2.%d", i+1), + "port": 80, + }) + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "total": 104, + "matches": matches, + }) + })) + defer shodanServer.Close() + + h := NewFofaHandler(&config.Config{ + Shodan: config.SpaceSearchConfig{ + BaseURL: shodanServer.URL, + APIKey: "test-shodan-key", + }, + }, zap.NewNop()) + + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + body := `{"provider":"shodan","query":"product:nginx","fields":"ip_str,port","size":1000,"page":1}` + ctx.Request = httptest.NewRequest(http.MethodPost, "/api/fofa/search", strings.NewReader(body)) + ctx.Request.Header.Set("Content-Type", "application/json") + + h.Search(ctx) + + if recorder.Code != http.StatusOK { + t.Fatalf("Search() status = %d, body = %s", recorder.Code, recorder.Body.String()) + } + var response fofaSearchResponse + if err := json.Unmarshal(recorder.Body.Bytes(), &response); err != nil { + t.Fatalf("decode response: %v", err) + } + if response.Total != 104 || response.ResultsCount != 103 { + t.Fatalf("counts: total=%d results_count=%d, want 104/103", response.Total, response.ResultsCount) + } + if response.ExpectedCount != 104 || response.Shortfall != 1 { + t.Fatalf("shortfall: expected=%d shortfall=%d, want 104/1", response.ExpectedCount, response.Shortfall) + } + if response.Warning == "" { + t.Fatal("warning should explain shortfall") + } +} + +func TestQuakeSearchHandlesStringErrorCode(t *testing.T) { + gin.SetMode(gin.TestMode) + t.Setenv("QUAKE_API_KEY", "") + + quakeServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if got := r.Header.Get("X-QuakeToken"); got != "test-quake-key" { + t.Fatalf("Quake token = %q, want test-quake-key", got) + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"code":"q5000","message":"查询语法错误"}`)) + })) + defer quakeServer.Close() + + h := NewFofaHandler(&config.Config{ + Quake: config.SpaceSearchConfig{ + BaseURL: quakeServer.URL, + APIKey: "test-quake-key", + }, + }, zap.NewNop()) + + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + body := `{"provider":"quake","query":"bad query","fields":"ip,port","size":10,"page":1}` + ctx.Request = httptest.NewRequest(http.MethodPost, "/api/fofa/search", strings.NewReader(body)) + ctx.Request.Header.Set("Content-Type", "application/json") + + h.Search(ctx) + + if recorder.Code != http.StatusBadGateway { + t.Fatalf("Search() status = %d, want %d, body = %s", recorder.Code, http.StatusBadGateway, recorder.Body.String()) + } + bodyText := recorder.Body.String() + if !strings.Contains(bodyText, "查询语法错误") { + t.Fatalf("response should include Quake error message, got %s", bodyText) + } + if strings.Contains(bodyText, "cannot unmarshal") { + t.Fatalf("response exposed JSON type decoding failure: %s", bodyText) + } +} + +func TestExtractInfoCollectJSONObject(t *testing.T) { + t.Parallel() + cases := []struct { + name string + in string + want string + }{ + { + name: "plain json", + in: `{"query":"title:\"CyberStrikeAI\"","warnings":[]}`, + want: `{"query":"title:\"CyberStrikeAI\"","warnings":[]}`, + }, + { + name: "fenced json", + in: "```json\n{\"query\":\"product:nginx\"}\n```", + want: `{"query":"product:nginx"}`, + }, + { + name: "prefixed explanation", + in: "解析结果如下：\n{\"query\":\"ssl.cert.subject.cn:example.com\",\"explanation\":\"ok\"}\n请确认。", + want: `{"query":"ssl.cert.subject.cn:example.com","explanation":"ok"}`, + }, + { + name: "braces inside string", + in: "结果：{\"query\":\"title:\\\"{admin}\\\"\",\"warnings\":[\"check\"]}", + want: `{"query":"title:\"{admin}\"","warnings":["check"]}`, + }, + } + for _, tc := range cases { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + got, err := extractInfoCollectJSONObject(tc.in) + if err != nil { + t.Fatalf("extractInfoCollectJSONObject() error = %v", err) + } + if got != tc.want { + t.Fatalf("extractInfoCollectJSONObject() = %q, want %q", got, tc.want) + } + }) + } +} diff --git a/internal/handler/group.go b/internal/handler/group.go new file mode 100644 index 00000000..b3df88ca --- /dev/null +++ b/internal/handler/group.go @@ -0,0 +1,438 @@ +package handler + +import ( + "errors" + "net/http" + "strings" + "time" + "unicode/utf8" + + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +// GroupHandler 分组处理器 +type GroupHandler struct { + db *database.DB + logger *zap.Logger +} + +const ( + maxGroupNameRunes = 64 + maxGroupIconRunes = 16 +) + +// NewGroupHandler 创建新的分组处理器 +func NewGroupHandler(db *database.DB, logger *zap.Logger) *GroupHandler { + return &GroupHandler{ + db: db, + logger: logger, + } +} + +func validateGroupTextField(field, value string, maxRunes int, required bool) (string, error) { + value = strings.TrimSpace(value) + if value == "" { + if required { + return "", errors.New(field + "不能为空") + } + return "", nil + } + if utf8.RuneCountInString(value) > maxRunes { + return "", errors.New(field + "过长") + } + for _, r := range value { + switch r { + case '<', '>', '"', '\'', '`': + return "", errors.New(field + "包含非法字符") + } + if r < 0x20 || r == 0x7f { + return "", errors.New(field + "包含非法控制字符") + } + } + return value, nil +} + +func validateGroupFields(name, icon string) (string, string, error) { + validName, err := validateGroupTextField("分组名称", name, maxGroupNameRunes, true) + if err != nil { + return "", "", err + } + validIcon, err := validateGroupTextField("分组图标", icon, maxGroupIconRunes, false) + if err != nil { + return "", "", err + } + return validName, validIcon, nil +} + +// CreateGroupRequest 创建分组请求 +type CreateGroupRequest struct { + Name string `json:"name"` + Icon string `json:"icon"` +} + +// CreateGroup 创建分组 +func (h *GroupHandler) CreateGroup(c *gin.Context) { + var req CreateGroupRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + name, icon, err := validateGroupFields(req.Name, req.Icon) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + session, _ := security.CurrentSession(c) + group, err := h.db.CreateGroup(name, icon, session.UserID) + if err != nil { + h.logger.Error("创建分组失败", zap.Error(err)) + // 如果是名称重复错误，返回400状态码 + if err.Error() == "分组名称已存在" { + c.JSON(http.StatusBadRequest, gin.H{"error": "分组名称已存在"}) + return + } + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + c.JSON(http.StatusOK, group) +} + +// ListGroups 列出所有分组 +func (h *GroupHandler) ListGroups(c *gin.Context) { + session, _ := security.CurrentSession(c) + groups, err := h.db.ListGroupsForAccess(session.UserID, session.Scope) + if err != nil { + h.logger.Error("获取分组列表失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + c.JSON(http.StatusOK, groups) +} + +// GetGroup 获取分组 +func (h *GroupHandler) GetGroup(c *gin.Context) { + id := c.Param("id") + if !h.groupAllowed(c, id) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + + group, err := h.db.GetGroup(id) + if err != nil { + h.logger.Error("获取分组失败", zap.Error(err)) + c.JSON(http.StatusNotFound, gin.H{"error": "分组不存在"}) + return + } + + c.JSON(http.StatusOK, group) +} + +// UpdateGroupRequest 更新分组请求 +type UpdateGroupRequest struct { + Name string `json:"name"` + Icon string `json:"icon"` +} + +// UpdateGroup 更新分组 +func (h *GroupHandler) UpdateGroup(c *gin.Context) { + id := c.Param("id") + if !h.groupAllowed(c, id) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + + var req UpdateGroupRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + name, icon, err := validateGroupFields(req.Name, req.Icon) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + if err := h.db.UpdateGroup(id, name, icon); err != nil { + h.logger.Error("更新分组失败", zap.Error(err)) + // 如果是名称重复错误，返回400状态码 + if err.Error() == "分组名称已存在" { + c.JSON(http.StatusBadRequest, gin.H{"error": "分组名称已存在"}) + return + } + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + group, err := h.db.GetGroup(id) + if err != nil { + h.logger.Error("获取更新后的分组失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + c.JSON(http.StatusOK, group) +} + +// DeleteGroup 删除分组 +func (h *GroupHandler) DeleteGroup(c *gin.Context) { + id := c.Param("id") + if !h.groupAllowed(c, id) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + + if err := h.db.DeleteGroup(id); err != nil { + h.logger.Error("删除分组失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + c.JSON(http.StatusOK, gin.H{"message": "删除成功"}) +} + +// AddConversationToGroupRequest 添加对话到分组请求 +type AddConversationToGroupRequest struct { + ConversationID string `json:"conversationId"` + GroupID string `json:"groupId"` +} + +// AddConversationToGroup 将对话添加到分组 +func (h *GroupHandler) AddConversationToGroup(c *gin.Context) { + var req AddConversationToGroupRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if !h.groupConversationAllowed(c, req.ConversationID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + if !h.groupAllowed(c, req.GroupID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该分组"}) + return + } + + if err := h.db.AddConversationToGroup(req.ConversationID, req.GroupID); err != nil { + h.logger.Error("添加对话到分组失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + c.JSON(http.StatusOK, gin.H{"message": "添加成功"}) +} + +// RemoveConversationFromGroup 从分组中移除对话 +func (h *GroupHandler) RemoveConversationFromGroup(c *gin.Context) { + conversationID := c.Param("conversationId") + groupID := c.Param("id") + if !h.groupAllowed(c, groupID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该分组"}) + return + } + if !h.groupConversationAllowed(c, conversationID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + + if err := h.db.RemoveConversationFromGroup(conversationID, groupID); err != nil { + h.logger.Error("从分组中移除对话失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + c.JSON(http.StatusOK, gin.H{"message": "移除成功"}) +} + +// GroupConversation 分组对话响应结构 +type GroupConversation struct { + ID string `json:"id"` + Title string `json:"title"` + Pinned bool `json:"pinned"` + GroupPinned bool `json:"groupPinned"` + CreatedAt time.Time `json:"createdAt"` + UpdatedAt time.Time `json:"updatedAt"` +} + +// GetGroupConversations 获取分组中的所有对话 +func (h *GroupHandler) GetGroupConversations(c *gin.Context) { + groupID := c.Param("id") + if !h.groupAllowed(c, groupID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该分组"}) + return + } + searchQuery := c.Query("search") // 获取搜索参数 + + var conversations []*database.Conversation + var err error + + // 如果有搜索关键词，使用搜索方法；否则使用普通方法 + if searchQuery != "" { + conversations, err = h.db.SearchConversationsByGroup(groupID, searchQuery) + } else { + conversations, err = h.db.GetConversationsByGroup(groupID) + } + + if err != nil { + h.logger.Error("获取分组对话失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + // 获取每个对话在分组中的置顶状态 + groupConvs := make([]GroupConversation, 0, len(conversations)) + for _, conv := range conversations { + if conv == nil || !h.groupConversationAllowed(c, conv.ID) { + continue + } + // 查询分组内置顶状态 + var groupPinned int + err := h.db.QueryRow( + "SELECT COALESCE(pinned, 0) FROM conversation_group_mappings WHERE conversation_id = ? AND group_id = ?", + conv.ID, groupID, + ).Scan(&groupPinned) + if err != nil { + h.logger.Warn("查询分组内置顶状态失败", zap.String("conversationId", conv.ID), zap.Error(err)) + groupPinned = 0 + } + + groupConvs = append(groupConvs, GroupConversation{ + ID: conv.ID, + Title: conv.Title, + Pinned: conv.Pinned, + GroupPinned: groupPinned != 0, + CreatedAt: conv.CreatedAt, + UpdatedAt: conv.UpdatedAt, + }) + } + + c.JSON(http.StatusOK, groupConvs) +} + +// GetAllMappings 批量获取所有分组映射（消除前端 N+1 请求） +func (h *GroupHandler) GetAllMappings(c *gin.Context) { + mappings, err := h.db.GetAllGroupMappings() + if err != nil { + h.logger.Error("获取分组映射失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + filtered := mappings[:0] + for _, mapping := range mappings { + if h.groupConversationAllowed(c, mapping.ConversationID) && h.groupAllowed(c, mapping.GroupID) { + filtered = append(filtered, mapping) + } + } + + c.JSON(http.StatusOK, filtered) +} + +// UpdateConversationPinnedRequest 更新对话置顶状态请求 +type UpdateConversationPinnedRequest struct { + Pinned bool `json:"pinned"` +} + +// UpdateConversationPinned 更新对话置顶状态 +func (h *GroupHandler) UpdateConversationPinned(c *gin.Context) { + conversationID := c.Param("id") + if !h.groupConversationAllowed(c, conversationID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + + var req UpdateConversationPinnedRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + if err := h.db.UpdateConversationPinned(conversationID, req.Pinned); err != nil { + h.logger.Error("更新对话置顶状态失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + c.JSON(http.StatusOK, gin.H{"message": "更新成功"}) +} + +// UpdateGroupPinnedRequest 更新分组置顶状态请求 +type UpdateGroupPinnedRequest struct { + Pinned bool `json:"pinned"` +} + +// UpdateGroupPinned 更新分组置顶状态 +func (h *GroupHandler) UpdateGroupPinned(c *gin.Context) { + groupID := c.Param("id") + if !h.groupAllowed(c, groupID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该分组"}) + return + } + + var req UpdateGroupPinnedRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + if err := h.db.UpdateGroupPinned(groupID, req.Pinned); err != nil { + h.logger.Error("更新分组置顶状态失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + c.JSON(http.StatusOK, gin.H{"message": "更新成功"}) +} + +// UpdateConversationPinnedInGroupRequest 更新分组对话置顶状态请求 +type UpdateConversationPinnedInGroupRequest struct { + Pinned bool `json:"pinned"` +} + +// UpdateConversationPinnedInGroup 更新对话在分组中的置顶状态 +func (h *GroupHandler) UpdateConversationPinnedInGroup(c *gin.Context) { + groupID := c.Param("id") + conversationID := c.Param("conversationId") + if !h.groupAllowed(c, groupID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该分组"}) + return + } + if !h.groupConversationAllowed(c, conversationID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + + var req UpdateConversationPinnedInGroupRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + if err := h.db.UpdateConversationPinnedInGroup(conversationID, groupID, req.Pinned); err != nil { + h.logger.Error("更新分组对话置顶状态失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + c.JSON(http.StatusOK, gin.H{"message": "更新成功"}) +} + +func (h *GroupHandler) groupConversationAllowed(c *gin.Context, conversationID string) bool { + session, ok := security.CurrentSession(c) + if !ok { + return false + } + return h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", conversationID) +} + +func (h *GroupHandler) groupAllowed(c *gin.Context, groupID string) bool { + session, ok := security.CurrentSession(c) + return ok && h.db.UserCanAccessGroup(session.UserID, session.Scope, groupID) +} diff --git a/internal/handler/group_test.go b/internal/handler/group_test.go new file mode 100644 index 00000000..1e7cbde8 --- /dev/null +++ b/internal/handler/group_test.go @@ -0,0 +1,41 @@ +package handler + +import ( + "strings" + "testing" +) + +func TestValidateGroupFieldsAllowsNormalNamesAndIcons(t *testing.T) { + name, icon, err := validateGroupFields(" 日常安全巡检 ", " 📁 ") + if err != nil { + t.Fatalf("validateGroupFields returned error: %v", err) + } + if name != "日常安全巡检" { + t.Fatalf("name = %q, want trimmed normal name", name) + } + if icon != "📁" { + t.Fatalf("icon = %q, want trimmed icon", icon) + } +} + +func TestValidateGroupFieldsRejectsStoredXSSPayloads(t *testing.T) { + tests := []struct { + name string + icon string + }{ + {name: ``, icon: "📁"}, + {name: "日常安全巡检", icon: ``}, + {name: "日常安全巡检`onmouseover=alert(1)", icon: "📁"}, + {name: "日常安全巡检\x00", icon: "📁"}, + {name: strings.Repeat("分", maxGroupNameRunes+1), icon: "📁"}, + {name: "日常安全巡检", icon: strings.Repeat("📁", maxGroupIconRunes+1)}, + } + + for _, tt := range tests { + t.Run(tt.name+"/"+tt.icon, func(t *testing.T) { + if _, _, err := validateGroupFields(tt.name, tt.icon); err == nil { + t.Fatal("validateGroupFields returned nil error for unsafe input") + } + }) + } +} diff --git a/internal/handler/hitl.go b/internal/handler/hitl.go new file mode 100644 index 00000000..c940a6ba --- /dev/null +++ b/internal/handler/hitl.go @@ -0,0 +1,1153 @@ +package handler + +import ( + "context" + "database/sql" + "encoding/json" + "errors" + "math" + "net/http" + "strconv" + "strings" + "sync" + "time" + + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/multiagent" + + "github.com/gin-gonic/gin" + "github.com/google/uuid" + "go.uber.org/zap" +) + +type hitlRuntimeConfig struct { + Enabled bool + Mode string + Reviewer string + SensitiveTools map[string]struct{} + Timeout time.Duration +} + +type hitlDecision struct { + Decision string + Comment string + EditedArguments map[string]interface{} +} + +type pendingInterrupt struct { + ConversationID string + InterruptID string + Mode string + ToolName string + ToolCallID string + decideCh chan hitlDecision +} + +type HITLManager struct { + db *database.DB + logger *zap.Logger + + mu sync.RWMutex + runtime map[string]hitlRuntimeConfig + pending map[string]*pendingInterrupt + // approvedExec 审批通过、待回写 tool_result 的队列（按会话 FIFO） + approvedExec map[string][]hitlApprovedExecTrack +} + +func NewHITLManager(db *database.DB, logger *zap.Logger) *HITLManager { + return &HITLManager{ + db: db, + logger: logger, + runtime: make(map[string]hitlRuntimeConfig), + pending: make(map[string]*pendingInterrupt), + } +} + +func (m *HITLManager) EnsureSchema() error { + if _, err := m.db.Exec(` +CREATE TABLE IF NOT EXISTS hitl_interrupts ( + id TEXT PRIMARY KEY, + conversation_id TEXT NOT NULL, + message_id TEXT, + mode TEXT NOT NULL, + tool_name TEXT NOT NULL, + tool_call_id TEXT, + payload TEXT, + status TEXT NOT NULL, + reviewer TEXT NOT NULL DEFAULT 'human', + decision TEXT, + decision_comment TEXT, + created_at DATETIME NOT NULL, + decided_at DATETIME +);`); err != nil { + return err + } + _, err := m.db.Exec(` +CREATE TABLE IF NOT EXISTS hitl_conversation_configs ( + conversation_id TEXT PRIMARY KEY, + enabled INTEGER NOT NULL DEFAULT 0, + mode TEXT NOT NULL DEFAULT 'off', + sensitive_tools TEXT NOT NULL DEFAULT '[]', + timeout_seconds INTEGER NOT NULL DEFAULT 0, + updated_at DATETIME NOT NULL +);`) + if err != nil { + return err + } + m.migrateHitlSchemaColumns() + + // On startup, cancel all orphaned pending interrupts from previous process. + // Their in-memory channels are gone, so they can never be resolved. + res, err := m.db.Exec(`UPDATE hitl_interrupts SET status='cancelled', decision='reject', + decision_comment='process restarted', decided_at=CURRENT_TIMESTAMP, decided_by='system' + WHERE status='pending'`) + if err != nil { + m.logger.Warn("failed to cancel orphaned HITL interrupts", zap.Error(err)) + } else if n, _ := res.RowsAffected(); n > 0 { + m.logger.Info("cancelled orphaned HITL interrupts from previous process", zap.Int64("count", n)) + } + if err := m.reconcileRestartInterruptedMessages(); err != nil { + m.logger.Warn("failed to finalize assistant messages interrupted by process restart", zap.Error(err)) + } + return nil +} + +// reconcileRestartInterruptedMessages completes durable terminal state for +// historical assistant placeholders that have explicit evidence of being over: +// a terminal HITL/process event, or a later message in the same conversation. +// The evidence requirement avoids rewriting a placeholder that could still be +// recoverable by another runtime. +func (m *HITLManager) reconcileRestartInterruptedMessages() error { + rows, err := m.db.Query(` +SELECT msg.id, msg.conversation_id, + COALESCE(( + SELECT pd.event_type + FROM process_details pd + WHERE pd.message_id = msg.id + AND pd.event_type IN ('cancelled', 'timeout', 'error') + ORDER BY pd.created_at DESC LIMIT 1 + ), '') AS terminal_event, + COALESCE(( + SELECT hi.status + FROM hitl_interrupts hi + WHERE hi.message_id = msg.id + ORDER BY COALESCE(hi.decided_at, hi.created_at) DESC LIMIT 1 + ), '') AS hitl_status, + COALESCE(( + SELECT hi.decision + FROM hitl_interrupts hi + WHERE hi.message_id = msg.id + ORDER BY COALESCE(hi.decided_at, hi.created_at) DESC LIMIT 1 + ), '') AS hitl_decision, + COALESCE(( + SELECT hi.decision_comment + FROM hitl_interrupts hi + WHERE hi.message_id = msg.id + ORDER BY COALESCE(hi.decided_at, hi.created_at) DESC LIMIT 1 + ), '') AS decision_comment, + COALESCE(( + SELECT MAX(COALESCE(hi.decided_at, hi.created_at)) + FROM hitl_interrupts hi + WHERE hi.message_id = msg.id + ), ( + SELECT MIN(later.created_at) + FROM messages later + WHERE later.conversation_id = msg.conversation_id + AND later.created_at > msg.created_at + ), ( + SELECT MAX(pd.created_at) + FROM process_details pd + WHERE pd.message_id = msg.id + ), msg.updated_at, msg.created_at) AS interrupted_at +FROM messages msg +WHERE msg.role = 'assistant' + AND TRIM(msg.content) IN ('处理中...', 'Processing...') + AND ( + EXISTS ( + SELECT 1 FROM hitl_interrupts hi + WHERE hi.message_id = msg.id + AND (hi.status IN ('cancelled', 'timeout') + OR (hi.status = 'decided' AND hi.decision = 'reject')) + ) + OR EXISTS ( + SELECT 1 FROM process_details pd + WHERE pd.message_id = msg.id + AND pd.event_type IN ('cancelled', 'timeout', 'error') + ) + OR EXISTS ( + SELECT 1 FROM messages later + WHERE later.conversation_id = msg.conversation_id + AND later.created_at > msg.created_at + ) + )`) + if err != nil { + return err + } + type interruptedMessage struct { + messageID string + conversationID string + terminalEvent string + hitlStatus string + hitlDecision string + decisionComment string + interruptedAt string + } + var interrupted []interruptedMessage + for rows.Next() { + var item interruptedMessage + if err := rows.Scan(&item.messageID, &item.conversationID, &item.terminalEvent, + &item.hitlStatus, &item.hitlDecision, &item.decisionComment, &item.interruptedAt); err != nil { + rows.Close() + return err + } + interrupted = append(interrupted, item) + } + if err := rows.Close(); err != nil { + return err + } + if len(interrupted) == 0 { + return nil + } + + tx, err := m.db.Begin() + if err != nil { + return err + } + defer func() { _ = tx.Rollback() }() + for _, item := range interrupted { + eventType := strings.ToLower(strings.TrimSpace(item.terminalEvent)) + decision := strings.ToLower(strings.TrimSpace(item.hitlDecision)) + comment := strings.ToLower(strings.TrimSpace(item.decisionComment)) + if eventType == "" { + if strings.EqualFold(strings.TrimSpace(item.hitlStatus), "timeout") || strings.Contains(comment, "timeout") { + eventType = "timeout" + } else { + eventType = "cancelled" + } + } + + notice := "任务因服务重启已中断。" + reason := "process_restarted" + switch eventType { + case "timeout": + notice = "任务等待审批超时，已自动拒绝。" + reason = "hitl_timeout" + case "error": + notice = "任务执行失败，已停止。" + reason = "execution_error" + case "cancelled": + if decision == "reject" && comment != "process restarted" { + notice = "任务审批已拒绝，执行已停止。" + reason = "hitl_rejected" + } else if comment == "process restarted" { + notice = "任务因服务重启已中断，审批已取消。" + } + default: + eventType = "cancelled" + } + detailData, _ := json.Marshal(map[string]string{"reason": reason, "status": eventType}) + result, err := tx.Exec(` +UPDATE messages +SET content = ?, updated_at = ? +WHERE id = ? AND TRIM(content) IN ('处理中...', 'Processing...')`, + notice, item.interruptedAt, item.messageID) + if err != nil { + return err + } + updated, _ := result.RowsAffected() + if updated == 0 { + continue + } + if _, err := tx.Exec(` +INSERT INTO process_details (id, message_id, conversation_id, event_type, message, data, created_at) +SELECT ?, ?, ?, ?, ?, ?, ? +WHERE NOT EXISTS ( + SELECT 1 FROM process_details + WHERE message_id = ? AND event_type IN ('cancelled', 'timeout', 'error') +)`, uuid.NewString(), item.messageID, item.conversationID, eventType, notice, string(detailData), + item.interruptedAt, item.messageID); err != nil { + return err + } + } + return tx.Commit() +} + +func normalizeHitlMode(mode string) string { + v := strings.ToLower(strings.TrimSpace(mode)) + if v == "" { + return "approval" + } + switch v { + case "off": + return "off" + case "feedback", "followup": + return "approval" + case "approval", "review_edit": + return v + default: + return "approval" + } +} + +func (m *HITLManager) ActivateConversation(conversationID string, req *HITLRequest) { + if req == nil || !req.Enabled { + m.DeactivateConversation(conversationID) + return + } + tools := make(map[string]struct{}) + for _, t := range req.SensitiveTools { + n := strings.ToLower(strings.TrimSpace(t)) + if n != "" { + tools[n] = struct{}{} + } + } + // timeout <= 0 means wait forever (no timeout). + timeout := time.Duration(0) + if req.TimeoutSeconds > 0 { + timeout = time.Duration(req.TimeoutSeconds) * time.Second + } + m.mu.Lock() + m.runtime[conversationID] = hitlRuntimeConfig{ + Enabled: true, + Mode: normalizeHitlMode(req.Mode), + Reviewer: normalizeHitlReviewer(req.Reviewer), + SensitiveTools: tools, + Timeout: timeout, + } + m.mu.Unlock() +} + +func (m *HITLManager) DeactivateConversation(conversationID string) { + m.mu.Lock() + delete(m.runtime, conversationID) + m.mu.Unlock() +} + +// hitlConfigGlobalToolWhitelist 来自 config.yaml hitl.tool_whitelist（去重、去空），并合并内置元工具免审批项。 +func (h *AgentHandler) hitlConfigGlobalToolWhitelist() []string { + if h == nil || h.config == nil { + return multiagent.MergeHitlExemptMetaTools(nil) + } + raw := h.config.Hitl.ToolWhitelist + seen := make(map[string]struct{}) + out := make([]string, 0, len(raw)+len(multiagent.HitlExemptMetaTools)) + for _, t := range raw { + n := strings.ToLower(strings.TrimSpace(t)) + if n == "" { + continue + } + if _, ok := seen[n]; ok { + continue + } + seen[n] = struct{}{} + out = append(out, strings.TrimSpace(t)) + } + return multiagent.MergeHitlExemptMetaTools(out) +} + +// hitlRequestWithMergedConfigWhitelist 将会话/API 中的白名单与 config.yaml 全局白名单及内置元工具免审批项合并（并集），仅用于运行时 Activate；不写入数据库。 +func (h *AgentHandler) hitlRequestWithMergedConfigWhitelist(req *HITLRequest) *HITLRequest { + if req == nil { + return nil + } + seen := make(map[string]struct{}) + union := make([]string, 0, len(req.SensitiveTools)+16) + add := func(t string) { + n := strings.ToLower(strings.TrimSpace(t)) + if n == "" { + return + } + if _, ok := seen[n]; ok { + return + } + seen[n] = struct{}{} + union = append(union, strings.TrimSpace(t)) + } + for _, t := range h.hitlConfigGlobalToolWhitelist() { + add(t) + } + for _, t := range req.SensitiveTools { + add(t) + } + out := *req + out.SensitiveTools = multiagent.MergeHitlExemptMetaTools(union) + return &out +} + +func (m *HITLManager) shouldInterrupt(conversationID, toolName string) (hitlRuntimeConfig, bool) { + m.mu.RLock() + cfg, ok := m.runtime[conversationID] + m.mu.RUnlock() + if !ok || !cfg.Enabled { + return hitlRuntimeConfig{}, false + } + // 语义：SensitiveTools 现在作为“白名单（免审批工具）” + // 空白名单 => 全部工具都需要审批 + if len(cfg.SensitiveTools) == 0 { + return cfg, true + } + _, inWhitelist := cfg.SensitiveTools[strings.ToLower(strings.TrimSpace(toolName))] + return cfg, !inWhitelist +} + +// NeedsToolApproval 与 Agent 工具层 shouldInterrupt 语义一致：仅当该会话已开启人机协同且工具不在免审批白名单时为 true。 +func (m *HITLManager) NeedsToolApproval(conversationID, toolName string) bool { + if m == nil { + return false + } + _, need := m.shouldInterrupt(conversationID, toolName) + return need +} + +func (m *HITLManager) CreatePendingInterrupt(conversationID, assistantMessageID, mode, toolName, toolCallID, payload, reviewer string) (*pendingInterrupt, error) { + now := time.Now() + id := "hitl_" + strings.ReplaceAll(uuid.New().String(), "-", "") + reviewer = normalizeHitlReviewer(reviewer) + if _, err := m.db.Exec(`INSERT INTO hitl_interrupts + (id, conversation_id, message_id, mode, tool_name, tool_call_id, payload, status, reviewer, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, 'pending', ?, ?)`, + id, conversationID, assistantMessageID, mode, toolName, toolCallID, payload, reviewer, now); err != nil { + return nil, err + } + // 刷新页面后侧栏依赖 DB 配置；若仅内存 Activate 未落库，会导致「有待审批却显示关闭」 + _ = m.ensureConversationHITLModePersisted(conversationID, mode) + p := &pendingInterrupt{ + ConversationID: conversationID, + InterruptID: id, + Mode: normalizeHitlMode(mode), + ToolName: toolName, + ToolCallID: toolCallID, + decideCh: make(chan hitlDecision, 1), + } + // Agent 审查不会等待人工决策，也不应进入人工审批的内存待办队列。 + if reviewer != "audit_agent" { + m.mu.Lock() + m.pending[id] = p + m.mu.Unlock() + } + return p, nil +} + +// ensureConversationHITLModePersisted 在产生待审批时把 mode 写入 hitl_conversation_configs，避免刷新后 GET 配置仍为关闭。 +func (m *HITLManager) ensureConversationHITLModePersisted(conversationID, interruptMode string) error { + if strings.TrimSpace(conversationID) == "" { + return nil + } + nm := normalizeHitlMode(interruptMode) + if nm == "off" { + return nil + } + cfg, err := m.LoadConversationConfig(conversationID) + if err != nil { + return err + } + if cfg.Enabled && normalizeHitlMode(cfg.Mode) == nm { + return nil + } + cfg.Enabled = true + cfg.Mode = nm + if cfg.TimeoutSeconds < 0 { + cfg.TimeoutSeconds = 0 + } + return m.SaveConversationConfig(conversationID, cfg) +} + +// PendingHITLInterruptMode 返回该会话最新一条 pending 中断的协同模式（用于 GET 配置时与库内「关闭」状态对齐）。 +func (m *HITLManager) PendingHITLInterruptMode(conversationID string) (string, bool) { + if strings.TrimSpace(conversationID) == "" { + return "", false + } + var mode string + err := m.db.QueryRow(`SELECT mode FROM hitl_interrupts WHERE conversation_id = ? AND status = 'pending' ORDER BY created_at DESC LIMIT 1`, conversationID). + Scan(&mode) + if err != nil { + if errors.Is(err, sql.ErrNoRows) { + return "", false + } + return "", false + } + mode = strings.TrimSpace(mode) + if mode == "" { + return "", false + } + return mode, true +} + +func hitlStoredConfigEffective(cfg *HITLRequest) bool { + if cfg == nil { + return false + } + if cfg.Enabled { + return true + } + return normalizeHitlMode(cfg.Mode) != "off" +} + +func (m *HITLManager) ResolveInterrupt(interruptID, decision, comment string, editedArguments map[string]interface{}) error { + decision = strings.ToLower(strings.TrimSpace(decision)) + if decision != "approve" && decision != "reject" { + return errors.New("decision must be approve/reject") + } + m.mu.RLock() + p, ok := m.pending[interruptID] + m.mu.RUnlock() + if !ok { + return errors.New("interrupt not found or already resolved") + } + d := hitlDecision{ + Decision: decision, + Comment: strings.TrimSpace(comment), + EditedArguments: editedArguments, + } + select { + case p.decideCh <- d: + return nil + default: + return errors.New("interrupt already resolved or decision channel busy") + } +} + +func (m *HITLManager) SaveConversationConfig(conversationID string, req *HITLRequest) error { + if strings.TrimSpace(conversationID) == "" { + return errors.New("conversationId is required") + } + if req == nil { + req = &HITLRequest{Enabled: false, Mode: "off", TimeoutSeconds: 0} + } + mode := normalizeHitlMode(req.Mode) + if !req.Enabled { + mode = "off" + } + tools, _ := json.Marshal(req.SensitiveTools) + timeout := req.TimeoutSeconds + if timeout < 0 { + timeout = 0 + } + _, err := m.db.Exec(`INSERT INTO hitl_conversation_configs + (conversation_id, enabled, mode, reviewer, sensitive_tools, timeout_seconds, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(conversation_id) DO UPDATE SET + enabled=excluded.enabled, mode=excluded.mode, reviewer=excluded.reviewer, sensitive_tools=excluded.sensitive_tools, timeout_seconds=excluded.timeout_seconds, updated_at=excluded.updated_at`, + conversationID, boolToInt(req.Enabled), mode, normalizeHitlReviewer(req.Reviewer), string(tools), timeout, time.Now()) + return err +} + +func (m *HITLManager) LoadConversationConfig(conversationID string) (*HITLRequest, error) { + var enabledInt int + var mode, reviewer, toolsJSON string + var timeout int + err := m.db.QueryRow(`SELECT enabled, mode, COALESCE(reviewer,'human'), sensitive_tools, timeout_seconds FROM hitl_conversation_configs WHERE conversation_id = ?`, conversationID). + Scan(&enabledInt, &mode, &reviewer, &toolsJSON, &timeout) + if errors.Is(err, sql.ErrNoRows) { + return &HITLRequest{Enabled: false, Mode: "off", Reviewer: "human", SensitiveTools: []string{}, TimeoutSeconds: 0}, nil + } + if err != nil { + return nil, err + } + if timeout < 0 { + timeout = 0 + } + tools := make([]string, 0) + _ = json.Unmarshal([]byte(toolsJSON), &tools) + return &HITLRequest{ + Enabled: enabledInt == 1, + Mode: mode, + Reviewer: normalizeHitlReviewer(reviewer), + SensitiveTools: tools, + TimeoutSeconds: timeout, + }, nil +} + +func (m *HITLManager) HasConversationConfig(conversationID string) (bool, error) { + if strings.TrimSpace(conversationID) == "" { + return false, nil + } + var one int + err := m.db.QueryRow(`SELECT 1 FROM hitl_conversation_configs WHERE conversation_id = ? LIMIT 1`, conversationID).Scan(&one) + if errors.Is(err, sql.ErrNoRows) { + return false, nil + } + return err == nil, err +} + +func (m *HITLManager) waitDecision(ctx context.Context, p *pendingInterrupt, timeout time.Duration) (hitlDecision, error) { + defer func() { + m.mu.Lock() + delete(m.pending, p.InterruptID) + m.mu.Unlock() + }() + var timeoutCh <-chan time.Time + if timeout > 0 { + timer := time.NewTimer(timeout) + defer timer.Stop() + timeoutCh = timer.C + } + select { + case d := <-p.decideCh: + // 只有 review_edit 模式允许改参；其他模式一律忽略 edited arguments + if p.Mode != "review_edit" && len(d.EditedArguments) > 0 { + d.EditedArguments = nil + } + _, _ = m.db.Exec(`UPDATE hitl_interrupts SET status='decided', decision=?, decision_comment=?, decided_at=?, decided_by='human' WHERE id=?`, + d.Decision, d.Comment, time.Now(), p.InterruptID) + return d, nil + case <-timeoutCh: + comment := "HITL timeout auto-reject for safety" + _, _ = m.db.Exec(`UPDATE hitl_interrupts SET status='timeout', decision='reject', decision_comment=?, decided_at=?, decided_by='system' WHERE id=?`, + comment, time.Now(), p.InterruptID) + return hitlDecision{Decision: "reject", Comment: comment}, nil + case <-ctx.Done(): + _, _ = m.db.Exec(`UPDATE hitl_interrupts SET status='cancelled', decision='reject', decision_comment='task cancelled', decided_at=?, decided_by='system' WHERE id=?`, + time.Now(), p.InterruptID) + return hitlDecision{Decision: "reject", Comment: "task cancelled"}, ctx.Err() + } +} + +func (h *AgentHandler) activateHITLForConversation(conversationID string, req *HITLRequest) { + if h.hitlManager == nil { + return + } + if req == nil { + cfg, err := h.loadHITLConversationConfig(conversationID) + if err == nil { + req = cfg + } + } + if req != nil && strings.TrimSpace(req.Reviewer) == "" { + req.Reviewer = h.hitlEffectiveDefaultReviewer() + } + h.hitlManager.ActivateConversation(conversationID, h.hitlRequestWithMergedConfigWhitelist(req)) +} + +func (h *AgentHandler) loadHITLConversationConfig(conversationID string) (*HITLRequest, error) { + cfg, err := h.hitlManager.LoadConversationConfig(conversationID) + if err != nil { + return nil, err + } + has, err := h.hitlManager.HasConversationConfig(conversationID) + if err != nil { + return nil, err + } + if !has { + cfg.Reviewer = h.hitlEffectiveDefaultReviewer() + } + return cfg, nil +} + +func (h *AgentHandler) waitHITLApproval(runCtx context.Context, cancelRun context.CancelCauseFunc, conversationID, assistantMessageID, toolName, toolCallID string, payload map[string]interface{}, sendEventFunc func(eventType, message string, data interface{})) (*hitlDecision, error) { + cfg, need := h.hitlManager.shouldInterrupt(conversationID, toolName) + if !need { + return nil, nil + } + h.enrichHitlApprovalPayload(conversationID, assistantMessageID, payload) + approvalStartedAt := time.Now().UTC() + timeoutSeconds := int(cfg.Timeout / time.Second) + var approvalExpiresAt *time.Time + if timeoutSeconds > 0 { + expiresAt := approvalStartedAt.Add(cfg.Timeout) + approvalExpiresAt = &expiresAt + } + payload["hitlApproval"] = map[string]interface{}{ + "createdAt": approvalStartedAt, + "timeoutSeconds": timeoutSeconds, + "expiresAt": approvalExpiresAt, + } + payloadRaw, _ := json.Marshal(payload) + p, err := h.hitlManager.CreatePendingInterrupt(conversationID, assistantMessageID, cfg.Mode, toolName, toolCallID, string(payloadRaw), cfg.Reviewer) + if err != nil { + h.logger.Warn("创建 HITL 中断失败", zap.Error(err)) + return nil, err + } + emitHITL := func(eventType, message string, eventData map[string]interface{}) { + clientData := enrichProgressEventData(eventData, conversationID, assistantMessageID) + if sendEventFunc != nil { + sendEventFunc(eventType, message, clientData) + } + if strings.TrimSpace(assistantMessageID) != "" && h.db != nil { + if err := h.db.AddProcessDetail(assistantMessageID, conversationID, eventType, message, clientData); err != nil { + h.logger.Warn("保存 HITL 过程详情失败", zap.Error(err), zap.String("eventType", eventType)) + } + } + } + + if cfg.Reviewer == "audit_agent" { + emitHITL("hitl_audit_agent_started", "审计 Agent 正在审查此请求", map[string]interface{}{ + "conversationId": conversationID, + "interruptId": p.InterruptID, + "toolName": toolName, + "toolCallId": toolCallID, + "mode": cfg.Mode, + "reviewer": "audit_agent", + "status": "audit_running", + "payload": payload, + }) + ad := h.auditAgentReview(runCtx, cfg.Mode, toolName, payload) + now := time.Now() + _, _ = h.db.Exec(`UPDATE hitl_interrupts SET status='decided', decision=?, decision_comment=?, decided_at=?, decided_by='audit_agent' WHERE id=?`, + ad.Decision, ad.Comment, now, p.InterruptID) + emitHITL("hitl_audit_agent", "审计 Agent 已裁决", map[string]interface{}{ + "conversationId": conversationID, + "interruptId": p.InterruptID, + "toolName": toolName, + "toolCallId": toolCallID, + "mode": cfg.Mode, + "status": "decided", + "decision": ad.Decision, + "comment": ad.Comment, + "editedArgs": ad.EditedArguments, + "decidedBy": "audit_agent", + "reviewer": "audit_agent", + }) + if ad.Decision == "reject" { + emitHITL("hitl_rejected", "审计 Agent 拒绝本次工具调用", map[string]interface{}{ + "conversationId": conversationID, + "interruptId": p.InterruptID, + "toolName": toolName, + "toolCallId": toolCallID, + "mode": cfg.Mode, + "decision": "reject", + "comment": ad.Comment, + "decidedBy": "audit_agent", + "reviewer": "audit_agent", + }) + return &ad, nil + } + emitHITL("hitl_resumed", "审计 Agent 已通过，继续执行", map[string]interface{}{ + "conversationId": conversationID, + "interruptId": p.InterruptID, + "toolName": toolName, + "toolCallId": toolCallID, + "mode": cfg.Mode, + "decision": "approve", + "comment": ad.Comment, + "editedArgs": ad.EditedArguments, + "decidedBy": "audit_agent", + "reviewer": "audit_agent", + }) + h.hitlManager.TrackApprovedHitlExecution(p.InterruptID, conversationID, toolName, toolCallID) + return &ad, nil + } + + emitHITL("hitl_interrupt", "命中人机协同审批", map[string]interface{}{ + "conversationId": conversationID, + "interruptId": p.InterruptID, + "mode": cfg.Mode, + "toolName": toolName, + "toolCallId": toolCallID, + "reviewer": "human", + "status": "pending", + "createdAt": approvalStartedAt, + "timeoutSeconds": timeoutSeconds, + "expiresAt": approvalExpiresAt, + "payload": payload, + }) + d, waitErr := h.hitlManager.waitDecision(runCtx, p, cfg.Timeout) + if waitErr != nil { + if cancelRun != nil && (errors.Is(waitErr, context.Canceled) || errors.Is(waitErr, context.DeadlineExceeded)) { + cause := context.Cause(runCtx) + switch { + case errors.Is(cause, ErrTaskCancelled): + cancelRun(ErrTaskCancelled) + case cause != nil: + cancelRun(cause) + case errors.Is(waitErr, context.DeadlineExceeded): + cancelRun(context.DeadlineExceeded) + default: + cancelRun(ErrTaskCancelled) + } + } + return nil, waitErr + } + if d.Decision == "reject" { + rejectMsg := "人工拒绝本次工具调用，模型将基于反馈继续迭代" + timedOut := strings.Contains(strings.ToLower(strings.TrimSpace(d.Comment)), "timeout") + if timedOut { + rejectMsg = "审批超时，安全起见已自动拒绝，模型将基于反馈继续迭代" + } + status := "decided" + decidedBy := "human" + if timedOut { + status = "timeout" + decidedBy = "system" + } + emitHITL("hitl_rejected", rejectMsg, map[string]interface{}{ + "conversationId": conversationID, + "interruptId": p.InterruptID, + "toolName": toolName, + "toolCallId": toolCallID, + "mode": cfg.Mode, + "status": status, + "decision": "reject", + "comment": d.Comment, + "decidedBy": decidedBy, + "reviewer": "human", + }) + return &d, nil + } + emitHITL("hitl_resumed", "人工确认通过，继续执行", map[string]interface{}{ + "conversationId": conversationID, + "interruptId": p.InterruptID, + "toolName": toolName, + "toolCallId": toolCallID, + "mode": cfg.Mode, + "decision": "approve", + "comment": d.Comment, + "editedArgs": d.EditedArguments, + "reviewer": "human", + }) + h.hitlManager.TrackApprovedHitlExecution(p.InterruptID, conversationID, toolName, toolCallID) + return &d, nil +} + +func (h *AgentHandler) handleHITLToolCall(runCtx context.Context, cancelRun context.CancelCauseFunc, conversationID, assistantMessageID string, data map[string]interface{}, sendEventFunc func(eventType, message string, data interface{})) { + if h.hitlManager == nil { + return + } + toolName, _ := data["toolName"].(string) + toolCallID, _ := data["toolCallId"].(string) + d, err := h.waitHITLApproval(runCtx, cancelRun, conversationID, assistantMessageID, toolName, toolCallID, data, sendEventFunc) + if err != nil || d == nil { + return + } + if len(d.EditedArguments) > 0 { + if argsObj, ok := data["argumentsObj"].(map[string]interface{}); ok { + for k := range argsObj { + delete(argsObj, k) + } + for k, v := range d.EditedArguments { + argsObj[k] = v + } + if b, mErr := json.Marshal(argsObj); mErr == nil { + data["arguments"] = string(b) + } + } + } +} + +func (h *AgentHandler) ListHITLPending(c *gin.Context) { + page, _ := strconv.Atoi(c.DefaultQuery("page", "1")) + if page < 1 { + page = 1 + } + pageSize, _ := strconv.Atoi(c.DefaultQuery("pageSize", "20")) + pageSize = int(math.Max(1, math.Min(float64(pageSize), 200))) + offset := (page - 1) * pageSize + q, args := h.buildHitlListQuery(false) + q, args = h.appendHitlListFilters(q, args, c) + q, args = appendConversationAccessSQL(q, args, "conversation_id", notificationAccessFromContext(c)) + total, err := h.countHitlQuery(q, args) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + q += " ORDER BY created_at DESC LIMIT ? OFFSET ?" + args = append(args, pageSize, offset) + rows, err := h.db.Query(q, args...) + if err != nil { + c.JSON(500, gin.H{"error": err.Error()}) + return + } + defer rows.Close() + items, err := h.scanHitlInterruptRows(rows) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"items": items, "page": page, "pageSize": pageSize, "total": total}) +} + +type hitlDecisionReq struct { + InterruptID string `json:"interruptId" binding:"required"` + Decision string `json:"decision" binding:"required"` + Comment string `json:"comment,omitempty"` + EditedArguments map[string]interface{} `json:"editedArguments,omitempty"` +} + +func (h *AgentHandler) DecideHITLInterrupt(c *gin.Context) { + var req hitlDecisionReq + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(400, gin.H{"error": err.Error()}) + return + } + if h.hitlManager == nil { + c.JSON(500, gin.H{"error": "hitl manager unavailable"}) + return + } + if !h.hitlInterruptAllowed(c, req.InterruptID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + if err := h.hitlManager.ResolveInterrupt(req.InterruptID, req.Decision, req.Comment, req.EditedArguments); err != nil { + c.JSON(http.StatusConflict, gin.H{"error": err.Error()}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "hitl", "decision", "HITL 审批决策", "hitl_interrupt", req.InterruptID, map[string]interface{}{ + "decision": req.Decision, + }) + } + c.JSON(http.StatusOK, gin.H{"ok": true}) +} + +func (h *AgentHandler) DismissHITLInterrupt(c *gin.Context) { + var req struct { + InterruptID string `json:"interruptId" binding:"required"` + } + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(400, gin.H{"error": err.Error()}) + return + } + if h.hitlManager == nil { + c.JSON(500, gin.H{"error": "hitl manager unavailable"}) + return + } + if !h.hitlInterruptAllowed(c, req.InterruptID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + res, err := h.db.Exec(`UPDATE hitl_interrupts SET status='cancelled', decision='reject', + decision_comment='dismissed by user', decided_at=CURRENT_TIMESTAMP, decided_by='human' + WHERE id=? AND status='pending'`, req.InterruptID) + if err != nil { + c.JSON(500, gin.H{"error": err.Error()}) + return + } + n, _ := res.RowsAffected() + if n == 0 { + c.JSON(404, gin.H{"error": "interrupt not found or already resolved"}) + return + } + // Also drain from in-memory map if present + h.hitlManager.mu.Lock() + if p, ok := h.hitlManager.pending[req.InterruptID]; ok { + delete(h.hitlManager.pending, req.InterruptID) + select { + case p.decideCh <- hitlDecision{Decision: "reject", Comment: "dismissed by user"}: + default: + } + } + h.hitlManager.mu.Unlock() + c.JSON(http.StatusOK, gin.H{"ok": true}) +} + +func (h *AgentHandler) interceptHITLForEinoTool(runCtx context.Context, cancelRun context.CancelCauseFunc, conversationID, assistantMessageID string, sendEventFunc func(eventType, message string, data interface{}), toolName, arguments string) (string, error) { + payload := map[string]interface{}{ + "toolName": toolName, + "arguments": arguments, + "source": "eino_middleware", + "toolCallId": "", + } + var argsObj map[string]interface{} + if strings.TrimSpace(arguments) != "" { + _ = json.Unmarshal([]byte(arguments), &argsObj) + if argsObj != nil { + payload["argumentsObj"] = argsObj + } + } + d, err := h.waitHITLApproval(runCtx, cancelRun, conversationID, assistantMessageID, toolName, "", payload, sendEventFunc) + if err != nil || d == nil { + return arguments, err + } + if d.Decision == "reject" { + return arguments, multiagent.NewHumanRejectError(d.Comment) + } + if len(d.EditedArguments) > 0 { + edited, mErr := json.Marshal(d.EditedArguments) + if mErr == nil { + return string(edited), nil + } + } + return arguments, nil +} + +type hitlConfigReq struct { + ConversationID string `json:"conversationId" binding:"required"` + HITLRequest +} + +func (h *AgentHandler) GetHITLConversationConfig(c *gin.Context) { + conversationID := strings.TrimSpace(c.Param("conversationId")) + if conversationID == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "conversationId is required"}) + return + } + if !h.hitlConversationAllowed(c, conversationID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + cfg, err := h.loadHITLConversationConfig(conversationID) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if !hitlStoredConfigEffective(cfg) { + if pendMode, ok := h.hitlManager.PendingHITLInterruptMode(conversationID); ok { + cfg2 := *cfg + cfg2.Enabled = true + cfg2.Mode = normalizeHitlMode(pendMode) + if cfg2.TimeoutSeconds < 0 { + cfg2.TimeoutSeconds = 0 + } + cfg = &cfg2 + } + } + c.JSON(http.StatusOK, gin.H{ + "conversationId": conversationID, + "hitl": cfg, + "defaultReviewer": h.hitlEffectiveDefaultReviewer(), + "hitlGlobalToolWhitelist": h.hitlConfigGlobalToolWhitelist(), + }) +} + +func (h *AgentHandler) UpsertHITLConversationConfig(c *gin.Context) { + var req hitlConfigReq + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if !h.hitlConversationAllowed(c, req.ConversationID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + req.Mode = normalizeHitlMode(req.Mode) + req.Reviewer = normalizeHitlReviewer(req.Reviewer) + if strings.TrimSpace(req.Reviewer) == "" { + req.Reviewer = h.hitlEffectiveDefaultReviewer() + } + if err := h.hitlManager.SaveConversationConfig(req.ConversationID, &req.HITLRequest); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if h.hitlWhitelistSaver != nil && len(req.SensitiveTools) > 0 { + if err := h.hitlWhitelistSaver.MergeHitlToolWhitelistIntoConfig(req.SensitiveTools); err != nil { + h.logger.Warn("HITL 会话配置已保存，但合并工具白名单到 config.yaml 失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{ + "error": "会话配置已保存，但写入 config.yaml 失败: " + err.Error(), + }) + return + } + } + h.hitlManager.ActivateConversation(req.ConversationID, h.hitlRequestWithMergedConfigWhitelist(&req.HITLRequest)) + c.JSON(http.StatusOK, gin.H{"ok": true}) +} + +type mergeHitlGlobalWhitelistReq struct { + SensitiveTools []string `json:"sensitiveTools"` +} + +type setHitlGlobalWhitelistReq struct { + ToolWhitelist []string `json:"toolWhitelist"` +} + +// GetHITLGlobalToolWhitelist 返回 config.yaml 中的全局免审批工具白名单。 +func (h *AgentHandler) GetHITLGlobalToolWhitelist(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{ + "toolWhitelist": h.hitlConfigGlobalToolWhitelist(), + "defaultReviewer": h.hitlEffectiveDefaultReviewer(), + }) +} + +type setHitlDefaultReviewerReq struct { + Reviewer string `json:"reviewer"` +} + +// GetHITLDefaultReviewer 返回 config.yaml 中的全局默认审批方。 +func (h *AgentHandler) GetHITLDefaultReviewer(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{ + "defaultReviewer": h.hitlEffectiveDefaultReviewer(), + }) +} + +// UpdateHITLDefaultReviewer 将全局默认审批方写入 config.yaml（未选会话时切换审批方）。 +func (h *AgentHandler) UpdateHITLDefaultReviewer(c *gin.Context) { + if h.hitlDefaultReviewerSaver == nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "HITL 配置持久化不可用"}) + return + } + var req setHitlDefaultReviewerReq + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + reviewer := normalizeHitlReviewer(req.Reviewer) + if err := h.hitlDefaultReviewerSaver.UpdateHitlDefaultReviewer(reviewer); err != nil { + h.logger.Warn("写入 HITL 默认审批方到 config.yaml 失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if h.config != nil { + h.config.Hitl.DefaultReviewer = reviewer + } + if h.audit != nil { + h.audit.RecordOK(c, "hitl", "default_reviewer_update", "HITL 全局默认审批方更新", "hitl_config", "default_reviewer", nil) + } + c.JSON(http.StatusOK, gin.H{ + "ok": true, + "defaultReviewer": reviewer, + }) +} + +// SetHITLGlobalToolWhitelist 整表替换 config.yaml 中的全局免审批工具白名单。 +func (h *AgentHandler) SetHITLGlobalToolWhitelist(c *gin.Context) { + if h.hitlWhitelistSaver == nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "HITL 配置持久化不可用"}) + return + } + var req setHitlGlobalWhitelistReq + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if err := h.hitlWhitelistSaver.SetHitlToolWhitelist(req.ToolWhitelist); err != nil { + h.logger.Warn("写入 HITL 工具白名单到 config.yaml 失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "hitl", "tool_whitelist_update", "HITL 全局白名单更新", "hitl_config", "tool_whitelist", nil) + } + c.JSON(http.StatusOK, gin.H{ + "ok": true, + "toolWhitelist": h.hitlConfigGlobalToolWhitelist(), + "hitlGlobalToolWhitelist": h.hitlConfigGlobalToolWhitelist(), + "hitlGlobalWhitelistMerged": false, + }) +} + +// MergeHITLGlobalToolWhitelist 无会话 ID 时将侧栏提交的免审批工具合并进 config.yaml（与 PUT /hitl/config 中白名单落盘规则一致）。 +func (h *AgentHandler) MergeHITLGlobalToolWhitelist(c *gin.Context) { + if h.hitlWhitelistSaver == nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "HITL 配置持久化不可用"}) + return + } + var req mergeHitlGlobalWhitelistReq + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if len(req.SensitiveTools) == 0 { + c.JSON(http.StatusOK, gin.H{ + "ok": true, + "hitlGlobalToolWhitelist": h.hitlConfigGlobalToolWhitelist(), + "hitlGlobalWhitelistMerged": false, + }) + return + } + if err := h.hitlWhitelistSaver.MergeHitlToolWhitelistIntoConfig(req.SensitiveTools); err != nil { + h.logger.Warn("合并 HITL 工具白名单到 config.yaml 失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{ + "ok": true, + "hitlGlobalToolWhitelist": h.hitlConfigGlobalToolWhitelist(), + "hitlGlobalWhitelistMerged": true, + }) +} + +func boolToInt(v bool) int { + if v { + return 1 + } + return 0 +} diff --git a/internal/handler/hitl_audit_agent.go b/internal/handler/hitl_audit_agent.go new file mode 100644 index 00000000..8bbd37b9 --- /dev/null +++ b/internal/handler/hitl_audit_agent.go @@ -0,0 +1,360 @@ +package handler + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "strings" + "time" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/openai" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +// auditAgentReview 在 reviewer=audit_agent 时由 LLM 代行审批。 +// 白名单工具在 shouldInterrupt 阶段已跳过，到达此处的一律需要裁决。 +func (h *AgentHandler) auditAgentReview(ctx context.Context, hitlMode, toolName string, payload map[string]interface{}) hitlDecision { + if h == nil { + return hitlDecision{Decision: "reject", Comment: "audit agent: handler unavailable"} + } + mode := normalizeHitlMode(hitlMode) + prompt := config.DefaultHitlAuditAgentPrompt() + if h.config != nil { + prompt = h.config.Hitl.EffectiveAuditAgentPromptForMode(mode) + } + llmCfg := h.auditLLMConfig() + if strings.TrimSpace(llmCfg.APIKey) == "" || strings.TrimSpace(llmCfg.Model) == "" { + return hitlDecision{Decision: "reject", Comment: "audit agent: LLM 未配置"} + } + if ctx == nil { + ctx = context.Background() + } + callCtx, cancel := context.WithTimeout(ctx, 90*time.Second) + defer cancel() + + userContent := buildAuditAgentReviewInput(mode, toolName, payload) + requestBody := map[string]interface{}{ + "model": strings.TrimSpace(llmCfg.Model), + "messages": []map[string]interface{}{ + {"role": "system", "content": prompt}, + {"role": "user", "content": userContent}, + }, + "temperature": 0.1, + "max_completion_tokens": 1024, + // 审计裁决需要结构化 JSON；关闭 thinking 避免 Qwen 等把正文放进 reasoning_content 导致解析失败。 + "thinking": map[string]interface{}{"type": "disabled"}, + } + + var apiResponse struct { + Choices []struct { + Message struct { + Content string `json:"content"` + ReasoningContent string `json:"reasoning_content"` + } `json:"message"` + } `json:"choices"` + } + client := openai.NewClient(&llmCfg, nil, h.logger) + if err := client.ChatCompletion(callCtx, requestBody, &apiResponse); err != nil { + h.logger.Warn("审计 Agent LLM 调用失败", zap.Error(err), zap.String("tool", toolName)) + return hitlDecision{ + Decision: "reject", + Comment: "audit agent: LLM 调用失败，保守拒绝", + } + } + if len(apiResponse.Choices) == 0 { + return hitlDecision{Decision: "reject", Comment: "audit agent: LLM 无有效响应，保守拒绝"} + } + msg := apiResponse.Choices[0].Message + raw := strings.TrimSpace(msg.Content) + if raw == "" { + raw = strings.TrimSpace(msg.ReasoningContent) + } + dec, err := parseAuditAgentLLMContent(raw) + if err != nil { + snippet := raw + if len(snippet) > 240 { + snippet = snippet[:240] + "..." + } + h.logger.Warn("审计 Agent 响应解析失败", + zap.Error(err), + zap.String("tool", toolName), + zap.String("mode", mode), + zap.String("snippet", snippet), + ) + return hitlDecision{Decision: "reject", Comment: "audit agent: 响应无法解析，保守拒绝"} + } + if mode != "review_edit" && len(dec.EditedArguments) > 0 { + h.logger.Warn("审计 Agent 在审批模式下返回 editedArguments，已忽略", + zap.String("tool", toolName), + ) + dec.EditedArguments = nil + } + if dec.Comment == "" { + dec.Comment = "audit agent: " + dec.Decision + } else if !strings.HasPrefix(strings.ToLower(dec.Comment), "audit agent") { + dec.Comment = "audit agent: " + dec.Comment + } + return dec +} + +func (h *AgentHandler) auditLLMConfig() config.OpenAIConfig { + if h != nil && h.config != nil { + return h.config.Hitl.AuditModelEffective(h.config.OpenAI) + } + return config.OpenAIConfig{} +} + +func buildAuditAgentReviewInput(hitlMode, toolName string, payload map[string]interface{}) string { + review := map[string]interface{}{ + "hitlMode": normalizeHitlMode(hitlMode), + "toolName": strings.TrimSpace(toolName), + } + if payload != nil { + for _, k := range []string{"arguments", "argumentsObj", "command", hitlPayloadUserMessage, hitlPayloadThinking, hitlPayloadReasoningChain, hitlPayloadPlanning} { + if v, ok := payload[k]; ok && v != nil && fmt.Sprint(v) != "" { + review[k] = v + } + } + } + b, err := json.MarshalIndent(review, "", " ") + if err != nil { + return fmt.Sprintf(`{"hitlMode":%q,"toolName":%q}`, normalizeHitlMode(hitlMode), toolName) + } + return string(b) +} + +func parseAuditAgentLLMContent(content string) (hitlDecision, error) { + s := strings.TrimSpace(content) + if s == "" { + return hitlDecision{}, errors.New("empty content") + } + for _, candidate := range auditAgentJSONCandidates(s) { + dec, comment, editedArgs, err := parseAuditAgentDecisionObject(candidate) + if err == nil { + return hitlDecision{ + Decision: dec, + Comment: comment, + EditedArguments: editedArgs, + }, nil + } + } + return hitlDecision{}, fmt.Errorf("no valid decision json in response") +} + +func auditAgentJSONCandidates(s string) []string { + out := make([]string, 0, 4) + seen := make(map[string]struct{}) + add := func(c string) { + c = strings.TrimSpace(c) + if c == "" { + return + } + if _, ok := seen[c]; ok { + return + } + seen[c] = struct{}{} + out = append(out, c) + } + add(s) + add(stripMarkdownCodeFence(s)) + if obj := extractFirstJSONObject(s); obj != "" { + add(obj) + } + if obj := extractFirstJSONObject(stripMarkdownCodeFence(s)); obj != "" { + add(obj) + } + return out +} + +func stripMarkdownCodeFence(s string) string { + s = strings.TrimSpace(s) + for _, fence := range []string{"```json", "```JSON", "```"} { + if strings.HasPrefix(s, fence) { + s = strings.TrimPrefix(s, fence) + } + } + s = strings.TrimSuffix(s, "```") + return strings.TrimSpace(s) +} + +func extractFirstJSONObject(s string) string { + start := strings.Index(s, "{") + if start < 0 { + return "" + } + depth := 0 + inStr := false + esc := false + for i := start; i < len(s); i++ { + ch := s[i] + if inStr { + if esc { + esc = false + continue + } + if ch == '\\' { + esc = true + continue + } + if ch == '"' { + inStr = false + } + continue + } + switch ch { + case '"': + inStr = true + case '{': + depth++ + case '}': + depth-- + if depth == 0 { + return s[start : i+1] + } + } + } + return "" +} + +func parseAuditAgentDecisionObject(jsonText string) (decision, comment string, editedArgs map[string]interface{}, err error) { + var parsed map[string]interface{} + if err := json.Unmarshal([]byte(jsonText), &parsed); err != nil { + return "", "", nil, err + } + rawDecision := auditAgentPickString(parsed, "decision", "Decision", "result", "action", "verdict", "决策", "决定") + decision = normalizeAuditAgentDecision(rawDecision) + if decision == "" { + return "", "", nil, fmt.Errorf("missing decision") + } + comment = auditAgentPickString(parsed, "comment", "Comment", "reason", "message", "rationale", "备注", "理由", "说明") + editedArgs = auditAgentPickObject(parsed, "editedArguments", "edited_arguments", "editedArgs") + return decision, strings.TrimSpace(comment), editedArgs, nil +} + +func auditAgentPickString(m map[string]interface{}, keys ...string) string { + for _, k := range keys { + if v, ok := m[k]; ok && v != nil { + s := strings.TrimSpace(fmt.Sprint(v)) + if s != "" { + return s + } + } + } + return "" +} + +func auditAgentPickObject(m map[string]interface{}, keys ...string) map[string]interface{} { + for _, k := range keys { + v, ok := m[k] + if !ok || v == nil { + continue + } + switch t := v.(type) { + case map[string]interface{}: + if len(t) > 0 { + return t + } + case string: + s := strings.TrimSpace(t) + if s == "" || s == "{}" { + continue + } + var obj map[string]interface{} + if err := json.Unmarshal([]byte(s), &obj); err == nil && len(obj) > 0 { + return obj + } + } + } + return nil +} + +func normalizeAuditAgentDecision(v string) string { + d := strings.ToLower(strings.TrimSpace(v)) + switch d { + case "approve", "approved", "pass", "passed", "allow", "allowed", "yes", "ok", "accept", "accepted": + return "approve" + case "reject", "rejected", "deny", "denied", "no", "block", "blocked", "refuse", "refused": + return "reject" + } + switch strings.TrimSpace(v) { + case "通过", "批准", "允许", "同意", "放行": + return "approve" + case "拒绝", "驳回", "禁止", "否决": + return "reject" + } + return "" +} + +type hitlAuditStrategyReq struct { + AuditAgentPrompt string `json:"auditAgentPrompt"` + AuditAgentPromptReviewEdit string `json:"auditAgentPromptReviewEdit"` +} + +func (h *AgentHandler) GetHITLAuditStrategy(c *gin.Context) { + approvalPrompt := config.DefaultHitlAuditAgentPrompt() + reviewEditPrompt := config.DefaultHitlAuditAgentPromptReviewEdit() + approvalCustom := false + reviewEditCustom := false + if h.config != nil { + approvalPrompt = h.config.Hitl.EffectiveAuditAgentPromptForMode("approval") + reviewEditPrompt = h.config.Hitl.EffectiveAuditAgentPromptForMode("review_edit") + approvalCustom = strings.TrimSpace(h.config.Hitl.AuditAgentPrompt) != "" + reviewEditCustom = strings.TrimSpace(h.config.Hitl.AuditAgentPromptReviewEdit) != "" + } + c.JSON(http.StatusOK, gin.H{ + "auditAgentPrompt": approvalPrompt, + "auditAgentPromptCustom": approvalCustom, + "auditAgentPromptReviewEdit": reviewEditPrompt, + "auditAgentPromptReviewEditCustom": reviewEditCustom, + "defaultAuditAgentPrompt": config.DefaultHitlAuditAgentPrompt(), + "defaultAuditAgentPromptReviewEdit": config.DefaultHitlAuditAgentPromptReviewEdit(), + }) +} + +func (h *AgentHandler) UpdateHITLAuditStrategy(c *gin.Context) { + if h.hitlStrategySaver == nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "HITL 策略持久化不可用"}) + return + } + var req hitlAuditStrategyReq + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + approvalPrompt := strings.TrimSpace(req.AuditAgentPrompt) + reviewEditPrompt := strings.TrimSpace(req.AuditAgentPromptReviewEdit) + if err := h.hitlStrategySaver.UpdateHitlAuditAgentStrategy(approvalPrompt, reviewEditPrompt); err != nil { + h.logger.Warn("保存审计 Agent 提示词失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "hitl", "audit_strategy_update", "HITL 审计策略更新", "hitl_config", "audit_agent_prompt", nil) + } + if h.config != nil { + h.config.Hitl.AuditAgentPrompt = approvalPrompt + h.config.Hitl.AuditAgentPromptReviewEdit = reviewEditPrompt + } + c.JSON(http.StatusOK, gin.H{ + "ok": true, + "auditAgentPrompt": config.HitlConfig{AuditAgentPrompt: approvalPrompt}.EffectiveAuditAgentPromptForMode("approval"), + "auditAgentPromptCustom": approvalPrompt != "", + "auditAgentPromptReviewEdit": config.HitlConfig{AuditAgentPromptReviewEdit: reviewEditPrompt}.EffectiveAuditAgentPromptForMode("review_edit"), + "auditAgentPromptReviewEditCustom": reviewEditPrompt != "", + }) +} + +// HitlAuditStrategySaver 持久化审计 Agent 提示词到 config.yaml。 +type HitlAuditStrategySaver interface { + UpdateHitlAuditAgentStrategy(approvalPrompt, reviewEditPrompt string) error +} + +// SetHitlAuditStrategySaver 设置审计策略落盘。 +func (h *AgentHandler) SetHitlAuditStrategySaver(s HitlAuditStrategySaver) { + h.hitlStrategySaver = s +} diff --git a/internal/handler/hitl_audit_agent_test.go b/internal/handler/hitl_audit_agent_test.go new file mode 100644 index 00000000..8a7d9a4c --- /dev/null +++ b/internal/handler/hitl_audit_agent_test.go @@ -0,0 +1,88 @@ +package handler + +import ( + "strings" + "testing" +) + +func TestParseAuditAgentLLMContentApprove(t *testing.T) { + d, err := parseAuditAgentLLMContent(`{"decision":"approve","comment":"与任务一致"}`) + if err != nil { + t.Fatal(err) + } + if d.Decision != "approve" || d.Comment != "与任务一致" { + t.Fatalf("unexpected %+v", d) + } +} + +func TestParseAuditAgentLLMContentReject(t *testing.T) { + d, err := parseAuditAgentLLMContent("```json\n{\"decision\":\"reject\",\"comment\":\"风险过高\"}\n```") + if err != nil { + t.Fatal(err) + } + if d.Decision != "reject" { + t.Fatalf("expected reject, got %s", d.Decision) + } +} + +func TestParseAuditAgentLLMContentInvalid(t *testing.T) { + _, err := parseAuditAgentLLMContent(`{"decision":"maybe"}`) + if err == nil { + t.Fatal("expected error for invalid decision") + } +} + +func TestParseAuditAgentLLMContentProseWrapped(t *testing.T) { + d, err := parseAuditAgentLLMContent("好的，裁决如下：\n```json\n{\"decision\":\"approve\",\"comment\":\"只读 ls\"}\n```\n以上。") + if err != nil { + t.Fatal(err) + } + if d.Decision != "approve" { + t.Fatalf("expected approve, got %s", d.Decision) + } +} + +func TestParseAuditAgentLLMContentChineseDecision(t *testing.T) { + d, err := parseAuditAgentLLMContent(`{"decision":"通过","comment":"风险低"}`) + if err != nil { + t.Fatal(err) + } + if d.Decision != "approve" { + t.Fatalf("expected approve, got %s", d.Decision) + } +} + +func TestParseAuditAgentLLMContentWithEditedArguments(t *testing.T) { + d, err := parseAuditAgentLLMContent(`{"decision":"approve","comment":"收窄路径","editedArguments":{"path":"/safe"}}`) + if err != nil { + t.Fatal(err) + } + if d.Decision != "approve" { + t.Fatalf("expected approve, got %s", d.Decision) + } + if d.EditedArguments == nil || d.EditedArguments["path"] != "/safe" { + t.Fatalf("unexpected edited args: %+v", d.EditedArguments) + } +} + +func TestBuildAuditAgentReviewInputIncludesMode(t *testing.T) { + s := buildAuditAgentReviewInput("review_edit", "execute", map[string]interface{}{ + "arguments": `{"command":"pwd"}`, + }) + if !strings.Contains(s, "review_edit") || !strings.Contains(s, "execute") { + t.Fatalf("unexpected input: %s", s) + } +} + +func TestBuildAuditAgentReviewInput(t *testing.T) { + s := buildAuditAgentReviewInput("approval", "nmap", map[string]interface{}{ + "arguments": `{"target":"10.0.0.1"}`, + "userMessage": "扫描内网", + }) + if s == "" { + t.Fatal("expected non-empty input") + } + if !strings.Contains(s, "nmap") || !strings.Contains(s, "10.0.0.1") || !strings.Contains(s, "扫描内网") { + t.Fatalf("unexpected input: %s", s) + } +} diff --git a/internal/handler/hitl_cognition.go b/internal/handler/hitl_cognition.go new file mode 100644 index 00000000..6b24cb57 --- /dev/null +++ b/internal/handler/hitl_cognition.go @@ -0,0 +1,97 @@ +package handler + +import ( + "strings" +) + +type hitlCognitionState struct { + AssistantMessageID string + UserMessage string + Thinking string + ReasoningChain string + Planning string +} + +// GetHitlCognition 返回当前运行任务上缓存的本轮 HITL 上下文（不含会话历史）。 +func (m *AgentTaskManager) GetHitlCognition(conversationID string) hitlCognitionFields { + conversationID = strings.TrimSpace(conversationID) + if m == nil || conversationID == "" { + return hitlCognitionFields{} + } + m.mu.RLock() + defer m.mu.RUnlock() + t, ok := m.tasks[conversationID] + if !ok || t == nil || t.hitlCognition == nil { + return hitlCognitionFields{} + } + c := t.hitlCognition + return hitlCognitionFields{ + UserMessage: c.UserMessage, + Thinking: c.Thinking, + ReasoningChain: c.ReasoningChain, + Planning: c.Planning, + } +} + +// ResetHitlCognition 新任务开始时重置本轮 HITL 上下文。 +func (m *AgentTaskManager) ResetHitlCognition(conversationID, userMessage string) { + conversationID = strings.TrimSpace(conversationID) + if m == nil || conversationID == "" { + return + } + m.mu.Lock() + defer m.mu.Unlock() + t, ok := m.tasks[conversationID] + if !ok || t == nil { + return + } + t.hitlCognition = &hitlCognitionState{UserMessage: strings.TrimSpace(userMessage)} +} + +// SetHitlAssistantMessageID 记录当前助手消息 ID，供 HITL 与 DB 回退对齐。 +func (m *AgentTaskManager) SetHitlAssistantMessageID(conversationID, assistantMessageID string) { + conversationID = strings.TrimSpace(conversationID) + assistantMessageID = strings.TrimSpace(assistantMessageID) + if m == nil || conversationID == "" || assistantMessageID == "" { + return + } + m.mu.Lock() + defer m.mu.Unlock() + t, ok := m.tasks[conversationID] + if !ok || t == nil { + return + } + if t.hitlCognition == nil { + t.hitlCognition = &hitlCognitionState{} + } + t.hitlCognition.AssistantMessageID = assistantMessageID +} + +// UpdateHitlCognitionSnapshot 从进行中的进度流快照更新 thinking / reasoning / planning。 +func (m *AgentTaskManager) UpdateHitlCognitionSnapshot(conversationID, assistantMessageID, thinking, reasoningChain, planning string) { + conversationID = strings.TrimSpace(conversationID) + if m == nil || conversationID == "" { + return + } + m.mu.Lock() + defer m.mu.Unlock() + t, ok := m.tasks[conversationID] + if !ok || t == nil { + return + } + if t.hitlCognition == nil { + t.hitlCognition = &hitlCognitionState{} + } + if id := strings.TrimSpace(assistantMessageID); id != "" { + t.hitlCognition.AssistantMessageID = id + } + if s := strings.TrimSpace(thinking); s != "" { + t.hitlCognition.Thinking = s + } + if s := strings.TrimSpace(reasoningChain); s != "" { + t.hitlCognition.ReasoningChain = s + } + if s := strings.TrimSpace(planning); s != "" { + t.hitlCognition.Planning = s + } +} diff --git a/internal/handler/hitl_context.go b/internal/handler/hitl_context.go new file mode 100644 index 00000000..e8d551ad --- /dev/null +++ b/internal/handler/hitl_context.go @@ -0,0 +1,102 @@ +package handler + +import ( + "strings" +) + +const ( + hitlPayloadUserMessage = "userMessage" + hitlPayloadThinking = "thinking" + hitlPayloadReasoningChain = "reasoningChain" + hitlPayloadPlanning = "planning" +) + +type hitlCognitionFields struct { + UserMessage string + Thinking string + ReasoningChain string + Planning string +} + +func (h *AgentHandler) enrichHitlApprovalPayload(conversationID, assistantMessageID string, payload map[string]interface{}) { + if h == nil || payload == nil { + return + } + cog := h.collectHitlCognition(conversationID, assistantMessageID) + if s := strings.TrimSpace(cog.UserMessage); s != "" { + payload[hitlPayloadUserMessage] = s + } + if s := strings.TrimSpace(cog.Thinking); s != "" { + payload[hitlPayloadThinking] = s + } + if s := strings.TrimSpace(cog.ReasoningChain); s != "" { + payload[hitlPayloadReasoningChain] = s + } + if s := strings.TrimSpace(cog.Planning); s != "" { + payload[hitlPayloadPlanning] = s + } +} + +func (h *AgentHandler) collectHitlCognition(conversationID, assistantMessageID string) hitlCognitionFields { + var out hitlCognitionFields + if h.tasks != nil { + out = h.tasks.GetHitlCognition(conversationID) + } + if strings.TrimSpace(out.UserMessage) == "" && h.db != nil { + if msg, err := h.db.GetTurnUserMessage(conversationID, assistantMessageID); err == nil { + out.UserMessage = msg + } + } + if h.db != nil && assistantMessageID != "" { + dbCog, err := h.db.GetAssistantCognitionTexts(assistantMessageID) + if err == nil { + if strings.TrimSpace(out.Thinking) == "" { + out.Thinking = dbCog.Thinking + } + if strings.TrimSpace(out.ReasoningChain) == "" { + out.ReasoningChain = dbCog.ReasoningChain + } + if strings.TrimSpace(out.Planning) == "" { + out.Planning = dbCog.Planning + } + } + } + return out +} + +func snapshotHitlCognitionFromStreams(thinkingStreams map[string]*thinkingBuf, respPlan *responsePlanAgg) (thinking, reasoningChain, planning string) { + if len(thinkingStreams) > 0 { + var thinkingParts, reasoningParts []string + for _, tb := range thinkingStreams { + if tb == nil { + continue + } + content := strings.TrimSpace(tb.b.String()) + if content == "" { + continue + } + if tb.persistAs == "reasoning_chain" { + reasoningParts = append(reasoningParts, content) + } else { + thinkingParts = append(thinkingParts, content) + } + } + thinking = strings.Join(thinkingParts, "\n\n") + reasoningChain = strings.Join(reasoningParts, "\n\n") + } + if respPlan != nil { + planning = strings.TrimSpace(respPlan.b.String()) + } + return thinking, reasoningChain, planning +} + +func (h *AgentHandler) syncHitlCognitionFromProgress(conversationID, assistantMessageID string, thinkingStreams map[string]*thinkingBuf, respPlan *responsePlanAgg) { + if h == nil || h.tasks == nil { + return + } + thinking, reasoning, planning := snapshotHitlCognitionFromStreams(thinkingStreams, respPlan) + if thinking == "" && reasoning == "" && planning == "" { + return + } + h.tasks.UpdateHitlCognitionSnapshot(conversationID, assistantMessageID, thinking, reasoning, planning) +} diff --git a/internal/handler/hitl_context_test.go b/internal/handler/hitl_context_test.go new file mode 100644 index 00000000..cdf3870c --- /dev/null +++ b/internal/handler/hitl_context_test.go @@ -0,0 +1,46 @@ +package handler + +import ( + "os" + "path/filepath" + "testing" + + "cyberstrike-ai/internal/database" + + "go.uber.org/zap" +) + +func TestEnrichHitlApprovalPayload(t *testing.T) { + tmp := t.TempDir() + db, err := database.NewDB(filepath.Join(tmp, "test.sqlite"), zap.NewNop()) + if err != nil { + t.Fatalf("db: %v", err) + } + defer os.RemoveAll(tmp) + + conv, err := db.CreateConversation("hitl ctx", database.ConversationCreateMeta{}) + if err != nil { + t.Fatalf("conv: %v", err) + } + if _, err := db.AddMessage(conv.ID, "user", "scan 10.0.0.1 please", nil); err != nil { + t.Fatalf("user msg: %v", err) + } + asst, err := db.AddMessage(conv.ID, "assistant", "", nil) + if err != nil { + t.Fatalf("asst msg: %v", err) + } + if err := db.AddProcessDetail(asst.ID, conv.ID, "thinking", "need port scan first", nil); err != nil { + t.Fatalf("detail: %v", err) + } + + h := &AgentHandler{db: db, tasks: NewAgentTaskManager()} + payload := map[string]interface{}{"toolName": "nmap", "arguments": "{}"} + h.enrichHitlApprovalPayload(conv.ID, asst.ID, payload) + + if got := payload["userMessage"]; got != "scan 10.0.0.1 please" { + t.Fatalf("userMessage=%v", got) + } + if got := payload["thinking"]; got != "need port scan first" { + t.Fatalf("thinking=%v", got) + } +} diff --git a/internal/handler/hitl_execution.go b/internal/handler/hitl_execution.go new file mode 100644 index 00000000..8d44b6d1 --- /dev/null +++ b/internal/handler/hitl_execution.go @@ -0,0 +1,132 @@ +package handler + +import ( + "encoding/json" + "strings" + "time" +) + +const hitlPayloadExecutionResult = "executionResult" + +type hitlExecutionResult struct { + Success bool `json:"success"` + Result string `json:"result,omitempty"` + ToolName string `json:"toolName,omitempty"` + ToolCallID string `json:"toolCallId,omitempty"` + RecordedAt time.Time `json:"recordedAt"` +} + +type hitlApprovedExecTrack struct { + InterruptID string + ConversationID string + ToolName string + ToolCallID string +} + +// TrackApprovedHitlExecution 审批通过后登记，待 tool_result 回写执行结果。 +func (m *HITLManager) TrackApprovedHitlExecution(interruptID, conversationID, toolName, toolCallID string) { + if m == nil { + return + } + interruptID = strings.TrimSpace(interruptID) + conversationID = strings.TrimSpace(conversationID) + if interruptID == "" || conversationID == "" { + return + } + m.mu.Lock() + defer m.mu.Unlock() + if m.approvedExec == nil { + m.approvedExec = make(map[string][]hitlApprovedExecTrack) + } + m.approvedExec[conversationID] = append(m.approvedExec[conversationID], hitlApprovedExecTrack{ + InterruptID: interruptID, + ConversationID: conversationID, + ToolName: strings.TrimSpace(toolName), + ToolCallID: strings.TrimSpace(toolCallID), + }) +} + +func (m *HITLManager) popApprovedInterruptForTool(conversationID, toolCallID, toolName string) string { + if m == nil { + return "" + } + conversationID = strings.TrimSpace(conversationID) + toolCallID = strings.TrimSpace(toolCallID) + toolName = strings.TrimSpace(toolName) + m.mu.Lock() + defer m.mu.Unlock() + queue := m.approvedExec[conversationID] + if len(queue) == 0 { + return "" + } + idx := -1 + if toolCallID != "" { + for i, t := range queue { + if t.ToolCallID == toolCallID { + idx = i + break + } + } + } + if idx < 0 && toolName != "" { + for i, t := range queue { + if strings.EqualFold(t.ToolName, toolName) { + idx = i + break + } + } + } + if idx < 0 { + return "" + } + id := queue[idx].InterruptID + queue = append(queue[:idx], queue[idx+1:]...) + if len(queue) == 0 { + delete(m.approvedExec, conversationID) + } else { + m.approvedExec[conversationID] = queue + } + return id +} + +func mergeHitlPayloadExecutionResult(payloadJSON string, exec hitlExecutionResult) (string, error) { + root := make(map[string]interface{}) + if strings.TrimSpace(payloadJSON) != "" { + _ = json.Unmarshal([]byte(payloadJSON), &root) + } + if root == nil { + root = make(map[string]interface{}) + } + root[hitlPayloadExecutionResult] = exec + out, err := json.Marshal(root) + if err != nil { + return payloadJSON, err + } + return string(out), nil +} + +func (h *AgentHandler) recordHitlToolExecutionResult(conversationID, toolCallID, toolName string, success bool, result string) { + if h == nil || h.hitlManager == nil || h.db == nil { + return + } + interruptID := h.hitlManager.popApprovedInterruptForTool(conversationID, toolCallID, toolName) + if interruptID == "" { + return + } + var payloadJSON string + err := h.db.QueryRow(`SELECT payload FROM hitl_interrupts WHERE id = ?`, interruptID).Scan(&payloadJSON) + if err != nil { + return + } + merged, err := mergeHitlPayloadExecutionResult(payloadJSON, hitlExecutionResult{ + Success: success, + Result: strings.TrimSpace(result), + ToolName: strings.TrimSpace(toolName), + ToolCallID: strings.TrimSpace(toolCallID), + RecordedAt: time.Now(), + }) + if err != nil { + return + } + _, _ = h.db.Exec(`UPDATE hitl_interrupts SET payload = ? WHERE id = ?`, merged, interruptID) +} diff --git a/internal/handler/hitl_execution_test.go b/internal/handler/hitl_execution_test.go new file mode 100644 index 00000000..1c620366 --- /dev/null +++ b/internal/handler/hitl_execution_test.go @@ -0,0 +1,39 @@ +package handler + +import ( + "encoding/json" + "testing" +) + +func TestMergeHitlPayloadExecutionResult(t *testing.T) { + merged, err := mergeHitlPayloadExecutionResult(`{"userMessage":"hi","toolName":"nmap"}`, hitlExecutionResult{ + Success: true, + Result: "open ports: 80", + }) + if err != nil { + t.Fatal(err) + } + var root map[string]interface{} + if err := json.Unmarshal([]byte(merged), &root); err != nil { + t.Fatal(err) + } + if root["userMessage"] != "hi" { + t.Fatalf("userMessage lost: %v", root["userMessage"]) + } + exec, ok := root["executionResult"].(map[string]interface{}) + if !ok || exec["success"] != true { + t.Fatalf("executionResult missing: %v", root["executionResult"]) + } +} + +func TestPopApprovedInterruptForTool(t *testing.T) { + m := NewHITLManager(nil, nil) + m.TrackApprovedHitlExecution("hitl_a", "conv1", "nmap", "tc1") + m.TrackApprovedHitlExecution("hitl_b", "conv1", "exec", "") + if id := m.popApprovedInterruptForTool("conv1", "tc1", "nmap"); id != "hitl_a" { + t.Fatalf("tc1 match=%q", id) + } + if id := m.popApprovedInterruptForTool("conv1", "", "exec"); id != "hitl_b" { + t.Fatalf("tool name match=%q", id) + } +} diff --git a/internal/handler/hitl_logs.go b/internal/handler/hitl_logs.go new file mode 100644 index 00000000..a6d787e0 --- /dev/null +++ b/internal/handler/hitl_logs.go @@ -0,0 +1,337 @@ +package handler + +import ( + "database/sql" + "errors" + "math" + "net/http" + "strconv" + "strings" + "time" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/security" + + "github.com/gin-gonic/gin" +) + +func normalizeHitlReviewer(v string) string { + switch strings.ToLower(strings.TrimSpace(v)) { + case "audit_agent", "agent", "ai": + return "audit_agent" + default: + return "human" + } +} + +func normalizeHitlDecidedBy(v string) string { + switch strings.ToLower(strings.TrimSpace(v)) { + case "audit_agent", "agent", "ai": + return "audit_agent" + case "system", "timeout": + return "system" + case "manual": + return "manual" + default: + return "human" + } +} + +func (m *HITLManager) migrateHitlSchemaColumns() { + _, _ = m.db.Exec(`ALTER TABLE hitl_interrupts ADD COLUMN decided_by TEXT NOT NULL DEFAULT 'human'`) + _, _ = m.db.Exec(`ALTER TABLE hitl_interrupts ADD COLUMN reviewer TEXT NOT NULL DEFAULT 'human'`) + _, _ = m.db.Exec(`UPDATE hitl_interrupts SET reviewer='audit_agent' + WHERE COALESCE(decided_by, '') IN ('audit_agent', 'agent', 'ai')`) + _, _ = m.db.Exec(`ALTER TABLE hitl_conversation_configs ADD COLUMN reviewer TEXT NOT NULL DEFAULT 'human'`) +} + +func hitlInterruptRowToMap( + id, cid, mode, toolName, toolCallID, payload, rowStatus, reviewer, decidedBy string, + messageID sql.NullString, + decision, comment sql.NullString, + createdAt time.Time, + decidedAt sql.NullTime, +) map[string]interface{} { + msgID := "" + if messageID.Valid { + msgID = messageID.String + } + return map[string]interface{}{ + "id": id, + "conversationId": cid, + "messageId": msgID, + "mode": mode, + "toolName": toolName, + "toolCallId": toolCallID, + "payload": payload, + "status": rowStatus, + "reviewer": reviewer, + "decision": decision.String, + "comment": comment.String, + "decidedBy": decidedBy, + "createdAt": createdAt, + "decidedAt": func() interface{} { + if decidedAt.Valid { + return decidedAt.Time + } + return nil + }(), + } +} + +func (h *AgentHandler) buildHitlListQuery(logs bool) (string, []interface{}) { + where, args := h.buildHitlLogsWhere(logs) + q := `SELECT id, conversation_id, message_id, mode, tool_name, tool_call_id, payload, status, COALESCE(reviewer,'human'), decision, decision_comment, COALESCE(decided_by,'human'), created_at, decided_at FROM hitl_interrupts` + where + return q, args +} + +func (h *AgentHandler) buildHitlLogsWhere(logs bool) (string, []interface{}) { + q := " WHERE 1=1" + args := []interface{}{} + if logs { + q += " AND status != 'pending'" + } else { + // 该接口只返回真正等待用户操作的人工审批。Agent 审查即使正在运行， + // 也不应触发弹窗、倒计时或项目待审批计数。 + q += " AND status = 'pending' AND COALESCE(reviewer,'human') = 'human'" + } + return q, args +} + +func (h *AgentHandler) appendHitlListFilters(q string, args []interface{}, c *gin.Context) (string, []interface{}) { + conversationID := strings.TrimSpace(c.Query("conversationId")) + toolName := strings.TrimSpace(c.Query("toolName")) + decision := strings.TrimSpace(c.Query("decision")) + decidedBy := strings.TrimSpace(c.Query("decidedBy")) + status := strings.TrimSpace(c.Query("status")) + search := strings.TrimSpace(c.Query("q")) + + if conversationID != "" { + q += " AND conversation_id = ?" + args = append(args, conversationID) + } + if toolName != "" { + q += " AND tool_name LIKE ?" + args = append(args, "%"+toolName+"%") + } + if decision != "" && decision != "all" { + q += " AND decision = ?" + args = append(args, decision) + } + if decidedBy != "" && decidedBy != "all" { + q += " AND COALESCE(decided_by,'human') = ?" + args = append(args, normalizeHitlDecidedBy(decidedBy)) + } + if status != "" && status != "all" { + q += " AND status = ?" + args = append(args, status) + } + if search != "" { + like := "%" + search + "%" + q += " AND (id LIKE ? OR conversation_id LIKE ? OR tool_name LIKE ? OR payload LIKE ? OR COALESCE(decision_comment,'') LIKE ?)" + args = append(args, like, like, like, like, like) + } + return q, args +} + +func (h *AgentHandler) scanHitlInterruptRows(rows *sql.Rows) ([]map[string]interface{}, error) { + items := make([]map[string]interface{}, 0) + for rows.Next() { + var id, cid, mode, toolName, toolCallID, payload, rowStatus, reviewer, decidedBy string + var messageID sql.NullString + var decision, comment sql.NullString + var createdAt time.Time + var decidedAt sql.NullTime + if err := rows.Scan(&id, &cid, &messageID, &mode, &toolName, &toolCallID, &payload, &rowStatus, &reviewer, &decision, &comment, &decidedBy, &createdAt, &decidedAt); err != nil { + continue + } + items = append(items, hitlInterruptRowToMap(id, cid, mode, toolName, toolCallID, payload, rowStatus, reviewer, decidedBy, messageID, decision, comment, createdAt, decidedAt)) + } + return items, nil +} + +func (h *AgentHandler) countHitlQuery(baseQ string, args []interface{}) (int, error) { + countQ := "SELECT COUNT(*) FROM (" + baseQ + ") AS hitl_cnt" + var total int + if err := h.db.QueryRow(countQ, args...).Scan(&total); err != nil { + return 0, err + } + return total, nil +} + +func (h *AgentHandler) ListHITLLogs(c *gin.Context) { + page, _ := strconv.Atoi(c.DefaultQuery("page", "1")) + if page < 1 { + page = 1 + } + pageSize, _ := strconv.Atoi(c.DefaultQuery("pageSize", "20")) + pageSize = int(math.Max(1, math.Min(float64(pageSize), 200))) + offset := (page - 1) * pageSize + + q, args := h.buildHitlListQuery(true) + q, args = h.appendHitlListFilters(q, args, c) + q, args = appendConversationAccessSQL(q, args, "conversation_id", notificationAccessFromContext(c)) + total, err := h.countHitlQuery(q, args) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + q += " ORDER BY COALESCE(decided_at, created_at) DESC LIMIT ? OFFSET ?" + args = append(args, pageSize, offset) + rows, err := h.db.Query(q, args...) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + defer rows.Close() + items, err := h.scanHitlInterruptRows(rows) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"items": items, "page": page, "pageSize": pageSize, "total": total, "retentionDays": h.hitlRetentionDays()}) +} + +func (h *AgentHandler) hitlRetentionDays() int { + if h.config != nil { + return h.config.Hitl.RetentionDaysEffective() + } + return config.HitlConfig{}.RetentionDaysEffective() +} + +// DeleteHITLLogs 批量删除或按筛选清空已决策的人机协同审计日志（不删除 pending）。 +func (h *AgentHandler) DeleteHITLLogs(c *gin.Context) { + var request struct { + IDs []string `json:"ids"` + All bool `json:"all"` + } + if err := c.ShouldBindJSON(&request); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "请求参数无效: " + err.Error()}) + return + } + + var deleted int64 + var err error + if request.All { + where, args := h.buildHitlLogsWhere(true) + where, args = h.appendHitlListFilters(where, args, c) + where, args = appendConversationAccessSQL(where, args, "conversation_id", notificationAccessFromContext(c)) + deleted, err = h.db.DeleteHitlInterruptLogsMatching(where, args) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "hitl", "logs_clear", "清空人机协同审计日志", "hitl_interrupt", "", map[string]interface{}{ + "deleted": deleted, + }) + } + } else { + if len(request.IDs) == 0 { + c.JSON(http.StatusBadRequest, gin.H{"error": "审计日志 ID 列表不能为空"}) + return + } + ids, filterErr := h.filterAllowedHitlInterruptIDs(c, request.IDs) + if filterErr != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": filterErr.Error()}) + return + } + deleted, err = h.db.DeleteHitlInterruptLogsByIDs(ids) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "hitl", "logs_delete_batch", "批量删除人机协同审计日志", "hitl_interrupt", "", map[string]interface{}{ + "count": len(request.IDs), + "deleted": deleted, + }) + } + } + + c.JSON(http.StatusOK, gin.H{"message": "删除成功", "deleted": deleted}) +} + +func (h *AgentHandler) GetHITLLog(c *gin.Context) { + id := strings.TrimSpace(c.Param("id")) + if id == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "id is required"}) + return + } + q := `SELECT id, conversation_id, message_id, mode, tool_name, tool_call_id, payload, status, COALESCE(reviewer,'human'), decision, decision_comment, COALESCE(decided_by,'human'), created_at, decided_at FROM hitl_interrupts WHERE id = ?` + var rowID, cid, mode, toolName, toolCallID, payload, rowStatus, reviewer, decidedBy string + var messageID sql.NullString + var decision, comment sql.NullString + var createdAt time.Time + var decidedAt sql.NullTime + err := h.db.QueryRow(q, id).Scan(&rowID, &cid, &messageID, &mode, &toolName, &toolCallID, &payload, &rowStatus, &reviewer, &decision, &comment, &decidedBy, &createdAt, &decidedAt) + if errors.Is(err, sql.ErrNoRows) { + c.JSON(http.StatusNotFound, gin.H{"error": "not found"}) + return + } + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if !h.hitlConversationAllowed(c, cid) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + c.JSON(http.StatusOK, hitlInterruptRowToMap(rowID, cid, mode, toolName, toolCallID, payload, rowStatus, reviewer, decidedBy, messageID, decision, comment, createdAt, decidedAt)) +} + +func (h *AgentHandler) filterAllowedHitlInterruptIDs(c *gin.Context, ids []string) ([]string, error) { + clean := make([]string, 0, len(ids)) + seen := map[string]struct{}{} + for _, id := range ids { + id = strings.TrimSpace(id) + if id == "" { + continue + } + if _, ok := seen[id]; ok { + continue + } + seen[id] = struct{}{} + clean = append(clean, id) + } + if len(clean) == 0 { + return clean, nil + } + query := `SELECT id, conversation_id FROM hitl_interrupts WHERE id IN (` + buildPlaceholders(len(clean)) + `)` + args := make([]interface{}, 0, len(clean)) + for _, id := range clean { + args = append(args, id) + } + rows, err := h.db.Query(query, args...) + if err != nil { + return nil, err + } + defer rows.Close() + allowed := make([]string, 0, len(clean)) + for rows.Next() { + var id, conversationID string + if err := rows.Scan(&id, &conversationID); err != nil { + continue + } + if h.hitlConversationAllowed(c, conversationID) { + allowed = append(allowed, id) + } + } + return allowed, rows.Err() +} + +func (h *AgentHandler) hitlInterruptAllowed(c *gin.Context, interruptID string) bool { + var conversationID string + if err := h.db.QueryRow(`SELECT conversation_id FROM hitl_interrupts WHERE id = ?`, strings.TrimSpace(interruptID)).Scan(&conversationID); err != nil { + return false + } + return h.hitlConversationAllowed(c, conversationID) +} + +func (h *AgentHandler) hitlConversationAllowed(c *gin.Context, conversationID string) bool { + session, ok := security.CurrentSession(c) + if !ok { + return false + } + return h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", conversationID) +} diff --git a/internal/handler/hitl_restart_test.go b/internal/handler/hitl_restart_test.go new file mode 100644 index 00000000..f5ec1475 --- /dev/null +++ b/internal/handler/hitl_restart_test.go @@ -0,0 +1,236 @@ +package handler + +import ( + "database/sql" + "path/filepath" + "strings" + "testing" + + "cyberstrike-ai/internal/database" + + "go.uber.org/zap" +) + +func TestEnsureSchemaCancelsPendingInterruptsAfterRestart(t *testing.T) { + db, err := database.NewDB(filepath.Join(t.TempDir(), "hitl-restart.db"), zap.NewNop()) + if err != nil { + t.Fatalf("open db: %v", err) + } + defer func() { _ = db.Close() }() + manager := NewHITLManager(db, zap.NewNop()) + if err := manager.EnsureSchema(); err != nil { + t.Fatalf("ensure schema: %v", err) + } + conversation, err := db.CreateConversation("restart interrupted", database.ConversationCreateMeta{}) + if err != nil { + t.Fatalf("create conversation: %v", err) + } + message, err := db.AddMessage(conversation.ID, "assistant", "处理中...", nil) + if err != nil { + t.Fatalf("create assistant placeholder: %v", err) + } + if _, err := db.Exec(`INSERT INTO hitl_interrupts + (id, conversation_id, message_id, mode, tool_name, tool_call_id, payload, status, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, 'pending', CURRENT_TIMESTAMP)`, + "restart-pending", conversation.ID, message.ID, "approval", "browser", "tool-call-1", `{}`); err != nil { + t.Fatalf("insert pending interrupt: %v", err) + } + + if err := manager.EnsureSchema(); err != nil { + t.Fatalf("reconcile restart: %v", err) + } + + var status, decision, comment, decidedBy string + var decidedAt sql.NullTime + if err := db.QueryRow(`SELECT status, decision, decision_comment, decided_by, decided_at + FROM hitl_interrupts WHERE id = ?`, "restart-pending"). + Scan(&status, &decision, &comment, &decidedBy, &decidedAt); err != nil { + t.Fatalf("query reconciled interrupt: %v", err) + } + if status != "cancelled" || decision != "reject" || comment != "process restarted" { + t.Fatalf("unexpected restart decision: status=%q decision=%q comment=%q", status, decision, comment) + } + if decidedBy != "system" { + t.Fatalf("decided_by=%q, want system", decidedBy) + } + if !decidedAt.Valid { + t.Fatal("decided_at should be set after restart reconciliation") + } + + var content string + var updatedAt sql.NullTime + if err := db.QueryRow(`SELECT content, updated_at FROM messages WHERE id = ?`, message.ID). + Scan(&content, &updatedAt); err != nil { + t.Fatalf("query reconciled assistant message: %v", err) + } + if content != "任务因服务重启已中断，审批已取消。" { + t.Fatalf("assistant content=%q, want restart interruption notice", content) + } + if !updatedAt.Valid { + t.Fatal("assistant updated_at should be set to the interruption time") + } + var eventType, eventMessage string + if err := db.QueryRow(`SELECT event_type, message FROM process_details WHERE message_id = ?`, message.ID). + Scan(&eventType, &eventMessage); err != nil { + t.Fatalf("query restart cancellation process detail: %v", err) + } + if eventType != "cancelled" || eventMessage != content { + t.Fatalf("unexpected terminal detail: type=%q message=%q", eventType, eventMessage) + } +} + +func TestEnsureSchemaFinalizesOnlyHistoricalPlaceholdersWithTerminalEvidence(t *testing.T) { + db, err := database.NewDB(filepath.Join(t.TempDir(), "hitl-history.db"), zap.NewNop()) + if err != nil { + t.Fatalf("open db: %v", err) + } + defer func() { _ = db.Close() }() + manager := NewHITLManager(db, zap.NewNop()) + if err := manager.EnsureSchema(); err != nil { + t.Fatalf("ensure schema: %v", err) + } + + supersededConversation, err := db.CreateConversation("superseded placeholder", database.ConversationCreateMeta{}) + if err != nil { + t.Fatalf("create superseded conversation: %v", err) + } + superseded, err := db.AddMessage(supersededConversation.ID, "assistant", "处理中...", nil) + if err != nil { + t.Fatalf("create superseded placeholder: %v", err) + } + if _, err := db.AddMessage(supersededConversation.ID, "user", "继续", nil); err != nil { + t.Fatalf("create later message: %v", err) + } + + timeoutConversation, err := db.CreateConversation("timeout placeholder", database.ConversationCreateMeta{}) + if err != nil { + t.Fatalf("create timeout conversation: %v", err) + } + timedOut, err := db.AddMessage(timeoutConversation.ID, "assistant", "处理中...", nil) + if err != nil { + t.Fatalf("create timeout placeholder: %v", err) + } + if _, err := db.Exec(`INSERT INTO hitl_interrupts + (id, conversation_id, message_id, mode, tool_name, status, decision, decision_comment, created_at, decided_at) + VALUES (?, ?, ?, 'approval', 'browser', 'timeout', 'reject', 'HITL timeout auto-reject for safety', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)`, + "timeout-interrupt", timeoutConversation.ID, timedOut.ID); err != nil { + t.Fatalf("insert timeout interrupt: %v", err) + } + + rejectedConversation, err := db.CreateConversation("rejected placeholder", database.ConversationCreateMeta{}) + if err != nil { + t.Fatalf("create rejected conversation: %v", err) + } + rejected, err := db.AddMessage(rejectedConversation.ID, "assistant", "处理中...", nil) + if err != nil { + t.Fatalf("create rejected placeholder: %v", err) + } + if _, err := db.Exec(`INSERT INTO hitl_interrupts + (id, conversation_id, message_id, mode, tool_name, status, decision, decision_comment, created_at, decided_at) + VALUES (?, ?, ?, 'approval', 'exec', 'decided', 'reject', 'user rejected', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)`, + "rejected-interrupt", rejectedConversation.ID, rejected.ID); err != nil { + t.Fatalf("insert rejected interrupt: %v", err) + } + + activeConversation, err := db.CreateConversation("potentially active placeholder", database.ConversationCreateMeta{}) + if err != nil { + t.Fatalf("create active conversation: %v", err) + } + potentiallyActive, err := db.AddMessage(activeConversation.ID, "assistant", "处理中...", nil) + if err != nil { + t.Fatalf("create potentially active placeholder: %v", err) + } + + if err := manager.EnsureSchema(); err != nil { + t.Fatalf("reconcile historical placeholders: %v", err) + } + + assertTerminal := func(messageID, wantContent, wantEvent string) { + t.Helper() + var content, eventType string + if err := db.QueryRow(`SELECT content FROM messages WHERE id = ?`, messageID).Scan(&content); err != nil { + t.Fatalf("query message %s: %v", messageID, err) + } + if content != wantContent { + t.Fatalf("message %s content=%q, want %q", messageID, content, wantContent) + } + if err := db.QueryRow(`SELECT event_type FROM process_details WHERE message_id = ? + AND event_type IN ('cancelled', 'timeout', 'error')`, messageID).Scan(&eventType); err != nil { + t.Fatalf("query terminal detail %s: %v", messageID, err) + } + if eventType != wantEvent { + t.Fatalf("message %s event=%q, want %q", messageID, eventType, wantEvent) + } + } + assertTerminal(superseded.ID, "任务因服务重启已中断。", "cancelled") + assertTerminal(timedOut.ID, "任务等待审批超时，已自动拒绝。", "timeout") + assertTerminal(rejected.ID, "任务审批已拒绝，执行已停止。", "cancelled") + + var activeContent string + if err := db.QueryRow(`SELECT content FROM messages WHERE id = ?`, potentiallyActive.ID).Scan(&activeContent); err != nil { + t.Fatalf("query potentially active message: %v", err) + } + if activeContent != "处理中..." { + t.Fatalf("potentially active message was rewritten to %q", activeContent) + } + var terminalCount int + if err := db.QueryRow(`SELECT COUNT(*) FROM process_details WHERE message_id = ? + AND event_type IN ('cancelled', 'timeout', 'error')`, potentiallyActive.ID).Scan(&terminalCount); err != nil { + t.Fatalf("count active terminal details: %v", err) + } + if terminalCount != 0 { + t.Fatalf("potentially active message got %d terminal details", terminalCount) + } +} + +func TestAuditAgentInterruptIsNotHumanPendingWork(t *testing.T) { + db, err := database.NewDB(filepath.Join(t.TempDir(), "hitl-reviewer.db"), zap.NewNop()) + if err != nil { + t.Fatalf("open db: %v", err) + } + defer func() { _ = db.Close() }() + manager := NewHITLManager(db, zap.NewNop()) + if err := manager.EnsureSchema(); err != nil { + t.Fatalf("ensure schema: %v", err) + } + + audit, err := manager.CreatePendingInterrupt("conversation-audit", "message-audit", "review_edit", "exec", "call-audit", `{}`, "audit_agent") + if err != nil { + t.Fatalf("create audit interrupt: %v", err) + } + human, err := manager.CreatePendingInterrupt("conversation-human", "message-human", "approval", "exec", "call-human", `{}`, "human") + if err != nil { + t.Fatalf("create human interrupt: %v", err) + } + + manager.mu.RLock() + _, auditWaitsForHuman := manager.pending[audit.InterruptID] + _, humanWaitsForHuman := manager.pending[human.InterruptID] + manager.mu.RUnlock() + if auditWaitsForHuman { + t.Fatal("audit-agent interrupt must not enter the human pending queue") + } + if !humanWaitsForHuman { + t.Fatal("human interrupt should enter the human pending queue") + } + + query, args := (&AgentHandler{}).buildHitlListQuery(false) + if len(args) != 0 { + t.Fatalf("unexpected pending query args: %v", args) + } + if !strings.Contains(query, "COALESCE(reviewer,'human') = 'human'") { + t.Fatalf("pending query must filter out audit-agent work: %s", query) + } + rows, err := db.Query(query) + if err != nil { + t.Fatalf("query human pending interrupts: %v", err) + } + defer rows.Close() + items, err := (&AgentHandler{}).scanHitlInterruptRows(rows) + if err != nil { + t.Fatalf("scan human pending interrupts: %v", err) + } + if len(items) != 1 || items[0]["id"] != human.InterruptID || items[0]["reviewer"] != "human" { + t.Fatalf("unexpected human pending result: %#v", items) + } +} diff --git a/internal/handler/hitl_whitelist_test.go b/internal/handler/hitl_whitelist_test.go new file mode 100644 index 00000000..5fcbfd4d --- /dev/null +++ b/internal/handler/hitl_whitelist_test.go @@ -0,0 +1,21 @@ +package handler + +import "testing" + +func TestHITLBuiltInWhitelistExemptsWriteFile(t *testing.T) { + h := &AgentHandler{} + req := h.hitlRequestWithMergedConfigWhitelist(&HITLRequest{ + Enabled: true, + Mode: "approval", + }) + + manager := NewHITLManager(nil, nil) + manager.ActivateConversation("conversation-1", req) + + if manager.NeedsToolApproval("conversation-1", "write_file") { + t.Fatal("write_file should use the built-in HITL exemption") + } + if !manager.NeedsToolApproval("conversation-1", "exec") { + t.Fatal("non-exempt tools should still require approval") + } +} diff --git a/internal/handler/knowledge.go b/internal/handler/knowledge.go new file mode 100644 index 00000000..f08882c8 --- /dev/null +++ b/internal/handler/knowledge.go @@ -0,0 +1,566 @@ +package handler + +import ( + "context" + "fmt" + "net/http" + "strings" + "time" + + "cyberstrike-ai/internal/audit" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/knowledge" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +// KnowledgeHandler 知识库处理器 +type KnowledgeHandler struct { + manager *knowledge.Manager + retriever *knowledge.Retriever + indexer *knowledge.Indexer + db *database.DB + logger *zap.Logger + audit *audit.Service +} + +// SetAudit wires platform audit logging. +func (h *KnowledgeHandler) SetAudit(s *audit.Service) { + h.audit = s +} + +// NewKnowledgeHandler 创建新的知识库处理器 +func NewKnowledgeHandler( + manager *knowledge.Manager, + retriever *knowledge.Retriever, + indexer *knowledge.Indexer, + db *database.DB, + logger *zap.Logger, +) *KnowledgeHandler { + return &KnowledgeHandler{ + manager: manager, + retriever: retriever, + indexer: indexer, + db: db, + logger: logger, + } +} + +// GetCategories 获取所有分类 +func (h *KnowledgeHandler) GetCategories(c *gin.Context) { + categories, err := h.manager.GetCategories() + if err != nil { + h.logger.Error("获取分类失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + c.JSON(http.StatusOK, gin.H{"categories": categories}) +} + +// GetItems 获取知识项列表（支持按分类分页和关键字搜索，默认不返回完整内容） +func (h *KnowledgeHandler) GetItems(c *gin.Context) { + category := c.Query("category") + searchKeyword := c.Query("search") // 搜索关键字 + + // 如果提供了搜索关键字，执行关键字搜索（在所有数据中搜索） + if searchKeyword != "" { + items, err := h.manager.SearchItemsByKeyword(searchKeyword, category) + if err != nil { + h.logger.Error("搜索知识项失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + // 按分类分组结果 + groupedByCategory := make(map[string][]*knowledge.KnowledgeItemSummary) + for _, item := range items { + cat := item.Category + if cat == "" { + cat = "未分类" + } + groupedByCategory[cat] = append(groupedByCategory[cat], item) + } + + // 转换为 CategoryWithItems 格式 + categoriesWithItems := make([]*knowledge.CategoryWithItems, 0, len(groupedByCategory)) + for cat, catItems := range groupedByCategory { + categoriesWithItems = append(categoriesWithItems, &knowledge.CategoryWithItems{ + Category: cat, + ItemCount: len(catItems), + Items: catItems, + }) + } + + // 按分类名称排序 + for i := 0; i < len(categoriesWithItems)-1; i++ { + for j := i + 1; j < len(categoriesWithItems); j++ { + if categoriesWithItems[i].Category > categoriesWithItems[j].Category { + categoriesWithItems[i], categoriesWithItems[j] = categoriesWithItems[j], categoriesWithItems[i] + } + } + } + + c.JSON(http.StatusOK, gin.H{ + "categories": categoriesWithItems, + "total": len(categoriesWithItems), + "search": searchKeyword, + "is_search": true, + }) + return + } + + // 分页模式：categoryPage=true 表示按分类分页，否则按项分页（向后兼容） + categoryPageMode := c.Query("categoryPage") != "false" // 默认使用分类分页 + + // 分页参数 + limit := 50 // 默认每页 50 条（分类分页时为分类数，项分页时为项数） + offset := 0 + if limitStr := c.Query("limit"); limitStr != "" { + if parsed, err := parseInt(limitStr); err == nil && parsed > 0 && parsed <= 500 { + limit = parsed + } + } + if offsetStr := c.Query("offset"); offsetStr != "" { + if parsed, err := parseInt(offsetStr); err == nil && parsed >= 0 { + offset = parsed + } + } + + // 如果指定了 category 参数，且使用分类分页模式，则只返回该分类 + if category != "" && categoryPageMode { + // 单分类模式：返回该分类的所有知识项（不分页） + items, total, err := h.manager.GetItemsSummary(category, 0, 0) + if err != nil { + h.logger.Error("获取知识项失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + // 包装成分类结构 + categoriesWithItems := []*knowledge.CategoryWithItems{ + { + Category: category, + ItemCount: total, + Items: items, + }, + } + + c.JSON(http.StatusOK, gin.H{ + "categories": categoriesWithItems, + "total": 1, // 只有一个分类 + "limit": limit, + "offset": offset, + }) + return + } + + if categoryPageMode { + // 按分类分页模式（默认） + // limit 表示每页分类数，推荐 5-10 个分类 + if limit <= 0 || limit > 100 { + limit = 10 // 默认每页 10 个分类 + } + + categoriesWithItems, totalCategories, err := h.manager.GetCategoriesWithItems(limit, offset) + if err != nil { + h.logger.Error("获取分类知识项失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + c.JSON(http.StatusOK, gin.H{ + "categories": categoriesWithItems, + "total": totalCategories, + "limit": limit, + "offset": offset, + }) + return + } + + // 按项分页模式（向后兼容） + // 是否包含完整内容（默认 false，只返回摘要） + includeContent := c.Query("includeContent") == "true" + + if includeContent { + // 返回完整内容（向后兼容） + items, err := h.manager.GetItemsWithOptions(category, limit, offset, true) + if err != nil { + h.logger.Error("获取知识项失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + // 获取总数 + total, err := h.manager.GetItemsCount(category) + if err != nil { + h.logger.Warn("获取知识项总数失败", zap.Error(err)) + total = len(items) + } + + c.JSON(http.StatusOK, gin.H{ + "items": items, + "total": total, + "limit": limit, + "offset": offset, + }) + } else { + // 返回摘要（不包含完整内容，推荐方式） + items, total, err := h.manager.GetItemsSummary(category, limit, offset) + if err != nil { + h.logger.Error("获取知识项失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + c.JSON(http.StatusOK, gin.H{ + "items": items, + "total": total, + "limit": limit, + "offset": offset, + }) + } +} + +// GetItem 获取单个知识项 +func (h *KnowledgeHandler) GetItem(c *gin.Context) { + id := c.Param("id") + + item, err := h.manager.GetItem(id) + if err != nil { + h.logger.Error("获取知识项失败", zap.Error(err)) + c.JSON(http.StatusNotFound, gin.H{"error": err.Error()}) + return + } + + c.JSON(http.StatusOK, item) +} + +// CreateItem 创建知识项 +func (h *KnowledgeHandler) CreateItem(c *gin.Context) { + var req struct { + Category string `json:"category" binding:"required"` + Title string `json:"title" binding:"required"` + Content string `json:"content" binding:"required"` + } + + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + item, err := h.manager.CreateItem(req.Category, req.Title, req.Content) + if err != nil { + h.logger.Error("创建知识项失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + // 异步索引 + go func() { + ctx := context.Background() + if err := h.indexer.IndexItem(ctx, item.ID); err != nil { + h.logger.Warn("索引知识项失败", zap.String("itemId", item.ID), zap.Error(err)) + } + }() + + c.JSON(http.StatusOK, item) +} + +// UpdateItem 更新知识项 +func (h *KnowledgeHandler) UpdateItem(c *gin.Context) { + id := c.Param("id") + + var req struct { + Category string `json:"category" binding:"required"` + Title string `json:"title" binding:"required"` + Content string `json:"content" binding:"required"` + } + + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + item, err := h.manager.UpdateItem(id, req.Category, req.Title, req.Content) + if err != nil { + h.logger.Error("更新知识项失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + // 异步重新索引 + go func() { + ctx := context.Background() + if err := h.indexer.IndexItem(ctx, item.ID); err != nil { + h.logger.Warn("重新索引知识项失败", zap.String("itemId", item.ID), zap.Error(err)) + } + }() + + c.JSON(http.StatusOK, item) +} + +// DeleteItem 删除知识项 +func (h *KnowledgeHandler) DeleteItem(c *gin.Context) { + id := c.Param("id") + + if err := h.manager.DeleteItem(id); err != nil { + h.logger.Error("删除知识项失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + if h.audit != nil { + h.audit.RecordOK(c, "knowledge", "item_delete", "删除知识项", "knowledge_item", id, nil) + } + c.JSON(http.StatusOK, gin.H{"message": "删除成功"}) +} + +// StartIndex 构建知识库向量索引。默认仅补齐尚无向量的知识项；mode=full 时全量重建。 +func (h *KnowledgeHandler) StartIndex(c *gin.Context) { + if err := h.indexer.TryBeginIndexRun(); err != nil { + c.JSON(http.StatusConflict, gin.H{"error": "已有索引任务正在进行，请等待完成"}) + return + } + + mode := strings.TrimSpace(c.Query("mode")) + if mode == "" { + mode = "missing" + } + if mode != "full" && mode != "missing" { + h.indexer.FinishIndexRun() + c.JSON(http.StatusBadRequest, gin.H{"error": "无效的 mode 参数，可选值：missing、full"}) + return + } + + fullRebuild := mode == "full" + message := "索引构建已开始，将在后台进行" + auditAction := "index_build" + auditDetail := "构建知识库索引" + if fullRebuild { + message = "全量索引重建已开始，将在后台进行" + auditAction = "index_rebuild_full" + auditDetail = "全量重建知识库索引" + } + + go func() { + defer h.indexer.FinishIndexRun() + ctx := context.Background() + var err error + if fullRebuild { + err = h.indexer.RunRebuildIndex(ctx) + } else { + err = h.indexer.RunIndexMissing(ctx) + } + if err != nil { + if fullRebuild { + h.logger.Error("全量重建索引失败", zap.Error(err)) + } else { + h.logger.Error("构建知识库索引失败", zap.Error(err)) + } + } + }() + + if h.audit != nil { + h.audit.RecordOK(c, "knowledge", auditAction, auditDetail, "knowledge", "", nil) + } + c.JSON(http.StatusOK, gin.H{"message": message, "mode": mode}) +} + +// ScanKnowledgeBase 扫描知识库 +func (h *KnowledgeHandler) ScanKnowledgeBase(c *gin.Context) { + itemsToIndex, err := h.manager.ScanKnowledgeBase() + if err != nil { + h.logger.Error("扫描知识库失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + if len(itemsToIndex) == 0 { + c.JSON(http.StatusOK, gin.H{"message": "扫描完成，没有需要索引的新项或更新项"}) + return + } + + // 异步索引新添加或更新的项（增量索引） + go func() { + ctx := context.Background() + h.logger.Info("开始增量索引", zap.Int("count", len(itemsToIndex))) + failedCount := 0 + consecutiveFailures := 0 + var firstFailureItemID string + var firstFailureError error + + for i, itemID := range itemsToIndex { + if err := h.indexer.IndexItem(ctx, itemID); err != nil { + failedCount++ + consecutiveFailures++ + + // 只在第一个失败时记录详细日志 + if consecutiveFailures == 1 { + firstFailureItemID = itemID + firstFailureError = err + h.logger.Warn("索引知识项失败", + zap.String("itemId", itemID), + zap.Int("totalItems", len(itemsToIndex)), + zap.Error(err), + ) + } + + // 如果连续失败 2 次，立即停止增量索引 + if consecutiveFailures >= 2 { + h.logger.Error("连续索引失败次数过多，立即停止增量索引", + zap.Int("consecutiveFailures", consecutiveFailures), + zap.Int("totalItems", len(itemsToIndex)), + zap.Int("processedItems", i+1), + zap.String("firstFailureItemId", firstFailureItemID), + zap.Error(firstFailureError), + ) + break + } + continue + } + + // 成功时重置连续失败计数 + if consecutiveFailures > 0 { + consecutiveFailures = 0 + firstFailureItemID = "" + firstFailureError = nil + } + + // 减少进度日志频率 + if (i+1)%10 == 0 || i+1 == len(itemsToIndex) { + h.logger.Info("索引进度", zap.Int("current", i+1), zap.Int("total", len(itemsToIndex)), zap.Int("failed", failedCount)) + } + } + h.logger.Info("增量索引完成", zap.Int("totalItems", len(itemsToIndex)), zap.Int("failedCount", failedCount)) + }() + + c.JSON(http.StatusOK, gin.H{ + "message": fmt.Sprintf("扫描完成，开始索引 %d 个新添加或更新的知识项", len(itemsToIndex)), + "items_to_index": len(itemsToIndex), + }) +} + +// GetRetrievalLogs 获取检索日志 +func (h *KnowledgeHandler) GetRetrievalLogs(c *gin.Context) { + conversationID := c.Query("conversationId") + messageID := c.Query("messageId") + limit := 50 // 默认 50 条 + + if limitStr := c.Query("limit"); limitStr != "" { + if parsed, err := parseInt(limitStr); err == nil && parsed > 0 { + limit = parsed + } + } + + logs, err := h.manager.GetRetrievalLogs(conversationID, messageID, limit) + if err != nil { + h.logger.Error("获取检索日志失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + c.JSON(http.StatusOK, gin.H{"logs": logs}) +} + +// DeleteRetrievalLog 删除检索日志 +func (h *KnowledgeHandler) DeleteRetrievalLog(c *gin.Context) { + id := c.Param("id") + + if err := h.manager.DeleteRetrievalLog(id); err != nil { + h.logger.Error("删除检索日志失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + c.JSON(http.StatusOK, gin.H{"message": "删除成功"}) +} + +// GetIndexStatus 获取索引状态 +func (h *KnowledgeHandler) GetIndexStatus(c *gin.Context) { + status, err := h.manager.GetIndexStatus() + if err != nil { + h.logger.Error("获取索引状态失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + // 获取索引器的错误信息 + if h.indexer != nil { + lastError, lastErrorTime := h.indexer.GetLastError() + if lastError != "" { + // 如果错误是最近发生的（5 分钟内），则返回错误信息 + if time.Since(lastErrorTime) < 5*time.Minute { + status["last_error"] = lastError + status["last_error_time"] = lastErrorTime.Format(time.RFC3339) + } + } + + // 获取重建索引状态 + isRebuilding, totalItems, current, failed, lastItemID, lastChunks, startTime := h.indexer.GetRebuildStatus() + if isRebuilding { + status["is_rebuilding"] = true + status["rebuild_total"] = totalItems + status["rebuild_current"] = current + status["rebuild_failed"] = failed + status["rebuild_start_time"] = startTime.Format(time.RFC3339) + if lastItemID != "" { + status["rebuild_last_item_id"] = lastItemID + } + if lastChunks > 0 { + status["rebuild_last_chunks"] = lastChunks + } + // 重建中时，is_complete 为 false + status["is_complete"] = false + // 计算重建进度百分比 + if totalItems > 0 { + status["progress_percent"] = float64(current) / float64(totalItems) * 100 + } + } + } + + c.JSON(http.StatusOK, status) +} + +// Search 搜索知识库（用于 API 调用，Agent 内部使用 Retriever） +func (h *KnowledgeHandler) Search(c *gin.Context) { + var req knowledge.SearchRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + // Retriever.Search 经 Eino VectorEinoRetriever，与 MCP 工具链一致。 + results, err := h.retriever.Search(c.Request.Context(), &req) + if err != nil { + h.logger.Error("搜索知识库失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + c.JSON(http.StatusOK, gin.H{"results": results}) +} + +// GetStats 获取知识库统计信息 +func (h *KnowledgeHandler) GetStats(c *gin.Context) { + totalCategories, totalItems, err := h.manager.GetStats() + if err != nil { + h.logger.Error("获取知识库统计信息失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + c.JSON(http.StatusOK, gin.H{ + "enabled": true, + "total_categories": totalCategories, + "total_items": totalItems, + }) +} + +// 辅助函数：解析整数 +func parseInt(s string) (int, error) { + var result int + _, err := fmt.Sscanf(s, "%d", &result) + return result, err +} diff --git a/internal/handler/markdown_agents.go b/internal/handler/markdown_agents.go new file mode 100644 index 00000000..70ba216d --- /dev/null +++ b/internal/handler/markdown_agents.go @@ -0,0 +1,333 @@ +package handler + +import ( + "fmt" + "net/http" + "os" + "path/filepath" + "regexp" + "strings" + + "cyberstrike-ai/internal/agents" + "cyberstrike-ai/internal/audit" + "cyberstrike-ai/internal/config" + + "github.com/gin-gonic/gin" +) + +var markdownAgentFilenameRe = regexp.MustCompile(`^[a-zA-Z0-9][a-zA-Z0-9_.-]*\.md$`) + +// MarkdownAgentsHandler 管理 agents 目录下子代理 Markdown（增删改查）。 +type MarkdownAgentsHandler struct { + dir string + audit *audit.Service +} + +// NewMarkdownAgentsHandler dir 须为已解析的绝对路径。 +func NewMarkdownAgentsHandler(dir string) *MarkdownAgentsHandler { + return &MarkdownAgentsHandler{dir: strings.TrimSpace(dir)} +} + +// SetAudit wires platform audit logging. +func (h *MarkdownAgentsHandler) SetAudit(s *audit.Service) { + h.audit = s +} + +func (h *MarkdownAgentsHandler) safeJoin(filename string) (string, error) { + filename = strings.TrimSpace(filename) + if filename == "" || !markdownAgentFilenameRe.MatchString(filename) { + return "", fmt.Errorf("非法文件名") + } + clean := filepath.Clean(filename) + if clean != filename || strings.Contains(clean, "..") { + return "", fmt.Errorf("非法文件名") + } + return filepath.Join(h.dir, clean), nil +} + +// existingOtherOrchestrator 若目录中已有同槽位的其他主代理文件，返回其文件名；writingBasename 为当前正在写入的文件名时不冲突。 +func existingOtherOrchestrator(dir, writingBasename string) (other string, err error) { + load, err := agents.LoadMarkdownAgentsDir(dir) + if err != nil { + return "", err + } + wb := filepath.Base(strings.TrimSpace(writingBasename)) + switch agents.OrchestratorMarkdownKind(wb) { + case "plan_execute": + if load.OrchestratorPlanExecute != nil && !strings.EqualFold(load.OrchestratorPlanExecute.Filename, wb) { + return load.OrchestratorPlanExecute.Filename, nil + } + case "supervisor": + if load.OrchestratorSupervisor != nil && !strings.EqualFold(load.OrchestratorSupervisor.Filename, wb) { + return load.OrchestratorSupervisor.Filename, nil + } + case "deep": + if load.Orchestrator != nil && !strings.EqualFold(load.Orchestrator.Filename, wb) { + return load.Orchestrator.Filename, nil + } + default: + if load.Orchestrator != nil && !strings.EqualFold(load.Orchestrator.Filename, wb) { + return load.Orchestrator.Filename, nil + } + } + return "", nil +} + +// ListMarkdownAgents GET /api/multi-agent/markdown-agents +func (h *MarkdownAgentsHandler) ListMarkdownAgents(c *gin.Context) { + if h.dir == "" { + c.JSON(http.StatusOK, gin.H{"agents": []any{}, "dir": "", "error": "未配置 agents 目录"}) + return + } + files, err := agents.LoadMarkdownAgentFiles(h.dir) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + out := make([]gin.H, 0, len(files)) + for _, fa := range files { + sub := fa.Config + out = append(out, gin.H{ + "filename": fa.Filename, + "id": sub.ID, + "name": sub.Name, + "description": sub.Description, + "is_orchestrator": fa.IsOrchestrator, + "kind": sub.Kind, + }) + } + c.JSON(http.StatusOK, gin.H{"agents": out, "dir": h.dir}) +} + +// GetMarkdownAgent GET /api/multi-agent/markdown-agents/:filename +func (h *MarkdownAgentsHandler) GetMarkdownAgent(c *gin.Context) { + filename := c.Param("filename") + path, err := h.safeJoin(filename) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + b, err := os.ReadFile(path) + if err != nil { + if os.IsNotExist(err) { + c.JSON(http.StatusNotFound, gin.H{"error": "文件不存在"}) + return + } + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + sub, err := agents.ParseMarkdownSubAgent(filename, string(b)) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + isOrch := agents.IsOrchestratorLikeMarkdown(filename, sub.Kind) + c.JSON(http.StatusOK, gin.H{ + "filename": filename, + "raw": string(b), + "id": sub.ID, + "name": sub.Name, + "description": sub.Description, + "tools": sub.RoleTools, + "instruction": sub.Instruction, + "bind_role": sub.BindRole, + "max_iterations": sub.MaxIterations, + "kind": sub.Kind, + "is_orchestrator": isOrch, + }) +} + +type markdownAgentBody struct { + Filename string `json:"filename"` + ID string `json:"id"` + Name string `json:"name"` + Description string `json:"description"` + Tools []string `json:"tools"` + Instruction string `json:"instruction"` + BindRole string `json:"bind_role"` + MaxIterations int `json:"max_iterations"` + Kind string `json:"kind"` + Raw string `json:"raw"` +} + +// CreateMarkdownAgent POST /api/multi-agent/markdown-agents +func (h *MarkdownAgentsHandler) CreateMarkdownAgent(c *gin.Context) { + if h.dir == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "未配置 agents 目录"}) + return + } + var body markdownAgentBody + if err := c.ShouldBindJSON(&body); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + filename := strings.TrimSpace(body.Filename) + if filename == "" { + if strings.EqualFold(strings.TrimSpace(body.Kind), "orchestrator") { + filename = agents.OrchestratorMarkdownFilename + } else { + base := agents.SlugID(body.Name) + if base == "" { + base = "agent" + } + filename = base + ".md" + } + } + path, err := h.safeJoin(filename) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if _, err := os.Stat(path); err == nil { + c.JSON(http.StatusConflict, gin.H{"error": "文件已存在"}) + return + } + sub := config.MultiAgentSubConfig{ + ID: strings.TrimSpace(body.ID), + Name: strings.TrimSpace(body.Name), + Description: strings.TrimSpace(body.Description), + Instruction: strings.TrimSpace(body.Instruction), + RoleTools: body.Tools, + BindRole: strings.TrimSpace(body.BindRole), + MaxIterations: body.MaxIterations, + Kind: strings.TrimSpace(body.Kind), + } + base := filepath.Base(path) + if (strings.EqualFold(base, agents.OrchestratorMarkdownFilename) || + strings.EqualFold(base, agents.OrchestratorPlanExecuteMarkdownFilename) || + strings.EqualFold(base, agents.OrchestratorSupervisorMarkdownFilename)) && sub.Kind == "" { + sub.Kind = "orchestrator" + } + if sub.ID == "" { + sub.ID = agents.SlugID(sub.Name) + } + if sub.Name == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "name 必填"}) + return + } + var out []byte + if strings.TrimSpace(body.Raw) != "" { + out = []byte(body.Raw) + } else { + out, err = agents.BuildMarkdownFile(sub) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + } + if want := agents.WantsMarkdownOrchestrator(filepath.Base(path), body.Kind, string(out)); want { + other, oerr := existingOtherOrchestrator(h.dir, filepath.Base(path)) + if oerr != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": oerr.Error()}) + return + } + if other != "" { + c.JSON(http.StatusConflict, gin.H{"error": fmt.Sprintf("已存在主代理定义：%s，请先删除或取消其主代理标记", other)}) + return + } + } + if err := os.MkdirAll(h.dir, 0755); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if err := os.WriteFile(path, out, 0644); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "agent", "markdown_create", "创建 Markdown 子代理", "markdown_agent", filepath.Base(path), nil) + } + c.JSON(http.StatusOK, gin.H{"filename": filepath.Base(path), "message": "已创建"}) +} + +// UpdateMarkdownAgent PUT /api/multi-agent/markdown-agents/:filename +func (h *MarkdownAgentsHandler) UpdateMarkdownAgent(c *gin.Context) { + filename := c.Param("filename") + path, err := h.safeJoin(filename) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + var body markdownAgentBody + if err := c.ShouldBindJSON(&body); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + sub := config.MultiAgentSubConfig{ + ID: strings.TrimSpace(body.ID), + Name: strings.TrimSpace(body.Name), + Description: strings.TrimSpace(body.Description), + Instruction: strings.TrimSpace(body.Instruction), + RoleTools: body.Tools, + BindRole: strings.TrimSpace(body.BindRole), + MaxIterations: body.MaxIterations, + Kind: strings.TrimSpace(body.Kind), + } + if (strings.EqualFold(filename, agents.OrchestratorMarkdownFilename) || + strings.EqualFold(filename, agents.OrchestratorPlanExecuteMarkdownFilename) || + strings.EqualFold(filename, agents.OrchestratorSupervisorMarkdownFilename)) && sub.Kind == "" { + sub.Kind = "orchestrator" + } + if sub.Name == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "name 必填"}) + return + } + if sub.ID == "" { + sub.ID = agents.SlugID(sub.Name) + } + var out []byte + if strings.TrimSpace(body.Raw) != "" { + out = []byte(body.Raw) + } else { + out, err = agents.BuildMarkdownFile(sub) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + } + if want := agents.WantsMarkdownOrchestrator(filename, body.Kind, string(out)); want { + other, oerr := existingOtherOrchestrator(h.dir, filename) + if oerr != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": oerr.Error()}) + return + } + if other != "" { + c.JSON(http.StatusConflict, gin.H{"error": fmt.Sprintf("已存在主代理定义：%s，请先删除或取消其主代理标记", other)}) + return + } + } + if err := os.WriteFile(path, out, 0644); err != nil { + if os.IsNotExist(err) { + c.JSON(http.StatusNotFound, gin.H{"error": "文件不存在"}) + return + } + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "agent", "markdown_update", "更新 Markdown 子代理", "markdown_agent", filename, nil) + } + c.JSON(http.StatusOK, gin.H{"message": "已保存"}) +} + +// DeleteMarkdownAgent DELETE /api/multi-agent/markdown-agents/:filename +func (h *MarkdownAgentsHandler) DeleteMarkdownAgent(c *gin.Context) { + filename := c.Param("filename") + path, err := h.safeJoin(filename) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if err := os.Remove(path); err != nil { + if os.IsNotExist(err) { + c.JSON(http.StatusNotFound, gin.H{"error": "文件不存在"}) + return + } + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "agent", "markdown_delete", "删除 Markdown 子代理", "markdown_agent", filename, nil) + } + c.JSON(http.StatusOK, gin.H{"message": "已删除"}) +} diff --git a/internal/handler/monitor.go b/internal/handler/monitor.go new file mode 100644 index 00000000..31f8a746 --- /dev/null +++ b/internal/handler/monitor.go @@ -0,0 +1,1033 @@ +package handler + +import ( + "encoding/json" + "errors" + "io" + "net/http" + "sort" + "strconv" + "strings" + "time" + + "cyberstrike-ai/internal/audit" + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/mcp" + "cyberstrike-ai/internal/monitor" + "cyberstrike-ai/internal/security" + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +// MonitorHandler 监控处理器 +type MonitorHandler struct { + mcpServer *mcp.Server + externalMCPMgr *mcp.ExternalMCPManager + taskManager *AgentTaskManager + agentHandler *AgentHandler + executor *security.Executor + db *database.DB + logger *zap.Logger + audit *audit.Service + monitorRetention *monitor.Service +} + +// SetMonitorRetention wires MCP execution retention settings. +func (h *MonitorHandler) SetMonitorRetention(s *monitor.Service) { + h.monitorRetention = s +} + +// SetAudit wires platform audit logging. +func (h *MonitorHandler) SetAudit(s *audit.Service) { + h.audit = s +} + +// NewMonitorHandler 创建新的监控处理器 +func NewMonitorHandler(mcpServer *mcp.Server, executor *security.Executor, db *database.DB, logger *zap.Logger) *MonitorHandler { + return &MonitorHandler{ + mcpServer: mcpServer, + externalMCPMgr: nil, // 将在创建后设置 + executor: executor, + db: db, + logger: logger, + } +} + +// SetExternalMCPManager 设置外部MCP管理器 +func (h *MonitorHandler) SetExternalMCPManager(mgr *mcp.ExternalMCPManager) { + h.externalMCPMgr = mgr +} + +// SetTaskManager 设置 Agent 任务管理器（用于 Eino execute 等按 executionId 终止）。 +func (h *MonitorHandler) SetTaskManager(mgr *AgentTaskManager) { + h.taskManager = mgr +} + +// SetAgentHandler 设置 Agent 处理器（MCP 监控终止与对话页「中断并继续」共用逻辑）。 +func (h *MonitorHandler) SetAgentHandler(ah *AgentHandler) { + h.agentHandler = ah +} + +const monitorPageTopTools = 6 + +// MonitorStatsSummary 工具调用汇总 +type MonitorStatsSummary struct { + TotalCalls int `json:"totalCalls"` + SuccessCalls int `json:"successCalls"` + FailedCalls int `json:"failedCalls"` + LastCallTime *time.Time `json:"lastCallTime,omitempty"` + ToolCount int `json:"toolCount"` +} + +// MonitorResponse 监控响应 +type MonitorResponse struct { + Executions []*mcp.ToolExecution `json:"executions"` + Summary *MonitorStatsSummary `json:"summary"` + TopTools []*mcp.ToolStats `json:"topTools"` + Timestamp time.Time `json:"timestamp"` + Total int `json:"total"` + Page int `json:"page"` + PageSize int `json:"pageSize"` + TotalPages int `json:"totalPages"` + RetentionDays int `json:"retentionDays"` +} + +// StatsResponse 统计信息响应（Dashboard 等） +type StatsResponse struct { + Summary *MonitorStatsSummary `json:"summary"` + TopTools []*mcp.ToolStats `json:"topTools"` +} + +// Monitor 获取监控信息 +func (h *MonitorHandler) Monitor(c *gin.Context) { + // 解析分页参数 + page := 1 + pageSize := 20 + if pageStr := c.Query("page"); pageStr != "" { + if p, err := strconv.Atoi(pageStr); err == nil && p > 0 { + page = p + } + } + if pageSizeStr := c.Query("page_size"); pageSizeStr != "" { + if ps, err := strconv.Atoi(pageSizeStr); err == nil && ps > 0 && ps <= 100 { + pageSize = ps + } + } + + // 解析状态筛选参数 + status := c.Query("status") + // 解析工具筛选参数（兼容 mcp__tool 与内部 mcp::tool） + toolName := normalizeToolNameFilter(c.Query("tool")) + + access := notificationAccessFromContext(c) + executions, total := h.loadExecutionListWithPagination(page, pageSize, status, toolName, access) + h.enrichExecutionsConversationID(executions) + var summary *MonitorStatsSummary + var topTools []*mcp.ToolStats + if access.Scope == database.RBACScopeAll { + summary, topTools = h.loadStatsSummary(monitorPageTopTools) + } else if h.db != nil { + if scoped, err := h.db.LoadToolStatsSummaryForAccess(monitorPageTopTools, access); err == nil { + summary, topTools = dbStatsSummaryToMonitor(scoped), scoped.TopTools + } else { + summary, topTools = summarizeAccessibleExecutionPage(executions, monitorPageTopTools) + } + } else { + summary, topTools = summarizeAccessibleExecutionPage(executions, monitorPageTopTools) + } + + totalPages := (total + pageSize - 1) / pageSize + if totalPages == 0 { + totalPages = 1 + } + + c.JSON(http.StatusOK, MonitorResponse{ + Executions: executions, + Summary: summary, + TopTools: topTools, + Timestamp: time.Now(), + Total: total, + Page: page, + PageSize: pageSize, + TotalPages: totalPages, + RetentionDays: h.monitorRetentionDays(), + }) +} + +func summarizeAccessibleExecutionPage(executions []*mcp.ToolExecution, topN int) (*MonitorStatsSummary, []*mcp.ToolStats) { + stats := map[string]*mcp.ToolStats{} + for _, exec := range executions { + if exec == nil { + continue + } + stat := stats[exec.ToolName] + if stat == nil { + stat = &mcp.ToolStats{ToolName: exec.ToolName} + stats[exec.ToolName] = stat + } + stat.TotalCalls++ + if monitorStatusCountsAsFailed(exec.Status) { + stat.FailedCalls++ + } else if exec.Status == "completed" { + stat.SuccessCalls++ + } + started := exec.StartTime + if stat.LastCallTime == nil || started.After(*stat.LastCallTime) { + stat.LastCallTime = &started + } + } + return summarizeToolStats(stats, topN) +} + +func monitorStatusCountsAsFailed(status string) bool { + switch strings.TrimSpace(strings.ToLower(status)) { + case "failed", "hard_timeout", "orphaned": + return true + default: + return false + } +} + +func (h *MonitorHandler) monitorRetentionDays() int { + if h.monitorRetention != nil { + return h.monitorRetention.RetentionDays() + } + return config.MonitorConfig{}.RetentionDaysEffective() +} + +func (h *MonitorHandler) loadExecutions() []*mcp.ToolExecution { + executions, _ := h.loadExecutionsWithPagination(1, 1000, "", "") + return executions +} + +func (h *MonitorHandler) loadExecutionListWithPagination(page, pageSize int, status, toolName string, access database.RBACListAccess) ([]*mcp.ToolExecution, int) { + if h.db == nil { + allExecutions := filterToolExecutionsForAccess(h.mcpServer.GetAllExecutions(), access, h.db) + if status != "" || toolName != "" { + filtered := make([]*mcp.ToolExecution, 0) + for _, exec := range allExecutions { + matchStatus := status == "" || exec.Status == status + matchTool := toolNameFilterMatches(exec.ToolName, toolName) + if matchStatus && matchTool { + filtered = append(filtered, exec) + } + } + allExecutions = filtered + } + total := len(allExecutions) + offset := (page - 1) * pageSize + end := offset + pageSize + if end > total { + end = total + } + if offset >= total { + return []*mcp.ToolExecution{}, total + } + pageSlice := allExecutions[offset:end] + out := make([]*mcp.ToolExecution, 0, len(pageSlice)) + for _, exec := range pageSlice { + if exec == nil { + continue + } + out = append(out, slimToolExecution(exec)) + } + return out, total + } + + offset := (page - 1) * pageSize + executions, err := h.db.LoadToolExecutionListPageForAccess(offset, pageSize, status, toolName, access) + if err != nil { + h.logger.Warn("从数据库加载执行记录列表失败，回退到内存数据", zap.Error(err)) + return h.loadExecutionListWithPaginationFromMemory(page, pageSize, status, toolName, access) + } + + total, err := h.db.CountToolExecutionsForAccess(status, toolName, access) + if err != nil { + h.logger.Warn("获取执行记录总数失败", zap.Error(err)) + total = offset + len(executions) + if len(executions) == pageSize { + total = offset + len(executions) + 1 + } + } + + return executions, total +} + +func (h *MonitorHandler) loadExecutionListWithPaginationFromMemory(page, pageSize int, status, toolName string, access database.RBACListAccess) ([]*mcp.ToolExecution, int) { + allExecutions := filterToolExecutionsForAccess(h.mcpServer.GetAllExecutions(), access, h.db) + if status != "" || toolName != "" { + filtered := make([]*mcp.ToolExecution, 0) + for _, exec := range allExecutions { + matchStatus := status == "" || exec.Status == status + matchTool := toolNameFilterMatches(exec.ToolName, toolName) + if matchStatus && matchTool { + filtered = append(filtered, exec) + } + } + allExecutions = filtered + } + total := len(allExecutions) + offset := (page - 1) * pageSize + end := offset + pageSize + if end > total { + end = total + } + if offset >= total { + return []*mcp.ToolExecution{}, total + } + pageSlice := allExecutions[offset:end] + out := make([]*mcp.ToolExecution, 0, len(pageSlice)) + for _, exec := range pageSlice { + if exec == nil { + continue + } + out = append(out, slimToolExecution(exec)) + } + return out, total +} + +func slimToolExecution(exec *mcp.ToolExecution) *mcp.ToolExecution { + if exec == nil { + return nil + } + slim := &mcp.ToolExecution{ + ID: exec.ID, + ToolName: exec.ToolName, + Status: exec.Status, + StartTime: exec.StartTime, + } + if exec.EndTime != nil { + end := *exec.EndTime + slim.EndTime = &end + } + if exec.Duration > 0 { + slim.Duration = exec.Duration + } + return slim +} + +func filterToolExecutionsForAccess(executions []*mcp.ToolExecution, access database.RBACListAccess, db *database.DB) []*mcp.ToolExecution { + if access.Scope == database.RBACScopeAll { + return executions + } + out := make([]*mcp.ToolExecution, 0, len(executions)) + for _, exec := range executions { + if toolExecutionVisible(exec, access, db) { + out = append(out, exec) + } + } + return out +} + +func toolExecutionVisible(exec *mcp.ToolExecution, access database.RBACListAccess, db *database.DB) bool { + if exec == nil || strings.TrimSpace(access.UserID) == "" { + return false + } + if access.Scope == database.RBACScopeAll || strings.TrimSpace(exec.OwnerUserID) == strings.TrimSpace(access.UserID) { + return true + } + conversationID := strings.TrimSpace(exec.ConversationID) + return conversationID != "" && db != nil && db.UserCanAccessResource(access.UserID, access.Scope, "conversation", conversationID) +} + +func (h *MonitorHandler) monitorExecutionAllowed(c *gin.Context, id string) bool { + access := notificationAccessFromContext(c) + if access.Scope == database.RBACScopeAll { + return true + } + id = strings.TrimSpace(id) + if id == "" { + return false + } + if exec, ok := h.mcpServer.GetExecution(id); ok { + return toolExecutionVisible(exec, access, h.db) + } + if h.externalMCPMgr != nil { + if exec, ok := h.externalMCPMgr.GetExecution(id); ok { + return toolExecutionVisible(exec, access, h.db) + } + } + return h.db != nil && h.db.UserCanAccessToolExecution(access.UserID, access.Scope, id) +} + +func (h *MonitorHandler) loadExecutionsWithPagination(page, pageSize int, status, toolName string) ([]*mcp.ToolExecution, int) { + if h.db == nil { + allExecutions := h.mcpServer.GetAllExecutions() + // 如果指定了状态筛选或工具筛选，先进行筛选 + if status != "" || toolName != "" { + filtered := make([]*mcp.ToolExecution, 0) + for _, exec := range allExecutions { + matchStatus := status == "" || exec.Status == status + // 支持部分匹配（模糊搜索） + matchTool := toolNameFilterMatches(exec.ToolName, toolName) + if matchStatus && matchTool { + filtered = append(filtered, exec) + } + } + allExecutions = filtered + } + total := len(allExecutions) + offset := (page - 1) * pageSize + end := offset + pageSize + if end > total { + end = total + } + if offset >= total { + return []*mcp.ToolExecution{}, total + } + return allExecutions[offset:end], total + } + + offset := (page - 1) * pageSize + executions, err := h.db.LoadToolExecutionsWithPagination(offset, pageSize, status, toolName) + if err != nil { + h.logger.Warn("从数据库加载执行记录失败，回退到内存数据", zap.Error(err)) + allExecutions := h.mcpServer.GetAllExecutions() + // 如果指定了状态筛选或工具筛选，先进行筛选 + if status != "" || toolName != "" { + filtered := make([]*mcp.ToolExecution, 0) + for _, exec := range allExecutions { + matchStatus := status == "" || exec.Status == status + // 支持部分匹配（模糊搜索） + matchTool := toolNameFilterMatches(exec.ToolName, toolName) + if matchStatus && matchTool { + filtered = append(filtered, exec) + } + } + allExecutions = filtered + } + total := len(allExecutions) + offset := (page - 1) * pageSize + end := offset + pageSize + if end > total { + end = total + } + if offset >= total { + return []*mcp.ToolExecution{}, total + } + return allExecutions[offset:end], total + } + + // 获取总数（考虑状态筛选和工具筛选） + total, err := h.db.CountToolExecutions(status, toolName) + if err != nil { + h.logger.Warn("获取执行记录总数失败", zap.Error(err)) + // 回退：使用已加载的记录数估算 + total = offset + len(executions) + if len(executions) == pageSize { + total = offset + len(executions) + 1 + } + } + + return executions, total +} + +func (h *MonitorHandler) loadStatsSummary(topN int) (*MonitorStatsSummary, []*mcp.ToolStats) { + if topN <= 0 { + topN = monitorPageTopTools + } + + if h.db != nil { + result, err := h.db.LoadToolStatsSummary(topN) + if err == nil { + return dbStatsSummaryToMonitor(result), result.TopTools + } + h.logger.Warn("从数据库加载统计汇总失败，回退到内存数据", zap.Error(err)) + } + + stats := h.loadStatsMap() + return summarizeToolStats(stats, topN) +} + +func dbStatsSummaryToMonitor(result *database.ToolStatsSummaryResult) *MonitorStatsSummary { + if result == nil { + return &MonitorStatsSummary{} + } + summary := &MonitorStatsSummary{ + TotalCalls: result.Summary.TotalCalls, + SuccessCalls: result.Summary.SuccessCalls, + FailedCalls: result.Summary.FailedCalls, + ToolCount: result.Summary.ToolCount, + } + if result.Summary.LastCallTime != nil { + t := *result.Summary.LastCallTime + summary.LastCallTime = &t + } + return summary +} + +func summarizeToolStats(stats map[string]*mcp.ToolStats, topN int) (*MonitorStatsSummary, []*mcp.ToolStats) { + summary := &MonitorStatsSummary{} + if len(stats) == 0 { + return summary, nil + } + + all := make([]*mcp.ToolStats, 0, len(stats)) + for _, stat := range stats { + if stat == nil { + continue + } + summary.ToolCount++ + summary.TotalCalls += stat.TotalCalls + summary.SuccessCalls += stat.SuccessCalls + summary.FailedCalls += stat.FailedCalls + if stat.LastCallTime != nil && (summary.LastCallTime == nil || stat.LastCallTime.After(*summary.LastCallTime)) { + t := *stat.LastCallTime + summary.LastCallTime = &t + } + if stat.TotalCalls > 0 { + statCopy := *stat + all = append(all, &statCopy) + } + } + + sort.Slice(all, func(i, j int) bool { + if all[i].TotalCalls == all[j].TotalCalls { + return all[i].ToolName < all[j].ToolName + } + return all[i].TotalCalls > all[j].TotalCalls + }) + if len(all) > topN { + all = all[:topN] + } + return summary, all +} + +func (h *MonitorHandler) loadStatsMap() map[string]*mcp.ToolStats { + // 合并内部MCP服务器和外部MCP管理器的统计信息 + stats := make(map[string]*mcp.ToolStats) + + // 加载内部MCP服务器的统计信息 + if h.db == nil { + internalStats := h.mcpServer.GetStats() + for k, v := range internalStats { + stats[k] = v + } + } else { + dbStats, err := h.db.LoadToolStats() + if err != nil { + h.logger.Warn("从数据库加载统计信息失败，回退到内存数据", zap.Error(err)) + internalStats := h.mcpServer.GetStats() + for k, v := range internalStats { + stats[k] = v + } + } else { + for k, v := range dbStats { + stats[k] = v + } + } + } + + // 合并外部MCP管理器的统计信息 + if h.externalMCPMgr != nil { + externalStats := h.externalMCPMgr.GetToolStats() + for k, v := range externalStats { + // 如果已存在，合并统计信息 + if existing, exists := stats[k]; exists { + existing.TotalCalls += v.TotalCalls + existing.SuccessCalls += v.SuccessCalls + existing.FailedCalls += v.FailedCalls + // 使用最新的调用时间 + if v.LastCallTime != nil && (existing.LastCallTime == nil || v.LastCallTime.After(*existing.LastCallTime)) { + existing.LastCallTime = v.LastCallTime + } + } else { + stats[k] = v + } + } + } + + return stats +} + +// GetExecution 获取特定执行记录 +func (h *MonitorHandler) GetExecution(c *gin.Context) { + id := c.Param("id") + if !h.monitorExecutionAllowed(c, id) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + + // 先从内部MCP服务器查找 + exec, exists := h.mcpServer.GetExecution(id) + if exists { + h.enrichExecutionsConversationID([]*mcp.ToolExecution{exec}) + c.JSON(http.StatusOK, exec) + return + } + + // 如果找不到，尝试从外部MCP管理器查找 + if h.externalMCPMgr != nil { + exec, exists = h.externalMCPMgr.GetExecution(id) + if exists { + h.enrichExecutionsConversationID([]*mcp.ToolExecution{exec}) + c.JSON(http.StatusOK, exec) + return + } + } + + // 如果都找不到，尝试从数据库查找（如果使用数据库存储） + if h.db != nil { + exec, err := h.db.GetToolExecution(id) + if err == nil && exec != nil { + h.enrichExecutionsConversationID([]*mcp.ToolExecution{exec}) + c.JSON(http.StatusOK, exec) + return + } + } + + c.JSON(http.StatusNotFound, gin.H{"error": "执行记录未找到"}) +} + +// CancelExecution 手动取消进行中的 MCP 工具调用（仅取消该次 tools/call 的上下文，不停止整条 Agent / 迭代任务） +// 请求体可选 JSON：{ "note": "用户说明" }，将与工具已返回输出合并交给模型（含「用户终止说明」标题块，与命令行原文区分）。 +func (h *MonitorHandler) CancelExecution(c *gin.Context) { + id := c.Param("id") + if id == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "执行记录ID不能为空"}) + return + } + if !h.monitorExecutionAllowed(c, id) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + note := "" + dec := json.NewDecoder(c.Request.Body) + var body struct { + Note string `json:"note"` + } + if err := dec.Decode(&body); err != nil && !errors.Is(err, io.EOF) { + c.JSON(http.StatusBadRequest, gin.H{"error": "请求体须为 JSON，例如 {\"note\":\"说明\"}，可为空对象"}) + return + } + note = strings.TrimSpace(body.Note) + + convID := h.conversationIDForRunningExecution(id) + if convID != "" && h.agentHandler != nil { + if ok, payload := h.agentHandler.cancelToolContinueAfter(convID, id, note); ok { + h.logger.Info("MCP 监控页终止工具（与对话中断并继续一致）", + zap.String("executionId", id), + zap.String("conversationId", convID), + zap.Bool("hasNote", note != ""), + ) + c.JSON(http.StatusOK, payload) + return + } + } + if h.mcpServer.CancelToolExecutionWithNote(id, note) { + h.logger.Info("已请求取消 MCP 工具执行", zap.String("executionId", id), zap.String("source", "internal"), zap.Bool("hasNote", note != "")) + c.JSON(http.StatusOK, gin.H{"message": "已发送终止信号", "executionId": id}) + return + } + if h.externalMCPMgr != nil && h.externalMCPMgr.CancelToolExecutionWithNote(id, note) { + h.logger.Info("已请求取消 MCP 工具执行", zap.String("executionId", id), zap.String("source", "external"), zap.Bool("hasNote", note != "")) + c.JSON(http.StatusOK, gin.H{"message": "已发送终止信号", "executionId": id}) + return + } + c.JSON(http.StatusNotFound, gin.H{"error": "未找到进行中的工具执行，或该任务已结束"}) +} + +func (h *MonitorHandler) enrichExecutionsConversationID(executions []*mcp.ToolExecution) { + for _, exec := range executions { + if exec == nil || exec.Status != "running" { + continue + } + exec.ConversationID = h.conversationIDForRunningExecution(exec.ID) + } +} + +func (h *MonitorHandler) conversationIDForRunningExecution(executionID string) string { + executionID = strings.TrimSpace(executionID) + if executionID == "" || h.taskManager == nil { + return "" + } + if conv := h.taskManager.ConversationIDForActiveMCPExecution(executionID); conv != "" { + return conv + } + exec := h.lookupExecution(executionID) + if exec == nil || exec.Status != "running" { + return "" + } + if strings.TrimSpace(exec.ToolName) == "execute" { + if onlyConv, ok := h.taskManager.ConversationIDForActiveEinoExecute(); ok { + return onlyConv + } + } + return "" +} + +func (h *MonitorHandler) lookupExecution(id string) *mcp.ToolExecution { + if exec, ok := h.mcpServer.GetExecution(id); ok { + return exec + } + if h.externalMCPMgr != nil { + if exec, ok := h.externalMCPMgr.GetExecution(id); ok { + return exec + } + } + if h.db != nil { + if exec, err := h.db.GetToolExecution(id); err == nil && exec != nil { + return exec + } + } + return nil +} + +// BatchGetToolNames 批量获取工具执行摘要（消除前端 N+1 请求） +func (h *MonitorHandler) BatchGetToolNames(c *gin.Context) { + var req struct { + IDs []string `json:"ids"` + } + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + type executionSummary struct { + ToolName string `json:"toolName"` + Status string `json:"status"` + } + + result := make(map[string]executionSummary, len(req.IDs)) + for _, id := range req.IDs { + if !h.monitorExecutionAllowed(c, id) { + continue + } + // 先从内部MCP服务器查找 + if exec, exists := h.mcpServer.GetExecution(id); exists { + result[id] = executionSummary{ToolName: exec.ToolName, Status: exec.Status} + continue + } + // 再从外部MCP管理器查找 + if h.externalMCPMgr != nil { + if exec, exists := h.externalMCPMgr.GetExecution(id); exists { + result[id] = executionSummary{ToolName: exec.ToolName, Status: exec.Status} + continue + } + } + // 最后从数据库查找 + if h.db != nil { + if exec, err := h.db.GetToolExecution(id); err == nil && exec != nil { + result[id] = executionSummary{ToolName: exec.ToolName, Status: exec.Status} + } + } + } + + c.JSON(http.StatusOK, result) +} + +// GetStats 获取统计信息 +func (h *MonitorHandler) GetStats(c *gin.Context) { + topN := 30 + if topStr := c.Query("top"); topStr != "" { + if t, err := strconv.Atoi(topStr); err == nil && t > 0 && t <= 100 { + topN = t + } + } + summary, topTools := h.loadStatsSummary(topN) + c.JSON(http.StatusOK, StatsResponse{ + Summary: summary, + TopTools: topTools, + }) +} + +// CallsTimelinePoint 调用趋势数据点 +type CallsTimelinePoint struct { + T time.Time `json:"t"` + Total int `json:"total"` + Failed int `json:"failed"` +} + +// CallsTimelineSummary 调用趋势汇总 +type CallsTimelineSummary struct { + TotalCalls int `json:"totalCalls"` + Peak int `json:"peak"` +} + +// CallsTimelineResponse 调用趋势响应 +type CallsTimelineResponse struct { + Range string `json:"range"` + Points []CallsTimelinePoint `json:"points"` + Summary CallsTimelineSummary `json:"summary"` +} + +type callsTimelineConfig struct { + rangeKey string + duration time.Duration + bucketSize time.Duration + dailyBuckets bool +} + +func parseCallsTimelineRange(raw string) (callsTimelineConfig, bool) { + switch strings.TrimSpace(raw) { + case "24h": + return callsTimelineConfig{rangeKey: "24h", duration: 24 * time.Hour, bucketSize: time.Hour, dailyBuckets: false}, true + case "30d": + return callsTimelineConfig{rangeKey: "30d", duration: 30 * 24 * time.Hour, bucketSize: 24 * time.Hour, dailyBuckets: true}, true + default: + return callsTimelineConfig{rangeKey: "7d", duration: 7 * 24 * time.Hour, bucketSize: time.Hour, dailyBuckets: false}, true + } +} + +func truncateToBucket(t time.Time, bucketSize time.Duration, dailyBuckets bool) time.Time { + if dailyBuckets { + y, m, d := t.Date() + return time.Date(y, m, d, 0, 0, 0, 0, t.Location()) + } + return t.Truncate(bucketSize) +} + +func buildCallsTimelinePoints(cfg callsTimelineConfig, buckets map[time.Time]struct{ total, failed int }) []CallsTimelinePoint { + now := time.Now() + start := truncateToBucket(now.Add(-cfg.duration), cfg.bucketSize, cfg.dailyBuckets) + end := truncateToBucket(now, cfg.bucketSize, cfg.dailyBuckets) + + points := make([]CallsTimelinePoint, 0) + for current := start; !current.After(end); current = current.Add(cfg.bucketSize) { + val := buckets[current] + points = append(points, CallsTimelinePoint{ + T: current, + Total: val.total, + Failed: val.failed, + }) + } + return points +} + +func (h *MonitorHandler) loadCallsTimeline(cfg callsTimelineConfig) []CallsTimelinePoint { + since := time.Now().Add(-cfg.duration) + bucketMap := make(map[time.Time]struct{ total, failed int }) + + if h.db != nil { + dbBuckets, err := h.db.LoadCallsTimeline(since, cfg.dailyBuckets) + if err != nil { + h.logger.Warn("从数据库加载调用趋势失败，回退到内存数据", zap.Error(err)) + } else { + for _, b := range dbBuckets { + key := truncateToBucket(b.BucketTime, cfg.bucketSize, cfg.dailyBuckets) + entry := bucketMap[key] + entry.total += b.Total + entry.failed += b.Failed + bucketMap[key] = entry + } + return buildCallsTimelinePoints(cfg, bucketMap) + } + } + + for _, exec := range h.mcpServer.GetAllExecutions() { + if exec == nil || exec.StartTime.Before(since) { + continue + } + key := truncateToBucket(exec.StartTime, cfg.bucketSize, cfg.dailyBuckets) + entry := bucketMap[key] + entry.total++ + if monitorStatusCountsAsFailed(exec.Status) { + entry.failed++ + } + bucketMap[key] = entry + } + return buildCallsTimelinePoints(cfg, bucketMap) +} + +// GetCallsTimeline 获取 MCP 工具调用趋势 +func (h *MonitorHandler) GetCallsTimeline(c *gin.Context) { + cfg, _ := parseCallsTimelineRange(c.Query("range")) + points := h.loadCallsTimeline(cfg) + + summary := CallsTimelineSummary{} + for _, p := range points { + summary.TotalCalls += p.Total + if p.Total > summary.Peak { + summary.Peak = p.Total + } + } + + c.JSON(http.StatusOK, CallsTimelineResponse{ + Range: cfg.rangeKey, + Points: points, + Summary: summary, + }) +} + +// DeleteExecution 删除执行记录 +func (h *MonitorHandler) DeleteExecution(c *gin.Context) { + id := c.Param("id") + if id == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "执行记录ID不能为空"}) + return + } + if !h.monitorExecutionAllowed(c, id) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + + // 如果使用数据库，先获取执行记录信息，然后删除并更新统计 + if h.db != nil { + // 先获取执行记录信息（用于更新统计） + exec, err := h.db.GetToolExecution(id) + if err != nil { + // 如果找不到记录，可能已经被删除，直接返回成功 + h.logger.Warn("执行记录不存在，可能已被删除", zap.String("executionId", id), zap.Error(err)) + c.JSON(http.StatusOK, gin.H{"message": "执行记录不存在或已被删除"}) + return + } + + // 删除执行记录 + err = h.db.DeleteToolExecution(id) + if err != nil { + h.logger.Error("删除执行记录失败", zap.Error(err), zap.String("executionId", id)) + c.JSON(http.StatusInternalServerError, gin.H{"error": "删除执行记录失败: " + err.Error()}) + return + } + + // 更新统计信息（减少相应的计数） + totalCalls := 1 + successCalls := 0 + failedCalls := 0 + if monitorStatusCountsAsFailed(exec.Status) { + failedCalls = 1 + } else if exec.Status == "completed" { + successCalls = 1 + } + + if exec.ToolName != "" { + if err := h.db.DecreaseToolStats(exec.ToolName, totalCalls, successCalls, failedCalls); err != nil { + h.logger.Warn("更新统计信息失败", zap.Error(err), zap.String("toolName", exec.ToolName)) + // 不返回错误，因为记录已经删除成功 + } + } + + h.logger.Info("执行记录已从数据库删除", zap.String("executionId", id), zap.String("toolName", exec.ToolName)) + if h.audit != nil { + h.audit.RecordOK(c, "tool", "execution_delete", "删除工具执行记录", "tool_execution", id, map[string]interface{}{ + "tool_name": exec.ToolName, + }) + } + c.JSON(http.StatusOK, gin.H{"message": "执行记录已删除"}) + return + } + + // 如果不使用数据库，尝试从内存中删除（内部MCP服务器） + // 注意：内存中的记录可能已经被清理，所以这里只记录日志 + h.logger.Info("尝试删除内存中的执行记录", zap.String("executionId", id)) + c.JSON(http.StatusOK, gin.H{"message": "执行记录已删除（如果存在）"}) +} + +// DeleteExecutions 批量删除执行记录 +func (h *MonitorHandler) DeleteExecutions(c *gin.Context) { + var request struct { + IDs []string `json:"ids"` + } + + if err := c.ShouldBindJSON(&request); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "请求参数无效: " + err.Error()}) + return + } + + if len(request.IDs) == 0 { + c.JSON(http.StatusBadRequest, gin.H{"error": "执行记录ID列表不能为空"}) + return + } + for _, id := range request.IDs { + if !h.monitorExecutionAllowed(c, id) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问一个或多个执行记录"}) + return + } + } + + // 如果使用数据库，先获取执行记录信息，然后删除并更新统计 + if h.db != nil { + // 先获取执行记录信息（用于更新统计） + executions, err := h.db.GetToolExecutionsByIds(request.IDs) + if err != nil { + h.logger.Error("获取执行记录失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": "获取执行记录失败: " + err.Error()}) + return + } + + // 按工具名称分组统计需要减少的数量 + toolStats := make(map[string]struct { + totalCalls int + successCalls int + failedCalls int + }) + + for _, exec := range executions { + if exec.ToolName == "" { + continue + } + + stats := toolStats[exec.ToolName] + stats.totalCalls++ + if monitorStatusCountsAsFailed(exec.Status) { + stats.failedCalls++ + } else if exec.Status == "completed" { + stats.successCalls++ + } + toolStats[exec.ToolName] = stats + } + + // 批量删除执行记录 + err = h.db.DeleteToolExecutions(request.IDs) + if err != nil { + h.logger.Error("批量删除执行记录失败", zap.Error(err), zap.Int("count", len(request.IDs))) + c.JSON(http.StatusInternalServerError, gin.H{"error": "批量删除执行记录失败: " + err.Error()}) + return + } + + // 更新统计信息（减少相应的计数） + for toolName, stats := range toolStats { + if err := h.db.DecreaseToolStats(toolName, stats.totalCalls, stats.successCalls, stats.failedCalls); err != nil { + h.logger.Warn("更新统计信息失败", zap.Error(err), zap.String("toolName", toolName)) + // 不返回错误，因为记录已经删除成功 + } + } + + h.logger.Info("批量删除执行记录成功", zap.Int("count", len(request.IDs))) + if h.audit != nil { + h.audit.RecordOK(c, "tool", "execution_delete_batch", "批量删除工具执行记录", "tool_execution", "", map[string]interface{}{ + "count": len(request.IDs), + }) + } + c.JSON(http.StatusOK, gin.H{"message": "成功删除执行记录", "deleted": len(executions)}) + return + } + + // 如果不使用数据库，尝试从内存中删除（内部MCP服务器） + // 注意：内存中的记录可能已经被清理，所以这里只记录日志 + h.logger.Info("尝试批量删除内存中的执行记录", zap.Int("count", len(request.IDs))) + c.JSON(http.StatusOK, gin.H{"message": "执行记录已删除（如果存在）"}) +} + +// normalizeToolNameFilter 将模型侧 mcp__tool 转为内部存储用的 mcp::tool。 +func normalizeToolNameFilter(name string) string { + name = strings.TrimSpace(name) + if name == "" { + return name + } + if strings.Contains(name, "::") { + return name + } + if idx := strings.Index(name, "__"); idx > 0 { + return name[:idx] + "::" + name[idx+2:] + } + return name +} + +func toolNameFilterMatches(storedName, filter string) bool { + filter = strings.TrimSpace(filter) + if filter == "" { + return true + } + storedLower := strings.ToLower(storedName) + filterLower := strings.ToLower(filter) + if strings.Contains(storedLower, filterLower) { + return true + } + normFilter := strings.ToLower(normalizeToolNameFilter(filter)) + if normFilter != filterLower && strings.Contains(storedLower, normFilter) { + return true + } + return strings.Contains(strings.ReplaceAll(storedLower, "::", "__"), filterLower) +} diff --git a/internal/handler/multi_agent.go b/internal/handler/multi_agent.go new file mode 100644 index 00000000..d68f0c28 --- /dev/null +++ b/internal/handler/multi_agent.go @@ -0,0 +1,628 @@ +package handler + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "strings" + "sync" + "time" + + "cyberstrike-ai/internal/agentfinalizer" + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/mcp" + "cyberstrike-ai/internal/multiagent" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +// MultiAgentLoopStream Eino DeepAgent 流式对话（需 config.multi_agent.enabled）。 +func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) { + c.Header("Content-Type", "text/event-stream; charset=utf-8") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + if h.config == nil || !h.config.MultiAgent.Enabled { + ev := StreamEvent{Type: "error", Message: "多代理未启用，请在设置或 config.yaml 中开启 multi_agent.enabled"} + b, _ := json.Marshal(ev) + fmt.Fprintf(c.Writer, "data: %s\n\n", b) + done := StreamEvent{Type: "done", Message: ""} + db, _ := json.Marshal(done) + fmt.Fprintf(c.Writer, "data: %s\n\n", db) + if flusher, ok := c.Writer.(http.Flusher); ok { + flusher.Flush() + } + return + } + + var req ChatRequest + if err := c.ShouldBindJSON(&req); err != nil { + event := StreamEvent{Type: "error", Message: "请求参数错误: " + err.Error()} + b, _ := json.Marshal(event) + fmt.Fprintf(c.Writer, "data: %s\n\n", b) + done := StreamEvent{Type: "done", Message: ""} + db, _ := json.Marshal(done) + fmt.Fprintf(c.Writer, "data: %s\n\n", db) + c.Writer.Flush() + return + } + + c.Header("X-Accel-Buffering", "no") + + // 用于在 sendEvent 中判断是否为用户主动停止导致的取消。 + // 注意：baseCtx 会在后面创建；该变量用于闭包提前捕获引用。 + var baseCtx context.Context + + clientDisconnected := false + // 与 sseKeepalive 共用：禁止并发写 ResponseWriter，否则会破坏 chunked 编码（ERR_INVALID_CHUNKED_ENCODING）。 + var sseWriteMu sync.Mutex + var ssePublishConversationID string + sendEvent := func(eventType, message string, data interface{}) { + // 用户主动停止时，Eino 可能仍会并发上报 eventType=="error"。 + // 为避免 UI 看到“取消错误 + cancelled 文案”两条回复，这里直接丢弃取消对应的 error。 + if eventType == "error" && baseCtx != nil { + cause := context.Cause(baseCtx) + if errors.Is(cause, ErrTaskCancelled) || errors.Is(cause, multiagent.ErrInterruptContinue) { + return + } + } + ev := StreamEvent{Type: eventType, Message: message, Data: data} + b, errMarshal := json.Marshal(ev) + if errMarshal != nil { + b = []byte(`{"type":"error","message":"marshal failed"}`) + } + sseLine := make([]byte, 0, len(b)+8) + sseLine = append(sseLine, []byte("data: ")...) + sseLine = append(sseLine, b...) + sseLine = append(sseLine, '\n', '\n') + if ssePublishConversationID != "" && h.taskEventBus != nil { + h.taskEventBus.Publish(ssePublishConversationID, sseLine) + } + if clientDisconnected { + return + } + select { + case <-c.Request.Context().Done(): + clientDisconnected = true + return + default: + } + sseWriteMu.Lock() + _, err := c.Writer.Write(sseLine) + if err != nil { + sseWriteMu.Unlock() + clientDisconnected = true + return + } + if flusher, ok := c.Writer.(http.Flusher); ok { + flusher.Flush() + } else { + c.Writer.Flush() + } + sseWriteMu.Unlock() + } + + h.logger.Info("收到 Eino DeepAgent 流式请求", + zap.String("conversationId", req.ConversationID), + ) + + prep, err := h.prepareMultiAgentSession(&req, c, "multi_agent_stream") + if err != nil { + sendEvent("error", err.Error(), nil) + sendEvent("done", "", nil) + return + } + ssePublishConversationID = prep.ConversationID + if prep.CreatedNew { + sendEvent("conversation", "会话已创建", map[string]interface{}{ + "conversationId": prep.ConversationID, + }) + } + + conversationID := prep.ConversationID + assistantMessageID := prep.AssistantMessageID + h.activateHITLForConversation(conversationID, req.Hitl) + if h.hitlManager != nil { + defer h.hitlManager.DeactivateConversation(conversationID) + } + + if prep.UserMessageID != "" { + sendEvent("message_saved", "", map[string]interface{}{ + "conversationId": conversationID, + "userMessageId": prep.UserMessageID, + }) + } + if h.runRoleWorkflowStreamIfBound(c, &req, prep, sendEvent) { + return + } + + var cancelWithCause context.CancelCauseFunc + curFinalMessage := prep.FinalMessage + curHistory := prep.History + roleTools := prep.RoleTools + orch := strings.TrimSpace(req.Orchestration) + + taskStatus := "completed" + // 仅在成功 StartTask 后再 FinishTask；避免「任务已存在」分支 return 时误删正在运行的同会话任务。 + taskOwned := false + defer func() { + if taskOwned { + h.tasks.FinishTask(conversationID, taskStatus) + } + }() + + sendEvent("progress", "正在启动 Eino 多代理...", map[string]interface{}{ + "conversationId": conversationID, + }) + + stopKeepalive := runSSEKeepalive(c, &sseWriteMu) + defer stopKeepalive() + runCfg, _, err := h.configForAIChannel(req.AIChannelID) + if err != nil { + sendEvent("error", err.Error(), nil) + sendEvent("done", "", map[string]interface{}{"conversationId": conversationID}) + return + } + + var result *multiagent.RunResult + var runErr error + + baseCtx, cancelWithCause = context.WithCancelCause(detachedAgentContext(c.Request.Context())) + taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute) + + if _, err := h.tasks.StartTask(conversationID, req.Message, cancelWithCause); err != nil { + var errorMsg string + if errors.Is(err, ErrTaskAlreadyRunning) { + errorMsg = "⚠️ 当前会话已有任务正在执行中，请等待当前任务完成或点击「停止任务」后再尝试。" + sendEvent("error", errorMsg, map[string]interface{}{ + "conversationId": conversationID, + "errorType": "task_already_running", + }) + } else { + errorMsg = "❌ 无法启动任务: " + err.Error() + sendEvent("error", errorMsg, nil) + } + if assistantMessageID != "" { + _, _ = h.db.Exec("UPDATE messages SET content = ?, updated_at = ? WHERE id = ?", errorMsg, time.Now(), assistantMessageID) + } + sendEvent("done", "", map[string]interface{}{"conversationId": conversationID}) + timeoutCancel() + return + } + taskOwned = true + + // 同一 HTTP 流内多段 Run（如中断并继续）合并 MCP execution id，供最终 response / 库表与工具芯片展示完整列表 + var cumulativeMCPExecutionIDs []string + // 同一请求内分段续跑时，主代理 iteration 事件按偏移累计，避免 UI 出现「第3轮 → 第1轮」回跳。 + var mainIterationOffset int + var emptyResponseContinueAttempt int + var finalizationAutoContinueAttempt int + effectiveOrch := config.NormalizeMultiAgentOrchestration(h.config.MultiAgent.Orchestration) + if o := strings.TrimSpace(req.Orchestration); o != "" { + effectiveOrch = config.NormalizeMultiAgentOrchestration(o) + } + agentMode := "eino_" + effectiveOrch + var decision agentfinalizer.Decision + + for { + segmentMainIterationMax := 0 + rawProgressCallback := h.createProgressCallback(taskCtx, cancelWithCause, conversationID, assistantMessageID, sendEvent) + progressCallback := func(eventType, message string, data interface{}) { + if eventType == "iteration" { + if m, ok := data.(map[string]interface{}); ok { + if scope, _ := m["einoScope"].(string); scope == "main" { + raw := 0 + switch v := m["iteration"].(type) { + case int: + raw = v + case int32: + raw = int(v) + case int64: + raw = int(v) + case float64: + raw = int(v) + case float32: + raw = int(v) + } + if raw > 0 { + if raw > segmentMainIterationMax { + segmentMainIterationMax = raw + } + m["iteration"] = raw + mainIterationOffset + } + } + } + } + rawProgressCallback(eventType, message, data) + } + taskCtxLoop := mcp.WithMCPConversationID(taskCtx, conversationID) + taskCtxLoop = mcp.WithToolRunRegistry(taskCtxLoop, h.tasks) + taskCtxLoop = mcp.WithEinoExecuteRunRegistry(taskCtxLoop, h.tasks) + taskCtxLoop = multiagent.WithAgentRuntimeCancelRegistrar(taskCtxLoop, func(cancel func(error) bool) func() { + return h.tasks.BindAgentRuntimeCancel(conversationID, cancel) + }) + taskCtxLoop = multiagent.WithAgentTurnLoopInterruptRegistrar(taskCtxLoop, func(push func(string) bool) func() { + return h.tasks.BindAgentTurnLoopInterrupt(conversationID, push) + }) + taskCtxLoop = multiagent.WithHITLToolInterceptor(taskCtxLoop, func(ctx context.Context, toolName, arguments string) (string, error) { + return h.interceptHITLForEinoTool(ctx, cancelWithCause, conversationID, assistantMessageID, sendEvent, toolName, arguments) + }) + + result, runErr = multiagent.RunDeepAgent( + taskCtxLoop, + runCfg, + &runCfg.MultiAgent, + h.agent, + h.db, + h.logger, + conversationID, + h.conversationProjectID(conversationID), + curFinalMessage, + curHistory, + roleTools, + progressCallback, + h.agentsMarkdownDir, + orch, + chatReasoningToClientIntent(req.Reasoning), + h.agentSessionContextBlock(conversationID), + ) + + if result != nil && len(result.MCPExecutionIDs) > 0 { + cumulativeMCPExecutionIDs = mergeMCPExecutionIDLists(cumulativeMCPExecutionIDs, result.MCPExecutionIDs) + } + + if runErr == nil { + mw := &h.config.MultiAgent.EinoMiddleware + if h.tryContinueOnEinoEmptyResponse(taskCtx, mw, conversationID, result, &emptyResponseContinueAttempt, &curHistory, &curFinalMessage, progressCallback) { + mainIterationOffset += segmentMainIterationMax + timeoutCancel() + baseCtx, cancelWithCause, taskCtx, timeoutCancel = h.rebindEinoRunningTask(taskCtx, conversationID, timeoutCancel) + continue + } + decision = h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, agentMode, result, cumulativeMCPExecutionIDs, requestRequiresExecutionEvidence(&req)) + if h.tryAutoContinueAfterFinalization(taskCtx, conversationID, result, decision, &finalizationAutoContinueAttempt, &curHistory, &curFinalMessage, progressCallback) { + mainIterationOffset += segmentMainIterationMax + timeoutCancel() + baseCtx, cancelWithCause, taskCtx, timeoutCancel = h.rebindEinoRunningTask(taskCtx, conversationID, timeoutCancel) + continue + } + timeoutCancel() + break + } + + cause := context.Cause(baseCtx) + if cause == nil { + switch { + case errors.Is(runErr, multiagent.ErrInterruptContinue): + cause = multiagent.ErrInterruptContinue + case errors.Is(runErr, ErrTaskCancelled): + cause = ErrTaskCancelled + } + } + if errors.Is(cause, multiagent.ErrInterruptContinue) { + if shouldPersistEinoAgentTraceAfterRunError(baseCtx) { + h.persistEinoAgentTraceForResume(conversationID, result) + } + note := h.tasks.TakeInterruptContinueNote(conversationID) + icSummary := interruptContinueTimelineSummary(note) + progressCallback("user_interrupt_continue", icSummary, map[string]interface{}{ + "conversationId": conversationID, + "rawReason": strings.TrimSpace(note), + "emptyReason": strings.TrimSpace(note) == "", + "kind": "no_active_mcp_tool", + }) + inject := formatInterruptContinueUserMessage(note) + // 不写入 messages 表为 user 气泡：避免主对话流出现大段模板；说明已由 user_interrupt_continue 记入助手 process_details（迭代详情）。 + if hist, err := h.loadHistoryFromAgentTrace(conversationID); err == nil && len(hist) > 0 { + curHistory = hist + } + curFinalMessage = inject + sendEvent("progress", "已合并用户补充与最新轨迹，正在继续推理…", map[string]interface{}{ + "conversationId": conversationID, + "source": "interrupt_continue", + }) + mainIterationOffset += segmentMainIterationMax + timeoutCancel() + baseCtx, cancelWithCause = context.WithCancelCause(detachedAgentContext(baseCtx)) + h.tasks.BindTaskCancel(conversationID, cancelWithCause) + taskCtx, timeoutCancel = context.WithTimeout(baseCtx, 600*time.Minute) + h.tasks.UpdateTaskStatus(conversationID, "running") + continue + } + + if shouldPersistEinoAgentTraceAfterRunError(baseCtx) { + h.persistEinoAgentTraceForResume(conversationID, result) + } + if errors.Is(cause, ErrTaskCancelled) { + taskStatus = "cancelled" + h.tasks.UpdateTaskStatus(conversationID, taskStatus) + cancelMsg := "任务已被用户取消，后续操作已停止。" + if assistantMessageID != "" { + if result != nil { + if err := h.mergeAssistantMessagePartialOnCancel(assistantMessageID, result.Response); err != nil { + h.logger.Warn("合并取消前的部分回复失败", zap.Error(err)) + } + } + if err := h.appendAssistantMessageNotice(assistantMessageID, cancelMsg); err != nil { + h.logger.Warn("更新取消后的助手消息失败", zap.Error(err)) + } + _ = h.db.AddProcessDetail(assistantMessageID, conversationID, "cancelled", cancelMsg, nil) + } + sendEvent("cancelled", cancelMsg, map[string]interface{}{ + "conversationId": conversationID, + "messageId": assistantMessageID, + }) + sendEvent("done", "", map[string]interface{}{"conversationId": conversationID}) + timeoutCancel() + return + } + + if errors.Is(runErr, context.DeadlineExceeded) || errors.Is(context.Cause(taskCtx), context.DeadlineExceeded) { + taskStatus = "timeout" + h.tasks.UpdateTaskStatus(conversationID, taskStatus) + timeoutMsg := "任务执行超时，已自动终止。" + if assistantMessageID != "" { + _, _ = h.db.Exec("UPDATE messages SET content = ?, updated_at = ? WHERE id = ?", timeoutMsg, time.Now(), assistantMessageID) + _ = h.db.AddProcessDetail(assistantMessageID, conversationID, "timeout", timeoutMsg, nil) + } + sendEvent("error", timeoutMsg, map[string]interface{}{ + "conversationId": conversationID, + "messageId": assistantMessageID, + "errorType": "timeout", + }) + sendEvent("done", "", map[string]interface{}{"conversationId": conversationID}) + timeoutCancel() + return + } + + h.logger.Error("Eino DeepAgent 执行失败", zap.Error(runErr)) + taskStatus = "failed" + h.tasks.UpdateTaskStatus(conversationID, taskStatus) + errMsg := "执行失败: " + runErr.Error() + if assistantMessageID != "" { + _, _ = h.db.Exec("UPDATE messages SET content = ?, updated_at = ? WHERE id = ?", errMsg, time.Now(), assistantMessageID) + _ = h.db.AddProcessDetail(assistantMessageID, conversationID, "error", errMsg, nil) + } + sendEvent("error", errMsg, map[string]interface{}{ + "conversationId": conversationID, + "messageId": assistantMessageID, + }) + sendEvent("done", "", map[string]interface{}{"conversationId": conversationID}) + timeoutCancel() + return + } + + timeoutCancel() + + if decision.CompletionReason == "" { + decision = h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, agentMode, result, cumulativeMCPExecutionIDs, requestRequiresExecutionEvidence(&req)) + } + h.persistFinalizationDecision(conversationID, assistantMessageID, agentMode, cumulativeMCPExecutionIDs, multiagent.AggregatedReasoningFromTraceJSON(result.LastAgentTraceInput), decision) + + if result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "" { + if err := h.db.SaveAgentTrace(conversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); err != nil { + h.logger.Warn("保存代理轨迹失败", zap.Error(err)) + } + } + + responseText := decision.FinalText + if !decision.Finalizable { + responseText = finalizationBlockedMessage(decision) + sendEvent("finalization_check", responseText, decision) + taskStatus = decision.Status + h.tasks.UpdateTaskStatus(conversationID, taskStatus) + } + sendEvent("response", responseText, finalizationResponsePayload(decision, map[string]interface{}{ + "mcpExecutionIds": cumulativeMCPExecutionIDs, + "conversationId": conversationID, + "messageId": assistantMessageID, + "agentMode": agentMode, + })) + sendEvent("done", "", map[string]interface{}{"conversationId": conversationID}) +} + +// MultiAgentLoop Eino DeepAgent 非流式对话（需 multi_agent.enabled）。 +func (h *AgentHandler) MultiAgentLoop(c *gin.Context) { + if h.config == nil || !h.config.MultiAgent.Enabled { + c.JSON(http.StatusNotFound, gin.H{"error": "多代理未启用，请在 config.yaml 中设置 multi_agent.enabled: true"}) + return + } + + var req ChatRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + h.logger.Info("收到 Eino DeepAgent 非流式请求", zap.String("conversationId", req.ConversationID)) + + prep, err := h.prepareMultiAgentSession(&req, c, "multi_agent") + if err != nil { + status, msg := multiAgentHTTPErrorStatus(err) + c.JSON(status, gin.H{"error": msg}) + return + } + h.activateHITLForConversation(prep.ConversationID, req.Hitl) + if h.hitlManager != nil { + defer h.hitlManager.DeactivateConversation(prep.ConversationID) + } + if h.runRoleWorkflowJSONIfBound(c, &req, prep) { + return + } + + baseCtx, cancelWithCause := context.WithCancelCause(c.Request.Context()) + defer cancelWithCause(nil) + taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute) + defer timeoutCancel() + progressCallback := h.createProgressCallback(taskCtx, cancelWithCause, prep.ConversationID, prep.AssistantMessageID, nil) + taskCtx = multiagent.WithHITLToolInterceptor(taskCtx, func(ctx context.Context, toolName, arguments string) (string, error) { + return h.interceptHITLForEinoTool(ctx, cancelWithCause, prep.ConversationID, prep.AssistantMessageID, nil, toolName, arguments) + }) + runCfg, _, err := h.configForAIChannel(req.AIChannelID) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + curHist := prep.History + curMsg := prep.FinalMessage + var result *multiagent.RunResult + var runErr error + var emptyResponseContinueAttempt int + var finalizationAutoContinueAttempt int + effectiveOrch := config.NormalizeMultiAgentOrchestration(h.config.MultiAgent.Orchestration) + if o := strings.TrimSpace(req.Orchestration); o != "" { + effectiveOrch = config.NormalizeMultiAgentOrchestration(o) + } + agentMode := "eino_" + effectiveOrch + var decision agentfinalizer.Decision + for { + result, runErr = multiagent.RunDeepAgent( + taskCtx, + runCfg, + &runCfg.MultiAgent, + h.agent, + h.db, + h.logger, + prep.ConversationID, + h.conversationProjectID(prep.ConversationID), + curMsg, + curHist, + prep.RoleTools, + progressCallback, + h.agentsMarkdownDir, + strings.TrimSpace(req.Orchestration), + chatReasoningToClientIntent(req.Reasoning), + h.agentSessionContextBlock(prep.ConversationID), + ) + if runErr != nil { + if shouldPersistEinoAgentTraceAfterRunError(baseCtx) { + h.persistEinoAgentTraceForResume(prep.ConversationID, result) + } + h.logger.Error("Eino DeepAgent 执行失败", zap.Error(runErr)) + errMsg := "执行失败: " + runErr.Error() + if prep.AssistantMessageID != "" { + _, _ = h.db.Exec("UPDATE messages SET content = ?, updated_at = ? WHERE id = ?", errMsg, time.Now(), prep.AssistantMessageID) + } + c.JSON(http.StatusInternalServerError, gin.H{"error": errMsg}) + return + } + mw := &h.config.MultiAgent.EinoMiddleware + if h.tryContinueOnEinoEmptyResponse(taskCtx, mw, prep.ConversationID, result, &emptyResponseContinueAttempt, &curHist, &curMsg, progressCallback) { + continue + } + decision = h.decideAgentRunForDeliveryWithPolicy(prep.ConversationID, prep.AssistantMessageID, agentMode, result, result.MCPExecutionIDs, requestRequiresExecutionEvidence(&req)) + if h.tryAutoContinueAfterFinalization(taskCtx, prep.ConversationID, result, decision, &finalizationAutoContinueAttempt, &curHist, &curMsg, progressCallback) { + continue + } + break + } + + h.persistFinalizationDecision(prep.ConversationID, prep.AssistantMessageID, agentMode, result.MCPExecutionIDs, multiagent.AggregatedReasoningFromTraceJSON(result.LastAgentTraceInput), decision) + + if result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "" { + if err := h.db.SaveAgentTrace(prep.ConversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); err != nil { + h.logger.Warn("保存代理轨迹失败", zap.Error(err)) + } + } + + responseText := decision.FinalText + if !decision.Finalizable { + responseText = finalizationBlockedMessage(decision) + } + c.JSON(http.StatusOK, ChatResponse{ + Response: responseText, + MCPExecutionIDs: result.MCPExecutionIDs, + ConversationID: prep.ConversationID, + Time: time.Now(), + Finalizable: decision.Finalizable, + Finalized: decision.Finalized, + Status: decision.Status, + CompletionReason: decision.CompletionReason, + EvidenceVerified: decision.EvidenceVerified, + EvidenceRefs: decision.EvidenceRefs, + PendingExecutionIDs: decision.PendingExecutionIDs, + MissingChecks: decision.MissingChecks, + }) +} + +// persistEinoAgentTraceForResume 在 Eino 运行异常结束时写入代理轨迹（库列 last_react_*），供下一请求 loadHistoryFromAgentTrace 软续跑。 +func (h *AgentHandler) persistEinoAgentTraceForResume(conversationID string, result *multiagent.RunResult) { + if h == nil || result == nil { + return + } + if result.LastAgentTraceInput == "" && result.LastAgentTraceOutput == "" { + return + } + if err := h.db.SaveAgentTrace(conversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); err != nil { + h.logger.Warn("保存 Eino 续跑上下文失败", zap.String("conversationId", conversationID), zap.Error(err)) + } +} + +// mergeMCPExecutionIDLists 去重合并多段 Run 的 MCP execution id（顺序：先 dst 后 more）。 +func mergeMCPExecutionIDLists(dst []string, more []string) []string { + seen := make(map[string]struct{}, len(dst)+len(more)) + out := make([]string, 0, len(dst)+len(more)) + add := func(ids []string) { + for _, id := range ids { + id = strings.TrimSpace(id) + if id == "" { + continue + } + if _, ok := seen[id]; ok { + continue + } + seen[id] = struct{}{} + out = append(out, id) + } + } + add(dst) + add(more) + return out +} + +// interruptContinueTimelineSummary 时间线 / process_details 中展示的简短正文（完整模板已写入另一条用户消息）。 +func interruptContinueTimelineSummary(note string) string { + note = strings.TrimSpace(note) + if note == "" { + return "用户选择「中断并继续」，未填写说明；已按默认渗透补充模板合并上下文并续跑。" + } + return "用户中断说明（原文）：\n\n" + note +} + +// formatInterruptContinueUserMessage 将「中断并继续」弹窗中的说明格式化为新一轮 user 消息（渗透场景下强调路径补充与端口复扫）。 +func formatInterruptContinueUserMessage(note string) string { + var b strings.Builder + b.WriteString("【用户补充 / 中断后继续】\n") + if s := strings.TrimSpace(note); s != "" { + b.WriteString(s) + b.WriteString("\n\n") + } + b.WriteString("【请在本轮落实】\n") + b.WriteString("- 将用户提供的接口路径、参数、业务变化纳入后续测试与推理。\n") + b.WriteString("- 若资产或目标信息有更新，请对目标重新执行端口/服务探测，再基于新结果规划下一步。\n") + b.WriteString("- 在已有轨迹基础上推进，避免无意义重复已完成的步骤。\n") + return strings.TrimSpace(b.String()) +} + +func multiAgentHTTPErrorStatus(err error) (int, string) { + msg := err.Error() + switch { + case strings.Contains(msg, "无权访问"): + return http.StatusForbidden, msg + case strings.Contains(msg, "对话不存在"): + return http.StatusNotFound, msg + case strings.Contains(msg, "未找到该 WebShell"): + return http.StatusBadRequest, msg + case strings.Contains(msg, "附件最多"): + return http.StatusBadRequest, msg + case strings.Contains(msg, "保存用户消息失败"), strings.Contains(msg, "创建对话失败"): + return http.StatusInternalServerError, msg + case strings.Contains(msg, "保存上传文件失败"): + return http.StatusInternalServerError, msg + default: + return http.StatusBadRequest, msg + } +} diff --git a/internal/handler/multi_agent_prepare.go b/internal/handler/multi_agent_prepare.go new file mode 100644 index 00000000..08564beb --- /dev/null +++ b/internal/handler/multi_agent_prepare.go @@ -0,0 +1,194 @@ +package handler + +import ( + "fmt" + "strings" + + "cyberstrike-ai/internal/agent" + "cyberstrike-ai/internal/audit" + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/mcp/builtin" + "cyberstrike-ai/internal/security" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +// multiAgentPrepared 多代理请求在调用 Eino 前的会话与消息准备结果。 +type multiAgentPrepared struct { + ConversationID string + CreatedNew bool + History []agent.ChatMessage + FinalMessage string + RoleTools []string + AssistantMessageID string + UserMessageID string +} + +func chatRequestAgentMode(req *ChatRequest, source string) string { + if strings.HasPrefix(strings.TrimSpace(source), "multi_agent") { + return config.NormalizeMultiAgentOrchestration(req.Orchestration) + } + return "eino_single" +} + +func (h *AgentHandler) prepareMultiAgentSession(req *ChatRequest, c *gin.Context, source string) (*multiAgentPrepared, error) { + if len(req.Attachments) > maxAttachments { + return nil, fmt.Errorf("附件最多 %d 个", maxAttachments) + } + + conversationID := strings.TrimSpace(req.ConversationID) + projectID := strings.TrimSpace(effectiveProjectID(h.config, req.ProjectID)) + webshellID := strings.TrimSpace(req.WebShellConnectionID) + session, hasSession := security.CurrentSession(c) + if !hasSession || !session.Permissions["chat:write"] { + return nil, fmt.Errorf("无权写入对话") + } + canAccess := func(resourceType, resourceID string) bool { + if !hasSession || h.db == nil || strings.TrimSpace(resourceID) == "" { + return false + } + return h.db.UserCanAccessResource(session.UserID, session.Scope, resourceType, resourceID) + } + if projectID != "" && (!session.Permissions["project:read"] || !canAccess("project", projectID)) { + return nil, fmt.Errorf("无权访问目标项目") + } + if webshellID != "" && (!session.Permissions["webshell:write"] || !canAccess("webshell", webshellID)) { + return nil, fmt.Errorf("无权访问该 WebShell 连接") + } + createdNew := false + if conversationID == "" { + title := safeTruncateString(req.Message, 50) + var conv *database.Conversation + var err error + meta := audit.ConversationCreateMetaFromGin(c, source) + meta.ProjectID = projectID + meta.RoleName = req.Role + meta.AgentMode = chatRequestAgentMode(req, source) + if webshellID != "" { + meta.Source = source + "_webshell" + meta.WebShellConnectionID = webshellID + conv, err = h.db.CreateConversationWithWebshell(meta.WebShellConnectionID, title, meta) + } else { + conv, err = h.db.CreateConversation(title, meta) + } + if err != nil { + return nil, fmt.Errorf("创建对话失败: %w", err) + } + conversationID = conv.ID + createdNew = true + if hasSession { + _ = h.db.SetResourceOwner("conversation", conversationID, session.UserID) + _ = h.db.AssignResourceToUser(session.UserID, "conversation", conversationID) + } + } else { + if _, err := h.db.GetConversation(conversationID); err != nil { + return nil, fmt.Errorf("对话不存在") + } + if !canAccess("conversation", conversationID) { + return nil, fmt.Errorf("无权访问该对话") + } + } + if err := h.db.SetConversationRoleName(conversationID, req.Role); err != nil { + h.logger.Warn("更新对话角色失败", zap.String("conversationId", conversationID), zap.String("role", req.Role), zap.Error(err)) + } + if err := h.db.SetConversationAgentMode(conversationID, chatRequestAgentMode(req, source)); err != nil { + h.logger.Warn("更新对话模式失败", zap.String("conversationId", conversationID), zap.String("source", source), zap.String("orchestration", req.Orchestration), zap.Error(err)) + } + + agentHistoryMessages, err := h.loadHistoryFromAgentTrace(conversationID) + if err != nil { + historyMessages, getErr := h.db.GetMessages(conversationID) + if getErr != nil { + agentHistoryMessages = []agent.ChatMessage{} + } else { + agentHistoryMessages = dbMessagesToAgentChatMessages(historyMessages) + } + } + + finalMessage := req.Message + var roleTools []string + if webshellID != "" { + conn, errConn := h.db.GetWebshellConnection(webshellID) + if errConn != nil || conn == nil { + h.logger.Warn("WebShell AI 助手：未找到连接", zap.String("id", req.WebShellConnectionID), zap.Error(errConn)) + return nil, fmt.Errorf("未找到该 WebShell 连接") + } + webshellContext := BuildWebshellAssistantContext(conn, WebshellSkillHintMultiAgent, req.Message) + // WebShell 模式下如果同时指定了角色，追加角色 user_prompt（工具集仍仅限 webshell 专用工具） + if req.Role != "" && req.Role != "默认" && h.config != nil && h.config.Roles != nil { + if role, exists := h.config.Roles[req.Role]; exists && role.Enabled && role.UserPrompt != "" { + finalMessage = role.UserPrompt + "\n\n" + webshellContext + h.logger.Info("WebShell + 角色: 应用角色提示词（多代理）", zap.String("role", req.Role)) + } else { + finalMessage = webshellContext + } + } else { + finalMessage = webshellContext + } + roleTools = []string{ + builtin.ToolWebshellExec, + builtin.ToolWebshellFileList, + builtin.ToolWebshellFileRead, + builtin.ToolWebshellFileWrite, + builtin.ToolRecordVulnerability, + builtin.ToolListVulnerabilities, + builtin.ToolGetVulnerability, + builtin.ToolUpsertProjectFact, + builtin.ToolGetProjectFact, + builtin.ToolListProjectFacts, + builtin.ToolSearchProjectFacts, + builtin.ToolDeprecateProjectFact, + builtin.ToolRestoreProjectFact, + builtin.ToolListKnowledgeRiskTypes, + builtin.ToolSearchKnowledgeBase, + } + } else if req.Role != "" && req.Role != "默认" && h.config != nil && 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 + } + roleTools = role.Tools + } + } + + var savedPaths []string + if len(req.Attachments) > 0 { + var aerr error + savedPaths, aerr = saveAttachmentsToDateAndConversationDir(req.Attachments, conversationID, h.logger) + if aerr != nil { + return nil, fmt.Errorf("保存上传文件失败: %w", aerr) + } + } + finalMessage = appendAttachmentsToMessage(finalMessage, req.Attachments, savedPaths) + + userContent := userMessageContentForStorage(req.Message, req.Attachments, savedPaths) + userMsgRow, uerr := h.db.AddMessage(conversationID, "user", userContent, nil) + if uerr != nil { + h.logger.Error("保存用户消息失败", zap.Error(uerr)) + return nil, fmt.Errorf("保存用户消息失败: %w", uerr) + } + userMessageID := "" + if userMsgRow != nil { + userMessageID = userMsgRow.ID + } + + assistantMsg, aerr := h.db.AddMessage(conversationID, "assistant", "处理中...", nil) + var assistantMessageID string + if aerr != nil { + h.logger.Warn("创建助手消息占位失败", zap.Error(aerr)) + } else if assistantMsg != nil { + assistantMessageID = assistantMsg.ID + } + + return &multiAgentPrepared{ + ConversationID: conversationID, + CreatedNew: createdNew, + History: agentHistoryMessages, + FinalMessage: finalMessage, + RoleTools: roleTools, + AssistantMessageID: assistantMessageID, + UserMessageID: userMessageID, + }, nil +} diff --git a/internal/handler/notification.go b/internal/handler/notification.go new file mode 100644 index 00000000..21f99d9b --- /dev/null +++ b/internal/handler/notification.go @@ -0,0 +1,837 @@ +package handler + +import ( + "fmt" + "net/http" + "sort" + "strconv" + "strings" + "time" + + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +// NotificationHandler 聚合通知（Phase 2：服务端统一计算） +type NotificationHandler struct { + db *database.DB + agentHandler *AgentHandler + logger *zap.Logger +} + +const notificationReadMaxRows = 150 + +// NotificationSummaryItem 通知项 +type NotificationSummaryItem struct { + ID string `json:"id"` + Level string `json:"level"` // p0/p1/p2 + Type string `json:"type"` + Title string `json:"title"` + Desc string `json:"desc"` + Ts string `json:"ts"` // RFC3339 + Count int `json:"count,omitempty"` + Actionable bool `json:"actionable"` + Read bool `json:"read"` + // 以下字段用于前端深链跳转（通知即入口） + ConversationID string `json:"conversationId,omitempty"` + VulnerabilityID string `json:"vulnerabilityId,omitempty"` + ExecutionID string `json:"executionId,omitempty"` + InterruptID string `json:"interruptId,omitempty"` + SessionID string `json:"sessionId,omitempty"` // C2 会话（如新会话上线） +} + +// NotificationSummaryResponse 聚合响应 +type NotificationSummaryResponse struct { + SinceMs int64 `json:"sinceMs"` + GeneratedAt string `json:"generatedAt"` + P0Count int `json:"p0Count"` + UnreadCount int `json:"unreadCount"` + Counts map[string]int `json:"counts"` + Items []NotificationSummaryItem `json:"items"` +} + +func NewNotificationHandler(db *database.DB, agentHandler *AgentHandler, logger *zap.Logger) *NotificationHandler { + return &NotificationHandler{ + db: db, + agentHandler: agentHandler, + logger: logger, + } +} + +func parseSinceMs(raw string) int64 { + v := strings.TrimSpace(raw) + if v == "" { + return 0 + } + if ms, err := strconv.ParseInt(v, 10, 64); err == nil && ms > 0 { + return ms + } + if t, err := time.Parse(time.RFC3339, v); err == nil { + return t.UnixMilli() + } + return 0 +} + +func unixSecToRFC3339(sec int64) string { + if sec <= 0 { + return time.Now().UTC().Format(time.RFC3339) + } + return time.Unix(sec, 0).UTC().Format(time.RFC3339) +} + +func normalizedSinceSec(sinceMs int64) int64 { + sec := sinceMs / 1000 + // SQLite 默认时间精度到秒；给 1s 回看窗口，避免“同秒内新增”被漏算。 + if sec > 0 { + return sec - 1 + } + return 0 +} + +func ptrTime(t time.Time) *time.Time { + return &t +} + +func normalizeSinceMs(raw int64) int64 { + if raw > 0 { + return raw + } + // 默认仅看最近 24 小时，避免首次打开拉全量历史噪音。 + return time.Now().Add(-24 * time.Hour).UnixMilli() +} + +func levelBySeverity(sev string) string { + switch strings.ToLower(strings.TrimSpace(sev)) { + case "critical", "high": + return "p0" + case "medium": + return "p1" + default: + return "p2" + } +} + +func requestWantsEnglish(c *gin.Context) bool { + if c == nil { + return false + } + lang := strings.ToLower(strings.TrimSpace(c.Query("lang"))) + if lang == "" { + lang = strings.ToLower(strings.TrimSpace(c.GetHeader("Accept-Language"))) + } + return strings.HasPrefix(lang, "en") +} + +func i18nText(english bool, zh string, en string) string { + if english { + return en + } + return zh +} + +func notificationAccessFromContext(c *gin.Context) database.RBACListAccess { + session, ok := security.CurrentSession(c) + if !ok { + return database.RBACListAccess{} + } + return database.RBACListAccess{UserID: session.UserID, Scope: session.Scope} +} + +func appendConversationAccessSQL(query string, args []interface{}, column string, access database.RBACListAccess) (string, []interface{}) { + userID := strings.TrimSpace(access.UserID) + if access.Scope == database.RBACScopeAll { + return query, args + } + if userID == "" { + return query + ` AND 1=0`, args + } + query += ` AND ` + column + ` IS NOT NULL AND ` + column + ` <> '' AND ( + EXISTS (SELECT 1 FROM conversations c WHERE c.id = ` + column + ` AND c.owner_user_id = ?) + OR EXISTS ( + SELECT 1 FROM rbac_resource_assignments ra + WHERE ra.user_id = ? AND ra.resource_type = 'conversation' AND ra.resource_id = ` + column + ` + ) + OR EXISTS ( + SELECT 1 FROM conversations c + JOIN projects p ON p.id = c.project_id + WHERE c.id = ` + column + ` AND p.owner_user_id = ? + ) + OR EXISTS ( + SELECT 1 FROM conversations c + JOIN rbac_resource_assignments pra ON pra.resource_id = c.project_id + WHERE c.id = ` + column + ` AND pra.user_id = ? AND pra.resource_type = 'project' + ) + )` + args = append(args, userID, userID, userID, userID) + return query, args +} + +func appendVulnerabilityNotificationAccessSQL(query string, args []interface{}, access database.RBACListAccess) (string, []interface{}) { + userID := strings.TrimSpace(access.UserID) + if access.Scope == database.RBACScopeAll { + return query, args + } + if userID == "" { + return query + ` AND 1=0`, args + } + query += ` AND ( + owner_user_id = ? + OR EXISTS ( + SELECT 1 FROM rbac_resource_assignments ra + WHERE ra.user_id = ? AND ra.resource_type = 'vulnerability' AND ra.resource_id = vulnerabilities.id + ) + OR ( + project_id IS NOT NULL AND project_id <> '' AND ( + EXISTS (SELECT 1 FROM projects p WHERE p.id = vulnerabilities.project_id AND p.owner_user_id = ?) + OR EXISTS ( + SELECT 1 FROM rbac_resource_assignments pra + WHERE pra.user_id = ? AND pra.resource_type = 'project' AND pra.resource_id = vulnerabilities.project_id + ) + ) + ) + OR ( + conversation_id IS NOT NULL AND conversation_id <> '' AND ( + EXISTS (SELECT 1 FROM conversations c WHERE c.id = vulnerabilities.conversation_id AND c.owner_user_id = ?) + OR EXISTS ( + SELECT 1 FROM rbac_resource_assignments cra + WHERE cra.user_id = ? AND cra.resource_type = 'conversation' AND cra.resource_id = vulnerabilities.conversation_id + ) + ) + ) + )` + args = append(args, userID, userID, userID, userID, userID, userID) + return query, args +} + +func (h *NotificationHandler) loadPendingHITLItems(limit int, english bool, access database.RBACListAccess) ([]NotificationSummaryItem, error) { + query := ` + SELECT + id, + conversation_id, + tool_name, + COALESCE(CAST(strftime('%s', created_at) AS INTEGER), 0) + FROM hitl_interrupts + WHERE status = 'pending' + ` + args := []interface{}{} + query, args = appendConversationAccessSQL(query, args, "conversation_id", access) + query += ` ORDER BY created_at DESC + LIMIT ? + ` + args = append(args, limit) + rows, err := h.db.Query(query, args...) + if err != nil { + return nil, err + } + defer rows.Close() + items := make([]NotificationSummaryItem, 0, limit) + for rows.Next() { + var id, conversationID, toolName string + var createdSec int64 + if err := rows.Scan(&id, &conversationID, &toolName, &createdSec); err != nil { + continue + } + desc := i18nText(english, "会话 "+conversationID+" 的审批中断待处理", "Conversation "+conversationID+" has pending HITL approval") + if strings.TrimSpace(toolName) != "" { + desc = i18nText(english, "工具 "+toolName+" 等待审批", "Tool "+toolName+" is waiting for approval") + } + items = append(items, NotificationSummaryItem{ + ID: "hitl:" + id, + Level: "p0", + Type: "hitl_pending", + Title: i18nText(english, "HITL 待审批", "HITL Pending Approval"), + Desc: desc, + Ts: unixSecToRFC3339(createdSec), + Count: 1, + Actionable: true, + Read: false, + ConversationID: conversationID, + InterruptID: id, + }) + } + return items, nil +} + +func (h *NotificationHandler) loadVulnerabilityItems(sinceMs int64, limit int, english bool, access database.RBACListAccess) ([]NotificationSummaryItem, map[string]int, error) { + sinceSec := normalizedSinceSec(sinceMs) + query := ` + SELECT + id, + title, + severity, + conversation_id, + COALESCE(CAST(strftime('%s', created_at) AS INTEGER), 0) + FROM vulnerabilities + WHERE CAST(strftime('%s', created_at) AS INTEGER) > ? + ` + args := []interface{}{sinceSec} + query, args = appendVulnerabilityNotificationAccessSQL(query, args, access) + query += ` + ORDER BY created_at DESC + LIMIT ? + ` + args = append(args, limit) + rows, err := h.db.Query(query, args...) + if err != nil { + return nil, nil, err + } + defer rows.Close() + items := make([]NotificationSummaryItem, 0, limit) + counts := map[string]int{ + "newCriticalVulns": 0, + "newHighVulns": 0, + "newMediumVulns": 0, + "newLowVulns": 0, + "newInfoVulns": 0, + } + for rows.Next() { + var id, title, severity, conversationID string + var createdSec int64 + if err := rows.Scan(&id, &title, &severity, &conversationID, &createdSec); err != nil { + continue + } + switch strings.ToLower(strings.TrimSpace(severity)) { + case "critical": + counts["newCriticalVulns"]++ + case "high": + counts["newHighVulns"]++ + case "medium": + counts["newMediumVulns"]++ + case "low": + counts["newLowVulns"]++ + default: + counts["newInfoVulns"]++ + } + sevUpper := strings.ToUpper(strings.TrimSpace(severity)) + if sevUpper == "" { + sevUpper = "INFO" + } + finalTitle := i18nText(english, "新漏洞（"+sevUpper+"）", "New Vulnerability ("+sevUpper+")") + finalDesc := strings.TrimSpace(title) + if finalDesc == "" { + finalDesc = i18nText(english, "（无标题）", "(Untitled)") + } + items = append(items, NotificationSummaryItem{ + ID: "vuln:" + id, + Level: levelBySeverity(severity), + Type: "vulnerability_created", + Title: finalTitle, + Desc: finalDesc, + Ts: unixSecToRFC3339(createdSec), + Count: 1, + Actionable: false, + Read: false, + ConversationID: conversationID, + VulnerabilityID: id, + }) + } + return items, counts, nil +} + +// loadC2SessionOnlineEvents 新会话上线（c2_events：session + critical，与 Manager.IngestCheckIn 一致） +func (h *NotificationHandler) loadC2SessionOnlineEvents(sinceMs int64, limit int, english bool, access database.RBACListAccess) ([]NotificationSummaryItem, int, error) { + sinceSec := normalizedSinceSec(sinceMs) + events, err := h.db.ListC2EventsForAccess(database.ListC2EventsFilter{ + Category: "session", + Level: "critical", + Since: ptrTime(time.Unix(sinceSec, 0)), + Limit: limit, + }, access) + if err != nil { + return nil, 0, err + } + items := make([]NotificationSummaryItem, 0, limit) + for _, e := range events { + if e == nil { + continue + } + desc := strings.TrimSpace(e.Message) + if len(desc) > 220 { + desc = desc[:200] + "…" + } + if desc == "" { + desc = i18nText(english, "新会话已建立", "A new session was created") + } + items = append(items, NotificationSummaryItem{ + ID: "c2evt:" + e.ID, + Level: "p0", + Type: "c2_session_online", + Title: i18nText(english, "C2 新会话上线", "C2 new session online"), + Desc: desc, + Ts: e.CreatedAt.UTC().Format(time.RFC3339), + Count: 1, + Actionable: false, + Read: false, + SessionID: e.SessionID, + }) + } + return items, len(items), nil +} + +func (h *NotificationHandler) loadFailedExecutionItems(sinceMs int64, limit int, english bool) ([]NotificationSummaryItem, int, error) { + sinceSec := normalizedSinceSec(sinceMs) + rows, err := h.db.Query(` + SELECT + id, + tool_name, + COALESCE(CAST(strftime('%s', start_time) AS INTEGER), 0) + FROM tool_executions + WHERE status = 'failed' + AND CAST(strftime('%s', start_time) AS INTEGER) > ? + ORDER BY start_time DESC + LIMIT ? + `, sinceSec, limit) + if err != nil { + return nil, 0, err + } + defer rows.Close() + items := make([]NotificationSummaryItem, 0, limit) + count := 0 + for rows.Next() { + var id, toolName string + var startSec int64 + if err := rows.Scan(&id, &toolName, &startSec); err != nil { + continue + } + count++ + if strings.TrimSpace(toolName) == "" { + toolName = i18nText(english, "未知工具", "unknown") + } + items = append(items, NotificationSummaryItem{ + ID: "exec_failed:" + id, + Level: "p0", + Type: "task_failed", + Title: i18nText(english, "任务执行失败", "Task Execution Failed"), + Desc: i18nText(english, "工具 "+toolName+" 执行失败", "Tool "+toolName+" execution failed"), + Ts: unixSecToRFC3339(startSec), + Count: 1, + Actionable: false, + Read: false, + ExecutionID: id, + }) + } + return items, count, nil +} + +func (h *NotificationHandler) summarizeLongRunningTasks(threshold time.Duration, english bool, access database.RBACListAccess) ([]NotificationSummaryItem, int) { + if h.agentHandler == nil || h.agentHandler.tasks == nil { + return nil, 0 + } + tasks := h.agentHandler.tasks.GetActiveTasks() + now := time.Now() + items := make([]NotificationSummaryItem, 0, len(tasks)) + for _, t := range tasks { + if t == nil { + continue + } + if !h.notificationConversationAllowed(access, t.ConversationID) { + continue + } + if now.Sub(t.StartedAt) >= threshold { + items = append(items, NotificationSummaryItem{ + ID: "task_long:" + t.ConversationID, + Level: "p1", + Type: "long_running_tasks", + Title: i18nText(english, "长时间运行任务", "Long Running Task"), + Desc: i18nText(english, "会话 "+t.ConversationID+" 运行超过 15 分钟", "Conversation "+t.ConversationID+" has been running over 15 minutes"), + Ts: t.StartedAt.UTC().Format(time.RFC3339), + Count: 1, + Actionable: true, + Read: false, + ConversationID: t.ConversationID, + }) + } + } + return items, len(items) +} + +func (h *NotificationHandler) summarizeCompletedTasksSince(sinceMs int64, limit int, english bool, access database.RBACListAccess) ([]NotificationSummaryItem, int) { + if h.agentHandler == nil || h.agentHandler.tasks == nil { + return nil, 0 + } + since := time.UnixMilli(sinceMs) + completed := h.agentHandler.tasks.GetCompletedTasks() + items := make([]NotificationSummaryItem, 0, limit) + for _, t := range completed { + if t == nil { + continue + } + if !h.notificationConversationAllowed(access, t.ConversationID) { + continue + } + if t.CompletedAt.After(since) { + items = append(items, NotificationSummaryItem{ + ID: "task_completed:" + t.ConversationID + ":" + strconv.FormatInt(t.CompletedAt.Unix(), 10), + Level: "p2", + Type: "task_completed", + Title: i18nText(english, "任务完成", "Task Completed"), + Desc: i18nText(english, "会话 "+t.ConversationID+" 已完成", "Conversation "+t.ConversationID+" completed"), + Ts: t.CompletedAt.UTC().Format(time.RFC3339), + Count: 1, + Actionable: false, + Read: false, + ConversationID: t.ConversationID, + }) + if len(items) >= limit { + break + } + } + } + return items, len(items) +} + +func buildPlaceholders(n int) string { + if n <= 0 { + return "" + } + out := make([]string, 0, n) + for i := 0; i < n; i++ { + out = append(out, "?") + } + return strings.Join(out, ",") +} + +func (h *NotificationHandler) readStatesByIDs(userID string, ids []string) (map[string]bool, error) { + result := make(map[string]bool, len(ids)) + userID = strings.TrimSpace(userID) + if len(ids) == 0 || userID == "" { + return result, nil + } + holders := buildPlaceholders(len(ids)) + query := "SELECT event_id FROM notification_reads_by_user WHERE user_id = ? AND event_id IN (" + holders + ")" + args := make([]interface{}, 0, len(ids)+1) + args = append(args, userID) + for _, id := range ids { + args = append(args, id) + } + rows, err := h.db.Query(query, args...) + if err != nil { + return result, err + } + defer rows.Close() + for rows.Next() { + var id string + if err := rows.Scan(&id); err != nil { + continue + } + result[id] = true + } + return result, nil +} + +func (h *NotificationHandler) applyReadStates(userID string, items []NotificationSummaryItem) ([]NotificationSummaryItem, error) { + markableIDs := make([]string, 0, len(items)) + for _, item := range items { + if item.Actionable { + continue + } + markableIDs = append(markableIDs, item.ID) + } + readMap, err := h.readStatesByIDs(userID, markableIDs) + if err != nil { + return items, err + } + for i := range items { + if items[i].Actionable { + items[i].Read = false + continue + } + items[i].Read = readMap[items[i].ID] + } + return items, nil +} + +func filterVisibleItems(items []NotificationSummaryItem) []NotificationSummaryItem { + out := make([]NotificationSummaryItem, 0, len(items)) + for _, item := range items { + if item.Actionable || !item.Read { + out = append(out, item) + } + } + return out +} + +func countP0(items []NotificationSummaryItem) int { + total := 0 + for _, item := range items { + if item.Level == "p0" { + if item.Count > 0 { + total += item.Count + } else { + total++ + } + } + } + return total +} + +func countUnread(items []NotificationSummaryItem) int { + total := 0 + for _, item := range items { + if item.Actionable || !item.Read { + if item.Count > 0 { + total += item.Count + } else { + total++ + } + } + } + return total +} + +func createNotificationReadTableIfNeeded(db *database.DB) error { + if db == nil { + return fmt.Errorf("db is nil") + } + _, err := db.Exec(` + CREATE TABLE IF NOT EXISTS notification_reads_by_user ( + user_id TEXT NOT NULL, + event_id TEXT NOT NULL, + read_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + PRIMARY KEY(user_id, event_id) + ); + `) + if err != nil { + return err + } + _, idxErr := db.Exec(`CREATE INDEX IF NOT EXISTS idx_notification_reads_user_read_at ON notification_reads_by_user(user_id, read_at DESC);`) + return idxErr +} + +func pruneNotificationReads(db *database.DB, userID string, maxRows int) error { + if db == nil { + return fmt.Errorf("db is nil") + } + userID = strings.TrimSpace(userID) + if maxRows <= 0 || userID == "" { + return nil + } + _, err := db.Exec(` + DELETE FROM notification_reads_by_user + WHERE user_id = ? AND event_id NOT IN ( + SELECT event_id + FROM notification_reads_by_user + WHERE user_id = ? + ORDER BY read_at DESC, rowid DESC + LIMIT ? + ) + `, userID, userID, maxRows) + return err +} + +type markReadRequest struct { + EventIDs []string `json:"eventIds"` +} + +func normalizeMarkableEventID(id string) (string, bool) { + v := strings.TrimSpace(id) + if v == "" { + return "", false + } + // 仅允许“可读后隐藏”的信息类事件；Actionable 事件不参与 read 标记。 + allowedPrefixes := []string{ + "vuln:", + "exec_failed:", + "task_completed:", + "c2evt:", + } + for _, prefix := range allowedPrefixes { + if strings.HasPrefix(v, prefix) { + return v, true + } + } + return "", false +} + +// MarkRead 按事件 ID 标记已读 +func (h *NotificationHandler) MarkRead(c *gin.Context) { + if err := createNotificationReadTableIfNeeded(h.db); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to prepare notification read table"}) + return + } + var req markReadRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"}) + return + } + if len(req.EventIDs) == 0 { + c.JSON(http.StatusOK, gin.H{"ok": true, "marked": 0}) + return + } + session, ok := security.CurrentSession(c) + if !ok || strings.TrimSpace(session.UserID) == "" { + c.JSON(http.StatusUnauthorized, gin.H{"error": "missing authenticated user"}) + return + } + tx, err := h.db.Begin() + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to begin transaction"}) + return + } + defer func() { + _ = tx.Rollback() + }() + stmt, err := tx.Prepare(` + INSERT INTO notification_reads_by_user(user_id, event_id, read_at) + VALUES(?, ?, CURRENT_TIMESTAMP) + ON CONFLICT(user_id, event_id) DO UPDATE SET read_at = CURRENT_TIMESTAMP + `) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to prepare statement"}) + return + } + defer stmt.Close() + marked := 0 + for _, raw := range req.EventIDs { + id, ok := normalizeMarkableEventID(raw) + if !ok { + continue + } + if _, err := stmt.Exec(session.UserID, id); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to mark read"}) + return + } + marked++ + } + if err := tx.Commit(); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to commit read marks"}) + return + } + if err := pruneNotificationReads(h.db, session.UserID, notificationReadMaxRows); err != nil { + h.logger.Warn("裁剪通知已读记录失败", zap.Error(err)) + } + c.JSON(http.StatusOK, gin.H{"ok": true, "marked": marked}) +} + +// GetSummary 返回通知聚合视图（用于头部铃铛） +func (h *NotificationHandler) GetSummary(c *gin.Context) { + if h.db == nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "database unavailable"}) + return + } + + if err := createNotificationReadTableIfNeeded(h.db); err != nil { + h.logger.Warn("初始化通知已读表失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to initialize notification read table"}) + return + } + + english := requestWantsEnglish(c) + sinceMs := normalizeSinceMs(parseSinceMs(c.Query("since"))) + limit, _ := strconv.Atoi(strings.TrimSpace(c.DefaultQuery("limit", "50"))) + if limit <= 0 { + limit = 50 + } + if limit > 200 { + limit = 200 + } + access := notificationAccessFromContext(c) + + hitlItems := []NotificationSummaryItem{} + if security.SessionHasPermission(c, "hitl:read") { + var err error + hitlItems, err = h.loadPendingHITLItems(limit, english, access) + if err != nil { + h.logger.Warn("加载 HITL 通知失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to summarize hitl notifications"}) + return + } + } + + vulnItems := []NotificationSummaryItem{} + vulnCounts := map[string]int{ + "newCriticalVulns": 0, + "newHighVulns": 0, + "newMediumVulns": 0, + "newLowVulns": 0, + "newInfoVulns": 0, + } + if security.SessionHasPermission(c, "vulnerability:read") { + var err error + vulnItems, vulnCounts, err = h.loadVulnerabilityItems(sinceMs, limit, english, access) + if err != nil { + h.logger.Warn("加载漏洞通知失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to summarize vulnerabilities"}) + return + } + } + + c2OnlineItems := []NotificationSummaryItem{} + c2OnlineCount := 0 + if security.SessionHasPermission(c, "c2:read") { + var err error + c2OnlineItems, c2OnlineCount, err = h.loadC2SessionOnlineEvents(sinceMs, limit, english, access) + if err != nil { + h.logger.Warn("加载 C2 会话上线通知失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to summarize c2 session events"}) + return + } + } + + longRunningItems := []NotificationSummaryItem{} + completedItems := []NotificationSummaryItem{} + longRunningCount := 0 + completedCount := 0 + if security.SessionHasPermission(c, "tasks:read") || security.SessionHasPermission(c, "chat:read") { + longRunningItems, longRunningCount = h.summarizeLongRunningTasks(15*time.Minute, english, access) + completedItems, completedCount = h.summarizeCompletedTasksSince(sinceMs, limit, english, access) + } + + items := make([]NotificationSummaryItem, 0, len(hitlItems)+len(vulnItems)+len(c2OnlineItems)+len(longRunningItems)+len(completedItems)) + items = append(items, hitlItems...) + items = append(items, vulnItems...) + items = append(items, c2OnlineItems...) + items = append(items, longRunningItems...) + items = append(items, completedItems...) + + session, _ := security.CurrentSession(c) + items, err := h.applyReadStates(session.UserID, items) + if err != nil { + h.logger.Warn("加载通知已读状态失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to load notification read states"}) + return + } + items = filterVisibleItems(items) + + sort.Slice(items, func(i, j int) bool { + ti, errI := time.Parse(time.RFC3339, items[i].Ts) + tj, errJ := time.Parse(time.RFC3339, items[j].Ts) + if errI != nil || errJ != nil { + return i < j + } + return ti.After(tj) + }) + + p0Count := countP0(items) + unreadCount := countUnread(items) + c.JSON(http.StatusOK, NotificationSummaryResponse{ + SinceMs: sinceMs, + GeneratedAt: time.Now().UTC().Format(time.RFC3339), + P0Count: p0Count, + UnreadCount: unreadCount, + Counts: map[string]int{ + "hitlPending": len(hitlItems), + "newCriticalVulns": vulnCounts["newCriticalVulns"], + "newHighVulns": vulnCounts["newHighVulns"], + "newMediumVulns": vulnCounts["newMediumVulns"], + "newLowVulns": vulnCounts["newLowVulns"], + "newInfoVulns": vulnCounts["newInfoVulns"], + "failedExecutions": 0, + "longRunningTasks": longRunningCount, + "completedTasks": completedCount, + "c2SessionOnline": c2OnlineCount, + }, + Items: items, + }) +} + +func (h *NotificationHandler) notificationConversationAllowed(access database.RBACListAccess, conversationID string) bool { + conversationID = strings.TrimSpace(conversationID) + if conversationID == "" { + return access.Scope == database.RBACScopeAll + } + return h.db.UserCanAccessResource(access.UserID, access.Scope, "conversation", conversationID) +} diff --git a/internal/handler/notification_rbac_test.go b/internal/handler/notification_rbac_test.go new file mode 100644 index 00000000..ad05edef --- /dev/null +++ b/internal/handler/notification_rbac_test.go @@ -0,0 +1,45 @@ +package handler + +import ( + "bytes" + "net/http" + "net/http/httptest" + "path/filepath" + "testing" + + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +func TestNotificationReadStateIsPerUser(t *testing.T) { + db, err := database.NewDB(filepath.Join(t.TempDir(), "notification-rbac.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = db.Close() }) + h := NewNotificationHandler(db, nil, zap.NewNop()) + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(security.ContextSessionKey, security.Session{UserID: "u1", Scope: database.RBACScopeAssigned}) + c.Next() + }) + router.POST("/notifications/read", h.MarkRead) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/notifications/read", bytes.NewBufferString(`{"eventIds":["vuln:v1"]}`)) + req.Header.Set("Content-Type", "application/json") + router.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("mark read status = %d: %s", w.Code, w.Body.String()) + } + u1, err := h.readStatesByIDs("u1", []string{"vuln:v1"}) + if err != nil || !u1["vuln:v1"] { + t.Fatalf("u1 read state = %#v, err=%v", u1, err) + } + u2, err := h.readStatesByIDs("u2", []string{"vuln:v1"}) + if err != nil || u2["vuln:v1"] { + t.Fatalf("u2 inherited u1 read state = %#v, err=%v", u2, err) + } +} diff --git a/internal/handler/openapi.go b/internal/handler/openapi.go new file mode 100644 index 00000000..fad523d0 --- /dev/null +++ b/internal/handler/openapi.go @@ -0,0 +1,6859 @@ +package handler + +import ( + "net/http" + + "cyberstrike-ai/internal/database" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +// OpenAPIHandler OpenAPI处理器 +type OpenAPIHandler struct { + db *database.DB + logger *zap.Logger + conversationHdlr *ConversationHandler + agentHdlr *AgentHandler +} + +// NewOpenAPIHandler 创建新的OpenAPI处理器 +func NewOpenAPIHandler(db *database.DB, logger *zap.Logger, conversationHdlr *ConversationHandler, agentHdlr *AgentHandler) *OpenAPIHandler { + return &OpenAPIHandler{ + db: db, + logger: logger, + conversationHdlr: conversationHdlr, + agentHdlr: agentHdlr, + } +} + +// GetOpenAPISpec 获取OpenAPI规范 +func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) { + host := c.Request.Host + scheme := "http" + if c.Request.TLS != nil { + scheme = "https" + } + + finalizationRequestSchema := map[string]interface{}{ + "type": "object", + "description": "最终回复交付策略。后端不会从自然语言内容推断执行意图；执行入口应显式声明是否要求 completed 工具证据。", + "properties": map[string]interface{}{ + "requireExecutionEvidence": map[string]interface{}{ + "type": "boolean", + "description": "为 true 时，缺少 completed 工具执行记录会触发无注入续跑或最终阻断；普通聊天可省略或设为 false。", + }, + }, + } + + spec := map[string]interface{}{ + "openapi": "3.0.0", + "info": map[string]interface{}{ + "title": "CyberStrikeAI API", + "description": "AI驱动的自动化安全测试平台API文档", + "version": "1.0.0", + "contact": map[string]interface{}{ + "name": "CyberStrikeAI", + }, + }, + "servers": []map[string]interface{}{ + { + "url": scheme + "://" + host, + "description": "当前服务器", + }, + }, + "components": map[string]interface{}{ + "securitySchemes": map[string]interface{}{ + "bearerAuth": map[string]interface{}{ + "type": "http", + "scheme": "bearer", + "bearerFormat": "JWT", + "description": "使用Bearer Token进行认证。Token通过 /api/auth/login 接口获取。", + }, + }, + "schemas": map[string]interface{}{ + "CreateConversationRequest": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "title": map[string]interface{}{ + "type": "string", + "description": "对话标题", + "example": "Web应用安全测试", + }, + "projectId": map[string]interface{}{ + "type": "string", + "description": "绑定的项目 ID（可选，共享事实黑板）", + }, + }, + }, + "SetConversationProjectRequest": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "projectId": map[string]interface{}{ + "type": "string", + "description": "项目 ID；空字符串表示解除绑定", + }, + }, + "required": []string{"projectId"}, + }, + "AgentChatResponse": map[string]interface{}{ + "type": "object", + "description": "Agent 非流式响应。response 只是交付文本；是否为成功最终回复必须以 finalized/finalizable/status 为准。", + "properties": map[string]interface{}{ + "response": map[string]interface{}{ + "type": "string", + "description": "交付给用户的文本。finalized=false 时为阻断/未完成说明，不是成功结论。", + }, + "conversationId": map[string]interface{}{ + "type": "string", + "description": "对话 ID", + }, + "assistantMessageId": map[string]interface{}{ + "type": "string", + "description": "助手消息 ID（部分接口返回）", + }, + "mcpExecutionIds": map[string]interface{}{ + "type": "array", + "description": "本轮关联的 MCP 工具执行 ID", + "items": map[string]interface{}{"type": "string"}, + }, + "agentMode": map[string]interface{}{ + "type": "string", + "description": "agent 模式，例如 eino_single、eino_deep、workflow", + }, + "finalized": map[string]interface{}{ + "type": "boolean", + "description": "是否已经通过最终回复检查。只有 true 才能当成功最终回复。", + }, + "finalizable": map[string]interface{}{ + "type": "boolean", + "description": "候选输出是否可提升为最终回复。", + }, + "status": map[string]interface{}{ + "type": "string", + "description": "最终化状态", + "enum": []string{"completed", "in_progress", "blocked", "failed", "cancelled", "awaiting_hitl"}, + }, + "completionReason": map[string]interface{}{ + "type": "string", + "description": "最终化或阻断原因，例如 verified、pending_tool_executions、missing_execution_evidence", + }, + "evidenceVerified": map[string]interface{}{ + "type": "boolean", + "description": "证据是否满足最终化要求", + }, + "evidenceRefs": map[string]interface{}{ + "type": "array", + "description": "证据引用，例如 mcp_execution:", + "items": map[string]interface{}{"type": "string"}, + }, + "pendingExecutionIds": map[string]interface{}{ + "type": "array", + "description": "仍处于 queued/running 的工具执行 ID", + "items": map[string]interface{}{"type": "string"}, + }, + "missingChecks": map[string]interface{}{ + "type": "array", + "description": "未通过最终化检查的原因列表", + "items": map[string]interface{}{"type": "string"}, + }, + }, + "required": []string{"response", "conversationId", "finalized", "finalizable", "status", "evidenceVerified"}, + }, + "Conversation": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "id": map[string]interface{}{ + "type": "string", + "description": "对话ID", + "example": "550e8400-e29b-41d4-a716-446655440000", + }, + "title": map[string]interface{}{ + "type": "string", + "description": "对话标题", + "example": "Web应用安全测试", + }, + "createdAt": map[string]interface{}{ + "type": "string", + "format": "date-time", + "description": "创建时间", + }, + "updatedAt": map[string]interface{}{ + "type": "string", + "format": "date-time", + "description": "更新时间", + }, + "projectId": map[string]interface{}{ + "type": "string", + "description": "绑定的项目 ID（可选）", + }, + }, + }, + "ConversationDetail": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "id": map[string]interface{}{ + "type": "string", + "description": "对话ID", + }, + "title": map[string]interface{}{ + "type": "string", + "description": "对话标题", + }, + "status": map[string]interface{}{ + "type": "string", + "description": "对话状态：active（进行中）、completed（已完成）、failed（失败）", + "enum": []string{"active", "completed", "failed"}, + }, + "createdAt": map[string]interface{}{ + "type": "string", + "format": "date-time", + "description": "创建时间", + }, + "updatedAt": map[string]interface{}{ + "type": "string", + "format": "date-time", + "description": "更新时间", + }, + "messages": map[string]interface{}{ + "type": "array", + "description": "消息列表", + "items": map[string]interface{}{ + "$ref": "#/components/schemas/Message", + }, + }, + "messageCount": map[string]interface{}{ + "type": "integer", + "description": "消息数量", + }, + }, + }, + "Message": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "id": map[string]interface{}{ + "type": "string", + "description": "消息ID", + }, + "conversationId": map[string]interface{}{ + "type": "string", + "description": "对话ID", + }, + "role": map[string]interface{}{ + "type": "string", + "description": "消息角色：user（用户）、assistant（助手）", + "enum": []string{"user", "assistant"}, + }, + "content": map[string]interface{}{ + "type": "string", + "description": "消息内容", + }, + "createdAt": map[string]interface{}{ + "type": "string", + "format": "date-time", + "description": "创建时间", + }, + }, + }, + "ConversationResults": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "conversationId": map[string]interface{}{ + "type": "string", + "description": "对话ID", + }, + "messages": map[string]interface{}{ + "type": "array", + "description": "消息列表", + "items": map[string]interface{}{ + "$ref": "#/components/schemas/Message", + }, + }, + "vulnerabilities": map[string]interface{}{ + "type": "array", + "description": "发现的漏洞列表", + "items": map[string]interface{}{ + "$ref": "#/components/schemas/Vulnerability", + }, + }, + "executionResults": map[string]interface{}{ + "type": "array", + "description": "执行结果列表", + "items": map[string]interface{}{ + "$ref": "#/components/schemas/ExecutionResult", + }, + }, + }, + }, + "Vulnerability": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "id": map[string]interface{}{ + "type": "string", + "description": "漏洞ID", + }, + "title": map[string]interface{}{ + "type": "string", + "description": "漏洞标题", + }, + "description": map[string]interface{}{ + "type": "string", + "description": "漏洞描述", + }, + "severity": map[string]interface{}{ + "type": "string", + "description": "严重程度", + "enum": []string{"critical", "high", "medium", "low", "info"}, + }, + "status": map[string]interface{}{ + "type": "string", + "description": "状态", + "enum": []string{"open", "confirmed", "fixed", "false_positive", "ignored"}, + }, + "target": map[string]interface{}{ + "type": "string", + "description": "受影响的目标", + }, + }, + }, + "AssetImportItem": map[string]interface{}{ + "type": "object", + "description": "待导入资产；host、ip、domain 至少一项非空", + "properties": map[string]interface{}{ + "project_id": map[string]interface{}{"type": "string", "description": "所属项目 ID；调用者必须有权访问"}, + "host": map[string]interface{}{"type": "string", "maxLength": 500, "example": "https://app.example.com:443"}, + "ip": map[string]interface{}{"type": "string", "example": "192.0.2.10"}, + "port": map[string]interface{}{"type": "integer", "minimum": 0, "maximum": 65535, "example": 443}, + "domain": map[string]interface{}{"type": "string", "example": "app.example.com"}, + "protocol": map[string]interface{}{"type": "string", "example": "https"}, + "title": map[string]interface{}{"type": "string", "maxLength": 500}, + "server": map[string]interface{}{"type": "string", "maxLength": 255, "example": "nginx"}, + "country": map[string]interface{}{"type": "string"}, + "province": map[string]interface{}{"type": "string"}, + "city": map[string]interface{}{"type": "string"}, + "responsible_person": map[string]interface{}{"type": "string", "maxLength": 255, "description": "资产负责人"}, + "department": map[string]interface{}{"type": "string", "maxLength": 255, "description": "所属部门"}, + "business_system": map[string]interface{}{"type": "string", "maxLength": 255, "description": "所属业务系统"}, + "environment": map[string]interface{}{"type": "string", "enum": []string{"production", "staging", "testing", "development", "other"}}, + "criticality": map[string]interface{}{"type": "string", "enum": []string{"critical", "high", "medium", "low"}}, + "source": map[string]interface{}{"type": "string"}, + "source_query": map[string]interface{}{"type": "string"}, + "status": map[string]interface{}{"type": "string", "enum": []string{"active", "inactive"}, "default": "active"}, + "tags": map[string]interface{}{ + "type": "array", + "maxItems": 30, + "items": map[string]interface{}{"type": "string", "maxLength": 64}, + }, + }, + }, + "AssetImportRequest": map[string]interface{}{ + "type": "object", + "required": []string{"assets"}, + "properties": map[string]interface{}{ + "assets": map[string]interface{}{ + "type": "array", + "minItems": 1, + "maxItems": 100000, + "items": map[string]interface{}{"$ref": "#/components/schemas/AssetImportItem"}, + }, + "source": map[string]interface{}{"type": "string", "description": "未在资产中填写来源时使用的默认来源"}, + "source_query": map[string]interface{}{"type": "string", "description": "默认来源查询或导入文件名"}, + }, + }, + "AssetImportResult": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "created": map[string]interface{}{"type": "integer", "description": "新建数量", "example": 120}, + "updated": map[string]interface{}{"type": "integer", "description": "去重合并数量", "example": 8}, + "skipped": map[string]interface{}{"type": "integer", "description": "跳过数量", "example": 2}, + }, + }, + "ExecutionResult": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "id": map[string]interface{}{ + "type": "string", + "description": "执行ID", + }, + "toolName": map[string]interface{}{ + "type": "string", + "description": "工具名称", + }, + "status": map[string]interface{}{ + "type": "string", + "description": "执行状态", + "enum": []string{"queued", "running", "completed", "failed", "cancelled", "hard_timeout", "orphaned"}, + }, + "result": map[string]interface{}{ + "type": "string", + "description": "执行结果", + }, + "createdAt": map[string]interface{}{ + "type": "string", + "format": "date-time", + "description": "创建时间", + }, + }, + }, + "Error": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "error": map[string]interface{}{ + "type": "string", + "description": "错误信息", + }, + }, + }, + "LoginRequest": map[string]interface{}{ + "type": "object", + "required": []string{"password"}, + "properties": map[string]interface{}{ + "password": map[string]interface{}{ + "type": "string", + "description": "登录密码", + }, + }, + }, + "LoginResponse": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "token": map[string]interface{}{ + "type": "string", + "description": "认证Token", + }, + "expires_at": map[string]interface{}{ + "type": "string", + "format": "date-time", + "description": "Token过期时间", + }, + "session_duration_hr": map[string]interface{}{ + "type": "integer", + "description": "会话持续时间（小时）", + }, + }, + }, + "ChangePasswordRequest": map[string]interface{}{ + "type": "object", + "required": []string{"oldPassword", "newPassword"}, + "properties": map[string]interface{}{ + "oldPassword": map[string]interface{}{ + "type": "string", + "description": "当前密码", + }, + "newPassword": map[string]interface{}{ + "type": "string", + "description": "新密码（至少8位）", + }, + }, + }, + "UpdateConversationRequest": map[string]interface{}{ + "type": "object", + "required": []string{"title"}, + "properties": map[string]interface{}{ + "title": map[string]interface{}{ + "type": "string", + "description": "对话标题", + }, + }, + }, + "Group": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "id": map[string]interface{}{ + "type": "string", + "description": "分组ID", + }, + "name": map[string]interface{}{ + "type": "string", + "description": "分组名称", + }, + "icon": map[string]interface{}{ + "type": "string", + "description": "分组图标", + }, + "createdAt": map[string]interface{}{ + "type": "string", + "format": "date-time", + "description": "创建时间", + }, + "updatedAt": map[string]interface{}{ + "type": "string", + "format": "date-time", + "description": "更新时间", + }, + }, + }, + "CreateGroupRequest": map[string]interface{}{ + "type": "object", + "required": []string{"name"}, + "properties": map[string]interface{}{ + "name": map[string]interface{}{ + "type": "string", + "description": "分组名称", + }, + "icon": map[string]interface{}{ + "type": "string", + "description": "分组图标（可选）", + }, + }, + }, + "UpdateGroupRequest": map[string]interface{}{ + "type": "object", + "required": []string{"name"}, + "properties": map[string]interface{}{ + "name": map[string]interface{}{ + "type": "string", + "description": "分组名称", + }, + "icon": map[string]interface{}{ + "type": "string", + "description": "分组图标", + }, + }, + }, + "AddConversationToGroupRequest": map[string]interface{}{ + "type": "object", + "required": []string{"conversationId", "groupId"}, + "properties": map[string]interface{}{ + "conversationId": map[string]interface{}{ + "type": "string", + "description": "对话ID", + }, + "groupId": map[string]interface{}{ + "type": "string", + "description": "分组ID", + }, + }, + }, + "BatchTaskRequest": map[string]interface{}{ + "type": "object", + "required": []string{"tasks"}, + "properties": map[string]interface{}{ + "title": map[string]interface{}{ + "type": "string", + "description": "任务标题（可选）", + }, + "tasks": map[string]interface{}{ + "type": "array", + "description": "任务列表，每行一个任务", + "items": map[string]interface{}{ + "type": "string", + }, + }, + "role": map[string]interface{}{ + "type": "string", + "description": "角色名称（可选）", + }, + "agentMode": map[string]interface{}{ + "type": "string", + "description": "代理模式：eino_single（Eino ADK 单代理，默认）| deep | plan_execute | supervisor", + "enum": []string{"eino_single", "deep", "plan_execute", "supervisor"}, + }, + "scheduleMode": map[string]interface{}{ + "type": "string", + "description": "调度方式（manual | cron）", + "enum": []string{"manual", "cron"}, + }, + "cronExpr": map[string]interface{}{ + "type": "string", + "description": "Cron 表达式（scheduleMode=cron 时必填）", + }, + "executeNow": map[string]interface{}{ + "type": "boolean", + "description": "是否创建后立即执行（默认 false）", + }, + }, + }, + "BatchQueue": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "id": map[string]interface{}{ + "type": "string", + "description": "队列ID", + }, + "title": map[string]interface{}{ + "type": "string", + "description": "队列标题", + }, + "status": map[string]interface{}{ + "type": "string", + "description": "队列状态", + "enum": []string{"pending", "running", "paused", "completed", "failed"}, + }, + "tasks": map[string]interface{}{ + "type": "array", + "description": "任务列表", + "items": map[string]interface{}{ + "type": "object", + }, + }, + "createdAt": map[string]interface{}{ + "type": "string", + "format": "date-time", + "description": "创建时间", + }, + }, + }, + "CancelAgentLoopRequest": map[string]interface{}{ + "type": "object", + "required": []string{"conversationId"}, + "properties": map[string]interface{}{ + "conversationId": map[string]interface{}{ + "type": "string", + "description": "对话ID", + }, + "reason": map[string]interface{}{ + "type": "string", + "description": "可选。与 MCP 监控页「终止并说明」一致：非空时合并进当前工具返回给模型的文本（含 USER INTERRUPT NOTE 块）", + }, + "continueAfter": map[string]interface{}{ + "type": "boolean", + "description": "为 true 时仅终止当前进行中的 MCP 工具调用（不取消整轮任务）；须已有工具在执行，否则 400", + }, + }, + }, + "AgentTask": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "conversationId": map[string]interface{}{ + "type": "string", + "description": "对话ID", + }, + "status": map[string]interface{}{ + "type": "string", + "description": "任务状态", + "enum": []string{"running", "completed", "failed", "cancelled", "timeout"}, + }, + "startedAt": map[string]interface{}{ + "type": "string", + "format": "date-time", + "description": "开始时间", + }, + }, + }, + "CreateVulnerabilityRequest": map[string]interface{}{ + "type": "object", + "required": []string{"conversation_id", "title", "description", "severity", "type", "target", "reproduction_steps", "evidence", "impact", "recommendation"}, + "properties": map[string]interface{}{ + "conversation_id": map[string]interface{}{ + "type": "string", + "description": "对话ID", + }, + "title": map[string]interface{}{ + "type": "string", + "description": "漏洞标题", + }, + "description": map[string]interface{}{ + "type": "string", + "description": "漏洞描述", + }, + "severity": map[string]interface{}{ + "type": "string", + "description": "严重程度", + "enum": []string{"critical", "high", "medium", "low", "info"}, + }, + "status": map[string]interface{}{ + "type": "string", + "description": "状态", + "enum": []string{"open", "closed", "fixed"}, + }, + "type": map[string]interface{}{ + "type": "string", + "description": "漏洞类型", + }, + "target": map[string]interface{}{ + "type": "string", + "description": "受影响的目标", + }, + "preconditions": map[string]interface{}{"type": "string", "description": "前置条件"}, + "reproduction_steps": map[string]interface{}{"type": "string", "description": "复现步骤"}, + "evidence": map[string]interface{}{"type": "string", "description": "证据/POC，包含请求响应、命令输出、截图说明、日志等"}, + "impact": map[string]interface{}{ + "type": "string", + "description": "影响", + }, + "recommendation": map[string]interface{}{ + "type": "string", + "description": "修复建议", + }, + "retest_notes": map[string]interface{}{"type": "string", "description": "复测方式"}, + }, + }, + "UpdateVulnerabilityRequest": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "title": map[string]interface{}{ + "type": "string", + "description": "漏洞标题", + }, + "description": map[string]interface{}{ + "type": "string", + "description": "漏洞描述", + }, + "severity": map[string]interface{}{ + "type": "string", + "description": "严重程度", + "enum": []string{"critical", "high", "medium", "low", "info"}, + }, + "status": map[string]interface{}{ + "type": "string", + "description": "状态", + "enum": []string{"open", "confirmed", "fixed", "false_positive", "ignored"}, + }, + "type": map[string]interface{}{ + "type": "string", + "description": "漏洞类型", + }, + "target": map[string]interface{}{ + "type": "string", + "description": "受影响的目标", + }, + "preconditions": map[string]interface{}{"type": "string", "description": "前置条件"}, + "reproduction_steps": map[string]interface{}{"type": "string", "description": "复现步骤"}, + "evidence": map[string]interface{}{"type": "string", "description": "证据/POC，包含请求响应、命令输出、截图说明、日志等"}, + "impact": map[string]interface{}{ + "type": "string", + "description": "影响", + }, + "recommendation": map[string]interface{}{ + "type": "string", + "description": "修复建议", + }, + "retest_notes": map[string]interface{}{"type": "string", "description": "复测方式"}, + }, + }, + "ListVulnerabilitiesResponse": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "vulnerabilities": map[string]interface{}{ + "type": "array", + "description": "漏洞列表", + "items": map[string]interface{}{ + "$ref": "#/components/schemas/Vulnerability", + }, + }, + "total": map[string]interface{}{ + "type": "integer", + "description": "总数", + }, + "page": map[string]interface{}{ + "type": "integer", + "description": "当前页", + }, + "page_size": map[string]interface{}{ + "type": "integer", + "description": "每页数量", + }, + "total_pages": map[string]interface{}{ + "type": "integer", + "description": "总页数", + }, + }, + }, + "VulnerabilityStats": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "total": map[string]interface{}{ + "type": "integer", + "description": "总漏洞数", + }, + "by_severity": map[string]interface{}{ + "type": "object", + "description": "按严重程度统计", + }, + "by_status": map[string]interface{}{ + "type": "object", + "description": "按状态统计", + }, + }, + }, + "RoleConfig": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "name": map[string]interface{}{ + "type": "string", + "description": "角色名称", + }, + "description": map[string]interface{}{ + "type": "string", + "description": "角色描述", + }, + "enabled": map[string]interface{}{ + "type": "boolean", + "description": "是否启用", + }, + "systemPrompt": map[string]interface{}{ + "type": "string", + "description": "系统提示词", + }, + "userPrompt": map[string]interface{}{ + "type": "string", + "description": "用户提示词", + }, + "tools": map[string]interface{}{ + "type": "array", + "description": "工具列表", + "items": map[string]interface{}{ + "type": "string", + }, + }, + }, + }, + "Skill": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "name": map[string]interface{}{ + "type": "string", + "description": "Skill名称", + }, + "description": map[string]interface{}{ + "type": "string", + "description": "Skill描述", + }, + "path": map[string]interface{}{ + "type": "string", + "description": "Skill路径", + }, + }, + }, + "CreateSkillRequest": map[string]interface{}{ + "type": "object", + "required": []string{"name", "description"}, + "properties": map[string]interface{}{ + "name": map[string]interface{}{ + "type": "string", + "description": "Skill名称", + }, + "description": map[string]interface{}{ + "type": "string", + "description": "Skill描述", + }, + }, + }, + "UpdateSkillRequest": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "description": map[string]interface{}{ + "type": "string", + "description": "Skill描述", + }, + }, + }, + "ToolExecution": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "id": map[string]interface{}{ + "type": "string", + "description": "执行ID", + }, + "toolName": map[string]interface{}{ + "type": "string", + "description": "工具名称", + }, + "status": map[string]interface{}{ + "type": "string", + "description": "执行状态", + "enum": []string{"queued", "running", "completed", "failed", "cancelled", "hard_timeout", "orphaned"}, + }, + "createdAt": map[string]interface{}{ + "type": "string", + "format": "date-time", + "description": "创建时间", + }, + }, + }, + "MonitorResponse": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "executions": map[string]interface{}{ + "type": "array", + "description": "执行记录列表（轻量字段，不含 arguments/result）", + "items": map[string]interface{}{ + "$ref": "#/components/schemas/ToolExecution", + }, + }, + "summary": map[string]interface{}{ + "type": "object", + "description": "工具调用汇总", + }, + "topTools": map[string]interface{}{ + "type": "array", + "description": "调用量 Top N 工具", + "items": map[string]interface{}{ + "type": "object", + }, + }, + "timestamp": map[string]interface{}{ + "type": "string", + "format": "date-time", + "description": "时间戳", + }, + "total": map[string]interface{}{ + "type": "integer", + "description": "执行记录总数", + }, + "page": map[string]interface{}{ + "type": "integer", + "description": "当前页", + }, + "pageSize": map[string]interface{}{ + "type": "integer", + "description": "每页数量", + }, + "totalPages": map[string]interface{}{ + "type": "integer", + "description": "总页数", + }, + "retentionDays": map[string]interface{}{ + "type": "integer", + "description": "执行记录保留天数", + }, + }, + }, + "ConfigResponse": map[string]interface{}{ + "type": "object", + "description": "配置信息（含 openai、vision、multi_agent 等）", + "properties": map[string]interface{}{ + "agent": map[string]interface{}{ + "$ref": "#/components/schemas/AgentConfig", + }, + "vision": map[string]interface{}{ + "$ref": "#/components/schemas/VisionConfig", + }, + }, + }, + "UpdateConfigRequest": map[string]interface{}{ + "type": "object", + "description": "更新配置请求", + "properties": map[string]interface{}{ + "agent": map[string]interface{}{ + "$ref": "#/components/schemas/AgentConfig", + }, + "vision": map[string]interface{}{ + "$ref": "#/components/schemas/VisionConfig", + }, + }, + }, + "AgentConfig": map[string]interface{}{ + "type": "object", + "description": "Agent 运行与外部 MCP 防卡死保护配置", + "properties": map[string]interface{}{ + "max_iterations": map[string]interface{}{"type": "integer", "description": "最大迭代次数"}, + "tool_timeout_minutes": map[string]interface{}{"type": "integer", "description": "单次工具执行硬超时（分钟）"}, + "tool_wait_timeout_seconds": map[string]interface{}{"type": "integer", "description": "工具单轮等待秒数；到时返回 execution_id，worker 继续后台执行"}, + "external_mcp_max_concurrent_per_server": map[string]interface{}{"type": "integer", "description": "单个外部 MCP server 并发上限；0=默认2；负数=不限制"}, + "external_mcp_max_concurrent_total": map[string]interface{}{"type": "integer", "description": "外部 MCP 全局并发上限；0=默认16；负数=不限制"}, + "external_mcp_circuit_failure_threshold": map[string]interface{}{"type": "integer", "description": "连续失败熔断阈值；0=默认3；负数=关闭熔断"}, + "external_mcp_circuit_cooldown_seconds": map[string]interface{}{"type": "integer", "description": "熔断冷却秒数；0=默认60"}, + "shell_no_output_timeout_seconds": map[string]interface{}{"type": "integer", "description": "execute/exec 连续无输出终止秒数"}, + "workspace_root_dir": map[string]interface{}{"type": "string", "description": "会话工作目录根路径"}, + "system_prompt_path": map[string]interface{}{"type": "string", "description": "单代理系统提示文件路径"}, + }, + }, + "VisionConfig": map[string]interface{}{ + "type": "object", + "description": "视觉分析（analyze_image MCP 工具）；enabled 且 model 非空时注册工具", + "properties": map[string]interface{}{ + "enabled": map[string]interface{}{"type": "boolean", "description": "是否启用 analyze_image"}, + "model": map[string]interface{}{"type": "string", "description": "视觉模型名（必填）", "example": "qwen-vl-max"}, + "api_key": map[string]interface{}{"type": "string", "description": "API Key；留空复用 openai.api_key"}, + "base_url": map[string]interface{}{"type": "string", "description": "Base URL；留空复用 openai.base_url"}, + "provider": map[string]interface{}{"type": "string", "description": "提供商；留空复用 openai.provider"}, + "timeout_seconds": map[string]interface{}{"type": "integer", "description": "VL 调用超时（秒）"}, + "max_image_bytes": map[string]interface{}{"type": "integer", "description": "原始文件大小上限（字节）"}, + "max_dimension": map[string]interface{}{"type": "integer", "description": "长边缩放像素"}, + "jpeg_quality": map[string]interface{}{"type": "integer", "description": "JPEG 质量 60-100"}, + "max_payload_bytes": map[string]interface{}{"type": "integer", "description": "送 API 体积上限（字节）"}, + "skip_preprocess_below_bytes": map[string]interface{}{"type": "integer", "description": "低于该字节且尺寸合规时可原图直传；0=始终压缩"}, + "detail": map[string]interface{}{"type": "string", "enum": []string{"low", "high", "auto"}, "description": "OpenAI 兼容 image detail"}, + }, + }, + "AnalyzeImageToolCall": map[string]interface{}{ + "type": "object", + "description": "内置 MCP 工具 analyze_image：分析服务器本地图片，返回纯文本（验证码/UI/报错等）", + "properties": map[string]interface{}{ + "path": map[string]interface{}{ + "type": "string", + "description": "图片绝对路径或相对于进程工作目录的路径", + }, + "question": map[string]interface{}{ + "type": "string", + "description": "可选：重点问题；验证码建议「只输出验证码字符」", + }, + }, + "required": []string{"path"}, + }, + "ExternalMCPConfig": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "enabled": map[string]interface{}{ + "type": "boolean", + "description": "是否启用", + }, + "command": map[string]interface{}{ + "type": "string", + "description": "命令", + }, + "args": map[string]interface{}{ + "type": "array", + "description": "参数列表", + "items": map[string]interface{}{ + "type": "string", + }, + }, + }, + }, + "ExternalMCPResponse": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "config": map[string]interface{}{ + "$ref": "#/components/schemas/ExternalMCPConfig", + }, + "status": map[string]interface{}{ + "type": "string", + "description": "状态", + "enum": []string{"connected", "disconnected", "error", "disabled"}, + }, + "toolCount": map[string]interface{}{ + "type": "integer", + "description": "工具数量", + }, + "error": map[string]interface{}{ + "type": "string", + "description": "错误信息", + }, + }, + }, + "AddOrUpdateExternalMCPRequest": map[string]interface{}{ + "type": "object", + "required": []string{"config"}, + "properties": map[string]interface{}{ + "config": map[string]interface{}{ + "$ref": "#/components/schemas/ExternalMCPConfig", + }, + }, + }, + "AttackChain": map[string]interface{}{ + "type": "object", + "description": "攻击链数据", + }, + "MCPMessage": map[string]interface{}{ + "type": "object", + "description": "MCP消息（符合JSON-RPC 2.0规范）", + "required": []string{"jsonrpc"}, + "properties": map[string]interface{}{ + "id": map[string]interface{}{ + "description": "消息ID，可以是字符串、数字或null。对于请求，必须提供；对于通知，可以省略", + "oneOf": []map[string]interface{}{ + {"type": "string"}, + {"type": "number"}, + {"type": "null"}, + }, + "example": "550e8400-e29b-41d4-a716-446655440000", + }, + "method": map[string]interface{}{ + "type": "string", + "description": "方法名。支持的方法：\n- `initialize`: 初始化MCP连接\n- `tools/list`: 列出所有可用工具\n- `tools/call`: 调用工具\n- `prompts/list`: 列出所有提示词模板\n- `prompts/get`: 获取提示词模板\n- `resources/list`: 列出所有资源\n- `resources/read`: 读取资源内容\n- `sampling/request`: 采样请求", + "enum": []string{ + "initialize", + "tools/list", + "tools/call", + "prompts/list", + "prompts/get", + "resources/list", + "resources/read", + "sampling/request", + }, + "example": "tools/list", + }, + "params": map[string]interface{}{ + "description": "方法参数（JSON对象），根据不同的method有不同的结构", + "type": "object", + }, + "jsonrpc": map[string]interface{}{ + "type": "string", + "description": "JSON-RPC版本，固定为\"2.0\"", + "enum": []string{"2.0"}, + "example": "2.0", + }, + }, + }, + "MCPInitializeParams": map[string]interface{}{ + "type": "object", + "required": []string{"protocolVersion", "capabilities", "clientInfo"}, + "properties": map[string]interface{}{ + "protocolVersion": map[string]interface{}{ + "type": "string", + "description": "协议版本", + "example": "2024-11-05", + }, + "capabilities": map[string]interface{}{ + "type": "object", + "description": "客户端能力", + }, + "clientInfo": map[string]interface{}{ + "type": "object", + "required": []string{"name", "version"}, + "properties": map[string]interface{}{ + "name": map[string]interface{}{ + "type": "string", + "description": "客户端名称", + "example": "MyClient", + }, + "version": map[string]interface{}{ + "type": "string", + "description": "客户端版本", + "example": "1.0.0", + }, + }, + }, + }, + }, + "MCPCallToolParams": map[string]interface{}{ + "type": "object", + "required": []string{"name", "arguments"}, + "properties": map[string]interface{}{ + "name": map[string]interface{}{ + "type": "string", + "description": "工具名称", + "example": "nmap", + }, + "arguments": map[string]interface{}{ + "type": "object", + "description": "工具参数（键值对），具体参数取决于工具定义", + "example": map[string]interface{}{ + "target": "192.168.1.1", + "ports": "80,443", + }, + }, + }, + }, + "MCPResponse": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "id": map[string]interface{}{ + "description": "消息ID（与请求中的id相同）", + "oneOf": []map[string]interface{}{ + {"type": "string"}, + {"type": "number"}, + {"type": "null"}, + }, + }, + "result": map[string]interface{}{ + "description": "方法执行结果（JSON对象），结构取决于调用的方法", + "type": "object", + }, + "error": map[string]interface{}{ + "type": "object", + "description": "错误信息（如果执行失败）", + "properties": map[string]interface{}{ + "code": map[string]interface{}{ + "type": "integer", + "description": "错误代码", + "example": -32600, + }, + "message": map[string]interface{}{ + "type": "string", + "description": "错误消息", + "example": "Invalid Request", + }, + "data": map[string]interface{}{ + "description": "错误详情（可选）", + }, + }, + }, + "jsonrpc": map[string]interface{}{ + "type": "string", + "description": "JSON-RPC版本", + "example": "2.0", + }, + }, + }, + }, + }, + "security": []map[string]interface{}{ + { + "bearerAuth": []string{}, + }, + }, + "paths": map[string]interface{}{ + "/api/auth/login": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"认证"}, + "summary": "用户登录", + "description": "使用密码登录获取认证Token", + "operationId": "login", + "security": []map[string]interface{}{}, + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/LoginRequest", + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "登录成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/LoginResponse", + }, + }, + }, + }, + "401": map[string]interface{}{ + "description": "密码错误", + }, + }, + }, + }, + "/api/auth/logout": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"认证"}, + "summary": "用户登出", + "description": "登出当前会话，使Token失效", + "operationId": "logout", + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "登出成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "message": map[string]interface{}{ + "type": "string", + "example": "已退出登录", + }, + }, + }, + }, + }, + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/auth/change-password": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"认证"}, + "summary": "修改密码", + "description": "修改登录密码，修改后所有会话将失效", + "operationId": "changePassword", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/ChangePasswordRequest", + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "密码修改成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "message": map[string]interface{}{ + "type": "string", + "example": "密码已更新，请使用新密码重新登录", + }, + }, + }, + }, + }, + }, + "400": map[string]interface{}{ + "description": "请求参数错误", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/auth/validate": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"认证"}, + "summary": "验证Token", + "description": "验证当前Token是否有效", + "operationId": "validateToken", + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "Token有效", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "token": map[string]interface{}{ + "type": "string", + "description": "Token", + }, + "expires_at": map[string]interface{}{ + "type": "string", + "format": "date-time", + "description": "过期时间", + }, + }, + }, + }, + }, + }, + "401": map[string]interface{}{ + "description": "Token无效或已过期", + }, + }, + }, + }, + "/api/conversations": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"对话管理"}, + "summary": "创建对话", + "description": "创建一个新的安全测试对话。\n**重要说明**：\n- ✅ 创建的对话会**立即保存到数据库**\n- ✅ 前端页面会**自动刷新**显示新对话\n- ✅ 与前端创建的对话**完全一致**\n**创建对话的两种方式**：\n**方式1（推荐）：** 直接使用 `/api/eino-agent` 发送消息，**不提供** `conversationId` 参数，系统会自动创建新对话并发送消息。这是最简单的方式，一步完成创建和发送。\n**方式2：** 先调用此端点创建空对话，然后使用返回的 `conversationId` 调用 `/api/eino-agent` 发送消息。适用于需要先创建对话，稍后再发送消息的场景。\n**示例**：\n```json\n{\n \"title\": \"Web应用安全测试\"\n}\n```", + "operationId": "createConversation", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/CreateConversationRequest", + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "对话创建成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/Conversation", + }, + }, + }, + }, + "400": map[string]interface{}{ + "description": "请求参数错误", + }, + "401": map[string]interface{}{ + "description": "未授权，需要有效的Token", + }, + "500": map[string]interface{}{ + "description": "服务器内部错误", + }, + }, + }, + "get": map[string]interface{}{ + "tags": []string{"对话管理"}, + "summary": "列出对话", + "description": "获取对话列表，支持分页和搜索", + "operationId": "listConversations", + "parameters": []map[string]interface{}{ + { + "name": "limit", + "in": "query", + "required": false, + "description": "返回数量限制", + "schema": map[string]interface{}{ + "type": "integer", + "default": 50, + "minimum": 1, + "maximum": 100, + }, + }, + { + "name": "offset", + "in": "query", + "required": false, + "description": "偏移量", + "schema": map[string]interface{}{ + "type": "integer", + "default": 0, + "minimum": 0, + }, + }, + { + "name": "search", + "in": "query", + "required": false, + "description": "搜索关键词", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + { + "name": "project_id", + "in": "query", + "required": false, + "description": "按项目筛选；传 __none__ 表示仅未绑定项目的对话", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + { + "name": "exclude_grouped", + "in": "query", + "required": false, + "description": "为 true 时排除已加入分组的对话（默认在未搜索且未按项目筛选时启用）", + "schema": map[string]interface{}{ + "type": "boolean", + }, + }, + { + "name": "sort_by", + "in": "query", + "required": false, + "description": "排序字段：updated_at（默认）或 created_at", + "schema": map[string]interface{}{ + "type": "string", + "enum": []string{"updated_at", "created_at"}, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "array", + "items": map[string]interface{}{ + "$ref": "#/components/schemas/Conversation", + }, + }, + }, + }, + }, + "401": map[string]interface{}{ + "description": "未授权，需要有效的Token", + }, + }, + }, + }, + "/api/conversations/{id}": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"对话管理"}, + "summary": "查看对话详情", + "description": "获取指定对话的详细信息，包括对话信息和消息列表", + "operationId": "getConversation", + "parameters": []map[string]interface{}{ + { + "name": "id", + "in": "path", + "required": true, + "description": "对话ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/ConversationDetail", + }, + }, + }, + }, + "404": map[string]interface{}{ + "description": "对话不存在", + }, + "401": map[string]interface{}{ + "description": "未授权，需要有效的Token", + }, + }, + }, + "put": map[string]interface{}{ + "tags": []string{"对话管理"}, + "summary": "更新对话", + "description": "更新对话标题", + "operationId": "updateConversation", + "parameters": []map[string]interface{}{ + { + "name": "id", + "in": "path", + "required": true, + "description": "对话ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/UpdateConversationRequest", + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "更新成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/Conversation", + }, + }, + }, + }, + "400": map[string]interface{}{ + "description": "请求参数错误", + }, + "404": map[string]interface{}{ + "description": "对话不存在", + }, + "401": map[string]interface{}{ + "description": "未授权，需要有效的Token", + }, + }, + }, + "delete": map[string]interface{}{ + "tags": []string{"对话管理"}, + "summary": "删除对话", + "description": "删除指定的对话及其会话数据（消息、攻击链等）。**漏洞记录会保留**，仅解除与会话的关联。**此操作不可恢复**。", + "operationId": "deleteConversation", + "parameters": []map[string]interface{}{ + { + "name": "id", + "in": "path", + "required": true, + "description": "对话ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "删除成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "message": map[string]interface{}{ + "type": "string", + "description": "成功消息", + "example": "删除成功", + }, + }, + }, + }, + }, + }, + "404": map[string]interface{}{ + "description": "对话不存在", + }, + "401": map[string]interface{}{ + "description": "未授权，需要有效的Token", + }, + "500": map[string]interface{}{ + "description": "服务器内部错误", + }, + }, + }, + }, + "/api/conversations/{id}/project": map[string]interface{}{ + "put": map[string]interface{}{ + "tags": []string{"对话管理"}, + "summary": "设置对话所属项目", + "description": "绑定或解除对话与项目的关联，用于共享事实黑板", + "operationId": "setConversationProject", + "parameters": []map[string]interface{}{ + { + "name": "id", "in": "path", "required": true, + "description": "对话ID", + "schema": map[string]interface{}{"type": "string"}, + }, + }, + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/SetConversationProjectRequest", + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{"description": "设置成功"}, + "400": map[string]interface{}{"description": "项目不存在或参数错误"}, + "404": map[string]interface{}{"description": "对话不存在"}, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + }, + "/api/conversations/{id}/results": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"对话管理"}, + "summary": "获取对话结果", + "description": "获取指定对话的执行结果，包括消息、漏洞信息和执行结果", + "operationId": "getConversationResults", + "parameters": []map[string]interface{}{ + { + "name": "id", + "in": "path", + "required": true, + "description": "对话ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/ConversationResults", + }, + }, + }, + }, + "404": map[string]interface{}{ + "description": "对话不存在或结果不存在", + }, + "401": map[string]interface{}{ + "description": "未授权，需要有效的Token", + }, + }, + }, + }, + "/api/eino-agent": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"对话交互"}, + "summary": "发送消息并获取 AI 回复（Eino ADK 单代理，非流式）", + "description": "向 AI 发送消息并获取回复（非流式）。由 **CloudWeGo Eino** `adk.NewChatModelAgent` + `adk.NewRunner.Run` 执行单代理 MCP 工具链。**不依赖** `multi_agent.enabled`；`multi_agent.eino_skills` / `eino_middleware` 等与多代理主代理一致时可生效。支持 `webshellConnectionId`、角色与附件。", + "operationId": "sendMessageEinoSingleAgent", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "message": map[string]interface{}{"type": "string"}, + "conversationId": map[string]interface{}{"type": "string"}, + "role": map[string]interface{}{"type": "string"}, + "webshellConnectionId": map[string]interface{}{"type": "string"}, + "finalization": finalizationRequestSchema, + }, + "required": []string{"message"}, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "成功。只有 finalized=true 表示成功最终回复；finalized=false 时 response 为未完成/阻断说明。", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{"$ref": "#/components/schemas/AgentChatResponse"}, + }, + }, + }, + "400": map[string]interface{}{"description": "参数错误"}, + "401": map[string]interface{}{"description": "未授权"}, + "500": map[string]interface{}{"description": "执行失败"}, + }, + }, + }, + "/api/eino-agent/stream": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"对话交互"}, + "summary": "发送消息并获取 AI 回复（Eino ADK 单代理，SSE）", + "description": "向 AI 发送消息并获取流式回复（SSE）。由 Eino **单代理** ADK 执行；事件类型与多代理流式一致（含 `tool_call` / `response_delta` / `thinking` 等）。`response_start` / `response_delta` 仅为候选/过程输出；只有 `type: response` 且 `data.finalized=true` 才表示成功最终回复。缺 completed 执行证据时可能先发送 `finalization_auto_continue`，表示服务端基于已有 trace 无注入续跑。`data.finalized=false` 时 message 为未完成/阻断说明。**不依赖** `multi_agent.enabled`。", + "operationId": "sendMessageEinoSingleAgentStream", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "message": map[string]interface{}{"type": "string"}, + "conversationId": map[string]interface{}{"type": "string"}, + "role": map[string]interface{}{"type": "string"}, + "webshellConnectionId": map[string]interface{}{"type": "string"}, + "finalization": finalizationRequestSchema, + }, + "required": []string{"message"}, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "text/event-stream（SSE）", + "content": map[string]interface{}{ + "text/event-stream": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "string", + "description": "SSE 流。终态 response 事件 data 包含 finalized、finalizable、status、completionReason、evidenceVerified、evidenceRefs、pendingExecutionIds、missingChecks；过程事件可能包含 finalization_auto_continue。", + }, + }, + }, + }, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + }, + "/api/multi-agent": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"对话交互"}, + "summary": "发送消息并获取 AI 回复（Eino 多代理，非流式）", + "description": "与 `POST /api/eino-agent` 请求体相同，但由 **CloudWeGo Eino** 多代理执行。编排由请求体 `orchestration`（`deep` | `plan_execute` | `supervisor`）指定，缺省为 `deep`。**前提**：`multi_agent.enabled: true`；未启用时返回 404 JSON。支持 `webshellConnectionId`。", + "operationId": "sendMessageMultiAgent", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "message": map[string]interface{}{ + "type": "string", + "description": "要发送的消息（必需）", + }, + "conversationId": map[string]interface{}{ + "type": "string", + "description": "对话 ID（可选，不提供则新建）", + }, + "role": map[string]interface{}{ + "type": "string", + "description": "角色名称（可选）", + }, + "webshellConnectionId": map[string]interface{}{ + "type": "string", + "description": "WebShell 连接 ID（可选，与 Eino 单/多代理流式行为一致）", + }, + "finalization": finalizationRequestSchema, + "orchestration": map[string]interface{}{ + "type": "string", + "description": "Eino 预置编排：deep | plan_execute | supervisor；缺省 deep", + "enum": []string{"deep", "plan_execute", "supervisor"}, + }, + }, + "required": []string{"message"}, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "成功。只有 finalized=true 表示成功最终回复；finalized=false 时 response 为未完成/阻断说明。", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{"$ref": "#/components/schemas/AgentChatResponse"}, + }, + }, + }, + "400": map[string]interface{}{"description": "参数错误"}, + "401": map[string]interface{}{"description": "未授权"}, + "404": map[string]interface{}{"description": "多代理未启用或对话不存在"}, + "500": map[string]interface{}{"description": "执行失败"}, + }, + }, + }, + "/api/multi-agent/stream": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"对话交互"}, + "summary": "发送消息并获取 AI 回复（Eino 多代理，SSE）", + "description": "与 `POST /api/eino-agent/stream` 类似；由 Eino 多代理执行。`orchestration` 指定 deep / plan_execute / supervisor，缺省 deep。`response_start` / `response_delta` 仅为候选/过程输出；只有 `type: response` 且 `data.finalized=true` 才表示成功最终回复。缺 completed 执行证据时可能先发送 `finalization_auto_continue`，表示服务端基于已有 trace 无注入续跑。**前提**：`multi_agent.enabled: true`；未启用时 SSE 内首条为 `type: error` 后接 `done`。支持 `webshellConnectionId`。", + "operationId": "sendMessageMultiAgentStream", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "message": map[string]interface{}{"type": "string"}, + "conversationId": map[string]interface{}{"type": "string"}, + "role": map[string]interface{}{"type": "string"}, + "webshellConnectionId": map[string]interface{}{"type": "string"}, + "finalization": finalizationRequestSchema, + "orchestration": map[string]interface{}{ + "type": "string", + "description": "deep | plan_execute | supervisor；缺省 deep", + "enum": []string{"deep", "plan_execute", "supervisor"}, + }, + }, + "required": []string{"message"}, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "text/event-stream（SSE）", + "content": map[string]interface{}{ + "text/event-stream": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "string", + "description": "SSE 流。终态 response 事件 data 包含 finalized、finalizable、status、completionReason、evidenceVerified、evidenceRefs、pendingExecutionIds、missingChecks；过程事件可能包含 finalization_auto_continue。", + }, + }, + }, + }, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + }, + "/api/agent-loop/cancel": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"对话交互"}, + "summary": "取消任务", + "description": "取消正在执行的Agent Loop任务", + "operationId": "cancelAgentLoop", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/CancelAgentLoopRequest", + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "取消请求已提交", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "status": map[string]interface{}{ + "type": "string", + "example": "cancelling", + }, + "conversationId": map[string]interface{}{ + "type": "string", + "description": "对话ID", + }, + "message": map[string]interface{}{ + "type": "string", + "example": "已提交取消请求，任务将在当前步骤完成后停止。", + }, + }, + }, + }, + }, + }, + "404": map[string]interface{}{ + "description": "未找到正在执行的任务", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/agent-loop/tasks": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"对话交互"}, + "summary": "列出运行中的任务", + "description": "获取所有正在运行的Agent Loop任务", + "operationId": "listAgentTasks", + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "tasks": map[string]interface{}{ + "type": "array", + "description": "任务列表", + "items": map[string]interface{}{ + "$ref": "#/components/schemas/AgentTask", + }, + }, + }, + }, + }, + }, + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/agent-loop/tasks/completed": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"对话交互"}, + "summary": "列出已完成的任务", + "description": "获取最近完成的Agent Loop任务历史", + "operationId": "listCompletedTasks", + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "tasks": map[string]interface{}{ + "type": "array", + "description": "已完成任务列表", + "items": map[string]interface{}{ + "$ref": "#/components/schemas/AgentTask", + }, + }, + }, + }, + }, + }, + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/batch-tasks": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"批量任务"}, + "summary": "创建批量任务队列", + "description": "创建一个批量任务队列，包含多个任务", + "operationId": "createBatchQueue", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/BatchTaskRequest", + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "创建成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "queueId": map[string]interface{}{ + "type": "string", + "description": "队列ID", + }, + "queue": map[string]interface{}{ + "$ref": "#/components/schemas/BatchQueue", + }, + "started": map[string]interface{}{ + "type": "boolean", + "description": "是否已立即启动执行", + }, + }, + }, + }, + }, + }, + "400": map[string]interface{}{ + "description": "请求参数错误", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + "get": map[string]interface{}{ + "tags": []string{"批量任务"}, + "summary": "列出批量任务队列", + "description": "获取所有批量任务队列", + "operationId": "listBatchQueues", + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "queues": map[string]interface{}{ + "type": "array", + "description": "队列列表", + "items": map[string]interface{}{ + "$ref": "#/components/schemas/BatchQueue", + }, + }, + }, + }, + }, + }, + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/batch-tasks/{queueId}": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"批量任务"}, + "summary": "获取批量任务队列", + "description": "获取指定批量任务队列的详细信息", + "operationId": "getBatchQueue", + "parameters": []map[string]interface{}{ + { + "name": "queueId", + "in": "path", + "required": true, + "description": "队列ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/BatchQueue", + }, + }, + }, + }, + "404": map[string]interface{}{ + "description": "队列不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + "delete": map[string]interface{}{ + "tags": []string{"批量任务"}, + "summary": "删除批量任务队列", + "description": "删除指定的批量任务队列", + "operationId": "deleteBatchQueue", + "parameters": []map[string]interface{}{ + { + "name": "queueId", + "in": "path", + "required": true, + "description": "队列ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "删除成功", + }, + "404": map[string]interface{}{ + "description": "队列不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/batch-tasks/{queueId}/start": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"批量任务"}, + "summary": "启动批量任务队列", + "description": "开始执行批量任务队列中的任务", + "operationId": "startBatchQueue", + "parameters": []map[string]interface{}{ + { + "name": "queueId", + "in": "path", + "required": true, + "description": "队列ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "启动成功", + }, + "404": map[string]interface{}{ + "description": "队列不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/batch-tasks/{queueId}/pause": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"批量任务"}, + "summary": "暂停批量任务队列", + "description": "暂停正在执行的批量任务队列", + "operationId": "pauseBatchQueue", + "parameters": []map[string]interface{}{ + { + "name": "queueId", + "in": "path", + "required": true, + "description": "队列ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "暂停成功", + }, + "404": map[string]interface{}{ + "description": "队列不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/batch-tasks/{queueId}/tasks": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"批量任务"}, + "summary": "添加任务到队列", + "description": "向批量任务队列添加新任务。任务会添加到队列末尾，按照队列顺序依次执行。每个任务会创建一个独立的对话，支持完整的状态跟踪。\n**任务格式**：\n任务内容是一个字符串，描述要执行的安全测试任务。例如：\n- \"扫描 http://example.com 的SQL注入漏洞\"\n- \"对 192.168.1.1 进行端口扫描\"\n- \"检测 https://target.com 的XSS漏洞\"\n**使用示例**：\n```json\n{\n \"task\": \"扫描 http://example.com 的SQL注入漏洞\"\n}\n```", + "operationId": "addBatchTask", + "parameters": []map[string]interface{}{ + { + "name": "queueId", + "in": "path", + "required": true, + "description": "队列ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "required": []string{"task"}, + "properties": map[string]interface{}{ + "task": map[string]interface{}{ + "type": "string", + "description": "任务内容，描述要执行的安全测试任务（必需）", + "example": "扫描 http://example.com 的SQL注入漏洞", + }, + }, + }, + "examples": map[string]interface{}{ + "sqlInjection": map[string]interface{}{ + "summary": "SQL注入扫描", + "description": "扫描目标网站的SQL注入漏洞", + "value": map[string]interface{}{ + "task": "扫描 http://example.com 的SQL注入漏洞", + }, + }, + "portScan": map[string]interface{}{ + "summary": "端口扫描", + "description": "对目标IP进行端口扫描", + "value": map[string]interface{}{ + "task": "对 192.168.1.1 进行端口扫描", + }, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "添加成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "taskId": map[string]interface{}{ + "type": "string", + "description": "新添加的任务ID", + }, + "message": map[string]interface{}{ + "type": "string", + "description": "成功消息", + "example": "任务已添加到队列", + }, + }, + }, + }, + }, + }, + "400": map[string]interface{}{ + "description": "请求参数错误（如task为空）", + }, + "404": map[string]interface{}{ + "description": "队列不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/batch-tasks/{queueId}/tasks/{taskId}": map[string]interface{}{ + "put": map[string]interface{}{ + "tags": []string{"批量任务"}, + "summary": "更新批量任务", + "description": "更新批量任务队列中的指定任务", + "operationId": "updateBatchTask", + "parameters": []map[string]interface{}{ + { + "name": "queueId", + "in": "path", + "required": true, + "description": "队列ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + { + "name": "taskId", + "in": "path", + "required": true, + "description": "任务ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "task": map[string]interface{}{ + "type": "string", + "description": "任务内容", + }, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "更新成功", + }, + "404": map[string]interface{}{ + "description": "任务不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + "delete": map[string]interface{}{ + "tags": []string{"批量任务"}, + "summary": "删除批量任务", + "description": "从批量任务队列中删除指定任务", + "operationId": "deleteBatchTask", + "parameters": []map[string]interface{}{ + { + "name": "queueId", + "in": "path", + "required": true, + "description": "队列ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + { + "name": "taskId", + "in": "path", + "required": true, + "description": "任务ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "删除成功", + }, + "404": map[string]interface{}{ + "description": "任务不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/groups": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"对话分组"}, + "summary": "创建分组", + "description": "创建一个新的对话分组", + "operationId": "createGroup", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/CreateGroupRequest", + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "创建成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/Group", + }, + }, + }, + }, + "400": map[string]interface{}{ + "description": "请求参数错误或分组名称已存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + "get": map[string]interface{}{ + "tags": []string{"对话分组"}, + "summary": "列出分组", + "description": "获取所有对话分组", + "operationId": "listGroups", + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "array", + "items": map[string]interface{}{ + "$ref": "#/components/schemas/Group", + }, + }, + }, + }, + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/groups/{id}": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"对话分组"}, + "summary": "获取分组", + "description": "获取指定分组的详细信息", + "operationId": "getGroup", + "parameters": []map[string]interface{}{ + { + "name": "id", + "in": "path", + "required": true, + "description": "分组ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/Group", + }, + }, + }, + }, + "404": map[string]interface{}{ + "description": "分组不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + "put": map[string]interface{}{ + "tags": []string{"对话分组"}, + "summary": "更新分组", + "description": "更新分组信息", + "operationId": "updateGroup", + "parameters": []map[string]interface{}{ + { + "name": "id", + "in": "path", + "required": true, + "description": "分组ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/UpdateGroupRequest", + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "更新成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/Group", + }, + }, + }, + }, + "400": map[string]interface{}{ + "description": "请求参数错误或分组名称已存在", + }, + "404": map[string]interface{}{ + "description": "分组不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + "delete": map[string]interface{}{ + "tags": []string{"对话分组"}, + "summary": "删除分组", + "description": "删除指定分组", + "operationId": "deleteGroup", + "parameters": []map[string]interface{}{ + { + "name": "id", + "in": "path", + "required": true, + "description": "分组ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "删除成功", + }, + "404": map[string]interface{}{ + "description": "分组不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/groups/{id}/conversations": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"对话分组"}, + "summary": "获取分组中的对话", + "description": "获取指定分组中的所有对话", + "operationId": "getGroupConversations", + "parameters": []map[string]interface{}{ + { + "name": "id", + "in": "path", + "required": true, + "description": "分组ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "array", + "items": map[string]interface{}{ + "$ref": "#/components/schemas/Conversation", + }, + }, + }, + }, + }, + "404": map[string]interface{}{ + "description": "分组不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/groups/conversations": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"对话分组"}, + "summary": "添加对话到分组", + "description": "将对话添加到指定分组", + "operationId": "addConversationToGroup", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/AddConversationToGroupRequest", + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "添加成功", + }, + "400": map[string]interface{}{ + "description": "请求参数错误", + }, + "404": map[string]interface{}{ + "description": "对话或分组不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/groups/{id}/conversations/{conversationId}": map[string]interface{}{ + "delete": map[string]interface{}{ + "tags": []string{"对话分组"}, + "summary": "从分组移除对话", + "description": "从指定分组中移除对话", + "operationId": "removeConversationFromGroup", + "parameters": []map[string]interface{}{ + { + "name": "id", + "in": "path", + "required": true, + "description": "分组ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + { + "name": "conversationId", + "in": "path", + "required": true, + "description": "对话ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "移除成功", + }, + "404": map[string]interface{}{ + "description": "对话或分组不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/assets/import": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"资产管理"}, + "summary": "批量导入资产", + "description": "新增或按“目标 + 端口 + 协议”去重更新资产。接收 JSON，不直接接收 XLSX/CSV 文件；单次最多 100000 条，需要 asset:write 权限。", + "operationId": "importAssets", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{"$ref": "#/components/schemas/AssetImportRequest"}, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "导入完成", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{"$ref": "#/components/schemas/AssetImportResult"}, + }, + }, + }, + "400": map[string]interface{}{"description": "数量或资产字段校验失败"}, + "401": map[string]interface{}{"description": "未授权"}, + "403": map[string]interface{}{"description": "缺少 asset:write 权限或无权访问指定项目"}, + "500": map[string]interface{}{"description": "导入事务失败"}, + }, + }, + }, + "/api/projects": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"项目管理"}, + "summary": "列出项目", + "operationId": "listProjects", + "parameters": []map[string]interface{}{ + {"name": "status", "in": "query", "schema": map[string]interface{}{"type": "string", "enum": []string{"active", "archived"}}}, + {"name": "limit", "in": "query", "schema": map[string]interface{}{"type": "integer", "default": 200}}, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{"description": "项目列表"}, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + "post": map[string]interface{}{ + "tags": []string{"项目管理"}, + "summary": "创建项目", + "operationId": "createProject", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "name": map[string]interface{}{"type": "string"}, + "description": map[string]interface{}{"type": "string"}, + "scope_json": map[string]interface{}{"type": "string"}, + }, + "required": []string{"name"}, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{"description": "创建成功"}, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + }, + "/api/projects/{id}": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"项目管理"}, "summary": "获取项目", "operationId": "getProject", + "parameters": []map[string]interface{}{ + {"name": "id", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}}, + }, + "responses": map[string]interface{}{"200": map[string]interface{}{"description": "项目详情"}}, + }, + "put": map[string]interface{}{ + "tags": []string{"项目管理"}, "summary": "更新项目", "operationId": "updateProject", + "parameters": []map[string]interface{}{ + {"name": "id", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}}, + }, + "responses": map[string]interface{}{"200": map[string]interface{}{"description": "更新成功"}}, + }, + "delete": map[string]interface{}{ + "tags": []string{"项目管理"}, "summary": "删除项目", "operationId": "deleteProject", + "parameters": []map[string]interface{}{ + {"name": "id", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}}, + }, + "responses": map[string]interface{}{"200": map[string]interface{}{"description": "删除成功"}}, + }, + }, + "/api/projects/{id}/facts": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"项目管理"}, "summary": "列出或按 key 获取事实", "operationId": "listProjectFacts", + "parameters": []map[string]interface{}{ + {"name": "id", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}}, + {"name": "fact_key", "in": "query", "schema": map[string]interface{}{"type": "string"}}, + {"name": "include_links", "in": "query", "schema": map[string]interface{}{"type": "boolean"}}, + {"name": "include_link_counts", "in": "query", "schema": map[string]interface{}{"type": "boolean"}}, + }, + "responses": map[string]interface{}{"200": map[string]interface{}{"description": "事实列表或单条（可含 link_counts / outgoing_links）"}}, + }, + "post": map[string]interface{}{ + "tags": []string{"项目管理"}, "summary": "创建/更新事实", "operationId": "upsertProjectFactREST", + "parameters": []map[string]interface{}{ + {"name": "id", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}}, + }, + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "fact_key": map[string]interface{}{"type": "string"}, + "summary": map[string]interface{}{"type": "string"}, + "links": map[string]interface{}{ + "type": "array", + "items": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "to": map[string]interface{}{"type": "string"}, + "type": map[string]interface{}{"type": "string"}, + }, + }, + }, + "links_text": map[string]interface{}{"type": "string", "description": "type: fact_key 每行一条"}, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{"200": map[string]interface{}{"description": "成功"}}, + }, + }, + "/api/projects/{id}/fact-graph": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"项目管理"}, "summary": "获取项目事实攻击路径图", "operationId": "getProjectFactGraph", + "parameters": []map[string]interface{}{ + {"name": "id", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}}, + {"name": "view", "in": "query", "schema": map[string]interface{}{"type": "string", "enum": []string{"path", "full"}, "default": "path"}}, + {"name": "exclude_deprecated", "in": "query", "schema": map[string]interface{}{"type": "boolean", "default": true}}, + }, + "responses": map[string]interface{}{"200": map[string]interface{}{"description": "nodes + edges"}}, + }, + }, + "/api/projects/{id}/fact-edges": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"项目管理"}, "summary": "列出项目全部事实边", "operationId": "listProjectFactEdges", + "parameters": []map[string]interface{}{ + {"name": "id", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}}, + }, + "responses": map[string]interface{}{"200": map[string]interface{}{"description": "边列表"}}, + }, + "post": map[string]interface{}{ + "tags": []string{"项目管理"}, "summary": "添加事实边", "operationId": "createProjectFactEdge", + "parameters": []map[string]interface{}{ + {"name": "id", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}}, + }, + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "required": []string{"source_fact_key", "target_fact_key", "edge_type"}, + "properties": map[string]interface{}{ + "source_fact_key": map[string]interface{}{"type": "string"}, + "target_fact_key": map[string]interface{}{"type": "string"}, + "edge_type": map[string]interface{}{"type": "string"}, + "confidence": map[string]interface{}{"type": "string"}, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{"200": map[string]interface{}{"description": "边已创建"}}, + }, + }, + "/api/projects/{id}/fact-edges/{edgeId}": map[string]interface{}{ + "delete": map[string]interface{}{ + "tags": []string{"项目管理"}, "summary": "删除事实边", "operationId": "deleteProjectFactEdge", + "parameters": []map[string]interface{}{ + {"name": "id", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}}, + {"name": "edgeId", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}}, + }, + "responses": map[string]interface{}{"200": map[string]interface{}{"description": "删除成功"}}, + }, + }, + "/api/projects/{id}/promote-attack-chain/{conversationId}": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"项目管理"}, "summary": "将对话攻击链沉淀到项目事实图", "operationId": "promoteAttackChainToProject", + "parameters": []map[string]interface{}{ + {"name": "id", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}}, + {"name": "conversationId", "in": "path", "required": true, "schema": map[string]interface{}{"type": "string"}}, + }, + "responses": map[string]interface{}{"200": map[string]interface{}{"description": "沉淀结果（facts/edges/graph）"}}, + }, + }, + "/api/vulnerabilities": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"漏洞管理"}, + "summary": "列出漏洞", + "description": "获取漏洞列表，支持分页和筛选", + "operationId": "listVulnerabilities", + "parameters": []map[string]interface{}{ + { + "name": "limit", + "in": "query", + "required": false, + "description": "每页数量", + "schema": map[string]interface{}{ + "type": "integer", + "default": 20, + "minimum": 1, + "maximum": 100, + }, + }, + { + "name": "offset", + "in": "query", + "required": false, + "description": "偏移量", + "schema": map[string]interface{}{ + "type": "integer", + "default": 0, + "minimum": 0, + }, + }, + { + "name": "page", + "in": "query", + "required": false, + "description": "页码（与offset二选一）", + "schema": map[string]interface{}{ + "type": "integer", + "minimum": 1, + }, + }, + { + "name": "id", + "in": "query", + "required": false, + "description": "漏洞ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + { + "name": "conversation_id", + "in": "query", + "required": false, + "description": "对话ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + { + "name": "project_id", + "in": "query", + "required": false, + "description": "项目ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + { + "name": "severity", + "in": "query", + "required": false, + "description": "严重程度", + "schema": map[string]interface{}{ + "type": "string", + "enum": []string{"critical", "high", "medium", "low", "info"}, + }, + }, + { + "name": "status", + "in": "query", + "required": false, + "description": "状态", + "schema": map[string]interface{}{ + "type": "string", + "enum": []string{"open", "closed", "fixed"}, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/ListVulnerabilitiesResponse", + }, + }, + }, + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + "post": map[string]interface{}{ + "tags": []string{"漏洞管理"}, + "summary": "创建漏洞", + "description": "创建一个新的漏洞记录", + "operationId": "createVulnerability", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/CreateVulnerabilityRequest", + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "创建成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/Vulnerability", + }, + }, + }, + }, + "400": map[string]interface{}{ + "description": "请求参数错误", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/vulnerabilities/stats": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"漏洞管理"}, + "summary": "获取漏洞统计", + "description": "获取漏洞统计信息", + "operationId": "getVulnerabilityStats", + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/VulnerabilityStats", + }, + }, + }, + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/vulnerabilities/{id}": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"漏洞管理"}, + "summary": "获取漏洞", + "description": "获取指定漏洞的详细信息", + "operationId": "getVulnerability", + "parameters": []map[string]interface{}{ + { + "name": "id", + "in": "path", + "required": true, + "description": "漏洞ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/Vulnerability", + }, + }, + }, + }, + "404": map[string]interface{}{ + "description": "漏洞不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + "put": map[string]interface{}{ + "tags": []string{"漏洞管理"}, + "summary": "更新漏洞", + "description": "更新漏洞信息", + "operationId": "updateVulnerability", + "parameters": []map[string]interface{}{ + { + "name": "id", + "in": "path", + "required": true, + "description": "漏洞ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/UpdateVulnerabilityRequest", + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "更新成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/Vulnerability", + }, + }, + }, + }, + "400": map[string]interface{}{ + "description": "请求参数错误", + }, + "404": map[string]interface{}{ + "description": "漏洞不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + "delete": map[string]interface{}{ + "tags": []string{"漏洞管理"}, + "summary": "删除漏洞", + "description": "删除指定漏洞", + "operationId": "deleteVulnerability", + "parameters": []map[string]interface{}{ + { + "name": "id", + "in": "path", + "required": true, + "description": "漏洞ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "删除成功", + }, + "404": map[string]interface{}{ + "description": "漏洞不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/roles": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"角色管理"}, + "summary": "列出角色", + "description": "获取所有安全测试角色", + "operationId": "getRoles", + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "roles": map[string]interface{}{ + "type": "array", + "description": "角色列表", + "items": map[string]interface{}{ + "$ref": "#/components/schemas/RoleConfig", + }, + }, + }, + }, + }, + }, + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + "post": map[string]interface{}{ + "tags": []string{"角色管理"}, + "summary": "创建角色", + "description": "创建一个新的安全测试角色", + "operationId": "createRole", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/RoleConfig", + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "创建成功", + }, + "400": map[string]interface{}{ + "description": "请求参数错误", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/roles/{name}": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"角色管理"}, + "summary": "获取角色", + "description": "获取指定角色的详细信息", + "operationId": "getRole", + "parameters": []map[string]interface{}{ + { + "name": "name", + "in": "path", + "required": true, + "description": "角色名称", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "role": map[string]interface{}{ + "$ref": "#/components/schemas/RoleConfig", + }, + }, + }, + }, + }, + }, + "404": map[string]interface{}{ + "description": "角色不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + "put": map[string]interface{}{ + "tags": []string{"角色管理"}, + "summary": "更新角色", + "description": "更新指定角色的配置", + "operationId": "updateRole", + "parameters": []map[string]interface{}{ + { + "name": "name", + "in": "path", + "required": true, + "description": "角色名称", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/RoleConfig", + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "更新成功", + }, + "400": map[string]interface{}{ + "description": "请求参数错误", + }, + "404": map[string]interface{}{ + "description": "角色不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + "delete": map[string]interface{}{ + "tags": []string{"角色管理"}, + "summary": "删除角色", + "description": "删除指定角色", + "operationId": "deleteRole", + "parameters": []map[string]interface{}{ + { + "name": "name", + "in": "path", + "required": true, + "description": "角色名称", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "删除成功", + }, + "404": map[string]interface{}{ + "description": "角色不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/skills": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"Skills管理"}, + "summary": "列出Skills", + "description": "获取所有Skills列表，支持分页和搜索", + "operationId": "getSkills", + "parameters": []map[string]interface{}{ + { + "name": "limit", + "in": "query", + "required": false, + "description": "每页数量", + "schema": map[string]interface{}{ + "type": "integer", + "default": 20, + }, + }, + { + "name": "offset", + "in": "query", + "required": false, + "description": "偏移量", + "schema": map[string]interface{}{ + "type": "integer", + "default": 0, + }, + }, + { + "name": "search", + "in": "query", + "required": false, + "description": "搜索关键词", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "skills": map[string]interface{}{ + "type": "array", + "description": "Skills列表", + "items": map[string]interface{}{ + "$ref": "#/components/schemas/Skill", + }, + }, + "total": map[string]interface{}{ + "type": "integer", + "description": "总数", + }, + }, + }, + }, + }, + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + "post": map[string]interface{}{ + "tags": []string{"Skills管理"}, + "summary": "创建Skill", + "description": "创建一个新的Skill", + "operationId": "createSkill", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/CreateSkillRequest", + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "创建成功", + }, + "400": map[string]interface{}{ + "description": "请求参数错误", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/skills/stats": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"Skills管理"}, + "summary": "获取Skill统计", + "description": "获取Skill调用统计信息", + "operationId": "getSkillStats", + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "description": "统计信息", + }, + }, + }, + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + "delete": map[string]interface{}{ + "tags": []string{"Skills管理"}, + "summary": "清空Skill统计", + "description": "清空所有Skill的调用统计", + "operationId": "clearSkillStats", + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "清空成功", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/skills/{name}": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"Skills管理"}, + "summary": "获取Skill", + "description": "获取指定Skill的详细信息", + "operationId": "getSkill", + "parameters": []map[string]interface{}{ + { + "name": "name", + "in": "path", + "required": true, + "description": "Skill名称", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/Skill", + }, + }, + }, + }, + "404": map[string]interface{}{ + "description": "Skill不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + "put": map[string]interface{}{ + "tags": []string{"Skills管理"}, + "summary": "更新Skill", + "description": "更新指定Skill的信息", + "operationId": "updateSkill", + "parameters": []map[string]interface{}{ + { + "name": "name", + "in": "path", + "required": true, + "description": "Skill名称", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/UpdateSkillRequest", + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "更新成功", + }, + "400": map[string]interface{}{ + "description": "请求参数错误", + }, + "404": map[string]interface{}{ + "description": "Skill不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + "delete": map[string]interface{}{ + "tags": []string{"Skills管理"}, + "summary": "删除Skill", + "description": "删除指定Skill", + "operationId": "deleteSkill", + "parameters": []map[string]interface{}{ + { + "name": "name", + "in": "path", + "required": true, + "description": "Skill名称", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "删除成功", + }, + "404": map[string]interface{}{ + "description": "Skill不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/skills/{name}/bound-roles": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"Skills管理"}, + "summary": "获取绑定角色", + "description": "获取使用指定Skill的所有角色", + "operationId": "getSkillBoundRoles", + "parameters": []map[string]interface{}{ + { + "name": "name", + "in": "path", + "required": true, + "description": "Skill名称", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "roles": map[string]interface{}{ + "type": "array", + "description": "角色列表", + "items": map[string]interface{}{ + "type": "string", + }, + }, + }, + }, + }, + }, + }, + "404": map[string]interface{}{ + "description": "Skill不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/skills/{name}/stats": map[string]interface{}{ + "delete": map[string]interface{}{ + "tags": []string{"Skills管理"}, + "summary": "清空Skill统计", + "description": "清空指定Skill的调用统计", + "operationId": "clearSkillStatsByName", + "parameters": []map[string]interface{}{ + { + "name": "name", + "in": "path", + "required": true, + "description": "Skill名称", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "清空成功", + }, + "404": map[string]interface{}{ + "description": "Skill不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/monitor": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"监控"}, + "summary": "获取监控信息", + "description": "获取工具执行监控信息，支持分页和筛选", + "operationId": "monitor", + "parameters": []map[string]interface{}{ + { + "name": "page", + "in": "query", + "required": false, + "description": "页码", + "schema": map[string]interface{}{ + "type": "integer", + "default": 1, + "minimum": 1, + }, + }, + { + "name": "page_size", + "in": "query", + "required": false, + "description": "每页数量", + "schema": map[string]interface{}{ + "type": "integer", + "default": 20, + "minimum": 1, + "maximum": 100, + }, + }, + { + "name": "status", + "in": "query", + "required": false, + "description": "状态筛选", + "schema": map[string]interface{}{ + "type": "string", + "enum": []string{"queued", "running", "completed", "failed", "cancelled", "hard_timeout", "orphaned"}, + }, + }, + { + "name": "tool", + "in": "query", + "required": false, + "description": "工具名称筛选（支持部分匹配）", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/MonitorResponse", + }, + }, + }, + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/monitor/execution/{id}": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"监控"}, + "summary": "获取执行记录", + "description": "获取指定执行记录的详细信息", + "operationId": "getExecution", + "parameters": []map[string]interface{}{ + { + "name": "id", + "in": "path", + "required": true, + "description": "执行ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/ToolExecution", + }, + }, + }, + }, + "404": map[string]interface{}{ + "description": "执行记录不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + "delete": map[string]interface{}{ + "tags": []string{"监控"}, + "summary": "删除执行记录", + "description": "删除指定的执行记录", + "operationId": "deleteExecution", + "parameters": []map[string]interface{}{ + { + "name": "id", + "in": "path", + "required": true, + "description": "执行ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "删除成功", + }, + "404": map[string]interface{}{ + "description": "执行记录不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/monitor/execution/{id}/cancel": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"监控"}, + "summary": "取消进行中的工具执行", + "description": "对当前进程内正在执行的 MCP 工具调用发送 context 取消信号；上层对话/多步任务可继续。若执行已结束或未在本进程内运行则返回 404。", + "operationId": "cancelExecution", + "parameters": []map[string]interface{}{ + { + "name": "id", + "in": "path", + "required": true, + "description": "执行ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "requestBody": map[string]interface{}{ + "required": false, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "note": map[string]interface{}{ + "type": "string", + "description": "可选。非空时与工具已返回输出合并交给大模型，并带有「用户终止说明」标题块以便与命令行原文区分", + }, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "已发送终止信号", + }, + "400": map[string]interface{}{ + "description": "请求体不是合法 JSON", + }, + "404": map[string]interface{}{ + "description": "未找到进行中的工具执行", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/monitor/executions": map[string]interface{}{ + "delete": map[string]interface{}{ + "tags": []string{"监控"}, + "summary": "批量删除执行记录", + "description": "批量删除执行记录", + "operationId": "deleteExecutions", + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "删除成功", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/monitor/stats": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"监控"}, + "summary": "获取统计信息", + "description": "获取工具执行统计信息", + "operationId": "getStats", + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "description": "统计信息", + }, + }, + }, + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/config": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"配置管理"}, + "summary": "获取配置", + "description": "获取系统配置信息", + "operationId": "getConfig", + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/ConfigResponse", + }, + }, + }, + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + "put": map[string]interface{}{ + "tags": []string{"配置管理"}, + "summary": "更新配置", + "description": "更新系统配置", + "operationId": "updateConfig", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/UpdateConfigRequest", + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "更新成功", + }, + "400": map[string]interface{}{ + "description": "请求参数错误", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/config/tools": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"配置管理"}, + "summary": "获取工具配置", + "description": "获取所有工具的配置信息", + "operationId": "getTools", + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "array", + "description": "工具配置列表", + }, + }, + }, + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/config/apply": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"配置管理"}, + "summary": "应用配置", + "description": "应用配置更改", + "operationId": "applyConfig", + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "应用成功", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/external-mcp": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"外部MCP管理"}, + "summary": "列出外部MCP", + "description": "获取所有外部MCP配置和状态", + "operationId": "getExternalMCPs", + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "servers": map[string]interface{}{ + "type": "object", + "description": "MCP服务器配置", + "additionalProperties": map[string]interface{}{ + "$ref": "#/components/schemas/ExternalMCPResponse", + }, + }, + "stats": map[string]interface{}{ + "type": "object", + "description": "统计信息", + }, + }, + }, + }, + }, + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/external-mcp/stats": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"外部MCP管理"}, + "summary": "获取外部MCP统计", + "description": "获取外部MCP统计信息", + "operationId": "getExternalMCPStats", + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "description": "统计信息", + }, + }, + }, + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/external-mcp/{name}": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"外部MCP管理"}, + "summary": "获取外部MCP", + "description": "获取指定外部MCP的配置和状态", + "operationId": "getExternalMCP", + "parameters": []map[string]interface{}{ + { + "name": "name", + "in": "path", + "required": true, + "description": "MCP名称", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/ExternalMCPResponse", + }, + }, + }, + }, + "404": map[string]interface{}{ + "description": "MCP不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + "put": map[string]interface{}{ + "tags": []string{"外部MCP管理"}, + "summary": "添加或更新外部MCP", + "description": "添加新的外部MCP配置或更新现有配置。\n**传输方式**：\n支持两种传输方式：\n**1. stdio（标准输入输出）**：\n```json\n{\n \"config\": {\n \"enabled\": true,\n \"command\": \"node\",\n \"args\": [\"/path/to/mcp-server.js\"],\n \"env\": {}\n }\n}\n```\n**2. sse（Server-Sent Events）**：\n```json\n{\n \"config\": {\n \"enabled\": true,\n \"transport\": \"sse\",\n \"url\": \"http://127.0.0.1:8082/sse\",\n \"timeout\": 30\n }\n}\n```\n**配置参数说明**：\n- `enabled`: 是否启用（boolean，必需）\n- `command`: 命令（stdio模式必需，如：\"node\", \"python\"）\n- `args`: 命令参数数组（stdio模式必需）\n- `env`: 环境变量（object，可选）\n- `transport`: 传输方式（\"stdio\" 或 \"sse\"，sse模式必需）\n- `url`: SSE端点URL（sse模式必需）\n- `timeout`: 超时时间（秒，可选，默认30）\n- `description`: 描述（可选）", + "operationId": "addOrUpdateExternalMCP", + "parameters": []map[string]interface{}{ + { + "name": "name", + "in": "path", + "required": true, + "description": "MCP名称（唯一标识符）", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/AddOrUpdateExternalMCPRequest", + }, + "examples": map[string]interface{}{ + "stdio": map[string]interface{}{ + "summary": "stdio模式配置", + "description": "使用标准输入输出方式连接外部MCP服务器", + "value": map[string]interface{}{ + "config": map[string]interface{}{ + "enabled": true, + "command": "node", + "args": []string{"/path/to/mcp-server.js"}, + "env": map[string]interface{}{}, + "timeout": 30, + "description": "Node.js MCP服务器", + }, + }, + }, + "sse": map[string]interface{}{ + "summary": "SSE模式配置", + "description": "使用Server-Sent Events方式连接外部MCP服务器", + "value": map[string]interface{}{ + "config": map[string]interface{}{ + "enabled": true, + "transport": "sse", + "url": "http://127.0.0.1:8082/sse", + "timeout": 30, + "description": "SSE MCP服务器", + }, + }, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "操作成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "message": map[string]interface{}{ + "type": "string", + "example": "外部MCP配置已保存", + }, + }, + }, + }, + }, + }, + "400": map[string]interface{}{ + "description": "请求参数错误（如配置格式不正确、缺少必需字段等）", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/Error", + }, + "example": map[string]interface{}{ + "error": "stdio模式需要提供command和args参数", + }, + }, + }, + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + "delete": map[string]interface{}{ + "tags": []string{"外部MCP管理"}, + "summary": "删除外部MCP", + "description": "删除指定的外部MCP配置", + "operationId": "deleteExternalMCP", + "parameters": []map[string]interface{}{ + { + "name": "name", + "in": "path", + "required": true, + "description": "MCP名称", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "删除成功", + }, + "404": map[string]interface{}{ + "description": "MCP不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/external-mcp/{name}/start": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"外部MCP管理"}, + "summary": "启动外部MCP", + "description": "启动指定的外部MCP服务器", + "operationId": "startExternalMCP", + "parameters": []map[string]interface{}{ + { + "name": "name", + "in": "path", + "required": true, + "description": "MCP名称", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "启动成功", + }, + "404": map[string]interface{}{ + "description": "MCP不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/external-mcp/{name}/stop": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"外部MCP管理"}, + "summary": "停止外部MCP", + "description": "停止指定的外部MCP服务器", + "operationId": "stopExternalMCP", + "parameters": []map[string]interface{}{ + { + "name": "name", + "in": "path", + "required": true, + "description": "MCP名称", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "停止成功", + }, + "404": map[string]interface{}{ + "description": "MCP不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/attack-chain/{conversationId}": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"攻击链"}, + "summary": "获取攻击链", + "description": "获取指定对话的攻击链可视化数据", + "operationId": "getAttackChain", + "parameters": []map[string]interface{}{ + { + "name": "conversationId", + "in": "path", + "required": true, + "description": "对话ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/AttackChain", + }, + }, + }, + }, + "404": map[string]interface{}{ + "description": "对话不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/attack-chain/{conversationId}/regenerate": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"攻击链"}, + "summary": "重新生成攻击链", + "description": "重新生成指定对话的攻击链可视化数据", + "operationId": "regenerateAttackChain", + "parameters": []map[string]interface{}{ + { + "name": "conversationId", + "in": "path", + "required": true, + "description": "对话ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "重新生成成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/AttackChain", + }, + }, + }, + }, + "404": map[string]interface{}{ + "description": "对话不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/conversations/{id}/pinned": map[string]interface{}{ + "put": map[string]interface{}{ + "tags": []string{"对话管理"}, + "summary": "设置对话置顶", + "description": "设置或取消对话的置顶状态", + "operationId": "updateConversationPinned", + "parameters": []map[string]interface{}{ + { + "name": "id", + "in": "path", + "required": true, + "description": "对话ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "required": []string{"pinned"}, + "properties": map[string]interface{}{ + "pinned": map[string]interface{}{ + "type": "boolean", + "description": "是否置顶", + }, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "更新成功", + }, + "404": map[string]interface{}{ + "description": "对话不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/groups/{id}/pinned": map[string]interface{}{ + "put": map[string]interface{}{ + "tags": []string{"对话分组"}, + "summary": "设置分组置顶", + "description": "设置或取消分组的置顶状态", + "operationId": "updateGroupPinned", + "parameters": []map[string]interface{}{ + { + "name": "id", + "in": "path", + "required": true, + "description": "分组ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "required": []string{"pinned"}, + "properties": map[string]interface{}{ + "pinned": map[string]interface{}{ + "type": "boolean", + "description": "是否置顶", + }, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "更新成功", + }, + "404": map[string]interface{}{ + "description": "分组不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/groups/{id}/conversations/{conversationId}/pinned": map[string]interface{}{ + "put": map[string]interface{}{ + "tags": []string{"对话分组"}, + "summary": "设置分组中对话的置顶", + "description": "设置或取消分组中对话的置顶状态", + "operationId": "updateConversationPinnedInGroup", + "parameters": []map[string]interface{}{ + { + "name": "id", + "in": "path", + "required": true, + "description": "分组ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + { + "name": "conversationId", + "in": "path", + "required": true, + "description": "对话ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "required": []string{"pinned"}, + "properties": map[string]interface{}{ + "pinned": map[string]interface{}{ + "type": "boolean", + "description": "是否置顶", + }, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "更新成功", + }, + "404": map[string]interface{}{ + "description": "对话或分组不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/knowledge/categories": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"知识库"}, + "summary": "获取分类", + "description": "获取知识库的所有分类", + "operationId": "getKnowledgeCategories", + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "categories": map[string]interface{}{ + "type": "array", + "description": "分类列表", + "items": map[string]interface{}{ + "type": "string", + }, + }, + "enabled": map[string]interface{}{ + "type": "boolean", + "description": "知识库是否启用", + }, + }, + }, + }, + }, + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/knowledge/items": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"知识库"}, + "summary": "列出知识项", + "description": "获取知识库中的所有知识项", + "operationId": "getKnowledgeItems", + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "items": map[string]interface{}{ + "type": "array", + "description": "知识项列表", + }, + "enabled": map[string]interface{}{ + "type": "boolean", + "description": "知识库是否启用", + }, + }, + }, + }, + }, + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + "post": map[string]interface{}{ + "tags": []string{"知识库"}, + "summary": "创建知识项", + "description": "创建新的知识项", + "operationId": "createKnowledgeItem", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "description": "知识项数据", + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "创建成功", + }, + "400": map[string]interface{}{ + "description": "请求参数错误", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/knowledge/items/{id}": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"知识库"}, + "summary": "获取知识项", + "description": "获取指定知识项的详细信息", + "operationId": "getKnowledgeItem", + "parameters": []map[string]interface{}{ + { + "name": "id", + "in": "path", + "required": true, + "description": "知识项ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + }, + "404": map[string]interface{}{ + "description": "知识项不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + "put": map[string]interface{}{ + "tags": []string{"知识库"}, + "summary": "更新知识项", + "description": "更新指定知识项", + "operationId": "updateKnowledgeItem", + "parameters": []map[string]interface{}{ + { + "name": "id", + "in": "path", + "required": true, + "description": "知识项ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "description": "知识项数据", + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "更新成功", + }, + "404": map[string]interface{}{ + "description": "知识项不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + "delete": map[string]interface{}{ + "tags": []string{"知识库"}, + "summary": "删除知识项", + "description": "删除指定知识项", + "operationId": "deleteKnowledgeItem", + "parameters": []map[string]interface{}{ + { + "name": "id", + "in": "path", + "required": true, + "description": "知识项ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "删除成功", + }, + "404": map[string]interface{}{ + "description": "知识项不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/knowledge/index-status": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"知识库"}, + "summary": "获取索引状态", + "description": "获取知识库索引的构建状态", + "operationId": "getIndexStatus", + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "enabled": map[string]interface{}{ + "type": "boolean", + "description": "知识库是否启用", + }, + "total_items": map[string]interface{}{ + "type": "integer", + "description": "总知识项数", + }, + "indexed_items": map[string]interface{}{ + "type": "integer", + "description": "已索引知识项数", + }, + "progress_percent": map[string]interface{}{ + "type": "number", + "description": "索引进度百分比", + }, + "is_complete": map[string]interface{}{ + "type": "boolean", + "description": "索引是否完成", + }, + }, + }, + }, + }, + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/knowledge/index": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"知识库"}, + "summary": "构建索引", + "description": "构建知识库向量索引。默认仅处理尚无向量的知识项；mode=full 时全量重建。", + "operationId": "startKnowledgeIndex", + "parameters": []map[string]interface{}{ + { + "name": "mode", + "in": "query", + "required": false, + "description": "索引模式：missing（默认，补齐缺失向量）或 full（全量重建）", + "schema": map[string]interface{}{ + "type": "string", + "enum": []string{"missing", "full"}, + "default": "missing", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "索引任务已启动", + }, + "400": map[string]interface{}{ + "description": "无效的 mode 参数", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + "409": map[string]interface{}{ + "description": "已有索引任务正在进行", + }, + }, + }, + }, + "/api/knowledge/scan": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"知识库"}, + "summary": "扫描知识库", + "description": "扫描知识库目录，导入新的知识文件", + "operationId": "scanKnowledgeBase", + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "扫描任务已启动", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/knowledge/search": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"知识库"}, + "summary": "搜索知识库", + "description": "在知识库中搜索相关内容。基于向量检索，按查询与知识片段的语义相似度（余弦）返回最相关结果。\n**搜索说明**：\n- 语义相似度搜索：嵌入向量 + 余弦相似度，可配置相似度阈值与 TopK\n- 可按风险类型等元数据过滤（如：SQL注入、XSS、文件上传等）\n- 建议先调用 `/api/knowledge/categories` 获取可用的风险类型列表\n**使用示例**：\n```json\n{\n \"query\": \"SQL注入漏洞的检测方法\",\n \"riskType\": \"SQL注入\",\n \"topK\": 5,\n \"threshold\": 0.7\n}\n```", + "operationId": "searchKnowledge", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "required": []string{"query"}, + "properties": map[string]interface{}{ + "query": map[string]interface{}{ + "type": "string", + "description": "搜索查询内容，描述你想要了解的安全知识主题（必需）", + "example": "SQL注入漏洞的检测方法", + }, + "riskType": map[string]interface{}{ + "type": "string", + "description": "可选：指定风险类型（如：SQL注入、XSS、文件上传等）。建议先调用 `/api/knowledge/categories` 获取可用的风险类型列表，然后使用正确的风险类型进行精确搜索，这样可以大幅减少检索时间。如果不指定则搜索所有类型。", + "example": "SQL注入", + }, + "topK": map[string]interface{}{ + "type": "integer", + "description": "可选：返回Top-K结果数量，默认5", + "default": 5, + "minimum": 1, + "maximum": 50, + "example": 5, + }, + "threshold": map[string]interface{}{ + "type": "number", + "format": "float", + "description": "可选：相似度阈值（0-1之间），默认0.7。只有相似度大于等于此值的结果才会返回", + "default": 0.7, + "minimum": 0, + "maximum": 1, + "example": 0.7, + }, + }, + }, + "examples": map[string]interface{}{ + "basic": map[string]interface{}{ + "summary": "基础搜索", + "description": "最简单的搜索，只提供查询内容", + "value": map[string]interface{}{ + "query": "SQL注入漏洞的检测方法", + }, + }, + "withRiskType": map[string]interface{}{ + "summary": "按风险类型搜索", + "description": "指定风险类型进行精确搜索", + "value": map[string]interface{}{ + "query": "SQL注入漏洞的检测方法", + "riskType": "SQL注入", + "topK": 5, + "threshold": 0.7, + }, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "搜索成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "results": map[string]interface{}{ + "type": "array", + "description": "搜索结果列表，每个结果包含：item（知识项信息）、chunks（匹配的知识片段）、score（相似度分数）", + "items": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "item": map[string]interface{}{ + "type": "object", + "description": "知识项信息", + }, + "chunks": map[string]interface{}{ + "type": "array", + "description": "匹配的知识片段列表", + }, + "score": map[string]interface{}{ + "type": "number", + "description": "相似度分数（0-1之间）", + }, + }, + }, + }, + "enabled": map[string]interface{}{ + "type": "boolean", + "description": "知识库是否启用", + }, + }, + }, + "example": map[string]interface{}{ + "results": []map[string]interface{}{ + { + "item": map[string]interface{}{ + "id": "item-1", + "title": "SQL注入漏洞检测", + "category": "SQL注入", + }, + "chunks": []map[string]interface{}{ + { + "text": "SQL注入漏洞的检测方法包括...", + }, + }, + "score": 0.85, + }, + }, + "enabled": true, + }, + }, + }, + }, + "400": map[string]interface{}{ + "description": "请求参数错误（如query为空）", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/Error", + }, + "example": map[string]interface{}{ + "error": "查询不能为空", + }, + }, + }, + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + "500": map[string]interface{}{ + "description": "服务器内部错误（如知识库未启用或检索失败）", + }, + }, + }, + }, + "/api/knowledge/retrieval-logs": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"知识库"}, + "summary": "获取检索日志", + "description": "获取知识库检索日志", + "operationId": "getRetrievalLogs", + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "logs": map[string]interface{}{ + "type": "array", + "description": "检索日志列表", + }, + "enabled": map[string]interface{}{ + "type": "boolean", + "description": "知识库是否启用", + }, + }, + }, + }, + }, + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + "/api/knowledge/retrieval-logs/{id}": map[string]interface{}{ + "delete": map[string]interface{}{ + "tags": []string{"知识库"}, + "summary": "删除检索日志", + "description": "删除指定的检索日志", + "operationId": "deleteRetrievalLog", + "parameters": []map[string]interface{}{ + { + "name": "id", + "in": "path", + "required": true, + "description": "日志ID", + "schema": map[string]interface{}{ + "type": "string", + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "删除成功", + }, + "404": map[string]interface{}{ + "description": "日志不存在", + }, + "401": map[string]interface{}{ + "description": "未授权", + }, + }, + }, + }, + // ==================== 对话交互 - 缺失端点 ==================== + "/api/conversations/{id}/delete-turn": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"对话交互"}, + "summary": "删除对话轮次", + "description": "删除指定消息所在的对话轮次（从该轮 user 消息到下一轮 user 消息之前的所有消息），并清空 last_react 状态。", + "operationId": "deleteConversationTurn", + "parameters": []map[string]interface{}{ + { + "name": "id", + "in": "path", + "required": true, + "description": "对话ID", + "schema": map[string]interface{}{"type": "string"}, + }, + }, + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "required": []string{"messageId"}, + "properties": map[string]interface{}{ + "messageId": map[string]interface{}{ + "type": "string", + "description": "锚点消息ID，标识要删除的轮次", + }, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "删除成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "deletedMessageIds": map[string]interface{}{ + "type": "array", + "items": map[string]interface{}{"type": "string"}, + "description": "被删除的消息ID列表", + }, + "message": map[string]interface{}{ + "type": "string", + "example": "ok", + }, + }, + }, + }, + }, + }, + "400": map[string]interface{}{"description": "参数错误或删除失败"}, + "401": map[string]interface{}{"description": "未授权"}, + "404": map[string]interface{}{"description": "对话不存在"}, + }, + }, + }, + "/api/messages/{id}/process-details": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"对话交互"}, + "summary": "获取消息过程详情", + "description": "按需分页加载指定消息的执行过程详情，包括工具调用、思考过程等事件。默认返回 50 条；导出或旧集成需要全量时可显式传 full=1。", + "operationId": "getMessageProcessDetails", + "parameters": []map[string]interface{}{ + { + "name": "id", + "in": "path", + "required": true, + "description": "消息ID", + "schema": map[string]interface{}{"type": "string"}, + }, + { + "name": "summary", + "in": "query", + "required": false, + "description": "仅返回过程详情摘要（total / iterationCount / maxIteration）", + "schema": map[string]interface{}{"type": "boolean"}, + }, + { + "name": "limit", + "in": "query", + "required": false, + "description": "分页大小，默认 50，最大 500", + "schema": map[string]interface{}{"type": "integer", "default": 50, "maximum": 500}, + }, + { + "name": "offset", + "in": "query", + "required": false, + "description": "分页偏移量，默认 0", + "schema": map[string]interface{}{"type": "integer", "default": 0}, + }, + { + "name": "full", + "in": "query", + "required": false, + "description": "显式返回全量过程详情；仅建议导出/兼容旧集成使用", + "schema": map[string]interface{}{"type": "boolean"}, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "processDetails": map[string]interface{}{ + "type": "array", + "items": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "id": map[string]interface{}{"type": "string", "description": "详情记录ID"}, + "messageId": map[string]interface{}{"type": "string", "description": "所属消息ID"}, + "conversationId": map[string]interface{}{"type": "string", "description": "所属对话ID"}, + "eventType": map[string]interface{}{"type": "string", "description": "事件类型（如tool_call, thinking等）"}, + "message": map[string]interface{}{"type": "string", "description": "事件消息"}, + "data": map[string]interface{}{"description": "事件附加数据（JSON对象）"}, + "createdAt": map[string]interface{}{"type": "string", "format": "date-time", "description": "创建时间"}, + }, + }, + }, + "total": map[string]interface{}{"type": "integer", "description": "过程详情总数"}, + "offset": map[string]interface{}{"type": "integer", "description": "当前分页偏移量"}, + "limit": map[string]interface{}{"type": "integer", "description": "当前分页大小"}, + "hasMore": map[string]interface{}{"type": "boolean", "description": "是否还有更多过程详情"}, + "summary": map[string]interface{}{ + "type": "object", + "description": "summary=1 时返回的摘要", + }, + }, + }, + }, + }, + }, + "400": map[string]interface{}{"description": "参数错误"}, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + }, + + // ==================== 批量任务 - 缺失端点 ==================== + "/api/batch-tasks/{queueId}/rerun": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"批量任务"}, + "summary": "重跑批量任务队列", + "description": "重置已完成或已取消的批量任务队列，重新开始执行所有任务。", + "operationId": "rerunBatchQueue", + "parameters": []map[string]interface{}{ + { + "name": "queueId", + "in": "path", + "required": true, + "description": "队列ID", + "schema": map[string]interface{}{"type": "string"}, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "重跑成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "message": map[string]interface{}{"type": "string", "example": "批量任务已重新开始执行"}, + "queueId": map[string]interface{}{"type": "string", "description": "队列ID"}, + }, + }, + }, + }, + }, + "400": map[string]interface{}{"description": "仅已完成或已取消的队列可以重跑"}, + "401": map[string]interface{}{"description": "未授权"}, + "404": map[string]interface{}{"description": "队列不存在"}, + }, + }, + }, + "/api/batch-tasks/{queueId}/metadata": map[string]interface{}{ + "put": map[string]interface{}{ + "tags": []string{"批量任务"}, + "summary": "修改队列元数据", + "description": "修改批量任务队列的标题、角色和代理模式。", + "operationId": "updateBatchQueueMetadata", + "parameters": []map[string]interface{}{ + { + "name": "queueId", + "in": "path", + "required": true, + "description": "队列ID", + "schema": map[string]interface{}{"type": "string"}, + }, + }, + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "title": map[string]interface{}{"type": "string", "description": "队列标题"}, + "role": map[string]interface{}{"type": "string", "description": "使用的角色名称"}, + "agentMode": map[string]interface{}{"type": "string", "description": "代理模式", "enum": []string{"eino_single", "deep", "plan_execute", "supervisor"}}, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "更新成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "queue": map[string]interface{}{"$ref": "#/components/schemas/BatchQueue"}, + }, + }, + }, + }, + }, + "400": map[string]interface{}{"description": "参数错误"}, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + }, + "/api/batch-tasks/{queueId}/schedule": map[string]interface{}{ + "put": map[string]interface{}{ + "tags": []string{"批量任务"}, + "summary": "修改队列调度配置", + "description": "修改批量任务队列的调度模式和Cron表达式。队列运行中无法修改。", + "operationId": "updateBatchQueueSchedule", + "parameters": []map[string]interface{}{ + { + "name": "queueId", + "in": "path", + "required": true, + "description": "队列ID", + "schema": map[string]interface{}{"type": "string"}, + }, + }, + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "scheduleMode": map[string]interface{}{"type": "string", "description": "调度模式", "enum": []string{"manual", "cron"}}, + "cronExpr": map[string]interface{}{"type": "string", "description": "Cron表达式（scheduleMode为cron时必填）", "example": "0 2 * * *"}, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "更新成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "queue": map[string]interface{}{"$ref": "#/components/schemas/BatchQueue"}, + }, + }, + }, + }, + }, + "400": map[string]interface{}{"description": "参数错误或队列正在运行中"}, + "401": map[string]interface{}{"description": "未授权"}, + "404": map[string]interface{}{"description": "队列不存在"}, + }, + }, + }, + "/api/batch-tasks/{queueId}/schedule-enabled": map[string]interface{}{ + "put": map[string]interface{}{ + "tags": []string{"批量任务"}, + "summary": "开关Cron自动调度", + "description": "开启或关闭批量任务队列的Cron自动调度功能，手工执行不受影响。", + "operationId": "setBatchQueueScheduleEnabled", + "parameters": []map[string]interface{}{ + { + "name": "queueId", + "in": "path", + "required": true, + "description": "队列ID", + "schema": map[string]interface{}{"type": "string"}, + }, + }, + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "required": []string{"scheduleEnabled"}, + "properties": map[string]interface{}{ + "scheduleEnabled": map[string]interface{}{"type": "boolean", "description": "是否启用自动调度"}, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "设置成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "queue": map[string]interface{}{"$ref": "#/components/schemas/BatchQueue"}, + }, + }, + }, + }, + }, + "401": map[string]interface{}{"description": "未授权"}, + "404": map[string]interface{}{"description": "队列不存在"}, + }, + }, + }, + + // ==================== 对话分组 - 缺失端点 ==================== + "/api/groups/mappings": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"对话分组"}, + "summary": "获取所有分组映射", + "description": "获取所有对话与分组之间的映射关系列表。", + "operationId": "getAllGroupMappings", + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "array", + "items": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "conversation_id": map[string]interface{}{"type": "string", "description": "对话ID"}, + "group_id": map[string]interface{}{"type": "string", "description": "分组ID"}, + "pinned": map[string]interface{}{"type": "boolean", "description": "是否置顶"}, + }, + }, + }, + }, + }, + }, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + }, + + // ==================== FOFA信息收集 ==================== + "/api/fofa/search": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"FOFA信息收集"}, + "summary": "FOFA搜索", + "description": "通过后端代理执行FOFA搜索查询，返回资产信息。", + "operationId": "fofaSearch", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "required": []string{"query"}, + "properties": map[string]interface{}{ + "query": map[string]interface{}{"type": "string", "description": "FOFA查询语法", "example": "domain=\"example.com\""}, + "size": map[string]interface{}{"type": "integer", "description": "返回数量（默认100，最大10000）", "default": 100}, + "page": map[string]interface{}{"type": "integer", "description": "页码（默认1）", "default": 1}, + "fields": map[string]interface{}{"type": "string", "description": "返回字段，逗号分隔", "example": "host,ip,port,title"}, + "full": map[string]interface{}{"type": "boolean", "description": "是否查询全部数据", "default": false}, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "搜索成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "query": map[string]interface{}{"type": "string", "description": "实际执行的查询"}, + "size": map[string]interface{}{"type": "integer"}, + "page": map[string]interface{}{"type": "integer"}, + "total": map[string]interface{}{"type": "integer", "description": "总匹配数"}, + "fields": map[string]interface{}{"type": "array", "items": map[string]interface{}{"type": "string"}}, + "results_count": map[string]interface{}{"type": "integer"}, + "results": map[string]interface{}{"type": "array", "items": map[string]interface{}{"type": "object"}, "description": "搜索结果列表"}, + }, + }, + }, + }, + }, + "400": map[string]interface{}{"description": "参数错误"}, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + }, + "/api/fofa/parse": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"FOFA信息收集"}, + "summary": "自然语言解析为FOFA语法", + "description": "使用AI将自然语言描述解析为FOFA查询语法，需人工确认后再执行查询。", + "operationId": "fofaParse", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "required": []string{"text"}, + "properties": map[string]interface{}{ + "text": map[string]interface{}{"type": "string", "description": "自然语言描述", "example": "查找使用WordPress的网站"}, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "解析成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "query": map[string]interface{}{"type": "string", "description": "生成的FOFA查询语法"}, + "explanation": map[string]interface{}{"type": "string", "description": "语法解释"}, + "warnings": map[string]interface{}{"type": "array", "items": map[string]interface{}{"type": "string"}, "description": "潜在风险或歧义提示"}, + }, + }, + }, + }, + }, + "400": map[string]interface{}{"description": "参数错误"}, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + }, + + // ==================== 配置管理 - 缺失端点 ==================== + "/api/config/test-vision": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"配置管理"}, + "summary": "测试视觉模型连接", + "description": "测试 Vision 模型 API 是否可用。vision.api_key/base_url 留空时可传 openai 段作回退。", + "operationId": "testVision", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "required": []string{"vision"}, + "properties": map[string]interface{}{ + "vision": map[string]interface{}{"$ref": "#/components/schemas/VisionConfig"}, + "openai": map[string]interface{}{ + "type": "object", + "description": "主 LLM 配置（vision 字段留空时用于 API Key/Base URL 回退）", + }, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "测试结果", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "success": map[string]interface{}{"type": "boolean"}, + "error": map[string]interface{}{"type": "string"}, + "model": map[string]interface{}{"type": "string"}, + "latency_ms": map[string]interface{}{"type": "number"}, + }, + }, + }, + }, + }, + "400": map[string]interface{}{"description": "参数错误"}, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + }, + "/api/config/test-openai": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"配置管理"}, + "summary": "测试OpenAI API连接", + "description": "测试指定的OpenAI/Claude API配置是否可用，发送一个最小请求验证连通性。", + "operationId": "testOpenAI", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "required": []string{"api_key", "model"}, + "properties": map[string]interface{}{ + "provider": map[string]interface{}{"type": "string", "description": "LLM提供商（openai/claude）", "example": "openai"}, + "base_url": map[string]interface{}{"type": "string", "description": "API基地址（可选，默认根据provider自动选择）"}, + "api_key": map[string]interface{}{"type": "string", "description": "API密钥"}, + "model": map[string]interface{}{"type": "string", "description": "模型名称", "example": "gpt-4"}, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "测试结果", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "success": map[string]interface{}{"type": "boolean", "description": "是否连接成功"}, + "error": map[string]interface{}{"type": "string", "description": "失败原因（success=false时）"}, + "model": map[string]interface{}{"type": "string", "description": "实际使用的模型（success=true时）"}, + "latency_ms": map[string]interface{}{"type": "number", "description": "延迟毫秒数（success=true时）"}, + }, + }, + }, + }, + }, + "400": map[string]interface{}{"description": "参数错误"}, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + }, + "/api/config/list-models": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"配置管理"}, + "summary": "获取模型列表", + "description": "代理调用 OpenAI 兼容 GET /models，返回可用模型 id 列表。Claude 不支持。", + "operationId": "listModels", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "required": []string{"api_key"}, + "properties": map[string]interface{}{ + "provider": map[string]interface{}{"type": "string", "description": "LLM提供商（openai/claude）", "example": "openai"}, + "base_url": map[string]interface{}{"type": "string", "description": "API基地址（可选）"}, + "api_key": map[string]interface{}{"type": "string", "description": "API密钥"}, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取结果", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "success": map[string]interface{}{"type": "boolean"}, + "supported": map[string]interface{}{"type": "boolean"}, + "error": map[string]interface{}{"type": "string"}, + "models": map[string]interface{}{"type": "array", "items": map[string]interface{}{"type": "string"}}, + "count": map[string]interface{}{"type": "integer"}, + }, + }, + }, + }, + }, + "400": map[string]interface{}{"description": "参数错误"}, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + }, + + // ==================== 终端 ==================== + "/api/terminal/run": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"终端"}, + "summary": "执行终端命令", + "description": "在服务器上执行Shell命令并返回结果。", + "operationId": "terminalRun", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "required": []string{"command"}, + "properties": map[string]interface{}{ + "command": map[string]interface{}{"type": "string", "description": "要执行的命令"}, + "shell": map[string]interface{}{"type": "string", "description": "Shell类型（默认sh/cmd）"}, + "cwd": map[string]interface{}{"type": "string", "description": "工作目录（可选）"}, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "执行完成", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "stdout": map[string]interface{}{"type": "string", "description": "标准输出"}, + "stderr": map[string]interface{}{"type": "string", "description": "标准错误"}, + "exit_code": map[string]interface{}{"type": "integer", "description": "退出码"}, + "error": map[string]interface{}{"type": "string", "description": "执行错误（可选）"}, + }, + }, + }, + }, + }, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + }, + "/api/terminal/run/stream": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"终端"}, + "summary": "流式执行终端命令", + "description": "以SSE流式方式执行Shell命令，实时返回输出。每个事件包含 JSON: {\"t\": \"out\"|\"err\"|\"exit\", \"d\": \"数据\", \"c\": 退出码}", + "operationId": "terminalRunStream", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "required": []string{"command"}, + "properties": map[string]interface{}{ + "command": map[string]interface{}{"type": "string", "description": "要执行的命令"}, + "shell": map[string]interface{}{"type": "string", "description": "Shell类型（默认sh/cmd）"}, + "cwd": map[string]interface{}{"type": "string", "description": "工作目录（可选）"}, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "SSE事件流", + "content": map[string]interface{}{ + "text/event-stream": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "string", + "description": "Server-Sent Events流，每个事件为JSON: {\"t\":\"out|err|exit\",\"d\":\"data\",\"c\":exitCode}", + }, + }, + }, + }, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + }, + "/api/terminal/ws": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"终端"}, + "summary": "WebSocket终端", + "description": "通过WebSocket建立交互式终端连接，支持PTY。客户端发送文本/二进制数据作为命令输入，也可发送JSON: {\"type\":\"resize\",\"cols\":80,\"rows\":24} 调整终端大小。服务端返回二进制PTY输出。", + "operationId": "terminalWS", + "responses": map[string]interface{}{ + "101": map[string]interface{}{"description": "WebSocket连接已建立"}, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + }, + + // ==================== WebShell管理 ==================== + "/api/webshell/connections": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"WebShell管理"}, + "summary": "列出WebShell连接", + "description": "获取所有已保存的WebShell连接配置列表。", + "operationId": "listWebshellConnections", + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "array", + "items": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "id": map[string]interface{}{"type": "string", "description": "连接ID"}, + "url": map[string]interface{}{"type": "string", "description": "WebShell URL"}, + "password": map[string]interface{}{"type": "string", "description": "连接密码"}, + "type": map[string]interface{}{"type": "string", "description": "Shell类型", "enum": []string{"php", "asp", "aspx", "jsp", "custom"}}, + "method": map[string]interface{}{"type": "string", "description": "请求方法", "enum": []string{"get", "post"}}, + "cmd_param": map[string]interface{}{"type": "string", "description": "命令参数名"}, + "remark": map[string]interface{}{"type": "string", "description": "备注"}, + "created_at": map[string]interface{}{"type": "string", "format": "date-time"}, + }, + }, + }, + }, + }, + }, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + "post": map[string]interface{}{ + "tags": []string{"WebShell管理"}, + "summary": "创建WebShell连接", + "description": "保存一个新的WebShell连接配置。", + "operationId": "createWebshellConnection", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "required": []string{"url"}, + "properties": map[string]interface{}{ + "url": map[string]interface{}{"type": "string", "description": "WebShell URL"}, + "password": map[string]interface{}{"type": "string", "description": "连接密码"}, + "type": map[string]interface{}{"type": "string", "description": "Shell类型", "enum": []string{"php", "asp", "aspx", "jsp", "custom"}}, + "method": map[string]interface{}{"type": "string", "description": "请求方法", "enum": []string{"get", "post"}}, + "cmd_param": map[string]interface{}{"type": "string", "description": "命令参数名"}, + "remark": map[string]interface{}{"type": "string", "description": "备注"}, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{"description": "创建成功"}, + "400": map[string]interface{}{"description": "参数错误"}, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + }, + "/api/webshell/connections/{id}": map[string]interface{}{ + "put": map[string]interface{}{ + "tags": []string{"WebShell管理"}, + "summary": "更新WebShell连接", + "description": "更新已有的WebShell连接配置。", + "operationId": "updateWebshellConnection", + "parameters": []map[string]interface{}{ + {"name": "id", "in": "path", "required": true, "description": "连接ID", "schema": map[string]interface{}{"type": "string"}}, + }, + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "url": map[string]interface{}{"type": "string"}, + "password": map[string]interface{}{"type": "string"}, + "type": map[string]interface{}{"type": "string", "enum": []string{"php", "asp", "aspx", "jsp", "custom"}}, + "method": map[string]interface{}{"type": "string", "enum": []string{"get", "post"}}, + "cmd_param": map[string]interface{}{"type": "string"}, + "remark": map[string]interface{}{"type": "string"}, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{"description": "更新成功"}, + "401": map[string]interface{}{"description": "未授权"}, + "404": map[string]interface{}{"description": "连接不存在"}, + }, + }, + "delete": map[string]interface{}{ + "tags": []string{"WebShell管理"}, + "summary": "删除WebShell连接", + "description": "删除指定的WebShell连接配置。", + "operationId": "deleteWebshellConnection", + "parameters": []map[string]interface{}{ + {"name": "id", "in": "path", "required": true, "description": "连接ID", "schema": map[string]interface{}{"type": "string"}}, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{"description": "删除成功"}, + "401": map[string]interface{}{"description": "未授权"}, + "404": map[string]interface{}{"description": "连接不存在"}, + }, + }, + }, + "/api/webshell/connections/{id}/state": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"WebShell管理"}, + "summary": "获取连接状态", + "description": "获取WebShell连接的保存状态数据。", + "operationId": "getWebshellConnectionState", + "parameters": []map[string]interface{}{ + {"name": "id", "in": "path", "required": true, "description": "连接ID", "schema": map[string]interface{}{"type": "string"}}, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "state": map[string]interface{}{"type": "object", "description": "状态数据（任意JSON）"}, + }, + }, + }, + }, + }, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + "put": map[string]interface{}{ + "tags": []string{"WebShell管理"}, + "summary": "保存连接状态", + "description": "保存WebShell连接的状态数据。", + "operationId": "saveWebshellConnectionState", + "parameters": []map[string]interface{}{ + {"name": "id", "in": "path", "required": true, "description": "连接ID", "schema": map[string]interface{}{"type": "string"}}, + }, + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "state": map[string]interface{}{"type": "object", "description": "状态数据（任意JSON）"}, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{"description": "保存成功"}, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + }, + "/api/webshell/connections/{id}/ai-history": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"WebShell管理"}, + "summary": "获取AI对话历史", + "description": "获取指定WebShell连接的AI辅助对话历史消息。", + "operationId": "getWebshellAIHistory", + "parameters": []map[string]interface{}{ + {"name": "id", "in": "path", "required": true, "description": "连接ID", "schema": map[string]interface{}{"type": "string"}}, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "conversationId": map[string]interface{}{"type": "string"}, + "messages": map[string]interface{}{ + "type": "array", + "items": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "id": map[string]interface{}{"type": "string"}, + "role": map[string]interface{}{"type": "string"}, + "content": map[string]interface{}{"type": "string"}, + "createdAt": map[string]interface{}{"type": "string", "format": "date-time"}, + }, + }, + }, + }, + }, + }, + }, + }, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + }, + "/api/webshell/connections/{id}/ai-conversations": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"WebShell管理"}, + "summary": "列出AI对话", + "description": "获取指定WebShell连接的所有AI辅助对话列表。", + "operationId": "listWebshellAIConversations", + "parameters": []map[string]interface{}{ + {"name": "id", "in": "path", "required": true, "description": "连接ID", "schema": map[string]interface{}{"type": "string"}}, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "array", + "items": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "id": map[string]interface{}{"type": "string"}, + "title": map[string]interface{}{"type": "string"}, + "createdAt": map[string]interface{}{"type": "string", "format": "date-time"}, + }, + }, + }, + }, + }, + }, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + }, + "/api/webshell/exec": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"WebShell管理"}, + "summary": "执行WebShell命令", + "description": "通过指定的WebShell连接执行远程命令。", + "operationId": "webshellExec", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "required": []string{"url", "command"}, + "properties": map[string]interface{}{ + "url": map[string]interface{}{"type": "string", "description": "WebShell URL"}, + "password": map[string]interface{}{"type": "string"}, + "type": map[string]interface{}{"type": "string", "enum": []string{"php", "asp", "aspx", "jsp", "custom"}}, + "method": map[string]interface{}{"type": "string", "enum": []string{"get", "post"}}, + "cmd_param": map[string]interface{}{"type": "string"}, + "command": map[string]interface{}{"type": "string", "description": "要执行的命令"}, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "执行结果", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "ok": map[string]interface{}{"type": "boolean"}, + "output": map[string]interface{}{"type": "string", "description": "命令输出"}, + "error": map[string]interface{}{"type": "string", "description": "错误信息"}, + "http_code": map[string]interface{}{"type": "integer", "description": "HTTP响应码"}, + }, + }, + }, + }, + }, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + }, + "/api/webshell/file": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"WebShell管理"}, + "summary": "WebShell文件操作", + "description": "通过WebShell执行远程文件操作（列目录、读写文件、创建目录、重命名、删除、上传等）。", + "operationId": "webshellFileOp", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "required": []string{"url", "action", "path"}, + "properties": map[string]interface{}{ + "url": map[string]interface{}{"type": "string", "description": "WebShell URL"}, + "password": map[string]interface{}{"type": "string"}, + "type": map[string]interface{}{"type": "string", "enum": []string{"php", "asp", "aspx", "jsp", "custom"}}, + "method": map[string]interface{}{"type": "string", "enum": []string{"get", "post"}}, + "cmd_param": map[string]interface{}{"type": "string"}, + "action": map[string]interface{}{"type": "string", "description": "操作类型", "enum": []string{"list", "read", "delete", "write", "mkdir", "rename", "upload", "upload_chunk"}}, + "path": map[string]interface{}{"type": "string", "description": "目标文件/目录路径"}, + "target_path": map[string]interface{}{"type": "string", "description": "目标路径（rename时使用）"}, + "content": map[string]interface{}{"type": "string", "description": "文件内容（write/upload时使用）"}, + "chunk_index": map[string]interface{}{"type": "integer", "description": "分块索引（upload_chunk时使用）"}, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "操作结果", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "ok": map[string]interface{}{"type": "boolean"}, + "output": map[string]interface{}{"type": "string"}, + "error": map[string]interface{}{"type": "string"}, + }, + }, + }, + }, + }, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + }, + + // ==================== 对话附件 ==================== + "/api/chat-uploads": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"对话附件"}, + "summary": "列出附件", + "description": "获取对话文件列表，包含手动上传附件、工具输出和会话产物，可按会话、项目、来源、文件名搜索和分页过滤。", + "operationId": "listChatUploads", + "parameters": []map[string]interface{}{ + {"name": "conversation", "in": "query", "required": false, "description": "按对话ID过滤", "schema": map[string]interface{}{"type": "string"}}, + {"name": "project", "in": "query", "required": false, "description": "按项目ID过滤", "schema": map[string]interface{}{"type": "string"}}, + {"name": "source", "in": "query", "required": false, "description": "按来源过滤：upload/reduction/workspace/conversation_artifact/all", "schema": map[string]interface{}{"type": "string", "enum": []string{"all", "upload", "reduction", "workspace", "conversation_artifact"}}}, + {"name": "search", "in": "query", "required": false, "description": "按文件名或子路径搜索", "schema": map[string]interface{}{"type": "string"}}, + {"name": "page", "in": "query", "required": false, "description": "页码，从1开始", "schema": map[string]interface{}{"type": "integer", "default": 1}}, + {"name": "pageSize", "in": "query", "required": false, "description": "每页数量，传 all 返回全部", "schema": map[string]interface{}{"oneOf": []map[string]interface{}{{"type": "integer"}, {"type": "string", "enum": []string{"all"}}}}}, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "files": map[string]interface{}{ + "type": "array", + "items": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "relativePath": map[string]interface{}{"type": "string"}, + "absolutePath": map[string]interface{}{"type": "string"}, + "name": map[string]interface{}{"type": "string"}, + "size": map[string]interface{}{"type": "integer"}, + "modifiedUnix": map[string]interface{}{"type": "integer"}, + "date": map[string]interface{}{"type": "string"}, + "conversationId": map[string]interface{}{"type": "string"}, + "conversationTitle": map[string]interface{}{"type": "string"}, + "projectId": map[string]interface{}{"type": "string"}, + "projectName": map[string]interface{}{"type": "string"}, + "subPath": map[string]interface{}{"type": "string"}, + "source": map[string]interface{}{"type": "string", "description": "upload/reduction/workspace/conversation_artifact"}, + }, + }, + }, + "folders": map[string]interface{}{"type": "array", "items": map[string]interface{}{"type": "string"}}, + "total": map[string]interface{}{"type": "integer"}, + "page": map[string]interface{}{"type": "integer"}, + "pageSize": map[string]interface{}{ + "type": "integer", + }, + }, + }, + }, + }, + }, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + "post": map[string]interface{}{ + "tags": []string{"对话附件"}, + "summary": "上传附件", + "description": "上传文件到对话附件目录（multipart/form-data）。", + "operationId": "uploadChatFile", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "multipart/form-data": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "required": []string{"file"}, + "properties": map[string]interface{}{ + "file": map[string]interface{}{"type": "string", "format": "binary", "description": "上传的文件"}, + "conversationId": map[string]interface{}{"type": "string", "description": "关联的对话ID（可选）"}, + "relativeDir": map[string]interface{}{"type": "string", "description": "目标目录相对路径（可选）"}, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "上传成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "ok": map[string]interface{}{"type": "boolean"}, + "relativePath": map[string]interface{}{"type": "string"}, + "absolutePath": map[string]interface{}{"type": "string"}, + "name": map[string]interface{}{"type": "string"}, + }, + }, + }, + }, + }, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + "delete": map[string]interface{}{ + "tags": []string{"对话附件"}, + "summary": "删除附件", + "description": "删除指定的对话附件文件。", + "operationId": "deleteChatUpload", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "required": []string{"path"}, + "properties": map[string]interface{}{ + "path": map[string]interface{}{"type": "string", "description": "文件相对路径"}, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{"description": "删除成功"}, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + }, + "/api/chat-uploads/export": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"对话附件"}, + "summary": "导出附件", + "description": "按当前过滤条件导出对话文件 ZIP，包含 manifest.json。", + "operationId": "exportChatUploads", + "parameters": []map[string]interface{}{ + {"name": "conversation", "in": "query", "required": false, "description": "按对话ID过滤", "schema": map[string]interface{}{"type": "string"}}, + {"name": "project", "in": "query", "required": false, "description": "按项目ID过滤", "schema": map[string]interface{}{"type": "string"}}, + {"name": "source", "in": "query", "required": false, "description": "按来源过滤：upload/reduction/workspace/conversation_artifact/all", "schema": map[string]interface{}{"type": "string", "enum": []string{"all", "upload", "reduction", "workspace", "conversation_artifact"}}}, + {"name": "search", "in": "query", "required": false, "description": "按文件名或子路径搜索", "schema": map[string]interface{}{"type": "string"}}, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "ZIP文件下载", + "content": map[string]interface{}{ + "application/zip": map[string]interface{}{ + "schema": map[string]interface{}{"type": "string", "format": "binary"}, + }, + }, + }, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + }, + "/api/chat-uploads/download": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"对话附件"}, + "summary": "下载附件", + "description": "下载指定的对话附件文件。", + "operationId": "downloadChatUpload", + "parameters": []map[string]interface{}{ + {"name": "path", "in": "query", "required": true, "description": "文件相对路径", "schema": map[string]interface{}{"type": "string"}}, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "文件下载", + "content": map[string]interface{}{ + "application/octet-stream": map[string]interface{}{ + "schema": map[string]interface{}{"type": "string", "format": "binary"}, + }, + }, + }, + "401": map[string]interface{}{"description": "未授权"}, + "404": map[string]interface{}{"description": "文件不存在"}, + }, + }, + }, + "/api/chat-uploads/path": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"对话附件"}, + "summary": "解析附件路径", + "description": "将文件管理中的相对路径或内部虚拟路径解析为服务器绝对路径，用于复制文件/目录路径。", + "operationId": "resolveChatUploadPath", + "parameters": []map[string]interface{}{ + {"name": "path", "in": "query", "required": true, "description": "相对路径或虚拟路径（如 __workspace__/projects//csv）", "schema": map[string]interface{}{"type": "string"}}, + {"name": "kind", "in": "query", "required": false, "description": "路径类型：file/directory，默认 file", "schema": map[string]interface{}{"type": "string", "enum": []string{"file", "directory"}}}, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "解析成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "absolutePath": map[string]interface{}{"type": "string"}, + "isDir": map[string]interface{}{"type": "boolean"}, + }, + }, + }, + }, + }, + "401": map[string]interface{}{"description": "未授权"}, + "403": map[string]interface{}{"description": "无权访问"}, + "404": map[string]interface{}{"description": "路径不存在"}, + }, + }, + }, + "/api/chat-uploads/content": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"对话附件"}, + "summary": "获取附件文本内容", + "description": "读取并返回文本文件的内容。", + "operationId": "getChatUploadContent", + "parameters": []map[string]interface{}{ + {"name": "path", "in": "query", "required": true, "description": "文件相对路径", "schema": map[string]interface{}{"type": "string"}}, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "content": map[string]interface{}{"type": "string", "description": "文件文本内容"}, + }, + }, + }, + }, + }, + "401": map[string]interface{}{"description": "未授权"}, + "404": map[string]interface{}{"description": "文件不存在"}, + }, + }, + "put": map[string]interface{}{ + "tags": []string{"对话附件"}, + "summary": "写入附件文本内容", + "description": "写入或覆盖文本文件的内容。", + "operationId": "putChatUploadContent", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "required": []string{"path", "content"}, + "properties": map[string]interface{}{ + "path": map[string]interface{}{"type": "string", "description": "文件相对路径"}, + "content": map[string]interface{}{"type": "string", "description": "文件文本内容"}, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{"description": "写入成功"}, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + }, + "/api/chat-uploads/mkdir": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"对话附件"}, + "summary": "创建附件目录", + "description": "在对话附件目录下创建子目录。", + "operationId": "mkdirChatUpload", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "required": []string{"name"}, + "properties": map[string]interface{}{ + "parent": map[string]interface{}{"type": "string", "description": "父目录相对路径"}, + "name": map[string]interface{}{"type": "string", "description": "目录名称"}, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "创建成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "ok": map[string]interface{}{"type": "boolean"}, + "relativePath": map[string]interface{}{"type": "string"}, + }, + }, + }, + }, + }, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + }, + "/api/chat-uploads/rename": map[string]interface{}{ + "put": map[string]interface{}{ + "tags": []string{"对话附件"}, + "summary": "重命名附件", + "description": "重命名对话附件文件或目录。", + "operationId": "renameChatUpload", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "required": []string{"path", "newName"}, + "properties": map[string]interface{}{ + "path": map[string]interface{}{"type": "string", "description": "当前文件相对路径"}, + "newName": map[string]interface{}{"type": "string", "description": "新名称"}, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "重命名成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "ok": map[string]interface{}{"type": "boolean"}, + "relativePath": map[string]interface{}{"type": "string"}, + }, + }, + }, + }, + }, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + }, + + // ==================== 机器人集成 ==================== + "/api/robot/wecom": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"机器人集成"}, + "summary": "企业微信回调验证", + "description": "企业微信服务器URL验证回调（用于配置消息接收地址时的验证）。无需认证。", + "operationId": "wecomCallbackVerify", + "security": []map[string]interface{}{}, + "parameters": []map[string]interface{}{ + {"name": "msg_signature", "in": "query", "required": true, "schema": map[string]interface{}{"type": "string"}}, + {"name": "timestamp", "in": "query", "required": true, "schema": map[string]interface{}{"type": "string"}}, + {"name": "nonce", "in": "query", "required": true, "schema": map[string]interface{}{"type": "string"}}, + {"name": "echostr", "in": "query", "required": true, "schema": map[string]interface{}{"type": "string"}}, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{"description": "验证成功，返回解密后的echostr"}, + }, + }, + "post": map[string]interface{}{ + "tags": []string{"机器人集成"}, + "summary": "企业微信消息回调", + "description": "接收企业微信推送的消息事件。无需认证，由企业微信服务器调用。", + "operationId": "wecomCallbackMessage", + "security": []map[string]interface{}{}, + "responses": map[string]interface{}{ + "200": map[string]interface{}{"description": "处理成功"}, + }, + }, + }, + "/api/robot/dingtalk": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"机器人集成"}, + "summary": "钉钉消息回调", + "description": "接收钉钉推送的消息事件。无需认证，由钉钉服务器调用。", + "operationId": "dingtalkCallback", + "security": []map[string]interface{}{}, + "responses": map[string]interface{}{ + "200": map[string]interface{}{"description": "处理成功"}, + }, + }, + }, + "/api/robot/lark": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"机器人集成"}, + "summary": "飞书消息回调", + "description": "接收飞书推送的消息事件。无需认证，由飞书服务器调用。", + "operationId": "larkCallback", + "security": []map[string]interface{}{}, + "responses": map[string]interface{}{ + "200": map[string]interface{}{"description": "处理成功"}, + }, + }, + }, + "/api/robot/test": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"机器人集成"}, + "summary": "测试机器人消息处理", + "description": "模拟机器人消息处理流程，用于调试和验证。需要登录认证。", + "operationId": "testRobot", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "required": []string{"platform", "text"}, + "properties": map[string]interface{}{ + "platform": map[string]interface{}{"type": "string", "description": "平台类型", "enum": []string{"dingtalk", "lark", "wecom"}}, + "user_id": map[string]interface{}{"type": "string", "description": "模拟用户ID", "example": "test"}, + "text": map[string]interface{}{"type": "string", "description": "消息文本", "example": "帮助"}, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{"description": "处理成功"}, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + }, + + // ==================== 多代理Markdown ==================== + "/api/multi-agent/markdown-agents": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"多代理Markdown"}, + "summary": "列出Markdown代理", + "description": "获取所有多代理Markdown定义文件列表。", + "operationId": "listMarkdownAgents", + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "agents": map[string]interface{}{ + "type": "array", + "items": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "filename": map[string]interface{}{"type": "string", "description": "文件名"}, + "id": map[string]interface{}{"type": "string", "description": "代理ID"}, + "name": map[string]interface{}{"type": "string", "description": "代理名称"}, + "description": map[string]interface{}{"type": "string", "description": "代理描述"}, + "is_orchestrator": map[string]interface{}{"type": "boolean", "description": "是否为编排器"}, + "kind": map[string]interface{}{"type": "string", "description": "编排类型"}, + }, + }, + }, + "dir": map[string]interface{}{"type": "string", "description": "代理定义目录路径"}, + }, + }, + }, + }, + }, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + "post": map[string]interface{}{ + "tags": []string{"多代理Markdown"}, + "summary": "创建Markdown代理", + "description": "创建新的多代理Markdown定义文件。", + "operationId": "createMarkdownAgent", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "required": []string{"name"}, + "properties": map[string]interface{}{ + "filename": map[string]interface{}{"type": "string", "description": "文件名（可选，自动生成）"}, + "id": map[string]interface{}{"type": "string", "description": "代理ID"}, + "name": map[string]interface{}{"type": "string", "description": "代理名称"}, + "description": map[string]interface{}{"type": "string", "description": "代理描述"}, + "tools": map[string]interface{}{"type": "array", "items": map[string]interface{}{"type": "string"}, "description": "可用工具列表"}, + "instruction": map[string]interface{}{"type": "string", "description": "代理指令"}, + "bind_role": map[string]interface{}{"type": "string", "description": "绑定角色"}, + "max_iterations": map[string]interface{}{"type": "integer", "description": "最大迭代次数"}, + "kind": map[string]interface{}{"type": "string", "description": "编排类型"}, + "raw": map[string]interface{}{"type": "string", "description": "原始Markdown内容"}, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "创建成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "filename": map[string]interface{}{"type": "string"}, + "message": map[string]interface{}{"type": "string", "example": "已创建"}, + }, + }, + }, + }, + }, + "400": map[string]interface{}{"description": "参数错误"}, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + }, + "/api/multi-agent/markdown-agents/{filename}": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"多代理Markdown"}, + "summary": "获取Markdown代理详情", + "description": "获取指定Markdown代理定义文件的详细内容。", + "operationId": "getMarkdownAgent", + "parameters": []map[string]interface{}{ + {"name": "filename", "in": "path", "required": true, "description": "文件名", "schema": map[string]interface{}{"type": "string"}}, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "filename": map[string]interface{}{"type": "string"}, + "raw": map[string]interface{}{"type": "string", "description": "原始Markdown内容"}, + "id": map[string]interface{}{"type": "string"}, + "name": map[string]interface{}{"type": "string"}, + "description": map[string]interface{}{"type": "string"}, + "tools": map[string]interface{}{"type": "array", "items": map[string]interface{}{"type": "string"}}, + "instruction": map[string]interface{}{"type": "string"}, + "bind_role": map[string]interface{}{"type": "string"}, + "max_iterations": map[string]interface{}{"type": "integer"}, + "kind": map[string]interface{}{"type": "string"}, + "is_orchestrator": map[string]interface{}{"type": "boolean"}, + }, + }, + }, + }, + }, + "401": map[string]interface{}{"description": "未授权"}, + "404": map[string]interface{}{"description": "代理不存在"}, + }, + }, + "put": map[string]interface{}{ + "tags": []string{"多代理Markdown"}, + "summary": "更新Markdown代理", + "description": "更新指定的Markdown代理定义。", + "operationId": "updateMarkdownAgent", + "parameters": []map[string]interface{}{ + {"name": "filename", "in": "path", "required": true, "description": "文件名", "schema": map[string]interface{}{"type": "string"}}, + }, + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "name": map[string]interface{}{"type": "string"}, + "description": map[string]interface{}{"type": "string"}, + "tools": map[string]interface{}{"type": "array", "items": map[string]interface{}{"type": "string"}}, + "instruction": map[string]interface{}{"type": "string"}, + "bind_role": map[string]interface{}{"type": "string"}, + "max_iterations": map[string]interface{}{"type": "integer"}, + "kind": map[string]interface{}{"type": "string"}, + "raw": map[string]interface{}{"type": "string"}, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "更新成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "message": map[string]interface{}{"type": "string", "example": "已保存"}, + }, + }, + }, + }, + }, + "401": map[string]interface{}{"description": "未授权"}, + "404": map[string]interface{}{"description": "代理不存在"}, + }, + }, + "delete": map[string]interface{}{ + "tags": []string{"多代理Markdown"}, + "summary": "删除Markdown代理", + "description": "删除指定的Markdown代理定义文件。", + "operationId": "deleteMarkdownAgent", + "parameters": []map[string]interface{}{ + {"name": "filename", "in": "path", "required": true, "description": "文件名", "schema": map[string]interface{}{"type": "string"}}, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "删除成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "message": map[string]interface{}{"type": "string", "example": "已删除"}, + }, + }, + }, + }, + }, + "401": map[string]interface{}{"description": "未授权"}, + "404": map[string]interface{}{"description": "代理不存在"}, + }, + }, + }, + + // ==================== Skills管理 - 缺失端点 ==================== + "/api/skills/{name}/files": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"Skills管理"}, + "summary": "列出技能包文件", + "description": "获取指定技能包目录下的所有文件列表。", + "operationId": "listSkillPackageFiles", + "parameters": []map[string]interface{}{ + {"name": "name", "in": "path", "required": true, "description": "技能名称/ID", "schema": map[string]interface{}{"type": "string"}}, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "files": map[string]interface{}{"type": "array", "items": map[string]interface{}{"type": "string"}, "description": "文件路径列表"}, + }, + }, + }, + }, + }, + "401": map[string]interface{}{"description": "未授权"}, + "404": map[string]interface{}{"description": "技能不存在"}, + }, + }, + }, + "/api/skills/{name}/file": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"Skills管理"}, + "summary": "获取技能包文件内容", + "description": "读取技能包中指定文件的内容。", + "operationId": "getSkillPackageFile", + "parameters": []map[string]interface{}{ + {"name": "name", "in": "path", "required": true, "description": "技能名称/ID", "schema": map[string]interface{}{"type": "string"}}, + {"name": "path", "in": "query", "required": true, "description": "文件相对路径", "schema": map[string]interface{}{"type": "string"}}, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "path": map[string]interface{}{"type": "string", "description": "文件路径"}, + "content": map[string]interface{}{"type": "string", "description": "文件内容"}, + }, + }, + }, + }, + }, + "401": map[string]interface{}{"description": "未授权"}, + "404": map[string]interface{}{"description": "文件不存在"}, + }, + }, + "put": map[string]interface{}{ + "tags": []string{"Skills管理"}, + "summary": "写入技能包文件", + "description": "写入或更新技能包中的文件内容。", + "operationId": "putSkillPackageFile", + "parameters": []map[string]interface{}{ + {"name": "name", "in": "path", "required": true, "description": "技能名称/ID", "schema": map[string]interface{}{"type": "string"}}, + }, + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "required": []string{"path"}, + "properties": map[string]interface{}{ + "path": map[string]interface{}{"type": "string", "description": "文件相对路径"}, + "content": map[string]interface{}{"type": "string", "description": "文件内容"}, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "保存成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "message": map[string]interface{}{"type": "string", "example": "saved"}, + "path": map[string]interface{}{"type": "string"}, + }, + }, + }, + }, + }, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + }, + + // ==================== 监控 - 缺失端点 ==================== + "/api/monitor/executions/names": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"监控"}, + "summary": "批量获取工具名称", + "description": "根据执行ID列表批量获取对应的工具名称，消除前端N+1请求问题。", + "operationId": "batchGetToolNames", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "required": []string{"ids"}, + "properties": map[string]interface{}{ + "ids": map[string]interface{}{ + "type": "array", + "items": map[string]interface{}{"type": "string"}, + "description": "执行记录ID列表", + }, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功，返回ID到工具名称的映射", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "additionalProperties": map[string]interface{}{"type": "string"}, + "description": "键为执行ID，值为工具名称", + "example": map[string]interface{}{"exec-001": "nmap", "exec-002": "sqlmap"}, + }, + }, + }, + }, + "400": map[string]interface{}{"description": "参数错误"}, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + }, + + // ==================== 知识库 - 缺失端点 ==================== + "/api/knowledge/stats": map[string]interface{}{ + "get": map[string]interface{}{ + "tags": []string{"知识库"}, + "summary": "获取知识库统计", + "description": "获取知识库的总体统计信息，包括分类数和条目数。", + "operationId": "getKnowledgeStats", + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "获取成功", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "enabled": map[string]interface{}{"type": "boolean", "description": "知识库是否启用"}, + "total_categories": map[string]interface{}{"type": "integer", "description": "分类总数"}, + "total_items": map[string]interface{}{"type": "integer", "description": "条目总数"}, + }, + }, + }, + }, + }, + "401": map[string]interface{}{"description": "未授权"}, + }, + }, + }, + + "/api/mcp": map[string]interface{}{ + "post": map[string]interface{}{ + "tags": []string{"MCP"}, + "summary": "MCP端点", + "description": "MCP (Model Context Protocol) 端点，用于处理MCP协议请求。\n**协议说明**：\n本端点遵循 JSON-RPC 2.0 规范，支持以下方法：\n**1. initialize** - 初始化MCP连接\n```json\n{\n \"jsonrpc\": \"2.0\",\n \"id\": \"init-1\",\n \"method\": \"initialize\",\n \"params\": {\n \"protocolVersion\": \"2024-11-05\",\n \"capabilities\": {},\n \"clientInfo\": {\n \"name\": \"MyClient\",\n \"version\": \"1.0.0\"\n }\n }\n}\n```\n**2. tools/list** - 列出所有可用工具\n```json\n{\n \"jsonrpc\": \"2.0\",\n \"id\": \"list-1\",\n \"method\": \"tools/list\",\n \"params\": {}\n}\n```\n**3. tools/call** - 调用工具\n```json\n{\n \"jsonrpc\": \"2.0\",\n \"id\": \"call-1\",\n \"method\": \"tools/call\",\n \"params\": {\n \"name\": \"nmap\",\n \"arguments\": {\n \"target\": \"192.168.1.1\",\n \"ports\": \"80,443\"\n }\n }\n}\n```\n**4. prompts/list** - 列出所有提示词模板\n```json\n{\n \"jsonrpc\": \"2.0\",\n \"id\": \"prompts-list-1\",\n \"method\": \"prompts/list\",\n \"params\": {}\n}\n```\n**5. prompts/get** - 获取提示词模板\n```json\n{\n \"jsonrpc\": \"2.0\",\n \"id\": \"prompt-get-1\",\n \"method\": \"prompts/get\",\n \"params\": {\n \"name\": \"prompt-name\",\n \"arguments\": {}\n }\n}\n```\n**6. resources/list** - 列出所有资源\n```json\n{\n \"jsonrpc\": \"2.0\",\n \"id\": \"resources-list-1\",\n \"method\": \"resources/list\",\n \"params\": {}\n}\n```\n**7. resources/read** - 读取资源内容\n```json\n{\n \"jsonrpc\": \"2.0\",\n \"id\": \"resource-read-1\",\n \"method\": \"resources/read\",\n \"params\": {\n \"uri\": \"resource://example\"\n }\n}\n```\n**错误代码说明**：\n- `-32700`: Parse error - JSON解析错误\n- `-32600`: Invalid Request - 无效请求\n- `-32601`: Method not found - 方法不存在\n- `-32602`: Invalid params - 参数无效\n- `-32603`: Internal error - 内部错误", + "operationId": "mcpEndpoint", + "requestBody": map[string]interface{}{ + "required": true, + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/MCPMessage", + }, + "examples": map[string]interface{}{ + "listTools": map[string]interface{}{ + "summary": "列出所有工具", + "description": "获取系统中所有可用的MCP工具列表", + "value": map[string]interface{}{ + "jsonrpc": "2.0", + "id": "list-tools-1", + "method": "tools/list", + "params": map[string]interface{}{}, + }, + }, + "callTool": map[string]interface{}{ + "summary": "调用工具", + "description": "调用指定的MCP工具", + "value": map[string]interface{}{ + "jsonrpc": "2.0", + "id": "call-tool-1", + "method": "tools/call", + "params": map[string]interface{}{ + "name": "nmap", + "arguments": map[string]interface{}{ + "target": "192.168.1.1", + "ports": "80,443", + }, + }, + }, + }, + "initialize": map[string]interface{}{ + "summary": "初始化连接", + "description": "初始化MCP连接，获取服务器能力", + "value": map[string]interface{}{ + "jsonrpc": "2.0", + "id": "init-1", + "method": "initialize", + "params": map[string]interface{}{ + "protocolVersion": "2024-11-05", + "capabilities": map[string]interface{}{}, + "clientInfo": map[string]interface{}{ + "name": "MyClient", + "version": "1.0.0", + }, + }, + }, + }, + }, + }, + }, + }, + "responses": map[string]interface{}{ + "200": map[string]interface{}{ + "description": "MCP响应（JSON-RPC 2.0格式）", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/MCPResponse", + }, + "examples": map[string]interface{}{ + "success": map[string]interface{}{ + "summary": "成功响应", + "description": "工具调用成功的响应示例", + "value": map[string]interface{}{ + "jsonrpc": "2.0", + "id": "call-tool-1", + "result": map[string]interface{}{ + "content": []map[string]interface{}{ + { + "type": "text", + "text": "工具执行结果...", + }, + }, + "isError": false, + }, + }, + }, + "error": map[string]interface{}{ + "summary": "错误响应", + "description": "工具调用失败的响应示例", + "value": map[string]interface{}{ + "jsonrpc": "2.0", + "id": "call-tool-1", + "error": map[string]interface{}{ + "code": -32601, + "message": "Tool not found", + "data": "工具 'unknown-tool' 不存在", + }, + }, + }, + }, + }, + }, + }, + "400": map[string]interface{}{ + "description": "请求格式错误（JSON解析失败）", + "content": map[string]interface{}{ + "application/json": map[string]interface{}{ + "schema": map[string]interface{}{ + "$ref": "#/components/schemas/MCPResponse", + }, + "example": map[string]interface{}{ + "id": nil, + "error": map[string]interface{}{ + "code": -32700, + "message": "Parse error", + "data": "unexpected end of JSON input", + }, + "jsonrpc": "2.0", + }, + }, + }, + }, + "401": map[string]interface{}{ + "description": "未授权，需要有效的Token", + }, + "405": map[string]interface{}{ + "description": "方法不允许（仅支持POST请求）", + }, + }, + }, + }, + }, + } + + enrichSpecWithI18nKeys(spec) + c.JSON(http.StatusOK, spec) +} + +// GetConversationResults 获取对话结果（OpenAPI端点） +// 注意：创建对话和获取对话详情直接使用标准的 /api/conversations 端点 +// 这个端点只是为了提供结果聚合功能 +func (h *OpenAPIHandler) GetConversationResults(c *gin.Context) { + conversationID := c.Param("id") + + // 验证对话是否存在 + conv, err := h.db.GetConversation(conversationID) + if err != nil { + h.logger.Error("获取对话失败", zap.Error(err)) + c.JSON(http.StatusNotFound, gin.H{"error": "对话不存在"}) + return + } + + // 获取消息列表 + messages, err := h.db.GetMessages(conversationID) + if err != nil { + h.logger.Error("获取消息失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + // 获取漏洞列表 + vulnList, err := h.db.ListVulnerabilities(1000, 0, database.VulnerabilityListFilter{ConversationID: conversationID}) + if err != nil { + h.logger.Warn("获取漏洞列表失败", zap.Error(err)) + vulnList = []*database.Vulnerability{} + } + vulnerabilities := make([]database.Vulnerability, len(vulnList)) + for i, v := range vulnList { + vulnerabilities[i] = *v + } + + // 获取执行结果（历史大结果由 Eino reduction 落盘，此处不再聚合文件存储） + executionResults := []map[string]interface{}{} + + response := map[string]interface{}{ + "conversationId": conv.ID, + "messages": messages, + "vulnerabilities": vulnerabilities, + "executionResults": executionResults, + } + + c.JSON(http.StatusOK, response) +} diff --git a/internal/handler/openapi_i18n.go b/internal/handler/openapi_i18n.go new file mode 100644 index 00000000..d0480c0f --- /dev/null +++ b/internal/handler/openapi_i18n.go @@ -0,0 +1,190 @@ +package handler + +// apiDocI18n 为 OpenAPI 文档提供 x-i18n-* 扩展键，供前端 apiDocs 国际化使用。 +// 前端通过 apiDocs.tags.* / apiDocs.summary.* / apiDocs.response.* 翻译。 + +var apiDocI18nTagToKey = map[string]string{ + "认证": "auth", "对话管理": "conversationManagement", "对话交互": "conversationInteraction", + "批量任务": "batchTasks", "对话分组": "conversationGroups", "漏洞管理": "vulnerabilityManagement", + "角色管理": "roleManagement", "Skills管理": "skillsManagement", "监控": "monitoring", + "配置管理": "configManagement", "外部MCP管理": "externalMCPManagement", "攻击链": "attackChain", + "知识库": "knowledgeBase", "MCP": "mcp", + "FOFA信息收集": "fofaRecon", "终端": "terminal", "WebShell管理": "webshellManagement", + "对话附件": "chatUploads", "机器人集成": "robotIntegration", "多代理Markdown": "markdownAgents", + "项目管理": "projectManagement", "资产管理": "assetManagement", +} + +var apiDocI18nSummaryToKey = map[string]string{ + "用户登录": "login", "用户登出": "logout", "修改密码": "changePassword", "验证Token": "validateToken", + "创建对话": "createConversation", "列出对话": "listConversations", "查看对话详情": "getConversationDetail", + "更新对话": "updateConversation", "删除对话": "deleteConversation", "获取对话结果": "getConversationResult", + "发送消息并获取AI回复（非流式）": "sendMessageNonStream", "发送消息并获取AI回复（流式）": "sendMessageStream", + "取消任务": "cancelTask", "列出运行中的任务": "listRunningTasks", "列出已完成的任务": "listCompletedTasks", + "创建批量任务队列": "createBatchQueue", "列出批量任务队列": "listBatchQueues", "获取批量任务队列": "getBatchQueue", + "删除批量任务队列": "deleteBatchQueue", "启动批量任务队列": "startBatchQueue", "暂停批量任务队列": "pauseBatchQueue", + "添加任务到队列": "addTaskToQueue", "SQL注入扫描": "sqlInjectionScan", "端口扫描": "portScan", + "更新批量任务": "updateBatchTask", "删除批量任务": "deleteBatchTask", + "创建分组": "createGroup", "列出分组": "listGroups", "获取分组": "getGroup", "更新分组": "updateGroup", + "删除分组": "deleteGroup", "获取分组中的对话": "getGroupConversations", "添加对话到分组": "addConversationToGroup", + "从分组移除对话": "removeConversationFromGroup", + "列出漏洞": "listVulnerabilities", "创建漏洞": "createVulnerability", "获取漏洞统计": "getVulnerabilityStats", + "获取漏洞": "getVulnerability", "更新漏洞": "updateVulnerability", "删除漏洞": "deleteVulnerability", + "列出角色": "listRoles", "创建角色": "createRole", "获取角色": "getRole", "更新角色": "updateRole", "删除角色": "deleteRole", + "获取可用Skills列表": "getAvailableSkills", "列出Skills": "listSkills", "创建Skill": "createSkill", + "获取Skill统计": "getSkillStats", "清空Skill统计": "clearSkillStats", "获取Skill": "getSkill", + "更新Skill": "updateSkill", "删除Skill": "deleteSkill", "获取绑定角色": "getBoundRoles", + "获取监控信息": "getMonitorInfo", "获取执行记录": "getExecutionRecords", "删除执行记录": "deleteExecutionRecord", + "批量删除执行记录": "batchDeleteExecutionRecords", "获取统计信息": "getStats", + "获取配置": "getConfig", "更新配置": "updateConfig", "获取工具配置": "getToolConfig", "应用配置": "applyConfig", + "列出外部MCP": "listExternalMCP", "获取外部MCP统计": "getExternalMCPStats", "获取外部MCP": "getExternalMCP", + "添加或更新外部MCP": "addOrUpdateExternalMCP", "stdio模式配置": "stdioModeConfig", "SSE模式配置": "sseModeConfig", + "删除外部MCP": "deleteExternalMCP", "启动外部MCP": "startExternalMCP", "停止外部MCP": "stopExternalMCP", + "获取攻击链": "getAttackChain", "重新生成攻击链": "regenerateAttackChain", + "设置对话置顶": "pinConversation", "设置分组置顶": "pinGroup", "设置分组中对话的置顶": "pinGroupConversation", + "获取分类": "getCategories", "列出知识项": "listKnowledgeItems", "创建知识项": "createKnowledgeItem", + "获取知识项": "getKnowledgeItem", "更新知识项": "updateKnowledgeItem", "删除知识项": "deleteKnowledgeItem", + "获取索引状态": "getIndexStatus", "构建索引": "startKnowledgeIndex", "扫描知识库": "scanKnowledgeBase", + "搜索知识库": "searchKnowledgeBase", "基础搜索": "basicSearch", "按风险类型搜索": "searchByRiskType", + "获取检索日志": "getRetrievalLogs", "删除检索日志": "deleteRetrievalLog", + "MCP端点": "mcpEndpoint", "列出所有工具": "listAllTools", "调用工具": "invokeTool", "初始化连接": "initConnection", + "成功响应": "successResponse", "错误响应": "errorResponse", + // 新增缺失端点 + "删除对话轮次": "deleteConversationTurn", "获取消息过程详情": "getMessageProcessDetails", + "重跑批量任务队列": "rerunBatchQueue", "修改队列元数据": "updateBatchQueueMetadata", + "修改队列调度配置": "updateBatchQueueSchedule", "开关Cron自动调度": "setBatchQueueScheduleEnabled", + "获取所有分组映射": "getAllGroupMappings", + "FOFA搜索": "fofaSearch", "自然语言解析为FOFA语法": "fofaParse", + "测试OpenAI API连接": "testOpenAI", + "执行终端命令": "terminalRun", "流式执行终端命令": "terminalRunStream", "WebSocket终端": "terminalWS", + "列出WebShell连接": "listWebshellConnections", "创建WebShell连接": "createWebshellConnection", + "更新WebShell连接": "updateWebshellConnection", "删除WebShell连接": "deleteWebshellConnection", + "获取连接状态": "getWebshellConnectionState", "保存连接状态": "saveWebshellConnectionState", + "获取AI对话历史": "getWebshellAIHistory", "列出AI对话": "listWebshellAIConversations", + "执行WebShell命令": "webshellExec", "WebShell文件操作": "webshellFileOp", + "列出附件": "listChatUploads", "导出附件": "exportChatUploads", "上传附件": "uploadChatFile", "删除附件": "deleteChatUpload", + "下载附件": "downloadChatUpload", "获取附件文本内容": "getChatUploadContent", + "写入附件文本内容": "putChatUploadContent", "创建附件目录": "mkdirChatUpload", "重命名附件": "renameChatUpload", + "企业微信回调验证": "wecomCallbackVerify", "企业微信消息回调": "wecomCallbackMessage", + "钉钉消息回调": "dingtalkCallback", "飞书消息回调": "larkCallback", "测试机器人消息处理": "testRobot", + "列出Markdown代理": "listMarkdownAgents", "创建Markdown代理": "createMarkdownAgent", + "获取Markdown代理详情": "getMarkdownAgent", "更新Markdown代理": "updateMarkdownAgent", "删除Markdown代理": "deleteMarkdownAgent", + "列出技能包文件": "listSkillPackageFiles", "获取技能包文件内容": "getSkillPackageFile", "写入技能包文件": "putSkillPackageFile", + "批量获取工具名称": "batchGetToolNames", + "获取知识库统计": "getKnowledgeStats", + "列出项目": "listProjects", "创建项目": "createProject", "获取项目": "getProject", + "更新项目": "updateProject", "删除项目": "deleteProject", + "批量导入资产": "importAssets", + "列出或按 key 获取事实": "listProjectFacts", "创建/更新事实": "upsertProjectFact", + "获取项目事实攻击路径图": "getProjectFactGraph", "列出项目全部事实边": "listProjectFactEdges", + "添加事实边": "createProjectFactEdge", "删除事实边": "deleteProjectFactEdge", + "将对话攻击链沉淀到项目事实图": "promoteAttackChainToProject", +} + +var apiDocI18nResponseDescToKey = map[string]string{ + "获取成功": "getSuccess", "未授权": "unauthorized", "未授权，需要有效的Token": "unauthorizedToken", + "创建成功": "createSuccess", "请求参数错误": "badRequest", "对话不存在": "conversationNotFound", + "对话不存在或结果不存在": "conversationOrResultNotFound", "请求参数错误（如task为空）": "badRequestTaskEmpty", + "请求参数错误或分组名称已存在": "badRequestGroupNameExists", "分组不存在": "groupNotFound", + "请求参数错误（如配置格式不正确、缺少必需字段等）": "badRequestConfig", + "请求参数错误（如query为空）": "badRequestQueryEmpty", "方法不允许（仅支持POST请求）": "methodNotAllowed", + "登录成功": "loginSuccess", "密码错误": "invalidPassword", "登出成功": "logoutSuccess", + "密码修改成功": "passwordChanged", "Token有效": "tokenValid", "Token无效或已过期": "tokenInvalid", + "对话创建成功": "conversationCreated", "服务器内部错误": "internalError", "更新成功": "updateSuccess", + "删除成功": "deleteSuccess", "队列不存在": "queueNotFound", "启动成功": "startSuccess", + "暂停成功": "pauseSuccess", "添加成功": "addSuccess", + "任务不存在": "taskNotFound", "对话或分组不存在": "conversationOrGroupNotFound", + "取消请求已提交": "cancelSubmitted", "未找到正在执行的任务": "noRunningTask", + "消息发送成功，返回AI回复": "messageSent", "流式响应（Server-Sent Events）": "streamResponse", + // 新增缺失端点响应 + "参数错误或删除失败": "badRequestOrDeleteFailed", + "参数错误": "paramError", "仅已完成或已取消的队列可以重跑": "onlyCompletedOrCancelledCanRerun", + "参数错误或队列正在运行中": "badRequestOrQueueRunning", "设置成功": "setSuccess", + "搜索成功": "searchSuccess", "解析成功": "parseSuccess", "测试结果": "testResult", + "执行完成": "executionDone", "SSE事件流": "sseEventStream", "WebSocket连接已建立": "wsEstablished", + "文件下载": "fileDownload", "文件不存在": "fileNotFound", "写入成功": "writeSuccess", + "重命名成功": "renameSuccess", "验证成功，返回解密后的echostr": "wecomVerifySuccess", + "处理成功": "processSuccess", "代理不存在": "agentNotFound", "保存成功": "saveSuccess", + "操作结果": "operationResult", "执行结果": "executionResult", "连接不存在": "connectionNotFound", + "项目列表": "projectList", "项目详情": "projectDetail", + "事实列表或单条（可含 link_counts / outgoing_links）": "projectFactList", + "成功": "success", "nodes + edges": "factGraphNodesEdges", + "边列表": "edgeList", "边已创建": "edgeCreated", + "沉淀结果（facts/edges/graph）": "promoteAttackChainResult", + "导入完成": "assetImportCompleted", "数量或资产字段校验失败": "assetImportValidationFailed", + "缺少 asset:write 权限或无权访问指定项目": "assetImportForbidden", + "导入事务失败": "assetImportTransactionFailed", +} + +// enrichSpecWithI18nKeys 在 spec 的每个 operation 上写入 x-i18n-tags、x-i18n-summary， +// 在每个 response 上写入 x-i18n-description，供前端按 key 做国际化。 +func enrichSpecWithI18nKeys(spec map[string]interface{}) { + paths, _ := spec["paths"].(map[string]interface{}) + if paths == nil { + return + } + for _, pathItem := range paths { + pm, _ := pathItem.(map[string]interface{}) + if pm == nil { + continue + } + for _, method := range []string{"get", "post", "put", "delete", "patch"} { + opVal, ok := pm[method] + if !ok { + continue + } + op, _ := opVal.(map[string]interface{}) + if op == nil { + continue + } + // x-i18n-tags: 与 tags 一一对应的 i18n 键数组（spec 中 tags 为 []string） + switch tags := op["tags"].(type) { + case []string: + if len(tags) > 0 { + keys := make([]string, 0, len(tags)) + for _, s := range tags { + if k := apiDocI18nTagToKey[s]; k != "" { + keys = append(keys, k) + } else { + keys = append(keys, s) + } + } + op["x-i18n-tags"] = keys + } + case []interface{}: + if len(tags) > 0 { + keys := make([]interface{}, 0, len(tags)) + for _, t := range tags { + if s, ok := t.(string); ok { + if k := apiDocI18nTagToKey[s]; k != "" { + keys = append(keys, k) + } else { + keys = append(keys, s) + } + } + } + if len(keys) > 0 { + op["x-i18n-tags"] = keys + } + } + } + // x-i18n-summary + if summary, _ := op["summary"].(string); summary != "" { + if k := apiDocI18nSummaryToKey[summary]; k != "" { + op["x-i18n-summary"] = k + } + } + // responses -> 每个 status -> x-i18n-description + if respMap, _ := op["responses"].(map[string]interface{}); respMap != nil { + for _, rv := range respMap { + if r, _ := rv.(map[string]interface{}); r != nil { + if desc, _ := r["description"].(string); desc != "" { + if k := apiDocI18nResponseDescToKey[desc]; k != "" { + r["x-i18n-description"] = k + } + } + } + } + } + } + } +} diff --git a/internal/handler/openapi_i18n_test.go b/internal/handler/openapi_i18n_test.go new file mode 100644 index 00000000..12f08a0d --- /dev/null +++ b/internal/handler/openapi_i18n_test.go @@ -0,0 +1,46 @@ +package handler + +import "testing" + +func TestEnrichSpecWithI18nKeysForAssetImport(t *testing.T) { + responses := map[string]interface{}{ + "200": map[string]interface{}{"description": "导入完成"}, + "400": map[string]interface{}{"description": "数量或资产字段校验失败"}, + "403": map[string]interface{}{"description": "缺少 asset:write 权限或无权访问指定项目"}, + "500": map[string]interface{}{"description": "导入事务失败"}, + } + operation := map[string]interface{}{ + "tags": []string{"资产管理"}, + "summary": "批量导入资产", + "responses": responses, + } + spec := map[string]interface{}{ + "paths": map[string]interface{}{ + "/api/assets/import": map[string]interface{}{ + "post": operation, + }, + }, + } + + enrichSpecWithI18nKeys(spec) + + tagKeys, ok := operation["x-i18n-tags"].([]string) + if !ok || len(tagKeys) != 1 || tagKeys[0] != "assetManagement" { + t.Fatalf("unexpected asset tag i18n keys: %#v", operation["x-i18n-tags"]) + } + if got := operation["x-i18n-summary"]; got != "importAssets" { + t.Fatalf("unexpected asset summary i18n key: %#v", got) + } + expectedResponseKeys := map[string]string{ + "200": "assetImportCompleted", + "400": "assetImportValidationFailed", + "403": "assetImportForbidden", + "500": "assetImportTransactionFailed", + } + for status, want := range expectedResponseKeys { + response := responses[status].(map[string]interface{}) + if got := response["x-i18n-description"]; got != want { + t.Errorf("unexpected asset response i18n key for %s: got %#v, want %q", status, got, want) + } + } +} diff --git a/internal/handler/project.go b/internal/handler/project.go new file mode 100644 index 00000000..c671ac65 --- /dev/null +++ b/internal/handler/project.go @@ -0,0 +1,657 @@ +package handler + +import ( + "fmt" + "net/http" + "strconv" + "strings" + + "cyberstrike-ai/internal/attackchain" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/project" + "cyberstrike-ai/internal/security" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +const maxProjectDescriptionRunes = 4000 + +func clampProjectDescription(s string) string { + r := []rune(s) + if len(r) <= maxProjectDescriptionRunes { + return s + } + return string(r[:maxProjectDescriptionRunes]) +} + +// ProjectHandler 项目管理处理器。 +type ProjectHandler struct { + db *database.DB + logger *zap.Logger +} + +// NewProjectHandler 创建项目管理处理器。 +func NewProjectHandler(db *database.DB, logger *zap.Logger) *ProjectHandler { + return &ProjectHandler{db: db, logger: logger} +} + +type createProjectRequest struct { + Name string `json:"name" binding:"required"` + Description string `json:"description"` + ScopeJSON string `json:"scope_json"` + Status string `json:"status"` +} + +// updateProjectRequest 部分更新：字段省略表示不修改；传 null 或 "" 可清空字符串字段。 +type updateProjectRequest struct { + Name *string `json:"name"` + Description *string `json:"description"` + ScopeJSON *string `json:"scope_json"` + Status *string `json:"status"` + Pinned *bool `json:"pinned"` +} + +// CreateProject POST /api/projects +func (h *ProjectHandler) CreateProject(c *gin.Context) { + var req createProjectRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + p := &database.Project{ + Name: strings.TrimSpace(req.Name), + Description: clampProjectDescription(req.Description), + ScopeJSON: req.ScopeJSON, + Status: strings.TrimSpace(req.Status), + } + created, err := h.db.CreateProject(p) + if err != nil { + h.logger.Error("创建项目失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if session, ok := security.CurrentSession(c); ok { + _ = h.db.SetResourceOwner("project", created.ID, session.UserID) + _ = h.db.AssignResourceToUser(session.UserID, "project", created.ID) + } + c.JSON(http.StatusOK, created) +} + +// GetDashboardSummary GET /api/projects/dashboard-summary +func (h *ProjectHandler) GetDashboardSummary(c *gin.Context) { + limit, _ := strconv.Atoi(strings.TrimSpace(c.DefaultQuery("fact_limit", "5"))) + if limit <= 0 { + limit = 5 + } + if limit > 50 { + limit = 50 + } + session, _ := security.CurrentSession(c) + summary, err := h.db.GetProjectDashboardSummaryForAccess(limit, session.UserID, session.Scope) + if err != nil { + h.logger.Error("获取项目仪表盘摘要失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if summary.RecentFacts == nil { + summary.RecentFacts = []database.ProjectDashboardFact{} + } + c.JSON(http.StatusOK, summary) +} + +// ListProjects GET /api/projects +func (h *ProjectHandler) ListProjects(c *gin.Context) { + status := c.Query("status") + search := c.Query("search") + limit, _ := strconv.Atoi(c.DefaultQuery("limit", "50")) + offset, _ := strconv.Atoi(c.Query("offset")) + if limit <= 0 { + limit = 50 + } + if limit > 500 { + limit = 500 + } + session, _ := security.CurrentSession(c) + list, err := h.db.ListProjectsForAccess(status, search, limit, offset, session.UserID, session.Scope) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if list == nil { + list = []*database.Project{} + } + total, err := h.db.CountProjectsForAccess(status, search, session.UserID, session.Scope) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{ + "projects": list, + "total": total, + "limit": limit, + "offset": offset, + }) +} + +// GetProjectStats GET /api/projects/:id/stats +func (h *ProjectHandler) GetProjectStats(c *gin.Context) { + stats, err := project.GetProjectStats(h.db, c.Param("id")) + if err != nil { + if strings.Contains(err.Error(), "不存在") { + c.JSON(http.StatusNotFound, gin.H{"error": "项目不存在"}) + return + } + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, stats) +} + +// ListProjectConversations GET /api/projects/:id/conversations +func (h *ProjectHandler) ListProjectConversations(c *gin.Context) { + projectID := c.Param("id") + if _, err := h.db.GetProject(projectID); err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": "项目不存在"}) + return + } + limit, _ := strconv.Atoi(c.DefaultQuery("limit", "100")) + offset, _ := strconv.Atoi(c.Query("offset")) + list, err := h.db.ListConversationsByProjectID(projectID, limit, offset) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if list == nil { + list = []*database.Conversation{} + } + total, _ := h.db.CountConversationsByProjectID(projectID) + c.JSON(http.StatusOK, gin.H{ + "conversations": list, + "total": total, + "limit": limit, + "offset": offset, + }) +} + +// GetProject GET /api/projects/:id +func (h *ProjectHandler) GetProject(c *gin.Context) { + p, err := h.db.GetProject(c.Param("id")) + if err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": "项目不存在"}) + return + } + c.JSON(http.StatusOK, p) +} + +// UpdateProject PUT /api/projects/:id +func (h *ProjectHandler) UpdateProject(c *gin.Context) { + id := c.Param("id") + p, err := h.db.GetProject(id) + if err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": "项目不存在"}) + return + } + var req updateProjectRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if req.Name != nil { + if s := strings.TrimSpace(*req.Name); s != "" { + p.Name = s + } + } + if req.Description != nil { + p.Description = clampProjectDescription(*req.Description) + } + if req.ScopeJSON != nil { + p.ScopeJSON = *req.ScopeJSON + } + if req.Status != nil { + if s := strings.TrimSpace(*req.Status); s != "" { + p.Status = s + } + } + if req.Pinned != nil { + p.Pinned = *req.Pinned + } + if err := h.db.UpdateProject(p); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, p) +} + +// DeleteProject DELETE /api/projects/:id +func (h *ProjectHandler) DeleteProject(c *gin.Context) { + if err := h.db.DeleteProject(c.Param("id")); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"success": true}) +} + +type factLinkRequest struct { + From string `json:"from"` + Type string `json:"type"` + Confidence string `json:"confidence,omitempty"` +} + +type upsertFactRequest struct { + FactKey string `json:"fact_key" binding:"required"` + Category string `json:"category"` + Summary string `json:"summary" binding:"required"` + Body string `json:"body"` + Confidence string `json:"confidence"` + Pinned bool `json:"pinned"` + RelatedVulnerabilityID string `json:"related_vulnerability_id"` + Links []factLinkRequest `json:"links"` + LinksText *string `json:"links_text"` +} + +// updateFactRequest 部分更新事实；指针字段省略=不修改，body 传 "" 可清空（仍走 merge 逻辑见 Upsert）。 +type updateFactRequest struct { + FactKey *string `json:"fact_key"` + Category *string `json:"category"` + Summary *string `json:"summary"` + Body *string `json:"body"` + Confidence *string `json:"confidence"` + Pinned *bool `json:"pinned"` + RelatedVulnerabilityID *string `json:"related_vulnerability_id"` + ClearBody bool `json:"clear_body"` + Links *[]factLinkRequest `json:"links"` + LinksText *string `json:"links_text"` +} + +func factLinksFromRequest(links []factLinkRequest, linksText *string) (*project.ParsedFactLinks, error) { + if len(links) > 0 { + parsed := &project.ParsedFactLinks{} + for i, l := range links { + from := strings.TrimSpace(l.From) + edgeType := strings.TrimSpace(l.Type) + if from == "" { + return nil, fmt.Errorf("links[%d] 须含 from", i) + } + if edgeType == "" { + return nil, fmt.Errorf("links[%d] 须含 type", i) + } + parsed.Incoming = append(parsed.Incoming, database.ProjectFactEdgeFromInput{ + From: from, Type: edgeType, Confidence: strings.TrimSpace(l.Confidence), + }) + } + return parsed, nil + } + if linksText != nil { + in, err := project.ParseFactLinksText(*linksText) + if err != nil { + return nil, err + } + return &project.ParsedFactLinks{Incoming: in}, nil + } + return &project.ParsedFactLinks{Incoming: []database.ProjectFactEdgeFromInput{}}, nil +} + +type factWithLinksResponse struct { + *database.ProjectFact + OutgoingLinks []*database.ProjectFactEdge `json:"outgoing_links,omitempty"` + IncomingLinks []*database.ProjectFactEdge `json:"incoming_links,omitempty"` + LinkCounts *project.LinkCounts `json:"link_counts,omitempty"` +} + +func (h *ProjectHandler) applyFactLinksAfterUpsert(projectID string, fact *database.ProjectFact, links []factLinkRequest, linksText *string, explicitLinks, parseBody bool) error { + if explicitLinks { + parsed, err := factLinksFromRequest(links, linksText) + if err != nil { + return err + } + return project.PersistFactLinksFromParsed(h.db, projectID, fact.FactKey, fact.SourceConversationID, parsed, true) + } + if parseBody { + inputs := project.ParseLinksFromBody(fact.Body) + if inputs == nil { + return nil + } + return project.PersistFactIncomingLinks(h.db, projectID, fact.FactKey, inputs, true) + } + return nil +} + +func (h *ProjectHandler) factResponseWithLinks(projectID string, f *database.ProjectFact, includeLinks bool) interface{} { + if !includeLinks || f == nil { + return f + } + out, _ := h.db.ListOutgoingProjectFactEdges(projectID, f.FactKey) + in, _ := h.db.ListIncomingProjectFactEdges(projectID, f.FactKey) + return &factWithLinksResponse{ + ProjectFact: f, + OutgoingLinks: out, + IncomingLinks: in, + } +} + +// ListFacts GET /api/projects/:id/facts （fact_key 查询参数可获取单条详情） +func (h *ProjectHandler) ListFacts(c *gin.Context) { + projectID := c.Param("id") + if key := strings.TrimSpace(c.Query("fact_key")); key != "" { + f, err := h.db.GetProjectFactByKey(projectID, key) + if err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": err.Error()}) + return + } + includeLinks := c.Query("include_links") == "1" || c.Query("include_links") == "true" + c.JSON(http.StatusOK, h.factResponseWithLinks(projectID, f, includeLinks)) + return + } + limit, _ := strconv.Atoi(c.DefaultQuery("limit", "100")) + offset, _ := strconv.Atoi(c.Query("offset")) + filter := database.ProjectFactListFilter{ + Category: c.Query("category"), + Confidence: c.Query("confidence"), + Search: c.Query("search"), + RelatedVulnerabilityID: c.Query("related_vulnerability_id"), + } + if c.Query("exclude_deprecated") == "1" || c.Query("exclude_deprecated") == "true" { + filter.ExcludeDeprecated = true + } + list, err := h.db.ListProjectFacts(projectID, filter, limit, offset) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if list == nil { + list = []*database.ProjectFact{} + } + if sparseOnly := c.Query("sparse_only"); sparseOnly == "1" || sparseOnly == "true" { + filtered := make([]*database.ProjectFact, 0, len(list)) + for _, f := range list { + if project.IsSparseFactBody(f.Category, f.FactKey, f.Body) { + filtered = append(filtered, f) + } + } + list = filtered + } + includeLinkCounts := c.Query("include_link_counts") == "1" || c.Query("include_link_counts") == "true" + if !includeLinkCounts { + c.JSON(http.StatusOK, list) + return + } + counts, err := project.LoadProjectFactLinkCounts(h.db, projectID) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + out := make([]factWithLinksResponse, 0, len(list)) + for _, f := range list { + item := factWithLinksResponse{ProjectFact: f} + if c, ok := counts[f.FactKey]; ok { + cc := c + item.LinkCounts = &cc + } + out = append(out, item) + } + c.JSON(http.StatusOK, out) +} + +// GetFactGraph GET /api/projects/:id/fact-graph?view=path|full +func (h *ProjectHandler) GetFactGraph(c *gin.Context) { + projectID := c.Param("id") + if _, err := h.db.GetProject(projectID); err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": "项目不存在"}) + return + } + view := c.DefaultQuery("view", "path") + excludeDeprecated := true + if v := c.Query("exclude_deprecated"); v == "0" || v == "false" { + excludeDeprecated = false + } + graph, err := project.BuildProjectFactGraph(h.db, projectID, view, excludeDeprecated) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if graph.Nodes == nil { + graph.Nodes = []database.ProjectFactGraphNode{} + } + if graph.Edges == nil { + graph.Edges = []database.ProjectFactGraphEdge{} + } + c.JSON(http.StatusOK, graph) +} + +// CreateFact POST /api/projects/:id/facts +func (h *ProjectHandler) CreateFact(c *gin.Context) { + var req upsertFactRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + projectID := c.Param("id") + f := &database.ProjectFact{ + ProjectID: projectID, + FactKey: req.FactKey, + Category: req.Category, + Summary: req.Summary, + Body: req.Body, + Confidence: req.Confidence, + Pinned: req.Pinned, + RelatedVulnerabilityID: req.RelatedVulnerabilityID, + } + created, err := h.db.UpsertProjectFact(f) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + explicitLinks := req.Links != nil || req.LinksText != nil + if err := h.applyFactLinksAfterUpsert(projectID, created, req.Links, req.LinksText, explicitLinks, !explicitLinks); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + created, _ = h.db.GetProjectFactByKey(projectID, created.FactKey) + c.JSON(http.StatusOK, h.factResponseWithLinks(projectID, created, true)) +} + +// UpdateFact PUT /api/projects/:id/facts/:factId +func (h *ProjectHandler) UpdateFact(c *gin.Context) { + projectID := c.Param("id") + existing, err := h.db.GetProjectFact(c.Param("factId")) + if err != nil || existing.ProjectID != projectID { + c.JSON(http.StatusNotFound, gin.H{"error": "事实不存在"}) + return + } + oldFactKey := existing.FactKey + var req updateFactRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if req.FactKey != nil { + if k := strings.TrimSpace(*req.FactKey); k != "" { + existing.FactKey = k + } + } + if req.Category != nil && strings.TrimSpace(*req.Category) != "" { + existing.Category = *req.Category + } + if req.Summary != nil && strings.TrimSpace(*req.Summary) != "" { + existing.Summary = *req.Summary + } + if req.ClearBody { + existing.Body = "" + } else if req.Body != nil { + existing.Body = *req.Body + } + if req.Confidence != nil && strings.TrimSpace(*req.Confidence) != "" { + existing.Confidence = *req.Confidence + } + if req.Pinned != nil { + existing.Pinned = *req.Pinned + } + if req.RelatedVulnerabilityID != nil { + existing.RelatedVulnerabilityID = *req.RelatedVulnerabilityID + } + updated, err := h.db.UpsertProjectFact(existing) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if oldFactKey != updated.FactKey { + if err := h.db.RenameProjectFactKeyEdges(projectID, oldFactKey, updated.FactKey); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + } + if req.Links != nil || req.LinksText != nil { + var links []factLinkRequest + if req.Links != nil { + links = *req.Links + } + if err := h.applyFactLinksAfterUpsert(projectID, updated, links, req.LinksText, true, false); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + } else if req.ClearBody || req.Body != nil { + if err := h.applyFactLinksAfterUpsert(projectID, updated, nil, nil, false, true); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + } + updated, _ = h.db.GetProjectFactByKey(projectID, updated.FactKey) + c.JSON(http.StatusOK, h.factResponseWithLinks(projectID, updated, true)) +} + +// DeleteFact DELETE /api/projects/:id/facts/:factId +func (h *ProjectHandler) DeleteFact(c *gin.Context) { + existing, err := h.db.GetProjectFact(c.Param("factId")) + if err != nil || existing.ProjectID != c.Param("id") { + c.JSON(http.StatusNotFound, gin.H{"error": "事实不存在"}) + return + } + if err := h.db.DeleteProjectFact(existing.ID); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"success": true}) +} + +type deprecateFactRequest struct { + FactKey string `json:"fact_key" binding:"required"` +} + +// DeprecateFact POST /api/projects/:id/facts/deprecate +func (h *ProjectHandler) DeprecateFact(c *gin.Context) { + var req deprecateFactRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if err := h.db.DeprecateProjectFact(c.Param("id"), req.FactKey); err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"success": true}) +} + +type restoreFactRequest struct { + FactKey string `json:"fact_key" binding:"required"` + Confidence string `json:"confidence"` // 可选：confirmed | tentative，默认 tentative +} + +// RestoreFact POST /api/projects/:id/facts/restore +func (h *ProjectHandler) RestoreFact(c *gin.Context) { + var req restoreFactRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if err := h.db.RestoreProjectFact(c.Param("id"), req.FactKey, req.Confidence); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"success": true}) +} + +type createFactEdgeRequest struct { + SourceFactKey string `json:"source_fact_key" binding:"required"` + TargetFactKey string `json:"target_fact_key" binding:"required"` + EdgeType string `json:"edge_type" binding:"required"` + Confidence string `json:"confidence"` +} + +// ListFactEdges GET /api/projects/:id/fact-edges +func (h *ProjectHandler) ListFactEdges(c *gin.Context) { + projectID := c.Param("id") + edges, err := h.db.ListProjectFactEdgesByProject(projectID) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if edges == nil { + edges = []*database.ProjectFactEdge{} + } + c.JSON(http.StatusOK, edges) +} + +// CreateFactEdge POST /api/projects/:id/fact-edges +func (h *ProjectHandler) CreateFactEdge(c *gin.Context) { + projectID := c.Param("id") + var req createFactEdgeRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + edge, err := h.db.AddProjectFactEdge(projectID, database.ProjectFactEdgeInput{ + To: req.TargetFactKey, + Type: req.EdgeType, + Confidence: req.Confidence, + }, req.SourceFactKey, "") + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if f, err := h.db.GetProjectFactByKey(projectID, req.TargetFactKey); err == nil { + in, _ := h.db.ListIncomingProjectFactEdges(projectID, req.TargetFactKey) + f.Body = project.SyncBodyLinksSection(f.Body, in) + _, _ = h.db.UpsertProjectFact(f) + } + c.JSON(http.StatusOK, edge) +} + +// DeleteFactEdge DELETE /api/projects/:id/fact-edges/:edgeId +func (h *ProjectHandler) DeleteFactEdge(c *gin.Context) { + projectID := c.Param("id") + edgeID := c.Param("edgeId") + edge, err := h.db.GetProjectFactEdge(edgeID) + if err != nil || edge.ProjectID != projectID { + c.JSON(http.StatusNotFound, gin.H{"error": "边不存在"}) + return + } + if err := h.db.DeleteProjectFactEdge(edgeID); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if f, err := h.db.GetProjectFactByKey(projectID, edge.TargetFactKey); err == nil { + in, _ := h.db.ListIncomingProjectFactEdges(projectID, edge.TargetFactKey) + f.Body = project.SyncBodyLinksSection(f.Body, in) + _, _ = h.db.UpsertProjectFact(f) + } + c.JSON(http.StatusOK, gin.H{"success": true}) +} + +// PromoteAttackChain POST /api/projects/:id/promote-attack-chain/:conversationId +func (h *ProjectHandler) PromoteAttackChain(c *gin.Context) { + projectID := c.Param("id") + conversationID := c.Param("conversationId") + session, ok := security.CurrentSession(c) + if !ok || !h.db.UserCanAccessResource(session.UserID, session.Scope, "project", projectID) || + !h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", conversationID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问目标项目或来源对话"}) + return + } + result, err := attackchain.PromoteToProject(h.db, projectID, conversationID) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, result) +} diff --git a/internal/handler/project_context.go b/internal/handler/project_context.go new file mode 100644 index 00000000..cc0e6830 --- /dev/null +++ b/internal/handler/project_context.go @@ -0,0 +1,85 @@ +package handler + +import ( + "strings" + + "cyberstrike-ai/internal/project" + "go.uber.org/zap" +) + +// agentSessionContextBlock 注入会话工作目录与项目黑板（用于 system prompt 追加块）。 +// 用户输入由 message history 承载；压缩后由 summarization 摘要指令保留关键约束。 +func (h *AgentHandler) agentSessionContextBlock(conversationID string) string { + var parts []string + if ws := h.buildWorkspaceBlock(conversationID); ws != "" { + parts = append(parts, ws) + } + if bb := h.projectBlackboardBlock(conversationID); bb != "" { + parts = append(parts, bb) + } + return strings.Join(parts, "\n\n") +} + +func (h *AgentHandler) buildWorkspaceBlock(conversationID string) string { + if h == nil || h.config == nil { + return "" + } + conversationID = strings.TrimSpace(conversationID) + if conversationID == "" { + return "" + } + projectID := h.conversationProjectID(conversationID) + rel := project.WorkspaceRootDir(h.config.Agent.WorkspaceRootDir, projectID, conversationID) + abs, err := project.EnsureWorkspace(rel) + if err != nil { + if h.logger != nil { + h.logger.Warn("创建会话工作目录失败", + zap.String("conversationId", conversationID), + zap.String("projectId", projectID), + zap.String("path", rel), + zap.Error(err)) + } + return "" + } + return project.BuildWorkspaceBlock(abs) +} + +// projectBlackboardBlock 根据对话 ID 构建项目事实索引块（用于注入 system prompt）。 +func (h *AgentHandler) projectBlackboardBlock(conversationID string) string { + if h == nil || h.db == nil || h.config == nil { + return "" + } + if !h.config.Project.Enabled { + return "" + } + conversationID = strings.TrimSpace(conversationID) + if conversationID == "" { + return "" + } + projectID, err := h.db.GetConversationProjectID(conversationID) + if err != nil || projectID == "" { + return "" + } + block, err := project.BuildProjectBlackboardBlock(h.db, projectID, h.config.Project) + if err != nil { + h.logger.Warn("构建项目黑板索引失败", zap.String("conversationId", conversationID), zap.Error(err)) + return "" + } + return strings.TrimSpace(block) +} + +// conversationProjectID 返回对话绑定的项目 ID；未绑定或查询失败时返回空字符串。 +func (h *AgentHandler) conversationProjectID(conversationID string) string { + if h == nil || h.db == nil { + return "" + } + conversationID = strings.TrimSpace(conversationID) + if conversationID == "" { + return "" + } + projectID, err := h.db.GetConversationProjectID(conversationID) + if err != nil { + return "" + } + return strings.TrimSpace(projectID) +} diff --git a/internal/handler/project_resolve.go b/internal/handler/project_resolve.go new file mode 100644 index 00000000..88885838 --- /dev/null +++ b/internal/handler/project_resolve.go @@ -0,0 +1,18 @@ +package handler + +import ( + "strings" + + "cyberstrike-ai/internal/config" +) + +// effectiveProjectID 请求/队列显式项目优先，否则使用 config.project.default_project_id。 +func effectiveProjectID(cfg *config.Config, explicit string) string { + if pid := strings.TrimSpace(explicit); pid != "" { + return pid + } + if cfg != nil { + return strings.TrimSpace(cfg.Project.DefaultProjectID) + } + return "" +} diff --git a/internal/handler/rbac.go b/internal/handler/rbac.go new file mode 100644 index 00000000..f4784e5e --- /dev/null +++ b/internal/handler/rbac.go @@ -0,0 +1,429 @@ +package handler + +import ( + "fmt" + "net/http" + "strconv" + "strings" + + "cyberstrike-ai/internal/audit" + "cyberstrike-ai/internal/authctx" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +type RBACHandler struct { + db *database.DB + logger *zap.Logger + audit *audit.Service + auth *security.AuthManager +} + +func NewRBACHandler(db *database.DB, logger *zap.Logger) *RBACHandler { + return &RBACHandler{db: db, logger: logger} +} + +func (h *RBACHandler) SetAudit(s *audit.Service) { + h.audit = s +} + +func (h *RBACHandler) SetAuthManager(m *security.AuthManager) { + h.auth = m +} + +func (h *RBACHandler) Me(c *gin.Context) { + session, _ := security.CurrentSession(c) + resolvedScope := session.Scope + permissionScopes := session.PermissionScopes + if principal, ok := authctx.PrincipalFromContext(c.Request.Context()); ok { + resolvedScope = principal.Scope + permissionScopes = principal.PermissionScopes + } + c.JSON(http.StatusOK, gin.H{ + "user": gin.H{ + "id": session.UserID, + "username": session.Username, + "display_name": session.DisplayName, + }, + "roles": session.Roles, + "permissions": permissionKeys(session.Permissions), + "scope": resolvedScope, + "permission_scopes": permissionScopes, + }) +} + +func (h *RBACHandler) Metadata(c *gin.Context) { + roles, err := h.db.ListRBACRoles() + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + rolePermissions := map[string][]string{} + for _, role := range roles { + keys, _ := h.db.ListRBACRolePermissionKeys(role.ID) + rolePermissions[role.ID] = keys + } + c.JSON(http.StatusOK, gin.H{ + "permissions": security.PermissionCatalog, + "roles": roles, + "role_permissions": rolePermissions, + "scopes": []string{database.RBACScopeAll, database.RBACScopeAssigned, database.RBACScopeOwn}, + }) +} + +func (h *RBACHandler) ListRoles(c *gin.Context) { + roles, err := h.db.ListRBACRoles() + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + out := make([]gin.H, 0, len(roles)) + for _, role := range roles { + keys, _ := h.db.ListRBACRolePermissionKeys(role.ID) + out = append(out, gin.H{ + "id": role.ID, + "name": role.Name, + "description": role.Description, + "scope": role.Scope, + "is_system": role.IsSystem, + "permissions": keys, + "created_at": role.CreatedAt, + "updated_at": role.UpdatedAt, + }) + } + c.JSON(http.StatusOK, gin.H{"roles": out}) +} + +type upsertRBACRoleRequest struct { + ID string `json:"id"` + Name string `json:"name" binding:"required"` + Description string `json:"description"` + Scope string `json:"scope"` + Permissions []string `json:"permissions"` +} + +func validateRBACPermissionKeys(keys []string) error { + for _, key := range keys { + key = strings.TrimSpace(key) + if key == "" { + continue + } + if _, ok := security.PermissionCatalog[key]; !ok { + return fmt.Errorf("未知权限: %s", key) + } + } + return nil +} + +func (h *RBACHandler) CreateRole(c *gin.Context) { + var req upsertRBACRoleRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if err := validateRBACPermissionKeys(req.Permissions); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + role, err := h.db.UpsertRBACRole("", req.Name, req.Description, req.Scope, req.Permissions) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "rbac", "create_role", "创建平台角色", "role", role.ID, nil) + } + c.JSON(http.StatusOK, gin.H{"role": role}) +} + +func (h *RBACHandler) UpdateRole(c *gin.Context) { + id := strings.TrimSpace(c.Param("id")) + existing, err := h.db.GetRBACRoleByID(id) + if err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": "角色不存在"}) + return + } + if existing.IsSystem { + c.JSON(http.StatusBadRequest, gin.H{"error": "系统内置角色不可修改，请创建自定义角色"}) + return + } + var req upsertRBACRoleRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if err := validateRBACPermissionKeys(req.Permissions); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + role, err := h.db.UpsertRBACRole(id, req.Name, req.Description, req.Scope, req.Permissions) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if existing.IsSystem { + role.IsSystem = true + } + if h.audit != nil { + h.audit.RecordOK(c, "rbac", "update_role", "更新平台角色", "role", id, nil) + } + if h.auth != nil { + h.auth.RevokeAllSessions() + } + c.JSON(http.StatusOK, gin.H{"role": role}) +} + +func (h *RBACHandler) DeleteRole(c *gin.Context) { + id := strings.TrimSpace(c.Param("id")) + if err := h.db.DeleteRBACRole(id); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "rbac", "delete_role", "删除平台角色", "role", id, nil) + } + if h.auth != nil { + h.auth.RevokeAllSessions() + } + c.JSON(http.StatusOK, gin.H{"success": true}) +} + +func (h *RBACHandler) ListUsers(c *gin.Context) { + users, err := h.db.ListRBACUsers() + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + out := make([]gin.H, 0, len(users)) + for _, user := range users { + roleIDs, _ := h.db.ListRBACUserRoleIDs(user.ID) + out = append(out, gin.H{ + "id": user.ID, + "username": user.Username, + "display_name": user.DisplayName, + "enabled": user.Enabled, + "is_builtin": user.IsBuiltin, + "roles": roleIDs, + "created_at": user.CreatedAt, + "updated_at": user.UpdatedAt, + }) + } + c.JSON(http.StatusOK, gin.H{"users": out}) +} + +type createRBACUserRequest struct { + Username string `json:"username" binding:"required"` + DisplayName string `json:"display_name"` + Password string `json:"password" binding:"required"` + Enabled *bool `json:"enabled"` + Roles []string `json:"roles"` +} + +func (h *RBACHandler) CreateUser(c *gin.Context) { + var req createRBACUserRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if len(strings.TrimSpace(req.Password)) < 8 { + c.JSON(http.StatusBadRequest, gin.H{"error": "密码长度至少需要 8 位"}) + return + } + hash, err := security.HashPassword(req.Password) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + enabled := true + if req.Enabled != nil { + enabled = *req.Enabled + } + user, err := h.db.CreateRBACUser(req.Username, req.DisplayName, hash, enabled, req.Roles) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "rbac", "create_user", "创建平台用户", "user", user.ID, map[string]interface{}{"username": user.Username}) + } + c.JSON(http.StatusOK, gin.H{"user": user}) +} + +type updateRBACUserRequest struct { + DisplayName *string `json:"display_name"` + Password *string `json:"password"` + Enabled *bool `json:"enabled"` + Roles *[]string `json:"roles"` +} + +func (h *RBACHandler) UpdateUser(c *gin.Context) { + id := strings.TrimSpace(c.Param("id")) + user, err := h.db.GetRBACUserByID(id) + if err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": "用户不存在"}) + return + } + var req updateRBACUserRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + displayName := user.DisplayName + if req.DisplayName != nil { + displayName = *req.DisplayName + } + if err := h.db.UpdateRBACUser(id, displayName, req.Enabled, req.Roles); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if req.Password != nil && strings.TrimSpace(*req.Password) != "" { + if len(strings.TrimSpace(*req.Password)) < 8 { + c.JSON(http.StatusBadRequest, gin.H{"error": "密码长度至少需要 8 位"}) + return + } + hash, err := security.HashPassword(*req.Password) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if err := h.db.UpdateRBACUserPassword(id, hash); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + } + if h.audit != nil { + h.audit.RecordOK(c, "rbac", "update_user", "更新平台用户", "user", id, nil) + } + if h.auth != nil { + h.auth.RevokeUserSessions(id) + } + updated, _ := h.db.GetRBACUserByID(id) + c.JSON(http.StatusOK, gin.H{"user": updated}) +} + +func (h *RBACHandler) DeleteUser(c *gin.Context) { + id := strings.TrimSpace(c.Param("id")) + if err := h.db.DeleteRBACUser(id); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "rbac", "delete_user", "删除平台用户", "user", id, nil) + } + if h.auth != nil { + h.auth.RevokeUserSessions(id) + } + c.JSON(http.StatusOK, gin.H{"success": true}) +} + +type assignResourceRequest struct { + UserID string `json:"user_id" binding:"required"` + ResourceType string `json:"resource_type" binding:"required"` + ResourceID string `json:"resource_id"` + ResourceIDs []string `json:"resource_ids"` + AutoDetect bool `json:"auto_detect"` +} + +func (h *RBACHandler) AssignResource(c *gin.Context) { + var req assignResourceRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + resourceIDs := append([]string(nil), req.ResourceIDs...) + if strings.TrimSpace(req.ResourceID) != "" { + resourceIDs = append(resourceIDs, req.ResourceID) + } + if len(resourceIDs) == 0 { + c.JSON(http.StatusBadRequest, gin.H{"error": "至少需要一个资源 ID"}) + return + } + var created int64 + var detectedTypes map[string]string + var err error + if req.AutoDetect { + created, detectedTypes, err = h.db.AssignResourcesToUserAuto(req.UserID, resourceIDs) + } else { + created, err = h.db.AssignResourcesToUser(req.UserID, req.ResourceType, resourceIDs) + } + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if h.audit != nil { + for _, resourceID := range resourceIDs { + resourceType := req.ResourceType + if detectedTypes != nil { + resourceType = detectedTypes[strings.TrimSpace(resourceID)] + } + h.audit.RecordOK(c, "rbac", "assign_resource", "授权资源访问", resourceType, strings.TrimSpace(resourceID), map[string]interface{}{"user_id": req.UserID}) + } + } + c.JSON(http.StatusOK, gin.H{ + "success": true, + "requested": len(resourceIDs), + "created": created, + "skipped": int64(len(resourceIDs)) - created, + "detected_types": detectedTypes, + }) +} + +func (h *RBACHandler) ListResourceAssignments(c *gin.Context) { + rows, err := h.db.ListRBACResourceAssignments(c.Query("user_id")) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"assignments": rows}) +} + +func (h *RBACHandler) ListAssignableResources(c *gin.Context) { + limit, _ := strconv.Atoi(c.DefaultQuery("limit", "50")) + if limit <= 0 || limit > 50 { + limit = 50 + } + offset, _ := strconv.Atoi(c.DefaultQuery("offset", "0")) + if offset < 0 { + offset = 0 + } + resources, err := h.db.ListAssignableRBACResourcesPage(c.Query("type"), c.Query("q"), limit+1, offset) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + hasMore := len(resources) > limit + if hasMore { + resources = resources[:limit] + } + total, err := h.db.CountAssignableRBACResources(c.Query("type"), c.Query("q")) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{ + "resources": resources, + "has_more": hasMore, + "limit": limit, + "offset": offset, + "total": total, + }) +} + +func (h *RBACHandler) DeleteResourceAssignment(c *gin.Context) { + id := strings.TrimSpace(c.Param("id")) + assignment, err := h.db.DeleteRBACResourceAssignmentWithDetails(id) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "rbac", "delete_resource_assignment", "撤销资源授权", assignment.ResourceType, assignment.ResourceID, map[string]interface{}{ + "user_id": assignment.UserID, + "assignment_id": assignment.ID, + }) + } + c.JSON(http.StatusOK, gin.H{"success": true}) +} diff --git a/internal/handler/rbac_boundary_test.go b/internal/handler/rbac_boundary_test.go new file mode 100644 index 00000000..d01e12e4 --- /dev/null +++ b/internal/handler/rbac_boundary_test.go @@ -0,0 +1,350 @@ +package handler + +import ( + "bytes" + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "net/url" + "os" + "path/filepath" + "testing" + "time" + + "cyberstrike-ai/internal/authctx" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/mcp" + "cyberstrike-ai/internal/security" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +func withTempWorkingDir(t *testing.T) string { + t.Helper() + old, err := os.Getwd() + if err != nil { + t.Fatal(err) + } + dir := t.TempDir() + if err := os.Chdir(dir); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _ = os.Chdir(old) + }) + return dir +} + +func TestDetachedAgentContextRetainsPrincipalWithoutParentCancellation(t *testing.T) { + parent, cancel := context.WithCancel(context.Background()) + parent = authctx.WithPrincipal(parent, authctx.NewPrincipal("u1", "user", database.RBACScopeAssigned, map[string]bool{"agent:execute": true})) + detached := detachedAgentContext(parent) + cancel() + if err := detached.Err(); err != nil { + t.Fatalf("detached context inherited cancellation: %v", err) + } + principal, ok := authctx.PrincipalFromContext(detached) + if !ok || principal.UserID != "u1" || !principal.HasPermission("agent:execute") { + t.Fatalf("detached context lost principal: %#v, ok=%v", principal, ok) + } +} + +func TestPromoteAttackChainRequiresSourceConversationAccess(t *testing.T) { + db, err := database.NewDB(filepath.Join(t.TempDir(), "promote-rbac.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = db.Close() }) + project, _ := db.CreateProject(&database.Project{Name: "owned"}) + conversation, _ := db.CreateConversation("foreign", database.ConversationCreateMeta{}) + _ = db.SetResourceOwner("project", project.ID, "u1") + _ = db.SetResourceOwner("conversation", conversation.ID, "u2") + h := NewProjectHandler(db, zap.NewNop()) + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(security.ContextSessionKey, security.Session{UserID: "u1", Scope: database.RBACScopeOwn}) + c.Next() + }) + router.POST("/api/projects/:id/promote-attack-chain/:conversationId", h.PromoteAttackChain) + w := httptest.NewRecorder() + router.ServeHTTP(w, httptest.NewRequest(http.MethodPost, "/api/projects/"+project.ID+"/promote-attack-chain/"+conversation.ID, nil)) + if w.Code != http.StatusForbidden { + t.Fatalf("status = %d, want 403: %s", w.Code, w.Body.String()) + } +} + +func TestVulnerabilityCannotBeReparentedToForeignProject(t *testing.T) { + db, err := database.NewDB(filepath.Join(t.TempDir(), "vuln-reparent-rbac.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = db.Close() }) + owned, _ := db.CreateProject(&database.Project{Name: "owned"}) + foreign, _ := db.CreateProject(&database.Project{Name: "foreign"}) + _ = db.SetResourceOwner("project", owned.ID, "u1") + _ = db.SetResourceOwner("project", foreign.ID, "u2") + vulnerability, err := db.CreateVulnerability(&database.Vulnerability{Title: "v", Severity: "high", ProjectID: owned.ID}) + if err != nil { + t.Fatal(err) + } + _ = db.SetResourceOwner("vulnerability", vulnerability.ID, "u1") + h := NewVulnerabilityHandler(db, zap.NewNop()) + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(security.ContextSessionKey, security.Session{UserID: "u1", Scope: database.RBACScopeOwn}) + c.Next() + }) + router.PUT("/api/vulnerabilities/:id", h.UpdateVulnerability) + body, _ := json.Marshal(map[string]interface{}{"project_id": foreign.ID}) + w := httptest.NewRecorder() + router.ServeHTTP(w, httptest.NewRequest(http.MethodPut, "/api/vulnerabilities/"+vulnerability.ID, bytes.NewReader(body))) + if w.Code != http.StatusForbidden { + t.Fatalf("status = %d, want 403: %s", w.Code, w.Body.String()) + } +} + +func TestAgentTaskEndpointsFilterAndRejectForeignConversations(t *testing.T) { + gin.SetMode(gin.TestMode) + db, user := setupConversationRBACTest(t) + allowed, _ := db.CreateConversation("allowed", database.ConversationCreateMeta{}) + hidden, _ := db.CreateConversation("hidden", database.ConversationCreateMeta{}) + if err := db.AssignResourceToUser(user.ID, "conversation", allowed.ID); err != nil { + t.Fatal(err) + } + tasks := NewAgentTaskManager() + if _, err := tasks.StartTask(allowed.ID, "visible", func(error) {}); err != nil { + t.Fatal(err) + } + if _, err := tasks.StartTask(hidden.ID, "secret", func(error) {}); err != nil { + t.Fatal(err) + } + h := &AgentHandler{db: db, tasks: tasks, logger: zap.NewNop()} + + w := performAssignedHandler(user, http.MethodGet, "/api/agent-loop/tasks", nil, h.ListAgentTasks) + if w.Code != http.StatusOK { + t.Fatalf("list status = %d: %s", w.Code, w.Body.String()) + } + var response struct { + Tasks []*AgentTask `json:"tasks"` + } + if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil { + t.Fatal(err) + } + if len(response.Tasks) != 1 || response.Tasks[0].ConversationID != allowed.ID { + t.Fatalf("tasks = %#v, want only %s", response.Tasks, allowed.ID) + } + + w = performAssignedHandler(user, http.MethodPost, "/api/agent-loop/cancel", map[string]string{"conversationId": hidden.ID}, h.CancelAgentLoop) + if w.Code != http.StatusForbidden { + t.Fatalf("cancel status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String()) + } +} + +func TestChatUploadPathAuthorizationFollowsConversationAccess(t *testing.T) { + db, user := setupConversationRBACTest(t) + allowed, _ := db.CreateConversation("allowed", database.ConversationCreateMeta{}) + hidden, _ := db.CreateConversation("hidden", database.ConversationCreateMeta{}) + if err := db.AssignResourceToUser(user.ID, "conversation", allowed.ID); err != nil { + t.Fatal(err) + } + h := NewChatUploadsHandler(zap.NewNop(), db) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Set(security.ContextSessionKey, security.Session{UserID: user.ID, Scope: database.RBACScopeAssigned, Permissions: map[string]bool{"chat:write": true}}) + + if !h.pathAllowed(c, filepath.ToSlash(filepath.Join("2026-07-10", allowed.ID, "a.txt"))) { + t.Fatal("assigned conversation attachment should be accessible") + } + if h.pathAllowed(c, filepath.ToSlash(filepath.Join("2026-07-10", hidden.ID, "secret.txt"))) { + t.Fatal("foreign conversation attachment should be denied") + } + if h.pathAllowed(c, "2026-07-10/_manual/secret.txt") { + t.Fatal("unowned manual attachment should fail closed") + } +} + +func TestChatUploadsListIncludesAuthorizedProjectWorkspaceFiles(t *testing.T) { + withTempWorkingDir(t) + db, user := setupConversationRBACTest(t) + fsBase := t.TempDir() + workspaceBase := filepath.Join(fsBase, "workspace") + reductionBase := filepath.Join(fsBase, "reduction") + db.SetEinoConversationDirs("", "", reductionBase, workspaceBase) + allowedProject, _ := db.CreateProject(&database.Project{Name: "allowed"}) + hiddenProject, _ := db.CreateProject(&database.Project{Name: "hidden"}) + conversation, _ := db.CreateConversation("project conversation", database.ConversationCreateMeta{ProjectID: allowedProject.ID}) + if err := db.AssignResourceToUser(user.ID, "project", allowedProject.ID); err != nil { + t.Fatal(err) + } + allowedFile := filepath.Join(workspaceBase, "projects", allowedProject.ID, "csv", "assets.csv") + hiddenFile := filepath.Join(workspaceBase, "projects", hiddenProject.ID, "csv", "secret.csv") + for _, path := range []string{allowedFile, hiddenFile} { + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(path, []byte("name\nexample\n"), 0o644); err != nil { + t.Fatal(err) + } + } + + h := NewChatUploadsHandler(zap.NewNop(), db) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodGet, "/api/chat-uploads?source=workspace&pageSize=all&conversation="+conversation.ID, nil) + c.Set(security.ContextSessionKey, security.Session{UserID: user.ID, Scope: database.RBACScopeAssigned}) + h.List(c) + if w.Code != http.StatusOK { + t.Fatalf("status = %d, want 200: %s", w.Code, w.Body.String()) + } + var response struct { + Files []ChatUploadFileItem `json:"files"` + Total int `json:"total"` + } + if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil { + t.Fatal(err) + } + if response.Total != 1 || len(response.Files) != 1 { + t.Fatalf("files = %#v, total = %d, want only authorized workspace file", response.Files, response.Total) + } + got := response.Files[0] + if got.Source != chatUploadSourceWorkspace || got.Name != "assets.csv" || got.ProjectID != allowedProject.ID { + t.Fatalf("workspace file = %#v", got) + } + if got.ProjectName != allowedProject.Name { + t.Fatalf("projectName = %q, want %q", got.ProjectName, allowedProject.Name) + } + if got.ConversationID != conversation.ID { + t.Fatalf("conversationId = %q, want %q", got.ConversationID, conversation.ID) + } + if got.ConversationTitle != conversation.Title { + t.Fatalf("conversationTitle = %q, want %q", got.ConversationTitle, conversation.Title) + } + if got.AbsolutePath != allowedFile { + t.Fatalf("absolutePath = %q, want %q", got.AbsolutePath, allowedFile) + } + + w = httptest.NewRecorder() + c, _ = gin.CreateTestContext(w) + resolveURL := "/api/chat-uploads/path?kind=directory&path=__workspace__%2Fprojects%2F" + allowedProject.ID + c.Request = httptest.NewRequest(http.MethodGet, resolveURL, nil) + c.Set(security.ContextSessionKey, security.Session{UserID: user.ID, Scope: database.RBACScopeAssigned}) + h.ResolvePath(c) + if w.Code != http.StatusOK { + t.Fatalf("resolve status = %d, want 200: %s", w.Code, w.Body.String()) + } + var resolved struct { + AbsolutePath string `json:"absolutePath"` + IsDir bool `json:"isDir"` + } + if err := json.Unmarshal(w.Body.Bytes(), &resolved); err != nil { + t.Fatal(err) + } + wantDir := filepath.Join(workspaceBase, "projects", allowedProject.ID) + if !resolved.IsDir || resolved.AbsolutePath != wantDir { + t.Fatalf("resolved = %#v, want dir %q", resolved, wantDir) + } + + w = httptest.NewRecorder() + c, _ = gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodGet, "/api/chat-uploads/path?kind=directory&path=__workspace__%2Fprojects", nil) + c.Set(security.ContextSessionKey, security.Session{UserID: user.ID, Scope: database.RBACScopeAssigned}) + h.ResolvePath(c) + if w.Code != http.StatusOK { + t.Fatalf("resolve projects container status = %d, want 200: %s", w.Code, w.Body.String()) + } + if err := json.Unmarshal(w.Body.Bytes(), &resolved); err != nil { + t.Fatal(err) + } + wantContainer := filepath.Join(workspaceBase, "projects") + if !resolved.IsDir || resolved.AbsolutePath != wantContainer { + t.Fatalf("resolved container = %#v, want dir %q", resolved, wantContainer) + } + + for _, tc := range []struct { + path string + want string + }{ + {"__workspace__/", workspaceBase}, + {"__reduction__/", reductionBase}, + {"__conversation_artifact__/", db.ConversationArtifactsBaseDir()}, + } { + if err := os.MkdirAll(tc.want, 0o755); err != nil { + t.Fatal(err) + } + w = httptest.NewRecorder() + c, _ = gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodGet, "/api/chat-uploads/path?kind=directory&path="+url.QueryEscape(tc.path), nil) + c.Set(security.ContextSessionKey, security.Session{UserID: user.ID, Scope: database.RBACScopeAssigned}) + h.ResolvePath(c) + if w.Code != http.StatusOK { + t.Fatalf("resolve root %q status = %d, want 200: %s", tc.path, w.Code, w.Body.String()) + } + if err := json.Unmarshal(w.Body.Bytes(), &resolved); err != nil { + t.Fatal(err) + } + wantAbs, _ := filepath.Abs(tc.want) + if !resolved.IsDir || resolved.AbsolutePath != wantAbs { + t.Fatalf("resolved root %q = %#v, want dir %q", tc.path, resolved, wantAbs) + } + } +} + +func TestPrepareMultiAgentSessionRejectsForeignConversation(t *testing.T) { + db, user := setupConversationRBACTest(t) + hidden, _ := db.CreateConversation("hidden", database.ConversationCreateMeta{}) + h := &AgentHandler{db: db, logger: zap.NewNop()} + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Set(security.ContextSessionKey, security.Session{UserID: user.ID, Scope: database.RBACScopeAssigned, Permissions: map[string]bool{"chat:write": true}}) + + _, err := h.prepareMultiAgentSession(&ChatRequest{ConversationID: hidden.ID, Message: "write"}, c, "test") + if err == nil || err.Error() != "无权访问该对话" { + t.Fatalf("err = %v, want unauthorized conversation", err) + } +} + +func TestMonitorExecutionDetailRejectsForeignOwner(t *testing.T) { + db, user := setupConversationRBACTest(t) + for _, exec := range []*mcp.ToolExecution{ + {ID: "exec-allowed", ToolName: "allowed", Status: "completed", StartTime: time.Now(), OwnerUserID: user.ID}, + {ID: "exec-hidden", ToolName: "hidden", Status: "completed", StartTime: time.Now(), OwnerUserID: "another-user"}, + } { + if err := db.SaveToolExecution(exec); err != nil { + t.Fatal(err) + } + } + h := NewMonitorHandler(mcp.NewServerWithStorage(zap.NewNop(), db), nil, db, zap.NewNop()) + + request := func(id string) *httptest.ResponseRecorder { + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodGet, "/api/monitor/execution/"+id, nil) + c.Params = gin.Params{{Key: "id", Value: id}} + c.Set(security.ContextSessionKey, security.Session{UserID: user.ID, Scope: database.RBACScopeAssigned}) + h.GetExecution(c) + return w + } + if w := request("exec-hidden"); w.Code != http.StatusForbidden { + t.Fatalf("hidden status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String()) + } + if w := request("exec-allowed"); w.Code != http.StatusOK { + t.Fatalf("allowed status = %d, want %d: %s", w.Code, http.StatusOK, w.Body.String()) + } +} + +func performAssignedHandler(user *database.RBACUser, method, path string, body interface{}, handler gin.HandlerFunc) *httptest.ResponseRecorder { + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + var req *http.Request + if body == nil { + req = httptest.NewRequest(method, path, nil) + } else { + payload, _ := json.Marshal(body) + req = httptest.NewRequest(method, path, bytes.NewReader(payload)) + req.Header.Set("Content-Type", "application/json") + } + c.Request = req + c.Set(security.ContextSessionKey, security.Session{UserID: user.ID, Scope: database.RBACScopeAssigned}) + handler(c) + return w +} diff --git a/internal/handler/rbac_context.go b/internal/handler/rbac_context.go new file mode 100644 index 00000000..f5b19c2e --- /dev/null +++ b/internal/handler/rbac_context.go @@ -0,0 +1,12 @@ +package handler + +import "context" + +// detachedAgentContext lets a long-running Agent survive an SSE disconnect +// while retaining immutable request values such as the authenticated Principal. +func detachedAgentContext(parent context.Context) context.Context { + if parent == nil { + parent = context.Background() + } + return context.WithoutCancel(parent) +} diff --git a/internal/handler/rbac_test.go b/internal/handler/rbac_test.go new file mode 100644 index 00000000..e53cfddd --- /dev/null +++ b/internal/handler/rbac_test.go @@ -0,0 +1,209 @@ +package handler + +import ( + "bytes" + "cyberstrike-ai/internal/audit" + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + "encoding/json" + "net/http" + "net/http/httptest" + "path/filepath" + "testing" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +func TestRBACAssignResourceBatchIsAtomicAndLegacyCompatible(t *testing.T) { + gin.SetMode(gin.TestMode) + db, err := database.NewDB(filepath.Join(t.TempDir(), "rbac-handler.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + defer db.Close() + + user, err := db.CreateRBACUser("api-member", "API Member", "hash", true, nil) + if err != nil { + t.Fatal(err) + } + p1, _ := db.CreateProject(&database.Project{Name: "p1"}) + p2, _ := db.CreateProject(&database.Project{Name: "p2"}) + p3, _ := db.CreateProject(&database.Project{Name: "p3"}) + + h := NewRBACHandler(db, zap.NewNop()) + router := gin.New() + router.POST("/api/rbac/resource-assignments", h.AssignResource) + + batch := performRBACJSONRequest(t, router, map[string]interface{}{ + "user_id": user.ID, "resource_type": "project", "resource_ids": []string{p1.ID, p2.ID}, + }) + if batch.Code != http.StatusOK { + t.Fatalf("batch status = %d, body = %s", batch.Code, batch.Body.String()) + } + var batchBody map[string]interface{} + if err := json.Unmarshal(batch.Body.Bytes(), &batchBody); err != nil { + t.Fatal(err) + } + if batchBody["created"] != float64(2) { + t.Fatalf("batch response = %#v, want created=2", batchBody) + } + + invalid := performRBACJSONRequest(t, router, map[string]interface{}{ + "user_id": user.ID, "resource_type": "project", "resource_ids": []string{p3.ID, "missing"}, + }) + if invalid.Code != http.StatusBadRequest { + t.Fatalf("invalid status = %d, body = %s", invalid.Code, invalid.Body.String()) + } + rows, err := db.ListRBACResourceAssignments(user.ID) + if err != nil { + t.Fatal(err) + } + if len(rows) != 2 { + t.Fatalf("failed batch persisted partial data: %#v", rows) + } + + legacy := performRBACJSONRequest(t, router, map[string]interface{}{ + "user_id": user.ID, "resource_type": "project", "resource_id": p3.ID, + }) + if legacy.Code != http.StatusOK { + t.Fatalf("legacy status = %d, body = %s", legacy.Code, legacy.Body.String()) + } +} + +func TestRBACAssignResourceAutoDetectsActualType(t *testing.T) { + gin.SetMode(gin.TestMode) + db, err := database.NewDB(filepath.Join(t.TempDir(), "rbac-auto-detect.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + defer db.Close() + user, err := db.CreateRBACUser("auto-member", "Auto Member", "hash", true, nil) + if err != nil { + t.Fatal(err) + } + project, err := db.CreateProject(&database.Project{Name: "auto-project"}) + if err != nil { + t.Fatal(err) + } + h := NewRBACHandler(db, zap.NewNop()) + router := gin.New() + router.POST("/api/rbac/resource-assignments", h.AssignResource) + + response := performRBACJSONRequest(t, router, map[string]interface{}{ + "user_id": user.ID, "resource_type": "conversation", "resource_ids": []string{project.ID}, "auto_detect": true, + }) + if response.Code != http.StatusOK { + t.Fatalf("auto-detect status = %d, body = %s", response.Code, response.Body.String()) + } + rows, err := db.ListRBACResourceAssignments(user.ID) + if err != nil { + t.Fatal(err) + } + if len(rows) != 1 || rows[0].ResourceType != "project" || rows[0].ResourceID != project.ID { + t.Fatalf("auto-detected assignments = %#v, want project/%s", rows, project.ID) + } +} + +func TestRBACAssignableResourcesArePaged(t *testing.T) { + gin.SetMode(gin.TestMode) + db, err := database.NewDB(filepath.Join(t.TempDir(), "rbac-picker.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + defer db.Close() + + for _, name := range []string{"p1", "p2", "p3"} { + if _, err := db.CreateProject(&database.Project{Name: name}); err != nil { + t.Fatal(err) + } + } + h := NewRBACHandler(db, zap.NewNop()) + router := gin.New() + router.GET("/api/rbac/resources", h.ListAssignableResources) + + request := httptest.NewRequest(http.MethodGet, "/api/rbac/resources?type=project&limit=2&offset=0", nil) + recorder := httptest.NewRecorder() + router.ServeHTTP(recorder, request) + if recorder.Code != http.StatusOK { + t.Fatalf("status = %d, body = %s", recorder.Code, recorder.Body.String()) + } + var body struct { + Resources []database.RBACResourceOption `json:"resources"` + HasMore bool `json:"has_more"` + } + if err := json.Unmarshal(recorder.Body.Bytes(), &body); err != nil { + t.Fatal(err) + } + if len(body.Resources) != 2 || !body.HasMore { + t.Fatalf("page = %#v, has_more = %v; want two rows and another page", body.Resources, body.HasMore) + } +} + +func TestRBACDeleteResourceAssignmentAuditsTargetResource(t *testing.T) { + gin.SetMode(gin.TestMode) + db, err := database.NewDB(filepath.Join(t.TempDir(), "rbac-revoke-audit.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + defer db.Close() + user, err := db.CreateRBACUser("audit-member", "Audit Member", "hash", true, nil) + if err != nil { + t.Fatal(err) + } + project, err := db.CreateProject(&database.Project{Name: "audit-project"}) + if err != nil { + t.Fatal(err) + } + if _, err := db.AssignResourcesToUser(user.ID, "project", []string{project.ID}); err != nil { + t.Fatal(err) + } + assignments, err := db.ListRBACResourceAssignments(user.ID) + if err != nil || len(assignments) != 1 { + t.Fatalf("assignments = %#v, err = %v", assignments, err) + } + + h := NewRBACHandler(db, zap.NewNop()) + h.SetAudit(audit.NewService(db, &config.Config{}, zap.NewNop())) + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(security.ContextUsernameKey, "operator-user") + c.Next() + }) + router.DELETE("/api/rbac/resource-assignments/:id", h.DeleteResourceAssignment) + request := httptest.NewRequest(http.MethodDelete, "/api/rbac/resource-assignments/"+assignments[0].ID, nil) + recorder := httptest.NewRecorder() + router.ServeHTTP(recorder, request) + if recorder.Code != http.StatusOK { + t.Fatalf("status = %d, body = %s", recorder.Code, recorder.Body.String()) + } + + logs, err := db.ListAuditLogs(database.ListAuditLogsFilter{Category: "rbac", RelatedUserID: user.ID, Limit: 10}) + if err != nil { + t.Fatal(err) + } + if len(logs) != 1 { + t.Fatalf("audit logs = %#v, want one member-related revoke", logs) + } + log := logs[0] + if log.Action != "delete_resource_assignment" || log.Actor != "operator-user" || log.ResourceType != "project" || log.ResourceID != project.ID { + t.Fatalf("audit log = %#v", log) + } + if log.Detail["user_id"] != user.ID || log.Detail["assignment_id"] != assignments[0].ID { + t.Fatalf("audit detail = %#v", log.Detail) + } +} + +func performRBACJSONRequest(t *testing.T, router http.Handler, payload map[string]interface{}) *httptest.ResponseRecorder { + t.Helper() + body, err := json.Marshal(payload) + if err != nil { + t.Fatal(err) + } + req := httptest.NewRequest(http.MethodPost, "/api/rbac/resource-assignments", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + recorder := httptest.NewRecorder() + router.ServeHTTP(recorder, req) + return recorder +} diff --git a/internal/handler/robot.go b/internal/handler/robot.go new file mode 100644 index 00000000..45f0ec36 --- /dev/null +++ b/internal/handler/robot.go @@ -0,0 +1,1977 @@ +package handler + +import ( + "bytes" + "context" + "crypto/aes" + "crypto/cipher" + "crypto/rand" + "crypto/sha1" + "crypto/sha256" + "encoding/base32" + "encoding/base64" + "encoding/binary" + "encoding/json" + "encoding/xml" + "errors" + "fmt" + "io" + "net/http" + "sort" + "strconv" + "strings" + "sync" + "time" + + "cyberstrike-ai/internal/audit" + "cyberstrike-ai/internal/authctx" + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +const ( + robotCmdHelp = "帮助" + robotCmdList = "列表" + robotCmdListAlt = "对话列表" + robotCmdSwitch = "切换" + robotCmdContinue = "继续" + robotCmdNew = "新对话" + robotCmdClear = "清空" + robotCmdStatus = "状态" + robotCmdStop = "停止" + robotCmdRoles = "角色" + robotCmdRolesList = "角色列表" + robotCmdSwitchRole = "切换角色" + robotCmdModes = "模式" + robotCmdModesList = "模式列表" + robotCmdSwitchMode = "切换模式" + robotCmdDelete = "删除" + robotCmdVersion = "版本" + robotCmdProjects = "项目" + robotCmdProjectsList = "项目列表" + robotCmdBindProject = "绑定项目" + robotCmdNewProject = "新建项目" + robotCmdUnbindProject = "解除项目" + robotCmdBindUser = "绑定" + robotCmdUnbindUser = "解绑" + robotCmdIdentity = "身份" + robotCmdTask = "任务" + robotCmdRename = "重命名" + robotCmdPermissions = "权限" + robotCmdDoctor = "诊断" + robotCmdConfirm = "确认" + robotCmdCancel = "取消" + robotCmdVulnAlerts = "漏洞提醒" + robotBindingCodeTTL = 5 * time.Minute +) + +type robotPendingConfirmation struct { + Action string + Target string + ExpiresAt time.Time +} + +// RobotHandler 企业微信/钉钉/飞书等机器人回调处理 +type RobotHandler struct { + config *config.Config + db *database.DB + agentHandler *AgentHandler + logger *zap.Logger + mu sync.RWMutex + sessions map[string]string // key: "platform_userID", value: conversationID + sessionRoles map[string]string // key: "platform_userID", value: roleName（默认"默认"） + sessionModes map[string]string // key: "platform_userID", value: agent mode + cancelMu sync.Mutex // 保护 runningCancels + runningCancels map[string]context.CancelFunc // key: "platform_userID", 用于停止命令中断任务 + wecomReplay map[string]time.Time + pendingConfirmations map[string]robotPendingConfirmation + alertWake chan struct{} + audit *audit.Service +} + +// NewRobotHandler 创建机器人处理器 +func NewRobotHandler(cfg *config.Config, db *database.DB, agentHandler *AgentHandler, logger *zap.Logger) *RobotHandler { + return &RobotHandler{ + config: cfg, + db: db, + agentHandler: agentHandler, + logger: logger, + sessions: make(map[string]string), + sessionRoles: make(map[string]string), + sessionModes: make(map[string]string), + runningCancels: make(map[string]context.CancelFunc), + wecomReplay: make(map[string]time.Time), + pendingConfirmations: make(map[string]robotPendingConfirmation), + alertWake: make(chan struct{}, 1), + } +} + +func (h *RobotHandler) SetAudit(s *audit.Service) { + h.audit = s +} + +func (h *RobotHandler) acceptFreshWecomRequest(timestamp, nonce, signature string) bool { + unixSeconds, err := strconv.ParseInt(strings.TrimSpace(timestamp), 10, 64) + if err != nil { + return false + } + now := time.Now() + requestTime := time.Unix(unixSeconds, 0) + if requestTime.Before(now.Add(-5*time.Minute)) || requestTime.After(now.Add(5*time.Minute)) { + return false + } + key := strings.TrimSpace(timestamp) + "\x00" + strings.TrimSpace(nonce) + "\x00" + strings.TrimSpace(signature) + if strings.TrimSpace(nonce) == "" || strings.TrimSpace(signature) == "" { + return false + } + h.mu.Lock() + defer h.mu.Unlock() + for replayKey, seenAt := range h.wecomReplay { + if now.Sub(seenAt) > 10*time.Minute { + delete(h.wecomReplay, replayKey) + } + } + if _, exists := h.wecomReplay[key]; exists { + return false + } + h.wecomReplay[key] = now + return true +} + +// sessionKey 生成会话 key +func (h *RobotHandler) sessionKey(platform, userID string) string { + return platform + "_" + userID +} + +func normalizeRobotBindingCode(code string) string { + return strings.ToUpper(strings.ReplaceAll(strings.TrimSpace(code), "-", "")) +} + +func hashRobotBindingCode(code string) string { + sum := sha256.Sum256([]byte(normalizeRobotBindingCode(code))) + return fmt.Sprintf("%x", sum[:]) +} + +func (h *RobotHandler) resolveRobotAccess(platform, userID string) (*database.RBACAccess, error) { + if h.db == nil { + return nil, fmt.Errorf("机器人鉴权服务不可用") + } + authorization := h.config.Robots.AuthorizationFor(platform) + var access *database.RBACAccess + var err error + switch authorization.EffectiveMode() { + case config.RobotAuthModeUserBinding: + access, err = h.db.ResolveRobotRBACAccess(platform, userID) + case config.RobotAuthModeServiceAccount: + if !authorization.ExternalUserAllowed(userID) { + return nil, fmt.Errorf("机器人发送者不在服务账号白名单中") + } + access, err = h.db.ResolveRBACAccess(strings.TrimSpace(authorization.ServiceUserID)) + default: + return nil, fmt.Errorf("机器人鉴权模式无效") + } + if err != nil { + return nil, err + } + if !access.User.Enabled { + return nil, fmt.Errorf("绑定的平台账号已被禁用") + } + return access, nil +} + +func robotPrincipal(access *database.RBACAccess) authctx.Principal { + if access == nil { + return authctx.Principal{} + } + return authctx.NewPrincipalWithScopes(access.User.ID, access.User.Username, access.Scope, access.Permissions, access.PermissionScopes) +} + +func (h *RobotHandler) robotAccessDeniedMessage(platform string) string { + if h.config.Robots.AuthorizationFor(platform).EffectiveMode() == config.RobotAuthModeServiceAccount { + return "当前平台账号不在该机器人的服务账号白名单中，或服务账号不可用。" + } + return "当前平台账号尚未绑定 CyberStrikeAI 用户。请先在网页端生成绑定码，然后发送：绑定 XXXX-XXXX" +} + +func (h *RobotHandler) loadSessionBinding(sk string) (convID, role, agentMode string) { + if h.db == nil || strings.TrimSpace(sk) == "" { + return "", "", "" + } + binding, err := h.db.GetRobotSessionBinding(sk) + if err != nil { + h.logger.Warn("读取机器人会话绑定失败", zap.String("session_key", sk), zap.Error(err)) + return "", "", "" + } + if binding == nil { + return "", "", "" + } + return binding.ConversationID, binding.RoleName, binding.AgentMode +} + +func (h *RobotHandler) persistSessionBinding(sk, convID, role, agentMode string) { + if h.db == nil || strings.TrimSpace(sk) == "" || strings.TrimSpace(convID) == "" { + return + } + if err := h.db.UpsertRobotSessionBinding(sk, convID, role, agentMode); err != nil { + h.logger.Warn("写入机器人会话绑定失败", zap.String("session_key", sk), zap.Error(err)) + } +} + +func (h *RobotHandler) deleteSessionBinding(sk string) { + if h.db == nil || strings.TrimSpace(sk) == "" { + return + } + if err := h.db.DeleteRobotSessionBinding(sk); err != nil { + h.logger.Warn("删除机器人会话绑定失败", zap.String("session_key", sk), zap.Error(err)) + } +} + +// getOrCreateConversation 获取或创建当前会话，title 用于新对话的标题（取用户首条消息前50字） +func (h *RobotHandler) getOrCreateConversation(platform, userID, title string, access *database.RBACAccess) (convID string, isNew bool) { + sk := h.sessionKey(platform, userID) + h.mu.RLock() + convID = h.sessions[sk] + h.mu.RUnlock() + ownerID := access.User.ID + readScope := robotPrincipal(access).ScopeFor("chat:read") + if convID != "" && access.Permissions["chat:read"] && h.db.UserCanAccessResource(ownerID, readScope, "conversation", convID) { + return convID, false + } + if persistedConvID, persistedRole, persistedMode := h.loadSessionBinding(sk); strings.TrimSpace(persistedConvID) != "" { + if !access.Permissions["chat:read"] || !h.db.UserCanAccessResource(ownerID, readScope, "conversation", persistedConvID) { + h.deleteSessionBinding(sk) + } else { + // 会话绑定持久化：服务重启后也可恢复当前对话和角色。 + h.mu.Lock() + h.sessions[sk] = persistedConvID + if strings.TrimSpace(persistedRole) != "" { + h.sessionRoles[sk] = persistedRole + } + if strings.TrimSpace(persistedMode) != "" { + h.sessionModes[sk] = config.NormalizeAgentMode(persistedMode) + } + h.mu.Unlock() + return persistedConvID, false + } + } + t := strings.TrimSpace(title) + if t == "" { + t = "新对话 " + time.Now().Format("01-02 15:04") + } else { + t = safeTruncateString(t, 50) + } + meta := database.ConversationCreateMeta{Source: "robot:" + platform} + if !access.Permissions["chat:write"] { + return "", false + } + meta.ProjectID = effectiveProjectID(h.config, "") + if meta.ProjectID != "" && (!access.Permissions["project:read"] || !h.db.UserCanAccessResource(ownerID, robotPrincipal(access).ScopeFor("project:read"), "project", meta.ProjectID)) { + meta.ProjectID = "" + } + conv, err := h.db.CreateConversation(t, meta) + if err != nil { + h.logger.Warn("创建机器人会话失败", zap.Error(err)) + return "", false + } + convID = conv.ID + _ = h.db.SetResourceOwner("conversation", convID, ownerID) + h.mu.Lock() + role := h.sessionRoles[sk] + agentMode := h.sessionModes[sk] + h.sessions[sk] = convID + h.mu.Unlock() + if agentMode == "" { + agentMode = config.NormalizeRobotAgentMode(h.config.MultiAgent) + } + h.persistSessionBinding(sk, convID, role, agentMode) + return convID, true +} + +// setConversation 切换当前会话 +func (h *RobotHandler) setConversation(platform, userID, convID string) { + sk := h.sessionKey(platform, userID) + h.mu.Lock() + role := h.sessionRoles[sk] + agentMode := h.sessionModes[sk] + h.sessions[sk] = convID + h.mu.Unlock() + h.persistSessionBinding(sk, convID, role, agentMode) +} + +// getRole 获取当前用户使用的角色，未设置时返回"默认" +func (h *RobotHandler) getRole(platform, userID string) string { + sk := h.sessionKey(platform, userID) + h.mu.RLock() + role := h.sessionRoles[sk] + h.mu.RUnlock() + if strings.TrimSpace(role) != "" { + return role + } + if _, persistedRole, _ := h.loadSessionBinding(sk); strings.TrimSpace(persistedRole) != "" { + h.mu.Lock() + h.sessionRoles[sk] = persistedRole + h.mu.Unlock() + return persistedRole + } + return "默认" +} + +// setRole 设置当前用户使用的角色 +func (h *RobotHandler) setRole(platform, userID, roleName string) { + sk := h.sessionKey(platform, userID) + h.mu.Lock() + h.sessionRoles[sk] = roleName + convID := h.sessions[sk] + agentMode := h.sessionModes[sk] + h.mu.Unlock() + h.persistSessionBinding(sk, convID, roleName, agentMode) +} + +func (h *RobotHandler) getAgentMode(platform, userID string) string { + sk := h.sessionKey(platform, userID) + h.mu.RLock() + mode := h.sessionModes[sk] + h.mu.RUnlock() + if mode != "" { + return config.NormalizeAgentMode(mode) + } + if _, _, persistedMode := h.loadSessionBinding(sk); persistedMode != "" { + mode = config.NormalizeAgentMode(persistedMode) + h.mu.Lock() + h.sessionModes[sk] = mode + h.mu.Unlock() + return mode + } + return config.NormalizeRobotAgentMode(h.config.MultiAgent) +} + +func (h *RobotHandler) setAgentMode(platform, userID, mode string) { + sk := h.sessionKey(platform, userID) + mode = config.NormalizeAgentMode(mode) + h.mu.Lock() + h.sessionModes[sk] = mode + convID := h.sessions[sk] + role := h.sessionRoles[sk] + h.mu.Unlock() + h.persistSessionBinding(sk, convID, role, mode) +} + +// clearConversation 清空当前会话（切换到新对话） +func (h *RobotHandler) clearConversation(platform, userID string, access *database.RBACAccess) (newConvID string) { + title := "新对话 " + time.Now().Format("01-02 15:04") + meta := database.ConversationCreateMeta{Source: "robot:" + platform + ":new"} + meta.ProjectID = effectiveProjectID(h.config, "") + ownerID := access.User.ID + if meta.ProjectID != "" && (!access.Permissions["project:read"] || !h.db.UserCanAccessResource(ownerID, robotPrincipal(access).ScopeFor("project:read"), "project", meta.ProjectID)) { + meta.ProjectID = "" + } + conv, err := h.db.CreateConversation(title, meta) + if err != nil { + h.logger.Warn("创建新对话失败", zap.Error(err)) + return "" + } + _ = h.db.SetResourceOwner("conversation", conv.ID, ownerID) + h.setConversation(platform, userID, conv.ID) + return conv.ID +} + +// HandleMessage 处理用户输入，返回回复文本（供各平台 webhook 调用） +func (h *RobotHandler) HandleMessage(platform, userID, text string) (reply string) { + platform = strings.TrimSpace(platform) + userID = strings.TrimSpace(userID) + text = strings.TrimSpace(text) + if platform == "" { + platform = "unknown" + } + if userID == "" { + h.logger.Warn("机器人消息缺少用户标识，已拒绝处理", zap.String("platform", platform)) + return "无法识别发送者身份，请检查机器人事件订阅权限（需返回可用的用户 ID）。" + } + if text == "" { + return "请输入内容或发送「帮助」/ help 查看命令。" + } + + // 先尝试作为命令处理（支持中英文） + if cmdReply, ok := h.handleRobotCommand(platform, userID, text); ok { + return cmdReply + } + access, err := h.resolveRobotAccess(platform, userID) + if err != nil { + return h.robotAccessDeniedMessage(platform) + } + if !access.Permissions["agent:execute"] || !access.Permissions["chat:read"] || !access.Permissions["chat:write"] { + return "权限不足：机器人对话需要 agent:execute、chat:read 和 chat:write 权限。" + } + if h.audit != nil && h.config.Robots.AuthorizationFor(platform).EffectiveMode() == config.RobotAuthModeServiceAccount { + hint := sha256.Sum256([]byte(userID)) + h.audit.RecordSystem(audit.Entry{ + Category: "robot", Action: "service_account_execute", Result: "success", Actor: access.User.Username, + ResourceType: "robot_sender", ResourceID: platform + ":" + fmt.Sprintf("%x", hint[:4]), + Message: "白名单平台发送者使用机器人服务账号执行 Agent", + }) + } + + // 普通消息：走 Agent + convID, _ := h.getOrCreateConversation(platform, userID, text, access) + if convID == "" { + return "无法创建或获取对话，请稍后再试。" + } + // 若对话标题为「新对话 xx:xx」格式（由「新对话」命令创建），将标题更新为首条消息内容，与 Web 端体验一致 + if conv, err := h.db.GetConversation(convID); err == nil && strings.HasPrefix(conv.Title, "新对话 ") { + newTitle := safeTruncateString(text, 50) + if newTitle != "" { + _ = h.db.UpdateConversationTitle(convID, newTitle) + } + } + ctx, cancel := context.WithTimeout(context.Background(), h.robotMessageTimeout()) + sk := h.sessionKey(platform, userID) + h.cancelMu.Lock() + h.runningCancels[sk] = cancel + h.cancelMu.Unlock() + defer func() { + cancel() + h.cancelMu.Lock() + delete(h.runningCancels, sk) + h.cancelMu.Unlock() + }() + role := h.getRole(platform, userID) + agentMode := h.getAgentMode(platform, userID) + resp, newConvID, err := h.agentHandler.ProcessMessageForRobot(ctx, platform, robotPrincipal(access), convID, text, role, agentMode) + if err != nil { + h.logger.Warn("机器人 Agent 执行失败", zap.String("platform", platform), zap.String("userID", userID), zap.Error(err)) + if errors.Is(err, context.Canceled) { + return "任务已取消。" + } + if errors.Is(err, context.DeadlineExceeded) { + return "任务执行超时，请稍后重试或精简本次请求范围。" + } + return "处理失败: " + err.Error() + } + if newConvID != convID { + h.setConversation(platform, userID, newConvID) + } + return resp +} + +func (h *RobotHandler) robotMessageTimeout() time.Duration { + // 机器人整次消息处理超时（与单次工具超时 agent.tool_timeout_minutes 解耦）。 + return 10 * time.Hour +} + +func (h *RobotHandler) cmdHelp(platform, userID string) string { + access, _ := h.resolveRobotAccess(platform, userID) + can := func(permission string) bool { + return access != nil && access.Permissions[permission] + } + var b strings.Builder + b.WriteString("【CyberStrikeAI 机器人命令】\n\n") + b.WriteString("【通用 General】\n") + b.WriteString("· 帮助 / help — 显示本帮助\n") + b.WriteString("· 版本 / version — 显示当前版本号\n") + b.WriteString("· 绑定 <绑定码> / bind — 绑定网页端 RBAC 用户\n") + b.WriteString("· 解绑 / unbind — 请求解除账号绑定（需确认）\n") + b.WriteString("· 身份 / whoami — 显示平台发送者、鉴权模式及当前实际 RBAC 身份\n") + if can("chat:read") || can("chat:write") || can("chat:delete") { + b.WriteString("\n【对话 Conversation】\n") + if can("chat:read") { + b.WriteString("· 列表 / list — 列出所有对话标题与 ID\n· 切换 / switch — 指定对话继续\n· 状态 / status — 汇总当前选择\n· 任务 / task — 查看当前任务状态\n") + } + if can("chat:write") { + b.WriteString("· 新对话 / new；清空 / clear — 开启新对话\n· 重命名 <名称> / rename — 修改当前对话标题\n") + } + if can("chat:delete") { + b.WriteString("· 删除 / delete — 删除指定对话（需确认）\n") + } + } + if can("roles:read") { + b.WriteString("\n【角色 Role】\n· 角色 / roles — 列出所有可用角色\n· 角色 <名> / role — 切换当前角色\n") + } + if can("agent:execute") { + b.WriteString("\n【模式 Mode】\n· 模式 / modes — 列出对话模式与当前选择\n· 模式 <名称> / mode — 切换对话模式\n· 停止 / stop — 中断当前任务\n") + } + if can("vulnerability:read") { + b.WriteString("\n【漏洞提醒 Vulnerability alerts】\n· 漏洞提醒 — 查看订阅状态\n· 漏洞提醒 开启 / vuln alerts on — 开启提醒\n· 漏洞提醒 仅严重|高危以上|中危以上 / vuln alerts critical|high|medium — 设置最低级别\n· 漏洞提醒 关闭 / vuln alerts off — 关闭提醒\n") + } + b.WriteString("\n【诊断 Diagnostics】\n") + b.WriteString("· 权限 / permissions — 查看当前业务权限\n") + if can("config:read") { + b.WriteString("· 诊断 / doctor — 检查机器人关键配置状态\n") + } + b.WriteString("· 确认 / confirm；取消 / cancel — 处理高风险操作确认\n") + if h.projectsEnabled() && (can("project:read") || can("project:write")) { + b.WriteString("\n【项目 Project】\n") + if can("project:read") { + b.WriteString("· 项目 / projects — 列出所有项目\n") + } + if can("project:write") { + b.WriteString("· 新建项目 <名称> / new project — 创建并绑定当前对话\n· 绑定项目 / bind project — 绑定已有项目\n· 解除项目 / unbind project — 解除项目绑定\n") + } + } + b.WriteString("\n──────────────\n") + b.WriteString("除以上命令外，直接输入内容将发送给 AI 进行渗透测试/安全分析。") + return b.String() +} + +func (h *RobotHandler) projectsEnabled() bool { + return h.config != nil && h.config.Project.Enabled +} + +func (h *RobotHandler) resolveProjectByIDOrName(access *database.RBACAccess, idOrName string) (*database.Project, string) { + idOrName = strings.TrimSpace(idOrName) + if idOrName == "" { + return nil, "请指定项目 ID 或名称，例如：绑定项目 xxx-xxx" + } + ownerID := access.User.ID + scope := robotPrincipal(access).ScopeFor("project:read") + if p, err := h.db.GetProject(idOrName); err == nil { + if h.db.UserCanAccessResource(ownerID, scope, "project", p.ID) { + return p, "" + } + return nil, "项目不存在或无权访问。" + } + list, err := h.db.ListProjectsForAccess("", "", 200, 0, ownerID, scope) + if err != nil { + return nil, "查询项目失败: " + err.Error() + } + var matches []*database.Project + for _, p := range list { + if p.Name == idOrName { + matches = append(matches, p) + } + } + switch len(matches) { + case 0: + return nil, fmt.Sprintf("项目「%s」不存在。发送「项目」查看列表。", idOrName) + case 1: + return matches[0], "" + default: + var b strings.Builder + b.WriteString(fmt.Sprintf("名称「%s」匹配到多个项目，请使用 ID 绑定：\n", idOrName)) + for _, p := range matches { + b.WriteString(fmt.Sprintf("· %s\n ID: %s\n", p.Name, p.ID)) + } + return nil, strings.TrimSuffix(b.String(), "\n") + } +} + +func (h *RobotHandler) formatProjectLabel(projectID string) string { + if strings.TrimSpace(projectID) == "" { + return "未绑定" + } + if p, err := h.db.GetProject(projectID); err == nil { + return fmt.Sprintf("「%s」 (%s)", p.Name, p.ID) + } + return projectID +} + +func (h *RobotHandler) cmdProjects(platform, userID string) string { + if !h.projectsEnabled() { + return "项目功能未启用（config.project.enabled）。" + } + access, err := h.resolveRobotAccess(platform, userID) + if err != nil { + return "当前平台账号尚未绑定。" + } + list, err := h.db.ListProjectsForAccess("", "", 50, 0, access.User.ID, robotPrincipal(access).ScopeFor("project:read")) + if err != nil { + return "获取项目列表失败: " + err.Error() + } + if len(list) == 0 { + return "暂无项目。发送「新建项目 <名称>」创建并绑定到当前对话。" + } + var b strings.Builder + b.WriteString("【项目列表】\n") + for i, p := range list { + if i >= 20 { + b.WriteString("… 仅显示前 20 条\n") + break + } + status := p.Status + if status == "" { + status = "active" + } + b.WriteString(fmt.Sprintf("· %s [%s]\n ID: %s\n", p.Name, status, p.ID)) + } + return strings.TrimSuffix(b.String(), "\n") +} + +func (h *RobotHandler) cmdBindProject(platform, userID, idOrName string) string { + if !h.projectsEnabled() { + return "项目功能未启用（config.project.enabled）。" + } + access, err := h.resolveRobotAccess(platform, userID) + if err != nil { + return "当前平台账号尚未绑定。" + } + p, errMsg := h.resolveProjectByIDOrName(access, idOrName) + if p == nil { + return errMsg + } + convID, _ := h.getOrCreateConversation(platform, userID, "", access) + if convID == "" { + return "无法获取当前对话，请稍后再试。" + } + if err := h.db.SetConversationProjectID(convID, p.ID); err != nil { + return "绑定失败: " + err.Error() + } + return fmt.Sprintf("已将当前对话绑定到项目：「%s」\nID: %s", p.Name, p.ID) +} + +func (h *RobotHandler) cmdNewProject(platform, userID, name string) string { + if !h.projectsEnabled() { + return "项目功能未启用（config.project.enabled）。" + } + name = strings.TrimSpace(name) + if name == "" { + return "请指定项目名称，例如：新建项目 某目标渗透" + } + access, accessErr := h.resolveRobotAccess(platform, userID) + if accessErr != nil { + return "当前平台账号尚未绑定。" + } + p := &database.Project{Name: name, Status: "active"} + created, err := h.db.CreateProject(p) + if err != nil { + return "创建项目失败: " + err.Error() + } + _ = h.db.SetResourceOwner("project", created.ID, access.User.ID) + convID, _ := h.getOrCreateConversation(platform, userID, name, access) + if convID == "" { + return fmt.Sprintf("项目已创建：「%s」\nID: %s\n（绑定当前对话失败，请手动发送「绑定项目 %s」）", created.Name, created.ID, created.ID) + } + if err := h.db.SetConversationProjectID(convID, created.ID); err != nil { + return fmt.Sprintf("项目已创建：「%s」\nID: %s\n绑定失败: %s", created.Name, created.ID, err.Error()) + } + return fmt.Sprintf("已创建项目并绑定当前对话：「%s」\nID: %s", created.Name, created.ID) +} + +func (h *RobotHandler) cmdUnbindProject(platform, userID string) string { + if !h.projectsEnabled() { + return "项目功能未启用（config.project.enabled）。" + } + sk := h.sessionKey(platform, userID) + h.mu.RLock() + convID := h.sessions[sk] + h.mu.RUnlock() + if convID == "" { + if persistedConvID, _, _ := h.loadSessionBinding(sk); persistedConvID != "" { + convID = persistedConvID + } + } + access, err := h.resolveRobotAccess(platform, userID) + if err != nil { + return "当前平台账号尚未绑定。" + } + if !h.db.UserCanAccessResource(access.User.ID, robotPrincipal(access).ScopeFor("chat:write"), "conversation", convID) { + return "当前对话不存在或无权访问。" + } + if convID == "" { + return "当前没有进行中的对话，无需解除绑定。" + } + projectID, err := h.db.GetConversationProjectID(convID) + if err != nil { + return "获取对话项目失败: " + err.Error() + } + if strings.TrimSpace(projectID) == "" { + return "当前对话未绑定项目。" + } + if err := h.db.SetConversationProjectID(convID, ""); err != nil { + return "解除绑定失败: " + err.Error() + } + return "已解除当前对话的项目绑定。" +} + +func (h *RobotHandler) cmdList(platform, userID string) string { + access, err := h.resolveRobotAccess(platform, userID) + if err != nil { + return "当前平台账号尚未绑定。" + } + convs, err := h.db.ListConversationsForAccess(50, 0, "", "", "", access.User.ID, robotPrincipal(access).ScopeFor("chat:read")) + if err != nil { + return "获取对话列表失败: " + err.Error() + } + if len(convs) == 0 { + return "暂无对话。发送任意内容将自动创建新对话。" + } + var b strings.Builder + b.WriteString("【对话列表】\n") + for i, c := range convs { + if i >= 20 { + b.WriteString("… 仅显示前 20 条\n") + break + } + b.WriteString(fmt.Sprintf("· %s\n ID: %s\n", c.Title, c.ID)) + } + return strings.TrimSuffix(b.String(), "\n") +} + +func (h *RobotHandler) cmdSwitch(platform, userID, convID string) string { + if convID == "" { + return "请指定对话 ID，例如：切换 xxx-xxx-xxx" + } + access, accessErr := h.resolveRobotAccess(platform, userID) + if accessErr != nil { + return "当前平台账号尚未绑定。" + } + conv, err := h.db.GetConversation(convID) + if err != nil || !h.db.UserCanAccessResource(access.User.ID, robotPrincipal(access).ScopeFor("chat:read"), "conversation", convID) { + return "对话不存在或 ID 错误。" + } + h.setConversation(platform, userID, conv.ID) + return fmt.Sprintf("已切换到对话：「%s」\nID: %s", conv.Title, conv.ID) +} + +func (h *RobotHandler) cmdNew(platform, userID string) string { + access, err := h.resolveRobotAccess(platform, userID) + if err != nil { + return "当前平台账号尚未绑定。" + } + newID := h.clearConversation(platform, userID, access) + if newID == "" { + return "创建新对话失败，请重试。" + } + return "已开启新对话，可直接发送内容。" +} + +func (h *RobotHandler) cmdClear(platform, userID string) string { + return h.cmdNew(platform, userID) +} + +func (h *RobotHandler) cmdStop(platform, userID string) string { + sk := h.sessionKey(platform, userID) + h.cancelMu.Lock() + cancel, ok := h.runningCancels[sk] + if ok { + delete(h.runningCancels, sk) + cancel() + } + h.cancelMu.Unlock() + if !ok { + return "当前没有正在执行的任务。" + } + return "已停止当前任务。" +} + +func (h *RobotHandler) cmdStatus(platform, userID string) string { + convID := h.currentConversationID(platform, userID) + if convID == "" { + return fmt.Sprintf("【当前状态】\n当前对话: 无\n当前角色: %s\n当前模式: %s\n当前项目: 无\n\n发送任意内容将创建新对话。", h.getRole(platform, userID), robotAgentModeLabel(h.getAgentMode(platform, userID))) + } + access, err := h.resolveRobotAccess(platform, userID) + if err != nil { + return "当前平台账号尚未绑定。" + } + if !h.db.UserCanAccessResource(access.User.ID, robotPrincipal(access).ScopeFor("chat:read"), "conversation", convID) { + return "当前对话不存在或无权访问。" + } + conv, err := h.db.GetConversation(convID) + if err != nil { + return "当前对话 ID: " + convID + "（获取标题失败）" + } + role := h.getRole(platform, userID) + reply := fmt.Sprintf("【当前状态】\n当前对话: %s\n对话 ID: %s\n当前模式: %s\n当前角色: %s", conv.Title, conv.ID, robotAgentModeLabel(h.getAgentMode(platform, userID)), role) + if h.projectsEnabled() { + projectID, _ := h.db.GetConversationProjectID(conv.ID) + reply += "\n当前项目: " + h.formatProjectLabel(projectID) + } else { + reply += "\n当前项目: 未启用" + } + return reply +} + +func (h *RobotHandler) currentConversationID(platform, userID string) string { + sk := h.sessionKey(platform, userID) + h.mu.RLock() + convID := h.sessions[sk] + h.mu.RUnlock() + if convID != "" { + return convID + } + persistedConvID, persistedRole, persistedMode := h.loadSessionBinding(sk) + if persistedConvID == "" { + return "" + } + h.mu.Lock() + h.sessions[sk] = persistedConvID + h.sessionRoles[sk] = persistedRole + h.sessionModes[sk] = config.NormalizeAgentMode(persistedMode) + h.mu.Unlock() + return persistedConvID +} + +func (h *RobotHandler) cmdTask(platform, userID string) string { + convID := h.currentConversationID(platform, userID) + if convID == "" { + return "【任务状态】\n当前没有对话，也没有正在执行的任务。" + } + if h.agentHandler == nil || h.agentHandler.tasks == nil { + return "任务状态服务不可用。" + } + task := h.agentHandler.tasks.GetTaskSnapshot(convID) + if task == nil { + return "【任务状态】\n状态: 空闲\n当前没有正在执行的任务。" + } + elapsed := time.Since(task.StartedAt).Round(time.Second) + return fmt.Sprintf("【任务状态】\n状态: %s\n已运行: %s\n对话 ID: %s\n模式: %s\n可用操作: 停止 / stop", task.Status, elapsed, convID, robotAgentModeLabel(h.getAgentMode(platform, userID))) +} + +func (h *RobotHandler) cmdRename(platform, userID, title string) string { + title = strings.TrimSpace(title) + if title == "" { + return "请指定新标题，例如：重命名 外网资产排查" + } + title = safeTruncateString(title, 100) + convID := h.currentConversationID(platform, userID) + if convID == "" { + return "当前没有对话，无法重命名。" + } + access, err := h.resolveRobotAccess(platform, userID) + if err != nil || !h.db.UserCanAccessResource(access.User.ID, robotPrincipal(access).ScopeFor("chat:write"), "conversation", convID) { + return "当前对话不存在或无权修改。" + } + if err := h.db.UpdateConversationTitle(convID, title); err != nil { + return "重命名失败: " + err.Error() + } + h.recordRobotCommandAudit(access, platform, "conversation_rename", "conversation", convID, "机器人重命名当前对话") + return fmt.Sprintf("已将当前对话重命名为：「%s」", title) +} + +func (h *RobotHandler) cmdPermissions(platform, userID string) string { + access, err := h.resolveRobotAccess(platform, userID) + if err != nil { + return h.robotAccessDeniedMessage(platform) + } + allowed := func(permission string) string { + if access.Permissions[permission] { + return "允许" + } + return "不允许" + } + return fmt.Sprintf("【当前权限】\n执行 Agent: %s\n读取对话: %s\n编辑对话: %s\n删除对话: %s\n读取角色: %s\n读取项目: %s\n编辑项目: %s\n资源范围: %s", allowed("agent:execute"), allowed("chat:read"), allowed("chat:write"), allowed("chat:delete"), allowed("roles:read"), allowed("project:read"), allowed("project:write"), access.Scope) +} + +func (h *RobotHandler) cmdDoctor() string { + configured := func(ok bool) string { + if ok { + return "正常" + } + return "未配置" + } + enabled := func(ok bool) string { + if ok { + return "已启用" + } + return "已关闭" + } + enabledInternalTools := 0 + for _, tool := range h.config.Security.Tools { + if tool.Enabled { + enabledInternalTools++ + } + } + enabledExternal := 0 + for _, server := range h.config.ExternalMCP.Servers { + if server.ExternalMCPEnable && !server.Disabled { + enabledExternal++ + } + } + return fmt.Sprintf("【配置诊断】\n主模型: %s\nEino 多代理: %s\n内置 MCP 工具: %d/%d 个已启用\nHTTP MCP 服务: %s\n外部 MCP: %d 个已启用\n知识库: %s\n项目功能: %s\n说明: 内置工具不依赖 HTTP MCP 服务；此命令只检查配置，不主动探测外部服务。", configured(strings.TrimSpace(h.config.OpenAI.Model) != "" && strings.TrimSpace(h.config.OpenAI.BaseURL) != ""), enabled(h.config.MultiAgent.Enabled), enabledInternalTools, len(h.config.Security.Tools), enabled(h.config.MCP.Enabled), enabledExternal, enabled(h.config.Knowledge.Enabled), enabled(h.config.Project.Enabled)) +} + +func (h *RobotHandler) recordRobotCommandAudit(access *database.RBACAccess, platform, action, resourceType, resourceID, message string) { + if h.audit == nil || access == nil { + return + } + h.audit.RecordSystem(audit.Entry{Category: "robot", Action: action, Result: "success", Actor: access.User.Username, ResourceType: resourceType, ResourceID: resourceID, Message: message + "（" + platform + "）"}) +} + +func (h *RobotHandler) cmdRoles() string { + if h.config.Roles == nil || len(h.config.Roles) == 0 { + return "暂无可用角色。" + } + names := make([]string, 0, len(h.config.Roles)) + for name, role := range h.config.Roles { + if role.Enabled { + names = append(names, name) + } + } + if len(names) == 0 { + return "暂无可用角色。" + } + sort.Slice(names, func(i, j int) bool { + if names[i] == "默认" { + return true + } + if names[j] == "默认" { + return false + } + return names[i] < names[j] + }) + var b strings.Builder + b.WriteString("【角色列表】\n") + for _, name := range names { + role := h.config.Roles[name] + desc := role.Description + if desc == "" { + desc = "无描述" + } + b.WriteString(fmt.Sprintf("· %s — %s\n", name, desc)) + } + return strings.TrimSuffix(b.String(), "\n") +} + +func (h *RobotHandler) cmdSwitchRole(platform, userID, roleName string) string { + if roleName == "" { + return "请指定角色名称，例如：角色 渗透测试" + } + if h.config.Roles == nil { + return "暂无可用角色。" + } + role, exists := h.config.Roles[roleName] + if !exists { + return fmt.Sprintf("角色「%s」不存在。发送「角色」查看可用角色。", roleName) + } + if !role.Enabled { + return fmt.Sprintf("角色「%s」已禁用。", roleName) + } + h.setRole(platform, userID, roleName) + return fmt.Sprintf("已切换到角色：「%s」\n%s", roleName, role.Description) +} + +func robotAgentModeLabel(mode string) string { + switch config.NormalizeAgentMode(mode) { + case "deep": + return "Deep" + case "plan_execute": + return "Plan-Execute" + case "supervisor": + return "Supervisor" + default: + return "Eino 单代理" + } +} + +func parseRobotAgentMode(input string) (string, bool) { + switch strings.ToLower(strings.TrimSpace(input)) { + case "eino_single", "eino-single", "single", "单代理", "eino单代理", "eino 单代理": + return "eino_single", true + case "deep": + return "deep", true + case "plan_execute", "plan-execute", "planexecute", "pe": + return "plan_execute", true + case "supervisor", "super", "sv": + return "supervisor", true + default: + return "", false + } +} + +func (h *RobotHandler) cmdModes(platform, userID string) string { + current := h.getAgentMode(platform, userID) + multiStatus := "可用" + if h.config == nil || !h.config.MultiAgent.Enabled { + multiStatus = "不可用（需在系统设置中启用 Eino 多代理）" + } + return fmt.Sprintf("【对话模式】\n· Eino 单代理 — 可用\n· Deep — %s\n· Plan-Execute — %s\n· Supervisor — %s\n\n当前模式: %s\n切换示例：模式 deep", multiStatus, multiStatus, multiStatus, robotAgentModeLabel(current)) +} + +func (h *RobotHandler) cmdSwitchMode(platform, userID, input string) string { + mode, ok := parseRobotAgentMode(input) + if !ok { + return fmt.Sprintf("不支持的对话模式「%s」。发送「模式」查看可用模式。", strings.TrimSpace(input)) + } + if mode != "eino_single" && (h.config == nil || !h.config.MultiAgent.Enabled) { + return fmt.Sprintf("无法切换到 %s：请先在系统设置中启用 Eino 多代理。", robotAgentModeLabel(mode)) + } + h.setAgentMode(platform, userID, mode) + return fmt.Sprintf("已切换对话模式：%s\n后续消息和新对话将使用该模式。", robotAgentModeLabel(mode)) +} + +func (h *RobotHandler) cmdDelete(platform, userID, convID string) string { + if convID == "" { + return "请指定对话 ID，例如：删除 xxx-xxx-xxx" + } + access, err := h.resolveRobotAccess(platform, userID) + if err != nil { + return "当前平台账号尚未绑定。" + } + if !h.db.UserCanAccessResource(access.User.ID, robotPrincipal(access).ScopeFor("chat:delete"), "conversation", convID) { + return "对话不存在或无权访问。" + } + h.setPendingConfirmation(platform, userID, "delete_conversation", convID) + return fmt.Sprintf("⚠️ 即将删除对话 ID: %s\n此操作不可撤销。请在 2 分钟内发送「确认」继续，或发送「取消」。", convID) +} + +func (h *RobotHandler) executeDelete(platform, userID, convID string) string { + access, err := h.resolveRobotAccess(platform, userID) + if err != nil || !h.db.UserCanAccessResource(access.User.ID, robotPrincipal(access).ScopeFor("chat:delete"), "conversation", convID) { + return "对话不存在或无权删除。" + } + sk := h.sessionKey(platform, userID) + h.mu.RLock() + currentConvID := h.sessions[sk] + h.mu.RUnlock() + if convID == currentConvID { + // 删除当前对话时，先清空会话绑定 + h.mu.Lock() + delete(h.sessions, sk) + delete(h.sessionRoles, sk) + delete(h.sessionModes, sk) + h.mu.Unlock() + h.deleteSessionBinding(sk) + } + if h.agentHandler != nil { + h.agentHandler.CancelRunningTaskForConversation(convID) + } + if err := h.db.DeleteConversation(convID); err != nil { + return "删除失败: " + err.Error() + } + h.recordRobotCommandAudit(access, platform, "conversation_delete", "conversation", convID, "机器人删除对话") + return fmt.Sprintf("已删除对话 ID: %s", convID) +} + +func (h *RobotHandler) cmdVersion() string { + v := h.config.Version + if v == "" { + v = "未知" + } + return "CyberStrikeAI " + v +} + +func (h *RobotHandler) cmdIdentity(platform, userID string) string { + authorization := h.config.Robots.AuthorizationFor(platform) + mode := authorization.EffectiveMode() + modeLabel := "逐用户绑定（user_binding）" + if mode == config.RobotAuthModeServiceAccount { + modeLabel = "专用服务账号（service_account）" + } + var b strings.Builder + b.WriteString("【机器人身份】\n") + b.WriteString("平台：" + platform + "\n") + b.WriteString("发送者 ID：" + userID + "\n") + b.WriteString("鉴权模式：" + modeLabel + "\n") + + access, err := h.resolveRobotAccess(platform, userID) + if err != nil { + if mode == config.RobotAuthModeServiceAccount { + b.WriteString("鉴权状态：拒绝（发送者不在白名单中，或服务账号不可用）") + } else { + b.WriteString("鉴权状态：未绑定\n") + b.WriteString("操作提示：请在 Web 端生成绑定码，然后发送“绑定 XXXX-XXXX”") + } + return b.String() + } + + name := strings.TrimSpace(access.User.DisplayName) + if name == "" { + name = access.User.Username + } + roleNames := make([]string, 0, len(access.Roles)) + for _, role := range access.Roles { + roleNames = append(roleNames, role.Name) + } + if len(roleNames) == 0 { + roleNames = append(roleNames, "未分配角色") + } + b.WriteString("鉴权状态：已授权\n") + b.WriteString("实际身份：" + name + " (" + access.User.Username + ")\n") + b.WriteString("RBAC User ID：" + access.User.ID + "\n") + b.WriteString("平台角色：" + strings.Join(roleNames, "、") + "\n") + b.WriteString("资源范围：" + access.Scope + "\n") + b.WriteString(fmt.Sprintf("有效权限：%d 项", len(access.Permissions))) + return b.String() +} + +func robotCommandPermission(text string) (string, bool) { + switch { + case text == robotCmdHelp || text == "help" || text == "？" || text == "?", text == robotCmdVersion || text == "version", text == robotCmdIdentity || text == "whoami": + return "", true + case text == robotCmdList || text == robotCmdListAlt || text == "list", + strings.HasPrefix(text, robotCmdSwitch+" "), strings.HasPrefix(text, robotCmdContinue+" "), + strings.HasPrefix(text, "switch "), strings.HasPrefix(text, "continue "), + text == robotCmdStatus || text == "status", text == robotCmdTask || text == "task": + return "chat:read", true + case text == robotCmdNew || text == "new", text == robotCmdClear || text == "clear", + strings.HasPrefix(text, robotCmdRename+" "), strings.HasPrefix(text, "rename "): + return "chat:write", true + case strings.HasPrefix(text, robotCmdDelete+" "), strings.HasPrefix(text, "delete "): + return "chat:delete", true + case text == robotCmdStop || text == "stop": + return "agent:execute", true + case text == robotCmdRoles || text == robotCmdRolesList || text == "roles", + strings.HasPrefix(text, robotCmdRoles+" "), strings.HasPrefix(text, robotCmdSwitchRole+" "), strings.HasPrefix(text, "role "): + return "roles:read", true + case text == robotCmdModes || text == robotCmdModesList || text == "modes", + strings.HasPrefix(text, robotCmdModes+" "), strings.HasPrefix(text, robotCmdSwitchMode+" "), strings.HasPrefix(text, "mode "): + return "agent:execute", true + case text == robotCmdPermissions || text == "permissions": + return "", true + case text == robotCmdConfirm || text == "confirm", text == robotCmdCancel || text == "cancel": + return "", true + case text == robotCmdDoctor || text == "doctor": + return "config:read", true + case text == robotCmdProjects || text == robotCmdProjectsList || text == "projects": + return "project:read", true + case text == robotCmdVulnAlerts || strings.HasPrefix(text, robotCmdVulnAlerts+" "), + text == "vuln alerts" || strings.HasPrefix(text, "vuln alerts "): + return "vulnerability:read", true + case text == robotCmdUnbindProject || text == "unbind project", + strings.HasPrefix(text, robotCmdNewProject+" "), strings.HasPrefix(text, "new project "), + strings.HasPrefix(text, robotCmdBindProject+" "), strings.HasPrefix(text, "bind project "): + return "project:write", true + default: + return "", false + } +} + +func (h *RobotHandler) cmdBindUser(platform, userID, code string) string { + if h.config.Robots.AuthorizationFor(platform).EffectiveMode() != config.RobotAuthModeUserBinding { + return "该机器人使用受控服务账号模式，不接受用户绑定。" + } + code = normalizeRobotBindingCode(code) + if code == "" { + return "请提供绑定码，例如：绑定 ABCD-1234" + } + user, err := h.db.ConsumeRobotBindingCode(platform, userID, hashRobotBindingCode(code)) + if err != nil { + return "绑定失败：绑定码无效、已使用或已过期。请在网页端重新生成。" + } + // Never carry an old synthetic-owner conversation into the RBAC identity. + sk := h.sessionKey(platform, userID) + h.mu.Lock() + delete(h.sessions, sk) + delete(h.sessionRoles, sk) + delete(h.sessionModes, sk) + h.mu.Unlock() + h.deleteSessionBinding(sk) + name := strings.TrimSpace(user.DisplayName) + if name == "" { + name = user.Username + } + if h.audit != nil { + hint := sha256.Sum256([]byte(userID)) + h.audit.RecordSystem(audit.Entry{ + Category: "auth", Action: "robot_bind", Result: "success", Actor: user.Username, + ResourceType: "robot_binding", ResourceID: platform + ":" + fmt.Sprintf("%x", hint[:4]), Message: "机器人平台账号绑定成功", + }) + } + return fmt.Sprintf("绑定成功，当前身份：%s。后续操作将实时使用该用户的 RBAC 权限。", name) +} + +func (h *RobotHandler) cmdUnbindUser(platform, userID string) string { + if h.config.Robots.AuthorizationFor(platform).EffectiveMode() != config.RobotAuthModeUserBinding { + return "该机器人使用受控服务账号模式，无需用户解绑。" + } + _, accessErr := h.resolveRobotAccess(platform, userID) + if accessErr != nil { + return "当前平台账号尚未绑定。" + } + h.setPendingConfirmation(platform, userID, "unbind_user", "") + return "⚠️ 即将解除当前平台账号绑定。请在 2 分钟内发送「确认」继续，或发送「取消」。" +} + +func (h *RobotHandler) executeUnbindUser(platform, userID string) string { + access, accessErr := h.resolveRobotAccess(platform, userID) + if accessErr != nil { + return "当前平台账号尚未绑定。" + } + if err := h.db.DeleteRobotIdentityBinding(platform, userID); err != nil { + return "解绑失败，请稍后重试。" + } + sk := h.sessionKey(platform, userID) + h.mu.Lock() + delete(h.sessions, sk) + delete(h.sessionRoles, sk) + delete(h.sessionModes, sk) + h.mu.Unlock() + h.deleteSessionBinding(sk) + if h.audit != nil { + hint := sha256.Sum256([]byte(userID)) + h.audit.RecordSystem(audit.Entry{ + Category: "auth", Action: "robot_unbind", Result: "success", Actor: access.User.Username, + ResourceType: "robot_binding", ResourceID: platform + ":" + fmt.Sprintf("%x", hint[:4]), Message: "机器人平台账号解绑成功", + }) + } + return "已解除当前平台账号与 CyberStrikeAI 用户的绑定。" +} + +func (h *RobotHandler) setPendingConfirmation(platform, userID, action, target string) { + sk := h.sessionKey(platform, userID) + now := time.Now() + h.mu.Lock() + for key, pending := range h.pendingConfirmations { + if now.After(pending.ExpiresAt) { + delete(h.pendingConfirmations, key) + } + } + h.pendingConfirmations[sk] = robotPendingConfirmation{Action: action, Target: target, ExpiresAt: now.Add(2 * time.Minute)} + h.mu.Unlock() +} + +func (h *RobotHandler) cmdConfirm(platform, userID string) string { + sk := h.sessionKey(platform, userID) + h.mu.Lock() + pending, ok := h.pendingConfirmations[sk] + delete(h.pendingConfirmations, sk) + h.mu.Unlock() + if !ok || time.Now().After(pending.ExpiresAt) { + return "当前没有待确认操作，或确认已超时。" + } + switch pending.Action { + case "delete_conversation": + return h.executeDelete(platform, userID, pending.Target) + case "unbind_user": + return h.executeUnbindUser(platform, userID) + default: + return "待确认操作无效，已取消。" + } +} + +func (h *RobotHandler) cmdCancelConfirmation(platform, userID string) string { + sk := h.sessionKey(platform, userID) + h.mu.Lock() + _, ok := h.pendingConfirmations[sk] + delete(h.pendingConfirmations, sk) + h.mu.Unlock() + if !ok { + return "当前没有待确认操作。" + } + return "已取消待确认操作。" +} + +// handleRobotCommand 处理机器人内置命令；若匹配到命令返回 (回复内容, true)，否则返回 ("", false) +func (h *RobotHandler) handleRobotCommand(platform, userID, text string) (string, bool) { + if (strings.HasPrefix(text, robotCmdBindUser+" ") || strings.HasPrefix(text, "bind ")) && !strings.HasPrefix(text, "bind project ") { + parts := strings.SplitN(text, " ", 2) + return h.cmdBindUser(platform, userID, strings.TrimSpace(parts[1])), true + } + if text == robotCmdUnbindUser || text == "unbind" { + return h.cmdUnbindUser(platform, userID), true + } + if permission, recognized := robotCommandPermission(text); recognized && permission != "" { + access, err := h.resolveRobotAccess(platform, userID) + if err != nil { + return h.robotAccessDeniedMessage(platform), true + } + if !access.Permissions[permission] { + return fmt.Sprintf("权限不足：缺少 %s 权限。", permission), true + } + } + switch { + case text == robotCmdVulnAlerts || text == "vuln alerts": + return h.cmdVulnerabilityAlerts(platform, userID, ""), true + case strings.HasPrefix(text, robotCmdVulnAlerts+" "): + return h.cmdVulnerabilityAlerts(platform, userID, strings.TrimSpace(text[len(robotCmdVulnAlerts)+1:])), true + case strings.HasPrefix(text, "vuln alerts "): + return h.cmdVulnerabilityAlerts(platform, userID, strings.TrimSpace(text[len("vuln alerts "):])), true + case text == robotCmdHelp || text == "help" || text == "？" || text == "?": + return h.cmdHelp(platform, userID), true + case text == robotCmdIdentity || text == "whoami": + return h.cmdIdentity(platform, userID), true + case text == robotCmdConfirm || text == "confirm": + return h.cmdConfirm(platform, userID), true + case text == robotCmdCancel || text == "cancel": + return h.cmdCancelConfirmation(platform, userID), true + case text == robotCmdList || text == robotCmdListAlt || text == "list": + return h.cmdList(platform, userID), true + case strings.HasPrefix(text, robotCmdSwitch+" ") || strings.HasPrefix(text, robotCmdContinue+" ") || strings.HasPrefix(text, "switch ") || strings.HasPrefix(text, "continue "): + var id string + switch { + case strings.HasPrefix(text, robotCmdSwitch+" "): + id = strings.TrimSpace(text[len(robotCmdSwitch)+1:]) + case strings.HasPrefix(text, robotCmdContinue+" "): + id = strings.TrimSpace(text[len(robotCmdContinue)+1:]) + case strings.HasPrefix(text, "switch "): + id = strings.TrimSpace(text[7:]) + default: + id = strings.TrimSpace(text[9:]) + } + return h.cmdSwitch(platform, userID, id), true + case text == robotCmdNew || text == "new": + return h.cmdNew(platform, userID), true + case text == robotCmdClear || text == "clear": + return h.cmdClear(platform, userID), true + case text == robotCmdStatus || text == "status": + return h.cmdStatus(platform, userID), true + case text == robotCmdTask || text == "task": + return h.cmdTask(platform, userID), true + case strings.HasPrefix(text, robotCmdRename+" ") || strings.HasPrefix(text, "rename "): + var title string + if strings.HasPrefix(text, robotCmdRename+" ") { + title = strings.TrimSpace(text[len(robotCmdRename)+1:]) + } else { + title = strings.TrimSpace(text[len("rename "):]) + } + return h.cmdRename(platform, userID, title), true + case text == robotCmdStop || text == "stop": + return h.cmdStop(platform, userID), true + case text == robotCmdRoles || text == robotCmdRolesList || text == "roles": + return h.cmdRoles(), true + case strings.HasPrefix(text, robotCmdRoles+" ") || strings.HasPrefix(text, robotCmdSwitchRole+" ") || strings.HasPrefix(text, "role "): + var roleName string + switch { + case strings.HasPrefix(text, robotCmdRoles+" "): + roleName = strings.TrimSpace(text[len(robotCmdRoles)+1:]) + case strings.HasPrefix(text, robotCmdSwitchRole+" "): + roleName = strings.TrimSpace(text[len(robotCmdSwitchRole)+1:]) + default: + roleName = strings.TrimSpace(text[5:]) + } + return h.cmdSwitchRole(platform, userID, roleName), true + case text == robotCmdModes || text == robotCmdModesList || text == "modes": + return h.cmdModes(platform, userID), true + case strings.HasPrefix(text, robotCmdModes+" ") || strings.HasPrefix(text, robotCmdSwitchMode+" ") || strings.HasPrefix(text, "mode "): + var mode string + switch { + case strings.HasPrefix(text, robotCmdModes+" "): + mode = strings.TrimSpace(text[len(robotCmdModes)+1:]) + case strings.HasPrefix(text, robotCmdSwitchMode+" "): + mode = strings.TrimSpace(text[len(robotCmdSwitchMode)+1:]) + default: + mode = strings.TrimSpace(text[5:]) + } + return h.cmdSwitchMode(platform, userID, mode), true + case text == robotCmdPermissions || text == "permissions": + return h.cmdPermissions(platform, userID), true + case text == robotCmdDoctor || text == "doctor": + return h.cmdDoctor(), true + case strings.HasPrefix(text, robotCmdDelete+" ") || strings.HasPrefix(text, "delete "): + var convID string + if strings.HasPrefix(text, robotCmdDelete+" ") { + convID = strings.TrimSpace(text[len(robotCmdDelete)+1:]) + } else { + convID = strings.TrimSpace(text[7:]) + } + return h.cmdDelete(platform, userID, convID), true + case text == robotCmdVersion || text == "version": + return h.cmdVersion(), true + case text == robotCmdProjects || text == robotCmdProjectsList || text == "projects": + return h.cmdProjects(platform, userID), true + case text == robotCmdUnbindProject || text == "unbind project": + return h.cmdUnbindProject(platform, userID), true + case strings.HasPrefix(text, robotCmdNewProject+" ") || strings.HasPrefix(text, "new project "): + var name string + if strings.HasPrefix(text, robotCmdNewProject+" ") { + name = strings.TrimSpace(text[len(robotCmdNewProject)+1:]) + } else { + name = strings.TrimSpace(text[len("new project "):]) + } + return h.cmdNewProject(platform, userID, name), true + case strings.HasPrefix(text, robotCmdBindProject+" ") || strings.HasPrefix(text, "bind project "): + var idOrName string + if strings.HasPrefix(text, robotCmdBindProject+" ") { + idOrName = strings.TrimSpace(text[len(robotCmdBindProject)+1:]) + } else { + idOrName = strings.TrimSpace(text[len("bind project "):]) + } + return h.cmdBindProject(platform, userID, idOrName), true + default: + return "", false + } +} + +// —————— 企业微信 —————— + +// wecomXML 企业微信回调 XML（明文模式下的简化结构；加密模式需先解密再解析） +type wecomXML struct { + ToUserName string `xml:"ToUserName"` + FromUserName string `xml:"FromUserName"` + CreateTime int64 `xml:"CreateTime"` + MsgType string `xml:"MsgType"` + Content string `xml:"Content"` + MsgID string `xml:"MsgId"` + AgentID int64 `xml:"AgentID"` + Encrypt string `xml:"Encrypt"` // 加密模式下消息在此 +} + +// wecomReplyXML 被动回复 XML（仅用于兼容，当前使用手动构造 XML） +type wecomReplyXML struct { + XMLName xml.Name `xml:"xml"` + ToUserName string `xml:"ToUserName"` + FromUserName string `xml:"FromUserName"` + CreateTime int64 `xml:"CreateTime"` + MsgType string `xml:"MsgType"` + Content string `xml:"Content"` +} + +// wecomRequireToken 企业微信回调必须配置 Token；未配置时拒绝请求，防止未授权触发 Agent。 +func (h *RobotHandler) wecomRequireToken(c *gin.Context) (string, bool) { + token := strings.TrimSpace(h.config.Robots.Wecom.Token) + if token == "" { + h.logger.Warn("企业微信已启用但未配置 token，已拒绝回调（请在配置中设置 robots.wecom.token）") + c.String(http.StatusForbidden, "") + return "", false + } + return token, true +} + +// HandleWecomGET 企业微信 URL 校验（GET） +func (h *RobotHandler) HandleWecomGET(c *gin.Context) { + if !h.config.Robots.Wecom.Enabled { + c.String(http.StatusNotFound, "") + return + } + token, ok := h.wecomRequireToken(c) + if !ok { + return + } + // Gin 的 Query() 会自动 URL 解码，拿到的就是正确的 base64 字符串 + echostr := c.Query("echostr") + msgSignature := c.Query("msg_signature") + timestamp := c.Query("timestamp") + nonce := c.Query("nonce") + + // 验证签名：将 token、timestamp、nonce、echostr 四个参数排序后拼接计算 SHA1 + signature := h.signWecomRequest(token, timestamp, nonce, echostr) + if signature != msgSignature { + h.logger.Warn("企业微信 URL 验证签名失败", zap.String("expected", msgSignature), zap.String("got", signature)) + c.String(http.StatusBadRequest, "invalid signature") + return + } + + if echostr == "" { + c.String(http.StatusBadRequest, "missing echostr") + return + } + + // 如果配置了 EncodingAESKey，说明是加密模式，需要解密 echostr + if h.config.Robots.Wecom.EncodingAESKey != "" { + decrypted, err := wecomDecrypt(h.config.Robots.Wecom.EncodingAESKey, echostr) + if err != nil { + h.logger.Warn("企业微信 echostr 解密失败", zap.Error(err)) + c.String(http.StatusBadRequest, "decrypt failed") + return + } + c.String(http.StatusOK, string(decrypted)) + return + } + + // 明文模式直接返回 echostr + c.String(http.StatusOK, echostr) +} + +// signWecomRequest 生成企业微信请求签名 +// 企业微信签名算法：将 token、timestamp、nonce、echostr 四个值排序后拼接成字符串，再计算 SHA1 +func (h *RobotHandler) signWecomRequest(token, timestamp, nonce, echostr string) string { + strs := []string{token, timestamp, nonce, echostr} + sort.Strings(strs) + s := strings.Join(strs, "") + hash := sha1.Sum([]byte(s)) + return fmt.Sprintf("%x", hash) +} + +// wecomDecrypt 企业微信消息解密（AES-256-CBC，PKCS7，明文格式：16字节随机+4字节长度+消息+corpID） +func wecomDecrypt(encodingAESKey, encryptedB64 string) ([]byte, error) { + key, err := base64.StdEncoding.DecodeString(encodingAESKey + "=") + if err != nil { + return nil, err + } + if len(key) != 32 { + return nil, fmt.Errorf("encoding_aes_key 解码后应为 32 字节") + } + ciphertext, err := base64.StdEncoding.DecodeString(encryptedB64) + if err != nil { + return nil, err + } + block, err := aes.NewCipher(key) + if err != nil { + return nil, err + } + iv := key[:16] + mode := cipher.NewCBCDecrypter(block, iv) + if len(ciphertext)%aes.BlockSize != 0 { + return nil, fmt.Errorf("密文长度不是块大小的倍数") + } + plain := make([]byte, len(ciphertext)) + mode.CryptBlocks(plain, ciphertext) + // 去除 PKCS7 填充 + n := int(plain[len(plain)-1]) + if n < 1 || n > 32 { + return nil, fmt.Errorf("无效的 PKCS7 填充") + } + plain = plain[:len(plain)-n] + // 企业微信格式：16 字节随机 + 4 字节长度(大端) + 消息 + corpID + if len(plain) < 20 { + return nil, fmt.Errorf("明文过短") + } + msgLen := binary.BigEndian.Uint32(plain[16:20]) + if int(20+msgLen) > len(plain) { + return nil, fmt.Errorf("消息长度越界") + } + return plain[20 : 20+msgLen], nil +} + +// wecomEncrypt 企业微信消息加密（AES-256-CBC，PKCS7，明文格式：16字节随机+4字节长度+消息+corpID） +func wecomEncrypt(encodingAESKey, message, corpID string) (string, error) { + key, err := base64.StdEncoding.DecodeString(encodingAESKey + "=") + if err != nil { + return "", err + } + if len(key) != 32 { + return "", fmt.Errorf("encoding_aes_key 解码后应为 32 字节") + } + // 构造明文：16 字节随机 + 4 字节长度 (大端) + 消息 + corpID + random := make([]byte, 16) + if _, err := rand.Read(random); err != nil { + // 降级方案：使用时间戳生成随机数 + for i := range random { + random[i] = byte(time.Now().UnixNano() % 256) + } + } + msgLen := len(message) + msgBytes := []byte(message) + corpBytes := []byte(corpID) + plain := make([]byte, 16+4+msgLen+len(corpBytes)) + copy(plain[:16], random) + binary.BigEndian.PutUint32(plain[16:20], uint32(msgLen)) + copy(plain[20:20+msgLen], msgBytes) + copy(plain[20+msgLen:], corpBytes) + // PKCS7 填充 + padding := aes.BlockSize - len(plain)%aes.BlockSize + pad := bytes.Repeat([]byte{byte(padding)}, padding) + plain = append(plain, pad...) + // AES-256-CBC 加密 + block, err := aes.NewCipher(key) + if err != nil { + return "", err + } + iv := key[:16] + ciphertext := make([]byte, len(plain)) + mode := cipher.NewCBCEncrypter(block, iv) + mode.CryptBlocks(ciphertext, plain) + return base64.StdEncoding.EncodeToString(ciphertext), nil +} + +// HandleWecomPOST 企业微信消息回调（POST），支持明文与加密模式 +func (h *RobotHandler) HandleWecomPOST(c *gin.Context) { + if !h.config.Robots.Wecom.Enabled { + h.logger.Debug("企业微信机器人未启用，跳过请求") + c.String(http.StatusOK, "") + return + } + // 从 URL 获取签名参数（加密模式回复时需要用到） + timestamp := c.Query("timestamp") + nonce := c.Query("nonce") + msgSignature := c.Query("msg_signature") + + // 先读取请求体，后续解析/签名验证都会用到 + bodyRaw, err := io.ReadAll(c.Request.Body) + if err != nil { + h.logger.Warn("企业微信 POST 读取请求体失败", zap.Error(err)) + c.String(http.StatusOK, "") + return + } + h.logger.Debug("企业微信 POST 收到请求", zap.String("body", string(bodyRaw))) + + // 验证请求签名防止伪造。企业微信签名算法同 URL 验证，使用 token、timestamp、nonce、 Encrypt 四个字段。 + // 启用企业微信时必须配置 token 并校验签名，避免未授权请求触发 Agent。 + token, ok := h.wecomRequireToken(c) + if !ok { + return + } + if msgSignature == "" { + h.logger.Warn("企业微信 POST 缺少签名，已拒绝（需确保回调携带 msg_signature）") + c.String(http.StatusOK, "") + return + } + var tmp wecomXML + if err := xml.Unmarshal(bodyRaw, &tmp); err != nil { + h.logger.Warn("企业微信 POST 签名验证前解析 XML 失败", zap.Error(err)) + c.String(http.StatusOK, "") + return + } + expected := h.signWecomRequest(token, timestamp, nonce, tmp.Encrypt) + if expected != msgSignature { + h.logger.Warn("企业微信 POST 签名验证失败", zap.String("expected", expected), zap.String("got", msgSignature)) + c.String(http.StatusOK, "") + return + } + if !h.acceptFreshWecomRequest(timestamp, nonce, msgSignature) { + h.logger.Warn("企业微信 POST 时间戳过期或请求重放，已拒绝") + c.String(http.StatusOK, "") + return + } + + var body wecomXML + if err := xml.Unmarshal(bodyRaw, &body); err != nil { + h.logger.Warn("企业微信 POST 解析 XML 失败", zap.Error(err)) + c.String(http.StatusOK, "") + return + } + h.logger.Debug("企业微信 XML 解析成功", zap.String("ToUserName", body.ToUserName), zap.String("FromUserName", body.FromUserName), zap.String("MsgType", body.MsgType), zap.String("Content", body.Content), zap.String("Encrypt", body.Encrypt)) + + // 保存企业 ID（用于明文模式回复） + enterpriseID := body.ToUserName + + // 配置了 EncodingAESKey 时必须走加密消息，拒绝明文 XML 绕过 + if strings.TrimSpace(h.config.Robots.Wecom.EncodingAESKey) != "" && strings.TrimSpace(body.Encrypt) == "" { + h.logger.Warn("企业微信已配置加密模式但收到明文消息，已拒绝") + c.String(http.StatusOK, "") + return + } + + // 加密模式：先解密再解析内层 XML + if body.Encrypt != "" && h.config.Robots.Wecom.EncodingAESKey != "" { + h.logger.Debug("企业微信进入加密模式解密流程") + decrypted, err := wecomDecrypt(h.config.Robots.Wecom.EncodingAESKey, body.Encrypt) + if err != nil { + h.logger.Warn("企业微信消息解密失败", zap.Error(err)) + c.String(http.StatusOK, "") + return + } + h.logger.Debug("企业微信解密成功", zap.String("decrypted", string(decrypted))) + if err := xml.Unmarshal(decrypted, &body); err != nil { + h.logger.Warn("企业微信解密后 XML 解析失败", zap.Error(err)) + c.String(http.StatusOK, "") + return + } + h.logger.Debug("企业微信内层 XML 解析成功", zap.String("FromUserName", body.FromUserName), zap.String("Content", body.Content)) + } + + tenantKey := strings.TrimSpace(enterpriseID) + if tenantKey == "" { + tenantKey = strings.TrimSpace(h.config.Robots.Wecom.CorpID) + } + if tenantKey == "" { + tenantKey = "default" + } + rawUserID := strings.TrimSpace(body.FromUserName) + replyUserID := rawUserID + userID := "" + if rawUserID != "" { + userID = "t:" + tenantKey + "|u:" + rawUserID + } + text := strings.TrimSpace(body.Content) + if userID == "" { + h.logger.Warn("企业微信消息缺少可用用户标识，已忽略") + c.String(http.StatusOK, "success") + return + } + + // 限制回复内容长度（企业微信限制 2048 字节） + maxReplyLen := 2000 + limitReply := func(s string) string { + if len(s) > maxReplyLen { + return s[:maxReplyLen] + "\n\n（内容过长，已截断）" + } + return s + } + + if body.MsgType != "text" { + h.logger.Debug("企业微信收到非文本消息", zap.String("MsgType", body.MsgType)) + h.sendWecomReply(c, replyUserID, enterpriseID, limitReply("暂仅支持文本消息，请发送文字。"), timestamp, nonce) + return + } + + // 文本消息：先判断是否为内置命令（如 帮助/列表/新对话 等），这类命令处理很快，可以直接走被动回复，避免依赖主动发送 API。 + if cmdReply, ok := h.handleRobotCommand("wecom", userID, text); ok { + h.logger.Debug("企业微信收到命令消息，走被动回复", zap.String("userID", userID), zap.String("text", text)) + h.sendWecomReply(c, replyUserID, enterpriseID, limitReply(cmdReply), timestamp, nonce) + return + } + + h.logger.Debug("企业微信开始处理消息（异步 AI）", zap.String("userID", userID), zap.String("text", text)) + + // 企业微信被动回复有 5 秒超时限制，而 AI 调用通常超过该时长。 + // 这里采用推荐做法：立即返回 success（或空串），然后通过主动发送接口推送完整回复。 + c.String(http.StatusOK, "success") + + // 异步处理消息并通过企业微信主动消息接口发送结果 + go func() { + reply := h.HandleMessage("wecom", userID, text) + reply = limitReply(reply) + h.logger.Debug("企业微信消息处理完成", zap.String("userID", userID), zap.String("reply", reply)) + // 调用企业微信 API 主动发送消息 + h.sendWecomMessageViaAPI(rawUserID, enterpriseID, reply) + }() +} + +// sendWecomReply 发送企业微信回复（加密模式自动加密） +// 参数：toUser=用户 ID, fromUser=企业 ID（明文模式）/CorpID（加密模式）, content=回复内容，timestamp/nonce=请求参数 +func (h *RobotHandler) sendWecomReply(c *gin.Context, toUser, fromUser, content, timestamp, nonce string) { + // 加密模式：判断 EncodingAESKey 是否配置 + if h.config.Robots.Wecom.EncodingAESKey != "" { + // 加密模式使用 CorpID 进行加密 + corpID := h.config.Robots.Wecom.CorpID + if corpID == "" { + h.logger.Warn("企业微信加密模式缺少 CorpID 配置") + c.String(http.StatusOK, "") + return + } + + // 构造完整的明文 XML 回复（格式严格按企业微信文档要求） + plainResp := fmt.Sprintf(` + + +%d + + +`, toUser, fromUser, time.Now().Unix(), content) + + encrypted, err := wecomEncrypt(h.config.Robots.Wecom.EncodingAESKey, plainResp, corpID) + if err != nil { + h.logger.Warn("企业微信回复加密失败", zap.Error(err)) + c.String(http.StatusOK, "") + return + } + // 使用请求中的 timestamp/nonce 生成签名（企业微信要求回复时使用与请求相同的 timestamp 和 nonce） + msgSignature := h.signWecomRequest(h.config.Robots.Wecom.Token, timestamp, nonce, encrypted) + + h.logger.Debug("企业微信发送加密回复", + zap.String("Encrypt", encrypted[:50]+"..."), + zap.String("MsgSignature", msgSignature), + zap.String("TimeStamp", timestamp), + zap.String("Nonce", nonce)) + + // 加密模式仅返回 4 个核心字段（企业微信官方要求） + xmlResp := fmt.Sprintf(``, encrypted, msgSignature, timestamp, nonce) + // also log the final response body so we can cross-check with the + // network traffic or developer console + h.logger.Debug("企业微信加密回复包", zap.String("xml", xmlResp)) + // for additional confidence, decrypt the payload ourselves and log it + if dec, err2 := wecomDecrypt(h.config.Robots.Wecom.EncodingAESKey, encrypted); err2 == nil { + h.logger.Debug("企业微信加密回复解密检查", zap.String("plain", string(dec))) + } else { + h.logger.Warn("企业微信加密回复解密检查失败", zap.Error(err2)) + } + + // 使用 c.Writer.Write 直接写入响应，避免 c.String 的转义问题 + c.Writer.WriteHeader(http.StatusOK) + // use text/xml as that's what WeCom examples show + c.Writer.Header().Set("Content-Type", "text/xml; charset=utf-8") + _, _ = c.Writer.Write([]byte(xmlResp)) + h.logger.Debug("企业微信加密回复已发送") + return + } + + // 明文模式 + h.logger.Debug("企业微信发送明文回复", zap.String("ToUserName", toUser), zap.String("FromUserName", fromUser), zap.String("Content", content[:50]+"...")) + + // 手动构造 XML 响应（使用 CDATA 包裹所有字段，并包含 AgentID） + xmlResp := fmt.Sprintf(` + + +%d + + +`, toUser, fromUser, time.Now().Unix(), content) + + // log the exact plaintext response for debugging + h.logger.Debug("企业微信明文回复包", zap.String("xml", xmlResp)) + + // use text/xml as recommended by WeCom docs + c.Header("Content-Type", "text/xml; charset=utf-8") + c.String(http.StatusOK, xmlResp) + h.logger.Debug("企业微信明文回复已发送") +} + +// —————— 测试接口（需登录，用于验证机器人逻辑，无需钉钉/飞书客户端） —————— + +// CreateRobotBindingCode creates a short-lived, single-use secret for the +// currently authenticated RBAC user. Only its hash is persisted. +func (h *RobotHandler) CreateRobotBindingCode(c *gin.Context) { + session, ok := security.CurrentSession(c) + if !ok || strings.TrimSpace(session.UserID) == "" { + c.JSON(http.StatusUnauthorized, gin.H{"error": "未授权访问"}) + return + } + random := make([]byte, 5) + if _, err := rand.Read(random); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "生成绑定码失败"}) + return + } + raw := base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(random) + code := raw[:4] + "-" + raw[4:] + expiresAt := time.Now().Add(robotBindingCodeTTL) + if err := h.db.CreateRobotBindingCode(session.UserID, hashRobotBindingCode(code), expiresAt); err != nil { + h.logger.Warn("创建机器人绑定码失败", zap.String("user_id", session.UserID), zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": "生成绑定码失败"}) + return + } + if h.audit != nil { + h.audit.Record(c, audit.Entry{Category: "auth", Action: "robot_binding_code_create", Result: "success", ResourceType: "user", ResourceID: session.UserID, Message: "生成机器人一次性绑定码"}) + } + c.Header("Cache-Control", "no-store") + c.JSON(http.StatusOK, gin.H{ + "code": code, "expires_at": expiresAt.UTC().Format(time.RFC3339), "expires_in_seconds": int(robotBindingCodeTTL.Seconds()), + }) +} + +func (h *RobotHandler) ListMyRobotBindings(c *gin.Context) { + session, ok := security.CurrentSession(c) + if !ok { + c.JSON(http.StatusUnauthorized, gin.H{"error": "未授权访问"}) + return + } + bindings, err := h.db.ListRobotUserBindings(session.UserID) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "获取机器人绑定失败"}) + return + } + items := make([]gin.H, 0, len(bindings)) + for _, binding := range bindings { + sum := sha256.Sum256([]byte(binding.ExternalUserID)) + items = append(items, gin.H{ + "id": binding.ID, "platform": binding.Platform, "external_user_hint": fmt.Sprintf("%x", sum[:4]), + "enabled": binding.Enabled, "created_at": binding.CreatedAt, "updated_at": binding.UpdatedAt, + }) + } + c.JSON(http.StatusOK, gin.H{"bindings": items}) +} + +func (h *RobotHandler) DeleteMyRobotBinding(c *gin.Context) { + session, ok := security.CurrentSession(c) + if !ok { + c.JSON(http.StatusUnauthorized, gin.H{"error": "未授权访问"}) + return + } + if err := h.db.DeleteRobotUserBindingForUser(c.Param("id"), session.UserID); err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": "绑定不存在"}) + return + } + if h.audit != nil { + h.audit.Record(c, audit.Entry{Category: "auth", Action: "robot_binding_revoke", Result: "success", ResourceType: "robot_binding", ResourceID: c.Param("id"), Message: "撤销机器人平台账号绑定"}) + } + c.Status(http.StatusNoContent) +} + +// RobotTestRequest 模拟机器人消息请求 +type RobotTestRequest struct { + Platform string `json:"platform"` // 如 "dingtalk"、"lark"、"wecom" + UserID string `json:"user_id"` + Text string `json:"text"` +} + +// HandleRobotTest 供本地验证：POST JSON { "platform", "user_id", "text" }，返回 { "reply": "..." } +func (h *RobotHandler) HandleRobotTest(c *gin.Context) { + var req RobotTestRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "请求体需为 JSON，包含 platform、user_id、text"}) + return + } + platform := strings.TrimSpace(req.Platform) + if platform == "" { + platform = "test" + } + userID := strings.TrimSpace(req.UserID) + if userID == "" { + userID = "test_user" + } + reply := h.HandleMessage(platform, userID, req.Text) + c.JSON(http.StatusOK, gin.H{"reply": reply}) +} + +// sendWecomMessageViaAPI 通过企业微信 API 主动发送消息（用于异步处理后的结果发送） +func (h *RobotHandler) sendWecomMessageViaAPI(toUser, toParty, content string) { + if !h.config.Robots.Wecom.Enabled { + return + } + + secret := h.config.Robots.Wecom.Secret + corpID := h.config.Robots.Wecom.CorpID + agentID := h.config.Robots.Wecom.AgentID + + if secret == "" || corpID == "" { + h.logger.Warn("企业微信主动 API 缺少 secret 或 corpID 配置") + return + } + + // 第 1 步：获取 access_token + tokenURL := fmt.Sprintf("https://qyapi.weixin.qq.com/cgi-bin/gettoken?corpid=%s&corpsecret=%s", corpID, secret) + resp, err := http.Get(tokenURL) + if err != nil { + h.logger.Warn("企业微信获取 token 失败", zap.Error(err)) + return + } + defer resp.Body.Close() + + var tokenResp struct { + AccessToken string `json:"access_token"` + ErrCode int `json:"errcode"` + ErrMsg string `json:"errmsg"` + } + if err := json.NewDecoder(resp.Body).Decode(&tokenResp); err != nil { + h.logger.Warn("企业微信 token 响应解析失败", zap.Error(err)) + return + } + if tokenResp.ErrCode != 0 { + h.logger.Warn("企业微信 token 获取错误", zap.String("errmsg", tokenResp.ErrMsg), zap.Int("errcode", tokenResp.ErrCode)) + return + } + + // 第 2 步：构造发送消息请求 + msgReq := map[string]interface{}{ + "touser": toUser, + "msgtype": "text", + "agentid": agentID, + "text": map[string]interface{}{ + "content": content, + }, + } + + msgBody, err := json.Marshal(msgReq) + if err != nil { + h.logger.Warn("企业微信消息序列化失败", zap.Error(err)) + return + } + + // 第 3 步：发送消息 + sendURL := fmt.Sprintf("https://qyapi.weixin.qq.com/cgi-bin/message/send?access_token=%s", tokenResp.AccessToken) + msgResp, err := http.Post(sendURL, "application/json", bytes.NewReader(msgBody)) + if err != nil { + h.logger.Warn("企业微信主动发送消息失败", zap.Error(err)) + return + } + defer msgResp.Body.Close() + + var sendResp struct { + ErrCode int `json:"errcode"` + ErrMsg string `json:"errmsg"` + InvalidUser string `json:"invaliduser"` + MsgID string `json:"msgid"` + } + if err := json.NewDecoder(msgResp.Body).Decode(&sendResp); err != nil { + h.logger.Warn("企业微信发送响应解析失败", zap.Error(err)) + return + } + + if sendResp.ErrCode == 0 { + h.logger.Debug("企业微信主动发送消息成功", zap.String("msgid", sendResp.MsgID)) + } else { + h.logger.Warn("企业微信主动发送消息失败", zap.String("errmsg", sendResp.ErrMsg), zap.Int("errcode", sendResp.ErrCode), zap.String("invaliduser", sendResp.InvalidUser)) + } +} + +// —————— 钉钉 —————— + +// HandleDingtalkPOST 钉钉事件回调（流式接入等）；当前为占位，返回 200 +func (h *RobotHandler) HandleDingtalkPOST(c *gin.Context) { + if !h.config.Robots.Dingtalk.Enabled { + c.JSON(http.StatusOK, gin.H{}) + return + } + // 钉钉流式/事件回调格式需按官方文档解析并异步回复，此处仅返回 200 + c.JSON(http.StatusOK, gin.H{"message": "ok"}) +} + +// —————— 飞书 —————— + +// HandleLarkPOST 飞书事件回调；当前为占位，返回 200；验证时需返回 challenge +func (h *RobotHandler) HandleLarkPOST(c *gin.Context) { + if !h.config.Robots.Lark.Enabled { + c.JSON(http.StatusOK, gin.H{}) + return + } + var body struct { + Challenge string `json:"challenge"` + } + if err := c.ShouldBindJSON(&body); err == nil && body.Challenge != "" { + c.JSON(http.StatusOK, gin.H{"challenge": body.Challenge}) + return + } + c.JSON(http.StatusOK, gin.H{}) +} diff --git a/internal/handler/robot_mode_test.go b/internal/handler/robot_mode_test.go new file mode 100644 index 00000000..ee9ab631 --- /dev/null +++ b/internal/handler/robot_mode_test.go @@ -0,0 +1,101 @@ +package handler + +import ( + "strings" + "testing" + + "cyberstrike-ai/internal/config" + + "go.uber.org/zap" +) + +func TestRobotModeSwitch(t *testing.T) { + h := NewRobotHandler(&config.Config{MultiAgent: config.MultiAgentConfig{Enabled: true}}, nil, nil, zap.NewNop()) + + if got := h.cmdSwitchMode("lark", "user-1", "plan-execute"); !strings.Contains(got, "Plan-Execute") { + t.Fatalf("unexpected switch response: %s", got) + } + if got := h.getAgentMode("lark", "user-1"); got != "plan_execute" { + t.Fatalf("mode = %q, want plan_execute", got) + } + if got := h.cmdModes("lark", "user-1"); !strings.Contains(got, "当前模式: Plan-Execute") { + t.Fatalf("unexpected modes response: %s", got) + } +} + +func TestRobotModeRejectsUnavailableMultiAgent(t *testing.T) { + h := NewRobotHandler(&config.Config{}, nil, nil, zap.NewNop()) + + if got := h.cmdSwitchMode("lark", "user-1", "deep"); !strings.Contains(got, "启用 Eino 多代理") { + t.Fatalf("unexpected rejection: %s", got) + } + if got := h.getAgentMode("lark", "user-1"); got != "eino_single" { + t.Fatalf("mode changed after rejection: %q", got) + } +} + +func TestParseRobotAgentModeRejectsUnknownMode(t *testing.T) { + if mode, ok := parseRobotAgentMode("unknown"); ok || mode != "" { + t.Fatalf("parseRobotAgentMode returned (%q, %v), want empty,false", mode, ok) + } +} + +func TestRobotStatusCommandPermission(t *testing.T) { + for _, command := range []string{"状态", "status"} { + permission, recognized := robotCommandPermission(command) + if !recognized || permission != "chat:read" { + t.Fatalf("command %q returned permission=%q recognized=%v", command, permission, recognized) + } + } + for _, removed := range []string{"当前", "current"} { + if _, recognized := robotCommandPermission(removed); recognized { + t.Fatalf("removed command %q is still recognized", removed) + } + } +} + +func TestRobotBestPracticeCommandPermissions(t *testing.T) { + cases := map[string]string{ + "任务": "chat:read", + "task": "chat:read", + "重命名 新标题": "chat:write", + "rename x": "chat:write", + "诊断": "config:read", + "doctor": "config:read", + } + for command, want := range cases { + permission, recognized := robotCommandPermission(command) + if !recognized || permission != want { + t.Fatalf("command %q returned permission=%q recognized=%v, want %q,true", command, permission, recognized, want) + } + } +} + +func TestRobotConfirmationCanBeCancelled(t *testing.T) { + h := NewRobotHandler(&config.Config{}, nil, nil, zap.NewNop()) + h.setPendingConfirmation("lark", "user-1", "delete_conversation", "conv-1") + if got := h.cmdCancelConfirmation("lark", "user-1"); got != "已取消待确认操作。" { + t.Fatalf("unexpected cancel response: %s", got) + } + if got := h.cmdConfirm("lark", "user-1"); !strings.Contains(got, "没有待确认操作") { + t.Fatalf("confirmation survived cancellation: %s", got) + } +} + +func TestRobotDoctorSeparatesInternalToolsFromHTTPMCP(t *testing.T) { + h := NewRobotHandler(&config.Config{ + Security: config.SecurityConfig{Tools: []config.ToolConfig{ + {Name: "enabled-tool", Enabled: true}, + {Name: "disabled-tool", Enabled: false}, + }}, + MCP: config.MCPConfig{Enabled: false}, + }, nil, nil, zap.NewNop()) + + got := h.cmdDoctor() + if !strings.Contains(got, "内置 MCP 工具: 1/2 个已启用") { + t.Fatalf("internal tool status missing: %s", got) + } + if !strings.Contains(got, "HTTP MCP 服务: 已关闭") { + t.Fatalf("HTTP MCP status missing: %s", got) + } +} diff --git a/internal/handler/robot_rbac_test.go b/internal/handler/robot_rbac_test.go new file mode 100644 index 00000000..d1a93bb8 --- /dev/null +++ b/internal/handler/robot_rbac_test.go @@ -0,0 +1,138 @@ +package handler + +import ( + "fmt" + "path/filepath" + "strings" + "testing" + "time" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + + "go.uber.org/zap" +) + +func TestRobotUsersAreResourceIsolated(t *testing.T) { + db, err := database.NewDB(filepath.Join(t.TempDir(), "robot-rbac.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = db.Close() }) + cfg := &config.Config{} + cfg.Project.Enabled = true + h := NewRobotHandler(cfg, db, nil, zap.NewNop()) + if err := db.BootstrapRBAC("hash", security.PermissionCatalog); err != nil { + t.Fatal(err) + } + alice, err := db.CreateRBACUser("robot-alice", "Robot Alice", "hash", true, []string{database.RBACSystemRoleOperator}) + if err != nil { + t.Fatal(err) + } + bob, err := db.CreateRBACUser("robot-bob", "Robot Bob", "hash", true, []string{database.RBACSystemRoleOperator}) + if err != nil { + t.Fatal(err) + } + for externalID, user := range map[string]*database.RBACUser{"alice": alice, "bob": bob} { + code := "TEST-" + strings.ToUpper(externalID) + if err := db.CreateRobotBindingCode(user.ID, hashRobotBindingCode(code), time.Now().Add(time.Minute)); err != nil { + t.Fatal(err) + } + if _, err := db.ConsumeRobotBindingCode("wecom", externalID, hashRobotBindingCode(code)); err != nil { + t.Fatal(err) + } + } + aliceAccess, err := db.ResolveRobotRBACAccess("wecom", "alice") + if err != nil { + t.Fatal(err) + } + if got := h.HandleMessage("wecom", "alice", "身份"); !strings.Contains(got, "Robot Alice") || !strings.Contains(got, alice.ID) || !strings.Contains(got, "user_binding") { + t.Fatalf("bound identity output is incomplete: %s", got) + } + + conversationID, _ := h.getOrCreateConversation("wecom", "alice", "alice conversation", aliceAccess) + if conversationID == "" { + t.Fatal("alice conversation was not created") + } + if got := h.cmdList("wecom", "bob"); strings.Contains(got, conversationID) { + t.Fatalf("bob listed alice conversation: %s", got) + } + if got := h.cmdSwitch("wecom", "bob", conversationID); !strings.Contains(got, "不存在") && !strings.Contains(got, "无权访问") { + t.Fatalf("bob switched to alice conversation: %s", got) + } + if got := h.cmdDelete("wecom", "bob", conversationID); !strings.Contains(got, "无权访问") { + t.Fatalf("bob deleted alice conversation: %s", got) + } + if _, err := db.GetConversation(conversationID); err != nil { + t.Fatalf("alice conversation was deleted: %v", err) + } + + createReply := h.cmdNewProject("wecom", "alice", "alice project") + if !strings.Contains(createReply, "已创建项目") { + t.Fatalf("create project reply: %s", createReply) + } + if got := h.cmdProjects("wecom", "bob"); strings.Contains(got, "alice project") { + t.Fatalf("bob listed alice project: %s", got) + } +} + +func TestWecomReplayGuardRequiresFreshUniqueRequest(t *testing.T) { + h := NewRobotHandler(&config.Config{}, nil, nil, zap.NewNop()) + timestamp := time.Now().Unix() + if !h.acceptFreshWecomRequest(fmt.Sprintf("%d", timestamp), "nonce", "signature") { + t.Fatal("fresh request was rejected") + } + if h.acceptFreshWecomRequest(fmt.Sprintf("%d", timestamp), "nonce", "signature") { + t.Fatal("duplicate request was accepted") + } + if h.acceptFreshWecomRequest(fmt.Sprintf("%d", timestamp-600), "old", "signature") { + t.Fatal("stale request was accepted") + } +} + +func TestRobotServiceAccountRequiresExactSenderAllowlist(t *testing.T) { + db, err := database.NewDB(filepath.Join(t.TempDir(), "robot-service.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = db.Close() }) + if err := db.BootstrapRBAC("hash", security.PermissionCatalog); err != nil { + t.Fatal(err) + } + serviceUser, err := db.CreateRBACUser("robot-service-user", "Robot Service", "hash", true, []string{database.RBACSystemRoleOperator}) + if err != nil { + t.Fatal(err) + } + cfg := &config.Config{} + cfg.Robots.Lark.Auth = config.RobotAuthorizationConfig{ + Mode: config.RobotAuthModeServiceAccount, ServiceUserID: serviceUser.ID, + AllowedExternalUsers: []string{"t:tenant|u:allowed"}, + } + h := NewRobotHandler(cfg, db, nil, zap.NewNop()) + if got := h.HandleMessage("lark", "t:tenant|u:allowed", "列表"); strings.Contains(got, "白名单") || strings.Contains(got, "尚未绑定") { + t.Fatalf("allowed service account sender was denied: %s", got) + } + if got := h.HandleMessage("lark", "t:tenant|u:denied", "列表"); !strings.Contains(got, "白名单") { + t.Fatalf("non-allowlisted sender was not denied: %s", got) + } + if got := h.HandleMessage("lark", "t:tenant|u:allowed", "绑定 ABCD-1234"); !strings.Contains(got, "服务账号模式") { + t.Fatalf("service-account robot accepted user binding: %s", got) + } + if got := h.HandleMessage("lark", "t:tenant|u:allowed", "whoami"); !strings.Contains(got, "Robot Service") || !strings.Contains(got, serviceUser.ID) || !strings.Contains(got, "service_account") { + t.Fatalf("service-account identity output is incomplete: %s", got) + } + if got := h.HandleMessage("lark", "t:tenant|u:denied", "whoami"); !strings.Contains(got, "鉴权状态：拒绝") || strings.Contains(got, serviceUser.ID) { + t.Fatalf("denied identity output leaked or omitted status: %s", got) + } + + adminCfg := &config.Config{} + adminCfg.Robots.Lark.Auth = config.RobotAuthorizationConfig{ + Mode: config.RobotAuthModeServiceAccount, ServiceUserID: "admin", + AllowedExternalUsers: []string{"t:tenant|u:owner"}, + } + adminHandler := NewRobotHandler(adminCfg, db, nil, zap.NewNop()) + if got := adminHandler.HandleMessage("lark", "t:tenant|u:owner", "身份"); !strings.Contains(got, "admin") || !strings.Contains(got, "鉴权状态：已授权") { + t.Fatalf("allowlisted admin service account was denied: %s", got) + } +} diff --git a/internal/handler/robot_wecom_test.go b/internal/handler/robot_wecom_test.go new file mode 100644 index 00000000..bb4a78ed --- /dev/null +++ b/internal/handler/robot_wecom_test.go @@ -0,0 +1,78 @@ +package handler + +import ( + "net/http" + "net/http/httptest" + "strings" + "testing" + + "cyberstrike-ai/internal/config" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +func newWecomTestHandler(token string, aesKey string) *RobotHandler { + return &RobotHandler{ + config: &config.Config{ + Robots: config.RobotsConfig{ + Wecom: config.RobotWecomConfig{ + Enabled: true, + Token: token, + EncodingAESKey: aesKey, + }, + }, + }, + logger: zap.NewNop(), + } +} + +func TestHandleWecomPOST_rejectsWhenTokenEmpty(t *testing.T) { + gin.SetMode(gin.TestMode) + + h := newWecomTestHandler("", "") + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + body := `attackertexthi` + c.Request = httptest.NewRequest(http.MethodPost, "/api/robot/wecom", strings.NewReader(body)) + + h.HandleWecomPOST(c) + + if w.Code != http.StatusForbidden { + t.Fatalf("status = %d, want %d", w.Code, http.StatusForbidden) + } + if w.Body.String() == "success" { + t.Fatal("expected rejection, got success") + } +} + +func TestHandleWecomPOST_rejectsPlaintextWhenEncryptionConfigured(t *testing.T) { + gin.SetMode(gin.TestMode) + + h := newWecomTestHandler("secret-token", "abcdefghijklmnopqrstuvwxyz0123456789ABCD") + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + body := `attackertexthi` + c.Request = httptest.NewRequest(http.MethodPost, "/api/robot/wecom?timestamp=1&nonce=2&msg_signature=fake", strings.NewReader(body)) + + h.HandleWecomPOST(c) + + if w.Body.String() == "success" { + t.Fatal("expected rejection for plaintext in encryption mode, got success") + } +} + +func TestHandleWecomGET_rejectsWhenTokenEmpty(t *testing.T) { + gin.SetMode(gin.TestMode) + + h := newWecomTestHandler("", "") + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodGet, "/api/robot/wecom?msg_signature=x×tamp=1&nonce=2&echostr=abc", nil) + + h.HandleWecomGET(c) + + if w.Code != http.StatusForbidden { + t.Fatalf("status = %d, want %d", w.Code, http.StatusForbidden) + } +} diff --git a/internal/handler/role.go b/internal/handler/role.go new file mode 100644 index 00000000..1c061256 --- /dev/null +++ b/internal/handler/role.go @@ -0,0 +1,469 @@ +package handler + +import ( + "fmt" + "net/http" + "os" + "path/filepath" + "regexp" + "strings" + + "cyberstrike-ai/internal/audit" + "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 + audit *audit.Service +} + +// SetAudit wires platform audit logging. +func (h *RoleHandler) SetAudit(s *audit.Service) { + h.audit = s +} + +// 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)) + if h.audit != nil { + h.audit.RecordOK(c, "role", "update", "更新角色", "role", finalKey, map[string]interface{}{"name": req.Name}) + } + c.JSON(http.StatusOK, gin.H{ + "message": "角色已更新", + "role": req, + }) +} + +// 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)) + if h.audit != nil { + h.audit.RecordOK(c, "role", "create", "创建角色", "role", req.Name, nil) + } + 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)) + if h.audit != nil { + h.audit.RecordOK(c, "role", "delete", "删除角色", "role", roleName, nil) + } + 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/handler/skills.go b/internal/handler/skills.go new file mode 100644 index 00000000..4246c297 --- /dev/null +++ b/internal/handler/skills.go @@ -0,0 +1,710 @@ +package handler + +import ( + "fmt" + "net/http" + "os" + "path/filepath" + "regexp" + "strings" + + "cyberstrike-ai/internal/audit" + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/skillpackage" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" + "gopkg.in/yaml.v3" +) + +// SkillsHandler Skills处理器（磁盘 + Eino 规范；运行时由 Eino ADK skill 中间件加载） +type SkillsHandler struct { + config *config.Config + configPath string + logger *zap.Logger + db *database.DB // 数据库连接（遗留统计；MCP list/read 已移除） + audit *audit.Service +} + +// SetAudit wires platform audit logging. +func (h *SkillsHandler) SetAudit(s *audit.Service) { + h.audit = s +} + +// NewSkillsHandler 创建新的Skills处理器 +func NewSkillsHandler(cfg *config.Config, configPath string, logger *zap.Logger) *SkillsHandler { + return &SkillsHandler{ + config: cfg, + configPath: configPath, + logger: logger, + } +} + +func (h *SkillsHandler) skillsRootAbs() string { + skillsDir := h.config.SkillsDir + if skillsDir == "" { + skillsDir = "skills" + } + configDir := filepath.Dir(h.configPath) + if !filepath.IsAbs(skillsDir) { + skillsDir = filepath.Join(configDir, skillsDir) + } + return skillsDir +} + +// SetDB 设置数据库连接（用于获取调用统计） +func (h *SkillsHandler) SetDB(db *database.DB) { + h.db = db +} + +// GetSkills 获取所有skills列表（支持分页和搜索） +func (h *SkillsHandler) GetSkills(c *gin.Context) { + allSummaries, err := skillpackage.ListSkillSummaries(h.skillsRootAbs()) + if err != nil { + h.logger.Error("获取skills列表失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + searchKeyword := strings.TrimSpace(c.Query("search")) + + allSkillsInfo := make([]map[string]interface{}, 0, len(allSummaries)) + for _, s := range allSummaries { + skillInfo := map[string]interface{}{ + "id": s.ID, + "name": s.Name, + "dir_name": s.DirName, + "description": s.Description, + "version": s.Version, + "path": s.Path, + "tags": s.Tags, + "triggers": s.Triggers, + "script_count": s.ScriptCount, + "file_count": s.FileCount, + "progressive": s.Progressive, + "file_size": s.FileSize, + "mod_time": s.ModTime, + } + allSkillsInfo = append(allSkillsInfo, skillInfo) + } + + filteredSkillsInfo := allSkillsInfo + if searchKeyword != "" { + keywordLower := strings.ToLower(searchKeyword) + filteredSkillsInfo = make([]map[string]interface{}, 0) + for _, skillInfo := range allSkillsInfo { + id := strings.ToLower(fmt.Sprintf("%v", skillInfo["id"])) + name := strings.ToLower(fmt.Sprintf("%v", skillInfo["name"])) + description := strings.ToLower(fmt.Sprintf("%v", skillInfo["description"])) + path := strings.ToLower(fmt.Sprintf("%v", skillInfo["path"])) + version := strings.ToLower(fmt.Sprintf("%v", skillInfo["version"])) + tagsJoined := "" + if tags, ok := skillInfo["tags"].([]string); ok { + tagsJoined = strings.ToLower(strings.Join(tags, " ")) + } + trigJoined := "" + if tr, ok := skillInfo["triggers"].([]string); ok { + trigJoined = strings.ToLower(strings.Join(tr, " ")) + } + if strings.Contains(id, keywordLower) || + strings.Contains(name, keywordLower) || + strings.Contains(description, keywordLower) || + strings.Contains(path, keywordLower) || + strings.Contains(version, keywordLower) || + strings.Contains(tagsJoined, keywordLower) || + strings.Contains(trigJoined, keywordLower) { + filteredSkillsInfo = append(filteredSkillsInfo, skillInfo) + } + } + } + + // 分页参数 + limit := 20 // 默认每页20条 + offset := 0 + if limitStr := c.Query("limit"); limitStr != "" { + if parsed, err := parseInt(limitStr); err == nil && parsed > 0 { + // 允许更大的limit用于搜索场景，但设置一个合理的上限（10000） + if parsed <= 10000 { + limit = parsed + } else { + limit = 10000 + } + } + } + if offsetStr := c.Query("offset"); offsetStr != "" { + if parsed, err := parseInt(offsetStr); err == nil && parsed >= 0 { + offset = parsed + } + } + + // 计算分页范围 + total := len(filteredSkillsInfo) + start := offset + end := offset + limit + if start > total { + start = total + } + if end > total { + end = total + } + + // 获取当前页的skill列表 + var paginatedSkillsInfo []map[string]interface{} + if start < end { + paginatedSkillsInfo = filteredSkillsInfo[start:end] + } else { + paginatedSkillsInfo = []map[string]interface{}{} + } + + c.JSON(http.StatusOK, gin.H{ + "skills": paginatedSkillsInfo, + "total": total, + "limit": limit, + "offset": offset, + }) +} + +// GetSkill 获取单个skill的详细信息 +func (h *SkillsHandler) GetSkill(c *gin.Context) { + skillName := c.Param("name") + if skillName == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "skill名称不能为空"}) + return + } + + resPath := strings.TrimSpace(c.Query("resource_path")) + if resPath == "" { + resPath = strings.TrimSpace(c.Query("skill_script_path")) + } + if resPath != "" { + content, err := skillpackage.ReadScriptText(h.skillsRootAbs(), skillName, resPath, 0) + if err != nil { + h.logger.Warn("读取skill资源失败", zap.String("skill", skillName), zap.String("path", resPath), zap.Error(err)) + c.JSON(http.StatusNotFound, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{ + "skill": map[string]interface{}{ + "id": skillName, + }, + "resource": map[string]interface{}{ + "path": resPath, + "content": content, + }, + }) + return + } + + depthStr := strings.ToLower(strings.TrimSpace(c.DefaultQuery("depth", "full"))) + section := strings.TrimSpace(c.Query("section")) + opt := skillpackage.LoadOptions{Section: section} + switch depthStr { + case "summary": + opt.Depth = "summary" + case "full", "": + opt.Depth = "full" + default: + c.JSON(http.StatusBadRequest, gin.H{"error": "depth 仅支持 summary 或 full"}) + return + } + + skill, err := skillpackage.LoadSkill(h.skillsRootAbs(), skillName, opt) + if err != nil { + h.logger.Warn("加载skill失败", zap.String("skill", skillName), zap.Error(err)) + c.JSON(http.StatusNotFound, gin.H{"error": "skill不存在: " + err.Error()}) + return + } + + skillPath := skill.Path + skillFile := filepath.Join(skillPath, "SKILL.md") + + fileInfo, _ := os.Stat(skillFile) + var fileSize int64 + var modTime string + if fileInfo != nil { + fileSize = fileInfo.Size() + modTime = fileInfo.ModTime().Format("2006-01-02 15:04:05") + } + + c.JSON(http.StatusOK, gin.H{ + "skill": map[string]interface{}{ + "id": skill.DirName, + "name": skill.Name, + "description": skill.Description, + "content": skill.Content, + "path": skill.Path, + "version": skill.Version, + "tags": skill.Tags, + "scripts": skill.Scripts, + "sections": skill.Sections, + "package_files": skill.PackageFiles, + "file_size": fileSize, + "mod_time": modTime, + "depth": depthStr, + "section": section, + }, + }) +} + +// ListSkillPackageFiles lists all files in a skill directory (Agent Skills layout). +func (h *SkillsHandler) ListSkillPackageFiles(c *gin.Context) { + skillID := c.Param("name") + files, err := skillpackage.ListPackageFiles(h.skillsRootAbs(), skillID) + if err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"files": files}) +} + +// GetSkillPackageFile returns one file by relative path (?path=). +func (h *SkillsHandler) GetSkillPackageFile(c *gin.Context) { + skillID := c.Param("name") + rel := strings.TrimSpace(c.Query("path")) + if rel == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "query path is required"}) + return + } + b, err := skillpackage.ReadPackageFile(h.skillsRootAbs(), skillID, rel, 0) + if err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"path": rel, "content": string(b)}) +} + +// PutSkillPackageFile writes a file inside the skill package. +func (h *SkillsHandler) PutSkillPackageFile(c *gin.Context) { + skillID := c.Param("name") + var req struct { + Path string `json:"path" binding:"required"` + Content string `json:"content"` + } + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()}) + return + } + if req.Path == "SKILL.md" { + if err := skillpackage.ValidateSkillMDPackage([]byte(req.Content), skillID); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + } + if err := skillpackage.WritePackageFile(h.skillsRootAbs(), skillID, req.Path, []byte(req.Content)); err != nil { + h.logger.Error("写入 skill 文件失败", zap.String("skill", skillID), zap.String("path", req.Path), zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"message": "saved", "path": req.Path}) +} + +// GetSkillBoundRoles 获取绑定指定skill的角色列表 +func (h *SkillsHandler) GetSkillBoundRoles(c *gin.Context) { + skillName := c.Param("name") + if skillName == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "skill名称不能为空"}) + return + } + + boundRoles := h.getRolesBoundToSkill(skillName) + c.JSON(http.StatusOK, gin.H{ + "skill": skillName, + "bound_roles": boundRoles, + "bound_count": len(boundRoles), + }) +} + +// getRolesBoundToSkill 预留：角色不再配置 skill 绑定，始终返回空列表。 +func (h *SkillsHandler) getRolesBoundToSkill(skillName string) []string { + _ = skillName + return nil +} + +// CreateSkill 创建新 skill（标准 Agent Skills：生成 SKILL.md + YAML front matter） +func (h *SkillsHandler) CreateSkill(c *gin.Context) { + var req struct { + Name string `json:"name" binding:"required"` + Description string `json:"description" binding:"required"` + Content string `json:"content" binding:"required"` + } + + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()}) + return + } + + if !isValidSkillName(req.Name) { + c.JSON(http.StatusBadRequest, gin.H{"error": "skill 目录名须为小写字母、数字、连字符（与 Agent Skills name 一致）"}) + return + } + + manifest := &skillpackage.SkillManifest{ + Name: req.Name, + Description: strings.TrimSpace(req.Description), + } + skillMD, err := skillpackage.BuildSkillMD(manifest, req.Content) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if err := skillpackage.ValidateSkillMDPackage(skillMD, req.Name); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + skillDir := filepath.Join(h.skillsRootAbs(), req.Name) + if err := os.MkdirAll(skillDir, 0755); err != nil { + h.logger.Error("创建skill目录失败", zap.String("skill", req.Name), zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": "创建skill目录失败: " + err.Error()}) + return + } + + if _, err := os.Stat(filepath.Join(skillDir, "SKILL.md")); err == nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "skill已存在"}) + return + } + + if err := os.WriteFile(filepath.Join(skillDir, "SKILL.md"), skillMD, 0644); err != nil { + h.logger.Error("创建 SKILL.md 失败", zap.String("skill", req.Name), zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": "创建 SKILL.md 失败: " + err.Error()}) + return + } + + h.logger.Info("创建skill成功", zap.String("skill", req.Name)) + if h.audit != nil { + h.audit.RecordOK(c, "skill", "create", "创建 Skill", "skill", req.Name, nil) + } + c.JSON(http.StatusOK, gin.H{ + "message": "skill已创建", + "skill": map[string]interface{}{ + "name": req.Name, + "path": skillDir, + }, + }) +} + +// UpdateSkill 更新 SKILL.md（保留 front matter 中除 description 外的字段；可选覆盖 description） +func (h *SkillsHandler) UpdateSkill(c *gin.Context) { + skillName := c.Param("name") + if skillName == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "skill名称不能为空"}) + return + } + + var req struct { + Description string `json:"description"` + Content string `json:"content" binding:"required"` + } + + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()}) + return + } + + mdPath := filepath.Join(h.skillsRootAbs(), skillName, "SKILL.md") + raw, err := os.ReadFile(mdPath) + if err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": "skill不存在: " + err.Error()}) + return + } + m, _, err := skillpackage.ParseSkillMD(raw) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if req.Description != "" { + m.Description = strings.TrimSpace(req.Description) + } + skillMD, err := skillpackage.BuildSkillMD(m, req.Content) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if err := skillpackage.ValidateSkillMDPackage(skillMD, skillName); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + skillDir := filepath.Join(h.skillsRootAbs(), skillName) + + if err := os.WriteFile(filepath.Join(skillDir, "SKILL.md"), skillMD, 0644); err != nil { + h.logger.Error("更新 SKILL.md 失败", zap.String("skill", skillName), zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": "更新 SKILL.md 失败: " + err.Error()}) + return + } + + h.logger.Info("更新skill成功", zap.String("skill", skillName)) + if h.audit != nil { + h.audit.RecordOK(c, "skill", "update", "更新 Skill", "skill", skillName, nil) + } + c.JSON(http.StatusOK, gin.H{ + "message": "skill已更新", + }) +} + +// DeleteSkill 删除skill +func (h *SkillsHandler) DeleteSkill(c *gin.Context) { + skillName := c.Param("name") + if skillName == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "skill名称不能为空"}) + return + } + + // 检查是否有角色绑定了该skill，如果有则自动移除绑定 + affectedRoles := h.removeSkillFromRoles(skillName) + if len(affectedRoles) > 0 { + h.logger.Info("从角色中移除skill绑定", + zap.String("skill", skillName), + zap.Strings("roles", affectedRoles)) + } + + skillDir := filepath.Join(h.skillsRootAbs(), skillName) + if err := os.RemoveAll(skillDir); err != nil { + h.logger.Error("删除skill失败", zap.String("skill", skillName), zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": "删除skill失败: " + err.Error()}) + return + } + responseMsg := "skill已删除" + if len(affectedRoles) > 0 { + responseMsg = fmt.Sprintf("skill已删除，已自动从 %d 个角色中移除绑定: %s", + len(affectedRoles), strings.Join(affectedRoles, ", ")) + } + + h.logger.Info("删除skill成功", zap.String("skill", skillName)) + if h.audit != nil { + h.audit.RecordOK(c, "skill", "delete", "删除 Skill", "skill", skillName, map[string]interface{}{ + "affected_roles": affectedRoles, + }) + } + c.JSON(http.StatusOK, gin.H{ + "message": responseMsg, + "affected_roles": affectedRoles, + }) +} + +// GetSkillStats 获取skills调用统计信息 +func (h *SkillsHandler) GetSkillStats(c *gin.Context) { + skillList, err := skillpackage.ListSkillDirNames(h.skillsRootAbs()) + if err != nil { + h.logger.Error("获取skills列表失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + skillsDir := h.skillsRootAbs() + + // 从数据库加载调用统计 + var skillStatsMap map[string]*database.SkillStats + if h.db != nil { + dbStats, err := h.db.LoadSkillStats() + if err != nil { + h.logger.Warn("从数据库加载Skills统计信息失败", zap.Error(err)) + skillStatsMap = make(map[string]*database.SkillStats) + } else { + skillStatsMap = dbStats + } + } else { + skillStatsMap = make(map[string]*database.SkillStats) + } + + // 构建统计信息（包含所有skills，即使没有调用记录） + statsList := make([]map[string]interface{}, 0, len(skillList)) + totalCalls := 0 + totalSuccess := 0 + totalFailed := 0 + + for _, skillName := range skillList { + stat, exists := skillStatsMap[skillName] + if !exists { + stat = &database.SkillStats{ + SkillName: skillName, + TotalCalls: 0, + SuccessCalls: 0, + FailedCalls: 0, + } + } + + totalCalls += stat.TotalCalls + totalSuccess += stat.SuccessCalls + totalFailed += stat.FailedCalls + + lastCallTimeStr := "" + if stat.LastCallTime != nil { + lastCallTimeStr = stat.LastCallTime.Format("2006-01-02 15:04:05") + } + + statsList = append(statsList, map[string]interface{}{ + "skill_name": stat.SkillName, + "total_calls": stat.TotalCalls, + "success_calls": stat.SuccessCalls, + "failed_calls": stat.FailedCalls, + "last_call_time": lastCallTimeStr, + }) + } + + c.JSON(http.StatusOK, gin.H{ + "total_skills": len(skillList), + "total_calls": totalCalls, + "total_success": totalSuccess, + "total_failed": totalFailed, + "skills_dir": skillsDir, + "stats": statsList, + }) +} + +// ClearSkillStats 清空所有Skills统计信息 +func (h *SkillsHandler) ClearSkillStats(c *gin.Context) { + if h.db == nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "数据库连接未配置"}) + return + } + + if err := h.db.ClearSkillStats(); err != nil { + h.logger.Error("清空Skills统计信息失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": "清空统计信息失败: " + err.Error()}) + return + } + + h.logger.Info("已清空所有Skills统计信息") + c.JSON(http.StatusOK, gin.H{ + "message": "已清空所有Skills统计信息", + }) +} + +// ClearSkillStatsByName 清空指定skill的统计信息 +func (h *SkillsHandler) ClearSkillStatsByName(c *gin.Context) { + skillName := c.Param("name") + if skillName == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "skill名称不能为空"}) + return + } + + if h.db == nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "数据库连接未配置"}) + return + } + + if err := h.db.ClearSkillStatsByName(skillName); err != nil { + h.logger.Error("清空指定skill统计信息失败", zap.String("skill", skillName), zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": "清空统计信息失败: " + err.Error()}) + return + } + + h.logger.Info("已清空指定skill统计信息", zap.String("skill", skillName)) + c.JSON(http.StatusOK, gin.H{ + "message": fmt.Sprintf("已清空skill '%s' 的统计信息", skillName), + }) +} + +// removeSkillFromRoles 预留：角色不再存储 skill 绑定，无操作。 +func (h *SkillsHandler) removeSkillFromRoles(skillName string) []string { + _ = skillName + return nil +} + +// saveRolesConfig 保存角色配置到文件（从SkillsHandler调用） +func (h *SkillsHandler) saveRolesConfig() 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 := sanitizeRoleFileName(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 +} + +// sanitizeRoleFileName 将角色名称转换为安全的文件名 +func sanitizeRoleFileName(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 +} + +// isValidSkillName 验证 skill 目录名（与 Agent Skills 的 name 字段一致：小写、数字、连字符） +func isValidSkillName(name string) bool { + if name == "" || len(name) > 100 { + return false + } + for _, r := range name { + if !((r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') || r == '-') { + return false + } + } + return true +} diff --git a/internal/handler/sse_keepalive.go b/internal/handler/sse_keepalive.go new file mode 100644 index 00000000..dac366d9 --- /dev/null +++ b/internal/handler/sse_keepalive.go @@ -0,0 +1,88 @@ +package handler + +import ( + "context" + "fmt" + "net/http" + "sync" + "time" + + "github.com/gin-gonic/gin" +) + +// sseInterval is how often we write on long SSE streams. Shorter intervals help NATs and +// some proxies that treat connections as idle; 10s is a reasonable balance with traffic. +const sseKeepaliveInterval = 10 * time.Second + +// runSSEKeepalive starts periodic SSE heartbeats in a background goroutine. +// The returned stop function must be deferred (or called) before the handler returns so the +// goroutine exits before Gin finalizes the ResponseWriter (avoids "Write called after Handler finished"). +// +// writeMu must be the same mutex used by the handler's event writes for this request: concurrent +// writes to http.ResponseWriter break chunked transfer encoding (browser: net::ERR_INVALID_CHUNKED_ENCODING). +func runSSEKeepalive(c *gin.Context, writeMu *sync.Mutex) func() { + if writeMu == nil { + return func() {} + } + stop := make(chan struct{}) + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + sseKeepaliveLoop(c, stop, writeMu) + }() + var once sync.Once + return func() { + once.Do(func() { + close(stop) + wg.Wait() + }) + } +} + +// sseKeepaliveLoop sends periodic SSE traffic so proxies (e.g. nginx proxy_read_timeout), NATs, +// and load balancers do not close long-running streams. Some intermediaries ignore comment-only +// lines, so we send both a comment and a minimal data frame (type heartbeat) per tick. +func sseKeepaliveLoop(c *gin.Context, stop <-chan struct{}, writeMu *sync.Mutex) { + ticker := time.NewTicker(sseKeepaliveInterval) + defer ticker.Stop() + ctx := c.Request.Context() + for { + select { + case <-stop: + return + case <-ctx.Done(): + return + case <-ticker.C: + writeMu.Lock() + if sseShuttingDown(stop, ctx) { + writeMu.Unlock() + return + } + if _, err := fmt.Fprintf(c.Writer, ": keepalive\n\n"); err != nil { + writeMu.Unlock() + return + } + // data: frame so strict proxies still see downstream bytes (comments alone may not reset timers) + if _, err := fmt.Fprintf(c.Writer, `data: {"type":"heartbeat"}`+"\n\n"); err != nil { + writeMu.Unlock() + return + } + if flusher, ok := c.Writer.(http.Flusher); ok { + flusher.Flush() + } + writeMu.Unlock() + } + } +} + +func sseShuttingDown(stop <-chan struct{}, ctx context.Context) bool { + select { + case <-stop: + return true + case <-ctx.Done(): + return true + default: + return false + } +} diff --git a/internal/handler/sse_keepalive_test.go b/internal/handler/sse_keepalive_test.go new file mode 100644 index 00000000..a468708d --- /dev/null +++ b/internal/handler/sse_keepalive_test.go @@ -0,0 +1,61 @@ +package handler + +import ( + "context" + "net/http/httptest" + "sync" + "testing" + "time" + + "github.com/gin-gonic/gin" +) + +func TestRunSSEKeepaliveStopsBeforeHandlerReturns(t *testing.T) { + gin.SetMode(gin.TestMode) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("GET", "/events", nil) + + var writeMu sync.Mutex + stop := runSSEKeepalive(c, &writeMu) + stop() + + // A second stop must be safe (channel already closed, goroutine already exited). + stop() +} + +func TestRunSSEKeepaliveExitsOnClientDisconnect(t *testing.T) { + gin.SetMode(gin.TestMode) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + ctx, cancel := context.WithCancel(context.Background()) + c.Request = httptest.NewRequest("GET", "/events", nil).WithContext(ctx) + + var writeMu sync.Mutex + stop := runSSEKeepalive(c, &writeMu) + cancel() + + done := make(chan struct{}) + go func() { + stop() + close(done) + }() + + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("keepalive stop did not complete after client disconnect") + } +} + +func TestRunSSEKeepaliveNilMutexIsNoop(t *testing.T) { + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("GET", "/events", nil) + + stop := runSSEKeepalive(c, nil) + stop() +} diff --git a/internal/handler/task_event_bus.go b/internal/handler/task_event_bus.go new file mode 100644 index 00000000..bf2ad880 --- /dev/null +++ b/internal/handler/task_event_bus.go @@ -0,0 +1,116 @@ +package handler + +import "sync" + +// TaskEventBus 将主 SSE 连接上的事件镜像给后订阅的客户端（例如刷新页面后、HITL 审批通过需继续收事件）。 +// 每个 payload 为完整 SSE 行： "data: {...}\n\n" +type TaskEventBus struct { + mu sync.RWMutex + subs map[string]map[*taskEventSub]struct{} +} + +type taskEventSub struct { + mu sync.Mutex + ch chan []byte + closed bool +} + +func (s *taskEventSub) sendNonBlocking(line []byte) bool { + if s == nil { + return false + } + s.mu.Lock() + defer s.mu.Unlock() + if s.closed { + return false + } + select { + case s.ch <- line: + return true + default: + return false + } +} + +func (s *taskEventSub) closeOnce() { + if s == nil { + return + } + s.mu.Lock() + defer s.mu.Unlock() + if s.closed { + return + } + s.closed = true + close(s.ch) +} + +func NewTaskEventBus() *TaskEventBus { + return &TaskEventBus{ + subs: make(map[string]map[*taskEventSub]struct{}), + } +} + +// Subscribe 注册订阅；cancel 时需调用 Unsubscribe。 +func (b *TaskEventBus) Subscribe(conversationID string) (sub *taskEventSub, ch <-chan []byte) { + chBuf := make(chan []byte, 256) + sub = &taskEventSub{ch: chBuf} + b.mu.Lock() + if b.subs[conversationID] == nil { + b.subs[conversationID] = make(map[*taskEventSub]struct{}) + } + b.subs[conversationID][sub] = struct{}{} + b.mu.Unlock() + return sub, chBuf +} + +func (b *TaskEventBus) Unsubscribe(conversationID string, sub *taskEventSub) { + if sub == nil { + return + } + b.mu.Lock() + m, ok := b.subs[conversationID] + if !ok { + b.mu.Unlock() + return + } + delete(m, sub) + if len(m) == 0 { + delete(b.subs, conversationID) + } + b.mu.Unlock() + sub.closeOnce() +} + +// Publish 非阻塞投递；慢消费者丢帧（HITL 场景以最新状态为准，丢帧可接受）。 +func (b *TaskEventBus) Publish(conversationID string, line []byte) { + if b == nil || conversationID == "" || len(line) == 0 { + return + } + b.mu.RLock() + m := b.subs[conversationID] + subs := make([]*taskEventSub, 0, len(m)) + for s := range m { + subs = append(subs, s) + } + b.mu.RUnlock() + + cp := append([]byte(nil), line...) + for _, s := range subs { + s.sendNonBlocking(cp) + } +} + +// CloseConversation 任务结束时关闭该会话所有订阅 channel。 +func (b *TaskEventBus) CloseConversation(conversationID string) { + if b == nil || conversationID == "" { + return + } + b.mu.Lock() + m := b.subs[conversationID] + delete(b.subs, conversationID) + b.mu.Unlock() + for sub := range m { + sub.closeOnce() + } +} diff --git a/internal/handler/task_manager.go b/internal/handler/task_manager.go new file mode 100644 index 00000000..1a55789e --- /dev/null +++ b/internal/handler/task_manager.go @@ -0,0 +1,627 @@ +package handler + +import ( + "context" + "errors" + "strings" + "sync" + "time" + + "cyberstrike-ai/internal/multiagent" +) + +// ErrTaskCancelled 用户取消任务的错误 +var ErrTaskCancelled = errors.New("agent task cancelled by user") + +// ErrTaskAlreadyRunning 会话已有任务正在执行 +var ErrTaskAlreadyRunning = errors.New("agent task already running for conversation") + +// shouldPersistEinoAgentTraceAfterRunError：Eino 相关 Run 非成功返回时，是否仍写入 last_react_* 供下轮 loadHistoryFromAgentTrace。 +// 当前策略：无论正常结束、异常结束或用户主动停止，都尽量保留最后可用轨迹， +// 以便在同一会话继续时可基于原始上下文续跑，而不是回退到仅消息文本历史。 +func shouldPersistEinoAgentTraceAfterRunError(baseCtx context.Context) bool { + return true +} + +// AgentTask 描述正在运行的Agent任务 +type AgentTask struct { + ConversationID string `json:"conversationId"` + Title string `json:"title,omitempty"` + Message string `json:"message,omitempty"` + StartedAt time.Time `json:"startedAt"` + Status string `json:"status"` + CancellingAt time.Time `json:"-"` // 进入 cancelling 状态的时间，用于清理长时间卡住的任务 + + // ActiveMCPExecutionID 当前正在执行的 MCP 工具 executionId（仅内存，供「中断并继续」= 仅掐当前工具） + ActiveMCPExecutionID string `json:"-"` + + // InterruptContinueNote 无 MCP 时「中断并继续」由用户在弹窗中填写的补充说明（Cancel 前写入，续跑轮次读取后清空） + InterruptContinueNote string `json:"-"` + + // activeEinoExecuteCancel 当前进行中的 Eino filesystem execute 取消函数（与 MCP 工具并行，供中断并继续） + activeEinoExecuteCancel context.CancelFunc + // activeEinoExecuteAbortNote AbortActiveEinoExecute 写入的用户说明，由 execute 收尾时合并进工具结果 + activeEinoExecuteAbortNote string + + // hitlCognition 本轮运行中供 HITL/审计 Agent 读取的上下文（用户原话 + 思考，不含会话历史） + hitlCognition *hitlCognitionState + + // agentRuntimeCancel 当前 Eino ADK 原生 AgentCancelFunc 包装；取消任务时先触发它，再走 context 兜底。 + agentRuntimeCancel func(error) bool + agentRuntimeCancelVersion uint64 + + // agentTurnLoopInterrupt 当前 Eino TurnLoop 用户补充 push hook；中断并继续时优先将补充作为新 turn item 入队。 + agentTurnLoopInterrupt func(string) bool + agentTurnLoopInterruptVersion uint64 + + cancel func(error) +} + +// RegisterRunningTool 实现 mcp.ToolRunRegistry：工具开始时登记本会话当前 executionId。 +func (m *AgentTaskManager) RegisterRunningTool(conversationID, executionID string) { + conversationID = strings.TrimSpace(conversationID) + executionID = strings.TrimSpace(executionID) + if conversationID == "" || executionID == "" { + return + } + m.mu.Lock() + defer m.mu.Unlock() + if t, ok := m.tasks[conversationID]; ok && t != nil { + t.ActiveMCPExecutionID = executionID + } +} + +// UnregisterRunningTool 工具结束时清除登记（仅当 id 仍匹配时清除，避免并发串单）。 +func (m *AgentTaskManager) UnregisterRunningTool(conversationID, executionID string) { + conversationID = strings.TrimSpace(conversationID) + executionID = strings.TrimSpace(executionID) + if conversationID == "" || executionID == "" { + return + } + m.mu.Lock() + defer m.mu.Unlock() + if t, ok := m.tasks[conversationID]; ok && t != nil { + if t.ActiveMCPExecutionID == executionID { + t.ActiveMCPExecutionID = "" + } + } +} + +// RegisterActiveEinoExecute 登记进行中的 Eino filesystem execute（每会话同时仅一条）。 +func (m *AgentTaskManager) RegisterActiveEinoExecute(conversationID string, cancel context.CancelFunc) { + conversationID = strings.TrimSpace(conversationID) + if conversationID == "" || cancel == nil { + return + } + m.mu.Lock() + defer m.mu.Unlock() + if t, ok := m.tasks[conversationID]; ok && t != nil { + t.activeEinoExecuteCancel = cancel + t.activeEinoExecuteAbortNote = "" + } +} + +// UnregisterActiveEinoExecute execute 正常结束或已取消后清除登记。 +func (m *AgentTaskManager) UnregisterActiveEinoExecute(conversationID string) { + conversationID = strings.TrimSpace(conversationID) + if conversationID == "" { + return + } + m.mu.Lock() + defer m.mu.Unlock() + if t, ok := m.tasks[conversationID]; ok && t != nil { + t.activeEinoExecuteCancel = nil + t.activeEinoExecuteAbortNote = "" + } +} + +// ConversationIDForActiveMCPExecution 根据当前登记的工具 executionId 反查会话 ID（供 MCP 监控页按 executionId 终止）。 +func (m *AgentTaskManager) ConversationIDForActiveMCPExecution(executionID string) string { + executionID = strings.TrimSpace(executionID) + if executionID == "" { + return "" + } + m.mu.Lock() + defer m.mu.Unlock() + for convID, t := range m.tasks { + if t != nil && t.ActiveMCPExecutionID == executionID { + return convID + } + } + return "" +} + +// ConversationIDForActiveEinoExecute 返回当前唯一进行 Eino execute 的会话 ID；多会话并行时返回空。 +func (m *AgentTaskManager) ConversationIDForActiveEinoExecute() (string, bool) { + m.mu.Lock() + defer m.mu.Unlock() + var found string + count := 0 + for convID, t := range m.tasks { + if t != nil && t.activeEinoExecuteCancel != nil { + found = convID + count++ + } + } + if count == 1 { + return found, true + } + return "", false +} + +// AbortActiveEinoExecute 终止当前 Eino execute 并暂存用户说明（与 MCP 工具终止一致）。 +func (m *AgentTaskManager) AbortActiveEinoExecute(conversationID, note string) bool { + conversationID = strings.TrimSpace(conversationID) + if conversationID == "" { + return false + } + m.mu.Lock() + t, ok := m.tasks[conversationID] + if !ok || t == nil || t.activeEinoExecuteCancel == nil { + m.mu.Unlock() + return false + } + t.activeEinoExecuteAbortNote = strings.TrimSpace(note) + cancel := t.activeEinoExecuteCancel + m.mu.Unlock() + cancel() + return true +} + +// TakeEinoExecuteAbortNote 读取并清空 execute 终止说明（execute 收尾时调用一次）。 +func (m *AgentTaskManager) TakeEinoExecuteAbortNote(conversationID string) string { + conversationID = strings.TrimSpace(conversationID) + if conversationID == "" { + return "" + } + m.mu.Lock() + defer m.mu.Unlock() + if t, ok := m.tasks[conversationID]; ok && t != nil { + n := t.activeEinoExecuteAbortNote + t.activeEinoExecuteAbortNote = "" + return n + } + return "" +} + +// SetInterruptContinueNote 在发起 ErrInterruptContinue 取消前写入用户补充说明（仅内存）。 +func (m *AgentTaskManager) SetInterruptContinueNote(conversationID, note string) { + conversationID = strings.TrimSpace(conversationID) + if conversationID == "" { + return + } + m.mu.Lock() + defer m.mu.Unlock() + if t, ok := m.tasks[conversationID]; ok && t != nil { + t.InterruptContinueNote = note + } +} + +// TakeInterruptContinueNote 读取并清空补充说明（续跑开始时调用一次）。 +func (m *AgentTaskManager) TakeInterruptContinueNote(conversationID string) string { + conversationID = strings.TrimSpace(conversationID) + if conversationID == "" { + return "" + } + m.mu.Lock() + defer m.mu.Unlock() + if t, ok := m.tasks[conversationID]; ok && t != nil { + n := t.InterruptContinueNote + t.InterruptContinueNote = "" + return n + } + return "" +} + +// BindTaskCancel 在同一运行任务内替换与 context 绑定的 cancel 函数（用于中断后继续时换新 baseCtx）。 +func (m *AgentTaskManager) BindTaskCancel(conversationID string, cancel context.CancelCauseFunc) { + conversationID = strings.TrimSpace(conversationID) + if conversationID == "" || cancel == nil { + return + } + m.mu.Lock() + defer m.mu.Unlock() + if t, ok := m.tasks[conversationID]; ok && t != nil { + t.cancel = func(err error) { + cancel(err) + } + } +} + +// BindAgentRuntimeCancel 登记当前运行段的 Eino 原生 cancel hook。 +func (m *AgentTaskManager) BindAgentRuntimeCancel(conversationID string, cancel func(error) bool) func() { + conversationID = strings.TrimSpace(conversationID) + if conversationID == "" || cancel == nil { + return func() {} + } + m.mu.Lock() + t, ok := m.tasks[conversationID] + if !ok || t == nil { + m.mu.Unlock() + return func() {} + } + t.agentRuntimeCancelVersion++ + version := t.agentRuntimeCancelVersion + t.agentRuntimeCancel = cancel + m.mu.Unlock() + + return func() { + m.mu.Lock() + defer m.mu.Unlock() + if cur, exists := m.tasks[conversationID]; exists && cur != nil && cur.agentRuntimeCancelVersion == version { + cur.agentRuntimeCancel = nil + } + } +} + +// BindAgentTurnLoopInterrupt 登记当前运行任务的 Eino TurnLoop 用户补充入队 hook。 +func (m *AgentTaskManager) BindAgentTurnLoopInterrupt(conversationID string, push func(string) bool) func() { + conversationID = strings.TrimSpace(conversationID) + if conversationID == "" || push == nil { + return func() {} + } + m.mu.Lock() + t, ok := m.tasks[conversationID] + if !ok || t == nil { + m.mu.Unlock() + return func() {} + } + t.agentTurnLoopInterruptVersion++ + version := t.agentTurnLoopInterruptVersion + t.agentTurnLoopInterrupt = push + m.mu.Unlock() + + return func() { + m.mu.Lock() + defer m.mu.Unlock() + if cur, exists := m.tasks[conversationID]; exists && cur != nil && cur.agentTurnLoopInterruptVersion == version { + cur.agentTurnLoopInterrupt = nil + } + } +} + +// ActiveMCPExecutionID 返回当前会话进行中的工具 executionId，无则空串。 +func (m *AgentTaskManager) ActiveMCPExecutionID(conversationID string) string { + conversationID = strings.TrimSpace(conversationID) + if conversationID == "" { + return "" + } + m.mu.RLock() + defer m.mu.RUnlock() + if t, ok := m.tasks[conversationID]; ok && t != nil { + return strings.TrimSpace(t.ActiveMCPExecutionID) + } + return "" +} + +// CompletedTask 已完成的任务（用于历史记录） +type CompletedTask struct { + ConversationID string `json:"conversationId"` + Title string `json:"title,omitempty"` + Message string `json:"message,omitempty"` + StartedAt time.Time `json:"startedAt"` + CompletedAt time.Time `json:"completedAt"` + Status string `json:"status"` +} + +// AgentTaskManager 管理正在运行的Agent任务 +type AgentTaskManager struct { + mu sync.RWMutex + tasks map[string]*AgentTask + completedTasks []*CompletedTask // 最近完成的任务历史 + maxHistorySize int // 最大历史记录数 + historyRetention time.Duration // 历史记录保留时间 + eventBus *TaskEventBus // 可选：任务结束时关闭镜像 SSE 订阅 + // toolCanceler 在用户整轮停止任务或会话结束时终止该会话仍在运行的 MCP 工具（非「中断并继续」）。 + toolCanceler func(conversationID string) +} + +const ( + // cancellingStuckThreshold 处于「取消中」超过此时长则强制从运行列表移除。正常取消会在当前步骤内返回， + // 超过则视为卡住，尽快释放会话。常见做法多为 30–60s 内释放。 + cancellingStuckThreshold = 45 * time.Second + // cancellingStuckThresholdLegacy 未记录 CancellingAt 时用 StartedAt 判断的兜底时长 + cancellingStuckThresholdLegacy = 2 * time.Minute + cleanupInterval = 15 * time.Second // 与上面阈值配合，最长约 60s 内移除 +) + +// NewAgentTaskManager 创建任务管理器 +func NewAgentTaskManager() *AgentTaskManager { + m := &AgentTaskManager{ + tasks: make(map[string]*AgentTask), + completedTasks: make([]*CompletedTask, 0), + maxHistorySize: 50, // 最多保留50条历史记录 + historyRetention: 24 * time.Hour, // 保留24小时 + } + go m.runStuckCancellingCleanup() + return m +} + +// SetTaskEventBus 设置任务事件总线（与 AgentHandler 共用同一实例）。 +func (m *AgentTaskManager) SetTaskEventBus(b *TaskEventBus) { + m.mu.Lock() + defer m.mu.Unlock() + m.eventBus = b +} + +// SetToolCanceler 设置整轮停止任务/会话结束时终止仍在运行 MCP 工具的回调（由 AgentHandler 注入）。 +func (m *AgentTaskManager) SetToolCanceler(fn func(conversationID string)) { + m.mu.Lock() + defer m.mu.Unlock() + m.toolCanceler = fn +} + +// GetTask 返回运行中任务（无则 nil）。 +func (m *AgentTaskManager) GetTask(conversationID string) *AgentTask { + m.mu.RLock() + defer m.mu.RUnlock() + return m.tasks[conversationID] +} + +// GetTaskSnapshot 返回运行任务的只读副本，供状态展示使用，避免锁外读取可变任务字段。 +func (m *AgentTaskManager) GetTaskSnapshot(conversationID string) *AgentTask { + m.mu.RLock() + defer m.mu.RUnlock() + task := m.tasks[conversationID] + if task == nil { + return nil + } + snapshot := *task + return &snapshot +} + +// runStuckCancellingCleanup 定期将长时间处于「取消中」的任务强制结束，避免卡住无法发新消息 +func (m *AgentTaskManager) runStuckCancellingCleanup() { + ticker := time.NewTicker(cleanupInterval) + defer ticker.Stop() + for range ticker.C { + m.cleanupStuckCancelling() + } +} + +func (m *AgentTaskManager) cleanupStuckCancelling() { + m.mu.Lock() + var toFinish []string + now := time.Now() + for id, task := range m.tasks { + if task.Status != "cancelling" { + continue + } + var elapsed time.Duration + if !task.CancellingAt.IsZero() { + elapsed = now.Sub(task.CancellingAt) + if elapsed < cancellingStuckThreshold { + continue + } + } else { + elapsed = now.Sub(task.StartedAt) + if elapsed < cancellingStuckThresholdLegacy { + continue + } + } + toFinish = append(toFinish, id) + } + m.mu.Unlock() + for _, id := range toFinish { + m.FinishTask(id, "cancelled") + } +} + +// StartTask 注册并开始一个新的任务 +func (m *AgentTaskManager) StartTask(conversationID, message string, cancel context.CancelCauseFunc) (*AgentTask, error) { + m.mu.Lock() + defer m.mu.Unlock() + + if _, exists := m.tasks[conversationID]; exists { + return nil, ErrTaskAlreadyRunning + } + + task := &AgentTask{ + ConversationID: conversationID, + Message: message, + StartedAt: time.Now(), + Status: "running", + cancel: func(err error) { + if cancel != nil { + cancel(err) + } + }, + } + + m.tasks[conversationID] = task + task.hitlCognition = &hitlCognitionState{UserMessage: strings.TrimSpace(message)} + return task, nil +} + +// CancelTask 取消指定会话的任务。若任务已在取消中，仍返回 (true, nil) 以便接口幂等、前端不报错。 +func (m *AgentTaskManager) CancelTask(conversationID string, cause error) (bool, error) { + m.mu.Lock() + task, exists := m.tasks[conversationID] + if !exists { + m.mu.Unlock() + return false, nil + } + + // 如果已经处于取消流程，视为成功（幂等），避免前端重复点击报「未找到任务」 + if task.Status == "cancelling" { + m.mu.Unlock() + return true, nil + } + + // ErrInterruptContinue：仅掐断当前推理步骤，随后由处理器续跑，不进入长时间「取消中」态。 + if cause != nil && errors.Is(cause, multiagent.ErrInterruptContinue) { + task.Status = "running" + } else { + task.Status = "cancelling" + task.CancellingAt = time.Now() + } + if cause != nil && errors.Is(cause, ErrTaskCancelled) { + task.InterruptContinueNote = "" + } + cancel := task.cancel + if cause == nil { + cause = ErrTaskCancelled + } + interruptPush := task.agentTurnLoopInterrupt + interruptNote := task.InterruptContinueNote + runtimeCancel := task.agentRuntimeCancel + var toolCanceler func(string) + if errors.Is(cause, ErrTaskCancelled) { + toolCanceler = m.toolCanceler + } + m.mu.Unlock() + + if errors.Is(cause, multiagent.ErrInterruptContinue) && interruptPush != nil && interruptPush(interruptNote) { + m.mu.Lock() + if cur, exists := m.tasks[conversationID]; exists && cur != nil { + cur.InterruptContinueNote = "" + } + m.mu.Unlock() + return true, nil + } + + runtimeHandled := false + if runtimeCancel != nil { + runtimeHandled = runtimeCancel(cause) + } + if cancel != nil && !runtimeHandled { + cancel(cause) + } + if toolCanceler != nil { + toolCanceler(conversationID) + } + return true, nil +} + +// UpdateTaskStatus 更新任务状态但不删除任务（用于在发送事件前更新状态） +func (m *AgentTaskManager) UpdateTaskStatus(conversationID string, status string) { + m.mu.Lock() + defer m.mu.Unlock() + + task, exists := m.tasks[conversationID] + if !exists { + return + } + + if status != "" { + task.Status = status + } +} + +// FinishTask 完成任务并从管理器中移除 +func (m *AgentTaskManager) FinishTask(conversationID string, finalStatus string) { + m.mu.Lock() + task, exists := m.tasks[conversationID] + if !exists { + m.mu.Unlock() + return + } + + if finalStatus != "" { + task.Status = finalStatus + } + toolCanceler := m.toolCanceler + activeEinoExecuteCancel := task.activeEinoExecuteCancel + + // 保存到历史记录 + completedTask := &CompletedTask{ + ConversationID: task.ConversationID, + Message: task.Message, + StartedAt: task.StartedAt, + CompletedAt: time.Now(), + Status: finalStatus, + } + + // 添加到历史记录 + m.completedTasks = append(m.completedTasks, completedTask) + + // 清理过期和过多的历史记录 + m.cleanupHistory() + + // 从运行任务中移除 + delete(m.tasks, conversationID) + bus := m.eventBus + m.mu.Unlock() + if toolCanceler != nil { + toolCanceler(conversationID) + } + if activeEinoExecuteCancel != nil { + activeEinoExecuteCancel() + } + if bus != nil { + bus.CloseConversation(conversationID) + } +} + +// cleanupHistory 清理过期的历史记录 +func (m *AgentTaskManager) cleanupHistory() { + now := time.Now() + cutoffTime := now.Add(-m.historyRetention) + + // 过滤掉过期的记录 + validTasks := make([]*CompletedTask, 0, len(m.completedTasks)) + for _, task := range m.completedTasks { + if task.CompletedAt.After(cutoffTime) { + validTasks = append(validTasks, task) + } + } + + // 如果仍然超过最大数量，只保留最新的 + if len(validTasks) > m.maxHistorySize { + // 按完成时间排序，保留最新的 + // 由于是追加的，最新的在最后，所以直接取最后N个 + start := len(validTasks) - m.maxHistorySize + validTasks = validTasks[start:] + } + + m.completedTasks = validTasks +} + +// GetActiveTasks 返回所有正在运行的任务 +func (m *AgentTaskManager) GetActiveTasks() []*AgentTask { + m.mu.RLock() + defer m.mu.RUnlock() + + result := make([]*AgentTask, 0, len(m.tasks)) + for _, task := range m.tasks { + result = append(result, &AgentTask{ + ConversationID: task.ConversationID, + Message: task.Message, + StartedAt: task.StartedAt, + Status: task.Status, + }) + } + return result +} + +// GetCompletedTasks 返回最近完成的任务历史 +func (m *AgentTaskManager) GetCompletedTasks() []*CompletedTask { + m.mu.RLock() + defer m.mu.RUnlock() + + // 清理过期记录（只读锁，不影响其他操作） + // 注意：这里不能直接调用cleanupHistory，因为需要写锁 + // 所以返回时过滤过期记录 + now := time.Now() + cutoffTime := now.Add(-m.historyRetention) + + result := make([]*CompletedTask, 0, len(m.completedTasks)) + for _, task := range m.completedTasks { + if task.CompletedAt.After(cutoffTime) { + result = append(result, task) + } + } + + // 按完成时间倒序排序（最新的在前） + // 由于是追加的，最新的在最后，需要反转 + for i, j := 0, len(result)-1; i < j; i, j = i+1, j-1 { + result[i], result[j] = result[j], result[i] + } + + // 限制返回数量 + if len(result) > m.maxHistorySize { + result = result[:m.maxHistorySize] + } + + return result +} diff --git a/internal/handler/task_manager_eino_execute_test.go b/internal/handler/task_manager_eino_execute_test.go new file mode 100644 index 00000000..d678aff6 --- /dev/null +++ b/internal/handler/task_manager_eino_execute_test.go @@ -0,0 +1,56 @@ +package handler + +import ( + "context" + "testing" + "time" +) + +func TestAbortActiveEinoExecute(t *testing.T) { + m := NewAgentTaskManager() + conv := "conv-eino-exec-abort" + ctx, cancel := context.WithCancel(context.Background()) + _, err := m.StartTask(conv, "test", func(error) {}) + if err != nil { + t.Fatalf("StartTask: %v", err) + } + m.RegisterActiveEinoExecute(conv, cancel) + + done := make(chan struct{}) + go func() { + <-ctx.Done() + close(done) + }() + + if !m.AbortActiveEinoExecute(conv, "跳过域名收集") { + t.Fatal("expected abort to succeed") + } + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("execute cancel did not propagate") + } + if got := m.TakeEinoExecuteAbortNote(conv); got != "跳过域名收集" { + t.Fatalf("abort note = %q, want 跳过域名收集", got) + } + m.UnregisterActiveEinoExecute(conv) + if m.AbortActiveEinoExecute(conv, "") { + t.Fatal("second abort should fail when no active execute") + } +} + +func TestConversationIDForActiveMCPExecution(t *testing.T) { + m := NewAgentTaskManager() + conv := "conv-mcp-exec" + _, err := m.StartTask(conv, "test", func(error) {}) + if err != nil { + t.Fatalf("StartTask: %v", err) + } + m.RegisterRunningTool(conv, "exec-123") + if got := m.ConversationIDForActiveMCPExecution("exec-123"); got != conv { + t.Fatalf("got %q, want %q", got, conv) + } + if got := m.ConversationIDForActiveMCPExecution("missing"); got != "" { + t.Fatalf("missing should be empty, got %q", got) + } +} diff --git a/internal/handler/task_manager_tool_cancel_test.go b/internal/handler/task_manager_tool_cancel_test.go new file mode 100644 index 00000000..51b411b0 --- /dev/null +++ b/internal/handler/task_manager_tool_cancel_test.go @@ -0,0 +1,246 @@ +package handler + +import ( + "context" + "errors" + "testing" + + "cyberstrike-ai/internal/multiagent" +) + +func TestCancelTaskInvokesToolCancelerOnFullStop(t *testing.T) { + tm := NewAgentTaskManager() + called := false + tm.SetToolCanceler(func(conversationID string) { + if conversationID == "conv-1" { + called = true + } + }) + + _, cancel := context.WithCancelCause(context.Background()) + _, err := tm.StartTask("conv-1", "hello", cancel) + if err != nil { + t.Fatalf("StartTask: %v", err) + } + + ok, err := tm.CancelTask("conv-1", ErrTaskCancelled) + if err != nil || !ok { + t.Fatalf("CancelTask: ok=%v err=%v", ok, err) + } + if !called { + t.Fatal("expected tool canceler to be invoked on full task cancel") + } +} + +func TestCancelTaskUsesAgentRuntimeCancelAsPrimaryPath(t *testing.T) { + tm := NewAgentTaskManager() + var order []string + tm.SetToolCanceler(func(conversationID string) { + if conversationID == "conv-native" { + order = append(order, "tool") + } + }) + + _, cancel := context.WithCancelCause(context.Background()) + if _, err := tm.StartTask("conv-native", "hello", func(err error) { + order = append(order, "context") + cancel(err) + }); err != nil { + t.Fatalf("StartTask: %v", err) + } + unregister := tm.BindAgentRuntimeCancel("conv-native", func(err error) bool { + if !errors.Is(err, ErrTaskCancelled) { + t.Fatalf("runtime cancel got %v", err) + } + order = append(order, "runtime") + return true + }) + defer unregister() + + ok, err := tm.CancelTask("conv-native", ErrTaskCancelled) + if err != nil || !ok { + t.Fatalf("CancelTask: ok=%v err=%v", ok, err) + } + want := []string{"runtime", "tool"} + if len(order) != len(want) { + t.Fatalf("order length got %d want %d: %#v", len(order), len(want), order) + } + for i := range want { + if order[i] != want[i] { + t.Fatalf("order[%d] got %q want %q; full=%#v", i, order[i], want[i], order) + } + } +} + +func TestCancelTaskFallsBackToContextWhenAgentRuntimeCancelMisses(t *testing.T) { + tm := NewAgentTaskManager() + var order []string + + _, cancel := context.WithCancelCause(context.Background()) + if _, err := tm.StartTask("conv-fallback", "hello", func(err error) { + order = append(order, "context") + cancel(err) + }); err != nil { + t.Fatalf("StartTask: %v", err) + } + unregister := tm.BindAgentRuntimeCancel("conv-fallback", func(err error) bool { + order = append(order, "runtime") + return false + }) + defer unregister() + + ok, err := tm.CancelTask("conv-fallback", ErrTaskCancelled) + if err != nil || !ok { + t.Fatalf("CancelTask: ok=%v err=%v", ok, err) + } + want := []string{"runtime", "context"} + if len(order) != len(want) { + t.Fatalf("order length got %d want %d: %#v", len(order), len(want), order) + } + for i := range want { + if order[i] != want[i] { + t.Fatalf("order[%d] got %q want %q; full=%#v", i, order[i], want[i], order) + } + } +} + +func TestCancelTaskSkipsToolCancelerOnInterruptContinue(t *testing.T) { + tm := NewAgentTaskManager() + called := false + tm.SetToolCanceler(func(conversationID string) { + called = true + }) + + _, cancel := context.WithCancelCause(context.Background()) + _, err := tm.StartTask("conv-1", "hello", cancel) + if err != nil { + t.Fatalf("StartTask: %v", err) + } + + ok, err := tm.CancelTask("conv-1", multiagent.ErrInterruptContinue) + if err != nil || !ok { + t.Fatalf("CancelTask: ok=%v err=%v", ok, err) + } + if called { + t.Fatal("tool canceler must not run for interrupt-continue") + } +} + +func TestCancelTaskPushesInterruptContinueToTurnLoopFirst(t *testing.T) { + tm := NewAgentTaskManager() + ctx, cancel := context.WithCancelCause(context.Background()) + if _, err := tm.StartTask("conv-turn", "hello", cancel); err != nil { + t.Fatalf("StartTask: %v", err) + } + tm.SetInterruptContinueNote("conv-turn", "focus ssh") + + var gotNote string + unregister := tm.BindAgentTurnLoopInterrupt("conv-turn", func(note string) bool { + gotNote = note + return true + }) + defer unregister() + + ok, err := tm.CancelTask("conv-turn", multiagent.ErrInterruptContinue) + if err != nil || !ok { + t.Fatalf("CancelTask: ok=%v err=%v", ok, err) + } + if gotNote != "focus ssh" { + t.Fatalf("turn loop note = %q, want focus ssh", gotNote) + } + if cause := context.Cause(ctx); cause != nil { + t.Fatalf("context should not be cancelled when turn loop accepted interrupt, got %v", cause) + } + if note := tm.TakeInterruptContinueNote("conv-turn"); note != "" { + t.Fatalf("interrupt note should be consumed after turn loop push, got %q", note) + } +} + +func TestCancelTaskFallsBackWhenTurnLoopInterruptRejects(t *testing.T) { + tm := NewAgentTaskManager() + var order []string + + _, cancel := context.WithCancelCause(context.Background()) + if _, err := tm.StartTask("conv-turn-fallback", "hello", func(err error) { + order = append(order, "context") + cancel(err) + }); err != nil { + t.Fatalf("StartTask: %v", err) + } + tm.SetInterruptContinueNote("conv-turn-fallback", "fallback note") + unregisterTurn := tm.BindAgentTurnLoopInterrupt("conv-turn-fallback", func(note string) bool { + order = append(order, "turn") + if note != "fallback note" { + t.Fatalf("turn loop note = %q, want fallback note", note) + } + return false + }) + defer unregisterTurn() + unregisterRuntime := tm.BindAgentRuntimeCancel("conv-turn-fallback", func(err error) bool { + order = append(order, "runtime") + return false + }) + defer unregisterRuntime() + + ok, err := tm.CancelTask("conv-turn-fallback", multiagent.ErrInterruptContinue) + if err != nil || !ok { + t.Fatalf("CancelTask: ok=%v err=%v", ok, err) + } + want := []string{"turn", "runtime", "context"} + if len(order) != len(want) { + t.Fatalf("order length got %d want %d: %#v", len(order), len(want), order) + } + for i := range want { + if order[i] != want[i] { + t.Fatalf("order[%d] got %q want %q; full=%#v", i, order[i], want[i], order) + } + } + if note := tm.TakeInterruptContinueNote("conv-turn-fallback"); note != "fallback note" { + t.Fatalf("interrupt note should remain for fallback rerun, got %q", note) + } +} + +func TestCancelTaskDefaultCauseIsTaskCancelled(t *testing.T) { + tm := NewAgentTaskManager() + var gotCause error + tm.SetToolCanceler(func(conversationID string) { + if conversationID == "conv-2" { + gotCause = ErrTaskCancelled + } + }) + + ctx, cancel := context.WithCancelCause(context.Background()) + if _, err := tm.StartTask("conv-2", "hello", cancel); err != nil { + t.Fatalf("StartTask: %v", err) + } + + if _, err := tm.CancelTask("conv-2", nil); err != nil { + t.Fatalf("CancelTask: %v", err) + } + if !errors.Is(context.Cause(ctx), ErrTaskCancelled) { + t.Fatalf("expected ErrTaskCancelled cause, got %v", context.Cause(ctx)) + } + if gotCause != ErrTaskCancelled { + t.Fatalf("expected tool canceler path for default cancel cause") + } +} + +func TestFinishTaskInvokesToolCancelerOnSessionEnd(t *testing.T) { + tm := NewAgentTaskManager() + calls := 0 + tm.SetToolCanceler(func(conversationID string) { + if conversationID == "conv-3" { + calls++ + } + }) + + _, cancel := context.WithCancelCause(context.Background()) + if _, err := tm.StartTask("conv-3", "hello", cancel); err != nil { + t.Fatalf("StartTask: %v", err) + } + + tm.FinishTask("conv-3", "completed") + if calls != 1 { + t.Fatalf("expected one tool cleanup on FinishTask, got %d", calls) + } +} diff --git a/internal/handler/terminal.go b/internal/handler/terminal.go new file mode 100644 index 00000000..3c3c53fb --- /dev/null +++ b/internal/handler/terminal.go @@ -0,0 +1,257 @@ +package handler + +import ( + "bytes" + "context" + "encoding/json" + "net/http" + "os" + "os/exec" + "path/filepath" + "runtime" + "strings" + "time" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +const ( + terminalMaxCommandLen = 4096 + terminalMaxOutputLen = 256 * 1024 // 256KB + terminalTimeout = 30 * time.Minute +) + +// TerminalHandler 处理系统设置中的终端命令执行 +type TerminalHandler struct { + logger *zap.Logger +} + +// maskTerminalCommand 对可能包含敏感信息的终端命令做脱敏，避免在日志中直接记录密码等内容 +func maskTerminalCommand(cmd string) string { + trimmed := strings.TrimSpace(cmd) + lower := strings.ToLower(trimmed) + if strings.Contains(lower, "sudo") || strings.Contains(lower, "password") { + return "[masked sensitive terminal command]" + } + if len(trimmed) > 256 { + return trimmed[:256] + "..." + } + return trimmed +} + +// NewTerminalHandler 创建终端处理器 +func NewTerminalHandler(logger *zap.Logger) *TerminalHandler { + return &TerminalHandler{logger: logger} +} + +// RunCommandRequest 执行命令请求 +type RunCommandRequest struct { + Command string `json:"command"` + Shell string `json:"shell,omitempty"` + Cwd string `json:"cwd,omitempty"` +} + +// RunCommandResponse 执行命令响应 +type RunCommandResponse struct { + Stdout string `json:"stdout"` + Stderr string `json:"stderr"` + ExitCode int `json:"exit_code"` + Error string `json:"error,omitempty"` +} + +// RunCommand 执行终端命令（需登录） +func (h *TerminalHandler) RunCommand(c *gin.Context) { + var req RunCommandRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "请求体无效，需要 command 字段"}) + return + } + + cmdStr := strings.TrimSpace(req.Command) + if cmdStr == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "command 不能为空"}) + return + } + if len(cmdStr) > terminalMaxCommandLen { + c.JSON(http.StatusBadRequest, gin.H{"error": "命令过长"}) + return + } + + shell := req.Shell + if shell == "" { + if runtime.GOOS == "windows" { + shell = "cmd" + } else { + shell = "sh" + } + } + + ctx, cancel := context.WithTimeout(c.Request.Context(), terminalTimeout) + defer cancel() + + var cmd *exec.Cmd + if runtime.GOOS == "windows" { + cmd = exec.CommandContext(ctx, "cmd", "/c", cmdStr) + } else { + cmd = exec.CommandContext(ctx, shell, "-c", cmdStr) + // 无 TTY 时设置 COLUMNS/TERM，使 ping 等工具的 usage 排版与真实终端一致 + cmd.Env = append(os.Environ(), "COLUMNS=256", "LINES=40", "TERM=xterm-256color") + } + + if req.Cwd != "" { + absCwd, err := filepath.Abs(req.Cwd) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "工作目录无效"}) + return + } + cur, _ := os.Getwd() + curAbs, _ := filepath.Abs(cur) + rel, err := filepath.Rel(curAbs, absCwd) + if err != nil || strings.HasPrefix(rel, "..") || rel == ".." { + c.JSON(http.StatusBadRequest, gin.H{"error": "工作目录必须在当前进程目录下"}) + return + } + cmd.Dir = absCwd + } + + var stdout, stderr bytes.Buffer + cmd.Stdout = &stdout + cmd.Stderr = &stderr + + err := cmd.Run() + stdoutBytes := stdout.Bytes() + stderrBytes := stderr.Bytes() + + // 限制输出长度，防止内存占用过大（复制后截断，避免修改原 buffer） + truncSuffix := []byte("\n...(输出已截断)\n") + if len(stdoutBytes) > terminalMaxOutputLen { + tmp := make([]byte, terminalMaxOutputLen+len(truncSuffix)) + n := copy(tmp, stdoutBytes[:terminalMaxOutputLen]) + copy(tmp[n:], truncSuffix) + stdoutBytes = tmp + } + if len(stderrBytes) > terminalMaxOutputLen { + tmp := make([]byte, terminalMaxOutputLen+len(truncSuffix)) + n := copy(tmp, stderrBytes[:terminalMaxOutputLen]) + copy(tmp[n:], truncSuffix) + stderrBytes = tmp + } + + exitCode := 0 + if err != nil { + if exitErr, ok := err.(*exec.ExitError); ok { + exitCode = exitErr.ExitCode() + } else { + exitCode = -1 + } + if ctx.Err() == context.DeadlineExceeded { + so := strings.ReplaceAll(string(stdoutBytes), "\r\n", "\n") + so = strings.ReplaceAll(so, "\r", "\n") + se := strings.ReplaceAll(string(stderrBytes), "\r\n", "\n") + se = strings.ReplaceAll(se, "\r", "\n") + resp := RunCommandResponse{ + Stdout: so, + Stderr: se, + ExitCode: -1, + Error: "命令执行超时（" + terminalTimeout.String() + "）", + } + c.JSON(http.StatusOK, resp) + return + } + h.logger.Debug("终端命令执行异常", zap.String("command", maskTerminalCommand(cmdStr)), zap.Error(err)) + } + + // 统一为 \n，避免前端因 \r 出现错位/对角线排版 + stdoutStr := strings.ReplaceAll(string(stdoutBytes), "\r\n", "\n") + stdoutStr = strings.ReplaceAll(stdoutStr, "\r", "\n") + stderrStr := strings.ReplaceAll(string(stderrBytes), "\r\n", "\n") + stderrStr = strings.ReplaceAll(stderrStr, "\r", "\n") + + resp := RunCommandResponse{ + Stdout: stdoutStr, + Stderr: stderrStr, + ExitCode: exitCode, + } + if err != nil && exitCode != 0 { + resp.Error = err.Error() + } + c.JSON(http.StatusOK, resp) +} + +// streamEvent SSE 事件 +type streamEvent struct { + T string `json:"t"` // "out" | "err" | "exit" + D string `json:"d,omitempty"` + C int `json:"c"` // exit code（不用 omitempty，否则 0 不序列化导致前端显示 [exit undefined]） +} + +// RunCommandStream 流式执行命令，输出实时推送到前端（SSE） +func (h *TerminalHandler) RunCommandStream(c *gin.Context) { + var req RunCommandRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "请求体无效，需要 command 字段"}) + return + } + cmdStr := strings.TrimSpace(req.Command) + if cmdStr == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "command 不能为空"}) + return + } + if len(cmdStr) > terminalMaxCommandLen { + c.JSON(http.StatusBadRequest, gin.H{"error": "命令过长"}) + return + } + shell := req.Shell + if shell == "" { + if runtime.GOOS == "windows" { + shell = "cmd" + } else { + shell = "sh" + } + } + ctx, cancel := context.WithTimeout(c.Request.Context(), terminalTimeout) + defer cancel() + + var cmd *exec.Cmd + if runtime.GOOS == "windows" { + cmd = exec.CommandContext(ctx, "cmd", "/c", cmdStr) + } else { + cmd = exec.CommandContext(ctx, shell, "-c", cmdStr) + cmd.Env = append(os.Environ(), "COLUMNS=256", "LINES=40", "TERM=xterm-256color") + } + if req.Cwd != "" { + absCwd, err := filepath.Abs(req.Cwd) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "工作目录无效"}) + return + } + cur, _ := os.Getwd() + curAbs, _ := filepath.Abs(cur) + rel, err := filepath.Rel(curAbs, absCwd) + if err != nil || strings.HasPrefix(rel, "..") || rel == ".." { + c.JSON(http.StatusBadRequest, gin.H{"error": "工作目录必须在当前进程目录下"}) + return + } + cmd.Dir = absCwd + } + + c.Header("Content-Type", "text/event-stream") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + c.Header("X-Accel-Buffering", "no") + c.Writer.WriteHeader(http.StatusOK) + flusher, ok := c.Writer.(http.Flusher) + if !ok { + cancel() + return + } + + sendEvent := func(ev streamEvent) { + body, _ := json.Marshal(ev) + c.SSEvent("", string(body)) + flusher.Flush() + } + + _ = runCommandStreamImpl(cmd, sendEvent, ctx) +} diff --git a/internal/handler/terminal_stream_unix.go b/internal/handler/terminal_stream_unix.go new file mode 100644 index 00000000..e8ab8c47 --- /dev/null +++ b/internal/handler/terminal_stream_unix.go @@ -0,0 +1,47 @@ +//go:build !windows + +package handler + +import ( + "bufio" + "context" + "os/exec" + "strings" + + "github.com/creack/pty" +) + +const ptyCols = 256 +const ptyRows = 40 + +// runCommandStreamImpl 在 Unix 下用 PTY 执行，使 ping 等命令按终端宽度排版（isatty 为真） +func runCommandStreamImpl(cmd *exec.Cmd, sendEvent func(streamEvent), ctx context.Context) int { + ptmx, err := pty.StartWithSize(cmd, &pty.Winsize{Cols: ptyCols, Rows: ptyRows}) + if err != nil { + sendEvent(streamEvent{T: "exit", C: -1}) + return -1 + } + defer ptmx.Close() + + normalize := func(s string) string { + s = strings.ReplaceAll(s, "\r\n", "\n") + return strings.ReplaceAll(s, "\r", "\n") + } + sc := bufio.NewScanner(ptmx) + for sc.Scan() { + sendEvent(streamEvent{T: "out", D: normalize(sc.Text())}) + } + exitCode := 0 + if err := cmd.Wait(); err != nil { + if exitErr, ok := err.(*exec.ExitError); ok { + exitCode = exitErr.ExitCode() + } else { + exitCode = -1 + } + } + if ctx.Err() == context.DeadlineExceeded { + exitCode = -1 + } + sendEvent(streamEvent{T: "exit", C: exitCode}) + return exitCode +} diff --git a/internal/handler/terminal_stream_windows.go b/internal/handler/terminal_stream_windows.go new file mode 100644 index 00000000..24e430a5 --- /dev/null +++ b/internal/handler/terminal_stream_windows.go @@ -0,0 +1,66 @@ +//go:build windows + +package handler + +import ( + "bufio" + "context" + "os/exec" + "strings" + "sync" +) + +// runCommandStreamImpl 在 Windows 下用 stdout/stderr 管道执行 +func runCommandStreamImpl(cmd *exec.Cmd, sendEvent func(streamEvent), ctx context.Context) int { + stdoutPipe, err := cmd.StdoutPipe() + if err != nil { + sendEvent(streamEvent{T: "exit", C: -1}) + return -1 + } + stderrPipe, err := cmd.StderrPipe() + if err != nil { + sendEvent(streamEvent{T: "exit", C: -1}) + return -1 + } + if err := cmd.Start(); err != nil { + sendEvent(streamEvent{T: "exit", C: -1}) + return -1 + } + + normalize := func(s string) string { + s = strings.ReplaceAll(s, "\r\n", "\n") + return strings.ReplaceAll(s, "\r", "\n") + } + + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + sc := bufio.NewScanner(stdoutPipe) + for sc.Scan() { + sendEvent(streamEvent{T: "out", D: normalize(sc.Text())}) + } + }() + go func() { + defer wg.Done() + sc := bufio.NewScanner(stderrPipe) + for sc.Scan() { + sendEvent(streamEvent{T: "err", D: normalize(sc.Text())}) + } + }() + + wg.Wait() + exitCode := 0 + if err := cmd.Wait(); err != nil { + if exitErr, ok := err.(*exec.ExitError); ok { + exitCode = exitErr.ExitCode() + } else { + exitCode = -1 + } + } + if ctx.Err() == context.DeadlineExceeded { + exitCode = -1 + } + sendEvent(streamEvent{T: "exit", C: exitCode}) + return exitCode +} diff --git a/internal/handler/terminal_ws_unix.go b/internal/handler/terminal_ws_unix.go new file mode 100644 index 00000000..0f446d83 --- /dev/null +++ b/internal/handler/terminal_ws_unix.go @@ -0,0 +1,111 @@ +//go:build !windows + +package handler + +import ( + "encoding/json" + "net/http" + "os" + "os/exec" + "time" + + "github.com/creack/pty" + "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" +) + +// terminalResize is sent by the frontend when the xterm.js terminal is resized. +type terminalResize struct { + Type string `json:"type"` + Cols uint16 `json:"cols"` + Rows uint16 `json:"rows"` +} + +// wsUpgrader 仅用于系统设置中的终端 WebSocket，会复用已有的登录保护（JWT 中间件在上层路由组） +var wsUpgrader = websocket.Upgrader{ + CheckOrigin: func(r *http.Request) bool { + // 由于已在 Gin 路由层做了认证，这里放宽 Origin，方便在同一域名下通过 HTTPS/WSS 访问 + return true + }, +} + +// RunCommandWS 提供真正交互式 Shell：基于 WebSocket + PTY 的长会话 +// 前端建立 WebSocket 连接后，所有键盘输入都会透传到 Shell，Shell 的输出也会实时写回前端。 +func (h *TerminalHandler) RunCommandWS(c *gin.Context) { + conn, err := wsUpgrader.Upgrade(c.Writer, c.Request, nil) + if err != nil { + return + } + defer conn.Close() + + // 启动交互式 Shell，这里优先使用 bash，找不到则退回 sh + shell := "bash" + if _, err := exec.LookPath(shell); err != nil { + shell = "sh" + } + cmd := exec.Command(shell) + cmd.Env = append(os.Environ(), + "COLUMNS=80", + "LINES=24", + "TERM=xterm-256color", + ) + + // Use 80x24 as a safe default; the frontend will send the actual size immediately after connecting. + ptmx, err := pty.StartWithSize(cmd, &pty.Winsize{Cols: 80, Rows: 24}) + if err != nil { + return + } + defer ptmx.Close() + + // Shell -> WebSocket：将 PTY 输出实时发给前端 + doneChan := make(chan struct{}) + go func() { + buf := make([]byte, 4096) + for { + n, err := ptmx.Read(buf) + if n > 0 { + _ = conn.WriteMessage(websocket.BinaryMessage, buf[:n]) + } + if err != nil { + break + } + } + close(doneChan) + }() + + // WebSocket -> Shell：将前端输入写入 PTY（包括 sudo 密码、Ctrl+C 等） + conn.SetReadLimit(64 * 1024) + _ = conn.SetReadDeadline(time.Now().Add(terminalTimeout)) + conn.SetPongHandler(func(string) error { + _ = conn.SetReadDeadline(time.Now().Add(terminalTimeout)) + return nil + }) + + for { + msgType, data, err := conn.ReadMessage() + if err != nil { + _ = cmd.Process.Kill() + break + } + if msgType != websocket.TextMessage && msgType != websocket.BinaryMessage { + continue + } + if len(data) == 0 { + continue + } + // Check if this is a resize message (JSON with type:"resize") + if msgType == websocket.TextMessage && len(data) > 0 && data[0] == '{' { + var resize terminalResize + if json.Unmarshal(data, &resize) == nil && resize.Type == "resize" && resize.Cols > 0 && resize.Rows > 0 { + _ = pty.Setsize(ptmx, &pty.Winsize{Cols: resize.Cols, Rows: resize.Rows}) + continue + } + } + if _, err := ptmx.Write(data); err != nil { + _ = cmd.Process.Kill() + break + } + } + + <-doneChan +} diff --git a/internal/handler/terminal_ws_windows.go b/internal/handler/terminal_ws_windows.go new file mode 100644 index 00000000..2d71fa6b --- /dev/null +++ b/internal/handler/terminal_ws_windows.go @@ -0,0 +1,16 @@ +//go:build windows + +package handler + +import ( + "net/http" + + "github.com/gin-gonic/gin" +) + +// RunCommandWS 交互式 PTY 终端依赖 Unix PTY（见 terminal_ws_unix.go）；Windows 暂不支持。 +func (h *TerminalHandler) RunCommandWS(c *gin.Context) { + c.JSON(http.StatusNotImplemented, gin.H{ + "error": "Interactive WebSocket terminal is not supported on Windows; use POST /terminal/run or /terminal/run/stream instead.", + }) +} diff --git a/internal/handler/vulnerability.go b/internal/handler/vulnerability.go new file mode 100644 index 00000000..ab1eef22 --- /dev/null +++ b/internal/handler/vulnerability.go @@ -0,0 +1,602 @@ +package handler + +import ( + "fmt" + "net/http" + "strconv" + "strings" + "time" + + "cyberstrike-ai/internal/audit" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +// VulnerabilityHandler 漏洞处理器 +type VulnerabilityHandler struct { + db *database.DB + logger *zap.Logger + audit *audit.Service +} + +// SetAudit wires platform audit logging. +func (h *VulnerabilityHandler) SetAudit(s *audit.Service) { + h.audit = s +} + +// NewVulnerabilityHandler 创建新的漏洞处理器 +func NewVulnerabilityHandler(db *database.DB, logger *zap.Logger) *VulnerabilityHandler { + return &VulnerabilityHandler{ + db: db, + logger: logger, + } +} + +// CreateVulnerabilityRequest 创建漏洞请求 +type CreateVulnerabilityRequest struct { + ConversationID string `json:"conversation_id" binding:"required"` + ProjectID string `json:"project_id"` + ConversationTag string `json:"conversation_tag"` + TaskTag string `json:"task_tag"` + Title string `json:"title" binding:"required"` + Description string `json:"description"` + Severity string `json:"severity" binding:"required"` + Status string `json:"status"` + Type string `json:"type"` + Target string `json:"target"` + Preconditions string `json:"preconditions"` + ReproSteps string `json:"reproduction_steps"` + Evidence string `json:"evidence"` + Impact string `json:"impact"` + Recommendation string `json:"recommendation"` + RetestNotes string `json:"retest_notes"` +} + +// CreateVulnerability 创建漏洞 +func (h *VulnerabilityHandler) CreateVulnerability(c *gin.Context) { + var req CreateVulnerabilityRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if session, ok := security.CurrentSession(c); ok && session.Scope != database.RBACScopeAll { + if strings.TrimSpace(req.ConversationID) != "" && !h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", strings.TrimSpace(req.ConversationID)) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权在该对话下创建漏洞"}) + return + } + if strings.TrimSpace(req.ProjectID) != "" && !h.db.UserCanAccessResource(session.UserID, session.Scope, "project", strings.TrimSpace(req.ProjectID)) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权在该项目下创建漏洞"}) + return + } + } + + vuln := &database.Vulnerability{ + ConversationID: req.ConversationID, + ProjectID: strings.TrimSpace(req.ProjectID), + ConversationTag: req.ConversationTag, + TaskTag: req.TaskTag, + Title: req.Title, + Description: req.Description, + Severity: req.Severity, + Status: req.Status, + Type: req.Type, + Target: req.Target, + Preconditions: req.Preconditions, + ReproSteps: req.ReproSteps, + Evidence: req.Evidence, + Impact: req.Impact, + Recommendation: req.Recommendation, + RetestNotes: req.RetestNotes, + } + + created, err := h.db.CreateVulnerability(vuln) + if err != nil { + h.logger.Error("创建漏洞失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if session, ok := security.CurrentSession(c); ok { + _ = h.db.SetResourceOwner("vulnerability", created.ID, session.UserID) + _ = h.db.AssignResourceToUser(session.UserID, "vulnerability", created.ID) + } + h.db.NotifyVulnerabilityCreated(created) + + if h.audit != nil { + h.audit.RecordOK(c, "vulnerability", "create", "创建漏洞记录", "vulnerability", created.ID, map[string]interface{}{ + "severity": created.Severity, "title": created.Title, + }) + } + c.JSON(http.StatusOK, created) +} + +// GetVulnerability 获取漏洞 +func (h *VulnerabilityHandler) GetVulnerability(c *gin.Context) { + id := c.Param("id") + + vuln, err := h.db.GetVulnerability(id) + if err != nil { + h.logger.Error("获取漏洞失败", zap.Error(err)) + c.JSON(http.StatusNotFound, gin.H{"error": "漏洞不存在"}) + return + } + + c.JSON(http.StatusOK, vuln) +} + +// ListVulnerabilitiesResponse 漏洞列表响应 +type ListVulnerabilitiesResponse struct { + Vulnerabilities []*database.Vulnerability `json:"vulnerabilities"` + Total int `json:"total"` + Page int `json:"page"` + PageSize int `json:"page_size"` + TotalPages int `json:"total_pages"` +} + +func parseVulnerabilityListFilter(c *gin.Context) database.VulnerabilityListFilter { + q := strings.TrimSpace(c.Query("q")) + if q == "" { + q = strings.TrimSpace(c.Query("search")) + } + return database.VulnerabilityListFilter{ + ProjectID: c.Query("project_id"), + ID: c.Query("id"), + Search: q, + ConversationID: c.Query("conversation_id"), + Severity: c.Query("severity"), + Status: c.Query("status"), + TaskID: c.Query("task_id"), + ConversationTag: c.Query("conversation_tag"), + TaskTag: c.Query("task_tag"), + } +} + +// ListVulnerabilities 列出漏洞 +func (h *VulnerabilityHandler) ListVulnerabilities(c *gin.Context) { + limitStr := c.DefaultQuery("limit", "20") + offsetStr := c.DefaultQuery("offset", "0") + pageStr := c.Query("page") + filter := parseVulnerabilityListFilter(c) + access := vulnerabilityAccessFromContext(c) + + limit, _ := strconv.Atoi(limitStr) + offset, _ := strconv.Atoi(offsetStr) + page := 1 + + // 如果提供了page参数，优先使用page计算offset + if pageStr != "" { + if p, err := strconv.Atoi(pageStr); err == nil && p > 0 { + page = p + offset = (page - 1) * limit + } + } + + if limit <= 0 || limit > 100 { + limit = 20 + } + if offset < 0 { + offset = 0 + } + + // 获取总数 + total, err := h.db.CountVulnerabilitiesForAccess(filter, access) + if err != nil { + h.logger.Error("获取漏洞总数失败", zap.Error(err)) + // 继续执行，使用0作为总数 + total = 0 + } + + // 获取漏洞列表 + vulnerabilities, err := h.db.ListVulnerabilitiesForAccess(limit, offset, filter, access) + if err != nil { + h.logger.Error("获取漏洞列表失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + // 计算总页数 + totalPages := (total + limit - 1) / limit + if totalPages == 0 { + totalPages = 1 + } + + // 如果使用offset计算page，需要重新计算 + if pageStr == "" { + page = (offset / limit) + 1 + } + + response := ListVulnerabilitiesResponse{ + Vulnerabilities: vulnerabilities, + Total: total, + Page: page, + PageSize: limit, + TotalPages: totalPages, + } + + c.JSON(http.StatusOK, response) +} + +// UpdateVulnerabilityRequest 更新漏洞请求 +type UpdateVulnerabilityRequest struct { + ProjectID *string `json:"project_id"` + ConversationTag *string `json:"conversation_tag"` + TaskTag *string `json:"task_tag"` + Title *string `json:"title"` + Description *string `json:"description"` + Severity *string `json:"severity"` + Status *string `json:"status"` + Type *string `json:"type"` + Target *string `json:"target"` + Preconditions *string `json:"preconditions"` + ReproSteps *string `json:"reproduction_steps"` + Evidence *string `json:"evidence"` + Impact *string `json:"impact"` + Recommendation *string `json:"recommendation"` + RetestNotes *string `json:"retest_notes"` +} + +// UpdateVulnerability 更新漏洞 +func (h *VulnerabilityHandler) UpdateVulnerability(c *gin.Context) { + id := c.Param("id") + + var req UpdateVulnerabilityRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + + // 获取现有漏洞 + existing, err := h.db.GetVulnerability(id) + if err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": "漏洞不存在"}) + return + } + + // 更新字段 + if req.ProjectID != nil { + targetProjectID := strings.TrimSpace(*req.ProjectID) + if targetProjectID != "" { + session, ok := security.CurrentSession(c) + if !ok || !h.db.UserCanAccessResource(session.UserID, session.Scope, "project", targetProjectID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权关联目标项目"}) + return + } + } + existing.ProjectID = targetProjectID + } + if req.ConversationTag != nil { + existing.ConversationTag = *req.ConversationTag + } + if req.TaskTag != nil { + existing.TaskTag = *req.TaskTag + } + if req.Title != nil { + existing.Title = *req.Title + } + if req.Description != nil { + existing.Description = *req.Description + } + if req.Severity != nil { + existing.Severity = *req.Severity + } + if req.Status != nil { + existing.Status = *req.Status + } + if req.Type != nil { + existing.Type = *req.Type + } + if req.Target != nil { + existing.Target = *req.Target + } + if req.Preconditions != nil { + existing.Preconditions = *req.Preconditions + } + if req.ReproSteps != nil { + existing.ReproSteps = *req.ReproSteps + } + if req.Evidence != nil { + existing.Evidence = *req.Evidence + } + if req.Impact != nil { + existing.Impact = *req.Impact + } + if req.Recommendation != nil { + existing.Recommendation = *req.Recommendation + } + if req.RetestNotes != nil { + existing.RetestNotes = *req.RetestNotes + } + + if err := h.db.UpdateVulnerability(id, existing); err != nil { + h.logger.Error("更新漏洞失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + // 返回更新后的漏洞 + updated, err := h.db.GetVulnerability(id) + if err != nil { + h.logger.Error("获取更新后的漏洞失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + if h.audit != nil { + h.audit.RecordOK(c, "vulnerability", "update", "更新漏洞记录", "vulnerability", id, map[string]interface{}{ + "severity": updated.Severity, "status": updated.Status, "project_id": updated.ProjectID, + }) + } + c.JSON(http.StatusOK, updated) +} + +// DeleteVulnerability 删除漏洞 +func (h *VulnerabilityHandler) DeleteVulnerability(c *gin.Context) { + id := c.Param("id") + + if err := h.db.DeleteVulnerability(id); err != nil { + h.logger.Error("删除漏洞失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + if h.audit != nil { + h.audit.Record(c, audit.Entry{ + Category: "vulnerability", + Action: "delete", + Result: "success", + ResourceType: "vulnerability", + ResourceID: id, + Message: "删除漏洞记录", + }) + } + + c.JSON(http.StatusOK, gin.H{"message": "删除成功"}) +} + +// BatchDeleteVulnerabilities 按当前筛选条件批量删除漏洞 +func (h *VulnerabilityHandler) BatchDeleteVulnerabilities(c *gin.Context) { + filter := parseVulnerabilityListFilter(c) + access := vulnerabilityAccessFromContext(c) + + total, err := h.db.CountVulnerabilitiesForAccess(filter, access) + if err != nil { + h.logger.Error("统计待删除漏洞失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if total == 0 { + c.JSON(http.StatusOK, gin.H{"message": "当前筛选条件下没有可删除的漏洞", "deleted": 0}) + return + } + + deleted, err := h.db.DeleteVulnerabilitiesByFilterForAccess(filter, access) + if err != nil { + h.logger.Error("批量删除漏洞失败", zap.Error(err), zap.Int("count", total)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + if h.audit != nil { + h.audit.RecordOK(c, "vulnerability", "delete_batch", "批量删除漏洞记录", "vulnerability", "", map[string]interface{}{ + "deleted": deleted, + "filter": filter, + }) + } + + c.JSON(http.StatusOK, gin.H{"message": "批量删除成功", "deleted": deleted}) +} + +// GetVulnerabilityStats 获取漏洞统计 +func (h *VulnerabilityHandler) GetVulnerabilityStats(c *gin.Context) { + filter := parseVulnerabilityListFilter(c) + access := vulnerabilityAccessFromContext(c) + + stats, err := h.db.GetVulnerabilityStatsForAccess(filter, access) + if err != nil { + h.logger.Error("获取漏洞统计失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + c.JSON(http.StatusOK, stats) +} + +// GetVulnerabilityFilterOptions 获取漏洞筛选建议项 +func (h *VulnerabilityHandler) GetVulnerabilityFilterOptions(c *gin.Context) { + options, err := h.db.GetVulnerabilityFilterOptionsForAccess(vulnerabilityAccessFromContext(c)) + if err != nil { + h.logger.Error("获取漏洞筛选建议失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, options) +} + +// ExportVulnerabilities 导出漏洞（支持按对话/任务分组，汇总或拆分） +func (h *VulnerabilityHandler) ExportVulnerabilities(c *gin.Context) { + groupBy := c.DefaultQuery("group_by", "conversation") + mode := c.DefaultQuery("mode", "summary") + if groupBy != "conversation" && groupBy != "task" { + c.JSON(http.StatusBadRequest, gin.H{"error": "group_by 仅支持 conversation 或 task"}) + return + } + if mode != "summary" && mode != "split" { + c.JSON(http.StatusBadRequest, gin.H{"error": "mode 仅支持 summary 或 split"}) + return + } + + filter := parseVulnerabilityListFilter(c) + access := vulnerabilityAccessFromContext(c) + + total, err := h.db.CountVulnerabilitiesForAccess(filter, access) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if total == 0 { + c.JSON(http.StatusOK, gin.H{"mode": mode, "group_by": groupBy, "total": 0, "files": []any{}}) + return + } + + items, err := h.db.ListVulnerabilitiesForAccess(total, 0, filter, access) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + + type exportFile struct { + FileName string `json:"filename"` + Content string `json:"content"` + } + grouped := map[string][]*database.Vulnerability{} + for _, v := range items { + key := v.ConversationID + if groupBy == "conversation" { + if strings.TrimSpace(v.ConversationTag) != "" { + key = strings.TrimSpace(v.ConversationTag) + } + } else { + key = firstNonEmpty(v.TaskTag, v.TaskID, v.TaskQueueID, "unassigned-task") + } + grouped[key] = append(grouped[key], v) + } + + files := make([]exportFile, 0) + nowStr := time.Now().Format("20060102-150405") + if mode == "summary" { + var b strings.Builder + b.WriteString("# 漏洞批量导出报告\n\n") + b.WriteString(fmt.Sprintf("- 导出时间: %s\n", time.Now().Format("2006-01-02 15:04:05"))) + b.WriteString(fmt.Sprintf("- 分组维度: %s\n", groupBy)) + b.WriteString(fmt.Sprintf("- 漏洞总数: %d\n", len(items))) + b.WriteString(fmt.Sprintf("- 分组数: %d\n\n", len(grouped))) + for group, list := range grouped { + b.WriteString(fmt.Sprintf("## %s (%d)\n\n", group, len(list))) + for _, v := range list { + appendVulnerabilityMarkdown(&b, v, "###") + } + } + files = append(files, exportFile{ + FileName: fmt.Sprintf("vulnerability-report-%s-%s.md", groupBy, nowStr), + Content: b.String(), + }) + } else { + for group, list := range grouped { + var b strings.Builder + b.WriteString(fmt.Sprintf("# 漏洞报告 - %s\n\n", group)) + b.WriteString(fmt.Sprintf("- 导出时间: %s\n", time.Now().Format("2006-01-02 15:04:05"))) + b.WriteString(fmt.Sprintf("- 漏洞数量: %d\n\n", len(list))) + for _, v := range list { + appendVulnerabilityMarkdown(&b, v, "##") + } + files = append(files, exportFile{ + FileName: fmt.Sprintf("vulnerability-%s-%s.md", sanitizeExportName(group), nowStr), + Content: b.String(), + }) + } + } + + c.JSON(http.StatusOK, gin.H{ + "mode": mode, + "group_by": groupBy, + "total": len(items), + "files": files, + }) +} + +func vulnerabilityAccessFromContext(c *gin.Context) database.RBACListAccess { + session, ok := security.CurrentSession(c) + if !ok { + return database.RBACListAccess{} + } + return database.RBACListAccess{UserID: session.UserID, Scope: session.Scope} +} + +// appendVulnerabilityMarkdown 单条漏洞的 Markdown 片段（与单文件下载字段对齐，缺省字段不写） +func appendVulnerabilityMarkdown(b *strings.Builder, v *database.Vulnerability, titleHeading string) { + b.WriteString(fmt.Sprintf("%s %s\n\n", titleHeading, v.Title)) + b.WriteString(fmt.Sprintf("- 漏洞ID: `%s`\n", v.ID)) + b.WriteString(fmt.Sprintf("- 严重程度: %s\n", v.Severity)) + b.WriteString(fmt.Sprintf("- 状态: %s\n", v.Status)) + if v.Type != "" { + b.WriteString(fmt.Sprintf("- 类型: %s\n", v.Type)) + } + if v.Target != "" { + b.WriteString(fmt.Sprintf("- 目标: %s\n", v.Target)) + } + b.WriteString(fmt.Sprintf("- 对话ID: `%s`\n", v.ConversationID)) + if v.ConversationTag != "" { + b.WriteString(fmt.Sprintf("- 对话标签: %s\n", v.ConversationTag)) + } + if v.TaskTag != "" { + b.WriteString(fmt.Sprintf("- 任务标签: %s\n", v.TaskTag)) + } + if v.TaskID != "" { + b.WriteString(fmt.Sprintf("- 任务ID: `%s`\n", v.TaskID)) + } + if v.TaskQueueID != "" { + b.WriteString(fmt.Sprintf("- 任务队列ID: `%s`\n", v.TaskQueueID)) + } + if !v.CreatedAt.IsZero() { + b.WriteString(fmt.Sprintf("- 创建时间: %s\n", v.CreatedAt.Format("2006-01-02 15:04:05"))) + } + if !v.UpdatedAt.IsZero() { + b.WriteString(fmt.Sprintf("- 更新时间: %s\n", v.UpdatedAt.Format("2006-01-02 15:04:05"))) + } + if v.Description != "" { + b.WriteString("\n#### 描述\n\n") + b.WriteString(v.Description) + b.WriteString("\n") + } + if v.Preconditions != "" { + b.WriteString("\n#### 前置条件\n\n") + b.WriteString(v.Preconditions) + b.WriteString("\n") + } + if v.ReproSteps != "" { + b.WriteString("\n#### 复现步骤\n\n") + b.WriteString(v.ReproSteps) + b.WriteString("\n") + } + if v.Evidence != "" { + b.WriteString("\n#### 证据 / POC\n\n```\n") + b.WriteString(v.Evidence) + b.WriteString("\n```\n") + } + if v.Impact != "" { + b.WriteString("\n#### 影响\n\n") + b.WriteString(v.Impact) + b.WriteString("\n") + } + if v.Recommendation != "" { + b.WriteString("\n#### 修复建议\n\n") + b.WriteString(v.Recommendation) + b.WriteString("\n") + } + if v.RetestNotes != "" { + b.WriteString("\n#### 复测方式\n\n") + b.WriteString(v.RetestNotes) + b.WriteString("\n") + } + b.WriteString("\n") +} + +func firstNonEmpty(values ...string) string { + for _, v := range values { + trimmed := strings.TrimSpace(v) + if trimmed != "" { + return trimmed + } + } + return "" +} + +func sanitizeExportName(raw string) string { + name := strings.TrimSpace(raw) + if name == "" { + return "unknown" + } + replacer := strings.NewReplacer("/", "-", "\\", "-", ":", "-", "*", "-", "?", "-", "\"", "-", "<", "-", ">", "-", "|", "-") + return replacer.Replace(name) +} diff --git a/internal/handler/vulnerability_alert.go b/internal/handler/vulnerability_alert.go new file mode 100644 index 00000000..941874c9 --- /dev/null +++ b/internal/handler/vulnerability_alert.go @@ -0,0 +1,227 @@ +package handler + +import ( + "context" + "fmt" + "net/http" + "strings" + "time" + + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/robot" + "cyberstrike-ai/internal/security" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +type updateVulnerabilityAlertRequest struct { + Enabled bool `json:"enabled"` + MinSeverity string `json:"min_severity"` +} + +func (h *VulnerabilityHandler) GetMyAlertSubscription(c *gin.Context) { + session, ok := security.CurrentSession(c) + if !ok { + c.JSON(http.StatusUnauthorized, gin.H{"error": "unauthorized"}) + return + } + sub, err := h.db.GetVulnerabilityAlertSubscription(session.UserID) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + bindings, _ := h.db.ListRobotUserBindings(session.UserID) + deliveryReady := false + for _, binding := range bindings { + if binding.Enabled && robot.SupportsProactive(binding.Platform) { + deliveryReady = true + break + } + } + c.JSON(http.StatusOK, gin.H{"subscription": sub, "bindings": bindings, "delivery_ready": deliveryReady}) +} + +func (h *VulnerabilityHandler) UpdateMyAlertSubscription(c *gin.Context) { + session, ok := security.CurrentSession(c) + if !ok { + c.JSON(http.StatusUnauthorized, gin.H{"error": "unauthorized"}) + return + } + var req updateVulnerabilityAlertRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + sub, err := h.db.UpsertVulnerabilityAlertSubscription(session.UserID, req.Enabled, req.MinSeverity) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "vulnerability", "alert_subscription_update", "更新漏洞提醒订阅", "user", session.UserID, map[string]interface{}{ + "enabled": sub.Enabled, "min_severity": sub.MinSeverity, + }) + } + c.JSON(http.StatusOK, sub) +} + +func vulnerabilityAlertSeverityLabel(value string) string { + switch value { + case "critical": + return "严重" + case "high": + return "高危" + case "medium": + return "中危" + case "low": + return "低危" + default: + return "信息" + } +} + +func formatVulnerabilityRobotAlert(v *database.Vulnerability) string { + return fmt.Sprintf("🚨 新漏洞提醒\n\n标题：%s\n严重程度：%s（%s）\n目标：%s\n影响：%s\n修复建议：%s\n漏洞 ID：%s\n\n请核实目标版本与实际暴露情况。", + strings.TrimSpace(v.Title), vulnerabilityAlertSeverityLabel(v.Severity), v.Severity, + fallbackAlertText(v.Target), fallbackAlertText(v.Impact), fallbackAlertText(v.Recommendation), v.ID) +} + +func fallbackAlertText(value string) string { + value = strings.TrimSpace(value) + if value == "" { + return "未填写" + } + if len([]rune(value)) > 300 { + return string([]rune(value)[:300]) + "…" + } + return value +} + +// NotifyNewVulnerability pushes an alert to every eligible bound robot identity. +// A failure on one platform never blocks vulnerability creation or other recipients. +func (h *RobotHandler) NotifyNewVulnerability(v *database.Vulnerability) { + if h == nil || h.db == nil || v == nil { + return + } + recipients, err := h.db.ListVulnerabilityAlertRecipients(v) + if err != nil { + h.logger.Warn("查询漏洞提醒接收人失败", zap.String("vulnerability_id", v.ID), zap.Error(err)) + return + } + eligible := recipients[:0] + for _, recipient := range recipients { + if robot.SupportsProactive(recipient.Platform) { + eligible = append(eligible, recipient) + } + } + if err := h.db.EnqueueVulnerabilityAlertDeliveries(v.ID, eligible); err != nil { + h.logger.Warn("写入漏洞提醒投递队列失败", zap.String("vulnerability_id", v.ID), zap.Error(err)) + return + } + select { + case h.alertWake <- struct{}{}: + default: + } +} + +// RunVulnerabilityAlertWorker drains the durable outbox. Failed sends use +// exponential backoff and remain retryable across application restarts. +func (h *RobotHandler) RunVulnerabilityAlertWorker(ctx context.Context) { + if h == nil || h.db == nil { + return + } + ticker := time.NewTicker(30 * time.Second) + defer ticker.Stop() + for { + h.processVulnerabilityAlertDeliveries(ctx) + select { + case <-ctx.Done(): + return + case <-ticker.C: + case <-h.alertWake: + } + } +} + +func (h *RobotHandler) processVulnerabilityAlertDeliveries(ctx context.Context) { + deliveries, err := h.db.ListDueVulnerabilityAlertDeliveries(50) + if err != nil { + h.logger.Warn("读取漏洞提醒投递队列失败", zap.Error(err)) + return + } + for _, delivery := range deliveries { + sendCtx, cancel := context.WithTimeout(ctx, 15*time.Second) + err := robot.SendProactive(sendCtx, h.config.Robots, delivery.Platform, delivery.ExternalUserID, formatVulnerabilityRobotAlert(delivery.Vulnerability)) + cancel() + if err != nil { + attempts := delivery.Attempts + 1 + _ = h.db.MarkVulnerabilityAlertDeliveryFailed(delivery.ID, attempts, err) + h.logger.Warn("发送漏洞机器人提醒失败", zap.String("platform", delivery.Platform), zap.String("user_id", delivery.UserID), zap.String("vulnerability_id", delivery.Vulnerability.ID), zap.Int("attempt", attempts), zap.Error(err)) + continue + } + _ = h.db.MarkVulnerabilityAlertDeliverySent(delivery.ID) + h.logger.Info("漏洞机器人提醒已发送", zap.String("platform", delivery.Platform), zap.String("user_id", delivery.UserID), zap.String("vulnerability_id", delivery.Vulnerability.ID)) + } +} + +func (h *RobotHandler) cmdVulnerabilityAlerts(platform, externalUserID, arg string) string { + access, err := h.resolveRobotAccess(platform, externalUserID) + if err != nil { + return h.robotAccessDeniedMessage(platform) + } + arg = strings.ToLower(strings.TrimSpace(arg)) + if arg == "" || arg == "status" || arg == "状态" { + sub, err := h.db.GetVulnerabilityAlertSubscription(access.User.ID) + if err != nil { + return "读取漏洞提醒设置失败，请稍后重试。" + } + state := "已关闭" + if sub.Enabled { + state = "已开启" + } + return fmt.Sprintf("漏洞提醒%s；最低提醒级别：%s（%s）。\n可在 Web「漏洞管理」页面同步修改。", state, vulnerabilityAlertSeverityLabel(sub.MinSeverity), sub.MinSeverity) + } + enabled, severity := true, "" + switch arg { + case "开启", "开", "on", "enable": + severity = "high" + case "关闭", "关", "off", "disable": + enabled, severity = false, "high" + case "仅严重", "严重", "critical": + severity = "critical" + case "高危以上", "高危", "high": + severity = "high" + case "中危以上", "中危", "medium": + severity = "medium" + case "低危以上", "低危", "low": + severity = "low" + case "全部", "all", "info": + severity = "info" + default: + return "用法：漏洞提醒 开启｜关闭｜仅严重｜高危以上｜中危以上" + } + current, _ := h.db.GetVulnerabilityAlertSubscription(access.User.ID) + if (arg == "开启" || arg == "开" || arg == "on" || arg == "enable" || !enabled) && current != nil { + severity = current.MinSeverity + } + sub, err := h.db.UpsertVulnerabilityAlertSubscription(access.User.ID, enabled, severity) + if err != nil { + return "更新漏洞提醒失败，请稍后重试。" + } + if !sub.Enabled { + return "已关闭漏洞提醒。Web 端设置已同步。" + } + bindings, _ := h.db.ListRobotUserBindings(access.User.ID) + deliveryReady := false + for _, binding := range bindings { + if binding.Enabled && robot.SupportsProactive(binding.Platform) { + deliveryReady = true + break + } + } + if !deliveryReady { + return fmt.Sprintf("已保存漏洞提醒：%s及以上。当前没有支持主动推送的已绑定账号；请绑定企业微信、飞书、Telegram、Slack 或 Discord。", vulnerabilityAlertSeverityLabel(sub.MinSeverity)) + } + return fmt.Sprintf("已开启漏洞提醒：%s及以上漏洞将通过已绑定机器人推送。Web 端设置已同步。", vulnerabilityAlertSeverityLabel(sub.MinSeverity)) +} diff --git a/internal/handler/vulnerability_alert_test.go b/internal/handler/vulnerability_alert_test.go new file mode 100644 index 00000000..b3596808 --- /dev/null +++ b/internal/handler/vulnerability_alert_test.go @@ -0,0 +1,51 @@ +package handler + +import ( + "path/filepath" + "strings" + "testing" + "time" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + + "go.uber.org/zap" +) + +func TestRobotVulnerabilityAlertCommandSharesSubscription(t *testing.T) { + db, err := database.NewDB(filepath.Join(t.TempDir(), "alert-command.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = db.Close() }) + if err := db.BootstrapRBAC("hash", security.PermissionCatalog); err != nil { + t.Fatal(err) + } + user, err := db.CreateRBACUser("alert-user", "Alert User", "hash", true, []string{database.RBACSystemRoleOperator}) + if err != nil { + t.Fatal(err) + } + if err := db.CreateRobotBindingCode(user.ID, "alert-code", time.Now().Add(time.Minute)); err != nil { + t.Fatal(err) + } + if _, err := db.ConsumeRobotBindingCode("wecom", "external-user", "alert-code"); err != nil { + t.Fatal(err) + } + h := NewRobotHandler(&config.Config{}, db, nil, zap.NewNop()) + + if got := h.HandleMessage("wecom", "external-user", "漏洞提醒 高危以上"); !strings.Contains(got, "已开启") { + t.Fatalf("enable reply: %s", got) + } + sub, err := db.GetVulnerabilityAlertSubscription(user.ID) + if err != nil || !sub.Enabled || sub.MinSeverity != "high" { + t.Fatalf("web subscription not updated: %#v %v", sub, err) + } + if got := h.HandleMessage("wecom", "external-user", "vuln alerts off"); !strings.Contains(got, "已关闭") { + t.Fatalf("disable reply: %s", got) + } + sub, _ = db.GetVulnerabilityAlertSubscription(user.ID) + if sub.Enabled { + t.Fatalf("subscription remained enabled: %#v", sub) + } +} diff --git a/internal/handler/webshell.go b/internal/handler/webshell.go new file mode 100644 index 00000000..90eb3ca8 --- /dev/null +++ b/internal/handler/webshell.go @@ -0,0 +1,1079 @@ +package handler + +import ( + "bytes" + "crypto/tls" + "database/sql" + "encoding/base64" + "encoding/json" + "io" + "net/http" + "net/url" + "strings" + "time" + "unicode/utf8" + + "cyberstrike-ai/internal/audit" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + + "github.com/gin-gonic/gin" + "github.com/google/uuid" + "go.uber.org/zap" + "golang.org/x/text/encoding/simplifiedchinese" + "golang.org/x/text/transform" +) + +// webshellSupportedEncodings 允许的 WebShell 响应编码取值（小写，含空串代表 auto） +// 仅暴露目前最常见的几种，其他需求可后续扩展（如 Big5、Shift_JIS 等）。 +var webshellSupportedEncodings = map[string]struct{}{ + "": {}, // 未配置，按 auto 处理 + "auto": {}, + "utf-8": {}, + "utf8": {}, + "gbk": {}, + "gb18030": {}, +} + +// normalizeWebshellEncoding 归一化编码标识：统一为小写，未知值回退为 auto，供持久化使用 +func normalizeWebshellEncoding(enc string) string { + enc = strings.ToLower(strings.TrimSpace(enc)) + if _, ok := webshellSupportedEncodings[enc]; !ok { + return "auto" + } + if enc == "" { + return "auto" + } + if enc == "utf8" { + return "utf-8" + } + return enc +} + +// decodeWebshellOutput 把 WebShell 返回的字节按指定编码转换为合法 UTF-8 字符串。 +// 约定： +// - "" / "auto"：若已是合法 UTF-8 原样返回，否则依次尝试 GB18030（GBK 超集）解码。 +// - "utf-8" / "utf8"：原样返回，非法字节交由 JSON 层按 U+FFFD 处理（保持原有行为）。 +// - "gbk" / "gb18030"：强制按对应编码解码；失败则回退原始字节。 +// +// 该函数对空输入直接返回空串，避免不必要的转换。 +func decodeWebshellOutput(raw []byte, encoding string) string { + if len(raw) == 0 { + return "" + } + enc := normalizeWebshellEncoding(encoding) + switch enc { + case "utf-8": + return string(raw) + case "gbk": + if out, _, err := transform.Bytes(simplifiedchinese.GBK.NewDecoder(), raw); err == nil { + return string(out) + } + return string(raw) + case "gb18030": + if out, _, err := transform.Bytes(simplifiedchinese.GB18030.NewDecoder(), raw); err == nil { + return string(out) + } + return string(raw) + default: // auto + if utf8.Valid(raw) { + return string(raw) + } + // GB18030 是 GBK 的超集，覆盖范围最广，auto 模式统一用它兜底 + if out, _, err := transform.Bytes(simplifiedchinese.GB18030.NewDecoder(), raw); err == nil { + return string(out) + } + return string(raw) + } +} + +// webshellSupportedOS 允许的 WebShell 目标操作系统（小写，空串代表 auto） +var webshellSupportedOS = map[string]struct{}{ + "": {}, + "auto": {}, + "linux": {}, + "windows": {}, +} + +// normalizeWebshellOS 归一化 OS 标识，未知值回退为 auto，供持久化使用 +func normalizeWebshellOS(osTag string) string { + osTag = strings.ToLower(strings.TrimSpace(osTag)) + if _, ok := webshellSupportedOS[osTag]; !ok { + return "auto" + } + if osTag == "" { + return "auto" + } + return osTag +} + +// resolveWebshellOS 根据连接的 os 与 shellType 推断最终目标 OS（仅返回 "linux" 或 "windows"）。 +// 规则： +// - 显式 linux / windows：按用户选择。 +// - auto 或未知：asp/aspx → windows，其他 → linux。保持历史行为，平滑向后兼容。 +func resolveWebshellOS(osTag, shellType string) string { + osTag = strings.ToLower(strings.TrimSpace(osTag)) + switch osTag { + case "linux": + return "linux" + case "windows": + return "windows" + } + t := strings.ToLower(strings.TrimSpace(shellType)) + if t == "asp" || t == "aspx" { + return "windows" + } + return "linux" +} + +// quoteCmdPath 把路径按 Windows cmd.exe 规则转义。 +// 使用双引号包裹，内部双引号转义为 ""（cmd 接受的写法）。 +func quoteCmdPath(p string) string { + if p == "" { + return "\".\"" + } + return "\"" + strings.ReplaceAll(p, "\"", "\"\"") + "\"" +} + +// normalizeWindowsCmdPath 把前端统一的 "/" 路径转换为 cmd 更稳定识别的 "\"。 +// 仅用于 Windows 命令构造，不改变语义（例如 "." / ".." 会保持不变）。 +func normalizeWindowsCmdPath(p string) string { + s := strings.TrimSpace(p) + if s == "" { + return s + } + return strings.ReplaceAll(s, "/", "\\") +} + +// quotePsSingle 把字符串按 PowerShell 单引号字符串规则转义（内部 ' → ”）。 +// 供 PowerShell 脚本参数使用，全脚本只用单引号，外层 cmd 再用双引号包裹即可安全传递。 +func quotePsSingle(s string) string { + return "'" + strings.ReplaceAll(s, "'", "''") + "'" +} + +// quoteShellSinglePosix 把路径按 POSIX sh 单引号规则转义（内部 ' → '\”） +func quoteShellSinglePosix(p string) string { + if p == "" { + return "." + } + return "'" + strings.ReplaceAll(p, "'", "'\\''") + "'" +} + +// quoteWebshellPath 按目标 OS 选择转义方案：Linux 用 POSIX 单引号，Windows 用 cmd 双引号 +func quoteWebshellPath(path, osTag string) string { + if resolveWebshellOS(osTag, "") == "windows" { + return quoteCmdPath(path) + } + return quoteShellSinglePosix(path) +} + +// buildWindowsPowerShellWrite 构造 Windows 端把 base64 内容一次性写入目标路径的 cmd 命令。 +// 外层走 cmd.exe 的 powershell 调用，PowerShell 脚本里只用单引号字符串，避免嵌套引号陷阱。 +func buildWindowsPowerShellWrite(path, b64 string) string { + script := "$b=[Convert]::FromBase64String(" + quotePsSingle(b64) + ");" + + "[IO.File]::WriteAllBytes(" + quotePsSingle(path) + ",$b)" + return "powershell -NoProfile -NonInteractive -Command \"" + script + "\"" +} + +// buildWindowsPowerShellAppend 构造 Windows 端把 base64 内容追加写入目标路径的 cmd 命令（用于分块上传） +func buildWindowsPowerShellAppend(path, b64 string) string { + script := "$b=[Convert]::FromBase64String(" + quotePsSingle(b64) + ");" + + "$f=[IO.File]::Open(" + quotePsSingle(path) + ",[IO.FileMode]::Append,[IO.FileAccess]::Write,[IO.FileShare]::None);" + + "try{$f.Write($b,0,$b.Length)}finally{$f.Close()}" + return "powershell -NoProfile -NonInteractive -Command \"" + script + "\"" +} + +// fileCommandInput 封装 buildFileCommand 的输入，避免长参数列表 +type fileCommandInput struct { + Action string + Path string + TargetPath string + Content string + ChunkIndex int + OS string + ShellType string +} + +// buildFileCommand 根据目标 OS 与文件操作类型生成具体的远端命令字符串。 +// 同一份实现供 HTTP 入口（FileOp）与 MCP 入口（FileOpWithConnection）共用，避免双份维护。 +// 返回值第二位是用户可见的业务错误（如 "path is required"）。 +func (h *WebShellHandler) buildFileCommand(in fileCommandInput) (string, error) { + targetOS := resolveWebshellOS(in.OS, in.ShellType) + action := strings.ToLower(strings.TrimSpace(in.Action)) + path := strings.TrimSpace(in.Path) + + switch action { + case "list": + p := path + if p == "" { + p = "." + } + if targetOS == "windows" { + p = normalizeWindowsCmdPath(p) + return "dir /a " + quoteCmdPath(p), nil + } + return "ls -la " + quoteShellSinglePosix(p), nil + + case "read": + if path == "" { + return "", errFileOpPathRequired + } + if targetOS == "windows" { + path = normalizeWindowsCmdPath(path) + return "type " + quoteCmdPath(path), nil + } + return "cat " + quoteShellSinglePosix(path), nil + + case "delete": + if path == "" { + return "", errFileOpPathRequired + } + if targetOS == "windows" { + path = normalizeWindowsCmdPath(path) + return "del /q /f " + quoteCmdPath(path), nil + } + return "rm -f " + quoteShellSinglePosix(path), nil + + case "mkdir": + if path == "" { + return "", errFileOpPathRequired + } + if targetOS == "windows" { + path = normalizeWindowsCmdPath(path) + // cmd 的 md 默认会自动创建中间目录（等价于 Linux 的 mkdir -p） + return "md " + quoteCmdPath(path), nil + } + return "mkdir -p " + quoteShellSinglePosix(path), nil + + case "rename": + oldPath := path + newPath := strings.TrimSpace(in.TargetPath) + if oldPath == "" || newPath == "" { + return "", errFileOpRenameNeedsBothPaths + } + if targetOS == "windows" { + oldPath = normalizeWindowsCmdPath(oldPath) + newPath = normalizeWindowsCmdPath(newPath) + return "move /y " + quoteCmdPath(oldPath) + " " + quoteCmdPath(newPath), nil + } + return "mv -f " + quoteShellSinglePosix(oldPath) + " " + quoteShellSinglePosix(newPath), nil + + case "write": + if path == "" { + return "", errFileOpPathRequired + } + // 统一策略：先把内容 base64 编码，再用目标平台对应方式解码写回， + // 这样既能写入任意二进制/含引号的文本，又避免各家 shell 的转义地狱。 + b64 := base64.StdEncoding.EncodeToString([]byte(in.Content)) + if targetOS == "windows" { + path = normalizeWindowsCmdPath(path) + return buildWindowsPowerShellWrite(path, b64), nil + } + return "echo '" + b64 + "' | base64 -d > " + quoteShellSinglePosix(path), nil + + case "upload": + if path == "" { + return "", errFileOpPathRequired + } + if len(in.Content) > 512*1024 { + return "", errFileOpUploadTooLarge + } + if targetOS == "windows" { + path = normalizeWindowsCmdPath(path) + return buildWindowsPowerShellWrite(path, in.Content), nil + } + return "echo '" + in.Content + "' | base64 -d > " + quoteShellSinglePosix(path), nil + + case "upload_chunk": + if path == "" { + return "", errFileOpPathRequired + } + if targetOS == "windows" { + path = normalizeWindowsCmdPath(path) + if in.ChunkIndex == 0 { + return buildWindowsPowerShellWrite(path, in.Content), nil + } + return buildWindowsPowerShellAppend(path, in.Content), nil + } + redir := ">>" + if in.ChunkIndex == 0 { + redir = ">" + } + return "echo '" + in.Content + "' | base64 -d " + redir + " " + quoteShellSinglePosix(path), nil + } + + return "", errFileOpUnsupportedAction(action) +} + +// 业务错误常量，便于上层统一返回用户可见提示 +var ( + errFileOpPathRequired = simpleError("path is required") + errFileOpRenameNeedsBothPaths = simpleError("path and target_path are required for rename") + errFileOpUploadTooLarge = simpleError("upload content too large (max 512KB base64)") +) + +func errFileOpUnsupportedAction(action string) error { + return simpleError("unsupported action: " + action) +} + +// simpleError 是不带堆栈的轻量错误类型，供 buildFileCommand 报可预期的参数校验错误 +type simpleError string + +func (e simpleError) Error() string { return string(e) } + +// WebShellHandler 代理执行 WebShell 命令（类似冰蝎/蚁剑），避免前端跨域并统一构建请求 +type WebShellHandler struct { + logger *zap.Logger + client *http.Client + db *database.DB + audit *audit.Service +} + +// SetAudit wires platform audit logging. +func (h *WebShellHandler) SetAudit(s *audit.Service) { + h.audit = s +} + +// NewWebShellHandler 创建 WebShell 处理器，db 可为 nil（连接配置接口将不可用） +func NewWebShellHandler(logger *zap.Logger, db *database.DB) *WebShellHandler { + return &WebShellHandler{ + logger: logger, + client: &http.Client{ + Timeout: 30 * time.Second, + Transport: &http.Transport{ + DisableKeepAlives: false, + // WebShell 场景常见自签证书或 IP 访问（证书无 IP SAN）；默认跳过校验，与蚁剑等客户端一致。 + TLSClientConfig: &tls.Config{InsecureSkipVerify: true}, //nolint:gosec // intentional for webshell proxy + }, + }, + db: db, + } +} + +// CreateConnectionRequest 创建连接请求 +type CreateConnectionRequest struct { + ProjectID string `json:"project_id"` + URL string `json:"url" binding:"required"` + Password string `json:"password"` + Type string `json:"type"` + Method string `json:"method"` + CmdParam string `json:"cmd_param"` + Remark string `json:"remark"` + Encoding string `json:"encoding"` + OS string `json:"os"` +} + +// UpdateConnectionRequest 更新连接请求 +type UpdateConnectionRequest struct { + ProjectID string `json:"project_id"` + URL string `json:"url" binding:"required"` + Password string `json:"password"` + Type string `json:"type"` + Method string `json:"method"` + CmdParam string `json:"cmd_param"` + Remark string `json:"remark"` + Encoding string `json:"encoding"` + OS string `json:"os"` +} + +// ListConnections 列出所有 WebShell 连接（GET /api/webshell/connections） +func (h *WebShellHandler) ListConnections(c *gin.Context) { + if h.db == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "database not available"}) + return + } + session, _ := security.CurrentSession(c) + list, err := h.db.ListWebshellConnectionsForAccess(session.UserID, session.Scope, c.Query("project_id")) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if list == nil { + list = []database.WebShellConnection{} + } + c.JSON(http.StatusOK, list) +} + +// CreateConnection 创建 WebShell 连接（POST /api/webshell/connections） +func (h *WebShellHandler) CreateConnection(c *gin.Context) { + if h.db == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "database not available"}) + return + } + var req CreateConnectionRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + req.URL = strings.TrimSpace(req.URL) + if req.URL == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "url is required"}) + return + } + if _, err := url.Parse(req.URL); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid url"}) + return + } + projectID := strings.TrimSpace(req.ProjectID) + if !h.canAccessProject(c, projectID) { + c.JSON(http.StatusForbidden, gin.H{"error": "project access denied"}) + return + } + method := strings.ToLower(strings.TrimSpace(req.Method)) + if method != "get" && method != "post" { + method = "post" + } + shellType := strings.ToLower(strings.TrimSpace(req.Type)) + if shellType == "" { + shellType = "php" + } + conn := &database.WebShellConnection{ + ID: "ws_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:12], + ProjectID: projectID, + URL: req.URL, + Password: strings.TrimSpace(req.Password), + Type: shellType, + Method: method, + CmdParam: strings.TrimSpace(req.CmdParam), + Remark: strings.TrimSpace(req.Remark), + Encoding: normalizeWebshellEncoding(req.Encoding), + OS: normalizeWebshellOS(req.OS), + CreatedAt: time.Now(), + } + if err := h.db.CreateWebshellConnection(conn); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if session, ok := security.CurrentSession(c); ok { + _ = h.db.SetResourceOwner("webshell", conn.ID, session.UserID) + _ = h.db.AssignResourceToUser(session.UserID, "webshell", conn.ID) + } + if h.audit != nil { + host := req.URL + if u, err := url.Parse(req.URL); err == nil { + host = u.Host + } + h.audit.RecordOK(c, "webshell", "connection_create", "创建 WebShell 连接", "webshell_connection", conn.ID, map[string]interface{}{ + "host": host, "type": shellType, + }) + } + c.JSON(http.StatusOK, conn) +} + +// UpdateConnection 更新 WebShell 连接（PUT /api/webshell/connections/:id） +func (h *WebShellHandler) UpdateConnection(c *gin.Context) { + if h.db == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "database not available"}) + return + } + id := strings.TrimSpace(c.Param("id")) + if id == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "id is required"}) + return + } + var req UpdateConnectionRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + req.URL = strings.TrimSpace(req.URL) + if req.URL == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "url is required"}) + return + } + if _, err := url.Parse(req.URL); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid url"}) + return + } + projectID := strings.TrimSpace(req.ProjectID) + if !h.canAccessProject(c, projectID) { + c.JSON(http.StatusForbidden, gin.H{"error": "project access denied"}) + return + } + method := strings.ToLower(strings.TrimSpace(req.Method)) + if method != "get" && method != "post" { + method = "post" + } + shellType := strings.ToLower(strings.TrimSpace(req.Type)) + if shellType == "" { + shellType = "php" + } + conn := &database.WebShellConnection{ + ID: id, + ProjectID: projectID, + URL: req.URL, + Password: strings.TrimSpace(req.Password), + Type: shellType, + Method: method, + CmdParam: strings.TrimSpace(req.CmdParam), + Remark: strings.TrimSpace(req.Remark), + Encoding: normalizeWebshellEncoding(req.Encoding), + OS: normalizeWebshellOS(req.OS), + } + if err := h.db.UpdateWebshellConnection(conn); err != nil { + if err == sql.ErrNoRows { + c.JSON(http.StatusNotFound, gin.H{"error": "connection not found"}) + return + } + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + updated, _ := h.db.GetWebshellConnection(id) + if updated != nil { + c.JSON(http.StatusOK, updated) + } else { + c.JSON(http.StatusOK, conn) + } +} + +// DeleteConnection 删除 WebShell 连接（DELETE /api/webshell/connections/:id） +func (h *WebShellHandler) DeleteConnection(c *gin.Context) { + if h.db == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "database not available"}) + return + } + id := strings.TrimSpace(c.Param("id")) + if id == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "id is required"}) + return + } + if err := h.db.DeleteWebshellConnection(id); err != nil { + if err == sql.ErrNoRows { + c.JSON(http.StatusNotFound, gin.H{"error": "connection not found"}) + return + } + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "webshell", "connection_delete", "删除 WebShell 连接", "webshell_connection", id, nil) + } + c.JSON(http.StatusOK, gin.H{"ok": true}) +} + +// GetConnectionState 获取 WebShell 连接关联的前端持久化状态（GET /api/webshell/connections/:id/state） +func (h *WebShellHandler) GetConnectionState(c *gin.Context) { + if h.db == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "database not available"}) + return + } + id := strings.TrimSpace(c.Param("id")) + if id == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "id is required"}) + return + } + conn, err := h.db.GetWebshellConnection(id) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if conn == nil { + c.JSON(http.StatusNotFound, gin.H{"error": "connection not found"}) + return + } + stateJSON, err := h.db.GetWebshellConnectionState(id) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + var state interface{} + if err := json.Unmarshal([]byte(stateJSON), &state); err != nil { + state = map[string]interface{}{} + } + c.JSON(http.StatusOK, gin.H{"state": state}) +} + +// SaveConnectionState 保存 WebShell 连接关联的前端持久化状态（PUT /api/webshell/connections/:id/state） +func (h *WebShellHandler) SaveConnectionState(c *gin.Context) { + if h.db == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "database not available"}) + return + } + id := strings.TrimSpace(c.Param("id")) + if id == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "id is required"}) + return + } + conn, err := h.db.GetWebshellConnection(id) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if conn == nil { + c.JSON(http.StatusNotFound, gin.H{"error": "connection not found"}) + return + } + var req struct { + State json.RawMessage `json:"state"` + } + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + raw := req.State + if len(raw) == 0 { + raw = json.RawMessage(`{}`) + } + if len(raw) > 2*1024*1024 { + c.JSON(http.StatusBadRequest, gin.H{"error": "state payload too large (max 2MB)"}) + return + } + var anyJSON interface{} + if err := json.Unmarshal(raw, &anyJSON); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "state must be valid json"}) + return + } + if err := h.db.UpsertWebshellConnectionState(id, string(raw)); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"ok": true}) +} + +// GetAIHistory 获取指定 WebShell 连接的 AI 助手对话历史（GET /api/webshell/connections/:id/ai-history） +func (h *WebShellHandler) GetAIHistory(c *gin.Context) { + if h.db == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "database not available"}) + return + } + id := strings.TrimSpace(c.Param("id")) + if id == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "id is required"}) + return + } + conv, err := h.db.GetConversationByWebshellConnectionID(id) + if err != nil { + h.logger.Warn("获取 WebShell AI 对话失败", zap.String("connectionId", id), zap.Error(err)) + c.JSON(http.StatusOK, gin.H{"conversationId": nil, "messages": []database.Message{}}) + return + } + if conv == nil { + c.JSON(http.StatusOK, gin.H{"conversationId": nil, "messages": []database.Message{}}) + return + } + c.JSON(http.StatusOK, gin.H{"conversationId": conv.ID, "messages": conv.Messages}) +} + +// ListAIConversations 列出该 WebShell 连接下的所有 AI 对话（供侧边栏） +func (h *WebShellHandler) ListAIConversations(c *gin.Context) { + if h.db == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "database not available"}) + return + } + id := strings.TrimSpace(c.Param("id")) + if id == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "id is required"}) + return + } + list, err := h.db.ListConversationsByWebshellConnectionID(id) + if err != nil { + h.logger.Warn("列出 WebShell AI 对话失败", zap.String("connectionId", id), zap.Error(err)) + c.JSON(http.StatusOK, []database.WebShellConversationItem{}) + return + } + if list == nil { + list = []database.WebShellConversationItem{} + } + c.JSON(http.StatusOK, list) +} + +// ExecRequest 执行命令请求（前端传入连接信息 + 命令） +type ExecRequest struct { + URL string `json:"url" binding:"required"` + Password string `json:"password"` + Type string `json:"type"` // php, asp, aspx, jsp, custom + Method string `json:"method"` // GET 或 POST，空则默认 POST + CmdParam string `json:"cmd_param"` // 命令参数名，如 cmd/xxx，空则默认 cmd + Encoding string `json:"encoding"` // 响应编码：auto / utf-8 / gbk / gb18030，空则 auto + OS string `json:"os"` // 目标操作系统：auto / linux / windows，当前 exec 不用它，保留字段便于未来扩展 + ConnectionID string `json:"connection_id,omitempty"` + Command string `json:"command" binding:"required"` +} + +// ExecResponse 执行命令响应 +type ExecResponse struct { + OK bool `json:"ok"` + Output string `json:"output"` + Error string `json:"error,omitempty"` + HTTPCode int `json:"http_code,omitempty"` +} + +// FileOpRequest 文件操作请求 +type FileOpRequest struct { + URL string `json:"url" binding:"required"` + Password string `json:"password"` + Type string `json:"type"` + Method string `json:"method"` // GET 或 POST，空则默认 POST + CmdParam string `json:"cmd_param"` // 命令参数名，如 cmd/xxx，空则默认 cmd + Encoding string `json:"encoding"` // 响应编码：auto / utf-8 / gbk / gb18030，空则 auto + OS string `json:"os"` // 目标操作系统：auto / linux / windows，空则按 shellType 推断 + ConnectionID string `json:"connection_id,omitempty"` // 可选：连接 ID；服务端探活出 OS 后会回写到此连接 + Action string `json:"action" binding:"required"` // list, read, delete, write, mkdir, rename, upload, upload_chunk + Path string `json:"path"` + TargetPath string `json:"target_path"` // rename 时目标路径 + Content string `json:"content"` // write/upload 时使用 + ChunkIndex int `json:"chunk_index"` // upload_chunk 时，0 表示首块 +} + +// FileOpResponse 文件操作响应 +type FileOpResponse struct { + OK bool `json:"ok"` + Output string `json:"output"` + Error string `json:"error,omitempty"` + DetectedOS string `json:"detected_os,omitempty"` // 仅在 auto 模式且探活成功时返回，前端应更新本地缓存 +} + +func (h *WebShellHandler) Exec(c *gin.Context) { + var req ExecRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + req.URL = strings.TrimSpace(req.URL) + req.Command = strings.TrimSpace(req.Command) + if req.URL == "" || req.Command == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "url and command are required"}) + return + } + // Pre-save connectivity tests send form credentials without connection_id. + // Saved connections must go through resource ACL; DB credentials are authoritative. + if cid := strings.TrimSpace(req.ConnectionID); cid != "" { + conn, allowed := h.authorizedWebshellConnection(c, cid, req.URL) + if !allowed { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + // Never let a caller pair an authorized ID with attacker-controlled + // transport credentials or a URL. + req.URL, req.Password, req.Type = conn.URL, conn.Password, conn.Type + req.Method, req.CmdParam, req.Encoding = conn.Method, conn.CmdParam, conn.Encoding + } else if !security.SessionHasPermission(c, "webshell:write") { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + + parsed, err := url.Parse(req.URL) + if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid url: only http(s) allowed"}) + return + } + + useGET := strings.ToUpper(strings.TrimSpace(req.Method)) == "GET" + cmdParam := strings.TrimSpace(req.CmdParam) + if cmdParam == "" { + cmdParam = "cmd" + } + var httpReq *http.Request + if useGET { + targetURL := h.buildExecURL(req.URL, req.Type, req.Password, cmdParam, req.Command) + httpReq, err = http.NewRequest(http.MethodGet, targetURL, nil) + } else { + body := h.buildExecBody(req.Type, req.Password, cmdParam, req.Command) + httpReq, err = http.NewRequest(http.MethodPost, req.URL, bytes.NewReader(body)) + httpReq.Header.Set("Content-Type", "application/x-www-form-urlencoded") + } + if err != nil { + h.logger.Warn("webshell exec NewRequest", zap.Error(err)) + c.JSON(http.StatusInternalServerError, ExecResponse{OK: false, Error: err.Error()}) + return + } + httpReq.Header.Set("User-Agent", "Mozilla/5.0 (compatible; CyberStrikeAI-WebShell/1.0)") + + resp, err := h.client.Do(httpReq) + if err != nil { + h.logger.Warn("webshell exec Do", zap.String("url", req.URL), zap.Error(err)) + c.JSON(http.StatusOK, ExecResponse{OK: false, Error: err.Error()}) + return + } + defer resp.Body.Close() + + out, readErr := io.ReadAll(resp.Body) + if readErr != nil { + h.logger.Warn("webshell exec read body", zap.Error(readErr)) + } + output := decodeWebshellOutput(out, req.Encoding) + httpCode := resp.StatusCode + + ok := resp.StatusCode == http.StatusOK + c.JSON(http.StatusOK, ExecResponse{ + OK: ok, + Output: output, + HTTPCode: httpCode, + }) +} + +// buildExecBody 按常见 WebShell 约定构建 POST 体（多数使用 pass + cmd，可配置命令参数名） +func (h *WebShellHandler) buildExecBody(shellType, password, cmdParam, command string) []byte { + form := h.execParams(shellType, password, cmdParam, command) + return []byte(form.Encode()) +} + +// buildExecURL 构建 GET 请求的完整 URL（baseURL + ?pass=xxx&cmd=yyy，cmd 可配置） +func (h *WebShellHandler) buildExecURL(baseURL, shellType, password, cmdParam, command string) string { + form := h.execParams(shellType, password, cmdParam, command) + if parsed, err := url.Parse(baseURL); err == nil { + parsed.RawQuery = form.Encode() + return parsed.String() + } + return baseURL + "?" + form.Encode() +} + +func (h *WebShellHandler) execParams(shellType, password, cmdParam, command string) url.Values { + shellType = strings.ToLower(strings.TrimSpace(shellType)) + if shellType == "" { + shellType = "php" + } + if strings.TrimSpace(cmdParam) == "" { + cmdParam = "cmd" + } + form := url.Values{} + form.Set("pass", password) + form.Set(cmdParam, command) + return form +} + +func (h *WebShellHandler) FileOp(c *gin.Context) { + var req FileOpRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + req.URL = strings.TrimSpace(req.URL) + req.Action = strings.ToLower(strings.TrimSpace(req.Action)) + if req.URL == "" || req.Action == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "url and action are required"}) + return + } + if cid := strings.TrimSpace(req.ConnectionID); cid != "" { + conn, allowed := h.authorizedWebshellConnection(c, cid, req.URL) + if !allowed { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + req.URL, req.Password, req.Type = conn.URL, conn.Password, conn.Type + req.Method, req.CmdParam, req.Encoding, req.OS = conn.Method, conn.CmdParam, conn.Encoding, conn.OS + } else if !security.SessionHasPermission(c, "webshell:write") { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + + parsed, err := url.Parse(req.URL) + if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid url: only http(s) allowed"}) + return + } + + // 若 OS 未显式配置，先发一次探活命令，识别出真实 OS 再构造文件操作命令。 + // 这解决了 "Windows + PHP + OS=auto" 场景下旧 fallback 错发 `ls -la` 导致目录列不出来的问题。 + osTag := req.OS + detectedOS := "" + if normalizeWebshellOS(osTag) == "auto" { + if probed := probeWebshellOSViaExec(h.newHTTPExecFn(req.URL, req.Password, req.Type, req.Method, req.CmdParam, req.Encoding)); probed != "" { + osTag = probed + detectedOS = probed + // 若前端带了 connection_id，顺带把探活结果持久化到该连接，后续刷新零成本 + if cid := strings.TrimSpace(req.ConnectionID); cid != "" { + h.persistDetectedOS(cid, probed) + } + } + } + + command, cmdErr := h.buildFileCommand(fileCommandInput{ + Action: req.Action, + Path: req.Path, + TargetPath: req.TargetPath, + Content: req.Content, + ChunkIndex: req.ChunkIndex, + OS: osTag, + ShellType: req.Type, + }) + if cmdErr != nil { + c.JSON(http.StatusBadRequest, FileOpResponse{OK: false, Error: cmdErr.Error()}) + return + } + + useGET := strings.ToUpper(strings.TrimSpace(req.Method)) == "GET" + cmdParam := strings.TrimSpace(req.CmdParam) + if cmdParam == "" { + cmdParam = "cmd" + } + var httpReq *http.Request + if useGET { + targetURL := h.buildExecURL(req.URL, req.Type, req.Password, cmdParam, command) + httpReq, err = http.NewRequest(http.MethodGet, targetURL, nil) + } else { + body := h.buildExecBody(req.Type, req.Password, cmdParam, command) + httpReq, err = http.NewRequest(http.MethodPost, req.URL, bytes.NewReader(body)) + httpReq.Header.Set("Content-Type", "application/x-www-form-urlencoded") + } + if err != nil { + c.JSON(http.StatusInternalServerError, FileOpResponse{OK: false, Error: err.Error()}) + return + } + httpReq.Header.Set("User-Agent", "Mozilla/5.0 (compatible; CyberStrikeAI-WebShell/1.0)") + + resp, err := h.client.Do(httpReq) + if err != nil { + c.JSON(http.StatusOK, FileOpResponse{OK: false, Error: err.Error()}) + return + } + defer resp.Body.Close() + + out, readErr := io.ReadAll(resp.Body) + if readErr != nil { + h.logger.Warn("webshell fileop read body", zap.Error(readErr)) + } + output := decodeWebshellOutput(out, req.Encoding) + + c.JSON(http.StatusOK, FileOpResponse{ + OK: resp.StatusCode == http.StatusOK, + Output: output, + DetectedOS: detectedOS, + }) +} + +func (h *WebShellHandler) authorizedWebshellConnection(c *gin.Context, connectionID, requestURL string) (*database.WebShellConnection, bool) { + connectionID = strings.TrimSpace(connectionID) + if connectionID == "" { + return nil, false + } + if h.db == nil { + return nil, false + } + session, ok := security.CurrentSession(c) + if !ok || !h.db.UserCanAccessResource(session.UserID, session.Scope, "webshell", connectionID) { + return nil, false + } + conn, err := h.db.GetWebshellConnection(connectionID) + if err != nil || conn == nil { + return nil, false + } + if requestURL = strings.TrimSpace(requestURL); requestURL != "" && strings.TrimSpace(conn.URL) != requestURL { + return nil, false + } + return conn, true +} + +func (h *WebShellHandler) canAccessProject(c *gin.Context, projectID string) bool { + projectID = strings.TrimSpace(projectID) + if projectID == "" || h.db == nil { + return true + } + session, ok := security.CurrentSession(c) + if !ok { + return false + } + if session.Scope == database.RBACScopeAll { + return true + } + return h.db.UserCanAccessResource(session.UserID, session.Scope, "project", projectID) +} + +// ExecWithConnection 在指定 WebShell 连接上执行命令（供 MCP/Agent 等非 HTTP 调用） +func (h *WebShellHandler) ExecWithConnection(conn *database.WebShellConnection, command string) (output string, ok bool, errMsg string) { + if conn == nil { + return "", false, "connection is nil" + } + command = strings.TrimSpace(command) + if command == "" { + return "", false, "command is required" + } + useGET := strings.ToUpper(strings.TrimSpace(conn.Method)) == "GET" + cmdParam := strings.TrimSpace(conn.CmdParam) + if cmdParam == "" { + cmdParam = "cmd" + } + var httpReq *http.Request + var err error + if useGET { + targetURL := h.buildExecURL(conn.URL, conn.Type, conn.Password, cmdParam, command) + httpReq, err = http.NewRequest(http.MethodGet, targetURL, nil) + } else { + body := h.buildExecBody(conn.Type, conn.Password, cmdParam, command) + httpReq, err = http.NewRequest(http.MethodPost, conn.URL, bytes.NewReader(body)) + httpReq.Header.Set("Content-Type", "application/x-www-form-urlencoded") + } + if err != nil { + return "", false, err.Error() + } + httpReq.Header.Set("User-Agent", "Mozilla/5.0 (compatible; CyberStrikeAI-WebShell/1.0)") + resp, err := h.client.Do(httpReq) + if err != nil { + return "", false, err.Error() + } + defer resp.Body.Close() + out, readErr := io.ReadAll(resp.Body) + if readErr != nil { + h.logger.Warn("webshell ExecWithConnection read body", zap.Error(readErr)) + } + return decodeWebshellOutput(out, conn.Encoding), resp.StatusCode == http.StatusOK, "" +} + +// FileOpWithConnection 在指定 WebShell 连接上执行文件操作（供 MCP/Agent 调用），支持 list / read / write +func (h *WebShellHandler) FileOpWithConnection(conn *database.WebShellConnection, action, path, content, targetPath string) (output string, ok bool, errMsg string) { + if conn == nil { + return "", false, "connection is nil" + } + action = strings.ToLower(strings.TrimSpace(action)) + // MCP 入口仅开放 list / read / write 三种动作，与工具文档的承诺保持一致 + switch action { + case "list", "read", "write": + // 支持的动作 + default: + return "", false, "unsupported action: " + action + " (supported: list, read, write)" + } + + // 若连接的 OS 为 auto，先探活并持久化，避免 AI/MCP 每次都对 Windows 发 `ls -la` + osTag := conn.OS + if normalizeWebshellOS(osTag) == "auto" { + if probed := probeWebshellOSViaExec(func(cmd string) (string, bool) { + out, exOk, _ := h.ExecWithConnection(conn, cmd) + return out, exOk + }); probed != "" { + osTag = probed + conn.OS = probed // 本次请求内使用探活结果 + h.persistDetectedOS(conn.ID, probed) + } + } + + command, cmdErr := h.buildFileCommand(fileCommandInput{ + Action: action, + Path: path, + TargetPath: targetPath, + Content: content, + OS: osTag, + ShellType: conn.Type, + }) + if cmdErr != nil { + return "", false, cmdErr.Error() + } + useGET := strings.ToUpper(strings.TrimSpace(conn.Method)) == "GET" + cmdParam := strings.TrimSpace(conn.CmdParam) + if cmdParam == "" { + cmdParam = "cmd" + } + var httpReq *http.Request + var err error + if useGET { + targetURL := h.buildExecURL(conn.URL, conn.Type, conn.Password, cmdParam, command) + httpReq, err = http.NewRequest(http.MethodGet, targetURL, nil) + } else { + body := h.buildExecBody(conn.Type, conn.Password, cmdParam, command) + httpReq, err = http.NewRequest(http.MethodPost, conn.URL, bytes.NewReader(body)) + httpReq.Header.Set("Content-Type", "application/x-www-form-urlencoded") + } + if err != nil { + return "", false, err.Error() + } + httpReq.Header.Set("User-Agent", "Mozilla/5.0 (compatible; CyberStrikeAI-WebShell/1.0)") + resp, err := h.client.Do(httpReq) + if err != nil { + return "", false, err.Error() + } + defer resp.Body.Close() + out, readErr := io.ReadAll(resp.Body) + if readErr != nil { + h.logger.Warn("webshell FileOpWithConnection read body", zap.Error(readErr)) + } + return decodeWebshellOutput(out, conn.Encoding), resp.StatusCode == http.StatusOK, "" +} diff --git a/internal/handler/webshell_context.go b/internal/handler/webshell_context.go new file mode 100644 index 00000000..6a29c908 --- /dev/null +++ b/internal/handler/webshell_context.go @@ -0,0 +1,106 @@ +package handler + +import ( + "strings" + + "cyberstrike-ai/internal/database" +) + +// WebshellSkillHintDefault 对话页 / Eino 单代理共用的 Skills 说明，放在 webshell 上下文末尾， +// 供 AI 选择 skill 加载入口时参考。 +const WebshellSkillHintDefault = "Skills 包请使用「多代理 / Eino DeepAgent」会话中的内置 `skill` 工具渐进加载。" + +// WebshellSkillHintMultiAgent 多代理 / Eino 多代理准备阶段使用的 Skills 说明 +const WebshellSkillHintMultiAgent = "Skills 包请使用 Eino 多代理内置 `skill` 工具。" + +// webshellAssistantToolList AI 助手在 WebShell 上下文下允许使用的工具清单（展示给模型用）。 +// 注意：此处只是展示字符串，真正的权限限制是在调用方设置的 roleTools 切片里。 +const webshellAssistantToolList = "webshell_exec、webshell_file_list、webshell_file_read、webshell_file_write、record_vulnerability、list_vulnerabilities、get_vulnerability、upsert_project_fact、get_project_fact、list_project_facts、search_project_facts、deprecate_project_fact、restore_project_fact、list_knowledge_risk_types、search_knowledge_base" + +// BuildWebshellAssistantContext 根据连接信息与用户原始消息组装 AI 助手的上下文提示词。 +// 上下文包含：连接 ID、备注、目标系统（及对应命令集建议）、响应编码、可用工具清单、Skills 加载入口、 +// 以及最终的用户请求。调用方只需要决定 skillHint 的文案（默认使用 WebshellSkillHintDefault）。 +// +// 之所以把这段逻辑抽到共享函数里，是为了避免 agent.go / multi_agent_prepare.go 等多处复制粘贴， +// 并确保当我们升级 OS / Encoding 文案时只需要改一处、测一处、同步生效。 +func BuildWebshellAssistantContext(conn *database.WebShellConnection, skillHint, userMsg string) string { + if conn == nil { + // 兜底：调用方已保证 conn 非 nil，这里只是防御性返回原消息 + return userMsg + } + remark := conn.Remark + if remark == "" { + remark = conn.URL + } + + targetOS := resolveWebshellOS(conn.OS, conn.Type) // 归一为 "linux" / "windows" + encoding := normalizeWebshellEncoding(conn.Encoding) + if skillHint == "" { + skillHint = WebshellSkillHintDefault + } + + var b strings.Builder + b.Grow(512 + len(userMsg)) + + b.WriteString("[WebShell 助手上下文] 连接 ID：") + b.WriteString(conn.ID) + b.WriteString("，备注：") + b.WriteString(remark) + b.WriteByte('\n') + + // 目标系统：明确告诉 AI 能用/不能用的命令集，避免它对着 Windows 发 ls/cat/rm + b.WriteString("- 目标系统：") + b.WriteString(describeTargetOSForPrompt(targetOS)) + b.WriteByte('\n') + + // 响应编码：仅在非 auto 时显式告知，auto 模式由后端自适应，不打扰模型 + if encHint := describeEncodingForPrompt(encoding); encHint != "" { + b.WriteString("- 响应编码：") + b.WriteString(encHint) + b.WriteByte('\n') + } + + // 工具清单 & connection_id 约束：保持旧有表达，AI 已熟悉 + b.WriteString("可用工具（仅在该连接上操作时使用，connection_id 填 \"") + b.WriteString(conn.ID) + b.WriteString("\"）：") + b.WriteString(webshellAssistantToolList) + b.WriteString("。边渗透边记录：每确认新认知即 upsert_project_fact，每验证漏洞即 record_vulnerability，勿等会话结束。") + b.WriteString(skillHint) + b.WriteString("\n\n用户请求：") + b.WriteString(userMsg) + + return b.String() +} + +// describeTargetOSForPrompt 返回某个 OS 对应的中文描述 + 推荐命令集 + 反例， +// 命令列表覆盖文件管理最常用的 6 类动作（查看/读/删/改名/建目录/查找），让 AI 能直接照抄。 +func describeTargetOSForPrompt(targetOS string) string { + switch targetOS { + case "windows": + return "Windows（推荐 cmd/PowerShell：dir /a、type、del /q /f、move /y、md、ren；" + + "查找文件用 `dir /s /b 过滤词` 或 PowerShell `Get-ChildItem -Recurse`；" + + "避免 ls / cat / rm / mv / find 等 Unix 命令，否则将返回 `不是内部或外部命令`）" + case "linux": + return "Linux/Unix（推荐 sh/bash：ls -la、cat、rm -f、mv、mkdir -p；" + + "查找文件用 `find /path -name '*pattern*'`；" + + "避免 dir、type、del、move 等 Windows 命令）" + default: + // 理论上不会走到这里，resolveWebshellOS 已经兜底 + return "未知（请先执行 `uname || ver` 探测再决定命令集）" + } +} + +// describeEncodingForPrompt 返回响应编码的人类可读描述；auto 返回空串以减少 token。 +func describeEncodingForPrompt(encoding string) string { + switch encoding { + case "utf-8": + return "UTF-8（目标原生 UTF-8，无需额外解码）" + case "gbk": + return "GBK（中文 Windows；后端已自动转码为 UTF-8 返回，若仍出现大量 \\uFFFD 替换字符说明命令失败或编码识别错误）" + case "gb18030": + return "GB18030（后端已自动转码为 UTF-8 返回）" + default: + return "" + } +} diff --git a/internal/handler/webshell_context_test.go b/internal/handler/webshell_context_test.go new file mode 100644 index 00000000..743c1a9e --- /dev/null +++ b/internal/handler/webshell_context_test.go @@ -0,0 +1,170 @@ +package handler + +import ( + "strings" + "testing" + + "cyberstrike-ai/internal/database" +) + +func TestBuildWebshellAssistantContext_WindowsExplicit(t *testing.T) { + conn := &database.WebShellConnection{ + ID: "ws_win01", + Remark: "IIS Windows 靶机", + URL: "http://example.com/shell.php", + Type: "php", + OS: "windows", + Encoding: "gbk", + } + got := BuildWebshellAssistantContext(conn, WebshellSkillHintDefault, "列出当前目录并告诉我 flag 在哪") + + mustContain(t, got, + "[WebShell 助手上下文]", + "ws_win01", + "IIS Windows 靶机", + "目标系统：Windows", + "dir /a", + "move /y", + "避免 ls / cat / rm", + "响应编码：GBK", + "后端已自动转码为 UTF-8", + "connection_id 填 \"ws_win01\"", + "webshell_exec、webshell_file_list", + WebshellSkillHintDefault, + "用户请求：列出当前目录并告诉我 flag 在哪", + ) + // Windows 场景下不应出现 Linux 命令推荐 + mustNotContain(t, got, "推荐 sh/bash") +} + +func TestBuildWebshellAssistantContext_LinuxAutoFromPHP(t *testing.T) { + conn := &database.WebShellConnection{ + ID: "ws_lnx01", + Remark: "", // 测试备注为空时 fallback URL + URL: "http://example.com/a.php", + Type: "php", + OS: "auto", // auto + php → linux + Encoding: "", // auto 编码不显式提示 + } + got := BuildWebshellAssistantContext(conn, WebshellSkillHintDefault, "看看 /etc/passwd") + + mustContain(t, got, + "连接 ID：ws_lnx01", + "备注：http://example.com/a.php", // 备注空时 fallback URL + "目标系统：Linux/Unix", + "ls -la", + "mkdir -p", + "避免 dir、type、del、move", + "用户请求：看看 /etc/passwd", + ) + // encoding=auto 不应出现"响应编码："这一行 + mustNotContain(t, got, "响应编码：") + // Linux 场景不应出现 Windows 命令 + mustNotContain(t, got, "推荐 cmd/PowerShell") +} + +func TestBuildWebshellAssistantContext_AutoFromASPDefaultsToWindows(t *testing.T) { + // 保留向后兼容：旧连接没配 os，shellType=asp 时应视为 Windows + conn := &database.WebShellConnection{ + ID: "ws_asp01", + Remark: "老 ASP 靶机", + Type: "asp", + OS: "", // 空串等同 auto + Encoding: "gb18030", + } + got := BuildWebshellAssistantContext(conn, WebshellSkillHintMultiAgent, "查当前用户") + + mustContain(t, got, + "目标系统：Windows", + "响应编码：GB18030", + "后端已自动转码为 UTF-8 返回", + WebshellSkillHintMultiAgent, + ) + // 多代理 skill 文案里没有 DeepAgent，不应混入 default 文案 + mustNotContain(t, got, "DeepAgent") +} + +func TestBuildWebshellAssistantContext_MultiAgentSkillHint(t *testing.T) { + conn := &database.WebShellConnection{ID: "ws_m1", Remark: "x", Type: "php", OS: "linux"} + got := BuildWebshellAssistantContext(conn, WebshellSkillHintMultiAgent, "hi") + mustContain(t, got, WebshellSkillHintMultiAgent) + mustNotContain(t, got, "DeepAgent") +} + +func TestBuildWebshellAssistantContext_DefaultSkillHintFallback(t *testing.T) { + conn := &database.WebShellConnection{ID: "ws_d1", Remark: "x", Type: "php", OS: "linux"} + // skillHint 传空字符串时应回退到 default + got := BuildWebshellAssistantContext(conn, "", "hi") + mustContain(t, got, WebshellSkillHintDefault) +} + +func TestBuildWebshellAssistantContext_UTF8EncodingIsAnnotated(t *testing.T) { + conn := &database.WebShellConnection{ + ID: "ws_u1", Remark: "u", Type: "jsp", OS: "linux", Encoding: "utf-8", + } + got := BuildWebshellAssistantContext(conn, WebshellSkillHintDefault, "hi") + mustContain(t, got, "响应编码：UTF-8", "目标原生 UTF-8") +} + +func TestBuildWebshellAssistantContext_NilConnReturnsUserMsg(t *testing.T) { + // 防御性：conn == nil 时不 panic，直接返回原消息 + got := BuildWebshellAssistantContext(nil, WebshellSkillHintDefault, "just the message") + if got != "just the message" { + t.Errorf("nil conn should return userMsg as-is, got %q", got) + } +} + +func TestDescribeTargetOSForPrompt(t *testing.T) { + cases := map[string][]string{ + "windows": {"Windows", "dir /a", "move /y", "PowerShell"}, + "linux": {"Linux/Unix", "ls -la", "mkdir -p"}, + "": {"未知", "uname"}, // 防御性分支 + } + for in, wants := range cases { + got := describeTargetOSForPrompt(in) + for _, w := range wants { + if !strings.Contains(got, w) { + t.Errorf("describeTargetOSForPrompt(%q) should contain %q, got: %s", in, w, got) + } + } + } +} + +func TestDescribeEncodingForPrompt(t *testing.T) { + cases := map[string]string{ + "utf-8": "UTF-8", + "gbk": "GBK", + "gb18030": "GB18030", + "auto": "", + "": "", + } + for in, want := range cases { + got := describeEncodingForPrompt(in) + if want == "" && got != "" { + t.Errorf("describeEncodingForPrompt(%q) should return empty string, got: %s", in, got) + } + if want != "" && !strings.Contains(got, want) { + t.Errorf("describeEncodingForPrompt(%q) should contain %q, got: %s", in, want, got) + } + } +} + +// ---- 小工具 ---- + +func mustContain(t *testing.T, text string, substrings ...string) { + t.Helper() + for _, s := range substrings { + if !strings.Contains(text, s) { + t.Errorf("expected text to contain %q\n--- text ---\n%s", s, text) + } + } +} + +func mustNotContain(t *testing.T, text string, substrings ...string) { + t.Helper() + for _, s := range substrings { + if strings.Contains(text, s) { + t.Errorf("text should not contain %q\n--- text ---\n%s", s, text) + } + } +} diff --git a/internal/handler/webshell_encoding_test.go b/internal/handler/webshell_encoding_test.go new file mode 100644 index 00000000..f246008a --- /dev/null +++ b/internal/handler/webshell_encoding_test.go @@ -0,0 +1,103 @@ +package handler + +import ( + "testing" + + "golang.org/x/text/encoding/simplifiedchinese" + "golang.org/x/text/transform" +) + +// mustEncode 使用指定编码对 UTF-8 字符串做编码，得到原始字节，用于构造测试输入 +func mustEncode(t *testing.T, s string, enc string) []byte { + t.Helper() + var tr transform.Transformer + switch enc { + case "gbk": + tr = simplifiedchinese.GBK.NewEncoder() + case "gb18030": + tr = simplifiedchinese.GB18030.NewEncoder() + default: + t.Fatalf("unsupported test encoding: %s", enc) + } + out, _, err := transform.Bytes(tr, []byte(s)) + if err != nil { + t.Fatalf("mustEncode(%s) failed: %v", enc, err) + } + return out +} + +func TestNormalizeWebshellEncoding(t *testing.T) { + cases := map[string]string{ + "": "auto", + " ": "auto", + "auto": "auto", + "AUTO": "auto", + "utf-8": "utf-8", + "UTF-8": "utf-8", + "utf8": "utf-8", + "gbk": "gbk", + "GBK": "gbk", + "gb18030": "gb18030", + "big5": "auto", // 未支持的回退到 auto + "anything": "auto", + } + for in, want := range cases { + if got := normalizeWebshellEncoding(in); got != want { + t.Errorf("normalizeWebshellEncoding(%q) = %q, want %q", in, got, want) + } + } +} + +func TestDecodeWebshellOutput_AutoDetectsGBK(t *testing.T) { + // 模拟 Windows 中文 cmd 输出的 GBK 字节流 + want := "用户名 SID 类型" + raw := mustEncode(t, want, "gbk") + + // auto 模式：UTF-8 校验失败后应当回退 GB18030 解码，得到原始中文 + got := decodeWebshellOutput(raw, "auto") + if got != want { + t.Errorf("decodeWebshellOutput(auto) = %q, want %q", got, want) + } + + // 显式 GBK 模式：同样应当正确解码 + got = decodeWebshellOutput(raw, "gbk") + if got != want { + t.Errorf("decodeWebshellOutput(gbk) = %q, want %q", got, want) + } + + // 显式 GB18030 模式：GBK 是 GB18030 子集，也应正确解码 + got = decodeWebshellOutput(raw, "gb18030") + if got != want { + t.Errorf("decodeWebshellOutput(gb18030) = %q, want %q", got, want) + } +} + +func TestDecodeWebshellOutput_PassthroughUTF8(t *testing.T) { + // 已经是 UTF-8 的中文字符串，各模式都应返回原串（不破坏） + want := "hello 世界" + for _, enc := range []string{"", "auto", "utf-8"} { + if got := decodeWebshellOutput([]byte(want), enc); got != want { + t.Errorf("decodeWebshellOutput(%q) passthrough = %q, want %q", enc, got, want) + } + } +} + +func TestDecodeWebshellOutput_ASCIIStable(t *testing.T) { + // 纯 ASCII 在任何模式下都必须保持原样 + want := "whoami\nAdministrator\n" + for _, enc := range []string{"", "auto", "utf-8", "gbk", "gb18030"} { + if got := decodeWebshellOutput([]byte(want), enc); got != want { + t.Errorf("decodeWebshellOutput(%q) ASCII = %q, want %q", enc, got, want) + } + } +} + +func TestDecodeWebshellOutput_EmptyInput(t *testing.T) { + // 空输入直接返回空串，不做额外分配 + if got := decodeWebshellOutput(nil, "gbk"); got != "" { + t.Errorf("decodeWebshellOutput(nil) = %q, want empty", got) + } + if got := decodeWebshellOutput([]byte{}, "auto"); got != "" { + t.Errorf("decodeWebshellOutput([]) = %q, want empty", got) + } +} diff --git a/internal/handler/webshell_os_test.go b/internal/handler/webshell_os_test.go new file mode 100644 index 00000000..5cf47b6b --- /dev/null +++ b/internal/handler/webshell_os_test.go @@ -0,0 +1,348 @@ +package handler + +import ( + "encoding/base64" + "strings" + "testing" + + "go.uber.org/zap" +) + +func newTestWebShellHandler() *WebShellHandler { + return NewWebShellHandler(zap.NewNop(), nil) +} + +func TestNormalizeWebshellOS(t *testing.T) { + cases := map[string]string{ + "": "auto", + " ": "auto", + "auto": "auto", + "AUTO": "auto", + "linux": "linux", + "Linux": "linux", + "windows": "windows", + "WINDOWS": "windows", + "macos": "auto", // 未支持的回退 auto + "solaris": "auto", + } + for in, want := range cases { + if got := normalizeWebshellOS(in); got != want { + t.Errorf("normalizeWebshellOS(%q) = %q, want %q", in, got, want) + } + } +} + +func TestResolveWebshellOS(t *testing.T) { + type testCase struct { + osTag string + shellType string + want string + } + cases := []testCase{ + // 显式 OS：按用户选择，忽略 shellType + {"linux", "asp", "linux"}, + {"windows", "php", "windows"}, + {"LINUX", "jsp", "linux"}, + + // auto + 各种 shellType：asp/aspx → windows，其他 → linux + {"auto", "asp", "windows"}, + {"auto", "aspx", "windows"}, + {"auto", "ASP", "windows"}, + {"auto", "php", "linux"}, + {"auto", "jsp", "linux"}, + {"auto", "custom", "linux"}, + {"auto", "", "linux"}, + + // 空/未知 OS 等价 auto + {"", "asp", "windows"}, + {"", "php", "linux"}, + {"unknown", "aspx", "windows"}, + } + for _, c := range cases { + got := resolveWebshellOS(c.osTag, c.shellType) + if got != c.want { + t.Errorf("resolveWebshellOS(%q,%q) = %q, want %q", c.osTag, c.shellType, got, c.want) + } + } +} + +func TestQuoteCmdPath(t *testing.T) { + cases := map[string]string{ + "": `"."`, + `C:\Windows\Temp`: `"C:\Windows\Temp"`, + `C:\Program Files\a`: `"C:\Program Files\a"`, + `C:\weird"name\f.txt`: `"C:\weird""name\f.txt"`, + `.`: `"."`, + } + for in, want := range cases { + if got := quoteCmdPath(in); got != want { + t.Errorf("quoteCmdPath(%q) = %q, want %q", in, got, want) + } + } +} + +func TestQuoteShellSinglePosix(t *testing.T) { + cases := map[string]string{ + "": ".", + "/tmp/a b": "'/tmp/a b'", + "/tmp/it's.txt": `'/tmp/it'\''s.txt'`, + } + for in, want := range cases { + if got := quoteShellSinglePosix(in); got != want { + t.Errorf("quoteShellSinglePosix(%q) = %q, want %q", in, got, want) + } + } +} + +// TestBuildFileCommand_LinuxBranch 覆盖 Linux 目标下每个 action 产出的命令 +func TestBuildFileCommand_LinuxBranch(t *testing.T) { + h := newTestWebShellHandler() + base := fileCommandInput{OS: "linux", ShellType: "php"} + + mustContain := func(t *testing.T, cmd string, substrings ...string) { + t.Helper() + for _, s := range substrings { + if !strings.Contains(cmd, s) { + t.Errorf("expected command to contain %q, got: %s", s, cmd) + } + } + } + mustNotContain := func(t *testing.T, cmd string, substrings ...string) { + t.Helper() + for _, s := range substrings { + if strings.Contains(cmd, s) { + t.Errorf("command should not contain %q, got: %s", s, cmd) + } + } + } + + // list with empty path defaults to '.' + in := base + in.Action = "list" + cmd, err := h.buildFileCommand(in) + if err != nil { + t.Fatalf("list linux: unexpected err: %v", err) + } + mustContain(t, cmd, "ls -la", "'.'") + + // list with path containing spaces + in.Path = "/tmp/my files" + cmd, _ = h.buildFileCommand(in) + mustContain(t, cmd, "ls -la ", "'/tmp/my files'") + + // read with path + in = base + in.Action = "read" + in.Path = "/etc/passwd" + cmd, _ = h.buildFileCommand(in) + mustContain(t, cmd, "cat ", "'/etc/passwd'") + + // read without path → error + in.Path = "" + if _, err := h.buildFileCommand(in); err != errFileOpPathRequired { + t.Errorf("read empty path: want errFileOpPathRequired, got %v", err) + } + + // delete + in = base + in.Action = "delete" + in.Path = "/tmp/a.txt" + cmd, _ = h.buildFileCommand(in) + mustContain(t, cmd, "rm -f ", "'/tmp/a.txt'") + mustNotContain(t, cmd, "del") + + // mkdir + in.Action = "mkdir" + in.Path = "/tmp/new/sub" + cmd, _ = h.buildFileCommand(in) + mustContain(t, cmd, "mkdir -p ", "'/tmp/new/sub'") + + // rename + in = base + in.Action = "rename" + in.Path = "/tmp/a" + in.TargetPath = "/tmp/b" + cmd, _ = h.buildFileCommand(in) + mustContain(t, cmd, "mv -f ", "'/tmp/a'", "'/tmp/b'") + + // rename missing target → error + in.TargetPath = "" + if _, err := h.buildFileCommand(in); err != errFileOpRenameNeedsBothPaths { + t.Errorf("rename empty target: want errFileOpRenameNeedsBothPaths, got %v", err) + } + + // write + in = base + in.Action = "write" + in.Path = "/tmp/w.txt" + in.Content = "hello 世界" + cmd, _ = h.buildFileCommand(in) + b64 := base64.StdEncoding.EncodeToString([]byte("hello 世界")) + mustContain(t, cmd, "echo '"+b64+"'", "| base64 -d", "> '/tmp/w.txt'") + + // upload + in = base + in.Action = "upload" + in.Path = "/tmp/bin" + in.Content = "YWJjZA==" // base64 of "abcd" + cmd, _ = h.buildFileCommand(in) + mustContain(t, cmd, "echo 'YWJjZA=='", "| base64 -d", "> '/tmp/bin'") + + // upload oversized content → error + in.Content = strings.Repeat("A", 513*1024) + if _, err := h.buildFileCommand(in); err != errFileOpUploadTooLarge { + t.Errorf("upload too large: want errFileOpUploadTooLarge, got %v", err) + } + + // upload_chunk with chunk_index=0 uses single redirect + in = base + in.Action = "upload_chunk" + in.Path = "/tmp/bin" + in.Content = "YWJj" + in.ChunkIndex = 0 + cmd, _ = h.buildFileCommand(in) + mustContain(t, cmd, "base64 -d > '/tmp/bin'") + mustNotContain(t, cmd, ">>") + + // upload_chunk with chunk_index>0 uses append redirect + in.ChunkIndex = 1 + cmd, _ = h.buildFileCommand(in) + mustContain(t, cmd, "base64 -d >> '/tmp/bin'") + + // unsupported action + in = base + in.Action = "nope" + if _, err := h.buildFileCommand(in); err == nil || !strings.Contains(err.Error(), "unsupported action") { + t.Errorf("unknown action: want unsupported action error, got %v", err) + } +} + +// TestBuildFileCommand_WindowsBranch 覆盖 Windows 目标下每个 action 产出的命令 +func TestBuildFileCommand_WindowsBranch(t *testing.T) { + h := newTestWebShellHandler() + base := fileCommandInput{OS: "windows", ShellType: "php"} + + mustContain := func(t *testing.T, cmd string, substrings ...string) { + t.Helper() + for _, s := range substrings { + if !strings.Contains(cmd, s) { + t.Errorf("expected command to contain %q, got: %s", s, cmd) + } + } + } + mustNotContain := func(t *testing.T, cmd string, substrings ...string) { + t.Helper() + for _, s := range substrings { + if strings.Contains(cmd, s) { + t.Errorf("command should not contain %q, got: %s", s, cmd) + } + } + } + + // list + in := base + in.Action = "list" + cmd, _ := h.buildFileCommand(in) + mustContain(t, cmd, "dir /a ", `"."`) + mustNotContain(t, cmd, "ls -la") + + in.Path = `C:\Users\Public Docs` + cmd, _ = h.buildFileCommand(in) + mustContain(t, cmd, "dir /a ", `"C:\Users\Public Docs"`) + + // read + in = base + in.Action = "read" + in.Path = `C:\flag.txt` + cmd, _ = h.buildFileCommand(in) + mustContain(t, cmd, "type ", `"C:\flag.txt"`) + + // delete + in.Action = "delete" + cmd, _ = h.buildFileCommand(in) + mustContain(t, cmd, "del /q /f ", `"C:\flag.txt"`) + mustNotContain(t, cmd, "rm -f") + + // mkdir + in.Action = "mkdir" + in.Path = `C:\a\b\c` + cmd, _ = h.buildFileCommand(in) + mustContain(t, cmd, "md ", `"C:\a\b\c"`) + + // rename + in = base + in.Action = "rename" + in.Path = `C:\a.txt` + in.TargetPath = `C:\b.txt` + cmd, _ = h.buildFileCommand(in) + mustContain(t, cmd, "move /y ", `"C:\a.txt"`, `"C:\b.txt"`) + + // write → PowerShell base64 one-liner + in = base + in.Action = "write" + in.Path = `C:\out.txt` + in.Content = "hello 世界" + cmd, _ = h.buildFileCommand(in) + wantB64 := base64.StdEncoding.EncodeToString([]byte("hello 世界")) + mustContain(t, cmd, + "powershell -NoProfile -NonInteractive -Command", + "[Convert]::FromBase64String('"+wantB64+"')", + "[IO.File]::WriteAllBytes('C:\\out.txt'", + ) + mustNotContain(t, cmd, "echo ", "base64 -d") + + // upload (chunk_index=0 equivalent) uses WriteAllBytes + in = base + in.Action = "upload" + in.Path = `C:\bin\f` + in.Content = "YWJjZA==" + cmd, _ = h.buildFileCommand(in) + mustContain(t, cmd, "WriteAllBytes('C:\\bin\\f'", "FromBase64String('YWJjZA==')") + + // upload_chunk index=0 → WriteAllBytes + in.Action = "upload_chunk" + in.ChunkIndex = 0 + cmd, _ = h.buildFileCommand(in) + mustContain(t, cmd, "WriteAllBytes(") + mustNotContain(t, cmd, "FileMode]::Append") + + // upload_chunk index>0 → append (Open with Append mode) + in.ChunkIndex = 1 + cmd, _ = h.buildFileCommand(in) + mustContain(t, cmd, "[IO.FileMode]::Append", "FromBase64String('YWJjZA==')") +} + +// TestBuildFileCommand_AutoFallbackMatchesLegacyBehavior 确保 os=auto 时与旧版 shellType 判定行为完全一致 +// asp/aspx 视为 Windows（旧行为），其他视为 Linux。 +func TestBuildFileCommand_AutoFallbackMatchesLegacyBehavior(t *testing.T) { + h := newTestWebShellHandler() + + // asp + auto → windows 命令 + cmd, _ := h.buildFileCommand(fileCommandInput{Action: "list", OS: "auto", ShellType: "asp"}) + if !strings.Contains(cmd, "dir /a") { + t.Errorf("auto + asp should use Windows cmd, got: %s", cmd) + } + + cmd, _ = h.buildFileCommand(fileCommandInput{Action: "list", OS: "auto", ShellType: "aspx"}) + if !strings.Contains(cmd, "dir /a") { + t.Errorf("auto + aspx should use Windows cmd, got: %s", cmd) + } + + // php/jsp/custom + auto → linux 命令（与历史行为一致） + for _, st := range []string{"php", "jsp", "custom", ""} { + cmd, _ = h.buildFileCommand(fileCommandInput{Action: "list", OS: "auto", ShellType: st}) + if !strings.Contains(cmd, "ls -la") { + t.Errorf("auto + %q should use Linux cmd, got: %s", st, cmd) + } + } + + // 显式 OS 覆盖 shellType + cmd, _ = h.buildFileCommand(fileCommandInput{Action: "list", OS: "windows", ShellType: "php"}) + if !strings.Contains(cmd, "dir /a") { + t.Errorf("explicit windows should override php shellType, got: %s", cmd) + } + cmd, _ = h.buildFileCommand(fileCommandInput{Action: "list", OS: "linux", ShellType: "asp"}) + if !strings.Contains(cmd, "ls -la") { + t.Errorf("explicit linux should override asp shellType, got: %s", cmd) + } +} diff --git a/internal/handler/webshell_probe.go b/internal/handler/webshell_probe.go new file mode 100644 index 00000000..75917206 --- /dev/null +++ b/internal/handler/webshell_probe.go @@ -0,0 +1,127 @@ +package handler + +import ( + "bytes" + "io" + "net/http" + "strings" + + "go.uber.org/zap" +) + +// webshellOSProbeCommand 探活命令：利用 Windows cmd 与 POSIX shell 对 `%OS%` 展开差异进行判定。 +// - Windows cmd：`%OS%` 被展开为 `Windows_NT`，回显 `:OSPROBE_Windows_NT:END` +// - POSIX sh/bash：`%OS%` 不是变量语法，作为字面量原样保留，回显 `:OSPROBE_%OS%:END` +// +// 一条命令即可得到明确的、互斥的信号，避免探活成本（相比发两次命令）。 +// 冒号包裹是为了避免部分 shell 输出多余空白/BOM 时字符串匹配失效。 +const webshellOSProbeCommand = "echo :OSPROBE_%OS%:END" + +// probeWebshellOSViaExec 通过一次命令执行的回显推断目标操作系统。 +// +// 返回值： +// - "windows" / "linux"：识别成功 +// - ""：无法判定（调用方应保留既有 fallback 逻辑） +// +// 入参 execFn 是一个"发命令并拿到回显"的闭包；让 HTTP 入口和 MCP 入口可以共用同一套探活逻辑 +// 而不必关心底层是如何发包的。 +func probeWebshellOSViaExec(execFn func(cmd string) (output string, ok bool)) string { + if execFn == nil { + return "" + } + out, ok := execFn(webshellOSProbeCommand) + if !ok { + return "" + } + return classifyWebshellOSProbeOutput(out) +} + +// classifyWebshellOSProbeOutput 纯函数：根据探活命令的回显判定 OS。 +// 抽出来是为了单测可直接覆盖所有分支，无需真实 HTTP 调用。 +func classifyWebshellOSProbeOutput(out string) string { + if out == "" { + return "" + } + lower := strings.ToLower(out) + + // Windows 强信号：cmd.exe 成功展开了 %OS% 变量 + if strings.Contains(out, "Windows_NT") { + return "windows" + } + // 容错：部分老版本 Windows 可能 `%OS%` 展开为其他字样（极少见），再看 PATH/OS 等次级线索 + if strings.Contains(lower, "microsoft windows") { + return "windows" + } + + // Linux/Unix 强信号：`%OS%` 字面量被原样回显，说明 shell 不是 cmd.exe + if strings.Contains(out, "%OS%") { + return "linux" + } + + // 次级线索：部分 webshell 在 Linux 上可能走了其他外壳（如 zsh/ash）， + // 但它们对 `%OS%` 同样不展开；若命中 OSPROBE 头部却没拿到 %OS% 字面量， + // 说明回显被中途截断或过滤，保守返回空让上层 fallback。 + return "" +} + +// newHTTPExecFn 为 HTTP FileOp 路径构造"发命令取回显"的闭包，供探活复用。 +// 参数来自 HTTP 请求，复用 buildExecURL / buildExecBody 两个已有的命令编排器， +// 确保探活包与实际文件操作包走完全一致的 webshell 协议（GET/POST、参数名、编码）。 +func (h *WebShellHandler) newHTTPExecFn(targetURL, password, shellType, method, cmdParam, encoding string) func(string) (string, bool) { + useGET := strings.ToUpper(strings.TrimSpace(method)) == "GET" + if strings.TrimSpace(cmdParam) == "" { + cmdParam = "cmd" + } + return func(cmd string) (string, bool) { + var ( + httpReq *http.Request + err error + ) + if useGET { + u := h.buildExecURL(targetURL, shellType, password, cmdParam, cmd) + httpReq, err = http.NewRequest(http.MethodGet, u, nil) + } else { + body := h.buildExecBody(shellType, password, cmdParam, cmd) + httpReq, err = http.NewRequest(http.MethodPost, targetURL, bytes.NewReader(body)) + if err == nil { + httpReq.Header.Set("Content-Type", "application/x-www-form-urlencoded") + } + } + if err != nil { + return "", false + } + httpReq.Header.Set("User-Agent", "Mozilla/5.0 (compatible; CyberStrikeAI-WebShell/1.0)") + resp, err := h.client.Do(httpReq) + if err != nil { + return "", false + } + defer resp.Body.Close() + raw, _ := io.ReadAll(resp.Body) + return decodeWebshellOutput(raw, encoding), resp.StatusCode == http.StatusOK + } +} + +// persistDetectedOS 把探活结果回写到连接表；失败只记日志不阻断主流程。 +// 设计上故意只触发 UPDATE，不会新建记录，因此即便 connectionID 不存在也只是悄悄放弃。 +func (h *WebShellHandler) persistDetectedOS(connectionID, detected string) { + connectionID = strings.TrimSpace(connectionID) + detected = normalizeWebshellOS(detected) + if connectionID == "" || detected == "" || detected == "auto" { + return + } + conn, err := h.db.GetWebshellConnection(connectionID) + if err != nil || conn == nil { + // 不是所有调用方都能提供有效 ID（比如临时测试），这里静默返回 + return + } + if normalizeWebshellOS(conn.OS) != "auto" { + // 用户已经显式选过 OS，尊重用户选择，不自动覆盖 + return + } + conn.OS = detected + if err := h.db.UpdateWebshellConnection(conn); err != nil { + h.logger.Warn("webshell 探活结果持久化失败", zap.String("id", connectionID), zap.String("os", detected), zap.Error(err)) + return + } + h.logger.Info("webshell auto OS 探活成功并持久化", zap.String("id", connectionID), zap.String("os", detected)) +} diff --git a/internal/handler/webshell_probe_test.go b/internal/handler/webshell_probe_test.go new file mode 100644 index 00000000..03917315 --- /dev/null +++ b/internal/handler/webshell_probe_test.go @@ -0,0 +1,68 @@ +package handler + +import "testing" + +func TestClassifyWebshellOSProbeOutput(t *testing.T) { + cases := []struct { + name string + in string + want string + }{ + {"Windows cmd 回显完整", ":OSPROBE_Windows_NT:END\r\n", "windows"}, + {"Windows cmd 回显带额外空行", "\r\n:OSPROBE_Windows_NT:END\r\n", "windows"}, + {"Windows 次级线索 - ver banner", "Microsoft Windows [版本 10.0.19045]\r\n", "windows"}, + {"Linux sh 字面量回显", ":OSPROBE_%OS%:END\n", "linux"}, + {"Linux 紧凑输出（无换行）", ":OSPROBE_%OS%:END", "linux"}, + {"空输出 - 无法判定", "", ""}, + {"被过滤的输出 - 无法判定", "something weird", ""}, + {"仅有 OSPROBE 前缀但被截断 - 保守返回空", ":OSPROBE_:END", ""}, + } + for _, c := range cases { + if got := classifyWebshellOSProbeOutput(c.in); got != c.want { + t.Errorf("case %q: got %q, want %q", c.name, got, c.want) + } + } +} + +func TestProbeWebshellOSViaExec_SendsOneCommandOnly(t *testing.T) { + var calls []string + fn := func(cmd string) (string, bool) { + calls = append(calls, cmd) + return ":OSPROBE_Windows_NT:END", true + } + got := probeWebshellOSViaExec(fn) + if got != "windows" { + t.Fatalf("want windows, got %q", got) + } + if len(calls) != 1 { + t.Fatalf("probe should issue exactly one exec call, got %d: %v", len(calls), calls) + } + if calls[0] != webshellOSProbeCommand { + t.Errorf("probe command mismatch: got %q", calls[0]) + } +} + +func TestProbeWebshellOSViaExec_NotOkReturnsEmpty(t *testing.T) { + // HTTP 非 200 的场景：execFn 返回 ok=false，探活应放弃 + fn := func(cmd string) (string, bool) { return "whatever", false } + if got := probeWebshellOSViaExec(fn); got != "" { + t.Errorf("want empty when exec not ok, got %q", got) + } +} + +func TestProbeWebshellOSViaExec_NilSafeguard(t *testing.T) { + if got := probeWebshellOSViaExec(nil); got != "" { + t.Errorf("nil execFn should return empty, got %q", got) + } +} + +func TestProbeWebshellOSViaExec_LinuxUname(t *testing.T) { + // 某些 webshell 对 `%OS%` 字面量也会过滤（例如安全规则）， + // 但主要路径是"%OS% 字面量被原样回显"。这里覆盖标准 Linux 场景。 + fn := func(cmd string) (string, bool) { + return ":OSPROBE_%OS%:END\n", true + } + if got := probeWebshellOSViaExec(fn); got != "linux" { + t.Errorf("Linux case: want linux, got %q", got) + } +} diff --git a/internal/handler/webshell_rbac_test.go b/internal/handler/webshell_rbac_test.go new file mode 100644 index 00000000..33b86e1e --- /dev/null +++ b/internal/handler/webshell_rbac_test.go @@ -0,0 +1,156 @@ +package handler + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "path/filepath" + "testing" + "time" + + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +func TestWebshellExecRequiresConnectionAccessWhenConnectionIDProvided(t *testing.T) { + gin.SetMode(gin.TestMode) + db, user, allowed, hidden := setupWebshellRBACTest(t) + handler := NewWebShellHandler(zap.NewNop(), db) + + w := performWebshellJSON(user, http.MethodPost, "/api/webshell/exec", map[string]interface{}{ + "url": hidden.URL, + "connection_id": hidden.ID, + "command": "id", + }, handler.Exec) + if w.Code != http.StatusForbidden { + t.Fatalf("hidden connection status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String()) + } + + w = performWebshellJSON(user, http.MethodPost, "/api/webshell/exec", map[string]interface{}{ + "url": hidden.URL, + "connection_id": allowed.ID, + "command": "id", + }, handler.Exec) + if w.Code != http.StatusForbidden { + t.Fatalf("mismatched URL status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String()) + } +} + +func TestWebshellExecAllowsAdHocURLWithoutConnectionID(t *testing.T) { + gin.SetMode(gin.TestMode) + _, user, _, _ := setupWebshellRBACTest(t) + handler := NewWebShellHandler(zap.NewNop(), nil) + // Ad-hoc probe (connectivity test before save) must not be rejected as "无权访问". + // The target URL will fail to connect; we only assert auth allows the request through. + w := performWebshellJSON(user, http.MethodPost, "/api/webshell/exec", map[string]interface{}{ + "url": "http://127.0.0.1:1/admin", "command": "id", + }, handler.Exec) + if w.Code == http.StatusForbidden { + t.Fatalf("ad-hoc URL status = %d, want non-403: %s", w.Code, w.Body.String()) + } + var resp ExecResponse + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatalf("decode response: %v body=%s", err, w.Body.String()) + } + if resp.OK { + t.Fatalf("expected connection failure for closed port, got ok=true") + } +} + +func TestWebshellExecRejectsAdHocWithoutWritePermission(t *testing.T) { + gin.SetMode(gin.TestMode) + user := &database.RBACUser{ID: "u_ro", Username: "readonly"} + handler := NewWebShellHandler(zap.NewNop(), nil) + payload, _ := json.Marshal(map[string]interface{}{ + "url": "http://127.0.0.1/admin", "command": "id", + }) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodPost, "/api/webshell/exec", bytes.NewReader(payload)) + c.Request.Header.Set("Content-Type", "application/json") + c.Set(security.ContextSessionKey, security.Session{ + UserID: user.ID, + Username: user.Username, + Permissions: map[string]bool{"webshell:read": true}, + Scope: database.RBACScopeAssigned, + }) + handler.Exec(c) + if w.Code != http.StatusForbidden { + t.Fatalf("status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String()) + } +} + +func TestWebshellFileOpRequiresConnectionAccessWhenConnectionIDProvided(t *testing.T) { + gin.SetMode(gin.TestMode) + db, user, _, hidden := setupWebshellRBACTest(t) + handler := NewWebShellHandler(zap.NewNop(), db) + + w := performWebshellJSON(user, http.MethodPost, "/api/webshell/file", map[string]interface{}{ + "url": hidden.URL, + "connection_id": hidden.ID, + "action": "list", + "path": ".", + }, handler.FileOp) + if w.Code != http.StatusForbidden { + t.Fatalf("status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String()) + } +} + +func setupWebshellRBACTest(t *testing.T) (*database.DB, *database.RBACUser, *database.WebShellConnection, *database.WebShellConnection) { + t.Helper() + db, err := database.NewDB(filepath.Join(t.TempDir(), "webshell-rbac.db"), zap.NewNop()) + if err != nil { + t.Fatalf("NewDB: %v", err) + } + t.Cleanup(func() { _ = db.Close() }) + user, err := db.CreateRBACUser("operator1", "Operator One", "hash", true, nil) + if err != nil { + t.Fatalf("CreateRBACUser: %v", err) + } + allowed := &database.WebShellConnection{ + ID: "ws_allowed", + URL: "http://127.0.0.1/allowed.php", + Type: "php", + Method: "post", + CmdParam: "cmd", + CreatedAt: time.Now(), + } + hidden := &database.WebShellConnection{ + ID: "ws_hidden", + URL: "http://127.0.0.1/hidden.php", + Type: "php", + Method: "post", + CmdParam: "cmd", + CreatedAt: time.Now(), + } + if err := db.CreateWebshellConnection(allowed); err != nil { + t.Fatalf("CreateWebshellConnection allowed: %v", err) + } + if err := db.CreateWebshellConnection(hidden); err != nil { + t.Fatalf("CreateWebshellConnection hidden: %v", err) + } + if err := db.AssignResourceToUser(user.ID, "webshell", allowed.ID); err != nil { + t.Fatalf("AssignResourceToUser: %v", err) + } + return db, user, allowed, hidden +} + +func performWebshellJSON(user *database.RBACUser, method, path string, body map[string]interface{}, handler gin.HandlerFunc) *httptest.ResponseRecorder { + payload, _ := json.Marshal(body) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(method, path, bytes.NewReader(payload)) + c.Request.Header.Set("Content-Type", "application/json") + c.Set(security.ContextSessionKey, security.Session{ + UserID: user.ID, + Username: user.Username, + Permissions: map[string]bool{"webshell:write": true}, + Scope: database.RBACScopeAssigned, + }) + handler(c) + return w +} diff --git a/internal/handler/wechat_robot.go b/internal/handler/wechat_robot.go new file mode 100644 index 00000000..93a5ea8f --- /dev/null +++ b/internal/handler/wechat_robot.go @@ -0,0 +1,293 @@ +package handler + +import ( + "context" + "net/http" + "strings" + "sync" + "time" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/robot/ilink" + + "github.com/gin-gonic/gin" + "github.com/google/uuid" + "go.uber.org/zap" +) + +const wechatLoginTTL = 5 * time.Minute + +// WechatConfigSaver 绑定成功后写入配置并重启机器人连接 +type WechatConfigSaver interface { + ApplyWechatRobotBinding(cfg config.RobotWechatConfig) error +} + +type wechatLoginSession struct { + QRCode string + QRCodeImgURL string + PendingVerify string + CurrentBaseURL string + StartedAt time.Time +} + +// WechatRobotHandler 微信 iLink 机器人（扫码绑定 + 配置） +type WechatRobotHandler struct { + config *config.Config + configSaver WechatConfigSaver + logger *zap.Logger + mu sync.Mutex + logins map[string]*wechatLoginSession +} + +// NewWechatRobotHandler 创建微信机器人处理器 +func NewWechatRobotHandler(cfg *config.Config, saver WechatConfigSaver, logger *zap.Logger) *WechatRobotHandler { + return &WechatRobotHandler{ + config: cfg, + configSaver: saver, + logger: logger, + logins: make(map[string]*wechatLoginSession), + } +} + +func (h *WechatRobotHandler) purgeExpiredLogins() { + now := time.Now() + for k, v := range h.logins { + if now.Sub(v.StartedAt) > wechatLoginTTL { + delete(h.logins, k) + } + } +} + +func (h *WechatRobotHandler) ilinkClient(baseURL string) *ilink.Client { + ver := h.config.Version + if ver == "" { + ver = "1.0.0" + } + ver = strings.TrimPrefix(strings.TrimSpace(ver), "v") + ver = strings.TrimPrefix(ver, "V") + wc := h.config.Robots.Wechat + return ilink.NewClient(baseURL, wc.BotToken, wc.BotAgent, ilink.BuildClientVersion(ver)) +} + +// HandleWechatQRCode POST /api/robot/wechat/qrcode — 生成绑定二维码 +func (h *WechatRobotHandler) HandleWechatQRCode(c *gin.Context) { + h.mu.Lock() + h.purgeExpiredLogins() + h.mu.Unlock() + + var req struct { + BotType string `json:"bot_type"` + } + _ = c.ShouldBindJSON(&req) + + botType := req.BotType + if botType == "" { + botType = h.config.Robots.Wechat.BotType + } + if botType == "" { + botType = ilink.DefaultBotType + } + baseURL := h.config.Robots.Wechat.BaseURL + if baseURL == "" { + baseURL = ilink.DefaultBaseURL + } + + var localTokens []string + if t := h.config.Robots.Wechat.BotToken; t != "" { + localTokens = []string{t} + } + + client := h.ilinkClient(baseURL) + ctx, cancel := context.WithTimeout(c.Request.Context(), 30*time.Second) + defer cancel() + + qr, err := client.GetBotQRCode(ctx, botType, localTokens) + if err != nil { + h.logger.Warn("获取微信二维码失败", zap.Error(err)) + c.JSON(http.StatusBadGateway, gin.H{"error": "获取二维码失败: " + err.Error()}) + return + } + if qr.QRCode == "" || qr.QRCodeImgContent == "" { + c.JSON(http.StatusBadGateway, gin.H{"error": "微信服务器未返回有效二维码"}) + return + } + + sessionKey := uuid.New().String() + h.mu.Lock() + h.logins[sessionKey] = &wechatLoginSession{ + QRCode: qr.QRCode, + QRCodeImgURL: qr.QRCodeImgContent, + CurrentBaseURL: baseURL, + StartedAt: time.Now(), + } + h.mu.Unlock() + + resp := gin.H{ + "session_key": sessionKey, + "qrcode": qr.QRCode, + "qrcode_open_url": qr.QRCodeImgContent, + "message": "请使用微信扫描二维码并确认绑定", + } + if dataURL, err := ilink.QRCodeDataURL(qr.QRCodeImgContent, 256); err != nil { + h.logger.Warn("生成二维码图片失败", zap.Error(err)) + } else { + resp["qrcode_image_data_url"] = dataURL + } + + c.JSON(http.StatusOK, resp) +} + +// HandleWechatQRCodeStatus GET /api/robot/wechat/qrcode/status — 轮询扫码状态 +func (h *WechatRobotHandler) HandleWechatQRCodeStatus(c *gin.Context) { + sessionKey := c.Query("session_key") + verifyCode := c.Query("verify_code") + if sessionKey == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "缺少 session_key"}) + return + } + + h.mu.Lock() + sess, ok := h.logins[sessionKey] + h.mu.Unlock() + if !ok { + c.JSON(http.StatusNotFound, gin.H{"error": "登录会话不存在或已过期，请重新生成二维码"}) + return + } + if time.Since(sess.StartedAt) > wechatLoginTTL { + h.mu.Lock() + delete(h.logins, sessionKey) + h.mu.Unlock() + c.JSON(http.StatusGone, gin.H{"error": "二维码已过期，请重新生成"}) + return + } + + baseURL := sess.CurrentBaseURL + if baseURL == "" { + baseURL = ilink.DefaultBaseURL + } + vc := verifyCode + if vc == "" { + vc = sess.PendingVerify + } + + client := h.ilinkClient(baseURL) + ctx, cancel := context.WithTimeout(c.Request.Context(), 40*time.Second) + defer cancel() + + st, err := client.GetQRCodeStatus(ctx, sess.QRCode, vc) + if err != nil { + h.logger.Warn("轮询微信二维码状态失败", zap.Error(err)) + c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()}) + return + } + + switch st.Status { + case "wait", "scaned": + c.JSON(http.StatusOK, gin.H{"status": st.Status}) + return + case "need_verifycode": + c.JSON(http.StatusOK, gin.H{ + "status": st.Status, + "message": "请在手机微信查看配对数字，并在下方输入", + }) + return + case "scaned_but_redirect": + if st.RedirectHost != "" { + h.mu.Lock() + if s, ok := h.logins[sessionKey]; ok { + s.CurrentBaseURL = "https://" + st.RedirectHost + } + h.mu.Unlock() + } + c.JSON(http.StatusOK, gin.H{"status": st.Status}) + return + case "binded_redirect": + h.mu.Lock() + delete(h.logins, sessionKey) + h.mu.Unlock() + c.JSON(http.StatusOK, gin.H{ + "status": st.Status, + "already_connected": true, + "message": "该微信已绑定过，无需重复绑定", + }) + return + case "confirmed": + if st.BotToken == "" || st.ILinkBotID == "" { + c.JSON(http.StatusBadGateway, gin.H{"error": "绑定确认成功但缺少 bot_token"}) + return + } + saveBase := st.BaseURL + if saveBase == "" { + saveBase = baseURL + } + wc := h.config.Robots.Wechat + wc.Enabled = true + wc.BotToken = st.BotToken + wc.ILinkBotID = st.ILinkBotID + wc.ILinkUserID = st.ILinkUserID + wc.BaseURL = saveBase + if wc.BotType == "" { + wc.BotType = ilink.DefaultBotType + } + if wc.BotAgent == "" { + wc.BotAgent = ilink.DefaultBotAgent + } + if h.configSaver != nil { + if err := h.configSaver.ApplyWechatRobotBinding(wc); err != nil { + h.logger.Warn("保存微信机器人配置失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": "保存配置失败: " + err.Error()}) + return + } + } else { + h.config.Robots.Wechat = wc + } + h.mu.Lock() + delete(h.logins, sessionKey) + h.mu.Unlock() + c.JSON(http.StatusOK, gin.H{ + "status": "confirmed", + "message": "绑定成功，微信机器人已启用", + "ilink_bot_id": st.ILinkBotID, + "ilink_user_id": st.ILinkUserID, + }) + return + default: + c.JSON(http.StatusOK, gin.H{"status": st.Status}) + } +} + +// HandleWechatVerifyCode POST /api/robot/wechat/qrcode/verify — 提交手机配对数字 +func (h *WechatRobotHandler) HandleWechatVerifyCode(c *gin.Context) { + var req struct { + SessionKey string `json:"session_key"` + VerifyCode string `json:"verify_code"` + } + if err := c.ShouldBindJSON(&req); err != nil || req.SessionKey == "" || req.VerifyCode == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "需要 session_key 与 verify_code"}) + return + } + h.mu.Lock() + sess, ok := h.logins[req.SessionKey] + if ok { + sess.PendingVerify = req.VerifyCode + } + h.mu.Unlock() + if !ok { + c.JSON(http.StatusNotFound, gin.H{"error": "登录会话不存在或已过期"}) + return + } + c.JSON(http.StatusOK, gin.H{"message": "已提交配对码，请继续等待绑定"}) +} + +// HandleWechatStatus GET /api/robot/wechat/status — 当前绑定状态（供前端展示） +func (h *WechatRobotHandler) HandleWechatStatus(c *gin.Context) { + wc := h.config.Robots.Wechat + bound := wc.BotToken != "" && wc.ILinkBotID != "" + c.JSON(http.StatusOK, gin.H{ + "enabled": wc.Enabled, + "bound": bound, + "ilink_bot_id": wc.ILinkBotID, + "ilink_user_id": wc.ILinkUserID, + "base_url": wc.BaseURL, + }) +} diff --git a/internal/handler/workflow.go b/internal/handler/workflow.go new file mode 100644 index 00000000..2b9a45c9 --- /dev/null +++ b/internal/handler/workflow.go @@ -0,0 +1,255 @@ +package handler + +import ( + "encoding/json" + "errors" + "net/http" + "strings" + + "cyberstrike-ai/internal/agent" + "cyberstrike-ai/internal/audit" + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + workflowrunner "cyberstrike-ai/internal/workflow" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +type WorkflowHandler struct { + db *database.DB + logger *zap.Logger + audit *audit.Service + agent *agent.Agent + cfg *config.Config +} + +func NewWorkflowHandler(db *database.DB, logger *zap.Logger) *WorkflowHandler { + return &WorkflowHandler{db: db, logger: logger} +} + +func (h *WorkflowHandler) SetAudit(s *audit.Service) { + h.audit = s +} + +type workflowSaveRequest struct { + ID string `json:"id"` + Name string `json:"name"` + Description string `json:"description,omitempty"` + Version int `json:"version,omitempty"` + Enabled *bool `json:"enabled,omitempty"` + Graph json.RawMessage `json:"graph,omitempty"` + GraphJSON json.RawMessage `json:"graph_json,omitempty"` +} + +type workflowDryRunRequest struct { + Graph json.RawMessage `json:"graph,omitempty"` + GraphJSON json.RawMessage `json:"graph_json,omitempty"` + Inputs map[string]interface{} `json:"inputs,omitempty"` +} + +type workflowGenerateDraftRequest struct { + Prompt string `json:"prompt"` + Options workflowrunner.DraftOptions `json:"options"` + AvailableTools []workflowrunner.DraftTool `json:"available_tools,omitempty"` +} + +func (h *WorkflowHandler) List(c *gin.Context) { + includeDisabled := strings.EqualFold(c.Query("includeDisabled"), "true") || c.Query("include_disabled") == "1" + items, err := h.db.ListWorkflowDefinitions(includeDisabled) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"workflows": items}) +} + +func (h *WorkflowHandler) Get(c *gin.Context) { + id := strings.TrimSpace(c.Param("id")) + wf, err := h.db.GetWorkflowDefinition(id) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if wf == nil { + c.JSON(http.StatusNotFound, gin.H{"error": "工作流不存在"}) + return + } + c.JSON(http.StatusOK, gin.H{"workflow": wf}) +} + +func (h *WorkflowHandler) Create(c *gin.Context) { + h.save(c, "") +} + +func (h *WorkflowHandler) Validate(c *gin.Context) { + var req workflowSaveRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"ok": false, "error": "无效的请求参数: " + err.Error()}) + return + } + graph := req.Graph + if len(graph) == 0 { + graph = req.GraphJSON + } + if len(graph) == 0 { + c.JSON(http.StatusBadRequest, gin.H{"ok": false, "error": "graph 不能为空"}) + return + } + if !json.Valid(graph) { + c.JSON(http.StatusBadRequest, gin.H{"ok": false, "error": "graph 必须是合法 JSON"}) + return + } + if err := workflowrunner.ValidateGraphJSON(c.Request.Context(), string(graph)); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"ok": false, "error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"ok": true}) +} + +func (h *WorkflowHandler) DryRun(c *gin.Context) { + var req workflowDryRunRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()}) + return + } + graph := req.Graph + if len(graph) == 0 { + graph = req.GraphJSON + } + if len(graph) == 0 || !json.Valid(graph) { + c.JSON(http.StatusBadRequest, gin.H{"error": "graph 必须是合法 JSON"}) + return + } + inputs := make(map[string]any, len(req.Inputs)) + for k, v := range req.Inputs { + inputs[k] = v + } + result, err := workflowrunner.DryRunGraphJSON(c.Request.Context(), string(graph), inputs) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"result": result}) +} + +func (h *WorkflowHandler) GenerateDraft(c *gin.Context) { + var req workflowGenerateDraftRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()}) + return + } + draftReq := workflowrunner.DraftRequest{ + Prompt: req.Prompt, + Options: req.Options, + AvailableTools: req.AvailableTools, + } + var result *workflowrunner.DraftResult + var llmErr error + if h.cfg != nil { + if llmCfg, _, ok := h.cfg.ResolveAIChannel(""); ok && strings.TrimSpace(llmCfg.APIKey) != "" && strings.TrimSpace(llmCfg.Model) != "" { + result, llmErr = workflowrunner.GenerateDraftFromLLM(c.Request.Context(), draftReq, llmCfg, h.logger) + } else { + llmErr = errors.New("AI 通道未配置 api_key 或 model") + } + } else { + llmErr = errors.New("工作流生成器未加载平台 AI 配置") + } + if llmErr != nil { + c.JSON(http.StatusBadGateway, gin.H{"error": "大模型生成失败: " + llmErr.Error()}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "workflow", "generate_draft", "自然语言生成工作流草稿", "", "", map[string]interface{}{ + "generator": result.Generator, + "nodes": result.Stats["nodes"], + "edges": result.Stats["edges"], + "high_risk": result.Audit.HighRisk, + "savable": result.Audit.Savable, + }) + } + c.JSON(http.StatusOK, gin.H{"result": result}) +} + +func (h *WorkflowHandler) Update(c *gin.Context) { + h.save(c, c.Param("id")) +} + +func (h *WorkflowHandler) save(c *gin.Context, pathID string) { + var req workflowSaveRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()}) + return + } + id := strings.TrimSpace(req.ID) + if strings.TrimSpace(pathID) != "" { + id = strings.TrimSpace(pathID) + } + name := strings.TrimSpace(req.Name) + if id == "" || name == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "工作流 id 和 name 不能为空"}) + return + } + graph := req.Graph + if len(graph) == 0 { + graph = req.GraphJSON + } + if len(graph) == 0 { + graph = []byte(`{"nodes":[],"edges":[],"config":{}}`) + } + if !json.Valid(graph) { + c.JSON(http.StatusBadRequest, gin.H{"error": "graph 必须是合法 JSON"}) + return + } + if err := workflowrunner.ValidateGraphJSON(c.Request.Context(), string(graph)); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "工作流图无法编译: " + err.Error()}) + return + } + var probe interface{} + if err := json.Unmarshal(graph, &probe); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "graph JSON 解析失败: " + err.Error()}) + return + } + enabled := true + if req.Enabled != nil { + enabled = *req.Enabled + } + wf := &database.WorkflowDefinition{ + ID: id, + Name: name, + Description: strings.TrimSpace(req.Description), + Version: req.Version, + GraphJSON: string(graph), + Enabled: enabled, + } + if err := h.db.UpsertWorkflowDefinition(wf); err != nil { + if h.logger != nil { + h.logger.Warn("保存工作流失败", zap.String("id", id), zap.Error(err)) + } + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + saved, _ := h.db.GetWorkflowDefinition(id) + workflowrunner.InvalidateCompiledCache(id) + if h.audit != nil { + h.audit.RecordOK(c, "workflow", "save", "保存工作流", "workflow", id, map[string]interface{}{"name": name}) + } + c.JSON(http.StatusOK, gin.H{"message": "工作流已保存", "workflow": saved}) +} + +func (h *WorkflowHandler) Delete(c *gin.Context) { + id := strings.TrimSpace(c.Param("id")) + if id == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "工作流 id 不能为空"}) + return + } + if err := h.db.DeleteWorkflowDefinition(id); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + workflowrunner.InvalidateCompiledCache(id) + if h.audit != nil { + h.audit.RecordOK(c, "workflow", "delete", "删除工作流", "workflow", id, nil) + } + c.JSON(http.StatusOK, gin.H{"message": "工作流已删除"}) +} diff --git a/internal/handler/workflow_integration.go b/internal/handler/workflow_integration.go new file mode 100644 index 00000000..bdbdc894 --- /dev/null +++ b/internal/handler/workflow_integration.go @@ -0,0 +1,304 @@ +package handler + +import ( + "context" + "errors" + "net/http" + "strings" + "time" + + "cyberstrike-ai/internal/config" + workflowrunner "cyberstrike-ai/internal/workflow" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +func (h *AgentHandler) roleForWorkflow(req *ChatRequest) (config.RoleConfig, bool) { + if h == nil || h.config == nil || h.config.Roles == nil || req == nil { + return config.RoleConfig{}, false + } + roleName := strings.TrimSpace(req.Role) + if roleName == "" { + return config.RoleConfig{}, false + } + role, ok := h.config.Roles[roleName] + if !ok || !role.Enabled { + return config.RoleConfig{}, false + } + if role.Name == "" { + role.Name = roleName + } + if !workflowrunner.ShouldAutoRunRoleWorkflow(role) { + return config.RoleConfig{}, false + } + return role, true +} + +func (h *AgentHandler) runRoleWorkflowStreamIfBound( + c *gin.Context, + req *ChatRequest, + prep *multiAgentPrepared, + sendEvent func(eventType, message string, data interface{}), +) bool { + role, ok := h.roleForWorkflow(req) + if !ok || prep == nil { + return false + } + + conversationID := prep.ConversationID + assistantMessageID := prep.AssistantMessageID + userMessage := "" + if req != nil { + userMessage = req.Message + } + + taskStatus := "completed" + taskOwned := false + defer func() { + if taskOwned { + h.tasks.FinishTask(conversationID, taskStatus) + } + }() + + if c == nil || c.Request == nil { + return false + } + baseCtx, cancelWithCause := context.WithCancelCause(detachedAgentContext(c.Request.Context())) + defer cancelWithCause(nil) + taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute) + defer timeoutCancel() + + if _, err := h.tasks.StartTask(conversationID, userMessage, cancelWithCause); err != nil { + var errorMsg string + if errors.Is(err, ErrTaskAlreadyRunning) { + errorMsg = "⚠️ 当前会话已有任务正在执行中，请等待当前任务完成或点击「停止任务」后再尝试。" + sendEvent("error", errorMsg, map[string]interface{}{ + "conversationId": conversationID, + "errorType": "task_already_running", + }) + } else { + errorMsg = "❌ 无法启动任务: " + err.Error() + sendEvent("error", errorMsg, nil) + } + if assistantMessageID != "" { + _, _ = h.db.Exec("UPDATE messages SET content = ?, updated_at = ? WHERE id = ?", errorMsg, time.Now(), assistantMessageID) + } + sendEvent("done", "", map[string]interface{}{"conversationId": conversationID}) + return true + } + taskOwned = true + + progress := h.createProgressCallback(taskCtx, cancelWithCause, conversationID, assistantMessageID, sendEvent) + result, err := workflowrunner.RunRoleBoundWorkflow(taskCtx, workflowrunner.RunArgs{ + DB: h.db, + Logger: h.logger, + Role: role, + AppCfg: h.config, + Agent: h.agent, + ConversationID: conversationID, + ProjectID: h.conversationProjectID(conversationID), + UserMessage: prep.FinalMessage, + History: prep.History, + RoleTools: prep.RoleTools, + AgentsMarkdownDir: h.agentsMarkdownDir, + SystemPromptExtra: h.agentSessionContextBlock(conversationID), + AssistantMessageID: assistantMessageID, + Progress: progress, + }) + if err != nil { + cause := context.Cause(baseCtx) + if errors.Is(cause, ErrTaskCancelled) { + taskStatus = "cancelled" + h.tasks.UpdateTaskStatus(conversationID, taskStatus) + cancelMsg := "任务已被用户取消，后续操作已停止。" + if assistantMessageID != "" { + if err := h.appendAssistantMessageNotice(assistantMessageID, cancelMsg); err != nil { + h.logger.Warn("更新取消后的助手消息失败", zap.Error(err)) + } + _ = h.db.AddProcessDetail(assistantMessageID, conversationID, "cancelled", cancelMsg, nil) + } + sendEvent("cancelled", cancelMsg, map[string]interface{}{ + "conversationId": conversationID, + "messageId": assistantMessageID, + }) + sendEvent("done", "", map[string]interface{}{"conversationId": conversationID}) + return true + } + if errors.Is(err, context.DeadlineExceeded) || errors.Is(context.Cause(taskCtx), context.DeadlineExceeded) { + taskStatus = "timeout" + h.tasks.UpdateTaskStatus(conversationID, taskStatus) + timeoutMsg := "任务执行超时，已自动终止。" + if assistantMessageID != "" { + _, _ = h.db.Exec("UPDATE messages SET content = ?, updated_at = ? WHERE id = ?", timeoutMsg, time.Now(), assistantMessageID) + _ = h.db.AddProcessDetail(assistantMessageID, conversationID, "timeout", timeoutMsg, nil) + } + sendEvent("error", timeoutMsg, map[string]interface{}{ + "conversationId": conversationID, + "messageId": assistantMessageID, + "errorType": "timeout", + }) + sendEvent("done", "", map[string]interface{}{"conversationId": conversationID}) + return true + } + errMsg := "执行角色绑定流程失败: " + err.Error() + taskStatus = "failed" + h.tasks.UpdateTaskStatus(conversationID, taskStatus) + if assistantMessageID != "" { + _, _ = h.db.Exec("UPDATE messages SET content = ?, updated_at = ? WHERE id = ?", errMsg, time.Now(), assistantMessageID) + _ = h.db.AddProcessDetail(assistantMessageID, conversationID, "error", errMsg, nil) + } + sendEvent("error", errMsg, map[string]interface{}{"conversationId": conversationID}) + sendEvent("done", "", map[string]interface{}{"conversationId": conversationID}) + return true + } + decision := h.finalizeCandidateForDeliveryWithPolicy( + prep.ConversationID, + prep.AssistantMessageID, + "workflow", + result.Response, + nil, + result.AwaitingHITL, + "", + true, + ) + responseText := decision.FinalText + if !decision.Finalizable { + responseText = finalizationBlockedMessage(decision) + taskStatus = decision.Status + h.tasks.UpdateTaskStatus(conversationID, taskStatus) + sendEvent("finalization_check", responseText, decision) + } + payload := finalizationResponsePayload(decision, map[string]interface{}{ + "conversationId": prep.ConversationID, + "messageId": prep.AssistantMessageID, + "agentMode": "workflow", + "workflowRunId": result.RunID, + }) + if result.AwaitingHITL { + payload["workflowStatus"] = "awaiting_hitl" + payload["awaitingHitl"] = true + } else { + payload["workflowStatus"] = result.Status + payload["awaitingHitl"] = false + } + sendEvent("response", responseText, payload) + sendEvent("done", "", map[string]interface{}{"conversationId": prep.ConversationID}) + return true +} + +func (h *AgentHandler) runRoleWorkflowJSONIfBound(c *gin.Context, req *ChatRequest, prep *multiAgentPrepared) bool { + role, ok := h.roleForWorkflow(req) + if !ok || prep == nil { + return false + } + + conversationID := prep.ConversationID + assistantMessageID := prep.AssistantMessageID + userMessage := "" + if req != nil { + userMessage = req.Message + } + + taskStatus := "completed" + taskOwned := false + defer func() { + if taskOwned { + h.tasks.FinishTask(conversationID, taskStatus) + } + }() + + baseCtx, cancelWithCause := context.WithCancelCause(c.Request.Context()) + defer cancelWithCause(nil) + taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute) + defer timeoutCancel() + + if _, err := h.tasks.StartTask(conversationID, userMessage, cancelWithCause); err != nil { + if errors.Is(err, ErrTaskAlreadyRunning) { + c.JSON(http.StatusConflict, gin.H{ + "error": "⚠️ 当前会话已有任务正在执行中，请等待当前任务完成或点击「停止任务」后再尝试。", + "conversationId": conversationID, + "errorType": "task_already_running", + }) + } else { + c.JSON(http.StatusInternalServerError, gin.H{"error": "❌ 无法启动任务: " + err.Error()}) + } + return true + } + taskOwned = true + + progress := h.createProgressCallback(taskCtx, cancelWithCause, conversationID, assistantMessageID, nil) + result, err := workflowrunner.RunRoleBoundWorkflow(taskCtx, workflowrunner.RunArgs{ + DB: h.db, + Logger: h.logger, + Role: role, + AppCfg: h.config, + Agent: h.agent, + ConversationID: conversationID, + ProjectID: h.conversationProjectID(conversationID), + UserMessage: prep.FinalMessage, + History: prep.History, + RoleTools: prep.RoleTools, + AgentsMarkdownDir: h.agentsMarkdownDir, + SystemPromptExtra: h.agentSessionContextBlock(conversationID), + AssistantMessageID: assistantMessageID, + Progress: progress, + }) + if err != nil { + cause := context.Cause(baseCtx) + if errors.Is(cause, ErrTaskCancelled) { + taskStatus = "cancelled" + cancelMsg := "任务已被用户取消，后续操作已停止。" + if assistantMessageID != "" { + _ = h.appendAssistantMessageNotice(assistantMessageID, cancelMsg) + _ = h.db.AddProcessDetail(assistantMessageID, conversationID, "cancelled", cancelMsg, nil) + } + c.JSON(http.StatusOK, gin.H{ + "status": "cancelled", + "message": cancelMsg, + "conversationId": conversationID, + }) + return true + } + errMsg := "执行角色绑定流程失败: " + err.Error() + taskStatus = "failed" + if assistantMessageID != "" { + _, _ = h.db.Exec("UPDATE messages SET content = ?, updated_at = ? WHERE id = ?", errMsg, time.Now(), assistantMessageID) + } + c.JSON(http.StatusInternalServerError, gin.H{"error": errMsg, "conversationId": conversationID}) + return true + } + decision := h.finalizeCandidateForDeliveryWithPolicy( + prep.ConversationID, + prep.AssistantMessageID, + "workflow", + result.Response, + nil, + result.AwaitingHITL, + "", + true, + ) + responseText := decision.FinalText + if !decision.Finalizable { + responseText = finalizationBlockedMessage(decision) + taskStatus = decision.Status + } + c.JSON(http.StatusOK, gin.H{ + "response": responseText, + "conversationId": prep.ConversationID, + "assistantMessageId": prep.AssistantMessageID, + "agentMode": "workflow", + "workflowRunId": result.RunID, + "workflowStatus": result.Status, + "awaitingHitl": result.AwaitingHITL, + "finalized": decision.Finalized, + "finalizable": decision.Finalizable, + "status": decision.Status, + "completionReason": decision.CompletionReason, + "evidenceVerified": decision.EvidenceVerified, + "evidenceRefs": decision.EvidenceRefs, + "pendingExecutionIds": decision.PendingExecutionIDs, + "missingChecks": decision.MissingChecks, + }) + return true +} diff --git a/internal/handler/workflow_package.go b/internal/handler/workflow_package.go new file mode 100644 index 00000000..ca2a5c5a --- /dev/null +++ b/internal/handler/workflow_package.go @@ -0,0 +1,319 @@ +package handler + +import ( + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "io" + "mime" + "net/http" + "strings" + "time" + + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + workflowrunner "cyberstrike-ai/internal/workflow" + workflowpkg "cyberstrike-ai/internal/workflow/package" + + "github.com/gin-gonic/gin" + "github.com/google/uuid" +) + +type workflowPackageResolution struct { + Action string `json:"action"` + NewWorkflowID string `json:"new_workflow_id"` +} +type workflowPackageImportRequest struct { + InspectionID string `json:"inspection_id"` + Resolution workflowPackageResolution `json:"resolution"` + ConfirmOverwrite bool `json:"confirm_overwrite"` +} + +func (h *WorkflowHandler) ExportPackage(c *gin.Context) { + wf, err := h.db.GetWorkflowDefinition(c.Param("id")) + if err != nil { + writeWorkflowPackageError(c, http.StatusInternalServerError, "WFPKG_EXPORT_FAILED", "导出工作流包失败", nil) + return + } + if wf == nil { + writeWorkflowPackageError(c, http.StatusNotFound, "WFPKG_WORKFLOW_NOT_FOUND", "工作流不存在", nil) + return + } + pkg, meta, err := workflowpkg.Export(workflowPackageDocument(wf)) + if err != nil { + writeWorkflowPackageError(c, http.StatusInternalServerError, "WFPKG_EXPORT_FAILED", "导出工作流包失败", nil) + return + } + c.Header("Content-Type", "application/zip") + c.Header("Content-Disposition", mime.FormatMediaType("attachment", map[string]string{"filename": meta.FileName})) + c.Header("ETag", `"`+meta.PackageHash+`"`) + c.Header("X-Workflow-Package-SHA256", meta.PackageHash) + if h.audit != nil { + h.audit.RecordOK(c, "workflow_package", "export", "导出工作流包", "workflow", wf.ID, map[string]interface{}{"package_hash": meta.PackageHash}) + } + c.Data(http.StatusOK, "application/zip", pkg) +} + +func (h *WorkflowHandler) CreatePackageInspection(c *gin.Context) { + session, ok := security.CurrentSession(c) + if !ok || strings.TrimSpace(session.UserID) == "" { + writeWorkflowPackageError(c, http.StatusUnauthorized, "WFPKG_INSPECTION_NOT_FOUND", "未授权访问", nil) + return + } + c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, workflowpkg.MaxArchiveBytes+1) + file, _, err := c.Request.FormFile("file") + if err != nil { + var maxErr *http.MaxBytesError + if errors.As(err, &maxErr) { + writeWorkflowPackageError(c, http.StatusUnprocessableEntity, "WFPKG_FILE_TOO_LARGE", "工作流包文件超过大小限制", nil) + return + } + writeWorkflowPackageError(c, http.StatusUnprocessableEntity, "WFPKG_FILE_REQUIRED", "必须上传工作流包文件", nil) + return + } + defer file.Close() + archive, err := io.ReadAll(io.LimitReader(file, workflowpkg.MaxArchiveBytes+1)) + if err != nil || len(archive) > workflowpkg.MaxArchiveBytes { + writeWorkflowPackageError(c, http.StatusUnprocessableEntity, "WFPKG_FILE_TOO_LARGE", "工作流包文件超过大小限制", nil) + return + } + inspected, err := workflowpkg.InspectArchive(c.Request.Context(), archive, workflowrunner.ValidateGraphJSON) + if err != nil { + code := workflowpkg.ErrorCode(err) + if code == "" { + code = "WFPKG_INVALID_ARCHIVE" + } + if h.audit != nil { + h.audit.RecordFail(c, "workflow_package", "inspect", "工作流包预检失败", map[string]interface{}{"code": code, "package_hash": workflowPackageHash(archive)}) + } + writeWorkflowPackageError(c, http.StatusUnprocessableEntity, code, "工作流包预检失败", nil) + return + } + state, local, err := h.workflowPackageConflict(inspected.Document.ID, inspected.ContentHash) + if err != nil { + writeWorkflowPackageError(c, http.StatusInternalServerError, "WFPKG_IMPORT_FAILED", "读取本地工作流失败", nil) + return + } + summary := workflowPackageInspectionSummary{ID: "wpi_" + strings.ReplaceAll(uuid.NewString(), "-", ""), Status: "ready", ExpiresAt: time.Now().UTC().Add(30 * time.Minute), Package: workflowPackagePackageSummary{PackageFormat: inspected.Manifest.PackageFormat, FormatVersion: inspected.Manifest.FormatVersion, PackageID: inspected.Manifest.PackageID, PackageHash: inspected.PackageHash}, Workflow: workflowPackageWorkflowSummary{SourceID: inspected.Document.ID, Name: inspected.Document.Name, Description: inspected.Document.Description, SourceRevision: inspected.Document.Version, Enabled: inspected.Document.Enabled, ContentHash: inspected.ContentHash, GraphHash: inspected.GraphHash, NodeCount: inspected.NodeCount, EdgeCount: inspected.EdgeCount}, Conflict: workflowPackageConflictSummary{State: state}, Warnings: []string{}} + if local != nil { + content, graph, _, _ := workflowpkg.DocumentHashes(workflowPackageDocument(local)) + summary.Conflict.LocalWorkflow = &workflowPackageLocalWorkflow{ID: local.ID, Version: local.Version, ContentHash: content, GraphHash: graph} + } + manifestJSON, _ := json.Marshal(inspected.Manifest) + payloadJSON, err := json.Marshal(inspected.Document) + if err != nil { + writeWorkflowPackageError(c, http.StatusInternalServerError, "WFPKG_IMPORT_FAILED", "保存预检失败", nil) + return + } + inspectionJSON, _ := json.Marshal(summary) + record := &database.WorkflowPackageInspection{ID: summary.ID, PackageHash: inspected.PackageHash, ManifestJSON: string(manifestJSON), WorkflowPayloadJSON: string(payloadJSON), InspectionJSON: string(inspectionJSON), SourceWorkflowID: inspected.Document.ID, SourceRevision: inspected.Document.Version, SourceContentHash: inspected.ContentHash, SourceGraphHash: inspected.GraphHash, LocalConflictState: state, CreatedBy: session.UserID, CreatedAt: time.Now().UTC(), ExpiresAt: summary.ExpiresAt} + if local != nil { + content, graph, _, _ := workflowpkg.DocumentHashes(workflowPackageDocument(local)) + record.LocalWorkflowID = local.ID + record.LocalContentHash = content + record.LocalGraphHash = graph + } + if err := h.db.CreateWorkflowPackageInspection(record); err != nil { + writeWorkflowPackageError(c, http.StatusInternalServerError, "WFPKG_IMPORT_FAILED", "保存预检失败", nil) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "workflow_package", "inspect", "工作流包预检成功", "inspection", record.ID, map[string]interface{}{"package_hash": record.PackageHash, "workflow_id": record.SourceWorkflowID}) + } + c.JSON(http.StatusCreated, gin.H{"inspection": summary}) +} + +func (h *WorkflowHandler) GetPackageInspection(c *gin.Context) { + session, ok := security.CurrentSession(c) + if !ok { + writeWorkflowPackageError(c, http.StatusUnauthorized, "WFPKG_INSPECTION_NOT_FOUND", "未授权访问", nil) + return + } + v, err := h.db.GetWorkflowPackageInspection(c.Param("inspectionId"), session.UserID) + if err != nil { + writeWorkflowPackageError(c, http.StatusInternalServerError, "WFPKG_IMPORT_FAILED", "读取预检失败", nil) + return + } + if v == nil { + writeWorkflowPackageError(c, http.StatusNotFound, "WFPKG_INSPECTION_NOT_FOUND", "预检不存在", nil) + return + } + if v.Status == "expired" { + writeWorkflowPackageError(c, http.StatusConflict, "WFPKG_INSPECTION_EXPIRED", "预检已过期", nil) + return + } + var summary any + _ = json.Unmarshal([]byte(v.InspectionJSON), &summary) + c.JSON(http.StatusOK, gin.H{"inspection": summary}) +} + +func (h *WorkflowHandler) ApplyPackageImport(c *gin.Context) { + session, ok := security.CurrentSession(c) + if !ok { + writeWorkflowPackageError(c, http.StatusUnauthorized, "WFPKG_INSPECTION_NOT_FOUND", "未授权访问", nil) + return + } + key := strings.TrimSpace(c.GetHeader("Idempotency-Key")) + if _, err := uuid.Parse(key); err != nil { + writeWorkflowPackageError(c, http.StatusBadRequest, "WFPKG_IDEMPOTENCY_KEY_REQUIRED", "必须提供 UUID 幂等键", nil) + return + } + var req workflowPackageImportRequest + if err := c.ShouldBindJSON(&req); err != nil { + writeWorkflowPackageError(c, http.StatusUnprocessableEntity, "WFPKG_INVALID_ACTION", "导入请求无效", nil) + return + } + req.InspectionID = strings.TrimSpace(req.InspectionID) + req.Resolution.Action = strings.TrimSpace(req.Resolution.Action) + req.Resolution.NewWorkflowID = strings.TrimSpace(req.Resolution.NewWorkflowID) + if req.Resolution.Action != "rename" && req.Resolution.NewWorkflowID != "" { + writeWorkflowPackageError(c, http.StatusUnprocessableEntity, "WFPKG_INVALID_ACTION", "当前导入动作不接受新工作流 ID", nil) + return + } + requestHash := workflowPackageRequestHash(req) + imp, replayed, err := h.db.ApplyWorkflowPackageImport(c.Request.Context(), database.WorkflowPackageApplyRequest{InspectionID: req.InspectionID, RequestHash: requestHash, IdempotencyKey: key, ActorUserID: session.UserID, Action: req.Resolution.Action, NewWorkflowID: req.Resolution.NewWorkflowID, ConfirmOverwrite: req.ConfirmOverwrite}) + if err != nil { + h.writeWorkflowPackageImportError(c, req.InspectionID, err) + return + } + wf, _ := h.db.GetWorkflowDefinition(imp.ResultingWorkflowID) + if !replayed && (imp.Result == "created" || imp.Result == "overwritten" || imp.Result == "renamed") { + workflowrunner.InvalidateCompiledCache(imp.ResultingWorkflowID) + } + response := h.workflowPackageImportResponse(imp, wf) + if !replayed && h.audit != nil { + h.audit.RecordOK(c, "workflow_package", "import", "工作流包导入成功", "workflow", imp.ResultingWorkflowID, map[string]interface{}{"inspection_id": imp.InspectionID, "action": imp.Action, "result": imp.Result}) + } + status := http.StatusCreated + if replayed { + status = http.StatusOK + } + c.JSON(status, gin.H{"import": response}) +} + +func (h *WorkflowHandler) GetPackageImport(c *gin.Context) { + session, ok := security.CurrentSession(c) + if !ok { + writeWorkflowPackageError(c, http.StatusUnauthorized, "WFPKG_INSPECTION_NOT_FOUND", "未授权访问", nil) + return + } + imp, err := h.db.GetWorkflowPackageImport(c.Param("importId"), session.UserID) + if err != nil { + writeWorkflowPackageError(c, http.StatusInternalServerError, "WFPKG_IMPORT_FAILED", "读取导入结果失败", nil) + return + } + if imp == nil { + writeWorkflowPackageError(c, http.StatusNotFound, "WFPKG_INSPECTION_NOT_FOUND", "导入结果不存在", nil) + return + } + wf, _ := h.db.GetWorkflowDefinition(imp.ResultingWorkflowID) + c.JSON(http.StatusOK, gin.H{"import": h.workflowPackageImportResponse(imp, wf)}) +} + +func (h *WorkflowHandler) workflowPackageConflict(id, sourceHash string) (string, *database.WorkflowDefinition, error) { + local, err := h.db.GetWorkflowDefinition(id) + if err != nil { + return "", nil, err + } + if local == nil { + return "none", nil, nil + } + localHash, _, _, err := workflowpkg.DocumentHashes(workflowPackageDocument(local)) + if err != nil { + return "", nil, err + } + if localHash == sourceHash { + return "identical", local, nil + } + return "id_conflict", local, nil +} +func workflowPackageDocument(w *database.WorkflowDefinition) workflowpkg.Document { + return workflowpkg.Document{ID: w.ID, Name: w.Name, Description: w.Description, Version: w.Version, GraphJSON: w.GraphJSON, Enabled: w.Enabled, UpdatedAt: w.UpdatedAt} +} +func workflowPackageRequestHash(req workflowPackageImportRequest) string { + value := struct { + ConfirmOverwrite bool `json:"confirm_overwrite"` + InspectionID string `json:"inspection_id"` + Resolution workflowPackageResolution `json:"resolution"` + }{req.ConfirmOverwrite, req.InspectionID, req.Resolution} + b, _ := json.Marshal(value) + sum := sha256.Sum256(b) + return "sha256:" + hex.EncodeToString(sum[:]) +} +func workflowPackageHash(b []byte) string { + sum := sha256.Sum256(b) + return "sha256:" + hex.EncodeToString(sum[:]) +} + +func (h *WorkflowHandler) writeWorkflowPackageImportError(c *gin.Context, inspectionID string, err error) { + var e *database.WorkflowPackageStoreError + if errors.As(err, &e) { + status := http.StatusConflict + if e.Code == "WFPKG_INVALID_ACTION" || e.Code == "WFPKG_INVALID_RENAME_ID" { + status = http.StatusUnprocessableEntity + } + if h.audit != nil { + h.audit.RecordFail(c, "workflow_package", "import", "工作流包导入失败", map[string]interface{}{"code": e.Code, "inspection_id": inspectionID}) + } + writeWorkflowPackageError(c, status, e.Code, e.Message, nil) + return + } + if h.audit != nil { + h.audit.RecordFail(c, "workflow_package", "import", "工作流包导入失败", map[string]interface{}{"code": "WFPKG_IMPORT_FAILED", "inspection_id": inspectionID}) + } + writeWorkflowPackageError(c, http.StatusInternalServerError, "WFPKG_IMPORT_FAILED", "导入工作流包失败", nil) +} +func writeWorkflowPackageError(c *gin.Context, status int, code, message string, details map[string]any) { + body := gin.H{"code": code, "message": message} + if len(details) > 0 { + body["details"] = details + } + c.JSON(status, gin.H{"error": body}) +} + +type workflowPackageInspectionSummary struct { + ID string `json:"id"` + Status string `json:"status"` + ExpiresAt time.Time `json:"expires_at"` + Package workflowPackagePackageSummary `json:"package"` + Workflow workflowPackageWorkflowSummary `json:"workflow"` + Conflict workflowPackageConflictSummary `json:"conflict"` + Warnings []string `json:"warnings"` +} +type workflowPackagePackageSummary struct { + PackageFormat string `json:"package_format"` + FormatVersion string `json:"format_version"` + PackageID string `json:"package_id"` + PackageHash string `json:"package_hash"` +} +type workflowPackageWorkflowSummary struct { + SourceID string `json:"source_id"` + Name string `json:"name"` + Description string `json:"description"` + SourceRevision int `json:"source_revision"` + Enabled bool `json:"enabled"` + ContentHash string `json:"content_hash"` + GraphHash string `json:"graph_hash"` + NodeCount int `json:"node_count"` + EdgeCount int `json:"edge_count"` +} +type workflowPackageConflictSummary struct { + State string `json:"state"` + LocalWorkflow *workflowPackageLocalWorkflow `json:"local_workflow,omitempty"` +} +type workflowPackageLocalWorkflow struct { + ID string `json:"id"` + Version int `json:"version"` + ContentHash string `json:"content_hash"` + GraphHash string `json:"graph_hash"` +} + +func (h *WorkflowHandler) workflowPackageImportResponse(imp *database.WorkflowPackageImport, wf *database.WorkflowDefinition) gin.H { + out := gin.H{"id": imp.ID, "inspection_id": imp.InspectionID, "status": "succeeded", "result": imp.Result, "action": imp.Action, "source_workflow_id": imp.SourceWorkflowID, "target_workflow_id": imp.TargetWorkflowID, "applied_at": imp.AppliedAt} + if wf != nil { + content, graph, _, _ := workflowpkg.DocumentHashes(workflowPackageDocument(wf)) + out["workflow"] = gin.H{"id": wf.ID, "version": wf.Version, "content_hash": content, "graph_hash": graph} + } + return out +} diff --git a/internal/handler/workflow_package_test.go b/internal/handler/workflow_package_test.go new file mode 100644 index 00000000..2fce24d6 --- /dev/null +++ b/internal/handler/workflow_package_test.go @@ -0,0 +1,104 @@ +package handler + +import ( + "bytes" + "encoding/json" + "mime/multipart" + "net/http" + "net/http/httptest" + "path/filepath" + "strings" + "testing" + "time" + + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + workflowrunner "cyberstrike-ai/internal/workflow" + workflowpkg "cyberstrike-ai/internal/workflow/package" + + "github.com/gin-gonic/gin" + "github.com/google/uuid" + "go.uber.org/zap" +) + +func TestWorkflowPackageHandlerInspectionAndCreateImport(t *testing.T) { + gin.SetMode(gin.TestMode) + workflowrunner.SetCheckpointDir(filepath.Join(t.TempDir(), "workflow-checkpoints")) + db, err := database.NewDB(filepath.Join(t.TempDir(), "workflow-package-handler.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = db.Close() }) + h := NewWorkflowHandler(db, zap.NewNop()) + pkg, _, err := workflowpkg.Export(workflowpkg.Document{ID: "wf-api", Name: "API workflow", Version: 4, Enabled: true, UpdatedAt: time.Now().UTC(), GraphJSON: `{"nodes":[{"id":"start-1","type":"start","label":"开始","position":{"x":0,"y":0},"config":{}},{"id":"out-1","type":"output","label":"输出","position":{"x":0,"y":120},"config":{"output_key":"result","source_binding":{"from":"inputs","field":"message"}}}],"edges":[{"id":"e1","source":"start-1","target":"out-1"}],"config":{"schema_version":1}}`}) + if err != nil { + t.Fatal(err) + } + var body bytes.Buffer + writer := multipart.NewWriter(&body) + part, err := writer.CreateFormFile("file", "wf-api.csapkg.zip") + if err != nil { + t.Fatal(err) + } + if _, err := part.Write(pkg); err != nil { + t.Fatal(err) + } + if err := writer.Close(); err != nil { + t.Fatal(err) + } + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodPost, "/api/workflow-package-inspections", &body) + c.Request.Header.Set("Content-Type", writer.FormDataContentType()) + c.Set(security.ContextSessionKey, security.Session{UserID: "user-1"}) + h.CreatePackageInspection(c) + if w.Code != http.StatusCreated { + t.Fatalf("inspection status=%d body=%s", w.Code, w.Body.String()) + } + var inspected struct { + Inspection struct { + ID string `json:"id"` + } `json:"inspection"` + } + if err := json.Unmarshal(w.Body.Bytes(), &inspected); err != nil { + t.Fatal(err) + } + applyBody := bytes.NewBufferString(`{"inspection_id":"` + inspected.Inspection.ID + `","resolution":{"action":"create","new_workflow_id":""},"confirm_overwrite":false}`) + w = httptest.NewRecorder() + c, _ = gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodPost, "/api/workflow-package-imports", applyBody) + c.Request.Header.Set("Content-Type", "application/json") + c.Request.Header.Set("Idempotency-Key", uuid.NewString()) + c.Set(security.ContextSessionKey, security.Session{UserID: "user-1"}) + h.ApplyPackageImport(c) + if w.Code != http.StatusCreated { + t.Fatalf("import status=%d body=%s", w.Code, w.Body.String()) + } + saved, _ := db.GetWorkflowDefinition("wf-api") + if saved == nil || saved.Version != 1 { + t.Fatalf("saved=%#v", saved) + } +} + +func TestWorkflowHandlerGenerateDraft(t *testing.T) { + gin.SetMode(gin.TestMode) + h := NewWorkflowHandler(nil, zap.NewNop()) + body := bytes.NewBufferString(`{"prompt":"对目标资产做端口扫描，如果发现高危端口就执行加固脚本，最后输出报告","options":{"include_objective":true},"available_tools":[{"key":"nmap","name":"nmap","enabled":true}]}`) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodPost, "/api/workflows/generate-draft", body) + c.Request.Header.Set("Content-Type", "application/json") + h.GenerateDraft(c) + if w.Code != http.StatusBadGateway { + t.Fatalf("status=%d body=%s", w.Code, w.Body.String()) + } + var resp struct { + Error string `json:"error"` + } + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatal(err) + } + if !strings.Contains(resp.Error, "大模型生成失败") { + t.Fatalf("unexpected error: %#v", resp.Error) + } +} diff --git a/internal/handler/workflow_run.go b/internal/handler/workflow_run.go new file mode 100644 index 00000000..1b18f05a --- /dev/null +++ b/internal/handler/workflow_run.go @@ -0,0 +1,200 @@ +package handler + +import ( + "encoding/json" + "net/http" + "strings" + "time" + + "cyberstrike-ai/internal/agent" + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + workflowrunner "cyberstrike-ai/internal/workflow" + + "github.com/gin-gonic/gin" +) + +func (h *WorkflowHandler) SetRuntime(agent *agent.Agent, cfg *config.Config) { + h.agent = agent + h.cfg = cfg +} + +func (h *WorkflowHandler) GetRun(c *gin.Context) { + runID := strings.TrimSpace(c.Param("runId")) + if !h.workflowRunAllowed(c, runID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + run, err := h.db.GetWorkflowRun(runID) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if run == nil { + c.JSON(http.StatusNotFound, gin.H{"error": "工作流运行不存在"}) + return + } + nodeRuns, err := h.db.ListWorkflowNodeRuns(runID) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"run": run, "nodeRuns": nodeRuns}) +} + +func (h *WorkflowHandler) ReplayRun(c *gin.Context) { + runID := strings.TrimSpace(c.Param("runId")) + if !h.workflowRunAllowed(c, runID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + nodeRuns, err := h.db.ListWorkflowNodeRuns(runID) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + steps := make([]gin.H, 0, len(nodeRuns)) + for i, nodeRun := range nodeRuns { + var input any + var output any + _ = json.Unmarshal([]byte(nodeRun.InputJSON), &input) + _ = json.Unmarshal([]byte(nodeRun.OutputJSON), &output) + steps = append(steps, gin.H{ + "step": i + 1, + "nodeRunId": nodeRun.ID, + "nodeId": nodeRun.NodeID, + "status": nodeRun.Status, + "input": input, + "output": output, + "error": nodeRun.Error, + "startedAt": nodeRun.StartedAt, + "finishedAt": nodeRun.FinishedAt, + }) + } + c.JSON(http.StatusOK, gin.H{"workflowRunId": runID, "steps": steps}) +} + +func (h *WorkflowHandler) ListPendingRuns(c *gin.Context) { + conversationID := strings.TrimSpace(c.Query("conversationId")) + if conversationID != "" && !h.workflowConversationAllowed(c, conversationID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + runs, err := h.db.ListWorkflowRunsAwaitingHITLFiltered(conversationID, 50) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + runs = filterSlice(runs, func(run *database.WorkflowRun) bool { + return run != nil && h.workflowConversationAllowed(c, run.ConversationID) + }) + c.JSON(http.StatusOK, gin.H{"runs": runs}) +} + +type workflowResumeRequest struct { + Approved bool `json:"approved"` + Comment string `json:"comment,omitempty"` +} + +func (h *WorkflowHandler) ResumeRun(c *gin.Context) { + if h.agent == nil || h.cfg == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "工作流运行时未初始化"}) + return + } + runID := strings.TrimSpace(c.Param("runId")) + if !h.workflowRunAllowed(c, runID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + var req workflowResumeRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()}) + return + } + run, err := h.db.GetWorkflowRun(runID) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if run == nil { + c.JSON(http.StatusNotFound, gin.H{"error": "工作流运行不存在"}) + return + } + role := config.RoleConfig{Name: strings.TrimSpace(run.RoleID)} + if role.Name != "" && h.cfg.Roles != nil { + if r, ok := h.cfg.Roles[role.Name]; ok { + role = r + if role.Name == "" { + role.Name = run.RoleID + } + } + } + if run.Status != "awaiting_hitl" { + c.JSON(http.StatusBadRequest, gin.H{"error": "工作流运行不在等待审批状态: " + run.Status}) + return + } + if err := h.db.RecordWorkflowRunHITLDecision(runID, req.Approved, req.Comment); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + decision := workflowrunner.HITLDecision{ + Approved: req.Approved, + Comment: strings.TrimSpace(req.Comment), + } + delegated := workflowrunner.NotifyHITLDecision(runID, decision) + if !delegated { + for i := 0; i < 10; i++ { + time.Sleep(50 * time.Millisecond) + if workflowrunner.NotifyHITLDecision(runID, decision) { + delegated = true + break + } + } + } + if delegated { + c.JSON(http.StatusOK, gin.H{ + "workflowRunId": runID, + "status": "delegated", + "streamResuming": true, + "approved": req.Approved, + }) + return + } + result, err := workflowrunner.ResumeWorkflowRun(c.Request.Context(), workflowrunner.RunArgs{ + DB: h.db, + Logger: h.logger, + Role: role, + AppCfg: h.cfg, + Agent: h.agent, + ConversationID: run.ConversationID, + ProjectID: run.ProjectID, + }, runID, req.Approved, req.Comment) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{ + "response": result.Response, + "workflowRunId": result.RunID, + "status": result.Status, + "awaitingHitl": result.AwaitingHITL, + }) +} + +func (h *WorkflowHandler) workflowConversationAllowed(c *gin.Context, conversationID string) bool { + session, ok := security.CurrentSession(c) + return ok && h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", strings.TrimSpace(conversationID)) +} + +func (h *WorkflowHandler) workflowRunAllowed(c *gin.Context, runID string) bool { + session, ok := security.CurrentSession(c) + if !ok { + return false + } + if session.Scope == database.RBACScopeAll { + return true + } + run, err := h.db.GetWorkflowRun(strings.TrimSpace(runID)) + return err == nil && run != nil && h.workflowConversationAllowed(c, run.ConversationID) +}