add options to specify User-Agent in headless template

v2/pkg/model/types/userAgent/user_agent.go

@@ -0,0 +1,95 @@
+package userAgent
+
+import (
+	"encoding/json"
+	"strings"
+
+	"github.com/alecthomas/jsonschema"
+	"github.com/pkg/errors"
+)
+
+type UserAgent int
+
+// name:UserAgent
+const (
+	// name:random
+	Random UserAgent = iota
+	// name:off
+	Off
+	// name:default
+	Default
+	// name:custom
+	Custom
+	limit
+)
+
+var userAgentMappings = map[UserAgent]string{
+	Random:  "random",
+	Off:     "off",
+	Default: "default",
+	Custom:  "custom",
+}
+
+func GetSupportedUserAgentOptions() []UserAgent {
+	var result []UserAgent
+	for index := UserAgent(1); index < limit; index++ {
+		result = append(result, index)
+	}
+	return result
+}
+
+func toUserAgent(valueToMap string) (UserAgent, error) {
+	normalizedValue := normalizeValue(valueToMap)
+	for key, currentValue := range userAgentMappings {
+		if normalizedValue == currentValue {
+			return key, nil
+		}
+	}
+	return -1, errors.New("Invalid userAgent: " + valueToMap)
+}
+
+func normalizeValue(value string) string {
+	return strings.TrimSpace(strings.ToLower(value))
+}
+
+func (userAgent UserAgent) String() string {
+	return userAgentMappings[userAgent]
+}
+
+// UserAgentHolder holds a UserAgent type. Required for un/marshalling purposes
+type UserAgentHolder struct {
+	Value UserAgent `mapping:"true"`
+}
+
+func (userAgentHolder UserAgentHolder) JSONSchemaType() *jsonschema.Type {
+	gotType := &jsonschema.Type{
+		Type:        "string",
+		Title:       "userAgent for the headless",
+		Description: "userAgent for the headless http request",
+	}
+	for _, userAgent := range GetSupportedUserAgentOptions() {
+		gotType.Enum = append(gotType.Enum, userAgent.String())
+	}
+	return gotType
+}
+
+func (userAgentHolder *UserAgentHolder) UnmarshalYAML(unmarshal func(interface{}) error) error {
+	var marshalledUserAgent string
+	if err := unmarshal(&marshalledUserAgent); err != nil {
+		return err
+	}
+	computedUserAgent, err := toUserAgent(marshalledUserAgent)
+	if err != nil {
+		return err
+	}
+	userAgentHolder.Value = computedUserAgent
+	return nil
+}
+
+func (userAgentHolder *UserAgentHolder) MarshalJSON() ([]byte, error) {
+	return json.Marshal(userAgentHolder.Value.String())
+}
+
+func (userAgentHolder UserAgentHolder) MarshalYAML() (interface{}, error) {
+	return userAgentHolder.Value.String(), nil
+}

v2/pkg/protocols/headless/engine/engine.go

@@ -8,7 +8,6 @@ import (
 	"runtime"
 	"strings"
 
-	"github.com/corpix/uarand"
 	"github.com/go-rod/rod"
 	"github.com/go-rod/rod/lib/launcher"
 	"github.com/pkg/errors"
@@ -89,9 +88,6 @@ func New(options *types.Options) (*Browser, error) {
 			customAgent = parts[1]
 		}
 	}
-	if customAgent == "" {
-		customAgent = uarand.GetRandom()
-	}
 
 	httpclient, err := newHttpClient(options)
 	if err != nil {
@@ -116,6 +112,16 @@ func MustDisableSandbox() bool {
 	return runtime.GOOS == "linux" && os.Geteuid() == 0
 }
 
+// SetUserAgent sets custom user agent to the browser
+func (b *Browser) SetUserAgent(customUserAgent string) {
+	b.customAgent = customUserAgent
+}
+
+// UserAgent fetch the currently set custom user agent
+func (b *Browser) UserAgent() string {
+	return b.customAgent
+}
+
 // Close closes the browser engine
 func (b *Browser) Close() {
 	b.engine.Close()

v2/pkg/protocols/headless/headless.go

@@ -1,9 +1,11 @@
 package headless
 
 import (
+	"github.com/corpix/uarand"
 	"github.com/pkg/errors"
 
 	"github.com/projectdiscovery/fileutil"
+	useragent "github.com/projectdiscovery/nuclei/v2/pkg/model/types/userAgent"
 	"github.com/projectdiscovery/nuclei/v2/pkg/operators"
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols"
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/generators"
@@ -33,6 +35,15 @@ type Request struct {
 	//   Steps is the list of actions to run for headless request
 	Steps []*engine.Action `yaml:"steps,omitempty" jsonschema:"title=list of actions for headless request,description=List of actions to run for headless request"`
 
+	// descriptions: |
+	// 	 User-Agent is the type of user-agent to use for the request.
+	UserAgent useragent.UserAgentHolder `yaml:"user_agent,omitempty" jsonschema:"title=user agent for the headless request,description=User agent for the headless request"`
+
+	// description: |
+	// 	 If UserAgent is set to custom, customUserAgent is the custom user-agent to use for the request.
+	CustomUserAgent   string `yaml:"custom_user_agent,omitempty" jsonschema:"title=custom user agent for the headless request,description=Custom user agent for the headless request"`
+	compiledUserAgent string
+
 	// Operators for the current request go here.
 	operators.Operators `yaml:",inline,omitempty"`
 	CompiledOperators   *operators.Operators `yaml:"-"`
@@ -90,6 +101,21 @@ func (request *Request) Compile(options *protocols.ExecuterOptions) error {
 		}
 	}
 
+	// Compile User-Agent
+	switch request.UserAgent.Value {
+	case useragent.Off:
+		request.compiledUserAgent = " "
+	case useragent.Default:
+		request.compiledUserAgent = ""
+	case useragent.Custom:
+		if request.CustomUserAgent == "" {
+			return errors.New("please set custom_user_agent in the template")
+		}
+		request.compiledUserAgent = request.CustomUserAgent
+	case useragent.Random:
+		request.compiledUserAgent = uarand.GetRandom()
+	}
+
 	if len(request.Matchers) > 0 || len(request.Extractors) > 0 {
 		compiled := &request.Operators
 		if err := compiled.Compile(); err != nil {

v2/pkg/protocols/headless/request.go

@@ -26,6 +26,9 @@ func (request *Request) Type() templateTypes.ProtocolType {
 
 // ExecuteWithResults executes the protocol requests and returns results instead of writing them.
 func (request *Request) ExecuteWithResults(inputURL string, metadata, previous output.InternalEvent /*TODO review unused parameter*/, callback protocols.OutputEventCallback) error {
+	if request.options.Browser.UserAgent() == "" {
+		request.options.Browser.SetUserAgent(request.compiledUserAgent)
+	}
 	payloads := generators.BuildPayloadFromOptions(request.options.Options)
 	if request.generator != nil {
 		iterator := request.generator.NewIterator()