#include <iostream>
#include <windows.h>
#include <imagehlp.h>

#include "Zydis/Zydis.h"

#define BCOMPAT_REGISTRY_KEY L"\\Registry\\Machine\\Software\\BCompat\\Functions"
#define FUNCTION_NAME_SIZE 64
#define PATCH_SECTION_SIZE 512

BYTE *buffer;

void printLastError() {
    DWORD error = GetLastError();
    if (error == ERROR_FILE_NOT_FOUND) {
        printf("File not found\n");
    } else if (error == ERROR_ACCESS_DENIED) {
        printf("Access to file denied\n");
    } else {
        printf("Error: %lu\n", error);
    }
}

int readInt(const DWORD pos) {
    return buffer[pos] | buffer[pos + 1] << 8;
}

void writeInt(const DWORD pos, const int value) {
    buffer[pos] = value & 0xff;
    buffer[pos + 1] = value >> 8;
}

void writeInt(const DWORD pos, LPDWORD offset, const int value) {
    writeInt(pos + *offset, value);
    *offset += 2;
}

DWORD readDWORD(const DWORD pos) {
    return buffer[pos] | buffer[pos + 1] << 8 | buffer[pos + 2] << 16 | buffer[pos + 3] << 24;
}

void writeDWORD(const DWORD pos, const DWORD value) {
    buffer[pos] = value & 0xff;
    buffer[pos + 1] = value >> 8 & 0xff;
    buffer[pos + 2] = value >> 16 & 0xff;
    buffer[pos + 3] = value >> 24 & 0xff;
}

void writeDWORD(const DWORD pos, LPDWORD offset, const DWORD value) {
    writeDWORD(pos + *offset, value);
    *offset += 4;
}

LPSTR readString(DWORD pos) {
    int size = 32;
    char* string = new char[size];
    int i = 0;
    while (i == 0 || string[i - 1] != '\0') {
        if (i >= size) {
            size += 32;
            char* newString = new char[size];
            memcpy(newString, string, size - 32);
            delete[] string;
            string = newString;
        }
        string[i] = buffer[pos + i];
        i++;
    }
    return string;
}

void writeString(const DWORD pos, LPDWORD offset, const char *string) {
    DWORD length = strlen(string);
    for (int i = 0; i <= length; i++) {
        buffer[pos + *offset + i] = string[i];
    }
    *offset += length + 1;
}

void writeUnicodeString(const DWORD pos, LPDWORD offset, const wchar_t *string) {
    DWORD length = wcslen(string) + 1;
    for (int i = 0; i < length; i++) {
        writeInt(pos, offset, string[i]);
    }
}

DWORD getRealAddress(DWORD *virtualAddresses, DWORD *realAddresses, DWORD sectionCount, DWORD address) {
    int section = 0;
    for (int i = 0; i < sectionCount; ++i) {
        if (virtualAddresses[i] <= address) {
            section = i;
        }
    }
    return address - virtualAddresses[section] + realAddresses[section];
}

DWORD readAddress(DWORD *virtualAddresses, DWORD *realAddresses, DWORD sectionCount, DWORD pos) {
    return getRealAddress(virtualAddresses, realAddresses, sectionCount, readDWORD(pos));
}

DWORD getExportedFunction(DWORD *virtualAddresses, DWORD *realAddresses, DWORD sectionCount, DWORD exportAddressTable, DWORD namePointerTable, DWORD ordinalTable, DWORD functionCount, const char *functionName) {
    printf("Locating %s... ", functionName);
    WORD ordinal = 0;
    for (DWORD i = 0; i < functionCount; i++) {
        DWORD namePtr = readAddress(virtualAddresses, realAddresses, sectionCount, namePointerTable + i * 4);
        LPSTR name = readString(namePtr);
        if (strcmp(name, functionName) == 0) {
            ordinal = readInt(ordinalTable + i * 2);
            delete[] name;
            break;
        }
        delete[] name;
    }
    if (ordinal == 0) {
        printf("Not found in export table");
        delete[] buffer;
        exit(1);
    }
    DWORD function = readDWORD(exportAddressTable + ordinal * 4);
    printf("0x%lx\n", function);
    return function;
}

DWORD getWrappedFunction(DWORD fileSize, DWORD wrapperFunction, const char *functionName) {
    printf("Locating %s... ", functionName);
    DWORD function = 0;
    {
        int i = 0;
        ZydisDisassembledInstruction instruction;
        while (ZYAN_SUCCESS(ZydisDisassembleIntel(ZYDIS_MACHINE_MODE_LEGACY_32, wrapperFunction + i, buffer + wrapperFunction + i, fileSize - wrapperFunction - i, &instruction))) {
            if (instruction.info.mnemonic == ZYDIS_MNEMONIC_CALL) {
                if (function != NULL) {
                    printf("Unexpected code structure");
                    delete[] buffer;
                    exit(1);
                }
                ZyanU64 address;
                ZydisCalcAbsoluteAddress(&instruction.info, &instruction.operands[0], wrapperFunction + i, &address);
                function = address;
            } else if (instruction.info.mnemonic == ZYDIS_MNEMONIC_RET) {
                if (function == NULL) {
                    printf("Could not find CALL");
                    delete[] buffer;
                    exit(1);
                }
                break;
            }
            i += instruction.info.length;
        }
    }
    printf("0x%lx\n", function);
    return function;
}

boolean writeFile(const LPTSTR name, const BYTE *buffer, const DWORD fileSize) {
    HANDLE file = CreateFile(name, GENERIC_WRITE, 0, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
    if (file == INVALID_HANDLE_VALUE) {
        printLastError();
        delete[] buffer;
        return false;
    }
    if (!WriteFile(file, buffer, fileSize, NULL, NULL)) {
        printLastError();
        CloseHandle(file);
        delete[] buffer;
        return false;
    }
    CloseHandle(file);
    return true;
}

void writeJumpAddress(const DWORD pos, const DWORD virtualPos, const DWORD target) {
    writeDWORD(pos, target - (virtualPos + 4));
}

void writeJumpAddress(const DWORD pos, const DWORD virtualPos, LPDWORD offset, const DWORD target) {
    writeJumpAddress(pos + *offset, virtualPos + *offset, target);
    *offset += 4;
}

void ensureAlignment(const DWORD pos, LPDWORD offset) {
    DWORD i = (pos + *offset) % 4;
    if (i > 0) {
        *offset += 4 - i;
    }
}

int main() {
    LPTSTR cmdLine = GetCommandLine();
    LPTSTR space = strrchr(cmdLine, ' ');
    if (space == NULL) {
        printf("Error: No Path provided\n");
        return 0;
    }
    LPTSTR originalFile = space + sizeof(TCHAR);
    if (strlen(originalFile) > MAX_PATH - 10 * sizeof(TCHAR)) {
        printf("Error: File name is too long\n");
        return 1;
    }
    printf("Copying original file... ");
    TCHAR patchedFile[MAX_PATH];
    sprintf(patchedFile, "%s.patched", originalFile);
    if (!CopyFile(originalFile, patchedFile, FALSE)) {
        printLastError();
        return 1;
    }
    printf("Done\n");

    printf("Opening file for patching... ");
    HANDLE file = CreateFile(patchedFile, GENERIC_READ, 0, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
    if (file == INVALID_HANDLE_VALUE) {
        printLastError();
        return 1;
    }
    DWORD fileSize = GetFileSize(file, NULL);
    if (fileSize == INVALID_FILE_SIZE) {
        printLastError();
        return 1;
    }
    buffer = new BYTE[fileSize + PATCH_SECTION_SIZE];
    DWORD bytesRead;
    ReadFile(file, buffer, fileSize, &bytesRead, NULL);
    CloseHandle(file);
    if (bytesRead != fileSize) {
        printLastError();
        delete[] buffer;
        return 1;
    }
    fileSize += PATCH_SECTION_SIZE ;
    printf("Done\n");

    printf("Locating PE Header... ");
    if (fileSize < 2 || buffer[0] != 'M' || buffer[1] != 'Z') {
        printf("Not a valid DLL (Invalid MZ Header)");
        delete[] buffer;
        return 1;
    }
    DWORD peHeader = readDWORD(0x3C);
    if (fileSize <= peHeader + 3 || buffer[peHeader] != 'P' || buffer[peHeader + 1] != 'E' || buffer[peHeader + 2] != '\0' || buffer[peHeader + 3] != '\0') {
        printf("Not a valid DLL (Invalid PE Header)");
        delete[] buffer;
        return 1;
    }
    if (fileSize <= peHeader + 0x19 || readInt(peHeader + 0x18) != 0x10b) {
        printf("Unsupported DLL (Not a normal executable)");
        delete[] buffer;
        return 1;
    }
    printf("0x%lx\n", peHeader);

    printf("Reading Base Image location... ");
    DWORD imageBase = readDWORD(peHeader + 0x18 + 28);
    printf("0x%lx\n", imageBase);

    printf("Locating Sections... ");
    DWORD sectionTable = peHeader + 0x18 + 224;
    WORD sectionCount = readInt(peHeader + 6);
    printf("%u\n", sectionCount);
    DWORD virtualAddresses[sectionCount];
    DWORD realAddresses[sectionCount];
    DWORD realSectionSizes[sectionCount];
    DWORD virtualSectionSizes[sectionCount];
    for (DWORD i = 0; i < sectionCount; i++) {
        LPSTR name = readString(sectionTable + i * 40);
        printf("Locating %s Section... ", name);
        delete[] name;
        DWORD virtualAddress = readDWORD(sectionTable + i * 40 + 12);
        DWORD realAddress = readDWORD(sectionTable + i * 40 + 20);
        virtualAddresses[i] = virtualAddress;
        realAddresses[i] = realAddress;
        realSectionSizes[i] = readDWORD(sectionTable + i * 40 + 16);
        virtualSectionSizes[i] = readDWORD(sectionTable + i * 40 + 8);
        printf("0x%lx / %lu\n", virtualAddress, realAddress);
    }

    printf("Locating Export Address Table... ");
    DWORD exportTable = readAddress(virtualAddresses, realAddresses, sectionCount, peHeader + 0x18 + 96);
    printf("0x%lx\n", exportTable);

    printf("Reading Export Address Table... ");
    DWORD functionCount = readDWORD(exportTable + 24);
    DWORD exportAddressTable = readAddress(virtualAddresses, realAddresses, sectionCount, exportTable + 28);
    DWORD namePointerTable = readAddress(virtualAddresses, realAddresses, sectionCount, exportTable + 32);
    DWORD ordinalTable = readAddress(virtualAddresses, realAddresses, sectionCount, exportTable + 36);
    printf("Done\n");

    DWORD ldrGetProcedureAddress = getExportedFunction(virtualAddresses, realAddresses, sectionCount,
        exportAddressTable, namePointerTable, ordinalTable, functionCount, "LdrGetProcedureAddress");

    DWORD ldrpGetProcedureAddress = getWrappedFunction(fileSize, ldrGetProcedureAddress, "LdrpGetProcedureAddress");

    DWORD rtlImageDirectoryEntryToData = getExportedFunction(virtualAddresses, realAddresses, sectionCount,
        exportAddressTable, namePointerTable, ordinalTable, functionCount, "RtlImageDirectoryEntryToData");

    printf("Locating LdrpSnapThunk... ");
    DWORD pSnapThunk = 0;
    {
        int i = 0;
        int found = 0;
        ZydisDisassembledInstruction instruction;
        while (ZYAN_SUCCESS(ZydisDisassembleIntel(ZYDIS_MACHINE_MODE_LEGACY_32, ldrpGetProcedureAddress + i, buffer + ldrpGetProcedureAddress + i, fileSize - ldrpGetProcedureAddress - i, &instruction))) {
            if (instruction.info.mnemonic == ZYDIS_MNEMONIC_CALL) {
                ZyanU64 address;
                ZydisCalcAbsoluteAddress(&instruction.info, &instruction.operands[0], ldrpGetProcedureAddress + i, &address);
                if (address == rtlImageDirectoryEntryToData) {
                    if (found > 0 || pSnapThunk != NULL) {
                        printf("Unexpected code structure");
                        delete[] buffer;
                        return 1;
                    }
                    found = 1;
                } else if (found == 1) { //__local_unwind2
                    found++;
                } else {
                    if (found == 2) {
                        pSnapThunk = address;
                    }
                    found = 0;
                }
            } else if (instruction.info.mnemonic == ZYDIS_MNEMONIC_RET) {
                if (pSnapThunk == NULL) {
                    printf("Could not find CALL");
                    delete[] buffer;
                    return 1;
                }
                break;
            }
            i += instruction.info.length;
        }
    }
    printf("0x%lx\n", pSnapThunk);

    printf("Locating patch point... ");
    DWORD patchPoint = 0;
    {
        int i = 0;
        ZydisDisassembledInstruction instruction;
        while (ZYAN_SUCCESS(ZydisDisassembleIntel(ZYDIS_MACHINE_MODE_LEGACY_32, pSnapThunk + i, buffer + pSnapThunk + i, fileSize - pSnapThunk - i, &instruction))) {
            if (instruction.info.mnemonic == ZYDIS_MNEMONIC_MOV) {
                if (instruction.operands[0].mem.disp.value == -20) { //ImportName = ...
                    patchPoint = instruction.runtime_address + instruction.info.length;
                    break;
                }
            } else if (instruction.info.mnemonic == ZYDIS_MNEMONIC_RET) {
                printf("Unexpected code structure");
                delete[] buffer;
                return 1;
            }
            i += instruction.info.length;
        }
    }
    printf("0x%lx\n", patchPoint);

    printf("Checking how many bytes to replace... ");
    DWORD targetSize = 5;
    DWORD patchSize = 0;
    {
        int i = 0;
        ZydisDisassembledInstruction instruction;
        while (ZYAN_SUCCESS(ZydisDisassembleIntel(ZYDIS_MACHINE_MODE_LEGACY_32, patchPoint + i, buffer + patchPoint + i, fileSize - patchPoint - i, &instruction))) {
            if (instruction.info.mnemonic == ZYDIS_MNEMONIC_RET) {
                printf("Unexpected end of function");
                delete[] buffer;
                return 1;
            }
            i += instruction.info.length;
            if (i >= targetSize) {
                patchSize = i;
                break;
            }
        }
    }
    printf("%lu\n", patchSize);

    printf("Locating LdrpSnapThunk end... ");
    DWORD ldrpSnapThunkEnd = 0;
    {
        int i = 0;
        ZydisDisassembledInstruction instruction;
        while (ZYAN_SUCCESS(ZydisDisassembleIntel(ZYDIS_MACHINE_MODE_LEGACY_32, pSnapThunk + i, buffer + pSnapThunk + i, fileSize - pSnapThunk - i, &instruction))) {
            if (instruction.info.mnemonic == ZYDIS_MNEMONIC_POP) {
                if (instruction.operands[0].reg.value == ZYDIS_REGISTER_EDI) { //POP EDI
                    if (ldrpSnapThunkEnd != NULL) {
                        printf("Unexpected code structure");
                        delete[] buffer;
                        return 1;
                    }
                    ldrpSnapThunkEnd = instruction.runtime_address;
                }
            } else if (instruction.info.mnemonic == ZYDIS_MNEMONIC_RET) {
                if (ldrpSnapThunkEnd == NULL) {
                    printf("Unexpected end of function");
                    delete[] buffer;
                    return 1;
                }
                break;
            }
            i += instruction.info.length;
        }
    }
    printf("0x%lx\n", ldrpSnapThunkEnd);

    printf("Adding .patch section... ");
    DWORD freeSpace = realAddresses[sectionCount - 1] + realSectionSizes[sectionCount - 1];
    DWORD virtualFreeSpace = virtualAddresses[sectionCount - 1] + virtualSectionSizes[sectionCount - 1];
    virtualFreeSpace = virtualFreeSpace + 0x1000 - virtualFreeSpace % 0x1000;
    {
        writeInt(peHeader + 6, sectionCount  + 1);
        writeDWORD(peHeader + 0x18 + 56, readDWORD(peHeader + 0x18 + 56) + PATCH_SECTION_SIZE);
        DWORD patchSectionHeader = sectionTable + sectionCount * 40;
        DWORD i = 0;
        writeString(patchSectionHeader, &i, ".patch");
        for (int x = 0; x < 7 - strlen(".patch"); x++) {
            buffer[patchSectionHeader + i++] = 0x00;
        }
        writeDWORD(patchSectionHeader, &i, PATCH_SECTION_SIZE);
        writeDWORD(patchSectionHeader, &i, virtualFreeSpace);
        writeDWORD(patchSectionHeader, &i, PATCH_SECTION_SIZE);
        writeDWORD(patchSectionHeader, &i, freeSpace);
        writeDWORD(patchSectionHeader, &i, 0);
        writeDWORD(patchSectionHeader, &i, 0);
        writeInt(patchSectionHeader, &i, 0);
        writeInt(patchSectionHeader, &i, 0);
        writeDWORD(patchSectionHeader, &i, IMAGE_SCN_CNT_CODE | IMAGE_SCN_MEM_EXECUTE | IMAGE_SCN_MEM_READ);
    }
    printf("0x%lx\n", freeSpace);

    DWORD strcmp = getExportedFunction(virtualAddresses, realAddresses, sectionCount,
        exportAddressTable, namePointerTable, ordinalTable, functionCount, "strcmp");

    DWORD ntOpenKey = getExportedFunction(virtualAddresses, realAddresses, sectionCount,
        exportAddressTable, namePointerTable, ordinalTable, functionCount, "NtOpenKey");

    DWORD strcpy = getExportedFunction(virtualAddresses, realAddresses, sectionCount,
        exportAddressTable, namePointerTable, ordinalTable, functionCount, "strcpy");

    DWORD strcat = getExportedFunction(virtualAddresses, realAddresses, sectionCount,
        exportAddressTable, namePointerTable, ordinalTable, functionCount, "strcat");

    DWORD rtlInitAnsiString = getExportedFunction(virtualAddresses, realAddresses, sectionCount,
        exportAddressTable, namePointerTable, ordinalTable, functionCount, "RtlInitAnsiString");

    DWORD rtlAnsiStringToUnicodeString = getExportedFunction(virtualAddresses, realAddresses, sectionCount,
        exportAddressTable, namePointerTable, ordinalTable, functionCount, "RtlAnsiStringToUnicodeString");

    DWORD ntQueryValueKey = getExportedFunction(virtualAddresses, realAddresses, sectionCount,
        exportAddressTable, namePointerTable, ordinalTable, functionCount, "NtQueryValueKey");

    DWORD ldrLoadDll = getExportedFunction(virtualAddresses, realAddresses, sectionCount,
        exportAddressTable, namePointerTable, ordinalTable, functionCount, "LdrLoadDll");

    DWORD ldrpLoadDll = getWrappedFunction(fileSize, ldrLoadDll, "LdrpLoadDll");

    DWORD rtlInitUnicodeString = getExportedFunction(virtualAddresses, realAddresses, sectionCount,
        exportAddressTable, namePointerTable, ordinalTable, functionCount, "RtlInitUnicodeString");

    printf("Locating LdrpWalkImportDescriptor... ");
    DWORD ldrpWalkImportDescriptor = 0;
    {
        int i = 0;
        int found = 0;
        ZydisDisassembledInstruction instruction;
        while (ZYAN_SUCCESS(ZydisDisassembleIntel(ZYDIS_MACHINE_MODE_LEGACY_32, ldrpLoadDll + i, buffer + ldrpLoadDll + i, fileSize - ldrpLoadDll - i, &instruction))) {
            if (instruction.info.mnemonic == ZYDIS_MNEMONIC_CALL) {
                ZyanU64 address;
                ZydisCalcAbsoluteAddress(&instruction.info, &instruction.operands[0], ldrpLoadDll + i, &address);
                if (address == rtlInitUnicodeString) {
                    if (found > 0 || ldrpWalkImportDescriptor != NULL) {
                        printf("Unexpected code structure");
                        delete[] buffer;
                        return 1;
                    }
                    found = 1;
                } else if (found == 1) { //LdrpCheckForLoadedDll
                    found++;
                } else if (found == 2) { //LdrpMapDll
                    found++;
                } else if (found == 3) { //__local_unwind2
                    found++;
                } else {
                    if (found == 4) {
                        ldrpWalkImportDescriptor = address;
                    }
                    found = 0;
                }
            } else if (instruction.info.mnemonic == ZYDIS_MNEMONIC_RET) {
                if (ldrpWalkImportDescriptor == NULL) {
                    printf("Could not find CALL");
                    delete[] buffer;
                    return 1;
                }
                break;
            }
            i += instruction.info.length;
        }
    }
    printf("0x%lx\n", ldrpWalkImportDescriptor);

    /*printf("Disabling bound imports... ");
    DWORD boundImportsPatch = 0;
    {
        int i = 0;
        ZydisDisassembledInstruction instruction;
        while (ZYAN_SUCCESS(ZydisDisassembleIntel(ZYDIS_MACHINE_MODE_LEGACY_32, ldrpWalkImportDescriptor + i, buffer + ldrpWalkImportDescriptor + i, fileSize - ldrpWalkImportDescriptor - i, &instruction))) {
            if (instruction.info.mnemonic == ZYDIS_MNEMONIC_MOV) {
                if (instruction.operands[0].mem.base == ZYDIS_REGISTER_EBP &&
                    instruction.operands[1].type == ZYDIS_OPERAND_TYPE_IMMEDIATE && instruction.operands[1].imm.value.u == 0) { //stale = false
                    if (boundImportsPatch != NULL) {
                        printf("Unexpected code structure");
                        delete[] buffer;
                        return 1;
                    }
                    boundImportsPatch = ldrpWalkImportDescriptor + i;
                }
            } else if (instruction.info.mnemonic == ZYDIS_MNEMONIC_RET) {
                if (boundImportsPatch == NULL) {
                    printf("Could not find patch point");
                    delete[] buffer;
                    return 1;
                }
                break;
            }
            i += instruction.info.length;
        }
    }
    buffer[boundImportsPatch + 3] = 0x01;
    printf("0x%lx\n", boundImportsPatch);*/
    /*printf("Disabling bound imports... ");
    DWORD boundImportsPatch = 0;
    {
        int i = 0;
        ZydisDisassembledInstruction instruction;
        while (ZYAN_SUCCESS(ZydisDisassembleIntel(ZYDIS_MACHINE_MODE_LEGACY_32, ldrpWalkImportDescriptor + i, buffer + ldrpWalkImportDescriptor + i, fileSize - ldrpWalkImportDescriptor - i, &instruction))) {
            if (instruction.info.mnemonic == ZYDIS_MNEMONIC_CMP) {
                if (instruction.operands[0].mem.base == ZYDIS_REGISTER_EBP && instruction.operands[0].mem.disp.value == -2 &&
                    instruction.operands[1].type == ZYDIS_OPERAND_TYPE_IMMEDIATE && instruction.operands[1].imm.value.u == 0) {//stale = false
                    if (boundImportsPatch != NULL) {
                        printf("Unexpected code structure");
                        delete[] buffer;
                        return 1;
                    }
                    boundImportsPatch = ldrpWalkImportDescriptor + i;
                }
            } else if (instruction.info.mnemonic == ZYDIS_MNEMONIC_RET) {
                if (boundImportsPatch == NULL) {
                    printf("Could not find patch point");
                    delete[] buffer;
                    return 1;
                }
                break;
            }
            i += instruction.info.length;
        }
    }
    buffer[boundImportsPatch + 6] = 0x90;
    buffer[boundImportsPatch + 7] = 0x90;
    buffer[boundImportsPatch + 8] = 0x90;
    buffer[boundImportsPatch + 9] = 0x90;
    buffer[boundImportsPatch + 10] = 0x90;
    buffer[boundImportsPatch + 11] = 0x90;
    printf("0x%lx\n", boundImportsPatch);*/

    DWORD dbgBreakPoint = getExportedFunction(virtualAddresses, realAddresses, sectionCount,
        exportAddressTable, namePointerTable, ordinalTable, functionCount, "DbgBreakPoint");

    printf("Locating LdrpLoadImportModule... ");
    DWORD ldrpLoadImportModule = 0;
    {
        int i = 0;
        bool found = false;
        ZydisDisassembledInstruction instruction;
        while (ZYAN_SUCCESS(ZydisDisassembleIntel(ZYDIS_MACHINE_MODE_LEGACY_32, ldrpWalkImportDescriptor + i, buffer + ldrpWalkImportDescriptor + i, fileSize - ldrpWalkImportDescriptor - i, &instruction))) {
            if (instruction.info.mnemonic == ZYDIS_MNEMONIC_CALL) {
                ZyanU64 address;
                ZydisCalcAbsoluteAddress(&instruction.info, &instruction.operands[0], ldrpWalkImportDescriptor + i, &address);
                if (address == dbgBreakPoint) {
                    if (found || ldrpLoadImportModule != NULL) {
                        printf("Unexpected code structure");
                        delete[] buffer;
                        return 1;
                    }
                    found = true;
                } else if (found) {
                    ldrpLoadImportModule = address;
                    found = false;
                }
            } else if (instruction.info.mnemonic == ZYDIS_MNEMONIC_RET) {
                if (ldrpLoadImportModule == NULL) {
                    printf("Could not find CALL");
                    delete[] buffer;
                    return 1;
                }
                break;
            }
            i += instruction.info.length;
        }
    }
    printf("0x%lx\n", ldrpLoadImportModule);

    DWORD ntCurrentTeb = getExportedFunction(virtualAddresses, realAddresses, sectionCount,
        exportAddressTable, namePointerTable, ordinalTable, functionCount, "NtCurrentTeb");

    printf("Patching free space... ");
    DWORD patchEntrypointAddress;
    {
        DWORD i = 0;
        //ADD ESP, 0x4
        buffer[freeSpace + i++] = 0x83;
        buffer[freeSpace + i++] = 0xc4;
        buffer[freeSpace + i++] = 0x04;
        //ADD ESP, 0x4
        buffer[freeSpace + i++] = 0x83;
        buffer[freeSpace + i++] = 0xc4;
        buffer[freeSpace + i++] = 0x04;
        DWORD patchEnd = virtualFreeSpace + i;
        //POP EDI
        buffer[freeSpace + i++] = 0x5f;
        //POP ESI
        buffer[freeSpace + i++] = 0x5e;
        //POP EAX
        buffer[freeSpace + i++] = 0x58;
        for (int x = 0; x < patchSize; x++) {
            buffer[freeSpace + i++] = buffer[patchPoint + x];
        }
        //JMP
        buffer[freeSpace + i++] = 0xe9;
        writeJumpAddress(freeSpace, virtualFreeSpace,  &i, patchPoint + patchSize);
        patchEntrypointAddress = virtualFreeSpace + i;

        //PUSH EAX
        buffer[freeSpace + i++] = 0x50;
        //PUSH ESI
        buffer[freeSpace + i++] = 0x56;
        //PUSH EDI
        buffer[freeSpace + i++] = 0x57;
        //MOV EAX, DllName
        buffer[freeSpace + i++] = 0x8b;
        buffer[freeSpace + i++] = 0x45;
        buffer[freeSpace + i++] = 0x24;
        //TEST EAX EAX
        buffer[freeSpace + i++] = 0x85;
        buffer[freeSpace + i++] = 0xc0;
        //JZ PatchEnd
        buffer[freeSpace + i++] = 0x0f;
        buffer[freeSpace + i++] = 0x84;
        writeJumpAddress(freeSpace, virtualFreeSpace,  &i, patchEnd);
        //JMP over "bcompat.dll"
        buffer[freeSpace + i++] = 0xe9;
        DWORD jumpPos = freeSpace + i;
        i += 4;
        DWORD dllString = virtualFreeSpace + i;
        writeString(freeSpace, &i, "bcompat.dll");
        writeJumpAddress(jumpPos, jumpPos, freeSpace + i);
        //PUSH "bcompat.dll"
        buffer[freeSpace + i++] = 0x68;
        writeDWORD(freeSpace, &i, imageBase + dllString);
        //PUSH EAX
        buffer[freeSpace + i++] = 0x50;
        //CALL strcmp
        buffer[freeSpace + i++] = 0xe8;
        writeJumpAddress(freeSpace, virtualFreeSpace, &i, strcmp);
        //ADD ESP, 8
        buffer[freeSpace + i++] = 0x83;
        buffer[freeSpace + i++] = 0xc4;
        buffer[freeSpace + i++] = 0x08;
        //TEST EAX EAX
        buffer[freeSpace + i++] = 0x85;
        buffer[freeSpace + i++] = 0xc0;
        //JZ PatchEnd
        buffer[freeSpace + i++] = 0x0f;
        buffer[freeSpace + i++] = 0x84;
        writeJumpAddress(freeSpace, virtualFreeSpace, &i, patchEnd);
        //JMP over
        buffer[freeSpace + i++] = 0xe9;
        jumpPos = freeSpace + i;
        i += 4;
        ensureAlignment(freeSpace, &i);
        DWORD keyName = virtualFreeSpace + i;
        writeUnicodeString(freeSpace, &i, BCOMPAT_REGISTRY_KEY);
        //UNICODE_STRING
        ensureAlignment(freeSpace, &i);
        DWORD keyString = virtualFreeSpace + i;
        buffer[freeSpace + i++] = wcslen(BCOMPAT_REGISTRY_KEY) * sizeof(wchar_t);       //Length
        buffer[freeSpace + i++] = 0x00;
        buffer[freeSpace + i++] = (wcslen(BCOMPAT_REGISTRY_KEY) + 1) * sizeof(wchar_t); //MaximumLength
        buffer[freeSpace + i++] = 0x00;
        writeDWORD(freeSpace, &i, imageBase + keyName); //Buffer
        //OBJECT_ATTRIBUTES
        DWORD objectAttributes = virtualFreeSpace + i;
        writeDWORD(freeSpace, &i, 0x18);      //Length
        writeDWORD(freeSpace, &i, 0x00);      //RootDirectory
        writeDWORD(freeSpace, &i, imageBase + keyString); //ObjectName
        writeDWORD(freeSpace, &i, 0x40);      //Attributes
        writeDWORD(freeSpace, &i, 0x00);      //SecurityDescriptor
        writeDWORD(freeSpace, &i, 0x00);      //SecurityQualityOfService
        writeJumpAddress(jumpPos, jumpPos, freeSpace + i);
        //SUB ESP, 4
        buffer[freeSpace + i++] = 0x83;
        buffer[freeSpace + i++] = 0xec;
        buffer[freeSpace + i++] = 0x04;
        //LEA EAX, [ESP]
        buffer[freeSpace + i++] = 0x8d;
        buffer[freeSpace + i++] = 0x04;
        buffer[freeSpace + i++] = 0x24;
        //PUSH POBJECT_ATTRIBUTES
        buffer[freeSpace + i++] = 0x68;
        writeDWORD(freeSpace, &i, imageBase + objectAttributes);
        //PUSH 0x20019
        buffer[freeSpace + i++] = 0x68;
        writeDWORD(freeSpace, &i, 0x20019);
        //PUSH EAX
        buffer[freeSpace + i++] = 0x50;
        //CALL NtOpenKey
        buffer[freeSpace + i++] = 0xe8;
        writeJumpAddress(freeSpace, virtualFreeSpace, &i, ntOpenKey);
        //TEST EAX EAX
        buffer[freeSpace + i++] = 0x85;
        buffer[freeSpace + i++] = 0xc0;
        //JS PatchEnd
        buffer[freeSpace + i++] = 0x0f;
        buffer[freeSpace + i++] = 0x88;
        writeJumpAddress(freeSpace, virtualFreeSpace, &i, patchEnd - 3);
        //SUB ESP, FUNCTION_NAME_SIZE
        buffer[freeSpace + i++] = 0x83;
        buffer[freeSpace + i++] = 0xec;
        buffer[freeSpace + i++] = FUNCTION_NAME_SIZE;
        //LEA ESI, [ESP]
        buffer[freeSpace + i++] = 0x8d;
        buffer[freeSpace + i++] = 0x34;
        buffer[freeSpace + i++] = 0x24;
        //MOV EAX, [DllName]
        buffer[freeSpace + i++] = 0x8b;
        buffer[freeSpace + i++] = 0x45;
        buffer[freeSpace + i++] = 0x24;
        //PUSH EAX
        buffer[freeSpace + i++] = 0x50;
        //PUSH ESI
        buffer[freeSpace + i++] = 0x56;
        //CALL strcpy
        buffer[freeSpace + i++] = 0xe8;
        writeJumpAddress(freeSpace, virtualFreeSpace, &i, strcpy);
        //ADD ESP, 8
        buffer[freeSpace + i++] = 0x83;
        buffer[freeSpace + i++] = 0xc4;
        buffer[freeSpace + i++] = 0x08;
        //MOV EAX, [ImportName]
        buffer[freeSpace + i++] = 0x8b;
        buffer[freeSpace + i++] = 0x45;
        buffer[freeSpace + i++] = 0xec;
        //PUSH EAX
        buffer[freeSpace + i++] = 0x50;
        //PUSH ESI
        buffer[freeSpace + i++] = 0x56;
        //CALL strcat
        buffer[freeSpace + i++] = 0xe8;
        writeJumpAddress(freeSpace, virtualFreeSpace, &i, strcat);
        //ADD ESP, 8
        buffer[freeSpace + i++] = 0x83;
        buffer[freeSpace + i++] = 0xc4;
        buffer[freeSpace + i++] = 0x08;
        //SUB ESP, 8
        buffer[freeSpace + i++] = 0x83;
        buffer[freeSpace + i++] = 0xec;
        buffer[freeSpace + i++] = 0x08;
        //LEA EDI, [ESP]
        buffer[freeSpace + i++] = 0x8d;
        buffer[freeSpace + i++] = 0x3c;
        buffer[freeSpace + i++] = 0x24;
        //PUSH ESI
        buffer[freeSpace + i++] = 0x56;
        //PUSH EDI
        buffer[freeSpace + i++] = 0x57;
        //CALL RtlInitAnsiString
        buffer[freeSpace + i++] = 0xe8;
        writeJumpAddress(freeSpace, virtualFreeSpace, &i, rtlInitAnsiString);
        //SUB ESP, 8
        buffer[freeSpace + i++] = 0x83;
        buffer[freeSpace + i++] = 0xec;
        buffer[freeSpace + i++] = 0x08;
        //LEA EAX, [ESP]
        buffer[freeSpace + i++] = 0x8d;
        buffer[freeSpace + i++] = 0x04;
        buffer[freeSpace + i++] = 0x24;
        //PUSH 1
        buffer[freeSpace + i++] = 0x6a;
        buffer[freeSpace + i++] = 0x01;
        //PUSH EDI
        buffer[freeSpace + i++] = 0x57;
        //PUSH EAX
        buffer[freeSpace + i++] = 0x50;
        //CALL RtlAnsiStringToUnicodeString
        buffer[freeSpace + i++] = 0xe8;
        writeJumpAddress(freeSpace, virtualFreeSpace, &i, rtlAnsiStringToUnicodeString);
        //MOV EDI, [ESP+0x50]
        buffer[freeSpace + i++] = 0x8b;
        buffer[freeSpace + i++] = 0x7c;
        buffer[freeSpace + i++] = 0x24;
        buffer[freeSpace + i++] = 0x50;
        //SUB ESP, 4
        buffer[freeSpace + i++] = 0x83;
        buffer[freeSpace + i++] = 0xec;
        buffer[freeSpace + i++] = 0x04;
        //LEA EAX, [ESP]
        buffer[freeSpace + i++] = 0x8d;
        buffer[freeSpace + i++] = 0x04;
        buffer[freeSpace + i++] = 0x24;
        //PUSH EAX
        buffer[freeSpace + i++] = 0x50;
        //PUSH 0
        buffer[freeSpace + i++] = 0x6a;
        buffer[freeSpace + i++] = 0x00;
        //PUSH 0
        buffer[freeSpace + i++] = 0x6a;
        buffer[freeSpace + i++] = 0x00;
        //PUSH 0
        buffer[freeSpace + i++] = 0x6a;
        buffer[freeSpace + i++] = 0x00;
        //LEA EAX, [ESP+0x14]
        buffer[freeSpace + i++] = 0x8d;
        buffer[freeSpace + i++] = 0x44;
        buffer[freeSpace + i++] = 0x24;
        buffer[freeSpace + i++] = 0x14;
        //PUSH EAX
        buffer[freeSpace + i++] = 0x50;
        //PUSH EDI
        buffer[freeSpace + i++] = 0x57;
        //CALL NtQueryValueKey
        buffer[freeSpace + i++] = 0xe8;
        writeJumpAddress(freeSpace, virtualFreeSpace, &i, ntQueryValueKey);
        //ADD ESP, 0x58
        buffer[freeSpace + i++] = 0x83;
        buffer[freeSpace + i++] = 0xc4;
        buffer[freeSpace + i++] = FUNCTION_NAME_SIZE + 0x18;
        //CMP EAX, 0xc0000023
        buffer[freeSpace + i++] = 0x3d;
        buffer[freeSpace + i++] = 0x23;
        buffer[freeSpace + i++] = 0x00;
        buffer[freeSpace + i++] = 0x00;
        buffer[freeSpace + i++] = 0xc0;
        //JNZ PatchEnd
        buffer[freeSpace + i++] = 0x0f;
        buffer[freeSpace + i++] = 0x85;
        writeJumpAddress(freeSpace, virtualFreeSpace, &i, patchEnd);
        //SUB ESP, 0x8
        buffer[freeSpace + i++] = 0x83;
        buffer[freeSpace + i++] = 0xec;
        buffer[freeSpace + i++] = 0x08;
        //LEA EDI, [ESP]
        buffer[freeSpace + i++] = 0x8d;
        buffer[freeSpace + i++] = 0x3c;
        buffer[freeSpace + i++] = 0x24;
        //LEA EAX, [ESP+0x4]
        buffer[freeSpace + i++] = 0x8d;
        buffer[freeSpace + i++] = 0x44;
        buffer[freeSpace + i++] = 0x24;
        buffer[freeSpace + i++] = 0x04;
        //PUSH EAX
        buffer[freeSpace + i++] = 0x50;
        //PUSH EDI
        buffer[freeSpace + i++] = 0x57;
        //MOV EAX, ImportBase
        buffer[freeSpace + i++] = 0x8b;
        buffer[freeSpace + i++] = 0x45;
        buffer[freeSpace + i++] = 0xc;
        //PUSH EAX
        buffer[freeSpace + i++] = 0x50;
        //PUSH "bcompat.dll"
        buffer[freeSpace + i++] = 0x68;
        writeDWORD(freeSpace, &i, imageBase + dllString);
        //PUSH 0
        buffer[freeSpace + i++] = 0x6a;
        buffer[freeSpace + i++] = 0x00;
        //CALL LdrpLoadImportModule
        buffer[freeSpace + i++] = 0xe8;
        writeJumpAddress(freeSpace, virtualFreeSpace, &i, ldrpLoadImportModule);
        //MOV EDI, [ESP]
        buffer[freeSpace + i++] = 0x8b;
        buffer[freeSpace + i++] = 0x3c;
        buffer[freeSpace + i++] = 0x24;
        //TEST EAX EAX
        buffer[freeSpace + i++] = 0x85;
        buffer[freeSpace + i++] = 0xc0;
        //JS PatchEnd
        buffer[freeSpace + i++] = 0x0f;
        buffer[freeSpace + i++] = 0x88;
        writeJumpAddress(freeSpace, virtualFreeSpace, &i, patchEnd - 6);
        //CMP byte ptr [ESP+0x4], 0
        buffer[freeSpace + i++] = 0x80;
        buffer[freeSpace + i++] = 0x7c;
        buffer[freeSpace + i++] = 0x24;
        buffer[freeSpace + i++] = 0x04;
        buffer[freeSpace + i++] = 0x00;
        //ADD ESP, 0x8
        buffer[freeSpace + i++] = 0x83;
        buffer[freeSpace + i++] = 0xc4;
        buffer[freeSpace + i++] = 0x08;
        //ADD EDI, 0x10
        buffer[freeSpace + i++] = 0x83;
        buffer[freeSpace + i++] = 0xc7;
        buffer[freeSpace + i++] = 0x10;
        //JZ SKIP_INIT
        buffer[freeSpace + i++] = 0x0f;
        buffer[freeSpace + i++] = 0x84;
        jumpPos = freeSpace + i;
        i += 4;
        //CALL NtCurrentTeb
        buffer[freeSpace + i++] = 0xe8;
        writeJumpAddress(freeSpace, virtualFreeSpace, &i, ntCurrentTeb);
        //MOV ESI, [EAX + 0x30] //PEB
        buffer[freeSpace + i++] = 0x8b;
        buffer[freeSpace + i++] = 0x70;
        buffer[freeSpace + i++] = 0x30;
        //MOV ESI, [ESI + 0xc] //Ldr
        buffer[freeSpace + i++] = 0x8b;
        buffer[freeSpace + i++] = 0x76;
        buffer[freeSpace + i++] = 0x0c;
        //ADD ESI, 0x1c //InInitializationOrderModuleList
        buffer[freeSpace + i++] = 0x83;
        buffer[freeSpace + i++] = 0xc6;
        buffer[freeSpace + i++] = 0x1c;
        //MOV EAX, [ESI + 0x4]
        buffer[freeSpace + i++] = 0x8b;
        buffer[freeSpace + i++] = 0x46;
        buffer[freeSpace + i++] = 0x04;
        //MOV [EDI], ESI
        buffer[freeSpace + i++] = 0x89;
        buffer[freeSpace + i++] = 0x37;
        //MOV [EDI + 0x4], EAX
        buffer[freeSpace + i++] = 0x89;
        buffer[freeSpace + i++] = 0x47;
        buffer[freeSpace + i++] = 0x04;
        //MOV [EAX], EDI
        buffer[freeSpace + i++] = 0x89;
        buffer[freeSpace + i++] = 0x38;
        //MOV [ESI + 0x4], EDI
        buffer[freeSpace + i++] = 0x89;
        buffer[freeSpace + i++] = 0x7e;
        buffer[freeSpace + i++] = 0x04;
        writeJumpAddress(jumpPos, jumpPos, freeSpace + i);
        //SUB ESP, 0x4
        buffer[freeSpace + i++] = 0x83;
        buffer[freeSpace + i++] = 0xec;
        buffer[freeSpace + i++] = 0x04;
        //LEA EAX, [ESP]
        buffer[freeSpace + i++] = 0x8d;
        buffer[freeSpace + i++] = 0x04;
        buffer[freeSpace + i++] = 0x24;
        //PUSH EAX
        buffer[freeSpace + i++] = 0x50;
        //PUSH 0
        buffer[freeSpace + i++] = 0x6a;
        buffer[freeSpace + i++] = 0x00;
        //PUSH 1
        buffer[freeSpace + i++] = 0x6a;
        buffer[freeSpace + i++] = 0x01;
        //ADD EDI, 0x8 //DllBase
        buffer[freeSpace + i++] = 0x83;
        buffer[freeSpace + i++] = 0xc7;
        buffer[freeSpace + i++] = 0x08;
        //mov EDI, [EDI]
        buffer[freeSpace + i++] = 0x8b;
        buffer[freeSpace + i++] = 0x3f;
        //PUSH EDI
        buffer[freeSpace + i++] = 0x57;
        //CALL RtlImageDirectoryEntryToData
        buffer[freeSpace + i++] = 0xe8;
        writeJumpAddress(freeSpace, virtualFreeSpace, &i, rtlImageDirectoryEntryToData);
        //MOV ESI, [ESP]
        buffer[freeSpace + i++] = 0x8b;
        buffer[freeSpace + i++] = 0x34;
        buffer[freeSpace + i++] = 0x24;
        //PUSH "bcompat.dll"
        buffer[freeSpace + i++] = 0x68;
        writeDWORD(freeSpace, &i, imageBase + dllString);
        //PUSH StaticSnap
        buffer[freeSpace + i++] = 0xff;
        buffer[freeSpace + i++] = 0x75;
        buffer[freeSpace + i++] = 0x20;
        //PUSH ESI
        buffer[freeSpace + i++] = 0x56,
        //PUSH EAX
        buffer[freeSpace + i++] = 0x50;
        //PUSH Thunk
        buffer[freeSpace + i++] = 0xff;
        buffer[freeSpace + i++] = 0x75;
        buffer[freeSpace + i++] = 0x14;
        //PUSH OriginalThunk
        buffer[freeSpace + i++] = 0xff;
        buffer[freeSpace + i++] = 0x75;
        buffer[freeSpace + i++] = 0x10;
        //PUSH ImageBase
        buffer[freeSpace + i++] = 0xff;
        buffer[freeSpace + i++] = 0x75;
        buffer[freeSpace + i++] = 0x0c;
        //PUSH EDI
        buffer[freeSpace + i++] = 0x57;
        //CALL LdrpSnapThunk
        buffer[freeSpace + i++] = 0xe8;
        writeJumpAddress(freeSpace, virtualFreeSpace, &i, pSnapThunk);
        //buffer[freeSpace + i++] = 0xcc;
        //ADD ESP, 0x4
        buffer[freeSpace + i++] = 0x83;
        buffer[freeSpace + i++] = 0xc4;
        buffer[freeSpace + i++] = 0x04;
        //POP EDI
        buffer[freeSpace + i++] = 0x5f;
        //POP ESI
        buffer[freeSpace + i++] = 0x5e;
        //POP EAX
        buffer[freeSpace + i++] = 0x58;
        //buffer[freeSpace + i++] = 0xcc;
        //JMP LdrpSnapThunk - End
        buffer[freeSpace + i++] = 0xe9;
        writeJumpAddress(freeSpace, virtualFreeSpace,  &i, ldrpSnapThunkEnd);
        printf("Done (%lu Bytes left)\n", PATCH_SECTION_SIZE - i);
    }

    printf("Patching LdrpSnapThunk... ");
    buffer[patchPoint] = 0xe9;
    writeJumpAddress(patchPoint + 1, patchPoint + 1, patchEntrypointAddress);
    for (int i = 5; i < patchSize; i++) {
        buffer[patchPoint + i] = 0x90;
    }
    printf("Done\n");

    printf("Writing to patched file... ");
    if (!writeFile(patchedFile, buffer, fileSize)) {
        return 1;
    }
    printf("Done\n");

    printf("Generating checksum... ");
    DWORD originalChecksum;
    DWORD checksum;
    if (MapFileAndCheckSum(patchedFile, &originalChecksum, &checksum) != CHECKSUM_SUCCESS) {
        printLastError();
        return 1;
    }
    printf("0x%lx\n", checksum);

    printf("Writing checksum to patched file... ");
    writeDWORD(peHeader + 0x18 + 64, checksum);
    if (!writeFile(patchedFile, buffer, fileSize)) {
        return 1;
    }
    printf("Finished\n");

    delete[] buffer;
    return 0;
}