diff --git a/v2/pkg/protocols/common/generators/maps.go b/v2/pkg/protocols/common/generators/maps.go new file mode 100644 index 00000000..da78af9d --- /dev/null +++ b/v2/pkg/protocols/common/generators/maps.go @@ -0,0 +1,38 @@ +package generators + +import "strings" + +// MergeMaps merges two maps into a new map +func MergeMaps(m1, m2 map[string]interface{}) map[string]interface{} { + m := make(map[string]interface{}, len(m1)+len(m2)) + for k, v := range m1 { + m[k] = v + } + for k, v := range m2 { + m[k] = v + } + return m +} + +// CopyMap creates a new copy of an existing map +func CopyMap(originalMap map[string]interface{}) map[string]interface{} { + newMap := make(map[string]interface{}) + for key, value := range originalMap { + newMap[key] = value + } + return newMap +} + +// CopyMapWithDefaultValue creates a new copy of an existing map and set a default value +func CopyMapWithDefaultValue(originalMap map[string][]string, defaultValue interface{}) map[string]interface{} { + newMap := make(map[string]interface{}) + for key := range originalMap { + newMap[key] = defaultValue + } + return newMap +} + +// TrimDelimiters removes trailing brackets +func TrimDelimiters(s string) string { + return strings.TrimSuffix(strings.TrimPrefix(s, "{{"), "}}") +} diff --git a/v2/pkg/protocols/common/replacer/replacer.go b/v2/pkg/protocols/common/replacer/replacer.go index 29f4aece..06e12c36 100644 --- a/v2/pkg/protocols/common/replacer/replacer.go +++ b/v2/pkg/protocols/common/replacer/replacer.go @@ -5,10 +5,11 @@ import ( "strings" ) +// Payload marker constants const ( - markerGeneral = "§" - markerParenthesisOpen = "{{" - markerParenthesisClose = "}}" + MarkerGeneral = "§" + MarkerParenthesisOpen = "{{" + MarkerParenthesisClose = "}}" ) // New creates a new replacer structure for values replacement on the fly. @@ -19,11 +20,11 @@ func New(values map[string]interface{}) *strings.Replacer { valueStr := fmt.Sprintf("%s", val) replacerItems = append(replacerItems, - fmt.Sprintf("%s%s%s", markerParenthesisOpen, key, markerParenthesisClose), + fmt.Sprintf("%s%s%s", MarkerParenthesisOpen, key, MarkerParenthesisClose), valueStr, ) replacerItems = append(replacerItems, - fmt.Sprintf("%s%s%s", markerGeneral, key, markerGeneral), + fmt.Sprintf("%s%s%s", MarkerGeneral, key, MarkerGeneral), valueStr, ) } diff --git a/v2/pkg/protocols/http/build_request.go b/v2/pkg/protocols/http/build_request.go index ac376b65..256031cd 100644 --- a/v2/pkg/protocols/http/build_request.go +++ b/v2/pkg/protocols/http/build_request.go @@ -1,13 +1,30 @@ package http import ( + "context" "fmt" + "io" + "io/ioutil" + "net" + "net/http" + "net/url" "regexp" + "strings" + "time" + "github.com/Knetic/govaluate" + "github.com/projectdiscovery/nuclei/v2/pkg/operators/common/dsl" "github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/generators" + "github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/replacer" + "github.com/projectdiscovery/nuclei/v2/pkg/protocols/http/race" + "github.com/projectdiscovery/nuclei/v2/pkg/protocols/http/raw" + "github.com/projectdiscovery/retryablehttp-go" ) -var urlWithPortRegex = regexp.MustCompile(`{{BaseURL}}:(\d+)`) +var ( + urlWithPortRegex = regexp.MustCompile(`{{BaseURL}}:(\d+)`) + templateExpressionRegex = regexp.MustCompile(`(?m)\{\{[^}]+\}\}`) +) // requestGenerator generates requests sequentially based on various // configurations for a http request template. @@ -65,7 +82,6 @@ func (r *requestGenerator) nextValue() (string, map[string]interface{}, bool) { } return "", nil, false } - fmt.Printf("index-last: %v\n", r.currentIndex) return r.request.Raw[r.currentIndex], payload, true } if item := r.request.Raw[r.currentIndex]; item != "" { @@ -76,11 +92,18 @@ func (r *requestGenerator) nextValue() (string, map[string]interface{}, bool) { return "", nil, false } -/* +// generatedRequest is a single wrapped generated request for a template request +type generatedRequest struct { + original *Request + rawRequest *raw.Request + meta map[string]interface{} + request *retryablehttp.Request +} + // Make creates a http request for the provided input. // It returns io.EOF as error when all the requests have been exhausted. -func (r *requestGenerator) Make(baseURL string, dynamicValues map[string]interface{}) (*HTTPRequest, error) { - data, ok := r.nextValue() +func (r *requestGenerator) Make(baseURL string, dynamicValues map[string]interface{}) (*generatedRequest, error) { + data, payloads, ok := r.nextValue() if !ok { return nil, io.EOF } @@ -92,7 +115,6 @@ func (r *requestGenerator) Make(baseURL string, dynamicValues map[string]interfa } hostname := parsed.Host - values := generators.MergeMaps(dynamicValues, map[string]interface{}{ "BaseURL": baseURLWithTemplatePrefs(data, parsed), "Hostname": hostname, @@ -101,7 +123,7 @@ func (r *requestGenerator) Make(baseURL string, dynamicValues map[string]interfa // If data contains \n it's a raw request, process it like that. Else // continue with the template based request flow. if strings.Contains(data, "\n") { - return r.makeHTTPRequestFromRaw(ctx, baseURL, data, values) + return r.makeHTTPRequestFromRaw(ctx, baseURL, data, values, payloads) } return r.makeHTTPRequestFromModel(ctx, data, values) } @@ -121,14 +143,12 @@ func baseURLWithTemplatePrefs(data string, parsedURL *url.URL) string { return parsedURL.String() } -/* - // MakeHTTPRequestFromModel creates a *http.Request from a request template -func (r *Request) makeHTTPRequestFromModel(ctx context.Context, data string, values map[string]interface{}) (*HTTPRequest, error) { +func (r *requestGenerator) makeHTTPRequestFromModel(ctx context.Context, data string, values map[string]interface{}) (*generatedRequest, error) { URL := replacer.New(values).Replace(data) // Build a request on the specified URL - req, err := http.NewRequestWithContext(ctx, r.Method, URL, nil) + req, err := http.NewRequestWithContext(ctx, r.request.Method, URL, nil) if err != nil { return nil, err } @@ -137,31 +157,27 @@ func (r *Request) makeHTTPRequestFromModel(ctx context.Context, data string, val if err != nil { return nil, err } - return &HTTPRequest{Request: request}, nil + return &generatedRequest{request: request}, nil } // makeHTTPRequestFromRaw creates a *http.Request from a raw request -func (r *Request) makeHTTPRequestFromRaw(ctx context.Context, baseURL, data string, values map[string]interface{}) (*HTTPRequest, error) { +func (r *requestGenerator) makeHTTPRequestFromRaw(ctx context.Context, baseURL, data string, values, payloads map[string]interface{}) (*generatedRequest, error) { // Add trailing line data += "\n" - // If we have payloads, handle them by creating a generator - if len(r.Payloads) > 0 { - r.gsfm.InitOrSkip(baseURL) - r.ReadOne(baseURL) - - payloads, err := r.getPayloadValues(baseURL) + // If we have payloads, handle them by evaluating them at runtime. + if len(r.request.Payloads) > 0 { + finalPayloads, err := r.getPayloadValues(baseURL, payloads) if err != nil { return nil, err } - return r.handleRawWithPaylods(ctx, data, baseURL, values, payloads) + return r.handleRawWithPaylods(ctx, data, baseURL, values, finalPayloads) } - - // otherwise continue with normal flow return r.handleRawWithPaylods(ctx, data, baseURL, values, nil) } -func (r *Request) handleRawWithPaylods(ctx context.Context, rawRequest, baseURL string, values, genValues map[string]interface{}) (*HTTPRequest, error) { +// handleRawWithPaylods handles raw requests along with paylaods +func (r *requestGenerator) handleRawWithPaylods(ctx context.Context, rawRequest, baseURL string, values, genValues map[string]interface{}) (*generatedRequest, error) { baseValues := generators.CopyMap(values) finValues := generators.MergeMaps(baseValues, genValues) @@ -169,12 +185,10 @@ func (r *Request) handleRawWithPaylods(ctx context.Context, rawRequest, baseURL rawRequest = replacer.New(finValues).Replace(rawRequest) dynamicValues := make(map[string]interface{}) - // find all potentials tokens between {{}} - var re = regexp.MustCompile(`(?m)\{\{[^}]+\}\}`) - for _, match := range re.FindAllString(rawRequest, -1) { + for _, match := range templateExpressionRegex.FindAllString(rawRequest, -1) { // check if the match contains a dynamic variable expr := generators.TrimDelimiters(match) - compiled, err := govaluate.NewEvaluableExpressionWithFunctions(expr, generators.HelperFunctions()) + compiled, err := govaluate.NewEvaluableExpressionWithFunctions(expr, dsl.HelperFunctions()) if err != nil { return nil, err @@ -189,40 +203,37 @@ func (r *Request) handleRawWithPaylods(ctx context.Context, rawRequest, baseURL // Replacer dynamic values if any in raw request and parse it rawRequest = replacer.New(dynamicValues).Replace(rawRequest) - rawRequestData, err := raw.Parse(rawRequest, baseURL, r.Unsafe) + rawRequestData, err := raw.Parse(rawRequest, baseURL, r.request.Unsafe) if err != nil { return nil, err } // rawhttp - if r.Unsafe { - unsafeReq := &HTTPRequest{ - RawRequest: rawRequest, - Meta: genValues, - AutomaticHostHeader: !r.DisableAutoHostname, - AutomaticContentLengthHeader: !r.DisableAutoContentLength, - Unsafe: true, - FollowRedirects: r.Redirects, + if r.request.Unsafe { + unsafeReq := &generatedRequest{ + rawRequest: rawRequestData, + meta: genValues, + original: r.request, } return unsafeReq, nil } // retryablehttp var body io.ReadCloser - body = ioutil.NopCloser(strings.NewReader(rawRequest.Data)) - if r.Race { + body = ioutil.NopCloser(strings.NewReader(rawRequestData.Data)) + if r.request.Race { // More or less this ensures that all requests hit the endpoint at the same approximated time // Todo: sync internally upon writing latest request byte body = race.NewOpenGateWithTimeout(body, time.Duration(2)*time.Second) } - req, err := http.NewRequestWithContext(ctx, rawRequest.Method, rawRequest.FullURL, body) + req, err := http.NewRequestWithContext(ctx, rawRequestData.Method, rawRequestData.FullURL, body) if err != nil { return nil, err } // copy headers - for key, value := range rawRequest.Headers { + for key, value := range rawRequestData.Headers { req.Header[key] = []string{value} } @@ -230,33 +241,33 @@ func (r *Request) handleRawWithPaylods(ctx context.Context, rawRequest, baseURL if err != nil { return nil, err } - - return &HTTPRequest{Request: request, Meta: genValues}, nil + return &generatedRequest{request: request, meta: genValues}, nil } -func (r *Request) fillRequest(req *http.Request, values map[string]interface{}) (*retryablehttp.Request, error) { - replacer := replacer.New(values) +// fillRequest fills various headers in the request with values +func (r *requestGenerator) fillRequest(req *http.Request, values map[string]interface{}) (*retryablehttp.Request, error) { // Set the header values requested - for header, value := range r.Headers { + replacer := replacer.New(values) + for header, value := range r.request.Headers { req.Header[header] = []string{replacer.Replace(value)} } // In case of multiple threads the underlying connection should remain open to allow reuse - if r.Threads <= 0 && req.Header.Get("Connection") == "" { + if r.request.Threads <= 0 && req.Header.Get("Connection") == "" { req.Close = true } // Check if the user requested a request body - if r.Body != "" { - req.Body = ioutil.NopCloser(strings.NewReader(r.Body)) + if r.request.Body != "" { + req.Body = ioutil.NopCloser(strings.NewReader(r.request.Body)) } setHeader(req, "User-Agent", "Nuclei - Open-source project (github.com/projectdiscovery/nuclei)") // raw requests are left untouched - if len(r.Raw) > 0 { + if len(r.request.Raw) > 0 { return retryablehttp.FromRequest(req) } - //setHeader(req, "Accept", "") + setHeader(req, "Accept", "*/*") setHeader(req, "Accept-Language", "en") return retryablehttp.FromRequest(req) @@ -269,28 +280,26 @@ func setHeader(req *http.Request, name, value string) { } } - // getPayloadValues returns current payload values for a request -func (r *Request) getPayloadValues(reqURL string) (map[string]interface{}, error) { +func (r *requestGenerator) getPayloadValues(reqURL string, templatePayloads map[string]interface{}) (map[string]interface{}, error) { payloadProcessedValues := make(map[string]interface{}) - payloadsFromTemplate := r.gsfm.Value(reqURL) - for k, v := range payloadsFromTemplate { + for k, v := range templatePayloads { kexp := v.(string) // if it doesn't containing markups, we just continue - if !hasMarker(kexp) { + if !strings.Contains(kexp, replacer.MarkerParenthesisOpen) || strings.Contains(kexp, replacer.MarkerParenthesisClose) || strings.Contains(kexp, replacer.MarkerGeneral) { payloadProcessedValues[k] = v continue } // attempts to expand expressions - compiled, err := govaluate.NewEvaluableExpressionWithFunctions(kexp, generators.HelperFunctions()) + compiled, err := govaluate.NewEvaluableExpressionWithFunctions(kexp, dsl.HelperFunctions()) if err != nil { // it is a simple literal payload => proceed with literal value payloadProcessedValues[k] = v continue } // it is an expression - try to solve it - expValue, err := compiled.Evaluate(payloadsFromTemplate) + expValue, err := compiled.Evaluate(templatePayloads) if err != nil { // an error occurred => proceed with literal value payloadProcessedValues[k] = v @@ -307,4 +316,3 @@ func (r *Request) getPayloadValues(reqURL string) (map[string]interface{}, error // ErrNoPayload error to avoid the additional base null request var ErrNoPayload = fmt.Errorf("no payload found") -*/ diff --git a/v2/pkg/protocols/http/build_request_test.go b/v2/pkg/protocols/http/build_request_test.go index 5793a23b..846b8896 100644 --- a/v2/pkg/protocols/http/build_request_test.go +++ b/v2/pkg/protocols/http/build_request_test.go @@ -1,13 +1,28 @@ package http import ( - "fmt" "testing" "github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/generators" "github.com/stretchr/testify/require" ) +func TestRequestGeneratorPaths(t *testing.T) { + req := &Request{ + Path: []string{"{{BaseURL}}/test", "{{BaseURL}}/test.php"}, + } + generator := req.newGenerator() + var payloads []string + for { + raw, _, ok := generator.nextValue() + if !ok { + break + } + payloads = append(payloads, raw) + } + require.Equal(t, req.Path, payloads, "Could not get correct paths") +} + func TestRequestGeneratorClusterSingle(t *testing.T) { var err error @@ -22,12 +37,11 @@ func TestRequestGeneratorClusterSingle(t *testing.T) { generator := req.newGenerator() var payloads []map[string]interface{} for { - raw, data, ok := generator.nextValue() + _, data, ok := generator.nextValue() if !ok { break } payloads = append(payloads, data) - fmt.Printf("%v %v\n", raw, data) } require.Equal(t, 9, len(payloads), "Could not get correct number of payloads") } @@ -46,12 +60,11 @@ func TestRequestGeneratorClusterMultipleRaw(t *testing.T) { generator := req.newGenerator() var payloads []map[string]interface{} for { - raw, data, ok := generator.nextValue() + _, data, ok := generator.nextValue() if !ok { break } payloads = append(payloads, data) - fmt.Printf("%v %v\n", raw, data) } require.Equal(t, 18, len(payloads), "Could not get correct number of payloads") } diff --git a/v2/pkg/protocols/http/http.go b/v2/pkg/protocols/http/http.go index 0c6be24b..fb40407e 100644 --- a/v2/pkg/protocols/http/http.go +++ b/v2/pkg/protocols/http/http.go @@ -1,8 +1,13 @@ package http import ( + "github.com/pkg/errors" + "github.com/projectdiscovery/nuclei/v2/pkg/operators" "github.com/projectdiscovery/nuclei/v2/pkg/protocols" "github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/generators" + "github.com/projectdiscovery/nuclei/v2/pkg/protocols/http/httpclientpool" + "github.com/projectdiscovery/rawhttp" + "github.com/projectdiscovery/retryablehttp-go" ) // Request contains a http request to be made from a template @@ -51,7 +56,44 @@ type Request struct { // The minimum number fof requests is determined by threads Race bool `yaml:"race"` - attackType generators.Type - generator *generators.Generator // optional, only enabled when using payloads - options *protocols.ExecuterOptions + // Operators for the current request go here. + *operators.Operators + + options *protocols.ExecuterOptions + attackType generators.Type + generator *generators.Generator // optional, only enabled when using payloads + httpClient *retryablehttp.Client + rawhttpClient *rawhttp.Client +} + +// Compile compiles the protocol request for further execution. +func (r *Request) Compile(options *protocols.ExecuterOptions) error { + client, err := httpclientpool.Get(options.Options, &httpclientpool.Configuration{ + Threads: r.Threads, + MaxRedirects: r.MaxRedirects, + FollowRedirects: r.Redirects, + }) + if err != nil { + return errors.Wrap(err, "could not get dns client") + } + r.httpClient = client + + if len(r.Raw) > 0 { + r.rawhttpClient = httpclientpool.GetRawHTTP() + } + if r.Operators != nil { + if err := r.Operators.Compile(); err != nil { + return errors.Wrap(err, "could not compile operators") + } + } + + if len(r.Payloads) > 0 { + r.attackType = generators.StringToType[r.AttackType] + r.generator, err = generators.New(r.Payloads, r.attackType) + if err != nil { + return errors.Wrap(err, "could not parse payloads") + } + } + r.options = options + return nil } diff --git a/v2/pkg/requests/bulk-http-request.go b/v2/pkg/requests/bulk-http-request.go deleted file mode 100644 index 0a1252e5..00000000 --- a/v2/pkg/requests/bulk-http-request.go +++ /dev/null @@ -1,60 +0,0 @@ -package requests - -import ( - "regexp" - - "github.com/projectdiscovery/nuclei/v2/pkg/generators" - "github.com/projectdiscovery/nuclei/v2/pkg/matchers" - "github.com/projectdiscovery/rawhttp" - retryablehttp "github.com/projectdiscovery/retryablehttp-go" -) - -const ( - two = 2 - three = 3 -) - -var urlWithPortRgx = regexp.MustCompile(`{{BaseURL}}:(\d+)`) - -// GetMatchersCondition returns the condition for the matcher -func (r *BulkHTTPRequest) GetMatchersCondition() matchers.ConditionType { - return r.matchersCondition -} - -// SetMatchersCondition sets the condition for the matcher -func (r *BulkHTTPRequest) SetMatchersCondition(condition matchers.ConditionType) { - r.matchersCondition = condition -} - -// GetAttackType returns the attack -func (r *BulkHTTPRequest) GetAttackType() generators.Type { - return r.attackType -} - -// SetAttackType sets the attack -func (r *BulkHTTPRequest) SetAttackType(attack generators.Type) { - r.attackType = attack -} - -// GetRequestCount returns the total number of requests the YAML rule will perform -func (r *BulkHTTPRequest) GetRequestCount() int64 { - return int64(r.gsfm.Total()) -} - -// HTTPRequest is the basic HTTP request -type HTTPRequest struct { - Request *retryablehttp.Request - RawRequest *RawRequest - Meta map[string]interface{} - - // flags - Unsafe bool - Pipeline bool - AutomaticHostHeader bool - AutomaticContentLengthHeader bool - AutomaticConnectionHeader bool - FollowRedirects bool - Rawclient *rawhttp.Client - Httpclient *retryablehttp.Client - PipelineClient *rawhttp.PipelineClient -} diff --git a/v2/pkg/requests/generator.go b/v2/pkg/requests/generator.go deleted file mode 100644 index 50e80e4a..00000000 --- a/v2/pkg/requests/generator.go +++ /dev/null @@ -1,273 +0,0 @@ -package requests - -import ( - "sync" - "time" - - "github.com/projectdiscovery/nuclei/v2/pkg/generators" -) - -type GeneratorState int - -const ( - fifteen = 15 - initial GeneratorState = iota - running - done -) - -type Generator struct { - sync.RWMutex - positionPath int - positionRaw int - gchan chan map[string]interface{} - currentGeneratorValue map[string]interface{} - state GeneratorState -} - -type GeneratorFSM struct { - sync.RWMutex - payloads map[string]interface{} - basePayloads map[string][]string - generator func(payloads map[string][]string) (out chan map[string]interface{}) - Generators map[string]*Generator - Type generators.Type - Paths []string - Raws []string -} - -func NewGeneratorFSM(typ generators.Type, payloads map[string]interface{}, paths, raws []string) *GeneratorFSM { - var gsfm GeneratorFSM - gsfm.payloads = payloads - gsfm.Paths = paths - gsfm.Raws = raws - gsfm.Type = typ - - if len(gsfm.payloads) > 0 { - // load payloads if not already done - if gsfm.basePayloads == nil { - gsfm.basePayloads = generators.LoadPayloads(gsfm.payloads) - } - - generatorFunc := generators.SniperGenerator - - switch typ { - case generators.PitchFork: - generatorFunc = generators.PitchforkGenerator - case generators.ClusterBomb: - generatorFunc = generators.ClusterbombGenerator - case generators.Sniper: - generatorFunc = generators.SniperGenerator - } - - gsfm.generator = generatorFunc - } - - gsfm.Generators = make(map[string]*Generator) - - return &gsfm -} - -func (gfsm *GeneratorFSM) Add(key string) { - gfsm.Lock() - defer gfsm.Unlock() - - if _, ok := gfsm.Generators[key]; !ok { - gfsm.Generators[key] = &Generator{state: initial} - } -} - -func (gfsm *GeneratorFSM) Has(key string) bool { - gfsm.RLock() - defer gfsm.RUnlock() - - _, ok := gfsm.Generators[key] - - return ok -} - -func (gfsm *GeneratorFSM) Delete(key string) { - gfsm.Lock() - defer gfsm.Unlock() - - delete(gfsm.Generators, key) -} - -func (gfsm *GeneratorFSM) ReadOne(key string) { - gfsm.RLock() - defer gfsm.RUnlock() - g, ok := gfsm.Generators[key] - - if !ok { - return - } - - for afterCh := time.After(fifteen * time.Second); ; { - select { - // got a value - case curGenValue, ok := <-g.gchan: - if !ok { - g.Lock() - g.gchan = nil - g.state = done - g.currentGeneratorValue = nil - g.Unlock() - - return - } - - g.currentGeneratorValue = curGenValue - - return - // timeout - case <-afterCh: - g.Lock() - g.gchan = nil - g.state = done - g.Unlock() - - return - } - } -} - -func (gfsm *GeneratorFSM) InitOrSkip(key string) { - gfsm.RLock() - defer gfsm.RUnlock() - - g, ok := gfsm.Generators[key] - if !ok { - return - } - - if len(gfsm.payloads) > 0 { - g.Lock() - defer g.Unlock() - - if g.gchan == nil { - g.gchan = gfsm.generator(gfsm.basePayloads) - g.state = running - } - } -} - -func (gfsm *GeneratorFSM) Value(key string) map[string]interface{} { - gfsm.RLock() - defer gfsm.RUnlock() - - g, ok := gfsm.Generators[key] - if !ok { - return nil - } - - return g.currentGeneratorValue -} - -func (gfsm *GeneratorFSM) Next(key string) bool { - gfsm.RLock() - defer gfsm.RUnlock() - - g, ok := gfsm.Generators[key] - if !ok { - return false - } - - if g.positionPath+g.positionRaw >= len(gfsm.Paths)+len(gfsm.Raws) { - return false - } - - return true -} - -func (gfsm *GeneratorFSM) Position(key string) int { - gfsm.RLock() - defer gfsm.RUnlock() - - g, ok := gfsm.Generators[key] - if !ok { - return 0 - } - - return g.positionPath + g.positionRaw -} - -func (gfsm *GeneratorFSM) Reset(key string) { - gfsm.Lock() - defer gfsm.Unlock() - - if !gfsm.Has(key) { - gfsm.Add(key) - } - - g, ok := gfsm.Generators[key] - if !ok { - return - } - - g.positionPath = 0 - g.positionRaw = 0 -} - -func (gfsm *GeneratorFSM) Current(key string) string { - gfsm.RLock() - defer gfsm.RUnlock() - - g, ok := gfsm.Generators[key] - if !ok { - return "" - } - - if g.positionPath < len(gfsm.Paths) && len(gfsm.Paths) != 0 { - return gfsm.Paths[g.positionPath] - } - - return gfsm.Raws[g.positionRaw] -} -func (gfsm *GeneratorFSM) Total() int { - estimatedRequestsWithPayload := 0 - if len(gfsm.basePayloads) > 0 { - switch gfsm.Type { - case generators.Sniper: - for _, kv := range gfsm.basePayloads { - estimatedRequestsWithPayload += len(kv) - } - case generators.PitchFork: - // Positional so it's equal to the length of one list - for _, kv := range gfsm.basePayloads { - estimatedRequestsWithPayload += len(kv) - break - } - case generators.ClusterBomb: - // Total of combinations => rule of product - prod := 1 - for _, kv := range gfsm.basePayloads { - prod *= len(kv) - } - estimatedRequestsWithPayload += prod - } - } - return len(gfsm.Paths) + len(gfsm.Raws) + estimatedRequestsWithPayload -} - -func (gfsm *GeneratorFSM) Increment(key string) { - gfsm.Lock() - defer gfsm.Unlock() - - g, ok := gfsm.Generators[key] - if !ok { - return - } - - if len(gfsm.Paths) > 0 && g.positionPath < len(gfsm.Paths) { - g.positionPath++ - return - } - - if len(gfsm.Raws) > 0 && g.positionRaw < len(gfsm.Raws) { - // if we have payloads increment only when the generators are done - if g.gchan == nil { - g.state = done - g.positionRaw++ - } - } -} diff --git a/v2/pkg/requests/util.go b/v2/pkg/requests/util.go index ce8ba795..ae4550d1 100644 --- a/v2/pkg/requests/util.go +++ b/v2/pkg/requests/util.go @@ -8,12 +8,6 @@ import ( "strings" ) -const ( - markerParenthesisOpen = "{{" - markerParenthesisClose = "}}" - markerGeneral = "§" -) - func newReplacer(values map[string]interface{}) *strings.Replacer { var replacerItems []string for key, val := range values { @@ -71,7 +65,3 @@ func ExpandMapValues(m map[string]string) (m1 map[string][]string) { } return } - -func hasMarker(s string) bool { - return strings.Contains(s, markerParenthesisOpen) || strings.Contains(s, markerParenthesisClose) || strings.Contains(s, markerGeneral) -}