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", "1.0.0",
) )
toolHandler := service.NewFscanMCPTool() toolHandler := service.NewFscanMCPTool()
// 添加提示词
// 添加工具处理器 // 添加工具处理器
s.AddTool( s.AddTool(
mcp.NewTool("StartScan", mcp.NewTool("StartScan",
mcp.WithDescription("开始扫描"), mcp.WithDescription("启动端口和服务扫描任务,适用于安全评估或资产排查场景。"),
mcp.WithString("target", mcp.WithString("target",
mcp.Required(), mcp.Required(),
mcp.Description("扫描目标"), mcp.Description("扫描目标地址支持IP、域名或CIDR格式如192.168.1.1、example.com、10.0.0.0/24"),
), ),
), ),
toolHandler.StartScan, toolHandler.StartScan,
) )
s.AddTool( s.AddTool(
mcp.NewTool("GetScanResults", mcp.NewTool("GetScanResults",
mcp.WithDescription("获取扫描结果"), mcp.WithDescription("获取当前扫描任务的执行进度和已完成部分的结果。若扫描尚未完成,也会返回当前阶段的中间结果,供用户分析或决策使用。"),
), ),
toolHandler.GetScanResults, toolHandler.GetScanResults,
) )

View File

@ -2,8 +2,10 @@ package service
import ( import (
"context" "context"
"encoding/json"
"errors" "errors"
"fmt" "fmt"
"strings"
"sync" "sync"
"sync/atomic" "sync/atomic"
"time" "time"
@ -17,6 +19,7 @@ type FscanMCPTool struct {
scanMutex sync.Mutex scanMutex sync.Mutex
isScanning int32 isScanning int32
scanStartTime time.Time scanStartTime time.Time
lastSentIndex int
} }
func NewFscanMCPTool() *FscanMCPTool { func NewFscanMCPTool() *FscanMCPTool {
@ -26,23 +29,28 @@ func NewFscanMCPTool() *FscanMCPTool {
func (s *FscanMCPTool) StartScan(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { func (s *FscanMCPTool) StartScan(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
target, ok := request.Params.Arguments["target"].(string) target, ok := request.Params.Arguments["target"].(string)
if !ok { 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) { if !atomic.CompareAndSwapInt32(&s.isScanning, 0, 1) {
return &mcp.CallToolResult{ return &mcp.CallToolResult{
Content: []mcp.Content{ Content: []mcp.Content{
mcp.TextContent{ mcp.TextContent{
Type: "text", 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 }, nil
} }
s.scanStartTime = time.Now() // 记录任务开始时间 s.scanStartTime = time.Now()
s.lastSentIndex = 0
go func() { go func() {
defer atomic.StoreInt32(&s.isScanning, 0) defer atomic.StoreInt32(&s.isScanning, 0)
@ -50,7 +58,7 @@ func (s *FscanMCPTool) StartScan(ctx context.Context, request mcp.CallToolReques
s.scanMutex.Lock() s.scanMutex.Lock()
defer s.scanMutex.Unlock() defer s.scanMutex.Unlock()
Common.LogDebug("异步执行扫描请求,目标: " + arg) Common.LogDebug("Starting scan asynchronously: " + arg)
var info Common.HostInfo var info Common.HostInfo
if err := Common.FlagFromRemote(&info, arg); err != nil { if err := Common.FlagFromRemote(&info, arg); err != nil {
@ -60,34 +68,94 @@ func (s *FscanMCPTool) StartScan(ctx context.Context, request mcp.CallToolReques
return return
} }
if err := Common.CloseOutput(); err != nil { if err := Common.CloseOutput(); err != nil {
Common.LogError(fmt.Sprintf("关闭输出系统失败: %v", err)) Common.LogError(fmt.Sprintf("Failed to close output: %v", err))
return return
} }
if err := Common.InitOutput(); err != nil { if err := Common.InitOutput(); err != nil {
Common.LogError(fmt.Sprintf("初始化输出系统失败: %v", err)) Common.LogError(fmt.Sprintf("Failed to initialize output: %v", err))
return return
} }
Core.Scan(info) Core.Scan(info)
Common.LogDebug("Scan completed.")
Common.LogDebug("扫描任务完成")
}() }()
return &mcp.CallToolResult{ return &mcp.CallToolResult{
Content: []mcp.Content{ Content: []mcp.Content{
mcp.TextContent{ mcp.TextContent{
Type: "text", Type: "text",
Text: fmt.Sprintf("扫描任务开始,扫描参数: %s", arg), Text: fmt.Sprintf("Scan started. Parameters: %s", arg),
}, },
}, },
}, nil }, nil
} }
func (s *FscanMCPTool) GetScanResults(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { 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{ return &mcp.CallToolResult{
Content: []mcp.Content{ Content: []mcp.Content{
mcp.TextContent{ mcp.TextContent{
Type: "text", Type: "text",
Text: Common.OutputFormat + Common.Outputfile, Text: progress,
},
mcp.TextContent{
Type: "text",
Text: sb.String(),
}, },
}, },
}, nil }, nil