Get lifetime info for NVMe devices (#6056)

* Get lifetime info for NVMe devices

* Fix lint and test issues

* Update tests/dbus_service_mocks/udisks2_manager.py

Co-authored-by: Copilot <[email protected]>

---------

Co-authored-by: Stefan Agner <[email protected]>
Co-authored-by: Copilot <[email protected]>

Author: 3 people

supervisor/dbus/const.py

@@ -32,6 +32,7 @@
 DBUS_IFACE_IP4CONFIG = "org.freedesktop.NetworkManager.IP4Config"
 DBUS_IFACE_IP6CONFIG = "org.freedesktop.NetworkManager.IP6Config"
 DBUS_IFACE_NM = "org.freedesktop.NetworkManager"
+DBUS_IFACE_NVME_CONTROLLER = "org.freedesktop.UDisks2.NVMe.Controller"
 DBUS_IFACE_PARTITION = "org.freedesktop.UDisks2.Partition"
 DBUS_IFACE_PARTITION_TABLE = "org.freedesktop.UDisks2.PartitionTable"
 DBUS_IFACE_RAUC_INSTALLER = "de.pengutronix.rauc.Installer"
@@ -87,6 +88,7 @@
 DBUS_ATTR_CURRENT_DEVICE = "CurrentDevice"
 DBUS_ATTR_CURRENT_DNS_SERVER = "CurrentDNSServer"
 DBUS_ATTR_CURRENT_DNS_SERVER_EX = "CurrentDNSServerEx"
+DBUS_ATTR_CONTROLLER_ID = "ControllerID"
 DBUS_ATTR_DEFAULT = "Default"
 DBUS_ATTR_DEPLOYMENT = "Deployment"
 DBUS_ATTR_DESCRIPTION = "Description"
@@ -111,6 +113,7 @@
 DBUS_ATTR_EJECTABLE = "Ejectable"
 DBUS_ATTR_FALLBACK_DNS = "FallbackDNS"
 DBUS_ATTR_FALLBACK_DNS_EX = "FallbackDNSEx"
+DBUS_ATTR_FGUID = "FGUID"
 DBUS_ATTR_FINISH_TIMESTAMP = "FinishTimestamp"
 DBUS_ATTR_FIRMWARE_TIMESTAMP_MONOTONIC = "FirmwareTimestampMonotonic"
 DBUS_ATTR_FREQUENCY = "Frequency"
@@ -147,6 +150,7 @@
 DBUS_ATTR_NTP = "NTP"
 DBUS_ATTR_NTPSYNCHRONIZED = "NTPSynchronized"
 DBUS_ATTR_NUMBER = "Number"
+DBUS_ATTR_NVME_REVISION = "NVMeRevision"
 DBUS_ATTR_OFFSET = "Offset"
 DBUS_ATTR_OPERATING_SYSTEM_PRETTY_NAME = "OperatingSystemPrettyName"
 DBUS_ATTR_OPERATION = "Operation"
@@ -161,15 +165,24 @@
 DBUS_ATTR_RESOLV_CONF_MODE = "ResolvConfMode"
 DBUS_ATTR_REVISION = "Revision"
 DBUS_ATTR_RCMANAGER = "RcManager"
+DBUS_ATTR_SANITIZE_PERCENT_REMAINING = "SanitizePercentRemaining"
+DBUS_ATTR_SANITIZE_STATUS = "SanitizeStatus"
 DBUS_ATTR_SEAT = "Seat"
 DBUS_ATTR_SERIAL = "Serial"
 DBUS_ATTR_SIZE = "Size"
+DBUS_ATTR_SMART_CRITICAL_WARNING = "SmartCriticalWarning"
+DBUS_ATTR_SMART_POWER_ON_HOURS = "SmartPowerOnHours"
+DBUS_ATTR_SMART_SELFTEST_PERCENT_REMAINING = "SmartSelftestPercentRemaining"
+DBUS_ATTR_SMART_SELFTEST_STATUS = "SmartSelftestStatus"
+DBUS_ATTR_SMART_TEMPERATURE = "SmartTemperature"
+DBUS_ATTR_SMART_UPDATED = "SmartUpdated"
 DBUS_ATTR_SSID = "Ssid"
 DBUS_ATTR_STATE = "State"
 DBUS_ATTR_STATE_FLAGS = "StateFlags"
 DBUS_ATTR_STATIC_HOSTNAME = "StaticHostname"
 DBUS_ATTR_STATIC_OPERATING_SYSTEM_CPE_NAME = "OperatingSystemCPEName"
 DBUS_ATTR_STRENGTH = "Strength"
+DBUS_ATTR_SUBSYSTEM_NQN = "SubsystemNQN"
 DBUS_ATTR_SUPPORTED_FILESYSTEMS = "SupportedFilesystems"
 DBUS_ATTR_SYMLINKS = "Symlinks"
 DBUS_ATTR_SWAP_SIZE = "SwapSize"
@@ -180,6 +193,7 @@
 DBUS_ATTR_TIMEZONE = "Timezone"
 DBUS_ATTR_TRANSACTION_STATISTICS = "TransactionStatistics"
 DBUS_ATTR_TYPE = "Type"
+DBUS_ATTR_UNALLOCATED_CAPACITY = "UnallocatedCapacity"
 DBUS_ATTR_USER_LED = "UserLED"
 DBUS_ATTR_USERSPACE_TIMESTAMP_MONOTONIC = "UserspaceTimestampMonotonic"
 DBUS_ATTR_UUID_UPPERCASE = "UUID"

supervisor/dbus/udisks2/__init__.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import logging
+from pathlib import Path
 from typing import Any
 
 from awesomeversion import AwesomeVersion
@@ -132,7 +133,10 @@ async def update(self, changed: dict[str, Any] | None = None) -> None:
                 for drive in drives
             }
 
-            # Update existing drives
+            # For existing drives, need to check their type and call update
+            await asyncio.gather(
+                *[self._drives[path].check_type() for path in unchanged_drives]
+            )
             await asyncio.gather(
                 *[self._drives[path].update() for path in unchanged_drives]
             )
@@ -160,20 +164,33 @@ def drives(self) -> list[UDisks2Drive]:
         return list(self._drives.values())
 
     @dbus_connected
-    def get_drive(self, drive_path: str) -> UDisks2Drive:
+    def get_drive(self, object_path: str) -> UDisks2Drive:
         """Get additional info on drive from object path."""
-        if drive_path not in self._drives:
-            raise DBusObjectError(f"Drive {drive_path} not found")
+        if object_path not in self._drives:
+            raise DBusObjectError(f"Drive {object_path} not found")
 
-        return self._drives[drive_path]
+        return self._drives[object_path]
 
     @dbus_connected
-    def get_block_device(self, device_path: str) -> UDisks2Block:
+    def get_block_device(self, object_path: str) -> UDisks2Block:
         """Get additional info on block device from object path."""
-        if device_path not in self._block_devices:
-            raise DBusObjectError(f"Block device {device_path} not found")
+        if object_path not in self._block_devices:
+            raise DBusObjectError(f"Block device {object_path} not found")
 
-        return self._block_devices[device_path]
+        return self._block_devices[object_path]
+
+    @dbus_connected
+    def get_block_device_by_path(self, device_path: Path) -> UDisks2Block:
+        """Get additional info on block device from device path.
+
+        Uses cache only. Use `resolve_device` to force a call for fresh data.
+        """
+        for device in self._block_devices.values():
+            if device.device == device_path:
+                return device
+        raise DBusObjectError(
+            f"Block device not found with device path {device_path.as_posix()}"
+        )
 
     @dbus_connected
     async def resolve_device(self, devspec: DeviceSpecification) -> list[UDisks2Block]:

supervisor/dbus/udisks2/drive.py

@@ -1,6 +1,7 @@
 """Interface to UDisks2 Drive over D-Bus."""
 
 from datetime import UTC, datetime
+from typing import Any
 
 from dbus_fast.aio import MessageBus
 
@@ -18,11 +19,13 @@
     DBUS_ATTR_VENDOR,
     DBUS_ATTR_WWN,
     DBUS_IFACE_DRIVE,
+    DBUS_IFACE_NVME_CONTROLLER,
     DBUS_NAME_UDISKS2,
 )
 from ..interface import DBusInterfaceProxy, dbus_property
 from ..utils import dbus_connected
 from .const import UDISKS2_DEFAULT_OPTIONS
+from .nvme_controller import UDisks2NVMeController
 
 
 class UDisks2Drive(DBusInterfaceProxy):
@@ -35,11 +38,18 @@ class UDisks2Drive(DBusInterfaceProxy):
     bus_name: str = DBUS_NAME_UDISKS2
     properties_interface: str = DBUS_IFACE_DRIVE
 
+    _nvme_controller: UDisks2NVMeController | None = None
+
     def __init__(self, object_path: str) -> None:
         """Initialize object."""
         self._object_path = object_path
         super().__init__()
 
+    async def connect(self, bus: MessageBus) -> None:
+        """Connect to bus."""
+        await super().connect(bus)
+        await self._reload_interfaces()
+
     @staticmethod
     async def new(object_path: str, bus: MessageBus) -> "UDisks2Drive":
         """Create and connect object."""
@@ -52,6 +62,11 @@ def object_path(self) -> str:
         """Object path for dbus object."""
         return self._object_path
 
+    @property
+    def nvme_controller(self) -> UDisks2NVMeController | None:
+        """NVMe controller interface if drive is one."""
+        return self._nvme_controller
+
     @property
     @dbus_property
     def vendor(self) -> str:
@@ -130,3 +145,40 @@ def ejectable(self) -> bool:
     async def eject(self) -> None:
         """Eject media from drive."""
         await self.connected_dbus.Drive.call("eject", UDISKS2_DEFAULT_OPTIONS)
+
+    @dbus_connected
+    async def update(self, changed: dict[str, Any] | None = None) -> None:
+        """Update properties via D-Bus."""
+        await super().update(changed)
+
+        if not changed and self.nvme_controller:
+            await self.nvme_controller.update()
+
+    @dbus_connected
+    async def check_type(self) -> None:
+        """Check if type of drive has changed and adjust interfaces if so."""
+        introspection = await self.connected_dbus.introspect()
+        interfaces = {intr.name for intr in introspection.interfaces}
+
+        # If interfaces changed, update the proxy from introspection and reload interfaces
+        if interfaces != set(self.connected_dbus.proxies.keys()):
+            await self.connected_dbus.init_proxy(introspection=introspection)
+            await self._reload_interfaces()
+
+    @dbus_connected
+    async def _reload_interfaces(self) -> None:
+        """Reload interfaces from introspection as necessary."""
+        # Check if drive is an nvme controller
+        if (
+            not self.nvme_controller
+            and DBUS_IFACE_NVME_CONTROLLER in self.connected_dbus.proxies
+        ):
+            self._nvme_controller = UDisks2NVMeController(self.object_path)
+            await self._nvme_controller.initialize(self.connected_dbus)
+
+        elif (
+            self.nvme_controller
+            and DBUS_IFACE_NVME_CONTROLLER not in self.connected_dbus.proxies
+        ):
+            self.nvme_controller.stop_sync_property_changes()
+            self._nvme_controller = None