From 088c8770cc03c373f5853d213bc11f43f96ced8c Mon Sep 17 00:00:00 2001
From: Ice3man543 <nizamulrana@gmail.com>
Date: Tue, 29 Dec 2020 16:33:25 +0530
Subject: [PATCH] Compiling templates + misc stuff

---
 v2/internal/runner/runner.go  | 21 +++++++++------------
 v2/pkg/templates/compile.go   | 25 ++++++++++++++++++-------
 v2/pkg/templates/templates.go |  2 ++
 3 files changed, 29 insertions(+), 19 deletions(-)

diff --git a/v2/internal/runner/runner.go b/v2/internal/runner/runner.go
index d7e8ddd7b..368616ac5 100644
--- a/v2/internal/runner/runner.go
+++ b/v2/internal/runner/runner.go
@@ -239,22 +239,19 @@ func (r *Runner) RunEnumeration() {
 
 		for _, t := range availableTemplates {
 			wgtemplates.Add()
-			go func(template interface{}) {
-				defer wgtemplates.Done()
-				switch tt := template.(type) {
-				case *templates.Template:
-					for _, request := range tt.RequestsDNS {
-						results.Or(r.processTemplateWithList(p, tt, request))
-					}
-					for _, request := range tt.BulkRequestsHTTP {
-						results.Or(r.processTemplateWithList(p, tt, request))
-					}
-				case *workflows.Workflow:
+			go func(template *templates.Template) {
+				if template.Workflow != nil {
 					results.Or(r.processWorkflowWithList(p, template.(*workflows.Workflow)))
+
+				}
+				for _, request := range template.RequestsDNS {
+					results.Or(r.processTemplateWithList(p, tt, request))
+				}
+				for _, request := range template.RequestsHTTP {
+					results.Or(r.processTemplateWithList(p, tt, request))
 				}
 			}(t)
 		}
-
 		wgtemplates.Wait()
 		p.Stop()
 	}
diff --git a/v2/pkg/templates/compile.go b/v2/pkg/templates/compile.go
index 4fc5d0bf6..be7e24df3 100644
--- a/v2/pkg/templates/compile.go
+++ b/v2/pkg/templates/compile.go
@@ -6,6 +6,8 @@ import (
 
 	"github.com/pkg/errors"
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols"
+	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/dns"
+	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/http"
 	"gopkg.in/yaml.v2"
 )
 
@@ -30,23 +32,32 @@ func Parse(file string, options *protocols.ExecuterOptions) (*Template, error) {
 	options.TemplateInfo = template.Info
 	options.TemplatePath = file
 
+	// We don't support both http and dns in a single template
+	if len(template.RequestsDNS) > 0 && len(template.RequestsHTTP) > 0 {
+		return nil, fmt.Errorf("both http and dns requests for %s", template.ID)
+	}
 	// If no requests, and it is also not a workflow, return error.
-	if len(template.RequestsDNS)+len(template.RequestsDNS)+len(template.Workflows) <= 0 {
+	if len(template.RequestsDNS)+len(template.RequestsDNS)+len(template.Workflows) == 0 {
 		return nil, fmt.Errorf("no requests defined for %s", template.ID)
 	}
 
 	// Compile the requests found
 	for _, request := range template.RequestsDNS {
-		if err := request.Compile(options); err != nil {
-			return nil, errors.Wrap(err, "could not compile dns request")
-		}
 		template.totalRequests += request.Requests()
 	}
 	for _, request := range template.RequestsHTTP {
-		if err := request.Compile(options); err != nil {
-			return nil, errors.Wrap(err, "could not compile dns request")
-		}
 		template.totalRequests += request.Requests()
 	}
+	if len(template.RequestsDNS) > 0 {
+		template.executer = dns.NewExecuter(template.RequestsDNS, options)
+		err = template.executer.Compile()
+	}
+	if len(template.RequestsHTTP) > 0 {
+		template.executer = http.NewExecuter(template.RequestsHTTP, options)
+		err = template.executer.Compile()
+	}
+	if err != nil {
+		return nil, errors.Wrap(err, "could not compile request")
+	}
 	return template, nil
 }
diff --git a/v2/pkg/templates/templates.go b/v2/pkg/templates/templates.go
index 41b7ff018..bbd679db4 100644
--- a/v2/pkg/templates/templates.go
+++ b/v2/pkg/templates/templates.go
@@ -1,6 +1,7 @@
 package templates
 
 import (
+	"github.com/projectdiscovery/nuclei/v2/pkg/protocols"
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/dns"
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/http"
 	"github.com/projectdiscovery/nuclei/v2/pkg/workflows"
@@ -22,6 +23,7 @@ type Template struct {
 
 	path          string
 	totalRequests int
+	executer      protocols.Executer
 }
 
 // GetPath returns the path of the template.