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 main YAML format validation tool
 20  // Used in CI pipelines to verify the format of YAML files under the data directory.
 21  // Usage: yamlcheck <path1> [path2] ...
 22  // Supports files or directories; directories are scanned recursively for .yaml/.yml files.
 23  package main
 24  
 25  import (
 26  	"fmt"
 27  	"os"
 28  	"path/filepath"
 29  	"strings"
 30  
 31  	"github.com/Tencent/AI-Infra-Guard/common/fingerprints/parser"
 32  	"github.com/Tencent/AI-Infra-Guard/pkg/vulstruct"
 33  )
 34  
 35  func main() {
 36  	if len(os.Args) < 2 {
 37  		fmt.Println("Usage: yamlcheck <path1> [path2] ...")
 38  		fmt.Println("  path can be a file or directory (directories are scanned recursively)")
 39  		fmt.Println("  Validates YAML files under data/fingerprints, data/vuln, data/vuln_en")
 40  		os.Exit(1)
 41  	}
 42  
 43  	// Collect all YAML files to check
 44  	var yamlFiles []string
 45  	for _, arg := range os.Args[1:] {
 46  		info, err := os.Stat(arg)
 47  		if err != nil {
 48  			fmt.Fprintf(os.Stderr, "❌ [path error] %s: %v\n", arg, err)
 49  			continue
 50  		}
 51  		if info.IsDir() {
 52  			files, err := walkYAMLFiles(arg)
 53  			if err != nil {
 54  				fmt.Fprintf(os.Stderr, "❌ [directory walk failed] %s: %v\n", arg, err)
 55  				continue
 56  			}
 57  			yamlFiles = append(yamlFiles, files...)
 58  		} else {
 59  			if isYAML(arg) {
 60  				yamlFiles = append(yamlFiles, arg)
 61  			}
 62  		}
 63  	}
 64  
 65  	hasError := false
 66  	checkedCount := 0
 67  	passCount := 0
 68  	failCount := 0
 69  
 70  	fmt.Println()
 71  	fmt.Println("╔══════════════════════════════════════════════╗")
 72  	fmt.Println("║         AIG YAML Validation Report          ║")
 73  	fmt.Println("╚══════════════════════════════════════════════╝")
 74  	fmt.Println()
 75  
 76  	for _, file := range yamlFiles {
 77  		category := categorizeFile(file)
 78  		if category == "" {
 79  			continue
 80  		}
 81  
 82  		checkedCount++
 83  		data, err := os.ReadFile(file)
 84  		if err != nil {
 85  			fmt.Fprintf(os.Stderr, "  ❌  FAIL  [read error]  %s\n      └─ %v\n", file, err)
 86  			hasError = true
 87  			failCount++
 88  			continue
 89  		}
 90  
 91  		switch category {
 92  		case "fingerprint":
 93  			_, err = parser.InitFingerPrintFromData(data)
 94  			if err != nil {
 95  				fmt.Fprintf(os.Stderr, "  ❌  FAIL  [fingerprint]  %s\n      └─ %v\n", file, err)
 96  				hasError = true
 97  				failCount++
 98  			} else {
 99  				fmt.Printf("  ✅  PASS  [fingerprint]  %s\n", file)
100  				passCount++
101  			}
102  		case "vuln":
103  			_, err = vulstruct.ReadVersionVul(data)
104  			if err != nil {
105  				fmt.Fprintf(os.Stderr, "  ❌  FAIL  [vuln rule]   %s\n      └─ %v\n", file, err)
106  				hasError = true
107  				failCount++
108  			} else {
109  				fmt.Printf("  ✅  PASS  [vuln rule]   %s\n", file)
110  				passCount++
111  			}
112  		}
113  	}
114  
115  	fmt.Println()
116  	fmt.Println("──────────────────────────────────────────────")
117  
118  	if checkedCount == 0 {
119  		fmt.Println("⚠️  No YAML files found to validate.")
120  		os.Exit(0)
121  	}
122  
123  	fmt.Printf("  Total checked : %d\n", checkedCount)
124  	fmt.Printf("  ✅ Passed     : %d\n", passCount)
125  	fmt.Printf("  ❌ Failed     : %d\n", failCount)
126  	fmt.Println("──────────────────────────────────────────────")
127  
128  	if hasError {
129  		fmt.Println()
130  		fmt.Fprintln(os.Stderr, "❌ Validation FAILED — please fix the errors listed above.")
131  		os.Exit(1)
132  	}
133  
134  	fmt.Println()
135  	fmt.Println("✅ All YAML files passed validation!")
136  }
137  
138  // walkYAMLFiles recursively walks a directory and returns all .yaml/.yml file paths.
139  func walkYAMLFiles(root string) ([]string, error) {
140  	var files []string
141  	err := filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
142  		if err != nil {
143  			return err
144  		}
145  		if !info.IsDir() && isYAML(path) {
146  			files = append(files, path)
147  		}
148  		return nil
149  	})
150  	return files, err
151  }
152  
153  // isYAML reports whether the file has a .yaml or .yml extension.
154  func isYAML(file string) bool {
155  	return strings.HasSuffix(file, ".yaml") || strings.HasSuffix(file, ".yml")
156  }
157  
158  // categorizeFile determines the category of a YAML file based on its path.
159  // Returns "fingerprint", "vuln", or "" (not a file that needs validation).
160  func categorizeFile(file string) string {
161  	normalized := filepath.ToSlash(file)
162  
163  	if strings.Contains(normalized, "data/fingerprints/") || strings.HasPrefix(normalized, "fingerprints/") {
164  		return "fingerprint"
165  	}
166  
167  	if strings.Contains(normalized, "data/vuln/") || strings.Contains(normalized, "data/vuln_en/") ||
168  		strings.HasPrefix(normalized, "vuln/") || strings.HasPrefix(normalized, "vuln_en/") {
169  		return "vuln"
170  	}
171  
172  	return ""
173  }