#include <openssl/evp.h>
#include <openssl/hmac.h>
#include <openssl/aes.h>
#include <openssl/rand.h>

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <stdbool.h>
#include <string.h>
#include <ctype.h>
#include <unistd.h>
#include <errno.h>
#include <termios.h>

#define PASS_MAXLEN	64

#define SALT_BYTES	8
#define KEY_BYTES	32
#define IV_BYTES	16

#define PBKDF2_DEFAULT_ROUNDS	1 << 20

#define AES_BUFFER(LEN) \
	(unsigned char *)( \
		calloc(AES_BLOCK_SIZE * (1 + ((LEN) / AES_BLOCK_SIZE)), \
			sizeof(unsigned char)));

#ifdef __APPLE__
void explcit_bzero(void *buf, size_t len)
{
	void *volatile const bufv = buf;
	memset(bufv, 0, len);
}
#endif

static bool chunk;

bool encrypt_file(const char *path_in, const char *path_out,
                  const char *passphrase, unsigned long iter)
{
	unsigned char salt[SALT_BYTES], key[KEY_BYTES], iv[IV_BYTES];
	bool result = true;

	FILE *file_in = fopen(path_in, "r");
	if (file_in == NULL) {
		fprintf(stderr, "Cannot open file %s: %s\n",
		        path_in, strerror(errno));
		return false;
	}

	FILE *file_out  = fopen(path_out, "wb");
	if (file_out == NULL) {
		fprintf(stderr, "Cannot open file %s: %s\n",
		        path_out, strerror(errno));
		fclose(file_in);
		return false;
	}

	RAND_bytes(salt, SALT_BYTES);
	fwrite(salt, SALT_BYTES, 1, file_out);

	explicit_bzero(&key, KEY_BYTES);
	explicit_bzero(&iv, IV_BYTES);

	EVP_BytesToKey(EVP_aes_256_cbc(), EVP_sha1(), salt,
	               (const unsigned char *)passphrase, strlen(passphrase),
	               iter ? iter : PBKDF2_DEFAULT_ROUNDS,
	               (unsigned char *)&key, (unsigned char *)&iv);

	EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();
	EVP_EncryptInit_ex(ctx, EVP_aes_256_cbc(), NULL,
	                   (unsigned char *)&key, (unsigned char *)&iv);

	unsigned char blocksize = EVP_CIPHER_CTX_block_size(ctx);

	unsigned char *rdbuf, *outbuf;
	int len, len_out;

	fseek(file_in, 0L, SEEK_END);
	size_t file_bytes = ftell(file_in);
	rewind(file_in);

	if (chunk) {
		rdbuf = malloc(BUFSIZ);
		outbuf = malloc(BUFSIZ + blocksize);

		do {
			len = fread(rdbuf, 1, BUFSIZ, file_in);
			if (!EVP_EncryptUpdate(ctx, outbuf, &len_out, rdbuf, len)) {
				result = false;
				break;
			}
			fwrite(outbuf, 1, len_out, file_out);
		} while (len == BUFSIZ);

		if (result && !EVP_EncryptFinal_ex(ctx, outbuf, &len_out)) {
			result = false;
		}
		fwrite(outbuf, 1, len_out, file_out);
	} else {
		rdbuf = AES_BUFFER(file_bytes);
		outbuf = AES_BUFFER(file_bytes + blocksize);

		fread(rdbuf, 1, file_bytes, file_in);

		EVP_EncryptUpdate(ctx, outbuf, &len, rdbuf, file_bytes);
		len_out = len;

		if (!EVP_EncryptFinal_ex(ctx, outbuf + len_out, &len)) {
			result = false;
		} else {
			len_out += len;
			fwrite(outbuf, len_out, 1, file_out);
		}
	}

	free(rdbuf);
	free(outbuf);

	EVP_CIPHER_CTX_free(ctx);

close_all:
	fclose(file_in);
	fclose(file_out);

	return result;
}

bool decrypt_file(const char *path_in, const char *path_out,
                  const char *passphrase, unsigned long iter)
{
	unsigned char salt[SALT_BYTES], key[KEY_BYTES], iv[IV_BYTES];
	bool result = true;

	FILE *file_in = fopen(path_in, "rb");
	if (file_in == NULL) {
		fprintf(stderr, "Cannot open file %s: %s\n",
		        path_in, strerror(errno));
		return false;
	}

	FILE *file_out = NULL;
	if (path_out[0] == '-' && strnlen(path_out, PATH_MAX) == 1) {
		file_out = stdout;
		if (chunk) {
			fprintf(stderr, "Warning: ignoring '-c' for decrypt to stdout.\n");
			chunk = false;
		}
	} else {
		file_out = fopen(path_out, "w");
		if (file_out == NULL) {
			fprintf(stderr, "Cannot open file %s: %s\n",
			       path_out, strerror(errno));
			fclose(file_in);
			return false;
		}
	}

	explicit_bzero(&key, KEY_BYTES);
	explicit_bzero(&iv, IV_BYTES);

	fseek(file_in, 0L, SEEK_END);
	size_t file_bytes = ftell(file_in) - SALT_BYTES;
	rewind(file_in);

	if (!fread(&salt, SALT_BYTES, 1, file_in)) {
		result = false;
		goto close_all;
	}

	EVP_BytesToKey(EVP_aes_256_cbc(), EVP_sha1(), salt,
	               (const unsigned char *)passphrase, strlen(passphrase),
	               iter ? iter : PBKDF2_DEFAULT_ROUNDS,
	               (unsigned char *)&key, (unsigned char *)&iv);

	EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();
	EVP_DecryptInit_ex(ctx, EVP_aes_256_cbc(), NULL,
                     (unsigned char *)&key, (unsigned char *)&iv);

	unsigned char blocksize = EVP_CIPHER_CTX_block_size(ctx);

	unsigned char *rdbuf, *outbuf;
	int len, len_out;

	if (chunk) {
		rdbuf = malloc(BUFSIZ);
		outbuf = malloc(BUFSIZ + blocksize);

		do {
			len = fread(rdbuf, 1, BUFSIZ, file_in);
			if (!EVP_DecryptUpdate(ctx, outbuf, &len_out, rdbuf, len)) {
				result = false;
				break;
			}
			fwrite(outbuf, 1, len_out, file_out);
		} while (len == BUFSIZ);

		if (result && !EVP_DecryptFinal_ex(ctx, outbuf, &len_out)) {
			result = false;
		}
		fwrite(outbuf, 1, len_out, file_out);
	} else {
		rdbuf = AES_BUFFER(file_bytes);
		outbuf = AES_BUFFER(file_bytes + blocksize);

		fread(rdbuf, 1, file_bytes, file_in);

		EVP_DecryptUpdate(ctx, outbuf, &len, rdbuf, file_bytes);
		len_out = len;

		if (!EVP_DecryptFinal_ex(ctx, outbuf + len_out, &len)) {
			result = false;
		} else {
			len_out += len;

			if (file_out == stdout) {
				for (unsigned char ch = 0; ch < file_bytes; ch++) {
					if (!isascii(*(outbuf + ch))) {
						fprintf(stderr,
                    "Invalid characters encountered"
                    " - not printing.\n");
						goto cleanup;
					}
				}
			}

			fwrite(outbuf, len_out, 1, file_out);
		}
	}

cleanup:
	free(rdbuf);
	free(outbuf);

	EVP_CIPHER_CTX_free(ctx);

close_all:
	fclose(file_in);
	fclose(file_out);

	return result;
}

void usage(void)
{
	printf("Minimal file encryption program (AES-CBC-256 + PBKDF2/SHA256)\n\n"
	       "Usage: encdec [ -edhs ] [ -io FILE ] [ -p PASSPHRASE ] [ -r ITER ]\n\n"
	       "  -c:\t Use chunking (ignored with '-o -')\n"
	       "  -d:\t Decrypt (mutually exclusive with encryption)\n"
	       "  -e:\t Encrypt (mutually exclusive with decryption)\n"
	       "  -h:\t Show this help printout\n"
	       "  -i:\t Input file name/path\n"
	       "  -o:\t Output file name/path ('-': decrypt to stdout)\n"
	       "  -p:\t Specify a passphrase (optional, less secure!)\n"
	       "  -r:\t Key derivation rounds (default: 2^20)\n");
}


int main(int argc, char *argv[])
{
	char path_in[PATH_MAX], path_out[PATH_MAX],
	     passphrase[PASS_MAXLEN + 1];

	bool encrypt, decrypt, prompt;
	encrypt = decrypt = false;
	prompt = true;

	unsigned long iter = 0;

	if (argc < 2) {
		fprintf(stderr, "ERROR: No arguments given.\n\n");
		usage();
		return 1;
	}

	explicit_bzero(&path_in, PATH_MAX);
	explicit_bzero(&path_out, PATH_MAX);
	explicit_bzero(&passphrase, PASS_MAXLEN);

	int opt, pos;
	while ((opt = getopt(argc, argv, "cdehi:o:p:r:")) != -1) {
		switch (opt) {
		case 'c':
			chunk = true;
			break;
		case 'd':
			decrypt = true;
			break;
		case 'e':
			encrypt = true;
			break;
		case 'i':
			strncpy(path_in, optarg, PATH_MAX);
			break;
		case 'o':
			strncpy(path_out, optarg, PATH_MAX);
			break;
		case 'p':
			strncpy(passphrase, optarg, PASS_MAXLEN);
			prompt = false;
			break;
		case 'r':
			if (sscanf(optarg, "%lu%n", &iter, &pos) != 1
			    || pos != strlen(optarg)) {
				fprintf(stderr, "-%c: Type mismatch: %s\n", opt, optarg);
				usage();
				return 1;
			}
			break;
		case 'h':
			usage();
			return 0;
		case '?':
			fprintf(stderr, "Unknown option: -%c\n", optopt);
			usage();
			return 1;
		case ':':
			fprintf(stderr, "Option -%c requires a value.\n", optopt);
			usage();
			break;
		}
	}

	if (optind < argc) {
		fprintf(stderr, "Warning, unused arguments:");
		for (; optind < argc; optind++)
			fprintf(stderr, " %s", argv[optind]);
		fprintf(stderr, "\n");
	}

	if (!strlen(path_in) || !strlen(path_out)) {
		fprintf(stderr, "File name not specified\n");
		usage();
		return 1;
	}

	if (encrypt && decrypt) {
		fprintf(stderr, "Cannot encrypt and decrypt at the same time\n");
		usage();
		return 1;
	} else if (!(encrypt || decrypt)) {
		fprintf(stderr, "Must either encrypt or decrypt\n");
		usage();
		return 1;
	}

	if (prompt) {
		struct termios term;
		tcgetattr(fileno(stdin), &term);

		term.c_lflag &= ~ECHO;
		tcsetattr(fileno(stdin), 0, &term);

		printf("Enter password ");
		fgets(passphrase, PASS_MAXLEN, stdin);
		passphrase[strcspn(passphrase, "\n")] = 0;
		printf("\n");

		if (encrypt) {
			char conf[PASS_MAXLEN];
			explicit_bzero(&conf, PASS_MAXLEN);

			printf("Confirm password ");
			fgets(conf, PASS_MAXLEN, stdin);
			conf[strcspn(conf, "\n")] = 0;
			printf("\n");

			if (strncmp(passphrase, conf, PASS_MAXLEN)) {
				printf("Passphrases do not match.\n");
				return 1;
			}
		}
		term.c_lflag |= ECHO;
		tcsetattr(fileno(stdin), 0, &term);
	}

	if (encrypt) {
		if (!encrypt_file(path_in, path_out, passphrase, iter)) {
			fprintf(stderr, "Failed to encrypt file\n");
			return 1;
		}
	}
	else if (decrypt) {
		if (!decrypt_file(path_in, path_out, passphrase, iter)) {
			fprintf(stderr, "Failed to decrypt file\n");
			return 1;
		}
	}

	return 0;
}