mirror of
https://github.com/shadow1ng/fscan.git
synced 2025-07-13 21:02:44 +08:00
fix: 更新扫描工具描述和参数提示,增强用户体验
This commit is contained in:
parent
d1e4a864a4
commit
226983a51c
@ -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,
|
||||||
)
|
)
|
||||||
|
@ -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
|
||||||
|
Loading…
Reference in New Issue
Block a user