Take the first valid known hosts domain name in generate command

cmd/generate.go

@@ -7,9 +7,11 @@ import (
 	"regexp"
 	"strings"
 
+	valid "github.com/asaskevich/govalidator"
 	"github.com/mitchellh/go-homedir"
 	"github.com/quantumsheep/sshconfig"
 	"github.com/spf13/cobra"
+	"github.com/spf13/pflag"
 	"github.com/spf13/viper"
 )
 
@@ -24,6 +26,7 @@ func init() {
 	flags := generateCmd.Flags()
 	flags.Bool("known-hosts", false, "Generate from known_hosts file")
 	flags.String("known-hosts-file", "~/.ssh/known_hosts", "Path of known_hosts file")
+	flags.Bool("known-hosts-allow-single-ip", false, "Allow single IP addresses (without hostname)")
 
 	viper.SetDefault("author", "quantumsheep <[email protected]>")
 	viper.SetDefault("license", "MIT")
@@ -32,10 +35,9 @@ func init() {
 func runGenerate(cmd *cobra.Command, args []string) {
 	flags := cmd.Flags()
 
-	knownHosts := false
-
-	if enabled, e := flags.GetBool("known-hosts"); e == nil {
-		knownHosts = enabled
+	knownHosts, e := flags.GetBool("known-hosts")
+	if e != nil {
+		log.Fatal(e)
 	}
 
 	if !knownHosts {
@@ -46,56 +48,88 @@ func runGenerate(cmd *cobra.Command, args []string) {
 	configs := make([]*KnownHostConfig, 0)
 
 	if knownHosts {
-		knownHostsFile := "~/.ssh/known_hosts"
+		configs = append(configs, generateFromKnownHosts(flags)...)
+	}
 
-		if str, e := flags.GetString("known-hosts-file"); e == nil && str != "" {
-			knownHostsFile = str
-		}
+	config := strings.Join(KnownHostConfigStrings(KnownHostConfigUniques(configs)), "\n\n")
+	fmt.Println(config)
+}
 
-		knownHostsFile, e := homedir.Expand(knownHostsFile)
-		if e != nil {
-			log.Fatal(e)
-		}
+func generateFromKnownHosts(flags *pflag.FlagSet) []*KnownHostConfig {
+	knownHostsFile := "~/.ssh/known_hosts"
 
-		// open file
-		bytes, e := os.ReadFile(knownHostsFile)
-		if e != nil {
-			log.Fatal(e)
+	if str, e := flags.GetString("known-hosts-file"); e == nil && str != "" {
+		knownHostsFile = str
+	}
+
+	knownHostsFile, e := homedir.Expand(knownHostsFile)
+	if e != nil {
+		log.Fatal(e)
+	}
+
+	// open file
+	bytes, e := os.ReadFile(knownHostsFile)
+	if e != nil {
+		log.Fatal(e)
+	}
+
+	data := string(bytes)
+	lines := strings.Split(data, "\n")
+
+	rx := regexp.MustCompile(`^(\[(?P<Host>.*?)\]:(?P<Port>\d+))|(?P<SingleHost>.*?)$`)
+
+	configs := make([]*KnownHostConfig, 0)
+
+	for _, line := range lines {
+		if line == "" {
+			continue
 		}
 
-		data := string(bytes)
-		rx := regexp.MustCompile(`^(\[(?P<Host>.*?)\]:(?P<Port>\d+))|(?P<SingleHost>.*?)$`)
+		lineConfigs := make([]*KnownHostConfig, 0)
 
-		lines := strings.Split(data, "\n")
-		for _, line := range lines {
-			if line == "" {
-				continue
-			}
+		targets := strings.Split(strings.Split(line, " ")[0], ",")
+		for _, target := range targets {
+			config := NewKnownHostConfig()
 
-			targets := strings.Split(strings.Split(line, " ")[0], ",")
-			for _, target := range targets {
-				config := NewKnownHostConfig()
+			matches := rx.FindStringSubmatch(target)
 
-				matches := rx.FindStringSubmatch(target)
+			if host := matches[rx.SubexpIndex("Host")]; host != "" {
+				port := matches[rx.SubexpIndex("Port")]
 
-				if host := matches[rx.SubexpIndex("Host")]; host != "" {
-					port := matches[rx.SubexpIndex("Port")]
+				config.Host = host + ":" + port
+				config.HostName = host
+				config.Port = port
+			} else if host := matches[rx.SubexpIndex("SingleHost")]; host != "" {
+				config.Host = host
+				config.HostName = host
+			}
 
-					config.Host = host + ":" + port
-					config.HostName = host
-					config.Port = port
-				} else if host := matches[rx.SubexpIndex("SingleHost")]; host != "" {
-					config.Host = host
-					config.HostName = host
-				}
+			lineConfigs = append(lineConfigs, config)
+		}
 
-				configs = append(configs, config)
+		allowSingleIp, e := flags.GetBool("known-hosts-allow-single-ip")
+		if e != nil {
+			log.Fatal(e)
+		}
+
+		var config *KnownHostConfig = nil
+
+		// Select the first config with a valid domain name (defaults to the first config)
+		for _, lineConfig := range lineConfigs {
+			if valid.IsDNSName(lineConfig.HostName) {
+				config = lineConfig
+				break
 			}
 		}
+
+		if config != nil {
+			configs = append(configs, config)
+		} else if allowSingleIp && len(lineConfigs) > 0 {
+			configs = append(configs, lineConfigs[0])
+		}
 	}
 
-	config := strings.Join(KnownHostConfigStrings(KnownHostConfigUniques(configs)), "\n\n")
-	fmt.Println(config)
+	return configs
 }
 
 type KnownHostConfig struct {

go.mod

@@ -9,6 +9,7 @@ require (
 	github.com/quantumsheep/sshconfig v1.1.0
 	github.com/rivo/tview v0.0.0-20220129131435-1f7581b67bd1
 	github.com/spf13/cobra v1.3.0
+	github.com/spf13/pflag v1.0.5
 	github.com/spf13/viper v1.10.1
 )
 
@@ -26,7 +27,6 @@ require (
 	github.com/spf13/afero v1.6.0 // indirect
 	github.com/spf13/cast v1.4.1 // indirect
 	github.com/spf13/jwalterweatherman v1.1.0 // indirect
-	github.com/spf13/pflag v1.0.5 // indirect
 	github.com/subosito/gotenv v1.2.0 // indirect
 	golang.org/x/sys v0.0.0-20211210111614-af8b64212486 // indirect
 	golang.org/x/term v0.0.0-20210220032956-6a3ed077a48d // indirect