add context vars in code and multi (#5051)

make the extracted variables available in subsequence templates when executing in a workflow

fix projectdiscovery/nuclei#4797

cmd/integration-test/workflow.go

@@ -3,12 +3,15 @@ package main
 import (
 	"fmt"
 	"io"
+	"log"
 	"net/http"
 	"net/http/httptest"
 	"strings"
 
 	"github.com/julienschmidt/httprouter"
 
+	"github.com/projectdiscovery/nuclei/v3/pkg/templates"
+	"github.com/projectdiscovery/nuclei/v3/pkg/templates/signer"
 	"github.com/projectdiscovery/nuclei/v3/pkg/testutils"
 )
 
@@ -20,9 +23,37 @@ var workflowTestcases = []TestCaseInfo{
 	{Path: "workflow/complex-conditions.yaml", TestCase: &workflowComplexConditions{}},
 	{Path: "workflow/http-value-share-workflow.yaml", TestCase: &workflowHttpKeyValueShare{}},
 	{Path: "workflow/dns-value-share-workflow.yaml", TestCase: &workflowDnsKeyValueShare{}},
+	{Path: "workflow/code-value-share-workflow.yaml", TestCase: &workflowCodeKeyValueShare{}, DisableOn: isCodeDisabled}, // isCodeDisabled declared in code.go
+	{Path: "workflow/multiprotocol-value-share-workflow.yaml", TestCase: &workflowMultiProtocolKeyValueShare{}},
 	{Path: "workflow/shared-cookie.yaml", TestCase: &workflowSharedCookies{}},
 }
 
+func init() {
+	// sign code templates (unless they are disabled)
+	if !isCodeDisabled() {
+		// allow local file access to load content of file references in template
+		// in order to sign them for testing purposes
+		templates.TemplateSignerLFA()
+
+		// testCertFile and testKeyFile are declared in code.go
+		tsigner, err := signer.NewTemplateSignerFromFiles(testCertFile, testKeyFile)
+		if err != nil {
+			panic(err)
+		}
+
+		// only the code templates are necessary to be signed
+		var templatesToSign = []string{
+			"workflow/code-template-1.yaml",
+			"workflow/code-template-2.yaml",
+		}
+		for _, templatePath := range templatesToSign {
+			if err := templates.SignTemplate(tsigner, templatePath); err != nil {
+				log.Fatalf("Could not sign template %v got: %s\n", templatePath, err)
+			}
+		}
+	}
+}
+
 type workflowBasic struct{}
 
 // Execute executes a test case and returns an error if occurred
@@ -159,6 +190,45 @@ func (h *workflowDnsKeyValueShare) Execute(filePath string) error {
 	return expectResultsCount(results, 1)
 }
 
+type workflowCodeKeyValueShare struct{}
+
+// Execute executes a test case and returns an error if occurred
+func (h *workflowCodeKeyValueShare) Execute(filePath string) error {
+	// provide the Certificate File that the code templates are signed with
+	certEnvVar := signer.CertEnvVarName + "=" + testCertFile
+
+	results, err := testutils.RunNucleiArgsWithEnvAndGetResults(debug, []string{certEnvVar}, "-workflows", filePath, "-target", "input", "-code")
+	if err != nil {
+		return err
+	}
+
+	return expectResultsCount(results, 1)
+}
+
+type workflowMultiProtocolKeyValueShare struct{}
+
+// Execute executes a test case and returns an error if occurred
+func (h *workflowMultiProtocolKeyValueShare) Execute(filePath string) error {
+	router := httprouter.New()
+	// the response of path1 contains a domain that will be extracted and shared with the second template
+	router.GET("/path1", func(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
+		fmt.Fprintf(w, "href=\"blog.projectdiscovery.io\"")
+	})
+	// path2 responds with the value of the "extracted" query parameter, e.g.: /path2?extracted=blog.projectdiscovery.io => blog.projectdiscovery.io
+	router.GET("/path2", func(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
+		fmt.Fprintf(w, "%s", r.URL.Query().Get("extracted"))
+	})
+	ts := httptest.NewServer(router)
+	defer ts.Close()
+
+	results, err := testutils.RunNucleiWorkflowAndGetResults(filePath, ts.URL, debug)
+	if err != nil {
+		return err
+	}
+
+	return expectResultsCount(results, 2)
+}
+
 type workflowSharedCookies struct{}
 
 // Execute executes a test case and returns an error if occurred

integration_tests/workflow/code-template-1.yaml

@@ -0,0 +1,22 @@
+id: code-template-1
+
+info:
+  name: code-template-1
+  author: tovask
+  severity: info
+  tags: code
+
+code:
+  - engine:
+      - py
+      - python3
+      - python
+    source: |
+      print("hello from first")
+    extractors:
+      - type: regex
+        name: extracted
+        regex:
+          - 'hello from (.*)'
+        group: 1
+# digest: 490a0046304402202c63d47bb0acdd40b3b852d95490d492ff5741b84071b2a8a40371be7797c13602202b6b977e157edf2ef70a402a2e57d4eb5a67c5ca91f0a2f9a10a966e8485ebaf:4a3eb6b4988d95847d4203be25ed1d46

integration_tests/workflow/code-template-2.yaml

@@ -0,0 +1,21 @@
+id: code-template-2
+
+info:
+  name: code-template-2
+  author: tovask
+  severity: info
+  tags: code
+
+code:
+  - engine:
+      - py
+      - python3
+      - python
+    source: |
+      import os
+      print("hello from " + os.getenv("extracted"))
+    matchers:
+      - type: word
+        words:
+          - "hello from first"
+# digest: 490a00463044022025661eab353b7f359c0d428a86b6287545d7f759375e8025cc8c9c77b616ca6502200bc2c019059622df3c88e7caa6dd7d1fb9b956010aa0de2ee2b9f7dd0a3c4954:4a3eb6b4988d95847d4203be25ed1d46

integration_tests/workflow/code-value-share-workflow.yaml

@@ -0,0 +1,12 @@
+id: code-value-sharing-workflow
+
+info:
+  name: Code Value Sharing Workflow
+  author: tovask
+  severity: info
+  tags: code
+
+workflows:
+  - template: workflow/code-template-1.yaml
+    subtemplates:
+      - template: workflow/code-template-2.yaml

integration_tests/workflow/multiprotocol-value-share-template.yaml

@@ -0,0 +1,22 @@
+id: multiprotocol-value-sharing-template
+
+info:
+  name: MultiProtocol Value Sharing Template
+  author: tovask
+  severity: info
+
+dns:
+  - name: "{{extracted}}"
+    type: PTR
+    matchers:
+      - type: word
+        words:
+          - "blog.projectdiscovery.io"
+
+http:
+  - path:
+      - "{{BaseURL}}/path2?extracted={{extracted}}"
+    matchers:
+      - type: word
+        words:
+          - "blog.projectdiscovery.io"

integration_tests/workflow/multiprotocol-value-share-workflow.yaml

@@ -0,0 +1,11 @@
+id: multiprotocol-value-sharing-workflow
+
+info:
+  name: MultiProtocol Value Sharing Workflow
+  author: tovask
+  severity: info
+
+workflows:
+  - template: workflow/http-value-share-template-1.yaml
+    subtemplates:
+      - template: workflow/multiprotocol-value-share-template.yaml

pkg/protocols/code/code.go

@@ -160,6 +160,8 @@ func (request *Request) ExecuteWithResults(input *contextargs.Context, dynamicVa
 	if request.options.HasTemplateCtx(input.MetaInput) {
 		allvars = generators.MergeMaps(allvars, request.options.GetTemplateCtx(input.MetaInput).GetAll())
 	}
+	// add dynamic and previous variables
+	allvars = generators.MergeMaps(allvars, dynamicValues, previous)
 	// optionvars are vars passed from CLI or env variables
 	optionVars := generators.BuildPayloadFromOptions(request.options.Options)
 	variablesMap := request.options.Variables.Evaluate(allvars)

pkg/tmplexec/multiproto/multi.go

@@ -46,6 +46,12 @@ func (m *MultiProtocol) Compile() error {
 func (m *MultiProtocol) ExecuteWithResults(ctx *scan.ScanContext) error {
 	// put all readonly args into template context
 	m.options.GetTemplateCtx(ctx.Input.MetaInput).Merge(m.readOnlyArgs)
+
+	// add all input args to template context
+	ctx.Input.ForEach(func(key string, value interface{}) {
+		m.options.GetTemplateCtx(ctx.Input.MetaInput).Set(key, value)
+	})
+
 	// callback to process results from all protocols
 	multiProtoCallback := func(event *output.InternalWrappedEvent) {
 		if event == nil {