[FIX] Fix disk type on Windows

This commit is contained in:
Maxim Romanko 2024-04-15 12:01:09 +03:00
parent 21324b62f1
commit 0ffe7d07df
2 changed files with 16 additions and 3 deletions

View File

@ -86,10 +86,13 @@ class App:
import wmi import wmi
c = wmi.WMI() c = wmi.WMI()
self.drives = [] self.drives = []
ws = wmi.WMI(namespace='root/Microsoft/Windows/Storage')
drives_mt = ws.MSFT_PhysicalDisk()
if disks := c.Win32_DiskDrive(): if disks := c.Win32_DiskDrive():
for disk in disks: for disk in disks:
self.drives.append(Drive(disk.Model, disk.Name, disk.InterfaceType, disk.DefaultBlockSize, int( for drive in drives_mt:
disk.Size), disk.SerialNumber, disk.Index)) 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": elif self.OS_TYPE == "Linux":
from diskinfo import DiskInfo from diskinfo import DiskInfo
di = DiskInfo() di = DiskInfo()

View File

@ -29,6 +29,13 @@ erasing_methods = {
} }
} }
win_types = {
0: "Unspecified",
3: "HDD",
4: "SSD",
5: "SCM"
}
def get_random_bytes(size): def get_random_bytes(size):
seed=os.urandom(32) seed=os.urandom(32)
backend = default_backend() backend = default_backend()
@ -56,6 +63,9 @@ class Drive():
def __init__(self, name, path, disk_type, block_size, capacity, dev_id, index): def __init__(self, name, path, disk_type, block_size, capacity, dev_id, index):
self.name = name self.name = name
if isinstance(disk_type, int):
self.disk_type = win_types[disk_type]
else:
self.disk_type = disk_type self.disk_type = disk_type
self.path = path self.path = path
if block_size is not None: if block_size is not None: