From ffe6ab04b37db52862724a02a75e577777552819 Mon Sep 17 00:00:00 2001
From: Ice3man <nizamulrana@gmail.com>
Date: Mon, 27 Jun 2022 18:09:29 +0530
Subject: [PATCH] Added include-templates force-loading for templates (#2232)

* Added include-templates force-loading for templates

* Fixed loader case with include-templates

* Added integration test for excluded-template in loader
---
 .../loader/excluded-template.yaml             | 16 +++++++++++++
 v2/cmd/integration-test/loader.go             | 24 +++++++++++++++++++
 v2/pkg/catalog/loader/filter/path_filter.go   |  6 +++++
 v2/pkg/catalog/loader/loader.go               | 14 +++++------
 4 files changed, 52 insertions(+), 8 deletions(-)
 create mode 100644 integration_tests/loader/excluded-template.yaml

diff --git a/integration_tests/loader/excluded-template.yaml b/integration_tests/loader/excluded-template.yaml
new file mode 100644
index 00000000..6c25a631
--- /dev/null
+++ b/integration_tests/loader/excluded-template.yaml
@@ -0,0 +1,16 @@
+id: excluded-template
+
+info:
+  name: Basic Excluded Template
+  author: pdteam
+  severity: info
+  tags: fuzz
+
+requests:
+  - method: GET
+    path:
+      - "{{BaseURL}}"
+    matchers:
+      - type: word
+        words:
+          - "This is test matcher text"
\ No newline at end of file
diff --git a/v2/cmd/integration-test/loader.go b/v2/cmd/integration-test/loader.go
index 9c0439c9..b2ce9a0e 100644
--- a/v2/cmd/integration-test/loader.go
+++ b/v2/cmd/integration-test/loader.go
@@ -16,6 +16,7 @@ import (
 var loaderTestcases = map[string]testutils.TestCase{
 	"loader/template-list.yaml":             &remoteTemplateList{},
 	"loader/workflow-list.yaml":             &remoteWorkflowList{},
+	"loader/excluded-template.yaml":         &excludedTemplate{},
 	"loader/nonexistent-template-list.yaml": &nonExistentTemplateList{},
 	"loader/nonexistent-workflow-list.yaml": &nonExistentWorkflowList{},
 	"loader/template-list-not-allowed.yaml": &remoteTemplateListNotAllowed{},
@@ -62,6 +63,29 @@ func (h *remoteTemplateList) Execute(templateList string) error {
 	return expectResultsCount(results, 2)
 }
 
+type excludedTemplate struct{}
+
+// Execute executes a test case and returns an error if occurred
+func (h *excludedTemplate) Execute(templateList string) error {
+	router := httprouter.New()
+
+	router.GET("/", func(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
+		fmt.Fprintf(w, "This is test matcher text")
+		if strings.EqualFold(r.Header.Get("test"), "nuclei") {
+			fmt.Fprintf(w, "This is test headers matcher text")
+		}
+	})
+	ts := httptest.NewServer(router)
+	defer ts.Close()
+
+	results, err := testutils.RunNucleiBareArgsAndGetResults(debug, "-target", ts.URL, "-t", templateList, "-include-templates", templateList)
+	if err != nil {
+		return err
+	}
+
+	return expectResultsCount(results, 1)
+}
+
 type remoteTemplateListNotAllowed struct{}
 
 // Execute executes a test case and returns an error if occurred
diff --git a/v2/pkg/catalog/loader/filter/path_filter.go b/v2/pkg/catalog/loader/filter/path_filter.go
index f938dcff..ba65bf71 100644
--- a/v2/pkg/catalog/loader/filter/path_filter.go
+++ b/v2/pkg/catalog/loader/filter/path_filter.go
@@ -43,3 +43,9 @@ func (p *PathFilter) Match(templates []string) map[string]struct{} {
 	}
 	return templatesMap
 }
+
+// MatchIncluded returns true if the template was included explicitly
+func (p *PathFilter) MatchIncluded(template string) bool {
+	_, found := p.alwaysIncludedTemplatesMap[template]
+	return found
+}
diff --git a/v2/pkg/catalog/loader/loader.go b/v2/pkg/catalog/loader/loader.go
index d776c6c6..18e7b333 100644
--- a/v2/pkg/catalog/loader/loader.go
+++ b/v2/pkg/catalog/loader/loader.go
@@ -261,10 +261,7 @@ func (store *Store) LoadTemplates(templatesList []string) []*templates.Template
 	loadedTemplates := make([]*templates.Template, 0, len(templatePathMap))
 	for templatePath := range templatePathMap {
 		loaded, err := parsers.LoadTemplate(templatePath, store.tagFilter, nil)
-		if err != nil {
-			gologger.Warning().Msgf("Could not load template %s: %s\n", templatePath, err)
-		}
-		if loaded {
+		if loaded || store.pathFilter.MatchIncluded(templatePath) {
 			parsed, err := templates.Parse(templatePath, store.preprocessor, store.config.ExecutorOptions)
 			if err != nil {
 				stats.Increment(parsers.RuntimeWarningsStats)
@@ -272,6 +269,8 @@ func (store *Store) LoadTemplates(templatesList []string) []*templates.Template
 			} else if parsed != nil {
 				loadedTemplates = append(loadedTemplates, parsed)
 			}
+		} else if err != nil {
+			gologger.Warning().Msgf("Could not load template %s: %s\n", templatePath, err)
 		}
 	}
 	return loadedTemplates
@@ -309,10 +308,7 @@ func (store *Store) LoadTemplatesWithTags(templatesList, tags []string) []*templ
 	loadedTemplates := make([]*templates.Template, 0, len(templatePathMap))
 	for templatePath := range templatePathMap {
 		loaded, err := parsers.LoadTemplate(templatePath, store.tagFilter, tags)
-		if err != nil {
-			gologger.Warning().Msgf("Could not load template %s: %s\n", templatePath, err)
-		}
-		if loaded {
+		if loaded || store.pathFilter.MatchIncluded(templatePath) {
 			parsed, err := templates.Parse(templatePath, store.preprocessor, store.config.ExecutorOptions)
 			if err != nil {
 				stats.Increment(parsers.RuntimeWarningsStats)
@@ -320,6 +316,8 @@ func (store *Store) LoadTemplatesWithTags(templatesList, tags []string) []*templ
 			} else if parsed != nil {
 				loadedTemplates = append(loadedTemplates, parsed)
 			}
+		} else if err != nil {
+			gologger.Warning().Msgf("Could not load template %s: %s\n", templatePath, err)
 		}
 	}
 	return loadedTemplates