From 1c394acd4a17a0b755a6839da52fbda5531cb2ba Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=85=AC=E6=98=8E?=
 <83812544+Ed1s0nZ@users.noreply.github.com>
Date: Wed, 25 Mar 2026 20:49:40 +0800
Subject: [PATCH] Add files via upload

---
 internal/app/app.go               | 365 ++++++++++++++++++++++++++++++
 internal/mcp/builtin/constants.go |  19 +-
 2 files changed, 383 insertions(+), 1 deletion(-)

diff --git a/internal/app/app.go b/internal/app/app.go
index bd58a0a5..fe0cf595 100644
--- a/internal/app/app.go
+++ b/internal/app/app.go
@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"os"
 	"path/filepath"
+	"strings"
 	"sync"
 	"time"
 
@@ -25,6 +26,7 @@ import (
 	"cyberstrike-ai/internal/storage"
 
 	"github.com/gin-gonic/gin"
+	"github.com/google/uuid"
 	"go.uber.org/zap"
 )
 
@@ -336,6 +338,7 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
 	webshellHandler := handler.NewWebShellHandler(log.Logger, db)
 	chatUploadsHandler := handler.NewChatUploadsHandler(log.Logger)
 	registerWebshellTools(mcpServer, db, webshellHandler, log.Logger)
+	registerWebshellManagementTools(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)
@@ -384,6 +387,7 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
 	// 设置 WebShell 工具注册器（ApplyConfig 时重新注册）
 	webshellRegistrar := func() error {
 		registerWebshellTools(mcpServer, db, webshellHandler, log.Logger)
+		registerWebshellManagementTools(mcpServer, db, webshellHandler, log.Logger)
 		return nil
 	}
 	configHandler.SetWebshellToolRegistrar(webshellRegistrar)
@@ -1270,6 +1274,367 @@ func registerWebshellTools(mcpServer *mcp.Server, db *database.DB, webshellHandl
 	logger.Info("WebShell 工具注册成功")
 }
 
+// registerWebshellManagementTools 注册 WebShell 连接管理 MCP 工具
+func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, webshellHandler *handler.WebShellHandler, logger *zap.Logger) {
+	if db == nil {
+		logger.Warn("跳过 WebShell 管理工具注册：db 为空")
+		return
+	}
+
+	// manage_webshell_list - 列出所有 webshell 连接
+	listTool := mcp.Tool{
+		Name:             builtin.ToolManageWebshellList,
+		Description:      "列出所有已保存的 WebShell 连接，返回连接ID、URL、类型、备注等信息。",
+		ShortDescription: "列出所有 WebShell 连接",
+		InputSchema: map[string]interface{}{
+			"type":       "object",
+			"properties": map[string]interface{}{},
+		},
+	}
+	listHandler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
+		connections, err := db.ListWebshellConnections()
+		if err != nil {
+			return &mcp.ToolResult{
+				Content: []mcp.Content{{Type: "text", Text: "获取连接列表失败: " + err.Error()}},
+				IsError: true,
+			}, nil
+		}
+		if len(connections) == 0 {
+			return &mcp.ToolResult{
+				Content: []mcp.Content{{Type: "text", Text: "暂无 WebShell 连接"}},
+				IsError: false,
+			}, nil
+		}
+		var sb strings.Builder
+		sb.WriteString(fmt.Sprintf("找到 %d 个 WebShell 连接：\n\n", len(connections)))
+		for _, conn := range connections {
+			sb.WriteString(fmt.Sprintf("ID: %s\n", conn.ID))
+			sb.WriteString(fmt.Sprintf("  URL: %s\n", conn.URL))
+			sb.WriteString(fmt.Sprintf("  类型: %s\n", conn.Type))
+			sb.WriteString(fmt.Sprintf("  请求方式: %s\n", conn.Method))
+			sb.WriteString(fmt.Sprintf("  命令参数: %s\n", conn.CmdParam))
+			if conn.Remark != "" {
+				sb.WriteString(fmt.Sprintf("  备注: %s\n", conn.Remark))
+			}
+			sb.WriteString(fmt.Sprintf("  创建时间: %s\n", conn.CreatedAt.Format("2006-01-02 15:04:05")))
+			sb.WriteString("\n")
+		}
+		return &mcp.ToolResult{
+			Content: []mcp.Content{{Type: "text", Text: sb.String()}},
+			IsError: false,
+		}, nil
+	}
+	mcpServer.RegisterTool(listTool, listHandler)
+
+	// manage_webshell_add - 添加新的 webshell 连接
+	addTool := mcp.Tool{
+		Name:        builtin.ToolManageWebshellAdd,
+		Description: "添加新的 WebShell 连接到管理系统。支持 PHP、ASP、ASPX、JSP 等类型的一句话木马。",
+		ShortDescription: "添加 WebShell 连接",
+		InputSchema: map[string]interface{}{
+			"type": "object",
+			"properties": map[string]interface{}{
+				"url": map[string]interface{}{
+					"type":        "string",
+					"description": "Shell 地址，如 http://target.com/shell.php（必填）",
+				},
+				"password": map[string]interface{}{
+					"type":        "string",
+					"description": "连接密码/密钥，如冰蝎/蚁剑的连接密码",
+				},
+				"type": map[string]interface{}{
+					"type":        "string",
+					"description": "Shell 类型：php、asp、aspx、jsp，默认为 php",
+					"enum":        []string{"php", "asp", "aspx", "jsp"},
+				},
+				"method": map[string]interface{}{
+					"type":        "string",
+					"description": "请求方式：GET 或 POST，默认为 POST",
+					"enum":        []string{"GET", "POST"},
+				},
+				"cmd_param": map[string]interface{}{
+					"type":        "string",
+					"description": "命令参数名，不填默认为 cmd",
+				},
+				"remark": map[string]interface{}{
+					"type":        "string",
+					"description": "备注，便于识别的备注名",
+				},
+			},
+			"required": []string{"url"},
+		},
+	}
+	addHandler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
+		urlStr, _ := args["url"].(string)
+		if urlStr == "" {
+			return &mcp.ToolResult{
+				Content: []mcp.Content{{Type: "text", Text: "错误: url 参数必填"}},
+				IsError: true,
+			}, nil
+		}
+
+		password, _ := args["password"].(string)
+		shellType, _ := args["type"].(string)
+		if shellType == "" {
+			shellType = "php"
+		}
+		method, _ := args["method"].(string)
+		if method == "" {
+			method = "post"
+		}
+		cmdParam, _ := args["cmd_param"].(string)
+		if cmdParam == "" {
+			cmdParam = "cmd"
+		}
+		remark, _ := args["remark"].(string)
+
+		// 生成连接ID
+		connID := "ws_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:12]
+		conn := &database.WebShellConnection{
+			ID:        connID,
+			URL:       urlStr,
+			Password:  password,
+			Type:      strings.ToLower(shellType),
+			Method:    strings.ToLower(method),
+			CmdParam:  cmdParam,
+			Remark:    remark,
+			CreatedAt: time.Now(),
+		}
+
+		if err := db.CreateWebshellConnection(conn); err != nil {
+			return &mcp.ToolResult{
+				Content: []mcp.Content{{Type: "text", Text: "添加 WebShell 连接失败: " + err.Error()}},
+				IsError: true,
+			}, nil
+		}
+
+		return &mcp.ToolResult{
+			Content: []mcp.Content{{
+				Type: "text",
+				Text: fmt.Sprintf("WebShell 连接添加成功！\n\n连接ID: %s\nURL: %s\n类型: %s\n请求方式: %s\n命令参数: %s", conn.ID, conn.URL, conn.Type, conn.Method, conn.CmdParam),
+			}},
+			IsError: false,
+		}, nil
+	}
+	mcpServer.RegisterTool(addTool, addHandler)
+
+	// manage_webshell_update - 更新 webshell 连接
+	updateTool := mcp.Tool{
+		Name:        builtin.ToolManageWebshellUpdate,
+		Description: "更新已存在的 WebShell 连接信息。",
+		ShortDescription: "更新 WebShell 连接",
+		InputSchema: map[string]interface{}{
+			"type": "object",
+			"properties": map[string]interface{}{
+				"connection_id": map[string]interface{}{
+					"type":        "string",
+					"description": "要更新的 WebShell 连接 ID（必填）",
+				},
+				"url": map[string]interface{}{
+					"type":        "string",
+					"description": "新的 Shell 地址",
+				},
+				"password": map[string]interface{}{
+					"type":        "string",
+					"description": "新的连接密码/密钥",
+				},
+				"type": map[string]interface{}{
+					"type":        "string",
+					"description": "新的 Shell 类型：php、asp、aspx、jsp",
+					"enum":        []string{"php", "asp", "aspx", "jsp"},
+				},
+				"method": map[string]interface{}{
+					"type":        "string",
+					"description": "新的请求方式：GET 或 POST",
+					"enum":        []string{"GET", "POST"},
+				},
+				"cmd_param": map[string]interface{}{
+					"type":        "string",
+					"description": "新的命令参数名",
+				},
+				"remark": map[string]interface{}{
+					"type":        "string",
+					"description": "新的备注",
+				},
+			},
+			"required": []string{"connection_id"},
+		},
+	}
+	updateHandler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
+		connID, _ := args["connection_id"].(string)
+		if connID == "" {
+			return &mcp.ToolResult{
+				Content: []mcp.Content{{Type: "text", Text: "错误: connection_id 参数必填"}},
+				IsError: true,
+			}, nil
+		}
+
+		// 获取现有连接
+		existing, err := db.GetWebshellConnection(connID)
+		if err != nil || existing == nil {
+			return &mcp.ToolResult{
+				Content: []mcp.Content{{Type: "text", Text: "未找到指定的 WebShell 连接: " + connID}},
+				IsError: true,
+			}, nil
+		}
+
+		// 更新字段（如果提供了新值）
+		if urlStr, ok := args["url"].(string); ok && urlStr != "" {
+			existing.URL = urlStr
+		}
+		if password, ok := args["password"].(string); ok {
+			existing.Password = password
+		}
+		if shellType, ok := args["type"].(string); ok && shellType != "" {
+			existing.Type = strings.ToLower(shellType)
+		}
+		if method, ok := args["method"].(string); ok && method != "" {
+			existing.Method = strings.ToLower(method)
+		}
+		if cmdParam, ok := args["cmd_param"].(string); ok && cmdParam != "" {
+			existing.CmdParam = cmdParam
+		}
+		if remark, ok := args["remark"].(string); ok {
+			existing.Remark = remark
+		}
+
+		if err := db.UpdateWebshellConnection(existing); err != nil {
+			return &mcp.ToolResult{
+				Content: []mcp.Content{{Type: "text", Text: "更新 WebShell 连接失败: " + err.Error()}},
+				IsError: true,
+			}, nil
+		}
+
+		return &mcp.ToolResult{
+			Content: []mcp.Content{{
+				Type: "text",
+				Text: fmt.Sprintf("WebShell 连接更新成功！\n\n连接ID: %s\nURL: %s\n类型: %s\n请求方式: %s\n命令参数: %s\n备注: %s", existing.ID, existing.URL, existing.Type, existing.Method, existing.CmdParam, existing.Remark),
+			}},
+			IsError: false,
+		}, nil
+	}
+	mcpServer.RegisterTool(updateTool, updateHandler)
+
+	// manage_webshell_delete - 删除 webshell 连接
+	deleteTool := mcp.Tool{
+		Name:        builtin.ToolManageWebshellDelete,
+		Description: "删除指定的 WebShell 连接。",
+		ShortDescription: "删除 WebShell 连接",
+		InputSchema: map[string]interface{}{
+			"type": "object",
+			"properties": map[string]interface{}{
+				"connection_id": map[string]interface{}{
+					"type":        "string",
+					"description": "要删除的 WebShell 连接 ID（必填）",
+				},
+			},
+			"required": []string{"connection_id"},
+		},
+	}
+	deleteHandler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
+		connID, _ := args["connection_id"].(string)
+		if connID == "" {
+			return &mcp.ToolResult{
+				Content: []mcp.Content{{Type: "text", Text: "错误: connection_id 参数必填"}},
+				IsError: true,
+			}, nil
+		}
+
+		if err := db.DeleteWebshellConnection(connID); err != nil {
+			return &mcp.ToolResult{
+				Content: []mcp.Content{{Type: "text", Text: "删除 WebShell 连接失败: " + err.Error()}},
+				IsError: true,
+			}, nil
+		}
+
+		return &mcp.ToolResult{
+			Content: []mcp.Content{{
+				Type: "text",
+				Text: fmt.Sprintf("WebShell 连接 %s 已成功删除", connID),
+			}},
+			IsError: false,
+		}, nil
+	}
+	mcpServer.RegisterTool(deleteTool, deleteHandler)
+
+	// manage_webshell_test - 测试 webshell 连接
+	testTool := mcp.Tool{
+		Name:        builtin.ToolManageWebshellTest,
+		Description: "测试指定的 WebShell 连接是否可用，会尝试执行一个简单的命令（如 whoami 或 dir）。",
+		ShortDescription: "测试 WebShell 连接",
+		InputSchema: map[string]interface{}{
+			"type": "object",
+			"properties": map[string]interface{}{
+				"connection_id": map[string]interface{}{
+					"type":        "string",
+					"description": "要测试的 WebShell 连接 ID（必填）",
+				},
+				"command": map[string]interface{}{
+					"type":        "string",
+					"description": "测试命令，默认为 whoami（Linux）或 dir（Windows）",
+				},
+			},
+			"required": []string{"connection_id"},
+		},
+	}
+	testHandler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
+		connID, _ := args["connection_id"].(string)
+		if connID == "" {
+			return &mcp.ToolResult{
+				Content: []mcp.Content{{Type: "text", Text: "错误: connection_id 参数必填"}},
+				IsError: true,
+			}, nil
+		}
+
+		// 获取连接
+		conn, err := db.GetWebshellConnection(connID)
+		if err != nil || conn == nil {
+			return &mcp.ToolResult{
+				Content: []mcp.Content{{Type: "text", Text: "未找到指定的 WebShell 连接: " + connID}},
+				IsError: true,
+			}, nil
+		}
+
+		// 确定测试命令
+		testCmd, _ := args["command"].(string)
+		if testCmd == "" {
+			// 根据 shell 类型选择默认命令
+			if conn.Type == "asp" || conn.Type == "aspx" {
+				testCmd = "dir"
+			} else {
+				testCmd = "whoami"
+			}
+		}
+
+		// 执行测试命令
+		output, ok, errMsg := webshellHandler.ExecWithConnection(conn, testCmd)
+		if errMsg != "" {
+			return &mcp.ToolResult{
+				Content: []mcp.Content{{Type: "text", Text: fmt.Sprintf("连接测试失败！\n\n连接ID: %s\nURL: %s\n错误: %s", connID, conn.URL, errMsg)}},
+				IsError: true,
+			}, nil
+		}
+
+		if !ok {
+			return &mcp.ToolResult{
+				Content: []mcp.Content{{Type: "text", Text: fmt.Sprintf("连接测试失败！HTTP 非 200\n\n连接ID: %s\nURL: %s\n输出: %s", connID, conn.URL, output)}},
+				IsError: true,
+			}, nil
+		}
+
+		return &mcp.ToolResult{
+			Content: []mcp.Content{{
+				Type: "text",
+				Text: fmt.Sprintf("连接测试成功！\n\n连接ID: %s\nURL: %s\n类型: %s\n\n测试命令: %s\n输出结果:\n%s", connID, conn.URL, conn.Type, testCmd, output),
+			}},
+			IsError: false,
+		}, nil
+	}
+	mcpServer.RegisterTool(testTool, testHandler)
+
+	logger.Info("WebShell 管理工具注册成功")
+}
+
 // initializeKnowledge 初始化知识库组件（用于动态初始化）
 func initializeKnowledge(
 	cfg *config.Config,
diff --git a/internal/mcp/builtin/constants.go b/internal/mcp/builtin/constants.go
index 13c01a6b..1ef554ae 100644
--- a/internal/mcp/builtin/constants.go
+++ b/internal/mcp/builtin/constants.go
@@ -19,6 +19,13 @@ const (
 	ToolWebshellFileList   = "webshell_file_list"
 	ToolWebshellFileRead   = "webshell_file_read"
 	ToolWebshellFileWrite  = "webshell_file_write"
+
+	// WebShell 连接管理工具（用于通过 MCP 管理 webshell 连接）
+	ToolManageWebshellList   = "manage_webshell_list"
+	ToolManageWebshellAdd    = "manage_webshell_add"
+	ToolManageWebshellUpdate = "manage_webshell_update"
+	ToolManageWebshellDelete = "manage_webshell_delete"
+	ToolManageWebshellTest   = "manage_webshell_test"
 )
 
 // IsBuiltinTool 检查工具名称是否是内置工具
@@ -32,7 +39,12 @@ func IsBuiltinTool(toolName string) bool {
 		ToolWebshellExec,
 		ToolWebshellFileList,
 		ToolWebshellFileRead,
-		ToolWebshellFileWrite:
+		ToolWebshellFileWrite,
+		ToolManageWebshellList,
+		ToolManageWebshellAdd,
+		ToolManageWebshellUpdate,
+		ToolManageWebshellDelete,
+		ToolManageWebshellTest:
 		return true
 	default:
 		return false
@@ -51,5 +63,10 @@ func GetAllBuiltinTools() []string {
 		ToolWebshellFileList,
 		ToolWebshellFileRead,
 		ToolWebshellFileWrite,
+		ToolManageWebshellList,
+		ToolManageWebshellAdd,
+		ToolManageWebshellUpdate,
+		ToolManageWebshellDelete,
+		ToolManageWebshellTest,
 	}
 }