fscan/RPC/service/fscan.go

package service

import (
	"context"
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"github.com/shadow1ng/fscan/Common"
	"github.com/shadow1ng/fscan/Core"
	pb "github.com/shadow1ng/fscan/RPC/lib"
	structpb "google.golang.org/protobuf/types/known/structpb"
)

type FscanService struct {
	pb.UnimplementedFscanServiceServer
	scanMutex     sync.Mutex
	isScanning    int32
	scanStartTime time.Time // 记录扫描开始时间
}

func (s *FscanService) StartScan(ctx context.Context, req *pb.StartScanRequest) (*pb.StartScanResponse, error) {
	if !atomic.CompareAndSwapInt32(&s.isScanning, 0, 1) {
		return &pb.StartScanResponse{
			TaskId:  "current",
			Message: "已有扫描任务正在运行，请稍后重试",
		}, nil
	}

	s.scanStartTime = time.Now() // 记录任务开始时间

	go func(req *pb.StartScanRequest) {
		defer atomic.StoreInt32(&s.isScanning, 0)

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

		Common.LogDebug("异步执行扫描请求，目标: " + req.Arg)

		var info Common.HostInfo
		if err := Common.FlagFromRemote(&info, req.Arg); err != nil {
			return
		}
		if err := Common.Parse(&info); err != nil {
			return
		}
		if err := Common.CloseOutput(); err != nil {
			Common.LogError(fmt.Sprintf("关闭输出系统失败: %v", err))
			return
		}
		if err := Common.InitOutput(); err != nil {
			Common.LogError(fmt.Sprintf("初始化输出系统失败: %v", err))
			return
		}

		Core.Scan(info)

		Common.LogDebug("扫描任务完成")
	}(req)

	return &pb.StartScanResponse{
		TaskId:  "current",
		Message: "成功启动扫描任务",
	}, nil
}

func (s *FscanService) GetScanResults(ctx context.Context, req *pb.TaskResultsRequest) (*pb.TaskResultsResponse, error) {
	// 获取扫描结果
	results, err := Common.GetResults()
	if err != nil {
		// 记录详细错误信息
		Common.LogError(fmt.Sprintf("读取结果失败: %v", err))
		return nil, fmt.Errorf("读取结果失败: %w", err)
	}

	// 创建一个用于存储转换后的pb扫描结果的切片
	pbResults := make([]*pb.ScanResult, 0, len(results))

	// 遍历每一项结果，进行转换
	for _, r := range results {
		// 尝试将详情转换为Struct
		detailsStruct, err := structpb.NewStruct(r.Details)
		if err != nil {
			// 记录转换失败的详细信息，并跳过当前项
			Common.LogError(fmt.Sprintf("转换为 Struct 失败 (Target: %s, Type: %s): %v", r.Target, r.Type, err))
			continue
		}

		// 将转换后的结果添加到 pbResults
		pbResults = append(pbResults, &pb.ScanResult{
			Time:        r.Time.Format(time.RFC3339), // 使用 RFC3339 格式化时间
			Type:        string(r.Type),
			Target:      r.Target,
			Status:      r.Status,
			DetailsJson: detailsStruct,
		})
	}

	// 通过原子操作判断扫描是否完成
	finished := atomic.LoadInt32(&s.isScanning) == 0

	// 如果任务未完成，计算 Total 和 End，仅在需要时计算
	var total, end int64
	if !finished {
		total = Common.Num
		end = Common.End
	}

	// 返回响应
	return &pb.TaskResultsResponse{
		Results:  pbResults,
		Finished: finished,
		Total:    total, // 返回 Total
		End:      end,   // 返回 End
	}, nil
}