Support deletion of SageMaker Notebook Instances

README.md

@@ -5,7 +5,7 @@
 This repo contains a CLI tool to delete all resources in an AWS account. cloud-nuke was created for situations when you might have an account you use for testing and need to clean up leftover resources so you're not charged for them. Also great for cleaning out accounts with redundant resources. Also great for removing unnecessary defaults like default VPCs and permissive ingress/egress rules in default security groups.
 
 In addition, cloud-nuke offers non-destructive inspecting functionality that can either be called via the command-line interface, or consumed as library methods, for scripting purposes.
- 
+
 The currently supported functionality includes:
 
 ## AWS
@@ -44,6 +44,7 @@ The currently supported functionality includes:
 - Inspecting and deleting all CloudWatch Log Groups in an AWS Account
 - Inspecting and deleting all GuardDuty Detectors in an AWS Account
 - Inspecting and deleting all Macie member accounts in an AWS account - as long as those accounts were created by Invitation - and not via AWS Organizations
+- Inspecting and deleting all SageMaker Notebook Instances in an AWS account
 
 ### BEWARE!
 
@@ -82,13 +83,13 @@ When using `cloud-nuke aws`, or `cloud-nuke inspect-aws`, you can use the `--reg
 cloud-nuke aws --region ap-south-1 --region ap-south-2
 ```
 
-Similarly, the following command will inspect resources only in `us-east-1` 
+Similarly, the following command will inspect resources only in `us-east-1`
 ```shell
 cloud-nuke inspect-aws --region us-east-1
 ```
 
-Including regions is available within: 
-- `cloud-nuke aws` 
+Including regions is available within:
+- `cloud-nuke aws`
 - `cloud-nuke defaults-aws`
 - `cloud-nuke inspect-aws`
 
@@ -108,8 +109,8 @@ cloud-nuke inspect-aws --exclude-region us-west-1
 
 `--region` and `--exclude-region` flags cannot be specified together i.e. they are mutually exclusive.
 
-Excluding regions is available within: 
-- `cloud-nuke aws` 
+Excluding regions is available within:
+- `cloud-nuke aws`
 - `cloud-nuke defaults-aws`
 - `cloud-nuke inspect-aws`
 
@@ -121,8 +122,8 @@ You can use the `--older-than` flag to only nuke resources that were created bef
 cloud-nuke aws --older-than 24h
 ```
 
-Excluding resources by age is available within: 
-- `cloud-nuke aws` 
+Excluding resources by age is available within:
+- `cloud-nuke aws`
 - `cloud-nuke inspect-aws`
 
 
@@ -134,8 +135,8 @@ You can use the `--list-resource-types` flag to list resource types whose termin
 cloud-nuke aws --list-resource-types
 ```
 
-Listing supported resource types is available within: 
-- `cloud-nuke aws` 
+Listing supported resource types is available within:
+- `cloud-nuke aws`
 - `cloud-nuke inspect-aws`
 
 
@@ -152,14 +153,14 @@ will search and target only `ec2` and `ami` resources. The specified resource ty
 i.e. it should be present in the `--list-resource-types` output. Using `--resource-type` also speeds up search because
 we are searching only for specific resource types.
 
-Similarly, the following command will inspect only ec2 instances: 
+Similarly, the following command will inspect only ec2 instances:
 
 ```shell
 cloud-nuke inspect-aws --resource-type ec2
 ```
 
-Specifying target resource types is available within: 
-- `cloud-nuke aws` 
+Specifying target resource types is available within:
+- `cloud-nuke aws`
 - `cloud-nuke inspect-aws`
 
 ### Exclude terminating specific resource types
@@ -175,8 +176,8 @@ This will terminate all resource types other than S3 and EC2.
 
 `--resource-type` and `--exclude-resource-type` flags cannot be specified together i.e. they are mutually exclusive.
 
-Specifying resource types to exclude is available within: 
-- `cloud-nuke aws` 
+Specifying resource types to exclude is available within:
+- `cloud-nuke aws`
 - `cloud-nuke inspect-aws`
 
 ### Dry run mode
@@ -188,14 +189,14 @@ If you want to check what resources are going to be targeted without actually te
 cloud-nuke aws --resource-type ec2 --dry-run
 ```
 
-Dry run mode is only available within: 
+Dry run mode is only available within:
 - `cloud-nuke aws`
 
 ### Using cloud-nuke as a library
 
-You can import cloud-nuke into other projects and use it as a library for programmatically inspecting and counting resources. 
+You can import cloud-nuke into other projects and use it as a library for programmatically inspecting and counting resources.
 
-```golang 
+```golang
 
 package main
 
@@ -333,6 +334,7 @@ The following resources support the Config file:
     - Config key: `CloudWatchLogGroup`
 - KMS customer keys
     - Resource type: `kmscustomerkeys`
+<<<<<<< HEAD
     - Config key: `KMSCustomerKeys`
 - Auto Scaling Groups
     - Resource type: `asg`
@@ -349,6 +351,13 @@ The following resources support the Config file:
 - EKS Clusters
     - Resource type: `ekscluster`
     - Config key: `EKSCluster`
+- SageMaker Notebook Instances
+  - Resource type: `sagemaker-notebook-instances`
+  - Config key: `SageMakerNotebook`
+
+Notes:
+  * no configuration options for KMS customer keys, since keys are created with auto-generated identifier
+
 
 #### Example
 
@@ -459,11 +468,10 @@ To find out what we options are supported in the config file today, consult this
 | eks                 | none | ✅   | none | none |
 | acmpca              | none | none | none | none |
 | iam role            | none | none | none | none |
+| sagemaker-notebook-instances| none| ✅   | none | none       |
 | ... (more to come)  | none | none | none | none |
 
 
-
-
 ### Log level
 
 You can set the log level by specifying the `--log-level` flag as per [logrus](https://github.com/sirupsen/logrus) log levels:

aws/aws.go

@@ -220,7 +220,6 @@ func GetAllResources(targetRegions []string, excludeAfter time.Time, resourceTyp
 		logging.Logger.Infof("Checking region [%d/%d]: %s", count, totalRegions, region)
 
 		cloudNukeSession := newSession(region)
-
 		resourcesInRegion := AwsRegionResource{}
 
 		// The order in which resources are nuked is important
@@ -717,6 +716,7 @@ func GetAllResources(targetRegions []string, excludeAfter time.Time, resourceTyp
 
 		}
 		// End GuardDuty detectors
+
 		// Macie member accounts
 		macieAccounts := MacieMember{}
 		if IsNukeable(macieAccounts.ResourceName(), resourceTypes) {
@@ -733,6 +733,20 @@ func GetAllResources(targetRegions []string, excludeAfter time.Time, resourceTyp
 		}
 		// End Macie member accounts
 
+		// Start SageMaker Notebook Instances
+		notebookInstances := SageMakerNotebookInstances{}
+		if IsNukeable(notebookInstances.ResourceName(), resourceTypes) {
+			instances, err := getAllNotebookInstances(cloudNukeSession, excludeAfter, configObj)
+			if err != nil {
+				return nil, errors.WithStackTrace(err)
+			}
+			if len(instances) > 0 {
+				notebookInstances.InstanceNames = awsgo.StringValueSlice(instances)
+				resourcesInRegion.Resources = append(resourcesInRegion.Resources, notebookInstances)
+			}
+		}
+		// End SageMaker Notebook Instances
+
 		if len(resourcesInRegion.Resources) > 0 {
 			account.Resources[region] = resourcesInRegion
 		}
@@ -846,6 +860,7 @@ func ListResourceTypes() []string {
 		CloudWatchLogGroups{}.ResourceName(),
 		GuardDuty{}.ResourceName(),
 		MacieMember{}.ResourceName(),
+		SageMakerNotebookInstances{}.ResourceName(),
 	}
 	sort.Strings(resourceTypes)
 	return resourceTypes

aws/sagemaker_notebook_instance.go

@@ -0,0 +1,99 @@
+package aws
+
+import (
+	"time"
+
+	awsgo "github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/session"
+	"github.com/aws/aws-sdk-go/service/sagemaker"
+	"github.com/gruntwork-io/cloud-nuke/config"
+	"github.com/gruntwork-io/cloud-nuke/logging"
+	"github.com/gruntwork-io/go-commons/errors"
+)
+
+func getAllNotebookInstances(session *session.Session, excludeAfter time.Time, configObj config.Config) ([]*string, error) {
+	svc := sagemaker.New(session)
+
+	result, err := svc.ListNotebookInstances(&sagemaker.ListNotebookInstancesInput{})
+
+	if err != nil {
+		return nil, errors.WithStackTrace(err)
+	}
+
+	var names []*string
+
+	for _, notebook := range result.NotebookInstances {
+		if notebook.CreationTime == nil {
+			continue
+		} 
+		if !excludeAfter.After(awsgo.TimeValue(notebook.CreationTime)) {
+			continue
+		}
+		if !config.ShouldInclude(awsgo.StringValue(notebook.NotebookInstanceName), configObj.S3.IncludeRule.NamesRegExp, configObj.S3.ExcludeRule.NamesRegExp){
+			continue
+		}
+		names = append(names, notebook.NotebookInstanceName)
+	}
+
+
+	return names, nil
+}
+
+func nukeAllNotebookInstances(session *session.Session, names []*string) error {
+	svc := sagemaker.New(session)
+
+	if len(names) == 0 {
+		logging.Logger.Infof("No Sagemaker Notebook Instance to nuke in region %s", *session.Config.Region)
+		return nil
+	}
+
+	logging.Logger.Infof("Deleting all Sagemaker Notebook Instances in region %s", *session.Config.Region)
+	deletedNames := []*string{}
+
+	for _, name := range names {
+		params := &sagemaker.DeleteNotebookInstanceInput{
+			NotebookInstanceName: name,
+		}
+
+		_, err := svc.StopNotebookInstance(&sagemaker.StopNotebookInstanceInput{
+			NotebookInstanceName: name,
+		})
+		if err != nil {
+			logging.Logger.Errorf("[Failed] %s: %s", *name, err)
+		}
+
+		err = svc.WaitUntilNotebookInstanceStopped(&sagemaker.DescribeNotebookInstanceInput{
+			NotebookInstanceName: name,
+		})
+
+		if err != nil {
+			logging.Logger.Errorf("[Failed] %s", err)
+		}
+
+		_, err = svc.DeleteNotebookInstance(params)
+
+		if err != nil {
+			logging.Logger.Errorf("[Failed] %s: %s", *name, err)
+		} else {
+			deletedNames = append(deletedNames, name)
+			logging.Logger.Infof("Deleted Sagemaker Notebook Instance: %s", awsgo.StringValue(name))
+		}
+	}
+
+	if len(deletedNames) > 0 {
+		for _, name := range deletedNames {
+
+			err := svc.WaitUntilNotebookInstanceDeleted(&sagemaker.DescribeNotebookInstanceInput{
+				NotebookInstanceName: name,
+			})
+
+			if err != nil {
+				logging.Logger.Errorf("[Failed] %s", err)
+				return errors.WithStackTrace(err)
+			}
+		}
+	}
+
+	logging.Logger.Infof("[OK] %d Sagemaker Notebook Instance(s) deleted in %s", len(deletedNames), *session.Config.Region)
+	return nil
+}

aws/sagemaker_notebook_instance_test.go

@@ -0,0 +1,102 @@
+package aws
+
+import (
+	"strings"
+	"testing"
+	"time"
+
+	awsgo "github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/session"
+	"github.com/aws/aws-sdk-go/service/sagemaker"
+	"github.com/gruntwork-io/cloud-nuke/config"
+	"github.com/gruntwork-io/cloud-nuke/logging"
+	"github.com/gruntwork-io/cloud-nuke/util"
+	"github.com/gruntwork-io/go-commons/errors"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// There's a built-in function WaitUntilDBInstanceAvailable but
+// the times that it was tested, it wasn't returning anything so we'll leave with the
+// custom one.
+
+func waitUntilNotebookInstanceCreated(svc *sagemaker.SageMaker, name *string) error {
+	input := &sagemaker.DescribeNotebookInstanceInput{
+		NotebookInstanceName: name,
+	}
+
+	for i := 0; i < 600; i++ {
+		instance, err := svc.DescribeNotebookInstance(input)
+		status := instance.NotebookInstanceStatus
+
+		if awsgo.StringValue(status) != "Pending" {
+			return nil
+		}
+
+		if err != nil {
+			return err
+		}
+
+		time.Sleep(1 * time.Second)
+		logging.Logger.Debug("Waiting for SageMaker Notebook Instance to be created")
+	}
+
+	return SageMakerNotebookInstanceDeleteError{name: *name}
+}
+
+func createTestNotebookInstance(t *testing.T, session *session.Session, name string, roleArn string) {
+	svc := sagemaker.New(session)
+
+	params := &sagemaker.CreateNotebookInstanceInput{
+		InstanceType:         awsgo.String("ml.t2.medium"),
+		NotebookInstanceName: awsgo.String(name),
+		RoleArn:              awsgo.String(roleArn),
+	}
+
+	_, err := svc.CreateNotebookInstance(params)
+	require.NoError(t, err)
+
+	waitUntilNotebookInstanceCreated(svc, &name)
+}
+
+func TestNukeNotebookInstance(t *testing.T) {
+	t.Parallel()
+
+	region, err := getRandomRegion()
+
+	require.NoError(t, errors.WithStackTrace(err))
+
+	session, err := session.NewSessionWithOptions(
+		session.Options{
+			SharedConfigState: session.SharedConfigEnable,
+			Config: awsgo.Config{
+				Region: awsgo.String(region),
+			},
+		},
+	)
+
+	notebookName := "cloud-nuke-test-" + util.UniqueID()
+	excludeAfter := time.Now().Add(1 * time.Hour)
+
+	role := createNotebookRole(t, session, notebookName+"-role")
+	defer deleteNotebookRole(session, role)
+
+	createTestNotebookInstance(t, session, notebookName, *role.Arn)
+
+	defer func() {
+		nukeAllNotebookInstances(session, []*string{&notebookName})
+
+		notebookNames, _ := getAllNotebookInstances(session, excludeAfter, config.Config{})
+
+		assert.NotContains(t, awsgo.StringValueSlice(notebookNames), strings.ToLower(notebookName))
+	}()
+
+	instances, err := getAllNotebookInstances(session, excludeAfter, config.Config{})
+
+	if err != nil {
+		assert.Failf(t, "Unable to fetch list of SageMaker Notebook Instances", errors.WithStackTrace(err).Error())
+	}
+
+	assert.Contains(t, awsgo.StringValueSlice(instances), notebookName)
+
+}