Adding domain name suffix

Also comes with the ability to now add more config variables later

cmd/reconf.go

@@ -17,7 +17,7 @@ which is a filename. In case of any errors, the preexisting file won't be touche
 		initConfig()
 	},
 	Run: func(cmd *cobra.Command, args []string) {
-		lib.Reconf(viper.GetStringSlice("profiles"), args[0], viper.GetBool("no-profile-prefix"))
+		lib.Reconf(viper.Get("profilesConfig").([]lib.ProfileConfig), args[0], viper.GetBool("no-profile-prefix"))
 	},
 }
 

cmd/root.go

@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"aws-ssh/lib"
 	"fmt"
 	"os"
 
@@ -54,11 +55,20 @@ func initSettings() {
 
 // initConfig reads in config file and ENV variables if set.
 func initConfig() {
+	profiles, err := getProfiles()
+	if err != nil {
+		log.WithError(err).Fatal("Profiles have not been provided and couldn't retrieve them from the config")
+	}
 	if len(viper.GetStringSlice("profiles")) == 0 {
-		profiles, err := getProfiles()
-		if err != nil {
-			log.WithError(err).Fatal("Profiles have not been provided and couldn't retrieve them from the config")
+		viper.Set("profilesConfig", profiles)
+	} else {
+		specifiedProfiles := viper.GetStringSlice("profiles")
+		filteredProfiles := make([]lib.ProfileConfig, 0, len(specifiedProfiles))
+		for _, profile := range profiles {
+			if contains(specifiedProfiles, profile.Name) {
+				filteredProfiles = append(filteredProfiles, profile)
+			}
 		}
-		viper.Set("profiles", profiles)
+		viper.Set("profilesConfig", filteredProfiles)
 	}
 }

cmd/test.go

@@ -20,7 +20,8 @@ Allows to identify permission issues early.
 		initConfig()
 	},
 	Run: func(cmd *cobra.Command, args []string) {
-		summaries, err := lib.TraverseProfiles(viper.GetStringSlice("profiles"))
+		profiles := viper.Get("profilesConfig").([]lib.ProfileConfig)
+		summaries, err := lib.TraverseProfiles(profiles)
 		if err != nil {
 			log.WithError(err).Fatal("Can't traverse through all profiles")
 		} else {

cmd/util.go

@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"aws-ssh/lib"
 	"os"
 	"path"
 	"strings"
@@ -11,13 +12,13 @@ import (
 )
 
 // gets profiles. The current Go AWS SDK doesn't have this function, whereas python boto3 has it. Why?
-func getProfiles() ([]string, error) {
+func getProfiles() ([]lib.ProfileConfig, error) {
 	// use a map here to get rid of duplicates
-	var profiles = make(map[string]struct{})
+	var profiles = make(map[string]lib.ProfileConfig)
 
 	home, err := homedir.Dir()
 	if err != nil {
-		return []string{}, err
+		return []lib.ProfileConfig{}, err
 	}
 
 	loadFile := func(envVariableName, defaultFileName, sectionPrefix string) error {
@@ -30,13 +31,24 @@ func getProfiles() ([]string, error) {
 			return err
 		}
 
-		for _, section := range config.SectionStrings() {
+		for _, section := range config.Sections() {
+			name := section.Name()
 			// skip the default section (top level keys/values)
-			if section == ini.DefaultSection {
+			if name == ini.DefaultSection {
 				continue
 			}
-			if strings.HasPrefix(section, sectionPrefix) {
-				profiles[section[len(sectionPrefix):]] = struct{}{}
+			if strings.HasPrefix(name, sectionPrefix) {
+				name = name[len(sectionPrefix):]
+				if _, exists := profiles[name]; !exists {
+					config := lib.ProfileConfig{Name: name}
+					if section.HasKey("aws-ssh-domain") {
+						config.Domain = section.Key("aws-ssh-domain").Value()
+					}
+					log.Debugf("Got profile - %s", name)
+					profiles[name] = config
+				} else {
+					log.Debugf("Skipping duplicate profile - %s", name)
+				}
 			}
 		}
 		return nil
@@ -50,10 +62,19 @@ func getProfiles() ([]string, error) {
 	}
 
 	// convert the map to a list
-	var profilesList []string
+	profilesList := make([]lib.ProfileConfig, 0, len(profiles))
 
-	for profile := range profiles {
+	for _, profile := range profiles {
 		profilesList = append(profilesList, profile)
 	}
 	return profilesList, nil
 }
+
+func contains(slice []string, element string) bool {
+	for _, item := range slice {
+		if item == element {
+			return true
+		}
+	}
+	return false
+}

lib/aws.go

@@ -19,6 +19,7 @@ type ProfileSummary struct {
 	Name      string
 	Region    string
 	Instances []*ec2.Instance
+	Domain    string
 }
 
 func makeSession(profile string) (*session.Session, error) {
@@ -37,14 +38,14 @@ func makeSession(profile string) (*session.Session, error) {
 }
 
 // TraverseProfiles goes through all profiles and returns a list of ProfileSummary
-func TraverseProfiles(profiles []string) ([]ProfileSummary, error) {
+func TraverseProfiles(profiles []ProfileConfig) ([]ProfileSummary, error) {
 	log.Debugf("Traversing through %d profiles", len(profiles))
 	var profileSummaryChan = make(chan ProfileSummary, len(profiles))
 	var errChan = make(chan error, len(profiles))
 
 	var profileSummaries []ProfileSummary
 	for _, profile := range profiles {
-		go func(profile string) {
+		go func(profile ProfileConfig) {
 			DescribeProfile(profile, profileSummaryChan, errChan)
 		}(profile)
 	}
@@ -66,16 +67,17 @@ func TraverseProfiles(profiles []string) ([]ProfileSummary, error) {
 }
 
 // DescribeProfile describes the specified profile
-func DescribeProfile(profile string, sum chan ProfileSummary, errChan chan error) {
-	awsSession, err := makeSession(profile)
+func DescribeProfile(profile ProfileConfig, sum chan ProfileSummary, errChan chan error) {
+	awsSession, err := makeSession(profile.Name)
 	if err != nil {
-		errChan <- fmt.Errorf("Couldn't create session for '%s': %s", profile, err)
+		errChan <- fmt.Errorf("Couldn't create session for '%s': %s", profile.Name, err)
 		return
 	}
 
 	profileSummary := ProfileSummary{
-		Name:   profile,
+		Name:   profile.Name,
 		Region: aws.StringValue(awsSession.Config.Region),
+		Domain: profile.Domain,
 	}
 
 	svc := ec2.New(awsSession)

lib/reconf.go

@@ -14,6 +14,12 @@ import (
 	linq "gopkg.in/ahmetb/go-linq.v3"
 )
 
+// ProfileConfig represents an entry in aws config
+type ProfileConfig struct {
+	Name,
+	Domain string
+}
+
 // SSHEntry represents an entry in ssh config
 type SSHEntry struct {
 	Address,
@@ -22,14 +28,20 @@ type SSHEntry struct {
 	ProxyJump,
 	Port,
 	User,
-	Profile string
+	Profile,
+	Domain string
 }
 
 // ConfigFormat returns formatted and stringified SSHEntry ready to use in ssh config
 func (e SSHEntry) ConfigFormat() string {
-	var output = []string{
-		fmt.Sprintf("Host %s %s %s.%s", e.Name, e.InstanceID, e.Address, e.Profile),
+	var output = []string{}
+
+	if e.Domain != "" {
+		output = append(output, fmt.Sprintf("Host %s %s %s.%s %s.%s", e.Name, e.InstanceID, e.Address, e.Profile, e.Name, e.Domain))
+	} else {
+		output = append(output, fmt.Sprintf("Host %s %s %s.%s", e.Name, e.InstanceID, e.Address, e.Profile))
 	}
+
 	if e.User != "" {
 		output = append(output, fmt.Sprintf("    User %s", e.User))
 	}
@@ -55,7 +67,7 @@ func instanceNameSorter(i interface{}) interface{} { // sort by instance name
 }
 
 // Reconf writes ssh config with profiles into the specified file
-func Reconf(profiles []string, filename string, noProfilePrefix bool) {
+func Reconf(profiles []ProfileConfig, filename string, noProfilePrefix bool) {
 	profileSummaries, err := TraverseProfiles(profiles)
 	if err != nil {
 		log.WithError(err).Error("got some errors")
@@ -124,6 +136,7 @@ func Reconf(profiles []string, filename string, noProfilePrefix bool) {
 						InstanceID: aws.StringValue(instance.InstanceId),
 						Name:       getInstanceCanonicalName(summary.Name, instanceName, instanceIndex),
 						Profile:    summary.Name,
+						Domain:     summary.Domain,
 					}
 					if noProfilePrefix {
 						entry.Name = getInstanceCanonicalName("", instanceName, instanceIndex)