#include <cstdio>
#include <ctime>
#include <Windows.h>
#include <algorithm>
#pragma warning(disable:28159)
#pragma warning(disable:4996)

DWORD pKeySink[623];
DWORD dwKeyLengthInDword = 623;
DWORD dwGenerateKeyCount;

BYTE cbDataSink[1000];
bool cbKeyGenerated;

void GenerateKeySinkWithRandom(_In_ DWORD Random) {
    DWORD dwCurrentValue;
    dwGenerateKeyCount = 0;
    cbKeyGenerated = true;
    pKeySink[0] = Random;
    for (DWORD index = 0; index < 622; ++index) {
        dwCurrentValue = pKeySink[index];
        pKeySink[index + 1] = index + 0x6C078965 * (dwCurrentValue ^ (dwCurrentValue >> 30));
    }
}

void updateKeySink() {
    bool v4; // zf

    for (auto i = 0; i < dwKeyLengthInDword; ++i) {
        v4 = (LOBYTE(pKeySink[i + 1] % 624) & 1) == 0;
        pKeySink[i] = ((pKeySink[(i + 1) % 624] & 0x7fffffffu) >> 1) ^ pKeySink[i % 624];
        if (!v4)
            pKeySink[i] ^= 0x9908B0DF;
    }
    return;
}

DWORD GenerateRandomInKeySink() {
    int v1; // eax
    DWORD v2; // ecx
    DWORD v3; // edx
    DWORD result; // eax

    if (!cbKeyGenerated) GenerateKeySinkWithRandom(time(nullptr));
    
    v1 = dwGenerateKeyCount;
    if (!dwGenerateKeyCount)
    {
        updateKeySink();
        v1 = dwGenerateKeyCount;
    }
    v2 = ((((pKeySink[v1] >> 11) ^ pKeySink[v1]) & 0xFF3A58AD) << 7) ^ (pKeySink[v1] >> 11) ^ pKeySink[v1];
    v3 = (v1 + 1) % 0x270u;
    result = ((v2 & 0xFFFFDF8C) << 15) ^ v2 ^ ((((v2 & 0xFFFFDF8C) << 15) ^ v2) >> 18);
    dwGenerateKeyCount = v3;
    return result;
}

__forceinline BOOLEAN NTAPI RtlIsValidImageHeader(_In_ LPCVOID lpBufferToImageHeader) {
    auto dos = PIMAGE_DOS_HEADER(lpBufferToImageHeader);
    return (dos->e_magic == IMAGE_DOS_SIGNATURE) &&
        (dos->e_lfanew >= sizeof(IMAGE_DOS_HEADER)) &&
        (dos->e_lfanew < 0x400) &&
        (PIMAGE_NT_HEADERS(SIZE_T(lpBufferToImageHeader) + dos->e_lfanew)->Signature == IMAGE_NT_SIGNATURE);
}

VOID NTAPI RtlUpdateBuffer(
    _In_ DWORD dwKey,
    _In_ BOOLEAN Encrypt,
    _In_ DWORD dwUserKeyLength,
    _In_ DWORD dwAutoKeyLength,
    _In_reads_bytes_(dwCiphertextLength) LPCVOID lpCiphertext,
    _In_ DWORD dwCiphertextLength,
    _Out_writes_bytes_all_(dwCiphertextLength) LPVOID lpResult,
    _Out_writes_opt_(dwUserKeyLength) LPDWORD lpKey) {
    LPDWORD pAutoKey = new DWORD[1ull + dwAutoKeyLength];
    DWORD dwTemp = 0;

    RtlMoveMemory(lpResult, lpCiphertext, dwCiphertextLength);

    GenerateKeySinkWithRandom(dwKey);
    for (auto i = 0; i < dwUserKeyLength; ++i)
        cbDataSink[i] = GenerateRandomInKeySink();
    for (auto i = 10; i; --i)
        for (auto j = 0; j < dwUserKeyLength; ++j)
            std::swap(cbDataSink[j], cbDataSink[(j + GenerateRandomInKeySink()) % dwUserKeyLength]);

    cbDataSink[dwUserKeyLength] = 0;

    if (dwUserKeyLength > 0)
        for (auto i = 0; i < dwUserKeyLength; ++i)
            dwTemp += cbDataSink[i];
    GenerateKeySinkWithRandom(dwTemp);
    if (dwAutoKeyLength > 0) {
        for (auto i = 0; i < dwAutoKeyLength; ++i)
            pAutoKey[i] = GenerateRandomInKeySink();
        for (auto i = 0; i < dwAutoKeyLength; ++i)
            std::swap(pAutoKey[i], pAutoKey[(i + GenerateRandomInKeySink()) % dwAutoKeyLength]);
    }
    GenerateKeySinkWithRandom(dwTemp * dwTemp);
    if (dwCiphertextLength > 0) {
        dwTemp = dwCiphertextLength % dwAutoKeyLength;
        for (auto i = 0; i < dwCiphertextLength; ++i) {
            if (i == dwTemp) GenerateKeySinkWithRandom(pAutoKey[1]);
            LPBYTE(lpResult)[i] += (Encrypt ? 1 : -1) * GenerateRandomInKeySink();
        }
    }
    if (lpKey) {
        for (auto i = 0; i < dwUserKeyLength; ++i)
            lpKey[i] = cbDataSink[i];
    }
    delete[]pAutoKey;
}

#define RtlEncryptBuffer(_dwKey_, _dwUserKeyLength_, _dwAutoKeyLength, _lpCiphertext_, _dwCiphertextLength_, _lpResult_, _lpKey_)\
    RtlUpdateBuffer(_dwKey_, TRUE, _dwUserKeyLength_, _dwAutoKeyLength, _lpCiphertext_, _dwCiphertextLength_, _lpResult_, _lpKey_)
#define RtlDecryptBuffer(_dwKey_, _dwUserKeyLength_, _dwAutoKeyLength, _lpCiphertext_, _dwCiphertextLength_, _lpResult_, _lpKey_)\
    RtlUpdateBuffer(_dwKey_, FALSE, _dwUserKeyLength_, _dwAutoKeyLength, _lpCiphertext_, _dwCiphertextLength_, _lpResult_, _lpKey_)

BOOLEAN NTAPI RtlCrackPeFileEncryption(
    _Out_writes_bytes_(21*2) LPSTR lpUserPassword,
    _In_opt_ DWORD dwBeginRandom,
    _In_ LPCVOID lpBufferToCrack,
    _In_ DWORD dwBufferSize) {
    lpUserPassword[0] = '\0';

    BOOLEAN result = FALSE;
    LPVOID lpData = nullptr;
    const auto Length = sizeof(IMAGE_DOS_HEADER) + sizeof(IMAGE_NT_HEADERS32);
    auto BeginRandom = dwBeginRandom;
    DWORD dwUserKey[2];

    if (dwBufferSize < Length)
        return FALSE;

    lpData = new char[dwBufferSize];
    while (BeginRandom & 0xffffffff) {
        RtlDecryptBuffer(BeginRandom, 2, 567, lpBufferToCrack, dwBufferSize, lpData, LPDWORD(&dwUserKey));
        if (RtlIsValidImageHeader(lpData)) {
            //found
            result = TRUE;
            sprintf(lpUserPassword, "%02X%02X", BYTE(dwUserKey[0]), BYTE(dwUserKey[1]));
            break;
        }
        --BeginRandom;
    }

    return result;
}

int main() {
    //BYTE test[4] = { 0x11,0x22,0x33,0x44 };
    //PDWORD key = nullptr;
    //calc(test, 4, key, &key, 2, 567);
    //calc(test, 4, key, nullptr, 2, 567);

    typedef struct _PE_HEADERS{
        IMAGE_DOS_HEADER dos;
        IMAGE_NT_HEADERS32 nt;
        _PE_HEADERS() {
            dos.e_magic = IMAGE_DOS_SIGNATURE;
            dos.e_lfanew = sizeof(dos);
            nt.Signature = IMAGE_NT_SIGNATURE;
        }
    }PE_HEADERS,*PPE_HEADERS;

    auto pImage = new PE_HEADERS;
    DWORD key[2];
    char szKey[42];
    RtlEncryptBuffer(GetTickCount(), 2, 567, pImage, sizeof(PE_HEADERS), pImage, LPDWORD(&key));
    if (!RtlCrackPeFileEncryption(szKey, GetTickCount(), pImage, sizeof(PE_HEADERS))) {
        //failed
        abort();
    }

    //check key[] and szKey

    //success
    delete pImage;
    return 0;
}