#include <openssl/err.h>
#include <openssl/evp.h>
#include <openssl/hmac.h>
#include <openssl/params.h>
#include <openssl/rand.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <termios.h>

#define KEY_SIZE 32
#define SALT_SIZE 16
#define ITERATIONS 80000

int aes_ctr(const char *input_file, const char *output_file, const unsigned char *key, int enc) {

	int BUF_SIZE = 1024;
	int cipher_block_size = EVP_CIPHER_block_size(EVP_aes_256_ctr());

	int input_size = BUF_SIZE;
	int output_size = input_size + (cipher_block_size - 1);
	int f_len = 0;

	unsigned char input_buf[input_size], output_buf[output_size];

	FILE *finput = fopen(input_file, "rb");
	if (finput == NULL) {
		fprintf(stderr, "Error opening input file\n");
		return 1;
	}

	// Open output file
	FILE *foutput = fopen(output_file, "ab");
	if (foutput == NULL) {
		fprintf(stderr, "Error opening output file\n");
		fclose(finput);
		return 1;
	}

	EVP_CIPHER_CTX *ctx = NULL;
	if (!(ctx = EVP_CIPHER_CTX_new())) {
		fprintf(stderr, "Error creating context\n");
		fclose(finput);
		fclose(foutput);
		return 1;
	}

	EVP_MAC *mac = EVP_MAC_fetch(NULL, "HMAC", NULL);
	if (!mac) {
		fprintf(stderr, "Error creating HMAC\n");
		fclose(finput);
		fclose(foutput);
		return 1;
	}

	EVP_MAC_CTX *hctx = NULL;
	if (!(hctx = EVP_MAC_CTX_new(mac))) {
		fprintf(stderr, "Error creating HMAC context\n");
		fclose(finput);
		fclose(foutput);
		return 1;
	}

	// Set parameters for HMAC algorithm
	OSSL_PARAM params[2], *p = params;
	const EVP_MD *md = EVP_sha256();
	*p++ = OSSL_PARAM_construct_utf8_string("digest", (char *)EVP_MD_name(md), 0);
	*p = OSSL_PARAM_construct_end();

	unsigned char *hmac = malloc(32);
	unsigned char *mac_input;
	if (enc) {
		mac_input = output_buf;
	} else {
		mac_input = input_buf;
	}

	fseek(finput, 0, SEEK_END);
	long file_size = ftell(finput);
	rewind(finput);
	if (!enc) {
		// Remove the size of the SALT, IV and HMAC from the file size
		fseek(finput, 32, SEEK_SET);
		file_size -= 64;
	}

	// If enc is 1, then we are encrypting, else we are decrypting
	// If we are encrypting, we need to generate an IV
	// If we are decrypting, we need to read the IV from the file
	unsigned char iv[16];
	if (enc) {
		if (RAND_bytes(iv, 16) != 1) {
			fprintf(stderr, "Error generating IV\n");
			return 1;
		}
		if (fwrite(iv, 1, 16, foutput) != 16) {
			fprintf(stderr, "Error writing IV to file\n");
			return 1;
		}
	} else {
		// Seek forward by 16 bytes to ignore the salt
		if (fseek(finput, 16, SEEK_SET) != 0) {

			fprintf(stderr, "Error seeking to IV position in input file\n");
			return 1;
		}
		if (fread(iv, 1, 16, finput) != 16) {
			fprintf(stderr, "Error reading IV from file\n");
			return 1;
		}
	}

	if (EVP_CipherInit_ex(ctx, EVP_aes_256_ctr(), NULL, key, iv, enc) != 1) {
		fprintf(stderr, "ERROR: EVP_CipherInit_ex failed. OpenSSL error: %s\n",
				ERR_error_string(ERR_get_error(), NULL));
		EVP_CIPHER_CTX_free(ctx);
		return 1;
	}

	if (EVP_MAC_init(hctx, key, 32, params) != 1) {
		fprintf(stderr, "ERROR: EVP_MAC_init failed. OpenSSL error: %s\n",
				ERR_error_string(ERR_get_error(), NULL));
		fclose(finput);
		fclose(foutput);
		return 1;
	}

	int len;
	while (file_size > 0) {
		if (file_size < BUF_SIZE) {
			BUF_SIZE = file_size;
		}
		size_t read_size = fread(input_buf, 1, BUF_SIZE, finput);
		if (read_size == 0) {
			break;
		}
		if (EVP_CipherUpdate(ctx, output_buf, &len, input_buf, read_size) != 1) {
			fprintf(stderr, "ERROR: EVP_CipherUpdate failed. OpenSSL error: %s\n",
					ERR_error_string(ERR_get_error(), NULL));
			fclose(finput);
			fclose(foutput);
			return 1;
		}
		if (EVP_MAC_update(hctx, mac_input, read_size) != 1) {
			fprintf(stderr, "ERROR: EVP_MAC_update failed. OpenSSL error: %s\n",
					ERR_error_string(ERR_get_error(), NULL));
			fclose(finput);
			fclose(foutput);
			return 1;
		}
		if (fwrite(output_buf, 1, len, foutput) != len) {
			fprintf(stderr, "Error writing to output file\n");
			return 1;
		}
		file_size -= read_size;
	}

	if (EVP_CipherFinal_ex(ctx, output_buf, &f_len) != 1) {
		fprintf(stderr, "ERROR: EVP_CipherFinal_ex failed. OpenSSL error: %s\n",
				ERR_error_string(ERR_get_error(), NULL));
		fclose(finput);
		fclose(foutput);
		return 1;
	}
	size_t m_len = 0;
	if (EVP_MAC_final(hctx, hmac, &m_len, 32) != 1) {
		fprintf(stderr, "ERROR: EVP_MAC_final failed. OpenSSL error: %s\n",
				ERR_error_string(ERR_get_error(), NULL));
		fclose(finput);
		fclose(foutput);
		return 1;
	}
	if (m_len != 32) {
		fprintf(stderr, "ERROR: HMAC length is not 32\n");
		fclose(finput);
		fclose(foutput);
		return 1;
	}

	if (f_len) {
		if (fwrite(output_buf, 1, f_len, foutput) != f_len) {
			fprintf(stderr, "Error writing to output file\n");
			fclose(finput);
			fclose(foutput);
			return 1;
		}
	}
	if (enc) {
		if (fwrite(hmac, 1, m_len, foutput) != m_len) {
			fprintf(stderr, "Error writing HMAC to file\n");
			fclose(finput);
			fclose(foutput);
			return 1;
		}
	} else {
		unsigned char *hmac2 = malloc(32);
		if (fread(hmac2, 1, 32, finput) != 32) {
			fprintf(stderr, "Error reading HMAC from file\n");
			fclose(finput);
			fclose(foutput);
			return 1;
		}
		int cmp = memcmp(hmac, hmac2, 32);
		if (cmp == 0) {
			printf("HMACs match\n");
		} else {
			fprintf(stderr, "HMACs do not match\n");
			fclose(finput);
			fclose(foutput);
			return 1;
		}
	}

	fclose(finput);
	fclose(foutput);

	return 0;
}

int encrypt(char *input_file, const char *passphrase) {
	unsigned char key[KEY_SIZE * 2];
	unsigned char salt[SALT_SIZE];
	// Derive key from passphrase using PBKDF2
	char *output_file = malloc(strlen(input_file) + 5);
	strcpy(output_file, input_file);
	strcat(output_file, ".enc");

	if (RAND_bytes(salt, SALT_SIZE) != 1) {
		fprintf(stderr, "Error generating salt\n");
		return 1;
	}
	// Write salt to output file
	FILE *foutput = fopen(output_file, "wb");
	if (foutput == NULL) {
		fprintf(stderr, "Error opening output file\n");
		return 1;
	}
	if (fwrite(salt, 1, 16, foutput) != 16) {
		fprintf(stderr, "Error writing salt to file\n");
		fclose(foutput);
		return 1;
	}
	fclose(foutput);

	int pass_len = strlen(passphrase);
	// Derive key from passphrase using PBKDF2
	if (PKCS5_PBKDF2_HMAC(passphrase, pass_len, salt, SALT_SIZE, ITERATIONS, EVP_sha256(),
						  KEY_SIZE * 2, key) != 1) {
		fprintf(stderr, "Error deriving key from passphrase\n");
		return 1;
	}

	aes_ctr(input_file, output_file, key, 1);
	return 0;
}

int decrypt(char *input_file, const char *passphrase) {
	unsigned char key[KEY_SIZE];
	unsigned char salt[SALT_SIZE];
	// Derive key from passphrase using PBKDF2
	char *output_file = malloc(strlen(input_file) + 5);
	strcpy(output_file, input_file);
	strcat(output_file, ".dec");
	// Read salt from input file
	FILE *finput = fopen(input_file, "rb");
	if (finput == NULL) {
		fprintf(stderr, "Error opening input file\n");
		return 1;
	}
	if (fread(salt, 1, 16, finput) != 16) {
		fprintf(stderr, "Error reading salt from file\n");
		fclose(finput);
		return 1;
	}
	fclose(finput);

	// Derive key from passphrase using PBKDF2
	if (PKCS5_PBKDF2_HMAC(passphrase, strlen(passphrase), salt, SALT_SIZE, ITERATIONS, EVP_sha256(),
						  KEY_SIZE, key) != 1) {
		fprintf(stderr, "Error deriving key from passphrase\n");
		return 1;
	}

	aes_ctr(input_file, output_file, key, 0);
	return 0;
}

void disableEcho() {
	struct termios oldTermios, newTermios;
	tcgetattr(0, &oldTermios);
	newTermios = oldTermios;
	newTermios.c_lflag &= ~(ECHO);
	tcsetattr(0, TCSANOW, &newTermios);
}

void enableEcho() {
	struct termios oldTermios;
	tcgetattr(0, &oldTermios);
	oldTermios.c_lflag |= ECHO;
	tcsetattr(0, TCSANOW, &oldTermios);
}

int main(int argc, char *argv[]) {
	if (argc < 3) {
		fprintf(stderr, "Usage: %s {enc|dec} [file_path]\n", argv[0]);
		return 1;
	}

	char *mode = argv[1];
	char *input_file = argv[2];

	if (!(strcmp(mode, "enc") == 0 || strcmp(mode, "dec") == 0)) {
		fprintf(stderr, "Invalid mode. Use 'enc' or 'dec'.\n");
	}

	if (argc != 3) {
		fprintf(stderr, "Usage: %s {enc|dec} [file_path]\n", argv[0]);
		return 1;
	}
	input_file = argv[2];

	char passphrase[256]; // Assuming maximum passphrase length of 255 characters
	printf("Enter passphrase: ");
	disableEcho();
	if (fgets(passphrase, sizeof(passphrase), stdin) == NULL) {
		fprintf(stderr, "Error reading passphrase from stdin\n");
		return 1;
	}
	enableEcho();
	putchar('\n');
	passphrase[strcspn(passphrase, "\n")] = '\0'; // Remove trailing newline
	int suc = 0;
	if (strcmp(mode, "enc") == 0) {
		suc = encrypt(input_file, passphrase);
	} else {
		suc = decrypt(input_file, passphrase);
	}
	if (suc == 0) {
		printf("Operation completed successfully\n");
	} else {
		printf("Operation failed\n");
	}
	return 0;
}