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

#define _CRT_SECURE_NO_WARNINGS 1

#include <mcrypt.h>
#include "utils.h"
#include "crc.h"
#include "test.h"
#include <string.h>
#include <time.h>
#include <inttypes.h>
#include <ctype.h>
#include <signal.h>

#define BUFF_SIZE 4096U

static const CHR *const ENVV_PASSWD_NAME = T("MCRYPT_PASSWD");

static char* read_passphrase(const CHR* const file_name)
{
	static const size_t buff_size = 512U;
	char *buffer = (char*) malloc(buff_size * sizeof(char));
	if (!buffer)
	{
		return NULL;
	}

	const int use_stdin = (STRICMP(file_name, T("-")) == 0);
	FILE *const file = use_stdin ? stdin : FOPEN(file_name, T("rb"));
	if (!file)
	{
		free(buffer);
		return NULL;
	}

	do
	{
		if (!fgets(buffer, (int)buff_size, file))
		{
			free(buffer);
			buffer = NULL;
			goto finish;
		}
		size_t length = strlen(buffer);
		while ((length > 0U) && ((buffer[length - 1U] == '\r') || (buffer[length - 1U] == '\n')))
		{
			buffer[--length] = '\0';
		}
	}
	while (!buffer[0U]);

finish:

	if ((!use_stdin) && file)
	{
		fclose(file);
	}

	return buffer;
}

static int weak_passphrase(const char *str)
{
	int flags[4U] = { 0, 0, 0, 0 };
	while (*str)
	{
		const CHR c = *str++;
		if (isalpha(c))
		{
			flags[isupper(c) ? 0U : 1U] = 1;
		}
		else 
		{
			flags[isdigit(c) ? 2U : 3U] = 1;
		}
	}
	const int strong = flags[0U] && flags[1U] && flags[2U] && flags[3U];
	return !strong;
}

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

	*file_out = FOPEN(output_path, T("wb"));
	if (!(*file_out))
	{
		FPUTS(T("Error: Failed to open output file for writing!\n\n"), stderr);
		fclose(*file_out);
		return 1;
	}

	return 0;
}

static int encrypt(const char* const passphrase, const CHR* const input_path, const CHR* const output_path)
{
	mcrypt_t ctx = MCRYPT_NULL;
	FILE * file_in = NULL, * file_out = NULL;
	int result = 1;

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

	const uint64_t file_size = get_file_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;
	}

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

	uint64_t seed;
	if (mcrypt_generate_seed(&seed) != MCRYPT_SUCCESS)
	{
		FPUTS(T("\n\nMCrypt error: Failed to generate seed!\n\n"), stderr);
		goto  clean_up;
	}

	if (fwrite(&seed, sizeof(uint64_t), 1U, file_out) < 1U)
	{
		FPUTS(T("\n\nI/O error: Failed to write seed value!\n\n"), stderr);
		goto clean_up;
	}

	ctx = mcrypt_alloc(seed, (const uint8_t*)passphrase, strlen(passphrase));
	if (!ctx)
	{
		FPUTS(g_macrypt_abort_flag ? T("\n\nProcess interrupted!\n\n") : T("\n\nMCrypt error: Failed to initialize encryption!\n\n"), stderr);
		goto clean_up;
	}

	clock_t clk_now, clk_update = clock();
	uint64_t crc_actual = CRC_INITIALIZER, bytes_read = 0U;
	uint8_t buffer[BUFF_SIZE];

	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 < BUFF_SIZE) ? ((size_t)bytes_remaining) : BUFF_SIZE;
		const size_t count = fread(buffer, sizeof(uint8_t), request_len, file_in);
		if (count > 0U)
		{
			crc_actual = crc64_update(crc_actual, buffer, count);
			bytes_read += count;
			const int status = mcrypt_encrypt_inplace(ctx, buffer, count);
			if (status != MCRYPT_SUCCESS)
			{
				FPUTS((status == MCRYPT_ABORTED) ? T("\n\nProcess interrupted!\n\n") : T("\n\nMCrypt 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*/
		}
		if ((clk_now = clock()) - clk_update > CLOCKS_PER_SEC)
		{
			FPRINTF(stderr, T("\b\b\b\b\b\b\b%5.1f%% "), (bytes_read / ((double)file_size)) * 100.0);
			fflush(stderr);
			clk_update = clk_now;
		}
	}

	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;
	}

	crc_actual = crc64_finish(crc_actual);
	FPRINTF(stderr, T("\b\b\b\b\b\b%5.1f%%\n\n"), 100.0);
	fflush(stderr);

	if (fwrite(&crc_actual, sizeof(uint64_t), 1U, file_out) < 1U)
	{
		FPUTS(T("I/O error: Failed to write CRC checksum!\n\n"), stderr);
		goto clean_up;
	}

	result = 0;

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

clean_up:

	if (ctx)
	{
		mcrypt_free(ctx);
	}

	if (file_out)
	{
		fclose(file_out);
	}

	if (file_in)
	{
		fclose(file_in);
	}

	return result;
}

static int decrypt(const char* const passphrase, const CHR* const input_path, const CHR* const output_path)
{
	mcrypt_t ctx = MCRYPT_NULL;
	FILE *file_in = NULL, *file_out = NULL;
	int result = 1;

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

	const uint64_t file_size = get_file_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 < 16U)
	{
		FPUTS(T("Error: Input file is too small! Truncated?\n\n"), stderr);
		goto clean_up;
	}

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

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

	ctx = mcrypt_alloc(seed, (const uint8_t*)passphrase, strlen(passphrase));
	if (!ctx)
	{
		FPUTS(g_macrypt_abort_flag ? T("\n\nProcess interrupted!\n\n") : T("\n\nMCrypt error: Failed to initialize decryption!\n\n"), stderr);
		goto clean_up;
	}
	
	clock_t clk_now, clk_update = clock();
	uint64_t crc_actual = CRC_INITIALIZER, bytes_read = sizeof(uint64_t);
	uint8_t buffer[BUFF_SIZE];
	const uint64_t read_limit = file_size - sizeof(uint64_t);

	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 < BUFF_SIZE) ? ((size_t)bytes_remaining) : BUFF_SIZE;
		const size_t count = fread(buffer, sizeof(uint8_t), request_len, file_in);
		if (count > 0U)
		{
			bytes_read += count;
			const int status = mcrypt_decrypt_inplace(ctx, buffer, count);
			if (status != MCRYPT_SUCCESS)
			{
				FPUTS((status == MCRYPT_ABORTED) ? T("\n\nProcess interrupted!\n\n") : T("\n\nMCrypt error: Failed to decrypt data!\n\n"), stderr);
				goto clean_up;
			}
			crc_actual = crc64_update(crc_actual, buffer, count);
			if (fwrite(buffer, sizeof(uint8_t), count, file_out) < count)
			{
				FPUTS(T("failed!\n\nI/O error: Failed to write decrypted data!\n\n"), stderr);
				goto clean_up;
			}
		}
		if (count < request_len)
		{
			break; /*EOF*/
		}
		if ((clk_now = clock()) - clk_update > CLOCKS_PER_SEC)
		{
			FPRINTF(stderr, T("\b\b\b\b\b\b\b%5.1f%% "), (bytes_read / ((double)read_limit)) * 100.0);
			fflush(stderr);
			clk_update = clk_now;
		}
	}

	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;
	}

	crc_actual = crc64_finish(crc_actual);
	FPRINTF(stderr, T("\b\b\b\b\b\b%5.1f%%\n\n"), 100.0);
	fflush(stderr);

	uint64_t crc_expected;
	if (fread(&crc_expected, sizeof(uint64_t), 1U, file_in) < 1U)
	{
		FPUTS(T("I/O error: Failed to read CRC checksum!\n\n"), stderr);
		goto clean_up;
	}

	if (crc_actual != crc_expected)
	{
		FPRINTF(stderr, T("CRC error: Checksum mismatch detected! [expected: 0x%016") T(PRIX64) T(", actual: 0x%016") T(PRIX64) T("]\n\n"), crc_actual, crc_expected);
		FPUTS(T("Wrong passphrase or corrupted file?\n\n"), stderr);
		goto clean_up;
	}

	result = 0;

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

clean_up:

	if (ctx)
	{
		mcrypt_free(ctx);
	}

	if (file_out)
	{
		fclose(file_out);
	}

	if (file_in)
	{
		fclose(file_in);
	}

	return result;
}

static int run_test(const char *const message)
{
	static const char* const passphrase = "OrpheanBeh0lderScry!Doubt";

	const size_t length = strlen(message) + 1U;
	int status, result = 1;
	mcrypt_t ctx = MCRYPT_NULL;

	uint64_t seed;
	if (mcrypt_generate_seed(&seed) != MCRYPT_SUCCESS)
	{
		FPUTS(T("\n\nWhoops: Failed to generate seed!\n\n"), stderr);
		return 1;
	}

	char* const text_temp = strdup(message);
	if (!text_temp)
	{
		FPUTS(T("\n\nWhoops: Failed to allocate text buffer!\n\n"), stderr);
		goto clean_up;
	}

	ctx = mcrypt_alloc(seed, (const uint8_t*)passphrase, strlen(passphrase));
	if (!ctx)
	{
		FPUTS(g_macrypt_abort_flag ? T("\n\nProcess interrupted!\n\n") : T("\n\nWhoops: Failed to initialize encoder!\n\n"), stderr);
		goto clean_up;
	}

	status = mcrypt_encrypt_inplace(ctx, (uint8_t*)text_temp, length);
	if (status != MCRYPT_SUCCESS)
	{
		FPUTS((status == MCRYPT_ABORTED) ? T("\n\nProcess interrupted!\n\n") : T("\n\nWhoops: Failed to encrypt the message!\n\n"), stderr);
		goto clean_up;
	}

	if (strncmp(message, text_temp, length) == 0)
	{
		FPUTS(T("\n\nWhoops: Encrypted message equals the original message!\n\n"), stderr);
		goto clean_up;
	}

	status = mcrypt_reset(ctx, seed, (const uint8_t*)passphrase, strlen(passphrase));
	if (status != MCRYPT_SUCCESS)
	{
		FPUTS((status == MCRYPT_ABORTED) ? T("\n\nProcess interrupted!\n\n") : T("\n\nWhoops: Failed to initialize decoder!\n\n"), stderr);
		goto clean_up;
	}

	status = mcrypt_decrypt_inplace(ctx, (uint8_t*)text_temp, length);
	if (status != MCRYPT_SUCCESS)
	{
		FPUTS((status == MCRYPT_ABORTED) ? T("\n\nProcess interrupted!\n\n") : T("\n\nWhoops: Failed to decrypt the message!\n\n"), stderr);
		goto clean_up;
	}

	if (strncmp(message, text_temp, length) != 0)
	{
		FPUTS(T("\n\nWhoops: Decrypted message does *not* match the original message!\n\n"), stderr);
		goto clean_up;
	}

	result = 0;

clean_up:

	if (ctx)
	{
		mcrypt_free(ctx);
	}

	if (text_temp)
	{
		mcrypt_bzero(text_temp, strlen(text_temp));
		free(text_temp);
	}

	return result;
}

static int self_test(void)
{
	static const size_t total = 32U;
	const char* const test_data[] = { TEST_DATA_0, TEST_DATA_1, TEST_DATA_2, TEST_DATA_3, NULL };

	size_t completed = 0U;
	FPRINTF(stderr, T("Self-test is running, please be patient... %2u/%2u "), (unsigned int)completed, (unsigned int)total);
	fflush(stderr);

	for (size_t i = 0U; i < 8U; ++i)
	{
		for (size_t j = 0U; test_data[j]; ++j)
		{
			FPRINTF(stderr, T("\b\b\b\b\b\b%2u/%2u "), (unsigned int)++completed, (unsigned int)total);
			fflush(stderr);
			if (run_test(test_data[j]))
			{
				return 1;
			}
		}
	}

	FPRINTF(stderr, T("\b\b\b\b\b\b%2u/%2u\n\nCompleted successfully.\n\n"), (unsigned int)total, (unsigned int)total);
	fflush(stderr);

	return 0;
}

static void sigint_handler(const int sig)
{
	if (sig == SIGINT)
	{
		g_macrypt_abort_flag = 1;
		signal(SIGINT, sigint_handler);
	}
}

int MAIN(int argc, CHR* argv[])
{
	init_terminal();
	signal(SIGINT, sigint_handler);

	FPRINTF(stderr, T("MCrypt Utility (%") T(PRIstr) T("-%") T(PRIstr) T("), by LoRd_MuldeR <MuldeR2@GMX.de>\n"), OS_TYPE, CPU_ARCH);
	FPRINTF(stderr, T("Using libMCrypt v%") T(PRIstr) T(" [%") T(PRIstr) T("]\n\n"), LIBMCRYPT_VERSION, LIBMCRYPT_BUILDNO);

	if (argc > 1)
	{
		if ((!STRICMP(argv[1U], T("-v"))) || (!STRICMP(argv[1U], T("--version"))))
		{
			return 0; /*exit right now*/
		}
		if ((!STRICMP(argv[1U], T("-h"))) || (!STRICMP(argv[1U], T("/?"))) || (!STRICMP(argv[1U], T("--help"))))
		{
			const CHR* const program = get_file_name(argv[0U]);
			FPUTS(T("====================================================================\n"), stderr);
			FPUTS(T("This software has been released under the CC0 1.0 Universal license:\n"), stderr);
			FPUTS(T("https://creativecommons.org/publicdomain/zero/1.0/legalcode\n"), stderr);
			FPUTS(T("====================================================================\n\n"), stderr);
			FPUTS(T("Synopsis:\n"), stderr);
			FPRINTF(stderr, T("  %") T(PRISTR) T(" --encrypt [[@][:]<passphrase>] <input.txt> <output.enc>\n"), program);
			FPRINTF(stderr, T("  %") T(PRISTR) T(" --decrypt [[@][:]<passphrase>] <input.enc> <output.txt>\n\n"), program);
			FPUTS(T("Remarks:\n"), stderr);
			FPUTS(T("- If <passphrase> is prefixed with a '@' character, then it specifies the file\n"), stderr);
			FPUTS(T("  to read the passphrase from; only the first line in that file is used!\n"), stderr);
			FPUTS(T("- If <passphrase> is prefixed with a ':' character, then the leading character\n"), stderr);
			FPUTS(T("  is skipped and the remainder of the argument is used as passphrase.\n"), stderr);
			FPUTS(T("- If the argument <passphrase> is *not* present, then the environment variable\n"), stderr);
			FPRINTF(stderr, T("  \"%") T(PRISTR) T("\" must be set; it specifies the passphrase to be used.\n"), ENVV_PASSWD_NAME);
			FPUTS(T("- Specify \"@-\" in order to read the passphrase from the standard input stream!\n\n"), stderr);
			return 0;
		}
		if ((!STRICMP(argv[1U], T("-t"))) || (!STRICMP(argv[1U], T("--self-test"))))
		{
			return self_test(); /*only self-test!*/
		}
	}

	if (argc < 4)
	{
		const CHR* const program = get_file_name(argv[0U]);
		FPRINTF(stderr, T("Error: Nothing to do. Please type '%") T(PRISTR) T(" --help' for details!\n\n"), program);
		return 1;
	}
	else if (argc > 5)
	{
		FPUTS(T("Warning: Excess command-line argument is ignored!\n\n"), stderr);
	}

	const CHR *const command = argv[1U], *const passphrase = (argc > 4) ? argv[2U] : GETENV(ENVV_PASSWD_NAME);
	const CHR *const input_file = argv[(argc > 4) ? 3U : 2U], *const output_file = argv[(argc > 4) ? 4U : 3U];

	if ((!passphrase) || (!passphrase[0U]) || (((passphrase[0U] == T('@')) || (passphrase[0U] == T(':'))) && (!passphrase[1U])))
	{
		FPUTS(T("Error: The passphrase (file name) must not be empty!\n\n"), stderr);
		return 1;
	}

	if ((!input_file[0U]) || (!output_file[0U]))
	{
		FPUTS(T("Error: The input file and/or output file must not be empty!\n\n"), stderr);
		return 1;
	}

	char *const passphrase_buffer = (passphrase[0U] == T('@')) ? read_passphrase(passphrase + 1U) : CHR_to_utf8((passphrase[0U] == T(':')) ? (passphrase + 1U) : passphrase);
	if (!passphrase_buffer)
	{
		FPUTS(T("Error: Failed to read the passphrase file!\n\n"), stderr);
		return 1;
	}

	mcrypt_bzero((CHR*)passphrase, STRLEN(passphrase) * sizeof(CHR));

	if (strlen(passphrase_buffer) < 12U)
	{
		FPUTS(T("Warning: Using a *short* passphrase; a length of 12 characters or more is recommended!\n\n"), stderr);
	}
	else if (weak_passphrase(passphrase_buffer))
	{
		FPUTS(T("Warning: Using a *weak* passphrase; a mix of upper-case letters, lower-case letters, digits and 'special' characters is recommended!\n\n"), stderr);
	}

	const clock_t clk_start = clock();
	int result = -1;

	if ((!STRICMP(command, T("-e"))) || (!STRICMP(command, T("--encrypt"))))
	{
		result = encrypt(passphrase_buffer, input_file, output_file);
	}
	else if ((!STRICMP(command, T("-d"))) || (!STRICMP(command, T("--decrypt"))))
	{
		result = decrypt(passphrase_buffer, input_file, output_file);
	}
	else
	{
		FPRINTF(stderr, T("Error: Command \"%") T(PRISTR) T("\" is unknown!\n\n"), command);
		goto exiting;
	}

	if (!g_macrypt_abort_flag)
	{
		FPUTS(T("--------\n\n"), stderr);
		fflush(stderr);
		const clock_t clk_end = clock();
		FPRINTF(stderr, T("Operation completed after %.1f seconds.\n\n"), (clk_end - clk_start) / ((double)CLOCKS_PER_SEC));
	}

exiting:
	
	if (passphrase_buffer)
	{
		mcrypt_bzero(passphrase_buffer, strlen(passphrase_buffer));
		free(passphrase_buffer);
	}

	return result;
}