Untitled

#define g_OverlayWnd OverlayWindow::Hwnd
using BitBlt_t = BOOL(WINAPI*)(HDC, int, int, int, int, HDC, int, int, DWORD);
using PrintWindow_t = BOOL(WINAPI*)(HWND, HDC, UINT);
using GetWindowTextA_t = int (WINAPI*)(HWND, LPSTR, int);
using GetClassNameA_t = int (WINAPI*)(HWND, LPSTR, int);
BitBlt_t oBitBlt = nullptr;
PrintWindow_t oPrintWindow = nullptr;
GetWindowTextA_t oGetWindowTextA = nullptr;
GetClassNameA_t oGetClassNameA = nullptr;
bool g_SuppressOverlayRendering = false;

bool IsProcessSuspicious()
{
    DWORD pid = 0;
    GetWindowThreadProcessId(GetForegroundWindow(), &pid);
    HANDLE hProc = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, pid);
    if (!hProc) return false;
    wchar_t buffer[MAX_PATH] = { 0 };
    GetModuleFileNameExW(hProc, NULL, buffer, MAX_PATH);
    CloseHandle(hProc);
    std::wstring path(buffer);
    return path.find(L"SnippingTool.exe") != std::wstring::npos ||
           path.find(L"GameBar") != std::wstring::npos ||
           path.find(L"obs64.exe") != std::wstring::npos ||
           path.find(L"ScreenShot") != std::wstring::npos;
}

BOOL WINAPI hkBitBlt(HDC hdcDest, int x, int y, int nWidth, int nHeight,
    HDC hdcSrc, int xSrc, int ySrc, DWORD dwRop)
{
    if (IsProcessSuspicious())
        return FALSE;
    return oBitBlt(hdcDest, x, y, nWidth, nHeight, hdcSrc, xSrc, ySrc, dwRop);
}

BOOL WINAPI hkPrintWindow(HWND hwnd, HDC hdcBlt, UINT nFlags)
{
    if (IsProcessSuspicious() && hwnd == g_OverlayWnd)
        return FALSE;
    return oPrintWindow(hwnd, hdcBlt, nFlags);
}

typedef BOOL(WINAPI* EnumWindows_t)(WNDENUMPROC, LPARAM);
EnumWindows_t oEnumWindows = nullptr;

BOOL CALLBACK EnumWindowsCallback(HWND hwnd, LPARAM lParamInner)
{
    if (hwnd == g_OverlayWnd)
        return TRUE;
    WNDENUMPROC originalCallback = reinterpret_cast<WNDENUMPROC>(lParamInner);
    return originalCallback(hwnd, lParamInner);
}

BOOL WINAPI hkEnumWindows(WNDENUMPROC lpEnumFunc, LPARAM lParam)
{
    return oEnumWindows(EnumWindowsCallback, reinterpret_cast<LPARAM>(lpEnumFunc));
}

int WINAPI hkGetWindowTextA(HWND hWnd, LPSTR lpString, int nMaxCount)
{
    if (hWnd == g_OverlayWnd)
    {
        if (lpString && nMaxCount > 0)
            lpString[0] = '\0';
        return 0;
    }
    return oGetWindowTextA(hWnd, lpString, nMaxCount);
}

int WINAPI hkGetClassNameA(HWND hWnd, LPSTR lpClassName, int nMaxCount)
{
    if (hWnd == g_OverlayWnd)
    {
        if (lpClassName && nMaxCount > 0)
            lpClassName[0] = '\0';
        return 0;
    }
    return oGetClassNameA(hWnd, lpClassName, nMaxCount);
}

StealthHook::HookContext ctxEnumWindows, ctxBitBlt, ctxNtBitBlt, ctxPrintWindow;
StealthHook::HookContext ctxGetWindowTextA, ctxGetClassNameA;

void InitScreenshotHooks()
{
    HMODULE user32 = GetModuleHandleW(ENC(L"user32"));
    HMODULE gdi32 = GetModuleHandleW(ENC(L"gdi32"));
    HMODULE win32u = GetModuleHandleW(ENC(L"win32u"));

    StealthHook::CreateHook(GetProcAddress(user32, ENC("EnumWindows")), hkEnumWindows, ctxEnumWindows);
    StealthHook::CreateHook(GetProcAddress(gdi32, ENC("BitBlt")), hkBitBlt, ctxBitBlt);
    StealthHook::CreateHook(GetProcAddress(win32u, ENC("NtGdiBitBlt")), hkBitBlt, ctxNtBitBlt);
    StealthHook::CreateHook(GetProcAddress(user32, ENC("PrintWindow")), hkPrintWindow, ctxPrintWindow);
    StealthHook::CreateHook(GetProcAddress(user32, ENC("GetWindowTextA")), hkGetWindowTextA, ctxGetWindowTextA);
    StealthHook::CreateHook(GetProcAddress(user32, ENC("GetClassNameA")), hkGetClassNameA, ctxGetClassNameA);
}

void EnableOverlayProtection()
{
    static bool last_state = false;
    bool current = g_secure_vars.get().allow_capture;

    if (current != last_state) {
        if (current)
            SetWindowDisplayAffinity(g_OverlayWnd, WDA_EXCLUDEFROMCAPTURE);
        else
            SetWindowDisplayAffinity(g_OverlayWnd, WDA_NONE);

        last_state = current;
    }
}