Refactor docgen: improve error handling and file operations (#6103)

* Refactor docgen: improve error handling and file operations

Moved the argument check to the start of main() so we don’t do any unnecessary work if the arguments aren’t provided. Also added proper error handling for encoder.Encode(jsonschemaData), since it was failing silently before.

Switched os.WriteFile to os.Create for better file handling and error reporting. To clean things up, I added a writeToFile() function to remove duplicate code.

For replacing schema paths, I used ReplaceAllString() instead of looping through matches—it’s cleaner and a bit more efficient. Also renamed r to t in the Namer function to make things clearer.

* Fix a Lint error

- Remove importing string

Author: Mehran-Seifalinia

cmd/docgen/docgen.go

@@ -6,53 +6,58 @@ import (
 	"os"
 	"reflect"
 	"regexp"
-	"strings"
 
 	"github.com/invopop/jsonschema"
-
 	"github.com/projectdiscovery/nuclei/v3/pkg/templates"
 	"github.com/projectdiscovery/nuclei/v3/pkg/utils/json"
 )
 
 var pathRegex = regexp.MustCompile(`github\.com/projectdiscovery/nuclei/v3/(?:internal|pkg)/(?:.*/)?([A-Za-z.]+)`)
 
-func main() {
-	// Generate yaml syntax documentation
-	data, err := templates.GetTemplateDoc().Encode()
+func writeToFile(filename string, data []byte) {
+	file, err := os.Create(filename)
 	if err != nil {
-		log.Fatalf("Could not encode docs: %s\n", err)
+		log.Fatalf("Could not create file %s: %s\n", filename, err)
 	}
+	defer file.Close()
 
+	_, err = file.Write(data)
+	if err != nil {
+		log.Fatalf("Could not write to file %s: %s\n", filename, err)
+	}
+}
+
+func main() {
 	if len(os.Args) < 3 {
 		log.Fatalf("syntax: %s md-docs-file jsonschema-file\n", os.Args[0])
 	}
 
-	err = os.WriteFile(os.Args[1], data, 0644)
+	// Generate YAML documentation
+	data, err := templates.GetTemplateDoc().Encode()
 	if err != nil {
-		log.Fatalf("Could not write docs: %s\n", err)
+		log.Fatalf("Could not encode docs: %s\n", err)
 	}
-
-	// Generate jsonschema
-	r := &jsonschema.Reflector{}
-	r.Namer = func(r reflect.Type) string {
-		if r.Kind() == reflect.Slice {
-			return ""
-		}
-		return r.String()
+	writeToFile(os.Args[1], data)
+
+	// Generate JSON Schema
+	r := &jsonschema.Reflector{
+		Namer: func(t reflect.Type) string {
+			if t.Kind() == reflect.Slice {
+				return ""
+			}
+			return t.String()
+		},
 	}
+
 	jsonschemaData := r.Reflect(&templates.Template{})
 
 	var buf bytes.Buffer
 	encoder := json.NewEncoder(&buf)
 	encoder.SetIndent("", "  ")
-	_ = encoder.Encode(jsonschemaData)
-
-	schema := buf.String()
-	for _, match := range pathRegex.FindAllStringSubmatch(schema, -1) {
-		schema = strings.ReplaceAll(schema, match[0], match[1])
-	}
-	err = os.WriteFile(os.Args[2], []byte(schema), 0644)
-	if err != nil {
-		log.Fatalf("Could not write jsonschema: %s\n", err)
+	if err := encoder.Encode(jsonschemaData); err != nil {
+		log.Fatalf("Could not encode JSON schema: %s\n", err)
 	}
+
+	schema := pathRegex.ReplaceAllString(buf.String(), "$1")
+	writeToFile(os.Args[2], []byte(schema))
 }