|
| 1 | +#!C:\Python27\python.exe |
| 2 | +# -*- coding: utf-8 -*- |
| 3 | +# The MIT License (MIT) |
| 4 | +# |
| 5 | +# Copyright © 2014-2016 Santoso Wijaya <santoso.wijaya@gmail.com> |
| 6 | +# |
| 7 | +# Permission is hereby granted, free of charge, to any person |
| 8 | +# obtaining a copy of this software and associated documentation files |
| 9 | +# (the "Software"), to deal in the Software without restriction, |
| 10 | +# including without limitation the rights to use, copy, modify, merge, |
| 11 | +# publish, distribute, sub-license, and/or sell copies of the Software, |
| 12 | +# and to permit persons to whom the Software is furnished to do so, |
| 13 | +# subject to the following conditions: |
| 14 | +# |
| 15 | +# The above copyright notice and this permission notice shall be |
| 16 | +# included in all copies or substantial portions of the Software. |
| 17 | +# |
| 18 | +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, |
| 19 | +# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF |
| 20 | +# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND |
| 21 | +# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS |
| 22 | +# BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN |
| 23 | +# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN |
| 24 | +# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE |
| 25 | +# SOFTWARE. |
| 26 | +# |
| 27 | +# Originally written by GitHub user @santa4nt (Santoso Wijaya) |
| 28 | +# |
| 29 | +# Source gist: https://gist.github.com/santa4nt/11068180 (rev 4) |
| 30 | +# |
| 31 | +# This implements primitives used to query device capacity on windows |
| 32 | +# machines. |
| 33 | + |
| 34 | +import ctypes |
| 35 | +import ctypes.wintypes as wintypes |
| 36 | +from ctypes import windll |
| 37 | + |
| 38 | + |
| 39 | +LPDWORD = ctypes.POINTER(wintypes.DWORD) |
| 40 | +LPOVERLAPPED = wintypes.LPVOID |
| 41 | +LPSECURITY_ATTRIBUTES = wintypes.LPVOID |
| 42 | + |
| 43 | +GENERIC_READ = 0x80000000 |
| 44 | +GENERIC_WRITE = 0x40000000 |
| 45 | +GENERIC_EXECUTE = 0x20000000 |
| 46 | +GENERIC_ALL = 0x10000000 |
| 47 | + |
| 48 | +CREATE_NEW = 1 |
| 49 | +CREATE_ALWAYS = 2 |
| 50 | +OPEN_EXISTING = 3 |
| 51 | +OPEN_ALWAYS = 4 |
| 52 | +TRUNCATE_EXISTING = 5 |
| 53 | + |
| 54 | +FILE_ATTRIBUTE_NORMAL = 0x00000080 |
| 55 | + |
| 56 | +INVALID_HANDLE_VALUE = -1 |
| 57 | + |
| 58 | +NULL = 0 |
| 59 | +FALSE = wintypes.BOOL(0) |
| 60 | +TRUE = wintypes.BOOL(1) |
| 61 | + |
| 62 | + |
| 63 | +def _CreateFile(filename, access, mode, creation, flags): |
| 64 | + """See: CreateFile function |
| 65 | +
|
| 66 | + http://msdn.microsoft.com/en-us/library/windows/desktop/aa363858(v=vs.85).aspx |
| 67 | +
|
| 68 | + """ |
| 69 | + CreateFile_Fn = windll.kernel32.CreateFileW |
| 70 | + CreateFile_Fn.argtypes = [ |
| 71 | + wintypes.LPWSTR, # _In_ LPCTSTR lpFileName |
| 72 | + wintypes.DWORD, # _In_ DWORD dwDesiredAccess |
| 73 | + wintypes.DWORD, # _In_ DWORD dwShareMode |
| 74 | + LPSECURITY_ATTRIBUTES, # _In_opt_ LPSECURITY_ATTRIBUTES lpSecurityAttributes |
| 75 | + wintypes.DWORD, # _In_ DWORD dwCreationDisposition |
| 76 | + wintypes.DWORD, # _In_ DWORD dwFlagsAndAttributes |
| 77 | + wintypes.HANDLE] # _In_opt_ HANDLE hTemplateFile |
| 78 | + CreateFile_Fn.restype = wintypes.HANDLE |
| 79 | + |
| 80 | + return wintypes.HANDLE(CreateFile_Fn(filename, |
| 81 | + access, |
| 82 | + mode, |
| 83 | + NULL, |
| 84 | + creation, |
| 85 | + flags, |
| 86 | + NULL)) |
| 87 | + |
| 88 | + |
| 89 | +def _DeviceIoControl(devhandle, ioctl, inbuf, inbufsiz, outbuf, outbufsiz): |
| 90 | + """See: DeviceIoControl function |
| 91 | +
|
| 92 | + http://msdn.microsoft.com/en-us/library/aa363216(v=vs.85).aspx |
| 93 | +
|
| 94 | + """ |
| 95 | + DeviceIoControl_Fn = windll.kernel32.DeviceIoControl |
| 96 | + DeviceIoControl_Fn.argtypes = [ |
| 97 | + wintypes.HANDLE, # _In_ HANDLE hDevice |
| 98 | + wintypes.DWORD, # _In_ DWORD dwIoControlCode |
| 99 | + wintypes.LPVOID, # _In_opt_ LPVOID lpInBuffer |
| 100 | + wintypes.DWORD, # _In_ DWORD nInBufferSize |
| 101 | + wintypes.LPVOID, # _Out_opt_ LPVOID lpOutBuffer |
| 102 | + wintypes.DWORD, # _In_ DWORD nOutBufferSize |
| 103 | + LPDWORD, # _Out_opt_ LPDWORD lpBytesReturned |
| 104 | + LPOVERLAPPED] # _Inout_opt_ LPOVERLAPPED lpOverlapped |
| 105 | + DeviceIoControl_Fn.restype = wintypes.BOOL |
| 106 | + |
| 107 | + # allocate a DWORD, and take its reference |
| 108 | + dwBytesReturned = wintypes.DWORD(0) |
| 109 | + lpBytesReturned = ctypes.byref(dwBytesReturned) |
| 110 | + |
| 111 | + status = DeviceIoControl_Fn(devhandle, |
| 112 | + ioctl, |
| 113 | + inbuf, |
| 114 | + inbufsiz, |
| 115 | + outbuf, |
| 116 | + outbufsiz, |
| 117 | + lpBytesReturned, |
| 118 | + None) |
| 119 | + |
| 120 | + return status, dwBytesReturned |
| 121 | + |
| 122 | + |
| 123 | +class DeviceIoControl(object): |
| 124 | + |
| 125 | + def __init__(self, path): |
| 126 | + self.path = path |
| 127 | + self._fhandle = None |
| 128 | + |
| 129 | + def _validate_handle(self): |
| 130 | + if self._fhandle is None: |
| 131 | + raise Exception('No file handle') |
| 132 | + if self._fhandle.value == wintypes.HANDLE(INVALID_HANDLE_VALUE).value: |
| 133 | + raise Exception('Failed to open %s. GetLastError(): %d' % |
| 134 | + (self.path, windll.kernel32.GetLastError())) |
| 135 | + |
| 136 | + def ioctl(self, ctl, inbuf, inbufsiz, outbuf, outbufsiz): |
| 137 | + self._validate_handle() |
| 138 | + return _DeviceIoControl(self._fhandle, ctl, inbuf, inbufsiz, outbuf, outbufsiz) |
| 139 | + |
| 140 | + def __enter__(self): |
| 141 | + self._fhandle = _CreateFile( |
| 142 | + self.path, |
| 143 | + GENERIC_READ | GENERIC_WRITE, |
| 144 | + 0, |
| 145 | + OPEN_EXISTING, |
| 146 | + FILE_ATTRIBUTE_NORMAL) |
| 147 | + self._validate_handle() |
| 148 | + return self |
| 149 | + |
| 150 | + def __exit__(self, typ, val, tb): |
| 151 | + try: |
| 152 | + self._validate_handle() |
| 153 | + except Exception: |
| 154 | + pass |
| 155 | + else: |
| 156 | + windll.kernel32.CloseHandle(self._fhandle) |
| 157 | + |
| 158 | + |
| 159 | +def query_device_capacity(device=r'\\.\PhysicalDrive0'): |
| 160 | + """Query a device and calculate its capacity on a Windows platform. |
| 161 | +
|
| 162 | + :return: the device capacity in bytes as a number type |
| 163 | + """ |
| 164 | + |
| 165 | + # sample code using \\.\PhysicalDrive0 |
| 166 | + # See: http://msdn.microsoft.com/en-us/library/windows/desktop/aa363147(v=vs.85).aspx |
| 167 | + |
| 168 | + # first, define the Structure in ctypes language |
| 169 | + class DISK_GEOMETRY(ctypes.Structure): |
| 170 | + """See: http://msdn.microsoft.com/en-us/library/aa363972(v=vs.85).aspx""" |
| 171 | + _fields_ = [ |
| 172 | + ('Cylinders', wintypes.LARGE_INTEGER), |
| 173 | + ('MediaType', wintypes.BYTE), # MEDIA_TYPE |
| 174 | + ('TracksPerCylinder', wintypes.DWORD), |
| 175 | + ('SectorsPerTrack', wintypes.DWORD), |
| 176 | + ('BytesPerSector', wintypes.DWORD) |
| 177 | + ] |
| 178 | + |
| 179 | + IOCTL_DISK_GET_DRIVE_GEOMETRY = 0x70000 |
| 180 | + |
| 181 | + disk_geometry = DISK_GEOMETRY() |
| 182 | + p_disk_geometry = ctypes.pointer(disk_geometry) |
| 183 | + |
| 184 | + with DeviceIoControl(device) as dctl: |
| 185 | + status, junk = dctl.ioctl(IOCTL_DISK_GET_DRIVE_GEOMETRY, |
| 186 | + None, 0, # no input buffer |
| 187 | + p_disk_geometry, ctypes.sizeof(DISK_GEOMETRY)) |
| 188 | + |
| 189 | + if status: |
| 190 | + cylinders = getattr(disk_geometry, 'Cylinders') |
| 191 | + trackspercylinder = getattr(disk_geometry, 'TracksPerCylinder') |
| 192 | + sectorspertrack = getattr(disk_geometry, 'SectorsPerTrack') |
| 193 | + bytespersector = getattr(disk_geometry, 'BytesPerSector') |
| 194 | + return cylinders * trackspercylinder * sectorspertrack * bytespersector |
| 195 | + else: |
| 196 | + print 'IOCTL returned failure. GetLastError(): %d' % windll.kernel32.GetLastError() |
0 commit comments