/* This code copies the target process in our own address space, performs relocation and executes new thread with a new function */

#include <Windows.h>
#include <stdio.h>
#include <string.h>

typedef BOOL(WINAPI* DllMain_t)(
    HINSTANCE hinstDLL,
    DWORD fdwReason,
    LPVOID lpReserved);

void loadDLL(const char *dllName)
{

    HANDLE hFile, hMapping;
    void *uiBaseAddress;

    hFile = CreateFileA((LPCSTR)dllName, GENERIC_READ, 0, 0, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, 0);

    if (hFile == INVALID_HANDLE_VALUE)
    {
        puts("Unable to open file\n");
        exit(1);
    }

    hMapping = CreateFileMapping(hFile, 0, PAGE_READONLY | SEC_COMMIT, 0, 0, 0);

    if (hMapping == INVALID_HANDLE_VALUE)
    {
        puts("Unable to map file\n");
        CloseHandle(hFile);
        exit(1);
    }

    uiBaseAddress = MapViewOfFile(hMapping, FILE_MAP_READ, 0, 0, 0);

    if (!uiBaseAddress)
    {
        puts("Unable mapping view of file to memory\n");
        UnmapViewOfFile(hMapping);
        CloseHandle(hFile);
        exit(1);
    }

    UINT_PTR pViewBase = (UINT_PTR)uiBaseAddress; // Mapping address
    IMAGE_NT_HEADERS *pNtHeader = (IMAGE_NT_HEADERS*)(pViewBase + ((IMAGE_DOS_HEADER*)pViewBase)->e_lfanew);


    /* Get location of NT headers */
    //IMAGE_NT_HEADERS *pNtHeader = (IMAGE_NT_HEADERS*)((PIMAGE_DOS_HEADER)pViewBase)->e_lfanew + (LONG)pViewBase;

    /* Allocate a memory block of pNtHeader.OptionalHeader.SizeOfImage bytes at position pNtHeader.OptionalHeader.ImageBase */
    void *pLibraryBase = VirtualAlloc((void *)pNtHeader->OptionalHeader.ImageBase, pNtHeader->OptionalHeader.SizeOfImage, MEM_RESERVE, PAGE_EXECUTE_READWRITE);

    if (!pLibraryBase)
    {
        puts("Unable to allocate memory for library in process\n");
        exit(1);
    }

    PIMAGE_SECTION_HEADER pSectionHeader = (PIMAGE_SECTION_HEADER)((unsigned long)pNtHeader + sizeof(IMAGE_NT_HEADERS));

    for (unsigned int i = 0; i < pNtHeader->FileHeader.NumberOfSections; i++)
    {
        /* Section start in memory */
        void *pSectionMemoryAddress = (void *)(pSectionHeader[i].VirtualAddress + (unsigned long)pLibraryBase);

        /* Section start in file */
        void *pSectionFileAddress = (void *)(pSectionHeader[i].PointerToRawData + (unsigned long)pViewBase);

        unsigned long ulSectionSize = pSectionHeader[i].SizeOfRawData;

        if (!ulSectionSize)
        {
            if (pSectionHeader[i].Characteristics & IMAGE_SCN_CNT_CODE)
                ulSectionSize = pNtHeader->OptionalHeader.SizeOfCode;

            else if (pSectionHeader[i].Characteristics & IMAGE_SCN_CNT_INITIALIZED_DATA)
                ulSectionSize = pNtHeader->OptionalHeader.SizeOfInitializedData;

            else if (pSectionHeader[i].Characteristics & IMAGE_SCN_CNT_UNINITIALIZED_DATA)
                ulSectionSize = pNtHeader->OptionalHeader.SizeOfUninitializedData;
        }

        if (VirtualAlloc(pSectionMemoryAddress, ulSectionSize, MEM_COMMIT, PAGE_READWRITE) != pSectionMemoryAddress)
        {
            puts("Error in commit memory\n");
            exit(1);
        }

        if (ulSectionSize > 0)
        {
            /* Copy sections from disk to memory */
            memcpy(pSectionMemoryAddress, pSectionFileAddress, ulSectionSize);
        }
    }

    /* Resolve imports and setup IAT */
    PIMAGE_IMPORT_DESCRIPTOR pImportDescriptor = (PIMAGE_IMPORT_DESCRIPTOR)(pNtHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress + (unsigned long)pLibraryBase);

    for (unsigned int i = 0; i < pImportDescriptor[i].FirstThunk; i++)
    {
        PIMAGE_THUNK_DATA pOFT = (PIMAGE_THUNK_DATA)(pImportDescriptor[i].OriginalFirstThunk + (unsigned long)pLibraryBase);
        PIMAGE_THUNK_DATA pIAT = (PIMAGE_THUNK_DATA)(pImportDescriptor[i].FirstThunk + (unsigned long)pLibraryBase);

        HMODULE hImportLib = LoadLibraryA((LPCSTR)((unsigned long)pLibraryBase + pImportDescriptor[i].Name));
        if (!hImportLib)
        {
            puts("Unable to find dependency library\n");
            exit(1);
        }

        for (unsigned int x = 0; x < pOFT[x].u1.Function; x++)
        {
            unsigned long ulImportNameOrdinal = 0;

            if (pOFT[x].u1.Function & (1 >> 31))
            {
                ulImportNameOrdinal = pIAT[x].u1.Function & ~(1 >> 31);
            }
            else
            {
                /* Address of the function called from the DLL */
                PIMAGE_IMPORT_BY_NAME pImport = PIMAGE_IMPORT_BY_NAME((unsigned long)pLibraryBase + pOFT[x].u1.Function);
                ulImportNameOrdinal = (unsigned long)pImport->Name;
            }

            /* Patch the IAT with the address of the function */
            if (!(pIAT[x].u1.Function = (unsigned long)(GetProcAddress(hImportLib, (const char *)ulImportNameOrdinal))))
            {
                puts("Error finding import\n");
                exit(1);
            }
        }
    }

    /* Do relocation as described in Relocation directory */
    DWORD_PTR difference = (DWORD)pLibraryBase - pNtHeader->OptionalHeader.ImageBase;
    if (difference != 0)
    {
        IMAGE_BASE_RELOCATION *relo_start = (IMAGE_BASE_RELOCATION*)((DWORD)pLibraryBase + pNtHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress);
        IMAGE_BASE_RELOCATION *relo_end = (IMAGE_BASE_RELOCATION*)((DWORD)relo_start + pNtHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].Size - sizeof(IMAGE_BASE_RELOCATION));

        for (; relo_start < relo_end; relo_start = (IMAGE_BASE_RELOCATION*)((DWORD)relo_start + relo_start->SizeOfBlock))
        {
            /* To determine the number of relocation in this block, substract the size of an
            IMAGE_BASE_RELOCATION (8 bytes) from the value of this field, and then divide
            by 2 (size of WORD).
            */
            WORD *reloc_item = (WORD*)(relo_start + 1);
            DWORD num_items = (relo_start->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION)) / sizeof(WORD);

            for (DWORD i = 0; i < num_items; ++i, ++reloc_item)
            {
                switch (*reloc_item >> 12)
                {
                case IMAGE_REL_BASED_ABSOLUTE:
                    break;
                case IMAGE_REL_BASED_HIGHLOW:
                    *(DWORD_PTR*)((DWORD)pLibraryBase + relo_start->VirtualAddress + (*reloc_item & 0xFFF)) += difference;
                }
            }
        }
    }


    /* Relocation have been done. Now apply proper page permissions */
    //  for (unsigned int i = 0; i < pNtHeader->FileHeader.NumberOfSections; i++)
    //  {
    //      void *pMemorySectionAddress = (void*)(pSectionHeader->VirtualAddress + (unsigned long)pLibraryBase);
    //      setSectionPermission(pMemorySectionAddress, pSectionHeader[i].Misc.VirtualSize, pSectionHeader[i].Characteristics);
    //  }

    DllMain_t entry = (DllMain_t)((DWORD)pLibraryBase + pNtHeader->OptionalHeader.AddressOfEntryPoint);
    entry((HINSTANCE)pLibraryBase, DLL_PROCESS_ATTACH, 0);

    UnmapViewOfFile(hMapping);
    CloseHandle(hFile);
}

DWORD WINAPI ThreadProc(PVOID p)
{
    MessageBox(NULL, "FUNCTION INJECTED", "HI", MB_ICONINFORMATION);
    char *dll = "C:\\Users\\root_sec\\Documents\\Visual Studio 2015\\Projects\\ManualMapping\\Debug\\TestFinal.dll";
    loadDLL(dll);
    return 0;
}

int main(int argc, char **argv)
{
    if (argc != 2)
    {
        printf("[!] Not enough arguments\n");
        return -1;
    }

    ULONG i, *p;
    WORD *relocItems;
    PIMAGE_DOS_HEADER pDosh;
    PIMAGE_NT_HEADERS pNth;
    PIMAGE_BASE_RELOCATION pBase;

    printf("[*] Opening the target process ...\n");

    HANDLE hProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE, atoi(argv[1]));
    if (hProcess == INVALID_HANDLE_VALUE)
    {
        puts("[!] Unable to open process");
        return -1;
    }

    HANDLE imageBase = GetModuleHandle(NULL);
    printf("[*] Image Base in current process: 0x%x\n", (unsigned int)imageBase);

    pDosh = (PIMAGE_DOS_HEADER)imageBase;
    pNth = (PIMAGE_NT_HEADERS)((PUCHAR)imageBase + pDosh->e_lfanew);

    printf("[*] Allocating memory in target process ...\n");

    void *mem = VirtualAllocEx(hProcess, NULL, pNth->OptionalHeader.SizeOfImage, MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE);

    if (!mem)
    {
        printf("[!] Unable to allocate memory. Error (%u)\n", GetLastError());
        CloseHandle(hProcess);
        return -1;
    }

    printf("[*] Memory allocated at 0x%x\n", (unsigned int)mem);

    void *Buffer = VirtualAlloc(NULL, pNth->OptionalHeader.SizeOfImage, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
    memcpy(Buffer, imageBase, pNth->OptionalHeader.SizeOfImage);

    printf("[*] Performing relocation ...\n");

    pBase = (PIMAGE_BASE_RELOCATION)((PUCHAR)Buffer + pNth->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress);
    ULONG difference = (ULONG)mem - (ULONG)imageBase;

    while (pBase->VirtualAddress)
    {
        if (pBase->SizeOfBlock >= sizeof(IMAGE_BASE_RELOCATION))
        {
            ULONG numItems = (pBase->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION)) / sizeof(WORD);
            relocItems = (WORD *)(pBase + 1);

            for (i = 0; i<numItems; i++)
            {
                if (relocItems[i])
                {
                    p = (PULONG)((PUCHAR)Buffer + pBase->VirtualAddress + (relocItems[i] & 0xFFF));
                    *p += difference;
                }
            }
        }

        pBase = (PIMAGE_BASE_RELOCATION)((PUCHAR)pBase + pBase->SizeOfBlock);
    }


    printf("[*] Writing relocated image into target process ...\n");
    if (!WriteProcessMemory(hProcess, mem, Buffer, pNth->OptionalHeader.SizeOfImage, NULL))
    {
        printf("[!] Unable to write to process memory. Error (%u)\n", GetLastError());
        VirtualFreeEx(hProcess, mem, 0, MEM_RELEASE); // size of memory should be 0 when using MEM_RELEASE
        CloseHandle(hProcess);
        return -1;
    }

    VirtualFree(Buffer, 0, MEM_RELEASE);

    printf("[*] Executing thread in remote process at : 0x%08x\n", (difference + (PUCHAR)loadDLL));

    char *dll = "C:\\Users\\root_sec\\Documents\\Visual Studio 2015\\Projects\\ManualMapping\\Debug\\TestFinal.dll";
    HANDLE hThread = CreateRemoteThread(hProcess, NULL, 0, ((LPTHREAD_START_ROUTINE)(difference + (PUCHAR)loadDLL)), dll, 0, NULL);

    if (hThread == INVALID_HANDLE_VALUE)
    {
        printf("[!] Could not create a remote thread.\n");
        VirtualFreeEx(hProcess, mem, 0, MEM_RELEASE);
        CloseHandle(hProcess);
        return -1;
    }

    printf("[*] Waiting for thread to terminate ...\n");
    WaitForSingleObject(hThread, INFINITE);

    printf("[*] Thread terminated. Freeing resources ...\n");
    VirtualFreeEx(hProcess, mem, 0, MEM_RELEASE);
    CloseHandle(hProcess);

    return 0;
}