implemented request count estimation

projectdiscovery · Oct 12, 2020 · 84c9373 · 84c9373
1 parent 500d7bb
commit 84c9373
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 4 deletions.
diff --git a/v2/pkg/requests/bulk-http-request.go b/v2/pkg/requests/bulk-http-request.go
@@ -74,7 +74,6 @@ type BulkHTTPRequest struct {
 	// DisableAutoContentLength Enable/Disable Content-Length header for unsafe raw requests
 	DisableAutoContentLength bool `yaml:"disable-automatic-content-length-header,omitempty"`
 	Threads                  int  `yaml:"threads,omitempty"`
-	RateLimit                int  `yaml:"rate-limit,omitempty"`
 
 	// Internal Finite State Machine keeping track of scan process
 	gsfm *GeneratorFSM
@@ -102,7 +101,7 @@ func (r *BulkHTTPRequest) SetAttackType(attack generators.Type) {
 
 // GetRequestCount returns the total number of requests the YAML rule will perform
 func (r *BulkHTTPRequest) GetRequestCount() int64 {
-	return int64(len(r.Raw) | len(r.Path))
+	return int64(r.gsfm.Total())
 }
 
 // MakeHTTPRequest makes the HTTP request
@@ -450,7 +449,7 @@ func (r *BulkHTTPRequest) Current(reqURL string) string {
 
 // Total is the total number of requests
 func (r *BulkHTTPRequest) Total() int {
-	return len(r.Path) + len(r.Raw)
+	return r.gsfm.Total()
 }
 
 // Increment increments the processed request

diff --git a/v2/pkg/requests/generator.go b/v2/pkg/requests/generator.go
@@ -41,6 +41,7 @@ func NewGeneratorFSM(typ generators.Type, payloads map[string]interface{}, paths
 	gsfm.payloads = payloads
 	gsfm.Paths = paths
 	gsfm.Raws = raws
+	gsfm.Type = typ
 
 	if len(gsfm.payloads) > 0 {
 		// load payloads if not already done
@@ -231,7 +232,30 @@ func (gfsm *GeneratorFSM) Current(key string) string {
 	return gfsm.Raws[g.positionRaw]
 }
 func (gfsm *GeneratorFSM) Total() int {
-	return len(gfsm.Paths) + len(gfsm.Raws)
+	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 = prod * len(kv)
+			}
+			estimatedRequestsWithPayload += prod
+		}
+	}
+
+	return len(gfsm.Paths) + len(gfsm.Raws) + estimatedRequestsWithPayload
 }
 
 func (gfsm *GeneratorFSM) Increment(key string) {