refactor and move encoding to MakeResultEventItem func

projectdiscovery · Nov 6, 2023 · 08eb802 · 08eb802
1 parent 3ebfb7a
commit 08eb802
Show file tree

Hide file tree

Showing 14 changed files with 35 additions and 18 deletions.
diff --git a/pkg/output/output.go b/pkg/output/output.go
@@ -1,7 +1,6 @@
 package output
 
 import (
-	"encoding/base64"
 	"fmt"
 	"io"
 	"os"
@@ -61,13 +60,9 @@ type StandardWriter struct {
 	severityColors   func(severity.Severity) string
 	storeResponse    bool
 	storeResponseDir string
-	omitTemplate     bool
 }
 
-var (
-	decolorizerRegex               = regexp.MustCompile(`\x1B\[[0-9;]*[a-zA-Z]`)
-	MaxTemplateFileSizeForEncoding = 1024 * 1024
-)
+var decolorizerRegex = regexp.MustCompile(`\x1B\[[0-9;]*[a-zA-Z]`)
 
 // InternalEvent is an internal output generation structure for nuclei.
 type InternalEvent map[string]interface{}
@@ -214,7 +209,6 @@ func NewStandardWriter(options *types.Options) (*StandardWriter, error) {
 		severityColors:   colorizer.New(auroraColorizer),
 		storeResponse:    options.StoreResponse,
 		storeResponseDir: options.StoreResponseDir,
-		omitTemplate:     options.OmitTemplate,
 	}
 	return writer, nil
 }
@@ -224,11 +218,6 @@ func (w *StandardWriter) Write(event *ResultEvent) error {
 	// Enrich the result event with extra metadata on the template-path and url.
 	if event.TemplatePath != "" {
 		event.Template, event.TemplateURL = utils.TemplatePathURL(types.ToString(event.TemplatePath), types.ToString(event.TemplateID))
-		if event.TemplateURL == "" && !w.omitTemplate {
-			if data, err := os.ReadFile(event.TemplatePath); err == nil && len(data) <= MaxTemplateFileSizeForEncoding {
-				event.TemplateEncoded = base64.StdEncoding.EncodeToString(data)
-			}
-		}
 	}
 
 	event.Timestamp = time.Now()

diff --git a/pkg/protocols/code/code.go b/pkg/protocols/code/code.go
@@ -254,6 +254,7 @@ func (request *Request) MakeResultEventItem(wrapped *output.InternalWrappedEvent
 		ExtractedResults: wrapped.OperatorsResult.OutputExtracts,
 		Timestamp:        time.Now(),
 		MatcherStatus:    true,
+		TemplateEncoded:  request.options.EncodeTemplate(),
 	}
 	return data
 }

diff --git a/pkg/protocols/dns/operators.go b/pkg/protocols/dns/operators.go
@@ -122,6 +122,7 @@ func (request *Request) MakeResultEventItem(wrapped *output.InternalWrappedEvent
 		Timestamp:        time.Now(),
 		Request:          types.ToString(wrapped.InternalEvent["request"]),
 		Response:         types.ToString(wrapped.InternalEvent["raw"]),
+		TemplateEncoded:  request.options.EncodeTemplate(),
 	}
 	return data
 }

diff --git a/pkg/protocols/file/operators.go b/pkg/protocols/file/operators.go
@@ -107,6 +107,7 @@ func (request *Request) MakeResultEventItem(wrapped *output.InternalWrappedEvent
 		ExtractedResults: wrapped.OperatorsResult.OutputExtracts,
 		Response:         types.ToString(wrapped.InternalEvent["raw"]),
 		Timestamp:        time.Now(),
+		TemplateEncoded:  request.options.EncodeTemplate(),
 	}
 	return data
 }
diff --git a/pkg/protocols/headless/operators.go b/pkg/protocols/headless/operators.go
@@ -138,6 +138,7 @@ func (request *Request) MakeResultEventItem(wrapped *output.InternalWrappedEvent
 		IP:               types.ToString(wrapped.InternalEvent["ip"]),
 		Request:          types.ToString(wrapped.InternalEvent["request"]),
 		Response:         types.ToString(wrapped.InternalEvent["data"]),
+		TemplateEncoded:  request.options.EncodeTemplate(),
 	}
 	return data
 }
diff --git a/pkg/protocols/http/operators.go b/pkg/protocols/http/operators.go
@@ -164,6 +164,7 @@ func (request *Request) MakeResultEventItem(wrapped *output.InternalWrappedEvent
 		Request:          types.ToString(wrapped.InternalEvent["request"]),
 		Response:         request.truncateResponse(wrapped.InternalEvent["response"]),
 		CURLCommand:      types.ToString(wrapped.InternalEvent["curl-command"]),
+		TemplateEncoded:  request.options.EncodeTemplate(),
 	}
 	return data
 }

diff --git a/pkg/protocols/javascript/js.go b/pkg/protocols/javascript/js.go
@@ -634,6 +634,7 @@ func (request *Request) MakeResultEventItem(wrapped *output.InternalWrappedEvent
 		Request:          types.ToString(wrapped.InternalEvent["request"]),
 		Response:         types.ToString(wrapped.InternalEvent["response"]),
 		IP:               types.ToString(wrapped.InternalEvent["ip"]),
+		TemplateEncoded:  request.options.EncodeTemplate(),
 	}
 	return data
 }

diff --git a/pkg/protocols/network/operators.go b/pkg/protocols/network/operators.go
@@ -108,6 +108,7 @@ func (request *Request) MakeResultEventItem(wrapped *output.InternalWrappedEvent
 		IP:               types.ToString(wrapped.InternalEvent["ip"]),
 		Request:          types.ToString(wrapped.InternalEvent["request"]),
 		Response:         types.ToString(wrapped.InternalEvent["data"]),
+		TemplateEncoded:  request.options.EncodeTemplate(),
 	}
 	return data
 }
diff --git a/pkg/protocols/offlinehttp/operators.go b/pkg/protocols/offlinehttp/operators.go
@@ -151,6 +151,7 @@ func (request *Request) MakeResultEventItem(wrapped *output.InternalWrappedEvent
 		IP:               types.ToString(wrapped.InternalEvent["ip"]),
 		Request:          types.ToString(wrapped.InternalEvent["request"]),
 		Response:         types.ToString(wrapped.InternalEvent["raw"]),
+		TemplateEncoded:  request.options.EncodeTemplate(),
 	}
 	return data
 }
diff --git a/pkg/protocols/protocols.go b/pkg/protocols/protocols.go
@@ -1,6 +1,7 @@
 package protocols
 
 import (
+	"encoding/base64"
 	"sync/atomic"
 
 	"github.com/projectdiscovery/ratelimit"
@@ -30,6 +31,8 @@ import (
 	"github.com/projectdiscovery/nuclei/v3/pkg/types"
 )
 
+var MaxTemplateFileSizeForEncoding = 1024 * 1024
+
 // Executer is an interface implemented any protocol based request executer.
 type Executer interface {
 	// Compile compiles the execution generators preparing any requests possible.
@@ -50,6 +53,8 @@ type ExecutorOptions struct {
 	TemplatePath string
 	// TemplateInfo contains information block of the template request
 	TemplateInfo model.Info
+	// RawTemplate is the raw template for the request
+	RawTemplate []byte
 	// Output is a writer interface for writing output events from executer.
 	Output output.Writer
 	// Options contains configuration options for the executer.
@@ -294,3 +299,10 @@ func MakeDefaultMatchFunc(data map[string]interface{}, matcher *matchers.Matcher
 	}
 	return false, nil
 }
+
+func (e *ExecutorOptions) EncodeTemplate() string {
+	if !e.Options.OmitTemplate && len(e.RawTemplate) <= MaxTemplateFileSizeForEncoding {
+		return base64.StdEncoding.EncodeToString(e.RawTemplate)
+	}
+	return ""
+}
diff --git a/pkg/protocols/ssl/ssl.go b/pkg/protocols/ssl/ssl.go
@@ -375,6 +375,7 @@ func (request *Request) MakeResultEventItem(wrapped *output.InternalWrappedEvent
 		Timestamp:        time.Now(),
 		MatcherStatus:    true,
 		IP:               types.ToString(wrapped.InternalEvent["ip"]),
+		TemplateEncoded:  request.options.EncodeTemplate(),
 	}
 	return data
 }
diff --git a/pkg/protocols/websocket/websocket.go b/pkg/protocols/websocket/websocket.go
@@ -409,6 +409,7 @@ func (request *Request) MakeResultEventItem(wrapped *output.InternalWrappedEvent
 		IP:               types.ToString(wrapped.InternalEvent["ip"]),
 		Request:          types.ToString(wrapped.InternalEvent["request"]),
 		Response:         types.ToString(wrapped.InternalEvent["response"]),
+		TemplateEncoded:  request.options.EncodeTemplate(),
 	}
 	return data
 }

diff --git a/pkg/protocols/whois/whois.go b/pkg/protocols/whois/whois.go
@@ -185,6 +185,7 @@ func (request *Request) MakeResultEventItem(wrapped *output.InternalWrappedEvent
 		MatcherStatus:    true,
 		Request:          types.ToString(wrapped.InternalEvent["request"]),
 		Response:         types.ToString(wrapped.InternalEvent["response"]),
+		TemplateEncoded:  request.options.EncodeTemplate(),
 	}
 	return data
 }

diff --git a/pkg/templates/compile.go b/pkg/templates/compile.go
@@ -128,7 +128,7 @@ func (template *Template) Requests() int {
 }
 
 // compileProtocolRequests compiles all the protocol requests for the template
-func (template *Template) compileProtocolRequests(options protocols.ExecutorOptions) error {
+func (template *Template) compileProtocolRequests(options *protocols.ExecutorOptions) error {
 	templateRequests := template.Requests()
 
 	if templateRequests == 0 {
@@ -180,7 +180,7 @@ func (template *Template) compileProtocolRequests(options protocols.ExecutorOpti
 			requests = append(requests, template.convertRequestToProtocolsRequest(template.RequestsJavascript)...)
 		}
 	}
-	template.Executer = tmplexec.NewTemplateExecuter(requests, &options)
+	template.Executer = tmplexec.NewTemplateExecuter(requests, options)
 	return nil
 }
 
@@ -206,7 +206,7 @@ func (template *Template) convertRequestToProtocolsRequest(requests interface{})
 // compileOfflineHTTPRequest iterates all requests if offline http mode is
 // specified and collects all matchers for all the base request templates
 // (those with URL {{BaseURL}} and it's slash variation.)
-func (template *Template) compileOfflineHTTPRequest(options protocols.ExecutorOptions) error {
+func (template *Template) compileOfflineHTTPRequest(options *protocols.ExecutorOptions) error {
 	operatorsList := []*operators.Operators{}
 
 mainLoop:
@@ -225,7 +225,7 @@ mainLoop:
 	}
 	if len(operatorsList) > 0 {
 		options.Operators = operatorsList
-		template.Executer = tmplexec.NewTemplateExecuter([]protocols.Request{&offlinehttp.Request{}}, &options)
+		template.Executer = tmplexec.NewTemplateExecuter([]protocols.Request{&offlinehttp.Request{}}, options)
 		return nil
 	}
 
@@ -360,7 +360,7 @@ func parseTemplate(data []byte, options protocols.ExecutorOptions) (*Template, e
 		return nil, errorutil.NewWithErr(err).Msgf("failed to load file refs for %s", template.ID)
 	}
 
-	if err := template.compileProtocolRequests(options); err != nil {
+	if err := template.compileProtocolRequests(template.Options); err != nil {
 		return nil, err
 	}
 
@@ -377,13 +377,18 @@ func parseTemplate(data []byte, options protocols.ExecutorOptions) (*Template, e
 
 	// check if the template is verified
 	// only valid templates can be verified or signed
-	for _, verifier := range signer.DefaultTemplateVerifiers {
+	var verifier *signer.TemplateSigner
+	for _, verifier = range signer.DefaultTemplateVerifiers {
 		template.Verified, _ = verifier.Verify(data, template)
 		if template.Verified {
 			SignatureStats[verifier.Identifier()].Add(1)
 			break
 		}
 	}
+
+	if !(template.Verified && verifier.Identifier() == "projectdiscovery/nuclei-templates") {
+		template.Options.RawTemplate = data
+	}
 	return template, nil
 }