#include <stdio.h>
#include <stdint.h>
#include <string.h>
#include <windows.h>

#define CASE(type) \
    case type: \
        fputs("Error: " #type "\n", stderr); \
        break;

static LONG WINAPI windows_exception_handler(EXCEPTION_POINTERS * ExceptionInfo)
{
  switch(ExceptionInfo->ExceptionRecord->ExceptionCode)
  {
    CASE(EXCEPTION_ACCESS_VIOLATION)
    CASE(EXCEPTION_ARRAY_BOUNDS_EXCEEDED)
    CASE(EXCEPTION_BREAKPOINT)
    CASE(EXCEPTION_DATATYPE_MISALIGNMENT)
    CASE(EXCEPTION_FLT_DENORMAL_OPERAND)
    CASE(EXCEPTION_FLT_DIVIDE_BY_ZERO)
    CASE(EXCEPTION_FLT_INEXACT_RESULT)
    CASE(EXCEPTION_FLT_INVALID_OPERATION)
    CASE(EXCEPTION_FLT_OVERFLOW)
    CASE(EXCEPTION_FLT_STACK_CHECK)
    CASE(EXCEPTION_FLT_UNDERFLOW)
    CASE(EXCEPTION_ILLEGAL_INSTRUCTION)
    CASE(EXCEPTION_IN_PAGE_ERROR)
    CASE(EXCEPTION_INT_DIVIDE_BY_ZERO)
    CASE(EXCEPTION_INT_OVERFLOW)
    CASE(EXCEPTION_INVALID_DISPOSITION)
    CASE(EXCEPTION_NONCONTINUABLE_EXCEPTION)
    CASE(EXCEPTION_PRIV_INSTRUCTION)
    CASE(EXCEPTION_SINGLE_STEP)
    CASE(EXCEPTION_STACK_OVERFLOW)
    default:
      fputs("Error: Unrecognized Exception\n", stderr);
      break;
  }
  fflush(stderr);

  return EXCEPTION_EXECUTE_HANDLER;
}

int
sym2hex(uint8_t *dest, uint8_t symbol)
{
    if ( (symbol >= 0x30) && (symbol <= 0x39) ) {
	    *dest |= symbol & 0x0F;
	    return 1;
    }

    if ( ( (symbol >= 0x41) && (symbol <= 0x46) ) ||
	     ( (symbol >= 0x61) && (symbol <= 0x66) ) ) {
	    *dest |= (symbol + 9) & 0x0F;
	    return 1;
    }

    return 0;
}

int
str2hex(uint8_t *dest, uint8_t *src, size_t size)
{
    int res;

    for (size_t i = 0, offset; i < size; i++) {
	    dest[i] = 0;
	    offset = i << 1;

	    res = sym2hex(&dest[i], src[i << 1]);
        if (src[(i << 1) + 1] == 0 || src[(i << 1) + 1] == '\n')
            break;

	    dest[i] <<= 4;
	    res = sym2hex(&dest[i], src[(i << 1) + 1]);

	    if (!res)
	        return -1;
    }

    return 0;
}

int
rc4(uint8_t *data, size_t data_size,
    const uint8_t *key, size_t key_size)
{
    uint8_t rc4_s[256];
    uint8_t key_item;
    int rc4_i, rc4_j;
    int tmp;

    if (strlen(key) > sizeof(rc4_s)) {
	    printf("Key must be under %ld bytes\n", sizeof(rc4_s));
	    return -1;
    }

    for (int i = 0; i < sizeof(rc4_s); i++)
	    rc4_s[i] = i;

    for (rc4_i = 0, rc4_j = 0; rc4_i < sizeof(rc4_s); rc4_i++) {
	    key_item = key[rc4_i % key_size];
	    rc4_j = (rc4_j + rc4_s[rc4_i] + key_item) % sizeof(rc4_s);

	    tmp = rc4_s[rc4_j];
	    rc4_s[rc4_j] = rc4_s[rc4_i];
	    rc4_s[rc4_i] = tmp;
    }

    rc4_i = 0;
    rc4_j = 0;
    for (int i = 0; i < data_size; i++) {
	    rc4_i = (rc4_i + 1) % sizeof(rc4_s);
	    rc4_j = (rc4_j + rc4_s[rc4_i]) % sizeof(rc4_s);

	    tmp = rc4_s[rc4_j];
	    rc4_s[rc4_j] = rc4_s[rc4_i];
	    rc4_s[rc4_i] = tmp;

	    tmp = (rc4_s[rc4_i] + rc4_s[rc4_j]) % sizeof(rc4_s);
	    data[i] ^= rc4_s[tmp];
    }

    return 0;
}

__asm__(".byte 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08");
volatile int
protected_main(void)
{
    WIN32_FIND_DATA fdFile;
    HANDLE hFind = NULL;

    puts("Call protected main");
    puts("Display all files from current directory");

    if((hFind = FindFirstFile(".\\*.*", &fdFile)) == INVALID_HANDLE_VALUE) {
        puts("Can not read first file from directory");
        return -1;
    }

    do {
        if (strcmp(fdFile.cFileName, ".") == 0)
            continue;

        if (strcmp(fdFile.cFileName, "..") == 0)
            continue;

        if (fdFile.dwFileAttributes &FILE_ATTRIBUTE_DIRECTORY) {
            printf("Directory: %s\n", fdFile.cFileName);
            continue;
        }
        printf("File: %s\n", fdFile.cFileName);
    } while (FindNextFile(hFind, &fdFile));

    FindClose(hFind);
    return 0;
}

volatile int
protected_main_end(void) { }

__asm__(".byte 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08");

int
main(void)
{
    uint8_t key_str[256 << 1];
    uint8_t key_hex[256];
    size_t key_str_len;
    size_t key_hex_len;
    size_t func_size;
    DWORD old_perm;
    int res;

    memset(key_str, 0, sizeof(key_str));
    memset(key_hex, 0, sizeof(key_hex));

    puts("Enter key:");
    fgets(key_str, sizeof(key_str), stdin);

    key_str_len = strlen(key_str);
    key_hex_len = (key_str_len >> 1) + (key_str_len & 1);
    res = str2hex(key_hex, key_str, key_hex_len);
    if (res) {
        puts("Invalid key value");
	    return res;
    }

    func_size = (size_t)protected_main_end - (size_t)protected_main;

    SetUnhandledExceptionFilter(windows_exception_handler);

    VirtualProtect(
        protected_main, func_size,
        PAGE_EXECUTE_READWRITE,
        &old_perm
    );
    rc4(
        (uint8_t *)protected_main - 8, func_size + 8,
        key_hex, key_hex_len
    );

    res = protected_main();

    rc4(
        (uint8_t *)protected_main - 8, func_size + 8,
        key_hex, key_hex_len
    );
    VirtualProtect(
        protected_main, func_size,
        old_perm,
        NULL
    );

    return res;