diff --git a/prowler/CHANGELOG.md b/prowler/CHANGELOG.md index 06599427e9..1fc1375568 100644 --- a/prowler/CHANGELOG.md +++ b/prowler/CHANGELOG.md @@ -47,6 +47,7 @@ All notable changes to the **Prowler SDK** are documented in this file. ### Fixed - Consolidate Azure Storage file service properties to the account level, improving the accuracy of the `storage_ensure_file_shares_soft_delete_is_enabled` check [(#8087)](https://github.com/prowler-cloud/prowler/pull/8087) +- Migrate Azure VM service and managed disk logic to Pydantic models for better serialization and type safety, and update all related tests to use the new models and fix UUID handling [(#https://github.com/prowler-cloud/prowler/pull/8151)](https://github.com/prowler-cloud/prowler/pull/https://github.com/prowler-cloud/prowler/pull/8151) ### Changed - Reworked `S3.test_connection` to match the AwsProvider logic [(#8088)](https://github.com/prowler-cloud/prowler/pull/8088) diff --git a/prowler/providers/azure/services/vm/vm_ensure_using_managed_disks/vm_ensure_using_managed_disks.py b/prowler/providers/azure/services/vm/vm_ensure_using_managed_disks/vm_ensure_using_managed_disks.py index 22735fa9b7..e771fc80e5 100644 --- a/prowler/providers/azure/services/vm/vm_ensure_using_managed_disks/vm_ensure_using_managed_disks.py +++ b/prowler/providers/azure/services/vm/vm_ensure_using_managed_disks/vm_ensure_using_managed_disks.py @@ -15,17 +15,17 @@ class vm_ensure_using_managed_disks(Check): using_managed_disks = ( True - if vm.storage_profile - and getattr(vm.storage_profile, "os_disk", False) - and getattr(vm.storage_profile.os_disk, "managed_disk", False) + if getattr( + getattr(getattr(vm, "storage_profile", None), "os_disk", None), + "managed_disk", + None, + ) else False ) - if using_managed_disks and getattr( - vm.storage_profile, "data_disks", False - ): + if using_managed_disks and getattr(vm, "storage_profile", None): for data_disk in vm.storage_profile.data_disks: - if not getattr(data_disk, "managed_disk", False): + if not getattr(data_disk, "managed_disk", None): using_managed_disks = False break diff --git a/prowler/providers/azure/services/vm/vm_service.py b/prowler/providers/azure/services/vm/vm_service.py index d5d0dde5c7..a30379d9bc 100644 --- a/prowler/providers/azure/services/vm/vm_service.py +++ b/prowler/providers/azure/services/vm/vm_service.py @@ -1,7 +1,9 @@ from dataclasses import dataclass +from enum import Enum from typing import List, Optional from azure.mgmt.compute import ComputeManagementClient +from pydantic import BaseModel from prowler.lib.logger import logger from prowler.providers.azure.azure_provider import AzureProvider @@ -30,24 +32,33 @@ class VirtualMachines(AzureService): 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, + managed_disk=ManagedDiskParameters( + id=( + getattr( + getattr(data_disk, "managed_disk", None), + "id", + None, + ) + if data_disk.managed_disk + else None + ) + ), ) 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", []) + for extension in vm.resources if extension ] @@ -60,8 +71,17 @@ class VirtualMachines(AzureService): StorageProfile( os_disk=OSDisk( name=getattr(os_disk, "name", None), - managed_disk=getattr( - os_disk, "managed_disk", None + operating_system_type=getattr( + os_disk, "os_type", None + ), + managed_disk=ManagedDiskParameters( + id=getattr( + getattr( + os_disk, "managed_disk", None + ), + "id", + None, + ) ), ), data_disks=data_disks, @@ -130,32 +150,37 @@ class SecurityProfile: uefi_settings: Optional[UefiSettings] -@dataclass -class OSDisk: - name: Optional[str] - managed_disk: Optional[bool] +class OperatingSystemType(Enum): + WINDOWS = "Windows" + LINUX = "Linux" -@dataclass -class DataDisk: +class ManagedDiskParameters(BaseModel): + id: str + + +class OSDisk(BaseModel): + name: str + operating_system_type: OperatingSystemType + managed_disk: Optional[ManagedDiskParameters] + + +class DataDisk(BaseModel): lun: int name: str - managed_disk: bool + managed_disk: Optional[ManagedDiskParameters] -@dataclass -class StorageProfile: +class StorageProfile(BaseModel): os_disk: Optional[OSDisk] data_disks: List[DataDisk] -@dataclass -class VirtualMachineExtension: +class VirtualMachineExtension(BaseModel): id: str -@dataclass -class VirtualMachine: +class VirtualMachine(BaseModel): resource_id: str resource_name: str location: str @@ -164,8 +189,7 @@ class VirtualMachine: storage_profile: Optional[StorageProfile] = None -@dataclass -class Disk: +class Disk(BaseModel): resource_id: str resource_name: str vms_attached: list[str] diff --git a/tests/providers/azure/services/vm/vm_ensure_attached_disks_encrypted_with_cmk/vm_ensure_attached_disks_encrypted_with_cmk_test.py b/tests/providers/azure/services/vm/vm_ensure_attached_disks_encrypted_with_cmk/vm_ensure_attached_disks_encrypted_with_cmk_test.py index ccf6e8f175..1eb8da64c4 100644 --- a/tests/providers/azure/services/vm/vm_ensure_attached_disks_encrypted_with_cmk/vm_ensure_attached_disks_encrypted_with_cmk_test.py +++ b/tests/providers/azure/services/vm/vm_ensure_attached_disks_encrypted_with_cmk/vm_ensure_attached_disks_encrypted_with_cmk_test.py @@ -54,15 +54,15 @@ class Test_vm_ensure_attached_disks_encrypted_with_cmk: assert len(result) == 0 def test_vm_subscription_one_disk_attached_encrypt_pk(self): - disk_id = uuid4() - resource_id = uuid4() + disk_id = str(uuid4()) + resource_id = str(uuid4()) vm_client = mock.MagicMock vm_client.disks = { AZURE_SUBSCRIPTION_ID: { disk_id: Disk( resource_id=resource_id, resource_name="test-disk", - vms_attached=[uuid4()], + vms_attached=[str(uuid4())], encryption_type="EncryptionAtRestWithPlatformKey", location="location", ) @@ -97,15 +97,15 @@ class Test_vm_ensure_attached_disks_encrypted_with_cmk: ) def test_vm_subscription_one_disk_attached_encrypt_cmk(self): - disk_id = uuid4() - resource_id = uuid4() + disk_id = str(uuid4()) + resource_id = str(uuid4()) vm_client = mock.MagicMock vm_client.disks = { AZURE_SUBSCRIPTION_ID: { disk_id: Disk( resource_id=resource_id, resource_name="test-disk", - vms_attached=[uuid4()], + vms_attached=[str(uuid4())], encryption_type="EncryptionAtRestWithCustomerKey", location="location", ) @@ -140,24 +140,24 @@ class Test_vm_ensure_attached_disks_encrypted_with_cmk: ) def test_vm_subscription_two_disk_attached_encrypt_cmk_and_pk(self): - disk_id_1 = uuid4() - resource_id_1 = uuid4() - disk_id_2 = uuid4() - resource_id_2 = uuid4() + disk_id_1 = str(uuid4()) + resource_id_1 = str(uuid4()) + disk_id_2 = str(uuid4()) + resource_id_2 = str(uuid4()) vm_client = mock.MagicMock vm_client.disks = { AZURE_SUBSCRIPTION_ID: { disk_id_1: Disk( resource_id=resource_id_1, resource_name="test-disk", - vms_attached=[uuid4()], + vms_attached=[str(uuid4())], encryption_type="EncryptionAtRestWithPlatformKey", location="location", ), disk_id_2: Disk( resource_id=resource_id_2, resource_name="test-disk-2", - vms_attached=[uuid4(), uuid4()], + vms_attached=[str(uuid4()), str(uuid4())], encryption_type="EncryptionAtRestWithCustomerKey", location="location2", ), @@ -200,8 +200,8 @@ class Test_vm_ensure_attached_disks_encrypted_with_cmk: ) def test_vm_unattached_disk_encrypt_cmk(self): - disk_id = uuid4() - resource_id = uuid4() + disk_id = str(uuid4()) + resource_id = str(uuid4()) vm_client = mock.MagicMock vm_client.disks = { AZURE_SUBSCRIPTION_ID: { @@ -209,8 +209,8 @@ class Test_vm_ensure_attached_disks_encrypted_with_cmk: resource_id=resource_id, resource_name="test-disk", vms_attached=[], - location="location", encryption_type="EncryptionAtRestWithCustomerKey", + location="location", ) } } diff --git a/tests/providers/azure/services/vm/vm_ensure_unattached_disks_encrypted_with_cmk/vm_ensure_unattached_disks_encrypted_with_cmk_test.py b/tests/providers/azure/services/vm/vm_ensure_unattached_disks_encrypted_with_cmk/vm_ensure_unattached_disks_encrypted_with_cmk_test.py index 44ec78b4ed..78d7920666 100644 --- a/tests/providers/azure/services/vm/vm_ensure_unattached_disks_encrypted_with_cmk/vm_ensure_unattached_disks_encrypted_with_cmk_test.py +++ b/tests/providers/azure/services/vm/vm_ensure_unattached_disks_encrypted_with_cmk/vm_ensure_unattached_disks_encrypted_with_cmk_test.py @@ -54,8 +54,8 @@ class Test_vm_ensure_unattached_disks_encrypted_with_cmk: assert len(result) == 0 def test_vm_one_unattached_disk_encrypt_pk(self): - disk_id = uuid4() - resource_id = uuid4() + disk_id = str(uuid4()) + resource_id = str(uuid4()) vm_client = mock.MagicMock vm_client.disks = { AZURE_SUBSCRIPTION_ID: { @@ -97,8 +97,8 @@ class Test_vm_ensure_unattached_disks_encrypted_with_cmk: ) def test_vm_one_unattached_disk_encrypt_cmk(self): - disk_id = uuid4() - resource_id = uuid4() + disk_id = str(uuid4()) + resource_id = str(uuid4()) vm_client = mock.MagicMock vm_client.disks = { AZURE_SUBSCRIPTION_ID: { @@ -140,10 +140,10 @@ class Test_vm_ensure_unattached_disks_encrypted_with_cmk: ) def test_vm_subscription_two_unattached_disk_encrypt_cmk_and_pk(self): - disk_id_1 = uuid4() - resource_id_1 = uuid4() - disk_id_2 = uuid4() - resource_id_2 = uuid4() + disk_id_1 = str(uuid4()) + resource_id_1 = str(uuid4()) + disk_id_2 = str(uuid4()) + resource_id_2 = str(uuid4()) vm_client = mock.MagicMock vm_client.disks = { AZURE_SUBSCRIPTION_ID: { @@ -200,17 +200,17 @@ class Test_vm_ensure_unattached_disks_encrypted_with_cmk: ) def test_vm_attached_disk_encrypt_cmk(self): - disk_id = uuid4() - resource_id = uuid4() + disk_id = str(uuid4()) + resource_id = str(uuid4()) vm_client = mock.MagicMock vm_client.disks = { AZURE_SUBSCRIPTION_ID: { disk_id: Disk( resource_id=resource_id, resource_name="test-disk", - location="location", - vms_attached=[uuid4()], + vms_attached=[str(uuid4())], encryption_type="EncryptionAtRestWithCustomerKey", + location="location", ) } } diff --git a/tests/providers/azure/services/vm/vm_ensure_using_managed_disks/vm_ensure_using_managed_disks_test.py b/tests/providers/azure/services/vm/vm_ensure_using_managed_disks/vm_ensure_using_managed_disks_test.py index dbc04dd53f..b73075337e 100644 --- a/tests/providers/azure/services/vm/vm_ensure_using_managed_disks/vm_ensure_using_managed_disks_test.py +++ b/tests/providers/azure/services/vm/vm_ensure_using_managed_disks/vm_ensure_using_managed_disks_test.py @@ -1,7 +1,15 @@ from unittest import mock from uuid import uuid4 -from prowler.providers.azure.services.vm.vm_service import VirtualMachine +from prowler.providers.azure.services.vm.vm_service import ( + DataDisk, + ManagedDiskParameters, + OSDisk, + SecurityProfile, + StorageProfile, + UefiSettings, + VirtualMachine, +) from tests.providers.azure.azure_fixtures import ( AZURE_SUBSCRIPTION_ID, set_mocked_azure_provider, @@ -62,18 +70,19 @@ class Test_vm_ensure_using_managed_disks: resource_id=vm_id, resource_name="VMTest", location="location", - security_profile=mock.MagicMock( + security_profile=SecurityProfile( security_type="TrustedLaunch", - uefi_settings=mock.MagicMock( + uefi_settings=UefiSettings( secure_boot_enabled=True, v_tpm_enabled=True, ), ), extensions=[], - storage_profile=mock.MagicMock( - os_disk=mock.MagicMock( - create_option="FromImage", - managed_disk=mock.MagicMock(id="managed_disk_id"), + storage_profile=StorageProfile( + os_disk=OSDisk( + name="os_disk_name", + operating_system_type="Linux", + managed_disk=ManagedDiskParameters(id="managed_disk_id"), ), data_disks=[], ), @@ -117,17 +126,18 @@ class Test_vm_ensure_using_managed_disks: resource_id=vm_id, resource_name="VMTest", location="location", - security_profile=mock.MagicMock( + security_profile=SecurityProfile( security_type="TrustedLaunch", - uefi_settings=mock.MagicMock( + uefi_settings=UefiSettings( secure_boot_enabled=True, v_tpm_enabled=True, ), ), extensions=[], - storage_profile=mock.MagicMock( - os_disk=mock.MagicMock( - create_option="FromImage", + storage_profile=StorageProfile( + os_disk=OSDisk( + name="os_disk_name", + operating_system_type="Linux", managed_disk=None, ), data_disks=[], @@ -172,20 +182,23 @@ class Test_vm_ensure_using_managed_disks: resource_id=vm_id, resource_name="VMTest", location="location", - security_profile=mock.MagicMock( + security_profile=SecurityProfile( security_type="TrustedLaunch", - uefi_settings=mock.MagicMock( + uefi_settings=UefiSettings( secure_boot_enabled=True, v_tpm_enabled=True, ), ), extensions=[], - storage_profile=mock.MagicMock( - os_disk=mock.MagicMock( - create_option="FromImage", - managed_disk=mock.MagicMock(id="managed_disk_id"), + storage_profile=StorageProfile( + os_disk=OSDisk( + name="os_disk_name", + operating_system_type="Linux", + managed_disk=ManagedDiskParameters(id="managed_disk_id"), ), - data_disks=[mock.MagicMock(managed_disk=None)], + data_disks=[ + DataDisk(lun=0, name="data_disk_1", managed_disk=None) + ], ), ) } diff --git a/tests/providers/azure/services/vm/vm_service_test.py b/tests/providers/azure/services/vm/vm_service_test.py index 6e0f53e2e5..78aaccda2a 100644 --- a/tests/providers/azure/services/vm/vm_service_test.py +++ b/tests/providers/azure/services/vm/vm_service_test.py @@ -1,10 +1,13 @@ -from unittest import mock from unittest.mock import patch -from azure.mgmt.compute.models import ManagedDiskParameters, OSDisk, StorageProfile - from prowler.providers.azure.services.vm.vm_service import ( Disk, + ManagedDiskParameters, + OperatingSystemType, + OSDisk, + SecurityProfile, + StorageProfile, + UefiSettings, VirtualMachine, VirtualMachines, ) @@ -21,9 +24,9 @@ def mock_vm_get_virtual_machines(_): resource_id="/subscriptions/resource_id", resource_name="VMTest", location="location", - security_profile=mock.MagicMock( + security_profile=SecurityProfile( security_type="TrustedLaunch", - uefi_settings=mock.MagicMock( + uefi_settings=UefiSettings( secure_boot_enabled=True, v_tpm_enabled=True, ), @@ -31,7 +34,8 @@ def mock_vm_get_virtual_machines(_): extensions=[], storage_profile=StorageProfile( os_disk=OSDisk( - create_option="FromImage", + name="os_disk_name", + operating_system_type=OperatingSystemType.LINUX, managed_disk=ManagedDiskParameters(id="managed_disk_id"), ), data_disks=[], @@ -49,7 +53,7 @@ def mock_vm_get_virtual_machines_with_none(_): resource_name="VMWithNoneValues", location="location", security_profile=None, - extensions=None, + extensions=[], storage_profile=None, ), "vm_id-2": VirtualMachine( @@ -57,10 +61,10 @@ def mock_vm_get_virtual_machines_with_none(_): resource_name="VMWithPartialNone", location="location", security_profile=None, - extensions=None, + extensions=[], storage_profile=StorageProfile( os_disk=None, - data_disks=None, + data_disks=[], ), ), } @@ -180,5 +184,5 @@ class Test_VirtualMachines_NoneCases: 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.storage_profile.data_disks == [] assert vm_2.resource_name == "VMWithPartialNone" diff --git a/tests/providers/azure/services/vm/vm_trusted_launch_enabled/vm_trusted_launch_enabled_test.py b/tests/providers/azure/services/vm/vm_trusted_launch_enabled/vm_trusted_launch_enabled_test.py index 05e96a5c1d..83ab63acce 100644 --- a/tests/providers/azure/services/vm/vm_trusted_launch_enabled/vm_trusted_launch_enabled_test.py +++ b/tests/providers/azure/services/vm/vm_trusted_launch_enabled/vm_trusted_launch_enabled_test.py @@ -1,6 +1,14 @@ from unittest import mock from uuid import uuid4 +from prowler.providers.azure.services.vm.vm_service import ( + ManagedDiskParameters, + OSDisk, + SecurityProfile, + StorageProfile, + UefiSettings, + VirtualMachine, +) from tests.providers.azure.azure_fixtures import ( AZURE_SUBSCRIPTION_ID, set_mocked_azure_provider, @@ -63,7 +71,6 @@ class Test_vm_trusted_launch_enabled: new=vm_client, ), ): - from prowler.providers.azure.services.vm.vm_service import VirtualMachine from prowler.providers.azure.services.vm.vm_trusted_launch_enabled.vm_trusted_launch_enabled import ( vm_trusted_launch_enabled, ) @@ -74,18 +81,21 @@ class Test_vm_trusted_launch_enabled: resource_id=vm_id, resource_name="VMTest", location="location", - security_profile=mock.MagicMock( + security_profile=SecurityProfile( security_type="TrustedLaunch", - uefi_settings=mock.MagicMock( + uefi_settings=UefiSettings( secure_boot_enabled=True, v_tpm_enabled=True, ), ), extensions=[], - storage_profile=mock.MagicMock( - os_disk=mock.MagicMock( - create_option="FromImage", - managed_disk=mock.MagicMock(id="managed_disk_id"), + storage_profile=StorageProfile( + os_disk=OSDisk( + name="os_disk_name", + operating_system_type="Linux", + managed_disk=ManagedDiskParameters( + id="managed_disk_id" + ), ), data_disks=[], ), @@ -117,7 +127,6 @@ class Test_vm_trusted_launch_enabled: new=vm_client, ), ): - from prowler.providers.azure.services.vm.vm_service import VirtualMachine from prowler.providers.azure.services.vm.vm_trusted_launch_enabled.vm_trusted_launch_enabled import ( vm_trusted_launch_enabled, ) @@ -128,18 +137,21 @@ class Test_vm_trusted_launch_enabled: resource_id=vm_id, resource_name="VMTest", location="location", - security_profile=mock.MagicMock( + security_profile=SecurityProfile( security_type="TrustedLaunch", - uefi_settings=mock.MagicMock( + uefi_settings=UefiSettings( secure_boot_enabled=False, v_tpm_enabled=False, ), ), extensions=[], - storage_profile=mock.MagicMock( - os_disk=mock.MagicMock( - create_option="FromImage", - managed_disk=mock.MagicMock(id="managed_disk_id"), + storage_profile=StorageProfile( + os_disk=OSDisk( + name="os_disk_name", + operating_system_type="Linux", + managed_disk=ManagedDiskParameters( + id="managed_disk_id" + ), ), data_disks=[], ), @@ -172,7 +184,6 @@ class Test_vm_trusted_launch_enabled: new=vm_client, ), ): - from prowler.providers.azure.services.vm.vm_service import VirtualMachine from prowler.providers.azure.services.vm.vm_trusted_launch_enabled.vm_trusted_launch_enabled import ( vm_trusted_launch_enabled, ) @@ -185,10 +196,13 @@ class Test_vm_trusted_launch_enabled: location="location", security_profile=None, extensions=[], - storage_profile=mock.MagicMock( - os_disk=mock.MagicMock( - create_option="FromImage", - managed_disk=mock.MagicMock(id="managed_disk_id"), + storage_profile=StorageProfile( + os_disk=OSDisk( + name="os_disk_name", + operating_system_type="Linux", + managed_disk=ManagedDiskParameters( + id="managed_disk_id" + ), ), data_disks=[], ),