From ae5f54dce98c91093a2166575cadff30d9617478 Mon Sep 17 00:00:00 2001
From: lc <corben@lynxsecurity.io>
Date: Sat, 2 May 2020 12:10:52 -0500
Subject: [PATCH 1/2] Add directory style support for templates

---
 internal/runner/runner.go | 56 ++++++++++++++++++++++++++++++++-------
 1 file changed, 46 insertions(+), 10 deletions(-)

diff --git a/internal/runner/runner.go b/internal/runner/runner.go
index b42025093..9a0073d3a 100644
--- a/internal/runner/runner.go
+++ b/internal/runner/runner.go
@@ -10,6 +10,7 @@ import (
 	"strings"
 	"sync"
 
+	"github.com/karrick/godirwalk"
 	"github.com/projectdiscovery/gologger"
 	"github.com/projectdiscovery/nuclei/pkg/executor"
 	"github.com/projectdiscovery/nuclei/pkg/requests"
@@ -67,12 +68,8 @@ func (r *Runner) Close() {
 // RunEnumeration sets up the input layer for giving input nuclei.
 // binary and runs the actual enumeration
 func (r *Runner) RunEnumeration() {
-	if !strings.HasSuffix(r.options.Templates, ".yaml") {
-		gologger.Fatalf("Could not run recognize template extension: %s\n", r.options.Templates)
-	}
-
 	// If the template path is a single template and not a glob, use that.
-	if !strings.Contains(r.options.Templates, "*") {
+	if !strings.Contains(r.options.Templates, "*") && strings.HasSuffix(r.options.Templates, ".yaml") {
 		template, err := templates.ParseTemplate(r.options.Templates)
 		if err != nil {
 			gologger.Errorf("Could not parse template file '%s': %s\n", r.options.Templates, err)
@@ -90,13 +87,51 @@ func (r *Runner) RunEnumeration() {
 		}
 		return
 	}
+	// If the template path is glob
+	if strings.Contains(r.options.Templates, "*") {
+		// Handle the glob, evaluate it and run all the template file checks
+		matches, err := filepath.Glob(r.options.Templates)
+		if err != nil {
+			gologger.Fatalf("Could not evaluate template path '%s': %s\n", r.options.Templates, err)
+		}
 
-	// Handle the glob, evaluate it and run all the template file checks
-	matches, err := filepath.Glob(r.options.Templates)
-	if err != nil {
-		gologger.Fatalf("Could not evaluate template path '%s': %s\n", r.options.Templates, err)
+		for _, match := range matches {
+			template, err := templates.ParseTemplate(match)
+			if err != nil {
+				gologger.Errorf("Could not parse template file '%s': %s\n", match, err)
+				return
+			}
+			for _, request := range template.RequestsDNS {
+				r.processTemplateRequest(template, request)
+			}
+			for _, request := range template.RequestsHTTP {
+				r.processTemplateRequest(template, request)
+			}
+		}
+		return
+	}
+	// If the template passed is a directory
+	matches := []string{}
+	// Recursively walk down the Templates directory and run all the template file checks
+	err := godirwalk.Walk(r.options.Templates, &godirwalk.Options{
+		Callback: func(path string, d *godirwalk.Dirent) error {
+			if !d.IsDir() && strings.HasSuffix(path, ".yaml") {
+				matches = append(matches, path)
+			}
+			return nil
+		},
+		ErrorCallback: func(path string, err error) godirwalk.ErrorAction {
+			return godirwalk.SkipNode
+		},
+		Unsorted: true,
+	})
+	if err != nil {
+		gologger.Fatalf("Error, walking directory: '%s': %s\n", r.options.Templates, err)
+	}
+	// 0 matches means no templates were found in directory
+	if len(matches) == 0 {
+		gologger.Fatalf("Error, no templates found in directory: '%s'\n", r.options.Templates)
 	}
-
 	for _, match := range matches {
 		template, err := templates.ParseTemplate(match)
 		if err != nil {
@@ -110,6 +145,7 @@ func (r *Runner) RunEnumeration() {
 			r.processTemplateRequest(template, request)
 		}
 	}
+	return
 }
 
 // processTemplate processes a template and runs the enumeration on all the targets

From 1b835085e645204847827cd274f7060dc31473ab Mon Sep 17 00:00:00 2001
From: lc <corben@lynxsecurity.io>
Date: Sat, 2 May 2020 12:22:05 -0500
Subject: [PATCH 2/2] Add directory style support for templates

---
 internal/runner/runner.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/internal/runner/runner.go b/internal/runner/runner.go
index 9a0073d3a..cff5a5a4e 100644
--- a/internal/runner/runner.go
+++ b/internal/runner/runner.go
@@ -126,7 +126,7 @@ func (r *Runner) RunEnumeration() {
 		Unsorted: true,
 	})
 	if err != nil {
-		gologger.Fatalf("Error, walking directory: '%s': %s\n", r.options.Templates, err)
+		gologger.Fatalf("Could not find templates in directory '%s': %s\n", r.options.Templates, err)
 	}
 	// 0 matches means no templates were found in directory
 	if len(matches) == 0 {