feature: custom timeout option (#70)

v-byte-cpu · May 12, 2021 · 5fc26bf · 5fc26bf
1 parent 9815930
commit 5fc26bf
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 9 deletions.
diff --git a/command/docker.go b/command/docker.go
@@ -6,7 +6,6 @@ import (
 	"os"
 	"os/signal"
 	"strings"
-	"time"
 
 	"github.com/spf13/cobra"
 	"github.com/v-byte-cpu/sx/command/log"
@@ -19,6 +18,7 @@ func init() {
 	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")
+	dockerCmd.Flags().DurationVarP(&cliTimeoutFlag, "timeout", "t", defaultTimeout, "set request timeout")
 	rootCmd.AddCommand(dockerCmd)
 }
 
@@ -61,8 +61,7 @@ var dockerCmd = &cobra.Command{
 }
 
 func newDockerScanEngine(ctx context.Context) scan.EngineResulter {
-	// TODO custom dataTimeout
-	scanner := docker.NewScanner(cliProtoFlag, docker.WithDataTimeout(10*time.Second))
+	scanner := docker.NewScanner(cliProtoFlag, docker.WithDataTimeout(cliTimeoutFlag))
 	results := scan.NewResultChan(ctx, 1000)
 	return scan.NewScanEngine(newIPPortGenerator(), scanner, results, scan.WithScanWorkerCount(cliWorkerCountFlag))
 }
diff --git a/command/elastic.go b/command/elastic.go
@@ -6,7 +6,6 @@ import (
 	"os"
 	"os/signal"
 	"strings"
-	"time"
 
 	"github.com/spf13/cobra"
 	"github.com/v-byte-cpu/sx/command/log"
@@ -19,6 +18,7 @@ func init() {
 	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")
+	elasticCmd.Flags().DurationVarP(&cliTimeoutFlag, "timeout", "t", defaultTimeout, "set request timeout")
 	rootCmd.AddCommand(elasticCmd)
 }
 
@@ -61,8 +61,7 @@ var elasticCmd = &cobra.Command{
 }
 
 func newElasticScanEngine(ctx context.Context) scan.EngineResulter {
-	// TODO custom dataTimeout
-	scanner := elastic.NewScanner(cliProtoFlag, elastic.WithDataTimeout(5*time.Second))
+	scanner := elastic.NewScanner(cliProtoFlag, elastic.WithDataTimeout(cliTimeoutFlag))
 	results := scan.NewResultChan(ctx, 1000)
 	return scan.NewScanEngine(newIPPortGenerator(), scanner, results, scan.WithScanWorkerCount(cliWorkerCountFlag))
 }
diff --git a/command/root.go b/command/root.go
@@ -114,6 +114,7 @@ var (
 	cliIPProtocolFlag   string
 	cliIPFlagsFlag      string
 	cliWorkerCountFlag  int
+	cliTimeoutFlag      time.Duration
 
 	cliInterface  *net.Interface
 	cliSrcIP      net.IP
@@ -135,6 +136,7 @@ const (
 	cliHTTPSProtoFlag = "https"
 
 	defaultWorkerCount = 100
+	defaultTimeout     = 5 * time.Second
 )
 
 var (

diff --git a/command/socks.go b/command/socks.go
@@ -17,6 +17,7 @@ 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")
+	socksCmd.Flags().DurationVarP(&cliTimeoutFlag, "timeout", "t", 2*time.Second, "set connect and data timeout")
 	rootCmd.AddCommand(socksCmd)
 }
 
@@ -53,10 +54,9 @@ var socksCmd = &cobra.Command{
 }
 
 func newSOCKSScanEngine(ctx context.Context) scan.EngineResulter {
-	// TODO custom dialTimeout, dataTimeout
 	scanner := socks5.NewScanner(
-		socks5.WithDialTimeout(2*time.Second),
-		socks5.WithDataTimeout(2*time.Second))
+		socks5.WithDialTimeout(cliTimeoutFlag),
+		socks5.WithDataTimeout(cliTimeoutFlag))
 	results := scan.NewResultChan(ctx, 1000)
 	return scan.NewScanEngine(newIPPortGenerator(), scanner, results, scan.WithScanWorkerCount(cliWorkerCountFlag))
 }