|
|
package websocket |
|
|
|
|
|
import ( |
|
|
"encoding/json" |
|
|
"fmt" |
|
|
"net/http" |
|
|
"os" |
|
|
"path/filepath" |
|
|
"regexp" |
|
|
"sort" |
|
|
"strconv" |
|
|
"strings" |
|
|
"time" |
|
|
|
|
|
"trpc.group/trpc-go/trpc-go/log" |
|
|
|
|
|
"github.com/Tencent/AI-Infra-Guard/common/fingerprints/parser" |
|
|
"github.com/Tencent/AI-Infra-Guard/pkg/vulstruct" |
|
|
"github.com/gin-gonic/gin" |
|
|
"gopkg.in/yaml.v3" |
|
|
) |
|
|
|
|
|
|
|
|
var validName = regexp.MustCompile(`^[a-zA-Z0-9\\._ -]+$`) |
|
|
|
|
|
func isValidName(name string) bool { |
|
|
if strings.Contains(name, "..") { |
|
|
return false |
|
|
} |
|
|
return validName.MatchString(name) |
|
|
} |
|
|
|
|
|
|
|
|
type FingerprintWithTime struct { |
|
|
FingerPrint parser.FingerPrint |
|
|
ModTime time.Time |
|
|
} |
|
|
|
|
|
|
|
|
type EvaluationDataItem struct { |
|
|
Prompt string `json:"prompt"` |
|
|
} |
|
|
|
|
|
type EvaluationDataset struct { |
|
|
Name string `json:"name"` |
|
|
Description string `json:"description"` |
|
|
DescriptionZh string `json:"description_zh,omitempty"` |
|
|
Author string `json:"author,omitempty"` |
|
|
Source []string `json:"source,omitempty"` |
|
|
Count int `json:"count"` |
|
|
Default bool `json:"default"` |
|
|
Tags []string `json:"tags,omitempty"` |
|
|
Recommendation int `json:"recommendation,omitempty"` |
|
|
Language string `json:"language,omitempty"` |
|
|
Data []EvaluationDataItem `json:"data"` |
|
|
} |
|
|
|
|
|
|
|
|
func HandleListFingerprints(c *gin.Context) { |
|
|
|
|
|
pageStr := c.DefaultQuery("page", "1") |
|
|
sizeStr := c.DefaultQuery("size", "20") |
|
|
page, _ := strconv.Atoi(pageStr) |
|
|
size, _ := strconv.Atoi(sizeStr) |
|
|
if page < 1 { |
|
|
page = 1 |
|
|
} |
|
|
if size < 1 { |
|
|
size = 10 |
|
|
} |
|
|
|
|
|
|
|
|
nameQuery := strings.ToLower(c.DefaultQuery("q", "")) |
|
|
|
|
|
|
|
|
var allFingerprintsWithTime []FingerprintWithTime |
|
|
root := "data/fingerprints" |
|
|
filepath.WalkDir(root, func(path string, d os.DirEntry, err error) error { |
|
|
if err != nil { |
|
|
return nil |
|
|
} |
|
|
if !d.IsDir() && strings.HasSuffix(d.Name(), ".yaml") { |
|
|
content, _ := os.ReadFile(path) |
|
|
fp, err := parser.InitFingerPrintFromData(content) |
|
|
if err == nil && fp != nil { |
|
|
|
|
|
fileInfo, _ := d.Info() |
|
|
modTime := time.Time{} |
|
|
if fileInfo != nil { |
|
|
modTime = fileInfo.ModTime() |
|
|
} |
|
|
allFingerprintsWithTime = append(allFingerprintsWithTime, FingerprintWithTime{ |
|
|
FingerPrint: *fp, |
|
|
ModTime: modTime, |
|
|
}) |
|
|
} |
|
|
} |
|
|
return nil |
|
|
}) |
|
|
|
|
|
|
|
|
var filteredFingerprintsWithTime []FingerprintWithTime |
|
|
if nameQuery == "" { |
|
|
filteredFingerprintsWithTime = allFingerprintsWithTime |
|
|
} else { |
|
|
for _, fpWithTime := range allFingerprintsWithTime { |
|
|
fp := fpWithTime.FingerPrint |
|
|
if strings.Contains(strings.ToLower(fp.Info.Name), nameQuery) { |
|
|
filteredFingerprintsWithTime = append(filteredFingerprintsWithTime, fpWithTime) |
|
|
continue |
|
|
} |
|
|
if strings.Contains(strings.ToLower(fp.Info.Desc), nameQuery) { |
|
|
filteredFingerprintsWithTime = append(filteredFingerprintsWithTime, fpWithTime) |
|
|
continue |
|
|
} |
|
|
if strings.Contains(strings.ToLower(fp.Info.Author), nameQuery) { |
|
|
filteredFingerprintsWithTime = append(filteredFingerprintsWithTime, fpWithTime) |
|
|
continue |
|
|
} |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
sort.Slice(filteredFingerprintsWithTime, func(i, j int) bool { |
|
|
fpI := filteredFingerprintsWithTime[i].FingerPrint |
|
|
fpJ := filteredFingerprintsWithTime[j].FingerPrint |
|
|
|
|
|
|
|
|
if fpI.Info.Recommendation > 0 && fpJ.Info.Recommendation > 0 { |
|
|
if fpI.Info.Recommendation != fpJ.Info.Recommendation { |
|
|
return fpI.Info.Recommendation > fpJ.Info.Recommendation |
|
|
} |
|
|
|
|
|
return filteredFingerprintsWithTime[i].ModTime.After(filteredFingerprintsWithTime[j].ModTime) |
|
|
} |
|
|
|
|
|
|
|
|
if fpI.Info.Recommendation > 0 && fpJ.Info.Recommendation <= 0 { |
|
|
return true |
|
|
} |
|
|
if fpI.Info.Recommendation <= 0 && fpJ.Info.Recommendation > 0 { |
|
|
return false |
|
|
} |
|
|
|
|
|
|
|
|
return filteredFingerprintsWithTime[i].ModTime.After(filteredFingerprintsWithTime[j].ModTime) |
|
|
}) |
|
|
|
|
|
|
|
|
var filteredFingerprints []parser.FingerPrint |
|
|
for _, fpWithTime := range filteredFingerprintsWithTime { |
|
|
filteredFingerprints = append(filteredFingerprints, fpWithTime.FingerPrint) |
|
|
} |
|
|
|
|
|
|
|
|
total := len(filteredFingerprints) |
|
|
start := (page - 1) * size |
|
|
end := start + size |
|
|
if start > total { |
|
|
start = total |
|
|
} |
|
|
if end > total { |
|
|
end = total |
|
|
} |
|
|
items := filteredFingerprints[start:end] |
|
|
|
|
|
|
|
|
c.JSON(http.StatusOK, gin.H{ |
|
|
"status": 0, |
|
|
"message": "success", |
|
|
"data": gin.H{ |
|
|
"total": total, |
|
|
"page": page, |
|
|
"size": size, |
|
|
"items": items, |
|
|
}, |
|
|
}) |
|
|
} |
|
|
|
|
|
|
|
|
func HandleCreateFingerprint(c *gin.Context) { |
|
|
|
|
|
type FingerprintUploadRequest struct { |
|
|
FileContent string `json:"file_content" binding:"required"` |
|
|
} |
|
|
var req FingerprintUploadRequest |
|
|
if err := c.ShouldBindJSON(&req); err != nil { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "参数解析失败"}) |
|
|
return |
|
|
} |
|
|
|
|
|
|
|
|
var fp parser.FingerPrint |
|
|
if err := yaml.Unmarshal([]byte(req.FileContent), &fp); err != nil { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "YAML解析失败: " + err.Error()}) |
|
|
return |
|
|
} |
|
|
if fp.Info.Name == "" { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "指纹名称不能为空"}) |
|
|
return |
|
|
} |
|
|
|
|
|
if _, err := parser.InitFingerPrintFromData([]byte(req.FileContent)); err != nil { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "指纹内容校验失败: " + err.Error()}) |
|
|
return |
|
|
} |
|
|
|
|
|
if !isValidName(fp.Info.Name) { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "指纹名称非法"}) |
|
|
return |
|
|
} |
|
|
yamlPath := filepath.Join("data/fingerprints", fp.Info.Name+".yaml") |
|
|
if _, err := os.Stat(yamlPath); err == nil { |
|
|
c.JSON(http.StatusConflict, gin.H{"status": 1, "message": "指纹已存在"}) |
|
|
return |
|
|
} |
|
|
|
|
|
if err := os.WriteFile(yamlPath, []byte(req.FileContent), 0644); err != nil { |
|
|
c.JSON(http.StatusInternalServerError, gin.H{"status": 1, "message": "文件写入失败: " + err.Error()}) |
|
|
return |
|
|
} |
|
|
|
|
|
c.JSON(http.StatusOK, gin.H{"status": 0, "message": "创建指纹成功"}) |
|
|
} |
|
|
|
|
|
|
|
|
type BatchDeleteRequest struct { |
|
|
Name []string `json:"name"` |
|
|
} |
|
|
|
|
|
func HandleDeleteFingerprint(c *gin.Context) { |
|
|
var req BatchDeleteRequest |
|
|
if err := c.ShouldBindJSON(&req); err != nil || len(req.Name) == 0 { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "参数错误", "data": nil}) |
|
|
return |
|
|
} |
|
|
|
|
|
var deleted []string |
|
|
var notFound []string |
|
|
var invalid []string |
|
|
|
|
|
for _, name := range req.Name { |
|
|
|
|
|
if !isValidName(name) { |
|
|
invalid = append(invalid, name) |
|
|
continue |
|
|
} |
|
|
yamlPath := filepath.Join("data/fingerprints", name+".yaml") |
|
|
if _, err := os.Stat(yamlPath); os.IsNotExist(err) { |
|
|
notFound = append(notFound, name) |
|
|
continue |
|
|
} |
|
|
if err := os.Remove(yamlPath); err == nil { |
|
|
deleted = append(deleted, name) |
|
|
} |
|
|
} |
|
|
|
|
|
msg := "删除完成" |
|
|
if len(notFound) > 0 { |
|
|
msg += ",部分指纹未找到: " + strings.Join(notFound, ", ") |
|
|
} |
|
|
if len(invalid) > 0 { |
|
|
msg += ",部分指纹名称非法: " + strings.Join(invalid, ", ") |
|
|
} |
|
|
|
|
|
c.JSON(http.StatusOK, gin.H{ |
|
|
"status": 0, |
|
|
"message": msg, |
|
|
"data": gin.H{ |
|
|
"deleted": deleted, |
|
|
"notFound": notFound, |
|
|
}, |
|
|
}) |
|
|
} |
|
|
|
|
|
|
|
|
func HandleEditFingerprint(c *gin.Context) { |
|
|
|
|
|
oldName := c.Param("name") |
|
|
if oldName == "" { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "指纹名称不能为空"}) |
|
|
return |
|
|
} |
|
|
|
|
|
type FingerprintUploadRequest struct { |
|
|
FileContent string `json:"file_content" binding:"required"` |
|
|
} |
|
|
var req FingerprintUploadRequest |
|
|
if err := c.ShouldBindJSON(&req); err != nil { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "参数解析失败"}) |
|
|
return |
|
|
} |
|
|
|
|
|
var fp parser.FingerPrint |
|
|
if err := yaml.Unmarshal([]byte(req.FileContent), &fp); err != nil { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "YAML解析失败: " + err.Error()}) |
|
|
return |
|
|
} |
|
|
if fp.Info.Name == "" { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "指纹名称不能为空"}) |
|
|
return |
|
|
} |
|
|
|
|
|
|
|
|
if _, err := parser.InitFingerPrintFromData([]byte(req.FileContent)); err != nil { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "指纹内容校验失败: " + err.Error()}) |
|
|
return |
|
|
} |
|
|
|
|
|
|
|
|
if !isValidName(oldName) || !isValidName(fp.Info.Name) { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "指纹名称非法"}) |
|
|
return |
|
|
} |
|
|
oldPath := filepath.Join("data/fingerprints", oldName+".yaml") |
|
|
if _, err := os.Stat(oldPath); os.IsNotExist(err) { |
|
|
c.JSON(http.StatusNotFound, gin.H{"status": 1, "message": "原指纹不存在"}) |
|
|
return |
|
|
} |
|
|
newPath := filepath.Join("data/fingerprints", fp.Info.Name+".yaml") |
|
|
|
|
|
|
|
|
if newPath != oldPath { |
|
|
if _, err := os.Stat(newPath); err == nil { |
|
|
c.JSON(http.StatusConflict, gin.H{"status": 1, "message": "新指纹名称已存在"}) |
|
|
return |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
if oldName != fp.Info.Name { |
|
|
_ = os.Remove(oldPath) |
|
|
} |
|
|
|
|
|
|
|
|
if err := os.WriteFile(newPath, []byte(req.FileContent), 0644); err != nil { |
|
|
c.JSON(http.StatusInternalServerError, gin.H{"status": 1, "message": "文件写入失败: " + err.Error()}) |
|
|
return |
|
|
} |
|
|
|
|
|
c.JSON(http.StatusOK, gin.H{"status": 0, "message": "修改指纹成功"}) |
|
|
} |
|
|
|
|
|
|
|
|
func HandleListVulnerabilities() gin.HandlerFunc { |
|
|
return func(c *gin.Context) { |
|
|
|
|
|
pageStr := c.DefaultQuery("page", "1") |
|
|
sizeStr := c.DefaultQuery("size", "20") |
|
|
query := strings.ToLower(c.DefaultQuery("q", "")) |
|
|
page, _ := strconv.Atoi(pageStr) |
|
|
size, _ := strconv.Atoi(sizeStr) |
|
|
if page < 1 { |
|
|
page = 1 |
|
|
} |
|
|
if size < 1 { |
|
|
size = 10 |
|
|
} |
|
|
|
|
|
engine := vulstruct.NewAdvisoryEngine() |
|
|
|
|
|
dir := "data/vuln" |
|
|
err := engine.LoadFromDirectory(dir) |
|
|
if err != nil { |
|
|
c.JSON(http.StatusInternalServerError, gin.H{"status": 1, "message": "加载漏洞库失败: " + err.Error()}) |
|
|
return |
|
|
} |
|
|
filteredVuls := make([]vulstruct.VersionVul, 0) |
|
|
if query == "" { |
|
|
filteredVuls = engine.GetAll() |
|
|
} else { |
|
|
for _, vul := range engine.GetAll() { |
|
|
if strings.Contains(strings.ToLower(vul.Info.CVEName), query) { |
|
|
filteredVuls = append(filteredVuls, vul) |
|
|
continue |
|
|
} |
|
|
if strings.Contains(strings.ToLower(vul.Info.Summary), query) { |
|
|
filteredVuls = append(filteredVuls, vul) |
|
|
continue |
|
|
} |
|
|
if strings.Contains(strings.ToLower(vul.Info.FingerPrintName), query) { |
|
|
filteredVuls = append(filteredVuls, vul) |
|
|
continue |
|
|
} |
|
|
if strings.Contains(strings.ToLower(vul.Info.Details), query) { |
|
|
filteredVuls = append(filteredVuls, vul) |
|
|
continue |
|
|
} |
|
|
for _, ref := range vul.References { |
|
|
if strings.Contains(strings.ToLower(ref), query) { |
|
|
filteredVuls = append(filteredVuls, vul) |
|
|
break |
|
|
} |
|
|
} |
|
|
} |
|
|
} |
|
|
|
|
|
total := len(filteredVuls) |
|
|
start := (page - 1) * size |
|
|
end := start + size |
|
|
if start > total { |
|
|
start = total |
|
|
} |
|
|
if end > total { |
|
|
end = total |
|
|
} |
|
|
items := filteredVuls[start:end] |
|
|
|
|
|
|
|
|
c.JSON(http.StatusOK, gin.H{ |
|
|
"status": 0, |
|
|
"message": "success", |
|
|
"data": gin.H{ |
|
|
"page": page, |
|
|
"size": size, |
|
|
"total": total, |
|
|
"items": items, |
|
|
}, |
|
|
}) |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
func HandleCreateVulnerability() gin.HandlerFunc { |
|
|
return func(c *gin.Context) { |
|
|
|
|
|
type VulnUploadRequest struct { |
|
|
FileContent string `json:"file_content" binding:"required"` |
|
|
} |
|
|
var req VulnUploadRequest |
|
|
if err := c.ShouldBindJSON(&req); err != nil { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "参数解析失败"}) |
|
|
return |
|
|
} |
|
|
|
|
|
|
|
|
var vul vulstruct.VersionVul |
|
|
if err := yaml.Unmarshal([]byte(req.FileContent), &vul); err != nil { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "YAML解析失败: " + err.Error()}) |
|
|
return |
|
|
} |
|
|
if vul.Info.CVEName == "" { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "CVE编号不能为空"}) |
|
|
return |
|
|
} |
|
|
if !isValidName(vul.Info.CVEName) { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "CVE编号非法"}) |
|
|
return |
|
|
} |
|
|
if vul.Info.FingerPrintName != "" && !isValidName(vul.Info.FingerPrintName) { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "指纹分类名称非法"}) |
|
|
return |
|
|
} |
|
|
|
|
|
|
|
|
_, err := vulstruct.ReadVersionVul([]byte(req.FileContent)) |
|
|
if err != nil { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "漏洞内容校验失败: " + err.Error()}) |
|
|
return |
|
|
} |
|
|
|
|
|
|
|
|
dir := "data/vuln" |
|
|
if vul.Info.FingerPrintName != "" { |
|
|
dir = filepath.Join(dir, vul.Info.FingerPrintName) |
|
|
} |
|
|
if err := os.MkdirAll(dir, 0755); err != nil { |
|
|
c.JSON(http.StatusInternalServerError, gin.H{"status": 1, "message": "创建目录失败: " + err.Error()}) |
|
|
return |
|
|
} |
|
|
fileName := strings.ToUpper(vul.Info.CVEName) + ".yaml" |
|
|
filePath := filepath.Join(dir, fileName) |
|
|
if _, err := os.Stat(filePath); err == nil { |
|
|
c.JSON(http.StatusConflict, gin.H{"status": 1, "message": "该CVE编号的漏洞已存在"}) |
|
|
return |
|
|
} |
|
|
data, err := yaml.Marshal(&vul) |
|
|
if err != nil { |
|
|
c.JSON(http.StatusInternalServerError, gin.H{"status": 1, "message": "YAML序列化失败: " + err.Error()}) |
|
|
return |
|
|
} |
|
|
if err := os.WriteFile(filePath, data, 0644); err != nil { |
|
|
c.JSON(http.StatusInternalServerError, gin.H{"status": 1, "message": "文件写入失败: " + err.Error()}) |
|
|
return |
|
|
} |
|
|
|
|
|
|
|
|
c.JSON(http.StatusOK, gin.H{"status": 0, "message": "创建漏洞库成功"}) |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
func HandleEditVulnerability(c *gin.Context) { |
|
|
|
|
|
oldCVE := c.Param("cve") |
|
|
if oldCVE == "" { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "CVE编号不能为空"}) |
|
|
return |
|
|
} |
|
|
if !isValidName(oldCVE) { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "原CVE编号非法"}) |
|
|
return |
|
|
} |
|
|
|
|
|
type VulnUploadRequest struct { |
|
|
FileContent string `json:"file_content" binding:"required"` |
|
|
} |
|
|
var req VulnUploadRequest |
|
|
if err := c.ShouldBindJSON(&req); err != nil { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "参数解析失败"}) |
|
|
return |
|
|
} |
|
|
|
|
|
var vul vulstruct.VersionVul |
|
|
if err := yaml.Unmarshal([]byte(req.FileContent), &vul); err != nil { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "YAML解析失败: " + err.Error()}) |
|
|
return |
|
|
} |
|
|
if vul.Info.CVEName == "" { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "CVE编号不能为空"}) |
|
|
return |
|
|
} |
|
|
if !isValidName(vul.Info.CVEName) { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "CVE编号非法"}) |
|
|
return |
|
|
} |
|
|
if vul.Info.FingerPrintName != "" && !isValidName(vul.Info.FingerPrintName) { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "指纹分类名称非法"}) |
|
|
return |
|
|
} |
|
|
|
|
|
_, err := vulstruct.ReadVersionVul([]byte(req.FileContent)) |
|
|
if err != nil { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "漏洞内容校验失败: " + err.Error()}) |
|
|
return |
|
|
} |
|
|
|
|
|
|
|
|
var oldPath string |
|
|
found := false |
|
|
baseDir := "data/vuln" |
|
|
_ = filepath.Walk(baseDir, func(path string, info os.FileInfo, err error) error { |
|
|
if err == nil && !info.IsDir() && strings.EqualFold(info.Name(), strings.ToUpper(oldCVE)+".yaml") { |
|
|
oldPath = path |
|
|
found = true |
|
|
return filepath.SkipDir |
|
|
} |
|
|
return nil |
|
|
}) |
|
|
if !found { |
|
|
c.JSON(http.StatusNotFound, gin.H{"status": 1, "message": "原漏洞不存在"}) |
|
|
return |
|
|
} |
|
|
|
|
|
|
|
|
newDir := "data/vuln" |
|
|
if vul.Info.FingerPrintName != "" { |
|
|
newDir = filepath.Join(newDir, vul.Info.FingerPrintName) |
|
|
} |
|
|
if err := os.MkdirAll(newDir, 0755); err != nil { |
|
|
c.JSON(http.StatusInternalServerError, gin.H{"status": 1, "message": "创建目录失败: " + err.Error()}) |
|
|
return |
|
|
} |
|
|
newPath := filepath.Join(newDir, strings.ToUpper(vul.Info.CVEName)+".yaml") |
|
|
|
|
|
|
|
|
if newPath != oldPath { |
|
|
if _, err := os.Stat(newPath); err == nil { |
|
|
c.JSON(http.StatusConflict, gin.H{"status": 1, "message": "新CVE编号的漏洞已存在"}) |
|
|
return |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
if err := os.Remove(oldPath); err != nil { |
|
|
c.JSON(http.StatusInternalServerError, gin.H{"status": 1, "message": "删除原文件失败: " + err.Error()}) |
|
|
return |
|
|
} |
|
|
|
|
|
|
|
|
data, err := yaml.Marshal(&vul) |
|
|
if err != nil { |
|
|
c.JSON(http.StatusInternalServerError, gin.H{"status": 1, "message": "YAML序列化失败: " + err.Error()}) |
|
|
return |
|
|
} |
|
|
if err := os.WriteFile(newPath, data, 0644); err != nil { |
|
|
c.JSON(http.StatusInternalServerError, gin.H{"status": 1, "message": "文件写入失败: " + err.Error()}) |
|
|
return |
|
|
} |
|
|
|
|
|
c.JSON(http.StatusOK, gin.H{"status": 0, "message": "修改漏洞成功"}) |
|
|
} |
|
|
|
|
|
|
|
|
type BatchDeleteVulnRequest struct { |
|
|
CVEs []string `json:"cves"` |
|
|
} |
|
|
|
|
|
func HandleBatchDeleteVulnerabilities(c *gin.Context) { |
|
|
var req BatchDeleteVulnRequest |
|
|
if err := c.ShouldBindJSON(&req); err != nil || len(req.CVEs) == 0 { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "参数解析失败或CVE列表为空"}) |
|
|
return |
|
|
} |
|
|
|
|
|
baseDir := "data/vuln" |
|
|
var notFound []string |
|
|
var failed []string |
|
|
|
|
|
for _, cve := range req.CVEs { |
|
|
if !isValidName(cve) { |
|
|
notFound = append(notFound, cve) |
|
|
continue |
|
|
} |
|
|
found := false |
|
|
_ = filepath.Walk(baseDir, func(path string, info os.FileInfo, err error) error { |
|
|
if err == nil && !info.IsDir() && strings.EqualFold(info.Name(), strings.ToUpper(cve)+".yaml") { |
|
|
|
|
|
if err := os.Remove(path); err != nil { |
|
|
failed = append(failed, cve) |
|
|
} |
|
|
found = true |
|
|
return filepath.SkipDir |
|
|
} |
|
|
return nil |
|
|
}) |
|
|
if !found { |
|
|
notFound = append(notFound, cve) |
|
|
} |
|
|
} |
|
|
|
|
|
if len(failed) > 0 { |
|
|
c.JSON(500, gin.H{"status": 1, "message": "部分删除失败", "failed": failed}) |
|
|
return |
|
|
} |
|
|
if len(notFound) > 0 { |
|
|
c.JSON(404, gin.H{"status": 1, "message": "部分CVE未找到", "not_found": notFound}) |
|
|
return |
|
|
} |
|
|
|
|
|
c.JSON(http.StatusOK, gin.H{"status": 0, "message": "批量删除成功"}) |
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
func HandleListEvaluations(c *gin.Context) { |
|
|
|
|
|
pageStr := c.DefaultQuery("page", "1") |
|
|
sizeStr := c.DefaultQuery("size", "20") |
|
|
detail := c.DefaultQuery("detail", "false") |
|
|
page, _ := strconv.Atoi(pageStr) |
|
|
size, _ := strconv.Atoi(sizeStr) |
|
|
if page < 1 { |
|
|
page = 1 |
|
|
} |
|
|
if size < 1 { |
|
|
size = 10 |
|
|
} |
|
|
|
|
|
|
|
|
nameQuery := strings.ToLower(c.DefaultQuery("q", "")) |
|
|
|
|
|
|
|
|
var allEvaluations []EvaluationDataset |
|
|
root := "data/eval" |
|
|
filepath.WalkDir(root, func(path string, d os.DirEntry, err error) error { |
|
|
if err != nil { |
|
|
return nil |
|
|
} |
|
|
if !d.IsDir() && strings.HasSuffix(d.Name(), ".json") { |
|
|
content, readErr := os.ReadFile(path) |
|
|
if readErr == nil { |
|
|
var eval EvaluationDataset |
|
|
err = json.Unmarshal(content, &eval) |
|
|
if err != nil { |
|
|
log.Error(path, err.Error()) |
|
|
return err |
|
|
} |
|
|
|
|
|
if detail != "true" { |
|
|
eval.Data = nil |
|
|
} |
|
|
allEvaluations = append(allEvaluations, eval) |
|
|
} |
|
|
} |
|
|
return nil |
|
|
}) |
|
|
|
|
|
|
|
|
var filteredEvaluations []EvaluationDataset |
|
|
if nameQuery == "" { |
|
|
filteredEvaluations = allEvaluations |
|
|
} else { |
|
|
for _, eval := range allEvaluations { |
|
|
if strings.Contains(strings.ToLower(eval.Name), nameQuery) { |
|
|
filteredEvaluations = append(filteredEvaluations, eval) |
|
|
continue |
|
|
} |
|
|
if strings.Contains(strings.ToLower(eval.Description), nameQuery) { |
|
|
filteredEvaluations = append(filteredEvaluations, eval) |
|
|
continue |
|
|
} |
|
|
if strings.Contains(strings.ToLower(eval.Author), nameQuery) { |
|
|
filteredEvaluations = append(filteredEvaluations, eval) |
|
|
continue |
|
|
} |
|
|
|
|
|
for _, tag := range eval.Tags { |
|
|
if strings.Contains(strings.ToLower(tag), nameQuery) { |
|
|
filteredEvaluations = append(filteredEvaluations, eval) |
|
|
break |
|
|
} |
|
|
} |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
sort.Slice(filteredEvaluations, func(i, j int) bool { |
|
|
return filteredEvaluations[i].Recommendation > filteredEvaluations[j].Recommendation |
|
|
}) |
|
|
|
|
|
|
|
|
total := len(filteredEvaluations) |
|
|
start := (page - 1) * size |
|
|
end := start + size |
|
|
if start > total { |
|
|
start = total |
|
|
} |
|
|
if end > total { |
|
|
end = total |
|
|
} |
|
|
items := filteredEvaluations[start:end] |
|
|
|
|
|
|
|
|
c.JSON(http.StatusOK, gin.H{ |
|
|
"status": 0, |
|
|
"message": "success", |
|
|
"data": gin.H{ |
|
|
"total": total, |
|
|
"page": page, |
|
|
"size": size, |
|
|
"items": items, |
|
|
}, |
|
|
}) |
|
|
} |
|
|
|
|
|
|
|
|
func HandleGetEvaluationDetail(c *gin.Context) { |
|
|
|
|
|
name := c.Param("name") |
|
|
if name == "" { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "评测集名称不能为空"}) |
|
|
return |
|
|
} |
|
|
|
|
|
|
|
|
if !isValidName(name) { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "评测集名称非法"}) |
|
|
return |
|
|
} |
|
|
|
|
|
|
|
|
var allEvaluations []EvaluationDataset |
|
|
root := "data/eval" |
|
|
filepath.WalkDir(root, func(path string, d os.DirEntry, err error) error { |
|
|
if err != nil { |
|
|
return nil |
|
|
} |
|
|
if !d.IsDir() && strings.HasSuffix(d.Name(), ".json") { |
|
|
content, readErr := os.ReadFile(path) |
|
|
if readErr == nil { |
|
|
var eval EvaluationDataset |
|
|
if parseErr := json.Unmarshal(content, &eval); parseErr == nil { |
|
|
allEvaluations = append(allEvaluations, eval) |
|
|
} |
|
|
} |
|
|
} |
|
|
return nil |
|
|
}) |
|
|
|
|
|
for _, eval := range allEvaluations { |
|
|
if eval.Name == name { |
|
|
c.JSON(http.StatusOK, gin.H{ |
|
|
"status": 0, |
|
|
"message": "success", |
|
|
"data": eval, |
|
|
}) |
|
|
return |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
c.JSON(http.StatusOK, gin.H{ |
|
|
"status": 0, |
|
|
"message": "success", |
|
|
"data": nil, |
|
|
}) |
|
|
} |
|
|
|
|
|
|
|
|
func HandleCreateEvaluation(c *gin.Context) { |
|
|
|
|
|
type EvaluationUploadRequest struct { |
|
|
FileContent string `json:"file_content" binding:"required"` |
|
|
} |
|
|
var req EvaluationUploadRequest |
|
|
if err := c.ShouldBindJSON(&req); err != nil { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "参数解析失败"}) |
|
|
return |
|
|
} |
|
|
|
|
|
|
|
|
var eval EvaluationDataset |
|
|
if err := json.Unmarshal([]byte(req.FileContent), &eval); err != nil { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "JSON解析失败: " + err.Error()}) |
|
|
return |
|
|
} |
|
|
if eval.Name == "" { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "评测集名称不能为空"}) |
|
|
return |
|
|
} |
|
|
|
|
|
|
|
|
if len(eval.Data) == 0 { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "评测数据不能为空"}) |
|
|
return |
|
|
} |
|
|
|
|
|
|
|
|
eval.Count = len(eval.Data) |
|
|
|
|
|
|
|
|
for i, item := range eval.Data { |
|
|
if item.Prompt == "" { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": fmt.Sprintf("第%d条数据的prompt不能为空", i+1)}) |
|
|
return |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
if !isValidName(eval.Name) { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "评测集名称非法,只允许字母、数字、下划线和横线"}) |
|
|
return |
|
|
} |
|
|
jsonPath := filepath.Join("data/eval", eval.Name+".json") |
|
|
if _, err := os.Stat(jsonPath); err == nil { |
|
|
c.JSON(http.StatusConflict, gin.H{"status": 1, "message": "评测集已存在"}) |
|
|
return |
|
|
} |
|
|
|
|
|
|
|
|
updatedContent, err := json.MarshalIndent(eval, "", " ") |
|
|
if err != nil { |
|
|
c.JSON(http.StatusInternalServerError, gin.H{"status": 1, "message": "JSON序列化失败: " + err.Error()}) |
|
|
return |
|
|
} |
|
|
|
|
|
if err := os.WriteFile(jsonPath, updatedContent, 0644); err != nil { |
|
|
c.JSON(http.StatusInternalServerError, gin.H{"status": 1, "message": "文件写入失败: " + err.Error()}) |
|
|
return |
|
|
} |
|
|
|
|
|
|
|
|
c.JSON(http.StatusOK, gin.H{"status": 0, "message": "创建评测集成功"}) |
|
|
} |
|
|
|
|
|
|
|
|
func HandleEditEvaluation(c *gin.Context) { |
|
|
|
|
|
oldName := c.Param("name") |
|
|
if oldName == "" { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "评测集名称不能为空"}) |
|
|
return |
|
|
} |
|
|
|
|
|
type EvaluationUploadRequest struct { |
|
|
FileContent string `json:"file_content" binding:"required"` |
|
|
} |
|
|
var req EvaluationUploadRequest |
|
|
if err := c.ShouldBindJSON(&req); err != nil { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "参数解析失败"}) |
|
|
return |
|
|
} |
|
|
|
|
|
|
|
|
var eval EvaluationDataset |
|
|
if err := json.Unmarshal([]byte(req.FileContent), &eval); err != nil { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "JSON解析失败: " + err.Error()}) |
|
|
return |
|
|
} |
|
|
if eval.Name == "" { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "评测集名称不能为空"}) |
|
|
return |
|
|
} |
|
|
|
|
|
|
|
|
if len(eval.Data) == 0 { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "评测数据不能为空"}) |
|
|
return |
|
|
} |
|
|
|
|
|
|
|
|
eval.Count = len(eval.Data) |
|
|
|
|
|
|
|
|
for i, item := range eval.Data { |
|
|
if item.Prompt == "" { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": fmt.Sprintf("第%d条数据的prompt不能为空", i+1)}) |
|
|
return |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
if !isValidName(oldName) || !isValidName(eval.Name) { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "评测集名称非法,只允许字母、数字、下划线和横线"}) |
|
|
return |
|
|
} |
|
|
oldPath := filepath.Join("data/eval", oldName+".json") |
|
|
if _, err := os.Stat(oldPath); os.IsNotExist(err) { |
|
|
c.JSON(http.StatusNotFound, gin.H{"status": 1, "message": "原评测集不存在"}) |
|
|
return |
|
|
} |
|
|
newPath := filepath.Join("data/eval", eval.Name+".json") |
|
|
|
|
|
|
|
|
if newPath != oldPath { |
|
|
if _, err := os.Stat(newPath); err == nil { |
|
|
c.JSON(http.StatusConflict, gin.H{"status": 1, "message": "新评测集名称已存在"}) |
|
|
return |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
if oldName != eval.Name { |
|
|
_ = os.Remove(oldPath) |
|
|
} |
|
|
|
|
|
|
|
|
updatedContent, err := json.MarshalIndent(eval, "", " ") |
|
|
if err != nil { |
|
|
c.JSON(http.StatusInternalServerError, gin.H{"status": 1, "message": "JSON序列化失败: " + err.Error()}) |
|
|
return |
|
|
} |
|
|
|
|
|
if err := os.WriteFile(newPath, updatedContent, 0644); err != nil { |
|
|
c.JSON(http.StatusInternalServerError, gin.H{"status": 1, "message": "文件写入失败: " + err.Error()}) |
|
|
return |
|
|
} |
|
|
|
|
|
c.JSON(http.StatusOK, gin.H{"status": 0, "message": "修改评测集成功"}) |
|
|
} |
|
|
|
|
|
|
|
|
type BatchDeleteEvaluationRequest struct { |
|
|
Names []string `json:"names"` |
|
|
} |
|
|
|
|
|
func HandleDeleteEvaluation(c *gin.Context) { |
|
|
var req BatchDeleteEvaluationRequest |
|
|
if err := c.ShouldBindJSON(&req); err != nil || len(req.Names) == 0 { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "参数错误", "data": nil}) |
|
|
return |
|
|
} |
|
|
|
|
|
for _, name := range req.Names { |
|
|
|
|
|
if !isValidName(name) { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "评测集名称非法,只允许字母、数字、下划线和横线"}) |
|
|
return |
|
|
} |
|
|
jsonPath := filepath.Join("data/eval", name+".json") |
|
|
if _, err := os.Stat(jsonPath); os.IsNotExist(err) { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "评测集不存在"}) |
|
|
return |
|
|
} |
|
|
if err := os.Remove(jsonPath); err != nil { |
|
|
c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "删除失败: " + err.Error()}) |
|
|
return |
|
|
} |
|
|
} |
|
|
c.JSON(http.StatusOK, gin.H{ |
|
|
"status": 0, |
|
|
"message": "ok", |
|
|
}) |
|
|
} |
|
|
|