taskbar-icon-changer.wh.cpp

// ==WindhawkMod==
// @id taskbar-icon-changer
// @name Change taskbar icon
// @description Assigns custom icons to app instances based on their command line arguments or window title.
// @version 0.5
// @author claude
// @include *
// @compilerOptions -lshlwapi -lcomctl32
// ==/WindhawkMod==

// ==WindhawkModReadme==
/*
# taskbar-icon-changer

If you are using the "Disable grouping on the taskbar" mod, then enable the "Use window icons" option.
*/
// ==/WindhawkModReadme==

// ==WindhawkModSettings==
/*
- rules:
  - - process: ""
      $name: Process name (firefox.exe)
    - cmdline: ""
      $name: Command line substring
    - title: ""
      $name: Window title substring (optional)
    - icon: ""
      $name: Icon path (.ico or exe,index)
*/
// ==/WindhawkModSettings==

#include <windows.h>
#include <shlwapi.h>
#include <commctrl.h>
#include <string>
#include <vector>

// ── Structs ────────────────────────────────────────────────────────────────────

struct Rule {
    std::wstring process, cmdline, title, iconPath;
    HICON hBig   = nullptr;
    HICON hSmall = nullptr;
};

struct WindowIconData {
    HICON hBig, hSmall;
};

static std::vector<Rule> g_rules;
static std::vector<Rule> g_matchedRules;
static bool  g_active = false;
static DWORD g_pid    = 0;

#define SUBCLASS_UID 0x49434F4Eul

// ── Settings ───────────────────────────────────────────────────────────────────

static std::vector<Rule> LoadRules() {
    std::vector<Rule> rules;
    for (int i = 0; ; i++) {
        PCWSTR proc  = Wh_GetStringSetting(L"rules[%d].process", i);
        PCWSTR c     = Wh_GetStringSetting(L"rules[%d].cmdline",  i);
        PCWSTR t     = Wh_GetStringSetting(L"rules[%d].title",    i);
        PCWSTR p     = Wh_GetStringSetting(L"rules[%d].icon",     i);
        bool atEnd   = (!proc || !*proc) && (!c || !*c)
                    && (!t   || !*t)   && (!p || !*p);
        if (p && *p)
            rules.push_back({
                proc ? proc : L"",
                c    ? c    : L"",
                t    ? t    : L"",
                p
            });
        Wh_FreeStringSetting(proc);
        Wh_FreeStringSetting(c);
        Wh_FreeStringSetting(t);
        Wh_FreeStringSetting(p);
        if (atEnd) break;
    }
    return rules;
}

// ── Icon loading ───────────────────────────────────────────────────────────────

static HICON LoadIconAt(const std::wstring& spec, int size) {
    std::wstring path = spec;
    int idx = 0;
    size_t comma = spec.rfind(L',');
    if (comma != std::wstring::npos) {
        try { idx = std::stoi(spec.substr(comma + 1)); } catch (...) {}
        path = spec.substr(0, comma);
    }
    if (StrStrIW(path.c_str(), L".ico")) {
        HICON h = (HICON)LoadImageW(nullptr, path.c_str(), IMAGE_ICON,
                                    size, size, LR_LOADFROMFILE);
        if (h) return h;
    }
    HICON big = nullptr, small = nullptr;
    ExtractIconExW(path.c_str(), idx, &big, &small, 1);
    if (size <= 20) { if (big)   DestroyIcon(big);   return small; }
    else            { if (small) DestroyIcon(small);  return big;   }
}

// ── Subclass ───────────────────────────────────────────────────────────────────

static LRESULT CALLBACK IconSubclassProc(
    HWND hwnd, UINT msg, WPARAM wp, LPARAM lp,
    UINT_PTR uid, DWORD_PTR data)
{
    auto* d = reinterpret_cast<WindowIconData*>(data);
    if (msg == WM_GETICON && d) {
        if (wp == ICON_BIG && d->hBig)                  return (LRESULT)d->hBig;
        if ((wp == ICON_SMALL || wp == 2) && d->hSmall) return (LRESULT)d->hSmall;
    }
    if (msg == WM_NCDESTROY) {
        delete d;
        RemoveWindowSubclass(hwnd, IconSubclassProc, uid);
    }
    return DefSubclassProc(hwnd, msg, wp, lp);
}

// ── Shell refresh ──────────────────────────────────────────────────────────────

static void NotifyTaskbarRedraw(HWND hwnd) {
    static UINT uMsg = RegisterWindowMessageW(L"SHELLHOOK");
    HWND hTaskbar = FindWindowW(L"Shell_TrayWnd", nullptr);
    if (hTaskbar && uMsg)
        PostMessageW(hTaskbar, uMsg, HSHELL_REDRAW, (LPARAM)hwnd);
}

// ── Rule matching per-window ───────────────────────────────────────────────────

static const Rule* FindRuleForWindow(HWND hwnd) {
    wchar_t title[1024] = {};
    bool fetched = false;

    // Сначала ищем по title
    for (const auto& r : g_matchedRules) {
        if (!r.title.empty()) {
            if (!fetched) {
                GetWindowTextW(hwnd, title, 1024);
                fetched = true;
            }
            if (StrStrIW(title, r.title.c_str()))
                return &r;
        }
    }
    // Если не нашли, берём первое правило без title
    for (const auto& r : g_matchedRules) {
        if (r.title.empty())
            return &r;
    }
    return nullptr;
}

// ── Apply icons ────────────────────────────────────────────────────────────────

static void ApplyIconToWindow(HWND hwnd) {
    if (!IsWindow(hwnd) || !IsWindowVisible(hwnd)) return;
    const Rule* r = FindRuleForWindow(hwnd);
    if (!r) return;

    // Уже применено — пропускаем
    WindowIconData* existing = nullptr;
    if (GetWindowSubclass(hwnd, IconSubclassProc, SUBCLASS_UID, (DWORD_PTR*)&existing) && existing)
        return;

    if (g_matchedRules.size() == 1 && g_matchedRules[0].title.empty()) {
        if (r->hBig)   SetClassLongPtrW(hwnd, GCLP_HICON,   (LONG_PTR)r->hBig);
        if (r->hSmall) SetClassLongPtrW(hwnd, GCLP_HICONSM, (LONG_PTR)r->hSmall);
    }

    SendMessageW(hwnd, WM_SETICON, ICON_BIG,   (LPARAM)r->hBig);
    SendMessageW(hwnd, WM_SETICON, ICON_SMALL, (LPARAM)r->hSmall);

    auto* d = new WindowIconData{r->hBig, r->hSmall};
    SetWindowSubclass(hwnd, IconSubclassProc, SUBCLASS_UID, (DWORD_PTR)d);

    SetWindowPos(hwnd, nullptr, 0, 0, 0, 0,
        SWP_NOSIZE | SWP_NOMOVE | SWP_NOZORDER | SWP_NOACTIVATE | SWP_FRAMECHANGED);
    NotifyTaskbarRedraw(hwnd);
    Wh_Log(L"[cmdline-icon] ApplyIconToWindow hwnd=%p rule=%s", hwnd, r->iconPath.c_str());
}

static BOOL CALLBACK EnumCb(HWND hwnd, LPARAM) {
    DWORD pid = 0;
    GetWindowThreadProcessId(hwnd, &pid);
    if (pid != g_pid)                                return TRUE;
    if (GetWindowLongW(hwnd, GWL_STYLE) & WS_CHILD) return TRUE;
    if (!IsWindowVisible(hwnd))                      return TRUE;
    ApplyIconToWindow(hwnd);
    return TRUE;
}

static void ApplyToAll() { EnumWindows(EnumCb, 0); }

static BOOL CALLBACK RemoveSubclassCb(HWND hwnd, LPARAM) {
    DWORD pid = 0;
    GetWindowThreadProcessId(hwnd, &pid);
    if (pid == g_pid) {
        WindowIconData* d = nullptr;
        GetWindowSubclass(hwnd, IconSubclassProc, SUBCLASS_UID, (DWORD_PTR*)&d);
        delete d;
        RemoveWindowSubclass(hwnd, IconSubclassProc, SUBCLASS_UID);
    }
    return TRUE;
}

static DWORD WINAPI ReapplyThread(LPVOID) {
    for (DWORD ms : {300u, 1000u, 3000u, 7000u}) {
        Sleep(ms);
        if (g_active) ApplyToAll();
    }
    return 0;
}

static DWORD WINAPI ApplyToHwndThread(LPVOID arg) {
    HWND hwnd = (HWND)arg;
    for (DWORD ms : {200u, 800u, 2500u, 6000u}) {
        Sleep(ms);
        if (!IsWindow(hwnd) || !g_active) return 0;
        if (IsWindowVisible(hwnd)) ApplyIconToWindow(hwnd);
    }
    return 0;
}

// ── Hook: CreateWindowExW ──────────────────────────────────────────────────────

using CreateWindowExW_t = HWND(WINAPI*)(
    DWORD, LPCWSTR, LPCWSTR, DWORD, int, int, int, int,
    HWND, HMENU, HINSTANCE, LPVOID);
CreateWindowExW_t pCreateWindowExW;

HWND WINAPI CreateWindowExW_Hook(
    DWORD ex, LPCWSTR cls, LPCWSTR name, DWORD style,
    int x, int y, int w, int h,
    HWND parent, HMENU menu, HINSTANCE inst, LPVOID param)
{
    HWND hwnd = pCreateWindowExW(ex, cls, name, style,
                                  x, y, w, h, parent, menu, inst, param);
    if (hwnd && !(style & WS_CHILD) && g_active)
        CloseHandle(CreateThread(nullptr, 0, ApplyToHwndThread,
                                 (LPVOID)hwnd, 0, nullptr));
    return hwnd;
}

// ── Activate / Deactivate ──────────────────────────────────────────────────────

static void Deactivate() {
    g_active = false;
    EnumWindows(RemoveSubclassCb, 0);
    for (auto& r : g_matchedRules) {
        if (r.hBig)   { DestroyIcon(r.hBig);   r.hBig   = nullptr; }
        if (r.hSmall) { DestroyIcon(r.hSmall);  r.hSmall = nullptr; }
    }
    g_matchedRules.clear();
}

static void Activate(std::vector<Rule> matched) {
    Deactivate();
    for (auto& r : matched) {
        r.hBig   = LoadIconAt(r.iconPath, 48); // 32
        r.hSmall = LoadIconAt(r.iconPath, 16);
        if (r.hBig || r.hSmall)
            g_matchedRules.push_back(r);
        else
            Wh_Log(L"[cmdline-icon] Failed to load: %s", r.iconPath.c_str());
    }
    if (g_matchedRules.empty()) return;
    g_active = true;
    ApplyToAll();
    CloseHandle(CreateThread(nullptr, 0, ReapplyThread, nullptr, 0, nullptr));
}

// ── Rule matching ──────────────────────────────────────────────────────────────

static std::wstring GetCurrentProcessName() {
    wchar_t path[MAX_PATH];
    GetModuleFileNameW(nullptr, path, MAX_PATH);
    std::wstring name = PathFindFileNameW(path);
    for (auto& ch : name) ch = towlower(ch);
    return name;
}

// Собирает ВСЕ подходящие правила по process+cmdline (title не проверяем — это per-window)
static std::vector<Rule> CollectMatchingRules(const std::vector<Rule>& rules) {
    std::wstring cmdline  = GetCommandLineW();
    std::wstring procName = GetCurrentProcessName();
    std::vector<Rule> matched;
    for (const auto& r : rules) {
        if (!r.process.empty()) {
            std::wstring rp = r.process;
            for (auto& ch : rp) ch = towlower(ch);
            if (procName != rp) continue;
        }
        if (!r.cmdline.empty() && cmdline.find(r.cmdline) == std::wstring::npos)
            continue;
        matched.push_back(r);
    }
    return matched;
}

// ── Windhawk callbacks ─────────────────────────────────────────────────────────

BOOL Wh_ModInit() {
    g_pid   = GetCurrentProcessId();
    g_rules = LoadRules();

    auto matched = CollectMatchingRules(g_rules);
    if (matched.empty()) return TRUE;

    Activate(std::move(matched));

    HMODULE u32 = GetModuleHandleW(L"user32.dll");
    if (u32) {
        void* pCreate = (void*)GetProcAddress(u32, "CreateWindowExW");
        if (pCreate)
            Wh_SetFunctionHook(pCreate, (void*)CreateWindowExW_Hook,
                               (void**)&pCreateWindowExW);
    }
    return TRUE;
}

void Wh_ModUninit() {
    Deactivate();
}

void Wh_ModSettingsChanged() {
    g_rules = LoadRules();
    auto matched = CollectMatchingRules(g_rules);
    if (!matched.empty()) Activate(std::move(matched));
    else                  Deactivate();
}