Working DNS and HTTP protocol implm

v2/pkg/output/output.go

@@ -130,6 +130,7 @@ func (w *StandardWriter) Write(event *ResultEvent) error {
 		return errors.Wrap(err, "could not format output")
 	}
 	_, _ = os.Stdout.Write(data)
+	_, _ = os.Stdout.Write([]byte("\n"))
 	if w.outputFile != nil {
 		if !w.json {
 			data = decolorizerRegex.ReplaceAll(data, []byte(""))
@@ -137,6 +138,7 @@ func (w *StandardWriter) Write(event *ResultEvent) error {
 		if writeErr := w.outputFile.Write(data); writeErr != nil {
 			return errors.Wrap(err, "could not write to output")
 		}
+		_ = w.outputFile.Write([]byte("\n"))
 	}
 	return nil
 }

v2/pkg/protocols/dns/dnsclientpool/clientpool.go

@@ -32,11 +32,7 @@ func Init(options *types.Options) error {
 	poolMutex = &sync.RWMutex{}
 	clientPool = make(map[string]*retryabledns.Client)
 
-	if client, err := Get(options, &Configuration{}); err != nil {
+	normalClient = retryabledns.New(defaultResolvers, 1)
-		return err
-	} else {
-		normalClient = client
-	}
 	return nil
 }
 
@@ -58,7 +54,7 @@ func (c *Configuration) Hash() string {
 
 // Get creates or gets a client for the protocol based on custom configuration
 func Get(options *types.Options, configuration *Configuration) (*retryabledns.Client, error) {
-	if !(configuration.Retries > 0) {
+	if !(configuration.Retries > 1) {
 		return normalClient, nil
 	}
 	hash := configuration.Hash()

v2/pkg/protocols/dns/executer_test.go

@@ -0,0 +1,73 @@
+package dns
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/projectdiscovery/nuclei/v2/internal/progress"
+	"github.com/projectdiscovery/nuclei/v2/pkg/operators"
+	"github.com/projectdiscovery/nuclei/v2/pkg/operators/matchers"
+	"github.com/projectdiscovery/nuclei/v2/pkg/output"
+	"github.com/projectdiscovery/nuclei/v2/pkg/protocols"
+	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/dns/dnsclientpool"
+	"github.com/projectdiscovery/nuclei/v2/pkg/types"
+	"github.com/stretchr/testify/require"
+)
+
+func TestRequest(t *testing.T) {
+	err := dnsclientpool.Init(&types.Options{})
+	require.Nil(t, err, "could not initialize dns client pool")
+
+	writer, err := output.NewStandardWriter(true, false, false, "", "")
+	require.Nil(t, err, "could not create standard output writer")
+
+	progress, err := progress.NewProgress(false, false, 0)
+	require.Nil(t, err, "could not create standard progress writer")
+
+	protocolOpts := &protocols.ExecuterOptions{
+		TemplateID:   "testing-dns",
+		TemplateInfo: map[string]string{"author": "test"},
+		Output:       writer,
+		Options:      &types.Options{},
+		Progress:     progress,
+	}
+	req := &Request{Name: "{{FQDN}}", Recursion: true, Class: "inet", Type: "CNAME", Retries: 5, Operators: &operators.Operators{
+		Matchers: []*matchers.Matcher{{Type: "word", Words: []string{"github.io"}, Part: "body"}},
+	}}
+	err = req.Compile(protocolOpts)
+	require.Nil(t, err, "could not compile request")
+
+	output, err := req.ExecuteWithResults("docs.hackerone.com.", nil)
+	require.Nil(t, err, "could not execute request")
+
+	for _, result := range output {
+		fmt.Printf("%+v\n", result)
+	}
+}
+
+func TestExecuter(t *testing.T) {
+	err := dnsclientpool.Init(&types.Options{})
+	require.Nil(t, err, "could not initialize dns client pool")
+
+	writer, err := output.NewStandardWriter(true, false, false, "", "")
+	require.Nil(t, err, "could not create standard output writer")
+
+	progress, err := progress.NewProgress(false, false, 0)
+	require.Nil(t, err, "could not create standard progress writer")
+
+	protocolOpts := &protocols.ExecuterOptions{
+		TemplateID:   "testing-dns",
+		TemplateInfo: map[string]string{"author": "test"},
+		Output:       writer,
+		Options:      &types.Options{},
+		Progress:     progress,
+	}
+	executer := NewExecuter([]*Request{&Request{Name: "{{FQDN}}", Recursion: true, Class: "inet", Type: "CNAME", Retries: 5, Operators: &operators.Operators{
+		Matchers: []*matchers.Matcher{{Type: "word", Words: []string{"github.io"}, Part: "body"}},
+	}}}, protocolOpts)
+	err = executer.Compile()
+	require.Nil(t, err, "could not compile request")
+
+	_, err = executer.Execute("docs.hackerone.com")
+	require.Nil(t, err, "could not execute request")
+}

v2/pkg/protocols/dns/operators.go

@@ -12,14 +12,9 @@ import (
 
 // Match matches a generic data response again a given matcher
 func (r *Request) Match(data map[string]interface{}, matcher *matchers.Matcher) bool {
-	part, ok := data[matcher.Part]
+	partString := matcher.Part
-	if !ok {
-		return false
-	}
-	partString := part.(string)
-
 	switch partString {
-	case "body", "all":
+	case "body", "all", "":
 		partString = "raw"
 	}
 
@@ -123,7 +118,7 @@ func (r *Request) responseToDSLMap(req, resp *dns.Msg, host, matched string) out
 
 // makeResultEvent creates a result event from internal wrapped event
 func (r *Request) makeResultEvent(wrapped *output.InternalWrappedEvent) []*output.ResultEvent {
-	results := make([]*output.ResultEvent, len(wrapped.OperatorsResult.Matches)+1)
+	results := make([]*output.ResultEvent, 0, len(wrapped.OperatorsResult.Matches)+1)
 
 	data := output.ResultEvent{
 		TemplateID:       r.options.TemplateID,

v2/pkg/protocols/http/build_request.go

@@ -168,7 +168,7 @@ func (r *requestGenerator) makeHTTPRequestFromModel(ctx context.Context, data st
 	if err != nil {
 		return nil, err
 	}
-	return &generatedRequest{request: request}, nil
+	return &generatedRequest{request: request, original: r.request}, nil
 }
 
 // makeHTTPRequestFromRaw creates a *http.Request from a raw request
@@ -221,11 +221,7 @@ func (r *requestGenerator) handleRawWithPaylods(ctx context.Context, rawRequest,
 
 	// rawhttp
 	if r.request.Unsafe {
-		unsafeReq := &generatedRequest{
+		unsafeReq := &generatedRequest{rawRequest: rawRequestData, meta: genValues, original: r.request}
-			rawRequest: rawRequestData,
-			meta:       genValues,
-			original:   r.request,
-		}
 		return unsafeReq, nil
 	}
 
@@ -252,7 +248,7 @@ func (r *requestGenerator) handleRawWithPaylods(ctx context.Context, rawRequest,
 	if err != nil {
 		return nil, err
 	}
-	return &generatedRequest{request: request, meta: genValues}, nil
+	return &generatedRequest{request: request, meta: genValues, original: r.request}, nil
 }
 
 // fillRequest fills various headers in the request with values

v2/pkg/protocols/http/executer.go

@@ -43,7 +43,7 @@ func (e *Executer) Execute(input string) (bool, error) {
 	var results bool
 
 	for _, req := range e.requests {
-		events, err := req.ExecuteHTTP(input, nil)
+		events, err := req.ExecuteWithResults(input, nil)
 		if err != nil {
 			return false, err
 		}
@@ -71,7 +71,7 @@ func (e *Executer) ExecuteWithResults(input string) ([]*output.InternalWrappedEv
 	var results []*output.InternalWrappedEvent
 
 	for _, req := range e.requests {
-		events, err := req.ExecuteHTTP(input, nil)
+		events, err := req.ExecuteWithResults(input, nil)
 		if err != nil {
 			return nil, err
 		}

v2/pkg/protocols/http/executer_test.go

@@ -0,0 +1,47 @@
+package http
+
+import (
+	"testing"
+
+	"github.com/projectdiscovery/nuclei/v2/internal/progress"
+	"github.com/projectdiscovery/nuclei/v2/pkg/operators"
+	"github.com/projectdiscovery/nuclei/v2/pkg/operators/matchers"
+	"github.com/projectdiscovery/nuclei/v2/pkg/output"
+	"github.com/projectdiscovery/nuclei/v2/pkg/protocols"
+	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/http/httpclientpool"
+	"github.com/projectdiscovery/nuclei/v2/pkg/types"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/ratelimit"
+)
+
+func TestRequest(t *testing.T) {
+	err := httpclientpool.Init(&types.Options{})
+	require.Nil(t, err, "could not initialize dns client pool")
+
+	writer, err := output.NewStandardWriter(true, false, false, "", "")
+	require.Nil(t, err, "could not create standard output writer")
+
+	progress, err := progress.NewProgress(false, false, 0)
+	require.Nil(t, err, "could not create standard progress writer")
+
+	protocolOpts := &protocols.ExecuterOptions{
+		TemplateID:   "testing-dns",
+		TemplateInfo: map[string]string{"author": "test"},
+		Output:       writer,
+		Options:      &types.Options{},
+		Progress:     progress,
+		RateLimiter:  ratelimit.New(100),
+	}
+	executer := NewExecuter([]*Request{&Request{Path: []string{"{{BaseURL}}"}, Method: "GET", Operators: &operators.Operators{
+		Matchers: []*matchers.Matcher{{Type: "dsl", DSL: []string{"!contains(tolower(all_headers), 'x-frame-options')"}, Part: "body"}},
+	}}}, protocolOpts)
+	err = executer.Compile()
+	require.Nil(t, err, "could not compile request")
+
+	_, err = executer.Execute("https://example.com")
+	require.Nil(t, err, "could not execute request")
+
+	//	for _, result := range output {
+	//		fmt.Printf("%+v\n", result)
+	//	}
+}

v2/pkg/protocols/http/httpclientpool/clientpool.go

@@ -37,11 +37,11 @@ func Init(options *types.Options) error {
 	poolMutex = &sync.RWMutex{}
 	clientPool = make(map[string]*retryablehttp.Client)
 
-	if client, err := Get(options, &Configuration{}); err != nil {
+	client, err := wrappedGet(options, &Configuration{})
+	if err != nil {
 		return err
-	} else {
-		normalClient = client
 	}
+	normalClient = client
 	return nil
 }
 
@@ -82,6 +82,11 @@ func Get(options *types.Options, configuration *Configuration) (*retryablehttp.C
 	if !(configuration.Threads > 0 && configuration.MaxRedirects > 0 && configuration.FollowRedirects) {
 		return normalClient, nil
 	}
+	return wrappedGet(options, configuration)
+}
+
+// wrappedGet wraps a get operation without normal cliet check
+func wrappedGet(options *types.Options, configuration *Configuration) (*retryablehttp.Client, error) {
 	var proxyURL *url.URL
 	var err error
 

v2/pkg/protocols/http/operators.go

@@ -14,12 +14,7 @@ import (
 
 // Match matches a generic data response again a given matcher
 func (r *Request) Match(data map[string]interface{}, matcher *matchers.Matcher) bool {
-	part, ok := data[matcher.Part]
+	partString := matcher.Part
-	if !ok {
-		return false
-	}
-	partString := part.(string)
-
 	switch partString {
 	case "header":
 		partString = "all_headers"
@@ -56,12 +51,7 @@ func (r *Request) Match(data map[string]interface{}, matcher *matchers.Matcher)
 
 // Extract performs extracting operation for a extractor on model and returns true or false.
 func (r *Request) Extract(data map[string]interface{}, extractor *extractors.Extractor) map[string]struct{} {
-	part, ok := data[extractor.Part]
+	partString := extractor.Part
-	if !ok {
-		return nil
-	}
-	partString := part.(string)
-
 	switch partString {
 	case "header":
 		partString = "all_headers"
@@ -85,12 +75,14 @@ func (r *Request) Extract(data map[string]interface{}, extractor *extractors.Ext
 }
 
 // responseToDSLMap converts a HTTP response to a map for use in DSL matching
-func (r *Request) responseToDSLMap(resp *http.Response, rawReq, rawResp, body, headers string, duration time.Duration, extra map[string]interface{}) map[string]interface{} {
+func (r *Request) responseToDSLMap(resp *http.Response, host, matched, rawReq, rawResp, body, headers string, duration time.Duration, extra map[string]interface{}) map[string]interface{} {
-	data := make(map[string]interface{}, len(extra)+6+len(resp.Header)+len(resp.Cookies()))
+	data := make(map[string]interface{}, len(extra)+8+len(resp.Header)+len(resp.Cookies()))
 	for k, v := range extra {
 		data[k] = v
 	}
 
+	data["host"] = host
+	data["matched"] = matched
 	if r.options.Options.JSONRequests {
 		data["request"] = rawReq
 		data["response"] = rawResp
@@ -119,7 +111,7 @@ func (r *Request) responseToDSLMap(resp *http.Response, rawReq, rawResp, body, h
 
 // makeResultEvent creates a result event from internal wrapped event
 func (r *Request) makeResultEvent(wrapped *output.InternalWrappedEvent) []*output.ResultEvent {
-	results := make([]*output.ResultEvent, len(wrapped.OperatorsResult.Matches)+1)
+	results := make([]*output.ResultEvent, 0, len(wrapped.OperatorsResult.Matches)+1)
 
 	data := output.ResultEvent{
 		TemplateID:       r.options.TemplateID,

v2/pkg/protocols/http/request.go

@@ -160,8 +160,8 @@ func (e *Request) executeTurboHTTP(reqURL string, dynamicValues map[string]inter
 	return outputs, requestErr
 }
 
-// ExecuteHTTP executes the HTTP request on a URL
+// ExecuteWithResults executes the final request on a URL
-func (e *Request) ExecuteHTTP(reqURL string, dynamicValues map[string]interface{}) ([]*output.InternalWrappedEvent, error) {
+func (e *Request) ExecuteWithResults(reqURL string, dynamicValues map[string]interface{}) ([]*output.InternalWrappedEvent, error) {
 	// verify if pipeline was requested
 	if e.Pipeline {
 		return e.executeTurboHTTP(reqURL, dynamicValues)
@@ -343,7 +343,14 @@ func (e *Request) executeRequest(reqURL string, request *generatedRequest, dynam
 	//	matchData = generators.MergeMaps(matchData, result.historyData)
 	//	result.Unlock()
 	//}
-	ouputEvent := e.responseToDSLMap(resp, unsafeToString(dumpedRequest), unsafeToString(dumpedResponse), unsafeToString(data), headersToString(resp.Header), duration, request.meta)
+	var matchedURL string
+	if request.rawRequest != nil {
+		matchedURL = request.rawRequest.FullURL
+	}
+	if request.request != nil {
+		matchedURL = request.request.URL.String()
+	}
+	ouputEvent := e.responseToDSLMap(resp, reqURL, matchedURL, unsafeToString(dumpedRequest), unsafeToString(dumpedResponse), unsafeToString(data), headersToString(resp.Header), duration, request.meta)
 
 	event := []*output.InternalWrappedEvent{{InternalEvent: ouputEvent}}
 	if e.Operators != nil {