fix(vm): handle Nonetype is not iterable for extensions (#7360)

Co-authored-by: Sergio Garcia <hello@mistercloudsec.com>
This commit is contained in:
Daniel Barranquero
2025-03-25 12:25:15 +01:00
committed by GitHub
parent dd1cc2d025
commit 5d6ed640f0
2 changed files with 83 additions and 21 deletions
@@ -24,6 +24,33 @@ class VirtualMachines(AzureService):
virtual_machines.update({subscription_name: {}})
for vm in virtual_machines_list:
storage_profile = getattr(vm, "storage_profile", None)
os_disk = (
getattr(storage_profile, "os_disk", None)
if storage_profile
else None
)
data_disks = []
if storage_profile and getattr(storage_profile, "data_disks", []):
data_disks = [
DataDisk(
lun=data_disk.lun,
name=data_disk.name,
managed_disk=data_disk.managed_disk,
)
for data_disk in getattr(storage_profile, "data_disks", [])
if data_disk
]
extensions = []
if getattr(vm, "resources", []):
extensions = [
VirtualMachineExtension(id=extension.id)
for extension in getattr(vm, "resources", [])
if extension
]
virtual_machines[subscription_name].update(
{
vm.id: VirtualMachine(
@@ -32,29 +59,19 @@ class VirtualMachines(AzureService):
storage_profile=(
StorageProfile(
os_disk=OSDisk(
name=vm.storage_profile.os_disk.name,
managed_disk=vm.storage_profile.os_disk.managed_disk,
name=getattr(os_disk, "name", None),
managed_disk=getattr(
os_disk, "managed_disk", None
),
),
data_disks=[
DataDisk(
lun=data_disk.lun,
name=data_disk.name,
managed_disk=data_disk.managed_disk,
)
for data_disk in getattr(
vm.storage_profile, "data_disks", []
)
],
data_disks=data_disks,
)
if getattr(vm, "storage_profile", None)
if storage_profile
else None
),
location=vm.location,
security_profile=getattr(vm, "security_profile", None),
extensions=[
VirtualMachineExtension(id=extension.id)
for extension in getattr(vm, "resources", [])
],
extensions=extensions,
)
}
)
@@ -110,13 +127,13 @@ class UefiSettings:
@dataclass
class SecurityProfile:
security_type: str
uefi_settings: UefiSettings
uefi_settings: Optional[UefiSettings]
@dataclass
class OSDisk:
name: str
managed_disk: bool
name: Optional[str]
managed_disk: Optional[bool]
@dataclass
@@ -128,7 +145,7 @@ class DataDisk:
@dataclass
class StorageProfile:
os_disk: OSDisk
os_disk: Optional[OSDisk]
data_disks: List[DataDisk]
@@ -41,6 +41,32 @@ def mock_vm_get_virtual_machines(_):
}
def mock_vm_get_virtual_machines_with_none(_):
return {
AZURE_SUBSCRIPTION_ID: {
"vm_id-1": VirtualMachine(
resource_id="/subscriptions/resource_id",
resource_name="VMWithNoneValues",
location="location",
security_profile=None,
extensions=None,
storage_profile=None,
),
"vm_id-2": VirtualMachine(
resource_id="/subscriptions/resource_id2",
resource_name="VMWithPartialNone",
location="location",
security_profile=None,
extensions=None,
storage_profile=StorageProfile(
os_disk=None,
data_disks=None,
),
),
}
}
def mock_vm_get_disks(_):
return {
AZURE_SUBSCRIPTION_ID: {
@@ -137,3 +163,22 @@ class Test_VirtualMachines_Service:
disks[AZURE_SUBSCRIPTION_ID]["disk_id-1"].encryption_type
== "EncryptionAtRestWithPlatformKey"
)
@patch(
"prowler.providers.azure.services.vm.vm_service.VirtualMachines._get_virtual_machines",
new=mock_vm_get_virtual_machines_with_none,
)
class Test_VirtualMachines_NoneCases:
def test_virtual_machine_with_none_storage_profile(self):
virtual_machines = VirtualMachines(set_mocked_azure_provider())
vm_1 = virtual_machines.virtual_machines[AZURE_SUBSCRIPTION_ID]["vm_id-1"]
assert vm_1.storage_profile is None
assert vm_1.resource_name == "VMWithNoneValues"
def test_virtual_machine_with_partial_none_storage_profile(self):
virtual_machines = VirtualMachines(set_mocked_azure_provider())
vm_2 = virtual_machines.virtual_machines[AZURE_SUBSCRIPTION_ID]["vm_id-2"]
assert vm_2.storage_profile.os_disk is None
assert vm_2.storage_profile.data_disks is None
assert vm_2.resource_name == "VMWithPartialNone"