Adding the ability to filter the account list. Closes #225

README.md

@@ -120,7 +120,9 @@ By supplying a `--mapping` flag with a comma-delimited list of key=value pairs c
 |--output|-o|Output file path (default: ./aws.config)|No|
 |--stdout||Write config to stdout instead of a file|No|
 |--sso-friendly-name||Alternative name for the SSO start URL|No|
-|--list-accounts|Lists all available AWS accounts|
+|--list-accounts|Lists all available AWS accounts|No|
+|--include-accounts|Includes only these comma-separated account IDs in the output|No|
+|--exclude-accounts|Exclude these comma-separated account IDs from the output|No|
 
 ## Generated Config Format
 

cmd/constants.go

@@ -28,6 +28,8 @@ const (
 	FlagSSOFriendlyName string = "sso-friendly-name"
 	FlagCheckUpdate     string = "check-update"
 	FlagListAccounts    string = "list-accounts"
+	FlagIncludeAccounts string = "include-accounts"
+	FlagExcludeAccounts string = "exclude-accounts"
 )
 
 // Default output filename if no filename is specified

cmd/flags.go

@@ -11,4 +11,6 @@ var (
 	permissions     bool   // Flag to print the permissions needed and exit
 	checkUpdate     bool   // Flag to check if an update is available
 	listAccounts    bool   // Only list AWS accounts found
+	includeAccounts string // Comma-separated list of accounts to include
+	excludeAccounts string // Comma-separated list of accounts to exclude
 )

cmd/init.go

@@ -72,6 +72,8 @@ func init() {
 	rootCmd.PersistentFlags().StringVar(&ssoFriendlyName, FlagSSOFriendlyName, "", "Use this instead of the identity store ID for the start URL")
 	rootCmd.PersistentFlags().BoolVar(&checkUpdate, FlagCheckUpdate, false, "Check if a newer version of the tool is available")
 	rootCmd.PersistentFlags().BoolVar(&listAccounts, FlagListAccounts, false, "List all available AWS accounts")
+	rootCmd.PersistentFlags().StringVarP(&includeAccounts, FlagIncludeAccounts, "", "", "Include only these comma-separated accounts")
+	rootCmd.PersistentFlags().StringVarP(&excludeAccounts, FlagExcludeAccounts, "", "", "Exclude these comma-separated accounts")
 
 	rootCmd.PreRunE = func(cmd *cobra.Command, args []string) error {
 		if permissions || checkUpdate || listAccounts {

cmd/rootCmd.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"os"
+	"strings"
 
 	"github.com/scottbrown/setlist"
 
@@ -130,12 +131,23 @@ func buildProfiles(
 ) ([]setlist.Profile, error) {
 	profiles := []setlist.Profile{}
 
+	includedAccounts := buildIncludedAccounts(accounts)
+	excludedAccounts := buildExcludedAccounts()
+
 	for _, account := range accounts {
 		if account.Id == nil {
 			fmt.Fprintf(os.Stderr, "Warning: Found account with nil ID, skipping\n")
 			continue
 		}
 
+		if !includedAccounts.Contains(*account.Id) {
+			continue
+		}
+
+		if excludedAccounts.Contains(*account.Id) {
+			continue
+		}
+
 		permissionSets, err := setlist.PermissionSets(ctx, ssoClient, *instance.InstanceArn, *account.Id)
 		if err != nil {
 			return nil, fmt.Errorf("failed to list permission sets for account %s: %w", *account.Id, err)
@@ -208,3 +220,35 @@ func displayAccounts(accounts []orgtypes.Account) error {
 
 	return nil
 }
+
+type AccountsFilter []string
+
+func buildIncludedAccounts(accounts []orgtypes.Account) AccountsFilter {
+	if includeAccounts == "" {
+		var x AccountsFilter
+		for _, i := range accounts {
+			x = append(x, *i.Id)
+		}
+		return x
+	}
+
+	return strings.Split(includeAccounts, ",")
+}
+
+func buildExcludedAccounts() AccountsFilter {
+	if includeAccounts == "" {
+		return AccountsFilter{}
+	}
+
+	return strings.Split(excludeAccounts, ",")
+}
+
+func (a AccountsFilter) Contains(id string) bool {
+	for _, i := range a {
+		if i == id {
+			return true
+		}
+	}
+
+	return false
+}