Merge pull request #385 from shinji-s/devel

Check physical disk capacity, not the size of the partition, before i…

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)
 

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: