// Copyright Epic Games, Inc. All Rights Reserved.

#include <zencore/windows.h>

#if ZEN_PLATFORM_WINDOWS

#	include <Aclapi.h>
#	include <sddl.h>
#	include <stdio.h>
#	include <userenv.h>

#	pragma comment(lib, "userenv.lib")
#	pragma comment(lib, "Advapi32.lib")

// Helper: add (or merge) an ACE for a SID on a directory
DWORD
GrantFolderAccessToSid(LPCWSTR path, PSID sid, DWORD accessMask)
{
	PACL				 oldDacl = nullptr, newDacl = nullptr;
	PSECURITY_DESCRIPTOR sd = nullptr;

	DWORD err = GetNamedSecurityInfoW(path, SE_FILE_OBJECT, DACL_SECURITY_INFORMATION, nullptr, nullptr, &oldDacl, nullptr, &sd);
	if (err != ERROR_SUCCESS)
		return err;

	EXPLICIT_ACCESSW ea		= {};
	ea.grfAccessMode		= GRANT_ACCESS;
	ea.grfAccessPermissions = accessMask;
	ea.grfInheritance		= OBJECT_INHERIT_ACE | CONTAINER_INHERIT_ACE;
	ea.Trustee.TrusteeForm	= TRUSTEE_IS_SID;
	ea.Trustee.TrusteeType	= TRUSTEE_IS_WELL_KNOWN_GROUP;	// fine for AppContainer SIDs
	ea.Trustee.ptstrName	= (LPWSTR)sid;

	err = SetEntriesInAclW(1, &ea, oldDacl, &newDacl);
	if (err == ERROR_SUCCESS)
	{
		err = SetNamedSecurityInfoW((LPWSTR)path, SE_FILE_OBJECT, DACL_SECURITY_INFORMATION, nullptr, nullptr, newDacl, nullptr);
	}

	if (newDacl)
		LocalFree(newDacl);

	if (sd)
		LocalFree(sd);

	return err;
}

int
zwmain()
{
	// 1) Choose a stable AppContainer name for your app/sandbox
	LPCWSTR containerName = L"com.epicgames.zenmaster.sandbox";
	LPCWSTR displayName	  = L"zenmaster sandbox";
	LPCWSTR description	  = L"Sandbox for running restricted zenserver processes";

	// 1a) Create the profile once (ok if it already exists)
	// Requires elevation the first time you create it. If it exists, creation will fail with ERROR_ALREADY_EXISTS—which is fine.
	HRESULT hr = CreateAppContainerProfile(containerName, displayName, description, nullptr, 0, nullptr);
	if (hr == HRESULT_FROM_WIN32(ERROR_ALREADY_EXISTS))
	{
		hr = S_OK;
	}
	if (FAILED(hr))
	{
		wprintf(L"CreateAppContainerProfile failed: 0x%08X\n", hr);
		return 1;
	}

	// 1b) Get the AppContainer SID
	PSID appContainerSid = nullptr;
	hr					 = DeriveAppContainerSidFromAppContainerName(containerName, &appContainerSid);
	if (FAILED(hr))
	{
		wprintf(L"DeriveAppContainerSidFromAppContainerName failed: 0x%08X\n", hr);
		return 1;
	}

	// 2) Grant access to specific folders (repeat for each allowed path)
	// Example: allow read/write/list/execute on D:\Allowed\WorkDir
	LPCWSTR allowedFolder = L"D:\\Allowed\\WorkDir";
	CreateDirectoryW(allowedFolder, nullptr);  // ensure it exists (optional)

	DWORD fsAccess = FILE_GENERIC_READ | FILE_GENERIC_WRITE | FILE_GENERIC_EXECUTE | READ_CONTROL | SYNCHRONIZE;

	DWORD err = GrantFolderAccessToSid(allowedFolder, appContainerSid, fsAccess);
	if (err != ERROR_SUCCESS)
	{
		wprintf(L"GrantFolderAccessToSid failed: %lu\n", err);
		FreeSid(appContainerSid);
		return 1;
	}

	// 3) Prepare SECURITY_CAPABILITIES with the AppContainer SID (no broad capabilities)
	SID_AND_ATTRIBUTES	  capSidAndAttrs = {};	// none for now
	SECURITY_CAPABILITIES secCaps		 = {};
	secCaps.AppContainerSid				 = appContainerSid;
	secCaps.Capabilities				 = nullptr;	 // you can add capability SIDs here if needed
	secCaps.CapabilityCount				 = 0;
	secCaps.Reserved					 = 0;

	// 4) Build attribute list with PROC_THREAD_ATTRIBUTE_SECURITY_CAPABILITIES
	SIZE_T attrListSize = 0;
	InitializeProcThreadAttributeList(nullptr, 1, 0, &attrListSize);
	LPPROC_THREAD_ATTRIBUTE_LIST attrList = (LPPROC_THREAD_ATTRIBUTE_LIST)HeapAlloc(GetProcessHeap(), 0, attrListSize);

	if (!attrList)
	{
		wprintf(L"HeapAlloc failed\n");
		FreeSid(appContainerSid);
		return 1;
	}

	if (!InitializeProcThreadAttributeList(attrList, 1, 0, &attrListSize))
	{
		wprintf(L"InitializeProcThreadAttributeList failed: %lu\n", GetLastError());
		HeapFree(GetProcessHeap(), 0, attrList);
		FreeSid(appContainerSid);
		return 1;
	}

	if (!UpdateProcThreadAttribute(attrList, 0, PROC_THREAD_ATTRIBUTE_SECURITY_CAPABILITIES, &secCaps, sizeof(secCaps), nullptr, nullptr))
	{
		wprintf(L"UpdateProcThreadAttribute failed: %lu\n", GetLastError());
		DeleteProcThreadAttributeList(attrList);
		HeapFree(GetProcessHeap(), 0, attrList);
		FreeSid(appContainerSid);
		return 1;
	}

	// 5) Launch the target EXE in the AppContainer
	STARTUPINFOEXW si  = {};
	si.StartupInfo.cb  = sizeof(si);
	si.lpAttributeList = attrList;

	PROCESS_INFORMATION pi		  = {};
	wchar_t				cmdline[] = L"C:\\Path\\To\\YourChild.exe";	 // child process
	BOOL				ok		  = CreateProcessW(nullptr,
							   cmdline,
							   nullptr,
							   nullptr,
							   FALSE,
							   EXTENDED_STARTUPINFO_PRESENT,
							   nullptr,		   // environment (null = inherit)
							   allowedFolder,  // set working dir to an allowed folder (optional but handy)
							   (LPSTARTUPINFOW)&si,
							   &pi);

	if (!ok)
	{
		wprintf(L"CreateProcessW failed: %lu\n", GetLastError());
	}
	else
	{
		wprintf(L"Launched pid=%lu in AppContainer.\n", pi.dwProcessId);
		CloseHandle(pi.hThread);
		CloseHandle(pi.hProcess);
	}

	// Cleanup
	DeleteProcThreadAttributeList(attrList);
	HeapFree(GetProcessHeap(), 0, attrList);
	FreeSid(appContainerSid);
	return ok ? 0 : 1;
}
#endif