fix proto prefix in template context

v2/pkg/protocols/code/code.go

@@ -32,6 +32,8 @@ type Request struct {
 	operators.Operators `yaml:",inline,omitempty"`
 	CompiledOperators   *operators.Operators `yaml:"-"`
 
+	// ID is the optional id of the request
+	ID string `yaml:"id,omitempty" json:"id,omitempty" jsonschema:"title=id of the dns request,description=ID is the optional ID of the DNS Request"`
 	// description: |
 	//   Engine type
 	Engine []string `yaml:"engine,omitempty" jsonschema:"title=engine,description=Engine,enum=python,enum=powershell,enum=command"`
@@ -96,7 +98,7 @@ func (request *Request) Requests() int {
 
 // GetID returns the ID for the request if any.
 func (request *Request) GetID() string {
-	return ""
+	return request.ID
 }
 
 // ExecuteWithResults executes the protocol requests and returns results instead of writing them.

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

@@ -42,7 +42,7 @@ func setup() {
 	executerOpts.WorkflowLoader = workflowLoader
 }
 
-func TestFlowTemplateWithID(t *testing.T) {
+func TestFlowTemplateWithIndex(t *testing.T) {
 	// test
 	setup()
 	Template, err := templates.Parse("testcases/nuclei-flow-dns.yaml", nil, executerOpts)
@@ -63,3 +63,43 @@ func TestFlowTemplateWithID(t *testing.T) {
 		require.True(t, len(value.([]string)) > 0)
 	}
 }
+
+func TestFlowTemplateWithID(t *testing.T) {
+	setup()
+	Template, err := templates.Parse("testcases/nuclei-flow-dns-id.yaml", nil, executerOpts)
+	require.Nil(t, err, "could not parse template")
+
+	require.True(t, Template.Flow != "", "not a flow template") // this is classifer if template is flow or not
+
+	err = Template.Executer.Compile()
+	require.Nil(t, err, "could not compile template")
+
+	gotresults, err := Template.Executer.Execute(contextargs.NewWithInput("hackerone.com"))
+	require.Nil(t, err, "could not execute template")
+	require.True(t, gotresults)
+
+	value, ok := Template.Options.TemplateCtx.Get("nameservers")
+	require.True(t, ok)
+	if value != nil {
+		require.True(t, len(value.([]string)) > 0)
+	}
+}
+
+func TestFlowWithProtoPrefix(t *testing.T) {
+	// test
+	setup()
+	Template, err := templates.Parse("testcases/nuclei-flow-dns.yaml", nil, executerOpts)
+	require.Nil(t, err, "could not parse template")
+
+	require.True(t, Template.Flow != "", "not a flow template") // this is classifer if template is flow or not
+
+	err = Template.Executer.Compile()
+	require.Nil(t, err, "could not compile template")
+
+	gotresults, err := Template.Executer.Execute(contextargs.NewWithInput("hackerone.com"))
+	require.Nil(t, err, "could not execute template")
+	require.True(t, gotresults)
+
+	// while there are lot of variables lets just look for only these
+	// protoVars := []string{"dns_"}
+}

v2/pkg/protocols/common/executer/testcases/nuclei-flow-dns-prefix.yaml

@@ -0,0 +1,41 @@
+id: nuclei-flow-dns
+
+info:
+  name: Nuclei flow dns
+  author: pdteam
+  severity: info
+  description: Description of the Template
+  reference: https://example-reference-link
+
+flow: |
+  dns("fetch-ns");
+  template["nameservers"].forEach(nameserver => {
+    set("nameserver",nameserver);
+    dns("1");
+  });
+
+dns:
+  - id: "fetch-ns"
+    name: "{{FQDN}}"
+    type: NS
+    matchers:
+      - type: word
+        words:
+          - "IN\tNS"
+    extractors:
+      - type: regex
+        internal: true
+        name: "nameservers"
+        group: 1
+        regex:
+          - "IN\tNS\t(.+)"
+
+  - name: "{{nameserver}}"
+    type: A
+    class: inet
+    retries: 3
+    recursion: true
+    extractors:
+      - type: dsl
+        dsl:
+          - "a"

v2/pkg/protocols/common/executer/testcases/nuclei-flow-dns.yaml

@@ -12,8 +12,8 @@ flow: |
   template["nameservers"].forEach(nameserver => {
     set("nameserver",nameserver);
     dns("1");
-    log(template)
   });
+  log(template)
 
 dns:
   - name: "{{FQDN}}"

v2/pkg/protocols/dns/request.go

@@ -151,7 +151,7 @@ func (request *Request) execute(domain string, metadata, previous output.Interna
 	outputEvent := request.responseToDSLMap(compiledRequest, response, domain, question, traceData)
 	// expose response variables in proto_var format
 	// this is no-op if the template is not a multi protocol template
-	request.options.AddTemplateVars(request.Type(), outputEvent)
+	request.options.AddTemplateVars(request.Type(), request.ID, outputEvent)
 	for k, v := range previous {
 		outputEvent[k] = v
 	}

v2/pkg/protocols/headless/request.go

@@ -156,7 +156,7 @@ func (request *Request) executeRequestWithPayloads(input *contextargs.Context, p
 
 	outputEvent := request.responseToDSLMap(responseBody, out["header"], out["status_code"], reqBuilder.String(), input.MetaInput.Input, input.MetaInput.Input, page.DumpHistory())
 	// add response fields to template context and merge templatectx variables to output event
-	request.options.AddTemplateVars(request.Type(), outputEvent)
+	request.options.AddTemplateVars(request.Type(), request.ID, outputEvent)
 	outputEvent = generators.MergeMaps(outputEvent, request.options.TemplateCtx.GetAll())
 	for k, v := range out {
 		outputEvent[k] = v

v2/pkg/protocols/http/request.go

@@ -721,7 +721,7 @@ func (request *Request) executeRequest(input *contextargs.Context, generatedRequ
 
 		outputEvent := request.responseToDSLMap(response.resp, input.MetaInput.Input, matchedURL, tostring.UnsafeToString(dumpedRequest), tostring.UnsafeToString(response.fullResponse), tostring.UnsafeToString(response.body), tostring.UnsafeToString(response.headers), duration, generatedRequest.meta)
 		// add response fields to template context and merge templatectx variables to output event
-		request.options.AddTemplateVars(request.Type(), outputEvent)
+		request.options.AddTemplateVars(request.Type(), request.ID, outputEvent)
 		outputEvent = generators.MergeMaps(outputEvent, request.options.TemplateCtx.GetAll())
 		if i := strings.LastIndex(hostname, ":"); i != -1 {
 			hostname = hostname[:i]

v2/pkg/protocols/network/request.go

@@ -281,7 +281,7 @@ func (request *Request) executeRequestWithPayloads(variables map[string]interfac
 	response := responseBuilder.String()
 	outputEvent := request.responseToDSLMap(reqBuilder.String(), string(final[:n]), response, input, actualAddress)
 	// add response fields to template context and merge templatectx variables to output event
-	request.options.AddTemplateVars(request.Type(), outputEvent)
+	request.options.AddTemplateVars(request.Type(), request.ID, outputEvent)
 	outputEvent = generators.MergeMaps(outputEvent, request.options.TemplateCtx.GetAll())
 	outputEvent["ip"] = request.dialer.GetDialedIP(hostname)
 	if request.options.StopAtFirstMatch {

v2/pkg/protocols/offlinehttp/request.go

@@ -88,7 +88,7 @@ func (request *Request) ExecuteWithResults(input *contextargs.Context, metadata
 
 			outputEvent := request.responseToDSLMap(resp, data, data, data, tostring.UnsafeToString(dumpedResponse), tostring.UnsafeToString(body), utils.HeadersToString(resp.Header), 0, nil)
 			// add response fields to template context and merge templatectx variables to output event
-			request.options.AddTemplateVars(request.Type(), outputEvent)
+			request.options.AddTemplateVars(request.Type(), request.GetID(), outputEvent)
 			outputEvent = generators.MergeMaps(outputEvent, request.options.TemplateCtx.GetAll())
 			outputEvent["ip"] = ""
 			for k, v := range previous {

v2/pkg/protocols/protocols.go

@@ -97,17 +97,19 @@ type ExecutorOptions struct {
 
 // AddTemplateVars adds vars to template context with given template type as prefix
 // this method is no-op if template is not multi protocol
-func (e *ExecutorOptions) AddTemplateVars(templateType templateTypes.ProtocolType, vars map[string]interface{}) {
+func (e *ExecutorOptions) AddTemplateVars(templateType templateTypes.ProtocolType, reqID string, vars map[string]interface{}) {
 	// if we wan't to disable adding response variables and other variables to template context
 	// this is the statement that does it . template context is currently only enabled for
 	// multiprotocol and flow templates
 	if e.ProtocolType != templateTypes.MultiProtocol && e.Flow == "" {
-		// no-op if not multi protocol template
+		// no-op if not multi protocol template or flow template
 		return
 	}
 	for k, v := range vars {
 		if !stringsutil.EqualFoldAny(k, "template-id", "template-info", "template-path") {
-			if templateType < templateTypes.InvalidProtocol {
+			if reqID != "" {
+				k = reqID + "_" + k
+			} else if templateType < templateTypes.InvalidProtocol {
 				k = templateType.String() + "_" + k
 			}
 			e.TemplateCtx.Set(k, v)
@@ -117,13 +119,15 @@ func (e *ExecutorOptions) AddTemplateVars(templateType templateTypes.ProtocolTyp
 
 // AddTemplateVar adds given var to template context with given template type as prefix
 // this method is no-op if template is not multi protocol
-func (e *ExecutorOptions) AddTemplateVar(prefix, key string, value interface{}) {
+func (e *ExecutorOptions) AddTemplateVar(templateType templateTypes.ProtocolType, reqID string, key string, value interface{}) {
 	if e.ProtocolType != templateTypes.MultiProtocol && e.Flow == "" {
-		// no-op if not multi protocol template
+		// no-op if not multi protocol template or flow template
 		return
 	}
-	if prefix != "" {
-		key = prefix + "_" + key
+	if reqID != "" {
+		key = reqID + "_" + key
+	} else if templateType < templateTypes.InvalidProtocol {
+		key = templateType.String() + "_" + key
 	}
 	e.TemplateCtx.Set(key, value)
 }

v2/pkg/protocols/ssl/ssl.go

@@ -270,7 +270,7 @@ func (request *Request) ExecuteWithResults(input *contextargs.Context, dynamicVa
 		if tag == "" || f.IsZero() {
 			continue
 		}
-		request.options.AddTemplateVar(request.Type().String(), tag, f.Value())
+		request.options.AddTemplateVar(request.Type(), request.ID, tag, f.Value())
 		data[tag] = f.Value()
 	}
 
@@ -289,7 +289,7 @@ func (request *Request) ExecuteWithResults(input *contextargs.Context, dynamicVa
 		if tag == "" || f.IsZero() {
 			continue
 		}
-		request.options.AddTemplateVar(request.Type().String(), tag, f.Value())
+		request.options.AddTemplateVar(request.Type(), request.ID, tag, f.Value())
 		data[tag] = f.Value()
 	}
 

v2/pkg/protocols/websocket/websocket.go

@@ -267,7 +267,7 @@ func (request *Request) executeRequestWithPayloads(input, hostname string, dynam
 	data["ip"] = request.dialer.GetDialedIP(hostname)
 
 	// add response fields to template context and merge templatectx variables to output event
-	request.options.AddTemplateVars(request.Type(), data)
+	request.options.AddTemplateVars(request.Type(), request.ID, data)
 	data = generators.MergeMaps(data, request.options.TemplateCtx.GetAll())
 
 	for k, v := range previous {

v2/pkg/protocols/whois/whois.go

@@ -137,7 +137,7 @@ func (request *Request) ExecuteWithResults(input *contextargs.Context, dynamicVa
 	data["response"] = jsonDataString
 
 	// add response fields to template context and merge templatectx variables to output event
-	request.options.AddTemplateVars(request.Type(), data)
+	request.options.AddTemplateVars(request.Type(), request.ID, data)
 	data = generators.MergeMaps(data, request.options.TemplateCtx.GetAll())
 
 	event := eventcreator.CreateEvent(request, data, request.options.Options.Debug || request.options.Options.DebugResponse)

v2/pkg/templates/compile.go

@@ -241,6 +241,11 @@ func ParseTemplateFromReader(reader io.Reader, preprocessor Preprocessor, option
 		options.Variables = template.Variables
 	}
 
+	// if more than 1 request per protocol exist we add request id to protocol request
+	// since in template context we have proto_prefix for each protocol it is overwritten
+	// if request id is not present
+	template.validateAllRequestIDs()
+
 	// TODO: we should add a syntax check here or somehow use a javascript linter
 	// simplest option for now seems to compile using goja and see if it fails
 	if strings.TrimSpace(template.Flow) != "" {

v2/pkg/templates/templates.go

@@ -3,6 +3,7 @@ package templates
 
 import (
 	"encoding/json"
+	"strconv"
 
 	validate "github.com/go-playground/validator/v10"
 	"github.com/projectdiscovery/nuclei/v2/pkg/model"
@@ -180,9 +181,84 @@ func (template *Template) Type() types.ProtocolType {
 	}
 }
 
+// validateAllRequestIDs check if that protocol already has given id if not
+// then is is manually set to proto_index
+func (template *Template) validateAllRequestIDs() {
+	// this is required in multiprotocol and flow where we save response variables
+	// and all other data in template context if template as two requests in a protocol
+	// then it is overwritten to avoid this we use proto_index as request ID
+	if len(template.RequestsCode) > 1 {
+		for i, req := range template.RequestsCode {
+			if req.ID == "" {
+				req.ID = template.Type().String() + "_" + strconv.Itoa(i)
+			}
+		}
+	}
+	if len(template.RequestsDNS) > 1 {
+		for i, req := range template.RequestsDNS {
+			if req.ID == "" {
+				req.ID = template.Type().String() + "_" + strconv.Itoa(i)
+			}
+		}
+	}
+	if len(template.RequestsFile) > 1 {
+		for i, req := range template.RequestsFile {
+			if req.ID == "" {
+				req.ID = template.Type().String() + "_" + strconv.Itoa(i)
+			}
+		}
+	}
+	if len(template.RequestsHTTP) > 1 {
+		for i, req := range template.RequestsHTTP {
+			if req.ID == "" {
+				req.ID = template.Type().String() + "_" + strconv.Itoa(i)
+			}
+		}
+	}
+	if len(template.RequestsHeadless) > 1 {
+		for i, req := range template.RequestsHeadless {
+			if req.ID == "" {
+				req.ID = template.Type().String() + "_" + strconv.Itoa(i)
+			}
+		}
+
+	}
+	if len(template.RequestsNetwork) > 1 {
+		for i, req := range template.RequestsNetwork {
+			if req.ID == "" {
+				req.ID = template.Type().String() + "_" + strconv.Itoa(i)
+			}
+		}
+	}
+	if len(template.RequestsSSL) > 1 {
+		for i, req := range template.RequestsSSL {
+			if req.ID == "" {
+				req.ID = template.Type().String() + "_" + strconv.Itoa(i)
+			}
+		}
+	}
+	if len(template.RequestsWebsocket) > 1 {
+		for i, req := range template.RequestsWebsocket {
+			if req.ID == "" {
+				req.ID = template.Type().String() + "_" + strconv.Itoa(i)
+			}
+		}
+	}
+	if len(template.RequestsWHOIS) > 1 {
+		for i, req := range template.RequestsWHOIS {
+			if req.ID == "" {
+				req.ID = template.Type().String() + "_" + strconv.Itoa(i)
+			}
+		}
+	}
+}
+
 // MarshalYAML forces recursive struct validation during marshal operation
 func (template *Template) MarshalYAML() ([]byte, error) {
 	out, marshalErr := yaml.Marshal(template)
+	// Review: we are adding requestIDs for templateContext
+	// if we are using this method then we might need to purge manually added IDS that start with `templatetype_`
+	// this is only applicable if there are more than 1 request fields in protocol
 	errValidate := validate.New().Struct(template)
 	return out, multierr.Append(marshalErr, errValidate)
 }