From b3ba75ea9ad4a00942bec76c056716ab61248c80 Mon Sep 17 00:00:00 2001
From: "David J. Allen" <allend@lanl.gov>
Date: Fri, 10 May 2024 13:41:08 -0600
Subject: [PATCH] Added format flag to convery output to JSON and slighly
 changed scanning logic

---
 cmd/list.go | 13 +++++++++--
 cmd/root.go |  1 +
 cmd/scan.go | 64 +++++++++++++++++++++++++++++++++--------------------
 3 files changed, 52 insertions(+), 26 deletions(-)

diff --git a/cmd/list.go b/cmd/list.go
index c5ce1c4..5c1a46c 100644
--- a/cmd/list.go
+++ b/cmd/list.go
@@ -1,7 +1,9 @@
 package cmd
 
 import (
+	"encoding/json"
 	"fmt"
+	"strings"
 
 	"github.com/OpenCHAMI/magellan/internal/db/sqlite"
 
@@ -17,12 +19,19 @@ var listCmd = &cobra.Command{
 		if err != nil {
 			logrus.Errorf("failed toget probe results: %v\n", err)
 		}
-		for _, r := range probeResults {
-			fmt.Printf("%s:%d (%s)\n", r.Host, r.Port, r.Protocol)
+		format = strings.ToLower(format)
+		if format == "json" {
+			b, _ := json.Marshal(probeResults)
+			fmt.Printf("%s\n", string(b))
+		} else {
+			for _, r := range probeResults {
+				fmt.Printf("%s:%d (%s)\n", r.Host, r.Port, r.Protocol)
+			}
 		}
 	},
 }
 
 func init() {
+	listCmd.Flags().StringVar(&format, "format", "", "set the output format")
 	rootCmd.AddCommand(listCmd)
 }
diff --git a/cmd/root.go b/cmd/root.go
index 41b0d20..0e1697c 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -15,6 +15,7 @@ import (
 var (
 	currentUser *user.User
 	accessToken string
+	format      string
 	timeout     int
 	concurrency int
 	ports       []int
diff --git a/cmd/scan.go b/cmd/scan.go
index e114749..557d061 100644
--- a/cmd/scan.go
+++ b/cmd/scan.go
@@ -1,10 +1,12 @@
 package cmd
 
 import (
+	"encoding/json"
 	"fmt"
 	"net"
 	"os"
 	"path"
+	"strings"
 
 	magellan "github.com/OpenCHAMI/magellan/internal"
 	"github.com/OpenCHAMI/magellan/internal/db/sqlite"
@@ -26,41 +28,56 @@ var scanCmd = &cobra.Command{
 	Use:   "scan",
 	Short: "Scan for BMC nodes on a network",
 	Run: func(cmd *cobra.Command, args []string) {
-		fmt.Printf("subnets in cmd: %v\n", subnets)
-		// set hosts to use for scanning
-		hostsToScan := []string{}
+		var (
+			hostsToScan []string
+			portsToScan []int
+		)
+
+		// start by adding `--host` supplied to scan
 		if len(hosts) > 0 {
 			hostsToScan = hosts
-		} else {
-			for i, subnet := range subnets {
-				if len(subnet) <= 0 {
-					return
-				}
-
-				if len(subnetMasks) < i+1 {
-					subnetMasks = append(subnetMasks, net.IP{255, 255, 255, 0})
-				}
-
-				hostsToScan = append(hostsToScan, magellan.GenerateHosts(subnet, &subnetMasks[i])...)
-			}
 		}
 
-		// set ports to use for scanning
-		portsToScan := []int{}
+		// add hosts from `--subnets` and `--subnet-mask`
+		for i, subnet := range subnets {
+			// subnet string is empty so nothing to do here
+			if subnet == "" {
+				continue
+			}
+
+			// NOTE: should we check if subnet is valid here or is it done elsewhere (maybe in GenerateHosts)?
+
+			// no subnet masks supplied so add a default one for class C private networks
+			if len(subnetMasks) < i+1 {
+				subnetMasks = append(subnetMasks, net.IP{255, 255, 255, 0})
+			}
+
+			// generate a slice of all hosts to scan from subnets
+			hostsToScan = append(hostsToScan, magellan.GenerateHosts(subnet, &subnetMasks[i])...)
+		}
+
+		// add ports to use for scanning
 		if len(ports) > 0 {
 			portsToScan = ports
 		} else {
-			portsToScan = append(magellan.GetDefaultPorts(), ports...)
+			// no ports supplied so only use defaults
+			portsToScan = magellan.GetDefaultPorts()
 		}
 
-		// scan and store probe data in dbPath
+		// scan and store scanned data in cache
 		if concurrency <= 0 {
 			concurrency = mathutil.Clamp(len(hostsToScan), 1, 255)
 		}
 		probeStates := magellan.ScanForAssets(hostsToScan, portsToScan, concurrency, timeout, disableProbing, verbose)
 		if verbose {
-			for _, r := range probeStates {
-				fmt.Printf("%s:%d (%s)\n", r.Host, r.Port, r.Protocol)
+			format = strings.ToLower(format)
+			if format == "json" {
+				b, _ := json.Marshal(probeStates)
+				fmt.Printf("%s\n", string(b))
+			} else {
+				for _, r := range probeStates {
+					fmt.Printf("%s:%d (%s)\n", r.Host, r.Port, r.Protocol)
+				}
 			}
 		}
 
@@ -77,10 +94,9 @@ 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().StringVar(&format, "format", "", "set the output format")
 	scanCmd.Flags().StringSliceVar(&subnets, "subnet", []string{}, "set additional subnets")
-	scanCmd.Flags().IPSliceVar(&subnetMasks, "subnet-mask", []net.IP{}, "set the subnet masks to use for network")
+	scanCmd.Flags().IPSliceVar(&subnetMasks, "subnet-mask", []net.IP{}, "set the subnet masks to use for network (must match number of subnets)")
 	scanCmd.Flags().BoolVar(&disableProbing, "disable-probing", false, "disable probing scanned results for BMC nodes")
 
 	viper.BindPFlag("scan.hosts", scanCmd.Flags().Lookup("host"))