Extending YAML to support include preprocessing

v2/pkg/parsers/parser.go

@@ -5,14 +5,13 @@ import (
 	"regexp"
 	"strings"
 
-	"gopkg.in/yaml.v2"
-
 	"github.com/projectdiscovery/nuclei/v2/pkg/catalog"
 	"github.com/projectdiscovery/nuclei/v2/pkg/catalog/loader/filter"
 	"github.com/projectdiscovery/nuclei/v2/pkg/templates"
 	"github.com/projectdiscovery/nuclei/v2/pkg/templates/cache"
 	"github.com/projectdiscovery/nuclei/v2/pkg/utils"
 	"github.com/projectdiscovery/nuclei/v2/pkg/utils/stats"
+	"gopkg.in/yaml.v2"
 )
 
 const (

v2/pkg/utils/utils.go

@@ -8,6 +8,8 @@ import (
 	"strings"
 
 	"github.com/projectdiscovery/nuclei/v2/pkg/catalog"
+	"github.com/projectdiscovery/nuclei/v2/pkg/utils/yaml"
+	fileutil "github.com/projectdiscovery/utils/file"
 )
 
 func IsBlank(value string) bool {
@@ -27,42 +29,42 @@ func UnwrapError(err error) error {
 
 // IsURL tests a string to determine if it is a well-structured url or not.
 func IsURL(input string) bool {
-	_, err := url.ParseRequestURI(input)
-	if err != nil {
-		return false
-	}
-
 	u, err := url.Parse(input)
-	if err != nil || u.Scheme == "" || u.Host == "" {
-		return false
-	}
-
-	return true
+	return err == nil && u.Scheme != "" && u.Host != ""
 }
 
 // ReadFromPathOrURL reads and returns the contents of a file or url.
 func ReadFromPathOrURL(templatePath string, catalog catalog.Catalog) (data []byte, err error) {
+	var reader io.Reader
 	if IsURL(templatePath) {
 		resp, err := http.Get(templatePath)
 		if err != nil {
 			return nil, err
 		}
 		defer resp.Body.Close()
-		data, err = io.ReadAll(resp.Body)
-		if err != nil {
-			return nil, err
-		}
+		reader = resp.Body
 	} else {
 		f, err := catalog.OpenFile(templatePath)
 		if err != nil {
 			return nil, err
 		}
 		defer f.Close()
-		data, err = io.ReadAll(f)
+		reader = f
+	}
+
+	data, err = io.ReadAll(reader)
+	if err != nil {
+		return nil, err
+	}
+
+	// pre-process directives only for local files
+	if fileutil.FileExists(templatePath) {
+		data, err = yaml.PreProcess(data)
 		if err != nil {
 			return nil, err
 		}
 	}
+
 	return
 }
 

v2/pkg/utils/yaml/preprocess.go

@@ -0,0 +1,65 @@
+package yaml
+
+import (
+	"bytes"
+	"os"
+	"regexp"
+	"strings"
+
+	fileutil "github.com/projectdiscovery/utils/file"
+	stringsutil "github.com/projectdiscovery/utils/strings"
+)
+
+var reImportsPattern = regexp.MustCompile(`(?m)# !include:(.+.yaml)`)
+
+// PreProcess all include directives
+func PreProcess(data []byte) ([]byte, error) {
+	// find all matches like !include:path\n
+	importMatches := reImportsPattern.FindAllSubmatch(data, -1)
+
+	var replaceItems []string
+
+	for _, match := range importMatches {
+		var (
+			matchString     string
+			includeFileName string
+		)
+		matchBytes := match[0]
+		matchString = string(matchBytes)
+		if len(match) > 0 {
+			includeFileName = string(match[1])
+		}
+
+		// gets the number of tabs/spaces between the last \n and the beginning of the match
+		matchIndex := bytes.Index(data, matchBytes)
+		lastNewLineIndex := bytes.LastIndex(data[:matchIndex], []byte("\n"))
+		padBytes := data[lastNewLineIndex:matchIndex]
+
+		// check if the file exists
+		if fileutil.FileExists(includeFileName) {
+			// and in case replace the comment with it
+			includeFileContent, err := os.ReadFile(includeFileName)
+			if err != nil {
+				return nil, err
+			}
+			// if it's yaml, tries to preprocess that too recursively
+			if stringsutil.HasSuffixAny(includeFileName, ".yaml") {
+				if subIncludedFileContent, err := PreProcess(includeFileContent); err == nil {
+					includeFileContent = subIncludedFileContent
+				} else {
+					return nil, err
+				}
+			}
+
+			// pad each line of file content with padBytes
+			includeFileContent = bytes.ReplaceAll(includeFileContent, []byte("\n"), padBytes)
+
+			replaceItems = append(replaceItems, matchString)
+			replaceItems = append(replaceItems, string(includeFileContent))
+		}
+	}
+
+	replacer := strings.NewReplacer(replaceItems...)
+
+	return []byte(replacer.Replace(string(data))), nil
+}