Adding dns trace support in dns templates (#1236)

* Adding dns trace support in dns templates + minor refactoring

v2/go.mod

@@ -36,7 +36,7 @@ require (
 	github.com/projectdiscovery/interactsh v0.0.6
 	github.com/projectdiscovery/nuclei-updatecheck-api v0.0.0-20211006155443-c0a8d610a4df
 	github.com/projectdiscovery/rawhttp v0.0.7
-	github.com/projectdiscovery/retryabledns v1.0.13-0.20210916165024-76c5b76fd59a
+	github.com/projectdiscovery/retryabledns v1.0.13-0.20211109182249-43d38df59660
 	github.com/projectdiscovery/retryablehttp-go v1.0.2
 	github.com/projectdiscovery/stringsutil v0.0.0-20210830151154-f567170afdd9
 	github.com/projectdiscovery/yamldoc-go v1.0.2

v2/go.sum

@@ -582,6 +582,7 @@ github.com/projectdiscovery/filekv v0.0.0-20210915124239-3467ef45dd08 h1:NwD1R/d
 github.com/projectdiscovery/filekv v0.0.0-20210915124239-3467ef45dd08/go.mod h1:paLCnwV8sL7ppqIwVQodQrk3F6mnWafwTDwRd7ywZwQ=
 github.com/projectdiscovery/fileutil v0.0.0-20210804142714-ebba15fa53ca/go.mod h1:U+QCpQnX8o2N2w0VUGyAzjM3yBAe4BKedVElxiImsx0=
 github.com/projectdiscovery/fileutil v0.0.0-20210914153648-31f843feaad4/go.mod h1:U+QCpQnX8o2N2w0VUGyAzjM3yBAe4BKedVElxiImsx0=
+github.com/projectdiscovery/fileutil v0.0.0-20210926202739-6050d0acf73c/go.mod h1:U+QCpQnX8o2N2w0VUGyAzjM3yBAe4BKedVElxiImsx0=
 github.com/projectdiscovery/fileutil v0.0.0-20210928100737-cab279c5d4b5 h1:2dbm7UhrAKnccZttr78CAmG768sSCd+MBn4ayLVDeqA=
 github.com/projectdiscovery/fileutil v0.0.0-20210928100737-cab279c5d4b5/go.mod h1:U+QCpQnX8o2N2w0VUGyAzjM3yBAe4BKedVElxiImsx0=
 github.com/projectdiscovery/goflags v0.0.7/go.mod h1:Jjwsf4eEBPXDSQI2Y+6fd3dBumJv/J1U0nmpM+hy2YY=
@@ -618,8 +619,9 @@ github.com/projectdiscovery/rawhttp v0.0.7 h1:5m4peVgjbl7gqDcRYMTVEuX+Xs/nh76ohT
 github.com/projectdiscovery/rawhttp v0.0.7/go.mod h1:PQERZAhAv7yxI/hR6hdDPgK1WTU56l204BweXrBec+0=
 github.com/projectdiscovery/retryabledns v1.0.11/go.mod h1:4sMC8HZyF01HXukRleSQYwz4870bwgb4+hTSXTMrkf4=
 github.com/projectdiscovery/retryabledns v1.0.12/go.mod h1:4sMC8HZyF01HXukRleSQYwz4870bwgb4+hTSXTMrkf4=
-github.com/projectdiscovery/retryabledns v1.0.13-0.20210916165024-76c5b76fd59a h1:WJQjr9qi/VjWhdNiGyNqcFi0967Gp0W3I769bCpHOJE=
 github.com/projectdiscovery/retryabledns v1.0.13-0.20210916165024-76c5b76fd59a/go.mod h1:tXaLDs4n3pRZHwfa8mdXpUWe/AYDNK3HlWDjldhRbjI=
+github.com/projectdiscovery/retryabledns v1.0.13-0.20211109182249-43d38df59660 h1:Ooa5htghPkdyfpzy6Y5KLdyv4w8ePZWmfzFSPQlJStQ=
+github.com/projectdiscovery/retryabledns v1.0.13-0.20211109182249-43d38df59660/go.mod h1:UfszkO3x+GLKVOpXB7boddJKbwNCr+tMPSkfgCSNhl4=
 github.com/projectdiscovery/retryablehttp-go v1.0.1/go.mod h1:SrN6iLZilNG1X4neq1D+SBxoqfAF4nyzvmevkTkWsek=
 github.com/projectdiscovery/retryablehttp-go v1.0.2 h1:LV1/KAQU+yeWhNVlvveaYFsjBYRwXlNEq0PvrezMV0U=
 github.com/projectdiscovery/retryablehttp-go v1.0.2/go.mod h1:dx//aY9V247qHdsRf0vdWHTBZuBQ2vm6Dq5dagxrDYI=

v2/pkg/protocols/common/executer/executer.go

@@ -47,8 +47,6 @@ func (e *Executer) Execute(input string) (bool, error) {
 	dynamicValues := make(map[string]interface{})
 	previous := make(map[string]interface{})
 	for _, req := range e.requests {
-		req := req
-
 		err := req.ExecuteWithResults(input, dynamicValues, previous, func(event *output.InternalWrappedEvent) {
 			ID := req.GetID()
 			if ID != "" {

v2/pkg/protocols/common/expressions/variables.go

@@ -10,47 +10,55 @@ var unresolvedVariablesRegex = regexp.MustCompile(`(?:%7[B|b]|\{){2}([^}]+)(?:%7
 
 // ContainsUnresolvedVariables returns an error with variable names if the passed
 // input contains unresolved {{<pattern-here>}} variables.
-func ContainsUnresolvedVariables(data string) error {
+func ContainsUnresolvedVariables(items ...string) error {
-	matches := unresolvedVariablesRegex.FindAllStringSubmatch(data, -1)
+	for _, data := range items {
-	if len(matches) == 0 {
+		matches := unresolvedVariablesRegex.FindAllStringSubmatch(data, -1)
-		return nil
+		if len(matches) == 0 {
-	}
+			return nil
-	errorString := &strings.Builder{}
-	errorString.WriteString("unresolved variables found: ")
-
-	for i, match := range matches {
-		if len(match) < 2 {
-			continue
 		}
-		errorString.WriteString(match[1])
+		errorString := &strings.Builder{}
-		if i != len(matches)-1 {
+		errorString.WriteString("unresolved variables found: ")
-			errorString.WriteString(",")
-		}
-	}
-	errorMessage := errorString.String()
-	return errors.New(errorMessage)
-}
 
-func ContainsVariablesWithNames(data string, names map[string]interface{}) error {
+		for i, match := range matches {
-	matches := unresolvedVariablesRegex.FindAllStringSubmatch(data, -1)
+			if len(match) < 2 {
-	if len(matches) == 0 {
+				continue
-		return nil
+			}
-	}
+			errorString.WriteString(match[1])
-	errorString := &strings.Builder{}
-	errorString.WriteString("unresolved variables with values found: ")
-
-	for i, match := range matches {
-		if len(match) < 2 {
-			continue
-		}
-		matchName := match[1]
-		if _, ok := names[matchName]; !ok {
-			errorString.WriteString(matchName)
 			if i != len(matches)-1 {
 				errorString.WriteString(",")
 			}
 		}
+		errorMessage := errorString.String()
+		return errors.New(errorMessage)
 	}
-	errorMessage := errorString.String()
+
-	return errors.New(errorMessage)
+	return nil
+}
+
+func ContainsVariablesWithNames(names map[string]interface{}, items ...string) error {
+	for _, data := range items {
+		matches := unresolvedVariablesRegex.FindAllStringSubmatch(data, -1)
+		if len(matches) == 0 {
+			return nil
+		}
+		errorString := &strings.Builder{}
+		errorString.WriteString("unresolved variables with values found: ")
+
+		for i, match := range matches {
+			if len(match) < 2 {
+				continue
+			}
+			matchName := match[1]
+			if _, ok := names[matchName]; !ok {
+				errorString.WriteString(matchName)
+				if i != len(matches)-1 {
+					errorString.WriteString(",")
+				}
+			}
+		}
+		errorMessage := errorString.String()
+		return errors.New(errorMessage)
+	}
+
+	return nil
 }

v2/pkg/protocols/dns/dns.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/projectdiscovery/nuclei/v2/pkg/operators"
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols"
+	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/expressions"
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/replacer"
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/dns/dnsclientpool"
 	"github.com/projectdiscovery/retryabledns"
@@ -62,6 +63,15 @@ type Request struct {
 	//   - name: Use a retry of 3 to 5 generally
 	//     value: 5
 	Retries int `yaml:"retries,omitempty" jsonschema:"title=retries for dns request,description=Retries is the number of retries for the DNS request"`
+	// description: |
+	//   Trace performs a trace operation for the target.
+	Trace bool `yaml:"trace,omitempty" jsonschema:"title=trace operation,description=Trace performs a trace operation for the target."`
+	// description: |
+	//   TraceMaxRecursion is the number of max recursion allowed for trace operations
+	// examples:
+	//   - name: Use a retry of 100 to 150 generally
+	//     value: 100
+	TraceMaxRecursion int `yaml:"trace-max-recursion,omitempty"  jsonschema:"title=trace-max-recursion level for dns request,description=TraceMaxRecursion is the number of max recursion allowed for trace operations"`
 
 	CompiledOperators *operators.Operators `yaml:"-"`
 	dnsClient         *retryabledns.Client
@@ -96,7 +106,7 @@ func (request *Request) Compile(options *protocols.ExecuterOptions) error {
 		dnsClientOptions.Resolvers = request.Resolvers
 	}
 	// Create a dns client for the class
-	client, err := dnsclientpool.Get(options.Options, dnsClientOptions)
+	client, err := request.getDnsClient(options, nil)
 	if err != nil {
 		return errors.Wrap(err, "could not get dns client")
 	}
@@ -115,6 +125,28 @@ func (request *Request) Compile(options *protocols.ExecuterOptions) error {
 	return nil
 }
 
+func (request *Request) getDnsClient(options *protocols.ExecuterOptions, metadata map[string]interface{}) (*retryabledns.Client, error) {
+	dnsClientOptions := &dnsclientpool.Configuration{
+		Retries: request.Retries,
+	}
+	if len(request.Resolvers) > 0 {
+		if len(request.Resolvers) > 0 {
+			for _, resolver := range request.Resolvers {
+				if expressions.ContainsUnresolvedVariables(resolver) != nil {
+					var err error
+					resolver, err = expressions.Evaluate(resolver, metadata)
+					if err != nil {
+						return nil, errors.Wrap(err, "could not resolve resolvers expressions")
+					}
+					dnsClientOptions.Resolvers = append(dnsClientOptions.Resolvers, resolver)
+				}
+			}
+		}
+		dnsClientOptions.Resolvers = request.Resolvers
+	}
+	return dnsclientpool.Get(options.Options, dnsClientOptions)
+}
+
 // Requests returns the total number of requests the YAML rule will perform
 func (request *Request) Requests() int {
 	return 1

v2/pkg/protocols/dns/operators.go

@@ -2,6 +2,8 @@ package dns
 
 import (
 	"bytes"
+	"fmt"
+	"strings"
 	"time"
 
 	"github.com/miekg/dns"
@@ -12,6 +14,7 @@ import (
 	"github.com/projectdiscovery/nuclei/v2/pkg/output"
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols"
 	"github.com/projectdiscovery/nuclei/v2/pkg/types"
+	"github.com/projectdiscovery/retryabledns"
 )
 
 // Match matches a generic data response against a given matcher
@@ -73,7 +76,7 @@ func (request *Request) getMatchPart(part string, data output.InternalEvent) (in
 }
 
 // responseToDSLMap converts a DNS response to a map for use in DSL matching
-func (request *Request) responseToDSLMap(req, resp *dns.Msg, host, matched string) output.InternalEvent {
+func (request *Request) responseToDSLMap(req, resp *dns.Msg, host, matched string, tracedata *retryabledns.TraceData) output.InternalEvent {
 	return output.InternalEvent{
 		"host":          host,
 		"matched":       matched,
@@ -87,6 +90,7 @@ func (request *Request) responseToDSLMap(req, resp *dns.Msg, host, matched strin
 		"template-id":   request.options.TemplateID,
 		"template-info": request.options.TemplateInfo,
 		"template-path": request.options.TemplatePath,
+		"trace":         traceToString(tracedata, false),
 	}
 }
 
@@ -126,3 +130,16 @@ func questionToString(resourceRecords []dns.Question) string {
 	}
 	return buffer.String()
 }
+
+func traceToString(tracedata *retryabledns.TraceData, withSteps bool) string {
+	buffer := &bytes.Buffer{}
+	if tracedata != nil {
+		for i, dnsRecord := range tracedata.DNSData {
+			if withSteps {
+				buffer.WriteString(fmt.Sprintf("request %d to resolver %s:\n", i, strings.Join(dnsRecord.Resolver, ",")))
+			}
+			buffer.WriteString(dnsRecord.Raw)
+		}
+	}
+	return buffer.String()
+}

v2/pkg/protocols/dns/operators_test.go

@@ -44,8 +44,8 @@ func TestResponseToDSLMap(t *testing.T) {
 	resp.Rcode = dns.RcodeSuccess
 	resp.Answer = append(resp.Answer, &dns.A{A: net.ParseIP("1.1.1.1"), Hdr: dns.RR_Header{Name: "one.one.one.one."}})
 
-	event := request.responseToDSLMap(req, resp, "one.one.one.one", "one.one.one.one")
+	event := request.responseToDSLMap(req, resp, "one.one.one.one", "one.one.one.one", nil)
-	require.Len(t, event, 12, "could not get correct number of items in dsl map")
+	require.Len(t, event, 13, "could not get correct number of items in dsl map")
 	require.Equal(t, dns.RcodeSuccess, event["rcode"], "could not get correct rcode")
 }
 
@@ -76,7 +76,7 @@ func TestDNSOperatorMatch(t *testing.T) {
 	resp.Rcode = dns.RcodeSuccess
 	resp.Answer = append(resp.Answer, &dns.A{A: net.ParseIP("1.1.1.1"), Hdr: dns.RR_Header{Name: "one.one.one.one."}})
 
-	event := request.responseToDSLMap(req, resp, "one.one.one.one", "one.one.one.one")
+	event := request.responseToDSLMap(req, resp, "one.one.one.one", "one.one.one.one", nil)
 
 	t.Run("valid", func(t *testing.T) {
 		matcher := &matchers.Matcher{
@@ -143,7 +143,7 @@ func TestDNSOperatorMatch(t *testing.T) {
 		resp.Rcode = dns.RcodeSuccess
 		resp.Answer = append(resp.Answer, &dns.A{A: net.ParseIP("1.1.1.1"), Hdr: dns.RR_Header{Name: "ONE.ONE.ONE.ONE."}})
 
-		event := request.responseToDSLMap(req, resp, "ONE.ONE.ONE.ONE", "ONE.ONE.ONE.ONE")
+		event := request.responseToDSLMap(req, resp, "ONE.ONE.ONE.ONE", "ONE.ONE.ONE.ONE", nil)
 
 		matcher := &matchers.Matcher{
 			Part:            "raw",
@@ -187,7 +187,7 @@ func TestDNSOperatorExtract(t *testing.T) {
 	resp.Rcode = dns.RcodeSuccess
 	resp.Answer = append(resp.Answer, &dns.A{A: net.ParseIP("1.1.1.1"), Hdr: dns.RR_Header{Name: "one.one.one.one."}})
 
-	event := request.responseToDSLMap(req, resp, "one.one.one.one", "one.one.one.one")
+	event := request.responseToDSLMap(req, resp, "one.one.one.one", "one.one.one.one", nil)
 
 	t.Run("extract", func(t *testing.T) {
 		extractor := &extractors.Extractor{
@@ -257,7 +257,7 @@ func TestDNSMakeResult(t *testing.T) {
 	resp.Rcode = dns.RcodeSuccess
 	resp.Answer = append(resp.Answer, &dns.A{A: net.ParseIP("1.1.1.1"), Hdr: dns.RR_Header{Name: "one.one.one.one."}})
 
-	event := request.responseToDSLMap(req, resp, "one.one.one.one", "one.one.one.one")
+	event := request.responseToDSLMap(req, resp, "one.one.one.one", "one.one.one.one", nil)
 	finalEvent := &output.InternalWrappedEvent{InternalEvent: event}
 	if request.CompiledOperators != nil {
 		result, ok := request.CompiledOperators.Execute(event, request.Match, request.Extract, false)

v2/pkg/protocols/dns/request.go

@@ -12,6 +12,7 @@ import (
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/expressions"
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/helpers/eventcreator"
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/helpers/responsehighlighter"
+	"github.com/projectdiscovery/retryabledns"
 	templateTypes "github.com/projectdiscovery/nuclei/v2/pkg/templates/types"
 )
 
@@ -40,6 +41,14 @@ func (request *Request) ExecuteWithResults(input string, metadata /*TODO review
 		return errors.Wrap(err, "could not build request")
 	}
 
+	dnsClient := request.dnsClient
+	if varErr := expressions.ContainsUnresolvedVariables(request.Resolvers...); varErr != nil {
+		if dnsClient, varErr = request.getDnsClient(request.options, metadata); varErr != nil {
+			gologger.Warning().Msgf("[%s] Could not make dns request for %s: %v\n", request.options.TemplateID, domain, varErr)
+			return nil
+		}
+	}
+
 	requestString := compiledRequest.String()
 	if varErr := expressions.ContainsUnresolvedVariables(requestString); varErr != nil {
 		gologger.Warning().Msgf("[%s] Could not make dns request for %s: %v\n", request.options.TemplateID, domain, varErr)
@@ -51,7 +60,7 @@ func (request *Request) ExecuteWithResults(input string, metadata /*TODO review
 	}
 
 	// Send the request to the target servers
-	response, err := request.dnsClient.Do(compiledRequest)
+	response, err := dnsClient.Do(compiledRequest)
 	if err != nil {
 		request.options.Output.Request(request.options.TemplatePath, domain, request.Type().String(), err)
 		request.options.Progress.IncrementFailedRequestsBy(1)
@@ -64,20 +73,33 @@ func (request *Request) ExecuteWithResults(input string, metadata /*TODO review
 	request.options.Output.Request(request.options.TemplatePath, domain, request.Type().String(), err)
 	gologger.Verbose().Msgf("[%s] Sent DNS request to %s\n", request.options.TemplateID, domain)
 
-	outputEvent := request.responseToDSLMap(compiledRequest, response, input, input)
+	// perform trace if necessary
+	var tracedata *retryabledns.TraceData
+	if request.Trace {
+		tracedata, err = request.dnsClient.Trace(domain, request.question, request.TraceMaxRecursion)
+		if err != nil {
+			request.options.Output.Request(request.options.TemplatePath, domain, "dns", err)
+		}
+	}
+
+	outputEvent := request.responseToDSLMap(compiledRequest, response, input, input, tracedata)
 	for k, v := range previous {
 		outputEvent[k] = v
 	}
 
 	event := eventcreator.CreateEvent(request, outputEvent, request.options.Options.Debug || request.options.Options.DebugResponse)
+	// TODO: dynamic values are not supported yet
 
 	dumpResponse(event, request.options, response.String(), domain)
+	if request.Trace {
+		dumpTraceData(event, request.options, traceToString(tracedata, true), domain)
+	}
 
 	callback(event)
 	return nil
 }
 
-func dumpResponse(event *output.InternalWrappedEvent, requestOptions *protocols.ExecuterOptions, response string, domain string) {
+func dumpResponse(event *output.InternalWrappedEvent, requestOptions *protocols.ExecuterOptions, response, domain string) {
 	cliOptions := requestOptions.Options
 	if cliOptions.Debug || cliOptions.DebugResponse {
 		hexDump := false
@@ -90,6 +112,19 @@ func dumpResponse(event *output.InternalWrappedEvent, requestOptions *protocols.
 	}
 }
 
+func dumpTraceData(event *output.InternalWrappedEvent, requestOptions *protocols.ExecuterOptions, tracedata, domain string) {
+	cliOptions := requestOptions.Options
+	if cliOptions.Debug || cliOptions.DebugResponse {
+		hexDump := false
+		if responsehighlighter.HasBinaryContent(tracedata) {
+			hexDump = true
+			tracedata = hex.Dump([]byte(tracedata))
+		}
+		highlightedResponse := responsehighlighter.Highlight(event.OperatorsResult, tracedata, cliOptions.NoColor, hexDump)
+		gologger.Debug().Msgf("[%s] Dumped DNS Trace data for %s\n\n%s", requestOptions.TemplateID, domain, highlightedResponse)
+	}
+}
+
 // isURL tests a string to determine if it is a well-structured url or not.
 func isURL(toTest string) bool {
 	if _, err := url.ParseRequestURI(toTest); err != nil {

v2/pkg/protocols/http/http.go

@@ -252,7 +252,7 @@ func (request *Request) Compile(options *protocols.ExecuterOptions) error {
 		}
 
 		for _, input := range inputs {
-			if expressions.ContainsVariablesWithNames(input, map[string]interface{}{name: payload}) == nil {
+			if expressions.ContainsVariablesWithNames(map[string]interface{}{name: payload}, input) == nil {
 				hasPayloadName = true
 				break
 			}

v2/pkg/protocols/network/network.go

@@ -170,7 +170,7 @@ func (request *Request) Compile(options *protocols.ExecuterOptions) error {
 			if input.Type.String() != "" {
 				continue
 			}
-			if expressions.ContainsVariablesWithNames(input.Data, map[string]interface{}{name: payload}) == nil {
+			if expressions.ContainsVariablesWithNames(map[string]interface{}{name: payload}, input.Data) == nil {
 				hasPayloadName = true
 				break
 			}