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

/* Internal */
#include "compiler.h"
#include "slunkcrypt.h"
#include "debug.h"
#include "keygen.h"
#include "initialize.h"
#include "thread.h"
#include "version.h"

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

/* Version */
const uint16_t SLUNKCRYPT_VERSION_MAJOR = LIB_VERSION_MAJOR;
const uint16_t SLUNKCRYPT_VERSION_MINOR = LIB_VERSION_MINOR;
const uint16_t SLUNKCRYPT_VERSION_PATCH = LIB_VERSION_PATCH;
const char *const SLUNKCRYPT_BUILD = __DATE__ ", " __TIME__;

/* Utilities */
#define BOOLIFY(X) (!!(X))
#define THREAD_COUNT(X) (((X)->thread_pool) ? slunkcrypt_thrdpl_count((X)->thread_pool) : 1U)

/* Configuration */
const int SLUNKCRYPT_HAVE_THREADS = BOOLIFY(MAX_THREADS > 1U);

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

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

typedef struct
{
	int reverse_mode;
	const uint8_t (*wheel)[256U];
	uint32_t counter;
	size_t index_off;
	rand_state_t random;
}
thread_state_t;

typedef struct
{
	uint8_t wheel[256U][256U];
	thread_state_t thread_data[MAX_THREADS];
}
crypt_data_t;

typedef struct
{
	int legacy_compat;
	int debug_logging;
	thrdpl_t *thread_pool;
	crypt_data_t data;
}
crypt_state_t;

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

volatile int g_slunkcrypt_abort_flag = 0;

#define CHECK_ABORTED() do \
{ \
	if (g_slunkcrypt_abort_flag) \
	{ \
		goto aborted; \
	} \
} \
while (0)

// ==========================================================================
// Deterministic random bit generator
// ==========================================================================

#define RANDOM_INIT(DST_LO, DST_HI, SRC) do \
{ \
	DST_LO = (uint32_t)(SRC); \
	DST_HI = (uint32_t)((SRC) >> 32U); \
} \
while (0)

static INLINE void random_init(rand_state_t *const state, const keydata_t *const key)
{
	RANDOM_INIT(state->x, state->y, key->a);
	RANDOM_INIT(state->z, state->w, key->b);
	RANDOM_INIT(state->v, state->d, key->c);
}

static INLINE 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 INLINE void random_skip(rand_state_t *const state, const size_t skip_count)
{
	size_t i;
	for (i = 0U; i < skip_count; ++i)
	{
		UNUSED /*volatile*/ uint32_t discard = random_next(state);
	}
}

static INLINE 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, const int legacy)
{
	keydata_t key;
	do
	{
		slunkcrypt_keygen(&key, salt++, pepper, passwd, passwd_len, legacy ? 99971U : 19997L);
		random_init(state, &key);
		slunkcrypt_bzero(&key, sizeof(keydata_t));
	}
	while (!(state->x || state->y || state->z || state->w || state->v));
	random_skip(state, legacy ? 97U : 23U);
}

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

static int initialize_state(crypt_data_t *const data, const size_t thread_count, const uint64_t nonce, const uint8_t *const passwd, const size_t passwd_len, const int mode, const int legacy, const int debug)
{
	uint8_t temp[256U][256U];
	size_t r, i;
	rand_state_t *const rand_state = &data->thread_data[0].random;
	const int reverse_mode = BOOLIFY(mode);

	/* initialize state */
	slunkcrypt_bzero(data, sizeof(crypt_data_t));

	/* initialize counter */
	random_seed(rand_state, nonce, legacy ? ((uint16_t)(-1)) : INITIALIZER[256U], passwd, passwd_len, legacy);
	data->thread_data[0].counter = random_next(rand_state);

	/* set up the wheel permutations */
	for (r = 0U; r < 256U; ++r)
	{
		random_seed(rand_state, nonce, legacy ? ((uint16_t)r) : INITIALIZER[r], passwd, passwd_len, legacy);
		for (i = 0U; i < 256U; ++i)
		{
			const size_t j = random_next(rand_state) % (i + 1U);
			if (j != i)
			{
				data->wheel[r][i] = data->wheel[r][j];
			}
			data->wheel[r][j] = (uint8_t)i;
		}
		CHECK_ABORTED();
	}

	/* reverse the wheels, if requested */
	if (reverse_mode)
	{
		for (r = 0U; r < 256U; ++r)
		{
			for (i = 0U; i < 256U; ++i)
			{
				temp[r][data->wheel[r][i]] = (uint8_t)i;
			}
		}
		for (r = 0U; r < 256U; ++r)
		{
			memcpy(data->wheel[255U - r], temp[r], 256U);
		}
		slunkcrypt_bzero(temp, sizeof(temp));
		CHECK_ABORTED();
	}

	/* initialize thread state */
	random_seed(rand_state, nonce, legacy ? 256U : INITIALIZER[257U], passwd, passwd_len, legacy);
	for (i = 0U; i < thread_count; ++i)
	{
		data->thread_data[i].reverse_mode = reverse_mode;
		data->thread_data[i].wheel = (const uint8_t(*)[256]) data->wheel;
		data->thread_data[i].index_off = i;
		if (i > 0U)
		{
			data->thread_data[i].counter = data->thread_data[0].counter + ((uint32_t)i);
			memcpy(&data->thread_data[i].random, &data->thread_data[0].random, sizeof(rand_state_t));
			random_skip(&data->thread_data[i].random, i * 63U);
		}
		CHECK_ABORTED();
	}

	/* dump the final configuration */
	if (debug)
	{
		slunkcrypt_debug_print("cntr = %08X", data->thread_data[0].counter);
		slunkcrypt_debug_print("drbg = %08X %08X %08X %08X %08X %08X", rand_state->d, rand_state->v, rand_state->w, rand_state->x, rand_state->y, rand_state->z);
		for (r = 0U; r < 256U; ++r)
		{
			char message[775U] = { '\0' };
			char *ptr = message;
			DBG_PRINTF(ptr, "[%02X] =", (unsigned)r);
			for (i = 0U; i < 256U; ++i)
			{
				DBG_PRINTF(ptr, " %02X", data->wheel[r][i]);
			}
			slunkcrypt_debug_write(message);
		}
	}

	return SLUNKCRYPT_SUCCESS;

	/* aborted */
aborted:
	slunkcrypt_bzero(data, sizeof(crypt_data_t));
	return SLUNKCRYPT_ABORTED;
}

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

static INLINE void update_offset(uint8_t *const offset, uint32_t seed, rand_state_t *const state, const int reverse)
{
	size_t i;
	for (i = 0U; i < 256U; ++i, seed >>= CHAR_BIT)
	{
		if (i && (!(i & 3U)))
		{
			seed = random_next(state);
		}
		offset[reverse ? (255U - i) : i] = (uint8_t)seed;
	}
}

static INLINE uint8_t process_next_symbol(thread_state_t *const state, uint8_t value)
{
	uint8_t offset[256U];
	size_t i;
	update_offset(offset, state->counter, &state->random, state->reverse_mode);
	for (i = 0U; i < 256U; ++i)
	{
		value = (state->wheel[i][(value + offset[i]) & 0xFF] - offset[i]) & 0xFF;
	}
	return value;
}

// ==========================================================================
// Thread entry point
// ==========================================================================

static INLINE void update_index(thread_state_t *const state, const size_t thread_count, const size_t length)
{
	const size_t remaining = thread_count - (length % thread_count);
	if (remaining != thread_count)
	{
		state->index_off = (state->index_off + remaining) % thread_count;
	}
}

static void thread_worker(const size_t thread_count, void *const context, uint8_t *const buffer, const size_t length)
{
	thread_state_t *const state = (thread_state_t*) context;
	size_t i;

	for (i = state->index_off; i < length; i += thread_count)
	{
		buffer[i] = process_next_symbol(state, buffer[i]);
		state->counter += (uint32_t)thread_count;
		random_skip(&state->random, 63U * (thread_count - 1U));
	}

	update_index(state, thread_count, length);
}

// ==========================================================================
// 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)
{
	const slunkparam_t param = { SLUNKCRYPT_PARAM_VERSION, 0U, SLUNKCRYPT_TRUE, SLUNKCRYPT_FALSE };
	return slunkcrypt_alloc_ext(nonce, passwd, passwd_len, mode, &param);
}

slunkcrypt_t slunkcrypt_alloc_ext(const uint64_t nonce, const uint8_t *const passwd, const size_t passwd_len, const int mode, const slunkparam_t *const param)
{
	crypt_state_t *state = NULL;

	if ((!passwd) || (passwd_len < SLUNKCRYPT_PWDLEN_MIN) || (passwd_len > SLUNKCRYPT_PWDLEN_MAX) ||
		(mode < SLUNKCRYPT_ENCRYPT) || (mode > SLUNKCRYPT_DECRYPT) || (!param) || (!param->version) || (param->version > SLUNKCRYPT_PARAM_VERSION))
	{
		return SLUNKCRYPT_NULL;
	}

	if (!(state = (crypt_state_t*) calloc(1U, sizeof(crypt_state_t))))
	{
		return SLUNKCRYPT_NULL;
	}

	if ((state->thread_pool = slunkcrypt_thrdpl_create(param->thread_count, thread_worker)))
	{
		const size_t thread_count = slunkcrypt_thrdpl_count(state->thread_pool);
		size_t i;
		for (i = 0U; i < thread_count; ++i)
		{
			slunkcrypt_thrdpl_init(state->thread_pool, i, &state->data.thread_data[i]);
		}
	}

	state->legacy_compat = (param->version > 1U) ? BOOLIFY(param->legacy_compat) : SLUNKCRYPT_FALSE;
	state->debug_logging = (param->version > 1U) ? BOOLIFY(param->debug_logging) : SLUNKCRYPT_FALSE;

	if (initialize_state(&state->data, THREAD_COUNT(state), nonce, passwd, passwd_len, mode, state->legacy_compat, state->debug_logging) == SLUNKCRYPT_SUCCESS)
	{
		return (slunkcrypt_t)state;
	}

	slunkcrypt_free((slunkcrypt_t)state);
	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->data, THREAD_COUNT(state), nonce, passwd, passwd_len, mode, state->legacy_compat, state->debug_logging)) != SLUNKCRYPT_SUCCESS)
	{
		slunkcrypt_bzero(&state->data, sizeof(crypt_data_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)
	{
		memcpy(output, input, length);
		return slunkcrypt_inplace(context, output, length);
	}

	return SLUNKCRYPT_SUCCESS;
}

int slunkcrypt_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)
	{
		if (THREAD_COUNT(state) > 1U)
		{
			slunkcrypt_thrdpl_exec(state->thread_pool, buffer, length);
		}
		else
		{
			thread_worker(1U, &state->data.thread_data[0U], buffer, length);
		}
	}

	CHECK_ABORTED();
	return SLUNKCRYPT_SUCCESS;

aborted:
	slunkcrypt_bzero(buffer, length);
	slunkcrypt_bzero(&state->data, sizeof(crypt_data_t));
	return SLUNKCRYPT_ABORTED;
}

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