Add files via upload

2026-05-31 19:41:42 +02:00 · 2026-03-14 00:49:25 +08:00
parent 226f9b79e2
commit 5e8fef0ad4
11 changed files with 1147 additions and 25 deletions
@@ -121,10 +121,11 @@ type ChatAttachment struct {

 // ChatRequest 聊天请求
 type ChatRequest struct {
-	Message        string            `json:"message" binding:"required"`
-	ConversationID string            `json:"conversationId,omitempty"`
-	Role           string            `json:"role,omitempty"` // 角色名称
-	Attachments    []ChatAttachment  `json:"attachments,omitempty"`
+	Message              string           `json:"message" binding:"required"`
+	ConversationID       string           `json:"conversationId,omitempty"`
+	Role                 string           `json:"role,omitempty"`                 // 角色名称
+	Attachments          []ChatAttachment  `json:"attachments,omitempty"`
+	WebShellConnectionID string           `json:"webshellConnectionId,omitempty"` // WebShell 管理 - AI 助手：当前选中的连接 ID，仅使用 webshell_* 工具
 }

 const (
@@ -316,7 +317,24 @@ func (h *AgentHandler) AgentLoop(c *gin.Context) {
 	finalMessage := req.Message
 	var roleTools []string // 角色配置的工具列表
 	var roleSkills []string // 角色配置的skills列表（用于提示AI，但不硬编码内容）
-	if req.Role != "" && req.Role != "默认" {
+
+	// WebShell AI 助手模式：绑定当前连接，仅开放 webshell_* 工具并注入 connection_id
+	if req.WebShellConnectionID != "" {
+		conn, err := h.db.GetWebshellConnection(strings.TrimSpace(req.WebShellConnectionID))
+		if err != nil || conn == nil {
+			h.logger.Warn("WebShell AI 助手：未找到连接", zap.String("id", req.WebShellConnectionID), zap.Error(err))
+			c.JSON(http.StatusBadRequest, gin.H{"error": "未找到该 WebShell 连接"})
+			return
+		}
+		remark := conn.Remark
+		if remark == "" {
+			remark = conn.URL
+		}
+		finalMessage = fmt.Sprintf("[WebShell 助手上下文] 当前连接 ID：%s，备注：%s。可用工具（仅在该连接上操作时使用，connection_id 填 \"%s\"）：webshell_exec、webshell_file_list、webshell_file_read、webshell_file_write。请根据用户输入决定下一步：若仅为问候、闲聊或简单问题，直接简短回复即可，不必调用工具；当用户明确需要执行命令、列目录、读写字件等操作时再调用上述工具。\n\n用户请求：%s",
+			conn.ID, remark, conn.ID, req.Message)
+		roleTools = []string{builtin.ToolWebshellExec, builtin.ToolWebshellFileList, builtin.ToolWebshellFileRead, builtin.ToolWebshellFileWrite}
+		roleSkills = nil
+	} else if req.Role != "" && req.Role != "默认" {
 		if h.config.Roles != nil {
 			if role, exists := h.config.Roles[req.Role]; exists && role.Enabled {
 				// 应用用户提示词
@@ -712,11 +730,17 @@ func (h *AgentHandler) AgentLoopStream(c *gin.Context) {
 		}
 	}

-	// 如果没有对话ID，创建新对话
+	// 如果没有对话ID，创建新对话（WebShell 助手模式下关联连接 ID 以便持久化展示）
 	conversationID := req.ConversationID
 	if conversationID == "" {
 		title := safeTruncateString(req.Message, 50)
-		conv, err := h.db.CreateConversation(title)
+		var conv *database.Conversation
+		var err error
+		if req.WebShellConnectionID != "" {
+			conv, err = h.db.CreateConversationWithWebshell(strings.TrimSpace(req.WebShellConnectionID), title)
+		} else {
+			conv, err = h.db.CreateConversation(title)
+		}
 		if err != nil {
 			h.logger.Error("创建对话失败", zap.Error(err))
 			sendEvent("error", "创建对话失败: "+err.Error(), nil)
@@ -769,7 +793,22 @@ func (h *AgentHandler) AgentLoopStream(c *gin.Context) {
 	// 应用角色用户提示词和工具配置
 	finalMessage := req.Message
 	var roleTools []string // 角色配置的工具列表
-	if req.Role != "" && req.Role != "默认" {
+	var roleSkills []string
+	if req.WebShellConnectionID != "" {
+		conn, errConn := h.db.GetWebshellConnection(strings.TrimSpace(req.WebShellConnectionID))
+		if errConn != nil || conn == nil {
+			h.logger.Warn("WebShell AI 助手：未找到连接", zap.String("id", req.WebShellConnectionID), zap.Error(errConn))
+			sendEvent("error", "未找到该 WebShell 连接", nil)
+			return
+		}
+		remark := conn.Remark
+		if remark == "" {
+			remark = conn.URL
+		}
+		finalMessage = fmt.Sprintf("[WebShell 助手上下文] 当前连接 ID：%s，备注：%s。可用工具（仅在该连接上操作时使用，connection_id 填 \"%s\"）：webshell_exec、webshell_file_list、webshell_file_read、webshell_file_write。请根据用户输入决定下一步：若仅为问候、闲聊或简单问题，直接简短回复即可，不必调用工具；当用户明确需要执行命令、列目录、读写字件等操作时再调用上述工具。\n\n用户请求：%s",
+			conn.ID, remark, conn.ID, req.Message)
+		roleTools = []string{builtin.ToolWebshellExec, builtin.ToolWebshellFileList, builtin.ToolWebshellFileRead, builtin.ToolWebshellFileWrite}
+	} else if req.Role != "" && req.Role != "默认" {
 		if h.config.Roles != nil {
 			if role, exists := h.config.Roles[req.Role]; exists && role.Enabled {
 				// 应用用户提示词
@@ -788,6 +827,7 @@ func (h *AgentHandler) AgentLoopStream(c *gin.Context) {
 				}
 				// 注意：角色配置的skills不再硬编码注入，AI可以通过list_skills和read_skill工具按需调用
 				if len(role.Skills) > 0 {
+					roleSkills = role.Skills
 					h.logger.Info("角色配置了skills，AI可通过工具按需调用", zap.String("role", req.Role), zap.Int("skillCount", len(role.Skills)), zap.Strings("skills", role.Skills))
 				}
 			}
@@ -886,17 +926,7 @@ func (h *AgentHandler) AgentLoopStream(c *gin.Context) {

 	// 执行Agent Loop，传入独立的上下文，确保任务不会因客户端断开而中断（使用包含角色提示词的finalMessage和角色工具列表）
 	sendEvent("progress", "正在分析您的请求...", nil)
-	// 注意：skills不会硬编码注入，但会在系统提示词中提示AI这个角色推荐使用哪些skills
-	var roleSkills []string // 角色配置的skills列表（用于提示AI，但不硬编码内容）
-	if req.Role != "" && req.Role != "默认" {
-		if h.config.Roles != nil {
-			if role, exists := h.config.Roles[req.Role]; exists && role.Enabled {
-				if len(role.Skills) > 0 {
-					roleSkills = role.Skills
-				}
-			}
-		}
-	}
+	// 注意：roleSkills 已在上方根据 req.Role 或 WebShell 模式设置
 	result, err := h.agent.AgentLoopWithProgress(taskCtx, finalMessage, agentHistoryMessages, conversationID, progressCallback, roleTools, roleSkills)
 	if err != nil {
 		h.logger.Error("Agent Loop执行失败", zap.Error(err))
@@ -28,6 +28,9 @@ type KnowledgeToolRegistrar func() error
 // VulnerabilityToolRegistrar 漏洞工具注册器接口
 type VulnerabilityToolRegistrar func() error

+// WebshellToolRegistrar WebShell 工具注册器接口（ApplyConfig 时重新注册）
+type WebshellToolRegistrar func() error
+
 // SkillsToolRegistrar Skills工具注册器接口
 type SkillsToolRegistrar func() error

@@ -60,6 +63,7 @@ type ConfigHandler struct {
 	externalMCPMgr             *mcp.ExternalMCPManager    // 外部MCP管理器
 	knowledgeToolRegistrar     KnowledgeToolRegistrar     // 知识库工具注册器（可选）
 	vulnerabilityToolRegistrar VulnerabilityToolRegistrar // 漏洞工具注册器（可选）
+	webshellToolRegistrar      WebshellToolRegistrar      // WebShell 工具注册器（可选）
 	skillsToolRegistrar        SkillsToolRegistrar        // Skills工具注册器（可选）
 	retrieverUpdater           RetrieverUpdater           // 检索器更新器（可选）
 	knowledgeInitializer       KnowledgeInitializer       // 知识库初始化器（可选）
@@ -120,6 +124,13 @@ func (h *ConfigHandler) SetVulnerabilityToolRegistrar(registrar VulnerabilityToo
 	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()
@@ -792,6 +803,16 @@ func (h *ConfigHandler) ApplyConfig(c *gin.Context) {
 		}
 	}

+	// 重新注册 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工具")
@@ -197,6 +197,53 @@ func (h *WebShellHandler) DeleteConnection(c *gin.Context) {
 	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"`
@@ -472,3 +519,108 @@ func (h *WebShellHandler) escapeForEcho(s string) string {
 	// 仅用于 write：base64 写入更安全，这里简单用单引号包裹
 	return "'" + strings.ReplaceAll(s, "'", "'\"'\"'") + "'"
 }
+
+// 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, _ := io.ReadAll(resp.Body)
+	return string(out), 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))
+	shellType := strings.ToLower(strings.TrimSpace(conn.Type))
+	if shellType == "" {
+		shellType = "php"
+	}
+	var command string
+	switch action {
+	case "list":
+		if path == "" {
+			path = "."
+		}
+		if shellType == "asp" || shellType == "aspx" {
+			command = "dir " + h.escapePath(strings.TrimSpace(path))
+		} else {
+			command = "ls -la " + h.escapePath(strings.TrimSpace(path))
+		}
+	case "read":
+		path = strings.TrimSpace(path)
+		if path == "" {
+			return "", false, "path is required for read"
+		}
+		if shellType == "asp" || shellType == "aspx" {
+			command = "type " + h.escapePath(path)
+		} else {
+			command = "cat " + h.escapePath(path)
+		}
+	case "write":
+		path = strings.TrimSpace(path)
+		if path == "" {
+			return "", false, "path is required for write"
+		}
+		command = "echo " + h.escapeForEcho(content) + " > " + h.escapePath(path)
+	default:
+		return "", false, "unsupported action: " + action + " (supported: list, read, write)"
+	}
+	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, _ := io.ReadAll(resp.Body)
+	return string(out), resp.StatusCode == http.StatusOK, ""
+}