github.com/projectdiscovery/nuclei/v2@v2.9.15/pkg/protocols/common/contextargs/contextargs.go (about)

     1  package contextargs
     2  
     3  import (
     4  	"net/http/cookiejar"
     5  	"strings"
     6  
     7  	mapsutil "github.com/projectdiscovery/utils/maps"
     8  	sliceutil "github.com/projectdiscovery/utils/slice"
     9  	stringsutil "github.com/projectdiscovery/utils/strings"
    10  	urlutil "github.com/projectdiscovery/utils/url"
    11  )
    12  
    13  var (
    14  	// reservedPorts contains list of reserved ports for non-network requests in nuclei
    15  	reservedPorts = []string{"80", "443", "8080", "8443", "8081", "53"}
    16  )
    17  
    18  // Context implements a shared context struct to share information across multiple templates within a workflow
    19  type Context struct {
    20  	// Meta is the target for the executor
    21  	MetaInput *MetaInput
    22  
    23  	// CookieJar shared within workflow's http templates
    24  	CookieJar *cookiejar.Jar
    25  
    26  	// Args is a workflow shared key-value store
    27  	args *mapsutil.SyncLockMap[string, interface{}]
    28  }
    29  
    30  // Create a new contextargs instance
    31  func New() *Context {
    32  	return &Context{MetaInput: &MetaInput{}}
    33  }
    34  
    35  // Create a new contextargs instance with input string
    36  func NewWithInput(input string) *Context {
    37  	return &Context{MetaInput: &MetaInput{Input: input}}
    38  }
    39  
    40  func (ctx *Context) initialize() {
    41  	ctx.args = &mapsutil.SyncLockMap[string, interface{}]{Map: mapsutil.Map[string, interface{}]{}}
    42  }
    43  
    44  // Set the specific key-value pair
    45  func (ctx *Context) Set(key string, value interface{}) {
    46  	if !ctx.isInitialized() {
    47  		ctx.initialize()
    48  	}
    49  
    50  	_ = ctx.args.Set(key, value)
    51  }
    52  
    53  func (ctx *Context) isInitialized() bool {
    54  	return ctx.args != nil
    55  }
    56  
    57  func (ctx *Context) hasArgs() bool {
    58  	return ctx.isInitialized() && !ctx.args.IsEmpty()
    59  }
    60  
    61  // UseNetworkPort updates input with required/default network port for that template
    62  // but is ignored if input/target contains non-http ports like 80,8080,8081 etc
    63  func (ctx *Context) UseNetworkPort(port string, excludePorts string) error {
    64  	ignorePorts := reservedPorts
    65  	if excludePorts != "" {
    66  		// TODO: add support for service names like http,https,ssh etc once https://github.com/projectdiscovery/netdb is ready
    67  		ignorePorts = sliceutil.Dedupe(strings.Split(excludePorts, ","))
    68  	}
    69  	if port == "" {
    70  		// if template does not contain port, do nothing
    71  		return nil
    72  	}
    73  	target, err := urlutil.Parse(ctx.MetaInput.Input)
    74  	if err != nil {
    75  		return err
    76  	}
    77  	inputPort := target.Port()
    78  	if inputPort == "" || stringsutil.EqualFoldAny(inputPort, ignorePorts...) {
    79  		// replace port with networkPort
    80  		target.UpdatePort(port)
    81  		ctx.MetaInput.Input = target.Host
    82  	}
    83  	return nil
    84  }
    85  
    86  // Get the value with specific key if exists
    87  func (ctx *Context) Get(key string) (interface{}, bool) {
    88  	if !ctx.hasArgs() {
    89  		return nil, false
    90  	}
    91  
    92  	return ctx.args.Get(key)
    93  }
    94  
    95  func (ctx *Context) GetAll() *mapsutil.SyncLockMap[string, interface{}] {
    96  	if !ctx.hasArgs() {
    97  		return nil
    98  	}
    99  
   100  	return ctx.args.Clone()
   101  }
   102  
   103  func (ctx *Context) ForEach(f func(string, interface{})) {
   104  	_ = ctx.args.Iterate(func(k string, v interface{}) error {
   105  		f(k, v)
   106  		return nil
   107  	})
   108  }
   109  
   110  // Has check if the key exists
   111  func (ctx *Context) Has(key string) bool {
   112  	return ctx.hasArgs() && ctx.args.Has(key)
   113  }
   114  
   115  func (ctx *Context) HasArgs() bool {
   116  	return ctx.hasArgs()
   117  }
   118  
   119  func (ctx *Context) Clone() *Context {
   120  	newCtx := &Context{
   121  		MetaInput: ctx.MetaInput.Clone(),
   122  		args:      ctx.args,
   123  		CookieJar: ctx.CookieJar,
   124  	}
   125  	return newCtx
   126  }