feat: support for multiple gimme-aws-creds profiles (#5)

cli/common.go

@@ -1,6 +1,7 @@
 package cli
 
 import (
+	"bufio"
 	"fmt"
 	"io"
 	"log"
@@ -21,6 +22,7 @@ type Axolotl struct {
 	autoGimmeAwsCreds  bool
 	defaultRegion      string
 	awsCredentialsFile *vault.CredentialsFile
+	profiles           map[string]string
 }
 
 func (a *Axolotl) AwsCredentialsFile() (*vault.CredentialsFile, error) {
@@ -35,7 +37,9 @@ func (a *Axolotl) AwsCredentialsFile() (*vault.CredentialsFile, error) {
 	return a.awsCredentialsFile, nil
 }
 
-func (a *Axolotl) MustGetProfileNames() []string {
+// MustGetAWSProfileNames returns a list of AWS profile names
+// based on the contents of the local credentials file
+func (a *Axolotl) MustGetAWSProfileNames() []string {
 	creds, err := a.AwsCredentialsFile()
 	if err != nil {
 		log.Fatalf("Error loading AWS credentials: %s", err.Error())
@@ -52,9 +56,35 @@ func (a *Axolotl) MustGetProfileNames() []string {
 	return profileNames
 }
 
+// MustGetGACProfileNames returns a list of gimme-aws-creds profile names
+func (a *Axolotl) MustGetGACProfileNames() []string {
+	// Verify .okta_aws_login_config exists
+	if _, err := os.Stat(filepath.Join(os.Getenv("HOME"), ".okta_aws_login_config")); os.IsNotExist(err) {
+		log.Fatalf(os.ExpandEnv("unable to locate .okta_aws_login_config in ${HOME}, please create it: %s\n\n\thttps://github.com/Nike-Inc/gimme-aws-creds#configuration"), err.Error())
+	}
+
+	// Parse .okta_aws_login_config for profile names
+	f, err := os.Open(filepath.Join(os.Getenv("HOME"), ".okta_aws_login_config"))
+	if err != nil {
+		log.Fatalf("unable to open .okta_aws_login_config: %s", err.Error())
+	}
+	defer f.Close()
+
+	profileNames := []string{}
+	scanner := bufio.NewScanner(f)
+	for scanner.Scan() {
+		line := scanner.Text()
+		if strings.HasPrefix(line, "[") && strings.HasSuffix(line, "]") {
+			profileNames = append(profileNames, strings.Trim(line, "[]"))
+		}
+	}
+
+	return profileNames
+}
+
 // AuthVerify checks if the user is authenticated and if not authenticates
 // with gimme-aws-creds
-func AuthVerify(enabled bool, profileName string) error {
+func AuthVerify(enabled bool, profile Profile) error {
 	if !enabled {
 		return nil
 	}
@@ -78,7 +108,7 @@ func AuthVerify(enabled bool, profileName string) error {
 			os.Setenv(strings.Split(e, "=")[0], strings.Split(e, "=")[1])
 		}
 	}
-	os.Setenv("AWS_PROFILE", profileName)
+	os.Setenv("AWS_PROFILE", profile.AWS)
 
 	// Check if we are authenticated by running aws sts get-caller-identity
 	// If we are not authenticated, we will get an error
@@ -98,11 +128,11 @@ func AuthVerify(enabled bool, profileName string) error {
 	}
 
 	// If we are not authenticated, we will run gimme-aws-creds
-	return AuthGimmeAwsCreds()
+	return AuthGimmeAwsCreds(profile.GimmeAWSCreds)
 }
 
 // AuthGimmeAwsCreds authenticates with gimme-aws-creds
-func AuthGimmeAwsCreds() error {
+func AuthGimmeAwsCreds(gacProfile string) error {
 	// Check if gimme-aws-creds is installed
 	if _, err := exec.LookPath("gimme-aws-creds"); err != nil {
 		return fmt.Errorf("unable to locate `gimme-aws-creds` in PATH, please install it: %w\n\n\thttps://github.com/Nike-Inc/gimme-aws-creds#installation", err)
@@ -114,7 +144,7 @@ func AuthGimmeAwsCreds() error {
 	}
 
 	// execute gimme-aws-creds
-	cmd := exec.Command("gimme-aws-creds")
+	cmd := exec.Command("gimme-aws-creds", "--profile", gacProfile)
 	cmd.Stdout = os.Stdout
 	cmd.Stderr = os.Stderr
 	cmd.Stdin = os.Stdin
@@ -134,10 +164,12 @@ func ConfigureGlobals(app *kingpin.Application) *Axolotl {
 
 	viper.GetBool("autoGimmeAwsCreds")
 	viper.GetString("defaultRegion")
+	viper.GetStringMapString("profiles")
 
 	a := &Axolotl{
 		autoGimmeAwsCreds: viper.GetBool("autoGimmeAwsCreds"),
 		defaultRegion:     viper.GetString("defaultRegion"),
+		profiles:          viper.GetStringMapString("profiles"),
 	}
 
 	var (
@@ -239,11 +271,22 @@ func restoreTermState() {
 	}
 }
 
-// profileCompleter returns a list of profile names
-func (a *Axolotl) profileCompleter() func(d prompt.Document) []prompt.Suggest {
+// awsProfileCompleter returns a list of AWS profile names
+func (a *Axolotl) awsProfileCompleter() func(d prompt.Document) []prompt.Suggest {
+	return func(d prompt.Document) []prompt.Suggest {
+		s := []prompt.Suggest{}
+		for _, p := range a.MustGetAWSProfileNames() {
+			s = append(s, prompt.Suggest{Text: p})
+		}
+		return prompt.FilterHasPrefix(s, d.GetWordBeforeCursor(), true)
+	}
+}
+
+// gacProfileCompleter returns a list of gimme-aws-creds profile names
+func (a *Axolotl) gacProfileCompleter() func(d prompt.Document) []prompt.Suggest {
 	return func(d prompt.Document) []prompt.Suggest {
 		s := []prompt.Suggest{}
-		for _, p := range a.MustGetProfileNames() {
+		for _, p := range a.MustGetGACProfileNames() {
 			s = append(s, prompt.Suggest{Text: p})
 		}
 		return prompt.FilterHasPrefix(s, d.GetWordBeforeCursor(), true)

cli/exec.go

@@ -2,21 +2,29 @@ package cli
 
 import (
 	"fmt"
+	"log"
 	"os"
 	"syscall"
 
 	"github.com/alecthomas/kingpin"
 	"github.com/c-bata/go-prompt"
+	"github.com/spf13/viper"
 	osexec "golang.org/x/sys/execabs"
 )
 
+// Profile represents a profile mapping between AWS and gimme-aws-creds profiles
+type Profile struct {
+	AWS           string // AWS profile name
+	GimmeAWSCreds string // gimme-aws-creds profile name
+}
+
 type ExecCommandInput struct {
-	ProfileName string
-	Region      string
-	Command     string
-	Args        []string
-	Verify      bool
-	AutoRegion  bool
+	Profile    Profile
+	Region     string
+	Command    string
+	Args       []string
+	Verify     bool
+	AutoRegion bool
 }
 
 func ConfigureExecCommand(app *kingpin.Application, a *Axolotl) {
@@ -26,8 +34,8 @@ func ConfigureExecCommand(app *kingpin.Application, a *Axolotl) {
 
 	app.Flag("profile", "The AWS profile to execute as").
 		Short('p').
-		HintAction(a.MustGetProfileNames).
-		StringVar(&input.ProfileName)
+		HintAction(a.MustGetAWSProfileNames).
+		StringVar(&input.Profile.AWS)
 
 	app.Flag("region", "The AWS region to execute to").
 		Default(a.defaultRegion).
@@ -47,26 +55,47 @@ func ConfigureExecCommand(app *kingpin.Application, a *Axolotl) {
 			return fmt.Errorf("ax sessions should be nested with care, unset AWS_AXOLOTL to force")
 		}
 
-		if input.ProfileName == "" {
+		if input.Profile.AWS == "" {
 			saveTermState()
-			fmt.Println("Please select profile.")
-			input.ProfileName = prompt.Input("> ", a.profileCompleter())
+			fmt.Println("Please select AWS profile.")
+			input.Profile.AWS = prompt.Input("> ", a.awsProfileCompleter())
 			restoreTermState()
 		}
 
+		var ok bool
+		input.Profile.GimmeAWSCreds, ok = a.profiles[input.Profile.AWS]
+		if !ok {
+			gacProfiles := a.MustGetGACProfileNames()
+			if len(gacProfiles) == 1 {
+				input.Profile.GimmeAWSCreds = gacProfiles[0]
+			} else {
+				saveTermState()
+				fmt.Println("Please select gimme-aws-creds profile.")
+				input.Profile.GimmeAWSCreds = prompt.Input("> ", a.gacProfileCompleter())
+				restoreTermState()
+			}
+
+			// save the mapping for next time
+			a.profiles[input.Profile.AWS] = input.Profile.GimmeAWSCreds
+			viper.Set("profiles", a.profiles)
+			if err := viper.WriteConfig(); err != nil {
+				log.Fatalf("error writing config file: %s", err.Error())
+			}
+		}
+
 		return ExecCommand(input)
 	})
 }
 
 func ExecCommand(input ExecCommandInput) error {
 	env := environ(os.Environ())
-	env.Set("AWS_DEFAULT_PROFILE", input.ProfileName)
-	env.Set("AWS_PROFILE", input.ProfileName)
+	env.Set("AWS_DEFAULT_PROFILE", input.Profile.AWS)
+	env.Set("AWS_PROFILE", input.Profile.AWS)
 	env.Set("AWS_DEFAULT_REGION", input.Region)
 	env.Set("AWS_REGION", input.Region)
 	env.Set("AWS_AXOLOTL", "42")
 
-	if err := AuthVerify(input.Verify, input.ProfileName); err != nil {
+	if err := AuthVerify(input.Verify, input.Profile); err != nil {
 		return err
 	}