// ==WindhawkMod==
// @id              classic-theme-explorer
// @name            Classic Theme Explorer
// @description     Classic Theme mitigations for Explorer as a Windhawk mod
// @version         0.1
// @author          Cynosphere
// @github          https://github.com/Cynosphere
// @homepage        https://c7.pm/
// @include         explorer.exe
// @compilerOptions -lcomctl32
// ==/WindhawkMod==

// ==WindhawkModReadme==
/*
# Classic Theme Explorer
Classic Theme mitigations for Explorer as a Windhawk mod.

Based on [SimpleClassicTheme.FileExplorerHook](https://github.com/AEAEAEAE4343/SimpleClassicTheme.FileExplorerHook) and the AutoHotkey scripts that came before it.
*/
// ==/WindhawkModReadme==

// ==WindhawkModSettings==
/*
- clientEdge: true
  $name: "Client Edge"
  $description: "Adds Client Edge styling"
- rebarBorder: true
  $name: "Bar Borders"
  $description: "Adds borders to RebarWindow32 controls"
- rebarFixedHeight: false
  $name: "Fixed Bar Height"
  $description: "Resizes all toolbars to the tallest bar's height"
- hideNavigation: true
  $name: "Hide Navigation"
  $description: "Hide default navigation for use with Classic Explorer Bar and Classic Address Bar/Quero"
- addressSync: true
  $name: "Classic Address Bar Sync"
  $description: "Syncs address bar location on folder change"
- addressSyncFullPath: true
  $name: "Address Sync: Use full path"
  $description: "Uses full path instead of folder name"
- addressHeight: true
  $name: "Classic Address Bar Height"
  $description: "Resizes address bar to 22px"
- hideRefresh: true
  $name: "Hide Refresh Button"
  $description: "Hides the refresh button from Classic Address Bar"
*/
// ==/WindhawkModSettings==

#include <windows.h>
#include <objbase.h>
#include <initguid.h>
#include <commctrl.h>
#include <dwmapi.h>
#include <string>
#include <regex>
#include <locale>
#include <codecvt>

struct {
    bool clientEdge;
    bool rebarBorder;
    bool rebarFixedHeight;
    bool hideNavigation;
    bool addressSync;
    bool addressSyncFullPath;
    bool addressHeight;
    bool hideRefresh;
} settings;

void LoadSettings() {
    settings.clientEdge = Wh_GetIntSetting(L"clientEdge");
    settings.rebarBorder = Wh_GetIntSetting(L"rebarBorder");
    settings.rebarFixedHeight = Wh_GetIntSetting(L"rebarFixedHeight");
    settings.hideNavigation = Wh_GetIntSetting(L"hideNavigation");
    settings.addressSync = Wh_GetIntSetting(L"addressSync");
    settings.addressSyncFullPath = Wh_GetIntSetting(L"addressSyncFullPath");
    settings.addressHeight = Wh_GetIntSetting(L"addressHeight");
    settings.hideRefresh = Wh_GetIntSetting(L"hideRefresh");
}

// Explorer modifications
typedef struct EnumChildProcData
{
    HWND result;
    int index;
    int instancesFound;
    const wchar_t* className;
};
BOOL CALLBACK EnumChildProc(_In_ HWND hwnd, _In_ LPARAM lParam)
{
    wchar_t buffer[256];
    EnumChildProcData* data = (EnumChildProcData*)lParam;
    GetClassNameW(hwnd, buffer, 256);
    if (lstrcmpW(buffer, data->className) == 0)
    {
        if (data->instancesFound + 1 == data->index)
        {
            data->result = hwnd;
            return FALSE;
        }
        else
        {
            data->instancesFound++;
        }
    }
    return TRUE;
};
HWND GetChildWindow(HWND parent, const wchar_t* className, int index)
{
    EnumChildProcData data = { 0 };
    data.className = className;
    data.index = index;
    EnumChildWindows(parent, EnumChildProc, (LPARAM)&data);
    return data.result;
}

VOID ClassicThemeExplorer(HWND hWnd, WPARAM wParam)
{
    wchar_t classNameBuffer[256];
    wchar_t pathBuffer[MAX_PATH];
    HKEY key;
    RECT rect;

    HWND NavBarParent = GetChildWindow((HWND)hWnd, L"WorkerW", 1);
    HWND NavBarAddressBandRoot = GetChildWindow(NavBarParent, L"Address Band Root", 1);
    HWND NavBarToolBar = GetChildWindow(NavBarAddressBandRoot, L"ToolbarWindow32", 2);
    HWND ClassicAddressBandRoot = GetChildWindow((HWND)hWnd, L"Address Band Root", 2);
    HWND ClassicProgressControl = GetChildWindow(ClassicAddressBandRoot, L"msctls_progress32", 1);
    HWND ClassicComboBox = GetChildWindow(ClassicAddressBandRoot, L"ComboBoxEx32", 1);
    HWND ClassicRebar = GetChildWindow((HWND)hWnd, L"ReBarWindow32", 2);

    if (settings.clientEdge && (wParam == 1 || wParam == 6)) {
        HWND TreeView = GetChildWindow((HWND)hWnd, L"SysTreeView32", 1);
        LONG TreeViewExtendedStyle = GetWindowLongPtrW(TreeView, GWL_EXSTYLE);
        TreeViewExtendedStyle |= WS_EX_CLIENTEDGE;
        SetWindowLongPtrW(TreeView, GWL_EXSTYLE, TreeViewExtendedStyle);

        HWND FolderView = GetChildWindow((HWND)hWnd, L"FolderView", 1);
        LONG FolderViewExtendedStyle = GetWindowLongPtrW(FolderView, GWL_EXSTYLE);
        FolderViewExtendedStyle |= WS_EX_CLIENTEDGE;
        SetWindowLongPtrW(FolderView, GWL_EXSTYLE, FolderViewExtendedStyle);

        HWND ShellTabWindowClass = GetChildWindow((HWND)hWnd, L"ShellTabWindowClass", 1);
        GetWindowRect(ShellTabWindowClass, &rect);
        SetWindowPos(ShellTabWindowClass, NULL, NULL, NULL, rect.right - rect.left + 1, rect.bottom - rect.top, SWP_NOMOVE | SWP_NOZORDER | SWP_DEFERERASE);
        SetWindowPos(ShellTabWindowClass, NULL, NULL, NULL, rect.right - rect.left, rect.bottom - rect.top, SWP_NOMOVE | SWP_NOZORDER | SWP_DEFERERASE);
    }

    if (settings.rebarBorder && (wParam == 1 || wParam == 6)) {
        LONG ClassicRebarStyle = GetWindowLongPtrW(ClassicRebar, GWL_STYLE);
        ClassicRebarStyle |= RBS_BANDBORDERS;
        ClassicRebarStyle |= WS_BORDER;
        SetWindowLongPtrW(ClassicRebar, GWL_STYLE, ClassicRebarStyle);
    }

    if (settings.rebarFixedHeight && (wParam == 1 || wParam == 6)) {
        LONG ClassicRebarStyle = GetWindowLongPtrW(ClassicRebar, GWL_STYLE);
        ClassicRebarStyle &= ~RBS_VARHEIGHT;
        SetWindowLongPtrW(ClassicRebar, GWL_STYLE, ClassicRebarStyle);
    }

    // Apply bar changes
    if ((settings.rebarBorder || settings.rebarFixedHeight) && (wParam == 1 || wParam == 6)) {
        GetWindowRect(ClassicRebar, &rect);
        SetWindowPos(ClassicRebar, NULL, NULL, NULL, rect.right - rect.left + 1, rect.bottom - rect.top, SWP_NOMOVE | SWP_NOZORDER | SWP_DEFERERASE);
        SetWindowPos(ClassicRebar, NULL, NULL, NULL, rect.right - rect.left, rect.bottom - rect.top, SWP_NOMOVE | SWP_NOZORDER | SWP_DEFERERASE);
    }

    if (settings.hideNavigation && wParam == 1) {
        // Save the location of the Rebar
        HWND NavBarRebar = GetChildWindow((HWND)hWnd, L"ReBarWindow32", 1);
        GetWindowRect(NavBarRebar, &rect);
        LONG xRebar = rect.left;
        LONG yRebar = rect.top;
        LONG cxRebar = rect.right - rect.left;
        LONG cyRebar = rect.bottom - rect.top;

        // Destroy the NC area of the rebar
        SendMessageW(NavBarRebar, WM_NCDESTROY, 0, 0);

        // Hide the WorkerW and the Rebar
        ShowWindow(NavBarParent, SW_HIDE);
        ShowWindow(NavBarRebar, SW_HIDE);

        // Save the location of the ShellTabWindowClass
        HWND ShellTabWnd = GetChildWindow((HWND)hWnd, L"ShellTabWindowClass", 1);
        GetWindowRect(NavBarRebar, &rect);
        LONG xTabWnd = rect.left;
        LONG yTabWnd = rect.top;
        LONG cxTabWnd = rect.right - rect.left;
        LONG cyTabWnd = rect.bottom - rect.top;

        // Move ShellTabWindow to (*; yRebar; *; yTabWnd - yRebar + cyTabWnd)
        SetWindowPos(ShellTabWnd, NULL, xTabWnd, yRebar, cxTabWnd, yTabWnd - yRebar + cyTabWnd, SWP_NOZORDER);

        // Move Rebar to (*; *; *; 0)
        SetWindowPos(NavBarRebar, NULL, xRebar, yRebar, cxRebar, 0, SWP_NOZORDER);

        // Resize the window to apply
        GetWindowRect((HWND)hWnd, &rect);
        SetWindowPos((HWND)hWnd, NULL, NULL, NULL, rect.right - rect.left, rect.bottom - rect.top + 1, SWP_NOMOVE | SWP_NOZORDER | SWP_DEFERERASE);
        SetWindowPos((HWND)hWnd, NULL, NULL, NULL, rect.right - rect.left, rect.bottom - rect.top, SWP_NOMOVE | SWP_NOZORDER | SWP_DEFERERASE);

        // Redraw the entire explorer window
        RedrawWindow((HWND)hWnd, NULL, NULL, RDW_INVALIDATE | RDW_ERASE | RDW_ERASENOW | RDW_UPDATENOW | RDW_ALLCHILDREN);
    }

    if (settings.addressSync) {
        if (settings.addressSyncFullPath) {
            HWND toolbarwindow = GetChildWindow(NavBarAddressBandRoot, L"ToolbarWindow32", 1);
            GetWindowTextW(toolbarwindow, pathBuffer, MAX_PATH);

            // remove "Address: "
            // FIXME: locale issues, codecvt deprecated as of c++17
            std::wstring_convert<std::codecvt_utf8_utf16<wchar_t>> converter;

            std::wstring _pathStr = pathBuffer;
            std::string pathStr(_pathStr.begin(), _pathStr.end());
            pathStr = std::regex_replace(pathStr, std::regex("Address: "), "");
            _pathStr = converter.from_bytes(pathStr);
            wcscpy(pathBuffer, _pathStr.c_str());
        } else {
            HWND shelltabwindow = GetChildWindow((HWND)hWnd, L"ShellTabWindowClass", 1);
            GetWindowTextW(shelltabwindow, pathBuffer, MAX_PATH);
        }

        HWND AddressBarEdit = GetChildWindow(ClassicAddressBandRoot, L"Edit", 1);
        for (int i = 0; i < 3; i++)
            SendMessageW(AddressBarEdit, WM_SETTEXT, 0, (LPARAM)pathBuffer);

        // TODO: figure out how to get and set folder icon
    }

    if (settings.hideRefresh && (wParam == 1 || wParam == 6)) {
        HWND GoButtonToolbar = GetChildWindow(ClassicAddressBandRoot, L"ToolbarWindow32", 1);
        SendMessageW(GoButtonToolbar, WM_CLOSE, 0, 0);
    }

    if (settings.addressHeight) {
        // Allocate memory inside Explorer
        DWORD count = SendMessage(GetParent(ClassicAddressBandRoot), RB_GETBANDCOUNT, 0, 0);
        SIZE_T bytesRead = 0;
        DWORD ExplorerPID = 0;
        GetWindowThreadProcessId(ClassicAddressBandRoot, &ExplorerPID);
        HANDLE ExplorerProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE, ExplorerPID);
        void* ExplorerMemoryRebar = VirtualAllocEx(ExplorerProcess, NULL, sizeof(REBARBANDINFO), MEM_COMMIT, PAGE_READWRITE);
        void* ExplorerMemoryComboBoxItem = VirtualAllocEx(ExplorerProcess, NULL, sizeof(COMBOBOXEXITEM), MEM_COMMIT, PAGE_READWRITE);
        void* ExplorerMemoryToolbarButton = VirtualAllocEx(ExplorerProcess, NULL, sizeof(TBBUTTON), MEM_COMMIT, PAGE_READWRITE);

        // Make the band that's 39 pixels high, 22 pixels high
        // MOD: Except this bar isn't always 39 pixels, so lets check between 39-23px instead
        REBARBANDINFO bandInfo = { 0 };
        for (int i = 0; i < count; i++)
        {
            bandInfo = { 0 };
            bandInfo.cbSize = sizeof(REBARBANDINFO);
            bandInfo.fMask = RBBIM_CHILDSIZE;
            WriteProcessMemory(ExplorerProcess, ExplorerMemoryRebar, &bandInfo, sizeof(REBARBANDINFO), &bytesRead);
            SendMessageW(GetParent(ClassicAddressBandRoot), RB_GETBANDINFO, i, (LPARAM)ExplorerMemoryRebar);
            ReadProcessMemory(ExplorerProcess, ExplorerMemoryRebar, &bandInfo, sizeof(REBARBANDINFO), &bytesRead);
            if (bandInfo.cyMinChild <= 39 && bandInfo.cyMinChild > 22) {
                bandInfo.cyMinChild = 22;
                WriteProcessMemory(ExplorerProcess, ExplorerMemoryRebar, &bandInfo, sizeof(REBARBANDINFO), &bytesRead);
                SendMessageW(GetParent(ClassicAddressBandRoot), RB_SETBANDINFO, i, (LPARAM)ExplorerMemoryRebar);
            }
        }

        // Free Explorer Memory
        VirtualFreeEx(ExplorerProcess, ExplorerMemoryRebar, sizeof(REBARBANDINFO), MEM_DECOMMIT);
        VirtualFreeEx(ExplorerProcess, ExplorerMemoryToolbarButton, sizeof(REBARBANDINFO), MEM_DECOMMIT);
        CloseHandle(ExplorerProcess);
        
        // Set ComboBox height
        SendMessageW(ClassicComboBox, CB_SETITEMHEIGHT, -1, 22 - 6);
        SetParent(ClassicComboBox, ClassicAddressBandRoot);
        
        // Remove ProgressBsr
        SendMessageW(ClassicProgressControl, WM_CLOSE, 0, 0);

        // Fix ComboBox
        GetWindowRect(ClassicComboBox, &rect);
        SetWindowPos(ClassicComboBox, NULL, NULL, NULL, rect.right - rect.left + 1, rect.bottom - rect.top, SWP_NOMOVE | SWP_NOZORDER | SWP_DEFERERASE);
        SetWindowPos(ClassicComboBox, NULL, NULL, NULL, rect.right - rect.left, rect.bottom - rect.top, SWP_NOMOVE | SWP_NOZORDER | SWP_DEFERERASE);

        // Redraw the ComboBox
        RedrawWindow(ClassicComboBox, NULL, NULL, RDW_INVALIDATE | RDW_ERASE | RDW_ERASENOW | RDW_UPDATENOW | RDW_ALLCHILDREN);
    }
}

// boilerplate from hide-search-bar
#if defined(__GNUC__) && __GNUC__ > 8
#define WINAPI_LAMBDA_RETURN(return_t) -> return_t WINAPI
#elif defined(__GNUC__)
#define WINAPI_LAMBDA_RETURN(return_t) WINAPI -> return_t
#else
#define WINAPI_LAMBDA_RETURN(return_t) -> return_t
#endif

DWORD g_uiThreadId;

inline LSTATUS SHRegGetValueFromHKCUHKLMWithOpt(PCWSTR pwszKey, PCWSTR pwszValue, REGSAM samDesired, void* pvData, DWORD* pcbData)
{
    LSTATUS lRes = ERROR_FILE_NOT_FOUND;
    HKEY hKey = NULL;

    RegOpenKeyExW(HKEY_CURRENT_USER, pwszKey, 0, samDesired, &hKey);
    if (hKey == NULL || hKey == INVALID_HANDLE_VALUE) hKey = NULL;
    if (hKey)
    {
        lRes = RegQueryValueExW(hKey, pwszValue, 0, NULL, (LPBYTE) pvData, (LPDWORD) pcbData);

        RegCloseKey(hKey);
        if (lRes == ERROR_SUCCESS || lRes == ERROR_MORE_DATA) return lRes;
    }

    RegOpenKeyExW(HKEY_LOCAL_MACHINE, pwszKey, 0, samDesired, &hKey);
    if (hKey == NULL || hKey == INVALID_HANDLE_VALUE) hKey = NULL;

    if (hKey)
    {
        lRes = RegQueryValueExW(hKey, pwszValue, 0, NULL, (LPBYTE) pvData, (LPDWORD) pcbData);
        RegCloseKey(hKey);

        if (lRes == ERROR_SUCCESS || lRes == ERROR_MORE_DATA) return lRes;
    }

    return lRes;
}

UINT g_subclassRegisteredMsg = RegisterWindowMessage(L"Windhawk_SetWindowSubclassFromAnyThread_classic-theme-explorer");

LRESULT CALLBACK ClassicThemeExplorerSubClass(HWND hWnd, UINT uMsg, WPARAM wParam, LPARAM lParam, UINT_PTR uIdSubclass, DWORD_PTR dwRefData) {
    if (uMsg == WM_DESTROY) {
        Wh_Log("[Destroy] uMsg: 0x%04x, wParam: 0x%x", uMsg, wParam);
        RemoveWindowSubclass(hWnd, ClassicThemeExplorerSubClass, 0);
    } else if (uMsg == g_subclassRegisteredMsg && !wParam) {
        Wh_Log("[Unsub] uMsg: 0x%04x, wParam: 0x%x", uMsg, wParam);
        RemoveWindowSubclass(hWnd, ClassicThemeExplorerSubClass, 0);
    } else if (uMsg == WM_PARENTNOTIFY || uMsg == WM_SIZE || uMsg == WM_GETICON) {
        Wh_Log("[Target] uMsg: 0x%04x, wParam: 0x%x", uMsg, wParam);
        ClassicThemeExplorer(hWnd, wParam);
    } else {
        Wh_Log("[Unknown] uMsg: 0x%04x, wParam: 0x%x", uMsg, wParam);
    }

    return DefSubclassProc(hWnd, uMsg, wParam, lParam);
}

struct SET_WINDOW_SUBCLASS_FROM_ANY_THREAD_PARAM {
    SUBCLASSPROC pfnSubclass;
    UINT_PTR uIdSubclass;
    DWORD_PTR dwRefData;
    BOOL result;
};

LRESULT CALLBACK CallWndProcForWindowSubclass(int nCode, WPARAM wParam, LPARAM lParam) {
    if (nCode == HC_ACTION) {
        const CWPSTRUCT* cwp = (const CWPSTRUCT*)lParam;
        if (cwp->message == g_subclassRegisteredMsg && cwp->wParam) {
            SET_WINDOW_SUBCLASS_FROM_ANY_THREAD_PARAM* param = (SET_WINDOW_SUBCLASS_FROM_ANY_THREAD_PARAM*)cwp->lParam;
            param->result = SetWindowSubclass(cwp->hwnd, param->pfnSubclass, param->uIdSubclass, param->dwRefData);
        }
    }

    return CallNextHookEx(nullptr, nCode, wParam, lParam);
}

BOOL SetWindowSubclassFromAnyThread(HWND hWnd, SUBCLASSPROC pfnSubclass, UINT_PTR uIdSubclass, DWORD_PTR dwRefData) {
    DWORD dwThreadId = GetWindowThreadProcessId(hWnd, nullptr);
    
    if (dwThreadId == 0) return FALSE;
    if (dwThreadId == GetCurrentThreadId()) return SetWindowSubclass(hWnd, pfnSubclass, uIdSubclass, dwRefData);

    HHOOK hook = SetWindowsHookEx(WH_CALLWNDPROC, CallWndProcForWindowSubclass, nullptr, dwThreadId);

    if (!hook) return FALSE;

    SET_WINDOW_SUBCLASS_FROM_ANY_THREAD_PARAM param;
    param.pfnSubclass = pfnSubclass;
    param.uIdSubclass = uIdSubclass;
    param.dwRefData = dwRefData;
    param.result = FALSE;
    
    SendMessage(hWnd, g_subclassRegisteredMsg, TRUE, (WPARAM)&param);

    UnhookWindowsHookEx(hook);

    return param.result;
}

BOOL CALLBACK EnumBrowserWindowsUnsubclassFunc(HWND hWnd, LPARAM lParam) {
    SendMessage(hWnd, g_subclassRegisteredMsg, FALSE, 0);

    return TRUE;
}


HWND(WINAPI *pOriginalSHCreateWorkerWindow)(WNDPROC wndProc, HWND hWndParent, DWORD dwExStyle, DWORD dwStyle, HMENU hMenu, LONG_PTR wnd_extra);

HWND WINAPI SHCreateWorkerWindowHook(WNDPROC wndProc, HWND hWndParent, DWORD dwExStyle, DWORD dwStyle, HMENU hMenu, LONG_PTR wnd_extra)
{
    HWND result;
    LSTATUS lRes = ERROR_FILE_NOT_FOUND;
    DWORD dwSize = 0;

    //Wh_Log("g_uiThreadId: %d, GetCurrentThreadId: %d", g_uiThreadId, GetCurrentThreadId());

    // is this even needed? it changes thread after some time anyways
    //if (g_uiThreadId && g_uiThreadId != GetCurrentThreadId()) return pOriginalSHCreateWorkerWindow(wndProc, hWndParent, dwExStyle, dwStyle, hMenu, wnd_extra);
    /*if (!g_uiThreadId)*/ g_uiThreadId = GetCurrentThreadId();


    if (SHRegGetValueFromHKCUHKLMWithOpt(
        TEXT("SOFTWARE\\Classes\\CLSID\\{056440FD-8568-48e7-A632-72157243B55B}\\InProcServer32"),
        TEXT(""),
        KEY_READ | KEY_WOW64_64KEY,
        NULL,
        (LPDWORD)(&dwSize)
    ) == ERROR_SUCCESS && (dwSize < 4) && dwExStyle == 0x10000 && dwStyle == 1174405120) result = 0;

    else result = pOriginalSHCreateWorkerWindow(wndProc, hWndParent, dwExStyle, dwStyle, hMenu, wnd_extra);

    //Wh_Log("dwExStyle: 0x%x, dwStyle: 0x%x, result: 0x%x", dwExStyle, dwStyle, result);

    if (dwExStyle == 0x10000 && dwStyle == 0x46000000 && result) SetWindowSubclassFromAnyThread(hWndParent, ClassicThemeExplorerSubClass, 0, 0);

    return result;
}

// Windhawk boilerplate
BOOL Wh_ModInit() {
    Wh_Log(L"Init");

    LoadSettings();

    HMODULE hShcore = GetModuleHandle(L"shcore.dll");

    void* origFunc = (void*)GetProcAddress(hShcore, (LPCSTR)188);
    Wh_SetFunctionHook(origFunc, (void*)SHCreateWorkerWindowHook, (void**)&pOriginalSHCreateWorkerWindow);

    return TRUE;
}

void Wh_ModUninit() {
    Wh_Log(L"Uninit");

    if (g_uiThreadId != 0) EnumThreadWindows(g_uiThreadId, EnumBrowserWindowsUnsubclassFunc, 0);
}

void Wh_ModSettingsChanged() {
    Wh_Log(L"SettingsChanged");

    LoadSettings();
}