From 32153f071da94065251d6d0d32422c616df0e108 Mon Sep 17 00:00:00 2001
From: v-byte-cpu <65545655+v-byte-cpu@users.noreply.github.com>
Date: Tue, 11 May 2021 02:12:20 +0300
Subject: [PATCH] feature: workers count option

---
 command/docker.go         | 4 ++--
 command/elastic.go        | 4 ++--
 command/root.go           | 6 ++++++
 command/socks.go          | 4 ++--
 pkg/scan/socks5/socks5.go | 2 +-
 5 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/command/docker.go b/command/docker.go
index 11902d7..60f92ab 100644
--- 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
index db6253d..c0964a7 100644
--- 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
index 8d6bb25..e1ecfa9 100644
--- 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
index de8ca19..63ccdaa 100644
--- 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
index 1669793..8c3a863 100644
--- 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"