Windows Drive Information query using Administator

win32-drive-devices.py

from ctypes import (
    POINTER,
    WinDLL,
    Structure,
    WinError,
    addressof,
    c_wchar,
    cast,
    create_string_buffer,
    get_last_error,
    create_unicode_buffer,
    byref,
    sizeof,
    string_at,
    wstring_at,
)
from ctypes.wintypes import (
    HANDLE,
    BOOL,
    BOOLEAN,
    DWORD,
    BYTE,
    LPWSTR,
)

kernel32 = WinDLL("kernel32", use_last_error=True)

INVALID_HANDLE_VALUE = HANDLE(-1).value

GENERIC_READ = 0x80000000

FILE_SHARE_READ = 1
FILE_SHARE_WRITE = 2

OPEN_EXISTING = 3

MAX_PATH = 260

IOCTL_STORAGE_GET_DEVICE_NUMBER = 0x2D1080
IOCTL_STORAGE_QUERY_PROPERTY = 0x2D1400

StorageDeviceProperty = 0
StorageDeviceSeekPenaltyProperty = 7

PropertyStandardQuery = 0


# ----------------------------------------------------------------------
# Structures
# ----------------------------------------------------------------------


class STORAGE_DEVICE_NUMBER(Structure):
    _fields_ = [
        ("DeviceType", DWORD),
        ("DeviceNumber", DWORD),
        ("PartitionNumber", DWORD),
    ]


class STORAGE_PROPERTY_QUERY(Structure):
    _fields_ = [
        ("PropertyId", DWORD),
        ("QueryType", DWORD),
        ("AdditionalParameters", BYTE * 1),
    ]


class STORAGE_DESCRIPTOR_HEADER(Structure):
    _fields_ = [
        ("Version", DWORD),
        ("Size", DWORD),
    ]


class STORAGE_DEVICE_DESCRIPTOR(Structure):
    _fields_ = [
        ("Version", DWORD),
        ("Size", DWORD),
        ("DeviceType", BYTE),
        ("DeviceTypeModifier", BYTE),
        ("RemovableMedia", BOOLEAN),
        ("CommandQueueing", BOOLEAN),
        ("VendorIdOffset", DWORD),
        ("ProductIdOffset", DWORD),
        ("ProductRevisionOffset", DWORD),
        ("SerialNumberOffset", DWORD),
        ("BusType", DWORD),
        ("RawPropertiesLength", DWORD),
        ("RawDeviceProperties", BYTE * 1),
    ]


class DEVICE_SEEK_PENALTY_DESCRIPTOR(Structure):
    _fields_ = [
        ("Version", DWORD),
        ("Size", DWORD),
        ("IncursSeekPenalty", BOOLEAN),
    ]


# ----------------------------------------------------------------------
# Prototypes
# ----------------------------------------------------------------------

FindFirstVolumeW = kernel32.FindFirstVolumeW
FindFirstVolumeW.argtypes = [LPWSTR, DWORD]
FindFirstVolumeW.restype = HANDLE

FindNextVolumeW = kernel32.FindNextVolumeW
FindNextVolumeW.argtypes = [HANDLE, LPWSTR, DWORD]
FindNextVolumeW.restype = BOOL

FindVolumeClose = kernel32.FindVolumeClose
FindVolumeClose.argtypes = [HANDLE]
FindVolumeClose.restype = BOOL

GetVolumePathNamesForVolumeNameW = kernel32.GetVolumePathNamesForVolumeNameW

CreateFileW = kernel32.CreateFileW
CreateFileW.restype = HANDLE

DeviceIoControl = kernel32.DeviceIoControl
CloseHandle = kernel32.CloseHandle


# ----------------------------------------------------------------------
# Helpers
# ----------------------------------------------------------------------


def open_device(path):
    h = CreateFileW(
        path,
        GENERIC_READ,
        FILE_SHARE_READ | FILE_SHARE_WRITE,
        None,
        OPEN_EXISTING,
        0,
        None,
    )

    if h == INVALID_HANDLE_VALUE:
        raise WinError(get_last_error())

    return h


# ----------------------------------------------------------------------
# Volume enumeration
# ----------------------------------------------------------------------


def enum_volumes():
    buf = create_unicode_buffer(MAX_PATH)

    h = FindFirstVolumeW(buf, MAX_PATH)

    if h == INVALID_HANDLE_VALUE:
        raise WinError(get_last_error())

    try:
        yield buf.value

        while FindNextVolumeW(h, buf, MAX_PATH):
            yield buf.value

    finally:
        FindVolumeClose(h)


# ----------------------------------------------------------------------
# Volume -> Mount points
# ----------------------------------------------------------------------


def get_mount_points(volume_guid):
    needed = DWORD()

    GetVolumePathNamesForVolumeNameW(
        volume_guid,
        None,
        0,
        byref(needed),
    )

    if needed.value == 0:
        return []

    buf = create_unicode_buffer(needed.value)

    if not GetVolumePathNamesForVolumeNameW(
        volume_guid,
        buf,
        needed.value,
        byref(needed),
    ):
        raise WinError(get_last_error())

    result = []

    offset = 0

    while True:
        s = wstring_at(addressof(buf) + offset * sizeof(c_wchar))

        if not s:
            break

        result.append(s)

        offset += len(s) + 1

    return result


# ----------------------------------------------------------------------
# Volume -> Disk Number
# ----------------------------------------------------------------------


def volume_to_disk_info(volume_guid) -> STORAGE_DEVICE_NUMBER:
    h = open_device(volume_guid.rstrip("\\"))

    try:
        result = STORAGE_DEVICE_NUMBER()
        returned = DWORD()

        # requires admin priviledge
        ok = DeviceIoControl(
            h,
            IOCTL_STORAGE_GET_DEVICE_NUMBER,
            None,
            0,
            byref(result),
            sizeof(result),
            byref(returned),
            None,
        )

        if not ok:
            raise WinError(get_last_error())

        return result

    finally:
        CloseHandle(h)


# ----------------------------------------------------------------------
# PhysicalDrive -> Model
# ----------------------------------------------------------------------


def get_disk_model(disk_number):
    h = open_device(rf"\\.\PhysicalDrive{disk_number}")

    try:
        query = STORAGE_PROPERTY_QUERY()
        query.PropertyId = StorageDeviceProperty
        query.QueryType = PropertyStandardQuery

        header = STORAGE_DESCRIPTOR_HEADER()
        returned = DWORD()

        DeviceIoControl(
            h,
            IOCTL_STORAGE_QUERY_PROPERTY,
            byref(query),
            sizeof(query),
            byref(header),
            sizeof(header),
            byref(returned),
            None,
        )

        buf = create_string_buffer(header.Size)

        ok = DeviceIoControl(
            h,
            IOCTL_STORAGE_QUERY_PROPERTY,
            byref(query),
            sizeof(query),
            buf,
            len(buf),
            byref(returned),
            None,
        )

        if not ok:
            raise WinError(get_last_error())

        desc = cast(buf, POINTER(STORAGE_DEVICE_DESCRIPTOR)).contents

        if desc.ProductIdOffset == 0:
            return "<unknown>"

        return (
            string_at(addressof(buf) + desc.ProductIdOffset)
            .decode("ascii", errors="ignore")
            .strip()
        )

    finally:
        CloseHandle(h)


# ----------------------------------------------------------------------
# PhysicalDrive -> SSD/HDD
# ----------------------------------------------------------------------


def get_disk_kind(disk_number):
    h = open_device(rf"\\.\PhysicalDrive{disk_number}")

    try:
        query = STORAGE_PROPERTY_QUERY()
        query.PropertyId = StorageDeviceSeekPenaltyProperty
        query.QueryType = PropertyStandardQuery

        result = DEVICE_SEEK_PENALTY_DESCRIPTOR()

        returned = DWORD()

        ok = DeviceIoControl(
            h,
            IOCTL_STORAGE_QUERY_PROPERTY,
            byref(query),
            sizeof(query),
            byref(result),
            sizeof(result),
            byref(returned),
            None,
        )

        if not ok:
            return "Unknown"

        return "HDD" if result.IncursSeekPenalty else "SSD"

    finally:
        CloseHandle(h)


# ----------------------------------------------------------------------
# Main
# ----------------------------------------------------------------------


def main():
    cache = {}

    for volume in enum_volumes():
        try:
            # otherwise raises error
            disk_info = volume_to_disk_info(volume)

            device_num = disk_info.DeviceNumber
            partition_num = disk_info.PartitionNumber

            if device_num not in cache:
                cache[device_num] = {
                    "model": get_disk_model(device_num),
                    "kind": get_disk_kind(device_num),
                }

            info = cache[device_num]

            print()
            print("Volume    :", volume)

            for mount in get_mount_points(volume):
                print("Mount     :", mount)

            print("Disk      :", device_num)
            print("Partition :", partition_num)
            print("Model     :", info["model"])
            print("Type      :", info["kind"])

        except Exception as e:
            print()
            print(volume)
            print("ERROR:", e)


if __name__ == "__main__":
    main()