#include <windows.h>
#include <userenv.h>
#include <sddl.h>
#include <iostream>
#include <string>
#include <vector>
#include <aclapi.h>

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

bool CreateAppContainerSid(const std::wstring& name, PSID* appContainerSid) {
    HRESULT hr = CreateAppContainerProfile(name.c_str(), name.c_str(), L"Test AppContainer Profile", nullptr, 0, appContainerSid);
    if (FAILED(hr)) {
        HRESULT hr = DeriveAppContainerSidFromAppContainerName(name.c_str(), appContainerSid);
        if (FAILED(hr)) {
            std::wcerr << L"Failed to derive AppContainer SID. Creating profile..." << std::endl;
            return false;
        }
        else
		{
			std::wcout << L"AppContainer SID derived successfully." << std::endl;
		}
    }
    else
    {
		std::wcout << L"AppContainer profile created successfully." << std::endl;
    }

    return true;
}

bool CreatePipes(HANDLE& readOut, HANDLE& writeOut) {
    SECURITY_ATTRIBUTES saAttr = {};
    saAttr.nLength = sizeof(SECURITY_ATTRIBUTES);
    saAttr.bInheritHandle = TRUE;
    saAttr.lpSecurityDescriptor = nullptr;

    if (!CreatePipe(&readOut, &writeOut, &saAttr, 0)) {
        std::cerr << "Stdout pipe creation failed\n";
        return false;
    }

    if (!SetHandleInformation(readOut, HANDLE_FLAG_INHERIT, 0)) {
        std::cerr << "Stdout SetHandleInformation failed\n";
        return false;
    }

    return true;
}

void ReadFromPipe(HANDLE readHandle) {
    CHAR buffer[4096];
    DWORD bytesRead;
    while (ReadFile(readHandle, buffer, sizeof(buffer) - 1, &bytesRead, nullptr) && bytesRead > 0) {
        buffer[bytesRead] = '\0';
        std::cout << buffer;
    }
}

bool GrantAccessToPathForAppContainer(const std::wstring& path, PSID appContainerSid, DWORD accessMask) {
    EXPLICIT_ACCESSW ea = {};
    ea.grfAccessPermissions = accessMask;
    ea.grfAccessMode = GRANT_ACCESS;
    ea.grfInheritance = OBJECT_INHERIT_ACE | CONTAINER_INHERIT_ACE;
	ea.Trustee.MultipleTrusteeOperation = NO_MULTIPLE_TRUSTEE;
    ea.Trustee.pMultipleTrustee = nullptr;
    ea.Trustee.TrusteeForm = TRUSTEE_IS_SID;
    ea.Trustee.TrusteeType = TRUSTEE_IS_GROUP;
    ea.Trustee.ptstrName = (LPWSTR)appContainerSid;

    PACL oldDacl = nullptr, newDacl = nullptr;
    PSECURITY_DESCRIPTOR sd = nullptr;

    DWORD result = GetNamedSecurityInfoW(
        path.c_str(), SE_FILE_OBJECT,
        DACL_SECURITY_INFORMATION,
        nullptr, nullptr, &oldDacl, nullptr, &sd);

    if (result != ERROR_SUCCESS) {
        std::wcerr << L"Failed to get security info for: " << path << L" Error: " << result << "\n";
        return false;
    }

    result = SetEntriesInAclW(1, &ea, oldDacl, &newDacl);
    if (result != ERROR_SUCCESS) {
        std::wcerr << L"SetEntriesInAclW failed: " << result << "\n";
        if (sd) LocalFree(sd);
        return false;
    }

    result = SetNamedSecurityInfoW(
        (LPWSTR)path.c_str(), SE_FILE_OBJECT,
        DACL_SECURITY_INFORMATION,
        nullptr, nullptr, newDacl, nullptr);

    if (result != ERROR_SUCCESS) {
        std::wcerr << L"SetNamedSecurityInfoW failed: " << result << "\n";
    }

	std::wcout << "Updated ACL for " << path << L" with AppContainer SID\n";

    if (sd) LocalFree(sd);
    if (newDacl) LocalFree(newDacl);

    return result == ERROR_SUCCESS;
}

std::wstring GetLastErrorMessage(DWORD errorCode = GetLastError()) {
    LPWSTR buffer = nullptr;
    DWORD size = FormatMessageW(
        FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
        nullptr,
        errorCode,
        MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
        (LPWSTR)&buffer,
        0,
        nullptr);

    std::wstring message(buffer, size);
    LocalFree(buffer);
    return message;
}

int wmain(int argc, wchar_t* argv[]) {
    if (argc < 2) {
        std::wcerr << L"Usage: AppContainerLauncher.exe <command line>\n";
        return 1;
    }

    std::vector<std::pair<std::wstring, DWORD>> accessRequests;
    std::vector<std::wstring> commandParts;
    std::wstring appExe;

    for (int i = 1; i < argc; ++i) {
        std::wstring arg = argv[i];
        if (arg == L"--grant-read" && i + 1 < argc) {
            accessRequests.emplace_back(argv[++i], FILE_GENERIC_READ | FILE_GENERIC_EXECUTE | GENERIC_READ | GENERIC_EXECUTE);
        }
        else if (arg == L"--grant-write" && i + 1 < argc) {
            accessRequests.emplace_back(argv[++i], FILE_ALL_ACCESS | GENERIC_ALL);
        }
        else {
            commandParts.emplace_back(arg);
        }
    }

    if (commandParts.empty()) {
        std::wcerr << L"No command provided.\n";
        return 1;
    }

    std::wstring commandLine;
	for (const auto& part : commandParts) {
		commandLine += part + L" ";
	}
	if (commandParts.size() > 0)
	    appExe = commandParts[0];

    std::wstring containerName = L"TestPythonAppContainer";
    PSID appContainerSid = nullptr;
    if (!CreateAppContainerSid(containerName, &appContainerSid)) {
        return 1;
    }

    for (const auto& p : accessRequests) {
        GrantAccessToPathForAppContainer(p.first, appContainerSid, p.second);
    }

    HANDLE readOut, writeOut;
    if (!CreatePipes(readOut, writeOut)) {
        return 1;
    }

    STARTUPINFOEXW siex = {};
    PROCESS_INFORMATION pi = {};
    SIZE_T attrListSize = 0;

    InitializeProcThreadAttributeList(nullptr, 2, 0, &attrListSize);
    siex.lpAttributeList = (LPPROC_THREAD_ATTRIBUTE_LIST)HeapAlloc(GetProcessHeap(), 0, attrListSize);
    if (!InitializeProcThreadAttributeList(siex.lpAttributeList, 2, 0, &attrListSize)) {
        std::cerr << "InitializeProcThreadAttributeList failed\n";
        return 1;
    }

    SECURITY_CAPABILITIES sc = { };
	sc.AppContainerSid = appContainerSid;
	sc.Capabilities = nullptr;
    sc.CapabilityCount = 0;
    sc.Reserved = 0;
    if (!UpdateProcThreadAttribute(siex.lpAttributeList, 0,
        PROC_THREAD_ATTRIBUTE_SECURITY_CAPABILITIES,
        &sc,
        sizeof(SECURITY_CAPABILITIES), nullptr, nullptr)) {
        std::cerr << "UpdateProcThreadAttribute failed\n";
        return 1;
    }

    HANDLE inheritHandles[] = { writeOut }; // Only pass the write end to the child
	if (!UpdateProcThreadAttribute(siex.lpAttributeList, 0,
		PROC_THREAD_ATTRIBUTE_HANDLE_LIST,
		inheritHandles,
		sizeof(inheritHandles), nullptr, nullptr)) {
		std::cerr << "UpdateProcThreadAttribute failed\n";
		return 1;
	}

    siex.StartupInfo.cb = sizeof(siex);
    siex.StartupInfo.hStdOutput = writeOut;
    siex.StartupInfo.hStdError = writeOut;
    siex.StartupInfo.dwFlags |= STARTF_USESTDHANDLES;

    BOOL success = CreateProcessW(
        nullptr,
        &commandLine[0],
        nullptr,
        nullptr,
        TRUE,
        EXTENDED_STARTUPINFO_PRESENT | CREATE_NO_WINDOW,
        nullptr,
        nullptr,
        &siex.StartupInfo,
        &pi
    );

    CloseHandle(writeOut); // The child process now owns the write end

    if (!success) {
        DWORD err = GetLastError();
        std::wcerr << L"CreateProcess failed (" << err << L"): " << GetLastErrorMessage(err) << "\n";
        return 1;
    }

    ReadFromPipe(readOut);

    WaitForSingleObject(pi.hProcess, INFINITE);

    CloseHandle(pi.hProcess);
    CloseHandle(pi.hThread);
    CloseHandle(readOut);
    if (appContainerSid) {
        FreeSid(appContainerSid);
    }
    DeleteProcThreadAttributeList(siex.lpAttributeList);
    HeapFree(GetProcessHeap(), 0, siex.lpAttributeList);

    return 0;
}