import collections import ctypes import io import pywintypes import struct import sys import time import win32api import win32con import win32file import winerror import winioctlcon import wmi from ctypes import wintypes from functools import reduce kernel32 = ctypes.WinDLL('kernel32') # , use_last_error=True) kernel32.FindFirstVolumeW.restype = wintypes.HANDLE kernel32.FindNextVolumeW.argtypes = (wintypes.HANDLE, wintypes.LPWSTR, wintypes.DWORD) kernel32.FindVolumeClose.argtypes = (wintypes.HANDLE,) def FindFirstVolume(): volume_name = ctypes.create_unicode_buffer(" " * 255) h = kernel32.FindFirstVolumeW(volume_name, 255) if h == win32file.INVALID_HANDLE_VALUE: raise RuntimeError("FindFirstVolume() returned an invalid handle.") return h, volume_name.value def FindNextVolume(hSearch): volume_name = ctypes.create_unicode_buffer(" " * 255) if kernel32.FindNextVolumeW(hSearch, volume_name, 255) != 0: return volume_name.value else: errno = ctypes.GetLastError() if errno == winerror.ERROR_NO_MORE_FILES: FindVolumeClose(hSearch) return None raise RuntimeError("FindNextVolume failed (%s)" % errno) def FindVolumeClose(hSearch): """Close a search handle opened by FindFirstVolume, typically after the last volume has been returned. """ if kernel32.FindVolumeClose(hSearch) == 0: raise RuntimeError("FindVolumeClose() failed.") def findAvailableDrives(): return [(d, win32file.GetDriveType(d)) for d in win32api.GetLogicalDriveStrings().rstrip('\0').split('\0')] def findNewDriveLetter(used_letters): all_letters = set([chr(i) for i in range(ord('C'), ord('Z')+1)]) return min(list(all_letters - set([s[0] for s in used_letters]))) def _openHandle(path, bWriteAccess, bWriteShare, logfunc = lambda s: None): TIMEOUT, NUM_RETRIES = 10, 20 for retry_count in range(6): try: access_flag = win32con.GENERIC_READ | \ (bWriteAccess and win32con.GENERIC_WRITE or 0) share_flag = win32con.FILE_SHARE_READ | \ (bWriteShare and win32con.FILE_SHARE_WRITE or 0) handle = win32file.CreateFile( path, access_flag, share_flag, None, win32con.OPEN_EXISTING, win32con.FILE_ATTRIBUTE_NORMAL, None) nth = { 0: 'first', 1:'second', 2:'third'} logfunc("Opening [%s]: success at the %s iteration" % (path, nth.get(retry_count, '%sth' % (retry_count+1)))) return handle except pywintypes.error as e: logfunc('Exception=>'+str(e)) if NUM_RETRIES/3 < retry_count: bWriteShare = True time.sleep(TIMEOUT / float(NUM_RETRIES)) else: raise RuntimeError("Couldn't open handle for %s." % path) def _closeHandle(h): x = win32file.CloseHandle(h) assert x != win32file.INVALID_HANDLE_VALUE return x class openHandle: def __init__(self, path, bWriteAccess, bWriteShare, logfunc = lambda s: None): self.path = path self.bWriteAccess = bWriteAccess self.bWriteShare = bWriteShare self.logfunc = logfunc self.h = None def __enter__(self): self.h = _openHandle(self.path, self.bWriteAccess, self.bWriteShare, self.logfunc) return self def __exit__(self, type_, value, traceback_): _closeHandle(self.h) def assert_physical_drive(self): if self.path.lower().find('physicaldrive')<0: raise RuntimeError("Handle is not one of a physical drive.") def LockPhysicalDrive(self): self.assert_physical_drive() lockPhysicalDrive(self.h, self.logfunc) self.logfunc("Successfully locked '%s'" % self.path) def ReadFile(self, size): return win32file.ReadFile(self.h, size, None) def WriteFile(self, b): return win32file.WriteFile(self.h, b, None) geometory_tuple = collections.namedtuple( 'DiskGeometory', ['number_of_cylinders', 'media_type', 'tracks_per_cylinder', 'sectors_per_track', 'bytes_per_sector', 'disk_size']) def DiskGeometory(self): self.assert_physical_drive() o = win32file.DeviceIoControl( self.h, winioctlcon.IOCTL_DISK_GET_DRIVE_GEOMETRY_EX, None, 256, None) return self.geometory_tuple(*struct.unpack('<qiiiiq', o[:32])) MAX_SECTORS_TO_CLEAR=128 def ZapMBRGPT(self, disk_size, sector_size, add1MB): self.assert_physical_drive() # Implementation borrowed from rufus: https://github.com/pbatard/rufus num_sectors_to_clear \ = (add1MB and 2048 or 0) + self.MAX_SECTORS_TO_CLEAR zeroBuf = b'\0' * sector_size for i in range(num_sectors_to_clear): self.WriteFile(zeroBuf) offset = disk_size - self.MAX_SECTORS_TO_CLEAR * sector_size win32file.SetFilePointer(self.h, offset, win32con.FILE_BEGIN) for i in range(num_sectors_to_clear): self.WriteFile(zeroBuf) # We need to append paddings as CREATE_DISK structure contains a union. param = struct.pack('<IIIHH8s', winioctlcon.PARTITION_STYLE_MBR, 0xdeadbeef, 0,0,0,b'abcdefgh') win32file.DeviceIoControl( self.h, winioctlcon.IOCTL_DISK_CREATE_DISK, param, 0, None) def CopyFrom(self, src_file, progress_cb): with openHandle(src_file, True, False, lambda s:sys.stdout.write(s+'\n')) as src: total_bytes = 0 hr, b = src.ReadFile(1024*1024) # win32file.ReadFile() seems to have a bug in the interpretation # of 'hr'. https://sourceforge.net/p/pywin32/bugs/689/ # The following loop condition is a workaround, which may not # work properly. while hr == 0 and len(b): win32file.WriteFile(self.h, b, None) total_bytes += len(b) progress_cb(total_bytes) hr, b = src.ReadFile(1024*1024) def lockPhysicalDrive(handle, logfunc=lambda s: None): try: win32file.DeviceIoControl( handle, winioctlcon.FSCTL_ALLOW_EXTENDED_DASD_IO, None, 0, None) except pywintypes.error as e: logfunc('IO boundary checks diabled.') for retry in range(20): try: win32file.DeviceIoControl(handle, winioctlcon.FSCTL_LOCK_VOLUME, None, 0, None) return except pywintypes.error as e: logfunc( str(e) ) time.sleep(1) raise RuntimeError("Couldn't lock the Volume.") def findVolumeGuids(): DiskExtent = collections.namedtuple( 'DiskExtent', ['DiskNumber', 'StartingOffset', 'ExtentLength']) Volume = collections.namedtuple( 'Volume', ['Guid', 'MediaType', 'DosDevice', 'Extents']) found = [] h, guid = FindFirstVolume() while h and guid: #print (guid) #print (guid, win32file.GetDriveType(guid), # win32file.QueryDosDevice(guid[4:-1])) hVolume = win32file.CreateFile( guid[:-1], win32con.GENERIC_READ, win32con.FILE_SHARE_READ | win32con.FILE_SHARE_WRITE, None, win32con.OPEN_EXISTING, win32con.FILE_ATTRIBUTE_NORMAL, None) extents = [] driveType = win32file.GetDriveType(guid) if driveType in [win32con.DRIVE_REMOVABLE, win32con.DRIVE_FIXED]: x = win32file.DeviceIoControl( hVolume, winioctlcon.IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS, None, 512, None) instream = io.BytesIO(x) numRecords = struct.unpack('<q', instream.read(8))[0] fmt = '<qqq' sz = struct.calcsize(fmt) while 1: b = instream.read(sz) if len(b) < sz: break rec = struct.unpack(fmt, b) extents.append( DiskExtent(*rec) ) vinfo = Volume(guid, driveType, win32file.QueryDosDevice(guid[4:-1]), extents) found.append(vinfo) guid = FindNextVolume(h) return found def ZapPhysicalDrive(target_drive, get_volume_info_func, log_func): with openHandle('\\\\.\\PhysicalDrive%d' % target_drive, True, False, lambda s:sys.stdout.write(s+'\n')) as hDrive: hDrive.LockPhysicalDrive() geom = hDrive.DiskGeometory() for v in get_volume_info_func(target_drive): volume_path = '\\\\.\\'+v.DeviceID log_func('Dismounting volume ' + volume_path) with openHandle(volume_path, False, False) as h: x = win32file.DeviceIoControl( h.h, winioctlcon.FSCTL_DISMOUNT_VOLUME, None, None) print ('FSCTL_DISMOUNT_VOLUME=>%s' % x) x = win32file.DeleteVolumeMountPoint(volume_path+'\\') log_func('DeleteVolumeMountPoint=>%s' % x) else: log_func('No volumes on %s' % target_drive) add1MB = False hDrive.ZapMBRGPT(geom.disk_size, geom.bytes_per_sector, add1MB) if __name__ == '__main__': # used_letters = [d for d in # win32api.GetLogicalDriveStrings().rstrip('\0').split('\0')] # print (used_letters) # print (findNewDriveLetter(used_letters)) # TargetDrive = 2 # vinfo_list = [x for x in findVolumeGuids() # if x.Extents and x.Extents[0].DiskNumber==TargetDrive] TargetDrive = 5 with openHandle('\\\\.\\PhysicalDrive%d' % TargetDrive, True, False, lambda s:sys.stdout.write(s+'\n')) as hDrive: hDrive.CopyFrom('c:/Users/shinj/Downloads/salitaz-rolling_iso', lambda b: None)