Desktop Icon Minder - Rearranges icons back to their proper position.

icon_minder.cpp

/*
 * This utility saves and restores windows desktop icon positions.
 *
 * NOTE: is probably quite brittle on how it interacts with the desktop and may break on different versions.
 * NOTE: the storage file is stored under %localappdata%\DesktopIconMinder\icon_positions.dat
 *
 * Compile with: cl.exe /W4 /O2 /std:c++20 /EHsc icon_minder.cpp
 */

#define UNICODE
#define _UNICODE

#include <combaseapi.h>
#include <commctrl.h>
#include <fileapi.h>
#include <handleapi.h>
#include <knownfolders.h>
#include <pathcch.h>
#include <shlobj_core.h>
#include <windows.h>
#include <winerror.h>
#include <winnt.h>

#include <cstddef>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <cwchar>
#include <clocale>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>

#pragma comment(lib, "Kernel32.lib")
#pragma comment(lib, "User32.lib")
#pragma comment(lib, "Shell32.lib")
#pragma comment(lib, "Ole32.lib")
#pragma comment(lib, "Pathcch.lib")

class AppException : public std::runtime_error {
   public:
    explicit AppException(const std::string &message) : std::runtime_error(message) {}
};

struct IconEntry {
    int index;
    POINT point;
    std::wstring text;
};

void WriteIconEntry(HANDLE handle, const IconEntry &entry) {
    // INDEX_u32 X_long, Y_long TEXT_LENGTH TEXT
    constexpr size_t initial = 4 + 4 + 4 + 4;
    char buffer[initial];
    memcpy(&buffer[0], &entry.index, 4);
    memcpy(&buffer[4], &entry.point.x, 4);
    memcpy(&buffer[8], &entry.point.y, 4);
    uint32_t numChars = static_cast<uint32_t>(entry.text.size());
    memcpy(&buffer[12], &numChars, 4);

    WriteFile(handle, buffer, initial, nullptr, nullptr);
    WriteFile(handle, entry.text.data(), numChars * sizeof(wchar_t), nullptr, nullptr);
}

IconEntry ReadIconEntry(HANDLE handle) {
    constexpr size_t initial = 4 + 4 + 4;
    char start[initial];

    IconEntry entry;
    BOOL status = TRUE;
    status &= ReadFile(handle, start, initial, nullptr, nullptr);
    memcpy(&entry.index, &start[0], 4);
    memcpy(&entry.point.x, &start[4], 4);
    memcpy(&entry.point.y, &start[8], 4);
    uint32_t numChars = 0;
    status &= ReadFile(handle, &numChars, 4, nullptr, nullptr);
    wchar_t buffer[MAX_PATH] = {0};
    status &= ReadFile(handle, &buffer, numChars * sizeof(wchar_t), nullptr, nullptr);
    if (status == FALSE) {
        throw AppException("Failed to read entry");
    }
    entry.text.assign(buffer, numChars);
    return entry;
}

struct DesktopMinder {
   private:
    static constexpr SIZE_T MEMORY_SIZE = 4096; // Take an entire page
    static constexpr SIZE_T MAX_STR_SIZE = 2048;
    static_assert(MAX_STR_SIZE % 16 == 0, "Must align to 16 bytes");

   public:
    DesktopMinder() = default;
    ~DesktopMinder() { DropMemorySpace(); }
    DesktopMinder(const DesktopMinder &other) = delete;
    DesktopMinder &operator=(const DesktopMinder &other) = delete;

    DesktopMinder(DesktopMinder &&other)
        : hProcess(std::exchange(other.hProcess, nullptr)),
          remoteMemory(std::exchange(other.remoteMemory, nullptr)),
          hwndListView(std::exchange(other.hwndListView, nullptr)) {}

    DesktopMinder &operator=(DesktopMinder &&other) {
        std::swap(hProcess, other.hProcess);
        std::swap(remoteMemory, other.remoteMemory);
        std::swap(hwndListView, other.hwndListView);
        return *this;
    }

    void AcquireHandle() {
        HWND programManager = FindWindow(L"Progman", NULL);
        HWND hwndDesktop = FindWindowEx(programManager, NULL, L"SHELLDLL_DefView", NULL);
        hwndListView = FindWindowEx(hwndDesktop, NULL, L"SysListView32", NULL);
        if (programManager == NULL || hwndDesktop == NULL || hwndListView == NULL) {
            throw AppException("Failed to find list view");
        }
    }

    void CreateRemoteMemory() {
        if (hasRemoteMemory()) {
            throw AppException("Misuse: memory space already injected");
        }
        if (hwndListView == NULL) {
            throw AppException("Need list view handle to find process id.");
        }
        DWORD pid = 0;
        if (GetWindowThreadProcessId(hwndListView, &pid) == 0) {
            throw AppException("Failed to find process id");
        }
        hProcess = OpenProcess(PROCESS_VM_OPERATION | PROCESS_VM_READ | PROCESS_VM_WRITE, FALSE, pid);
        if (hProcess == 0) {
            throw AppException("Failed to 'borrow' process");
        }

        remoteMemory = VirtualAllocEx(hProcess, NULL, MEMORY_SIZE, MEM_COMMIT, PAGE_READWRITE);
    }

    int GetIconCount() const { return static_cast<int>(SendMessage(hwndListView, LVM_GETITEMCOUNT, 0, 0)); }

    POINT GetIconPosition(int index) const {
        checkRemoteMemory();

        if (index < 0) {
            throw AppException("Bad index");
        }
        LPVOID start = remoteMemory;
        POINT point = {0};
        // We probably don't need to actually initialize the memory
        if (!WriteProcessMemory(hProcess, remoteMemory, &point, sizeof(point), NULL)) {
            throw AppException("Bad remote memory write");
        }
        if (!SendMessageW(hwndListView, LVM_GETITEMPOSITION, (WPARAM)index, (LPARAM)start)) {
            throw AppException("Failed to retrieve position");
        }
        if (!ReadProcessMemory(hProcess, remoteMemory, &point, sizeof(point), NULL)) {
            throw AppException("Bad remote memory read");
        }
        return point;
    }

    void SetIconPosition(int index, POINT position) {
        checkRemoteMemory();

        if (!WriteProcessMemory(hProcess, remoteMemory, &position, sizeof(position), NULL)) {
            throw AppException("Bad remote memory write");
        }
        if (!SendMessageW(hwndListView, LVM_SETITEMPOSITION32, (WPARAM)index, (LPARAM)remoteMemory)) {
            throw AppException("Failed to set position");
        }
    }

    std::wstring GetIconText(int index) const {
        checkRemoteMemory();

        if (index < 0) {
            throw AppException("Bad index");
        }
        LPVOID start = remoteMemory;
        LPVOID pRemoteLvItem = (char *)start + MAX_STR_SIZE;

        // LVM_GETITEMTEXT
        LVITEMW lvItem = {0};
        lvItem.mask = LVIF_TEXT;
        lvItem.iItem = index;
        lvItem.iSubItem = 0;
        lvItem.pszText = (LPWSTR)start;
        lvItem.cchTextMax = MAX_STR_SIZE;

        if (!WriteProcessMemory(hProcess, pRemoteLvItem, &lvItem, sizeof(lvItem), NULL)) {
            throw AppException("Bad remote memory write");
        }
        LRESULT numChars = SendMessage(hwndListView, LVM_GETITEMTEXTW, (WPARAM)index, (LPARAM)pRemoteLvItem);
        wchar_t buffer[MAX_STR_SIZE] = {0};
        if (!ReadProcessMemory(hProcess, start, buffer, sizeof(TCHAR) * numChars, NULL)) {
            throw AppException("Bad remote memory read");
        }
        return std::wstring(buffer, numChars);
    }

    void DropMemorySpace() {
        if (remoteMemory != NULL) {
            // For MEM_RELEASE dwSize = 0
            VirtualFreeEx(hProcess, remoteMemory, 0, MEM_RELEASE);
            CloseHandle(hProcess);
            remoteMemory = NULL;
            hProcess = NULL;
        }
    }

   private:
    bool hasRemoteMemory() const { return remoteMemory != NULL; }
    void checkRemoteMemory() const {
        if (!hasRemoteMemory()) {
            throw AppException("DesktopManager has not setup remote space");
        }
    }

    HANDLE hProcess = NULL;
    LPVOID remoteMemory = NULL;
    HWND hwndListView = NULL;
};

std::vector<IconEntry> GetIconEntries(DesktopMinder &minder) {
    std::vector<IconEntry> results{};
    int iconCount = minder.GetIconCount();
    for (int i = 0; i < iconCount; ++i) {
        auto text = minder.GetIconText(i);
        auto p = minder.GetIconPosition(i);
        results.emplace_back(i, p, text);
    }
    return results;
}

void PrintIconEntries(const std::vector<IconEntry> &entries) {
    wprintf(L"There are %zu icons\n", entries.size());
    for (const auto &entry : entries) {
        const auto &p = entry.point;
        wprintf(L"Icon[%d]: '%s' at (%ld, %ld)\n", entry.index, entry.text.data(), p.x, p.y);
    }
}

int CmdView() {
    DesktopMinder minder{};
    minder.AcquireHandle();
    minder.CreateRemoteMemory();
    auto results = GetIconEntries(minder);
    PrintIconEntries(results);
    return 0;
}

std::wstring GetStoragePath() {
    std::wstring result{};
    {
        PWSTR knownPath = nullptr;
        // SHGetKnownFolderPath works without CoInitialize if HANDLE hToken is NULL.
        if (SHGetKnownFolderPath(FOLDERID_LocalAppData, 0, NULL, &knownPath) != S_OK) {
            // We must free regardless of returned status.
            CoTaskMemFree(knownPath);
            throw AppException("Failed to retrieve storage path");
        }
        result = knownPath;
        CoTaskMemFree(knownPath);
    }
    if (result.back() != L'\\') {
        result += L'\\';
    }
    result += L"DesktopIconMinder\\icon_positions.dat";

    return result;
}

HANDLE OpenStorageFile(DWORD dwDesiredAccess, DWORD dwDisposition) {
    std::wstring path = GetStoragePath();
    wchar_t pwstrFolder[MAX_PATH] = {0};
    {
        path._Copy_s(pwstrFolder, MAX_PATH, path.length());
        HRESULT hr = PathCchRemoveFileSpec(pwstrFolder, MAX_PATH);
        if (hr != S_OK) {
            throw AppException("Failed to process storage directory path");
        }
    }
    int directoryResult = SHCreateDirectoryEx(NULL, pwstrFolder, NULL);
    if (directoryResult != ERROR_SUCCESS && directoryResult != ERROR_ALREADY_EXISTS) {
        throw AppException("Failed to create directory for storage file");
    }

    HANDLE result =
        CreateFile(path.data(), dwDesiredAccess, 0, nullptr, dwDisposition, FILE_ATTRIBUTE_NORMAL, nullptr);
    if (result == INVALID_HANDLE_VALUE) {
        throw AppException("Failed to open storage file");
    }

    return result;
}

void WriteToStorage(HANDLE handle, const std::vector<IconEntry> &entries) {
    // We just little-endian it out.
    std::string magic{"DesktopIconMinder|"};
    WriteFile(handle, magic.data(), static_cast<DWORD>(sizeof(char) * magic.length()), nullptr, nullptr);
    uint32_t count = static_cast<uint32_t>(entries.size());
    WriteFile(handle, &count, sizeof(count), nullptr, nullptr);
    for (const auto &entry : entries) {
        WriteIconEntry(handle, entry);
    }
}

std::vector<IconEntry> ReadFromStorage(HANDLE handle) {
    // Expect header
    std::vector<IconEntry> entries{};
    constexpr size_t BUF_SIZE = 64;
    char buffer[BUF_SIZE];
    // NOTE: We really expect it to not give partial read here.
    if (!ReadFile(handle, buffer, 18, nullptr, nullptr)) {
        throw AppException("Failed to read header");
    }
    if (std::string{buffer, 18} != "DesktopIconMinder|") {
        throw AppException("Header magic incorrect");
    }
    int32_t count = 0;
    if (!ReadFile(handle, &count, 4, nullptr, nullptr)) {
        throw AppException("Failed to read count");
    }
    entries.reserve(count);
    for (int i = 0; i < count; ++i) {
        entries.push_back(ReadIconEntry(handle));
    }
    return entries;
}

int CmdSave() {
    // Fail early on strange cases
    // TODO: RAII for Handle.
    HANDLE storageFile = OpenStorageFile(GENERIC_WRITE, CREATE_ALWAYS);

    DesktopMinder minder{};
    minder.AcquireHandle();
    minder.CreateRemoteMemory();
    auto entries = GetIconEntries(minder);

    {
        // Check and warn that duplicates cannot be properly handled.
        std::unordered_map<std::wstring, IconEntry> seenNames{};
        for (const auto &entry : entries) {
            const auto [iter, wasInserted] = seenNames.insert({entry.text, entry});
            if (!wasInserted) {
                wprintf(L"WARNING: index %d and %d collides on text: %s\n", iter->second.index, entry.index,
                    entry.text.data());
            }
        }
    }

    WriteToStorage(storageFile, entries);

    return 0;
}

int CmdRead() {
    HANDLE storageFile = OpenStorageFile(GENERIC_READ, OPEN_EXISTING);
    if (storageFile == INVALID_HANDLE_VALUE) {
        throw AppException("Failed to open storage file");
    }
    std::vector<IconEntry> entries = ReadFromStorage(storageFile);
    PrintIconEntries(entries);
    return 0;
}

int CmdRestore() {
    HANDLE storageFile = OpenStorageFile(GENERIC_READ, OPEN_EXISTING);
    if (storageFile == INVALID_HANDLE_VALUE) {
        throw AppException("Failed to open storage file");
    }
    auto storedEntries = ReadFromStorage(storageFile);

    DesktopMinder minder{};
    minder.AcquireHandle();
    minder.CreateRemoteMemory();
    auto desktopEntries = GetIconEntries(minder);

    std::unordered_map<std::wstring, const IconEntry *> mapStored{};
    for (const auto &entry : storedEntries) mapStored.emplace(entry.text, &entry);
    {
        std::unordered_map<std::wstring, const IconEntry *> mapDesktop{};
        for (const auto &entry : desktopEntries) mapDesktop.emplace(entry.text, &entry);

        // Write current icons that have no stored positions
        // and write stored icon positions that are no longer found.
        for (const auto &[text, _entry] : mapStored) {
            if (mapDesktop.find(text) == mapDesktop.end()) {
                wprintf(L"Note: '%s' is stored but not found on desktop\n", text.data());
            }
        }
        for (const auto &[text, _entry] : mapDesktop) {
            if (mapStored.find(text) == mapStored.end()) {
                wprintf(L"Warning: '%s' is on desktop but has no entry in storage\n", text.data());
            }
        }
    }

    // Modify desktop entries positions only.
    std::vector<IconEntry> moving{};
    for (auto &entry : desktopEntries) {
        auto iter = mapStored.find(entry.text);
        if (iter != mapStored.end()) {
            const auto &sEntry = *iter->second;
            long eX = entry.point.x;
            long eY = entry.point.y;
            long sX = sEntry.point.x;
            long sY = sEntry.point.y;
            if (sX == eX && sY == eY) continue;
            wprintf(L"Moving '%s' from (%ld, %ld) to (%ld, %ld))\n", entry.text.data(), eX, eY, sX, sY);
            IconEntry moved = entry;
            moved.point = sEntry.point;
            moving.push_back(moved);
        }
    }

    if (moving.empty()) {
        wprintf(L"\nDone. Nothing to move.\n");
    } else {
        for (const auto &toMove : moving) {
            minder.SetIconPosition(toMove.index, toMove.point);
        }
        wprintf(L"\nDone moving.\n");
    }

    return 0;
}

int wmain(int argc, wchar_t *argv[]) {
    SetConsoleOutputCP(CP_UTF8);
    _wsetlocale(LC_ALL, L".UTF8");

    if (argc != 2) {
        wprintf(
            L"Usage: %s COMMAND\n\nWhere COMMAND can be one of:\n"
            L"    view: Shows current desktop icon positions\n"
            L"    save: Stores the desktop icon positions\n"
            L"    read: Reads and outputs the stored positions (does not restore)\n"
            L"    restore: Restores the desktop icon positions\n",
            argv[0]);
        return 1;
    }

    int returnCode = 1;
    try {
        std::wstring command{argv[1]};
        if (command == L"view") {
            returnCode = CmdView();
        } else if (command == L"save") {
            returnCode = CmdSave();
        } else if (command == L"read") {
            returnCode = CmdRead();
        } else if (command == L"restore") {
            returnCode = CmdRestore();
        } else {
            returnCode = 1;
            fwprintf(stderr, L"Unrecognized command '%s'\n", command.data());
        }
    } catch (const AppException &exception) {
        returnCode = 2;
        fwprintf(stderr, L"*** ERROR: %S\n", exception.what());
    }

    return returnCode;
}