feature: workers count option (#69)

v-byte-cpu · web-flow · commit 9815930becdc · 2021-05-11T02:21:33.000+03:00
diff --git a/command/docker.go b/command/docker.go
@@ -18,6 +18,7 @@ func init() {
 	dockerCmd.Flags().StringVarP(&cliPortsFlag, "ports", "p", "", "set ports to scan")
 	dockerCmd.Flags().StringVarP(&cliIPPortFileFlag, "file", "f", "", "set JSONL file with ip/port pairs to scan")
 	dockerCmd.Flags().StringVar(&cliProtoFlag, "proto", "", "set protocol to use, http is used by default; only http or https are valid")
+	dockerCmd.Flags().IntVarP(&cliWorkerCountFlag, "workers", "w", defaultWorkerCount, "set workers count")
 	rootCmd.AddCommand(dockerCmd)
 }
 
@@ -63,6 +64,5 @@ func newDockerScanEngine(ctx context.Context) scan.EngineResulter {
 	// TODO custom dataTimeout
 	scanner := docker.NewScanner(cliProtoFlag, docker.WithDataTimeout(10*time.Second))
 	results := scan.NewResultChan(ctx, 1000)
-	// TODO custom workerCount
-	return scan.NewScanEngine(newIPPortGenerator(), scanner, results, scan.WithScanWorkerCount(50))
+	return scan.NewScanEngine(newIPPortGenerator(), scanner, results, scan.WithScanWorkerCount(cliWorkerCountFlag))
 }
diff --git a/command/elastic.go b/command/elastic.go
@@ -18,6 +18,7 @@ func init() {
 	elasticCmd.Flags().StringVarP(&cliPortsFlag, "ports", "p", "", "set ports to scan")
 	elasticCmd.Flags().StringVarP(&cliIPPortFileFlag, "file", "f", "", "set JSONL file with ip/port pairs to scan")
 	elasticCmd.Flags().StringVar(&cliProtoFlag, "proto", "", "set protocol to use, http is used by default; only http or https are valid")
+	elasticCmd.Flags().IntVarP(&cliWorkerCountFlag, "workers", "w", defaultWorkerCount, "set workers count")
 	rootCmd.AddCommand(elasticCmd)
 }
 
@@ -63,6 +64,5 @@ func newElasticScanEngine(ctx context.Context) scan.EngineResulter {
 	// TODO custom dataTimeout
 	scanner := elastic.NewScanner(cliProtoFlag, elastic.WithDataTimeout(5*time.Second))
 	results := scan.NewResultChan(ctx, 1000)
-	// TODO custom workerCount
-	return scan.NewScanEngine(newIPPortGenerator(), scanner, results, scan.WithScanWorkerCount(50))
+	return scan.NewScanEngine(newIPPortGenerator(), scanner, results, scan.WithScanWorkerCount(cliWorkerCountFlag))
 }
diff --git a/command/root.go b/command/root.go
@@ -90,6 +90,9 @@ var rootCmd = &cobra.Command{
 				return
 			}
 		}
+		if cliWorkerCountFlag <= 0 {
+			return errors.New("invalid workers count")
+		}
 		return
 	},
 }
@@ -110,6 +113,7 @@ var (
 	cliIPTotalLenFlag   string
 	cliIPProtocolFlag   string
 	cliIPFlagsFlag      string
+	cliWorkerCountFlag  int
 
 	cliInterface  *net.Interface
 	cliSrcIP      net.IP
@@ -129,6 +133,8 @@ var (
 const (
 	cliHTTPProtoFlag  = "http"
 	cliHTTPSProtoFlag = "https"
+
+	defaultWorkerCount = 100
 )
 
 var (
diff --git a/command/socks.go b/command/socks.go
@@ -16,6 +16,7 @@ import (
 func init() {
 	socksCmd.Flags().StringVarP(&cliPortsFlag, "ports", "p", "", "set ports to scan")
 	socksCmd.Flags().StringVarP(&cliIPPortFileFlag, "file", "f", "", "set JSONL file with ip/port pairs to scan")
+	socksCmd.Flags().IntVarP(&cliWorkerCountFlag, "workers", "w", defaultWorkerCount, "set workers count")
 	rootCmd.AddCommand(socksCmd)
 }
 
@@ -57,6 +58,5 @@ func newSOCKSScanEngine(ctx context.Context) scan.EngineResulter {
 		socks5.WithDialTimeout(2*time.Second),
 		socks5.WithDataTimeout(2*time.Second))
 	results := scan.NewResultChan(ctx, 1000)
-	// TODO custom workerCount
-	return scan.NewScanEngine(newIPPortGenerator(), scanner, results, scan.WithScanWorkerCount(50))
+	return scan.NewScanEngine(newIPPortGenerator(), scanner, results, scan.WithScanWorkerCount(cliWorkerCountFlag))
 }
diff --git a/pkg/scan/socks5/socks5.go b/pkg/scan/socks5/socks5.go
@@ -2,7 +2,7 @@ package socks5
 
 import (
 	"context"
-	json "encoding/json"
+	"encoding/json"
 	"fmt"
 	"net"
 	"time"

Original file line number	Diff line number	Diff line change
`@@ -18,6 +18,7 @@ func init() {`
`18`	`18`	`dockerCmd.Flags().StringVarP(&cliPortsFlag, "ports", "p", "", "set ports to scan")`
`19`	`19`	`dockerCmd.Flags().StringVarP(&cliIPPortFileFlag, "file", "f", "", "set JSONL file with ip/port pairs to scan")`
`20`	`20`	`dockerCmd.Flags().StringVar(&cliProtoFlag, "proto", "", "set protocol to use, http is used by default; only http or https are valid")`
	`21`	`+ dockerCmd.Flags().IntVarP(&cliWorkerCountFlag, "workers", "w", defaultWorkerCount, "set workers count")`
`21`	`22`	`rootCmd.AddCommand(dockerCmd)`
`22`	`23`	`}`
`23`	`24`
`@@ -63,6 +64,5 @@ func newDockerScanEngine(ctx context.Context) scan.EngineResulter {`
`63`	`64`	`// TODO custom dataTimeout`
`64`	`65`	`scanner := docker.NewScanner(cliProtoFlag, docker.WithDataTimeout(10*time.Second))`
`65`	`66`	`results := scan.NewResultChan(ctx, 1000)`
`66`		`- // TODO custom workerCount`
`67`		`- return scan.NewScanEngine(newIPPortGenerator(), scanner, results, scan.WithScanWorkerCount(50))`
	`67`	`+ return scan.NewScanEngine(newIPPortGenerator(), scanner, results, scan.WithScanWorkerCount(cliWorkerCountFlag))`
`68`	`68`	`}`
Original file line number	Diff line number	Diff line change
`@@ -90,6 +90,9 @@ var rootCmd = &cobra.Command{`
`90`	`90`	`return`
`91`	`91`	`}`
`92`	`92`	`}`
	`93`	`+ if cliWorkerCountFlag <= 0 {`
	`94`	`+ return errors.New("invalid workers count")`
	`95`	`+ }`
`93`	`96`	`return`
`94`	`97`	`},`
`95`	`98`	`}`
`@@ -110,6 +113,7 @@ var (`
`110`	`113`	`cliIPTotalLenFlag string`
`111`	`114`	`cliIPProtocolFlag string`
`112`	`115`	`cliIPFlagsFlag string`
	`116`	`+ cliWorkerCountFlag int`
`113`	`117`
`114`	`118`	`cliInterface *net.Interface`
`115`	`119`	`cliSrcIP net.IP`
`@@ -129,6 +133,8 @@ var (`
`129`	`133`	`const (`
`130`	`134`	`cliHTTPProtoFlag = "http"`
`131`	`135`	`cliHTTPSProtoFlag = "https"`
	`136`	`+`
	`137`	`+ defaultWorkerCount = 100`
`132`	`138`	`)`
`133`	`139`
`134`	`140`	`var (`
Original file line number	Diff line number	Diff line change
`@@ -16,6 +16,7 @@ import (`
`16`	`16`	`func init() {`
`17`	`17`	`socksCmd.Flags().StringVarP(&cliPortsFlag, "ports", "p", "", "set ports to scan")`
`18`	`18`	`socksCmd.Flags().StringVarP(&cliIPPortFileFlag, "file", "f", "", "set JSONL file with ip/port pairs to scan")`
	`19`	`+ socksCmd.Flags().IntVarP(&cliWorkerCountFlag, "workers", "w", defaultWorkerCount, "set workers count")`
`19`	`20`	`rootCmd.AddCommand(socksCmd)`
`20`	`21`	`}`
`21`	`22`
`@@ -57,6 +58,5 @@ func newSOCKSScanEngine(ctx context.Context) scan.EngineResulter {`
`57`	`58`	`socks5.WithDialTimeout(2*time.Second),`
`58`	`59`	`socks5.WithDataTimeout(2*time.Second))`
`59`	`60`	`results := scan.NewResultChan(ctx, 1000)`
`60`		`- // TODO custom workerCount`
`61`		`- return scan.NewScanEngine(newIPPortGenerator(), scanner, results, scan.WithScanWorkerCount(50))`
	`61`	`+ return scan.NewScanEngine(newIPPortGenerator(), scanner, results, scan.WithScanWorkerCount(cliWorkerCountFlag))`
`62`	`62`	`}`