Add SMART status to disk infos

debian/control

@@ -17,7 +17,9 @@ Depends: python3-all (>= 3.11),
  , python3-ldap, python3-zeroconf (>= 0.47), python3-lexicon,
  , python3-cryptography, python3-jwt, python3-passlib, python3-magic
  , python-is-python3, python3-pydantic, python3-email-validator
+ , python3-sortedcollections
  , udisks2, udisks2-bcache, udisks2-btrfs, udisks2-lvm2, udisks2-zram
+ , smartmontools
  , python3-zmq
  , nginx, nginx-extras (>=1.22)
  , apt, apt-transport-https, apt-utils, aptitude, dirmngr

src/disks.py

@@ -1,44 +1,77 @@
+import enum
+
 import dbus
+from sdbus import sd_bus_open_system, DbusInterfaceCommon
 from yunohost.utils.system import binary_to_human
 
+__all__ = ["info", "do_list"]
 
-__all__ = ["info", "list"]
+from yunohost.utils.udisks2_interfaces import Udisks2Manager
 
 
 UDISK_DRIVE_PATH = "/org/freedesktop/UDisks2/drives/"
 UDISK_DRIVE_IFC = "org.freedesktop.UDisks2.Drive"
-
-
-def _query_udisks() -> list[tuple[str, dict]]:
-    bus = dbus.SystemBus()
-    manager = dbus.Interface(
-        bus.get_object("org.freedesktop.UDisks2", "/org/freedesktop/UDisks2"),
-        "org.freedesktop.DBus.ObjectManager",
-    )
-
-    return sorted(
-        (
-            (name.removeprefix(UDISK_DRIVE_PATH), dev[UDISK_DRIVE_IFC])
-            for name, dev in manager.GetManagedObjects().items()
-            if name.startswith(UDISK_DRIVE_PATH)
-        ),
-        key=lambda item: item[1]["SortKey"],
-    )
+UDISK_DRIVE_ATA_IFC = "org.freedesktop.UDisks2.Drive.Ata"
+UDISK_DRIVE_NVME_IFC = "org.freedesktop.UDisks2.Manager.NVMe"
+
+
+class DiskState(enum.StrEnum):
+    @staticmethod
+    def _generate_next_value_(name, start, count, last_values):
+        return name.upper()
+
+    TEST_SANE = enum.auto()
+    TEST_IN_PROGRESS = enum.auto()
+    TEST_FAILURE = enum.auto()
+    """Test result is unavailable for some reason"""
+    TEST_UNKNOWN = enum.auto()
+    """No data. Either because not run or previously aborted"""
+    UNKNOWN = enum.auto()
+
+    @staticmethod
+    def parse(status: str | None):
+        """
+        https://storaged.org/doc/udisks2-api/latest/gdbus-org.freedesktop.UDisks2.Drive.Ata.html#gdbus-property-org-freedesktop-UDisks2-Drive-Ata.SmartSelftestStatus
+        https://storaged.org/doc/udisks2-api/latest/gdbus-org.freedesktop.UDisks2.NVMe.Controller.html#gdbus-property-org-freedesktop-UDisks2-NVMe-Controller.SmartSelftestStatus
+        """
+        if not status:
+            return DiskState.UNKNOWN
+
+        match status:
+            case "success":
+                return DiskState.TEST_SANE
+            case "inprogress":
+                return DiskState.TEST_IN_PROGRESS
+            case (
+            "fatal_error"
+            | "known_seg_fail"
+            | "unknown_seg_fail"
+            | "fatal"
+            | "error_unknown"
+            | "error_electrical"
+            | "error_servo"
+            | "error_read"
+            | "error_handling"
+            ):
+                return DiskState.TEST_FAILURE
+            case _:
+                return DiskState.UNKNOWN
 
 
 def _disk_infos(name: str, drive: dict, human_readable=False):
     result = {
         "name": name,
-        "model": drive["Model"],
-        "serial": drive["Serial"],
-        "removable": bool(drive["MediaRemovable"]),
-        "size": binary_to_human(drive["Size"]) if human_readable else drive["Size"],
+        "model": drive["model"],
+        "serial": drive["serial"],
+        "removable": bool(drive["media_removable"]),
+        "size": binary_to_human(drive["size"]) if human_readable else drive["size"],
+        "smartStatus": DiskState.parse(drive.get("smart_selftest_status")),
     }
 
-    if connection_bus := drive["ConnectionBus"]:
-        result["connection_bus"] = connection_bus
+    if "connection_bus" in drive:
+        result["connectionBus"] = drive["connection_bus"]
 
-    if (rotation_rate := drive["RotationRate"]) == -1:
+    if (rotation_rate := drive["rotation_rate"]) == -1:
         result.update(
             {
                 "type": "HDD",
@@ -58,25 +91,24 @@
     return result
 
 
-def list(with_info=False, human_readable=False):
+def do_list(with_info=False, human_readable=False):
+    bus = sd_bus_open_system()
+    disks = Udisks2Manager(bus).get_disks()
     if not with_info:
-        return [name for name, _ in _query_udisks()]
+        return list(disks.keys())
 
-    result = []
-
-    for name, drive in _query_udisks():
-        result.append(_disk_infos(name, drive, human_readable))
+    result = [
+        _disk_infos(name, disk.props, human_readable) for name, disk in disks.items()
+    ]
 
     return {"disks": result}
 
 
 def info(name, human_readable=False):
-    bus = dbus.SystemBus()
-    drive = dbus.Interface(
-        bus.get_object(
-            "org.freedesktop.UDisks2", f"/org/freedesktop/UDisks2/drives/{name}"
-        ),
-        dbus.PROPERTIES_IFACE,
-    ).GetAll("org.freedesktop.UDisks2.Drive")
-
-    return _disk_infos(name, drive, human_readable)
+    bus = sd_bus_open_system()
+    disk = Udisks2Manager(bus).get_disks().get(name)
+
+    if not disk:
+        return f"Unknown disk with name {name}" if human_readable else None
+
+    return _disk_infos(name, disk.props, human_readable)

src/storage.py

@@ -1,7 +1,7 @@
 def storage_disk_list(with_info=False, human_readable=False):
-    from yunohost.disks import list
+    from yunohost.disks import do_list
 
-    return list(with_info=with_info, human_readable=human_readable)
+    return do_list(with_info=with_info, human_readable=human_readable)
 
 
 def storage_disk_info(name, human_readable=False):