/******************************************************************************/
/* MCrypt, by LoRd_MuldeR <MuldeR2@GMX.de>                                    */
/* This work has been released under the CC0 1.0 Universal license!           */
/******************************************************************************/

#ifdef _WIN32
#define  _CRT_RAND_S 1
#endif

#include <mcrypt.h>
#include <string.h>

#ifdef _MSC_VER
#define FORCE_INLINE __forceinline
#else
#ifdef __GNUC__
#define FORCE_INLINE __attribute__((always_inline)) inline
#else
#define FORCE_INLINE inline
#endif
#endif

const char* const LIBMCRYPT_VERSION = "1.0.0";
const char* const LIBMCRYPT_BUILDNO = __DATE__", "__TIME__;

typedef struct
{
	uint64_t a, b;
}
key_data_t;

typedef struct
{
	uint8_t wheel_fwd[256U][256U];
	uint8_t wheel_bwd[256U][256U];
	uint8_t step_fwd[256U];
	uint8_t step_bwd[256U];
	uint8_t rotation_fwd[2U][256U];
	uint8_t rotation_bwd[2U][256U];
	uint8_t counter;
}
crypt_state_t;

typedef struct
{
	uint32_t a, b, c, d;
	uint32_t counter;
}
rand_state_t;

// ==========================================================================
// Hash function
// ==========================================================================

static FORCE_INLINE void hash_update(uint64_t *const h, const uint8_t* const data, const size_t data_len)
{
	for (size_t i = 0U; i < data_len; ++i)
	{
		*h ^= data[i];
		*h *= 0x00000100000001B3ull;
	}
}

static uint64_t hash_code(const uint64_t salt, const uint16_t pepper, const uint8_t* const data, const size_t data_len)
{
	uint64_t h = 0xCBF29CE484222325ull;
	hash_update(&h, (uint8_t*)&salt,   sizeof(uint64_t));
	hash_update(&h, (uint8_t*)&pepper, sizeof(uint16_t));
	hash_update(&h, data, data_len);
	return h;
}

// ==========================================================================
// Key derivation
// ==========================================================================

static FORCE_INLINE uint64_t keygen_loop(uint64_t value, const uint16_t pepper, const uint8_t* const passwd, const size_t passwd_len)
{
	for (size_t i = 0U; i < 99971U; ++i)
	{
		value ^= hash_code(value, pepper, passwd, passwd_len);
	}
	return value;
}

static void generate_key(key_data_t *const key, const uint64_t salt, const uint16_t pepper, const uint8_t* const passwd, const size_t passwd_len)
{
	key->a = keygen_loop(salt, pepper & 0x7FFF, passwd, passwd_len);
	key->b = keygen_loop(salt, pepper | 0x8000, passwd, passwd_len);
}

// ==========================================================================
// PRNG
// ==========================================================================

static void random_init(rand_state_t* const state, const uint64_t seed_0, const uint64_t seed_1)
{
	mcrypt_bzero(state, sizeof(rand_state_t));
	state->a = (uint32_t)( seed_0        & 0xFFFFFFFF);
	state->b = (uint32_t)((seed_0 >> 32) & 0xFFFFFFFF);
	state->c = (uint32_t)( seed_1        & 0xFFFFFFFF);
	state->d = (uint32_t)((seed_1 >> 32) & 0xFFFFFFFF);
}

static uint32_t random_next(rand_state_t *const state)
{
	uint32_t t = state->d;
	const uint32_t s = state->a;
	state->d = state->c;
	state->c = state->b;
	state->b = s;
	t ^= t >> 2;
	t ^= t << 1;
	t ^= s ^ (s << 4);
	state->a = t;
	return t + (state->counter += 362437U);
}

static void random_seed(rand_state_t* const state, const uint64_t salt, const uint16_t pepper, const uint8_t *const passwd, const size_t passwd_len)
{
	key_data_t key;
	generate_key(&key, salt, pepper, passwd, passwd_len);
	random_init(state, key.a, key.b);
	mcrypt_bzero(&key, sizeof(key_data_t));
	for (size_t i = 0U; i < 97U; ++i)
	{
		volatile uint32_t u = random_next(state);
	}
}

// ==========================================================================
// Initialization
// ==========================================================================

static void initialize_state(crypt_state_t* const crypt_state, const uint64_t salt, const uint8_t* const passwd, const size_t passwd_len)
{
	mcrypt_bzero(crypt_state, sizeof(crypt_state_t));

	/* set up wheels and initial rotation */
	rand_state_t rand_state;
	for (size_t r = 0U; r < 256U; ++r)
	{
		random_seed(&rand_state, salt, (uint16_t)r, passwd, passwd_len);
		crypt_state->rotation_bwd[0U][255U - r] = crypt_state->rotation_fwd[0U][r] = (uint8_t)random_next(&rand_state);
		crypt_state->rotation_bwd[1U][255U - r] = crypt_state->rotation_fwd[1U][r] = 0U;
		for (size_t i = 0U; i < 256U; ++i)
		{
			const size_t j = random_next(&rand_state) % (i + 1U);
			if (j != i)
			{
				crypt_state->wheel_fwd[r][i] = crypt_state->wheel_fwd[r][j];
			}
			crypt_state->wheel_fwd[r][j] = (uint8_t)i;
		}
		for (size_t i = 0U; i < 256U; ++i)
		{
			const size_t j = crypt_state->wheel_fwd[r][i];
			crypt_state->wheel_bwd[255U - r][j] = (uint8_t)i;
		}
	}

	/* set up stepping */
	random_seed(&rand_state, salt, 256U, passwd, passwd_len);
	for (size_t i = 0U; i < 256U; ++i)
	{
		const size_t j = random_next(&rand_state) % (i + 1U);
		if (j != i)
		{
			crypt_state->step_fwd[i] = crypt_state->step_fwd[j];
			crypt_state->step_bwd[i] = crypt_state->step_bwd[j];
		}
		crypt_state->step_fwd[j] = (uint8_t)i;
		crypt_state->step_bwd[j] = (uint8_t)(255U - i);
	}

	mcrypt_bzero(&rand_state, sizeof(rand_state_t));
}

// ==========================================================================
// Encrypt / Decrypt
// ==========================================================================

static FORCE_INLINE void increment(uint8_t *const arr, const int rev)
{
	for (size_t i = 0U; i < 256U; ++i)
	{
		if (++arr[rev ? (255U - i) : i] != 0U)
		{
			break;
		}
	}
}

static FORCE_INLINE uint8_t process_enc(crypt_state_t* const crypt_state, uint8_t value)
{
	for (size_t i = 0U; i < 256U; ++i)
	{
		const uint8_t offset = crypt_state->rotation_fwd[0U][i] + crypt_state->rotation_fwd[1U][i];
		value = crypt_state->wheel_fwd[i][(value + offset) & 0xFF];
	}
	++crypt_state->rotation_fwd[0U][crypt_state->step_fwd[crypt_state->counter++]];
	increment(crypt_state->rotation_fwd[1U], 0);
	return value;
}

static FORCE_INLINE uint8_t process_dec(crypt_state_t* const crypt_state, uint8_t value)
{
	for (size_t i = 0U; i < 256U; ++i)
	{
		const uint8_t offset = crypt_state->rotation_bwd[0U][i] + crypt_state->rotation_bwd[1U][i];
		value = (crypt_state->wheel_bwd[i][value] - offset) & 0xFF;
	}
	++crypt_state->rotation_bwd[0U][crypt_state->step_bwd[crypt_state->counter++]];
	increment(crypt_state->rotation_bwd[1U], 1);
	return value;
}

// ==========================================================================
// Public API
// ==========================================================================

int mcrypt_generate_seed(uint64_t* const seed)
{
	if (!seed)
	{
		return MCRYPT_FAILURE;
	}
	do
	{
		if (mcrypt_random_bytes((uint8_t*)seed, sizeof(uint64_t)) != 0)
		{
			return MCRYPT_FAILURE;
		}
	}
	while (!(*seed));
	return MCRYPT_SUCCESS;
}

mcrypt_t mcrypt_alloc(const uint64_t salt, const uint8_t *const passwd, const size_t passwd_len)
{
	if ((!passwd) || (passwd_len < 1U))
	{
		return ((mcrypt_t)NULL);
	}
	crypt_state_t* const state = (crypt_state_t*)malloc(sizeof(crypt_state_t));
	if (!state)
	{
		return ((mcrypt_t)NULL);
	}
	initialize_state(state, salt, passwd, passwd_len);
	return ((mcrypt_t)state);
}

int mcrypt_reset(const mcrypt_t context, const uint64_t salt, const uint8_t *const passwd, const size_t passwd_len)
{
	crypt_state_t* const state = (crypt_state_t*)context;
	if ((!state) || (!passwd) || (passwd_len < 1U))
	{
		return MCRYPT_FAILURE;
	}
	initialize_state(state, salt, passwd, passwd_len);
	return MCRYPT_SUCCESS;
}

int mcrypt_encrypt(const mcrypt_t context, const uint8_t* const input, uint8_t* const output, size_t length)
{
	crypt_state_t* const state = (crypt_state_t*)context;
	if (!state)
	{
		return MCRYPT_FAILURE;
	}
	for (size_t i = 0; i < length; ++i)
	{
		output[i] = process_enc(state, input[i]);
	}
	return MCRYPT_SUCCESS;
}

int mcrypt_encrypt_inplace(const mcrypt_t context, uint8_t* const buffer, size_t length)
{
	crypt_state_t* const state = (crypt_state_t*)context;
	if (!state)
	{
		return MCRYPT_FAILURE;
	}
	for (size_t i = 0; i < length; ++i)
	{
		buffer[i] = process_enc(state, buffer[i]);
	}
	return MCRYPT_SUCCESS;
}


int mcrypt_decrypt(const mcrypt_t context, const uint8_t* const input, uint8_t* const output, size_t length)
{
	crypt_state_t* const state = (crypt_state_t*)context;
	if (!state)
	{
		return MCRYPT_FAILURE;
	}
	for (size_t i = 0; i < length; ++i)
	{
		output[i] = process_dec(state, input[i]);
	}
	return MCRYPT_SUCCESS;
}

int mcrypt_decrypt_inplace(const mcrypt_t context, uint8_t* const buffer, size_t length)
{
	crypt_state_t* const state = (crypt_state_t*)context;
	if (!state)
	{
		return MCRYPT_FAILURE;
	}
	for (size_t i = 0; i < length; ++i)
	{
		buffer[i] = process_dec(state, buffer[i]);
	}
	return MCRYPT_SUCCESS;
}

void mcrypt_free(const mcrypt_t context)
{
	crypt_state_t* const state = (crypt_state_t*)context;
	if (state)
	{
		mcrypt_bzero(state, sizeof(crypt_state_t));
		free(state);
	}
}