/******************************************************************************/
/* 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 <string.h>
#include <time.h>
#include <inttypes.h>
#include <ctype.h>

static void print_string(const char *const text, const size_t length)
{
	fputc('"', stderr);
	for (size_t i = 0; i < length; ++i)
	{
		fputc((text[i] >= 0x20) ? text[i] : '?', stderr);
	}
	fputs("\"\n\n", stderr);
}

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

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

	return 0;
}

static int weak_passphrase(const CHR *str)
{
	int flags[4U] = { 0, 0, 0, 0 };
	while (*str)
	{
		const CHR c = *str++;
		if      (ISUPPER(c)) flags[0U] = 1;
		else if (ISLOWER(c)) flags[1U] = 1;
		else if (ISDIGIT(c)) flags[2U] = 1;
		else if (ISGRAPH(c)) flags[3U] = 1;
	}
	const int strong = flags[0U] && flags[1U] && flags[2U] && flags[3U];
	return !strong;
}

static int encrypt(const char* const passphrase, const CHR* const input, const CHR* const output)
{
	mcrypt_t ctx = NULL;
	FILE *fin = NULL, *fout = NULL;
	int result = -1;

	if (open_files(&fin, &fout, input, output) != 0)
	{
		goto clean_up;;
	}

	const uint64_t file_size = get_file_size(fin);
	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 < 1LL)
	{
		FPUTS(T("Error: Input file is empty!\n\n"), stderr);
		goto clean_up;
	}

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

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

	ctx = mcrypt_alloc(seed, passphrase);
	if (!ctx)
	{
		FPUTS(T("MCrypt error: Failed to initialize encryption!\n\n"), stderr);
		goto clean_up;
	}

	FPRINTF(stderr, T("Encrypting file contents, please be patient... %5.1f%%"), 0.0);

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

	while ((!feof(fin)) && (bytes_read < file_size))
	{
		const size_t count = fread(buffer, sizeof(uint8_t), 1024U, fin);
		if (ferror(fin))
		{
			FPUTS(T("\n\nI/O error: Failed to read input data!\n\n"), stderr);
			goto clean_up;
		}
		if (count > 0U)
		{
			crc_actual = crc64_update(crc_actual, buffer, count);
			bytes_read += count;
			if (mcrypt_enc_process_inplace(ctx, buffer, count) != 0)
			{
				FPUTS(T("\n\nMCrypt error: Failed to encrypt data!\n\n"), stderr);
				goto  clean_up;
			}
			if (fwrite(buffer, sizeof(uint8_t), count, fout) < count)
			{
				FPUTS(T("\n\nI/O error: Failed to write encrypted data!\n\n"), stderr);
				goto clean_up;
			}
		}
		if ((clk_now = clock()) - clk_update > CLOCKS_PER_SEC)
		{
			FPRINTF(stderr, T("\b\b\b\b\b\b%5.1f%%"), (bytes_read / ((double)file_size)) * 100.0);
			fflush(stderr);
			clk_update = clk_now;
		}
	}

	crc_actual = crc64_finish(crc_actual);

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

	FPRINTF(stderr, T("\b\b\b\b\b\b%5.1f%%\n\n"), 100.0);
	
	if (fwrite(&crc_actual, sizeof(uint64_t), 1U, fout) < 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 (fout)
	{
		fclose(fout);
	}

	if (fin)
	{
		fclose(fin);
	}

	return result;
}

static int decrypt(const char* const passphrase, const CHR* const input, const CHR* const output)
{
	mcrypt_t ctx = NULL;
	FILE *fin = NULL, *fout = NULL;
	int result = -1;

	if (open_files(&fin, &fout, input, output) != 0)
	{
		goto clean_up;
	}

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

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

	ctx = mcrypt_alloc(seed, passphrase);
	if (!ctx)
	{
		FPUTS(T("MCrypt error: Failed to initialize decryption!\n\n"), stderr);
		goto clean_up;
	}

	FPRINTF(stderr, T("Decrypting file contents, please be patient... %5.1f%%"), 0.0);
	
	clock_t clk_now, clk_update = clock();
	uint64_t crc_actual = CRC_INITIALIZER, bytes_read = sizeof(uint64_t);
	uint8_t buffer[1024U];
	const uint64_t read_limit = file_size - sizeof(uint64_t);

	while ((!feof(fin)) && (bytes_read < read_limit))
	{
		const uint64_t bytes_remaining = read_limit - bytes_read;
		const size_t read_len = (bytes_remaining < 1024U) ? ((size_t)bytes_remaining) : 1024U;
		const size_t count = fread(buffer, sizeof(uint8_t), read_len, fin);
		if (ferror(fin))
		{
			FPUTS(T("\n\nI/O error: Failed to read encrypted data!\n\n"), stderr);
			goto clean_up;
		}
		if (count > 0U)
		{
			bytes_read += count;
			if (mcrypt_dec_process_inplace(ctx, buffer, count) != 0)
			{
				FPUTS(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, fout) < count)
			{
				FPUTS(T("failed!\n\nI/O error: Failed to write decrypted data!\n\n"), stderr);
				goto clean_up;
			}
		}
		if ((clk_now = clock()) - clk_update > CLOCKS_PER_SEC)
		{
			FPRINTF(stderr, T("\b\b\b\b\b\b%5.1f%%"), (bytes_read / ((double)read_limit)) * 100.0);
			fflush(stderr);
			clk_update = clk_now;
		}
	}

	crc_actual = crc64_finish(crc_actual);

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

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

	uint64_t crc_expected;
	if (fread(&crc_expected, sizeof(uint64_t), 1U, fin) < 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 (fout)
	{
		fclose(fout);
	}

	if (fin)
	{
		fclose(fin);
	}

	return result;
}

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

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

	const int help_requested = (argc > 1) && ((!STRICMP(argv[1U], T("/?"))) || (!STRICMP(argv[1U], T("--help"))) || (!STRICMP(argv[1U], T("--version"))));
	if ((argc < 5) || help_requested)
	{
		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);
		if (!help_requested)
		{
			FPUTS(T("Error: Required argument is missing!\n\n"), stderr);
		}
		FPUTS(T("Usage:\n"), stderr);
		FPRINTF(stderr, T("  %") T(PRISTR) T(" --encrypt <passphrase> <input.txt> <output.enc>\n"),   argv[0U]);
		FPRINTF(stderr, T("  %") T(PRISTR) T(" --decrypt <passphrase> <input.enc> <output.txt>\n\n"), argv[0U]);
		return help_requested ? 0 : 1;
	}

	const CHR *const command = argv[1U], *const passphrase = argv[2U], *const input_file = argv[3U], *const output_file = argv[4U];
	if (!passphrase[0U])
	{
		FPUTS(T("Error: The passphrase must not be empty!\n\n"), stderr);
		return 1;
	}
	else
	{
		if (STRLEN(passphrase) < 12U)
		{
			FPUTS(T("Warning: Using a *short* passphrase; a length of 12 characters or more is recommended!\n\n"), stderr);
		}
		if (weak_passphrase(passphrase))
		{
			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);
		}
	}

	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_utf8 = CHR_to_utf8(passphrase);
	if (!passphrase_utf8)
	{
		FPUTS(T("Error: Failed to convert passphrase to the UTF-8 format!\n\n"), stderr);
		return 1;
	}

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

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

	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_utf8)
	{
		mcrypt_bzero(passphrase_utf8, strlen(passphrase_utf8));
		free(passphrase_utf8);
	}

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