#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdbool.h>
#include <unistd.h>
#include <limits.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sodium.h>
#include <pawstd.h>
#define KEY_SIZE crypto_secretbox_KEYBYTES
#define SALT_SIZE crypto_pwhash_SALTBYTES


int encryptMessage(const char* inputFile, const char* outputFile, const char* password) {
	// Open files
	FILE *inFile = fopen(inputFile,"rb");
	FILE *outFile = fopen(outputFile,"wb");
	if (inFile == NULL || outFile == NULL) {
		fprintf(stderr,ERR"Could not open files!\n");
		return 1;
	}
	// Variables
	unsigned char salt[SALT_SIZE];
	unsigned char nonce[crypto_secretbox_NONCEBYTES];
	unsigned char key[KEY_SIZE];

	// Read input file

	fprintf(stderr,NOR"Reading input filesize...\n");
	fseek(inFile,0,SEEK_END);
	size_t inLen = ftell(inFile);
	fseek(inFile,0,SEEK_SET);
	fprintf(stderr,OK"File is %d bytes.\n",inLen);
	unsigned char decrypted[inLen];

	// Generate salt
	randombytes_buf(salt, SALT_SIZE);
	fprintf(stderr,NOR"Beginning key generation...\n");
	
	// Key generation from password
	if (crypto_pwhash(
				key,KEY_SIZE,
				password,strlen(password),
				salt,crypto_pwhash_OPSLIMIT_INTERACTIVE,
				crypto_pwhash_MEMLIMIT_INTERACTIVE,
				crypto_pwhash_ALG_DEFAULT) != 0) {
		fprintf(stderr,ERR"Out of memory! Key could not be generated.\n");
		fclose(inFile);
		fclose(outFile);
		return 1;
	}

	fprintf(stderr,OK"Key generation successful!\n");
	fprintf(stderr,NOR"Reading file...\n");

	// Read the input file
	if(fread(decrypted,1,inLen,inFile) != inLen) {
		fprintf(stderr,ERR"Error reading from file.\n");
		fclose(inFile);
		fclose(outFile);
		return 1;
	}
	fprintf(stderr,OK"File read successfully!\n");
	fprintf(stderr,NOR"Beginning encryption...\n");

	// Generate a nonce, write headers and encrypted message
	randombytes_buf(nonce,sizeof(nonce));

	if (fwrite("gneurshk",1,strlen("gneurshk"),outFile) != strlen("gneurshk")) {
		fprintf(stderr,ERR"Error writing header to file.\n");
		fclose(inFile);
		fclose(outFile);
		return 1;
	}

	if (fwrite(salt,1,sizeof(salt),outFile) != sizeof(salt)) {
		fprintf(stderr,ERR"Error writing salt to file.\n");
		fclose(inFile);
		fclose(outFile);
		return 1;
	}

	if (fwrite(nonce,1,sizeof(nonce),outFile) != sizeof(nonce)) {
		fprintf(stderr,ERR"Error writing nonce to file.\n");
		fclose(inFile);
		fclose(outFile);
		return 1;
	}

	fprintf(stderr,OK"Headers written!\n");
	fprintf(stderr,NOR"Writing encrypted data...\n");

	unsigned char encrypted[inLen + crypto_secretbox_MACBYTES];

	crypto_secretbox_easy(encrypted, decrypted, inLen, nonce, key);

	if(fwrite(encrypted,1,sizeof(encrypted),outFile) != sizeof(encrypted)) {
		fprintf(stderr,ERR"Error writing encrypted data to file.\n");
		fclose(inFile);
		fclose(outFile);
		return 1;
	}


	fclose(inFile);
	fclose(outFile);

fprintf(stderr,"\n"OK"File encrypted!\n");

	return 0;
}

int decryptMessage(const char* inputFile, const char* outputFile, const char* password) {
	// Open files
	FILE *inFile = fopen(inputFile,"rb");
	FILE *outFile = fopen(outputFile,"wb");
	if (inFile == NULL || outFile == NULL) {
		fprintf(stderr,ERR"Could not open files.\n");
		return 1;
	}

	// Variables
	unsigned char salt[SALT_SIZE];
	unsigned char nonce[crypto_secretbox_NONCEBYTES];
	unsigned char key[KEY_SIZE];

	// Get the size of the file and assume the length of the content
	fseek(inFile,0,SEEK_END);
	size_t fileLen = ftell(inFile);
	fseek(inFile,0,SEEK_SET);

	size_t encLen = fileLen - sizeof(salt) - sizeof(nonce) - 8;

	fprintf(stderr,DBG"Size of encrypted file: %d bytes.\n",fileLen);
	fprintf(stderr,DBG"This means the encrypted content is %d bytes.\n",encLen);

	unsigned char encrypted[encLen];

	// Verify file
	unsigned char header[8];
	if (fread(header,1,sizeof(header),inFile) != sizeof(header)) {
		fprintf(stderr,ERR"Error reading header from file.\n");
		fclose(inFile);
		fclose(outFile);
		return 1;
	}
	
	if (memcmp(header,"gneurshk",sizeof(header)) != 0) {
		fprintf(stderr,ERR"Invalid header.\n");
		fprintf(stderr,ERR"Make sure the input file was signed by this program!\n");
		fclose(inFile);
		fclose(outFile);
		return 1;
	}
	fprintf(stderr,OK"Valid header! Reading the rest...\n");

	// Read salt and nonce
	if(fread(salt,1,sizeof(salt),inFile) != sizeof(salt)) {
		fprintf(stderr,ERR"Error reading salt from file.\n");
		fclose(inFile);
		fclose(outFile);
		return 1;
	}
	fprintf(stderr,OK"Salt retrieved.\n");

	if(fread(nonce,1,sizeof(nonce),inFile) != sizeof(nonce)) {
		fprintf(stderr,ERR"Error reading nonce from file.\n");
		fclose(inFile);
		fclose(outFile);
		return 1;
	}
	fprintf(stderr,OK"Nonce retrieved.\n");

	if(fread(encrypted,1,encLen,inFile) != encLen) {
		fprintf(stderr,ERR"Error reading data from file.\n");
		fclose(inFile);
		fclose(outFile);
		return 1;
	}
	fprintf(stderr,OK"Data retrieved.\n");
	fprintf(stderr,OK"Deriving key from password...\n");

	// Key derivation
	if (crypto_pwhash(
				key,KEY_SIZE,
				password,strlen(password),
				salt,crypto_pwhash_OPSLIMIT_INTERACTIVE,
				crypto_pwhash_MEMLIMIT_INTERACTIVE,
				crypto_pwhash_ALG_DEFAULT) != 0) {
		fprintf(stderr,ERR"Out of memory! Key could not be generated.\n");
		return 1;
	}
	fprintf(stderr,OK"Key derivation successful!\n");
	fprintf(stderr,OK"Proceeding to decrypt file...\n");

	size_t decLen = encLen - crypto_secretbox_MACBYTES;
	unsigned char decrypted[decLen];

	if (crypto_secretbox_open_easy(decrypted,encrypted,encLen,nonce,key) < 0) {
		fprintf(stderr,ERR"Error decrypting file.\n");
		fclose(inFile);
		fclose(outFile);
		return 1;
	}

	if (fwrite(decrypted,1,decLen,outFile) != decLen) {
		fprintf(stderr,ERR"Error writing data to file.\n");
		fclose(inFile);
		fclose(outFile);
		return 1;
	}

	fclose(inFile);
	fclose(outFile);
	
	fprintf(stderr,"\n"OK"File decrypted!\n");

	return 0;

}

bool isFile(const char* filename) {
	struct stat buffer;
	return (stat(filename, &buffer) == 0);
}

int main(int argc, char *argv[]) {
	fprintf(stderr,"\n- NetPaws Salty - File Encryption Program\n");
	fprintf(stderr,"- 2024 Ignacio Rivero\n");
	fprintf(stderr,"- FOR INTERNAL USE ONLY\n\n");
	// Initialize libsodium
	if (sodium_init() < 0) {
		fprintf(stderr,ERR"libsodium initialization failed! It is not safe to proceed.\n");
		return 1;
	}

	// Handle input, output and password arguments
	int in = -1, out = -1, pass = -1;
	bool decrypt = false;
	
	int i = 1, exit = 0;
	while (i < argc && exit == 0) {
		if (strcmp (argv[i],"-in") == 0) {
			if (i+1 >= argc) {
				fprintf(stderr,ERR"Argument -in requires a file.\n");
				exit = 1;
			} else if (!isFile(argv[i+1])) {
				fprintf(stderr,ERR"File not found error: %s\n",argv[i+1]);
				exit = 1;
			} else if (in == -1) {
				in = ++i;
			} else {
				fprintf(stderr,ERR"Too many arguments.\n");
				exit = 1;
			}
		} else if (strcmp(argv[i],"-out") == 0) {
			if (i+1 >= argc) {
				fprintf(stderr,ERR"Argument -out requires a file.\n");
				exit = 1;
			} else if (strcmp(argv[i+1],"-") == 0) {
				out = -2;
				i++;
			} else if (out == -1) {
				out = ++i;
			} else {
				fprintf(stderr,ERR"Too many arguments.\n");
				exit = 1;
			}
		} else if (strcmp (argv[i],"-key") == 0) {
			if (i+1 >= argc) {
				fprintf(stderr,ERR"Argument -key requires a password.\n");
				exit = 1;
			} else if (pass == -1) {
				if (strlen(argv[i+1]) > 128) {
					fprintf(stderr,ERR"Password too long.");
					exit = 1;
				} else {
					pass = ++i;
				}
			} else {
				fprintf(stderr,ERR"Too many arguments.\n");
				exit = 1;
			}
		} else if (strcmp (argv[i],"-d") == 0) {
			decrypt = true;
		} else {
			fprintf(stderr,ERR"Invalid argument: %s\n",argv[i]);
			exit = 1;
		}
		i++;
	}

	if (exit != 0)
		return exit;

	// Set input and output filenames, or go to standard I/O if none
	char input[FILENAME_MAX];
	char output[FILENAME_MAX];
	char psw[129];

	if (in == -1) {
		fprintf(stderr,ERR"No input file specified.\n");
		return 1;
	} else {
		snprintf(input,sizeof(input),argv[in]);
	}

	if (out == -1) {
		fprintf(stderr,NOR"No output file, writing to stdout.\n");
		snprintf(output,sizeof(output),"/dev/stdout");
	} else if (out == -2) {
		fprintf(stderr,NOR"Writing to stdout.\n");
		snprintf(output,sizeof(output),"/dev/stdout");
	} else {
		snprintf(output,sizeof(output),argv[out]);
	}

	char* password;

	// Set password and run!
	if (pass == -1) {
		password = getpass("Enter your encryption password: ");
		printf(DBG"Password : %s\n",password);
	} else {
		password = argv[pass];
		printf(DBG"Password : %s\n",password);
	}

	if (decrypt) {
			return decryptMessage(input,output,password);
		} else {
			return encryptMessage(input,output,password);
		}

}