Port ranges

command/arp.go

@@ -61,9 +61,9 @@ var arpCmd = &cobra.Command{
 }
 
 func newARPScanMethod(ctx context.Context) *arp.ScanMethod {
-	var reqgen scan.RequestGenerator = scan.RequestGeneratorFunc(scan.Requests)
+	var reqgen scan.RequestGenerator = scan.NewIPRequestGenerator(scan.NewIPGenerator())
 	if arpLiveModeFlag {
-		reqgen = scan.NewLiveRequestGenerator(1 * time.Second)
+		reqgen = scan.NewLiveRequestGenerator(reqgen, 1*time.Second)
 	}
 	pktgen := scan.NewPacketMultiGenerator(arp.NewPacketFiller(), runtime.NumCPU())
 	psrc := scan.NewPacketSource(reqgen, pktgen)

command/root.go

@@ -65,7 +65,7 @@ func parseScanConfig(scanName, subnet, ports string) (c *scanConfig, err error)
 	if r, err = parseScanRange(subnet); err != nil {
 		return
 	}
-	if r.StartPort, r.EndPort, err = parsePortRange(ports); err != nil {
+	if r.Ports, err = parsePortRanges(ports); err != nil {
 		return
 	}
 
@@ -133,22 +133,31 @@ func parseScanRange(subnet string) (*scan.Range, error) {
 		SrcMAC:    srcMAC}, nil
 }
 
-// TODO port ranges with tests
-func parsePortRange(portsRange string) (startPort, endPort uint16, err error) {
+func parsePortRange(portsRange string) (r *scan.PortRange, err error) {
 	ports := strings.Split(portsRange, "-")
 	var port uint64
 	if port, err = strconv.ParseUint(ports[0], 10, 16); err != nil {
 		return
 	}
-	startPort = uint16(port)
+	result := &scan.PortRange{StartPort: uint16(port), EndPort: uint16(port)}
 	if len(ports) < 2 {
-		endPort = startPort
-		return
+		return result, nil
 	}
 	if port, err = strconv.ParseUint(ports[1], 10, 16); err != nil {
 		return
 	}
-	endPort = uint16(port)
+	result.EndPort = uint16(port)
+	return result, nil
+}
+
+func parsePortRanges(portsRanges string) (result []*scan.PortRange, err error) {
+	var ports *scan.PortRange
+	for _, portsRange := range strings.Split(portsRanges, ",") {
+		if ports, err = parsePortRange(portsRange); err != nil {
+			return
+		}
+		result = append(result, ports)
+	}
 	return
 }
 

command/root_test.go

@@ -0,0 +1,143 @@
+package command
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"github.com/v-byte-cpu/sx/pkg/scan"
+)
+
+func TestParsePortRangeError(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name       string
+		portsRange string
+	}{
+		{
+			name:       "EmptyPortRange",
+			portsRange: "",
+		},
+		{
+			name:       "EmptyStartPort",
+			portsRange: "-22",
+		},
+		{
+			name:       "EmptyEndPort",
+			portsRange: "22-",
+		},
+		{
+			name:       "InvalidLargePort",
+			portsRange: "65536",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			_, err := parsePortRange(tt.portsRange)
+			require.Error(t, err)
+		})
+	}
+}
+
+func TestParsePortRange(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name       string
+		portsRange string
+		expected   *scan.PortRange
+	}{
+		{
+			name:       "OnePort",
+			portsRange: "22",
+			expected: &scan.PortRange{
+				StartPort: 22,
+				EndPort:   22,
+			},
+		},
+		{
+			name:       "TwoPorts",
+			portsRange: "22-23",
+			expected: &scan.PortRange{
+				StartPort: 22,
+				EndPort:   23,
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ports, err := parsePortRange(tt.portsRange)
+			require.NoError(t, err)
+			require.Equal(t, tt.expected, ports)
+		})
+	}
+}
+
+func TestParsePortRanges(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name       string
+		portsRange string
+		expected   []*scan.PortRange
+	}{
+		{
+			name:       "OneRangeOnePort",
+			portsRange: "22",
+			expected: []*scan.PortRange{
+				{
+					StartPort: 22,
+					EndPort:   22,
+				},
+			},
+		},
+		{
+			name:       "OneRangeTwoPorts",
+			portsRange: "22-23",
+			expected: []*scan.PortRange{
+				{
+					StartPort: 22,
+					EndPort:   23,
+				},
+			},
+		},
+		{
+			name:       "TwoRangesOnePort",
+			portsRange: "22,23",
+			expected: []*scan.PortRange{
+				{
+					StartPort: 22,
+					EndPort:   22,
+				},
+				{
+					StartPort: 23,
+					EndPort:   23,
+				},
+			},
+		},
+		{
+			name:       "TwoRangesTwoPorts",
+			portsRange: "22-23,26-27",
+			expected: []*scan.PortRange{
+				{
+					StartPort: 22,
+					EndPort:   23,
+				},
+				{
+					StartPort: 26,
+					EndPort:   27,
+				},
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ports, err := parsePortRanges(tt.portsRange)
+			require.NoError(t, err)
+			require.Equal(t, tt.expected, ports)
+		})
+	}
+}

command/tcp.go

@@ -158,8 +158,10 @@ func newTCPScanMethod(ctx context.Context, conf *scanConfig, opts ...tcpScanConf
 	for _, o := range opts {
 		o(c)
 	}
+	portgen := scan.NewPortGenerator()
+	ipgen := scan.NewIPGenerator()
 	reqgen := arp.NewCacheRequestGenerator(
-		scan.RequestGeneratorFunc(scan.Requests), conf.gatewayIP, conf.cache)
+		scan.NewIPPortRequestGenerator(ipgen, portgen), conf.gatewayIP, conf.cache)
 	pktgen := scan.NewPacketMultiGenerator(c.packetFiller, runtime.NumCPU())
 	psrc := scan.NewPacketSource(reqgen, pktgen)
 	results := scan.NewResultChan(ctx, 1000)

command/udp.go

@@ -51,8 +51,10 @@ var udpCmd = &cobra.Command{
 }
 
 func newUDPScanMethod(ctx context.Context, conf *scanConfig) *udp.ScanMethod {
+	portgen := scan.NewPortGenerator()
+	ipgen := scan.NewIPGenerator()
 	reqgen := arp.NewCacheRequestGenerator(
-		scan.RequestGeneratorFunc(scan.Requests), conf.gatewayIP, conf.cache)
+		scan.NewIPPortRequestGenerator(ipgen, portgen), conf.gatewayIP, conf.cache)
 	pktgen := scan.NewPacketMultiGenerator(udp.NewPacketFiller(), runtime.NumCPU())
 	psrc := scan.NewPacketSource(reqgen, pktgen)
 	results := scan.NewResultChan(ctx, 1000)

pkg/packet/receiver.go

@@ -25,15 +25,15 @@ type Receiver interface {
 	ReceivePackets(ctx context.Context) <-chan error
 }
 
+func NewReceiver(sr Reader, p Processor) Receiver {
+	return &receiver{sr, p}
+}
+
 type receiver struct {
 	sr Reader
 	p  Processor
 }
 
-func NewReceiver(sr Reader, p Processor) Receiver {
-	return &receiver{sr, p}
-}
-
 func isTemporaryError(err error) bool {
 	if err == syscall.EAGAIN {
 		return true

pkg/packet/sender.go

@@ -21,14 +21,14 @@ type Writer interface {
 	WritePacketData(pkt []byte) error
 }
 
-type sender struct {
-	w Writer
-}
-
 func NewSender(w Writer) Sender {
 	return &sender{w}
 }
 
+type sender struct {
+	w Writer
+}
+
 func (s *sender) SendPackets(ctx context.Context, in <-chan *BufferData) (<-chan interface{}, <-chan error) {
 	done := make(chan interface{})
 	errc := make(chan error, 100)

pkg/scan/engine.go

@@ -10,6 +10,11 @@ import (
 	"github.com/v-byte-cpu/sx/pkg/packet"
 )
 
+type PortRange struct {
+	StartPort uint16
+	EndPort   uint16
+}
+
 type Range struct {
 	Interface *net.Interface
 	DstSubnet *net.IPNet
@@ -18,21 +23,22 @@ type Range struct {
 	SrcMAC    net.HardwareAddr
 	StartPort uint16
 	EndPort   uint16
+	Ports     []*PortRange
 }
 
 type PacketSource interface {
 	Packets(ctx context.Context, r *Range) <-chan *packet.BufferData
 }
 
+func NewPacketSource(reqgen RequestGenerator, pktgen PacketGenerator) PacketSource {
+	return &packetSource{reqgen, pktgen}
+}
+
 type packetSource struct {
 	reqgen RequestGenerator
 	pktgen PacketGenerator
 }
 
-func NewPacketSource(reqgen RequestGenerator, pktgen PacketGenerator) PacketSource {
-	return &packetSource{reqgen, pktgen}
-}
-
 func (s *packetSource) Packets(ctx context.Context, r *Range) <-chan *packet.BufferData {
 	requests, err := s.reqgen.GenerateRequests(ctx, r)
 	if err != nil {

pkg/scan/engine_test.go

@@ -106,8 +106,12 @@ func TestEngineStartCollectsAllErrors(t *testing.T) {
 			IP:   net.IPv4(192, 168, 0, 1),
 			Mask: net.CIDRMask(32, 32),
 		},
-		StartPort: 888,
-		EndPort:   888,
+		Ports: []*PortRange{
+			{
+				StartPort: 888,
+				EndPort:   888,
+			},
+		},
 	})
 
 	result := chanToSlice(t, chanErrToGeneric(out), 2)
@@ -129,10 +133,14 @@ func TestPacketSourceReturnsError(t *testing.T) {
 		pktgen := NewMockPacketGenerator(ctrl)
 
 		expectedScanRange := &Range{
-			SrcIP:     net.IPv4(192, 168, 0, 1),
-			SrcMAC:    net.HardwareAddr{0x1, 0x2, 0x3, 0x4, 0x5, 0x6},
-			StartPort: 22,
-			EndPort:   22,
+			SrcIP:  net.IPv4(192, 168, 0, 1),
+			SrcMAC: net.HardwareAddr{0x1, 0x2, 0x3, 0x4, 0x5, 0x6},
+			Ports: []*PortRange{
+				{
+					StartPort: 22,
+					EndPort:   22,
+				},
+			},
 		}
 		var scanRange Range
 		err := copier.Copy(&scanRange, expectedScanRange)
@@ -166,10 +174,14 @@ func TestPacketSourceReturnsData(t *testing.T) {
 		pktgen := NewMockPacketGenerator(ctrl)
 
 		scanRange := &Range{
-			SrcIP:     net.IPv4(192, 168, 0, 1),
-			SrcMAC:    net.HardwareAddr{0x1, 0x2, 0x3, 0x4, 0x5, 0x6},
-			StartPort: 22,
-			EndPort:   22,
+			SrcIP:  net.IPv4(192, 168, 0, 1),
+			SrcMAC: net.HardwareAddr{0x1, 0x2, 0x3, 0x4, 0x5, 0x6},
+			Ports: []*PortRange{
+				{
+					StartPort: 22,
+					EndPort:   22,
+				},
+			},
 		}
 		requests := make(chan *Request)
 		close(requests)

pkg/scan/generator.go

@@ -18,14 +18,14 @@ type PacketGenerator interface {
 	Packets(ctx context.Context, in <-chan *Request) <-chan *packet.BufferData
 }
 
-type packetGenerator struct {
-	filler PacketFiller
-}
-
 func NewPacketGenerator(filler PacketFiller) PacketGenerator {
 	return &packetGenerator{filler}
 }
 
+type packetGenerator struct {
+	filler PacketFiller
+}
+
 func (g *packetGenerator) Packets(ctx context.Context, in <-chan *Request) <-chan *packet.BufferData {
 	out := make(chan *packet.BufferData)
 	go func() {
@@ -63,16 +63,16 @@ func writeBufToChan(ctx context.Context, out chan *packet.BufferData, buf *packe
 	}
 }
 
-type packetMultiGenerator struct {
-	gen        *packetGenerator
-	numWorkers int
-}
-
 func NewPacketMultiGenerator(filler PacketFiller, numWorkers int) PacketGenerator {
 	gen := &packetGenerator{filler}
 	return &packetMultiGenerator{gen, numWorkers}
 }
 
+type packetMultiGenerator struct {
+	gen        *packetGenerator
+	numWorkers int
+}
+
 func (g *packetMultiGenerator) Packets(ctx context.Context, in <-chan *Request) <-chan *packet.BufferData {
 	workers := make([]<-chan *packet.BufferData, g.numWorkers)
 	for i := 0; i < g.numWorkers; i++ {