From af631b7ad81d6b50aa57e0e09ff8a5f24341714d Mon Sep 17 00:00:00 2001 From: Jan Grulich Date: Fri, 4 Oct 2024 16:12:21 +0200 Subject: [PATCH] WIP: rewrite Windows code --- .github/workflows/clang-format-check.yml | 2 +- src/app/CMakeLists.txt | 2 +- src/app/main.cpp | 1 + src/app/windrivemanager.cpp | 298 +++--------- src/app/windrivemanager.h | 12 +- src/helper/win/CMakeLists.txt | 1 + src/helper/win/main.cpp | 24 +- src/helper/win/restorejob.cpp | 4 +- src/helper/win/restorejob.h | 2 +- src/helper/win/writejob.cpp | 595 ++++++++++++----------- src/helper/win/writejob.h | 43 +- src/lib/CMakeLists.txt | 3 + src/lib/libwmi/CMakeLists.txt | 9 + src/lib/libwmi/libwmi.cpp | 361 ++++++++++++++ src/lib/libwmi/libwmi.h | 81 +++ 15 files changed, 875 insertions(+), 563 deletions(-) create mode 100644 src/lib/libwmi/CMakeLists.txt create mode 100644 src/lib/libwmi/libwmi.cpp create mode 100644 src/lib/libwmi/libwmi.h diff --git a/.github/workflows/clang-format-check.yml b/.github/workflows/clang-format-check.yml index 5d45c7e1..9c856529 100644 --- a/.github/workflows/clang-format-check.yml +++ b/.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' diff --git a/src/app/CMakeLists.txt b/src/app/CMakeLists.txt index 9449894a..06613cf7 100644 --- a/src/app/CMakeLists.txt +++ b/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) diff --git a/src/app/main.cpp b/src/app/main.cpp index e7512088..eb09fd7a 100644 --- a/src/app/main.cpp +++ b/src/app/main.cpp @@ -1,5 +1,6 @@ /* * Fedora Media Writer + * Copyright (C) 2024 Jan Grulich * Copyright (C) 2016 Martin Bříza * * This program is free software; you can redistribute it and/or diff --git a/src/app/windrivemanager.cpp b/src/app/windrivemanager.cpp index 7fb1035c..1ae38fe5 100644 --- a/src/app/windrivemanager.cpp +++ b/src/app/windrivemanager.cpp @@ -1,6 +1,6 @@ /* * Fedora Media Writer - * Copyright (C) 2022 Jan Grulich + * Copyright (C) 2022-2024 Jan Grulich * Copyright (C) 2011-2022 Pete Batard * Copyright (C) 2016 Martin Bříza * @@ -22,262 +22,98 @@ #include "windrivemanager.h" #include "notifications.h" -#include #include #include -#include -#define INITGUID -#include +#include -#include -#include - -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(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(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 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 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 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(); + 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 +215,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 +296,8 @@ 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); + mWarning() << "Written: " << bytes; if (ok) { if (bytes < 0) m_progress->setValue(NAN); diff --git a/src/app/windrivemanager.h b/src/app/windrivemanager.h index 3d1ac43d..c99b18d9 100644 --- a/src/app/windrivemanager.h +++ b/src/app/windrivemanager.h @@ -21,27 +21,28 @@ #define WINDRIVEMANAGER_H #include "drivemanager.h" +#include "libwmi/libwmi.h" +#include #include 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 findPhysicalDrive(char driveLetter); - bool describeDrive(int driveNumber, bool verbose); - bool isMountable(int driveNumber); - QMap m_drives; + std::unique_ptr 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; diff --git a/src/helper/win/CMakeLists.txt b/src/helper/win/CMakeLists.txt index 50e8d5f7..26e8c4fb 100644 --- a/src/helper/win/CMakeLists.txt +++ b/src/helper/win/CMakeLists.txt @@ -26,6 +26,7 @@ remove_definitions(-std=c++17) target_link_libraries(helper Qt6::Core isomd5 + libwmi ${LIBLZMA_LIBRARIES} ) diff --git a/src/helper/win/main.cpp b/src/helper/win/main.cpp index 4152eec7..0053533f 100644 --- a/src/helper/win/main.cpp +++ b/src/helper/win/main.cpp @@ -1,5 +1,6 @@ /* * Fedora Media Writer + * Copyright (C) 2024 Jan Grulich * Copyright (C) 2016 Martin Bříza * * This program is free software; you can redistribute it and/or @@ -18,27 +19,28 @@ */ #include +#include #include #include -#include #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; diff --git a/src/helper/win/restorejob.cpp b/src/helper/win/restorejob.cpp index b4633756..78881ed3 100644 --- a/src/helper/win/restorejob.cpp +++ b/src/helper/win/restorejob.cpp @@ -22,8 +22,8 @@ #include #include -RestoreJob::RestoreJob(const QString &where) - : QObject(nullptr) +RestoreJob::RestoreJob(const QString &where, QObject *parent) + : QObject(parent) { bool ok = false; m_where = where.toInt(&ok); diff --git a/src/helper/win/restorejob.h b/src/helper/win/restorejob.h index a46ada9b..7b52533d 100644 --- a/src/helper/win/restorejob.h +++ b/src/helper/win/restorejob.h @@ -28,7 +28,7 @@ class RestoreJob : public QObject { Q_OBJECT public: - explicit RestoreJob(const QString &where); + explicit RestoreJob(const QString &where, QObject *parent); signals: diff --git a/src/helper/win/writejob.cpp b/src/helper/win/writejob.cpp index 30947f30..2c57b87c 100644 --- a/src/helper/win/writejob.cpp +++ b/src/helper/win/writejob.cpp @@ -1,5 +1,6 @@ /* * Fedora Media Writer + * Copyright (C) 2024 Jan Grulich * Copyright (C) 2016 Martin Bříza * * This program is free software; you can redistribute it and/or @@ -20,83 +21,92 @@ #include "writejob.h" #include -#include -#include -#include #include -#include +#include #include -#include +#include +#include #include -#include #include #include "isomd5/libcheckisomd5.h" +static constexpr qint64 BLOCK_SIZE{512 * 128}; -WriteJob::WriteJob(const QString &what, const QString &where) - : QObject(nullptr), what(what) +WriteJob::WriteJob(const QString &image, const QString &driveNumber, QObject *parent) + : QObject(parent) + , m_image(image) { - bool ok = false; - this->where = where.toInt(&ok); + const int wmiDriveNumber = driveNumber.toInt(); - if (what.endsWith(".part")) { - connect(&watcher, &QFileSystemWatcher::fileChanged, this, &WriteJob::onFileChanged); - watcher.addPath(what); - } - else { + m_wmi = std::make_unique(this); + m_wmiDiskDrive = m_wmi->getDiskDriveInformation(wmiDriveNumber); + + if (m_image.endsWith(".part")) { + connect(&m_watcher, &QFileSystemWatcher::fileChanged, this, &WriteJob::onFileChanged); + m_watcher.addPath(m_image); + } else { QTimer::singleShot(0, this, &WriteJob::work); } } -int WriteJob::staticOnMediaCheckAdvanced(void *data, long long offset, long long total) { - return ((WriteJob*)data)->onMediaCheckAdvanced(offset, total); +int WriteJob::staticOnMediaCheckAdvanced(void *data, long long offset, long long total) +{ + return ((WriteJob *)data)->onMediaCheckAdvanced(offset, total); } -int WriteJob::onMediaCheckAdvanced(long long offset, long long total) { +int WriteJob::onMediaCheckAdvanced(long long offset, long long total) +{ Q_UNUSED(total); - out << offset << "\n"; - out.flush(); + m_out << offset << "\n"; + m_out.flush(); return 0; } -HANDLE WriteJob::openDrive(int physicalDriveNumber) { - HANDLE hVol; - QString drivePath = QString("\\\\.\\PhysicalDrive%0").arg(physicalDriveNumber); - - hVol = CreateFile(drivePath.toStdWString().c_str(), GENERIC_READ | GENERIC_WRITE, FILE_SHARE_READ | FILE_SHARE_WRITE, 0, OPEN_EXISTING, 0, 0); +void WriteJob::cleanDrive(HANDLE driveHandle) +{ + QFile drive; + drive.open(_open_osfhandle(reinterpret_cast(driveHandle), 0), QIODevice::WriteOnly | QIODevice::Unbuffered, QFile::AutoCloseHandle); + if (!drive.isOpen()) { + m_err << tr("Failed to open the device for writing"); + m_err.flush(); + qApp->exit(1); + return; + } - if( hVol == INVALID_HANDLE_VALUE ) { - TCHAR message[256]; - FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM, NULL, GetLastError(), MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), message, 255, NULL); - err << tr("Couldn't open the drive for writing") << " (" << QString::fromWCharArray(message).trimmed() << ")\n"; - err.flush(); - return hVol; + void *buffer = NULL; + buffer = VirtualAlloc(NULL, BLOCK_SIZE, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE); + if (buffer == NULL) { + m_err << tr("Failed to allocate the buffer"); + m_err.flush(); + qApp->exit(1); + return; } - return hVol; + memset(buffer, 0, BLOCK_SIZE); + drive.write(static_cast(buffer), BLOCK_SIZE); + drive.close(); } -bool WriteJob::lockDrive(HANDLE drive) { +bool WriteJob::lockDrive(HANDLE driveHandle) +{ int attempts = 0; DWORD status; while (true) { - if (!DeviceIoControl(drive, FSCTL_LOCK_VOLUME, NULL, 0, NULL, 0, &status, NULL)) { + if (!DeviceIoControl(driveHandle, FSCTL_LOCK_VOLUME, NULL, 0, NULL, 0, &status, NULL)) { attempts++; - } - else { + } else { return true; } if (attempts == 10) { TCHAR message[256]; FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM, NULL, GetLastError(), MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), message, 255, NULL); - - err << tr("Couldn't lock the drive") << " (" << QString::fromWCharArray(message).trimmed() << ")\n"; - err.flush(); + m_err << tr("Couldn't lock the drive") << " (" << QString::fromWCharArray(message).trimmed() << ")\n"; + m_err.flush(); break; } @@ -106,118 +116,63 @@ bool WriteJob::lockDrive(HANDLE drive) { return false; } -bool WriteJob::removeMountPoints(uint diskNumber) { - DWORD drives = ::GetLogicalDrives(); - - for (char i = 0; i < 26; i++) { - if (drives & (1 << i)) { - char currentDrive = 'A' + i; - QString drivePath = QString("\\\\.\\%1:").arg(currentDrive); - - HANDLE hDevice = ::CreateFile(drivePath.toStdWString().c_str(), GENERIC_READ | GENERIC_WRITE, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL); - - DWORD bytesReturned; - VOLUME_DISK_EXTENTS vde; // TODO FIXME: handle ERROR_MORE_DATA (this is an extending structure) - BOOL bResult = DeviceIoControl(hDevice, IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS, NULL, 0, &vde, sizeof(vde), &bytesReturned, NULL); - - if (bResult) { - for (uint j = 0; j < vde.NumberOfDiskExtents; j++) { - if (vde.Extents[j].DiskNumber == diskNumber) { - QString volumePath = QString("%1:\\").arg(currentDrive); - - CloseHandle(hDevice); - hDevice = nullptr; - - if (!DeleteVolumeMountPointA(volumePath.toStdString().c_str())) { - TCHAR message[256]; - FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM, NULL, GetLastError(), MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), message, 255, NULL); - err << tr("Couldn't remove the drive %1:").arg(currentDrive) << " (" << QString::fromWCharArray(message).trimmed() << "\n"; - err.flush(); - return false; - } - - break; - } - } - } - if (hDevice) - CloseHandle(hDevice); - } - } - - return true; -} - -bool WriteJob::cleanDrive(uint driveNumber) { - QProcess diskpart; - diskpart.setProgram("diskpart.exe"); - diskpart.setProcessChannelMode(QProcess::ForwardedChannels); - - diskpart.start(QIODevice::ReadWrite); - - diskpart.write(qPrintable(QString("select disk %0\r\n").arg(driveNumber))); - diskpart.write("clean\r\n"); - // for some reason this works (tm) - diskpart.write("create part pri\r\n"); - diskpart.write("clean\r\n"); - diskpart.write("exit\r\n"); - - diskpart.waitForFinished(); - - if (diskpart.exitCode() == 0) { - // as advised in the diskpart documentation - QThread::sleep(15); +HANDLE WriteJob::openDrive() +{ + HANDLE hVol; + hVol = CreateFile(m_wmiDiskDrive->deviceID().toStdWString().c_str(), GENERIC_WRITE, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, FILE_FLAG_WRITE_THROUGH | FILE_FLAG_NO_BUFFERING, NULL); - return true; + if (hVol == INVALID_HANDLE_VALUE) { + TCHAR message[256]; + FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM, NULL, GetLastError(), MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), message, 255, NULL); + m_err << tr("Couldn't open the drive for writing") << " (" << QString::fromWCharArray(message).trimmed() << ")\n"; + m_err.flush(); + return hVol; } - return false; + return hVol; } -bool WriteJob::writeBlock(HANDLE drive, OVERLAPPED *overlap, char *data, uint size) { - DWORD bytesWritten; - - if (!WriteFile(drive, data, size, &bytesWritten, overlap)) { - DWORD Errorcode = GetLastError(); - if (Errorcode == ERROR_IO_PENDING) { - WaitForSingleObject(overlap->hEvent, INFINITE); - } - else { - TCHAR message[256]; - FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM, NULL, GetLastError(), MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), message, 255, NULL); - err << tr("Destination drive is not writable") << " (" << QString::fromWCharArray(message).trimmed() << ")\n"; - err.flush(); - return false; +void WriteJob::unmountVolumes() +{ + m_out << "Unmounting volumes"; + m_out.flush(); + + QStringList partitions = m_wmi->getDevicePartitions(m_wmiDiskDrive->index()); + for (const QString &partition : partitions) { + QStringList volumes = m_wmi->getLogicalDisks(partition); + for (const QString &volumeID : volumes) { + const QString volumePath = QString("\\\\.\\%1").arg(volumeID); + m_out << "Unmounting volume: " << volumePath; + m_out.flush(); + + DWORD ret; + HANDLE volume = CreateFile(volumePath.toStdWString().c_str(), GENERIC_READ | GENERIC_WRITE, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL); + if (volume == INVALID_HANDLE_VALUE) { + TCHAR message[256]; + FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM, NULL, GetLastError(), MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), message, 255, NULL); + m_err << tr("Failed to unmount the drive %1:").arg(volumeID) << " (" << QString::fromWCharArray(message).trimmed() << ")\n"; + m_err.flush(); + return; + } + if (!DeviceIoControl(volume, FSCTL_DISMOUNT_VOLUME, NULL, 0, NULL, 0, &ret, NULL)) { + TCHAR message[256]; + FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM, NULL, GetLastError(), MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), message, 255, NULL); + m_err << tr("Failed to unmount the drive %1:").arg(volumeID) << " (" << QString::fromWCharArray(message).trimmed() << ")\n"; + m_err.flush(); + } + CloseHandle(volume); } } - - if (bytesWritten != size) { - err << tr("Destination drive is not writable") << "\n"; - err.flush(); - return false; - } - - return true; } - -void WriteJob::unlockDrive(HANDLE drive) { - DWORD status; - if (!DeviceIoControl(drive, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &status, NULL)) { - TCHAR message[256]; - FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM, NULL, GetLastError(), MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), message, 255, NULL); - err << tr("Couldn't unlock the drive") << " (" << QString::fromWCharArray(message).trimmed() << ")\n"; - err.flush(); - } -} - -void WriteJob::work() { +void WriteJob::work() +{ if (!write()) { - out << "0\n"; - out.flush(); - QThread::sleep(5); - if (!write()) - return; + // m_out << "0\n"; + // m_out.flush(); + // QThread::sleep(5); + // if (!write()) + return; } if (!check()) @@ -226,193 +181,253 @@ void WriteJob::work() { qApp->exit(0); } -void WriteJob::onFileChanged(const QString &path) { +void WriteJob::onFileChanged(const QString &path) +{ if (QFile::exists(path)) return; QRegularExpression reg("[.]part$"); - what = what.replace(reg, ""); + m_image = m_image.replace(reg, ""); - out << "WRITE\n"; - out.flush(); + m_out << "WRITE\n"; + m_out.flush(); work(); } -bool WriteJob::write() { - removeMountPoints(where); - cleanDrive(where); +bool WriteJob::write() +{ + unmountVolumes(); - HANDLE drive = openDrive(where); + HANDLE drive = openDrive(); if (!lockDrive(drive)) { qApp->exit(1); return false; } + // cleanDrive(drive); + + // drive = openDrive(); + // if (!lockDrive(drive)) { + // qApp->exit(1); + // return false; + // } + // if (m_image.endsWith(".xz")) + // return writeCompressed(drive); + // else + bool ret = writePlain(drive); + if (!ret) { + qApp->exit(1); + } + CloseHandle(drive); - if (what.endsWith(".xz")) - return writeCompressed(drive); - else - return writePlain(drive); + return ret; } -bool WriteJob::writeCompressed(HANDLE drive) { - qint64 totalRead = 0; - - lzma_stream strm = LZMA_STREAM_INIT; - lzma_ret ret; - - uint8_t *inBuffer = new uint8_t[BLOCK_SIZE]; - uint8_t *outBuffer = new uint8_t[BLOCK_SIZE]; +bool WriteJob::writeCompressed(HANDLE drive) +{ + // qint64 totalRead = 0; + + // lzma_stream strm = LZMA_STREAM_INIT; + // lzma_ret ret; + + // uint8_t *inBuffer = new uint8_t[BLOCK_SIZE]; + // uint8_t *outBuffer = new uint8_t[BLOCK_SIZE]; + + // QFile file(m_image); + // file.open(QIODevice::ReadOnly); + + // ret = lzma_stream_decoder(&strm, MEDIAWRITER_LZMA_LIMIT, LZMA_CONCATENATED); + // if (ret != LZMA_OK) { + // m_err << tr("Failed to start decompressing."); + // return false; + // } + + // strm.next_in = inBuffer; + // strm.avail_in = 0; + // strm.next_out = outBuffer; + // strm.avail_out = BLOCK_SIZE; + + // OVERLAPPED osWrite; + // memset(&osWrite, 0, sizeof(osWrite)); + // osWrite.hEvent = 0; + + // while (true) { + // if (strm.avail_in == 0) { + // qint64 len = file.read((char*) inBuffer, BLOCK_SIZE); + // totalRead += len; + + // strm.next_in = inBuffer; + // strm.avail_in = len; + + // m_out << totalRead << "\n"; + // m_out.flush(); + // } + + // ret = lzma_code(&strm, strm.avail_in == 0 ? LZMA_FINISH : LZMA_RUN); + // if (ret == LZMA_STREAM_END) { + // if (!writeBlock(drive, &osWrite, (char *) outBuffer, BLOCK_SIZE - strm.avail_out)) { + // qApp->exit(1); + // CloseHandle(drive); + // return false; + // } + + // if (osWrite.Offset + BLOCK_SIZE < osWrite.Offset) + // osWrite.OffsetHigh++; + // osWrite.Offset += BLOCK_SIZE; + + // CloseHandle(drive); + + // return true; + // } + // if (ret != LZMA_OK) { + // switch (ret) { + // case LZMA_MEM_ERROR: + // m_err << tr("There is not enough memory to decompress the file."); + // break; + // case LZMA_FORMAT_ERROR: + // case LZMA_DATA_ERROR: + // case LZMA_BUF_ERROR: + // m_err << tr("The downloaded compressed file is corrupted."); + // break; + // case LZMA_OPTIONS_ERROR: + // m_err << tr("Unsupported compression options."); + // break; + // default: + // m_err << tr("Unknown decompression error."); + // break; + // } + // qApp->exit(4); + // CloseHandle(drive); + // return false; + // } + + // if (strm.avail_out == 0) { + // if (!writeBlock(drive, &osWrite, (char *) outBuffer, BLOCK_SIZE - strm.avail_out)) { + // qApp->exit(1); + // CloseHandle(drive); + // return false; + // } + + // if (osWrite.Offset + BLOCK_SIZE < osWrite.Offset) + // osWrite.OffsetHigh++; + // osWrite.Offset += BLOCK_SIZE; + + // strm.next_out = outBuffer; + // strm.avail_out = BLOCK_SIZE; + // } + // } + return false; +} - QFile file(what); - file.open(QIODevice::ReadOnly); +bool WriteJob::writePlain(HANDLE driveHandle) +{ + const qint64 blockSize = m_wmiDiskDrive->sectorSize() * 128; - ret = lzma_stream_decoder(&strm, MEDIAWRITER_LZMA_LIMIT, LZMA_CONCATENATED); - if (ret != LZMA_OK) { - err << tr("Failed to start decompressing."); + QFile isoFile(m_image); + isoFile.open(QIODevice::ReadOnly); + if (!isoFile.isOpen()) { + m_err << tr("Source image is not readable"); + m_err.flush(); return false; } + auto isoCleanup = qScopeGuard([&isoFile] { + isoFile.close(); + }); + + QFile drive; + drive.open(_open_osfhandle(reinterpret_cast(driveHandle), 0), QIODevice::WriteOnly | QIODevice::Unbuffered, QFile::AutoCloseHandle); + if (!drive.isOpen()) { + m_err << tr("Failed to open the device for writing"); + m_err.flush(); + qApp->exit(1); + return false; + } + auto driveCleanup = qScopeGuard([&drive] {drive.close();}); - strm.next_in = inBuffer; - strm.avail_in = 0; - strm.next_out = outBuffer; - strm.avail_out = BLOCK_SIZE; - - OVERLAPPED osWrite; - memset(&osWrite, 0, sizeof(osWrite)); - osWrite.hEvent = 0; - + void *buffer = NULL; + buffer = VirtualAlloc(NULL, blockSize, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE); + if (buffer == NULL) { + m_err << tr("Failed to allocate the buffer"); + m_err.flush(); + return false; + } + auto bufferCleanup = qScopeGuard([buffer, blockSize] { + VirtualFree(buffer, blockSize, MEM_DECOMMIT | MEM_RELEASE); + }); + + qint64 sectorSize = m_wmiDiskDrive->sectorSize(); + qint64 totalBytes = 0; + qint64 readBytes; + qint64 writtenBytes; + // DWORD writtenBytes = 0; while (true) { - if (strm.avail_in == 0) { - qint64 len = file.read((char*) inBuffer, BLOCK_SIZE); - totalRead += len; - - strm.next_in = inBuffer; - strm.avail_in = len; - - out << totalRead << "\n"; - out.flush(); + if ((readBytes = isoFile.read(static_cast(buffer), blockSize)) <= 0) { + break; } - ret = lzma_code(&strm, strm.avail_in == 0 ? LZMA_FINISH : LZMA_RUN); - if (ret == LZMA_STREAM_END) { - if (!writeBlock(drive, &osWrite, (char *) outBuffer, BLOCK_SIZE - strm.avail_out)) { - qApp->exit(1); - CloseHandle(drive); - return false; - } - - if (osWrite.Offset + BLOCK_SIZE < osWrite.Offset) - osWrite.OffsetHigh++; - osWrite.Offset += BLOCK_SIZE; - - CloseHandle(drive); - - return true; - } - if (ret != LZMA_OK) { - switch (ret) { - case LZMA_MEM_ERROR: - err << tr("There is not enough memory to decompress the file."); - break; - case LZMA_FORMAT_ERROR: - case LZMA_DATA_ERROR: - case LZMA_BUF_ERROR: - err << tr("The downloaded compressed file is corrupted."); - break; - case LZMA_OPTIONS_ERROR: - err << tr("Unsupported compression options."); - break; - default: - err << tr("Unknown decompression error."); - break; - } - qApp->exit(4); - CloseHandle(drive); + readBytes = ((readBytes + sectorSize - 1) / sectorSize) * sectorSize; + writtenBytes = drive.write(static_cast(buffer), readBytes); + // BOOL result = ::WriteFile(driveHandle, buffer, readBytes, &writtenBytes, nullptr); + if (writtenBytes < 0) { + // if (!result) { + m_err << tr("Failed to write to the device: ") << drive.errorString(); + // TCHAR message[256]; + // FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM, NULL, GetLastError(), MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), message, 255, NULL); + // m_err << tr("Failed to write to the device:") << " (" << QString::fromWCharArray(message).trimmed() << ")\n"; + m_err.flush(); return false; } - if (strm.avail_out == 0) { - if (!writeBlock(drive, &osWrite, (char *) outBuffer, BLOCK_SIZE - strm.avail_out)) { - qApp->exit(1); - CloseHandle(drive); - return false; - } - - if (osWrite.Offset + BLOCK_SIZE < osWrite.Offset) - osWrite.OffsetHigh++; - osWrite.Offset += BLOCK_SIZE; - - strm.next_out = outBuffer; - strm.avail_out = BLOCK_SIZE; - } - } -} - -bool WriteJob::writePlain(HANDLE drive) { - OVERLAPPED osWrite; - memset(&osWrite, 0, sizeof(osWrite)); - osWrite.hEvent = 0; - - uint64_t cnt = 0; - QByteArray buffer; - QFile isoFile(what); - isoFile.open(QIODevice::ReadOnly); - if (!isoFile.isOpen()) { - err << tr("Source image is not readable"); - err.flush(); - qApp->exit(1); - return false; - } - - while (true) { - buffer = isoFile.read(BLOCK_SIZE); - if (!writeBlock(drive, &osWrite, buffer.data(), buffer.size())) { - qApp->exit(1); + if (writtenBytes != readBytes) { + m_err << tr("The last block was not fully written"); + m_err.flush(); return false; } - if (osWrite.Offset + BLOCK_SIZE < osWrite.Offset) - osWrite.OffsetHigh++; - osWrite.Offset += BLOCK_SIZE; - cnt += buffer.size(); - out << cnt << "\n"; - out.flush(); + totalBytes += readBytes; + m_out << totalBytes << "\n"; + m_out.flush(); - if (buffer.size() != BLOCK_SIZE || isoFile.atEnd()) + if (readBytes != blockSize || isoFile.atEnd()) { break; + } } - CloseHandle(drive); + if (readBytes < 0) { + m_err << tr("Failed to read the image file: ") << isoFile.errorString(); + m_err.flush(); + return false; + } return true; } -bool WriteJob::check() { - out << "CHECK\n"; - out.flush(); - - HANDLE drive = openDrive(where); - - switch (mediaCheckFD(_open_osfhandle(reinterpret_cast(drive), 0), &WriteJob::staticOnMediaCheckAdvanced, this)) { - case ISOMD5SUM_CHECK_NOT_FOUND: - case ISOMD5SUM_CHECK_PASSED: - out << "DONE\n"; - out.flush(); - err << "OK\n"; - err.flush(); - qApp->exit(0); - break; - case ISOMD5SUM_CHECK_FAILED: - err << tr("Your drive is probably damaged.") << "\n"; - err.flush(); - qApp->exit(1); - return false; - default: - err << tr("Unexpected error occurred during media check.") << "\n"; - err.flush(); - qApp->exit(1); - return false; - } +bool WriteJob::check() +{ + // m_out << "CHECK\n"; + // m_out.flush(); + + // HANDLE drive = openDrive(m_drive); + + // switch (mediaCheckFD(_open_osfhandle(reinterpret_cast(drive), 0), &WriteJob::staticOnMediaCheckAdvanced, this)) { + // case ISOMD5SUM_CHECK_NOT_FOUND: + // case ISOMD5SUM_CHECK_PASSED: + // m_out << "DONE\n"; + // m_out.flush(); + // m_err << "OK\n"; + // m_err.flush(); + // qApp->exit(0); + // break; + // case ISOMD5SUM_CHECK_FAILED: + // m_err << tr("Your drive is probably damaged.") << "\n"; + // m_err.flush(); + // qApp->exit(1); + // return false; + // default: + // m_err << tr("Unexpected error occurred during media check.") << "\n"; + // m_err.flush(); + // qApp->exit(1); + // return false; + // } return true; } diff --git a/src/helper/win/writejob.h b/src/helper/win/writejob.h index a5aa222b..b530bae8 100644 --- a/src/helper/win/writejob.h +++ b/src/helper/win/writejob.h @@ -1,5 +1,6 @@ /* * Fedora Media Writer + * Copyright (C) 2024 Jan Grulich * Copyright (C) 2016 Martin Bříza * * This program is free software; you can redistribute it and/or @@ -20,58 +21,52 @@ #ifndef WRITEJOB_H #define WRITEJOB_H -#include -#include -#include +#include #include +#include +#include #include #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); - + void cleanDrive(HANDLE driveHandle); + bool lockDrive(HANDLE driveHandle); + HANDLE openDrive(); + void unmountVolumes(); private slots: void work(); void onFileChanged(const QString &path); bool write(); - bool writeCompressed(HANDLE drive); - bool writePlain(HANDLE drive); + bool writeCompressed(HANDLE driveHandle); + bool writePlain(HANDLE driveHandle); bool check(); -private: - QString what; - uint where; - QTextStream out { stdout }; - QTextStream err { stderr }; +private: + QString m_image; + std::unique_ptr m_wmi; + std::unique_ptr m_wmiDiskDrive; - QFileSystemWatcher watcher { }; + QTextStream m_out{stdout}; + QTextStream m_err{stderr}; - const int BLOCK_SIZE { 512 * 128 }; + QFileSystemWatcher m_watcher; }; #endif // WRITEJOB_H diff --git a/src/lib/CMakeLists.txt b/src/lib/CMakeLists.txt index 3a352a7a..2142e16d 100644 --- a/src/lib/CMakeLists.txt +++ b/src/lib/CMakeLists.txt @@ -1 +1,4 @@ add_subdirectory(isomd5) +if (WIN32) + add_subdirectory(libwmi) +endif() diff --git a/src/lib/libwmi/CMakeLists.txt b/src/lib/libwmi/CMakeLists.txt new file mode 100644 index 00000000..221bff5c --- /dev/null +++ b/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 +) diff --git a/src/lib/libwmi/libwmi.cpp b/src/lib/libwmi/libwmi.cpp new file mode 100644 index 00000000..92674395 --- /dev/null +++ b/src/lib/libwmi/libwmi.cpp @@ -0,0 +1,361 @@ +/* + * Fedora Media Writer + * Copyright (C) 2024 Jan Grulich + * + * 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 + +#include +#include + +#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(&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 LibWMI::getUSBDeviceList() +{ + QMap 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 LibWMI::getDiskDriveInformation(quint32 index, const QString &deviceID) +{ + std::unique_ptr result = std::make_unique(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; +} diff --git a/src/lib/libwmi/libwmi.h b/src/lib/libwmi/libwmi.h new file mode 100644 index 00000000..fa6ef485 --- /dev/null +++ b/src/lib/libwmi/libwmi.h @@ -0,0 +1,81 @@ +/* + * Fedora Media Writer + * Copyright (C) 2024 Jan Grulich + * + * 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 +#include + +class LibWMIDiskDrive; + +class LibWMI : public QObject +{ + Q_OBJECT +public: + LibWMI(QObject *parent); + ~LibWMI(); + + // Returns a map with list of devices + QMap 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 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