#include <array>
#include <ntdll.h>
#include <Windows.h>
#include <shlobj.h>
#include <Shlwapi.h>
#include "structs.h"
#include "misc.h"
#include "wow64ext.h"
#include "syscall64.h"
#include "get_syscall64_ids.h"

void WINAPI RtlInitUnicodeString(
	unsigned char* target,
	PCWSTR source)
{
	SET_MEMBER_PTR(_UNICODE_STRING_T, Buffer, target, source);
	if(source)
	{
		unsigned int length = lstrlenW(source) * sizeof(WCHAR);
		if(length > 0xfffc)
			length = 0xfffc;
		SET_MEMBER(_UNICODE_STRING_T, Length, target, length);
		SET_MEMBER(_UNICODE_STRING_T, MaximumLength, target, length + sizeof(WCHAR));
	}
	else
	{
		SET_MEMBER(_UNICODE_STRING_T, Length, target, 0);
		SET_MEMBER(_UNICODE_STRING_T, MaximumLength, target, 0);
	}
}


VOID initialize_object_attributes(unsigned char* p, LPVOID n, ULONG a)
{
	SET_MEMBER(_OBJECT_ATTRIBUTES_T, uLength, p, is_WOW64() ? sizeof(_OBJECT_ATTRIBUTES_T<DWORD64>) : sizeof(_OBJECT_ATTRIBUTES_T<DWORD>));
	SET_MEMBER(_OBJECT_ATTRIBUTES_T, hRootDirectory, p, NULL);
	SET_MEMBER(_OBJECT_ATTRIBUTES_T, uAttributes, p, a);
	SET_MEMBER_PTR(_OBJECT_ATTRIBUTES_T, pObjectName, p, n);
	SET_MEMBER(_OBJECT_ATTRIBUTES_T, pSecurityDescriptor, p, NULL); // Actually a ptr but we set a NULL
	SET_MEMBER(_OBJECT_ATTRIBUTES_T, pSecurityQualityOfService, p, NULL);
}

BOOL file_test()
{
	ALLOC_STRUCTURE(_UNICODE_STRING_T, filename);
	wchar_t desktopPath[MAX_PATH];
	lstrcpyW(desktopPath, L"\\??\\");
	if(!SHGetSpecialFolderPathW(HWND_DESKTOP, &desktopPath[lstrlenW(desktopPath)], CSIDL_DESKTOPDIRECTORY, FALSE))
		return FALSE;
	PathAppendW(desktopPath, L"WaitAMinute.HowDidThisGetHere.txt");
	RtlInitUnicodeString(filename, desktopPath);
	printf("Desktop: %ws\n", desktopPath);

	ALLOC_STRUCTURE(_OBJECT_ATTRIBUTES_T, obja);
	initialize_object_attributes(obja, (void*)filename, OBJ_CASE_INSENSITIVE);

	ALLOC_STRUCTURE(_HANDLE_T, fileHandle);
	ALLOC_STRUCTURE(_IO_STATUS_BLOCK_T, iostatusblock);

	NTSTATUS stat = DO_SYSCALL(get_syscall_ID(NTCREATEFILE), fileHandle, FILE_READ_DATA | FILE_WRITE_DATA | SYNCHRONIZE,
		obja, iostatusblock, NULL, FILE_ATTRIBUTE_NORMAL, FILE_SHARE_READ, FILE_OVERWRITE_IF,
		FILE_NON_DIRECTORY_FILE | FILE_SYNCHRONOUS_IO_NONALERT, NULL, 0);
	if(STATUS_SUCCESS != stat)
	{
		printf("NTSTATUS: %X\n", stat);
		return FALSE;
	}
	printf("Handle: %X\tIO.uInformation:%X\n", *(PHANDLE)GET_MEMBER(_HANDLE_T, h, fileHandle), *(PDWORD)GET_MEMBER(_IO_STATUS_BLOCK_T, uInformation, iostatusblock));

	FREE_STRUCTURE(_UNICODE_STRING_T, filename);
	FREE_STRUCTURE(_OBJECT_ATTRIBUTES_T, obja);

	const char TEST_STRING[] = {"0123456789ABCDEF"};
	if(STATUS_SUCCESS != (stat = DO_SYSCALL(get_syscall_ID(NTWRITEFILE), *(PHANDLE)GET_MEMBER(_HANDLE_T, h, fileHandle),
		NULL, NULL, NULL, iostatusblock, TEST_STRING, lstrlenA(TEST_STRING), NULL, NULL)))
	{
		printf("Write: %X\t\n", stat);
		return FALSE;
	}

	FlushFileBuffers(*(PHANDLE)GET_MEMBER(_HANDLE_T, h, fileHandle));

	char buffer[sizeof(TEST_STRING)+1] = {0};
	LARGE_INTEGER offset = {0};
	if(STATUS_SUCCESS != (stat = DO_SYSCALL(get_syscall_ID(NTREADFILE), *(PHANDLE)GET_MEMBER(_HANDLE_T, h, fileHandle),
		NULL, NULL, NULL, iostatusblock, buffer, sizeof(buffer), &offset, NULL)))
	{
		printf("Reading failed: %X\t%X != %X\n", stat, sizeof(buffer), offset.LowPart);
		return FALSE;
	}
	buffer[*(PDWORD)GET_MEMBER(_IO_STATUS_BLOCK_T, uInformation, iostatusblock)] = 0;
	printf("io stat: %X\tio inf: %X\tread: %X\n", *(PDWORD)GET_MEMBER(_IO_STATUS_BLOCK_T, Status, iostatusblock), *(PDWORD)GET_MEMBER(_IO_STATUS_BLOCK_T, uInformation, iostatusblock), offset.LowPart);

	DO_SYSCALL(get_syscall_ID(NTCLOSE), *(PHANDLE)GET_MEMBER(_HANDLE_T, h, fileHandle));
	FREE_STRUCTURE(_HANDLE_T, fileHandle);

	printf("'%s' == '%s'\n", TEST_STRING, buffer);
	return 0 == lstrcmpA(TEST_STRING, buffer);

}