mirror of
https://github.com/shadow1ng/fscan.git
synced 2025-07-13 12:52:44 +08:00
118 lines
3.1 KiB
Go
118 lines
3.1 KiB
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
|
||
}
|