From d8e2be5d0a31d6b960992fba910557ae1f3032c1 Mon Sep 17 00:00:00 2001
From: CllaudiaB <claudia.bura@numberly.com>
Date: Wed, 12 Feb 2025 16:30:55 +0100
Subject: [PATCH] fix: find storage

---
 netbox_agent/lshw.py           | 29 +++++++++++++++--------------
 netbox_agent/virtualmachine.py |  2 +-
 2 files changed, 16 insertions(+), 15 deletions(-)

diff --git a/netbox_agent/lshw.py b/netbox_agent/lshw.py
index 381c5414..3ff866d0 100644
--- a/netbox_agent/lshw.py
+++ b/netbox_agent/lshw.py
@@ -94,20 +94,7 @@ def find_network(self, obj):
         )
 
     def find_storage(self, obj):
-        if "children" in obj:
-            for device in obj["children"]:
-                self.disks.append(
-                    {
-                        "logicalname": device.get("logicalname"),
-                        "product": device.get("product"),
-                        "serial": device.get("serial"),
-                        "version": device.get("version"),
-                        "size": device.get("size"),
-                        "description": device.get("description"),
-                        "type": device.get("description"),
-                    }
-                )
-        elif "driver" in obj["configuration"] and "nvme" in obj["configuration"]["driver"]:
+        if obj.get("configuration") and "driver" in obj["configuration"] and "nvme" in obj["configuration"]["driver"]:
             if not is_tool("nvme"):
                 logging.error("nvme-cli >= 1.0 does not seem to be installed")
                 return
@@ -131,6 +118,20 @@ def find_storage(self, obj):
                     self.disks.append(d)
             except Exception:
                 pass
+        elif "children" in obj:
+            for device in obj["children"]:
+                if device.get("class") == "disk":
+                    self.disks.append(
+                        {
+                            "logicalname": device.get("logicalname"),
+                            "product": device.get("product"),
+                            "serial": device.get("serial"),
+                            "version": device.get("version"),
+                            "size": device.get("size"),
+                            "description": device.get("description"),
+                            "type": device.get("description"),
+                        }
+                    )
 
     def find_cpus(self, obj):
         if "product" in obj:
diff --git a/netbox_agent/virtualmachine.py b/netbox_agent/virtualmachine.py
index 00fa5ccc..7564d6e7 100644
--- a/netbox_agent/virtualmachine.py
+++ b/netbox_agent/virtualmachine.py
@@ -125,7 +125,7 @@ def netbox_create_or_update(self, config):
                 vm.memory = memory
                 updated += 1
             if vm.disk != disk:
-                vm.disk = disk
+                vm.disk = round(disk)
                 updated += 1
 
             vm_tags = sorted(set([x.name for x in vm.tags]))