🧹 Rework azure snapshot connection to use device connection. Drop all…

… flags that were used as a temporary workaround. (#4085)

Signed-off-by: Preslav <[email protected]>

providers/azure/config/config.go

@@ -98,27 +98,6 @@ 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,
-				},
-				{
-					Long:    "skip-snapshot-setup",
-					Type:    plugin.FlagType_Bool,
-					Default: "",
-					Desc:    "If set, no setup will be performed for the snapshot connection. It is expected that the target's disk is already attached. Use together with --lun.",
-					Option:  plugin.FlagOption_Hidden,
-				},
-				{
-					Long:    "lun",
-					Type:    plugin.FlagType_Int,
-					Default: "",
-					Desc:    "The logical unit number of the attached disk that should be scanned. Use together with --skip-snapshot-setup.",
-					Option:  plugin.FlagOption_Hidden,
-				},
 			},
 		},
 	},

providers/azure/connection/azureinstancesnapshot/instance.go

@@ -0,0 +1,84 @@
+// Copyright (c) Mondoo, Inc.
+// SPDX-License-Identifier: BUSL-1.1
+
+package azureinstancesnapshot
+
+import (
+	"context"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
+	compute "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute"
+	"github.com/cockroachdb/errors"
+	"github.com/rs/zerolog/log"
+)
+
+type instanceInfo struct {
+	subscriptionId string
+	resourceGroup  string
+	instanceName   string
+	location       string
+	bootDiskId     string
+	zones          []*string
+	// Attach the entire VM response as well
+	vm compute.VirtualMachine
+}
+
+func (s *instanceInfo) getFirstAvailableLun() (int, error) {
+	takenLuns := []int{}
+
+	if s.vm.Properties.StorageProfile == nil {
+		return 0, errors.New("instance storage profile not found")
+	}
+	if s.vm.Properties.StorageProfile.DataDisks == nil {
+		return 0, errors.New("instance data disks not found")
+	}
+
+	for _, disk := range s.vm.Properties.StorageProfile.DataDisks {
+		takenLuns = append(takenLuns, int(*disk.Lun))
+	}
+
+	availableLuns := []int{}
+	for i := 0; i < 64; i++ {
+		// exclude the taken LUNs
+		available := true
+		for _, d := range takenLuns {
+			if i == d {
+				available = false
+				break
+			}
+		}
+		if available {
+			availableLuns = append(availableLuns, i)
+		} else {
+			// log just for visibility
+			log.Debug().Int("LUN", i).Msg("azure snapshot> LUN is taken, skipping")
+		}
+	}
+	if len(availableLuns) == 0 {
+		return 0, errors.New("no available LUNs")
+	}
+	return availableLuns[0], nil
+}
+
+func GetInstanceInfo(resourceGroup, instanceName, subId string, token azcore.TokenCredential) (instanceInfo, error) {
+	ctx := context.Background()
+	ii := instanceInfo{}
+
+	computeSvc, err := computeClient(token, subId, nil)
+	if err != nil {
+		return ii, err
+	}
+
+	instance, err := computeSvc.Get(ctx, resourceGroup, instanceName, &compute.VirtualMachinesClientGetOptions{})
+	if err != nil {
+		return ii, err
+	}
+	ii.resourceGroup = resourceGroup
+	ii.instanceName = *instance.Name
+	ii.bootDiskId = *instance.Properties.StorageProfile.OSDisk.ManagedDisk.ID
+	ii.location = *instance.Location
+	ii.subscriptionId = subId
+	ii.zones = instance.Zones
+	ii.vm = instance.VirtualMachine
+	return ii, nil
+}

providers/azure/connection/azureinstancesnapshot/instance_test.go

@@ -0,0 +1,71 @@
+// Copyright (c) Mondoo, Inc.
+// SPDX-License-Identifier: BUSL-1.1
+
+package azureinstancesnapshot
+
+import (
+	"testing"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
+	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute"
+	"github.com/stretchr/testify/require"
+)
+
+func TestGetFirstAvailableLun(t *testing.T) {
+	t.Run("no luns available", func(t *testing.T) {
+		instanceInfo := &instanceInfo{
+			vm: armcompute.VirtualMachine{
+				Properties: &armcompute.VirtualMachineProperties{
+					StorageProfile: &armcompute.StorageProfile{
+						DataDisks: []*armcompute.DataDisk{},
+					},
+				},
+			},
+		}
+		// fill in all available LUNs for the scanner
+		for i := int32(0); i < 64; i++ {
+			instanceInfo.vm.Properties.StorageProfile.DataDisks = append(instanceInfo.vm.Properties.StorageProfile.DataDisks, &armcompute.DataDisk{
+				Lun: to.Ptr(i),
+			})
+		}
+
+		_, err := instanceInfo.getFirstAvailableLun()
+		require.Error(t, err)
+	})
+	t.Run("first available lun on a scanner with no disks", func(t *testing.T) {
+		instanceInfo := instanceInfo{
+			vm: armcompute.VirtualMachine{
+				Properties: &armcompute.VirtualMachineProperties{
+					StorageProfile: &armcompute.StorageProfile{
+						DataDisks: []*armcompute.DataDisk{},
+					},
+				},
+			},
+		}
+
+		lun, err := instanceInfo.getFirstAvailableLun()
+		require.NoError(t, err)
+		require.Equal(t, 0, lun)
+	})
+	t.Run("first available lun on a scanner with some disks", func(t *testing.T) {
+		instanceInfo := instanceInfo{
+			vm: armcompute.VirtualMachine{
+				Properties: &armcompute.VirtualMachineProperties{
+					StorageProfile: &armcompute.StorageProfile{
+						DataDisks: []*armcompute.DataDisk{},
+					},
+				},
+			},
+		}
+		// fill in 15 luns
+		for i := int32(0); i < 16; i++ {
+			instanceInfo.vm.Properties.StorageProfile.DataDisks = append(instanceInfo.vm.Properties.StorageProfile.DataDisks, &armcompute.DataDisk{
+				Lun: to.Ptr(i),
+			})
+		}
+
+		lun, err := instanceInfo.getFirstAvailableLun()
+		require.NoError(t, err)
+		require.Equal(t, 16, lun)
+	})
+}