From 88cc39ed218286d54aaaf34ba7a6cf24a9e90aa1 Mon Sep 17 00:00:00 2001
From: v-byte-cpu <65545655+v-byte-cpu@users.noreply.github.com>
Date: Sat, 20 Mar 2021 17:58:10 +0300
Subject: [PATCH] Refactor ARP scan

---
 command/arp.go           | 17 ++++++++++-----
 pkg/scan/arp/arp.go      | 46 ++++++++--------------------------------
 pkg/scan/arp/arp_test.go |  2 +-
 3 files changed, 22 insertions(+), 43 deletions(-)

diff --git a/command/arp.go b/command/arp.go
index 504eb79..098f08e 100644
--- a/command/arp.go
+++ b/command/arp.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"os"
 	"os/signal"
+	"runtime"
 	"strings"
 	"time"
 
@@ -48,11 +49,7 @@ var arpCmd = &cobra.Command{
 			logger = log.NewUniqueLogger(logger)
 		}
 
-		var opts []arp.ScanMethodOption
-		if arpLiveModeFlag {
-			opts = append(opts, arp.LiveMode(1*time.Second))
-		}
-		m := arp.NewScanMethod(ctx, opts...)
+		m := newARPScanMethod(ctx)
 
 		return startEngine(ctx, &engineConfig{
 			logger:     logger,
@@ -62,3 +59,13 @@ var arpCmd = &cobra.Command{
 		})
 	},
 }
+
+func newARPScanMethod(ctx context.Context) *arp.ScanMethod {
+	var reqgen scan.RequestGenerator = scan.RequestGeneratorFunc(scan.Requests)
+	if arpLiveModeFlag {
+		reqgen = scan.NewLiveRequestGenerator(1 * time.Second)
+	}
+	pktgen := scan.NewPacketMultiGenerator(arp.NewPacketFiller(), runtime.NumCPU())
+	psrc := scan.NewPacketSource(reqgen, pktgen)
+	return arp.NewScanMethod(ctx, psrc)
+}
diff --git a/pkg/scan/arp/arp.go b/pkg/scan/arp/arp.go
index eaa3a20..b2214d0 100644
--- a/pkg/scan/arp/arp.go
+++ b/pkg/scan/arp/arp.go
@@ -6,19 +6,15 @@ import (
 	"context"
 	"fmt"
 	"net"
-	"runtime"
-	"time"
 
 	"github.com/google/gopacket"
 	"github.com/google/gopacket/layers"
 	"github.com/google/gopacket/macs"
-	"github.com/v-byte-cpu/sx/pkg/packet"
 	"github.com/v-byte-cpu/sx/pkg/scan"
 )
 
 type ScanMethod struct {
-	reqgen  scan.RequestGenerator
-	pktgen  scan.PacketGenerator
+	scan.PacketSource
 	parser  *gopacket.DecodingLayerParser
 	results *scan.ResultChan
 	ctx     context.Context
@@ -47,28 +43,15 @@ func (r *ScanResult) ID() string {
 	return r.IP
 }
 
-type ScanMethodOption func(sm *ScanMethod)
-
-func LiveMode(rescanTimeout time.Duration) ScanMethodOption {
-	return func(sm *ScanMethod) {
-		sm.reqgen = scan.NewLiveRequestGenerator(rescanTimeout)
-	}
-}
-
-func NewScanMethod(ctx context.Context, opts ...ScanMethodOption) *ScanMethod {
+func NewScanMethod(ctx context.Context, psrc scan.PacketSource) *ScanMethod {
 	sm := &ScanMethod{
-		ctx:     ctx,
-		results: scan.NewResultChan(ctx, 1000),
-		reqgen:  scan.RequestGeneratorFunc(scan.Requests),
-		pktgen:  scan.NewPacketMultiGenerator(newPacketFiller(), runtime.NumCPU()),
+		PacketSource: psrc,
+		ctx:          ctx,
+		results:      scan.NewResultChan(ctx, 1000),
 	}
 	parser := gopacket.NewDecodingLayerParser(layers.LayerTypeEthernet, &sm.rcvEth, &sm.rcvARP)
 	parser.IgnoreUnsupported = true
 	sm.parser = parser
-
-	for _, o := range opts {
-		o(sm)
-	}
 	return sm
 }
 
@@ -76,17 +59,6 @@ func (s *ScanMethod) Results() <-chan scan.Result {
 	return s.results.Chan()
 }
 
-func (s *ScanMethod) Packets(ctx context.Context, r *scan.Range) <-chan *packet.BufferData {
-	requests, err := s.reqgen.GenerateRequests(ctx, r)
-	if err != nil {
-		out := make(chan *packet.BufferData, 1)
-		out <- &packet.BufferData{Err: err}
-		close(out)
-		return out
-	}
-	return s.pktgen.Packets(ctx, requests)
-}
-
 func (s *ScanMethod) ProcessPacketData(data []byte, _ *gopacket.CaptureInfo) error {
 	// try to exit as early as possible
 	select {
@@ -115,13 +87,13 @@ func (s *ScanMethod) ProcessPacketData(data []byte, _ *gopacket.CaptureInfo) err
 	return nil
 }
 
-type packetFiller struct{}
+type PacketFiller struct{}
 
-func newPacketFiller() *packetFiller {
-	return &packetFiller{}
+func NewPacketFiller() *PacketFiller {
+	return &PacketFiller{}
 }
 
-func (*packetFiller) Fill(packet gopacket.SerializeBuffer, pair *scan.Request) error {
+func (*PacketFiller) Fill(packet gopacket.SerializeBuffer, pair *scan.Request) error {
 	eth := &layers.Ethernet{
 		SrcMAC:       pair.SrcMAC,
 		DstMAC:       net.HardwareAddr{0xff, 0xff, 0xff, 0xff, 0xff, 0xff},
diff --git a/pkg/scan/arp/arp_test.go b/pkg/scan/arp/arp_test.go
index 141f659..6be6763 100644
--- a/pkg/scan/arp/arp_test.go
+++ b/pkg/scan/arp/arp_test.go
@@ -22,7 +22,7 @@ func TestProcessPacketData(t *testing.T) {
 
 		ctx, cancel := context.WithCancel(context.Background())
 		defer cancel()
-		sm := NewScanMethod(ctx)
+		sm := NewScanMethod(ctx, nil)
 
 		// generate packet data
 		packet := gopacket.NewSerializeBuffer()