Add live mode

command/arp.go

@@ -3,6 +3,7 @@ package command
 import (
 	"context"
 	"errors"
+	"io"
 	"net"
 	"os"
 	"os/signal"
@@ -23,11 +24,13 @@ var errSrcIP = errors.New("invalid source IP")
 var interfaceFlag string
 var srcIPFlag string
 var srcMACFlag string
+var liveModeFlag bool
 
 func init() {
 	arpCmd.Flags().StringVarP(&interfaceFlag, "iface", "i", "", "set interface to send/receive packets")
 	arpCmd.Flags().StringVar(&srcIPFlag, "srcip", "", "set source IP address for generated packets")
 	arpCmd.Flags().StringVar(&srcMACFlag, "srcmac", "", "set source MAC address for generated packets")
+	arpCmd.Flags().BoolVar(&liveModeFlag, "live", false, "enable live mode")
 	rootCmd.AddCommand(arpCmd)
 }
 
@@ -76,23 +79,35 @@ var arpCmd = &cobra.Command{
 			}
 		}
 
+		ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
+		defer cancel()
+
 		var logger log.Logger
-		if jsonFlag {
-			logger, err = log.NewJSONLogger(os.Stdout, "arp")
-		} else {
-			logger, err = log.NewPlainLogger(os.Stdout, "arp")
-		}
-		if err != nil {
+		if logger, err = getLogger(ctx, os.Stdout); err != nil {
 			return err
 		}
 
 		r := &scan.Range{Subnet: dstSubnet, Interface: iface, SrcIP: srcIP.To4(), SrcMAC: srcMAC}
-		return startEngine(logger, r)
+		return startEngine(ctx, logger, r)
 	},
 }
 
-func startEngine(logger log.Logger, r *scan.Range) error {
-	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
+func getLogger(ctx context.Context, w io.Writer) (logger log.Logger, err error) {
+	opts := []log.LoggerOption{log.FlushInterval(1 * time.Second)}
+	if jsonFlag {
+		opts = append(opts, log.JSON())
+	}
+	if logger, err = log.NewLogger(w, "arp", opts...); err != nil {
+		return
+	}
+	if liveModeFlag {
+		logger = log.NewUniqueLogger(ctx, logger)
+	}
+	return
+}
+
+func startEngine(ctx context.Context, logger log.Logger, r *scan.Range) error {
+	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
 	// setup network interface to read/write packets
@@ -106,7 +121,11 @@ func startEngine(logger log.Logger, r *scan.Range) error {
 		return err
 	}
 
-	m := arp.NewScanMethod(ctx)
+	var opts []arp.ScanMethodOption
+	if liveModeFlag {
+		opts = append(opts, arp.LiveMode(1*time.Second))
+	}
+	m := arp.NewScanMethod(ctx, opts...)
 
 	// setup result logging
 	var wg sync.WaitGroup

command/log/logger.go

@@ -1,8 +1,99 @@
 package log
 
-import "github.com/v-byte-cpu/sx/pkg/scan/arp"
+import (
+	"bufio"
+	"io"
+	"time"
+
+	"github.com/v-byte-cpu/sx/pkg/scan/arp"
+	"go.uber.org/zap"
+)
 
 type Logger interface {
 	Error(err error)
 	LogResults(results <-chan *arp.ScanResult)
 }
+
+type FlushWriter interface {
+	io.Writer
+	Flush() error
+}
+
+type ResultWriter interface {
+	Write(w io.Writer, result *arp.ScanResult) error
+}
+
+type logger struct {
+	zapl  *zap.Logger
+	label string
+
+	w             io.Writer
+	rw            ResultWriter
+	flushInterval time.Duration
+}
+
+type LoggerOption func(*logger)
+
+func JSON() LoggerOption {
+	return func(l *logger) {
+		l.rw = &JSONResultWriter{}
+	}
+}
+
+func Plain() LoggerOption {
+	return func(l *logger) {
+		l.rw = &PlainResultWriter{}
+	}
+}
+
+func FlushInterval(interval time.Duration) LoggerOption {
+	return func(l *logger) {
+		l.flushInterval = interval
+	}
+}
+
+func NewLogger(w io.Writer, label string, opts ...LoggerOption) (Logger, error) {
+	zapl, err := zap.NewProduction()
+	if err != nil {
+		return nil, err
+	}
+
+	l := &logger{
+		zapl:          zapl,
+		label:         label,
+		rw:            &PlainResultWriter{},
+		w:             w,
+		flushInterval: 1 * time.Second,
+	}
+	for _, o := range opts {
+		o(l)
+	}
+	return l, nil
+}
+
+func (l *logger) Error(err error) {
+	l.zapl.Error(l.label, zap.Error(err))
+}
+
+func (l *logger) LogResults(results <-chan *arp.ScanResult) {
+	bw := bufio.NewWriter(l.w)
+	defer bw.Flush()
+	var err error
+	timec := time.After(l.flushInterval)
+	for {
+		select {
+		case result, ok := <-results:
+			if !ok {
+				return
+			}
+			if err := l.rw.Write(l.w, result); err != nil {
+				l.Error(err)
+			}
+		case <-timec:
+			if err = bw.Flush(); err != nil {
+				l.Error(err)
+			}
+			timec = time.After(l.flushInterval)
+		}
+	}
+}