#include <cstdio>
#include <ntdll.h>
#include <cstdint>

#include "structs.h"
#include "wow64ext.h"
#include "misc.h"
#include "syscall64.h"
#include "get_syscall64_ids.h"

/**
\file
*/

/**
\brief All functions doing direct syscalls are saved with a hash of their name
and the accompying syscall ID in the table (API_TO_INDEX*) ID_table.
To get the ID for a given hash the table can be searched.
(To speed up the search the table is sorted and a binary search is used)
*/
struct API_TO_INDEX
{
	DWORD hash;
	DWORD id;
} *ID_table = NULL;
DWORD ID_table_count = 0; //!< Count of entries in ID_table

/**
	\brief Parses the ID out of a x64 function.
	\return Returns the ID or INVALID_SYSCALL_ID on err
*/
static DWORD get_syscall_id(LPBYTE function)
{
	if(!function)
		return INVALID_SYSCALL_ID;

	/*
	00000000`77b61800 4c8bd1          mov     r10,rcx
	00000000`77b61803 b852000000      mov     eax,52h
	00000000`77b61808 0f05            syscall

	On Windows 10 it's:
	0:000> u NtOpenFile
	ntdll!NtOpenFile:
	00007ffe`62bf6720 4c8bd1          mov     r10,rcx
	00007ffe`62bf6723 b833000000      mov     eax,33h
	00007ffe`62bf6728 f604250803fe7f01 test    byte ptr [SharedUserData+0x308 (00000000`7ffe0308)],1
	00007ffe`62bf6730 7503            jne     ntdll!NtOpenFile+0x15 (00007ffe`62bf6735)
	00007ffe`62bf6732 0f05            syscall
	00007ffe`62bf6734 c3              ret
	00007ffe`62bf6735 cd2e            int     2Eh
	00007ffe`62bf6737 c3              ret
	*/

	const unsigned char MOV_R10_RCX_OPCODE[] = {0x4c, 0x8b, 0xd1};
	if(memcmp(function, MOV_R10_RCX_OPCODE, sizeof(MOV_R10_RCX_OPCODE)))
		return INVALID_SYSCALL_ID;

	const unsigned char SYSCALL_OPCODE[] = {0x0f, 0x05};
	const size_t SYSCALL_OFFSET_OLD = 8,
		SYSCALL_OFFSET_10 = 0x12;
	uint8_t* old = function + SYSCALL_OFFSET_OLD;
	uint8_t* win10 = function + SYSCALL_OFFSET_10;
	if (memcmp(old, SYSCALL_OPCODE, sizeof(SYSCALL_OPCODE)) &&
		memcmp(win10, SYSCALL_OPCODE, sizeof(SYSCALL_OPCODE))) 
	{
		return INVALID_SYSCALL_ID;
	}

	const unsigned char MOV_EAX_OPCODE = 0xB8;
	const size_t MOV_EAX_ID_OFFSET = 3;
	if(MOV_EAX_OPCODE != *(function + MOV_EAX_ID_OFFSET))
		return INVALID_SYSCALL_ID;

	return *(PDWORD)(function + MOV_EAX_ID_OFFSET + 1);
}


/**
	\brief Parses the given image to get the ID or just count them.
	\return Number of entries found / populated
	\param imageBase Image base of the DLL the info will be taken from
	\param ids The struct the info is put into. If NULL - won't be touched
*/
static DWORD get_syscall_ids(LPVOID imageBase, API_TO_INDEX* ids)
{
	if(!imageBase)
		return 0;

	PIMAGE_DOS_HEADER dos = (PIMAGE_DOS_HEADER)imageBase;
	// not a valid DOS header
	if(IMAGE_DOS_SIGNATURE != dos->e_magic)
		return 0;

	PIMAGE_NT_HEADERS64 nt = (PIMAGE_NT_HEADERS64)((LPBYTE)imageBase + dos->e_lfanew);
	// not a valid PE or not the correct architecture
	if(IMAGE_NT_SIGNATURE != nt->Signature || IMAGE_FILE_MACHINE_AMD64 != nt->FileHeader.Machine)
	{
		return 0;
	}

	// No exports?
	if(!nt->OptionalHeader.DataDirectory->Size)
		return 0;

	DWORD count = 0; //! Number of functions who do direct syscalls
	IMAGE_EXPORT_DIRECTORY *exportDir = (IMAGE_EXPORT_DIRECTORY *)((LPBYTE)imageBase + nt->OptionalHeader.DataDirectory->VirtualAddress);
	PDWORD nameRef = (DWORD *)((LPBYTE)imageBase + exportDir->AddressOfNames);
	WORD* ordinal = (WORD *)((LPBYTE)imageBase + exportDir->AddressOfNameOrdinals);
	DWORD* addressOfFunctions = (DWORD*)((LPBYTE)imageBase + exportDir->AddressOfFunctions);

	printf("Total functions: %d\n", exportDir->NumberOfNames);
	for(DWORD i = 0; i < exportDir->NumberOfNames; i++, nameRef++)
	{
		const char* name = (const char*)((LPBYTE)imageBase + (*nameRef));
		LPBYTE address = (LPBYTE)imageBase + addressOfFunctions[ordinal[i]];

		DWORD id = get_syscall_id(address);
		if(INVALID_SYSCALL_ID == id)
			continue;

		// Put it into the table
		if(ids)
		{
			ids[count].hash = hash(name);
			ids[count].id = id;
		}
		// Print info
		else
		{
			printf("%s (%X) = %X\n", name, hash(name), id);
		}
		count++;
	}

	return count;
}

/**
\brief (Bubble)Sorts the ID_table so it can be searched via binary search
*/
static void sort_ID_table()
{
	DWORD n = ID_table_count;

	bool swapped = false;
	do
	{
		for(unsigned int i = 0; i < n - 1; ++i)
		{
			if(ID_table[i].hash > ID_table[i + 1].hash)
			{
				API_TO_INDEX tmp = {ID_table[i].hash, ID_table[i].id};

				ID_table[i].hash = ID_table[i + 1].hash;
				ID_table[i].id = ID_table[i + 1].id;

				ID_table[i + 1].hash = tmp.hash;
				ID_table[i + 1].id = tmp.id;

				swapped = true;
			}
		}
		n = n - 1;
	} while(swapped == true && n);
}

/**
\brief Gets the path to the x64 ntdll in native format
\param path Output buffer
*/
static void get_ntdll_path(wchar_t* path)
{
	wchar_t systemDirectory[MAX_PATH];

	if(!path)
		return;

	GetSystemDirectory(systemDirectory, _countof(systemDirectory));

	wsprintfW(path, L"\\??\\%s\\ntdll.dll", systemDirectory);
}
/**
*/
static BOOL read_ntdll(LPBYTE& content, DWORD& size)
{
	wchar_t ntdllPath[MAX_PATH]; //! Path in the native format 
	get_ntdll_path(ntdllPath);
	printf("%ls\n", ntdllPath);

	_UNICODE_STRING_T<DWORD64> filename = {lstrlenW(ntdllPath) * sizeof(WCHAR), MAX_PATH  * sizeof(WCHAR), (DWORD64)ntdllPath};
	_OBJECT_ATTRIBUTES_T<DWORD64> obja = {sizeof(_OBJECT_ATTRIBUTES_T<DWORD64>), NULL, (DWORD64)&filename, OBJ_CASE_INSENSITIVE, NULL, NULL};
	_IO_STATUS_BLOCK_T<DWORD64> iostatusblock = {0};

	_HANDLE_T<DWORD64> fileHandle = {(DWORD64)INVALID_HANDLE_VALUE};
	NTSTATUS stat = DO_SYSCALL(get_basic_syscall_ID(NTOPENFILE), &fileHandle,
		FILE_READ_DATA | SYNCHRONIZE, &obja, &iostatusblock, FILE_SHARE_READ,
		FILE_NON_DIRECTORY_FILE | FILE_SYNCHRONOUS_IO_NONALERT);
	if(STATUS_SUCCESS != stat)
	{
		printf("NTSTATUS: %X\tHandle: %llX\n", stat, fileHandle.h);
		return FALSE;
	}
	printf("Handle: %llX\n", fileHandle.h);

	if(!(size = GetFileSize((HANDLE)fileHandle.h, NULL)))
	{
		DO_SYSCALL(get_basic_syscall_ID(NTCLOSE), fileHandle.h);
		return FALSE;
	}
	printf("Size of ntdll: %dkb\n", size / 1024);

	/* As no code is actually executed in the new ntdll - but only
	a few bytes are read from it - rw is enough */
	if(!(content = (LPBYTE)VirtualAlloc(NULL, size, MEM_COMMIT | MEM_RESERVE,
		PAGE_READWRITE)))
	{
		DO_SYSCALL(get_basic_syscall_ID(NTCLOSE), fileHandle.h);
		return FALSE;
	}
	printf("content %p\n", content);

	printf("ios ptr %X", &iostatusblock);
	LARGE_INTEGER offset = {0};
	stat = DO_SYSCALL(get_basic_syscall_ID(NTREADFILE), 
		(HANDLE)fileHandle.h, 
		NULL, // event
		NULL, // ApcRoutine
		NULL, // ApcContext
		&iostatusblock,
		content,
		size,
		&offset,
		NULL); // key
	if(STATUS_SUCCESS != stat)
	{
		VirtualFree(content, 0, MEM_RELEASE);
		content = NULL;
		DO_SYSCALL(get_basic_syscall_ID(NTCLOSE), fileHandle.h);
		printf("Reading failed: %X\t%X != %X\n", stat, size, offset.LowPart);
		return FALSE;
	}
	printf("ios: %X\toffset: %X\n", iostatusblock.Status, offset.LowPart);

	DO_SYSCALL(get_basic_syscall_ID(NTCLOSE), fileHandle.h);
	return TRUE;
}

/**
	\brief (Partly) maps the image.

	This does NOT fix relocations, imports, TLS, sets section protections yadda yadda.
	As this image is only used to parse out the IDs that is no problem.
	\param content Raw data of the file to be mapped
	\param mappedImage The finished image
*/
static BOOL map_PE(LPBYTE content, LPBYTE* mappedImage)
{
	PIMAGE_DOS_HEADER dos = (PIMAGE_DOS_HEADER)content;
	if(IMAGE_DOS_SIGNATURE != dos->e_magic)
	{
		printf("Not a valid DOS header: %s\n", content);
		return FALSE;
	}

	PIMAGE_NT_HEADERS64 nt = (PIMAGE_NT_HEADERS64)(content + dos->e_lfanew);
	if(IMAGE_NT_SIGNATURE != nt->Signature)
	{
		printf("Not a valid PE header\n");
		return FALSE;
	}

	/* Must be mapped properly (Actually almost -
	only the code section would be needed) so the IDs can be be
	parsed properly */
	if(!(*mappedImage = (LPBYTE)VirtualAlloc(NULL, nt->OptionalHeader.SizeOfImage, MEM_COMMIT | MEM_RESERVE,
		PAGE_READWRITE)))
	{
		printf("Can't alloc\n");
		return FALSE;
	}

	printf("Let's map that shit @ %p. %X sections. %X imageSize\n", *mappedImage,
		nt->FileHeader.NumberOfSections, nt->OptionalHeader.SizeOfImage);

	// Copy up to the first section (Size of optional header + delta optional header to start of image)
	memcpy(*mappedImage, content, nt->FileHeader.SizeOfOptionalHeader + ((LPBYTE)&nt->OptionalHeader - (LPBYTE)content));

	// Copy sections
#define IMAGE_FIRST_SECTION64( ntheader ) ((PIMAGE_SECTION_HEADER)        \
	((ULONG_PTR)(ntheader)+\
	FIELD_OFFSET(IMAGE_NT_HEADERS64, OptionalHeader) + \
	((ntheader))->FileHeader.SizeOfOptionalHeader   \
	))
	PIMAGE_SECTION_HEADER pish = IMAGE_FIRST_SECTION64(nt);
	for(unsigned int i = 0; i < nt->FileHeader.NumberOfSections; i++)
	{
		printf("Copy %Xbytes @ RVA %X From %X\n", pish->SizeOfRawData, pish->VirtualAddress, pish->PointerToRawData);
		memcpy((*mappedImage + pish->VirtualAddress), (content + pish->PointerToRawData),
			pish->SizeOfRawData);
		pish++;
	}

	return TRUE;
}

/**
\brief Initalize the ID_table.
*/
BOOL initalize_ID_table()
{
	LPBYTE ntdll = NULL, imageBase = NULL;
	DWORD ntdllSize = 0;
	if(!read_ntdll(ntdll, ntdllSize))
		return FALSE;
	if(!map_PE(ntdll, &imageBase))
		return FALSE;

	ID_table_count = get_syscall_ids(imageBase, NULL);
	printf("%i functions who do direct syscalls\n", ID_table_count);
	if(!ID_table_count)
		return FALSE;

	if(!(ID_table = (API_TO_INDEX*)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, ID_table_count * sizeof(API_TO_INDEX))))
		return FALSE;

	if(!get_syscall_ids(imageBase, ID_table))
		return FALSE;

	VirtualFree(ntdll, 0, MEM_RELEASE);
	VirtualFree(imageBase, 0, MEM_RELEASE);

	sort_ID_table();

	return TRUE;
}

/**
	\brief Has a few hardcoded IDs and returns the correct one for the current OS

	Which IDs exactly are hardcoded can be seen in the enum.
	Table made mostly with data from:
	x86: http://j00ru.vexillium.org/ntapi/
	WOW64: http://j00ru.vexillium.org/ntapi_64/

	\fixme >=Win8 Support?
*/
DWORD get_basic_syscall_ID(SYSCALL_IDS func)
{
	_KUSER_SHARED_DATA* _kuser_s_d = GET_KUSER_SHARED_DATA();
	ULONG majorVersion = _kuser_s_d->NtMajorVersion;
	ULONG minorVersion = _kuser_s_d->NtMinorVersion;
	NT_PRODUCT_TYPE productType = _kuser_s_d->NtProductType;
	_PEB* p = (_PEB*)__readfsdword(0x30);
	ULONG buildID = p->NtBuildNumber;

	switch(majorVersion)
	{
	case 5:
		// XP
		if(1 == minorVersion ||
			(2 == minorVersion && VER_NT_WORKSTATION == productType))
		{
			if(NTOPENFILE == func)
				return 0x30;
			else if(NTCREATEFILE == func)
				return 0x52;
			else if(NTREADFILE == func)
				return 0x03;
			else if(NTCLOSE == func)
				return 0x0C;
		}
		// Server 2003
		else
		{
			printf("SRV03 unsupported\n"); // fixme
			if(NTOPENFILE == func)
				return 0x7a;
			else if(NTCREATEFILE == func)
				return 0x26;
			else if(NTREADFILE == func)
				return 0xbf;
			else if(NTCLOSE == func)
				return 0x1b;
		}
		break;
	case 6:
		switch(minorVersion)
		{
		// Vista SP0-2 & Server 2008
		case 0:
			printf("Vista unsupported\n"); // fixme
			if(NTOPENFILE == func)
				return 0xba;
			else if(NTCREATEFILE == func)
				return 0x3c;
			else if(NTREADFILE == func)
				return 0x102;
			else if(NTCLOSE == func)
				return 0x30;
			break;
		// Win7
		case 1:
			if(NTOPENFILE == func)
				return 0x30;
			else if(NTCREATEFILE == func)
				return 0x52;
			else if(NTREADFILE == func)
				return 0x03;
			else if(NTCLOSE == func)
				return 0x0C;
		// Win 8
		case 2:
			printf("Win8 unsupported\n"); // fixme
			if(NTOPENFILE == func)
				return 0xe7;
			else if(NTCREATEFILE == func)
				return 0x15f;
			else if(NTREADFILE == func)
				return 0x86;
			else if(NTCLOSE == func)
				return 0x0171;
			break;
		}
		break;
	// Win 10
	case 10:
		if (NTOPENFILE == func)
			return 0x0033;
		else if (NTCREATEFILE == func)
			return 0x55;
		else if (NTREADFILE == func)
			return 0x06;
		else if (NTCLOSE == func)
			return 0x0f;
	}

	return INVALID_SYSCALL_ID;
}

DWORD get_syscall_ID(DWORD func)
{
	DWORD left = 0, right = ID_table_count;

	while(left != right)
	{
		DWORD middle = (left + right) / 2;

		if(func == ID_table[middle].hash)
			return ID_table[middle].id;
		if(func > ID_table[middle].hash)
			left = middle + 1;
		else
			right = middle - 1;
	}

	if(func == ID_table[left].hash)
	{
		return ID_table[left].id;
	}

	return INVALID_SYSCALL_ID;
}

VOID destroy_ID_table()
{
	HeapFree(GetProcessHeap(), 0, ID_table);
}

#if 0
BOOL test_id_table()
{
	LPVOID imageBase = GetModuleHandle(TEXT("ntdll.dll"));

	PIMAGE_DOS_HEADER dos = (PIMAGE_DOS_HEADER)imageBase;
	// not a valid DOS header
	if(IMAGE_DOS_SIGNATURE != dos->e_magic)
		return FALSE;

	PIMAGE_NT_HEADERS nt = (PIMAGE_NT_HEADERS)((LPBYTE)imageBase + dos->e_lfanew);
	// not a valid PE or not the correct architecture
	if(IMAGE_NT_SIGNATURE != nt->Signature || IMAGE_FILE_MACHINE_I386 != nt->FileHeader.Machine)
	{
		return FALSE;
	}

	// No exports?
	if(!nt->OptionalHeader.DataDirectory->Size)
		return FALSE;

	IMAGE_EXPORT_DIRECTORY *exportDir = (IMAGE_EXPORT_DIRECTORY *)((LPBYTE)imageBase + nt->OptionalHeader.DataDirectory->VirtualAddress);
	PDWORD nameRef = (DWORD *)((LPBYTE)imageBase + exportDir->AddressOfNames);
	WORD* ordinal = (WORD *)((LPBYTE)imageBase + exportDir->AddressOfNameOrdinals);
	DWORD* addressOfFunctions = (DWORD*)((LPBYTE)imageBase + exportDir->AddressOfFunctions);

	for(DWORD i = 0; i < exportDir->NumberOfNames; i++, nameRef++)
	{
		const char* name = (const char*)((LPBYTE)imageBase + (*nameRef));
		LPBYTE address = (LPBYTE)imageBase + addressOfFunctions[ordinal[i]];

		DWORD id = get_syscall_id(address);
		if(INVALID_SYSCALL_ID == id)
			continue;

		// At address is a function doing direct syscalls - check if can be found
		if(INVALID_SYSCALL_ID == get_syscall_ID(hash(name)))
		{
			printf("%s not found", name);
			return FALSE;
		}
	}

	return TRUE;
}

#endif