//Emil Fihlman
//License TBD
//gcc -Wall -Werror -Wextra -O3 -flto -o sha256crypt sha256crypt.c

#include <stdtypes.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <termios.h>
#include <unistd.h>
#include <signal.h>
#include <sys/syscall.h>
#include "sha256.c"

void restoreEcho(void)
{
	struct termios terminal={0};
	tcgetattr(STDIN_FILENO, &terminal);
	terminal.c_lflag|=ECHO;
	tcsetattr(STDIN_FILENO, TCSANOW, &terminal);
	write(STDIN_FILENO, "\e[?25h", 6);
}

void signalHandler(s32 signal)
{
	s8 *message;
	switch(signal)
	{
		case SIGINT:
			message="\r\nGot SIGINT\r\n";
			write(2, (u8 *)message, strlen(message));
			restoreEcho();
			_Exit(-1);
		case SIGPIPE:
			message="\r\nGot SIGPIPE\r\n";
			break;
		default:
			message="\r\nGot some other SIGNAL\r\n";
	}
	write(2, (u8 *)message, strlen(message));
}

s32 getRandom(void *buffer, u64 length)
{
	return(syscall(SYS_getrandom, buffer, length, 0));
}

s32 main(s32 argc, s8 **argv)
{
	s32 rs32;
	u64 ru64, wrt;
	FILE *inputFile=NULL, *outputFile=NULL, *keyFile=NULL;
	u8 buffer[SHA256_DIGEST_SIZE], key[256], digest[SHA256_DIGEST_SIZE], mac[SHA256_DIGEST_SIZE], failFlag=0, keySize=0;
	s8 *action=NULL, *inputPath=NULL, *outputPath=NULL, *keyPath=NULL, *rs8p=NULL;
	struct sha256cryptcontext context;
	struct termios terminal;
	if(argc<4)
	{
		fprintf(stderr, "Usage: %s --encrypt|--decrypt --input-file input_path --output-file output_path --key-file key_path\r\n", argv[0]);
		return(0);
	}
	--argc;
	++argv;
	while(argc)
	{
		if(!strcmp(*argv, "--encrypt") || !strcmp(*argv, "--decrypt"))
		{
			action=*argv;
		}
		else if(argc>1)
		{
			if(!strcmp(*argv, "--input-file"))
			{
				++argv;
				--argc;
				inputPath=*argv;
			}
			else if(!strcmp(*argv, "--output-file"))
			{
				++argv;
				--argc;
				outputPath=*argv;
			}
			else if(!strcmp(*argv, "--key-file"))
			{
				++argv;
				--argc;
				keyPath=*argv;
			}
			else
			{
				goto unknown;
			}
		}
		else
		{
			unknown:
			failFlag=1;
			fprintf(stderr, "Unknown argument: %s\r\n", *argv);
		}
		--argc;
		++argv;
	}
	if(!action)
	{
		fprintf(stderr, "Action required\r\n");
		failFlag=1;
	}
	if(!inputPath && !keyPath)
	{
		fprintf(stderr, "At least one of --input-file and --key-file required\r\n");
		failFlag=1;
	}
	if(failFlag)
	{
		fprintf(stderr, "Failed\r\n");
		return(-1);
	}
	if(inputPath)
	{
		inputFile=fopen(inputPath, "r");
		if(!inputFile)
		{
			fprintf(stderr, "Couldn't open \"%s\" for input\r\n", inputPath);
			return(-1);
		}
	}
	else
	{
		inputFile=stdin;
	}
	if(outputPath)
	{
		outputFile=fopen(outputPath, "w");
		if(!outputFile)
		{
			fprintf(stderr, "Couldn't open \"%s\" for output\r\n", outputPath);
			return(-1);
		}
	}
	else
	{
		outputFile=stdout;
	}
	if(keyPath)
	{
		keyFile=fopen(keyPath, "r");
		if(!keyFile)
		{
			fprintf(stderr, "Couldn't open \"%s\" for key\r\n", keyPath);
			return(-1);
		}
		ru64=fread(key, 1, 256, keyFile);
		if(ru64==256)
		{
			rs32=fgetc(keyFile);
			if(rs32!=EOF)
			{
				fprintf(stderr, "Key file \"%s\" larger than 256 bytes, truncating key\r\n", keyPath);
			}
		}
		else if(ru64==0)
		{
			fprintf(stderr, "Couldn't read key from \"%s\"\r\n", keyPath);
			memset(key, 0, 256);
			return(-1);
		}
		fclose(keyFile);
		keySize=ru64;
	}
	else
	{
		if(signal(SIGINT, signalHandler)==SIG_ERR)
		{
			return(-1);
		}
		fprintf(stderr, "\e[?25lEnter password:");
		fflush(stderr);
		tcgetattr(STDIN_FILENO, &terminal);
		terminal.c_lflag&=~ECHO;
		tcsetattr(STDIN_FILENO, TCSANOW, &terminal);
		rs32=atexit(restoreEcho);
		if(rs32)
		{
			return(-1);
		}
		rs8p=fgets((s8 *)key, 256, stdin);
		if(!rs8p)
		{
			return(-1);
		}
		if(strchr((s8 *)key, '\r'))
		{
			*strchr((s8 *)key, '\r')=0;
		}
		else if(strchr((s8 *)key, '\n'))
		{
			*strchr((s8 *)key, '\n')=0;
		}
		keySize=strlen((s8 *)key);
		fprintf(stderr, "\r\n");
	}
	if(!strcmp(action, "--encrypt"))
	{
		rs32=getRandom(buffer, SHA256_DIGEST_SIZE);
		if(rs32!=SHA256_DIGEST_SIZE)
		{
			fprintf(stderr, "Vector initialisation failed\r\n");
			memset(key, 0, 256);
			return(-1);
		}
		ru64=fwrite(buffer, 1, SHA256_DIGEST_SIZE, outputFile);
		if(ru64!=SHA256_DIGEST_SIZE)
		{
			fprintf(stderr, "Couldn't write vector\r\n");
			memset(key, 0, 256);
			return(-1);
		}
		sha256cryptinit(&context, key, keySize, buffer, SHA256_DIGEST_SIZE);
		memset(key, 0, 256);
		while(1)
		{
			ru64=fread(buffer, 1, SHA256_DIGEST_SIZE, inputFile);
			wrt=ru64;
			sha256cryptupdate(&context, buffer, ru64, 1);
			ru64=fwrite(buffer, 1, ru64, outputFile);
			if(ru64!=wrt)
			{
				fprintf(stderr, "Couldn't write data\r\n");
				return(-1);
			}
			if(ru64<SHA256_DIGEST_SIZE)
			{
				break;
			}
		}
		sha256cryptfinal(&context, mac);
		ru64=fwrite(mac, 1, SHA256_DIGEST_SIZE, outputFile);
		if(ru64!=SHA256_DIGEST_SIZE)
		{
			fprintf(stderr, "Couldn't write mac\r\n");
			return(-1);
		}
	}
	else if(!strcmp(action, "--decrypt"))
	{
		ru64=fread(buffer, 1, SHA256_DIGEST_SIZE, inputFile);
		if(ru64!=SHA256_DIGEST_SIZE)
		{
			fprintf(stderr, "Couldn't read vector\r\n");
			memset(key, 0, 256);
			return(-1);
		}
		sha256cryptinit(&context, key, keySize, buffer, SHA256_DIGEST_SIZE);
		memset(key, 0, 256);
		ru64=fread(buffer, 1, SHA256_DIGEST_SIZE, inputFile);
		if(ru64!=SHA256_DIGEST_SIZE)
		{
			fprintf(stderr, "Couldn't read mac\r\n");
			return(-1);
		}
		while(1)
		{
			ru64=fread(mac, 1, SHA256_DIGEST_SIZE, inputFile);
			sha256cryptupdate(&context, buffer, ru64, 0);
			wrt=fwrite(buffer, 1, ru64, outputFile);
			if(wrt!=ru64)
			{
				fprintf(stderr, "Couldn't write data\r\n");
				return(-1);
			}
			if(ru64<SHA256_DIGEST_SIZE)
			{
				memmove(mac+(SHA256_DIGEST_SIZE-ru64), mac, ru64);
				memmove(mac, buffer+ru64, SHA256_DIGEST_SIZE-ru64);
				break;
			}
			else
			{
				memmove(buffer, mac, SHA256_DIGEST_SIZE);
			}
		}
		sha256cryptfinal(&context, digest);
		ru64=sha256compare(mac, digest);
		if(ru64)
		{
			fprintf(stderr, "Mac mismatch, broken/tampered data!\r\n");
			return(-1);
		}
	}
	else
	{
		fprintf(stderr, "Action error\r\n");
		return(-1);
	}
	fclose(outputFile);
	fclose(inputFile);
	return(0);
}