From 1493c998d5361f0a81551df7fb94b6fb622bcff7 Mon Sep 17 00:00:00 2001
From: v-byte-cpu <65545655+v-byte-cpu@users.noreply.github.com>
Date: Sun, 28 Mar 2021 17:55:43 +0300
Subject: [PATCH] ARP cache file option

---
 README.md       | 12 ++++++++++++
 command/root.go | 49 ++++++++++++++++++++++++++++++++++++++-----------
 command/tcp.go  |  2 ++
 command/udp.go  |  2 ++
 4 files changed, 54 insertions(+), 11 deletions(-)

diff --git a/README.md b/README.md
index 6e2e029..d026fb8 100644
--- a/README.md
+++ b/README.md
@@ -128,6 +128,18 @@ or individual ports:
 cat arp.cache | ./sx tcp -p 22,443 192.168.0.171
 ```
 
+It is possible to specify the ARP cache file using the `-a` or `--arp-cache` options:
+
+```
+./sx tcp -a arp.cache -p 22,443 192.168.0.171
+```
+
+or stdin redirect:
+
+```
+./sx tcp -p 22,443 192.168.0.171 < arp.cache
+```
+
 You can also use the `tcp syn` subcommand instead of the `tcp`:
 
 ```
diff --git a/command/root.go b/command/root.go
index 3909ce5..be11e3d 100644
--- a/command/root.go
+++ b/command/root.go
@@ -1,6 +1,7 @@
 package command
 
 import (
+	"bufio"
 	"context"
 	"errors"
 	"io"
@@ -63,13 +64,14 @@ var rootCmd = &cobra.Command{
 }
 
 var (
-	cliJSONFlag      bool
-	cliInterfaceFlag string
-	cliSrcIPFlag     string
-	cliSrcMACFlag    string
-	cliPortsFlag     string
-	cliRateLimitFlag string
-	cliExitDelayFlag string
+	cliJSONFlag         bool
+	cliInterfaceFlag    string
+	cliSrcIPFlag        string
+	cliSrcMACFlag       string
+	cliPortsFlag        string
+	cliRateLimitFlag    string
+	cliExitDelayFlag    string
+	cliARPCacheFileFlag string
 
 	cliInterface  *net.Interface
 	cliSrcIP      net.IP
@@ -85,6 +87,7 @@ var (
 	errSrcMAC       = errors.New("invalid source MAC")
 	errSrcInterface = errors.New("invalid source interface")
 	errRateLimit    = errors.New("invalid ratelimit")
+	errStdin        = errors.New("stdin is from a terminal")
 )
 
 func init() {
@@ -128,10 +131,8 @@ func parseScanConfig(scanName, subnet string) (c *scanConfig, err error) {
 		return
 	}
 
-	// TODO file argument
-	// TODO handle pipes
-	cache := arp.NewCache()
-	if err = arp.FillCache(cache, os.Stdin); err != nil {
+	var cache *arp.Cache
+	if cache, err = parseARPCache(); err != nil {
 		return
 	}
 
@@ -148,6 +149,32 @@ func parseScanConfig(scanName, subnet string) (c *scanConfig, err error) {
 	return
 }
 
+func parseARPCache() (cache *arp.Cache, err error) {
+	var r io.Reader
+	if len(cliARPCacheFileFlag) > 0 {
+		var f *os.File
+		if f, err = os.Open(cliARPCacheFileFlag); err != nil {
+			return
+		}
+		defer f.Close()
+		r = bufio.NewReader(f)
+	} else {
+		var info os.FileInfo
+		if info, err = os.Stdin.Stat(); err != nil {
+			return
+		}
+		// only data being piped to stdin is valid
+		if (info.Mode() & os.ModeCharDevice) != 0 {
+			// stdin from terminal is not valid
+			return nil, errStdin
+		}
+		r = os.Stdin
+	}
+	cache = arp.NewCache()
+	err = arp.FillCache(cache, r)
+	return
+}
+
 func parseScanRange(subnet string) (*scan.Range, error) {
 	dstSubnet, err := ip.ParseIPNet(subnet)
 	if err != nil {
diff --git a/command/tcp.go b/command/tcp.go
index 4501b10..bc44884 100644
--- a/command/tcp.go
+++ b/command/tcp.go
@@ -38,6 +38,8 @@ func init() {
 	if err := tcpCmd.MarkPersistentFlagRequired("ports"); err != nil {
 		golog.Fatalln(err)
 	}
+	tcpCmd.PersistentFlags().StringVarP(&cliARPCacheFileFlag, "arp-cache", "a", "",
+		strings.Join([]string{"set ARP cache file", "reads from stdin by default"}, "\n"))
 	tcpCmd.Flags().StringVar(&cliTCPPacketFlags, "flags", "", "set TCP flags")
 	rootCmd.AddCommand(tcpCmd)
 }
diff --git a/command/udp.go b/command/udp.go
index 6659058..58d78d0 100644
--- a/command/udp.go
+++ b/command/udp.go
@@ -17,6 +17,8 @@ import (
 
 func init() {
 	udpCmd.Flags().StringVarP(&cliPortsFlag, "ports", "p", "", "set ports to scan")
+	udpCmd.Flags().StringVarP(&cliARPCacheFileFlag, "arp-cache", "a", "",
+		strings.Join([]string{"set ARP cache file", "reads from stdin by default"}, "\n"))
 	rootCmd.AddCommand(udpCmd)
 }