From dc46bd263b97de6d0fa50d2e7829e718ed366289 Mon Sep 17 00:00:00 2001 From: Sajad Parra Date: Thu, 24 Feb 2022 12:31:08 +0530 Subject: [PATCH] add options to specify User-Agent in headless template --- v2/pkg/model/types/userAgent/user_agent.go | 95 ++++++++++++++++++++++ v2/pkg/protocols/headless/engine/engine.go | 14 +++- v2/pkg/protocols/headless/headless.go | 26 ++++++ v2/pkg/protocols/headless/request.go | 3 + 4 files changed, 134 insertions(+), 4 deletions(-) create mode 100644 v2/pkg/model/types/userAgent/user_agent.go diff --git a/v2/pkg/model/types/userAgent/user_agent.go b/v2/pkg/model/types/userAgent/user_agent.go new file mode 100644 index 00000000..998bb67a --- /dev/null +++ b/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 +} diff --git a/v2/pkg/protocols/headless/engine/engine.go b/v2/pkg/protocols/headless/engine/engine.go index 6c1046b6..8d72f85b 100644 --- a/v2/pkg/protocols/headless/engine/engine.go +++ b/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() diff --git a/v2/pkg/protocols/headless/headless.go b/v2/pkg/protocols/headless/headless.go index dc017470..5133cfdd 100644 --- a/v2/pkg/protocols/headless/headless.go +++ b/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 { diff --git a/v2/pkg/protocols/headless/request.go b/v2/pkg/protocols/headless/request.go index 7422447e..66179090 100644 --- a/v2/pkg/protocols/headless/request.go +++ b/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()