updates to support priv sessions (#41)

aws/guardduty.go

@@ -129,10 +129,10 @@ func enableGuardDutyAutoEnable(client *guardduty.GuardDuty, autoEnableS3Protecti
 // EnableGuardDutyAdministratorAccount enables the GuardDuty Administrator account within the AWS Organization
 func EnableGuardDutyAdministratorAccount(region string, administratorAccountRole string, rootRole string, autoEnableS3Protection bool) error {
 	rootSession := GetSession()
-	rootAccountID := GetAccountID(rootSession, rootRole)
+	rootAccountID := GetAccountIDWithRole(rootSession, rootRole)
 
 	adminAcctSession := GetSession()
-	adminAccountID := GetAccountID(adminAcctSession, administratorAccountRole)
+	adminAccountID := GetAccountIDWithRole(adminAcctSession, administratorAccountRole)
 
 	enabledRegions := GetEnabledRegions(region, rootRole, false)
 

aws/securityhub.go

@@ -19,8 +19,10 @@ package aws
 import (
 	"errors"
 	"fmt"
+	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/awserr"
 	"github.com/aws/aws-sdk-go/service/securityhub"
 	common "github.com/cloudposse/turf/common/error"
 	"github.com/sirupsen/logrus"
@@ -85,17 +87,33 @@ func (hub SecurityHub) addSecurityHubMemberAccounts(memberAccounts []string, adm
 	}
 }
 
+func (hub SecurityHub) disableControl(currentControl string) error {
+	_, err := hub.currentAccountClient.UpdateStandardsControl(&securityhub.UpdateStandardsControlInput{
+		ControlStatus:       aws.String("DISABLED"),
+		DisabledReason:      aws.String("Global Resources are not collected in this region"),
+		StandardsControlArn: aws.String(currentControl),
+	})
+
+	if err != nil {
+		if aerr, ok := err.(awserr.Error); ok {
+			if aerr.Code() == "TooManyRequestsException" {
+				logrus.Warnf("    received too many requests error. Sleeping then trying again while disabling control %s", currentControl)
+
+				time.Sleep(2 * time.Second)
+				return hub.disableControl(currentControl)
+			}
+		}
+	}
+
+	return err
+}
+
 func (hub SecurityHub) disableControls(region string, accountID string, controls []string) {
 	for i := range controls {
 		currentControl := fmt.Sprintf(controls[i], region, accountID)
 
 		logrus.Infof("    disabling control %s", currentControl)
-
-		_, err := hub.currentAccountClient.UpdateStandardsControl(&securityhub.UpdateStandardsControlInput{
-			ControlStatus:       aws.String("DISABLED"),
-			DisabledReason:      aws.String("Global Resources are not collected in this region"),
-			StandardsControlArn: aws.String(currentControl),
-		})
+		err := hub.disableControl(currentControl)
 
 		if err != nil {
 			logrus.Error(err)
@@ -149,7 +167,7 @@ func getCIS120Controls(isGlobalCollectionRegion bool, isCloudTrailAccount bool)
 		}...)
 	}
 
-	if !isCloudTrailAccount {
+	if !isCloudTrailAccount || (isCloudTrailAccount && !isGlobalCollectionRegion) {
 		controls = append(controls, []string{
 			"arn:aws:securityhub:%s:%s:control/cis-aws-foundations-benchmark/v/1.2.0/2.7",
 			"arn:aws:securityhub:%s:%s:control/cis-aws-foundations-benchmark/v/1.2.0/3.1",
@@ -207,10 +225,10 @@ func logSecurityHubMemberAccounts(memberAccounts []string) {
 // EnableSecurityHubAdministratorAccount enables the Security Hub Administrator account within the AWS Organization
 func EnableSecurityHubAdministratorAccount(region string, administratorAccountRole string, rootRole string) error {
 	rootSession := GetSession()
-	rootAccountID := GetAccountID(rootSession, rootRole)
+	rootAccountID := GetAccountIDWithRole(rootSession, rootRole)
 
 	adminAcctSession := GetSession()
-	adminAccountID := GetAccountID(adminAcctSession, administratorAccountRole)
+	adminAccountID := GetAccountIDWithRole(adminAcctSession, administratorAccountRole)
 
 	enabledRegions := GetEnabledRegions(region, rootRole, false)
 
@@ -269,8 +287,15 @@ func DisableSecurityHubGlobalResourceControls(globalCollectionRegion string, rol
 	}
 
 	session := GetSession()
-	accountID := GetAccountID(session, role)
-	enabledRegions := GetEnabledRegions("us-east-1", role, false)
+	var accountID string
+
+	if isPrivileged {
+		accountID = GetAccountID(session)
+	} else {
+		accountID = GetAccountIDWithRole(session, role)
+	}
+
+	enabledRegions := GetEnabledRegions("us-east-1", role, isPrivileged)
 
 	if !validateRegion(enabledRegions, globalCollectionRegion) {
 		return fmt.Errorf("%s is not a valid enabled region in this account", globalCollectionRegion)
@@ -284,16 +309,16 @@ func DisableSecurityHubGlobalResourceControls(globalCollectionRegion string, rol
 		var currentAccountClient *securityhub.SecurityHub
 
 		if isPrivileged {
-			currentAccountClient = getSecurityHubClientWithRole(currentRegion, role)
-		} else {
 			currentAccountClient = getSecurityHubClient(currentRegion)
+		} else {
+			currentAccountClient = getSecurityHubClientWithRole(currentRegion, role)
 		}
 
 		hub := SecurityHub{
 			currentAccountClient: currentAccountClient,
 		}
 
-		isGlobalCollectionRegion := currentRegion != globalCollectionRegion
+		isGlobalCollectionRegion := currentRegion == globalCollectionRegion
 
 		if isGlobalCollectionRegion {
 			logrus.Infof("  processing global collection region %s", currentRegion)

aws/sts.go

@@ -46,14 +46,20 @@ func GetCreds(sess *session.Session, role string) *credentials.Credentials {
 }
 
 // GetAccountID returns the AWS Account ID of the session
-func GetAccountID(sess *session.Session, role string) string {
-	var client *sts.STS
-	if role == "" {
-		client = getStsClient(sess)
-	} else {
-		creds := GetCreds(sess, role)
-		client = getStsClientWithCreds(sess, creds)
-	}
+func GetAccountID(sess *session.Session) string {
+	client := getStsClient(sess)
+
+	input := sts.GetCallerIdentityInput{}
+	ident, err := client.GetCallerIdentity(&input)
+
+	common.AssertErrorNil(err)
+	return *ident.Account
+}
+
+// GetAccountIDWithRole returns the AWS Account ID of the session after assuming a role
+func GetAccountIDWithRole(sess *session.Session, role string) string {
+	creds := GetCreds(sess, role)
+	client := getStsClientWithCreds(sess, creds)
 
 	input := sts.GetCallerIdentityInput{}
 	ident, err := client.GetCallerIdentity(&input)

go.mod

@@ -3,14 +3,13 @@ module github.com/cloudposse/turf
 go 1.15
 
 require (
-	github.com/aws/aws-sdk-go v1.38.18
+	github.com/aws/aws-sdk-go v1.38.21
 	github.com/fsnotify/fsnotify v1.4.9 // indirect
-	github.com/magefile/mage v1.11.0 // indirect
 	github.com/magiconair/properties v1.8.4 // indirect
 	github.com/mitchellh/go-homedir v1.1.0
 	github.com/mitchellh/mapstructure v1.4.1 // indirect
 	github.com/pelletier/go-toml v1.8.1 // indirect
-	github.com/sirupsen/logrus v1.8.0
+	github.com/sirupsen/logrus v1.8.1
 	github.com/spf13/afero v1.5.1 // indirect
 	github.com/spf13/cast v1.3.1 // indirect
 	github.com/spf13/cobra v1.1.3

go.sum

@@ -19,8 +19,8 @@ github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf/go.mod h1:ybxpYRF
 github.com/armon/circbuf v0.0.0-20150827004946-bbbad097214e/go.mod h1:3U/XgcO3hCbHZ8TKRvWD2dDTCfh9M9ya+I9JpbB7O8o=
 github.com/armon/go-metrics v0.0.0-20180917152333-f0300d1749da/go.mod h1:Q73ZrmVTwzkszR9V5SSuryQ31EELlFMUz1kKyl939pY=
 github.com/armon/go-radix v0.0.0-20180808171621-7fddfc383310/go.mod h1:ufUuZ+zHj4x4TnLV4JWEpy2hxWSpsRywHrMgIH9cCH8=
-github.com/aws/aws-sdk-go v1.38.18 h1:Y5W5siOHoiZbCHQcoEiih267L21sbw1FW+2Vdlq9dlU=
-github.com/aws/aws-sdk-go v1.38.18/go.mod h1:hcU610XS61/+aQV88ixoOzUoG7v3b31pl2zKMmprdro=
+github.com/aws/aws-sdk-go v1.38.21 h1:D08DXWI4QRaawLaW+OtsIEClOI90I6eheJs1GwXTQVI=
+github.com/aws/aws-sdk-go v1.38.21/go.mod h1:hcU610XS61/+aQV88ixoOzUoG7v3b31pl2zKMmprdro=
 github.com/beorn7/perks v0.0.0-20180321164747-3a771d992973/go.mod h1:Dwedo/Wpr24TaqPxmxbtue+5NUziq4I4S80YR8gNf3Q=
 github.com/beorn7/perks v1.0.0/go.mod h1:KWe93zE9D1o94FZ5RNwFwVgaQK1VOXiVxmqh+CedLV8=
 github.com/bgentry/speakeasy v0.1.0/go.mod h1:+zsyZBPWlz7T6j88CTgSN5bM796AkVf0kBD4zp0CCIs=
@@ -119,10 +119,6 @@ github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORN
 github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
 github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE=
 github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
-github.com/magefile/mage v1.10.0 h1:3HiXzCUY12kh9bIuyXShaVe529fJfyqoVM42o/uom2g=
-github.com/magefile/mage v1.10.0/go.mod h1:z5UZb/iS3GoOSn0JgWuiw7dxlurVYTu+/jHXqQg881A=
-github.com/magefile/mage v1.11.0 h1:C/55Ywp9BpgVVclD3lRnSYCwXTYxmSppIgLeDYlNuls=
-github.com/magefile/mage v1.11.0/go.mod h1:z5UZb/iS3GoOSn0JgWuiw7dxlurVYTu+/jHXqQg881A=
 github.com/magiconair/properties v1.8.1 h1:ZC2Vc7/ZFkGmsVC9KvOjumD+G5lXy2RtTKyzRKO2BQ4=
 github.com/magiconair/properties v1.8.1/go.mod h1:PppfXfuXeibc/6YijjN8zIbojt8czPbwD3XqdrwzmxQ=
 github.com/magiconair/properties v1.8.4 h1:8KGKTcQQGm0Kv7vEbKFErAoAOFyyacLStRtQSeYtvkY=
@@ -176,8 +172,8 @@ github.com/sean-/seed v0.0.0-20170313163322-e2103e2c3529/go.mod h1:DxrIzT+xaE7yg
 github.com/shurcooL/sanitized_anchor_name v1.0.0/go.mod h1:1NzhyTcUVG4SuEtjjoZeVRXNmyL/1OwPU0+IJeTBvfc=
 github.com/sirupsen/logrus v1.2.0 h1:juTguoYk5qI21pwyTXY3B3Y5cOTH3ZUyZCg1v/mihuo=
 github.com/sirupsen/logrus v1.2.0/go.mod h1:LxeOpSwHxABJmUn/MG1IvRgCAasNZTLOkJPxbbu5VWo=
-github.com/sirupsen/logrus v1.8.0 h1:nfhvjKcUMhBMVqbKHJlk5RPrrfYr/NMo3692g0dwfWU=
-github.com/sirupsen/logrus v1.8.0/go.mod h1:4GuYW9TZmE769R5STWrRakJc4UqQ3+QQ95fyz7ENv1A=
+github.com/sirupsen/logrus v1.8.1 h1:dJKuHgqk1NNQlqoA6BTlM1Wf9DOH3NBjQyu0h9+AZZE=
+github.com/sirupsen/logrus v1.8.1/go.mod h1:yWOB1SBYBC5VeMP7gHvWumXLIWorT60ONWic61uBYv0=
 github.com/smartystreets/assertions v0.0.0-20180927180507-b2de0cb4f26d h1:zE9ykElWQ6/NYmHa3jpm/yHnI4xSofP+UP6SpjHcSeM=
 github.com/smartystreets/assertions v0.0.0-20180927180507-b2de0cb4f26d/go.mod h1:OnSkiWE9lh6wB0YB77sQom3nweQdgAjqCqsofrRNTgc=
 github.com/smartystreets/goconvey v1.6.4 h1:fv0U8FUIMPNf1L9lnHLvLhgicrIVChEkdzIKYqbNC9s=