Golang 'script' to query lambda runtimes across AWS accounts

lambdaruntimequery.go

package main

import (
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/credentials"
	"github.com/aws/aws-sdk-go-v2/service/lambda"
	"github.com/aws/aws-sdk-go-v2/service/sso"
	"github.com/aws/aws-sdk-go-v2/service/sso/types"
	"github.com/aws/aws-sdk-go-v2/service/sts"
	"github.com/briandowns/spinner"
)

type Lambda struct {
	AccountName  string `json:"account_name"`
	AccountID    string `json:"account_id"`
	FunctionName string `json:"function_name"`
	Runtime      string `json:"runtime"`
}

func getRuntimesToFilterFor() []string {
	// return []string{"python3.6"}
	return []string{"python3.6", "python3.7", "python3.8"}
}

func ssoToken() (string, error) {
	homeDir, _ := os.UserHomeDir()
	awsSSOCacheDir := homeDir + "/.aws/sso/cache"

	files, _ := os.ReadDir(awsSSOCacheDir)
	for _, file := range files {
		if strings.HasSuffix(file.Name(), ".json") {
			if !strings.HasPrefix(file.Name(), "botocore-client") {
				plan, _ := os.ReadFile(awsSSOCacheDir + string(os.PathSeparator) + file.Name())
				ssoCache := struct {
					AccessToken string `json:"accessToken"`
				}{}
				_ = json.Unmarshal(plan, &ssoCache)

				return ssoCache.AccessToken, nil
			}
		}
	}

	return "", errors.New("could not get access token from SSO cache")
}

func getAccounts(ssoClient *sso.Client) ([]types.AccountInfo, error) {
	var accounts []types.AccountInfo

	token, _ := ssoToken()

	accountsInput := sso.ListAccountsInput{AccessToken: &token}

	for {
		response, err := ssoClient.ListAccounts(context.TODO(), &accountsInput)
		if err != nil {
			return []types.AccountInfo{}, err
		}

		accounts = append(accounts, response.AccountList...)

		token := response.NextToken
		if token == nil || *token == "" {
			break
		} else {
			accountsInput.NextToken = response.NextToken
		}
	}

	return accounts, nil
}

func main() {
	formatFlag := flag.String("format", "text", "Output format (text|json)")
	flag.Parse()

	const duration = 80 * time.Millisecond
	progressSpinner := spinner.New(spinner.CharSets[11], duration)
	progressSpinner.Suffix = " Starting..."
	progressSpinner.Start()

	token, _ := ssoToken()

	cfg, _ := config.LoadDefaultConfig(context.TODO())

	ssoClient := sso.NewFromConfig(cfg)
	_ = sts.NewFromConfig(cfg)

	var lambdas []Lambda

	accountsList, _ := getAccounts(ssoClient)
	for i, account := range accountsList {
		accountID := account.AccountId
		accountName := account.AccountName
		progressSpinner.Suffix = fmt.Sprintf(" [%d/%d] Querying lambdas in %v... ", i, len(accountsList), *accountName)

		rolesOutput, err := ssoClient.ListAccountRoles(context.TODO(),
			&sso.ListAccountRolesInput{
				AccessToken: &token, AccountId: accountID,
			},
		)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}

		role := "AdministratorAccess"

		for _, r := range rolesOutput.RoleList {
			if *r.RoleName == "ViewOnlyAccess" {
				role = "ViewOnlyAccess"
			}
		}

		roleCredentials, err := ssoClient.GetRoleCredentials(context.TODO(), &sso.GetRoleCredentialsInput{
			AccessToken: &token, AccountId: accountID, RoleName: &role,
		},
		)
		if err != nil {
			progressSpinner.Stop()
			fmt.Printf("Could not get credentails: %v", err)
			os.Exit(1)
		}

		accountCredentials := credentials.NewStaticCredentialsProvider(
			*roleCredentials.RoleCredentials.AccessKeyId,
			*roleCredentials.RoleCredentials.SecretAccessKey,
			*roleCredentials.RoleCredentials.SessionToken,
		)

		accountConfig := aws.Config{Region: "eu-central-1", Credentials: accountCredentials}
		lambdaClient := lambda.NewFromConfig(accountConfig)

		lambdaFunctions, err := lambdaClient.ListFunctions(context.TODO(), &lambda.ListFunctionsInput{})
		if err != nil {
			progressSpinner.Stop()
			fmt.Printf("Could not list lambda functions: %v", err)
			os.Exit(1)
		}

		isInFilterScope := func(runtimeName string) bool {
			for _, r := range getRuntimesToFilterFor() {
				if r == runtimeName {
					return true
				}
			}

			return false
		}

		for _, function := range lambdaFunctions.Functions {
			lambdaRuntime := string(function.Runtime)
			if isInFilterScope(lambdaRuntime) {
				lambdas = append(lambdas, Lambda{
					AccountName:  *accountName,
					AccountID:    *accountID,
					FunctionName: *function.FunctionName,
					Runtime:      string(function.Runtime),
				})
			}
		}
	}

	progressSpinner.Stop()

	if *formatFlag == "json" {
		b, _ := json.MarshalIndent(lambdas, "", "  ") // encoding/json
		fmt.Println(string(b))
	} else {
		for _, l := range lambdas {
			fmt.Printf("account: %-45v id: %-15v function: %v [%v]\n", l.AccountName, l.AccountID, l.FunctionName, l.Runtime)
		}
		fmt.Printf("Found %v lambda functions with runtime(progressSpinner) %v\n", len(lambdas), getRuntimesToFilterFor())
	}
}