fix: 更新扫描工具描述和参数提示,增强用户体验

This commit is contained in:
tongque 2025-05-08 12:39:56 +08:00
parent d1e4a864a4
commit 226983a51c
2 changed files with 84 additions and 15 deletions

View File

@ -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,
)

View File

@ -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