Merge branch 'region' (pr remind101#30)

main.go

@@ -71,6 +71,7 @@ func main() {
 
 	// Load credentials from configFilePath if it exists, else use regular AWS config
 	var creds *credentials.Value
+	var region *string
 	var err error
 	if roleArnRe.MatchString(role) {
 		creds, err = assumeRole(role, "", *duration)
@@ -88,31 +89,31 @@ func main() {
 
 		creds, err = assumeRole(roleConfig.Role, roleConfig.MFA, *duration)
 	} else {
-		creds, err = assumeProfile(role)
+		creds, region, err = assumeProfile(role)
 	}
 
 	must(err)
 
 	if len(args) == 0 {
 		switch *format {
 		case "powershell":
-			printPowerShellCredentials(role, creds)
+			printPowerShellCredentials(role, creds, region)
 		case "bash":
-			printCredentials(role, creds)
+			printCredentials(role, creds, region)
 		case "fish":
-			printFishCredentials(role, creds)
+			printFishCredentials(role, creds, region)
 		default:
 			flag.Usage()
 			os.Exit(1)
 		}
 		return
 	}
 
-	err = execWithCredentials(role, args, creds)
+	err = execWithCredentials(role, args, creds, region)
 	must(err)
 }
 
-func execWithCredentials(role string, argv []string, creds *credentials.Value) error {
+func execWithCredentials(role string, argv []string, creds *credentials.Value, region *string) error {
 	argv0, err := exec.LookPath(argv[0])
 	if err != nil {
 		return err
@@ -123,62 +124,79 @@ func execWithCredentials(role string, argv []string, creds *credentials.Value) e
 	os.Setenv("AWS_SESSION_TOKEN", creds.SessionToken)
 	os.Setenv("AWS_SECURITY_TOKEN", creds.SessionToken)
 	os.Setenv("ASSUMED_ROLE", role)
+	if region != nil {
+		os.Setenv("AWS_DEFAULT_REGION", *region)
+		os.Setenv("AWS_REGION", *region)
+	}
 
 	env := os.Environ()
 	return syscall.Exec(argv0, argv, env)
 }
 
 // printCredentials prints the credentials in a way that can easily be sourced
 // with bash.
-func printCredentials(role string, creds *credentials.Value) {
+func printCredentials(role string, creds *credentials.Value, region *string) {
 	fmt.Printf("export AWS_ACCESS_KEY_ID=\"%s\"\n", creds.AccessKeyID)
 	fmt.Printf("export AWS_SECRET_ACCESS_KEY=\"%s\"\n", creds.SecretAccessKey)
 	fmt.Printf("export AWS_SESSION_TOKEN=\"%s\"\n", creds.SessionToken)
 	fmt.Printf("export AWS_SECURITY_TOKEN=\"%s\"\n", creds.SessionToken)
 	fmt.Printf("export ASSUMED_ROLE=\"%s\"\n", role)
+	if region != nil {
+		fmt.Printf("export AWS_DEFAULT_REGION=\"%s\"\n", *region)
+		fmt.Printf("export AWS_REGION=\"%s\"\n", *region)
+	}
 	fmt.Printf("# Run this to configure your shell:\n")
 	fmt.Printf("# eval $(%s)\n", strings.Join(os.Args, " "))
 }
 
 // printFishCredentials prints the credentials in a way that can easily be sourced
 // with fish.
-func printFishCredentials(role string, creds *credentials.Value) {
+func printFishCredentials(role string, creds *credentials.Value, region *string) {
 	fmt.Printf("set -gx AWS_ACCESS_KEY_ID \"%s\";\n", creds.AccessKeyID)
 	fmt.Printf("set -gx AWS_SECRET_ACCESS_KEY \"%s\";\n", creds.SecretAccessKey)
 	fmt.Printf("set -gx AWS_SESSION_TOKEN \"%s\";\n", creds.SessionToken)
 	fmt.Printf("set -gx AWS_SECURITY_TOKEN \"%s\";\n", creds.SessionToken)
 	fmt.Printf("set -gx ASSUMED_ROLE \"%s\";\n", role)
+	if region != nil {
+		fmt.Printf("set -gx AWS_DEFAULT_REGION \"%s\";\n", *region)
+		fmt.Printf("set -gx AWS_REGION \"%s\";\n", *region)
+	}
 	fmt.Printf("# Run this to configure your shell:\n")
 	fmt.Printf("# eval (%s)\n", strings.Join(os.Args, " "))
 }
 
 // printPowerShellCredentials prints the credentials in a way that can easily be sourced
 // with Windows powershell using Invoke-Expression.
-func printPowerShellCredentials(role string, creds *credentials.Value) {
+func printPowerShellCredentials(role string, creds *credentials.Value, region *string) {
 	fmt.Printf("$env:AWS_ACCESS_KEY_ID=\"%s\"\n", creds.AccessKeyID)
 	fmt.Printf("$env:AWS_SECRET_ACCESS_KEY=\"%s\"\n", creds.SecretAccessKey)
 	fmt.Printf("$env:AWS_SESSION_TOKEN=\"%s\"\n", creds.SessionToken)
 	fmt.Printf("$env:AWS_SECURITY_TOKEN=\"%s\"\n", creds.SessionToken)
 	fmt.Printf("$env:ASSUMED_ROLE=\"%s\"\n", role)
+	if region != nil {
+		fmt.Printf("$env:AWS_DEFAULT_REGION=\"%s\"\n", *region)
+		fmt.Printf("$env:AWS_REGION=\"%s\"\n", *region)
+	}
 	fmt.Printf("# Run this to configure your shell:\n")
 	fmt.Printf("# %s | Invoke-Expression \n", strings.Join(os.Args, " "))
 }
 
 // assumeProfile assumes the named profile which must exist in ~/.aws/config
 // (https://docs.aws.amazon.com/cli/latest/userguide/cli-roles.html) and returns the temporary STS
 // credentials.
-func assumeProfile(profile string) (*credentials.Value, error) {
+func assumeProfile(profile string) (*credentials.Value, *string, error) {
 	sess := session.Must(session.NewSessionWithOptions(session.Options{
 		Profile:                 profile,
 		SharedConfigState:       session.SharedConfigEnable,
 		AssumeRoleTokenProvider: readTokenCode,
 	}))
 
+	region := sess.Config.Region
 	creds, err := sess.Config.Credentials.Get()
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
-	return &creds, nil
+	return &creds, region, nil
 }
 
 // assumeRole assumes the given role and returns the temporary STS credentials.