162 lines
6.7 KiB
Python
162 lines
6.7 KiB
Python
"""Fast Windows card discovery straight from the storage drivers (ctypes).
|
|
|
|
Why not Get-Disk / WMI: every fresh PowerShell/WMI query enumerates ALL
|
|
disks, which wakes sleeping hard drives - measured 1.6 to 21 s per listing
|
|
on a PC with two SATA HDDs, freezing the UI each time. Here each physical
|
|
drive is opened with access 0 (no admin needed, no I/O to the medium) and
|
|
only its cached device descriptor is read; size and volumes are queried only
|
|
for USB/SD/MMC disks. A full listing takes milliseconds.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import concurrent.futures
|
|
import ctypes
|
|
import string
|
|
from ctypes import wintypes
|
|
|
|
from .disks import Disk, Volume
|
|
|
|
IOCTL_STORAGE_QUERY_PROPERTY = 0x002D1400
|
|
IOCTL_STORAGE_GET_DEVICE_NUMBER = 0x002D1080
|
|
IOCTL_DISK_GET_DRIVE_GEOMETRY_EX = 0x000700A0
|
|
BUS_TYPES = {1: "SCSI", 2: "ATAPI", 3: "ATA", 4: "1394", 7: "USB", 8: "RAID", 10: "SAS", 11: "SATA",
|
|
12: "SD", 13: "MMC", 14: "Virtual", 15: "FileBackedVirtual", 17: "NVMe", 18: "SCM", 19: "UFS"}
|
|
SHARE_RW, OPEN_EXISTING = 0x3, 3
|
|
DRIVE_REMOVABLE, DRIVE_FIXED = 2, 3
|
|
MAX_DRIVES = 32
|
|
|
|
_k32 = ctypes.WinDLL("kernel32", use_last_error=True)
|
|
_k32.CreateFileW.restype = wintypes.HANDLE
|
|
_k32.CreateFileW.argtypes = [wintypes.LPCWSTR, wintypes.DWORD, wintypes.DWORD, wintypes.LPVOID,
|
|
wintypes.DWORD, wintypes.DWORD, wintypes.HANDLE]
|
|
_k32.DeviceIoControl.argtypes = [wintypes.HANDLE, wintypes.DWORD, wintypes.LPVOID, wintypes.DWORD,
|
|
wintypes.LPVOID, wintypes.DWORD, ctypes.POINTER(wintypes.DWORD), wintypes.LPVOID]
|
|
_k32.GetDriveTypeW.argtypes = [wintypes.LPCWSTR]
|
|
_k32.GetVolumeInformationW.argtypes = [wintypes.LPCWSTR, wintypes.LPWSTR, wintypes.DWORD,
|
|
ctypes.POINTER(wintypes.DWORD), ctypes.POINTER(wintypes.DWORD),
|
|
ctypes.POINTER(wintypes.DWORD), wintypes.LPWSTR, wintypes.DWORD]
|
|
_INVALID = wintypes.HANDLE(-1).value
|
|
|
|
|
|
class _Query(ctypes.Structure): # STORAGE_PROPERTY_QUERY (StorageDeviceProperty, standard)
|
|
_fields_ = [("PropertyId", wintypes.DWORD), ("QueryType", wintypes.DWORD), ("Extra", ctypes.c_ubyte * 4)]
|
|
|
|
|
|
class _DevNumber(ctypes.Structure): # STORAGE_DEVICE_NUMBER
|
|
_fields_ = [("DeviceType", wintypes.DWORD), ("DeviceNumber", wintypes.DWORD),
|
|
("PartitionNumber", wintypes.DWORD)]
|
|
|
|
|
|
def _open(path: str):
|
|
h = _k32.CreateFileW(path, 0, SHARE_RW, None, OPEN_EXISTING, 0, None)
|
|
return None if h in (None, _INVALID) else h
|
|
|
|
|
|
def _ioctl(h, code: int, inbuf=None, insize: int = 0, outsize: int = 1024) -> bytes | None:
|
|
out = ctypes.create_string_buffer(outsize)
|
|
got = wintypes.DWORD()
|
|
ok = _k32.DeviceIoControl(h, code, inbuf, insize, out, outsize, ctypes.byref(got), None)
|
|
return out.raw[:got.value] if ok else None
|
|
|
|
|
|
def _descriptor(h) -> tuple[str, str] | None:
|
|
"""(bus, "vendor product") from STORAGE_DEVICE_DESCRIPTOR - served from the driver's cache."""
|
|
q = _Query(0, 0)
|
|
raw = _ioctl(h, IOCTL_STORAGE_QUERY_PROPERTY, ctypes.byref(q), ctypes.sizeof(q))
|
|
if not raw or len(raw) < 32:
|
|
return None
|
|
vendor_off, product_off = int.from_bytes(raw[12:16], "little"), int.from_bytes(raw[16:20], "little")
|
|
bus = int.from_bytes(raw[28:32], "little")
|
|
|
|
def text(off: int) -> str:
|
|
if not off or off >= len(raw):
|
|
return ""
|
|
return raw[off:].split(b"\0", 1)[0].decode("ascii", "replace").strip()
|
|
|
|
return BUS_TYPES.get(bus, f"bus{bus}"), " ".join(t for t in (text(vendor_off), text(product_off)) if t) or "Disk"
|
|
|
|
|
|
def _size(h) -> int:
|
|
raw = _ioctl(h, IOCTL_DISK_GET_DRIVE_GEOMETRY_EX, outsize=256) # fails when no card is inserted
|
|
return int.from_bytes(raw[24:32], "little") if raw and len(raw) >= 32 else 0
|
|
|
|
|
|
def _volumes_by_disk() -> dict[int, list[Volume]]:
|
|
found: dict[int, list[Volume]] = {}
|
|
mask = _k32.GetLogicalDrives()
|
|
for i, letter in enumerate(string.ascii_uppercase):
|
|
if not mask & (1 << i) or _k32.GetDriveTypeW(f"{letter}:\\") not in (DRIVE_REMOVABLE, DRIVE_FIXED):
|
|
continue
|
|
h = _open(f"\\\\.\\{letter}:")
|
|
if h is None:
|
|
continue
|
|
try:
|
|
raw = _ioctl(h, IOCTL_STORAGE_GET_DEVICE_NUMBER, outsize=ctypes.sizeof(_DevNumber))
|
|
finally:
|
|
_k32.CloseHandle(h)
|
|
if raw and len(raw) >= 8:
|
|
num = int.from_bytes(raw[4:8], "little")
|
|
found.setdefault(num, []).append(Volume(f"{letter}:\\", "", ""))
|
|
return found
|
|
|
|
|
|
def _fill_volume_info(v: Volume) -> None:
|
|
name, fs = ctypes.create_unicode_buffer(64), ctypes.create_unicode_buffer(32)
|
|
serial, maxlen, flags = wintypes.DWORD(), wintypes.DWORD(), wintypes.DWORD()
|
|
if _k32.GetVolumeInformationW(v.mount, name, 64, ctypes.byref(serial), ctypes.byref(maxlen),
|
|
ctypes.byref(flags), fs, 32):
|
|
v.filesystem, v.label = fs.value, name.value
|
|
|
|
|
|
MEDIA_TIMEOUT_S = 1.5
|
|
# Slot queries still running from an earlier listing (an empty card-reader
|
|
# slot can block ~13 s on its first media query). Never start a second one.
|
|
_pending: dict[int, concurrent.futures.Future] = {}
|
|
_pool = concurrent.futures.ThreadPoolExecutor(max_workers=8, thread_name_prefix="sd-media")
|
|
|
|
|
|
def _media_size(n: int) -> int:
|
|
h = _open(f"\\\\.\\PhysicalDrive{n}")
|
|
if h is None:
|
|
return 0
|
|
try:
|
|
return _size(h)
|
|
finally:
|
|
_k32.CloseHandle(h)
|
|
|
|
|
|
def list_disks(buses: set[str]) -> list[Disk]:
|
|
"""Disks on the given buses (e.g. {"USB", "SD", "MMC"}), with size and volumes.
|
|
|
|
Media queries run in parallel with a short timeout: a slot that does not
|
|
answer in time is reported as "no card" instead of freezing the caller.
|
|
"""
|
|
candidates: list[tuple[int, tuple[str, str]]] = []
|
|
for n in range(MAX_DRIVES):
|
|
h = _open(f"\\\\.\\PhysicalDrive{n}")
|
|
if h is None:
|
|
continue
|
|
try:
|
|
desc = _descriptor(h)
|
|
finally:
|
|
_k32.CloseHandle(h)
|
|
if desc is not None and desc[0] in buses: # internal disks: only the cached descriptor is read
|
|
candidates.append((n, desc))
|
|
for n, _ in candidates:
|
|
if n not in _pending or _pending[n].done():
|
|
_pending[n] = _pool.submit(_media_size, n)
|
|
concurrent.futures.wait([_pending[n] for n, _ in candidates], timeout=MEDIA_TIMEOUT_S)
|
|
volumes: dict[int, list[Volume]] | None = None
|
|
disks = []
|
|
for n, (bus, name) in candidates:
|
|
fut = _pending[n]
|
|
size = fut.result() if fut.done() and fut.exception() is None else 0
|
|
if size and volumes is None:
|
|
volumes = _volumes_by_disk()
|
|
vols = (volumes or {}).get(n, []) if size else []
|
|
for v in vols:
|
|
_fill_volume_info(v)
|
|
disks.append(Disk(str(n), name, size, bus, vols, has_media=size > 0))
|
|
return disks
|