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

/* Internal */
#include "../include/slunkcrypt.h"
#include "version.h"

/* CRT */
#include <string.h>
#include <limits.h>
#include <assert.h>

/* Compiler compatibility */
#if defined(_MSC_VER)
#	define FORCE_INLINE __forceinline
#	define UNUSED __pragma(warning(suppress: 4189))
#elif defined(__GNUC__)
#	define FORCE_INLINE __attribute__((always_inline)) inline
#	define UNUSED __attribute__((unused))
#else
#	define FORCE_INLINE inline
#	define UNUSED
#endif

/* Version info */
const uint16_t SLUNKCRYPT_VERSION_MAJOR = MY_VERSION_MAJOR;
const uint16_t SLUNKCRYPT_VERSION_MINOR = MY_VERSION_MINOR;
const uint16_t SLUNKCRYPT_VERSION_PATCH = MY_VERSION_PATCH;
const char* const SLUNKCRYPT_BUILD = __DATE__ " " __TIME__;

/* Const */
#define HASH_MAGIC_PRIME 0x00000100000001B3ull
#define HASH_OFFSET_BASE 0xCBF29CE484222325ull

/* Utils */
#define BOOLIFY(X) (!!(X))
#define LOGICAL_XOR(X,Y) ((Y) ? (!(X)) : (X))

// ==========================================================================
// Data structures
// ==========================================================================

typedef struct
{
	uint64_t a, b, c;
}
key_data_t;

typedef struct
{
	uint32_t x, y, z, w, v, d;
}
rand_state_t;

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

// ==========================================================================
// Abort flag
// ==========================================================================

volatile int g_slunkcrypt_abort_flag = 0;

#define CHECK_ABORTED() do \
{ \
	if (g_slunkcrypt_abort_flag) \
	{ \
		return SLUNKCRYPT_ABORTED; \
	} \
} \
while (0)

// ==========================================================================
// Byte access (endianness agnostic)
// ==========================================================================

static FORCE_INLINE uint32_t lower_u64(const uint64_t value)
{
	return (uint32_t)(value & 0xFFFFFFFF);
}

static FORCE_INLINE uint32_t upper_u64(const uint64_t value)
{
	return (uint32_t)((value >> 32U) & 0xFFFFFFFF);
}

static FORCE_INLINE uint8_t byte_u16(const uint16_t value, const size_t off)
{
	assert(off < sizeof(uint16_t));
	return (uint8_t)((value >> (CHAR_BIT * off)) & 0xFF);
}

static FORCE_INLINE uint8_t byte_u64(const uint64_t value, const size_t off)
{
	assert(off < sizeof(uint64_t));
	return (uint8_t)((value >> (CHAR_BIT * off)) & 0xFF);
}

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

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

static FORCE_INLINE void hash_update_u64(uint64_t* const hash, const uint64_t value)
{
	size_t i;
	for (i = 0U; i < sizeof(uint64_t); ++i)
	{
		*hash = ((*hash) ^ byte_u64(value, i)) * HASH_MAGIC_PRIME;
	}
}

static FORCE_INLINE void hash_update_u16(uint64_t* const hash, const uint16_t value)
{
	size_t i;
	for (i = 0U; i < sizeof(uint16_t); ++i)
	{
		*hash = ((*hash) ^ byte_u16(value, i)) * HASH_MAGIC_PRIME;
	}
}

static uint64_t hash_code_init(const uint64_t salt, const uint16_t i, const uint8_t* const data, const size_t data_len)
{
	uint64_t hash = HASH_OFFSET_BASE;
	hash_update_u64(&hash, salt);
	hash_update_u16(&hash, i);
	hash_update_str(&hash, data, data_len);
	return hash;
}

static uint64_t hash_code_next(const uint64_t salt, const uint8_t* const data, const size_t data_len)
{
	uint64_t hash = HASH_OFFSET_BASE;
	hash_update_u64(&hash, salt);
	hash_update_str(&hash, data, data_len);
	return hash;
}

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

static FORCE_INLINE uint64_t keygen_loop(uint64_t salt, const uint16_t i, const uint8_t* const passwd, const size_t passwd_len)
{
	size_t u;
	uint64_t result = salt = hash_code_init(salt, i, passwd, passwd_len);
	for (u = 1U; u < 99971U; ++u)
	{
		result ^= salt = hash_code_next(salt, passwd, passwd_len);
	}
	return result;
}

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 & 0x3FFF) | 0x0000, passwd, passwd_len);
	key->b = keygen_loop(salt, (pepper & 0x3FFF) | 0x4000, passwd, passwd_len);
	key->c = keygen_loop(salt, (pepper & 0x3FFF) | 0x8000, passwd, passwd_len);
}

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

static void random_init(rand_state_t *const state, const key_data_t *const key)
{
	slunkcrypt_bzero(state, sizeof(rand_state_t));
	state->x = lower_u64(key->a);
	state->y = upper_u64(key->a);
	state->z = lower_u64(key->b);
	state->w = upper_u64(key->b);
	state->v = lower_u64(key->c);
	state->d = upper_u64(key->c);
}

static uint32_t random_next(rand_state_t *const state)
{
	const uint32_t t = state->x ^ (state->x >> 2);
	state->x = state->y;
	state->y = state->z;
	state->z = state->w;
	state->w = state->v;
	state->v ^= (state->v << 4) ^ t ^ (t << 1);
	return (state->d += 0x000587C5) + state->v;
}

static void random_seed(rand_state_t *const state, uint64_t salt, const uint16_t pepper, const uint8_t *const passwd, const size_t passwd_len)
{
	size_t i;
	key_data_t key;
	do
	{
		generate_key(&key, salt++, pepper, passwd, passwd_len);
		random_init(state, &key);
		slunkcrypt_bzero(&key, sizeof(key_data_t));
	}
	while (!(state->x || state->y || state->z || state->w || state->v));
	for (i = 0U; i < 97U; ++i)
	{
		UNUSED volatile uint32_t q = random_next(state);
	}
}

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

static int initialize_state(crypt_state_t *const crypt_state, const uint64_t nonce, const uint8_t *const passwd, const size_t passwd_len, const int reverse)
{
	rand_state_t rand_state;
	size_t r, i;

	/* initialize state */
	slunkcrypt_bzero(crypt_state, sizeof(crypt_state_t));
	crypt_state->reverse_mode = BOOLIFY(reverse);

	/* set up wheels and initial rotation */
	for (r = 0U; r < 256U; ++r)
	{
		random_seed(&rand_state, nonce, (uint16_t)r, passwd, passwd_len);
		crypt_state->rotation_bwd[255U - r] = crypt_state->rotation_fwd[r] = (uint8_t)random_next(&rand_state);
		for (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 (i = 0U; i < 256U; ++i)
		{
			crypt_state->wheel_bwd[255U - r][crypt_state->wheel_fwd[r][i]] = (uint8_t)i;
		}
		CHECK_ABORTED();
	}

	/* set up stepping */
	random_seed(&rand_state, nonce, 256U, passwd, passwd_len);
	for (i = 0U; i < 241U; ++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)(  6U + i);
		crypt_state->step_bwd[j] = (uint8_t)(249U - i);
	}

	/* final clean-up */
	slunkcrypt_bzero(&rand_state, sizeof(rand_state_t));
	return SLUNKCRYPT_SUCCESS;
}

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

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

static FORCE_INLINE void odometer_step(uint8_t *const arr, const int bwd)
{
	increment(arr, 0U, 6U, LOGICAL_XOR(bwd, 0));
	increment(arr, 0U, 3U, LOGICAL_XOR(bwd, 1));
	increment(arr, 3U, 6U, LOGICAL_XOR(bwd, 1));
	increment(arr, 6U, 9U, LOGICAL_XOR(bwd, 1));
}

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

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

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

int slunkcrypt_generate_nonce(uint64_t* const nonce)
{
	if (!nonce)
	{
		return SLUNKCRYPT_FAILURE;
	}
	do
	{
		if (slunkcrypt_random_bytes((uint8_t*)nonce, sizeof(uint64_t)) != sizeof(uint64_t))
		{
			return SLUNKCRYPT_FAILURE;
		}
	}
	while (!(*nonce));
	return SLUNKCRYPT_SUCCESS;
}

slunkcrypt_t slunkcrypt_alloc(const uint64_t nonce, const uint8_t *const passwd, const size_t passwd_len, const int mode)
{
	crypt_state_t* state = NULL;
	if ((!passwd) || (passwd_len < SLUNKCRYPT_PWDLEN_MIN) || (passwd_len > SLUNKCRYPT_PWDLEN_MAX) || (mode < SLUNKCRYPT_ENCRYPT) || (mode > SLUNKCRYPT_DECRYPT))
	{
		return SLUNKCRYPT_NULL;
	}
	if (!(state = (crypt_state_t*)malloc(sizeof(crypt_state_t))))
	{
		return SLUNKCRYPT_NULL;
	}
	if (initialize_state(state, nonce, passwd, passwd_len, mode) == SLUNKCRYPT_SUCCESS)
	{
		return ((slunkcrypt_t)state);
	}
	else
	{
		slunkcrypt_bzero(state, sizeof(crypt_state_t));
		return SLUNKCRYPT_NULL;
	}
}

int slunkcrypt_reset(const slunkcrypt_t context, const uint64_t nonce, const uint8_t *const passwd, const size_t passwd_len, const int mode)
{
	crypt_state_t* const state = (crypt_state_t*)context;
	int result = SLUNKCRYPT_FAILURE;
	if ((!state) || (!passwd) || (passwd_len < SLUNKCRYPT_PWDLEN_MIN) || (passwd_len > SLUNKCRYPT_PWDLEN_MAX) || (mode < SLUNKCRYPT_ENCRYPT) || (mode > SLUNKCRYPT_DECRYPT))
	{
		return SLUNKCRYPT_FAILURE;
	}
	if ((result = initialize_state(state, nonce, passwd, passwd_len, mode)) != SLUNKCRYPT_SUCCESS)
	{
		slunkcrypt_bzero(state, sizeof(crypt_state_t));
	}
	return result;
}

int slunkcrypt_process(const slunkcrypt_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 SLUNKCRYPT_FAILURE;
	}
	if (length > 0U)
	{
		size_t i;
		for (i = 0; i < length; ++i)
		{
			output[i] = state->reverse_mode ? process_decrypt(state, input[i]) : process_encrypt(state, input[i]);
			CHECK_ABORTED();
		}
	}
	return SLUNKCRYPT_SUCCESS;
}

int slunkcrypt_process_inplace(const slunkcrypt_t context, uint8_t* const buffer, size_t length)
{
	crypt_state_t* const state = (crypt_state_t*)context;
	if (!state)
	{
		return SLUNKCRYPT_FAILURE;
	}
	if (length > 0U)
	{
		size_t i;
		for (i = 0; i < length; ++i)
		{
			buffer[i] = state->reverse_mode ? process_decrypt(state, buffer[i]) : process_encrypt(state, buffer[i]);
			CHECK_ABORTED();
		}
	}
	return SLUNKCRYPT_SUCCESS;
}

void slunkcrypt_free(const slunkcrypt_t context)
{
	crypt_state_t* const state = (crypt_state_t*)context;
	if (state)
	{
		slunkcrypt_bzero(state, sizeof(crypt_state_t));
		free(state);
	}
}