feature: workers count option

v-byte-cpu · May 10, 2021 · 32153f0 · 32153f0
1 parent 8845651
commit 32153f0
Show file tree

Hide file tree

Showing 5 changed files with 13 additions and 7 deletions.
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"