Adding variables support for headless templates (#2064)

v2/pkg/protocols/common/generators/maps.go

@@ -47,15 +47,14 @@ func MergeMapsMany(maps ...interface{}) map[string][]string {
 }
 
 // MergeMaps merges two maps into a new map
-func MergeMaps(m1, m2 map[string]interface{}) map[string]interface{} {
-	m := make(map[string]interface{}, len(m1)+len(m2))
-	for k, v := range m1 {
-		m[k] = v
+func MergeMaps(maps ...map[string]interface{}) map[string]interface{} {
+	merged := make(map[string]interface{})
+	for _, m := range maps {
+		for k, v := range m {
+			merged[k] = v
 		}
-	for k, v := range m2 {
-		m[k] = v
 	}
-	return m
+	return merged
 }
 
 // ExpandMapValues converts values from flat string to string slice

v2/pkg/protocols/common/generators/maps_test.go

@@ -7,10 +7,10 @@ import (
 )
 
 func TestMergeMapsMany(t *testing.T) {
-	got := MergeMapsMany(map[string]interface{}{"a": []string{"1", "2"}, "c": "5"}, map[string][]string{"b": []string{"3", "4"}})
+	got := MergeMapsMany(map[string]interface{}{"a": []string{"1", "2"}, "c": "5"}, map[string][]string{"b": {"3", "4"}})
 	require.Equal(t, map[string][]string{
-		"a": []string{"1", "2"},
-		"b": []string{"3", "4"},
-		"c": []string{"5"},
+		"a": {"1", "2"},
+		"b": {"3", "4"},
+		"c": {"5"},
 	}, got, "could not get correct merged map")
 }

v2/pkg/protocols/dns/dns.go

@@ -245,8 +245,8 @@ func classToInt(class string) uint16 {
 	return uint16(result)
 }
 
-// GenerateDNSVariables from a dns name
-func GenerateDNSVariables(domain string) map[string]interface{} {
+// GenerateVariables from a dns name
+func GenerateVariables(domain string) map[string]interface{} {
 	parsed, err := publicsuffix.Parse(strings.TrimSuffix(domain, "."))
 	if err != nil {
 		return map[string]interface{}{"FQDN": domain}

v2/pkg/protocols/dns/dns_test.go

@@ -11,7 +11,7 @@ import (
 )
 
 func TestGenerateDNSVariables(t *testing.T) {
-	vars := GenerateDNSVariables("www.projectdiscovery.io")
+	vars := GenerateVariables("www.projectdiscovery.io")
 	require.Equal(t, map[string]interface{}{
 		"FQDN": "www.projectdiscovery.io",
 		"RDN":  "projectdiscovery.io",

v2/pkg/protocols/dns/request.go

@@ -43,7 +43,7 @@ func (request *Request) ExecuteWithResults(input string, metadata /*TODO review
 	if err != nil {
 		return errors.Wrap(err, "could not build request")
 	}
-	vars := GenerateDNSVariables(domain)
+	vars := GenerateVariables(domain)
 	variablesMap := request.options.Variables.Evaluate(vars)
 	vars = generators.MergeMaps(variablesMap, vars)
 

v2/pkg/protocols/headless/request.go

@@ -14,6 +14,7 @@ import (
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/helpers/eventcreator"
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/helpers/responsehighlighter"
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/interactsh"
+	httpProtocol "github.com/projectdiscovery/nuclei/v2/pkg/protocols/http"
 	templateTypes "github.com/projectdiscovery/nuclei/v2/pkg/templates/types"
 )
 
@@ -31,9 +32,11 @@ func (request *Request) ExecuteWithResults(inputURL string, metadata, previous o
 	if request.options.Browser.UserAgent() == "" {
 		request.options.Browser.SetUserAgent(request.compiledUserAgent)
 	}
-	payloads := generators.BuildPayloadFromOptions(request.options.Options)
 
-	variablesMap := request.options.Variables.Evaluate(generators.MergeMaps(metadata, payloads))
+	vars := GenerateVariables(inputURL)
+	payloads := generators.BuildPayloadFromOptions(request.options.Options)
+	values := generators.MergeMaps(vars, metadata, payloads)
+	variablesMap := request.options.Variables.Evaluate(values)
 	payloads = generators.MergeMaps(variablesMap, payloads)
 
 	if request.generator != nil {
@@ -141,3 +144,13 @@ func dumpResponse(event *output.InternalWrappedEvent, requestOptions *protocols.
 		gologger.Debug().Msgf("[%s] Dumped Headless response for %s\n\n%s", requestOptions.TemplateID, input, highlightedResponse)
 	}
 }
+
+// GenerateVariables will create default variables
+func GenerateVariables(URL string) map[string]interface{} {
+	parsed, err := url.Parse(URL)
+	if err != nil {
+		return nil
+	}
+
+	return httpProtocol.GenerateVariables(parsed, false)
+}

v2/pkg/protocols/http/build_request.go

@@ -93,7 +93,7 @@ func (r *requestGenerator) Make(baseURL, data string, payloads, dynamicValues ma
 	}
 
 	values := generators.MergeMaps(
-		generators.MergeMaps(dynamicValues, generateVariables(parsed, trailingSlash)),
+		generators.MergeMaps(dynamicValues, GenerateVariables(parsed, trailingSlash)),
 		generators.BuildPayloadFromOptions(r.request.options.Options),
 	)
 
@@ -156,7 +156,7 @@ func (r *requestGenerator) makeSelfContainedRequest(data string, payloads, dynam
 			return nil, fmt.Errorf("could not parse request URL: %w", err)
 		}
 		values = generators.MergeMaps(
-			generators.MergeMaps(dynamicValues, generateVariables(parsed, false)),
+			generators.MergeMaps(dynamicValues, GenerateVariables(parsed, false)),
 			values,
 		)
 
@@ -377,8 +377,8 @@ func setHeader(req *http.Request, name, value string) {
 	}
 }
 
-// generateVariables will create default variables after parsing a url
-func generateVariables(parsed *url.URL, trailingSlash bool) map[string]interface{} {
+// GenerateVariables will create default variables after parsing a url
+func GenerateVariables(parsed *url.URL, trailingSlash bool) map[string]interface{} {
 	domain := parsed.Host
 	if strings.Contains(parsed.Host, ":") {
 		domain = strings.Split(parsed.Host, ":")[0]
@@ -416,5 +416,5 @@ func generateVariables(parsed *url.URL, trailingSlash bool) map[string]interface
 		"File":     base,
 		"Scheme":   parsed.Scheme,
 	}
-	return generators.MergeMaps(httpVariables, dns.GenerateDNSVariables(domain))
+	return generators.MergeMaps(httpVariables, dns.GenerateVariables(domain))
 }

v2/pkg/protocols/http/build_request_test.go

@@ -27,7 +27,7 @@ func TestBaseURLWithTemplatePrefs(t *testing.T) {
 func TestVariables(t *testing.T) {
 	baseURL := "http://localhost:9001/test/123"
 	parsed, _ := url.Parse(baseURL)
-	values := generateVariables(parsed, true)
+	values := GenerateVariables(parsed, true)
 
 	require.Equal(t, values["BaseURL"], parsed.String(), "incorrect baseurl")
 	require.Equal(t, values["RootURL"], "http://localhost:9001", "incorrect rootURL")
@@ -40,7 +40,7 @@ func TestVariables(t *testing.T) {
 
 	baseURL = "https://example.com"
 	parsed, _ = url.Parse(baseURL)
-	values = generateVariables(parsed, false)
+	values = GenerateVariables(parsed, false)
 
 	require.Equal(t, values["BaseURL"], parsed.String(), "incorrect baseurl")
 	require.Equal(t, values["Host"], "example.com", "incorrect domain name")
@@ -52,7 +52,7 @@ func TestVariables(t *testing.T) {
 
 	baseURL = "ftp://foobar.com/"
 	parsed, _ = url.Parse(baseURL)
-	values = generateVariables(parsed, true)
+	values = GenerateVariables(parsed, true)
 
 	require.Equal(t, values["BaseURL"], parsed.String(), "incorrect baseurl")
 	require.Equal(t, values["Host"], "foobar.com", "incorrect domain name")