🧹 Attach data disks with a delete opt. Add an option to skip detach/d…

…elete steps (#3276) * 🧹 Attach data disks with a delete opt. Add an option to skip detach/delete steps. Signed-off-by: Preslav <[email protected]> * Add a flag to skip cleanup. * Better flag assertion. --------- Signed-off-by: Preslav <[email protected]>
mondoohq · Feb 14, 2024 · 114413a · 114413a
1 parent 5aeb69d
commit 114413a
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 4 deletions.
diff --git a/providers/azure/config/config.go b/providers/azure/config/config.go
@@ -97,6 +97,13 @@ var Config = plugin.Provider{
 					Desc:    "Comma-separated list of Azure subscriptions to exclude.",
 					Option:  plugin.FlagOption_Hidden,
 				},
+				{
+					Long:    "skip-snapshot-cleanup",
+					Type:    plugin.FlagType_Bool,
+					Default: "",
+					Desc:    "If set, no cleanup will be performed for the snapshot connection.",
+					Option:  plugin.FlagOption_Hidden,
+				},
 			},
 		},
 	},

diff --git a/providers/azure/connection/azureinstancesnapshot/provider.go b/providers/azure/connection/azureinstancesnapshot/provider.go
@@ -29,6 +29,7 @@ const (
 	DiskTargetType         string                = "disk"
 	SnapshotTargetType     string                = "snapshot"
 	InstanceTargetType     string                = "instance"
+	SkipCleanup            string                = "skip-snapshot-cleanup"
 )
 
 // the instance from which we're performing the scan
@@ -307,8 +308,9 @@ func (c *AzureSnapshotConnection) Close() {
 			log.Error().Err(err).Msg("unable to unmount volume")
 		}
 	}
-
-	if c.snapshotCreator != nil {
+	if c.skipCleanup() {
+		log.Debug().Msgf("skipping azure snapshot cleanup, %s flag is set to true", SkipCleanup)
+	} else if c.snapshotCreator != nil {
 		if c.mountInfo.diskName != "" {
 			err := c.snapshotCreator.detachDisk(c.mountInfo.diskName, c.scanner.instanceInfo)
 			if err != nil {
@@ -332,6 +334,10 @@ func (c *AzureSnapshotConnection) Close() {
 	}
 }
 
+func (c *AzureSnapshotConnection) skipCleanup() bool {
+	return c.opts[SkipCleanup] == "true"
+}
+
 func (c *AzureSnapshotConnection) Kind() string {
 	return "api"
 }

diff --git a/providers/azure/connection/azureinstancesnapshot/snapshot.go b/providers/azure/connection/azureinstancesnapshot/snapshot.go
@@ -214,12 +214,14 @@ func (sc *snapshotCreator) attachDisk(targetInstance instanceInfo, diskName, dis
 		return err
 	}
 	attachOpt := compute.DiskCreateOptionTypesAttach
+	deleteOpt := compute.DiskDeleteOptionTypesDelete
 	// the Azure API requires all disks to be specified, even the already attached ones.
 	// we simply attach the new disk to the end of the already present list of data disks
 	disks := targetInstance.vm.Properties.StorageProfile.DataDisks
 	disks = append(disks, &compute.DataDisk{
 		Name:         &diskName,
 		CreateOption: &attachOpt,
+		DeleteOption: &deleteOpt,
 		Lun:          &lun,
 		ManagedDisk: &compute.ManagedDiskParameters{
 			ID: &diskId,

diff --git a/providers/azure/provider/provider.go b/providers/azure/provider/provider.go
@@ -42,7 +42,7 @@ func (s *Service) ParseCLI(req *plugin.ParseCLIReq) (*plugin.ParseCLIRes, error)
 	subscriptionsToExclude := flags["subscriptions-exclude"]
 	certificatePath := flags["certificate-path"]
 	certificateSecret := flags["certificate-secret"]
-
+	skipSnapshotCleanup := flags["skip-snapshot-cleanup"]
 	opts := map[string]string{}
 	creds := []*vault.Credential{}
 
@@ -57,7 +57,10 @@ func (s *Service) ParseCLI(req *plugin.ParseCLIReq) (*plugin.ParseCLIRes, error)
 	if len(subscriptionsToExclude.Value) > 0 {
 		opts["subscriptions-exclude"] = string(subscriptionsToExclude.Value)
 	}
-
+	// the presence of the flag indicates that we should skip cleanup
+	if present := skipSnapshotCleanup.RawData().Value.(bool); present {
+		opts[azureinstancesnapshot.SkipCleanup] = "true"
+	}
 	if len(clientSecret.Value) > 0 {
 		creds = append(creds, &vault.Credential{
 			Type:   vault.CredentialType_password,