diff --git a/scripts/imager.py b/scripts/imager.py index e113529..7ae4b56 100644 --- a/scripts/imager.py +++ b/scripts/imager.py @@ -165,6 +165,10 @@ class Imager(QtWidgets.QMainWindow, Ui_MainWindow): if not model: model = "Unknown" else: + if partition == 0: + dinfo = osdriver.wmi_get_physicaldrive_info(usb_disk) + return _ntuple_diskusage(*[dinfo[a] for a in [ + 'Size', 'MediaType', 'Model']]) try: selected_usb_part = str(usb_disk) oFS = win32com.client.Dispatch("Scripting.FileSystemObject") @@ -178,6 +182,7 @@ class Imager(QtWidgets.QMainWindow, Ui_MainWindow): model = label except: log("Error detecting USB details.") + raise return _ntuple_diskusage(total_size, usb_type, model) diff --git a/scripts/osdriver.py b/scripts/osdriver.py index 702c15b..962fc98 100644 --- a/scripts/osdriver.py +++ b/scripts/osdriver.py @@ -83,6 +83,28 @@ def wmi_get_drive_info(usb_disk): return (partition, disk) raise RuntimeError('Failed to obtain drive information ' + usb_disk) +def wmi_get_physicaldrive_info(usb_disk): + "Return information about the drive that contains 'usb_disk'." + partition, disk = wmi_get_drive_info(usb_disk) + import wmi + c = wmi.WMI() + drv_list = [d for d in c.Win32_DiskDrive() + if d.Index == partition.DiskIndex] + assert len(drv_list)==1 + d = drv_list[0] + rdict = {} + for attrname, convfunc in [ + ('BytesPerSector', int), + ('DeviceID', str), + ('MediaType', str), + ('Model', str), + ('Partitions', int), + ('SerialNumber', str), + ('Size', int), + ('TotalSectors', int), + ]: + rdict[attrname] = getattr(d, attrname) + return rdict def wmi_get_drive_info_ex(usb_disk): assert platform.system() == 'Windows' @@ -149,7 +171,7 @@ def dd_win_clean_usb(usb_disk_no, status_update): with open('\\\\.\\PhysicalDrive%d' % usb_disk_no, 'rb'): return raise RuntimeError('PhysicalDrive%d is now gone!' % usb_disk_no) - raise RuntimeError("Execution of dd(.exe) has failed.") + raise RuntimeError("Execution of diskpart.exe has failed.") class Base: