support assume role (#130)

* support assume role * fix config
similarweb · Aug 13, 2020 · 0f47880 · 0f47880
1 parent 5c2f18d
commit 0f47880
Show file tree

Hide file tree

Showing 7 changed files with 29 additions and 16 deletions.
diff --git a/collector/aws/session.go → collector/aws/auth.go b/collector/aws/session.go → collector/aws/auth.go
@@ -3,12 +3,13 @@ package aws
 import (
 	awsClient "github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/credentials"
+	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
 	"github.com/aws/aws-sdk-go/aws/session"
 	log "github.com/sirupsen/logrus"
 )
 
-// CreateNewSession return new AWS session
-func CreateNewSession(accessKey, secretKey, sessionToken, region string) *session.Session {
+// CreateAuthConfiguration return aws auth configuration
+func CreateAuthConfiguration(accessKey, secretKey, sessionToken, role, region string) (*session.Session, *awsClient.Config) {
 	var credentialsAWS *credentials.Credentials
 
 	// Use separate call for AWS credentials defined in config.yaml
@@ -18,11 +19,16 @@ func CreateNewSession(accessKey, secretKey, sessionToken, region string) *sessio
 		credentialsAWS = credentials.NewStaticCredentials(accessKey, secretKey, sessionToken)
 	}
 
-	sess := session.Must(session.NewSession(&awsClient.Config{
+	config := &awsClient.Config{
 		Region:      &region,
 		Credentials: credentialsAWS,
-	}))
+	}
 
-	return sess
+	sess := session.Must(session.NewSession(config))
 
+	if role != "" {
+		log.WithField("role", role).Info("assume role provided")
+		config.Credentials = stscreds.NewCredentials(sess, role, func(p *stscreds.AssumeRoleProvider) {})
+	}
+	return sess, config
 }
diff --git a/collector/aws/common/structs.go b/collector/aws/common/structs.go
@@ -6,6 +6,7 @@ import (
 	"finala/collector/aws/pricing"
 	"finala/collector/config"
 
+	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/sts"
 )
@@ -25,7 +26,7 @@ type AWSManager interface {
 	GetCloudWatchClient() *cloudwatch.CloudwatchManager
 	GetPricingClient() *pricing.PricingManager
 	GetRegion() string
-	GetSession() *session.Session
+	GetSession() (*session.Session, *aws.Config)
 	GetAccountIdentity() *sts.GetCallerIdentityOutput
 	SetGlobal(resourceName collector.ResourceIdentifier)
 	IsGlobalSet(resourceName collector.ResourceIdentifier) bool

diff --git a/collector/aws/detector.go b/collector/aws/detector.go
@@ -7,6 +7,7 @@ import (
 	"finala/collector/config"
 	"fmt"
 
+	awsClient "github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/session"
 	awsCloudwatch "github.com/aws/aws-sdk-go/service/cloudwatch"
 	awsPricing "github.com/aws/aws-sdk-go/service/pricing"
@@ -20,7 +21,7 @@ type DetectorDescriptor interface {
 	GetCloudWatchClient() *cloudwatch.CloudwatchManager
 	GetPricingClient() *pricing.PricingManager
 	GetRegion() string
-	GetSession() *session.Session
+	GetSession() (*session.Session, *awsClient.Config)
 	GetAccountIdentity() *sts.GetCallerIdentityOutput
 }
 
@@ -35,6 +36,7 @@ type DetectorManager struct {
 	cloudWatchClient *cloudwatch.CloudwatchManager
 	pricing          *pricing.PricingManager
 	session          *session.Session
+	awsConfig        *awsClient.Config
 	accountIdentity  *sts.GetCallerIdentityOutput
 	region           string
 	global           map[string]struct{}
@@ -43,11 +45,11 @@ type DetectorManager struct {
 // NewDetectorManager create new instance of detector manager
 func NewDetectorManager(collector collector.CollectorDescriber, account config.AWSAccount, stsManager *STSManager, global map[string]struct{}, region string) *DetectorManager {
 
-	priceSession := CreateNewSession(account.AccessKey, account.SecretKey, account.SessionToken, defaultRegionPrice)
+	priceSession, _ := CreateAuthConfiguration(account.AccessKey, account.SecretKey, account.SessionToken, account.Role, defaultRegionPrice)
 	pricingManager := pricing.NewPricingManager(awsPricing.New(priceSession), defaultRegionPrice)
 
-	regionSession := CreateNewSession(account.AccessKey, account.SecretKey, account.SessionToken, region)
-	cloudWatchCLient := cloudwatch.NewCloudWatchManager(awsCloudwatch.New(regionSession))
+	regionSession, regionConfig := CreateAuthConfiguration(account.AccessKey, account.SecretKey, account.SessionToken, account.Role, region)
+	cloudWatchCLient := cloudwatch.NewCloudWatchManager(awsCloudwatch.New(regionSession, regionConfig))
 
 	callerIdentityOutput, _ := stsManager.client.GetCallerIdentity(&sts.GetCallerIdentityInput{})
 	return &DetectorManager{
@@ -56,6 +58,7 @@ func NewDetectorManager(collector collector.CollectorDescriber, account config.A
 		pricing:          pricingManager,
 		region:           region,
 		session:          regionSession,
+		awsConfig:        regionConfig,
 		accountIdentity:  callerIdentityOutput,
 		global:           global,
 	}
@@ -87,8 +90,8 @@ func (dm *DetectorManager) GetRegion() string {
 }
 
 // GetSession return the aws session
-func (dm *DetectorManager) GetSession() *session.Session {
-	return dm.session
+func (dm *DetectorManager) GetSession() (*session.Session, *awsClient.Config) {
+	return dm.session, dm.awsConfig
 }
 
 // GetAccountIdentity return the caller identity

diff --git a/collector/aws/run.go b/collector/aws/run.go
@@ -38,8 +38,8 @@ func (app *Analyze) All() {
 
 	for _, account := range app.awsAccounts {
 
-		globalsession := CreateNewSession(account.AccessKey, account.SecretKey, account.SessionToken, "")
-		stsManager := NewSTSManager(sts.New(globalsession))
+		globalsession, globalConfig := CreateAuthConfiguration(account.AccessKey, account.SecretKey, account.SessionToken, account.Role, "")
+		stsManager := NewSTSManager(sts.New(globalsession, globalConfig))
 
 		for _, region := range account.Regions {
 			resourcesDetection := NewDetectorManager(app.cl, account, stsManager, app.global, region)

diff --git a/collector/aws/testutils/detector.go b/collector/aws/testutils/detector.go
@@ -6,6 +6,7 @@ import (
 	"finala/collector/aws/pricing"
 	"fmt"
 
+	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/sts"
 )
@@ -57,8 +58,8 @@ func (dm *MockAWSManager) GetRegion() string {
 	return dm.region
 }
 
-func (dm *MockAWSManager) GetSession() *session.Session {
-	return dm.session
+func (dm *MockAWSManager) GetSession() (*session.Session, *aws.Config) {
+	return dm.session, &aws.Config{}
 }
 
 func (dm *MockAWSManager) GetAccountIdentity() *sts.GetCallerIdentityOutput {

diff --git a/collector/config/config.go b/collector/config/config.go
@@ -15,6 +15,7 @@ type AWSAccount struct {
 	Name         string   `yaml:"name"`
 	AccessKey    string   `yaml:"access_key"`
 	SecretKey    string   `yaml:"secret_key"`
+	Role         string   `yaml:"role"`
 	SessionToken string   `yaml:"session_token"`
 	Regions      []string `yaml:"regions"`
 }

diff --git a/configuration/collector.yaml b/configuration/collector.yaml
@@ -11,6 +11,7 @@ providers:
       - name: <account_name>
         # access_key: <access_key>
         # secret_key: <secret_key>
+        # role: 
         regions:
           - us-east-1
           - us-west-2