fscan/Remote/service/mcp.go

package service

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/mark3labs/mcp-go/mcp"
	"github.com/shadow1ng/fscan/Common"
	"github.com/shadow1ng/fscan/Core"
)

type FscanMCPTool struct {
	scanMutex     sync.Mutex
	isScanning    int32
	scanStartTime time.Time
	lastSentIndex int
}

func NewFscanMCPTool() *FscanMCPTool {
	return &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("target must be a string")
	}

	arg := fmt.Sprintf("-h %s", target)
	if !atomic.CompareAndSwapInt32(&s.isScanning, 0, 1) {
		return &mcp.CallToolResult{
			Content: []mcp.Content{
				mcp.TextContent{
					Type: "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.lastSentIndex = 0

	go func() {
		defer atomic.StoreInt32(&s.isScanning, 0)

		s.scanMutex.Lock()
		defer s.scanMutex.Unlock()

		Common.LogDebug("Starting scan asynchronously: " + arg)

		var info Common.HostInfo
		if err := Common.FlagFromRemote(&info, arg); err != nil {
			return
		}
		if err := Common.Parse(&info); err != nil {
			return
		}
		if err := Common.CloseOutput(); err != nil {
			Common.LogError(fmt.Sprintf("Failed to close output: %v", err))
			return
		}
		if err := Common.InitOutput(); err != nil {
			Common.LogError(fmt.Sprintf("Failed to initialize output: %v", err))
			return
		}

		Core.Scan(info)
		Common.LogDebug("Scan completed.")
	}()

	return &mcp.CallToolResult{
		Content: []mcp.Content{
			mcp.TextContent{
				Type: "text",
				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: progress,
			},
			mcp.TextContent{
				Type: "text",
				Text: sb.String(),
			},
		},
	}, nil
}