#define KEY "79c86fb12b36dfa33d1a537c9af100b4c7928a9c"
#include "openwechat.h"
#include <string>
#include <wchar.h>
#include <cstring>
#include <vector>
#include <tlhelp32.h>
ZWQUERYSYSTEMINFORMATION ZwQuerySystemInformation = (ZWQUERYSYSTEMINFORMATION)GetProcAddress(GetModuleHandleA("ntdll.dll"), "ZwQuerySystemInformation");
NTQUERYOBJECT NtQueryObject = (NTQUERYOBJECT)GetProcAddress(
	GetModuleHandleA("ntdll.dll"), "NtQueryObject");
BOOL IsTargetPid(DWORD Pid, std::vector<DWORD> Pids, int num)
{
	for (auto pid:Pids)
	{
		if (Pid == pid)
		{
			return TRUE;
		}
	}
	return FALSE;
}
HANDLE DuplicateHandleEx(DWORD pid, HANDLE h, DWORD flags)
{
	HANDLE hHandle = NULL;

	HANDLE hProc = OpenProcess(PROCESS_ALL_ACCESS, FALSE, pid);
	if (hProc)
	{
		if (!DuplicateHandle(hProc,
			(HANDLE)h, GetCurrentProcess(),
			&hHandle, 0, FALSE, /*DUPLICATE_SAME_ACCESS*/flags))
		{
			hHandle = NULL;
		}
	}

	CloseHandle(hProc);
	return hHandle;
}

typedef struct _SYSTEM_HANDLE_TABLE_ENTRY_INFO {
	USHORT UniqueProcessId;
	USHORT CreatorBackTraceIndex;
	UCHAR ObjectTypeIndex;
	UCHAR HandleAttributes;
	USHORT HandleValue;
	PVOID Object;
	ULONG GrantedAccess;
} SYSTEM_HANDLE_TABLE_ENTRY_INFO, *PSYSTEM_HANDLE_TABLE_ENTRY_INFO;
typedef struct
{
	USHORT Length;
	USHORT MaxLen;
	USHORT *Buffer;
}UNICODE_STRING, *PUNICODE_STRING;
typedef struct _OBJECT_NAME_INFORMATION {
	UNICODE_STRING Name;
} OBJECT_NAME_INFORMATION, *POBJECT_NAME_INFORMATION;

typedef struct _SYSTEM_HANDLE_INFORMATION1 {
	ULONG NumberOfHandles;
	SYSTEM_HANDLE_TABLE_ENTRY_INFO Handles[1];
} SYSTEM_HANDLE_INFORMATION1, *PSYSTEM_HANDLE_INFORMATION1;

DWORD qureyProcessId(std::wstring name, std::vector<DWORD> *pids) {
	DWORD pid;
	PROCESSENTRY32 entry;
	entry.dwSize = sizeof(PROCESSENTRY32);

	HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, NULL);

	if (Process32First(snapshot, &entry) == TRUE)
	{
		while (Process32Next(snapshot, &entry) == TRUE)
		{
			if (std::wstring(entry.szExeFile) == name) {
				HANDLE hProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE, entry.th32ProcessID);
				pid = GetProcessId(hProcess);
				CloseHandle(hProcess);
				pids->emplace_back(pid);
			}
		}
	}
	CloseHandle(snapshot);
	return pid;
}
BOOL ElevatePrivileges()
{
	HANDLE hToken;
	TOKEN_PRIVILEGES tkp;
	tkp.PrivilegeCount = 1;
	if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken))
		return FALSE;
	LookupPrivilegeValue(NULL, SE_DEBUG_NAME, &tkp.Privileges[0].Luid);
	tkp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
	if (!AdjustTokenPrivileges(hToken, FALSE, &tkp, sizeof(TOKEN_PRIVILEGES), NULL, NULL))
	{
		return FALSE;
	}

	return TRUE;
}
int PatchWeChat()
{
	DWORD dwSize = 0;
	POBJECT_NAME_INFORMATION pNameInfo;
	POBJECT_NAME_INFORMATION pNameType;
	PVOID pbuffer = NULL;
	NTSTATUS Status;
	int nIndex = 0;
	DWORD dwFlags = 0;
	char szType[128] = { 0 };
	char szName[512] = { 0 };
	PSYSTEM_HANDLE_INFORMATION1 pHandleInfo = NULL;
	int ret = -1;

	ElevatePrivileges();
	std::vector<DWORD> pids;
	DWORD Num = qureyProcessId(L"WeChat.exe", &pids);
	if (Num == 0)
	{
		return ret;
	}

	if (!ZwQuerySystemInformation)
	{
		goto Exit0;
	}

	pbuffer = VirtualAlloc(NULL, 0x1000, MEM_COMMIT, PAGE_READWRITE);

	if (!pbuffer)
	{
		goto Exit0;
	}

	Status = ZwQuerySystemInformation(SystemHandleInformation, pbuffer, 0x1000, &dwSize);

	if (Status<0)
	{
		if ((LONG)0xC0000004L != Status)
		{
			goto Exit0;
		}
		else
		{
			// 这里大家可以保证程序的正确性使用循环分配稍好
			if (NULL != pbuffer)
			{
				VirtualFree(pbuffer, 0, MEM_RELEASE);
			}

			if (dwSize * 2 > 0x4000000)  // MAXSIZE
			{
				goto Exit0;
			}

			pbuffer = VirtualAlloc(NULL, dwSize * 2, MEM_COMMIT, PAGE_READWRITE);

			if (!pbuffer)
			{
				goto Exit0;
			}

			Status = ZwQuerySystemInformation(SystemHandleInformation, pbuffer, dwSize * 2, NULL);

			if (Status<0)
			{
				goto Exit0;
			}
		}
	}

	pHandleInfo = (PSYSTEM_HANDLE_INFORMATION1)pbuffer;

	for (nIndex = 0; nIndex < pHandleInfo->NumberOfHandles; nIndex++)
	{
		if (IsTargetPid(pHandleInfo->Handles[nIndex].UniqueProcessId, pids, Num))
		{
			//
			HANDLE hHandle = DuplicateHandleEx(pHandleInfo->Handles[nIndex].UniqueProcessId,
				(HANDLE)pHandleInfo->Handles[nIndex].HandleValue,
				DUPLICATE_SAME_ACCESS
			);
			if (hHandle == NULL) continue;

			Status = NtQueryObject(hHandle, ObjectNameInformation, szName, 512, &dwFlags);

			if (Status<0)
			{
				CloseHandle(hHandle);
				continue;
			}

			Status = NtQueryObject(hHandle, ObjectTypeInformation, szType, 128, &dwFlags);

			if (Status<0)
			{
				CloseHandle(hHandle);
				continue;
			}

			pNameInfo = (POBJECT_NAME_INFORMATION)szName;
			pNameType = (POBJECT_NAME_INFORMATION)szType;

			WCHAR TypName[1024] = { 0 };
			WCHAR Name[1024] = { 0 };

			wcsncpy_s(TypName, (WCHAR*)pNameType->Name.Buffer, pNameType->Name.Length / 2);
			wcsncpy_s(Name, (WCHAR*)pNameInfo->Name.Buffer, pNameInfo->Name.Length / 2);

			// 匹配是否为需要关闭的句柄名称
			if (0 == wcscmp(TypName, L"Mutant"))
			{
				//WeChat_aj5r8jpxt_Instance_Identity_Mutex_Name
				//if (wcsstr(Name, L"_WeChat_App_Instance_Identity_Mutex_Name"))
				if (wcsstr(Name, L"_WeChat_") &&
					wcsstr(Name, L"_Instance_Identity_Mutex_Name"))
				{
					CloseHandle(hHandle);

					hHandle = DuplicateHandleEx(pHandleInfo->Handles[nIndex].UniqueProcessId,
						(HANDLE)pHandleInfo->Handles[nIndex].HandleValue,
						DUPLICATE_CLOSE_SOURCE
					);

					if (hHandle)
					{
						ret = ERROR_SUCCESS;
						CloseHandle(hHandle);
					}
					else
					{
						ret = GetLastError();
					}

					//goto Exit0;
				}
			}

			CloseHandle(hHandle);
		}

	}

Exit0:
	if (NULL != pbuffer)
	{
		VirtualFree(pbuffer, 0, MEM_RELEASE);
	}

	return ret;
}
int WINAPI WinMain(
	HINSTANCE hInstance,
	HINSTANCE hPrevInstance,
	PSTR nCmdLine,
	int iCmdShow
) {
	if (strcmp(nCmdLine, KEY) != 0){
		return 0;
	}
	else {
		int ret = 1;
		PatchWeChat();
		WCHAR Path[1024];
		HKEY hKey = NULL;
		if (ERROR_SUCCESS != RegOpenKey(HKEY_CURRENT_USER, L"Software\\Tencent\\WeChat", &hKey))
		{
			
			return ret;
		}

		DWORD Type = REG_SZ;
		// WCHAR Path[MAX_PATH] = { 0 };
		DWORD cbData = MAX_PATH * sizeof(WCHAR);
		if (ERROR_SUCCESS != RegQueryValueEx(hKey, L"InstallPath", 0, &Type, (LPBYTE)Path, &cbData))
		{
			ret = GetLastError();
			if (hKey)
			{
				RegCloseKey(hKey);
			}
			return ret;
		}
		WCHAR exe[1024];
		wcscpy_s(exe, Path);
		wcscat_s(exe, L"\\WeChat.exe");

		ShellExecute(GetDesktopWindow(), L"open", exe, L"", Path, SW_SHOW);

	}
	return 0;
}