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

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

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

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

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

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

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];
    DWORD bytesRead;
    ReadFile(file, buffer, fileSize, &bytesRead, NULL);
    CloseHandle(file);
    if (bytesRead != fileSize) {
        printLastError();
        delete[] buffer;
        return 1;
    }
    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("Locating Sections... ");
    DWORD sectionTable = peHeader + 0x18 + 224;
    WORD sectionCount = readInt(peHeader + 6);
    printf("%u\n", sectionCount);
    DWORD virtualAddresses[sectionCount];
    DWORD realAddresses[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;
        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 exportsCount = readDWORD(exportTable + 20);
    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);
    if (exportsCount != functionCount) {
        printf("DLL already contains ordinal exports");
        delete[] buffer;
        return 1;
    }
    printf("Done\n");

    printf("Patching Ordinals... ");
    DWORD ordinalBase = -1;
    DWORD ordinalEnd = 0;
    DWORD exportAddresses[exportsCount];
    for (int i = 0; i < functionCount; ++i) {
        DWORD namePtr = readAddress(virtualAddresses, realAddresses, sectionCount, namePointerTable + i * 4);
        LPSTR name = readString(namePtr);
        if (strstr(name, "ordinal") == name) {
            DWORD ordinal = std::stoi(name + 7);
            if (ordinal < ordinalBase) {
                ordinalBase = ordinal;
            }
            if (ordinal > ordinalEnd) {
                ordinalEnd = ordinal;
            }
        }
        exportAddresses[i] = readDWORD(exportAddressTable + i * 4);
        writeDWORD(exportAddressTable + i * 4, NULL);
        delete[] name;
    }

    for (int i = 0; i < functionCount; ++i) {
        DWORD namePtr = readAddress(virtualAddresses, realAddresses, sectionCount, namePointerTable + i * 4);
        LPSTR name = readString(namePtr);
        DWORD realOrdinal;
        if (strstr(name, "ordinal") == name) {
            DWORD ordinal = std::stoi(name + 7);

            DWORD stringAddr = /*readAddress(virtualAddresses, realAddresses, sectionCount, */exportAddresses[i] - (readDWORD(peHeader + 0x18 + 96) - exportTable);//);
            LPSTR forwardString = readString(stringAddr);
            char* delimiter = strchr(forwardString, '.');
            if (delimiter == NULL) {
                printf("Ordinal Forward String has unexpected format");
                delete[] buffer;
                return 1;
            }
            buffer[stringAddr + (delimiter - forwardString) + 1] = '#';
            delete[] forwardString;

            realOrdinal = ordinal - ordinalBase;
        } else if (strstr(name, "_ordFiller") == name) {
            writeInt(ordinalTable + i * 2, 0);
            delete[] name;
            continue;
        } else {
            ordinalEnd++;
            realOrdinal = ordinalEnd - ordinalBase;
        }
        writeInt(ordinalTable + i * 2, realOrdinal);
        writeDWORD(exportAddressTable + realOrdinal * 4, exportAddresses[i]);
        delete[] name;
    }
    writeDWORD(exportTable + 16, ordinalBase);
    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;
}