OSDN Git Service

enable command notify on x64
[yamy/yamy.git] / fixscancodemap.cpp
index 5fece9d..249ba1a 100755 (executable)
@@ -3,13 +3,16 @@
 #include "windowstool.h"\r
 #include <tchar.h>\r
 #include <tlhelp32.h>\r
+#include <process.h>\r
 \r
 #pragma runtime_checks( "", off )\r
-static DWORD invokeFunc(InjectInfo *info)\r
+static DWORD WINAPI invokeFunc(InjectInfo *info)\r
 {\r
        BOOL ret;\r
        HANDLE hToken;\r
        HMODULE hAdvapi32;\r
+       DWORD result = 0;\r
+\r
        FpImpersonateLoggedOnUser pImpersonateLoggedOnUser;\r
        FpRevertToSelf pRevertToSelf;\r
        FpOpenProcessToken pOpenProcessToken;\r
@@ -22,29 +25,44 @@ static DWORD invokeFunc(InjectInfo *info)
 \r
        HANDLE hProcess = info->pOpenProcess(PROCESS_QUERY_INFORMATION, FALSE, info->pid_);\r
        if (hProcess == NULL) {\r
-               return 1;\r
+               result = YAMY_ERROR_ON_OPEN_YAMY_PROCESS;\r
+               goto exit;\r
        }\r
 \r
        ret = pOpenProcessToken(hProcess, TOKEN_QUERY | TOKEN_DUPLICATE , &hToken);\r
        if (ret == FALSE) {\r
-               return 2;\r
+               result = YAMY_ERROR_ON_OPEN_YAMY_TOKEN;\r
+               goto exit;\r
        }\r
 \r
        ret = pImpersonateLoggedOnUser(hToken);\r
        if (ret == FALSE) {\r
-               return 3;\r
+               result = YAMY_ERROR_ON_IMPERSONATE;\r
+               goto exit;\r
        }\r
 \r
-       info->pUpdate(0, 1);\r
+       if (info->isVistaOrLater_) {\r
+               info->pUpdate4(1);\r
+       } else {\r
+               info->pUpdate8(0, 1);\r
+       }\r
 \r
        ret = pRevertToSelf();\r
        if (ret == FALSE) {\r
-               return 4;\r
+               result = YAMY_ERROR_ON_REVERT_TO_SELF;\r
+               goto exit;\r
        }\r
 \r
-       info->pCloseHandle(hToken);\r
-       info->pCloseHandle(hProcess);\r
-       return 0;\r
+exit:\r
+       if (hToken != NULL) {\r
+               info->pCloseHandle(hToken);\r
+       }\r
+\r
+       if (hProcess != NULL) {\r
+               info->pCloseHandle(hProcess);\r
+       }\r
+\r
+       return result;\r
 }\r
 static int afterFunc(int arg)\r
 {\r
@@ -70,13 +88,13 @@ int FixScancodeMap::acquirePrivileges()
        HANDLE hToken = NULL;\r
 \r
        if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES, &hToken)) {\r
-               ret = 5;\r
+               ret = YAMY_ERROR_ON_OPEN_CURRENT_PROCESS;\r
                goto exit;\r
        }\r
 \r
        LUID luid;\r
        if (!LookupPrivilegeValue(NULL, SE_DEBUG_NAME, &luid)) {\r
-               ret = 6;\r
+               ret = YAMY_ERROR_ON_LOOKUP_PRIVILEGE;\r
                goto exit;\r
        }\r
 \r
@@ -86,7 +104,7 @@ int FixScancodeMap::acquirePrivileges()
        tk_priv.Privileges[0].Luid = luid;\r
 \r
        if (!AdjustTokenPrivileges(hToken, FALSE, &tk_priv, 0, NULL, NULL)) {\r
-               ret = 7;\r
+               ret = YAMY_ERROR_ON_ADJUST_PRIVILEGE;\r
                goto exit;\r
        }\r
 \r
@@ -117,7 +135,7 @@ DWORD FixScancodeMap::getWinLogonPid()
 \r
     BOOL bResult = Process32First(hSnap, &pe);\r
        while (bResult){\r
-               if (!_tcscmp(pe.szExeFile, _T("winlogon.exe"))) {\r
+               if (!_tcsicmp(pe.szExeFile, _T("winlogon.exe"))) {\r
                        DWORD sessionId;\r
 \r
                        if (ProcessIdToSessionId(pe.th32ProcessID, &sessionId) != FALSE) {\r
@@ -135,37 +153,34 @@ DWORD FixScancodeMap::getWinLogonPid()
 }\r
 \r
 \r
-int FixScancodeMap::clean()\r
+bool FixScancodeMap::clean(WlInfo wl)\r
 {\r
        int ret = 0;\r
 \r
-       if (WaitForSingleObject(m_hThread, 5000) == WAIT_TIMEOUT) {\r
-               ret = 15;\r
-               goto dirty_exit;\r
-       }\r
-       DWORD result = -1;\r
-       GetExitCodeThread(m_hThread, &result);\r
-       ret = result;\r
-       CloseHandle(m_hThread);\r
-       m_hThread = NULL;\r
+       if (wl.m_hThread != NULL) {\r
+               DWORD result;\r
 \r
-       if (m_remoteMem != NULL && m_hProcess != NULL) {\r
-               VirtualFreeEx(m_hProcess, m_remoteMem, 0, MEM_RELEASE);\r
-               m_remoteMem = NULL;\r
-       }\r
+               if (WaitForSingleObject(wl.m_hThread, 5000) == WAIT_TIMEOUT) {\r
+                       return false;\r
+               }\r
 \r
-       if (m_remoteInfo != NULL && m_hProcess != NULL) {\r
-               VirtualFreeEx(m_hProcess, m_remoteInfo, 0, MEM_RELEASE);\r
-               m_remoteInfo = NULL;\r
-       }\r
+               GetExitCodeThread(wl.m_hThread, &result);\r
+               CloseHandle(wl.m_hThread);\r
+\r
+               if (wl.m_remoteMem != NULL && wl.m_hProcess != NULL) {\r
+                       VirtualFreeEx(wl.m_hProcess, wl.m_remoteMem, 0, MEM_RELEASE);\r
+               }\r
+\r
+               if (wl.m_remoteInfo != NULL && wl.m_hProcess != NULL) {\r
+                       VirtualFreeEx(wl.m_hProcess, wl.m_remoteInfo, 0, MEM_RELEASE);\r
+               }\r
 \r
-       if (m_hProcess != NULL) {\r
-               CloseHandle(m_hProcess);\r
-               m_hProcess = NULL;\r
+               if (wl.m_hProcess != NULL) {\r
+                       CloseHandle(wl.m_hProcess);\r
+               }\r
        }\r
 \r
-dirty_exit:\r
-       return ret;\r
+       return true;\r
 }\r
 \r
 \r
@@ -174,97 +189,80 @@ int FixScancodeMap::injectThread(DWORD dwPID)
        int ret = 0;\r
        DWORD err = 0;\r
        BOOL wFlag;\r
+       WlInfo wi;\r
 \r
-       DWORD invokeFuncAddr = (DWORD)invokeFunc;\r
-       DWORD afterFuncAddr = (DWORD)afterFunc;\r
-       DWORD memSize =  afterFuncAddr - invokeFuncAddr;\r
+       wi.m_hProcess = NULL;\r
+       wi.m_remoteMem = NULL;\r
+       wi.m_remoteInfo = NULL;\r
+       wi.m_hThread = NULL;\r
 \r
-       if (m_hThread != NULL) {\r
-               ret = clean();\r
-               if (ret) {\r
-                       return ret;\r
-               }\r
-       }\r
+       ULONG_PTR invokeFuncAddr = (ULONG_PTR)invokeFunc;\r
+       ULONG_PTR afterFuncAddr = (ULONG_PTR)afterFunc;\r
+       SIZE_T memSize =  afterFuncAddr - invokeFuncAddr;\r
 \r
-       if ((m_hProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE, dwPID)) == NULL) {\r
-               ret = 8;\r
+       if ((wi.m_hProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE, dwPID)) == NULL) {\r
+               ret = YAMY_ERROR_ON_OPEN_WINLOGON_PROCESS;\r
                goto exit;\r
        }\r
 \r
-       m_remoteMem = VirtualAllocEx(m_hProcess, NULL, memSize, MEM_COMMIT, PAGE_EXECUTE_READWRITE);\r
-       if (m_remoteMem == NULL) {\r
-               ret = 9;\r
+       wi.m_remoteMem = VirtualAllocEx(wi.m_hProcess, NULL, memSize, MEM_COMMIT, PAGE_EXECUTE_READWRITE);\r
+       if (wi.m_remoteMem == NULL) {\r
+               ret = YAMY_ERROR_ON_VIRTUALALLOCEX;\r
                err = GetLastError();\r
                goto exit;\r
        }\r
 \r
-       wFlag = WriteProcessMemory(m_hProcess, m_remoteMem, (char*)invokeFunc, memSize, (SIZE_T*)0);\r
+       wFlag = WriteProcessMemory(wi.m_hProcess, wi.m_remoteMem, (char*)invokeFunc, memSize, (SIZE_T*)0);\r
        if (wFlag == FALSE) {\r
-               ret = 10;\r
+               ret = YAMY_ERROR_ON_WRITEPROCESSMEMORY;\r
                goto exit;\r
        }\r
 \r
-       m_remoteInfo = VirtualAllocEx(m_hProcess, NULL, sizeof(m_info), MEM_COMMIT, PAGE_READWRITE);\r
-       if (m_remoteInfo == NULL) {\r
-               ret = 11;\r
+       wi.m_remoteInfo = VirtualAllocEx(wi.m_hProcess, NULL, sizeof(m_info), MEM_COMMIT, PAGE_READWRITE);\r
+       if (wi.m_remoteInfo == NULL) {\r
+               ret = YAMY_ERROR_ON_VIRTUALALLOCEX;\r
                err = GetLastError();\r
                goto exit;\r
        }\r
 \r
-       wFlag = WriteProcessMemory(m_hProcess, m_remoteInfo, (char*)&m_info, sizeof(m_info), (SIZE_T*)0);\r
+       wFlag = WriteProcessMemory(wi.m_hProcess, wi.m_remoteInfo, (char*)&m_info, sizeof(m_info), (SIZE_T*)0);\r
        if (wFlag == FALSE) {\r
-               ret = 12;\r
+               ret = YAMY_ERROR_ON_WRITEPROCESSMEMORY;\r
                goto exit;\r
        }\r
 \r
-#if 0\r
-       TCHAR buf[1024];\r
-\r
-       _stprintf_s(buf, sizeof(buf)/sizeof(buf[0]),\r
-               _T("execute UpdatePerUserSystemParameters(), inject code to winlogon.exe?\r\n")\r
-               _T("invokeFunc=0x%p\r\n")\r
-               _T("afterFunc=0x%p\r\n")\r
-               _T("afterFunc - invokeFunc=%d\r\n")\r
-               _T("remoteMem=0x%p\r\n")\r
-               _T("remoteInfo=0x%p(size: %d)\r\n"),\r
-               invokeFunc, afterFunc, memSize, m_remoteMem, m_remoteInfo, sizeof(m_info));\r
-       if (MessageBox((HWND)NULL, buf, _T("upusp"), MB_OKCANCEL | MB_ICONSTOP) == IDCANCEL) {\r
-               (m_info.pUpdate)(0, 1);\r
+       wi.m_hThread = CreateRemoteThread(wi.m_hProcess, NULL, 0, \r
+               (LPTHREAD_START_ROUTINE)wi.m_remoteMem, wi.m_remoteInfo, 0, NULL);\r
+       if (wi.m_hThread == NULL) {\r
+               ret = YAMY_ERROR_ON_CREATEREMOTETHREAD;\r
                goto exit;\r
        }\r
-#endif\r
 \r
-       m_hThread = CreateRemoteThread(m_hProcess, NULL, 0, \r
-               (LPTHREAD_START_ROUTINE)m_remoteMem, m_remoteInfo, 0, NULL);\r
-       if (m_hThread == NULL) {\r
-               ret = 13;\r
-               goto exit;\r
-       }\r
-\r
-       if (WaitForSingleObject(m_hThread, 5000) == WAIT_TIMEOUT) {\r
-               ret = 14;\r
+       if (WaitForSingleObject(wi.m_hThread, 5000) == WAIT_TIMEOUT) {\r
+               ret = YAMY_ERROR_TIMEOUT_INJECTION;\r
+               m_wlTrash.push_back(wi);\r
                goto dirty_exit;\r
        }\r
        DWORD result = -1;\r
-       GetExitCodeThread(m_hThread, &result);\r
+       GetExitCodeThread(wi.m_hThread, &result);\r
        ret = result;\r
-       CloseHandle(m_hThread);\r
-       m_hThread = NULL;\r
+       CloseHandle(wi.m_hThread);\r
+       wi.m_hThread = NULL;\r
 \r
 exit:\r
-       if (m_remoteMem != NULL && m_hProcess != NULL) {\r
-               VirtualFreeEx(m_hProcess, m_remoteMem, 0, MEM_RELEASE);\r
-               m_remoteMem = NULL;\r
+       if (wi.m_remoteMem != NULL && wi.m_hProcess != NULL) {\r
+               VirtualFreeEx(wi.m_hProcess, wi.m_remoteMem, 0, MEM_RELEASE);\r
+               wi.m_remoteMem = NULL;\r
        }\r
 \r
-       if (m_remoteInfo != NULL && m_hProcess != NULL) {\r
-               VirtualFreeEx(m_hProcess, m_remoteInfo, 0, MEM_RELEASE);\r
-               m_remoteInfo = NULL;\r
+       if (wi.m_remoteInfo != NULL && wi.m_hProcess != NULL) {\r
+               VirtualFreeEx(wi.m_hProcess, wi.m_remoteInfo, 0, MEM_RELEASE);\r
+               wi.m_remoteInfo = NULL;\r
        }\r
 \r
-       if (m_hProcess != NULL) {\r
-               CloseHandle(m_hProcess);\r
-               m_hProcess = NULL;\r
+       if (wi.m_hProcess != NULL) {\r
+               CloseHandle(wi.m_hProcess);\r
+               wi.m_hProcess = NULL;\r
        }\r
 \r
 dirty_exit:\r
@@ -276,26 +274,27 @@ int FixScancodeMap::update()
        MINIMIZEDMETRICS mm;\r
        int result = 0;\r
 \r
-       result = acquirePrivileges();\r
-       if (result) {\r
+       if (m_errorOnConstruct) {\r
+               result = m_errorOnConstruct;\r
                goto exit;\r
        }\r
 \r
-       DWORD dwPID;\r
-       if ((dwPID = getWinLogonPid()) == 0) {\r
-               result = 15;\r
-               goto exit;\r
-       }\r
+       m_wlTrash.erase(remove_if(m_wlTrash.begin(), m_wlTrash.end(), FixScancodeMap::clean), m_wlTrash.end());\r
 \r
        memset(&mm, 0, sizeof(mm));\r
        mm.cbSize = sizeof(mm);\r
        SystemParametersInfo(SPI_GETMINIMIZEDMETRICS, sizeof(mm), &mm, 0);\r
 \r
-       result = injectThread(dwPID);\r
-       if (result && m_hThread == NULL) {\r
-               goto exit;\r
+       result = injectThread(m_winlogonPid);\r
+       if (result == YAMY_ERROR_TIMEOUT_INJECTION) {\r
+               // retry once\r
+               result = injectThread(m_winlogonPid);\r
+               if (result == YAMY_SUCCESS) {\r
+                       result = YAMY_ERROR_RETRY_INJECTION_SUCCESS;\r
+               }\r
        }\r
 \r
+       mm.iArrange = ARW_HIDE;\r
        SystemParametersInfo(SPI_SETMINIMIZEDMETRICS, sizeof(mm), &mm, 0);\r
 \r
 exit:\r
@@ -314,20 +313,20 @@ int FixScancodeMap::fix()
        if (ret) {\r
                origMap = reinterpret_cast<ScancodeMap*>(malloc(origSize));\r
                if (origMap == NULL) {\r
-                       result = 16;\r
+                       result = YAMY_ERROR_NO_MEMORY;\r
                        goto exit;\r
                }\r
 \r
                ret = m_pReg->read(_T("Scancode Map"), reinterpret_cast<BYTE*>(origMap), &origSize, NULL, 0);\r
                if (ret == false) {\r
-                       result = 17;\r
+                       result = YAMY_ERROR_ON_READ_SCANCODE_MAP;\r
                        goto exit;\r
                }\r
 \r
                fixSize = origSize;\r
                fixMap = reinterpret_cast<ScancodeMap*>(malloc(origSize + s_fixEntryNum * sizeof(s_fixEntry[0])));\r
                if (fixMap == NULL) {\r
-                       result = 18;\r
+                       result = YAMY_ERROR_NO_MEMORY;\r
                        goto exit;\r
                }\r
 \r
@@ -339,7 +338,7 @@ int FixScancodeMap::fix()
                fixSize = sizeof(ScancodeMap);\r
                fixMap = reinterpret_cast<ScancodeMap*>(malloc(sizeof(ScancodeMap) + s_fixEntryNum * sizeof(s_fixEntry[0])));\r
                if (fixMap == NULL) {\r
-                       result = 19;\r
+                       result = YAMY_ERROR_NO_MEMORY;\r
                        goto exit;\r
                }\r
 \r
@@ -374,7 +373,7 @@ int FixScancodeMap::fix()
 \r
        ret = m_pReg->write(_T("Scancode Map"), reinterpret_cast<BYTE*>(fixMap), fixSize);\r
        if (ret == false) {\r
-               result = 20;\r
+               result = YAMY_ERROR_ON_WRITE_SCANCODE_MAP;\r
                goto exit;\r
        }\r
 \r
@@ -386,7 +385,7 @@ int FixScancodeMap::fix()
                ret = m_pReg->remove(_T("Scancode Map"));\r
        }\r
        if (ret == false) {\r
-               result = 21;\r
+               result = YAMY_ERROR_ON_WRITE_SCANCODE_MAP;\r
                goto exit;\r
        }\r
 \r
@@ -405,11 +404,58 @@ int FixScancodeMap::restore()
        return update();\r
 }\r
 \r
+int FixScancodeMap::escape(bool i_escape)\r
+{\r
+       if (i_escape) {\r
+               SetEvent(m_hFixEvent);\r
+       } else {\r
+               SetEvent(m_hRestoreEvent);\r
+       }\r
+       return 0;\r
+}\r
+\r
+unsigned int WINAPI FixScancodeMap::threadLoop(void *i_this)\r
+{\r
+       int err;\r
+       DWORD ret;\r
+       FixScancodeMap *This = reinterpret_cast<FixScancodeMap*>(i_this);\r
+       HANDLE handles[] = {This->m_hFixEvent, This->m_hRestoreEvent, This->m_hQuitEvent};\r
+       while ((ret = MsgWaitForMultipleObjects(NUMBER_OF(handles), &handles[0],\r
+               FALSE, INFINITE, QS_POSTMESSAGE)) != WAIT_FAILED) {\r
+               switch (ret) {\r
+               case WAIT_OBJECT_0:                     // m_hFixEvent\r
+                       ResetEvent(This->m_hFixEvent);\r
+                       err = This->fix();\r
+                       PostMessage(This->m_hwnd, This->m_messageOnFail, err, 1);\r
+                       break;\r
+               case WAIT_OBJECT_0 + 1:         // m_hRestoreEvent\r
+                       ResetEvent(This->m_hRestoreEvent);\r
+                       err = This->restore();\r
+                       PostMessage(This->m_hwnd, This->m_messageOnFail, err, 0);\r
+                       break;\r
+               case WAIT_OBJECT_0 + 2:         // m_hQuiteEvent\r
+                       ResetEvent(This->m_hQuitEvent);\r
+                       // through below\r
+               default:\r
+                       return 0;\r
+                       break;\r
+               }\r
+       }\r
+       return 1;\r
+}\r
+\r
+int FixScancodeMap::init(HWND i_hwnd, UINT i_messageOnFail)\r
+{\r
+       m_hwnd = i_hwnd;\r
+       m_messageOnFail = i_messageOnFail;\r
+       return 0;\r
+}\r
+\r
 FixScancodeMap::FixScancodeMap() :\r
-       m_hProcess(NULL),\r
-       m_remoteMem(NULL),\r
-       m_remoteInfo(NULL),\r
-       m_hThread(NULL),\r
+       m_hwnd(NULL),\r
+       m_messageOnFail(WM_NULL),\r
+       m_errorOnConstruct(0),\r
+       m_winlogonPid(0),\r
        m_regHKCU(HKEY_CURRENT_USER, _T("Keyboard Layout")),\r
        m_regHKLM(HKEY_LOCAL_MACHINE, _T("SYSTEM\\CurrentControlSet\\Control\\Keyboard Layout")),\r
        m_pReg(NULL)\r
@@ -423,10 +469,20 @@ FixScancodeMap::FixScancodeMap() :
        memcpy(&m_info.revertToSelf_, "RevertToSelf", sizeof(m_info.revertToSelf_));\r
        memcpy(&m_info.openProcessToken_, "OpenProcessToken", sizeof(m_info.openProcessToken_));\r
 \r
+       m_hFixEvent = CreateEvent(NULL, TRUE, FALSE, NULL);\r
+       ASSERT(m_hFixEvent);\r
+       m_hRestoreEvent = CreateEvent(NULL, TRUE, FALSE, NULL);\r
+       ASSERT(m_hRestoreEvent);\r
+       m_hQuitEvent = CreateEvent(NULL, TRUE, FALSE, NULL);\r
+       ASSERT(m_hQuitEvent);\r
+\r
+       m_hThread = (HANDLE)_beginthreadex(NULL, 0, threadLoop, this, 0, &m_threadId);\r
+\r
        hMod = GetModuleHandle(_T("user32.dll"));\r
        if (hMod != NULL) {\r
-               m_info.pUpdate = (FpUpdatePerUserSystemParameters)GetProcAddress(hMod, "UpdatePerUserSystemParameters");\r
-               if (m_info.pUpdate == NULL) {\r
+               m_info.pUpdate4 = (FpUpdatePerUserSystemParameters4)GetProcAddress(hMod, "UpdatePerUserSystemParameters");\r
+               m_info.pUpdate8 = (FpUpdatePerUserSystemParameters8)m_info.pUpdate4;\r
+               if (m_info.pUpdate4 == NULL) {\r
                        return;\r
                }\r
        }\r
@@ -460,8 +516,30 @@ FixScancodeMap::FixScancodeMap() :
        } else {\r
                m_pReg = &m_regHKLM; // Windows7 or later\r
        }\r
+\r
+       // prototype of UpdatePerUserSystemParameters() differ vista or earlier\r
+       if (checkWindowsVersion(6, 0) == FALSE) {\r
+               m_info.isVistaOrLater_ = 0; // before Vista\r
+       } else {\r
+               m_info.isVistaOrLater_ = 1; // Vista or later\r
+       }\r
+\r
+       m_errorOnConstruct = acquirePrivileges();\r
+       if (m_errorOnConstruct) {\r
+               goto exit;\r
+       }\r
+\r
+       if ((m_winlogonPid = getWinLogonPid()) == 0) {\r
+               m_errorOnConstruct = YAMY_ERROR_ON_GET_WINLOGON_PID;\r
+               goto exit;\r
+       }\r
+\r
+exit:\r
+       ;\r
 }\r
 \r
 FixScancodeMap::~FixScancodeMap()\r
 {\r
+       SetEvent(m_hQuitEvent);\r
+       WaitForSingleObject(m_hThread, INFINITE);\r
 }\r