github.com/projectdiscovery/nuclei/v2@v2.9.15/pkg/templates/types/types.go (about)

     1  package types
     2  
     3  import (
     4  	"encoding/json"
     5  	"fmt"
     6  	"strings"
     7  
     8  	"github.com/alecthomas/jsonschema"
     9  	"github.com/pkg/errors"
    10  	"github.com/projectdiscovery/goflags"
    11  	"github.com/projectdiscovery/nuclei/v2/pkg/model/types/stringslice"
    12  )
    13  
    14  // ProtocolType is the type of the request protocol specified
    15  type ProtocolType int
    16  
    17  // Supported values for the ProtocolType
    18  // name:ProtocolType
    19  const (
    20  	// name:dns
    21  	DNSProtocol ProtocolType = iota + 1
    22  	// name:file
    23  	FileProtocol
    24  	// name:http
    25  	HTTPProtocol
    26  	OfflineHTTPProtocol
    27  	// name:headless
    28  	HeadlessProtocol
    29  	// name:network
    30  	NetworkProtocol
    31  	// name:workflow
    32  	WorkflowProtocol
    33  	// name:ssl
    34  	SSLProtocol
    35  	// name:websocket
    36  	WebsocketProtocol
    37  	// name:whois
    38  	WHOISProtocol
    39  	limit
    40  	InvalidProtocol
    41  )
    42  
    43  // ExtractorTypes is a table for conversion of extractor type from string.
    44  var protocolMappings = map[ProtocolType]string{
    45  	InvalidProtocol:   "invalid",
    46  	DNSProtocol:       "dns",
    47  	FileProtocol:      "file",
    48  	HTTPProtocol:      "http",
    49  	HeadlessProtocol:  "headless",
    50  	NetworkProtocol:   "tcp",
    51  	WorkflowProtocol:  "workflow",
    52  	SSLProtocol:       "ssl",
    53  	WebsocketProtocol: "websocket",
    54  	WHOISProtocol:     "whois",
    55  }
    56  
    57  func GetSupportedProtocolTypes() ProtocolTypes {
    58  	var result []ProtocolType
    59  	for index := ProtocolType(1); index < limit; index++ {
    60  		result = append(result, index)
    61  	}
    62  	return result
    63  }
    64  
    65  func toProtocolType(valueToMap string) (ProtocolType, error) {
    66  	normalizedValue := normalizeValue(valueToMap)
    67  	for key, currentValue := range protocolMappings {
    68  		if normalizedValue == currentValue {
    69  			return key, nil
    70  		}
    71  	}
    72  	return -1, errors.New("Invalid protocol type: " + valueToMap)
    73  }
    74  
    75  func normalizeValue(value string) string {
    76  	return strings.TrimSpace(strings.ToLower(value))
    77  }
    78  
    79  func (t ProtocolType) String() string {
    80  	return protocolMappings[t]
    81  }
    82  
    83  // TypeHolder is used to hold internal type of the protocol
    84  type TypeHolder struct {
    85  	ProtocolType ProtocolType `mapping:"true"`
    86  }
    87  
    88  func (holder TypeHolder) JSONSchemaType() *jsonschema.Type {
    89  	gotType := &jsonschema.Type{
    90  		Type:        "string",
    91  		Title:       "type of the protocol",
    92  		Description: "Type of the protocol",
    93  	}
    94  	for _, types := range GetSupportedProtocolTypes() {
    95  		gotType.Enum = append(gotType.Enum, types.String())
    96  	}
    97  	return gotType
    98  }
    99  
   100  func (holder *TypeHolder) UnmarshalYAML(unmarshal func(interface{}) error) error {
   101  	var marshalledTypes string
   102  	if err := unmarshal(&marshalledTypes); err != nil {
   103  		return err
   104  	}
   105  
   106  	computedType, err := toProtocolType(marshalledTypes)
   107  	if err != nil {
   108  		return err
   109  	}
   110  
   111  	holder.ProtocolType = computedType
   112  	return nil
   113  }
   114  
   115  func (holder *TypeHolder) MarshalJSON() ([]byte, error) {
   116  	return json.Marshal(holder.ProtocolType.String())
   117  }
   118  
   119  func (holder TypeHolder) MarshalYAML() (interface{}, error) {
   120  	return holder.ProtocolType.String(), nil
   121  }
   122  
   123  type ProtocolTypes []ProtocolType
   124  
   125  func (protocolTypes *ProtocolTypes) Set(values string) error {
   126  	inputTypes, err := goflags.ToStringSlice(values, goflags.FileNormalizedStringSliceOptions)
   127  	if err != nil {
   128  		return err
   129  	}
   130  
   131  	for _, inputType := range inputTypes {
   132  		if err := setProtocolType(protocolTypes, inputType); err != nil {
   133  			return err
   134  		}
   135  	}
   136  	return nil
   137  }
   138  
   139  func (protocolTypes *ProtocolTypes) UnmarshalYAML(unmarshal func(interface{}) error) error {
   140  	var stringSliceValue stringslice.StringSlice
   141  	if err := unmarshal(&stringSliceValue); err != nil {
   142  		return err
   143  	}
   144  
   145  	stringSLice := stringSliceValue.ToSlice()
   146  	var result = make(ProtocolTypes, 0, len(stringSLice))
   147  	for _, typeString := range stringSLice {
   148  		if err := setProtocolType(&result, typeString); err != nil {
   149  			return err
   150  		}
   151  	}
   152  	*protocolTypes = result
   153  	return nil
   154  }
   155  
   156  func (protocolTypes ProtocolTypes) MarshalJSON() ([]byte, error) {
   157  	var stringProtocols = make([]string, 0, len(protocolTypes))
   158  	for _, protocol := range protocolTypes {
   159  		stringProtocols = append(stringProtocols, protocol.String())
   160  	}
   161  	return json.Marshal(stringProtocols)
   162  }
   163  
   164  func (protocolTypes ProtocolTypes) String() string {
   165  	var stringTypes []string
   166  	for _, t := range protocolTypes {
   167  		protocolMapping := t.String()
   168  		if protocolMapping != "" {
   169  			stringTypes = append(stringTypes, protocolMapping)
   170  		}
   171  
   172  	}
   173  	return strings.Join(stringTypes, ", ")
   174  }
   175  
   176  func setProtocolType(protocolTypes *ProtocolTypes, value string) error {
   177  	computedType, err := toProtocolType(value)
   178  	if err != nil {
   179  		return fmt.Errorf("'%s' is not a valid extract type", value)
   180  	}
   181  	*protocolTypes = append(*protocolTypes, computedType)
   182  	return nil
   183  }