code: fix variables merge order (#4623)

* fix variables merge order

* format screen: quote and trim extracted result

* code: interpret env vars in debug mode

* update integration test

cmd/integration-test/headless.go

@@ -113,7 +113,7 @@ func (h *headlessExtractValues) Execute(filePath string) error {
 		return err
 	}
 
-	return expectResultsCount(results, 3)
+	return expectResultsCount(results, 1)
 }
 
 type headlessPayloads struct{}

cmd/integration-test/http.go

@@ -368,6 +368,7 @@ func (h *httpDSLFunctions) Execute(filePath string) error {
 	}
 
 	for _, header := range extracted {
+		header = strings.Trim(header, `"`)
 		parts := strings.Split(header, ": ")
 		index, err := strconv.Atoi(parts[0])
 		if err != nil {

pkg/output/format_screen.go

@@ -3,6 +3,7 @@ package output
 import (
 	"bytes"
 	"strconv"
+	"strings"
 
 	"github.com/projectdiscovery/nuclei/v3/pkg/types"
 	mapsutil "github.com/projectdiscovery/utils/maps"
@@ -57,6 +58,9 @@ func (w *StandardWriter) formatScreen(output *ResultEvent) []byte {
 		builder.WriteString(" [")
 
 		for i, item := range output.ExtractedResults {
+			// trim trailing space
+			item = strings.TrimSpace(item)
+			item = strconv.QuoteToASCII(item)
 			builder.WriteString(w.aurora.BrightCyan(item).String())
 
 			if i != len(output.ExtractedResults)-1 {

pkg/protocols/code/code.go

@@ -3,6 +3,7 @@ package code
 import (
 	"context"
 	"fmt"
+	"regexp"
 	"strings"
 	"time"
 
@@ -28,6 +29,14 @@ import (
 	errorutil "github.com/projectdiscovery/utils/errors"
 )
 
+const (
+	pythonEnvRegex = `os\.getenv\(['"]([^'"]+)['"]\)`
+)
+
+var (
+	pythonEnvRegexCompiled = regexp.MustCompile(pythonEnvRegex)
+)
+
 // Request is a request for the SSL protocol
 type Request struct {
 	// Operators for the current request go here.
@@ -125,17 +134,20 @@ func (request *Request) ExecuteWithResults(input *contextargs.Context, dynamicVa
 
 	var interactshURLs []string
 
-	// inject all template context values as gozero env variables
-	variables := protocolutils.GenerateVariables(input.MetaInput.Input, false, nil)
+	// inject all template context values as gozero env allvars
+	allvars := protocolutils.GenerateVariables(input.MetaInput.Input, false, nil)
 	// add template context values
-	variables = generators.MergeMaps(variables, request.options.GetTemplateCtx(input.MetaInput).GetAll())
+	allvars = generators.MergeMaps(allvars, request.options.GetTemplateCtx(input.MetaInput).GetAll())
 	// optionvars are vars passed from CLI or env variables
 	optionVars := generators.BuildPayloadFromOptions(request.options.Options)
-	variablesMap := request.options.Variables.Evaluate(variables)
-	variables = generators.MergeMaps(variablesMap, variables, optionVars, request.options.Constants)
-	for name, value := range variables {
+	variablesMap := request.options.Variables.Evaluate(allvars)
+	// since we evaluate variables using allvars, give precedence to variablesMap
+	allvars = generators.MergeMaps(allvars, variablesMap, optionVars, request.options.Constants)
+	for name, value := range allvars {
 		v := fmt.Sprint(value)
 		v, interactshURLs = request.options.Interactsh.Replace(v, interactshURLs)
+		// if value is updated by interactsh, update allvars to reflect the change downstream
+		allvars[name] = v
 		metaSrc.AddVariable(gozerotypes.Variable{Name: name, Value: v})
 	}
 	gOutput, err := request.gozero.Eval(context.Background(), request.src, metaSrc)
@@ -145,11 +157,11 @@ func (request *Request) ExecuteWithResults(input *contextargs.Context, dynamicVa
 	gologger.Verbose().Msgf("[%s] Executed code on local machine %v", request.options.TemplateID, input.MetaInput.Input)
 
 	if vardump.EnableVarDump {
-		gologger.Debug().Msgf("Code Protocol request variables: \n%s\n", vardump.DumpVariables(variables))
+		gologger.Debug().Msgf("Code Protocol request variables: \n%s\n", vardump.DumpVariables(allvars))
 	}
 
 	if request.options.Options.Debug || request.options.Options.DebugRequests {
-		gologger.Debug().Msgf("[%s] Dumped Executed Source Code for %v\n\n%v\n", request.options.TemplateID, input.MetaInput.Input, request.Source)
+		gologger.Debug().Msgf("[%s] Dumped Executed Source Code for %v\n\n%v\n", request.options.TemplateID, input.MetaInput.Input, interpretEnvVars(request.Source, allvars))
 	}
 
 	dataOutputString := fmtStdout(gOutput.Stdout.String())
@@ -272,3 +284,24 @@ func (request *Request) MakeResultEventItem(wrapped *output.InternalWrappedEvent
 func fmtStdout(data string) string {
 	return strings.Trim(data, " \n\r\t")
 }
+
+// interpretEnvVars replaces environment variables in the input string
+func interpretEnvVars(source string, vars map[string]interface{}) string {
+	// bash mode
+	if strings.Contains(source, "$") {
+		for k, v := range vars {
+			source = strings.ReplaceAll(source, "$"+k, fmt.Sprintf("'%s'", v))
+		}
+	}
+	// python mode
+	if strings.Contains(source, "os.getenv") {
+		matches := pythonEnvRegexCompiled.FindAllStringSubmatch(source, -1)
+		for _, match := range matches {
+			if len(match) == 0 {
+				continue
+			}
+			source = strings.ReplaceAll(source, fmt.Sprintf("os.getenv('%s')", match), fmt.Sprintf("'%s'", vars[match[0]]))
+		}
+	}
+	return source
+}