fonchain-fiee/pkg/service/check/security.go

package check

import (
	"errors"
	"fmt"
	"mime/multipart"
	"path/filepath"
	"strconv"
	"time"

	modelSecurity "fonchain-fiee/pkg/model/security"
	pkgSecurity "fonchain-fiee/pkg/security"

	"github.com/fonchain/utils/security"
)

// ImageCheckUrlValid 图片黄疸检测 true 是通过
func ImageCheckUrlValid(imgUrl string) (bool, error) {
	resp, err := pkgSecurity.ImageScanner.ScanImageByURL(imgUrl, fmt.Sprint(time.Now().UnixMicro()), security.BaselineCheckGlobal)
	if err != nil {
		err = errors.New("图片检测请求失败")
		return false, err
	}
	if resp.Code != 200 {
		err = errors.New("图片检测失败，错误码")
		return false, err
	}
	if len(resp.Data) == 0 || len(resp.Data[0].Results) == 0 {
		return false, errors.New("图片检测结果异常")
	}
	riskLevel := resp.Data[0].Results[0].RiskLevel
	if *riskLevel == "none" {
		return true, nil
	}
	return false, nil
}

func ImageCheckByte(file *multipart.FileHeader) (bool, error) {
	//imageScanner, err := security.NewImageScanner(&security.Config{
	//	RAMAccessKeyID:      "LTAI5tNBzbeEbG1yCitvHsMb",
	//	RAMAccessKeySecret:  "G1xAUB8G6WDVo0SLr6DJaJjNWIlpmO",
	//	RAMRoleArn:          "acs:ram::5828544250383902:role/content-secret",
	//	Region:              "ap-southeast-1",
	//	Endpoint:            "green-cip.ap-southeast-1.aliyuncs.com",
	//	TempAccessKeyID:     "",
	//	TempAccessKeySecret: "",
	//	SecurityToken:       "",
	//})
	//if err != nil {
	//	return false, err
	//}
	//resp, err := imageScanner.ScanImageByFileByte(file, fmt.Sprint(time.Now().UnixMicro()), security.BaselineCheckGlobal)
	//fmt.Println(resp)
	return false, nil
}

func SecurityFile(textVal string) (bool, error) {
	if textVal == "" {
		return true, nil
	}
	var fileInfo modelSecurity.FileInfo
	fileInfo.FileName = textVal
	fileInfo.FileUrl = textVal
	//获取文件名称后缀
	extension := filepath.Ext(textVal)
	switch extension {
	case ".jpg", ".jpeg", ".png", ".gif", ".bmp", ".webp":
		fileInfo.FileType = "image"
	case ".mp4", ".avi", ".mov", ".wmv", ".flv", ".mkv":
		fileInfo.FileType = "video"
	default:
		return true, nil
	}
	err := securityScan(&fileInfo)
	if err != nil {
		return false, err
	}
	return true, nil
}
func SecurityText(textVal string) (bool, error) {
	aliConfig, err := pkgSecurity.GetGlobalConfig("./data/alibabacloud.env")
	if err != nil {
		return false, err
	}
	fmt.Println("开始获取STS Token")
	err = aliConfig.GetSTSToken()
	if err != nil {
		return false, errors.New("获取STS Token失败")
	}
	if err = handleTextScan(aliConfig, textVal); err != nil {
		return false, err
	}
	return true, nil
}

func securityScan(fileInfo *modelSecurity.FileInfo) (err error) {
	//加载阿里云配置获取临时token
	aliConfig, err := pkgSecurity.GetGlobalConfig("./data/alibabacloud.env")
	if err != nil {
		return err
	}
	fmt.Println("开始获取STS Token")
	err = aliConfig.GetSTSToken()
	if err != nil {
		return errors.New("获取STS Token失败")
	}
	fmt.Println("获取STS Token成功")
	switch fileInfo.FileType {
	case "image":
		return handleImageScan(aliConfig, fileInfo)
	case "video":
		return handleVideoScan(aliConfig, fileInfo)
	default:
		return errors.New("不支持的审核类型")
	}
}

func handleVideoScan(config *security.Config, fileInfo *modelSecurity.FileInfo) (err error) {
	fmt.Println("\n=== 视频内容安全审核 ===")

	scanner, err := security.NewVideoScanner(config)
	if err != nil {
		fmt.Printf("创建视频扫描器失败: %v\n", err)
		return errors.New("创建视频扫描器失败")
	}

	if fileInfo.FileUrl == "" {
		fmt.Println("视频不能为空")
		return errors.New("视频不能为空")
	}
	serviceType := security.VideoPostCheckByVLGlobal
	dataID := "video_" + time.Now().Format("20060102150405")

	fmt.Println("正在扫描视频（这可能需要几分钟）...")
	result, err := scanner.ScanVideoAndWait(fileInfo.FileUrl, dataID, serviceType, 10*time.Minute)
	if err != nil {
		fmt.Printf("扫描失败: %v\n", err)
		return errors.New("扫描失败")
	}
	scanner.PrintResult(result)
	for _, v := range result.Data {
		fileInfo.SecurityStatus = *v.Results.RiskLevel
		if v.Results.FrameResult != nil {
			frameResult := *v.Results.FrameResult
			if len(frameResult.FrameSummarys) > 0 {
				for _, summary := range frameResult.FrameSummarys {
					fileInfo.Describe = fileInfo.Describe + *summary.Description + ";"
				}
			}
			if len(frameResult.Frames) > 0 {
				for _, frame := range frameResult.Frames {
					fileInfo.ProblemTimeFrame = fileInfo.ProblemTimeFrame + strconv.FormatFloat(float64(*frame.Offset), 'f', 2, 64) + "秒;"
				}
			}
		}
		if v.Results.RiskLevel != nil {
			lv := v.Results.RiskLevel
			if *lv == "none" {
				return nil
			}
			return errors.New("图片违规")
		}

		return errors.New("风险等级检测未成功")
	}
	return nil
}

func handleImageScan(config *security.Config, fileInfo *modelSecurity.FileInfo) (err error) {
	fmt.Println("\n=== 图片内容安全审核 ===")

	scanner, err := security.NewImageScanner(config)
	if err != nil {
		fmt.Printf("创建图片扫描器失败:%v", err)
		return errors.New("创建图片扫描器失败")
	}

	if fileInfo.FileUrl == "" {
		fmt.Println("文件不能为空")
		return errors.New("文件不能为空")
	}

	serviceType := security.PostImageCheckByVLGlobal

	fmt.Println("正在扫描图片...")
	result, err := scanner.ScanImageByURL(fileInfo.FileUrl, "image_"+time.Now().Format("20060102150405"), serviceType)
	if err != nil {
		fmt.Printf("扫描失败: %v\n", err)
		return errors.New("扫描失败")
	}
	scanner.PrintResult(result)
	for _, v := range result.Data {
		for _, v2 := range v.Results {
			fileInfo.SecurityStatus = *v2.RiskLevel
			if v2.RiskLevel != nil {
				lv := v2.RiskLevel
				if *lv == "none" {
					return nil
				}
				return errors.New("图片违规")
			}
			return errors.New("风险等级检测未成功")
		}
	}
	return nil
}

func handleTextScan(config *security.Config, textVal string) (err error) {
	fmt.Println("\n=== 图片内容安全审核 ===")

	scanner, err := security.NewTextScanner(config)
	if err != nil {
		fmt.Printf("创建扫描器失败:%v", err)
		return errors.New("创建扫描器失败")
	}

	serviceType := security.TextBaselineCheckGlobal

	fmt.Println("正在扫描图片...")
	result, err := scanner.ScanText(textVal, "image_"+time.Now().Format("20060102150405"), serviceType)
	if err != nil {
		fmt.Printf("扫描失败: %v\n", err)
		return errors.New("扫描失败")
	}
	scanner.PrintResult(result)
	for _, v := range result.Data {
		for _, vv := range v.Results {
			label := *vv.Label
			if label != "nonLabel" {
				return errors.New("风险等级过高")
			}
		}
	}
	fmt.Println(result)
	return nil
}