From 7ab9b0302fe6a32637a8f850717eacf16e333d18 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?V=C3=ADctor=20Zamanillo?= <bellzebu@hotmail.com>
Date: Fri, 28 Aug 2020 11:17:37 +0200
Subject: [PATCH] Added option to list available templates

---
 internal/runner/options.go  |  2 ++
 internal/runner/runner.go   | 63 ++++++++++++++++++++++++++++++++-----
 internal/runner/validate.go | 14 +++++----
 3 files changed, 66 insertions(+), 13 deletions(-)

diff --git a/internal/runner/options.go b/internal/runner/options.go
index 8b2bacfa..97a936df 100644
--- a/internal/runner/options.go
+++ b/internal/runner/options.go
@@ -20,6 +20,7 @@ type Options struct {
 	JSON              bool // JSON writes json output to files
 	JSONRequests      bool // write requests/responses for matches in JSON output
 	EnableProgressBar bool // Enable progrss bar
+	ListTemplates     bool // List available templates
 
 	Stdin              bool                   // Stdin specifies whether stdin input was given to the process
 	Templates          multiStringFlag        // Signature specifies the template/templates to use
@@ -74,6 +75,7 @@ func ParseOptions() *Options {
 	flag.BoolVar(&options.JSON, "json", false, "Write json output to files")
 	flag.BoolVar(&options.JSONRequests, "json-requests", false, "Write requests/responses for matches in JSON output")
 	flag.BoolVar(&options.EnableProgressBar, "pbar", false, "Enable the progress bar")
+	flag.BoolVar(&options.ListTemplates, "lt", false, "List available templates")
 
 	flag.Parse()
 
diff --git a/internal/runner/runner.go b/internal/runner/runner.go
index 1e558bf8..71952d4e 100644
--- a/internal/runner/runner.go
+++ b/internal/runner/runner.go
@@ -68,6 +68,11 @@ func New(options *Options) (*Runner, error) {
 		gologger.Warningf("Could not update templates: %s\n", err)
 	}
 
+	if options.ListTemplates {
+		runner.ListAvailableTemplates()
+		os.Exit(0)
+	}
+
 	if (len(options.Templates) == 0 || (options.Targets == "" && !options.Stdin && options.Target == "")) && options.UpdateTemplates {
 		os.Exit(0)
 	}
@@ -348,8 +353,9 @@ func (r *Runner) getTemplatesFor(definitions []string) []string {
 				matches := []string{}
 
 				// Recursively walk down the Templates directory and run all the template file checks
-				err = godirwalk.Walk(absPath, &godirwalk.Options{
-					Callback: func(path string, d *godirwalk.Dirent) error {
+				err := directoryWalker(
+					absPath,
+					func(path string, d *godirwalk.Dirent) error {
 						if !d.IsDir() && strings.HasSuffix(path, ".yaml") {
 							if !r.checkIfInNucleiIgnore(path) && isNewPath(path, processed) {
 								matches = append(matches, path)
@@ -358,11 +364,7 @@ func (r *Runner) getTemplatesFor(definitions []string) []string {
 						}
 						return nil
 					},
-					ErrorCallback: func(path string, err error) godirwalk.ErrorAction {
-						return godirwalk.SkipNode
-					},
-					Unsorted: true,
-				})
+				)
 
 				// directory couldn't be walked
 				if err != nil {
@@ -789,3 +791,50 @@ func (r *Runner) parse(file string) (interface{}, error) {
 
 	return nil, errors.New("unknown error occurred")
 }
+
+func directoryWalker(path string, callback func(path string, d *godirwalk.Dirent) error) error {
+	err := godirwalk.Walk(path, &godirwalk.Options{
+		Callback: callback,
+		ErrorCallback: func(path string, err error) godirwalk.ErrorAction {
+			return godirwalk.SkipNode
+		},
+		Unsorted: true,
+	})
+
+	// directory couldn't be walked
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// ListAvailableTemplates prints available templates to stdout
+func (r *Runner) ListAvailableTemplates() {
+	gologger.Infof("Listing available templates...")
+	if r.templatesConfig != nil {
+		r.colorizer = aurora.NewAurora(true)
+		err := directoryWalker(
+			r.templatesConfig.TemplatesDirectory,
+			func(path string, d *godirwalk.Dirent) error {
+				if !d.IsDir() && strings.HasSuffix(path, ".yaml") {
+					t, err := r.parse(path)
+					switch tp := t.(type) {
+					case *templates.Template:
+						r.logTemplateLoaded(tp.ID, tp.Info.Name, tp.Info.Author, tp.Info.Severity)
+					case *workflows.Workflow:
+						r.logTemplateLoaded(tp.ID, tp.Info.Name, tp.Info.Author, tp.Info.Severity)
+					default:
+						gologger.Errorf("Could not parse file '%s': %s\n", path, err)
+					}
+				}
+				return nil
+			},
+		)
+
+		// directory couldn't be walked
+		if err != nil {
+			gologger.Labelf("Could not find templates in directory '%s': %s\n", r.templatesConfig.TemplatesDirectory, err)
+		}
+	}
+}
diff --git a/internal/runner/validate.go b/internal/runner/validate.go
index e25aaa96..0002ab51 100644
--- a/internal/runner/validate.go
+++ b/internal/runner/validate.go
@@ -14,13 +14,15 @@ func (options *Options) validateOptions() error {
 		return errors.New("both verbose and silent mode specified")
 	}
 
-	// Check if a list of templates was provided and it exists
-	if len(options.Templates) == 0 && !options.UpdateTemplates {
-		return errors.New("no template/templates provided")
-	}
+	if !options.ListTemplates {
+		// Check if a list of templates was provided and it exists
+		if len(options.Templates) == 0 && !options.UpdateTemplates {
+			return errors.New("no template/templates provided")
+		}
 
-	if options.Targets == "" && !options.Stdin && options.Target == "" && !options.UpdateTemplates {
-		return errors.New("no target input provided")
+		if options.Targets == "" && !options.Stdin && options.Target == "" && !options.UpdateTemplates {
+			return errors.New("no target input provided")
+		}
 	}
 
 	// Validate proxy options if provided