#include <windows/dbghelpProtos.h>
#include "dbghelpFuncs.h"

BOOL CALLBACK EnumProc( 
    PSYMBOL_INFOW pSymInfo,   
    ULONG SymbolSize,      
    PVOID UserContext)
{
	PVOID pAddr = (PVOID)pSymInfo->Address;
	std::wcout << L"Address: 0x" << pAddr << L", Size: " << SymbolSize << L", Name: " << pSymInfo->Name << L'\n';
    return TRUE;
}

#include <map>
#include <string>
#include <algorithm>
#include <cctype>

HANDLE hProc = GetCurrentProcess();

void ListWin32KSysTable(const ULONG64& baseAddress, HMODULE hMod)
{
	SYMBOL_INFOW tableSym = {sizeof(tableSym), 0};
	if(symFromName(hProc, L"W32pServiceTable", &tableSym))
	{
		char* modIter = (char*)hMod;
		ULONG64 tableStart = tableSym.Address;
		ULONG64 offset = tableStart - baseAddress;
		DWORD* tableAddress = (DWORD*)(modIter + offset);
		DWORD64 symOffset = 0;
		SYMBOL_INFO_PACKAGEW function = {sizeof(function.si), 0};
		function.si.MaxNameLen = MAX_SYM_NAME;
		DWORD index = 1;
		DWORD sysCall = 0x1000;
		while(symFromAddr(hProc, *tableAddress, &symOffset, &function.si))
		{
			std::wcout << std::dec << index++ << L". syscall(" << std::hex << sysCall++ << L") " << function.si.Name << '\n';
			++tableAddress;
		}
	}
}

struct ZWFuncParams
{
	HMODULE hMod;
	std::map<DWORD, std::wstring>& serviceMap;
};

#ifndef _WIN64
BOOL CALLBACK ZWFuncProc(PSYMBOL_INFOW pSymInfo, ULONG SymbolSize, PVOID UserContext)
{
	ZWFuncParams* pParams = (ZWFuncParams*)UserContext;
	ULONG64 offset = pSymInfo->Address - pSymInfo->ModBase;
	char* pServiceNumber = ((char*)pParams->hMod + offset) + 1;  // move to the function, and then past the mov prefix
	DWORD dwService = *(DWORD*)pServiceNumber;
	if(dwService < 0x1000) // Gdi/User services start at 0x1000
	{
		std::map<DWORD, std::wstring>& serviceMap = pParams->serviceMap;
		if(serviceMap.find(dwService) == serviceMap.end())
		{
			serviceMap.insert(std::map<DWORD, std::wstring>::value_type(dwService, pSymInfo->Name));
		}
	}
    return TRUE;
}
#endif

void ListNTSysTable(const ULONG64& loadedBase, HMODULE hMod)
{
	std::map<DWORD, std::wstring> serviceMap;
	ZWFuncParams params = {hMod, serviceMap};
	symEnumSymbols(hProc, loadedBase, L"Nt*", &ZWFuncProc, &params);
	symEnumSymbols(hProc, loadedBase, L"Zw*", &ZWFuncProc, &params);
	DWORD count = 0;
	for(std::map<DWORD, std::wstring>::const_iterator iter = serviceMap.begin();
		iter != serviceMap.end();
		++iter)
	{
		std::wcout << std::dec << ++count << L". Syscall(" << std::hex << iter->first << L") " << iter->second << L'\n';
	}
}

int __cdecl wmain(int argc, wchar_t** argv)
//int WINAPI WinMain(HINSTANCE, HINSTANCE, LPSTR, int)
{
	if(argc > 1)
	{
		void (*parserFunc)(const ULONG64&, HMODULE) = NULL;
		std::wstring temp(argv[1]);
		std::transform(temp.begin(), temp.end(), temp.begin(), &towlower);
		if(temp.find(L"win32k.sys") != std::wstring::npos)
		{
			parserFunc = &ListWin32KSysTable;
		}
		else if(temp.find(L"ntdll.dll") != std::wstring::npos)
		{
			parserFunc = &ListNTSysTable;
		}
		else
		{
			std::wcout << L"Argument must be a path to ntdll.dll or win32k.sys";
		}
		if(parserFunc)
		{
			LoadDbghelp();
			HMODULE hMod = LoadLibraryEx(argv[1], NULL, DONT_RESOLVE_DLL_REFERENCES);
			if(init(hProc, NULL, FALSE))
			{
				ULONG64 baseAddress = 0;
				if((baseAddress = loader(hProc, NULL, argv[1], NULL, 0, 0, NULL, 0)) != 0)
				{
					parserFunc(baseAddress, hMod);
				}
				uninit(hProc);
			}
			FreeLibrary(hMod);
		}
	}
	else
	{
		std::wcout << L"Expected path to \"win32k.sys\" or \"ntdll.dll\" file as argument\n";
	}
	return 0;
}