knowledge_api.go
1 // Copyright (c) 2024-2026 Tencent Zhuque Lab. All rights reserved. 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 // 15 // Requirement: Any integration or derivative work must explicitly attribute 16 // Tencent Zhuque Lab (https://github.com/Tencent/AI-Infra-Guard) in its 17 // documentation or user interface, as detailed in the NOTICE file. 18 19 package websocket 20 21 import ( 22 "encoding/json" 23 "fmt" 24 "net/http" 25 "os" 26 "path/filepath" 27 "regexp" 28 "sort" 29 "strconv" 30 "strings" 31 "time" 32 33 "trpc.group/trpc-go/trpc-go/log" 34 35 "github.com/Tencent/AI-Infra-Guard/common/fingerprints/parser" 36 "github.com/Tencent/AI-Infra-Guard/pkg/vulstruct" 37 "github.com/gin-gonic/gin" 38 "gopkg.in/yaml.v3" 39 ) 40 41 // 合法性校验 42 var validName = regexp.MustCompile(`^[a-zA-Z0-9\\._ -]+$`) 43 44 func isValidName(name string) bool { 45 if strings.Contains(name, "..") { 46 return false 47 } 48 return validName.MatchString(name) 49 } 50 51 // safeJoinPath joins base and elem using filepath.Join, then verifies the result 52 // is still within base to prevent path traversal attacks (defense in depth). 53 func safeJoinPath(base, elem string) (string, error) { 54 cleanBase := filepath.Clean(base) 55 joined := filepath.Clean(filepath.Join(cleanBase, elem)) 56 if joined != cleanBase && !strings.HasPrefix(joined, cleanBase+string(os.PathSeparator)) { 57 return "", fmt.Errorf("path traversal detected: %q is outside base %q", elem, base) 58 } 59 return joined, nil 60 } 61 62 // FingerprintWithTime 包含指纹数据和文件修改时间的结构体 63 type FingerprintWithTime struct { 64 FingerPrint parser.FingerPrint 65 ModTime time.Time 66 } 67 68 // 评测集数据结构定义 69 type EvaluationDataItem struct { 70 Prompt string `json:"prompt"` 71 } 72 73 type EvaluationDataset struct { 74 Name string `json:"name"` 75 Description string `json:"description"` 76 DescriptionZh string `json:"description_zh,omitempty"` 77 Author string `json:"author,omitempty"` 78 Source []string `json:"source,omitempty"` 79 Count int `json:"count"` 80 Default bool `json:"default"` 81 Tags []string `json:"tags,omitempty"` 82 Recommendation int `json:"recommendation,omitempty"` 83 Language string `json:"language,omitempty"` 84 Data []EvaluationDataItem `json:"data"` 85 } 86 87 // 获取指纹列表，支持分页和名字模糊 88 func HandleListFingerprints(c *gin.Context) { 89 // 1. 解析分页参数 90 pageStr := c.DefaultQuery("page", "1") 91 sizeStr := c.DefaultQuery("size", "20") 92 page, _ := strconv.Atoi(pageStr) 93 size, _ := strconv.Atoi(sizeStr) 94 if page < 1 { 95 page = 1 96 } 97 if size < 1 { 98 size = 10 99 } 100 101 // 2. 获取查询参数 102 nameQuery := strings.ToLower(c.DefaultQuery("q", "")) 103 104 // 3. 读取 data/fingerprints/ 下所有分类和YAML文件 105 var allFingerprintsWithTime []FingerprintWithTime 106 root := "data/fingerprints" 107 filepath.WalkDir(root, func(path string, d os.DirEntry, err error) error { 108 if err != nil { 109 return nil // 忽略错误 110 } 111 if !d.IsDir() && strings.HasSuffix(d.Name(), ".yaml") { 112 content, _ := os.ReadFile(path) 113 fp, err := parser.InitFingerPrintFromData(content) 114 if err == nil && fp != nil { 115 // 获取文件修改时间 116 fileInfo, _ := d.Info() 117 modTime := time.Time{} 118 if fileInfo != nil { 119 modTime = fileInfo.ModTime() 120 } 121 allFingerprintsWithTime = append(allFingerprintsWithTime, FingerprintWithTime{ 122 FingerPrint: *fp, 123 ModTime: modTime, 124 }) 125 } 126 } 127 return nil 128 }) 129 130 // 4. 条件过滤 131 var filteredFingerprintsWithTime []FingerprintWithTime 132 if nameQuery == "" { 133 filteredFingerprintsWithTime = allFingerprintsWithTime 134 } else { 135 for _, fpWithTime := range allFingerprintsWithTime { 136 fp := fpWithTime.FingerPrint 137 if strings.Contains(strings.ToLower(fp.Info.Name), nameQuery) { 138 filteredFingerprintsWithTime = append(filteredFingerprintsWithTime, fpWithTime) 139 continue 140 } 141 if strings.Contains(strings.ToLower(fp.Info.Desc), nameQuery) { 142 filteredFingerprintsWithTime = append(filteredFingerprintsWithTime, fpWithTime) 143 continue 144 } 145 if strings.Contains(strings.ToLower(fp.Info.Author), nameQuery) { 146 filteredFingerprintsWithTime = append(filteredFingerprintsWithTime, fpWithTime) 147 continue 148 } 149 } 150 } 151 152 // 5. 复合排序：如果有Recommendation，先按Recommendation降序；否则按文件修改时间降序 153 sort.Slice(filteredFingerprintsWithTime, func(i, j int) bool { 154 fpI := filteredFingerprintsWithTime[i].FingerPrint 155 fpJ := filteredFingerprintsWithTime[j].FingerPrint 156 157 // 如果两个都有Recommendation（>0），按Recommendation降序 158 if fpI.Info.Recommendation > 0 && fpJ.Info.Recommendation > 0 { 159 if fpI.Info.Recommendation != fpJ.Info.Recommendation { 160 return fpI.Info.Recommendation > fpJ.Info.Recommendation 161 } 162 // Recommendation相同时，按文件修改时间降序 163 return filteredFingerprintsWithTime[i].ModTime.After(filteredFingerprintsWithTime[j].ModTime) 164 } 165 166 // 如果只有一个有Recommendation，有Recommendation的排前面 167 if fpI.Info.Recommendation > 0 && fpJ.Info.Recommendation <= 0 { 168 return true 169 } 170 if fpI.Info.Recommendation <= 0 && fpJ.Info.Recommendation > 0 { 171 return false 172 } 173 174 // 如果两个都没有Recommendation，按文件修改时间降序 175 return filteredFingerprintsWithTime[i].ModTime.After(filteredFingerprintsWithTime[j].ModTime) 176 }) 177 178 // 提取指纹数据用于分页和返回 179 var filteredFingerprints []parser.FingerPrint 180 for _, fpWithTime := range filteredFingerprintsWithTime { 181 filteredFingerprints = append(filteredFingerprints, fpWithTime.FingerPrint) 182 } 183 184 // 6. 分页 185 total := len(filteredFingerprints) 186 start := (page - 1) * size 187 end := start + size 188 if start > total { 189 start = total 190 } 191 if end > total { 192 end = total 193 } 194 items := filteredFingerprints[start:end] 195 196 // 7. 返回 197 c.JSON(http.StatusOK, gin.H{ 198 "status": 0, 199 "message": "success", 200 "data": gin.H{ 201 "total": total, 202 "page": page, 203 "size": size, 204 "items": items, 205 }, 206 }) 207 } 208 209 // 创建指纹 210 func HandleCreateFingerprint(c *gin.Context) { 211 // 1. 解析请求体，获取file_content字段 212 type FingerprintUploadRequest struct { 213 FileContent string `json:"file_content" binding:"required"` 214 } 215 var req FingerprintUploadRequest 216 if err := c.ShouldBindJSON(&req); err != nil { 217 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "参数解析失败"}) 218 return 219 } 220 221 // 2. 解析YAML为parser.FingerPrint结构体 222 var fp parser.FingerPrint 223 if err := yaml.Unmarshal([]byte(req.FileContent), &fp); err != nil { 224 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "YAML解析失败: " + err.Error()}) 225 return 226 } 227 if fp.Info.Name == "" { 228 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "指纹名称不能为空"}) 229 return 230 } 231 232 if _, err := parser.InitFingerPrintFromData([]byte(req.FileContent)); err != nil { 233 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "指纹内容校验失败: " + err.Error()}) 234 return 235 } 236 237 if !isValidName(fp.Info.Name) { 238 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "指纹名称非法"}) 239 return 240 } 241 yamlPath, err := safeJoinPath("data/fingerprints", fp.Info.Name+".yaml") 242 if err != nil { 243 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "非法路径"}) 244 return 245 } 246 if _, err := os.Stat(yamlPath); err == nil { 247 c.JSON(http.StatusConflict, gin.H{"status": 1, "message": "指纹已存在"}) 248 return 249 } 250 251 if err := os.WriteFile(yamlPath, []byte(req.FileContent), 0644); err != nil { 252 c.JSON(http.StatusInternalServerError, gin.H{"status": 1, "message": "文件写入失败: " + err.Error()}) 253 return 254 } 255 256 c.JSON(http.StatusOK, gin.H{"status": 0, "message": "创建指纹成功"}) 257 } 258 259 // 批量删除指纹处理函数 260 type BatchDeleteRequest struct { 261 Name []string `json:"name"` 262 } 263 264 func HandleDeleteFingerprint(c *gin.Context) { 265 var req BatchDeleteRequest 266 if err := c.ShouldBindJSON(&req); err != nil || len(req.Name) == 0 { 267 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "参数错误", "data": nil}) 268 return 269 } 270 271 var deleted []string 272 var notFound []string 273 var invalid []string 274 275 for _, name := range req.Name { 276 // 使用已存在的合法性校验函数防止路径遍历攻击 277 if !isValidName(name) { 278 invalid = append(invalid, name) 279 continue 280 } 281 yamlPath, pathErr := safeJoinPath("data/fingerprints", name+".yaml") 282 if pathErr != nil { 283 invalid = append(invalid, name) 284 continue 285 } 286 if _, err := os.Stat(yamlPath); os.IsNotExist(err) { 287 notFound = append(notFound, name) 288 continue 289 } 290 if err := os.Remove(yamlPath); err == nil { 291 deleted = append(deleted, name) 292 } 293 } 294 295 msg := "删除完成" 296 if len(notFound) > 0 { 297 msg += "，部分指纹未找到: " + strings.Join(notFound, ", ") 298 } 299 if len(invalid) > 0 { 300 msg += "，部分指纹名称非法: " + strings.Join(invalid, ", ") 301 } 302 303 c.JSON(http.StatusOK, gin.H{ 304 "status": 0, 305 "message": msg, 306 "data": gin.H{ 307 "deleted": deleted, 308 "notFound": notFound, 309 }, 310 }) 311 } 312 313 // 编辑指纹处理函数 314 func HandleEditFingerprint(c *gin.Context) { 315 // 1. 获取原指纹名称 316 oldName := c.Param("name") 317 if oldName == "" { 318 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "指纹名称不能为空"}) 319 return 320 } 321 322 type FingerprintUploadRequest struct { 323 FileContent string `json:"file_content" binding:"required"` 324 } 325 var req FingerprintUploadRequest 326 if err := c.ShouldBindJSON(&req); err != nil { 327 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "参数解析失败"}) 328 return 329 } 330 // 2. 解析YAML为parser.FingerPrint结构体 331 var fp parser.FingerPrint 332 if err := yaml.Unmarshal([]byte(req.FileContent), &fp); err != nil { 333 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "YAML解析失败: " + err.Error()}) 334 return 335 } 336 if fp.Info.Name == "" { 337 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "指纹名称不能为空"}) 338 return 339 } 340 341 // 新增：用和读取时一致的解析逻辑做一次完整校验 342 if _, err := parser.InitFingerPrintFromData([]byte(req.FileContent)); err != nil { 343 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "指纹内容校验失败: " + err.Error()}) 344 return 345 } 346 347 // 3. 校验原文件是否存在 348 if !isValidName(oldName) || !isValidName(fp.Info.Name) { 349 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "指纹名称非法"}) 350 return 351 } 352 oldPath, err := safeJoinPath("data/fingerprints", oldName+".yaml") 353 if err != nil { 354 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "非法路径"}) 355 return 356 } 357 if _, err := os.Stat(oldPath); os.IsNotExist(err) { 358 c.JSON(http.StatusNotFound, gin.H{"status": 1, "message": "原指纹不存在"}) 359 return 360 } 361 newPath, err := safeJoinPath("data/fingerprints", fp.Info.Name+".yaml") 362 if err != nil { 363 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "非法路径"}) 364 return 365 } 366 367 // 4. 校验新文件名是否已存在（且不是原文件） 368 if newPath != oldPath { 369 if _, err := os.Stat(newPath); err == nil { 370 c.JSON(http.StatusConflict, gin.H{"status": 1, "message": "新指纹名称已存在"}) 371 return 372 } 373 } 374 375 // 5. 如果新旧文件名不同，删除原文件 376 if oldName != fp.Info.Name { 377 _ = os.Remove(oldPath) // 删除老文件 378 } 379 380 // 6. 写入新内容（新文件名） 381 if err := os.WriteFile(newPath, []byte(req.FileContent), 0644); err != nil { 382 c.JSON(http.StatusInternalServerError, gin.H{"status": 1, "message": "文件写入失败: " + err.Error()}) 383 return 384 } 385 386 c.JSON(http.StatusOK, gin.H{"status": 0, "message": "修改指纹成功"}) 387 } 388 389 // 漏洞库分页+条件查询接口 390 func HandleListVulnerabilities() gin.HandlerFunc { 391 return func(c *gin.Context) { 392 // 1. 解析分页和查询参数 393 pageStr := c.DefaultQuery("page", "1") 394 sizeStr := c.DefaultQuery("size", "20") 395 query := strings.ToLower(c.DefaultQuery("q", "")) 396 page, _ := strconv.Atoi(pageStr) 397 size, _ := strconv.Atoi(sizeStr) 398 if page < 1 { 399 page = 1 400 } 401 if size < 1 { 402 size = 10 403 } 404 405 engine := vulstruct.NewAdvisoryEngine() 406 // load from directory 407 dir := "data/vuln" 408 err := engine.LoadFromDirectory(dir) 409 if err != nil { 410 c.JSON(http.StatusInternalServerError, gin.H{"status": 1, "message": "加载漏洞库失败: " + err.Error()}) 411 return 412 } 413 filteredVuls := make([]vulstruct.VersionVul, 0) 414 if query == "" { 415 filteredVuls = engine.GetAll() 416 } else { 417 for _, vul := range engine.GetAll() { 418 if strings.Contains(strings.ToLower(vul.Info.CVEName), query) { 419 filteredVuls = append(filteredVuls, vul) 420 continue 421 } 422 if strings.Contains(strings.ToLower(vul.Info.Summary), query) { 423 filteredVuls = append(filteredVuls, vul) 424 continue 425 } 426 if strings.Contains(strings.ToLower(vul.Info.FingerPrintName), query) { 427 filteredVuls = append(filteredVuls, vul) 428 continue 429 } 430 if strings.Contains(strings.ToLower(vul.Info.Details), query) { 431 filteredVuls = append(filteredVuls, vul) 432 continue 433 } 434 for _, ref := range vul.References { 435 if strings.Contains(strings.ToLower(ref), query) { 436 filteredVuls = append(filteredVuls, vul) 437 break 438 } 439 } 440 } 441 } 442 // 5. 分页 443 total := len(filteredVuls) 444 start := (page - 1) * size 445 end := start + size 446 if start > total { 447 start = total 448 } 449 if end > total { 450 end = total 451 } 452 items := filteredVuls[start:end] 453 454 // 6. 返回 455 c.JSON(http.StatusOK, gin.H{ 456 "status": 0, 457 "message": "success", 458 "data": gin.H{ 459 "page": page, 460 "size": size, 461 "total": total, 462 "items": items, 463 }, 464 }) 465 } 466 } 467 468 // 添加漏洞信息（带严格校验） 469 func HandleCreateVulnerability() gin.HandlerFunc { 470 return func(c *gin.Context) { 471 // 1. 解析请求体，获取file_content 472 type VulnUploadRequest struct { 473 FileContent string `json:"file_content" binding:"required"` 474 } 475 var req VulnUploadRequest 476 if err := c.ShouldBindJSON(&req); err != nil { 477 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "参数解析失败"}) 478 return 479 } 480 481 // 2. 反序列化为vulstruct.VersionVul，校验CVE编号等必填字段 482 var vul vulstruct.VersionVul 483 if err := yaml.Unmarshal([]byte(req.FileContent), &vul); err != nil { 484 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "YAML解析失败: " + err.Error()}) 485 return 486 } 487 if vul.Info.CVEName == "" { 488 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "CVE编号不能为空"}) 489 return 490 } 491 if !isValidName(vul.Info.CVEName) { 492 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "CVE编号非法"}) 493 return 494 } 495 if vul.Info.FingerPrintName != "" && !isValidName(vul.Info.FingerPrintName) { 496 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "指纹分类名称非法"}) 497 return 498 } 499 500 // 4. 用vulstruct.NewAdvisoryEngine加载临时文件做完整业务校验 501 _, err := vulstruct.ReadVersionVul([]byte(req.FileContent)) 502 if err != nil { 503 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "漏洞内容校验失败: " + err.Error()}) 504 return 505 } 506 507 // 5. 校验通过后，正式写入到目标目录（如已存在则报冲突） 508 dir := "data/vuln" 509 if vul.Info.FingerPrintName != "" { 510 var pathErr error 511 dir, pathErr = safeJoinPath(dir, vul.Info.FingerPrintName) 512 if pathErr != nil { 513 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "非法路径"}) 514 return 515 } 516 } 517 if err := os.MkdirAll(dir, 0755); err != nil { 518 c.JSON(http.StatusInternalServerError, gin.H{"status": 1, "message": "创建目录失败: " + err.Error()}) 519 return 520 } 521 fileName := strings.ToUpper(vul.Info.CVEName) + ".yaml" 522 filePath, err := safeJoinPath(dir, fileName) 523 if err != nil { 524 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "非法路径"}) 525 return 526 } 527 if _, err := os.Stat(filePath); err == nil { 528 c.JSON(http.StatusConflict, gin.H{"status": 1, "message": "该CVE编号的漏洞已存在"}) 529 return 530 } 531 data, err := yaml.Marshal(&vul) 532 if err != nil { 533 c.JSON(http.StatusInternalServerError, gin.H{"status": 1, "message": "YAML序列化失败: " + err.Error()}) 534 return 535 } 536 if err := os.WriteFile(filePath, data, 0644); err != nil { 537 c.JSON(http.StatusInternalServerError, gin.H{"status": 1, "message": "文件写入失败: " + err.Error()}) 538 return 539 } 540 541 // 6. 返回结果 542 c.JSON(http.StatusOK, gin.H{"status": 0, "message": "创建漏洞库成功"}) 543 } 544 } 545 546 // 编辑漏洞处理函数 547 func HandleEditVulnerability(c *gin.Context) { 548 // 1. 获取原CVE编号 549 oldCVE := c.Param("cve") 550 if oldCVE == "" { 551 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "CVE编号不能为空"}) 552 return 553 } 554 if !isValidName(oldCVE) { 555 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "原CVE编号非法"}) 556 return 557 } 558 559 type VulnUploadRequest struct { 560 FileContent string `json:"file_content" binding:"required"` 561 } 562 var req VulnUploadRequest 563 if err := c.ShouldBindJSON(&req); err != nil { 564 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "参数解析失败"}) 565 return 566 } 567 // 2. 反序列化为vulstruct.VersionVul，校验CVE编号等必填字段 568 var vul vulstruct.VersionVul 569 if err := yaml.Unmarshal([]byte(req.FileContent), &vul); err != nil { 570 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "YAML解析失败: " + err.Error()}) 571 return 572 } 573 if vul.Info.CVEName == "" { 574 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "CVE编号不能为空"}) 575 return 576 } 577 if !isValidName(vul.Info.CVEName) { 578 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "CVE编号非法"}) 579 return 580 } 581 if vul.Info.FingerPrintName != "" && !isValidName(vul.Info.FingerPrintName) { 582 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "指纹分类名称非法"}) 583 return 584 } 585 // 4. 用vulstruct.NewAdvisoryEngine加载临时文件做完整业务校验 586 _, err := vulstruct.ReadVersionVul([]byte(req.FileContent)) 587 if err != nil { 588 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "漏洞内容校验失败: " + err.Error()}) 589 return 590 } 591 592 // 5. 在所有分类目录下查找原文件 593 var oldPath string 594 found := false 595 baseDir := "data/vuln" 596 _ = filepath.Walk(baseDir, func(path string, info os.FileInfo, err error) error { 597 if err == nil && !info.IsDir() && strings.EqualFold(info.Name(), strings.ToUpper(oldCVE)+".yaml") { 598 oldPath = path 599 found = true 600 return filepath.SkipDir // 找到就停止遍历 601 } 602 return nil 603 }) 604 if !found { 605 c.JSON(http.StatusNotFound, gin.H{"status": 1, "message": "原漏洞不存在"}) 606 return 607 } 608 609 // 6. 生成新文件路径 610 newDir := "data/vuln" 611 if vul.Info.FingerPrintName != "" { 612 newDir, err = safeJoinPath(newDir, vul.Info.FingerPrintName) 613 if err != nil { 614 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "非法路径"}) 615 return 616 } 617 } 618 if err := os.MkdirAll(newDir, 0755); err != nil { 619 c.JSON(http.StatusInternalServerError, gin.H{"status": 1, "message": "创建目录失败: " + err.Error()}) 620 return 621 } 622 newPath, err := safeJoinPath(newDir, strings.ToUpper(vul.Info.CVEName)+".yaml") 623 if err != nil { 624 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "非法路径"}) 625 return 626 } 627 628 // 7. 校验新文件名是否已存在（且不是原文件） 629 if newPath != oldPath { 630 if _, err := os.Stat(newPath); err == nil { 631 c.JSON(http.StatusConflict, gin.H{"status": 1, "message": "新CVE编号的漏洞已存在"}) 632 return 633 } 634 } 635 636 // 8. 删除原文件 637 if err := os.Remove(oldPath); err != nil { 638 c.JSON(http.StatusInternalServerError, gin.H{"status": 1, "message": "删除原文件失败: " + err.Error()}) 639 return 640 } 641 642 // 9. 写入新内容（新文件名/新目录） 643 data, err := yaml.Marshal(&vul) 644 if err != nil { 645 c.JSON(http.StatusInternalServerError, gin.H{"status": 1, "message": "YAML序列化失败: " + err.Error()}) 646 return 647 } 648 if err := os.WriteFile(newPath, data, 0644); err != nil { 649 c.JSON(http.StatusInternalServerError, gin.H{"status": 1, "message": "文件写入失败: " + err.Error()}) 650 return 651 } 652 653 c.JSON(http.StatusOK, gin.H{"status": 0, "message": "修改漏洞成功"}) 654 } 655 656 // 批量删除漏洞处理函数 657 type BatchDeleteVulnRequest struct { 658 CVEs []string `json:"cves"` 659 } 660 661 func HandleBatchDeleteVulnerabilities(c *gin.Context) { 662 var req BatchDeleteVulnRequest 663 if err := c.ShouldBindJSON(&req); err != nil || len(req.CVEs) == 0 { 664 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "参数解析失败或CVE列表为空"}) 665 return 666 } 667 668 baseDir := "data/vuln" 669 var notFound []string 670 var failed []string 671 672 for _, cve := range req.CVEs { 673 if !isValidName(cve) { 674 notFound = append(notFound, cve) 675 continue 676 } 677 found := false 678 _ = filepath.Walk(baseDir, func(path string, info os.FileInfo, err error) error { 679 if err == nil && !info.IsDir() && strings.EqualFold(info.Name(), strings.ToUpper(cve)+".yaml") { 680 // 找到就删除 681 if err := os.Remove(path); err != nil { 682 failed = append(failed, cve) 683 } 684 found = true 685 return filepath.SkipDir 686 } 687 return nil 688 }) 689 if !found { 690 notFound = append(notFound, cve) 691 } 692 } 693 694 if len(failed) > 0 { 695 c.JSON(500, gin.H{"status": 1, "message": "部分删除失败", "failed": failed}) 696 return 697 } 698 if len(notFound) > 0 { 699 c.JSON(404, gin.H{"status": 1, "message": "部分CVE未找到", "not_found": notFound}) 700 return 701 } 702 703 c.JSON(http.StatusOK, gin.H{"status": 0, "message": "批量删除成功"}) 704 } 705 706 // ================== 评测集管理接口 ================== 707 708 // 获取评测集列表，支持分页和名字模糊搜索 709 func HandleListEvaluations(c *gin.Context) { 710 // 1. 解析分页参数 711 pageStr := c.DefaultQuery("page", "1") 712 sizeStr := c.DefaultQuery("size", "20") 713 detail := c.DefaultQuery("detail", "false") 714 page, _ := strconv.Atoi(pageStr) 715 size, _ := strconv.Atoi(sizeStr) 716 if page < 1 { 717 page = 1 718 } 719 if size < 1 { 720 size = 10 721 } 722 723 // 2. 获取查询参数 724 nameQuery := strings.ToLower(c.DefaultQuery("q", "")) 725 726 // 3. 读取 data/eval/ 下所有JSON文件 727 var allEvaluations []EvaluationDataset 728 root := "data/eval" 729 filepath.WalkDir(root, func(path string, d os.DirEntry, err error) error { 730 if err != nil { 731 return nil // 忽略错误 732 } 733 if !d.IsDir() && strings.HasSuffix(d.Name(), ".json") { 734 content, readErr := os.ReadFile(path) 735 if readErr == nil { 736 var eval EvaluationDataset 737 err = json.Unmarshal(content, &eval) 738 if err != nil { 739 log.Error(path, err.Error()) 740 return err 741 } 742 // 转换为摘要格式（不包含data字段） 743 if detail != "true" { 744 eval.Data = nil 745 } 746 allEvaluations = append(allEvaluations, eval) 747 } 748 } 749 return nil 750 }) 751 752 // 4. 条件过滤 753 var filteredEvaluations []EvaluationDataset 754 if nameQuery == "" { 755 filteredEvaluations = allEvaluations 756 } else { 757 for _, eval := range allEvaluations { 758 if strings.Contains(strings.ToLower(eval.Name), nameQuery) { 759 filteredEvaluations = append(filteredEvaluations, eval) 760 continue 761 } 762 if strings.Contains(strings.ToLower(eval.Description), nameQuery) { 763 filteredEvaluations = append(filteredEvaluations, eval) 764 continue 765 } 766 if strings.Contains(strings.ToLower(eval.Author), nameQuery) { 767 filteredEvaluations = append(filteredEvaluations, eval) 768 continue 769 } 770 // 搜索标签 771 for _, tag := range eval.Tags { 772 if strings.Contains(strings.ToLower(tag), nameQuery) { 773 filteredEvaluations = append(filteredEvaluations, eval) 774 break 775 } 776 } 777 } 778 } 779 780 // 5. 默认排序：按照Recommendation降序排列 781 sort.Slice(filteredEvaluations, func(i, j int) bool { 782 return filteredEvaluations[i].Recommendation > filteredEvaluations[j].Recommendation 783 }) 784 785 // 6. 分页 786 total := len(filteredEvaluations) 787 start := (page - 1) * size 788 end := start + size 789 if start > total { 790 start = total 791 } 792 if end > total { 793 end = total 794 } 795 items := filteredEvaluations[start:end] 796 797 // 7. 返回 798 c.JSON(http.StatusOK, gin.H{ 799 "status": 0, 800 "message": "success", 801 "data": gin.H{ 802 "total": total, 803 "page": page, 804 "size": size, 805 "items": items, 806 }, 807 }) 808 } 809 810 // 获取评测集详情，返回包含data的完整信息 811 func HandleGetEvaluationDetail(c *gin.Context) { 812 // 1. 获取评测集名称 813 name := c.Param("name") 814 if name == "" { 815 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "评测集名称不能为空"}) 816 return 817 } 818 819 // 2. 验证名称合法性 820 if !isValidName(name) { 821 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "评测集名称非法"}) 822 return 823 } 824 825 // 3. 读取评测集文件 826 var allEvaluations []EvaluationDataset 827 root := "data/eval" 828 filepath.WalkDir(root, func(path string, d os.DirEntry, err error) error { 829 if err != nil { 830 return nil // 忽略错误 831 } 832 if !d.IsDir() && strings.HasSuffix(d.Name(), ".json") { 833 content, readErr := os.ReadFile(path) 834 if readErr == nil { 835 var eval EvaluationDataset 836 if parseErr := json.Unmarshal(content, &eval); parseErr == nil { 837 allEvaluations = append(allEvaluations, eval) 838 } 839 } 840 } 841 return nil 842 }) 843 844 for _, eval := range allEvaluations { 845 if eval.Name == name { 846 c.JSON(http.StatusOK, gin.H{ 847 "status": 0, 848 "message": "success", 849 "data": eval, 850 }) 851 return 852 } 853 } 854 855 // 5. 返回完整的评测集信息（包含data字段） 856 c.JSON(http.StatusOK, gin.H{ 857 "status": 0, 858 "message": "success", 859 "data": nil, 860 }) 861 } 862 863 // 创建评测集 864 func HandleCreateEvaluation(c *gin.Context) { 865 // 1. 解析请求体，获取file_content字段 866 type EvaluationUploadRequest struct { 867 FileContent string `json:"file_content" binding:"required"` 868 } 869 var req EvaluationUploadRequest 870 if err := c.ShouldBindJSON(&req); err != nil { 871 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "参数解析失败"}) 872 return 873 } 874 875 // 2. 解析JSON为EvaluationDataset结构体 876 var eval EvaluationDataset 877 if err := json.Unmarshal([]byte(req.FileContent), &eval); err != nil { 878 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "JSON解析失败: " + err.Error()}) 879 return 880 } 881 if eval.Name == "" { 882 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "评测集名称不能为空"}) 883 return 884 } 885 886 // 3. 验证数据完整性 887 if len(eval.Data) == 0 { 888 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "评测数据不能为空"}) 889 return 890 } 891 892 // 更新count字段为实际数据条数 893 eval.Count = len(eval.Data) 894 895 // 验证数据项 896 for i, item := range eval.Data { 897 if item.Prompt == "" { 898 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": fmt.Sprintf("第%d条数据的prompt不能为空", i+1)}) 899 return 900 } 901 } 902 903 // 4. 检查评测集名称是否已存在 904 if !isValidName(eval.Name) { 905 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "评测集名称非法，只允许字母、数字、下划线和横线"}) 906 return 907 } 908 jsonPath, err := safeJoinPath("data/eval", eval.Name+".json") 909 if err != nil { 910 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "非法路径"}) 911 return 912 } 913 if _, err := os.Stat(jsonPath); err == nil { 914 c.JSON(http.StatusConflict, gin.H{"status": 1, "message": "评测集已存在"}) 915 return 916 } 917 918 // 5. 序列化并写入JSON文件 919 updatedContent, err := json.MarshalIndent(eval, "", " ") 920 if err != nil { 921 c.JSON(http.StatusInternalServerError, gin.H{"status": 1, "message": "JSON序列化失败: " + err.Error()}) 922 return 923 } 924 925 if err := os.WriteFile(jsonPath, updatedContent, 0644); err != nil { 926 c.JSON(http.StatusInternalServerError, gin.H{"status": 1, "message": "文件写入失败: " + err.Error()}) 927 return 928 } 929 930 // 6. 返回精简响应 931 c.JSON(http.StatusOK, gin.H{"status": 0, "message": "创建评测集成功"}) 932 } 933 934 // 编辑评测集处理函数 935 func HandleEditEvaluation(c *gin.Context) { 936 // 1. 获取原评测集名称 937 oldName := c.Param("name") 938 if oldName == "" { 939 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "评测集名称不能为空"}) 940 return 941 } 942 943 type EvaluationUploadRequest struct { 944 FileContent string `json:"file_content" binding:"required"` 945 } 946 var req EvaluationUploadRequest 947 if err := c.ShouldBindJSON(&req); err != nil { 948 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "参数解析失败"}) 949 return 950 } 951 952 // 2. 解析JSON为EvaluationDataset结构体 953 var eval EvaluationDataset 954 if err := json.Unmarshal([]byte(req.FileContent), &eval); err != nil { 955 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "JSON解析失败: " + err.Error()}) 956 return 957 } 958 if eval.Name == "" { 959 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "评测集名称不能为空"}) 960 return 961 } 962 963 // 验证数据完整性 964 if len(eval.Data) == 0 { 965 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "评测数据不能为空"}) 966 return 967 } 968 969 // 更新count字段为实际数据条数 970 eval.Count = len(eval.Data) 971 972 // 验证数据项 973 for i, item := range eval.Data { 974 if item.Prompt == "" { 975 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": fmt.Sprintf("第%d条数据的prompt不能为空", i+1)}) 976 return 977 } 978 } 979 980 // 3. 校验原文件是否存在 981 if !isValidName(oldName) || !isValidName(eval.Name) { 982 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "评测集名称非法，只允许字母、数字、下划线和横线"}) 983 return 984 } 985 oldPath, err := safeJoinPath("data/eval", oldName+".json") 986 if err != nil { 987 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "非法路径"}) 988 return 989 } 990 if _, err := os.Stat(oldPath); os.IsNotExist(err) { 991 c.JSON(http.StatusNotFound, gin.H{"status": 1, "message": "原评测集不存在"}) 992 return 993 } 994 newPath, err := safeJoinPath("data/eval", eval.Name+".json") 995 if err != nil { 996 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "非法路径"}) 997 return 998 } 999 1000 // 4. 校验新文件名是否已存在（且不是原文件） 1001 if newPath != oldPath { 1002 if _, err := os.Stat(newPath); err == nil { 1003 c.JSON(http.StatusConflict, gin.H{"status": 1, "message": "新评测集名称已存在"}) 1004 return 1005 } 1006 } 1007 1008 // 5. 如果新旧文件名不同，删除原文件 1009 if oldName != eval.Name { 1010 _ = os.Remove(oldPath) // 删除老文件 1011 } 1012 1013 // 6. 序列化并写入新内容（新文件名） 1014 updatedContent, err := json.MarshalIndent(eval, "", " ") 1015 if err != nil { 1016 c.JSON(http.StatusInternalServerError, gin.H{"status": 1, "message": "JSON序列化失败: " + err.Error()}) 1017 return 1018 } 1019 1020 if err := os.WriteFile(newPath, updatedContent, 0644); err != nil { 1021 c.JSON(http.StatusInternalServerError, gin.H{"status": 1, "message": "文件写入失败: " + err.Error()}) 1022 return 1023 } 1024 1025 c.JSON(http.StatusOK, gin.H{"status": 0, "message": "修改评测集成功"}) 1026 } 1027 1028 // 批量删除评测集处理函数 1029 type BatchDeleteEvaluationRequest struct { 1030 Names []string `json:"names"` 1031 } 1032 1033 func HandleDeleteEvaluation(c *gin.Context) { 1034 var req BatchDeleteEvaluationRequest 1035 if err := c.ShouldBindJSON(&req); err != nil || len(req.Names) == 0 { 1036 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "参数错误", "data": nil}) 1037 return 1038 } 1039 1040 for _, name := range req.Names { 1041 // 使用已存在的合法性校验函数防止路径遍历攻击 1042 if !isValidName(name) { 1043 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "评测集名称非法，只允许字母、数字、下划线和横线"}) 1044 return 1045 } 1046 jsonPath, pathErr := safeJoinPath("data/eval", name+".json") 1047 if pathErr != nil { 1048 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "非法路径"}) 1049 return 1050 } 1051 if _, err := os.Stat(jsonPath); os.IsNotExist(err) { 1052 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "评测集不存在"}) 1053 return 1054 } 1055 if err := os.Remove(jsonPath); err != nil { 1056 c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "删除失败: " + err.Error()}) 1057 return 1058 } 1059 } 1060 c.JSON(http.StatusOK, gin.H{ 1061 "status": 0, 1062 "message": "ok", 1063 }) 1064 }