WIP: rewrite Windows code

.github/workflows/clang-format-check.yml

@@ -11,4 +11,4 @@ jobs:
       with:
         clang-format-version: '17'
         check-path: 'src'
-        exclude-regex: 'src\/helper\/win|src\/app\/crashhandler.cpp'
+        exclude-regex: 'src\/app\/crashhandler.cpp'

src/app/CMakeLists.txt

@@ -90,7 +90,7 @@ if (UNIX AND NOT APPLE)
 endif()
 
 if (WIN32)
-    target_link_libraries(mediawriter dbghelp)
+    target_link_libraries(mediawriter dbghelp libwmi)
 endif()
 
 if (APPLE)

src/app/main.cpp

@@ -1,5 +1,6 @@
 /*
  * Fedora Media Writer
+ * Copyright (C) 2024 Jan Grulich <jgrulichredhat.com>
  * Copyright (C) 2016 Martin Bříza <[email protected]>
  *
  * This program is free software; you can redistribute it and/or

src/app/windrivemanager.cpp

@@ -1,6 +1,6 @@
 /*
  * Fedora Media Writer
- * Copyright (C) 2022 Jan Grulich <[email protected]>
+ * Copyright (C) 2022-2024 Jan Grulich <[email protected]>
  * Copyright (C) 2011-2022 Pete Batard <[email protected]>
  * Copyright (C) 2016 Martin Bříza <[email protected]>
  *
@@ -22,262 +22,99 @@
 #include "windrivemanager.h"
 #include "notifications.h"
 
-#include <QDebug>
 #include <QStringList>
 #include <QTimer>
 
-#include <windows.h>
-#define INITGUID
-#include <guiddef.h>
+#include <dbt.h>
 
-#include <cmath>
-#include <cstring>
-
-const int maxPartitionCount = 16;
-
-DEFINE_GUID(PARTITION_MICROSOFT_DATA, 0xEBD0A0A2, 0xB9E5, 0x4433, 0x87, 0xC0, 0x68, 0xB6, 0xB7, 0x26, 0x99, 0xC7);
+#pragma comment(lib, "wbemuuid.lib")
 
 WinDriveProvider::WinDriveProvider(DriveManager *parent)
     : DriveProvider(parent)
+    , m_wmi(std::make_unique<LibWMI>(this))
 {
     mDebug() << this->metaObject()->className() << "construction";
+    qApp->installNativeEventFilter(this);
     QTimer::singleShot(0, this, &WinDriveProvider::checkDrives);
 }
 
-void WinDriveProvider::checkDrives()
+bool WinDriveProvider::nativeEventFilter(const QByteArray &eventType, void *message, qintptr *result)
 {
-    static bool firstRun = true;
-
-    if (firstRun)
-        mDebug() << this->metaObject()->className() << "Looking for the drives for the first time";
-
-    for (int i = 0; i < 64; i++) {
-        bool present = describeDrive(i, firstRun);
-        if (!present && m_drives.contains(i)) {
-            emit driveRemoved(m_drives[i]);
-            m_drives[i]->deleteLater();
-            m_drives.remove(i);
+    Q_UNUSED(eventType);
+
+    MSG *msg = static_cast<MSG *>(message);
+    if ((msg->message == WM_DEVICECHANGE) && ((msg->wParam == DBT_DEVICEARRIVAL) || (msg->wParam == DBT_DEVICEREMOVECOMPLETE))) {
+        mDebug() << "Recieved device change event";
+        *result = TRUE;
+        for (const WinDrive *drive : m_drives) {
+            // Ignore device change events when we are restoring or writting and schedule
+            // re-check once we are done
+            if (drive->busy()) {
+                return true;
+            }
         }
+        QTimer::singleShot(0, this, &WinDriveProvider::checkDrives);
+        return true;
     }
-
-    if (firstRun)
-        mDebug() << this->metaObject()->className() << "Finished looking for the drives for the first time";
-    firstRun = false;
-    QTimer::singleShot(2500, this, &WinDriveProvider::checkDrives);
-}
-
-QString getPhysicalName(int driveNumber)
-{
-    return QString("\\\\.\\PhysicalDrive%0").arg(driveNumber);
-}
-
-HANDLE getPhysicalHandle(int driveNumber)
-{
-    HANDLE physicalHandle = INVALID_HANDLE_VALUE;
-    QString physicalPath = getPhysicalName(driveNumber);
-    physicalHandle = CreateFileA(physicalPath.toStdString().c_str(), GENERIC_READ, FILE_SHARE_WRITE, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
-    return physicalHandle;
+    return false;
 }
 
-bool WinDriveProvider::isMountable(int driveNumber)
+void WinDriveProvider::checkDrives()
 {
-    mDebug() << this->metaObject()->className() << "Checking whether " << getPhysicalName(driveNumber) << " is mountable";
-
-    HANDLE physicalHandle = getPhysicalHandle(driveNumber);
-    if (physicalHandle == INVALID_HANDLE_VALUE) {
-        mDebug() << this->metaObject()->className() << "Could not get physical handle for drive " << getPhysicalName(driveNumber);
-        return false;
-    }
-
-    DWORD size;
-    BYTE geometry[256];
-    bool ret = DeviceIoControl(physicalHandle, IOCTL_DISK_GET_DRIVE_GEOMETRY_EX, NULL, 0, geometry, sizeof(geometry), &size, NULL);
-    if (!ret || size <= 0) {
-        mDebug() << this->metaObject()->className() << "Could not get geometry for drive " << getPhysicalName(driveNumber);
-        CloseHandle(physicalHandle);
-        return false;
-    }
-
-    PDISK_GEOMETRY_EX diskGeometry = (PDISK_GEOMETRY_EX)(void *)geometry;
-    // Drive info
-    LONGLONG diskSize;
-    DWORD sectorSize;
-    DWORD sectorsPerTrack;
-    DWORD firstDataSector;
-    MEDIA_TYPE mediaType;
-
-    diskSize = diskGeometry->DiskSize.QuadPart;
-    sectorSize = diskGeometry->Geometry.BytesPerSector;
-    firstDataSector = MAXDWORD;
-    if (sectorSize < 512) {
-        mDebug() << this->metaObject()->className() << "Warning: Drive " << getPhysicalName(driveNumber) << " reports a sector size of " << sectorSize << " - Correcting to 512 bytes.";
-        sectorSize = 512;
-    }
-    sectorsPerTrack = diskGeometry->Geometry.SectorsPerTrack;
-    mediaType = diskGeometry->Geometry.MediaType;
-
-    BYTE layout[4096] = {0};
-    ret = DeviceIoControl(physicalHandle, IOCTL_DISK_GET_DRIVE_LAYOUT_EX, NULL, 0, layout, sizeof(layout), &size, NULL);
-    if (!ret || size <= 0) {
-        mDebug() << this->metaObject()->className() << "Could not get layout for drive " << getPhysicalName(driveNumber);
-        CloseHandle(physicalHandle);
-        return false;
-    }
-
-    PDRIVE_LAYOUT_INFORMATION_EX driveLayout = (PDRIVE_LAYOUT_INFORMATION_EX)(void *)layout;
-
-    switch (driveLayout->PartitionStyle) {
-    case PARTITION_STYLE_MBR:
-        mDebug() << this->metaObject()->className() << "MBR partition style";
-        for (int i = 0; i < driveLayout->PartitionCount; i++) {
-            if (driveLayout->PartitionEntry[i].Mbr.PartitionType != PARTITION_ENTRY_UNUSED) {
-                QVector<uint8_t> mbrMountable = {0x01, 0x04, 0x06, 0x07, 0x0b, 0x0c, 0x0e};
-                BYTE partType = driveLayout->PartitionEntry[i].Mbr.PartitionType;
-                mDebug() << this->metaObject()->className() << "Partition type: " << partType;
-                if (!mbrMountable.contains(partType)) {
-                    CloseHandle(physicalHandle);
-                    mDebug() << this->metaObject()->className() << getPhysicalName(driveNumber) << " is not mountable";
-                    return false;
-                }
-            }
-        }
-        break;
-    case PARTITION_STYLE_GPT:
-        mDebug() << this->metaObject()->className() << "GPT partition style";
-        for (int i = 0; i < driveLayout->PartitionCount; i++) {
-            if (memcmp(&driveLayout->PartitionEntry[i].Gpt.PartitionType, &PARTITION_MICROSOFT_DATA, sizeof(GUID)) != 0) {
-                CloseHandle(physicalHandle);
-                mDebug() << this->metaObject()->className() << getPhysicalName(driveNumber) << " is not mountable";
-                return false;
+    mDebug() << this->metaObject()->className() << "Looking for the drives";
+
+    QMap<int, WinDrive *> drives;
+    auto usbDeviceList = m_wmi->getUSBDeviceList();
+    for (auto it = usbDeviceList.cbegin(); it != usbDeviceList.cend(); it++) {
+        bool noLogicalDevice = false;
+        QStringList partitionList = m_wmi->getDevicePartitions(it.key());
+        for (const QString &partition : partitionList) {
+            if (m_wmi->getLogicalDisks(partition).isEmpty()) {
+                noLogicalDevice = true;
+                break;
             }
         }
-        break;
-    default:
-        mDebug() << this->metaObject()->className() << "Partition type: RAW";
-        break;
-    }
-
-    mDebug() << this->metaObject()->className() << getPhysicalName(driveNumber) << " is mountable";
 
-    CloseHandle(physicalHandle);
-    return true;
-}
-
-bool WinDriveProvider::describeDrive(int nDriveNumber, bool verbose)
-{
-    BOOL removable;
-    QString productVendor;
-    QString productId;
-    QString serialNumber;
-    uint64_t deviceBytes;
-    STORAGE_BUS_TYPE storageBus;
-
-    BOOL bResult = FALSE; // results flag
-    // DWORD dwRet = NO_ERROR;
-
-    // Format physical drive path (may be '\\.\PhysicalDrive0', '\\.\PhysicalDrive1' and so on).
-    QString strDrivePath = getPhysicalName(nDriveNumber);
-
-    // Get a handle to physical drive
-    HANDLE hDevice = ::CreateFile(strDrivePath.toStdWString().c_str(), 0, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL);
-
-    if (hDevice == INVALID_HANDLE_VALUE)
-        return false; //::GetLastError();
-
-    if (verbose)
-        mDebug() << this->metaObject()->className() << strDrivePath << "is present";
-
-    // Set the input data structure
-    STORAGE_PROPERTY_QUERY storagePropertyQuery;
-    ZeroMemory(&storagePropertyQuery, sizeof(STORAGE_PROPERTY_QUERY));
-    storagePropertyQuery.PropertyId = StorageDeviceProperty;
-    storagePropertyQuery.QueryType = PropertyStandardQuery;
-
-    // Get the necessary output buffer size
-    STORAGE_DESCRIPTOR_HEADER storageDescriptorHeader;
-    ZeroMemory(&storageDescriptorHeader, sizeof(STORAGE_DESCRIPTOR_HEADER));
-    DWORD dwBytesReturned = 0;
-    if (verbose)
-        mDebug() << this->metaObject()->className() << strDrivePath << "IOCTL_STORAGE_QUERY_PROPERTY";
-    if (!::DeviceIoControl(hDevice, IOCTL_STORAGE_QUERY_PROPERTY, &storagePropertyQuery, sizeof(STORAGE_PROPERTY_QUERY), &storageDescriptorHeader, sizeof(STORAGE_DESCRIPTOR_HEADER), &dwBytesReturned, NULL)) {
-        // dwRet = ::GetLastError();
-        ::CloseHandle(hDevice);
-        return false; // dwRet;
-    }
+        auto diskDrive = m_wmi->getDiskDriveInformation(it.key(), it.value());
+        if (diskDrive->model().isEmpty() || !diskDrive->size() || diskDrive->serialNumber().isEmpty()) {
+            continue;
+        }
 
-    // Alloc the output buffer
-    const DWORD dwOutBufferSize = storageDescriptorHeader.Size;
-    BYTE *pOutBuffer = new BYTE[dwOutBufferSize];
-    ZeroMemory(pOutBuffer, dwOutBufferSize);
-
-    if (verbose)
-        mDebug() << this->metaObject()->className() << strDrivePath << "IOCTL_STORAGE_QUERY_PROPERTY with a bigger buffer";
-    // Get the storage device descriptor
-    if (!(bResult = ::DeviceIoControl(hDevice, IOCTL_STORAGE_QUERY_PROPERTY, &storagePropertyQuery, sizeof(STORAGE_PROPERTY_QUERY), pOutBuffer, dwOutBufferSize, &dwBytesReturned, NULL))) {
-        // dwRet = ::GetLastError();
-        delete[] pOutBuffer;
-        ::CloseHandle(hDevice);
-        return false; // dwRet;
+        WinDrive *currentDrive = new WinDrive(this, diskDrive->model(), diskDrive->size(), noLogicalDevice, diskDrive->index(), diskDrive->serialNumber());
+        drives[diskDrive->index()] = currentDrive;
     }
 
-    // Now, the output buffer points to a STORAGE_DEVICE_DESCRIPTOR structure
-    // followed by additional info like vendor ID, product ID, serial number, and so on.
-    STORAGE_DEVICE_DESCRIPTOR *pDeviceDescriptor = (STORAGE_DEVICE_DESCRIPTOR *)pOutBuffer;
-    removable = pDeviceDescriptor->RemovableMedia;
-    if (pDeviceDescriptor->ProductIdOffset != 0)
-        productId = QString((char *)pOutBuffer + pDeviceDescriptor->ProductIdOffset).trimmed();
-    if (pDeviceDescriptor->VendorIdOffset != 0)
-        productVendor = QString((char *)pOutBuffer + pDeviceDescriptor->VendorIdOffset).trimmed();
-    if (pDeviceDescriptor->SerialNumberOffset != 0)
-        serialNumber = QString((char *)pOutBuffer + pDeviceDescriptor->SerialNumberOffset).trimmed();
-    storageBus = pDeviceDescriptor->BusType;
-
-    if (verbose)
-        mDebug() << this->metaObject()->className() << strDrivePath << "detected:" << productVendor << productId << (removable ? ", removable" : ", nonremovable") << (storageBus == BusTypeUsb ? "USB" : "notUSB");
-
-    if (!removable && storageBus != BusTypeUsb)
-        return false;
-
-    DISK_GEOMETRY pdg;
-    DWORD junk = 0; // discard results
-
-    if (verbose)
-        mDebug() << this->metaObject()->className() << strDrivePath << "IOCTL_DISK_GET_DRIVE_GEOMETRY";
-    bResult = DeviceIoControl(hDevice, // device to be queried
-                              IOCTL_DISK_GET_DRIVE_GEOMETRY, // operation to perform
-                              NULL,
-                              0, // no input buffer
-                              &pdg,
-                              sizeof(pdg), // output buffer
-                              &junk, // # bytes returned
-                              (LPOVERLAPPED)NULL); // synchronous I/O
-
-    if (!bResult || pdg.MediaType == Unknown)
-        return false;
+    // Update our list of drives and notify about added and removed drives
+    QList<int> driveIndexes = m_drives.keys();
+    for (auto it = drives.constBegin(); it != drives.constEnd(); it++) {
+        if (m_drives.contains(it.key()) && *m_drives[it.key()] == *it.value()) {
+            mDebug() << "Drive " << it.key() << " already exists";
+            it.value()->deleteLater();
+            driveIndexes.removeAll(it.key());
+            continue;
+        }
 
-    deviceBytes = pdg.Cylinders.QuadPart * pdg.TracksPerCylinder * pdg.SectorsPerTrack * pdg.BytesPerSector;
+        if (m_drives.contains(it.key())) {
+            mDebug() << "Replacing old drive in the list on index " << it.key();
+            emit driveRemoved(m_drives[it.key()]);
+            m_drives[it.key()]->deleteLater();
+            m_drives.remove(it.key());
+        }
 
-    // Do cleanup and return
-    if (verbose)
-        mDebug() << this->metaObject()->className() << strDrivePath << "cleanup, adding to the list";
-    delete[] pOutBuffer;
-    ::CloseHandle(hDevice);
+        mDebug() << "Adding new drive to the list with index " << it.key();
+        m_drives[it.key()] = it.value();
+        emit driveConnected(it.value());
 
-    WinDrive *currentDrive = new WinDrive(this, productVendor + " " + productId, deviceBytes, !isMountable(nDriveNumber), nDriveNumber, serialNumber);
-    if (m_drives.contains(nDriveNumber) && *m_drives[nDriveNumber] == *currentDrive) {
-        currentDrive->deleteLater();
-        return true;
+        driveIndexes.removeAll(it.key());
     }
 
-    if (m_drives.contains(nDriveNumber)) {
-        emit driveRemoved(m_drives[nDriveNumber]);
-        m_drives[nDriveNumber]->deleteLater();
+    // Remove our previously stored drives that were not present in the last check
+    for (int index : driveIndexes) {
+        mDebug() << "Removing old drive with index" << index;
+        emit driveRemoved(m_drives[index]);
+        m_drives[index]->deleteLater();
+        m_drives.remove(index);
     }
-
-    m_drives[nDriveNumber] = currentDrive;
-    emit driveConnected(currentDrive);
-
-    return true;
 }
 
 WinDrive::WinDrive(WinDriveProvider *parent, const QString &name, uint64_t size, bool containsLive, int device, const QString &serialNumber)
@@ -379,6 +216,11 @@ void WinDrive::restore()
     m_child->start(QIODevice::ReadOnly);
 }
 
+bool WinDrive::busy() const
+{
+    return (m_child && m_child->state() == QProcess::Running);
+}
+
 QString WinDrive::serialNumber() const
 {
     return m_serialNo;
@@ -455,7 +297,7 @@ void WinDrive::onReadyRead()
             Notifications::notify(tr("Finished!"), tr("Writing %1 was successful").arg(m_image->fullName()));
         } else {
             bool ok;
-            qreal bytes = line.toLongLong(&ok);
+            qint64 bytes = line.toLongLong(&ok);
             if (ok) {
                 if (bytes < 0)
                     m_progress->setValue(NAN);

src/app/windrivemanager.h

@@ -21,27 +21,28 @@
 #define WINDRIVEMANAGER_H
 
 #include "drivemanager.h"
+#include "libwmi/libwmi.h"
 
+#include <QAbstractNativeEventFilter>
 #include <QProcess>
 
 class WinDriveProvider;
 class WinDrive;
 
-class WinDriveProvider : public DriveProvider
+class WinDriveProvider : public DriveProvider, public QAbstractNativeEventFilter
 {
     Q_OBJECT
 public:
     WinDriveProvider(DriveManager *parent);
 
+    bool nativeEventFilter(const QByteArray &eventType, void *message, qintptr *result) override;
+
 public slots:
     void checkDrives();
 
 private:
-    QSet<int> findPhysicalDrive(char driveLetter);
-    bool describeDrive(int driveNumber, bool verbose);
-    bool isMountable(int driveNumber);
-
     QMap<int, WinDrive *> m_drives;
+    std::unique_ptr<LibWMI> m_wmi;
 };
 
 class WinDrive : public Drive
@@ -55,6 +56,7 @@ class WinDrive : public Drive
     Q_INVOKABLE virtual void cancel() override;
     Q_INVOKABLE virtual void restore() override;
 
+    bool busy() const;
     QString serialNumber() const;
 
     bool operator==(const WinDrive &o) const;

src/helper/win/CMakeLists.txt

@@ -26,6 +26,7 @@ remove_definitions(-std=c++17)
 target_link_libraries(helper
     Qt6::Core
     isomd5
+    libwmi
     ${LIBLZMA_LIBRARIES}
 )
 

src/helper/win/main.cpp

@@ -1,5 +1,6 @@
 /*
  * Fedora Media Writer
+ * Copyright (C) 2024 Jan Grulich <jgrulich@redhat.com>
  * Copyright (C) 2016 Martin Bříza <mbriza@redhat.com>
  *
  * This program is free software; you can redistribute it and/or
@@ -18,27 +19,28 @@
  */
 
 #include <QCoreApplication>
+#include <QLocale>
 #include <QTextStream>
 #include <QTranslator>
-#include <QLocale>
 
 #include "restorejob.h"
 #include "writejob.h"
 
-int main(int argc, char *argv[]) {
+int main(int argc, char *argv[])
+{
     QCoreApplication app(argc, argv);
 
     QTranslator translator;
-    if (translator.load(QLocale(QLocale().language(), QLocale().country()), QLatin1String(), QLatin1String(), ":/translations"))
-        app.installTranslator(&translator);   
-
-    if (app.arguments().count() == 3 && app.arguments()[1] == "restore") {
-        new RestoreJob(app.arguments()[2]);
+    if (translator.load(QLocale(), QLatin1String(), QLatin1String(), ":/translations")) {
+        app.installTranslator(&translator);
     }
-    else if (app.arguments().count() == 4 && app.arguments()[1] == "write") {
-        new WriteJob(app.arguments()[2], app.arguments()[3]);
-    }
-    else {
+
+    const QStringList args = app.arguments();
+    if (args.count() == 3 && args[1] == "restore") {
+        new RestoreJob(args[2], &app);
+    } else if (args.count() == 4 && args[1] == "write") {
+        new WriteJob(args[2], args[3], &app);
+    } else {
         QTextStream err(stderr);
         err << "Helper: Wrong arguments entered\n";
         return 1;

src/helper/win/restorejob.cpp

@@ -22,8 +22,8 @@
 #include <QTextStream>
 #include <QTimer>
 
-RestoreJob::RestoreJob(const QString &where)
-    : QObject(nullptr)
+RestoreJob::RestoreJob(const QString &where, QObject *parent)
+    : QObject(parent)
 {
     bool ok = false;
     m_where = where.toInt(&ok);
@@ -33,7 +33,8 @@ RestoreJob::RestoreJob(const QString &where)
         QTimer::singleShot(0, this, &RestoreJob::work);
 }
 
-void RestoreJob::work() {
+void RestoreJob::work()
+{
     m_diskpart.setProgram("diskpart.exe");
     m_diskpart.setProcessChannelMode(QProcess::ForwardedChannels);
 
@@ -50,8 +51,7 @@ void RestoreJob::work() {
 
     if (m_diskpart.waitForFinished()) {
         qApp->exit(0);
-    }
-    else {
+    } else {
         err << m_diskpart.readAllStandardError();
         err.flush();
         qApp->exit(1);

src/helper/win/restorejob.h

@@ -28,16 +28,16 @@ class RestoreJob : public QObject
 {
     Q_OBJECT
 public:
-    explicit RestoreJob(const QString &where);
+    explicit RestoreJob(const QString &where, QObject *parent);
 
 signals:
 
 private slots:
     void work();
 
 private:
-    QTextStream out { stdout };
-    QTextStream err { stderr };
+    QTextStream out{stdout};
+    QTextStream err{stderr};
 
     QProcess m_diskpart;
     int m_where;

src/helper/win/writejob.h

@@ -1,5 +1,6 @@
 /*
  * Fedora Media Writer
+ * Copyright (C) 2024 Jan Grulich <jgrulich@redhat.com>
  * Copyright (C) 2016 Martin Bříza <mbriza@redhat.com>
  *
  * This program is free software; you can redistribute it and/or
@@ -20,58 +21,48 @@
 #ifndef WRITEJOB_H
 #define WRITEJOB_H
 
-#include <QObject>
-#include <QTextStream>
-#include <QProcess>
+#include <libwmi/libwmi.h>
 
 #include <QFileSystemWatcher>
+#include <QObject>
+#include <QTextStream>
 
 #include <windows.h>
 
 #ifndef MEDIAWRITER_LZMA_LIMIT
 // 256MB memory limit for the decompressor
-# define MEDIAWRITER_LZMA_LIMIT (1024*1024*256)
+#define MEDIAWRITER_LZMA_LIMIT (1024 * 1024 * 256)
 #endif
 
 class WriteJob : public QObject
 {
     Q_OBJECT
 public:
-    explicit WriteJob(const QString &what, const QString &where);
+    explicit WriteJob(const QString &image, const QString &driveNumber, QObject *parent);
 
     static int staticOnMediaCheckAdvanced(void *data, long long offset, long long total);
     int onMediaCheckAdvanced(long long offset, long long total);
 
 private:
-    HANDLE openDrive(int physicalDriveNumber);
-    bool lockDrive(HANDLE drive);
-    bool removeMountPoints(uint diskNumber);
-    // bool dismountDrive(HANDLE drive, int diskNumber);
-    bool cleanDrive(uint diskNumber);
-
-    bool writeBlock(HANDLE drive, OVERLAPPED *overlap, char *data, uint size);
-
-    void unlockDrive(HANDLE drive);
-
+    bool check();
+    bool lockDrive(HANDLE driveHandle);
+    bool write();
+    bool writeCompressed(HANDLE driveHandle);
+    bool writePlain(HANDLE driveHandle);
 
 private slots:
-    void work();
     void onFileChanged(const QString &path);
+    void work();
 
-    bool write();
-    bool writeCompressed(HANDLE drive);
-    bool writePlain(HANDLE drive);
-    bool check();
 private:
-    QString what;
-    uint where;
-
-    QTextStream out { stdout };
-    QTextStream err { stderr };
+    QString m_image;
+    std::unique_ptr<LibWMI> m_wmi;
+    std::unique_ptr<LibWMIDiskDrive> m_wmiDiskDrive;
 
-    QFileSystemWatcher watcher { };
+    QTextStream m_out{stdout};
+    QTextStream m_err{stderr};
 
-    const int BLOCK_SIZE { 512 * 128 };
+    QFileSystemWatcher m_watcher;
 };
 
 #endif // WRITEJOB_H

src/lib/CMakeLists.txt

@@ -1 +1,4 @@
 add_subdirectory(isomd5)
+if (WIN32)
+    add_subdirectory(libwmi)
+endif()

src/lib/libwmi/CMakeLists.txt

@@ -0,0 +1,9 @@
+set(LIBWMI_SRCS
+    libwmi.cpp
+)
+
+add_library(libwmi STATIC ${LIBWMI_SRCS})
+
+target_link_libraries(libwmi
+    Qt6::Core
+)

src/lib/libwmi/libwmi.cpp

@@ -0,0 +1,361 @@
+/*
+ * Fedora Media Writer
+ * Copyright (C) 2024 Jan Grulich <jgrulich@redhat.com>
+ *
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU General Public License
+ * as published by the Free Software Foundation; either version 2
+ * of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+ */
+
+#include "libwmi.h"
+
+#include <QDebug>
+
+#include <comdef.h>
+#include <dbt.h>
+
+#pragma comment(lib, "wbemuuid.lib")
+
+LibWMI::LibWMI(QObject *parent)
+    : QObject(parent)
+{
+    HRESULT res = S_OK;
+    // This needs to be initialized here before any RPC communication occurs
+    res = CoInitializeSecurity(NULL, -1, NULL, NULL, RPC_C_AUTHN_LEVEL_DEFAULT, RPC_C_IMP_LEVEL_IMPERSONATE, NULL, EOAC_NONE, 0);
+    if (FAILED(res)) {
+        _com_error err(res);
+        qWarning() << "Failed to initialize security. Error = " << err.ErrorMessage();
+        return;
+    }
+
+    res = CoCreateInstance(CLSID_WbemAdministrativeLocator, NULL, CLSCTX_INPROC_SERVER, IID_IWbemLocator, reinterpret_cast<void **>(&m_IWbemLocator));
+    if (FAILED(res)) {
+        _com_error err(res);
+        qWarning() << "Failed to create IWbemLocator object. Error = " << QString::fromWCharArray(err.ErrorMessage());
+        return;
+    }
+
+    res = m_IWbemLocator->ConnectServer(_bstr_t(L"root\\cimv2"), NULL, NULL, NULL, 0, NULL, NULL, &m_IWbemServices);
+    if (FAILED(res)) {
+        _com_error err(res);
+        qWarning() << "Could not connect to WMI. Error = " << QString::fromWCharArray(err.ErrorMessage());
+        return;
+    }
+
+    initialized = true;
+}
+
+LibWMI::~LibWMI()
+{
+    if (m_IWbemLocator) {
+        m_IWbemLocator->Release();
+    }
+    if (m_IWbemServices) {
+        m_IWbemServices->Release();
+    }
+    CoUninitialize();
+}
+
+QMap<quint32, QString> LibWMI::getUSBDeviceList()
+{
+    QMap<quint32, QString> result;
+    if (!initialized) {
+        return result;
+    }
+
+    HRESULT res = S_OK;
+    IEnumWbemClassObject *pEnumDiskObjects = NULL;
+
+    res = m_IWbemServices->ExecQuery(_bstr_t(L"WQL"), _bstr_t(L"SELECT * FROM Win32_DiskDrive WHERE InterfaceType = 'USB'"), WBEM_FLAG_RETURN_IMMEDIATELY, NULL, &pEnumDiskObjects);
+    if (FAILED(res)) {
+        _com_error err(res);
+        qWarning() << "WMI query failed. Error = " << QString::fromWCharArray(err.ErrorMessage());
+        return result;
+    }
+
+    while (true) {
+        IWbemClassObject *pDiskObject = NULL;
+        ULONG uReturn = 0;
+        pEnumDiskObjects->Next(WBEM_INFINITE, 1, &pDiskObject, &uReturn);
+        if (uReturn == 0) {
+            break;
+        }
+
+        // Fetch disk information
+        quint32 index;
+        QString deviceID;
+        VARIANT var;
+
+        if ((pDiskObject->Get(_bstr_t(L"Index"), 0, &var, 0, 0)) == WBEM_S_NO_ERROR) {
+            if (var.vt == VT_I4) {
+                index = var.intVal;
+                qDebug() << "Disk Index: " << index;
+            } else if (var.vt == VT_UI4) {
+                index = var.uintVal;
+                qDebug() << "Disk Index: " << index;
+            }
+            VariantClear(&var);
+        }
+
+        if ((pDiskObject->Get(_bstr_t(L"DeviceID"), 0, &var, 0, 0)) == WBEM_S_NO_ERROR) {
+            if (var.vt == VT_BSTR) {
+                deviceID = QString::fromWCharArray(var.bstrVal);
+                qDebug() << "Device ID" << deviceID;
+            }
+            VariantClear(&var);
+        }
+        pDiskObject->Release();
+
+        result.insert(index, deviceID);
+    }
+    pEnumDiskObjects->Release();
+
+    return result;
+}
+
+QStringList LibWMI::getDevicePartitions(quint32 index)
+{
+    QStringList result;
+    if (!initialized) {
+        return result;
+    }
+
+    HRESULT res = S_OK;
+    IEnumWbemClassObject *pPartitionObjects = NULL;
+    std::wstring partitionQuery = L"SELECT * FROM Win32_DiskPartition WHERE DiskIndex = " + std::to_wstring(index);
+
+    res = m_IWbemServices->ExecQuery(_bstr_t(L"WQL"), _bstr_t(partitionQuery.c_str()), WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY, NULL, &pPartitionObjects);
+    if (FAILED(res)) {
+        _com_error err(res);
+        qWarning() << "Query for disk partitions failed. Error = " << QString::fromWCharArray(err.ErrorMessage());
+        return result;
+    }
+
+    while (true) {
+        IWbemClassObject *pPartitionObject = NULL;
+        ULONG uReturnPartition = 0;
+        pPartitionObjects->Next(WBEM_INFINITE, 1, &pPartitionObject, &uReturnPartition);
+        if (uReturnPartition == 0) {
+            break;
+        }
+
+        QString deviceID;
+        VARIANT partitionDeviceID;
+        if ((pPartitionObject->Get(_bstr_t(L"DeviceID"), 0, &partitionDeviceID, 0, 0)) == WBEM_S_NO_ERROR) {
+            if (partitionDeviceID.vt == VT_BSTR) {
+                deviceID = QString::fromWCharArray(partitionDeviceID.bstrVal);
+                qDebug() << "Partition device ID " << deviceID;
+            }
+            VariantClear(&partitionDeviceID);
+        }
+        pPartitionObject->Release();
+        if (!deviceID.isEmpty()) {
+            result << deviceID;
+        }
+    }
+    pPartitionObjects->Release();
+
+    return result;
+}
+
+QStringList LibWMI::getLogicalDisks(const QString &partitionID)
+{
+    QStringList result;
+    if (!initialized) {
+        return result;
+    }
+
+    HRESULT res = S_OK;
+    std::wstring partitionToLogicalQuery = L"ASSOCIATORS OF {Win32_DiskPartition.DeviceID='" + partitionID.toStdWString() + L"'} WHERE AssocClass = Win32_LogicalDiskToPartition";
+    IEnumWbemClassObject *pLogicalDiskObjects = NULL;
+    res = m_IWbemServices->ExecQuery(_bstr_t(L"WQL"), _bstr_t(partitionToLogicalQuery.c_str()), WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY, NULL, &pLogicalDiskObjects);
+    if (FAILED(res)) {
+        _com_error err(res);
+        qWarning() << "Query for logical disks failed. Error = " << QString::fromWCharArray(err.ErrorMessage());
+        return result;
+    }
+
+    while (true) {
+        IWbemClassObject *pLogicalDiskObject = NULL;
+        ULONG uReturnLogicalDisk = 0;
+        pLogicalDiskObjects->Next(WBEM_INFINITE, 1, &pLogicalDiskObject, &uReturnLogicalDisk);
+        if (uReturnLogicalDisk == 0) {
+            break;
+        }
+
+        QString deviceID;
+        VARIANT logicalDiskDeviceID;
+        if ((pLogicalDiskObject->Get(_bstr_t(L"DeviceID"), 0, &logicalDiskDeviceID, 0, 0)) == WBEM_S_NO_ERROR) {
+            if (logicalDiskDeviceID.vt == VT_BSTR) {
+                deviceID = QString::fromWCharArray(logicalDiskDeviceID.bstrVal);
+                qDebug() << "Logical disk device ID " << deviceID;
+            }
+            VariantClear(&logicalDiskDeviceID);
+        }
+        pLogicalDiskObject->Release();
+        if (!deviceID.isEmpty()) {
+            result << deviceID;
+        }
+    }
+    pLogicalDiskObjects->Release();
+
+    return result;
+}
+
+std::unique_ptr<LibWMIDiskDrive> LibWMI::getDiskDriveInformation(quint32 index, const QString &deviceID)
+{
+    std::unique_ptr<LibWMIDiskDrive> result = std::make_unique<LibWMIDiskDrive>(index, deviceID);
+    if (!initialized) {
+        return result;
+    }
+
+    HRESULT res = S_OK;
+    IEnumWbemClassObject *pEnumDiskObjects = NULL;
+
+    std::wstring deviceQuery = L"SELECT * FROM Win32_DiskDrive WHERE Index = " + std::to_wstring(index);
+    res = m_IWbemServices->ExecQuery(_bstr_t(L"WQL"), _bstr_t(deviceQuery.c_str()), WBEM_FLAG_RETURN_IMMEDIATELY, NULL, &pEnumDiskObjects);
+    if (FAILED(res)) {
+        _com_error err(res);
+        qWarning() << "WMI query failed. Error = " << QString::fromWCharArray(err.ErrorMessage());
+        return result;
+    }
+
+    while (true) {
+        IWbemClassObject *pDiskObject = NULL;
+        ULONG uReturn = 0;
+        pEnumDiskObjects->Next(WBEM_INFINITE, 1, &pDiskObject, &uReturn);
+        if (uReturn == 0) {
+            break;
+        }
+
+        VARIANT var;
+        if (result->deviceID().isEmpty()) {
+            if ((pDiskObject->Get(_bstr_t(L"DeviceID"), 0, &var, 0, 0)) == WBEM_S_NO_ERROR) {
+                if (var.vt == VT_BSTR) {
+                    result->setDeviceID(QString::fromWCharArray(var.bstrVal));
+                    qDebug() << "DeviceID " << result->deviceID();
+                }
+                VariantClear(&var);
+            }
+        }
+
+        if ((pDiskObject->Get(_bstr_t(L"Model"), 0, &var, 0, 0)) == WBEM_S_NO_ERROR) {
+            if (var.vt == VT_BSTR) {
+                result->setModel(QString::fromWCharArray(var.bstrVal));
+                qDebug() << "Disk model: " << result->model();
+            }
+            VariantClear(&var);
+        }
+
+        if ((pDiskObject->Get(_bstr_t(L"Size"), 0, &var, 0, 0)) == WBEM_S_NO_ERROR) {
+            if (var.vt == VT_BSTR) {
+                result->setSize(QString::fromWCharArray(var.bstrVal).toULongLong());
+                qDebug() << "Size " << result->size();
+            } else if (var.vt == VT_I4) {
+                result->setSize(var.intVal);
+                qDebug() << "Size " << result->size();
+            } else if (var.vt == VT_UI4) {
+                result->setSize(var.uintVal);
+                qDebug() << "Size " << result->size();
+            }
+            VariantClear(&var);
+        }
+
+        if ((pDiskObject->Get(_bstr_t(L"BytesPerSector"), 0, &var, 0, 0)) == WBEM_S_NO_ERROR) {
+            if (var.vt == VT_BSTR) {
+                result->setSectorSize(QString::fromWCharArray(var.bstrVal).toULongLong());
+                qDebug() << "Sector size " << result->sectorSize();
+            } else if (var.vt == VT_I4) {
+                result->setSectorSize(var.intVal);
+                qDebug() << "Sector size " << result->sectorSize();
+            } else if (var.vt == VT_UI4) {
+                result->setSectorSize(var.uintVal);
+                qDebug() << "Sector size " << result->sectorSize();
+            }
+            VariantClear(&var);
+        }
+
+        if ((pDiskObject->Get(_bstr_t(L"SerialNumber"), 0, &var, 0, 0)) == WBEM_S_NO_ERROR) {
+            if (var.vt == VT_BSTR) {
+                result->setSerialNumber(QString::fromWCharArray(var.bstrVal));
+                qDebug() << "Serial number " << result->serialNumber();
+            }
+            VariantClear(&var);
+        }
+        pDiskObject->Release();
+    }
+    pEnumDiskObjects->Release();
+
+    return result;
+}
+
+LibWMIDiskDrive::LibWMIDiskDrive(quint32 index, const QString &deviceID)
+    : m_index(index)
+    , m_deviceID(deviceID)
+{
+}
+
+quint32 LibWMIDiskDrive::index() const
+{
+    return m_index;
+}
+
+QString LibWMIDiskDrive::deviceID() const
+{
+    return m_deviceID;
+}
+
+void LibWMIDiskDrive::setDeviceID(const QString &deviceID)
+{
+    m_deviceID = deviceID;
+}
+
+QString LibWMIDiskDrive::model() const
+{
+    return m_model;
+}
+
+void LibWMIDiskDrive::setModel(const QString &model)
+{
+    m_model = model;
+}
+
+quint64 LibWMIDiskDrive::size() const
+{
+    return m_size;
+}
+
+void LibWMIDiskDrive::setSize(quint64 size)
+{
+    m_size = size;
+}
+
+QString LibWMIDiskDrive::serialNumber() const
+{
+    return m_serialNumber;
+}
+
+void LibWMIDiskDrive::setSerialNumber(const QString &serialNumber)
+{
+    m_serialNumber = serialNumber;
+}
+
+quint32 LibWMIDiskDrive::sectorSize()
+{
+    return m_sectorSize;
+}
+
+void LibWMIDiskDrive::setSectorSize(quint32 sectorSize)
+{
+    m_sectorSize = sectorSize;
+}

src/lib/libwmi/libwmi.h

@@ -0,0 +1,81 @@
+/*
+ * Fedora Media Writer
+ * Copyright (C) 2024 Jan Grulich <jgrulich@redhat.com>
+ *
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU General Public License
+ * as published by the Free Software Foundation; either version 2
+ * of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+ */
+
+#ifndef LIBWMI_H
+#define LIBWMI_H
+
+#include <QObject>
+#include <wbemidl.h>
+
+class LibWMIDiskDrive;
+
+class LibWMI : public QObject
+{
+    Q_OBJECT
+public:
+    LibWMI(QObject *parent);
+    ~LibWMI();
+
+    // Returns a map<index, deviceID> with list of devices
+    QMap<quint32, QString> getUSBDeviceList();
+    // Returns a list of partition IDs for device with given index
+    QStringList getDevicePartitions(quint32 index);
+    // Returns a list of logical disks IDs for given partition
+    QStringList getLogicalDisks(const QString &partitionID);
+    // Returns information about disk drive
+    std::unique_ptr<LibWMIDiskDrive> getDiskDriveInformation(quint32 index, const QString &deviceID = QString());
+
+private:
+    bool initialized = false;
+    IWbemLocator *m_IWbemLocator = NULL;
+    IWbemServices *m_IWbemServices = NULL;
+};
+
+class LibWMIDiskDrive
+{
+public:
+    LibWMIDiskDrive(quint32 index, const QString &deviceID = QString());
+
+    quint32 index() const;
+
+    QString deviceID() const;
+    void setDeviceID(const QString &deviceID);
+
+    QString model() const;
+    void setModel(const QString &model);
+
+    quint64 size() const;
+    void setSize(quint64 size);
+
+    QString serialNumber() const;
+    void setSerialNumber(const QString &serialNumber);
+
+    quint32 sectorSize();
+    void setSectorSize(quint32 sectorSize);
+
+private:
+    quint32 m_index = 0;
+    quint32 m_sectorSize = 0;
+    quint64 m_size = 0;
+    QString m_deviceID;
+    QString m_model;
+    QString m_serialNumber;
+};
+
+#endif // LIBWMI_H