From 0ffe7d07df1a37ca9c85980c16f082d3ce1bbf1c Mon Sep 17 00:00:00 2001
From: Maxim Romanko <mromanko@list.ru>
Date: Mon, 15 Apr 2024 12:01:09 +0300
Subject: [PATCH] [FIX] Fix disk type on Windows

---
 app/__init__.py    |  7 +++++--
 app/utils_drive.py | 12 +++++++++++-
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/app/__init__.py b/app/__init__.py
index 026bd9e..c10b37b 100644
--- a/app/__init__.py
+++ b/app/__init__.py
@@ -86,10 +86,13 @@ class App:
             import wmi
             c = wmi.WMI()
             self.drives = []
+            ws = wmi.WMI(namespace='root/Microsoft/Windows/Storage')
+            drives_mt = ws.MSFT_PhysicalDisk()
             if disks := c.Win32_DiskDrive():
                 for disk in disks:
-                    self.drives.append(Drive(disk.Model, disk.Name, disk.InterfaceType, disk.DefaultBlockSize, int(
-                        disk.Size), disk.SerialNumber, disk.Index))
+                    for drive in drives_mt:
+                        if disk.Index == int(drive.DeviceId):
+                            self.drives.append(Drive(disk.Model, disk.Name, drive.MediaType, disk.DefaultBlockSize, int(disk.Size), disk.SerialNumber, disk.Index))
         elif self.OS_TYPE == "Linux":
             from diskinfo import DiskInfo
             di = DiskInfo()
diff --git a/app/utils_drive.py b/app/utils_drive.py
index edee52a..96bf511 100644
--- a/app/utils_drive.py
+++ b/app/utils_drive.py
@@ -29,6 +29,13 @@ erasing_methods = {
     }
 }
 
+win_types = {
+    0: "Unspecified",
+    3: "HDD",
+    4: "SSD",
+    5: "SCM"
+}
+
 def get_random_bytes(size):
     seed=os.urandom(32)
     backend = default_backend()
@@ -56,7 +63,10 @@ class Drive():
 
     def __init__(self, name, path, disk_type, block_size, capacity, dev_id, index):
         self.name = name
-        self.disk_type = disk_type
+        if isinstance(disk_type, int):
+            self.disk_type = win_types[disk_type]
+        else:
+            self.disk_type = disk_type
         self.path = path
         if block_size is not None:
             self.block_size = block_size