Exit delay option

v-byte-cpu · Mar 27, 2021 · 71c3d1b · 71c3d1b
1 parent 9a0726f
commit 71c3d1b
Show file tree

Hide file tree

Showing 8 changed files with 95 additions and 59 deletions.
diff --git a/command/arp.go b/command/arp.go
@@ -51,14 +51,11 @@ var arpCmd = &cobra.Command{
 
 		m := newARPScanMethod(ctx)
 
-		return startEngine(ctx, &engineConfig{
-			logger:     logger,
-			scanRange:  r,
-			scanMethod: m,
-			bpfFilter:  arp.BPFFilter,
-			rateCount:  cliRateCount,
-			rateWindow: cliRateWindow,
-		})
+		return startEngine(ctx, newEngineConfig(
+			withLogger(logger),
+			withScanRange(r),
+			withScanMethod(m),
+			withBPFFilter(arp.BPFFilter)))
 	},
 }
 

diff --git a/command/root.go b/command/root.go
@@ -53,6 +53,11 @@ var rootCmd = &cobra.Command{
 				return
 			}
 		}
+		if len(cliExitDelayFlag) > 0 {
+			if cliExitDelay, err = time.ParseDuration(cliExitDelayFlag); err != nil {
+				return
+			}
+		}
 		return
 	},
 }
@@ -64,13 +69,15 @@ var (
 	cliSrcMACFlag    string
 	cliPortsFlag     string
 	cliRateLimitFlag string
+	cliExitDelayFlag string
 
 	cliInterface  *net.Interface
 	cliSrcIP      net.IP
 	cliSrcMAC     net.HardwareAddr
 	cliPortRanges []*scan.PortRange
 	cliRateCount  int
 	cliRateWindow time.Duration
+	cliExitDelay  = 300 * time.Millisecond
 )
 
 var (
@@ -85,7 +92,16 @@ func init() {
 	rootCmd.PersistentFlags().StringVarP(&cliInterfaceFlag, "iface", "i", "", "set interface to send/receive packets")
 	rootCmd.PersistentFlags().StringVar(&cliSrcIPFlag, "srcip", "", "set source IP address for generated packets")
 	rootCmd.PersistentFlags().StringVar(&cliSrcMACFlag, "srcmac", "", "set source MAC address for generated packets")
-	rootCmd.PersistentFlags().StringVarP(&cliRateLimitFlag, "rate", "r", "", "set rate limit for generated packets")
+	rootCmd.PersistentFlags().StringVarP(&cliRateLimitFlag, "rate", "r", "",
+		strings.Join([]string{
+			"set rate limit for generated packets",
+			`format: "rateCount/rateWindow"`,
+			"where rateCount is a number of packets, rateWindow is the time interval",
+			"e.g. 1000/s -- 1000 packets per second", "500/7s -- 500 packets per 7 seconds\n"}, "\n"))
+	rootCmd.PersistentFlags().StringVar(&cliExitDelayFlag, "exit-delay", "",
+		strings.Join([]string{
+			"set exit delay to wait for response packets",
+			"any expression accepted by time.ParseDuration is valid (300ms by default)"}, "\n"))
 }
 
 func Main() {
@@ -258,13 +274,54 @@ func getGatewayIP(r *scan.Range) (gatewayIP net.IP, err error) {
 	return
 }
 
+type bpfFilterFunc func(r *scan.Range) (filter string, maxPacketLength int)
+
 type engineConfig struct {
 	logger     log.Logger
 	scanRange  *scan.Range
 	scanMethod resultScanMethod
-	bpfFilter  func(r *scan.Range) (filter string, maxPacketLength int)
+	bpfFilter  bpfFilterFunc
 	rateCount  int
 	rateWindow time.Duration
+	exitDelay  time.Duration
+}
+
+type engineConfigOption func(c *engineConfig)
+
+func withLogger(logger log.Logger) engineConfigOption {
+	return func(c *engineConfig) {
+		c.logger = logger
+	}
+}
+
+func withScanRange(r *scan.Range) engineConfigOption {
+	return func(c *engineConfig) {
+		c.scanRange = r
+	}
+}
+
+func withScanMethod(sm resultScanMethod) engineConfigOption {
+	return func(c *engineConfig) {
+		c.scanMethod = sm
+	}
+}
+
+func withBPFFilter(bpfFilter bpfFilterFunc) engineConfigOption {
+	return func(c *engineConfig) {
+		c.bpfFilter = bpfFilter
+	}
+}
+
+func newEngineConfig(opts ...engineConfigOption) *engineConfig {
+	c := &engineConfig{
+		rateCount:  cliRateCount,
+		rateWindow: cliRateWindow,
+		exitDelay:  cliExitDelay,
+	}
+	for _, o := range opts {
+		o(c)
+	}
+	return c
 }
 
 type resultScanMethod interface {
@@ -311,7 +368,7 @@ func startEngine(ctx context.Context, conf *engineConfig) error {
 	go func() {
 		defer cancel()
 		<-done
-		<-time.After(300 * time.Millisecond)
+		<-time.After(conf.exitDelay)
 	}()
 
 	// error logging

diff --git a/command/tcp.go b/command/tcp.go
@@ -86,14 +86,11 @@ var tcpCmd = &cobra.Command{
 			withTCPPacketFlags(tcp.AllFlags),
 		)
 
-		return startEngine(ctx, &engineConfig{
-			logger:     conf.logger,
-			scanRange:  conf.scanRange,
-			scanMethod: m,
-			bpfFilter:  tcp.BPFFilter,
-			rateCount:  cliRateCount,
-			rateWindow: cliRateWindow,
-		})
+		return startEngine(ctx, newEngineConfig(
+			withLogger(conf.logger),
+			withScanRange(conf.scanRange),
+			withScanMethod(m),
+			withBPFFilter(tcp.BPFFilter)))
 	},
 }
 

diff --git a/command/tcp_fin.go b/command/tcp_fin.go
@@ -43,13 +43,10 @@ var tcpfinCmd = &cobra.Command{
 			withTCPPacketFlags(tcp.AllFlags),
 		)
 
-		return startEngine(ctx, &engineConfig{
-			logger:     conf.logger,
-			scanRange:  conf.scanRange,
-			scanMethod: m,
-			bpfFilter:  tcp.BPFFilter,
-			rateCount:  cliRateCount,
-			rateWindow: cliRateWindow,
-		})
+		return startEngine(ctx, newEngineConfig(
+			withLogger(conf.logger),
+			withScanRange(conf.scanRange),
+			withScanMethod(m),
+			withBPFFilter(tcp.BPFFilter)))
 	},
 }
diff --git a/command/tcp_null.go b/command/tcp_null.go
@@ -43,13 +43,10 @@ var tcpnullCmd = &cobra.Command{
 			withTCPPacketFlags(tcp.AllFlags),
 		)
 
-		return startEngine(ctx, &engineConfig{
-			logger:     conf.logger,
-			scanRange:  conf.scanRange,
-			scanMethod: m,
-			bpfFilter:  tcp.BPFFilter,
-			rateCount:  cliRateCount,
-			rateWindow: cliRateWindow,
-		})
+		return startEngine(ctx, newEngineConfig(
+			withLogger(conf.logger),
+			withScanRange(conf.scanRange),
+			withScanMethod(m),
+			withBPFFilter(tcp.BPFFilter)))
 	},
 }
diff --git a/command/tcp_syn.go b/command/tcp_syn.go
@@ -51,13 +51,10 @@ func startTCPSYNScan(ctx context.Context, subnet string) (err error) {
 		withTCPPacketFlags(tcp.EmptyFlags),
 	)
 
-	return startEngine(ctx, &engineConfig{
-		logger:     conf.logger,
-		scanRange:  conf.scanRange,
-		scanMethod: m,
+	return startEngine(ctx, newEngineConfig(
+		withLogger(conf.logger),
+		withScanRange(conf.scanRange),
+		withScanMethod(m),
 		// TODO SYN,ACK filter
-		bpfFilter:  tcp.BPFFilter,
-		rateCount:  cliRateCount,
-		rateWindow: cliRateWindow,
-	})
+		withBPFFilter(tcp.BPFFilter)))
 }
diff --git a/command/tcp_xmas.go b/command/tcp_xmas.go
@@ -43,13 +43,10 @@ var tcpxmasCmd = &cobra.Command{
 			withTCPPacketFlags(tcp.AllFlags),
 		)
 
-		return startEngine(ctx, &engineConfig{
-			logger:     conf.logger,
-			scanRange:  conf.scanRange,
-			scanMethod: m,
-			bpfFilter:  tcp.BPFFilter,
-			rateCount:  cliRateCount,
-			rateWindow: cliRateWindow,
-		})
+		return startEngine(ctx, newEngineConfig(
+			withLogger(conf.logger),
+			withScanRange(conf.scanRange),
+			withScanMethod(m),
+			withBPFFilter(tcp.BPFFilter)))
 	},
 }
diff --git a/command/udp.go b/command/udp.go
@@ -41,14 +41,11 @@ var udpCmd = &cobra.Command{
 
 		m := newUDPScanMethod(ctx, conf)
 
-		return startEngine(ctx, &engineConfig{
-			logger:     conf.logger,
-			scanRange:  conf.scanRange,
-			scanMethod: m,
-			bpfFilter:  icmp.BPFFilter,
-			rateCount:  cliRateCount,
-			rateWindow: cliRateWindow,
-		})
+		return startEngine(ctx, newEngineConfig(
+			withLogger(conf.logger),
+			withScanRange(conf.scanRange),
+			withScanMethod(m),
+			withBPFFilter(icmp.BPFFilter)))
 	},
 }