ben.carrasco/prowler
1
0
Fork 0
mirror of https://github.com/prowler-cloud/prowler.git synced 2026-07-23 12:31:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(azure): use Pydantic models in VM service and fix managed disk logic (#8151)

Browse Source
This commit is contained in:
Rubén De la Torre Vico
2025-07-02 10:23:51 +02:00
committed by GitHub
parent c719d705e0
commit f78a29206c
8 changed files with 159 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files
prowler/CHANGELOG.md
+1
View File
@@ -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)
prowler/providers/azure/services/vm/vm_ensure_using_managed_disks/vm_ensure_using_managed_disks.py
+7 -7
View File
@@ -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
prowler/providers/azure/services/vm/vm_service.py
+45 -21
View File
@@ -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]
tests/providers/azure/services/vm/vm_ensure_attached_disks_encrypted_with_cmk/vm_ensure_attached_disks_encrypted_with_cmk_test.py
+15 -15
View File
@@ -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",
)
}
}
tests/providers/azure/services/vm/vm_ensure_unattached_disks_encrypted_with_cmk/vm_ensure_unattached_disks_encrypted_with_cmk_test.py
+12 -12
View File
@@ -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",
)
}
}
tests/providers/azure/services/vm/vm_ensure_using_managed_disks/vm_ensure_using_managed_disks_test.py
+32 -19
View File
@@ -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)
],
),
)
}
tests/providers/azure/services/vm/vm_service_test.py
+14 -10
View File
@@ -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"
tests/providers/azure/services/vm/vm_trusted_launch_enabled/vm_trusted_launch_enabled_test.py
+33 -19
View File
@@ -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=[],
),