From 5d6ed640f082172acbd6b29bb74228f010d765f9 Mon Sep 17 00:00:00 2001
From: Daniel Barranquero <74871504+danibarranqueroo@users.noreply.github.com>
Date: Tue, 25 Mar 2025 12:25:15 +0100
Subject: [PATCH] fix(vm): handle `Nonetype is not iterable` for extensions
 (#7360)

Co-authored-by: Sergio Garcia <hello@mistercloudsec.com>
---
 .../providers/azure/services/vm/vm_service.py | 59 ++++++++++++-------
 .../azure/services/vm/vm_service_test.py      | 45 ++++++++++++++
 2 files changed, 83 insertions(+), 21 deletions(-)

diff --git a/prowler/providers/azure/services/vm/vm_service.py b/prowler/providers/azure/services/vm/vm_service.py
index 2ef3c9e30a..d5d0dde5c7 100644
--- a/prowler/providers/azure/services/vm/vm_service.py
+++ b/prowler/providers/azure/services/vm/vm_service.py
@@ -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]
 
 
diff --git a/tests/providers/azure/services/vm/vm_service_test.py b/tests/providers/azure/services/vm/vm_service_test.py
index 5bd225e1bf..6e0f53e2e5 100644
--- a/tests/providers/azure/services/vm/vm_service_test.py
+++ b/tests/providers/azure/services/vm/vm_service_test.py
@@ -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"