Skip to content
This repository has been archived by the owner on Jul 31, 2023. It is now read-only.

Commit

Permalink
Fetch accounts and various roles in parallel (#8)
Browse files Browse the repository at this point in the history
* Add time display

* Fetch accounts and various roles in parallel

Baseline (3393bfc): `10.501797317s`
This commit: `3.150612075s`
  • Loading branch information
nihalgonsalves committed Mar 18, 2021
1 parent 113aaef commit 033af5f
Show file tree
Hide file tree
Showing 2 changed files with 116 additions and 33 deletions.
8 changes: 7 additions & 1 deletion aws-cfg-generator/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@ package main
*/

import (
"time"

"github.com/alecthomas/kong"
"github.com/rs/zerolog"
"github.com/rs/zerolog/log"
Expand All @@ -22,6 +24,8 @@ import (
)

func main() {
start := time.Now()

var cli cmd.CLI
ctx := kong.Parse(&cli)

Expand All @@ -30,9 +34,11 @@ func main() {
} else {
zerolog.SetGlobalLevel(zerolog.InfoLevel)
}

err := ctx.Run(cli)
if err != nil {
log.Panic().Err(err).Msgf("unexpected CLI error")
}

elapsed := time.Since(start)
log.Info().Msgf("Done in %s", elapsed)
}
141 changes: 109 additions & 32 deletions pkg/util/aws.go
Original file line number Diff line number Diff line change
Expand Up @@ -57,26 +57,19 @@ func GetAWSContext() (client *AWSContext) {
}

func (ctx *AWSContext) GetRolesAndAccounts() (roleArns []string, accountMap map[string]string) {
accountMap = ctx.getAccountNames()
cRoles := make(chan []string)
cAccount := make(chan map[string]string)

gcio, err := ctx.sts.GetCallerIdentity(&sts.GetCallerIdentityInput{})
if err != nil {
log.Panic().Err(err).Msg("could not get caller identity")
}
go func() {
cRoles <- ctx.getRoles()
}()

log.Info().Str("user-arn", *gcio.Arn).Msg("Found user")
go func() {
cAccount <- ctx.getAccountNames()
}()

lgfuo, err := ctx.iam.ListGroupsForUser(&iam.ListGroupsForUserInput{
UserName: getUser(gcio.Arn),
})
if err != nil {
log.Panic().Err(err).Str("user", *getUser(gcio.Arn)).Msg("could not list groups for user")
}

for _, group := range lgfuo.Groups {
log.Debug().Str("group", *group.GroupName).Msg("Finding roles for group")
roleArns = append(roleArns, ctx.getRoleArnsForGroup(group)...)
}
roleArns = <-cRoles
accountMap = <-cAccount

return
}
Expand Down Expand Up @@ -132,6 +125,44 @@ func getUser(userArn *string) *string {
return &arnParts[1]
}

func (ctx *AWSContext) getRoles() (roleArns []string) {
log.Debug().Msg("getting caller identity")

gcio, err := ctx.sts.GetCallerIdentity(&sts.GetCallerIdentityInput{})
if err != nil {
log.Panic().Err(err).Msg("could not get caller identity")
}

log.Info().Str("user-arn", *gcio.Arn).Msg("Found user")

lgfuo, err := ctx.iam.ListGroupsForUser(&iam.ListGroupsForUserInput{
UserName: getUser(gcio.Arn),
})
if err != nil {
log.Panic().Err(err).Str("user", *getUser(gcio.Arn)).Msg("could not list groups for user")
}

log.Debug().Msgf("Found %d groups", len(lgfuo.Groups))

c := make(chan []string)

for _, group := range lgfuo.Groups {
go func(g iam.Group) {
log.Debug().Str("group", *g.GroupName).Msg("Finding roles for group")
c <- ctx.getRoleArnsForGroup(&g)
}(*group)
}

for range lgfuo.Groups {
roleArns = append(roleArns, (<-c)...)
}

log.Info().Msgf("Found %d roles", len(roleArns))
log.Debug().Strs("roles", roleArns).Msgf("Roles")

return
}

func (ctx *AWSContext) getAccountNames() map[string]string {
accIDToName := map[string]string{}

Expand All @@ -145,9 +176,13 @@ func (ctx *AWSContext) getAccountNames() map[string]string {
break
}

log.Debug().Msgf("found %d member accounts", len(lao.Accounts))

for _, acc := range lao.Accounts {
accIDToName[*acc.Id] = *acc.Name
log.Debug().Str("account-id", *acc.Id).Str("account-name", *acc.Name).
log.Debug().
Str("account-id", *acc.Id).
Str("account-name", *acc.Name).
Msg("found organization member account")
}

Expand All @@ -162,35 +197,51 @@ func (ctx *AWSContext) getAccountNames() map[string]string {
}

func (ctx *AWSContext) getRoleArnsForGroup(group *iam.Group) (roles []string) {
c := make(chan []string)

go func() {
c <- ctx.listInlinePolicyAndGetRoles(group)
}()
go func() {
c <- ctx.listAttachedPolicyAndGetRoles(group)
}()

roles = append(roles, (<-c)...)
roles = append(roles, (<-c)...)

return
}

func (ctx *AWSContext) listInlinePolicyAndGetRoles(group *iam.Group) (roleArns []string) {
log.Debug().Str("group", *group.GroupName).Msg("finding roles from group inline policies")

lgpo, err := ctx.iam.ListGroupPolicies(&iam.ListGroupPoliciesInput{
GroupName: group.GroupName,
})
if err != nil {
log.Panic().Err(err).Str("group", *group.GroupName).Msg("could not list inline group policies")
}

c := make(chan []string)

for _, policy := range lgpo.PolicyNames {
log.Debug().Str("policy", *policy).Msg("Finding roles for inlined policy")
roles = append(roles, ctx.getRoleArnsForInlinePolicy(group.GroupName, policy)...)
go func(p string) {
log.Debug().Str("policy", p).Msg("Finding roles for inlined policy")
c <- ctx.getRoleArnsForInlinePolicy(*group.GroupName, p)
}(*policy)
}

lagpo, err := ctx.iam.ListAttachedGroupPolicies(&iam.ListAttachedGroupPoliciesInput{
GroupName: group.GroupName,
})
if err != nil {
log.Panic().Err(err).Str("group", *group.GroupName).Msg("could not list attached group policies")
}
for _, policy := range lagpo.AttachedPolicies {
log.Debug().Str("policy ARN", *policy.PolicyArn).Msg("Finding roles for attached policy")
roles = append(roles, ctx.getRoleArnsForAttachedPolicy(policy)...)
for range lgpo.PolicyNames {
roleArns = append(roleArns, (<-c)...)
}

return
}

func (ctx *AWSContext) getRoleArnsForInlinePolicy(group, policyName *string) []string {
func (ctx *AWSContext) getRoleArnsForInlinePolicy(group, policyName string) []string {
ggpo, err := ctx.iam.GetGroupPolicy(&iam.GetGroupPolicyInput{
GroupName: group,
PolicyName: policyName,
GroupName: &group,
PolicyName: &policyName,
})
if err != nil {
log.Panic().Err(err).Msg("could not get group policy")
Expand All @@ -199,6 +250,32 @@ func (ctx *AWSContext) getRoleArnsForInlinePolicy(group, policyName *string) []s
return getRolesArnsFromPolicy(ggpo.PolicyDocument)
}

func (ctx *AWSContext) listAttachedPolicyAndGetRoles(group *iam.Group) (roleArns []string) {
log.Debug().Str("group", *group.GroupName).Msg("finding roles from group attached policies")

lagpo, err := ctx.iam.ListAttachedGroupPolicies(&iam.ListAttachedGroupPoliciesInput{
GroupName: group.GroupName,
})
if err != nil {
log.Panic().Err(err).Str("group", *group.GroupName).Msg("could not list attached group policies")
}

c := make(chan []string)

for _, policy := range lagpo.AttachedPolicies {
go func(p iam.AttachedPolicy) {
log.Debug().Str("policy ARN", *p.PolicyArn).Msg("Finding roles for attached policy")
c <- ctx.getRoleArnsForAttachedPolicy(&p)
}(*policy)
}

for range lagpo.AttachedPolicies {
roleArns = append(roleArns, (<-c)...)
}

return
}

func (ctx *AWSContext) getRoleArnsForAttachedPolicy(policy *iam.AttachedPolicy) []string {
gpio, err := ctx.iam.GetPolicy(&iam.GetPolicyInput{
PolicyArn: policy.PolicyArn,
Expand Down

0 comments on commit 033af5f

Please sign in to comment.