Show phyiscal disks on Windows too.
Unmount partitions and lock physical drive when dd-ing iso on Windows Reset progres before starting dd-iso thread.pull/388/head
shinji-s
6 years ago
parent
9d349c7c64
commit
30128a7933
@ -0,0 +1,258 @@
|
||||
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)
|
||||
Loading…
Reference in New Issue