From 5e8fef0ad4af524bab9838580f64f4eb81b2f496 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=85=AC=E6=98=8E?= <83812544+Ed1s0nZ@users.noreply.github.com> Date: Sat, 14 Mar 2026 00:49:25 +0800 Subject: [PATCH] Add files via upload --- internal/app/app.go | 162 +++++++++++++ internal/database/conversation.go | 101 +++++++- internal/database/database.go | 15 ++ internal/handler/agent.go | 68 ++++-- internal/handler/config.go | 21 ++ internal/handler/webshell.go | 152 ++++++++++++ internal/mcp/builtin/constants.go | 16 +- web/static/css/style.css | 241 +++++++++++++++++++ web/static/i18n/en-US.json | 8 + web/static/i18n/zh-CN.json | 8 + web/static/js/webshell.js | 380 +++++++++++++++++++++++++++++- 11 files changed, 1147 insertions(+), 25 deletions(-) diff --git a/internal/app/app.go b/internal/app/app.go index 13e3dad9..65e8e715 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -320,6 +320,7 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) { attackChainHandler := handler.NewAttackChainHandler(db, &cfg.OpenAI, log.Logger) vulnerabilityHandler := handler.NewVulnerabilityHandler(db, log.Logger) webshellHandler := handler.NewWebShellHandler(log.Logger, db) + registerWebshellTools(mcpServer, db, webshellHandler, log.Logger) configHandler := handler.NewConfigHandler(configPath, cfg, mcpServer, executor, agent, attackChainHandler, externalMCPMgr, log.Logger) externalMCPHandler := handler.NewExternalMCPHandler(externalMCPMgr, cfg, configPath, log.Logger) roleHandler := handler.NewRoleHandler(cfg, configPath, log.Logger) @@ -365,6 +366,13 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) { } configHandler.SetVulnerabilityToolRegistrar(vulnerabilityRegistrar) + // 设置 WebShell 工具注册器(ApplyConfig 时重新注册) + webshellRegistrar := func() error { + registerWebshellTools(mcpServer, db, webshellHandler, log.Logger) + return nil + } + configHandler.SetWebshellToolRegistrar(webshellRegistrar) + // 设置Skills工具注册器(内置工具,必须设置) skillsRegistrar := func() error { // 创建一个适配器,将database.DB适配为SkillStatsStorage接口 @@ -823,6 +831,8 @@ func setupRoutes( // WebShell 管理(代理执行 + 连接配置存 SQLite) protected.GET("/webshell/connections", webshellHandler.ListConnections) protected.POST("/webshell/connections", webshellHandler.CreateConnection) + protected.GET("/webshell/connections/:id/ai-history", webshellHandler.GetAIHistory) + protected.GET("/webshell/connections/:id/ai-conversations", webshellHandler.ListAIConversations) protected.PUT("/webshell/connections/:id", webshellHandler.UpdateConnection) protected.DELETE("/webshell/connections/:id", webshellHandler.DeleteConnection) protected.POST("/webshell/exec", webshellHandler.Exec) @@ -1067,6 +1077,158 @@ func registerVulnerabilityTool(mcpServer *mcp.Server, db *database.DB, logger *z logger.Info("漏洞记录工具注册成功") } +// registerWebshellTools 注册 WebShell 相关 MCP 工具,供 AI 助手在指定连接上执行命令与文件操作 +func registerWebshellTools(mcpServer *mcp.Server, db *database.DB, webshellHandler *handler.WebShellHandler, logger *zap.Logger) { + if db == nil || webshellHandler == nil { + logger.Warn("跳过 WebShell 工具注册:db 或 webshellHandler 为空") + return + } + + // webshell_exec + execTool := mcp.Tool{ + Name: builtin.ToolWebshellExec, + Description: "在指定的 WebShell 连接上执行一条系统命令,返回命令的标准输出。connection_id 由用户在 AI 助手上下文中选定。", + ShortDescription: "在 WebShell 连接上执行命令", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "connection_id": map[string]interface{}{ + "type": "string", + "description": "WebShell 连接 ID(如 ws_xxx)", + }, + "command": map[string]interface{}{ + "type": "string", + "description": "要执行的系统命令", + }, + }, + "required": []string{"connection_id", "command"}, + }, + } + execHandler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + cid, _ := args["connection_id"].(string) + cmd, _ := args["command"].(string) + if cid == "" || cmd == "" { + return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "connection_id 和 command 均为必填"}}, IsError: true}, nil + } + conn, err := db.GetWebshellConnection(cid) + if err != nil || conn == nil { + return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "未找到该 WebShell 连接或查询失败"}}, IsError: true}, nil + } + output, ok, errMsg := webshellHandler.ExecWithConnection(conn, cmd) + if errMsg != "" { + return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: errMsg}}, IsError: true}, nil + } + if !ok { + return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "HTTP 非 200,输出:\n" + output}}, IsError: false}, nil + } + return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: output}}, IsError: false}, nil + } + mcpServer.RegisterTool(execTool, execHandler) + + // webshell_file_list + listTool := mcp.Tool{ + Name: builtin.ToolWebshellFileList, + Description: "在指定 WebShell 连接上列出目录内容。path 默认为当前目录(.)。", + ShortDescription: "在 WebShell 上列出目录", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "connection_id": map[string]interface{}{"type": "string", "description": "WebShell 连接 ID"}, + "path": map[string]interface{}{"type": "string", "description": "目录路径,默认 ."}, + }, + "required": []string{"connection_id"}, + }, + } + listHandler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + cid, _ := args["connection_id"].(string) + path, _ := args["path"].(string) + if cid == "" { + return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "connection_id 必填"}}, IsError: true}, nil + } + conn, err := db.GetWebshellConnection(cid) + if err != nil || conn == nil { + return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "未找到该 WebShell 连接"}}, IsError: true}, nil + } + output, ok, errMsg := webshellHandler.FileOpWithConnection(conn, "list", path, "", "") + if errMsg != "" { + return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: errMsg}}, IsError: true}, nil + } + return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: output}}, IsError: !ok}, nil + } + mcpServer.RegisterTool(listTool, listHandler) + + // webshell_file_read + readTool := mcp.Tool{ + Name: builtin.ToolWebshellFileRead, + Description: "在指定 WebShell 连接上读取文件内容。", + ShortDescription: "在 WebShell 上读取文件", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "connection_id": map[string]interface{}{"type": "string", "description": "WebShell 连接 ID"}, + "path": map[string]interface{}{"type": "string", "description": "文件路径"}, + }, + "required": []string{"connection_id", "path"}, + }, + } + readHandler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + cid, _ := args["connection_id"].(string) + path, _ := args["path"].(string) + if cid == "" || path == "" { + return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "connection_id 和 path 必填"}}, IsError: true}, nil + } + conn, err := db.GetWebshellConnection(cid) + if err != nil || conn == nil { + return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "未找到该 WebShell 连接"}}, IsError: true}, nil + } + output, ok, errMsg := webshellHandler.FileOpWithConnection(conn, "read", path, "", "") + if errMsg != "" { + return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: errMsg}}, IsError: true}, nil + } + return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: output}}, IsError: !ok}, nil + } + mcpServer.RegisterTool(readTool, readHandler) + + // webshell_file_write + writeTool := mcp.Tool{ + Name: builtin.ToolWebshellFileWrite, + Description: "在指定 WebShell 连接上写入文件内容(会覆盖已有文件)。", + ShortDescription: "在 WebShell 上写入文件", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "connection_id": map[string]interface{}{"type": "string", "description": "WebShell 连接 ID"}, + "path": map[string]interface{}{"type": "string", "description": "文件路径"}, + "content": map[string]interface{}{"type": "string", "description": "要写入的内容"}, + }, + "required": []string{"connection_id", "path", "content"}, + }, + } + writeHandler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + cid, _ := args["connection_id"].(string) + path, _ := args["path"].(string) + content, _ := args["content"].(string) + if cid == "" || path == "" { + return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "connection_id 和 path 必填"}}, IsError: true}, nil + } + conn, err := db.GetWebshellConnection(cid) + if err != nil || conn == nil { + return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "未找到该 WebShell 连接"}}, IsError: true}, nil + } + output, ok, errMsg := webshellHandler.FileOpWithConnection(conn, "write", path, content, "") + if errMsg != "" { + return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: errMsg}}, IsError: true}, nil + } + if !ok { + return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "写入可能失败,输出:\n" + output}}, IsError: false}, nil + } + return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "写入成功\n" + output}}, IsError: false}, nil + } + mcpServer.RegisterTool(writeTool, writeHandler) + + logger.Info("WebShell 工具注册成功") +} + // initializeKnowledge 初始化知识库组件(用于动态初始化) func initializeKnowledge( cfg *config.Config, diff --git a/internal/database/conversation.go b/internal/database/conversation.go index f6a9e0fe..432f870d 100644 --- a/internal/database/conversation.go +++ b/internal/database/conversation.go @@ -33,13 +33,26 @@ type Message struct { // CreateConversation 创建新对话 func (db *DB) CreateConversation(title string) (*Conversation, error) { + return db.CreateConversationWithWebshell("", title) +} + +// CreateConversationWithWebshell 创建新对话,可选绑定 WebShell 连接 ID(为空则普通对话) +func (db *DB) CreateConversationWithWebshell(webshellConnectionID, title string) (*Conversation, error) { id := uuid.New().String() now := time.Now() - _, err := db.Exec( - "INSERT INTO conversations (id, title, created_at, updated_at) VALUES (?, ?, ?, ?)", - id, title, now, now, - ) + var err error + if webshellConnectionID != "" { + _, err = db.Exec( + "INSERT INTO conversations (id, title, created_at, updated_at, webshell_connection_id) VALUES (?, ?, ?, ?, ?)", + id, title, now, now, webshellConnectionID, + ) + } else { + _, err = db.Exec( + "INSERT INTO conversations (id, title, created_at, updated_at) VALUES (?, ?, ?, ?)", + id, title, now, now, + ) + } if err != nil { return nil, fmt.Errorf("创建对话失败: %w", err) } @@ -52,6 +65,86 @@ func (db *DB) CreateConversation(title string) (*Conversation, error) { }, nil } +// GetConversationByWebshellConnectionID 根据 WebShell 连接 ID 获取该连接下最近一条对话(用于 AI 助手持久化) +func (db *DB) GetConversationByWebshellConnectionID(connectionID string) (*Conversation, error) { + if connectionID == "" { + return nil, fmt.Errorf("connectionID is empty") + } + var conv Conversation + var createdAt, updatedAt string + var pinned int + err := db.QueryRow( + "SELECT id, title, pinned, created_at, updated_at FROM conversations WHERE webshell_connection_id = ? ORDER BY updated_at DESC LIMIT 1", + connectionID, + ).Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt) + if err != nil { + if err == sql.ErrNoRows { + return nil, nil + } + return nil, fmt.Errorf("查询对话失败: %w", err) + } + conv.Pinned = pinned != 0 + if t, e := time.Parse("2006-01-02 15:04:05.999999999-07:00", createdAt); e == nil { + conv.CreatedAt = t + } else if t, e := time.Parse("2006-01-02 15:04:05", createdAt); e == nil { + conv.CreatedAt = t + } else { + conv.CreatedAt, _ = time.Parse(time.RFC3339, createdAt) + } + if t, e := time.Parse("2006-01-02 15:04:05.999999999-07:00", updatedAt); e == nil { + conv.UpdatedAt = t + } else if t, e := time.Parse("2006-01-02 15:04:05", updatedAt); e == nil { + conv.UpdatedAt = t + } else { + conv.UpdatedAt, _ = time.Parse(time.RFC3339, updatedAt) + } + messages, err := db.GetMessages(conv.ID) + if err != nil { + return nil, fmt.Errorf("加载消息失败: %w", err) + } + conv.Messages = messages + return &conv, nil +} + +// WebShellConversationItem 用于侧边栏列表,不含消息 +type WebShellConversationItem struct { + ID string `json:"id"` + Title string `json:"title"` + UpdatedAt time.Time `json:"updatedAt"` +} + +// ListConversationsByWebshellConnectionID 列出该 WebShell 连接下的所有对话(按更新时间倒序),供侧边栏展示 +func (db *DB) ListConversationsByWebshellConnectionID(connectionID string) ([]WebShellConversationItem, error) { + if connectionID == "" { + return nil, nil + } + rows, err := db.Query( + "SELECT id, title, updated_at FROM conversations WHERE webshell_connection_id = ? ORDER BY updated_at DESC", + connectionID, + ) + if err != nil { + return nil, fmt.Errorf("查询对话列表失败: %w", err) + } + defer rows.Close() + var list []WebShellConversationItem + for rows.Next() { + var item WebShellConversationItem + var updatedAt string + if err := rows.Scan(&item.ID, &item.Title, &updatedAt); err != nil { + continue + } + if t, e := time.Parse("2006-01-02 15:04:05.999999999-07:00", updatedAt); e == nil { + item.UpdatedAt = t + } else if t, e := time.Parse("2006-01-02 15:04:05", updatedAt); e == nil { + item.UpdatedAt = t + } else { + item.UpdatedAt, _ = time.Parse(time.RFC3339, updatedAt) + } + list = append(list, item) + } + return list, rows.Err() +} + // GetConversation 获取对话 func (db *DB) GetConversation(id string) (*Conversation, error) { var conv Conversation diff --git a/internal/database/database.go b/internal/database/database.go index c228edc8..1c48f5e0 100644 --- a/internal/database/database.go +++ b/internal/database/database.go @@ -415,6 +415,21 @@ func (db *DB) migrateConversationsTable() error { } } + // 检查 webshell_connection_id 字段是否存在(WebShell AI 助手对话关联) + err = db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('conversations') WHERE name='webshell_connection_id'").Scan(&count) + if err != nil { + if _, addErr := db.Exec("ALTER TABLE conversations ADD COLUMN webshell_connection_id TEXT"); addErr != nil { + errMsg := strings.ToLower(addErr.Error()) + if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") { + db.logger.Warn("添加webshell_connection_id字段失败", zap.Error(addErr)) + } + } + } else if count == 0 { + if _, err := db.Exec("ALTER TABLE conversations ADD COLUMN webshell_connection_id TEXT"); err != nil { + db.logger.Warn("添加webshell_connection_id字段失败", zap.Error(err)) + } + } + return nil } diff --git a/internal/handler/agent.go b/internal/handler/agent.go index 99fa0803..f41b9803 100644 --- a/internal/handler/agent.go +++ b/internal/handler/agent.go @@ -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)) diff --git a/internal/handler/config.go b/internal/handler/config.go index 80f2c9e2..47487b6e 100644 --- a/internal/handler/config.go +++ b/internal/handler/config.go @@ -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工具") diff --git a/internal/handler/webshell.go b/internal/handler/webshell.go index 34356c68..b32bcb05 100644 --- a/internal/handler/webshell.go +++ b/internal/handler/webshell.go @@ -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, "" +} diff --git a/internal/mcp/builtin/constants.go b/internal/mcp/builtin/constants.go index feec7dc3..13c01a6b 100644 --- a/internal/mcp/builtin/constants.go +++ b/internal/mcp/builtin/constants.go @@ -13,6 +13,12 @@ const ( // Skills工具 ToolListSkills = "list_skills" ToolReadSkill = "read_skill" + + // WebShell 助手工具(AI 在 WebShell 管理 - AI 助手 中使用) + ToolWebshellExec = "webshell_exec" + ToolWebshellFileList = "webshell_file_list" + ToolWebshellFileRead = "webshell_file_read" + ToolWebshellFileWrite = "webshell_file_write" ) // IsBuiltinTool 检查工具名称是否是内置工具 @@ -22,7 +28,11 @@ func IsBuiltinTool(toolName string) bool { ToolListKnowledgeRiskTypes, ToolSearchKnowledgeBase, ToolListSkills, - ToolReadSkill: + ToolReadSkill, + ToolWebshellExec, + ToolWebshellFileList, + ToolWebshellFileRead, + ToolWebshellFileWrite: return true default: return false @@ -37,5 +47,9 @@ func GetAllBuiltinTools() []string { ToolSearchKnowledgeBase, ToolListSkills, ToolReadSkill, + ToolWebshellExec, + ToolWebshellFileList, + ToolWebshellFileRead, + ToolWebshellFileWrite, } } diff --git a/web/static/css/style.css b/web/static/css/style.css index aa2b88cd..f71c7e0b 100644 --- a/web/static/css/style.css +++ b/web/static/css/style.css @@ -9092,6 +9092,247 @@ header { min-width: 72px; } +/* WebShell AI 助手 Tab:仅在此 pane 激活时显示;左右布局 = 左侧栏 + 右侧主区(与对话页一致) */ +#webshell-pane-ai { + flex-direction: row; + flex: 1; + min-height: 0; + overflow: hidden; +} +/* 不设 display,避免覆盖 .webshell-pane 的 display:none(否则终端/文件管理页会露出 AI 面板)。左右布局由 #webshell-pane-ai 的 flex-direction:row 提供 */ +.webshell-pane-ai-with-sidebar { + flex: 1; + min-height: 0; + overflow: hidden; + width: 100%; +} +.webshell-ai-sidebar { + flex-shrink: 0; + width: 280px; + min-width: 200px; + display: flex; + flex-direction: column; + border-right: 1px solid var(--border-color); + background: var(--bg-secondary); +} +.webshell-ai-new-btn { + margin: 10px 12px; + width: calc(100% - 24px); +} +.webshell-ai-conv-list { + flex: 1; + min-height: 0; + overflow-y: auto; + padding: 4px 0; +} +.webshell-ai-conv-item { + display: flex; + align-items: center; + gap: 6px; + padding: 8px 12px; + cursor: pointer; + font-size: 0.9rem; + border-left: 3px solid transparent; +} +.webshell-ai-conv-item:hover { + background: var(--border-color); +} +.webshell-ai-conv-item.active { + background: var(--accent-light, rgba(59, 130, 246, 0.1)); + border-left-color: var(--accent-color); +} +.webshell-ai-conv-item-title { + flex: 1; + min-width: 0; + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; +} +.webshell-ai-conv-item-date { + flex-shrink: 0; + font-size: 0.75rem; + color: var(--text-secondary); +} +.webshell-ai-conv-del { + flex-shrink: 0; + padding: 2px 6px; + opacity: 0.6; +} +.webshell-ai-conv-item:hover .webshell-ai-conv-del { + opacity: 1; +} +.webshell-ai-main { + flex: 1; + min-width: 0; + display: flex; + flex-direction: column; + min-height: 0; +} +.webshell-ai-hint { + flex-shrink: 0; + padding: 10px 14px; + font-size: 0.9rem; + color: var(--text-secondary); + border-bottom: 1px solid var(--border-color); +} +.webshell-ai-toolbar { + flex-shrink: 0; + padding: 6px 14px 8px; + display: flex; + gap: 8px; + border-bottom: 1px solid var(--border-color); +} +.webshell-ai-progress { + font-size: 0.85rem; + color: var(--text-secondary); + padding: 6px 12px; + margin: 0 0 4px 0; + border-radius: 8px; + background: var(--bg-secondary); + border: 1px solid var(--border-color); + align-self: flex-start; +} +.webshell-ai-timeline { + display: none; + flex-direction: column; + gap: 6px; + width: 100%; + margin-bottom: 8px; + padding: 8px 12px; + border-radius: 8px; + background: var(--bg-secondary); + border: 1px solid var(--border-color); +} +.webshell-ai-timeline.has-items { + display: flex; +} +.webshell-ai-timeline-item { + font-size: 0.85rem; + color: var(--text-primary); + padding: 4px 0; + border-bottom: 1px solid var(--border-color); +} +.webshell-ai-timeline-item:last-child { + border-bottom: none; +} +.webshell-ai-timeline-title { + display: block; + font-weight: 500; + color: var(--text-secondary); +} +.webshell-ai-timeline-msg { + margin-top: 4px; + padding-left: 0; + font-size: 0.8rem; + color: var(--text-secondary); + white-space: pre-wrap; + word-break: break-word; + max-height: 120px; + overflow-y: auto; +} +.webshell-ai-old-conv { + width: 100%; + margin-bottom: 8px; + border: 1px solid var(--border-color); + border-radius: 10px; + background: var(--bg-secondary); + overflow: hidden; +} +.webshell-ai-old-conv-label { + display: block; + width: 100%; + padding: 8px 12px; + text-align: left; + font-size: 0.9rem; + color: var(--text-secondary); + background: transparent; + border: none; + cursor: pointer; +} +.webshell-ai-old-conv-label:hover { + background: var(--border-color); + color: var(--text-primary); +} +.webshell-ai-old-conv-body { + padding: 0 12px 12px; + display: flex; + flex-direction: column; + gap: 12px; +} +.webshell-ai-messages { + flex: 1; + min-height: 0; + overflow-y: auto; + padding: 12px 14px; + display: flex; + flex-direction: column; + gap: 12px; +} +.webshell-ai-msg { + padding: 10px 14px; + border-radius: 10px; + max-width: 90%; + white-space: pre-wrap; + word-break: break-word; +} +.webshell-ai-msg.user { + align-self: flex-end; + background: var(--accent-color); + color: #fff; +} +.webshell-ai-msg.assistant { + align-self: flex-start; + background: var(--bg-secondary); + border: 1px solid var(--border-color); +} +.webshell-ai-input-row { + flex-shrink: 0; + display: flex; + gap: 10px; + padding: 8px 14px; + border-top: 1px solid var(--border-color); + align-items: center; +} +.webshell-ai-input { + flex: 1; + height: 36px; + min-height: 36px; + resize: none; + padding: 8px 12px; + border-radius: 10px; + border: 1px solid var(--border-color); + font-size: 0.95rem; + /* 更柔和的滚动条样式 */ + scrollbar-width: thin; + scrollbar-color: rgba(15, 23, 42, 0.25) transparent; +} +.webshell-ai-input:focus { + outline: none; + border-color: var(--accent-color); +} +.webshell-ai-input::-webkit-scrollbar { + width: 6px; +} +.webshell-ai-input::-webkit-scrollbar-track { + background: transparent; +} +.webshell-ai-input::-webkit-scrollbar-thumb { + background: rgba(15, 23, 42, 0.25); + border-radius: 999px; +} +.webshell-ai-input::-webkit-scrollbar-thumb:hover { + background: rgba(15, 23, 42, 0.4); +} +.webshell-ai-input-row .btn-primary { + flex-shrink: 0; + height: 36px; + min-width: 72px; + padding: 0 16px; + display: flex; + align-items: center; + justify-content: center; +} + /* 仪表盘页面样式(最佳实践布局 + 视觉增强) */ .dashboard-page { height: 100%; diff --git a/web/static/i18n/en-US.json b/web/static/i18n/en-US.json index 3802d33e..68dda4ac 100644 --- a/web/static/i18n/en-US.json +++ b/web/static/i18n/en-US.json @@ -355,6 +355,14 @@ "editConnectionTitle": "Edit connection", "tabTerminal": "Virtual terminal", "tabFileManager": "File manager", + "tabAiAssistant": "AI Assistant", + "aiSystemReadyMessage": "System is ready. Please enter your test requirements, and the system will automatically perform the corresponding security tests.", + "aiNewConversation": "New conversation", + "aiPreviousConversation": "Previous conversation", + "aiDeleteConversation": "Delete conversation", + "aiDeleteConversationConfirm": "Delete this conversation?", + "aiPlaceholder": "e.g. List files in the current directory", + "aiSend": "Send", "quickCommands": "Quick commands", "downloadFile": "Download", "terminalWelcome": "WebShell virtual terminal — type a command and press Enter (Ctrl+L clear)", diff --git a/web/static/i18n/zh-CN.json b/web/static/i18n/zh-CN.json index e266c3b7..33af68e9 100644 --- a/web/static/i18n/zh-CN.json +++ b/web/static/i18n/zh-CN.json @@ -355,6 +355,14 @@ "editConnectionTitle": "编辑连接", "tabTerminal": "虚拟终端", "tabFileManager": "文件管理", + "tabAiAssistant": "AI 助手", + "aiSystemReadyMessage": "系统已就绪。请输入您的测试需求,系统将自动执行相应的安全测试。", + "aiNewConversation": "新对话", + "aiPreviousConversation": "之前的对话", + "aiDeleteConversation": "删除对话", + "aiDeleteConversationConfirm": "确定删除当前对话记录?", + "aiPlaceholder": "例如:列出当前目录下的文件", + "aiSend": "发送", "quickCommands": "快捷命令", "downloadFile": "下载", "terminalWelcome": "WebShell 虚拟终端 — 输入命令后按回车执行(Ctrl+L 清屏)", diff --git a/web/static/js/webshell.js b/web/static/js/webshell.js index bc961404..6f8b81ef 100644 --- a/web/static/js/webshell.js +++ b/web/static/js/webshell.js @@ -20,6 +20,11 @@ let webshellHistoryIndex = -1; const WEBSHELL_HISTORY_MAX = 100; // 清屏防重入:一次点击只执行一次(避免多次绑定或重复触发导致多个 shell>) let webshellClearInProgress = false; +// AI 助手:按连接 ID 保存对话 ID,便于多轮对话 +let webshellAiConvMap = {}; +let webshellAiSending = false; +// 流式打字机效果:当前会话的 response 序号,用于中止过期的打字 +let webshellStreamingTypingId = 0; // 从服务端(SQLite)拉取连接列表 function getWebshellConnections() { @@ -61,6 +66,10 @@ function wsT(key) { 'webshell.editConnectionTitle': '编辑连接', 'webshell.tabTerminal': '虚拟终端', 'webshell.tabFileManager': '文件管理', + 'webshell.tabAiAssistant': 'AI 助手', + 'webshell.aiSystemReadyMessage': '系统已就绪。请输入您的测试需求,系统将自动执行相应的安全测试。', + 'webshell.aiPlaceholder': '例如:列出当前目录下的文件', + 'webshell.aiSend': '发送', 'webshell.terminalWelcome': 'WebShell 虚拟终端 — 输入命令后按回车执行(Ctrl+L 清屏)', 'webshell.quickCommands': '快捷命令', 'webshell.downloadFile': '下载', @@ -268,6 +277,95 @@ function escapeHtml(s) { return div.innerHTML; } +function formatWebshellAiConvDate(updatedAt) { + if (!updatedAt) return ''; + var d = typeof updatedAt === 'string' ? new Date(updatedAt) : updatedAt; + if (isNaN(d.getTime())) return ''; + var now = new Date(); + var sameDay = d.getDate() === now.getDate() && d.getMonth() === now.getMonth() && d.getFullYear() === now.getFullYear(); + if (sameDay) return d.getHours() + ':' + String(d.getMinutes()).padStart(2, '0'); + return (d.getMonth() + 1) + '/' + d.getDate(); +} + +function fetchAndRenderWebshellAiConvList(conn, listEl) { + if (!conn || !conn.id || !listEl || typeof apiFetch !== 'function') return Promise.resolve(); + return apiFetch('/api/webshell/connections/' + encodeURIComponent(conn.id) + '/ai-conversations', { method: 'GET' }) + .then(function (r) { return r.json(); }) + .then(function (list) { + if (!Array.isArray(list)) list = []; + listEl.innerHTML = ''; + list.forEach(function (item) { + var row = document.createElement('div'); + row.className = 'webshell-ai-conv-item'; + row.dataset.convId = item.id; + var title = (item.title || '').trim() || item.id.slice(0, 8); + var dateStr = item.updatedAt ? formatWebshellAiConvDate(item.updatedAt) : ''; + row.innerHTML = '' + escapeHtml(title) + '' + escapeHtml(dateStr) + ''; + if (webshellAiConvMap[conn.id] === item.id) row.classList.add('active'); + row.addEventListener('click', function () { + webshellAiConvListSelect(conn, item.id, document.getElementById('webshell-ai-messages'), listEl); + }); + var delBtn = document.createElement('button'); + delBtn.type = 'button'; + delBtn.className = 'btn-ghost btn-sm webshell-ai-conv-del'; + delBtn.textContent = '×'; + delBtn.title = wsT('webshell.aiDeleteConversation') || '删除对话'; + delBtn.addEventListener('click', function (e) { + e.stopPropagation(); + if (!confirm(wsT('webshell.aiDeleteConversationConfirm') || '确定删除该对话?')) return; + apiFetch('/api/conversations/' + encodeURIComponent(item.id), { method: 'DELETE' }) + .then(function (r) { + if (r.ok) { + if (webshellAiConvMap[conn.id] === item.id) { + delete webshellAiConvMap[conn.id]; + var msgs = document.getElementById('webshell-ai-messages'); + if (msgs) msgs.innerHTML = ''; + } + fetchAndRenderWebshellAiConvList(conn, listEl); + } + }) + .catch(function (e) { console.warn('删除对话失败', e); }); + }); + row.appendChild(delBtn); + listEl.appendChild(row); + }); + }) + .catch(function (e) { console.warn('加载对话列表失败', e); }); +} + +function webshellAiConvListSelect(conn, convId, messagesContainer, listEl) { + if (!conn || !convId || !messagesContainer) return; + webshellAiConvMap[conn.id] = convId; + if (listEl) listEl.querySelectorAll('.webshell-ai-conv-item').forEach(function (el) { + el.classList.toggle('active', el.dataset.convId === convId); + }); + if (typeof apiFetch !== 'function') return; + apiFetch('/api/conversations/' + encodeURIComponent(convId), { method: 'GET' }) + .then(function (r) { return r.json(); }) + .then(function (data) { + messagesContainer.innerHTML = ''; + var list = data.messages || []; + list.forEach(function (msg) { + var role = (msg.role || '').toLowerCase(); + var content = (msg.content || '').trim(); + if (!content && role !== 'assistant') return; + var div = document.createElement('div'); + div.className = 'webshell-ai-msg ' + (role === 'user' ? 'user' : 'assistant'); + div.textContent = content; + messagesContainer.appendChild(div); + }); + if (list.length === 0) { + var readyMsg = wsT('webshell.aiSystemReadyMessage') || '系统已就绪。请输入您的测试需求,系统将自动执行相应的安全测试。'; + var readyDiv = document.createElement('div'); + readyDiv.className = 'webshell-ai-msg assistant'; + readyDiv.textContent = readyMsg; + messagesContainer.appendChild(readyDiv); + } + messagesContainer.scrollTop = messagesContainer.scrollHeight; + }) + .catch(function (e) { console.warn('加载对话失败', e); }); +} + // 选择连接:渲染终端 + 文件管理 Tab,并初始化终端 function selectWebshell(id) { currentWebshellId = id; @@ -287,6 +385,7 @@ function selectWebshell(id) { '