selftest/selftest.py

#!/usr/bin/env python3

# --------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See LICENSE in the project root for license information.
# --------------------------------------------------------------------------

"""Azure VM utilities self-tests script."""
import argparse
import glob
import json
import logging
import os
import re
import subprocess
import time
import urllib.error
import urllib.request
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Literal, Optional

logger = logging.getLogger("selftest")

# pylint: disable=broad-exception-caught
# pylint: disable=line-too-long
# pylint: disable=too-many-lines
# pylint: disable=too-many-instance-attributes
# pylint: disable=too-many-lines
# pylint: disable=too-many-locals

SYS_CLASS_NET = "/sys/class/net"


@dataclass(eq=True, repr=True)
class SkuConfig:
    """VM sku-specific configuration related to disks."""

    vm_size: str
    vm_size_type: Literal["arm64", "x64"] = "x64"
    nvme_controller_toggle_supported: bool = (
        False  # whether the sku supports NVMe controller toggle (Eb[d]s_v5)
    )
    nvme_only: bool = False  # NVMe-only skus (v6+)
    nvme_id_enabled_local: bool = False  # whether the sku supports NVMe ID locally
    nvme_id_enabled_remote: bool = False  # whether the sku supports NVMe ID remotely
    nvme_local_disk_count: int = 0
    nvme_local_disk_size_gib: int = 0
    temp_disk_size_gib: int = 0  # SCSI temp/resource disk size in GiB


@dataclass(eq=True, repr=True)
class V6SkuConfig(SkuConfig):
    """V6 VM sku-specific configuration related to disks."""

    nvme_only: bool = True
    nvme_id_enabled_local: bool = True
    nvme_id_enabled_remote: bool = False


def gb_to_gib(size_gb: int) -> int:
    """Roughly convert GB to GiB as sizes are documented in both ways."""
    return int(size_gb * (1000**3) / (1024**3))


SKU_CONFIGS = {
    "Standard_B2ts_v2": SkuConfig(vm_size="Standard_B2ts_v2"),
    # "Standard_D2s_v3": SkuConfig(vm_size="Standard_D2s_v3", temp_disk_size_gib=16),
    "Standard_D2s_v4": SkuConfig(vm_size="Standard_D2s_v4"),
    "Standard_D2ds_v4": SkuConfig(vm_size="Standard_D2ds_v4", temp_disk_size_gib=75),
    "Standard_D2s_v5": SkuConfig(vm_size="Standard_D2s_v5"),
    "Standard_D2ds_v5": SkuConfig(vm_size="Standard_D2ds_v5", temp_disk_size_gib=75),
    "Standard_D2ads_v5": SkuConfig(vm_size="Standard_D2ads_v5", temp_disk_size_gib=75),
    "Standard_D16ads_v5": SkuConfig(
        vm_size="Standard_D16ads_v5", temp_disk_size_gib=600
    ),
    "Standard_L8s_v2": SkuConfig(
        vm_size="Standard_L8s_v2",
        temp_disk_size_gib=80,
        nvme_local_disk_count=1,
        nvme_local_disk_size_gib=gb_to_gib(1920),
    ),
    "Standard_L8s_v3": SkuConfig(
        vm_size="Standard_L8s_v3",
        temp_disk_size_gib=80,
        nvme_local_disk_count=1,
        nvme_local_disk_size_gib=gb_to_gib(1920),
    ),
    "Standard_L80s_v3": SkuConfig(
        vm_size="Standard_L80s_v3",
        nvme_controller_toggle_supported=True,
        temp_disk_size_gib=800,
        nvme_local_disk_count=10,
        nvme_local_disk_size_gib=gb_to_gib(1920),
    ),
    "Standard_E2bs_v5": SkuConfig(
        vm_size="Standard_E2bs_v5", nvme_controller_toggle_supported=True
    ),
    "Standard_E2bds_v5": SkuConfig(
        vm_size="Standard_E2bds_v5",
        nvme_controller_toggle_supported=True,
        temp_disk_size_gib=75,
    ),
    "Standard_D2s_v6": V6SkuConfig(vm_size="Standard_D2s_v6"),
    "Standard_D2ds_v6": V6SkuConfig(
        vm_size="Standard_D2ds_v6",
        nvme_local_disk_count=1,
        nvme_local_disk_size_gib=110,
    ),
    "Standard_D16ds_v6": V6SkuConfig(
        vm_size="Standard_D16ds_v6",
        nvme_local_disk_count=2,
        nvme_local_disk_size_gib=440,
    ),
    "Standard_D32ds_v6": V6SkuConfig(
        vm_size="Standard_D32ds_v6",
        nvme_local_disk_count=4,
        nvme_local_disk_size_gib=440,
    ),
    "Standard_D2as_v6": V6SkuConfig(vm_size="Standard_D2as_v6"),
    "Standard_D2ads_v6": V6SkuConfig(
        vm_size="Standard_D2ads_v6",
        nvme_local_disk_count=1,
        nvme_local_disk_size_gib=110,
    ),
    "Standard_D16ads_v6": V6SkuConfig(
        vm_size="Standard_D16ads_v6",
        nvme_local_disk_count=2,
        nvme_local_disk_size_gib=440,
    ),
    "Standard_D32ads_v6": V6SkuConfig(
        vm_size="Standard_D32ads_v6",
        nvme_local_disk_count=4,
        nvme_local_disk_size_gib=440,
    ),
    "Standard_D2pls_v5": SkuConfig(
        vm_size="Standard_D2pls_v5",
        vm_size_type="arm64",
    ),
    "Standard_D2plds_v5": SkuConfig(
        vm_size="Standard_D2plds_v5",
        vm_size_type="arm64",
        temp_disk_size_gib=75,
    ),
    "Standard_D8pls_v5": SkuConfig(
        vm_size="Standard_D8pls_v5",
        vm_size_type="arm64",
    ),
    "Standard_D8plds_v5": SkuConfig(
        vm_size="Standard_D8plds_v5",
        vm_size_type="arm64",
        temp_disk_size_gib=300,
    ),
    "Standard_D2pls_v6": SkuConfig(
        vm_size="Standard_D2pls_v6",
        vm_size_type="arm64",
    ),
    "Standard_D2plds_v6": SkuConfig(
        vm_size="Standard_D2plds_v6",
        vm_size_type="arm64",
        nvme_local_disk_count=1,
        nvme_local_disk_size_gib=110,
    ),
    "Standard_D16pls_v6": SkuConfig(
        vm_size="Standard_D16pls_v6",
        vm_size_type="arm64",
    ),
    "Standard_D16plds_v6": SkuConfig(
        vm_size="Standard_D16plds_v6",
        vm_size_type="arm64",
        nvme_local_disk_count=2,
        nvme_local_disk_size_gib=440,
    ),
}


def device_sort(devices: List[str]) -> List[str]:
    """Natural sort for devices."""

    def natural_sort_key(s: str):
        # Natural sort by turning a string into a list of string and number chunks.
        # e.g. "nvme0n10" -> ["nvme", 0, "n", 10]
        return [
            int(text) if text.isdigit() else text for text in re.split("([0-9]+)", s)
        ]

    return sorted(devices, key=natural_sort_key)


def get_disk_size_gb(disk_path: str) -> int:
    """Get the size of the disk in GB."""
    try:
        proc = subprocess.run(
            ["lsblk", "-b", "-n", "-o", "SIZE", "-d", disk_path],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            check=True,
        )
        logger.debug("lsblk output: %r", proc)
        size_bytes = int(proc.stdout.strip())
        size_gib = size_bytes // (1000**3)
        return size_gib
    except subprocess.CalledProcessError as error:
        logger.error("error while fetching disk size: %r", error)
        raise
    except FileNotFoundError:
        logger.error("lsblk command not found")
        raise


def get_disk_size_gib(disk_path: str) -> int:
    """Get the size of the disk in GiB."""
    try:
        proc = subprocess.run(
            ["lsblk", "-b", "-n", "-o", "SIZE", "-d", disk_path],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            check=True,
        )
        logger.debug("lsblk output: %r", proc)
        size_bytes = int(proc.stdout.strip())
        size_gib = size_bytes // (1024**3)
        return size_gib
    except subprocess.CalledProcessError as error:
        logger.error("error while fetching disk size: %r", error)
        raise
    except FileNotFoundError:
        logger.error("lsblk command not found")
        raise


def get_imds_metadata() -> Dict:
    """Fetch IMDS metadata using urllib."""
    url = "http://169.254.169.254/metadata/instance?api-version=2021-02-01"
    headers = {"Metadata": "true"}

    req = urllib.request.Request(url, headers=headers)

    last_error = None
    deadline = time.time() + 300
    while time.time() < deadline:
        try:
            with urllib.request.urlopen(req, timeout=60) as response:
                if response.status != 200:
                    raise urllib.error.HTTPError(
                        url,
                        response.status,
                        "Failed to fetch metadata",
                        response.headers,
                        None,
                    )
                metadata = json.load(response)
                logger.debug("fetched IMDS metadata: %r", metadata)
                return metadata
        except urllib.error.URLError as error:
            last_error = error
            logger.error("error fetching IMDS metadata: %r", error)
            time.sleep(1)

    raise RuntimeError(f"failed to fetch IMDS metadata: {last_error}")


def get_local_nvme_disks() -> List[str]:
    """Get all local NVMe disks."""
    local_disk_controllers = get_nvme_controllers_with_model(
        "Microsoft NVMe Direct Disk"
    )
    local_disk_controllers_v2 = get_nvme_controllers_with_model(
        "Microsoft NVMe Direct Disk v2"
    )

    return device_sort(
        [
            namespace
            for controller in local_disk_controllers + local_disk_controllers_v2
            for namespace in get_nvme_namespace_devices(controller)
        ]
    )


def get_remote_nvme_disks() -> List[str]:
    """Get all remote NVMe disks."""
    remote_disk_controllers = get_nvme_controllers_with_model(
        "MSFT NVMe Accelerator v1.0"
    )

    assert (
        len(remote_disk_controllers) <= 1
    ), f"unexpected number of remote controllers {remote_disk_controllers}"
    return device_sort(
        [
            namespace
            for controller in remote_disk_controllers
            for namespace in get_nvme_namespace_devices(controller)
        ]
    )


def get_nvme_controllers_with_model(model: str) -> List[str]:
    """Get a list of all NVMe controllers with the specified model."""
    nvme_controllers = []
    nvme_path = "/sys/class/nvme"

    for controller in glob.glob(os.path.join(nvme_path, "nvme*")):
        logger.debug("checking controller: %s", controller)
        model_path = os.path.join(controller, "model")
        try:
            with open(model_path, "r", encoding="utf-8") as file:
                controller_model = file.read().strip()
                logger.debug("controller: %s model: %s", controller, controller_model)
                if controller_model == model:
                    controller_name = controller.split("/")[-1]
                    nvme_controllers.append(controller_name)
        except FileNotFoundError:
            logger.debug("model file not found: %s", model_path)
            continue

    return device_sort(nvme_controllers)


def get_nvme_namespace_devices_with_model(model: str) -> List[str]:
    """Get all NVMe namespace devices for a given NVMe controller model."""
    controllers = get_nvme_controllers_with_model(model)
    logger.debug("controllers found for model=%s: %r", model, controllers)
    return device_sort(
        [
            namespace
            for controller in controllers
            for namespace in get_nvme_namespace_devices(controller)
        ]
    )


def get_nvme_namespace_devices(controller: str) -> List[str]:
    """Get all NVMe namespace devices for a given NVMe controller."""
    namespace_devices = []
    controller_name = controller.split("/")[-1]
    nvme_path = f"/sys/class/nvme/{controller_name}"

    logger.debug("checking namespaces under %s", nvme_path)
    for namespace in glob.glob(os.path.join(nvme_path, "nvme*")):
        logger.debug("checking namespace device: %s", namespace)
        if os.path.isdir(namespace):
            device_name = namespace.split("/")[-1]
            namespace_devices.append(device_name)

    return device_sort(namespace_devices)


def get_root_block_device() -> str:
    """Get the root block device using findmnt."""
    try:
        proc = subprocess.run(
            ["findmnt", "-n", "-o", "SOURCE", "/"],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            check=True,
        )
        logger.debug("findmnt output: %r", proc)
        return proc.stdout.strip()
    except subprocess.CalledProcessError as error:
        logger.error("error while fetching root block device: %r", error)
        raise
    except FileNotFoundError:
        logger.error("findmnt command not found")
        raise


def get_scsi_resource_disk() -> Optional[str]:
    """Get the SCSI resource disk device."""
    paths = [
        # cloud-init udev rules
        "/dev/disk/cloud/azure_resource",
        # gen2
        "/dev/disk/by-path/acpi-VMBUS:00-vmbus-f8b3781a1e824818a1c363d806ec15bb-lun-1",
        # gen1
        "/dev/disk/by-path/acpi-VMBUS:01-vmbus-000000000001*-lun-0",
    ]

    for path in paths:
        if "*" in path:
            matched_paths = glob.glob(path)
            for matched_path in matched_paths:
                resolved_path = os.path.realpath(matched_path)
                if os.path.exists(resolved_path):
                    return resolved_path.split("/")[-1]
        else:
            if os.path.exists(path):
                resolved_path = os.path.realpath(path)
                if os.path.exists(resolved_path):
                    return resolved_path.split("/")[-1]

    logger.info("no SCSI resource disk found")
    return None


DEV_DISK_AZURE_RESOURCE = "/dev/disk/azure/resource"


@dataclass(eq=True, repr=True)
class DiskInfo:
    """Information about different types of disks present."""

    root_device: str
    dev_disk_azure_links: List[str] = field(default_factory=list)
    dev_disk_azure_resource_disk: Optional[str] = None  # resolved path
    dev_disk_azure_resource_disk_size_gib: int = 0
    nvme_local_disk_size_gib: int = 0
    nvme_local_disks_v1: List[str] = field(default_factory=list)
    nvme_local_disks_v2: List[str] = field(default_factory=list)
    nvme_local_disks: List[str] = field(default_factory=list)
    nvme_remote_data_disks: List[str] = field(default_factory=list)
    nvme_remote_disks: List[str] = field(default_factory=list)
    nvme_remote_os_disk: Optional[str] = None
    root_device_is_nvme: bool = False
    scsi_resource_disk: Optional[str] = None
    scsi_resource_disk_size_gib: int = 0

    @classmethod
    def gather(cls) -> "DiskInfo":
        """Gather disk information and return an instance of DiskInfo."""
        dev_disk_azure_links = device_sort(
            [
                link
                for link in glob.glob(
                    os.path.join("/dev/disk/azure", "**"), recursive=True
                )
                if os.path.islink(link)
            ]
        )

        dev_disk_azure_resource_disk = None
        dev_disk_azure_resource_disk_size_gib = 0
        if os.path.exists(DEV_DISK_AZURE_RESOURCE):
            dev_disk_azure_resource_disk = os.path.realpath(DEV_DISK_AZURE_RESOURCE)
            dev_disk_azure_resource_disk_size_gib = get_disk_size_gib(
                dev_disk_azure_resource_disk
            )

        nvme_local_disks_v1 = get_nvme_namespace_devices_with_model(
            "Microsoft NVMe Direct Disk"
        )
        nvme_local_disks_v2 = get_nvme_namespace_devices_with_model(
            "Microsoft NVMe Direct Disk v2"
        )
        nvme_local_disks = device_sort(nvme_local_disks_v1 + nvme_local_disks_v2)
        nvme_local_disk_size_gib = 0
        if nvme_local_disks:
            nvme_local_disk_size_gib = min(
                get_disk_size_gib(f"/dev/{disk}") for disk in nvme_local_disks
            )
            local_disk_max_size = max(
                get_disk_size_gib(f"/dev/{disk}") for disk in nvme_local_disks
            )
            assert (
                nvme_local_disk_size_gib == local_disk_max_size
            ), f"local disk size mismatch: {nvme_local_disk_size_gib} != {local_disk_max_size} for {nvme_local_disks}"

        nvme_remote_disks = get_remote_nvme_disks()
        if nvme_remote_disks:
            nvme_remote_os_disk = nvme_remote_disks.pop(0)
            nvme_remote_data_disks = nvme_remote_disks
        else:
            nvme_remote_os_disk = None
            nvme_remote_data_disks = []

        root_device = get_root_block_device()
        root_device_is_nvme = root_device.startswith("/dev/nvme")
        root_device = root_device.split("/")[-1]

        scsi_resource_disk = get_scsi_resource_disk()
        scsi_resource_disk_size_gib = (
            get_disk_size_gib(f"/dev/{scsi_resource_disk}") if scsi_resource_disk else 0
        )

        disk_info = cls(
            dev_disk_azure_links=dev_disk_azure_links,
            dev_disk_azure_resource_disk=dev_disk_azure_resource_disk,
            dev_disk_azure_resource_disk_size_gib=dev_disk_azure_resource_disk_size_gib,
            nvme_local_disk_size_gib=nvme_local_disk_size_gib,
            nvme_local_disks_v1=nvme_local_disks_v1,
            nvme_local_disks_v2=nvme_local_disks_v2,
            nvme_local_disks=nvme_local_disks,
            nvme_remote_os_disk=nvme_remote_os_disk,
            nvme_remote_data_disks=nvme_remote_data_disks,
            nvme_remote_disks=nvme_remote_disks,
            root_device=root_device,
            root_device_is_nvme=root_device_is_nvme,
            scsi_resource_disk=scsi_resource_disk,
            scsi_resource_disk_size_gib=scsi_resource_disk_size_gib,
        )

        logger.info("disks info: %r", disk_info)
        return disk_info


@dataclass
class AzureNvmeIdDevice:
    """Azure NVMe ID device."""

    device: str
    model: Optional[str]
    nvme_id: str
    type: Optional[str]
    index: Optional[int]
    lun: Optional[int]
    name: Optional[str]
    extra: Dict[str, str]


@dataclass(repr=True)
class AzureNvmeIdInfo:
    """Azure NVMe ID."""

    azure_nvme_id_stdout: str
    azure_nvme_id_stderr: str
    azure_nvme_id_returncode: int
    azure_nvme_id_disks: Dict[str, AzureNvmeIdDevice]

    azure_nvme_id_json_stdout: str
    azure_nvme_id_json_stderr: str
    azure_nvme_id_json_returncode: int
    azure_nvme_id_json_disks: Dict[str, AzureNvmeIdDevice]

    azure_nvme_id_help_stdout: str
    azure_nvme_id_help_stderr: str
    azure_nvme_id_help_returncode: int

    azure_nvme_id_version_stdout: str
    azure_nvme_id_version_stderr: str
    azure_nvme_id_version_returncode: int
    azure_nvme_id_version: str

    azure_nvme_id_zzz_stdout: str
    azure_nvme_id_zzz_stderr: str
    azure_nvme_id_zzz_returncode: int

    def _validate_azure_nvme_disks(
        self, azure_nvme_id_disks: Dict[str, AzureNvmeIdDevice], disk_info: DiskInfo
    ) -> None:
        disk_cfg: Optional[AzureNvmeIdDevice] = None
        for device_name, disk_cfg in azure_nvme_id_disks.items():
            assert f"/dev/{device_name}" == disk_cfg.device
            assert disk_cfg.device.startswith(
                "/dev/nvme"
            ), f"unexpected device: {disk_cfg}"

        for device_name in disk_info.nvme_local_disks_v2:
            assert (
                device_name in azure_nvme_id_disks
            ), f"missing azure-nvme-id for {device_name}"
            disk_cfg = azure_nvme_id_disks.get(device_name)
            assert disk_cfg, f"failed to find azure-nvme-id for {device_name}"
            assert disk_cfg.type == "local", f"unexpected local disk type {disk_cfg}"
            assert disk_cfg.name, f"unexpected local disk name {disk_cfg}"
            assert disk_cfg.index, f"unexpected local disk index {disk_cfg}"
            assert disk_cfg.lun is None, f"unexpected local disk lun {disk_cfg}"
            assert disk_cfg.nvme_id, f"unexpected local disk id {disk_cfg}"
            assert not disk_cfg.extra, f"unexpected local disk extra {disk_cfg}"

        for device_name in disk_info.nvme_local_disks_v1:
            assert (
                device_name in azure_nvme_id_disks
            ), f"missing azure-nvme-id for {device_name}"
            disk_cfg = azure_nvme_id_disks.get(device_name)
            assert disk_cfg, f"failed to find azure-nvme-id for {device_name}"
            assert disk_cfg.type == "local", f"unexpected disk type {disk_cfg}"
            assert not disk_cfg.name, f"unexpected disk name {disk_cfg}"
            assert not disk_cfg.index, f"unexpected disk index {disk_cfg}"
            assert disk_cfg.lun is None, f"unexpected local disk lun {disk_cfg}"
            assert disk_cfg.nvme_id, f"unexpected disk id {disk_cfg}"
            assert not disk_cfg.extra, f"unexpected disk extra {disk_cfg}"

        for device_name in disk_info.nvme_remote_disks:
            assert (
                device_name in azure_nvme_id_disks
            ), f"missing azure-nvme-id for {device_name}"
            disk_cfg = azure_nvme_id_disks.get(device_name)
            assert disk_cfg, f"failed to find azure-nvme-id for {device_name}"
            assert disk_cfg.type in (
                "os",
                "data",
            ), f"unexpected remote disk type {disk_cfg}"
            if disk_cfg.type == "data":
                assert (
                    disk_cfg.lun is not None and disk_cfg.lun >= 0
                ), f"unexpected remote disk index {disk_cfg}"
            else:
                assert disk_cfg.lun is None, f"unexpected remote disk index {disk_cfg}"
            assert not disk_cfg.name, f"unexpected remote disk name {disk_cfg}"
            assert disk_cfg.nvme_id, f"unexpected remote disk id {disk_cfg}"
            assert not disk_cfg.extra, f"unexpected remote disk extra {disk_cfg}"

        logger.info("validate_azure_nvme_disks OK: %r", self.azure_nvme_id_disks)

    def validate_azure_nvme_id(self, disk_info: DiskInfo) -> None:
        """Validate azure-nvme-id outputs."""
        assert self.azure_nvme_id_returncode == 0, "azure-nvme-id failed"
        if not os.path.exists("/sys/class/nvme"):
            assert (
                self.azure_nvme_id_stderr
                == "no NVMe devices in /sys/class/nvme: No such file or directory\n"
            ), f"unexpected azure-nvme-id stderr without /sys/class/nvme: {self.azure_nvme_id_stderr}"
        else:
            assert (
                self.azure_nvme_id_stderr == ""
            ), f"unexpected azure-nvme-id stderr: {self.azure_nvme_id_stderr}"

        self._validate_azure_nvme_disks(self.azure_nvme_id_disks, disk_info)
        logger.info("validate_azure_nvmve_id OK: %r", self.azure_nvme_id_stdout)

    def validate_azure_nvme_id_help(self) -> None:
        """Validate azure-nvme-id --help outputs."""
        assert self.azure_nvme_id_help_returncode == 0, "azure-nvme-id --help failed"
        assert (
            self.azure_nvme_id_help_stderr == ""
        ), f"unexpected azure-nvme-id --help stderr: {self.azure_nvme_id_help_stderr!r}"
        assert (
            self.azure_nvme_id_help_stdout
            and self.azure_nvme_id_help_stdout.startswith("Usage: azure-nvme-id ")
        ), "unexpected azure-nvme-id --help stdout: {self.azure_nvme_id_help_stdout!r}"

        logger.info(
            "validate_azure_nvme_id_help OK: %r", self.azure_nvme_id_help_stdout
        )

    def validate_azure_nvme_id_json(self, disk_info: DiskInfo) -> None:
        """Validate azure-nvme-id --format json outputs."""
        assert self.azure_nvme_id_json_returncode == 0, "azure-nvme-id failed"
        if not os.path.exists("/sys/class/nvme"):
            assert (
                self.azure_nvme_id_json_stderr
                == "no NVMe devices in /sys/class/nvme: No such file or directory\n"
            ), f"unexpected azure-nvme-id stderr without /sys/class/nvme: {self.azure_nvme_id_json_stderr}"
        else:
            assert (
                self.azure_nvme_id_json_stderr == ""
            ), f"unexpected azure-nvme-id stderr: {self.azure_nvme_id_json_stderr}"

        self._validate_azure_nvme_disks(self.azure_nvme_id_disks, disk_info)

        assert all(
            disk.model
            in (
                "MSFT NVMe Accelerator v1.0",
                "Microsoft NVMe Direct Disk",
                "Microsoft NVMe Direct Disk v2",
            )
            for disk in self.azure_nvme_id_json_disks.values()
        ), "missing model in azure-nvme-id --format json"
        logger.info(
            "validate_azure_nvmve_id_json OK: %r", self.azure_nvme_id_json_stdout
        )

    def validate_azure_nvme_id_version(self) -> None:
        """Validate azure-nvme-id --version outputs."""
        assert (
            self.azure_nvme_id_version_returncode == 0
        ), "azure-nvme-id --version failed"
        assert (
            self.azure_nvme_id_version_stderr == ""
        ), f"unexpected azure-nvme-id stderr: {self.azure_nvme_id_stderr}"
        assert self.azure_nvme_id_version_stdout, "missing azure-nvme-id version stdout"
        assert re.match(
            r"azure-nvme-id [0v]\.*", self.azure_nvme_id_version_stdout.strip()
        ), f"unexpected azure-nvme-id version stdout: {self.azure_nvme_id_version_stdout}"
        assert re.match(
            r"[0v]\.*", self.azure_nvme_id_version
        ), f"unexpected azure-nvme-id version: {self.azure_nvme_id_version}"

        logger.info("validate_azure_nvme_id_version OK: %s", self.azure_nvme_id_version)

    def validate_azure_nvme_id_zzz_invalid_arg(self) -> None:
        """Validate azure-nvme-id handles invalid arguments."""
        assert (
            self.azure_nvme_id_zzz_returncode == 1
        ), f"azure-nvme-id zzz rc={self.azure_nvme_id_zzz_returncode}"
        assert (
            self.azure_nvme_id_zzz_stderr == "invalid argument: zzz\n"
        ), f"unexpected azure-nvme-id zzz stderr: {self.azure_nvme_id_zzz_stderr!r}"
        assert (
            self.azure_nvme_id_zzz_stdout
            and self.azure_nvme_id_zzz_stdout.startswith("Usage: azure-nvme-id ")
        ), (f"unexpected azure-nvme-id zzz stdout: {self.azure_nvme_id_zzz_stdout!r}")

        logger.info(
            "validate_azure_nvme_id_invalid_arg OK: %r", self.azure_nvme_id_zzz_stdout
        )

    def validate(self, disk_info: DiskInfo) -> None:
        """Validate Azure NVMe ID output."""
        self.validate_azure_nvme_id_help()
        self.validate_azure_nvme_id_version()
        self.validate_azure_nvme_id_zzz_invalid_arg()
        self.validate_azure_nvme_id(disk_info)
        self.validate_azure_nvme_id_json(disk_info)

    @classmethod
    def gather(cls) -> "AzureNvmeIdInfo":
        """Gather Azure NVMe ID information."""
        proc = subprocess.run(["azure-nvme-id"], capture_output=True, check=False)
        azure_nvme_id_stdout = proc.stdout.decode("utf-8")
        azure_nvme_id_stderr = proc.stderr.decode("utf-8")
        azure_nvme_id_returncode = proc.returncode
        azure_nvme_id_disks = cls.parse_azure_nvme_id_output(azure_nvme_id_stdout)

        proc = subprocess.run(
            ["azure-nvme-id", "--format", "json"], capture_output=True, check=False
        )
        azure_nvme_id_json_stdout = proc.stdout.decode("utf-8")
        azure_nvme_id_json_stderr = proc.stderr.decode("utf-8")
        azure_nvme_id_json_returncode = proc.returncode
        azure_nvme_id_json_disks = cls.parse_azure_nvme_id_json_output(
            azure_nvme_id_json_stdout
        )

        proc = subprocess.run(
            ["azure-nvme-id", "--help"], capture_output=True, check=False
        )
        azure_nvme_id_help_stdout = proc.stdout.decode("utf-8")
        azure_nvme_id_help_stderr = proc.stderr.decode("utf-8")
        azure_nvme_id_help_returncode = proc.returncode

        proc = subprocess.run(
            ["azure-nvme-id", "--version"], capture_output=True, check=False
        )
        azure_nvme_id_version_stdout = proc.stdout.decode("utf-8")
        azure_nvme_id_version_stderr = proc.stderr.decode("utf-8")
        azure_nvme_id_version_returncode = proc.returncode
        azure_nvme_id_version = cls.parse_azure_nvme_id_version(
            azure_nvme_id_version_stdout
        )

        proc = subprocess.run(
            ["azure-nvme-id", "zzz"], capture_output=True, check=False
        )
        azure_nvme_id_zzz_stdout = proc.stdout.decode("utf-8")
        azure_nvme_id_zzz_stderr = proc.stderr.decode("utf-8")
        azure_nvme_id_zzz_returncode = proc.returncode

        azure_nvme_id_info = cls(
            azure_nvme_id_stdout=azure_nvme_id_stdout,
            azure_nvme_id_stderr=azure_nvme_id_stderr,
            azure_nvme_id_returncode=azure_nvme_id_returncode,
            azure_nvme_id_help_stdout=azure_nvme_id_help_stdout,
            azure_nvme_id_help_stderr=azure_nvme_id_help_stderr,
            azure_nvme_id_help_returncode=azure_nvme_id_help_returncode,
            azure_nvme_id_disks=azure_nvme_id_disks,
            azure_nvme_id_json_stdout=azure_nvme_id_json_stdout,
            azure_nvme_id_json_stderr=azure_nvme_id_json_stderr,
            azure_nvme_id_json_returncode=azure_nvme_id_json_returncode,
            azure_nvme_id_json_disks=azure_nvme_id_json_disks,
            azure_nvme_id_version_stdout=azure_nvme_id_version_stdout,
            azure_nvme_id_version_stderr=azure_nvme_id_version_stderr,
            azure_nvme_id_version_returncode=azure_nvme_id_version_returncode,
            azure_nvme_id_version=azure_nvme_id_version,
            azure_nvme_id_zzz_returncode=azure_nvme_id_zzz_returncode,
            azure_nvme_id_zzz_stdout=azure_nvme_id_zzz_stdout,
            azure_nvme_id_zzz_stderr=azure_nvme_id_zzz_stderr,
        )
        logger.info("azure-nvme-id info: %r", azure_nvme_id_info)
        return azure_nvme_id_info

    @staticmethod
    def parse_azure_nvme_id_json_output(output: str) -> Dict[str, AzureNvmeIdDevice]:
        """Parse azure-nvme-id --format json output.
        Example output:
        [
            {
                "path": "/dev/nvme0n33",
                "model": "MSFT NVMe Accelerator v1.0",
                "properties": {
                    "type": "data",
                    "lun": 31
                },
                "vs": ""
            },
            {
                "path": "/dev/nvme1n1",
                "model": "Microsoft NVMe Direct Disk v2",
                "properties": {
                    "type": "local",
                    "index": 1,
                    "name": "nvme-440G-1"
                },
                "vs": "type=local,index=1,name=nvme-440G-1"
            }
        ]
        """
        devices = {}

        for device in json.loads(output):
            device_path = device["path"]
            model = device["model"]
            properties = device["properties"]
            device_type = properties.pop("type", None)
            device_index = (
                int(properties.pop("index")) if "index" in properties else None
            )
            device_lun = int(properties.pop("lun")) if "lun" in properties else None
            device_name = properties.pop("name", None)
            azure_nvme_id_device = AzureNvmeIdDevice(
                device=device_path,
                model=model,
                nvme_id=",".join([f"{k}={v}" for k, v in properties.items()]),
                type=device_type,
                index=device_index,
                lun=device_lun,
                name=device_name,
                extra=properties,
            )

            key = device_path.split("/")[-1]
            devices[key] = azure_nvme_id_device

        return devices

    @staticmethod
    def parse_azure_nvme_id_output(output: str) -> Dict[str, AzureNvmeIdDevice]:
        """Parse azure-nvme-id output.

        Example output:
        /dev/nvme0n1: type=os
        /dev/nvme0n2: type=data,lun=0
        /dev/nvme0n3: type=data,lun=1
        /dev/nvme1n1: type=local,index=1,name=nvme-440G-1
        /dev/nvme2n1: type=local,index=2,name=nvme-440G-2
        /dev/nvme3n1:
        """
        devices = {}

        for line in output.splitlines():
            parts = line.strip().split(":", 1)
            if parts[-1] == "":
                parts.pop()

            device = parts[0].strip()
            if len(parts) == 2:
                nvme_id = parts[1].strip()
                properties = dict(kv.split("=", 1) for kv in nvme_id.split(","))
            elif len(parts) == 1:
                nvme_id = ""
                properties = {}
            else:
                raise ValueError(f"unexpected azure-nvme-id output: {line}")

            device_type = properties.pop("type", None)
            device_index = (
                int(properties.pop("index")) if "index" in properties else None
            )
            device_lun = int(properties.pop("lun")) if "lun" in properties else None
            device_name = properties.pop("name", None)
            azure_nvme_id_device = AzureNvmeIdDevice(
                device=device,
                model=None,
                nvme_id=nvme_id,
                type=device_type,
                index=device_index,
                lun=device_lun,
                name=device_name,
                extra=properties,
            )

            key = device.split("/")[-1]
            devices[key] = azure_nvme_id_device

        return devices

    @staticmethod
    def parse_azure_nvme_id_version(azure_nvme_id_version_output: str) -> str:
        """Parse azure-nvme-id version output and return version info."""
        parts = azure_nvme_id_version_output.strip().split(" ")
        assert (
            len(parts) == 2
        ), f"unexpected azure-nvme-id version output: {azure_nvme_id_version_output}"
        return parts[1]


@dataclass
class NetworkInterface:
    """Network interface."""

    name: str
    driver: str
    mac: str
    ipv4_addrs: List[str]
    udev_properties: Dict[str, str]


@dataclass(eq=True, repr=True)
class NetworkInfo:
    """Network information."""

    interfaces: Dict[str, NetworkInterface] = field(default_factory=dict)

    @classmethod
    def enumerate_interfaces(cls) -> Dict[str, NetworkInterface]:
        """Retrieve all Ethernet interfaces on the system."""
        interfaces: Dict[str, NetworkInterface] = {}
        interface_names = [
            interface
            for interface in os.listdir(SYS_CLASS_NET)
            if os.path.exists(os.path.join(SYS_CLASS_NET, interface, "device"))
        ]

        for interface_name in interface_names:
            sys_path = Path(SYS_CLASS_NET, interface_name)
            udev_properties = cls.query_udev_properties(interface_name)
            driver_path = Path(sys_path, "device", "driver")
            if not driver_path.is_symlink():
                logger.debug(
                    "ignoring interface %s without driver symlink", interface_name
                )
                continue

            link = os.readlink(driver_path)
            driver = os.path.basename(link)
            mac = (sys_path / "address").read_text().strip()
            ipv4_addrs = cls.get_ipv4_addresses(interface_name)
            interfaces[interface_name] = NetworkInterface(
                name=interface_name,
                driver=driver,
                mac=mac,
                ipv4_addrs=ipv4_addrs,
                udev_properties=udev_properties,
            )

        return interfaces

    @staticmethod
    def get_ipv4_addresses(interface_name: str) -> List[str]:
        """Get the IPv4 addresses of a given network interface using `ip addr`."""
        try:
            result = subprocess.run(
                ["ip", "-4", "addr", "show", interface_name],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                check=True,
            )

            ipv4_addresses = re.findall(r"inet (\d+\.\d+\.\d+\.\d+)", result.stdout)
            return ipv4_addresses
        except subprocess.CalledProcessError as error:
            logger.error("failed to get IPv4 address for %s: %r", interface_name, error)
            raise

    @staticmethod
    def query_udev_properties(interface_name: str) -> Dict[str, str]:
        """Query all udev properties for a given interface using udevadm."""
        try:
            result = subprocess.run(
                [
                    "udevadm",
                    "info",
                    "--query=property",
                    f"--path={SYS_CLASS_NET}/{interface_name}",
                ],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                check=True,
            )
            properties: Dict[str, str] = {}
            for line in result.stdout.splitlines():
                if "=" in line:
                    key, value = line.split("=", 1)
                    properties[key] = value
            return properties
        except subprocess.CalledProcessError as error:
            logger.error(
                "Failed to query udev properties for %s: %r", interface_name, error
            )
            return {}

    def _validate_interface(self, interface: NetworkInterface) -> None:
        """Ensure the required properties are set for hv_netvsc, mlx4, mlx5, and mana devices."""
        if interface.driver in ["mlx4_core", "mlx5_core", "mana"]:
            assert (
                interface.udev_properties.get("NM_UNMANAGED") == "1"
                and interface.udev_properties.get("AZURE_UNMANAGED_SRIOV") == "1"
                and interface.udev_properties.get("ID_NET_MANAGED_BY") == "unmanaged"
            ), f"missing required properties for network interface: {interface}"
        elif interface.driver == "hv_netvsc":
            assert (
                "AZURE_UNMANAGED_SRIOV" not in interface.udev_properties
            ), f"unexpected AZURE_UNMANAGED_SRIOV property: {interface}"
            assert (
                interface.udev_properties.get("ID_NET_MANAGED_BY") != "unmanaged"
            ), f"hv_netvsc interface should be managed: {interface}"

        mana_has_synthetic_netvsc = interface.driver == "mana" and any(
            i.driver == "hv_netvsc" and i.mac == interface.mac
            for i in self.interfaces.values()
        )
        if interface.driver == "hv_netvsc" or (
            interface.driver == "mana" and not mana_has_synthetic_netvsc
        ):
            assert interface.ipv4_addrs, f"missing IPv4 addresses for {interface}"
        else:
            assert (
                not interface.ipv4_addrs
            ), f"unexpected IPv4 addresses for {interface}"

        logger.info("validate_interface %s OK: %r", interface.name, interface)

    def validate(self) -> None:
        """Validate network configuration."""
        for _, interface in self.interfaces.items():
            self._validate_interface(interface)

    @classmethod
    def gather(cls) -> "NetworkInfo":
        """Gather networking information."""
        return NetworkInfo(interfaces=cls.enumerate_interfaces())


class AzureVmUtilsValidator:
    """Validate Azure VM utilities."""

    def __init__(
        self,
        *,
        skip_imds_validation: bool = False,
        skip_network_validation: bool = False,
        skip_symlink_validation: bool = False,
    ) -> None:
        self.azure_nvme_id_info = AzureNvmeIdInfo.gather()
        self.disk_info = DiskInfo.gather()
        self.net_info = NetworkInfo.gather()
        self.skip_imds_validation = skip_imds_validation
        self.skip_network_validation = skip_network_validation
        self.skip_symlink_validation = skip_symlink_validation

        try:
            self.imds_metadata = get_imds_metadata()
        except Exception as error:
            logger.error("failed to fetch IMDS metadata: %r", error)
            if not self.skip_imds_validation:
                raise
            self.imds_metadata = {}

        self.vm_size = self.imds_metadata.get("compute", {}).get("vmSize")
        self.sku_config = SKU_CONFIGS.get(self.vm_size)

        logger.info("sku config: %r", self.sku_config)

    def validate_dev_disk_azure_links_data(self) -> None:
        """Validate /dev/disk/azure/data links.

        All data disks should have by-lun if azure-vm-utils is installed.
        Future variants of remote disks will include by-name.
        """
        imds_data_disks = (
            self.imds_metadata.get("compute", {})
            .get("storageProfile", {})
            .get("dataDisks", [])
        )
        expected_data_disks = len(imds_data_disks)
        data_disks = [
            link
            for link in self.disk_info.dev_disk_azure_links
            if link.startswith("/dev/disk/azure/data/by-lun")
        ]
        if self.disk_info.nvme_remote_disks:
            assert len(data_disks) == len(
                self.disk_info.nvme_remote_data_disks
            ), f"unexpected number of data disks: {data_disks} configured={self.disk_info.nvme_remote_data_disks}"

        assert (
            len(data_disks) == expected_data_disks
        ), f"unexpected number of data disks: {data_disks} IMDS configured={imds_data_disks} (note that IMDS may not be accurate)"

        # Verify disk sizes match up with IMDS configuration.
        for imds_disk in imds_data_disks:
            lun = imds_disk.get("lun")
            # Disk size is actually reported in GiB not GB.
            expected_size_gib = int(imds_disk.get("diskSizeGB"))
            disk_path = f"/dev/disk/azure/data/by-lun/{lun}"
            actual_size_gib = get_disk_size_gib(disk_path)
            assert (
                actual_size_gib == expected_size_gib
            ), f"disk size mismatch for {disk_path}: expected {expected_size_gib} GiB, found {actual_size_gib} GiB"

        logger.info("validate_dev_disk_azure_links_data OK: %r", data_disks)

    def validate_dev_disk_azure_links_local(self) -> None:
        """Validate /dev/disk/azure/local links.

        All local disks should have by-serial if azure-vm-utils is installed.
        If NVMe id is supported, by-index and by-name will be available as well.
        """
        local_disks = sorted(
            [
                link
                for link in self.disk_info.dev_disk_azure_links
                if link.startswith("/dev/disk/azure/local")
            ]
        )

        for key in ["index", "name", "serial"]:
            local_disks_by_key = sorted(
                [
                    link
                    for link in self.disk_info.dev_disk_azure_links
                    if link.startswith(f"/dev/disk/azure/local/by-{key}")
                ]
            )
            if key == "serial":
                expected_count = len(self.disk_info.nvme_local_disks)
            else:
                expected_count = len(self.disk_info.nvme_local_disks_v2)

            assert (
                len(local_disks_by_key) == expected_count
            ), f"unexpected number of local disks by-{key}: {local_disks_by_key} (expected {expected_count})"
            assert (
                not self.sku_config
                or not self.sku_config.nvme_id_enabled_local
                or len(local_disks_by_key) == self.sku_config.nvme_local_disk_count
            ), f"unexpected number of local disks by sku for by-{key}: {local_disks_by_key} (expected {expected_count})"

            if key == "name":
                for disk in local_disks_by_key:
                    name = disk.split("/")[-1]
                    assert name.startswith(
                        "nvme-"
                    ), f"unexpected local disk name: {name}"
                    match = re.match(r"nvme-(\d+)G-(\d+)", name)
                    assert (
                        match
                    ), f"local disk name does not conform to expected pattern: {name}"
                    size, index = match.groups()
                    assert (
                        size.isdigit() and index.isdigit()
                    ), f"invalid size or index in local disk name: {name}"

                    # Cross-check by-index links with by-name links.
                    by_index_path = f"/dev/disk/azure/local/by-index/{index}"
                    assert os.path.realpath(by_index_path) == os.path.realpath(
                        disk
                    ), f"mismatch between by-index and by-name links: {by_index_path} != {disk}"

        logger.info("validate_dev_disk_azure_links_local OK: %r", local_disks)

    def validate_dev_disk_azure_links_os(self) -> None:
        """Validate /dev/disk/azure/os link."""
        os_disk = "/dev/disk/azure/os"
        assert os_disk in self.disk_info.dev_disk_azure_links, f"missing {os_disk}"

        logger.info("validate_dev_disk_azure_links_os OK: %r", os_disk)

    def validate_dev_disk_azure_links_resource(self) -> None:
        """Validate /dev/disk/azure/resource link."""
        resource_disk = "/dev/disk/azure/resource"
        expected = (self.sku_config and self.sku_config.temp_disk_size_gib) or bool(
            self.disk_info.scsi_resource_disk
        )
        if expected:
            assert (
                resource_disk in self.disk_info.dev_disk_azure_links
            ), f"missing {resource_disk}"
        else:
            assert (
                resource_disk not in self.disk_info.dev_disk_azure_links
            ), f"unexpected {resource_disk}"

        logger.info("validate_dev_disk_azure_links_resource OK: %r", resource_disk)

    def validate_networking(self) -> None:
        """Validate networking configuration."""
        self.net_info.validate()
        logger.info("validate_networking OK: %r", self.net_info)

    def validate_nvme_local_disks(self) -> None:
        """Validate NVMe local disks."""
        logger.info("validate_nvme_local_disks OK: %r", self.disk_info.nvme_local_disks)

    def validate_scsi_resource_disk(self) -> None:
        """Validate SCSI resource disk symlink and size."""
        assert (
            self.disk_info.scsi_resource_disk_size_gib
            == self.disk_info.dev_disk_azure_resource_disk_size_gib
        ), f"resource disk size mismatch: {self.disk_info}"
        if self.disk_info.scsi_resource_disk:
            assert (
                f"/dev/{self.disk_info.scsi_resource_disk}"
                == self.disk_info.dev_disk_azure_resource_disk
            ), f"unexpected resource disk path: {self.disk_info}"
        else:
            assert (
                self.disk_info.scsi_resource_disk is None
                and self.disk_info.dev_disk_azure_resource_disk is None
            ), f"unexpected resource disk path: {self.disk_info}"

        logger.info(
            "validate_scsi_resource_disk OK: /dev/disk/azure/resource => %s",
            self.disk_info.dev_disk_azure_resource_disk,
        )

    def validate_sku_config(self) -> None:
        """Validate SKU config."""
        if not self.sku_config:
            logger.warning(
                "validate_sku_config SKIPPED: no sku configuration for VM size %r",
                self.vm_size,
            )
            return

        assert (
            self.sku_config.vm_size == self.vm_size
        ), f"vm size mismatch: {self.sku_config.vm_size} != {self.vm_size}"
        assert (
            len(self.disk_info.nvme_local_disks)
            == self.sku_config.nvme_local_disk_count
        ), f"local disk count mismatch: {len(self.disk_info.nvme_local_disks)} != {self.sku_config.nvme_local_disk_count}"
        assert (
            self.disk_info.nvme_local_disk_size_gib
            == self.sku_config.nvme_local_disk_size_gib
        ), f"local disk size mismatch: {self.disk_info.nvme_local_disk_size_gib} != {self.sku_config.nvme_local_disk_size_gib}"
        assert (
            self.disk_info.scsi_resource_disk_size_gib
            == self.sku_config.temp_disk_size_gib
        ), f"temp disk size mismatch: {self.disk_info.scsi_resource_disk_size_gib} != {self.sku_config.temp_disk_size_gib}"
        assert (
            self.disk_info.dev_disk_azure_resource_disk_size_gib
            == self.sku_config.temp_disk_size_gib
        ), f"temp disk size mismatch: {self.disk_info.dev_disk_azure_resource_disk_size_gib} != {self.sku_config.temp_disk_size_gib}"

        logger.info("validate_sku_config OK: %r", self.sku_config)

    def validate(self) -> None:
        """Run validations."""
        self.azure_nvme_id_info.validate(self.disk_info)

        if self.skip_symlink_validation:
            logger.info("validate_dev_disk_azure_links_data SKIPPED")
            logger.info("validate_dev_disk_azure_links_local SKIPPED")
            logger.info("validate_dev_disk_azure_links_os SKIPPED")
            logger.info("validate_dev_disk_azure_links_resource SKIPPED")
            logger.info("validate_scsi_resource_disk SKIPPED")
        else:
            self.validate_dev_disk_azure_links_data()
            self.validate_dev_disk_azure_links_local()
            self.validate_dev_disk_azure_links_os()
            self.validate_dev_disk_azure_links_resource()
            self.validate_scsi_resource_disk()

        if self.skip_network_validation:
            logger.info("validate_networking SKIPPED")
        else:
            self.validate_networking()

        self.validate_sku_config()

        logger.info("success!")


def main() -> None:
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description="Azure VM utilities self-tests script."
    )
    parser.add_argument(
        "--debug",
        action="store_true",
        help="Enable debug logging",
    )
    parser.add_argument(
        "--skip-imds-validation",
        action="store_true",
        help="Skip imds validation (allow for running tests outside Azure VM)",
    )
    parser.add_argument(
        "--skip-network-validation",
        action="store_true",
        help="Skip network validation (allow for running test without reboot after install)",
    )
    parser.add_argument(
        "--skip-symlink-validation",
        action="store_true",
        help="Skip symlink validation (allow for running test without reboot after install)",
    )
    args = parser.parse_args()

    if args.debug:
        logging.basicConfig(format="[%(asctime)s] %(message)s", level=logging.DEBUG)
    else:
        logging.basicConfig(format="[%(asctime)s] %(message)s", level=logging.INFO)

    validator = AzureVmUtilsValidator(
        skip_imds_validation=args.skip_imds_validation,
        skip_network_validation=args.skip_network_validation,
        skip_symlink_validation=args.skip_symlink_validation,
    )
    validator.validate()


if __name__ == "__main__":
    main()