diff --git a/Remote/server/mcp.go b/Remote/server/mcp.go index ddc9f6a..5793285 100644 --- a/Remote/server/mcp.go +++ b/Remote/server/mcp.go @@ -34,21 +34,22 @@ func NewFscanMCPServer() *server.MCPServer { "1.0.0", ) toolHandler := service.NewFscanMCPTool() + // 添加提示词 // 添加工具处理器 s.AddTool( mcp.NewTool("StartScan", - mcp.WithDescription("开始扫描"), + mcp.WithDescription("启动端口和服务扫描任务,适用于安全评估或资产排查场景。"), mcp.WithString("target", mcp.Required(), - mcp.Description("扫描目标"), + mcp.Description("待扫描的目标地址,支持IP、域名或CIDR格式(如192.168.1.1、example.com、10.0.0.0/24)。"), ), ), toolHandler.StartScan, ) s.AddTool( mcp.NewTool("GetScanResults", - mcp.WithDescription("获取扫描结果"), + mcp.WithDescription("获取当前扫描任务的执行进度和已完成部分的结果。若扫描尚未完成,也会返回当前阶段的中间结果,供用户分析或决策使用。"), ), toolHandler.GetScanResults, ) diff --git a/Remote/service/mcp.go b/Remote/service/mcp.go index e4d50a0..33c8ef3 100644 --- a/Remote/service/mcp.go +++ b/Remote/service/mcp.go @@ -2,8 +2,10 @@ package service import ( "context" + "encoding/json" "errors" "fmt" + "strings" "sync" "sync/atomic" "time" @@ -17,6 +19,7 @@ type FscanMCPTool struct { scanMutex sync.Mutex isScanning int32 scanStartTime time.Time + lastSentIndex int } func NewFscanMCPTool() *FscanMCPTool { @@ -26,23 +29,28 @@ func NewFscanMCPTool() *FscanMCPTool { func (s *FscanMCPTool) StartScan(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { target, ok := request.Params.Arguments["target"].(string) if !ok { - return nil, errors.New("name must be a string") + return nil, errors.New("target must be a string") } - //构造扫描字符串 - arg := fmt.Sprintf("-h %s", target) + arg := fmt.Sprintf("-h %s", target) if !atomic.CompareAndSwapInt32(&s.isScanning, 0, 1) { return &mcp.CallToolResult{ Content: []mcp.Content{ mcp.TextContent{ Type: "text", - Text: "已有扫描任务正在运行,请稍后重试", + Text: fmt.Sprintf("A scan is already in progress. Progress: %s%%", func() string { + if Common.Num == 0 || Common.End == 0 { + return "initializing" + } + return fmt.Sprintf("%.2f", (float64(Common.End)/float64(Common.Num))*100) + }()), }, }, }, nil } - s.scanStartTime = time.Now() // 记录任务开始时间 + s.scanStartTime = time.Now() + s.lastSentIndex = 0 go func() { defer atomic.StoreInt32(&s.isScanning, 0) @@ -50,7 +58,7 @@ func (s *FscanMCPTool) StartScan(ctx context.Context, request mcp.CallToolReques s.scanMutex.Lock() defer s.scanMutex.Unlock() - Common.LogDebug("异步执行扫描请求,目标: " + arg) + Common.LogDebug("Starting scan asynchronously: " + arg) var info Common.HostInfo if err := Common.FlagFromRemote(&info, arg); err != nil { @@ -60,34 +68,94 @@ func (s *FscanMCPTool) StartScan(ctx context.Context, request mcp.CallToolReques return } if err := Common.CloseOutput(); err != nil { - Common.LogError(fmt.Sprintf("关闭输出系统失败: %v", err)) + Common.LogError(fmt.Sprintf("Failed to close output: %v", err)) return } if err := Common.InitOutput(); err != nil { - Common.LogError(fmt.Sprintf("初始化输出系统失败: %v", err)) + Common.LogError(fmt.Sprintf("Failed to initialize output: %v", err)) return } Core.Scan(info) - - Common.LogDebug("扫描任务完成") + Common.LogDebug("Scan completed.") }() + return &mcp.CallToolResult{ Content: []mcp.Content{ mcp.TextContent{ Type: "text", - Text: fmt.Sprintf("扫描任务开始,扫描参数: %s", arg), + Text: fmt.Sprintf("Scan started. Parameters: %s", arg), }, }, }, nil } func (s *FscanMCPTool) GetScanResults(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { + // immediate, _ := request.Params.Arguments["immediate"].(bool) + // server := server.ServerFromContext(ctx) + + // if !immediate { + // for { + // select { + // case <-ctx.Done(): + // return nil, ctx.Err() + // default: + // if atomic.LoadInt32(&s.isScanning) == 0 { + // break + // } + // _ = server.SendNotificationToClient( + // ctx, + // "Fscan/progress", + // map[string]interface{}{ + // "end": Common.End, + // "total": Common.Num, + // }, + // ) + // time.Sleep(1000 * time.Millisecond) + // } + // } + // } + results, err := Common.GetResults() + if err != nil { + Common.LogError(fmt.Sprintf("Failed to read results: %v", err)) + return nil, fmt.Errorf("failed to read results: %w", err) + } + + var sb strings.Builder + for _, r := range results { + detailsJSON, err := json.Marshal(r.Details) + if err != nil { + Common.LogError(fmt.Sprintf("Failed to encode result details (Target: %s, Type: %s): %v", r.Target, r.Type, err)) + continue + } + + sb.WriteString(fmt.Sprintf( + "--- Result ---\nTime: %s\nType: %s\nTarget: %s\nStatus: %s\nDetails: %s\n\n", + r.Time.Format(time.RFC3339), + r.Type, + r.Target, + r.Status, + string(detailsJSON), + )) + } + progress := fmt.Sprintf("Scan Progress: %.2f%%\n", func() float64 { + if Common.Num == 0 || Common.End == 0 { + return 0 + } + if s.isScanning == 0 { + return 100 + } + return (float64(Common.End) / float64(Common.Num)) * 100 + }()) return &mcp.CallToolResult{ Content: []mcp.Content{ mcp.TextContent{ Type: "text", - Text: Common.OutputFormat + Common.Outputfile, + Text: progress, + }, + mcp.TextContent{ + Type: "text", + Text: sb.String(), }, }, }, nil