Merge pull request #185 from edoardottt/devel

Add ratelimit option

Author: edoardottt

Makefile

@@ -16,7 +16,7 @@ lint:
 	@golangci-lint run
 
 linux:
-	@go build -o csprecon ./main.go
+	@go build -o csprecon ./cmd/csprecon
 	@sudo mv csprecon /usr/local/bin/
 	@echo "Done."
 

README.md

@@ -41,13 +41,13 @@ Install 📡
 
 ### Using Snap
 
-```bash
+```console
 sudo snap install csprecon
 ```
 
 ### Using Go
 
-```
+```console
 go install github.com/edoardottt/csprecon/cmd/csprecon@latest
 ```
 
@@ -68,6 +68,7 @@ CONFIGURATIONS:
    -d, -domain string[]  Filter results belonging to these domains (comma separated)
    -c, -concurrency int  Concurrency level (default 50)
    -t, -timeout int      Connection timeout in seconds (default 10)
+   -rl, -rate-limit int  Set a rate limit (per second)
 
 OUTPUT:
    -o, -output string  File to write output results
@@ -110,6 +111,12 @@ Grab all possible results from single CIDR
 csprecon -u 192.168.1.0/24 -cidr
 ```
 
+Set a rate limit of 10 requests per second
+
+```bash
+cat targets.txt | csprecon -rl 10
+```
+
 Changelog 📌
 -------
 

go.mod

@@ -12,6 +12,8 @@ require (
 	github.com/stretchr/testify v1.8.4
 )
 
+require github.com/benbjohnson/clock v1.3.0 // indirect
+
 require (
 	github.com/andybalholm/brotli v1.0.6 // indirect
 	github.com/andybalholm/cascadia v1.3.2 // indirect
@@ -39,6 +41,7 @@ require (
 	github.com/saintfish/chardet v0.0.0-20230101081208-5e3ef4b5456d // indirect
 	github.com/ulikunitz/xz v0.5.11 // indirect
 	github.com/xi2/xz v0.0.0-20171230120015-48954b6210f8 // indirect
+	go.uber.org/ratelimit v0.3.0
 	golang.org/x/exp v0.0.0-20231006140011-7918f672742d // indirect
 	golang.org/x/mod v0.13.0 // indirect
 	golang.org/x/net v0.17.0 // indirect

go.sum

@@ -10,6 +10,8 @@ github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2 h1:DklsrG3d
 github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2/go.mod h1:WaHUgvxTVq04UNunO+XhnAqY/wQc+bxr74GqbsZ/Jqw=
 github.com/aymerick/douceur v0.2.0 h1:Mv+mAeH1Q+n9Fr+oyamOlAkUNPWPlA8PPGR0QAaYuPk=
 github.com/aymerick/douceur v0.2.0/go.mod h1:wlT5vV2O3h55X9m7iVYN0TBM0NH/MmbLnd30/FjWUq4=
+github.com/benbjohnson/clock v1.3.0 h1:ip6w0uFQkncKQ979AypyG0ER7mqUSBdKLOgAle/AT8A=
+github.com/benbjohnson/clock v1.3.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
 github.com/cnf/structhash v0.0.0-20201127153200-e1b16c1ebc08 h1:ox2F0PSMlrAAiAdknSRMDrAr8mfxPCfSZolH+/qQnyQ=
 github.com/cnf/structhash v0.0.0-20201127153200-e1b16c1ebc08/go.mod h1:pCxVEbcm3AMg7ejXyorUXi6HQCzOIBf7zEDVPtw0/U4=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
@@ -81,6 +83,8 @@ github.com/ulikunitz/xz v0.5.11/go.mod h1:nbz6k7qbPmH4IRqmfOplQw/tblSgqTqBwxkY0o
 github.com/xi2/xz v0.0.0-20171230120015-48954b6210f8 h1:nIPpBwaJSVYIxUFsDv3M8ofmx9yWTog9BfvIu0q41lo=
 github.com/xi2/xz v0.0.0-20171230120015-48954b6210f8/go.mod h1:HUYIGzjTL3rfEspMxjDjgmT5uz5wzYJKVo23qUhYTos=
 github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=
+go.uber.org/ratelimit v0.3.0 h1:IdZd9wqvFXnvLvSEBo0KPcGfkoBGNkpTHlrE3Rcjkjw=
+go.uber.org/ratelimit v0.3.0/go.mod h1:So5LG7CV1zWpY1sHe+DXTJqQvOx+FFPFaAs2SnoyBaI=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
 golang.org/x/exp v0.0.0-20231006140011-7918f672742d h1:jtJma62tbqLibJ5sFQz8bKtEM8rJBtfilJ2qTU199MI=

pkg/csprecon/csprecon.go

@@ -127,6 +127,7 @@ func execute(r *Runner) {
 	defer r.InWg.Done()
 
 	dregex := CompileRegex(DomainRegex)
+	rl := rateLimiter(r)
 
 	for i := 0; i < r.Options.Concurrency; i++ {
 		r.InWg.Add(1)
@@ -144,6 +145,8 @@ func execute(r *Runner) {
 					return
 				}
 
+				rl.Take()
+
 				client := customClient(r.Options.Timeout)
 
 				result, err := CheckCSP(targetURL, r.UserAgent, dregex, client)

pkg/csprecon/utils.go renamed to pkg/csprecon/net.go

@@ -0,0 +1,20 @@
+/*
+csprecon - Discover new target domains using Content Security Policy
+
+This repository is under MIT License https://github.com/edoardottt/csprecon/blob/main/LICENSE
+*/
+
+package csprecon
+
+import "go.uber.org/ratelimit"
+
+func rateLimiter(r *Runner) ratelimit.Limiter {
+	var ratelimiter ratelimit.Limiter
+	if r.Options.RateLimit > 0 {
+		ratelimiter = ratelimit.New(r.Options.RateLimit)
+	} else {
+		ratelimiter = ratelimit.NewUnlimited()
+	}
+
+	return ratelimiter
+}

pkg/csprecon/utils_test.go renamed to pkg/csprecon/net_test.go

@@ -31,7 +31,11 @@ func (options *Options) validateOptions() error {
 	}
 
 	if options.Concurrency <= 0 {
-		return fmt.Errorf("%w", ErrNegativeValue)
+		return fmt.Errorf("concurrency: %w", ErrNegativeValue)
+	}
+
+	if options.RateLimit != 0 && options.RateLimit <= 0 {
+		return fmt.Errorf("rate limit: %w", ErrNegativeValue)
 	}
 
 	return nil

pkg/csprecon/ratelimit.go

@@ -20,6 +20,7 @@ import (
 const (
 	DefaultTimeout     = 10
 	DefaultConcurrency = 50
+	DefaultRateLimit   = 0
 )
 
 type Options struct {
@@ -33,6 +34,7 @@ type Options struct {
 	Concurrency int
 	Timeout     int
 	Cidr        bool
+	RateLimit   int
 }
 
 // configureOutput configures the output on the screen.
@@ -62,6 +64,7 @@ func ParseOptions() *Options {
 		flagSet.StringSliceVarP(&options.Domain, "domain", "d", nil, `Filter results belonging to these domains (comma separated)`, goflags.CommaSeparatedStringSliceOptions),
 		flagSet.IntVarP(&options.Concurrency, "concurrency", "c", DefaultConcurrency, `Concurrency level`),
 		flagSet.IntVarP(&options.Timeout, "timeout", "t", DefaultTimeout, `Connection timeout in seconds`),
+		flagSet.IntVarP(&options.RateLimit, "rate-limit", "rl", DefaultRateLimit, `Set a rate limit (per second)`),
 	)
 
 	// Output