nuclei/pkg/protocols/websocket/websocket.go

package websocket

import (
	"context"
	"crypto/tls"
	"fmt"
	"io"
	"maps"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/gobwas/ws"
	"github.com/gobwas/ws/wsutil"
	"github.com/pkg/errors"

	"github.com/projectdiscovery/fastdialer/fastdialer"
	"github.com/projectdiscovery/gologger"
	"github.com/projectdiscovery/nuclei/v3/pkg/operators"
	"github.com/projectdiscovery/nuclei/v3/pkg/operators/extractors"
	"github.com/projectdiscovery/nuclei/v3/pkg/operators/matchers"
	"github.com/projectdiscovery/nuclei/v3/pkg/output"
	"github.com/projectdiscovery/nuclei/v3/pkg/protocols"
	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/contextargs"
	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/generators"
	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/helpers/eventcreator"
	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/helpers/responsehighlighter"
	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/render"
	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/utils/vardump"
	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/network/networkclientpool"
	protocolutils "github.com/projectdiscovery/nuclei/v3/pkg/protocols/utils"
	templateTypes "github.com/projectdiscovery/nuclei/v3/pkg/templates/types"
	"github.com/projectdiscovery/nuclei/v3/pkg/types"
	urlutil "github.com/projectdiscovery/utils/url"
)

// Request is a request for the Websocket protocol
type Request struct {
	// Operators for the current request go here.
	operators.Operators `yaml:",inline,omitempty" json:",inline,omitempty"`
	CompiledOperators   *operators.Operators `yaml:"-" json:"-"`

	// ID is the optional id of the request
	ID string `yaml:"id,omitempty" json:"id,omitempty" jsonschema:"title=id of the request,description=ID of the network request"`
	// description: |
	//   Address contains address for the request
	Address string `yaml:"address,omitempty" json:"address,omitempty" jsonschema:"title=address for the websocket request,description=Address contains address for the request"`
	// description: |
	//   Inputs contains inputs for the websocket protocol
	Inputs []*Input `yaml:"inputs,omitempty" json:"inputs,omitempty" jsonschema:"title=inputs for the websocket request,description=Inputs contains any input/output for the current request"`
	// description: |
	//   Headers contains headers for the request.
	Headers map[string]string `yaml:"headers,omitempty" json:"headers,omitempty" jsonschema:"title=headers contains the request headers,description=Headers contains headers for the request"`

	// description: |
	//   Attack is the type of payload combinations to perform.
	//
	//   Sniper is each payload once, pitchfork combines multiple payload sets and clusterbomb generates
	//   permutations and combinations for all payloads.
	AttackType generators.AttackTypeHolder `yaml:"attack,omitempty" json:"attack,omitempty" jsonschema:"title=attack is the payload combination,description=Attack is the type of payload combinations to perform,enum=sniper,enum=pitchfork,enum=clusterbomb"`
	// description: |
	//   Payloads contains any payloads for the current request.
	//
	//   Payloads support both key-values combinations where a list
	//   of payloads is provided, or optionally a single file can also
	//   be provided as payload which will be read on run-time.
	Payloads map[string]interface{} `yaml:"payloads,omitempty" json:"payloads,omitempty" jsonschema:"title=payloads for the websocket request,description=Payloads contains any payloads for the current request"`

	generator *generators.PayloadGenerator

	// cache any variables that may be needed for operation.
	dialer  *fastdialer.Dialer
	options *protocols.ExecutorOptions
}

// Input is an input for the websocket protocol
type Input struct {
	// description: |
	//   Data is the data to send as the input.
	//
	//   It supports DSL Helper Functions as well as normal expressions.
	// examples:
	//   - value: "\"TEST\""
	//   - value: "\"hex_decode('50494e47')\""
	Data string `yaml:"data,omitempty" json:"data,omitempty" jsonschema:"title=data to send as input,description=Data is the data to send as the input"`
	// description: |
	//   Name is the optional name of the data read to provide matching on.
	// examples:
	//   - value: "\"prefix\""
	Name string `yaml:"name,omitempty" json:"name,omitempty" jsonschema:"title=optional name for data read,description=Optional name of the data read to provide matching on"`
}

const (
	parseUrlErrorMessage                   = "could not parse input url"
	evaluateTemplateExpressionErrorMessage = "could not evaluate template expressions"
)

// Compile compiles the request generators preparing any requests possible.
func (request *Request) Compile(options *protocols.ExecutorOptions) error {
	request.options = options

	client, err := networkclientpool.Get(options.Options, &networkclientpool.Configuration{
		CustomDialer: options.CustomFastdialer,
	})
	if err != nil {
		return errors.Wrap(err, "could not get network client")
	}
	request.dialer = client

	if len(request.Payloads) > 0 {
		request.generator, err = generators.New(request.Payloads, request.AttackType.Value, request.options.TemplatePath, options.Catalog, options.Options.AttackType, options.Options)
		if err != nil {
			return errors.Wrap(err, "could not parse payloads")
		}
	}

	if len(request.Matchers) > 0 || len(request.Extractors) > 0 {
		compiled := &request.Operators
		compiled.ExcludeMatchers = options.ExcludeMatchers
		compiled.TemplateID = options.TemplateID
		if err := compiled.Compile(); err != nil {
			return errors.Wrap(err, "could not compile operators")
		}
		// llm operators only evaluate once they hold the scan's client; without
		// it every llm path returns a negative result instead of an error.
		protocols.BindLLMOperators(compiled, options)
		request.CompiledOperators = compiled
	}
	return nil
}

// Requests returns the total number of requests the rule will perform
func (request *Request) Requests() int {
	if request.generator != nil {
		return request.generator.NewIterator().Total()
	}
	return 1
}

// GetID returns the ID for the request if any.
func (request *Request) GetID() string {
	return request.ID
}

// ExecuteWithResults executes the protocol requests and returns results instead of writing them.
func (request *Request) ExecuteWithResults(input *contextargs.Context, dynamicValues, previous output.InternalEvent, callback protocols.OutputEventCallback) error {
	hostname, err := getAddress(input.MetaInput.Input)
	if err != nil {
		return err
	}

	if request.generator != nil {
		iterator := request.generator.NewIterator()

		for {
			value, ok := iterator.Value()
			if !ok {
				break
			}
			if err := request.executeRequestWithPayloads(input, hostname, value, previous, callback); err != nil {
				return err
			}
		}
	} else {
		value := make(map[string]interface{})
		if err := request.executeRequestWithPayloads(input, hostname, value, previous, callback); err != nil {
			return err
		}
	}
	return nil
}

// ExecuteWithResults executes the protocol requests and returns results instead of writing them.
func (request *Request) executeRequestWithPayloads(target *contextargs.Context, hostname string, dynamicValues, previous output.InternalEvent, callback protocols.OutputEventCallback) error {
	header := http.Header{}
	input := target.MetaInput.Input

	parsed, err := urlutil.Parse(input)
	if err != nil {
		return errors.Wrap(err, parseUrlErrorMessage)
	}

	defaultVars := protocolutils.GenerateVariables(parsed, false, nil)
	optionVars := generators.BuildPayloadFromOptions(request.options.Options)
	// add templatecontext variables to varMap
	scope := request.options.NewVariablesScope(defaultVars, optionVars, dynamicValues, previous)

	request.options.AddTemplateCtxToVariablesScope(target.MetaInput, scope)
	scope.AddData(request.options.Constants)

	variables := request.options.Variables.EvaluateScope(scope).Values
	payloadValues := generators.MergeMaps(variables, defaultVars, optionVars, dynamicValues, previous)
	if request.options.HasTemplateCtx(target.MetaInput) {
		payloadValues = generators.MergeMaps(payloadValues, request.options.GetTemplateCtx(target.MetaInput).GetAll())
	}
	payloadValues = generators.MergeMaps(payloadValues, request.options.Constants)

	requestOptions := request.options
	// subprotocols requested via a Sec-WebSocket-Protocol header are negotiated
	// through the dialer rather than sent as a raw header: gobwas/ws validates the
	// server-selected protocol against Dialer.Protocols and rejects one that was
	// only set as a header, so a server that echoes the requested subprotocol
	// (e.g. Vite's "vite-hmr") would otherwise fail the handshake.
	var subprotocols []string
	for key, value := range request.Headers {
		result, dataErr := render.Render(render.Input{Options: request.options.GetOptions(), Text: value, Values: payloadValues})
		if dataErr != nil {
			requestOptions.Output.Request(requestOptions.TemplateID, input, request.Type().String(), dataErr)
			requestOptions.Progress.IncrementFailedRequestsBy(1)

			return errors.Wrap(dataErr, evaluateTemplateExpressionErrorMessage)
		}

		if http.CanonicalHeaderKey(key) == "Sec-Websocket-Protocol" {
			for _, proto := range strings.Split(result.Text, ",") {
				// Subprotocol values are written into the handshake request line by
				// gobwas/ws, so a CR or LF would smuggle extra handshake headers;
				// reject any such value instead of negotiating it.
				if trimmed := strings.TrimSpace(proto); trimmed != "" && !strings.ContainsAny(trimmed, "\r\n") {
					subprotocols = append(subprotocols, trimmed)
				}
			}
			continue
		}

		header.Set(key, result.Text)
	}

	tlsConfig := &tls.Config{
		InsecureSkipVerify: true,
		ServerName:         hostname,
		MinVersion:         tls.VersionTLS10,
	}

	if requestOptions.Options.SNI != "" {
		tlsConfig.ServerName = requestOptions.Options.SNI
	}

	// Dial reports both a refused TCP connection and a rejected HTTP upgrade.
	// Only the TCP result belongs to host backoff: an upgrade failure means
	// the port already accepted the connection.
	var dialed bool
	var dialErr error
	websocketDialer := ws.Dialer{
		Header:    ws.HandshakeHeaderHTTP(header),
		Protocols: subprotocols,
		Timeout:   time.Duration(requestOptions.Options.Timeout) * time.Second,
		NetDial: func(ctx context.Context, network, address string) (net.Conn, error) {
			dialed = true
			conn, err := request.dialer.Dial(ctx, network, address)
			dialErr = err
			return conn, err
		},
		TLSConfig: tlsConfig,
	}

	if vardump.EnableVarDump {
		gologger.Debug().Msgf("WebSocket Protocol request variables: %s\n", vardump.DumpVariables(payloadValues))
	}

	addressResult, dataErr := render.Render(render.Input{Options: request.options.GetOptions(), Text: request.Address, Values: payloadValues})
	if dataErr != nil {
		requestOptions.Output.Request(requestOptions.TemplateID, input, request.Type().String(), dataErr)
		requestOptions.Progress.IncrementFailedRequestsBy(1)

		return errors.Wrap(dataErr, evaluateTemplateExpressionErrorMessage)
	}

	addressToDial := addressResult.Text

	parsedAddress, err := url.Parse(addressToDial)
	if err != nil {
		requestOptions.Output.Request(requestOptions.TemplateID, input, request.Type().String(), err)
		requestOptions.Progress.IncrementFailedRequestsBy(1)

		return errors.Wrap(err, parseUrlErrorMessage)
	}

	if parsedAddress.Path == "" || parsedAddress.Path == "/" {
		parsedAddress.Path = parsed.Path
	}

	addressToDial = parsedAddress.String()

	if err := requestOptions.HostBackoffWait(target.Context(), addressToDial); err != nil {
		return err
	}
	if err := requestOptions.RateLimitTakeContext(target.Context(), addressToDial); err != nil {
		requestOptions.Output.Request(requestOptions.TemplateID, input, request.Type().String(), err)
		requestOptions.Progress.IncrementFailedRequestsBy(1)

		return errors.Wrap(err, "could not take rate limit")
	}

	timeStart := time.Now()
	conn, readBuffer, _, err := websocketDialer.Dial(target.Context(), addressToDial)
	handshakeDuration := time.Since(timeStart)
	if dialed {
		requestOptions.HostBackoffObserveDial(addressToDial, dialErr)
	}
	if err != nil {
		requestOptions.Output.Request(requestOptions.TemplateID, input, request.Type().String(), err)
		requestOptions.Progress.IncrementFailedRequestsBy(1)

		return errors.Wrap(err, "could not connect to server")
	}

	defer func() {
		_ = conn.Close()
	}()

	responseBuilder := &strings.Builder{}
	if readBuffer != nil {
		_, _ = io.Copy(responseBuilder, readBuffer) // Copy initial response
	}

	events, requestOutput, stepDurations, err := request.readWriteInputWebsocket(conn, payloadValues, input, responseBuilder)
	if err != nil {
		requestOptions.Output.Request(requestOptions.TemplateID, input, request.Type().String(), err)
		requestOptions.Progress.IncrementFailedRequestsBy(1)

		return errors.Wrap(err, "could not read write response")
	}

	requestOptions.Progress.IncrementRequests()

	if requestOptions.Options.Debug || requestOptions.Options.DebugRequests {
		gologger.Debug().Str("address", input).Msgf("[%s] Dumped Websocket request for %s", requestOptions.TemplateID, input)
		gologger.Print().Msgf("%s", requestOutput)
	}

	requestOptions.Output.Request(requestOptions.TemplateID, input, request.Type().String(), err)
	gologger.Verbose().Msgf("Sent Websocket request to %s", input)

	data := make(map[string]interface{})

	data["type"] = request.Type().String()
	data["success"] = "true"
	data["request"] = requestOutput
	data["response"] = responseBuilder.String()
	data["host"] = input
	data["matched"] = addressToDial
	if target.MetaInput.CustomIP != "" {
		data["ip"] = target.MetaInput.CustomIP
	} else {
		data["ip"] = request.dialer.GetDialedIP(hostname)
	}
	if len(stepDurations) > 0 {
		addWebsocketDurationFields(data, stepDurations)
	} else {
		data["duration"] = handshakeDuration.Seconds()
	}

	// add response fields to template context and merge templatectx variables to output event
	request.options.AddTemplateVars(target.MetaInput, request.Type(), request.ID, data)
	data = generators.MergeMaps(data, request.options.GetTemplateCtx(target.MetaInput).GetAll())

	maps.Copy(data, previous)
	maps.Copy(data, events)

	event := eventcreator.CreateEventWithAdditionalOptions(request, data, requestOptions.Options.Debug || requestOptions.Options.DebugResponse, func(internalWrappedEvent *output.InternalWrappedEvent) {
		internalWrappedEvent.OperatorsResult.PayloadValues = payloadValues
	})
	if requestOptions.Options.Debug || requestOptions.Options.DebugResponse {
		responseOutput := responseBuilder.String()
		gologger.Debug().Msgf("[%s] Dumped Websocket response for %s", requestOptions.TemplateID, input)
		gologger.Print().Msgf("%s", responsehighlighter.Highlight(event.OperatorsResult, responseOutput, requestOptions.Options.NoColor, false))
	}

	callback(event)

	return nil
}

func (request *Request) readWriteInputWebsocket(conn net.Conn, payloadValues map[string]interface{}, input string, respBuilder *strings.Builder) (events map[string]interface{}, req string, durations []time.Duration, err error) {
	reqBuilder := &strings.Builder{}
	inputEvents := make(map[string]interface{})

	requestOptions := request.options
	for _, req := range request.Inputs {
		reqBuilder.Grow(len(req.Data))

		result, dataErr := render.Render(render.Input{Options: request.options.GetOptions(), Text: req.Data, Values: payloadValues})
		if dataErr != nil {
			requestOptions.Output.Request(requestOptions.TemplateID, input, request.Type().String(), dataErr)
			requestOptions.Progress.IncrementFailedRequestsBy(1)
			return nil, "", nil, errors.Wrap(dataErr, evaluateTemplateExpressionErrorMessage)
		}
		reqBuilder.WriteString(result.Text)

		timeStart := time.Now()
		err = wsutil.WriteClientMessage(conn, ws.OpText, []byte(result.Text))
		if err != nil {
			requestOptions.Output.Request(requestOptions.TemplateID, input, request.Type().String(), err)
			requestOptions.Progress.IncrementFailedRequestsBy(1)
			return nil, "", nil, errors.Wrap(err, "could not write request to server")
		}

		msg, opCode, err := wsutil.ReadServerData(conn)
		if err != nil {
			requestOptions.Output.Request(requestOptions.TemplateID, input, request.Type().String(), err)
			requestOptions.Progress.IncrementFailedRequestsBy(1)
			return nil, "", nil, errors.Wrap(err, "could not read response from server")
		}
		durations = append(durations, time.Since(timeStart))
		// Only perform matching and writes in case we receive
		// text or binary opcode from the websocket server.
		if opCode != ws.OpText && opCode != ws.OpBinary {
			continue
		}

		respBuilder.Write(msg)
		if req.Name != "" {
			bufferStr := string(msg)
			inputEvents[req.Name] = bufferStr

			// Run any internal extractors for the request here and add found values to map.
			if request.CompiledOperators != nil {
				values := request.CompiledOperators.ExecuteInternalExtractors(map[string]interface{}{req.Name: bufferStr}, request.Extract)
				maps.Copy(inputEvents, values)
			}
		}
	}
	return inputEvents, reqBuilder.String(), durations, nil
}

func addWebsocketDurationFields(event map[string]interface{}, durations []time.Duration) {
	for i, duration := range durations {
		seconds := duration.Seconds()
		event["duration_"+strconv.Itoa(i+1)] = seconds
		event["duration"] = seconds
	}
}

// getAddress returns the address of the host to make request to
func getAddress(toTest string) (string, error) {
	parsed, err := url.Parse(toTest)
	if err != nil {
		return "", errors.Wrap(err, parseUrlErrorMessage)
	}
	scheme := strings.ToLower(parsed.Scheme)

	if scheme != "ws" && scheme != "wss" {
		return "", fmt.Errorf("invalid url scheme provided: %s", scheme)
	}
	if parsed != nil && parsed.Host != "" {
		return parsed.Host, nil
	}
	return "", nil
}

// Match performs matching operation for a matcher on model and returns:
// true and a list of matched snippets if the matcher type is supports it
// otherwise false and an empty string slice
func (request *Request) Match(data map[string]interface{}, matcher *matchers.Matcher) (bool, []string) {
	return protocols.MakeDefaultMatchFuncWithExecutorOptions(data, matcher, request.options)
}

// Extract performs extracting operation for an extractor on model and returns true or false.
func (request *Request) Extract(data map[string]interface{}, matcher *extractors.Extractor) map[string]struct{} {
	return protocols.MakeDefaultExtractFuncWithExecutorOptions(data, matcher, request.options)
}

// MakeResultEvent creates a result event from internal wrapped event
func (request *Request) MakeResultEvent(wrapped *output.InternalWrappedEvent) []*output.ResultEvent {
	return protocols.MakeDefaultResultEvent(request, wrapped)
}

// GetCompiledOperators returns a list of the compiled operators
func (request *Request) GetCompiledOperators() []*operators.Operators {
	return []*operators.Operators{request.CompiledOperators}
}

// RequestPartDefinitions contains a mapping of request part definitions and their
// description. Multiple definitions are separated by commas.
// Definitions not having a name (generated on runtime) are prefixed & suffixed by <>.
var RequestPartDefinitions = map[string]string{
	"type":     "Type is the type of request made",
	"success":  "Success specifies whether websocket connection was successful",
	"request":  "Websocket request made to the server",
	"duration": "Latest measured operation duration in seconds",
	"response": "Websocket response received from the server",
	"host":     "Host is the input to the template",
	"matched":  "Matched is the input which was matched upon",
}

func (request *Request) MakeResultEventItem(wrapped *output.InternalWrappedEvent) *output.ResultEvent {
	fields := protocolutils.GetJsonFieldsFromURL(types.ToString(wrapped.InternalEvent["host"]))
	if types.ToString(wrapped.InternalEvent["ip"]) != "" {
		fields.Ip = types.ToString(wrapped.InternalEvent["ip"])
	}
	data := &output.ResultEvent{
		TemplateID:       types.ToString(request.options.TemplateID),
		TemplatePath:     types.ToString(request.options.TemplatePath),
		Info:             request.options.TemplateInfo,
		TemplateVerifier: request.options.TemplateVerifier,
		Type:             types.ToString(wrapped.InternalEvent["type"]),
		Host:             fields.Host,
		Port:             fields.Port,
		Matched:          types.ToString(wrapped.InternalEvent["matched"]),
		Metadata:         wrapped.OperatorsResult.PayloadValues,
		ExtractedResults: wrapped.OperatorsResult.OutputExtracts,
		Timestamp:        time.Now(),
		MatcherStatus:    true,
		IP:               fields.Ip,
		Request:          types.ToString(wrapped.InternalEvent["request"]),
		Response:         types.ToString(wrapped.InternalEvent["response"]),
		TemplateEncoded:  request.options.EncodeTemplate(),
		Error:            types.ToString(wrapped.InternalEvent["error"]),
	}
	return data
}

// Type returns the type of the protocol request
func (request *Request) Type() templateTypes.ProtocolType {
	return templateTypes.WebsocketProtocol
}

// UpdateOptions replaces this request's options with a new copy
func (r *Request) UpdateOptions(opts *protocols.ExecutorOptions) {
	r.Operators = protocols.CloneOperators(r.Operators)
	r.CompiledOperators = nil
	if r.options == nil {
		if opts != nil {
			r.options = opts.Copy()
		}
		return
	}
	r.options = r.options.Copy()
	r.options.ApplyNewEngineOptions(opts)
}