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

/* Internal */
#include "crypt.h"
#include "utils.h"
#include "blake2.h"

/* Library */
#include <slunkcrypt.h>

/* CRT */
#include <time.h>
#include <inttypes.h>
#include <ctype.h>
#include <assert.h>
#include <errno.h>

// ==========================================================================
// Constants
// ==========================================================================

static const uint64_t MAGIC_NUMBER = 0x243F6A8885A308D3ull;

#define MAX_BUFFER_SIZE 1048576U // 1 MiB

// ==========================================================================
// Auxiliary functions
// ==========================================================================

static int open_files(FILE **const file_in, FILE **const file_out, const CHR *const input_path, const CHR *const output_path)
{
	if (!(*file_in = FOPEN(input_path, "rb")))
	{
		FPRINTF(stderr, T("Error: Failed to open input file \"%") T(PRISTR) T("\" for reading!\n\n%") T(PRISTR) T("\n\n"), input_path, STRERROR(errno));
		*file_out = NULL;
		return EXIT_FAILURE;
	}

	if (!(*file_out = FOPEN(output_path, "wb")))
	{
		FPRINTF(stderr, T("Error: Failed to open output file \"%") T(PRISTR) T("\" for writing!\n\n%") T(PRISTR) T("\n\n"), output_path, STRERROR(errno));
		return EXIT_FAILURE;
	}

	//setvbuf(*file_in,  NULL, _IOFBF, (((8U * BUFFER_SIZE) + (BUFSIZ - 1U)) / BUFSIZ) * BUFSIZ);
	//setvbuf(*file_out, NULL, _IOFBF, (((8U * BUFFER_SIZE) + (BUFSIZ - 1U)) / BUFSIZ) * BUFSIZ);

	return EXIT_SUCCESS;
}

static size_t get_buffer_size(const uint64_t file_size)
{
	const uint64_t upper_limit = file_size / 101U;
	size_t pwr2, result = 0U;
	for (pwr2 = 1U; (pwr2 <= MAX_BUFFER_SIZE) && (pwr2 <= upper_limit); pwr2 <<= 1)
	{
		result = pwr2;
	}
	return BOUND(BUFSIZ, result, MAX_BUFFER_SIZE);
}

static void init_slunk_param(slunkparam_t *const param, const crypt_options_t *const options)
{
	slunkcrypt_bzero(param, sizeof(slunkparam_t));
	param->version = SLUNKCRYPT_PARAM_VERSION;
	param->thread_count = options->thread_count;
}

#define UPDATE_PROGRESS_INDICATOR(CLK_UPDATE, CURRENT, TOTAL) do \
{ \
	const uint64_t clk_now = clock_read(); \
	if ((clk_now < (CLK_UPDATE)) || (clk_now - (CLK_UPDATE) > update_interval)) \
	{ \
		FPRINTF(stderr, T("\b\b\b\b\b\b\b%5.1f%% "), ((CURRENT) / ((double)(TOTAL))) * 100.0); \
		fflush(stderr); \
		CLK_UPDATE = clk_now; \
	} \
} \
while(0)

// ==========================================================================
// Encrypt
// ==========================================================================

int encrypt(const char *const passphrase, const CHR *const input_path, const CHR *const output_path, const crypt_options_t *const options)
{
	slunkcrypt_t ctx = SLUNKCRYPT_NULL;
	slunkparam_t param;
	FILE *file_in = NULL, *file_out = NULL;
	uint8_t *buffer = NULL;
	size_t buffer_size = BUFSIZ;
	int result = EXIT_FAILURE, status = -1;

	if (open_files(&file_in, &file_out, input_path, output_path) != EXIT_SUCCESS)
	{
		goto clean_up;
	}

	const uint64_t file_size = get_size(file_in);
	if (file_size == UINT64_MAX)
	{
		FPUTS(T("I/O error: Failed to determine size of input file!\n\n"), stderr);
		goto clean_up;
	}
	else if (file_size < 1U)
	{
		FPUTS(T("Error: Input file is empty or an unsupported type!\n\n"), stderr);
		goto clean_up;
	}

	if (!(buffer = malloc((buffer_size = get_buffer_size(file_size)) * sizeof(uint8_t))))
	{
		FPUTS(T("Error: Failed to allocate the I/O buffer!\n\n"), stderr);
		goto clean_up;
	}

	FPUTS(T("Encrypting file contents, please be patient... "), stderr);
	fflush(stderr);

	uint64_t nonce;
	if (slunkcrypt_generate_nonce(&nonce) != SLUNKCRYPT_SUCCESS)
	{
		FPUTS(T("\n\nSlunkCrypt error: Failed to generate nonce!\n\n"), stderr);
		goto  clean_up;
	}

	init_slunk_param(&param, options);
	ctx = slunkcrypt_alloc_ext(nonce, (const uint8_t*)passphrase, strlen(passphrase), SLUNKCRYPT_ENCRYPT, &param);
	if (!ctx)
	{
		FPUTS(g_slunkcrypt_abort_flag ? T("\n\nProcess interrupted!\n\n") : T("\n\nSlunkCrypt error: Failed to initialize encryption!\n\n"), stderr);
		goto clean_up;
	}

	if (fwrite_ui64(nonce ^ MAGIC_NUMBER, file_out) < 1U)
	{
		FPUTS(T("\n\nI/O error: Failed to write nonce value!\n\n"), stderr);
		goto clean_up;
	}

	uint64_t bytes_read = 0U, clk_update = clock_read();
	const uint64_t update_interval = (uint64_t)(clock_freq() * 1.414);

	blake2s_t blake2s_state;
	blake2s_init(&blake2s_state);

	FPRINTF(stderr, T("%5.1f%% "), 0.0);
	fflush(stderr);

	while (bytes_read < file_size)
	{
		const uint64_t bytes_remaining = file_size - bytes_read;
		const size_t request_len = (bytes_remaining < buffer_size) ? ((size_t)bytes_remaining) : buffer_size;
		const size_t count = fread(buffer, sizeof(uint8_t), request_len, file_in);
		if (count > 0U)
		{
			blake2s_update(&blake2s_state, buffer, count);
			bytes_read += count;
			if ((status = slunkcrypt_inplace(ctx, buffer, count)) != SLUNKCRYPT_SUCCESS)
			{
				FPUTS((status == SLUNKCRYPT_ABORTED) ? T("\n\nProcess interrupted!\n\n") : T("\n\nSlunkCrypt error: Failed to encrypt data!\n\n"), stderr);
				goto clean_up;
			}
			if (fwrite(buffer, sizeof(uint8_t), count, file_out) < count)
			{
				FPUTS(T("\n\nI/O error: Failed to write encrypted data!\n\n"), stderr);
				goto clean_up;
			}
		}
		if (count < request_len)
		{
			break; /*EOF*/
		}
		UPDATE_PROGRESS_INDICATOR(clk_update, bytes_read, file_size);
	}

	if (ferror(file_in))
	{
		FPUTS(T("\n\nI/O error: Failed to read input data!\n\n"), stderr);
		goto clean_up;
	}

	if (bytes_read != file_size)
	{
		FPUTS(T("\n\nI/O error: Input file could not be fully read!\n\n"), stderr);
		goto clean_up;
	}

	const size_t padding = sizeof(uint64_t) - (file_size % sizeof(uint64_t));
	assert(padding && (padding <= sizeof(uint64_t)));
	if (slunkcrypt_random_bytes(buffer, padding) < padding)
	{
		FPUTS(T("\n\nSlunkCrypt error: Failed to generate random data!\n\n"), stderr);
		goto clean_up;
	}

	SET_LOWBITS(buffer[padding - 1U], padding - 1U);
	if ((status = slunkcrypt_inplace(ctx, buffer, padding)) != SLUNKCRYPT_SUCCESS)
	{
		FPUTS((status == SLUNKCRYPT_ABORTED) ? T("\n\nProcess interrupted!\n\n") : T("\n\nSlunkCrypt error: Failed to encrypt data!\n\n"), stderr);
		goto clean_up;
	}

	if (fwrite(buffer, sizeof(uint8_t), padding, file_out) < padding)
	{
		FPUTS(T("\n\nI/O error: Failed to write padding data!\n\n"), stderr);
		goto clean_up;
	}

	uint8_t checksum_buffer[sizeof(uint64_t)];
	store_ui64(checksum_buffer, blake2s_final(&blake2s_state));

	if ((status = slunkcrypt_inplace(ctx, checksum_buffer, sizeof(uint64_t))) != SLUNKCRYPT_SUCCESS)
	{
		FPUTS((status == SLUNKCRYPT_ABORTED) ? T("\n\nProcess interrupted!\n\n") : T("\n\nSlunkCrypt error: Failed to encrypt checksum!\n\n"), stderr);
		goto  clean_up;
	}

	if (fwrite(checksum_buffer, sizeof(uint8_t), sizeof(uint64_t), file_out) < sizeof(uint64_t))
	{
		FPUTS(T("\n\nI/O error: Failed to write the checksum!\n\n"), stderr);
		goto clean_up;
	}

	FPRINTF(stderr, T("\b\b\b\b\b\b\b%5.1f%%\n\n"), 100.0);

	result = EXIT_SUCCESS;

	FPUTS(T("All is done.\n\n"), stderr);
	fflush(stderr);

clean_up:

	SLUNKCRYPT_SAFE_FREE(ctx);

	if (file_out)
	{
		fclose(file_out);
		if ((result != EXIT_SUCCESS) && (!options->keep_incomplete))
		{
			if (REMOVE(output_path))
			{
				FPUTS(T("Warning: Failed to remove incomplete output file!\n\n"), stderr);
			}
		}
	}

	if (file_in)
	{
		fclose(file_in);
	}

	if (buffer)
	{
		slunkcrypt_bzero(buffer, buffer_size * sizeof(uint8_t));
		free(buffer);
	}

	slunkcrypt_bzero(checksum_buffer, sizeof(uint64_t));
	slunkcrypt_bzero(&blake2s_state, sizeof(blake2s_t));
	slunkcrypt_bzero(&nonce, sizeof(uint64_t));

	return result;
}

// ==========================================================================
// Decrypt
// ==========================================================================

int decrypt(const char *const passphrase, const CHR *const input_path, const CHR *const output_path, const crypt_options_t *const options)
{
	slunkcrypt_t ctx = SLUNKCRYPT_NULL;
	slunkparam_t param;
	FILE *file_in = NULL, *file_out = NULL;
	uint8_t *buffer = NULL;
	size_t buffer_size = BUFSIZ;
	int result = EXIT_FAILURE, status = -1;

	if (open_files(&file_in, &file_out, input_path, output_path) != EXIT_SUCCESS)
	{
		goto clean_up;
	}

	const uint64_t file_size = get_size(file_in);
	if (file_size == UINT64_MAX)
	{
		FPUTS(T("I/O error: Failed to determine size of input file!\n\n"), stderr);
		goto clean_up;
	}
	else if (file_size < (3U * sizeof(uint64_t)))
	{
		FPUTS(T("Error: Input file is too small! Truncated?\n\n"), stderr);
		goto clean_up;
	}
	else if ((file_size % sizeof(uint64_t)) != 0)
	{
		FPRINTF(stderr, T("Warning: File size is *not* an integer multiple of %u, ignoring excess bytes!\n\n"), (unsigned)sizeof(uint64_t));
	}

	if (!(buffer = malloc((buffer_size = get_buffer_size(file_size)) * sizeof(uint8_t))))
	{
		FPUTS(T("Error: Failed to allocate the I/O buffer!\n\n"), stderr);
		goto clean_up;
	}

	FPUTS(T("Decrypting file contents, please be patient... "), stderr);
	fflush(stderr);

	uint64_t nonce;
	if (fread_ui64(&nonce, file_in) < 1U)
	{
		FPUTS(T("\n\nI/O error: Failed to read nonce value!\n\n"), stderr);
		goto clean_up;
	}

	init_slunk_param(&param, options);
	ctx = slunkcrypt_alloc_ext(nonce ^ MAGIC_NUMBER, (const uint8_t*)passphrase, strlen(passphrase), SLUNKCRYPT_DECRYPT, &param);
	if (!ctx)
	{
		FPUTS(g_slunkcrypt_abort_flag ? T("\n\nProcess interrupted!\n\n") : T("\n\nSlunkCrypt error: Failed to initialize decryption!\n\n"), stderr);
		goto clean_up;
	}

	uint64_t bytes_read = sizeof(uint64_t), clk_update = clock_read();
	const uint64_t update_interval = (uint64_t)(clock_freq() * 1.414);
	const uint64_t read_limit = round_down(file_size, sizeof(uint64_t)) - (2U * sizeof(uint64_t));

	blake2s_t blake2s_state;
	blake2s_init(&blake2s_state);

	FPRINTF(stderr, T("%5.1f%% "), 0.0);
	fflush(stderr);

	while (bytes_read < read_limit)
	{
		const uint64_t bytes_remaining = read_limit - bytes_read;
		const size_t request_len = (bytes_remaining < buffer_size) ? ((size_t)bytes_remaining) : buffer_size;
		const size_t count = fread(buffer, sizeof(uint8_t), request_len, file_in);
		if (count > 0U)
		{
			bytes_read += count;
			if ((status = slunkcrypt_inplace(ctx, buffer, count)) != SLUNKCRYPT_SUCCESS)
			{
				FPUTS((status == SLUNKCRYPT_ABORTED) ? T("\n\nProcess interrupted!\n\n") : T("\n\nSlunkCrypt error: Failed to decrypt data!\n\n"), stderr);
				goto clean_up;
			}
			blake2s_update(&blake2s_state, buffer, count);
			if (fwrite(buffer, sizeof(uint8_t), count, file_out) < count)
			{
				FPUTS(T("\n\nI/O error: Failed to write decrypted data!\n\n"), stderr);
				goto clean_up;
			}
		}
		if (count < request_len)
		{
			break; /*EOF*/
		}
		UPDATE_PROGRESS_INDICATOR(clk_update, bytes_read, read_limit);
	}

	if (ferror(file_in))
	{
		FPUTS(T("\n\nI/O error: Failed to read input data!\n\n"), stderr);
		goto clean_up;
	}

	if (bytes_read != read_limit)
	{
		FPUTS(T("\n\nI/O error: Input file could not be fully read!\n\n"), stderr);
		goto clean_up;
	}

	if (fread(buffer, sizeof(uint8_t), sizeof(uint64_t), file_in) < sizeof(uint64_t))
	{
		FPUTS(T("\n\nI/O error: Failed to read final block!\n\n"), stderr);
		goto clean_up;
	}

	if ((status = slunkcrypt_inplace(ctx, buffer, sizeof(uint64_t))) != SLUNKCRYPT_SUCCESS)
	{
		FPUTS((status == SLUNKCRYPT_ABORTED) ? T("\n\nProcess interrupted!\n\n") : T("\n\nSlunkCrypt error: Failed to decrypt data!\n\n"), stderr);
		goto  clean_up;
	}

	const size_t padding = GET_LOWBITS(buffer[sizeof(uint64_t) - 1U]) + 1U;
	assert(padding && (padding <= sizeof(uint64_t)));
	if (padding != sizeof(uint64_t))
	{
		const size_t count = sizeof(uint64_t) - padding;
		if (fwrite(buffer, sizeof(uint8_t), count, file_out) < count)
		{
			FPUTS(T("\n\nI/O error: Failed to write decrypted data!\n\n"), stderr);
			goto clean_up;
		}
		blake2s_update(&blake2s_state, buffer, count);
	}

	const uint64_t checksum_actual = blake2s_final(&blake2s_state);

	uint8_t checksum_buffer[sizeof(uint64_t)];
	if (fread(checksum_buffer, sizeof(uint8_t), sizeof(uint64_t), file_in) < sizeof(uint64_t))
	{
		FPUTS(T("\n\nI/O error: Failed to read the checksum!\n\n"), stderr);
		goto clean_up;
	}

	if ((status = slunkcrypt_inplace(ctx, checksum_buffer, sizeof(uint64_t))) != SLUNKCRYPT_SUCCESS)
	{
		FPUTS((status == SLUNKCRYPT_ABORTED) ? T("\n\nProcess interrupted!\n\n") : T("\n\nSlunkCrypt error: Failed to decrypt checksum!\n\n"), stderr);
		goto clean_up;
	}

	FPRINTF(stderr, T("\b\b\b\b\b\b\b%5.1f%%\n\n"), 100.0);

	const uint64_t checksum_stored = load_ui64(checksum_buffer);
	if (checksum_actual != checksum_stored)
	{
		FPRINTF(stderr, T("Error: Checksum mismatch detected! [expected: 0x%016") T(PRIX64) T(", actual: 0x%016") T(PRIX64) T("]\n\n"), checksum_stored, checksum_actual);
		FPUTS(T("Wrong passphrase or corrupted file?\n\n"), stderr);
		goto clean_up;
	}

	result = EXIT_SUCCESS;

	FPUTS(T("Checksum is correct.\n\n"), stderr);
	fflush(stderr);

clean_up:

	SLUNKCRYPT_SAFE_FREE(ctx);

	if (file_out)
	{
		fclose(file_out);
		if ((result != EXIT_SUCCESS) && (!options->keep_incomplete))
		{
			if (REMOVE(output_path))
			{
				FPUTS(T("Warning: Failed to remove incomplete output file!\n\n"), stderr);
			}
		}
	}

	if (file_in)
	{
		fclose(file_in);
	}

	if (buffer)
	{
		slunkcrypt_bzero(buffer, buffer_size * sizeof(uint8_t));
		free(buffer);
	}

	slunkcrypt_bzero(checksum_buffer, sizeof(uint64_t));
	slunkcrypt_bzero(&blake2s_state, sizeof(blake2s_t));
	slunkcrypt_bzero(&nonce, sizeof(uint64_t));
	slunkcrypt_bzero((void*)&checksum_stored, sizeof(uint64_t));
	slunkcrypt_bzero((void*)&checksum_actual, sizeof(uint64_t));

	return result;
}
