From a1be3e9931377bda9248d36315462690a2ca8d9f Mon Sep 17 00:00:00 2001
From: Parth Malhotra <28601533+parthmalhotra@users.noreply.github.com>
Date: Tue, 26 Sep 2023 10:07:37 +0530
Subject: [PATCH] Update yaml2json.go - fix #8172

---
 .github/scripts/yaml2json.go | 135 +++++++++++++++++++----------------
 1 file changed, 75 insertions(+), 60 deletions(-)

diff --git a/.github/scripts/yaml2json.go b/.github/scripts/yaml2json.go
index 914cb5cd0a4..19d7b6a6a23 100644
--- a/.github/scripts/yaml2json.go
+++ b/.github/scripts/yaml2json.go
@@ -16,78 +16,93 @@ type Classification struct {
 }
 
 type Info struct {
-	Name        string         `yaml:"name"`
-	Severity    string         `yaml:"severity"`
-	Description string         `yaml:"description"`
+	Name           string         `yaml:"name"`
+	Severity       string         `yaml:"severity"`
+	Description    string         `yaml:"description"`
 	Classification Classification `yaml:"classification,omitempty"`
 }
 
 type Data struct {
-	ID          string `yaml:"id"`
-	Info        Info   `yaml:"info"`
-	FilePath    string `json:"file_path"`
+	ID       string `yaml:"id"`
+	Info     Info   `yaml:"info"`
+	FilePath string `json:"file_path"`
 }
 
 func main() {
 	if len(os.Args) != 3 {
-		fmt.Println("Usage: go run main.go <directory> <output_file>")
+		fmt.Println("Usage: go run main.go <directory1[,directory2,...]> <output_file>")
 		os.Exit(1)
 	}
 
-	directory := os.Args[1]
+	input := os.Args[1]
 	outputFile := os.Args[2]
+	var directories []string
+
+	// Check if the input contains a comma
+	if strings.Contains(input, ",") {
+		directories = strings.Split(input, ",")
+	} else {
+		directories = []string{input}
+	}
 
 	var data []Data
 
-	err := filepath.Walk(directory, func(path string, info os.FileInfo, err error) error {
-		if strings.HasSuffix(path, ".yaml") || strings.HasSuffix(path, ".yml") {
-			yamlFile, err := ioutil.ReadFile(path)
-                if err != nil {
-                    fmt.Printf("Error reading YAML file %s: %v\n", path, err)
-                    return err
-                }
-    
-                var d Data
-                err = yaml.Unmarshal(yamlFile, &d)
-                if err != nil {
-                    fmt.Printf("Error unmarshalling YAML file %s: %v\n", path, err)
-                    return err
-                }
-                if d.Info.Classification.CVSSScore == "" {
-                    d.Info.Classification.CVSSScore = "N/A"
-                }
-                if d.Info.Classification == (Classification{}) {
-                    d.Info.Classification.CVSSScore = "N/A"
-                }
-		fpath := strings.Replace(path, "/home/runner/work/nuclei-templates/nuclei-templates/", "", 1)
-                d.FilePath = fpath
-    
-                data = append(data, d)
-            }
-            return nil
-        })
-    
-        if err != nil {
-            fmt.Printf("Error reading directory: %v\n", err)
-            os.Exit(1)
-        }
-    
-        var jsonData []byte
-        for _, d := range data {
-            temp, err := json.Marshal(d)
-            if err != nil {
-                fmt.Printf("Error marshalling JSON: %v\n", err)
-                os.Exit(1)
-            }
-            jsonData = append(jsonData, temp...)
-            jsonData = append(jsonData, byte('\n'))
-        }
-        err = ioutil.WriteFile(outputFile, jsonData, 0644)
-        if err != nil {
-            fmt.Printf("Error writing JSON data to file: %v\n", err)
-            os.Exit(1)
-        }
-    
-        fmt.Println("JSON data written to", outputFile)
-    }
-    
+	for _, directory := range directories {
+		fmt.Println("Generating data for", directory)
+
+		err := filepath.Walk(directory, func(path string, info os.FileInfo, err error) error {
+			if err != nil {
+				fmt.Printf("Error accessing path %s: %v\n", path, err)
+				return err
+			}
+			if strings.HasSuffix(path, ".yaml") || strings.HasSuffix(path, ".yml") {
+				yamlFile, err := ioutil.ReadFile(path)
+				if err != nil {
+					fmt.Printf("Error reading YAML file %s: %v\n", path, err)
+					return err
+				}
+
+				var d Data
+				err = yaml.Unmarshal(yamlFile, &d)
+				if err != nil {
+					fmt.Printf("Error unmarshalling YAML file %s: %v\n", path, err)
+					return err
+				}
+				if d.Info.Classification.CVSSScore == "" {
+					d.Info.Classification.CVSSScore = "N/A"
+				}
+				if d.Info.Classification == (Classification{}) {
+					d.Info.Classification.CVSSScore = "N/A"
+				}
+				fpath := strings.Replace(path, "/home/runner/work/nuclei-templates/nuclei-templates/", "", 1)
+				d.FilePath = fpath
+
+				data = append(data, d)
+			}
+			return nil
+		})
+
+		if err != nil {
+			fmt.Printf("Error reading directory: %v\n", err)
+			os.Exit(1)
+		}
+	}
+
+	var jsonData []byte
+	for _, d := range data {
+		temp, err := json.Marshal(d)
+		if err != nil {
+			fmt.Printf("Error marshalling JSON: %v\n", err)
+			os.Exit(1)
+		}
+		jsonData = append(jsonData, temp...)
+		jsonData = append(jsonData, byte('\n'))
+	}
+	err := ioutil.WriteFile(outputFile, jsonData, 0644)
+	if err != nil {
+		fmt.Printf("Error writing JSON data to file: %v\n", err)
+		os.Exit(1)
+	}
+
+	fmt.Println("JSON data written to", outputFile)
+}