From 1c3ab5a8596cd412ff8da15976243c4f774b6872 Mon Sep 17 00:00:00 2001
From: Nathan Beneke <14852919+nathan-beneke@users.noreply.github.com>
Date: Mon, 15 Mar 2021 10:24:19 -0500
Subject: [PATCH] Use `include_profile` to inherit fields from the source
 profile (#3)

* Use `include_profile` to make tool MFA agnostic

* Update README.md with include_profile info

* Only start session after completing command parsing

* Use `ne` operator

* Correct flag information
---
 README.md                 | 12 +++-----
 aws-cfg-generator/main.go | 59 +++++++++++++++++++++++++++------------
 2 files changed, 45 insertions(+), 26 deletions(-)

diff --git a/README.md b/README.md
index e03eea2..192e4b7 100644
--- a/README.md
+++ b/README.md
@@ -26,25 +26,22 @@ should look something like this:
 
 ```
 [profile 123456789098]
-region=eu-central-1
 role_arn=arn:aws:iam::123456789098:role/role-name
 source_profile=default
-mfa_serial=arn:aws:iam::111111111111:mfa/user-name@example.com
+include_profile=default
 
 [profile 098765432123]
-region=eu-central-1
 role_arn=arn:aws:iam::098765432123:role/role-name-two
 source_profile=default
-mfa_serial=arn:aws:iam::111111111111:mfa/user-name@example.com
+include_profile=default
 
-. . .
+[. . .]
 ```
 
 #### Flags
 
 - `source-profile` Can be used to specify a source profile other than `default`
-- `region` Can be used to specify a region other than `eu-central-1`
-
+- `region` Can be used to specify a region other than the region in your source profile
 
 ### aws-extend-switch-roles
 
@@ -74,7 +71,6 @@ color = 00ff7f
 - No human-readable profile names
 - Only recognizes policies that are attached to groups
 - Can only recognize explicit permissions (i.e. it doesn't work when the `Resource` is not a role ARN)
-- Assumes there's an MFA
   
 ## Planned features
 
diff --git a/aws-cfg-generator/main.go b/aws-cfg-generator/main.go
index 8eb3152..ff02efd 100644
--- a/aws-cfg-generator/main.go
+++ b/aws-cfg-generator/main.go
@@ -14,10 +14,12 @@ package main
 */
 
 import (
+	"bytes"
 	"encoding/json"
 	"fmt"
 	"net/url"
 	"strings"
+	"text/template"
 
 	"github.com/alecthomas/kong"
 	"github.com/aws/aws-sdk-go/aws/arn"
@@ -30,11 +32,18 @@ const (
 	assumeAction = "sts:AssumeRole"
 )
 
-var vaultTemplate = `[profile %s]
-region=%s
-role_arn=%s
-source_profile=%s
-mfa_serial=arn:aws:iam::%s:mfa/%s
+type VaultModel struct {
+	ProfileName   string
+	Region        string
+	SourceProfile string
+	RoleArn       string
+}
+
+const vaultTemplate = `[profile {{.ProfileName}}]
+{{if (ne .Region "") }}region={{.Region}}
+{{end}}role_arn={{.RoleArn}}
+source_profile={{.SourceProfile}}
+input_profile={{.SourceProfile}}
 
 `
 
@@ -53,7 +62,7 @@ type CLI struct {
 
 type VaultCmd struct {
 	SourceProfile string `help:"The profile that your credentials should come from" default:"default"`
-	Region        string `help:"The AWS region each profile should set as default" default:"eu-central-1"`
+	Region        string `help:"Override the region configured with your source profile"`
 }
 
 type SwitchRolesCmd struct {
@@ -87,7 +96,7 @@ func generateSwitchRolesProfile(role, color string) {
 	fmt.Printf(switchRolesTemplate, roleArn.AccountID, roleArn.AccountID, roleSplit[1], color)
 }
 
-func generateVaultProfile(role, region, sourceProfile, rootAccount, user string) {
+func generateVaultProfile(role, region, sourceProfile string) {
 	// skip creating this profile if the role isn't a valid ARN (e.g. `*`)
 	if !arn.IsARN(role) {
 		return
@@ -95,7 +104,20 @@ func generateVaultProfile(role, region, sourceProfile, rootAccount, user string)
 
 	roleArn, _ := arn.Parse(role)
 
-	fmt.Printf(vaultTemplate, roleArn.AccountID, region, role, sourceProfile, rootAccount, user)
+	t := template.Must(template.New("vaultText").Parse(vaultTemplate))
+
+	var b bytes.Buffer
+	err := t.Execute(&b, VaultModel{
+		ProfileName:   roleArn.AccountID,
+		Region:        region,
+		SourceProfile: sourceProfile,
+		RoleArn:       role,
+	})
+	if err != nil {
+		panic(err)
+	}
+
+	fmt.Print(b.String())
 }
 
 func checkAction(action interface{}) bool {
@@ -210,21 +232,12 @@ func main() {
 	var cli CLI
 	ctx := kong.Parse(&cli)
 
-	sess := session.Must(session.NewSession())
-	stsClient := sts.New(sess)
-
-	gcio, err := stsClient.GetCallerIdentity(&sts.GetCallerIdentityInput{})
-	if err != nil {
-		panic(err)
-	}
-
-	user := getUser(gcio.Arn)
 	var generatorFunc func(role string)
 
 	switch ctx.Command() {
 	case "vault":
 		generatorFunc = func(role string) {
-			generateVaultProfile(role, cli.Vault.Region, cli.Vault.SourceProfile, *gcio.Account, *user)
+			generateVaultProfile(role, cli.Vault.Region, cli.Vault.SourceProfile)
 		}
 	case "switch-roles":
 		generatorFunc = func(role string) {
@@ -234,6 +247,16 @@ func main() {
 		panic(fmt.Errorf("unsupported command '%s'", ctx.Command()))
 	}
 
+	sess := session.Must(session.NewSession())
+	stsClient := sts.New(sess)
+
+	gcio, err := stsClient.GetCallerIdentity(&sts.GetCallerIdentityInput{})
+	if err != nil {
+		panic(err)
+	}
+
+	user := getUser(gcio.Arn)
+
 	iamClient := iam.New(sess)
 
 	cfgCreator := &configCreator{