From 555ecf679d3622e315cec349553fe71b64199b1f Mon Sep 17 00:00:00 2001
From: "David J. Allen" <allend@lanl.gov>
Date: Mon, 16 Oct 2023 14:37:23 -0600
Subject: [PATCH] Added function implementations to generate hosts by subnet
 and CIDR

---
 cmd/scan.go           | 14 ++++++---
 internal/scan.go      | 72 +++++++++++++++++++++++++++++++++++++++++--
 internal/util/util.go | 12 ++++++++
 3 files changed, 92 insertions(+), 6 deletions(-)

diff --git a/cmd/scan.go b/cmd/scan.go
index d03e734..9e41c74 100644
--- a/cmd/scan.go
+++ b/cmd/scan.go
@@ -16,6 +16,7 @@ var (
 	begin   uint8
 	end     uint8
 	subnets []string
+	subnetMasks []string
 	disableProbing bool
 )
 
@@ -28,8 +29,12 @@ var scanCmd = &cobra.Command{
 		if len(hosts) > 0 {
 			hostsToScan = hosts
 		} else {
-			for _, subnet := range subnets {
-				hostsToScan = append(hostsToScan, magellan.GenerateHosts(subnet, begin, end)...)
+			for i, subnet := range subnets {
+				if len(subnetMasks) > 0 {
+					hostsToScan = append(hostsToScan, magellan.GenerateHostsWithSubnet(subnet, subnetMasks[i])...)
+				} else {
+					hostsToScan = append(hostsToScan, magellan.GenerateHosts(subnet)...)
+				}
 			}
 		}
 
@@ -63,9 +68,10 @@ var scanCmd = &cobra.Command{
 func init() {
 	scanCmd.Flags().StringSliceVar(&hosts, "host", []string{}, "set additional hosts to scan")
 	scanCmd.Flags().IntSliceVar(&ports, "port", []int{}, "set the ports to scan")
-	scanCmd.Flags().Uint8Var(&begin, "begin", 0, "set the starting point for range of IP addresses")
-	scanCmd.Flags().Uint8Var(&end, "end", 255, "set the ending point for range of IP addresses")
+	// scanCmd.Flags().Uint8Var(&begin, "begin", 0, "set the starting point for range of IP addresses")
+	// scanCmd.Flags().Uint8Var(&end, "end", 255, "set the ending point for range of IP addresses")
 	scanCmd.Flags().StringSliceVar(&subnets, "subnet", []string{}, "set additional subnets")
+	scanCmd.Flags().StringSliceVar(&subnetMasks, "subnet-mask", []string{}, "set the subnet masks to use for network")
 	scanCmd.Flags().BoolVar(&disableProbing, "disable-probing", false, "disable probing scanned results for BMC nodes")
 
 	rootCmd.AddCommand(scanCmd)
diff --git a/internal/scan.go b/internal/scan.go
index beee775..2031cf0 100644
--- a/internal/scan.go
+++ b/internal/scan.go
@@ -50,11 +50,79 @@ func rawConnect(host string, ports []int, timeout int, keepOpenOnly bool) []Scan
 	return results
 }
 
-func GenerateHosts(subnet string, begin uint8, end uint8) []string {
+func GenerateHosts(subnet string, mask string, begin uint8, end uint8) []string {
 	hosts := []string{}
 	ip := net.ParseIP(subnet).To4()
 	for i := begin; i < end; i++ {
-		ip[3] = byte(i)
+		ip = util.GetNextIP(ip, 1)
+		hosts = append(hosts, fmt.Sprintf("%d.%d.%d.%d", ip[0], ip[1], ip[2], ip[3]))
+	}
+	return hosts
+}
+
+func GenerateHostsWithCIDR(subnet string) []string {
+	// check for network with valid CIDR
+	ip, network, err := net.ParseCIDR(subnet)
+	if err != nil && (network != nil || ip != nil) {
+		network.Mask = ip.DefaultMask()
+	}
+
+	// check for IP with no CIDR
+	if ip == nil {
+		ip = net.ParseIP(subnet)
+		if ip == nil {
+			return nil
+		}
+	}
+
+	if network == nil {
+		network = &net.IPNet{
+			Mask: ip.DefaultMask(),
+		}
+	}
+	
+	// get all IP addresses in network
+	return generateHosts(ip, network.Mask)
+}
+
+func GenerateHostsWithSubnet(subnet string, subnetMask string) []string {
+	// if no subnet mask, use a default 24-bit mask (for now)
+	if subnetMask != "" {
+		ip, network, err := net.ParseCIDR(subnet)
+		if err != nil && (network != nil || ip != nil) {
+			network.Mask = ip.DefaultMask()
+		}
+		// check for IP with no CIDR
+		if ip == nil {
+			ip = net.ParseIP(subnet)
+			if ip == nil {
+				return nil
+			}
+		}
+
+		if network == nil {
+			network = &net.IPNet{
+				Mask: ip.DefaultMask(),
+			}
+		}
+		return generateHosts(ip, network.Mask)
+	} else {
+		ip := net.ParseIP(subnetMask)
+		if ip != nil {
+			return []string{}
+		}
+		return generateHosts(ip, ip.DefaultMask())
+	}
+}
+
+func generateHosts(ip net.IP, mask net.IPMask) []string {
+	// get all IP addresses in network
+	ones, _ := mask.Size()
+	hosts := []string{}
+	fmt.Printf("ones: %d\n", ones)
+	for i := 0; i < 32-ones; i++ {
+		// ip[3] = byte(i)
+		ip = util.GetNextIP(ip, 1)
 		hosts = append(hosts, fmt.Sprintf("%d.%d.%d.%d", ip[0], ip[1], ip[2], ip[3]))
 	}
 	return hosts
diff --git a/internal/util/util.go b/internal/util/util.go
index 990bda9..668c473 100644
--- a/internal/util/util.go
+++ b/internal/util/util.go
@@ -5,6 +5,7 @@ import (
 	"crypto/tls"
 	"fmt"
 	"io"
+	"net"
 	"net/http"
 	"os"
 	"time"
@@ -17,6 +18,17 @@ func PathExists(path string) (bool, error) {
 	return false, err
 }
 
+func GetNextIP(ip net.IP, inc uint) net.IP {
+	i := ip.To4()
+	v := uint(i[0])<<24 + uint(i[1])<<16 + uint(i[2])<<8 + uint(i[3])
+	v += inc
+	v3 := byte(v & 0xFF)
+	v2 := byte((v >> 8) & 0xFF)
+	v1 := byte((v >> 16) & 0xFF)
+	v0 := byte((v >> 24) & 0xFF)
+	return net.IPv4(v0, v1, v2, v3)
+}
+
 func MakeRequest(url string, httpMethod string, body []byte, headers map[string]string) (*http.Response, []byte, error) {
 	http.DefaultTransport.(*http.Transport).TLSClientConfig = &tls.Config{InsecureSkipVerify: true}
 	req, err := http.NewRequest(httpMethod, url, bytes.NewBuffer(body))