#include <windows.h>
#include <iostream>
#include <vector>
#include <algorithm>

struct SYSTEM_HANDLE_INFORMATION
{
	ULONG	ProcessId;
	BYTE	ObjectTypeNumber;
	BYTE	Flags;
	USHORT	Handle;
	PVOID	Object;
	ACCESS_MASK	GrantedAccess;
};

typedef NTSTATUS (WINAPI*QueryObject)(HANDLE, ULONG, PVOID, ULONG, PULONG);
typedef NTSTATUS (WINAPI*QuerySysInfo)(ULONG, PVOID, ULONG, PULONG);
typedef NTSTATUS (WINAPI*OpenDirectory)(PHANDLE, ACCESS_MASK, PVOID);
typedef NTSTATUS (WINAPI*QueryDirectory)(HANDLE, PVOID, ULONG, BOOL, BOOL, PULONG, PULONG);

BYTE GetObjectTypeNumber(LPCWSTR objectName)
{
	HMODULE hNtDll = GetModuleHandle(L"ntdll.dll");
	OpenDirectory openDir = reinterpret_cast<OpenDirectory>(GetProcAddress(hNtDll, "NtOpenDirectoryObject"));
	QueryDirectory queryDir = reinterpret_cast<QueryDirectory>(GetProcAddress(hNtDll, "NtQueryDirectoryObject"));
	BYTE objectNumber = 0;
	UNICODE_STRING us = {0};
	us.Buffer = L"\\ObjectTypes";
	us.Length = wcslen(us.Buffer) * sizeof(WCHAR);
	us.MaximumLength = us.Length + sizeof(WCHAR);
	OBJECT_ATTRIBUTES oa = {sizeof(oa), 0};
	oa.Attributes = OBJ_CASE_INSENSITIVE;
	oa.ObjectName = &us;
	HANDLE hDir = NULL;	
    NTSTATUS stat = openDir(&hDir, DIRECTORY_QUERY, &oa);
	if(stat == STATUS_SUCCESS)
	{
		std::vector<BYTE> buffer(32000);
		ULONG len = 0, ctx = 0;
		while((stat = queryDir(hDir, &buffer[0], buffer.size(), FALSE, TRUE, &ctx, &len)) == STATUS_INFO_LENGTH_MISMATCH)
		{
			if(len)
			{
				buffer.resize(len);
			}
			else
			{
				buffer.resize(buffer.size() * 2);
			}
		}
		if(stat == STATUS_SUCCESS)
		{
			POBJECT_DIRECTORY_INFORMATION pObjInf = reinterpret_cast<POBJECT_DIRECTORY_INFORMATION>(&buffer[0]);
			for(ULONG i = 0; i < ctx; ++i)
			{
				OBJECT_DIRECTORY_INFORMATION& obj = pObjInf[i];
				std::wcout << std::dec << i + 1 << ',' << std::hex << i + 1 << ',' << obj.Name.Buffer << '\n';
				if(!objectNumber)
				{
					if(_wcsicmp(objectName, obj.Name.Buffer) == 0)
					{
						objectNumber = i + 1;
						//break;
					}
				}
			}
		}
		CloseHandle(hDir);
	}
	return objectNumber;
}

bool GetSystemHandleTable(std::vector<SYSTEM_HANDLE_INFORMATION>& handles)
{
	struct SYS_HANDLE_INFO_EX
	{
		ULONG size;
		SYSTEM_HANDLE_INFORMATION handles[1];
	};

	__if_not_exists(SystemHandleInformation)
	{
		const UINT SystemHandleInformation = 16;
	}
	std::vector<BYTE> buffer(32000);
	DWORD size = 0;
	NTSTATUS status = 0;
	QuerySysInfo qsi = reinterpret_cast<QuerySysInfo>(GetProcAddress(GetModuleHandle(L"ntdll.dll"), "NtQuerySystemInformation"));
	while((status = qsi(SystemHandleInformation, &buffer[0], buffer.size(), &size)) == STATUS_INFO_LENGTH_MISMATCH)
	{
		if(size)
		{
			buffer.resize(size);
		}
		else
		{
			buffer.resize(buffer.size() * 2);
		}
	}
	if(status == STATUS_SUCCESS)
	{
		SYS_HANDLE_INFO_EX* handleInf = reinterpret_cast<SYS_HANDLE_INFO_EX*>(&buffer[0]);
		handles.resize(handleInf->size);
		std::copy(&(handleInf->handles[0]), &(handleInf->handles[handleInf->size]), handles.begin());
	}
	return status == STATUS_SUCCESS;
}

namespace detail
{
	struct ProcessIDMatcher
	{
		DWORD id;
		ProcessIDMatcher(DWORD id) : id(id) {}
		bool operator()(SYSTEM_HANDLE_INFORMATION& handle) const
		{
			return handle.ProcessId != id;
		}
	};
}

void GetProcessHandleTable(DWORD procID, std::vector<SYSTEM_HANDLE_INFORMATION>& procHandles)
{
	std::vector<SYSTEM_HANDLE_INFORMATION> handles;
	if(GetSystemHandleTable(handles))
	{
		std::remove_copy_if(
			handles.begin(),
			handles.end(),
			std::back_inserter(procHandles),
			detail::ProcessIDMatcher(procID)
		);
	}
}

BYTE GetFileTypeNumber()
{
	BYTE objectNumber = 0;
	HANDLE hThisFile = CreateFileW(_wtmpnam(NULL), GENERIC_WRITE, 0, NULL, OPEN_ALWAYS, FILE_ATTRIBUTE_TEMPORARY | FILE_FLAG_DELETE_ON_CLOSE, NULL);
	if(hThisFile != INVALID_HANDLE_VALUE)
	{
		std::vector<SYSTEM_HANDLE_INFORMATION> procHandles;
		GetProcessHandleTable(GetCurrentProcessId(), procHandles);
		for(ULONG i = 0; i < procHandles.size(); ++i)
		{
			if(reinterpret_cast<HANDLE>(procHandles[i].Handle) == hThisFile)
			{
				objectNumber = procHandles[i].ObjectTypeNumber;
				break;
			}
		}
		CloseHandle(hThisFile);
	}
	return objectNumber;
}

wchar_t* LowerCase(LPCWSTR str)
{
	wchar_t* lowercase = _wcsdup(str);
	if(lowercase)
	{
		_wcslwr(lowercase);
	}
	return lowercase;
}

#include <Cfgmgr32.h>
#pragma comment(lib, "Cfgmgr32.lib")

void MatchName(ULONG procID, USHORT handle, LPCWSTR matchText)
{
	static QueryObject queryObject = NULL;
	if(!queryObject)
	{
		HMODULE hNtDll = GetModuleHandle(L"ntdll.dll");
		queryObject = reinterpret_cast<QueryObject>(GetProcAddress(hNtDll, "NtQueryObject"));
	}
	HANDLE hProc = OpenProcess(PROCESS_DUP_HANDLE, FALSE, procID);
	if(hProc)
	{
		HANDLE hOurCopy = NULL;
		if(DuplicateHandle(hProc, reinterpret_cast<HANDLE>(handle), GetCurrentProcess(), &hOurCopy, 0, FALSE, 0))
		{
			std::vector<BYTE> buffer(4096);
			ULONG required = 0;
			NTSTATUS stat = STATUS_SUCCESS;
			while((stat = queryObject(hOurCopy, ObjectNameInformation, &buffer[0], buffer.size(), &required)) == STATUS_INFO_LENGTH_MISMATCH)
			{
				if(required)
				{
					buffer.resize(required);
				}
				else
				{
					buffer.resize(buffer.size() * 2);
				}
			}
			if(stat == STATUS_SUCCESS)
			{
				UNICODE_STRING* pName = reinterpret_cast<UNICODE_STRING*>(&buffer[0]);
				if(pName->Buffer)
				{
					wchar_t* handleNameCopy = LowerCase(pName->Buffer);
					if(handleNameCopy)
					{
						if(wcsstr(handleNameCopy, matchText))
						{
							std::wcout << L"Matched with: " << handleNameCopy << L'\n';
						}
						free(handleNameCopy);
					}
				}
			}
			CloseHandle(hOurCopy);
		}
		CloseHandle(hProc);
	}
}

void EnableDebugPrivilege()
{
     HANDLE hToken;
     TOKEN_PRIVILEGES tp = {0};
     
     OpenProcessToken(GetCurrentProcess(), TOKEN_WRITE, &hToken);
                    
     LookupPrivilegeValue(NULL, SE_DEBUG_NAME, &tp.Privileges[0].Luid);
     tp.PrivilegeCount = 1;
     tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
                    
     AdjustTokenPrivileges(hToken, FALSE, &tp, 0, NULL, 0);
     CloseHandle(hToken);
}

int __cdecl wmain(int argc, wchar_t** argv)
{
	if(argc >= 3)
	{
		wchar_t* matchTextLower = LowerCase(argv[2]);
		if(matchTextLower)
		{
			EnableDebugPrivilege();
			GetObjectTypeNumber(argv[1]);
			BYTE objID = GetFileTypeNumber();
			if(objID)
			{
				std::vector<SYSTEM_HANDLE_INFORMATION> handles;
				GetSystemHandleTable(handles);
				for(UINT i = 0; i < handles.size(); ++i)
				{
					SYSTEM_HANDLE_INFORMATION& curHandle = handles[i];
					if(curHandle.ObjectTypeNumber == objID)
					{
						MatchName(curHandle.ProcessId, curHandle.Handle, matchTextLower);
					}
				}
			}
			free(matchTextLower);
		}
	}
	else
	{
		;
	}
    return 0;
}