ConsoleLogonHook - Windhawk edition

// ==WindhawkMod==
// @id           consolelogonhook
// @name         Console Logon Hook - Windhawk edition
// @description  Windows 7 login screen for W10 1903-22H2
// @version      1.0
// @author       teknixstuff
// @github       https://github.com/teknixstuff
// @include      logonui.exe
// ==/WindhawkMod==

#include <libloaderapi.h>
#include <windows.h>
#include <winnt.h>
#include <string>

const WCHAR hookpath[] = L"C:\\ProgramData\\Windhawk\\Engine\\ModsWritable\\ConsoleLogonHook.dll";
const WCHAR uipath[] = L"C:\\ProgramData\\Windhawk\\Engine\\ModsWritable\\ConsoleLogonUI.dll";

typedef LONG NTSTATUS;

#ifndef STATUS_SUCCESS
#define STATUS_SUCCESS ((NTSTATUS)0x00000000L)
#endif

#ifndef STATUS_BUFFER_TOO_SMALL
#define STATUS_BUFFER_TOO_SMALL ((NTSTATUS)0xC0000023L)
#endif

std::wstring GetKeyPathFromHKEY(HKEY key)
{
    std::wstring keyPath;
    if (key != NULL)
    {
        HMODULE dll = LoadLibrary(L"ntdll.dll");
        if (dll != NULL) {
            typedef DWORD (__stdcall *NtQueryKeyType)(
                HANDLE  KeyHandle,
                int KeyInformationClass,
                PVOID  KeyInformation,
                ULONG  Length,
                PULONG  ResultLength);

            NtQueryKeyType func = reinterpret_cast<NtQueryKeyType>(::GetProcAddress(dll, "NtQueryKey"));

            if (func != NULL) {
                DWORD size = 0;
                NTSTATUS result = 0;
                result = func(key, 3, 0, 0, &size);
                if (result == STATUS_BUFFER_TOO_SMALL)
                {
                    size = size + 2;
                    wchar_t* buffer = new (std::nothrow) wchar_t[size/sizeof(wchar_t)]; // size is in bytes
                    if (buffer != NULL)
                    {
                        result = func(key, 3, buffer, size, &size);
                        if (result == STATUS_SUCCESS)
                        {
                            buffer[size / sizeof(wchar_t)] = L'\0';
                            keyPath = std::wstring(buffer + 2);
                        }

                        delete[] buffer;
                    }
                }
            }

            FreeLibrary(dll);
        }
    }
    return keyPath;
}

typedef LONG (WINAPI *REGQUERYVALUEEXW)(HKEY hKey, LPCWSTR lpValueName, LPDWORD lpReserved, LPDWORD lpType, LPBYTE lpData, LPDWORD lpcbData);

REGQUERYVALUEEXW pOriginalRegQueryValueExW;

LONG WINAPI RegQueryValueExWHook(HKEY hKey, LPCWSTR lpValueName, LPDWORD lpReserved, LPDWORD lpType, LPBYTE lpData, LPDWORD lpcbData)
{   
    if ((lstrcmpiW(lpValueName, L"DllPath") == 0) && (
        (GetKeyPathFromHKEY(hKey) == L"\\REGISTRY\\MACHINE\\SOFTWARE\\Microsoft\\WindowsRuntime\\ActivatableClassId\\Windows.Internal.UI.Logon.Controller.ConsoleBlockedShutdownResolver") ||
        (GetKeyPathFromHKEY(hKey) == L"\\REGISTRY\\MACHINE\\SOFTWARE\\Microsoft\\WindowsRuntime\\ActivatableClassId\\Windows.Internal.UI.Logon.Controller.ConsoleLockScreen") ||
        (GetKeyPathFromHKEY(hKey) == L"\\REGISTRY\\MACHINE\\SOFTWARE\\Microsoft\\WindowsRuntime\\ActivatableClassId\\Windows.Internal.UI.Logon.Controller.ConsoleLogonUX") ||
        (GetKeyPathFromHKEY(hKey) == L"\\REGISTRY\\MACHINE\\SOFTWARE\\Microsoft\\WindowsRuntime\\ActivatableClassId\\Windows.Internal.UI.Logon.Controller.LogonUX") ||
        (GetKeyPathFromHKEY(hKey) == L"\\REGISTRY\\MACHINE\\SOFTWARE\\Microsoft\\WindowsRuntime\\ActivatableClassId\\Windows.Internal.Shell.PlatformExtensions.ConsoleCredUX")
        ))    
    { 
            if (lpType)
                *lpType = REG_SZ;
            if (lpData && lpcbData && *lpcbData >= sizeof(hookpath))
            {
                memcpy(lpData, hookpath, sizeof(hookpath));
                *lpcbData = sizeof(hookpath);
            }
            return ERROR_SUCCESS;
    }

    return pOriginalRegQueryValueExW(hKey, lpValueName, lpReserved, lpType, lpData, lpcbData);
}

HMODULE (WINAPI* pOriginalLoadLibraryExW)(LPCWSTR lpLibFileName, HANDLE hFile, DWORD dwFlags);
HMODULE LoadLibraryExWHook(LPCWSTR lpLibFileName, HANDLE hFile, DWORD dwFlags) {
    if (lstrcmpiW(lpLibFileName, L"ConsoleLogonUI.dll")) {
        return pOriginalLoadLibraryExW(uipath, hFile, dwFlags);
    } else {
        return pOriginalLoadLibraryExW(lpLibFileName, hFile, dwFlags);
    }
}

BOOL Wh_ModInit(void)
{
    Wh_Log(L"Init");

    if (GetFileAttributes(hookpath) == INVALID_FILE_ATTRIBUTES) {
        Wh_Log(L"Downloading ConsoleLogonHook");
        const WH_URL_CONTENT* HookFile = Wh_GetUrlContent(L"https://github.com/teknixstuff/misc-files/releases/download/ConsoleLogonHook-Windhawk/ConsoleLogonHook.dll", NULL);
        HANDLE hHookFile = CreateFile(hookpath, GENERIC_WRITE, 0, NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
        DWORD lpNumberOfBytesWritten;
        WriteFile(hHookFile, HookFile->data, HookFile->length, &lpNumberOfBytesWritten, NULL);
        CloseHandle(hHookFile);
        Wh_FreeUrlContent(HookFile);
    }

    if (GetFileAttributes(uipath) == INVALID_FILE_ATTRIBUTES) {
        Wh_Log(L"Downloading ConsoleLogonUI");
        const WH_URL_CONTENT* UIFile = Wh_GetUrlContent(L"https://github.com/teknixstuff/misc-files/releases/download/ConsoleLogonHook-Windhawk/ConsoleLogonUI.dll", NULL);
        HANDLE hUIFile = CreateFile(uipath, GENERIC_WRITE, 0, NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
        DWORD lpNumberOfBytesWritten;
        WriteFile(hUIFile, UIFile->data, UIFile->length, &lpNumberOfBytesWritten, NULL);
        CloseHandle(hUIFile);
        Wh_FreeUrlContent(UIFile);
    }

    Wh_SetFunctionHook((void*)GetProcAddress(LoadLibrary(L"kernelbase.dll"), "RegQueryValueExW"), (void*)RegQueryValueExWHook, (void**)&pOriginalRegQueryValueExW);
    Wh_SetFunctionHook((void*)LoadLibraryExW, (void*)LoadLibraryExWHook, (void**)&pOriginalLoadLibraryExW);

    return TRUE;
}