#define _POSIX_C_SOURCE 200809L
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdbool.h>
#include <sys/stat.h>
#include <sodium.h>
#include <termios.h>
#include "pawstd.h"
#define KEY_SIZE crypto_secretbox_KEYBYTES
#define SALT_SIZE crypto_pwhash_SALTBYTES
#define BUF_SIZE 1024
#define VER "1.00"

bool stdinput = false;

size_t fileSize(FILE *file) {
	fseek(file,0,SEEK_END);
	size_t size = ftell(file);
	fseek(file,0,SEEK_SET);
	return size;
}

bool keyGen(const unsigned char* salt, const char* password, unsigned char* key) {
	fprintf(stderr,NOR"Beginning key generation...\n");
	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 false;
	}
	fprintf(stderr,OK"Key generation successful!\n");
	return true;
}

int encryptFile(const char* inputFile, const char* outputFile, const char* password) {

	// Generate salt
	unsigned char salt[SALT_SIZE];
	randombytes_buf(salt, SALT_SIZE);
	
	// Key generation from password
	unsigned char key[KEY_SIZE];
	if (!keyGen(salt,password,key)) {
		return 1;
	}

	fprintf(stderr,NOR"Reading file...\n");

	FILE *outFile = fopen(outputFile,"wb");
	if (outFile == NULL) {
		fprintf(stderr,ERR"Could not open output file!\n");
		return 1;
	}

	unsigned char* decrypted;
	size_t decLen = 0;

	// Check if input is stdin, and read the input
	if (stdinput) {
		// Begin reading into buffer
		unsigned char buffer[BUF_SIZE];
		unsigned char* input = malloc(BUF_SIZE);
		if (input == NULL) {
			fprintf(stderr,ERR"Memory allocation error.\n");
			fclose(outFile);
			return 1;
		}
        while(true) {
			size_t readBytes = fread(buffer,1,BUF_SIZE,stdin);
			decLen += readBytes;
			// Prevents reallocating to zero bytes.
			if (decLen > 0) {
				unsigned char* old = input;
				input = realloc(input, decLen);
				if (input == NULL) {
                    fprintf(stderr, ERR"Memory allocation error.\n");
                    fclose(outFile);
                    free(old);
                    return 1;
                }
			}
			memcpy(input + decLen - readBytes,buffer,readBytes);
			if (readBytes == 0 && !ferror(stdin))
				break;
			else if (readBytes == 0 && ferror(stdin)) {
				fprintf(stderr,ERR"Error reading from stdin.\n");
				fclose(outFile);
				free(input);
				return 1;
			}
		}
		decrypted = input;
	} else {
		// Open input file
		FILE *inFile = fopen(inputFile,"rb");
		if (inFile == NULL) {
			fprintf(stderr,ERR"Could not open input file!\n");
			return 1;
		}

		// Read input file
		fprintf(stderr,NOR"Reading input filesize...\n");
		decLen = fileSize(inFile);
		decrypted = malloc(decLen);

		if (fread(decrypted,1,decLen,inFile) != decLen) {
			fprintf(stderr,ERR"Error reading from file.\n");
			fclose(inFile);
			fclose(outFile);
			free(decrypted);
			return 1;
		}
		fclose(inFile);
	}

	if (decLen == 0) {
		fprintf(stderr,ERR"No bytes read from file.\n");
		fclose(outFile);
		return 1;
	}

	size_t encLen = decLen + crypto_secretbox_MACBYTES;
	fprintf(stderr,OK"Input file is %zu bytes.\n",decLen);
	fprintf(stderr,OK"Encrypting file to %zu bytes.\n", encLen);
	fprintf(stderr,NOR"Beginning encryption...\n");

	// Generate a nonce, write headers and encrypted message
	unsigned char nonce[crypto_secretbox_NONCEBYTES];
	randombytes_buf(nonce,sizeof(nonce));

	// Encrypt the data
	unsigned char* encrypted = malloc(encLen);
    if (encrypted == NULL) {
        fprintf(stderr,ERR"Memory allocation error: File too large.\n");
        free(decrypted);
        fclose(outFile);
        return 1;
    }
	crypto_secretbox_easy(encrypted, decrypted, decLen, nonce, key);

	// Patchy but correct-ish attempt to solve annoying stdout issue
	size_t totalSize = strlen("gneurshk") + sizeof(salt) + sizeof(nonce) + encLen;
	size_t offset = 0;
	unsigned char* outputBuffer = malloc(totalSize);
	if (outputBuffer == NULL) {
		fprintf(stderr,ERR"Memory allocation error.\n");
		free(decrypted);
		fclose(outFile);
		return 1;
	}
	memcpy(outputBuffer + offset, "gneurshk", strlen("gneurshk"));
	offset += strlen("gneurshk");

	memcpy(outputBuffer + offset, salt, sizeof(salt));
	offset += sizeof(salt);

	memcpy(outputBuffer + offset, nonce, sizeof(nonce));
	offset += sizeof(nonce);

	memcpy(outputBuffer + offset, encrypted, encLen);

	fprintf(stderr,NOR"Writing encrypted data...\n");

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

	fclose(outFile);
	free(decrypted);
	free(outputBuffer);

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

	return 0;
}

int decryptFile(const char* inputFile, const char* outputFile, const char* password) {
	// Open output file
	FILE *outFile = fopen(outputFile,"wb");
	if (outFile == NULL) {
		fprintf(stderr,ERR"Could not open output file!\n");
		return 1;
	}

	// Variables
	unsigned char salt[SALT_SIZE];
	unsigned char nonce[crypto_secretbox_NONCEBYTES];
	unsigned char* fullInput;
	size_t inLen = 0;

	// Check if input is stdin, and read the input
	if (stdinput) {
		// Begin reading into buffer
		unsigned char buffer[BUF_SIZE];
		unsigned char* input = malloc(BUF_SIZE);
		if (input == NULL) {
			fprintf(stderr,ERR"Memory allocation error.");
			fclose(outFile);
			return 1;
		}
		while(true) {
			size_t readBytes = fread(buffer,1,BUF_SIZE,stdin);
			inLen += readBytes;
			// Prevents reallocating to zero bytes.
			if (inLen > 0) {
				unsigned char* old = input;
				input = realloc(input, inLen);
				if (input == NULL) {
					fprintf(stderr,ERR"Memory allocation error.\n");
					fclose(outFile);
					free(old);
					return 1;
				}
			}
			memcpy(input + inLen - readBytes,buffer,readBytes);
			if (readBytes == 0 && !ferror(stdin))
				break;
			else if (readBytes == 0 && ferror(stdin)) {
				fprintf(stderr,ERR"Error reading from stdin.\n");
				fclose(outFile);
				free(input);
				return 1;
			}
		}
		fullInput = input;
	} else {	
		FILE *inFile = fopen(inputFile,"rb");
		
		// Get the size of the file and read it
		inLen = fileSize(inFile);
		unsigned char* input = malloc(inLen);
        if (input == NULL) {
            fprintf(stderr, ERR"Memory allocation error: File too large.\n");
            fclose(inFile);
            fclose(outFile);
            return 1;
        }

		if (fread(input,1,inLen,inFile) != inLen) {
			fprintf(stderr,ERR"Error reading from input file.\n");
			fclose(inFile);
			fclose(outFile);
			return 1;
		}
		
		fclose(inFile);
		fullInput = input;
	}

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

	if (inLen == 0) {
		fprintf(stderr,ERR"No bytes read from file.\n");
		fclose(outFile);
		return 1;
	}

	fprintf(stderr,OK"Size of input file: %zu bytes.\n",inLen);
	fprintf(stderr,OK"Encrypted content is %zu bytes.\n",encLen);

	unsigned char* encrypted = malloc(encLen);
    if (encrypted == NULL) {
        fprintf(stderr,ERR"Memory allocation error.");
        free(fullInput);
        fclose(outFile);
        return 1;
    }

	// Verify file
	size_t offset = 0;
	if (memcmp(fullInput,"gneurshk",strlen("gneurshk")) != 0) {
		fprintf(stderr,ERR"Invalid header.\n");
		fprintf(stderr,ERR"Make sure the input file was signed by this program!\n");
		free(fullInput);
		fclose(outFile);
		return 1;
	}
	fprintf(stderr,OK"Valid header! Reading the rest...\n");
	offset += strlen("gneurshk");

	// Read salt and nonce
	memcpy(salt, fullInput + offset, sizeof(salt));
	offset += sizeof(salt);

	memcpy(nonce, fullInput + offset, sizeof(nonce));
	offset += sizeof(nonce);

	memcpy(encrypted, fullInput + offset, encLen);

	fprintf(stderr,OK"Data retrieved.\n");

    free(fullInput);

	// Key
	unsigned char key[KEY_SIZE];
	
	if (!keyGen(salt,password,key)) {
		free(encrypted);
		fclose(outFile);
		return 1;
	}

	fprintf(stderr,OK"Proceeding to decrypt file...\n");

	size_t decLen = encLen - crypto_secretbox_MACBYTES;
	unsigned char* decrypted = malloc(decLen);

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

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

	fclose(outFile);
    free(encrypted);
    free(decrypted);

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

	return 0;
}

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

bool getPassword(char *pw, int size) {
    struct termios term;
    tcgetattr(fileno(stdin), &term);
    int i = 0;
    char c;
    pw[0] = '\0';
    // Hide input
    term.c_lflag &= ~ECHO;
    tcsetattr(fileno(stdin), 0, &term);
    // Get password from stdin
    while (true) {
        c = fgetc(stdin);
        if (c == '\r' || c == '\n' || feof(stdin)) {
            break;
        }
        // Return false on excess characters
        if (i < size - 1) {
            pw[i] = c;
            pw[i + 1] = '\0';
        } else
            return false;
        i++;
    }
    // Show input
    term.c_lflag |= ECHO;
    tcsetattr(fileno(stdin), 0, &term);
    return true;
}

void printUsage() {
	fprintf(stderr,"\nUsage: salty [-d] -in [INPUT] -out [OUTPUT] [OPTIONS]\n");
	fprintf(stderr,"Encrypts or decrypts a file with a password. When INPUT or OUTPUT are not set,\n");
	fprintf(stderr,"or set to -, standard input or output is used respectively.\n\n");
	fprintf(stderr,"Options:\n");
	fprintf(stderr,"    -in\t\tTakes an input file\n");
	fprintf(stderr,"    -out\tTakes an output file\n");
	fprintf(stderr,"    -d\t\tEnable decryption mode\n");
	fprintf(stderr,"    -key\tTakes a password, needed if reading from stdin\n");
	fprintf(stderr,"    \t\tOtherwise this password is read from stdin itself\n");
	fprintf(stderr,"    -v\t\tDisplays version information and exits\n");
	fprintf(stderr,"    -h\t\tDisplays this help and exits\n");
}

int main(int argc, char *argv[]) {
    setvbuf(stdout, NULL, _IONBF, 0);
	fprintf(stderr,NOR"NetPaws Salty - File Encryption Program\n");
	fprintf(stderr,NOR"2024 Ignacio Rivero\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, password, and other arguments
	int in = -1, out = -1, pass = -1;
	bool decrypt = false;
	
	int i = 1, exit = -1;
	while (i < argc && exit == -1) {
		if (strcmp (argv[i],"-in") == 0) {
			if (i+1 >= argc) {
				fprintf(stderr,ERR"Argument -in requires a file.\n");
				exit = 1;
			} else if (strcmp(argv[i+1],"-") == 0) {
				in = -2;
				i++;
			} 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 is 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 if (strcmp (argv[i],"-v") == 0) {
			fprintf(stderr,NOR"Version %s\n",VER);
			exit = 0;
		} else if (strcmp (argv[i],"-h") == 0) {
			fprintf(stderr,NOR"Version %s\n",VER);
			printUsage();
			exit = 0;
		} else {
			fprintf(stderr,ERR"Invalid argument: %s\n",argv[i]);
			exit = 1;
		}
		i++;
	}

	if (exit == 1) {
		printUsage();
		return exit;
	} else if (exit != -1)
		return exit;

	fprintf(stderr,"\n");

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

	if (in == -1) {
		fprintf(stderr,NOR"No input file, reading from stdin.\n");
		stdinput = true;
	} else if (in == -2) {
		fprintf(stderr,NOR"Reading from stdin.\n");
		stdinput = true;
	} else {
		snprintf(input,sizeof(input),"%s", 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),"%s", argv[out]);
	}

	char* password;

	// Set password and run!
	if (pass > 0) {
		password = argv[pass];
	} else if (in < 0) {
        fprintf(stderr, ERR"Cannot read from stdin without a password.\n");
        return 1;
    } else {
        password = malloc(129);
        printf(NOR"Enter your encryption password: ");
        if (!getPassword(password,130)) {
            fprintf(stderr,"\n"ERR"Password is too long.");
            return 1;
        }
        printf("\n");
    }

    if (decrypt) {
			return decryptFile(input, output, password);
		} else {
			return encryptFile(input, output, password);
		}
}