reworked the role assumption, works better now for non-role assumptio…

…n usage (#17)
capitalone · Apr 3, 2017 · 95db0a6 · 95db0a6
1 parent fca49d3
commit 95db0a6
Show file tree

Hide file tree

Showing 5 changed files with 79 additions and 27 deletions.
diff --git a/cmd/images.go b/cmd/images.go
@@ -76,8 +76,8 @@ func init() {
 		viper.SetDefault("ImageFilterRegex", map[string]string{
 			"name": `^COF-[0-9A-Za-z_.\-]+-x64-HVM-Enc-[0-9A-Za-z_.\-]+`,
 		})
-
 	}
+
 	viper.SetDefault("VersionRegex", `-([\d]+(-[\d]+)*)$`)
 	viper.SetDefault("PlatformRegex", `^[0-9A-Za-z_.]+-([0-9A-Za-z_.]+)`)
 

diff --git a/cmd/root.go b/cmd/root.go
@@ -24,9 +24,10 @@ import (
 )
 
 const (
-	projectName = "stack-deployment-tool"
-	dryModeFlag = "drymode"
-	debugFlag   = "debug"
+	projectName    = "stack-deployment-tool"
+	dryModeFlag    = "drymode"
+	debugFlag      = "debug"
+	assumeRoleFlag = "assume-role"
 )
 
 var (
@@ -76,6 +77,10 @@ func init() {
 	RootCmd.PersistentFlags().BoolP(dryModeFlag, "q", false, "enable dry mode")
 	viper.BindPFlag(dryModeFlag, RootCmd.PersistentFlags().Lookup(dryModeFlag))
 
+	RootCmd.PersistentFlags().Bool(assumeRoleFlag, false,
+		"force assume role, if AWS_ROLE_ARN is not provided it will be guessed based on AWS_PROFILE and ~/.aws/config")
+	viper.BindPFlag(assumeRoleFlag, RootCmd.PersistentFlags().Lookup(assumeRoleFlag))
+
 	RootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is $HOME/."+projectName+".yaml)")
 	// local flags
 	RootCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle")

diff --git a/images/ami.go b/images/ami.go
@@ -102,6 +102,9 @@ func (a *AmiFinder) describeImages(params *ec2.DescribeImagesInput) (*ec2.Descri
 		resp, err = a.DescribeImages(params)
 	}
 	log.Debugf("DescribeImages Err: %#v Response: %#v", err, resp)
+	if err != nil {
+		log.Errorf("Error finding image: %v", err)
+	}
 	return resp, err
 }
 

diff --git a/images/ami_template_helper.go b/images/ami_template_helper.go
@@ -15,6 +15,8 @@
 package images
 
 import (
+	"sync"
+
 	"github.com/capitalone/stack-deployment-tool/stacks"
 
 	"github.com/aymerick/raymond"
@@ -26,8 +28,29 @@ const (
 )
 
 func init() {
-	stacks.RegisterTemplateHelper(amiTemplCmd, NewAmiFinder(), amiHelper)
-	stacks.RegisterTemplateHelper(latestAmiTemplCmd, NewAmiFinder(), latestAmiHelper)
+	lazyImageFinder := &LazyImageFinder{}
+	stacks.RegisterTemplateHelper(amiTemplCmd, lazyImageFinder, amiHelper)
+	stacks.RegisterTemplateHelper(latestAmiTemplCmd, lazyImageFinder, latestAmiHelper)
+}
+
+type LazyImageFinder struct {
+	finderInit sync.Once
+	finderInst *AmiFinder
+}
+
+func (l *LazyImageFinder) FindImageId(platform, version string) string {
+	return l.finder().FindImageId(platform, version)
+}
+
+func (l *LazyImageFinder) FindLatestImageId(platform string) string {
+	return l.finder().FindLatestImageId(platform)
+}
+
+func (l *LazyImageFinder) finder() *AmiFinder {
+	l.finderInit.Do(func() {
+		l.finderInst = NewAmiFinder()
+	})
+	return l.finderInst
 }
 
 func amiHelper(options *raymond.Options) raymond.SafeString {

diff --git a/providers/aws.go b/providers/aws.go
@@ -26,6 +26,7 @@ import (
 
 	"github.com/capitalone/stack-deployment-tool/sdt"
 	"github.com/capitalone/stack-deployment-tool/utils"
+	"github.com/spf13/viper"
 
 	log "github.com/Sirupsen/logrus"
 	"github.com/aws/aws-sdk-go/aws"
@@ -43,6 +44,10 @@ const (
 	retry_max = 5
 )
 
+var (
+	roleArn *string
+)
+
 type AWSApi struct {
 	Session *session.Session
 	cfsrvc  *cloudformation.CloudFormation
@@ -63,26 +68,38 @@ func sessionName() string {
 	return fmt.Sprintf("%s-%d", utils.GetenvWithDefault("USER", ""), time.Now().UTC().Unix())
 }
 
+func EnvRoleArn() string {
+	return utils.GetenvWithDefault("AWS_ROLE_ARN", "")
+}
+
+func EnvProfile() string {
+	return utils.GetenvWithDefault("AWS_PROFILE", "")
+}
+
 func RoleArn() string {
-	role := utils.GetenvWithDefault("AWS_ROLE_ARN", "")
-	if len(role) == 0 { // maybe try $HOME/.aws/config
-		profile := utils.GetenvWithDefault("AWS_PROFILE", "")
-		configFile := filepath.Join(os.Getenv("HOME"), ".aws", "config")
-		if len(profile) > 0 && utils.FileExists(configFile) {
-			cfg, err := ini.Load(configFile)
-			if err != nil {
-				log.Debugf("Error loading %s: %+v", configFile, err)
-			} else {
-				role = cfg.Section("profile " + profile).Key("saml_role").Value()
+	if roleArn == nil {
+		role := EnvRoleArn()
+		if len(role) == 0 { // maybe try $HOME/.aws/config
+			profile := EnvProfile()
+			configFile := filepath.Join(os.Getenv("HOME"), ".aws", "config")
+			if len(profile) > 0 && utils.FileExists(configFile) {
+				cfg, err := ini.Load(configFile)
+				if err != nil {
+					log.Debugf("Error loading %s: %+v", configFile, err)
+				} else {
+					role = cfg.Section("profile " + profile).Key("saml_role").Value()
+				}
 			}
 		}
+		if len(role) == 0 { // check if it is still empty..
+			log.Debugf("AWS_ROLE_ARN is empty")
+		} else {
+			log.Debugf("Found AWS ROLE: %s", role)
+		}
+		roleArn = &role
 	}
-	if len(role) == 0 { // check if it is still empty..
-		log.Infof("AWS_ROLE_ARN is empty")
-	} else {
-		log.Infof("Using AWS ROLE: %s", role)
-	}
-	return role
+
+	return *roleArn
 }
 
 func Region() string {
@@ -104,11 +121,15 @@ func createSession(httpClient *http.Client) *session.Session {
 
 	sess := session.New(config)
 
-	role := RoleArn()
-	if len(role) > 0 {
-		sess.Config.Credentials = stscreds.NewCredentials(sess,
-			role,
-			func(p *stscreds.AssumeRoleProvider) { p.RoleSessionName = sessionName() })
+	forceAssume := viper.GetBool("assume-role")
+	if len(EnvRoleArn()) == 0 && forceAssume {
+		role := RoleArn()
+		if len(role) > 0 {
+			log.Infof("Assuming Role: %s", role)
+			sess.Config.Credentials = stscreds.NewCredentials(sess,
+				role,
+				func(p *stscreds.AssumeRoleProvider) { p.RoleSessionName = sessionName() })
+		}
 	}
 
 	sess.Handlers.Build.PushFrontNamed(addNameAndVersionToUserAgent)