Update yaml2json.go - fix #8172

patch-1
Parth Malhotra 2023-09-26 10:07:37 +05:30 committed by GitHub
parent 6082650d6c
commit a1be3e9931
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 75 additions and 60 deletions

View File

@ -16,78 +16,93 @@ type Classification struct {
} }
type Info struct { type Info struct {
Name string `yaml:"name"` Name string `yaml:"name"`
Severity string `yaml:"severity"` Severity string `yaml:"severity"`
Description string `yaml:"description"` Description string `yaml:"description"`
Classification Classification `yaml:"classification,omitempty"` Classification Classification `yaml:"classification,omitempty"`
} }
type Data struct { type Data struct {
ID string `yaml:"id"` ID string `yaml:"id"`
Info Info `yaml:"info"` Info Info `yaml:"info"`
FilePath string `json:"file_path"` FilePath string `json:"file_path"`
} }
func main() { func main() {
if len(os.Args) != 3 { 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) os.Exit(1)
} }
directory := os.Args[1] input := os.Args[1]
outputFile := os.Args[2] 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 var data []Data
err := filepath.Walk(directory, func(path string, info os.FileInfo, err error) error { for _, directory := range directories {
if strings.HasSuffix(path, ".yaml") || strings.HasSuffix(path, ".yml") { fmt.Println("Generating data for", directory)
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 := filepath.Walk(directory, func(path string, info os.FileInfo, err error) error {
err = yaml.Unmarshal(yamlFile, &d) if err != nil {
if err != nil { fmt.Printf("Error accessing path %s: %v\n", path, err)
fmt.Printf("Error unmarshalling YAML file %s: %v\n", path, err) return err
return err }
} if strings.HasSuffix(path, ".yaml") || strings.HasSuffix(path, ".yml") {
if d.Info.Classification.CVSSScore == "" { yamlFile, err := ioutil.ReadFile(path)
d.Info.Classification.CVSSScore = "N/A" if err != nil {
} fmt.Printf("Error reading YAML file %s: %v\n", path, err)
if d.Info.Classification == (Classification{}) { return err
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) var d Data
} err = yaml.Unmarshal(yamlFile, &d)
return nil 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
if err != nil { data = append(data, d)
fmt.Printf("Error reading directory: %v\n", err) }
os.Exit(1) return nil
} })
var jsonData []byte if err != nil {
for _, d := range data { fmt.Printf("Error reading directory: %v\n", err)
temp, err := json.Marshal(d) os.Exit(1)
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) 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)
}