#280 Regexp expressions in KMS Customer key removal (#284)

* Updated KMS customer key removal to support aliases matching by regexp

* Updated imports

* Code cleanup

* Updated empty config creation

README.md

@@ -203,10 +203,7 @@ The following resources support the Config file:
     - Config key: `CloudWatchLogGroup`
 - KMS customer keys
     - Resource type: `kmscustomerkeys`
-
-Notes:
-  * no configuration options for KMS customer keys, since keys are created with auto-generated identifier
-
+    - Config key: `KMSCustomerKeys`
 
 #### Example
 

aws/aws.go

@@ -709,7 +709,7 @@ func GetAllResources(targetRegions []string, excludeAfter time.Time, resourceTyp
 		// KMS Customer managed keys
 		customerKeys := KmsCustomerKeys{}
 		if IsNukeable(customerKeys.ResourceName(), resourceTypes) {
-			keys, err := getAllKmsUserKeys(session, customerKeys.MaxBatchSize(), excludeAfter)
+			keys, err := getAllKmsUserKeys(session, customerKeys.MaxBatchSize(), excludeAfter, configObj)
 			if err != nil {
 				return nil, errors.WithStackTrace(err)
 			}

aws/kms_customer_key.go

@@ -4,6 +4,8 @@ import (
 	"sync"
 	"time"
 
+	"github.com/gruntwork-io/cloud-nuke/config"
+
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/kms"
@@ -12,83 +14,130 @@ import (
 	"github.com/hashicorp/go-multierror"
 )
 
-func getAllKmsUserKeys(session *session.Session, batchSize int, excludeAfter time.Time) ([]*string, error) {
+func getAllKmsUserKeys(session *session.Session, batchSize int, excludeAfter time.Time, configObj config.Config) ([]*string, error) {
 	svc := kms.New(session)
-	var kmsIds []*string
-	input := &kms.ListKeysInput{
-		Limit: aws.Int64(int64(batchSize)),
+	// collect all aliases for each key
+	keyAliases, err := listKeyAliases(svc, batchSize)
+	if err != nil {
+		return nil, errors.WithStackTrace(err)
 	}
-	listPage := 1
-	err := svc.ListKeysPages(input, func(page *kms.ListKeysOutput, lastPage bool) bool {
-		logging.Logger.Debugf("Loading User Key from page %d", listPage)
-
-		wg := new(sync.WaitGroup)
-		wg.Add(len(page.Keys))
-		keyChans := make([]chan *kms.KeyListEntry, len(page.Keys))
-		errChans := make([]chan error, len(page.Keys))
-		for i, key := range page.Keys {
-			keyChans[i] = make(chan *kms.KeyListEntry, 1)
-			errChans[i] = make(chan error, 1)
-			go shouldIncludeKmsUserKey(wg, svc, key, excludeAfter, keyChans[i], errChans[i])
-		}
-		wg.Wait()
 
-		// collect errors
-		for _, errChan := range errChans {
-			if err := <-errChan; err != nil {
-				logging.Logger.Errorf("[Failed] %s", err)
-			}
+	// checking in parallel if keys can be considered for removal
+	var wg sync.WaitGroup
+	wg.Add(len(keyAliases))
+	resultsChan := make([]chan *KmsCheckIncludeResult, len(keyAliases))
+	var id = 0
+	for key, aliases := range keyAliases {
+		resultsChan[id] = make(chan *KmsCheckIncludeResult, 1)
+		go shouldIncludeKmsUserKey(&wg, resultsChan[id], svc, key, aliases, excludeAfter, configObj)
+		id++
+	}
+	wg.Wait()
+
+	var kmsIds []*string
+	for _, channel := range resultsChan {
+		result := <-channel
+		if result.Error != nil {
+			logging.Logger.Warnf("Can't read KMS key %s", result.Error)
+
+			continue
 		}
-		// collect keys
-		for _, keyChan := range keyChans {
-			if key := <-keyChan; key != nil {
-				kmsIds = append(kmsIds, key.KeyId)
-			}
+		if result.KeyId != "" {
+			kmsIds = append(kmsIds, &result.KeyId)
 		}
+	}
+	return kmsIds, nil
+}
 
-		listPage++
-		return true
-	})
-
-	return kmsIds, errors.WithStackTrace(err)
+// KmsCheckIncludeResult - structure used results of evaluation: not null KeyId - key should be included
+type KmsCheckIncludeResult struct {
+	KeyId string
+	Error error
 }
 
-func shouldIncludeKmsUserKey(wg *sync.WaitGroup, svc *kms.KMS, key *kms.KeyListEntry, excludeAfter time.Time, keyChan chan *kms.KeyListEntry, errChan chan error) {
+func shouldIncludeKmsUserKey(wg *sync.WaitGroup, resultsChan chan *KmsCheckIncludeResult, svc *kms.KMS, key string,
+	aliases []string, excludeAfter time.Time, configObj config.Config) {
 	defer wg.Done()
+	var includedByName = false
+	// verify if key aliases matches configurations
+	for _, alias := range aliases {
+		v := config.ShouldInclude(alias, configObj.KMSCustomerKeys.IncludeRule.NamesRegExp,
+			configObj.KMSCustomerKeys.ExcludeRule.NamesRegExp)
+		if v {
+			includedByName = true
+
+			break
+		}
+	}
+	if !includedByName {
+		resultsChan <- &KmsCheckIncludeResult{KeyId: ""}
+		return
+	}
 	// additional request to describe key and get information about creation date, removal status
-	details, err := svc.DescribeKey(&kms.DescribeKeyInput{KeyId: key.KeyId})
+	details, err := svc.DescribeKey(&kms.DescribeKeyInput{KeyId: &key})
+
 	if err != nil {
-		errChan <- err
-		keyChan <- nil
+		resultsChan <- &KmsCheckIncludeResult{Error: err}
 		return
 	}
 	metadata := details.KeyMetadata
 	// evaluate only user keys
 	if *metadata.KeyManager != kms.KeyManagerTypeCustomer {
-		keyChan <- nil
-		errChan <- nil
+		resultsChan <- &KmsCheckIncludeResult{KeyId: ""}
 		return
 	}
 	// skip keys already scheduled for removal
 	if metadata.DeletionDate != nil {
-		keyChan <- nil
-		errChan <- nil
+		resultsChan <- &KmsCheckIncludeResult{KeyId: ""}
 		return
 	}
 	if metadata.PendingDeletionWindowInDays != nil {
-		keyChan <- nil
-		errChan <- nil
+		resultsChan <- &KmsCheckIncludeResult{KeyId: ""}
 		return
 	}
 	var referenceTime = *metadata.CreationDate
 	if referenceTime.After(excludeAfter) {
-		keyChan <- nil
-		errChan <- nil
+		resultsChan <- &KmsCheckIncludeResult{KeyId: ""}
 		return
 	}
 	// put key in channel to be considered for removal
-	keyChan <- key
-	errChan <- nil
+	resultsChan <- &KmsCheckIncludeResult{KeyId: key}
+}
+
+func listKeyAliases(svc *kms.KMS, batchSize int) (map[string][]string, error) {
+	// map key - KMS key id, value list of aliases
+	aliases := map[string][]string{}
+	var next *string
+
+	for {
+		list, err := svc.ListAliases(&kms.ListAliasesInput{
+			Marker: next,
+			Limit:  aws.Int64(int64(batchSize)),
+		})
+		if err != nil {
+			return nil, errors.WithStackTrace(err)
+		}
+
+		// collect key aliases to map
+		for _, alias := range list.Aliases {
+			key := alias.TargetKeyId
+			if key == nil {
+				continue
+			}
+			list, found := aliases[*key]
+			if !found {
+				list = make([]string, 0)
+			}
+			list = append(list, *alias.AliasName)
+			aliases[*key] = list
+		}
+
+		if list.NextMarker == nil || len(*list.NextMarker) == 0 {
+			break
+		}
+		next = list.NextMarker
+	}
+	return aliases, nil
 }
 
 func nukeAllCustomerManagedKmsKeys(session *session.Session, keyIds []*string) error {

aws/kms_customer_key_test.go

@@ -1,9 +1,14 @@
 package aws
 
 import (
+	"fmt"
+	"regexp"
 	"testing"
 	"time"
 
+	"github.com/gruntwork-io/cloud-nuke/config"
+	"github.com/gruntwork-io/cloud-nuke/util"
+
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/kms"
@@ -20,16 +25,45 @@ func TestListKmsUserKeys(t *testing.T) {
 	session, err := session.NewSession(&aws.Config{Region: aws.String(region)})
 	require.NoError(t, err)
 
-	createdKeyId := createKmsCustomerManagedKey(t, session, err)
+	aliasName := "cloud-nuke-test-" + util.UniqueID()
+	keyAlias := fmt.Sprintf("alias/%s", aliasName)
+	createdKeyId := createKmsCustomerManagedKey(t, session, keyAlias, err)
 
 	// test if listing of keys will return new key
-	keys, err := getAllKmsUserKeys(session, KmsCustomerKeys{}.MaxBatchSize(), time.Now())
+	keys, err := getAllKmsUserKeys(session, KmsCustomerKeys{}.MaxBatchSize(), time.Now(), config.Config{})
 	require.NoError(t, err)
 	assert.Contains(t, aws.StringValueSlice(keys), createdKeyId)
 
 	// test if time shift works
 	olderThan := time.Now().Add(-1 * time.Hour)
-	keys, err = getAllKmsUserKeys(session, KmsCustomerKeys{}.MaxBatchSize(), olderThan)
+	keys, err = getAllKmsUserKeys(session, KmsCustomerKeys{}.MaxBatchSize(), olderThan, config.Config{})
+	require.NoError(t, err)
+	assert.NotContains(t, aws.StringValueSlice(keys), createdKeyId)
+
+	// test if matching by regexp works
+	keys, err = getAllKmsUserKeys(session, KmsCustomerKeys{}.MaxBatchSize(), time.Now(), config.Config{
+		KMSCustomerKeys: config.ResourceType{
+			IncludeRule: config.FilterRule{
+				NamesRegExp: []config.Expression{
+					{RE: *regexp.MustCompile(fmt.Sprintf("^%s", keyAlias))},
+				},
+			},
+		},
+	})
+	require.NoError(t, err)
+	assert.Contains(t, aws.StringValueSlice(keys), createdKeyId)
+	assert.Equal(t, 1, len(keys))
+
+	// test if exclusion by regexp works
+	keys, err = getAllKmsUserKeys(session, KmsCustomerKeys{}.MaxBatchSize(), time.Now(), config.Config{
+		KMSCustomerKeys: config.ResourceType{
+			ExcludeRule: config.FilterRule{
+				NamesRegExp: []config.Expression{
+					{RE: *regexp.MustCompile(fmt.Sprintf("^%s", keyAlias))},
+				},
+			},
+		},
+	})
 	require.NoError(t, err)
 	assert.NotContains(t, aws.StringValueSlice(keys), createdKeyId)
 }
@@ -43,22 +77,28 @@ func TestRemoveKmsUserKeys(t *testing.T) {
 	session, err := session.NewSession(&aws.Config{Region: aws.String(region)})
 	require.NoError(t, err)
 
-	createdKeyId := createKmsCustomerManagedKey(t, session, err)
+	keyAlias := "alias/cloud-nuke-test-" + util.UniqueID()
+	createdKeyId := createKmsCustomerManagedKey(t, session, keyAlias, err)
 
 	err = nukeAllCustomerManagedKmsKeys(session, []*string{&createdKeyId})
 	require.NoError(t, err)
 
 	// test if key is not included for removal second time
-	keys, err := getAllKmsUserKeys(session, KmsCustomerKeys{}.MaxBatchSize(), time.Now())
+	keys, err := getAllKmsUserKeys(session, KmsCustomerKeys{}.MaxBatchSize(), time.Now(), config.Config{})
 	require.NoError(t, err)
 	assert.NotContains(t, aws.StringValueSlice(keys), createdKeyId)
 }
 
-func createKmsCustomerManagedKey(t *testing.T, session *session.Session, err error) string {
+func createKmsCustomerManagedKey(t *testing.T, session *session.Session, alias string, err error) string {
 	svc := kms.New(session)
 	input := &kms.CreateKeyInput{}
 	result, err := svc.CreateKey(input)
 	require.NoError(t, err)
 	createdKeyId := *result.KeyMetadata.KeyId
+
+	aliasInput := &kms.CreateAliasInput{AliasName: &alias, TargetKeyId: &createdKeyId}
+	_, err = svc.CreateAlias(aliasInput)
+	require.NoError(t, err)
+
 	return createdKeyId
 }

config/config.go

@@ -27,6 +27,7 @@ type Config struct {
 	VPC                   ResourceType `yaml:"VPC"`
 	OIDCProvider          ResourceType `yaml:"OIDCProvider"`
 	CloudWatchLogGroup    ResourceType `yaml:"CloudWatchLogGroup"`
+	KMSCustomerKeys       ResourceType `yaml:"KMSCustomerKeys"`
 }
 
 type ResourceType struct {

config/config_test.go

@@ -28,6 +28,7 @@ func emptyConfig() *Config {
 		ResourceType{FilterRule{}, FilterRule{}},
 		ResourceType{FilterRule{}, FilterRule{}},
 		ResourceType{FilterRule{}, FilterRule{}},
+		ResourceType{FilterRule{}, FilterRule{}},
 	}
 }