/*
 * Copyright (C) 2016 Guido Berhoerster <guido+sencrypt@berhoerster.name>
 *
 * Permission is hereby granted, free of charge, to any person obtaining
 * a copy of this software and associated documentation files (the
 * "Software"), to deal in the Software without restriction, including
 * without limitation the rights to use, copy, modify, merge, publish,
 * distribute, sublicense, and/or sell copies of the Software, and to
 * permit persons to whom the Software is furnished to do so, subject to
 * the following conditions:
 *
 * The above copyright notice and this permission notice shall be included
 * in all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
 * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
 * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
 * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

#include <stdio.h>
#include <string.h>
#include <stdint.h>
#include <stdbool.h>
#include <unistd.h>
#include <limits.h>
#include <libgen.h>
#include <arpa/inet.h>
#include <sys/stat.h>
#include <openssl/conf.h>
#include <openssl/rand.h>
#include <openssl/evp.h>
#include <openssl/err.h>

#ifdef HAVE_ERR_H
#include <err.h>
#endif /* HAVE_ERR_H */
#include "compat.h"

#define	MAX(a, b)	(((a) > (b)) ? (a) : (b))

#define	EXIT_USAGE	2

#define	SENCRYPT_FORMAT_VERSION	1
#define	PBKDF2_ITERATIONS	50000
#define	SALT_LEN		16
#define	BUFFER_SIZE		(16 * 1024)
#define	MAX_PASSWORD_LEN	256

enum {
	CMD_SENCRYPT,
	CMD_SDECRYPT
};

static void
openssl_warn(void) {
	unsigned long	errcode;

	while ((errcode = ERR_get_error()) != 0) {
		warnx("%s", ERR_error_string(errcode, NULL));
	}
}

static size_t
read_keyfile(const char *filename, unsigned char *key, size_t key_size_max)
{
	size_t		keyfile_size = 0;
	FILE		*fp = NULL;
	struct		stat statbuf;

	fp = fopen(filename, "r");
	if (fp == NULL) {
		warn("could not open key file \"%s\"", filename);
		goto out;
	}

	if (fstat(fileno(fp), &statbuf) == -1) {
		warn("could not stat key file \"%s\"", filename);
		goto out;
	}

	if (!S_ISREG(statbuf.st_mode)) {
		warnx("key file \"%s\" is not a regular file", filename);
		goto out;
	}

	if ((uintmax_t)statbuf.st_size > SIZE_MAX) {
		warnx("key file \"%s\" is too large", filename);
		goto out;
	}
	keyfile_size = (size_t)statbuf.st_size;
	if ((keyfile_size > key_size_max) ||
	    (keyfile_size == 0)) {
		warnx("invalid key size");
		goto out;
	}

	if (fread(key, 1, keyfile_size, fp) != keyfile_size) {
		warnx("could not read key file \"%s\"", filename);
		goto out;
	}

out:
	if (fp != NULL) {
		fclose(fp);
	}

	return (keyfile_size);
}

static int
find_algorithm(const char *algo_name, const EVP_CIPHER **cipher_ptr,
    size_t *key_len_ptr)
{
	int	retval = 0;
	const EVP_CIPHER	*cipher = NULL;
	size_t	key_len = *key_len_ptr;

	if (strcmp(algo_name, "aes") == 0) {
		switch (key_len) {
		case 0:
			key_len = 16;
		case 16:
			cipher = EVP_aes_128_cbc();
			break;
		case 24:
			cipher = EVP_aes_192_cbc();
			break;
		case 32:
			cipher = EVP_aes_256_cbc();
			break;
		default:
			warnx("invalid key length %zu", key_len);
			retval = -1;
		}
	} else if (strcmp(algo_name, "arcfour") == 0) {
		if (key_len == 0) {
			key_len = 16;
			cipher = EVP_rc4();
		} else if (key_len <= EVP_MAX_KEY_LENGTH) {
			/*
			 * for RC4 keys are not used verbatim but dervied using
			 * PBKDF2 with a hardcoded key length of 128 bit
			 */
			key_len = 16;
			cipher = EVP_rc4();
		} else {
			warnx("invalid key length %zu", key_len);
			retval = -1;
		}
	} else if (strcmp(algo_name, "des") == 0) {
		if (key_len == 0) {
			key_len = 8;
			cipher = EVP_des_cbc();
		} else if (key_len == 8) {
			cipher = EVP_des_cbc();
		} else {
			warnx("invalid key length %zu", key_len);
			retval = -1;
		}
	} else if (strcmp(algo_name, "3des") == 0) {
		if (key_len == 0) {
			key_len = 24;
			cipher = EVP_des_ede3_cbc();
		} else if (key_len == 24) {
			cipher = EVP_des_ede3_cbc();
		} else {
			warnx("invalid key length %zu", key_len);
			retval = -1;
		}
	} else {
		warnx("unknown algorithm \"%s\"", algo_name);
		retval = -1;
	}

	*cipher_ptr = cipher;
	*key_len_ptr = key_len;

	return (retval);
}

static int
read_header(BIO *bio_in, uint32_t *iterations, unsigned char *iv, int iv_len,
    unsigned char *salt, int salt_len)
{
	int		read_len;
	uint32_t	version;
	int		retval = 0;

	read_len = BIO_read(bio_in, &version, sizeof (version));
	if (read_len != sizeof (version)) {
		warnx("failed to read version from input file");
		if (read_len < 0) {
			openssl_warn();
		}
		retval = -1;
		goto out;
	}
	version = htonl(version);
	if (version != SENCRYPT_FORMAT_VERSION) {
		warnx("unknown format version %d", version);
		retval = -1;
		goto out;
	}

	read_len = BIO_read(bio_in, iterations, sizeof (*iterations));
	if (read_len != sizeof (*iterations)) {
		warnx("failed to read iterations from input file");
		if (read_len < 0) {
			openssl_warn();
		}
		retval = -1;
		goto out;
	}
	*iterations = htonl(*iterations);
	if ((*iterations == 0) || ((sizeof (int) <= sizeof (uint32_t)) &&
	    (*iterations > INT_MAX))) {
		warnx("invalid number of iterations");
		retval = -1;
		goto out;
	}

	if (iv_len > 0) {
		read_len = BIO_read(bio_in, iv, iv_len);
		if (read_len != iv_len) {
			warnx("failed to read IV from input file");
			if (read_len < 0) {
				openssl_warn();
			}
			retval = -1;
			goto out;
		}
	}

	read_len = BIO_read(bio_in, salt, salt_len);
	if (read_len != salt_len) {
		warnx("failed to read salt from input file");
		if (read_len < 0) {
			openssl_warn();
		}
		retval = -1;
		goto out;
	}

out:
	return (retval);
}

static int
sencrypt(const EVP_CIPHER *cipher, BIO *bio_in, BIO *bio_out,
    const unsigned char *key, size_t key_len, const unsigned char *iv,
    const unsigned char *salt)
{
	int		retval = 0;
	uint32_t	version;
	uint32_t	iterations;
	int		iv_len;
	int		write_len;
	int		read_len;
	BIO		*bio_cipher = NULL;
	char		*buf = NULL;
	EVP_CIPHER_CTX	*cipher_ctx;

	/* set up cipher filter */
	bio_cipher = BIO_new(BIO_f_cipher());
	BIO_set_cipher(bio_cipher, cipher, NULL, NULL, 1);
	BIO_get_cipher_ctx(bio_cipher, &cipher_ctx);
	if (EVP_CIPHER_CTX_set_key_length(cipher_ctx, (int)key_len) != 1) {
		warnx("failed to set key length");
		openssl_warn();
		retval = 1;
		goto out;
	}
	if (EVP_CipherInit_ex(cipher_ctx, NULL, NULL, key, iv, 1) != 1) {
		warnx("failed to initialize cipher");
		openssl_warn();
		retval = 1;
		goto out;
	}
	BIO_push(bio_cipher, bio_out);

	/* write header */
	version = htonl(SENCRYPT_FORMAT_VERSION);
	write_len = BIO_write(bio_out, &version, sizeof (version));
	if (write_len != sizeof (version)) {
		warnx("failed to write version to output file");
		if (write_len < 0) {
			openssl_warn();
		}
		retval = 1;
		goto out;
	}

	iterations = htonl(PBKDF2_ITERATIONS);
	write_len = BIO_write(bio_out, &iterations, sizeof (iterations));
	if (write_len != sizeof (iterations)) {
		warnx("failed to write iterations to output file");
		if (write_len < 0) {
			openssl_warn();
		}
		retval = 1;
		goto out;
	}

	iv_len = EVP_CIPHER_iv_length(cipher);
	if (iv_len > 0) {
		write_len = BIO_write(bio_out, iv, iv_len);
		if (write_len != iv_len) {
			warnx("failed to write IV to output file");
			if (write_len < 0) {
				openssl_warn();
			}
			retval = 1;
			goto out;
		}
	}

	write_len = BIO_write(bio_out, salt, SALT_LEN);
	if (write_len != SALT_LEN) {
		warnx("failed to write salt to output file");
		if (write_len < 0) {
			openssl_warn();
		}
		retval = 1;
		goto out;
	}

	if (BIO_flush(bio_out) < 1) {
		warnx("failed to flush output file");
		openssl_warn();
		retval = 1;
		goto out;
	}

	buf = malloc(BUFFER_SIZE);
	if (buf == NULL) {
		warn(NULL);
		retval = 1;
		goto out;
	}

	/* encrypt data */
	while ((read_len = BIO_read(bio_in, buf, BUFFER_SIZE)) > 0) {
		if ((write_len = BIO_write(bio_cipher, buf, read_len)) !=
		    read_len) {
			warnx("failed to write to output file");
			if (write_len < 0) {
				openssl_warn();
			}
			retval = 1;
			goto out;
		}
	}
	if (read_len < 0) {
		warnx("failed to read from input file");
		openssl_warn();
		retval = 1;
		goto out;
	}

	if (BIO_flush(bio_cipher) < 1) {
		warnx("failed to flush output file");
		openssl_warn();
		retval = 1;
		goto out;
	}

out:
	free(buf);

	if (bio_cipher != NULL) {
		BIO_pop(bio_cipher);
		BIO_free(bio_cipher);
	}

	return (retval);
}

static int
sdecrypt(const EVP_CIPHER *cipher, BIO *bio_in, BIO *bio_out,
    const unsigned char *key, size_t key_len, const unsigned char *iv)
{
	int		read_len;
	BIO		*bio_cipher = NULL;
	int		write_len;
	char		*buf = NULL;
	EVP_CIPHER_CTX	*cipher_ctx;
	int		retval = 0;

	buf = malloc(BUFFER_SIZE);
	if (buf == NULL) {
		warn(NULL);
		retval = 1;
		goto out;
	}

	/* set up cipher filter */
	bio_cipher = BIO_new(BIO_f_cipher());
	BIO_set_cipher(bio_cipher, cipher, NULL, NULL, 0);
	BIO_get_cipher_ctx(bio_cipher, &cipher_ctx);
	if (EVP_CIPHER_CTX_set_key_length(cipher_ctx, (int)key_len) != 1) {
		warnx("failed to set key length");
		openssl_warn();
		retval = 1;
		goto out;
	}
	if (EVP_CipherInit_ex(cipher_ctx, NULL, NULL, key, iv, 0) != 1) {
		warnx("failed to initialize cipher");
		openssl_warn();
		retval = 1;
		goto out;
	}
	BIO_push(bio_cipher, bio_in);

	/* decrypt data */
	while ((read_len = BIO_read(bio_cipher, buf, BUFFER_SIZE)) > 0) {
		if ((write_len = BIO_write(bio_out, buf, read_len)) !=
		    read_len) {
			warnx("failed to write to to output file");
			if (write_len < 0) {
				openssl_warn();
			}
			retval = 1;
			goto out;
		}
	}
	if (read_len < 0) {
		warnx("failed to read from input file");
		openssl_warn();
		retval = 1;
		goto out;
	}

	if (BIO_flush(bio_out) < 1) {
		warnx("failed to flush output file");
		openssl_warn();
		retval = 1;
		goto out;
	}

	if (BIO_get_cipher_status(bio_cipher) == 0) {
		warnx("decryption failed");
		openssl_warn();
		retval = 1;
		goto out;
	}

out:
	free(buf);

	if (bio_cipher != NULL) {
		BIO_pop(bio_cipher);
		BIO_free(bio_cipher);
	}

	return (retval);
}

static void
list_algorithms(void)
{
	printf("Algorithm       Keysize:  Min   Max (bits)\n"
	    "------------------------------------------\n");
	printf("%-15s         %5u %5u\n", "aes", 128, 256);
	printf("%-15s         %5u %5u\n", "arcfour", 8,
	    EVP_MAX_KEY_LENGTH * 8);
	printf("%-15s         %5u %5u\n", "des", 64, 64);
	printf("%-15s         %5u %5u\n", "3des", 192, 192);
}

static void
usage(int cmd)
{
	if (cmd == CMD_SENCRYPT) {
		fprintf(stderr, "usage: sencrypt -l | [-v] -a algorithm "
		    "[-k key_file] [-i input_file] [-o output_file]\n");
	} else if (cmd == CMD_SDECRYPT) {
		fprintf(stderr, "usage: sdecrypt -l | [-v] -a algorithm "
		    "[-k key_file] [-i input_file] [-o output_file]\n");
	}
}

int
main(int argc, char *argv[])
{
	char		*progname;
	int		cmd;
	int		c;
	bool		aflag = false;
	char		*algo_name = NULL;
	bool		is_algo_rc4 = false;
	bool		iflag = false;
	char		*in_filename = NULL;
	bool		kflag = false;
	char		*key_filename = NULL;
	bool		lflag = false;
	bool		oflag = false;
	char		*out_filename = NULL;
	bool		vflag = false;
	bool		errflag = false;
	unsigned char	key[EVP_MAX_KEY_LENGTH];
	size_t		key_len = 0;
	size_t		key_file_len;
	const EVP_CIPHER	*cipher;
	BIO		*bio_in = NULL;
	uint32_t	iterations = PBKDF2_ITERATIONS;
	unsigned char	iv[EVP_MAX_IV_LENGTH];
	unsigned char	salt[SALT_LEN];
	BIO		*bio_out = NULL;
	int		need_tmpfile = 0;
	FILE		*fp_in;
	struct stat	statbuf_in;
	struct stat	statbuf_out;
	int		fd_tmp = -1;
	FILE		*fp_tmp = NULL;
	char		*out_filename_tmp = NULL;
	char		*out_dir = NULL;
	char		*tmp_filename = NULL;
	int		len;
	mode_t		old_mode;
	char		pwdata[MAX(MAX_PASSWORD_LEN, EVP_MAX_KEY_LENGTH)];
	size_t		pwdata_len = 0;
	int		status = EXIT_SUCCESS;

	/* initialize OpenSSL */
	OpenSSL_add_all_algorithms();
	ERR_load_crypto_strings();
	OPENSSL_config(NULL);

	progname = strrchr(argv[0], '/');
	progname = (progname != NULL) ? progname + 1 : argv[0];
	if ((strcmp(progname, "sencrypt") == 0) ||
	    (strcmp(progname, "encrypt") == 0)) {
		cmd = CMD_SENCRYPT;
	} else if ((strcmp(progname, "sdecrypt") == 0) ||
	    (strcmp(progname, "decrypt") == 0)) {
		cmd = CMD_SDECRYPT;
	} else {
		fprintf(stderr, "invalid command name");
		status = EXIT_FAILURE;
		goto out;
	}

	while (!errflag && (c = getopt(argc, argv, "a:i:k:lo:v")) != -1) {
		switch (c) {
		case 'a':
			aflag = true;
			algo_name = optarg;
			is_algo_rc4 = (strcmp(algo_name, "arcfour") == 0);
			break;
		case 'i':
			iflag = true;
			in_filename = optarg;
			break;
		case 'k':
			kflag = true;
			key_filename = optarg;
			break;
		case 'l':
			lflag = true;
			break;
		case 'o':
			oflag = true;
			out_filename = optarg;
			break;
		case 'v':
			vflag = true;
			break;
		default:
			errflag = true;
		}
	}
	if (errflag || (!lflag && !aflag) || (lflag && aflag) ||
	    (argc > optind)) {
		usage(cmd);
		status = EXIT_USAGE;
		goto out;
	}

	if (lflag) {
		list_algorithms();
		goto out;
	}

	if (kflag) {
		key_file_len = read_keyfile(key_filename, key,
		    (off_t)sizeof (key));
		if (key_file_len < 1) {
			status = EXIT_FAILURE;
			goto out;
		}
		key_len = key_file_len;
	} else {
		if (EVP_read_pw_string(pwdata, sizeof (pwdata), "Enter key:",
		    (cmd == CMD_SENCRYPT) ? 1 : 0) != 0) {
			warnx("could not read passphrase");
			openssl_warn();
			status = EXIT_FAILURE;
			goto out;
		}
		pwdata_len = strlen(pwdata);
		if (pwdata_len < 1) {
			warnx("invalid passphrase");
			status = EXIT_FAILURE;
			goto out;
		}
	}

	/* the cipher is determined based on name and length of the key file */
	if (find_algorithm(algo_name, &cipher, &key_len) == -1) {
		status = EXIT_FAILURE;
		goto out;
	}
	if ((cmd == CMD_SENCRYPT) && ((cipher != EVP_aes_128_cbc()) &&
	    (cipher != EVP_aes_192_cbc()) && (cipher != EVP_aes_256_cbc()))) {
		fprintf(stderr, "warning: the %s algorithm is no longer "
		    "considered secure", algo_name);
	}

	if (iflag) {
		bio_in = BIO_new_file(in_filename, "r");
	} else {
		bio_in = BIO_new_fp(stdin, BIO_NOCLOSE);
	}
	if (bio_in == NULL) {
		warnx("could not open input file");
		openssl_warn();
		status = EXIT_FAILURE;
		goto out;
	}

	if (cmd == CMD_SENCRYPT) {
		/* generate random salt and IV */
		if ((RAND_bytes(salt, sizeof (salt)) != 1) ||
		    (RAND_bytes(iv, EVP_CIPHER_iv_length(cipher)) != 1)) {
			/* not enough entropy or unknown error */
			warnx("failed to generate random data");
			status = EXIT_FAILURE;
			goto out;
		}
	} else {
		read_header(bio_in, &iterations, iv,
		    EVP_CIPHER_iv_length(cipher), salt, (int)sizeof (salt));
	}

	/*
	 * if no keyfile was given or the RC4 cipher is used, derive the key
	 * from the password and salt
	 */
	if (kflag && is_algo_rc4) {
		memcpy(pwdata, key, key_file_len);
		pwdata_len = key_file_len;
	}
	if (!kflag || is_algo_rc4) {
		if (PKCS5_PBKDF2_HMAC_SHA1(pwdata, (int)pwdata_len, salt,
		    sizeof (salt), (int)iterations, (int)key_len, key) == 0) {
			warnx("failed to generate key");
			status = EXIT_FAILURE;
			goto out;
		}
	}

	if (oflag) {
		/*
		 * if input and output files are identical, create and write the
		 * output to a temporary file for the output which is then
		 * renamed to out_filename
		 */
		if (iflag) {
			BIO_get_fp(bio_in, &fp_in);
			if (fstat(fileno(fp_in), &statbuf_in) == -1) {
				warn("could not stat input file");
				status = EXIT_FAILURE;
				goto out;
			}
			if (stat(out_filename, &statbuf_out) == -1) {
				if (errno != ENOENT) {
					warn("could not stat output file");
					status = EXIT_FAILURE;
					goto out;
				}
			} else if ((statbuf_in.st_ino == statbuf_out.st_ino) &&
				    (statbuf_in.st_dev == statbuf_out.st_dev)) {
				need_tmpfile = 1;
			}
		}

		if (need_tmpfile) {
			out_filename_tmp = strdup(out_filename);
			if (out_filename_tmp == NULL) {
				warn(NULL);
				status = EXIT_FAILURE;
				goto out;
			}
			out_dir = dirname(out_filename_tmp);
			len = snprintf(NULL, 0, "%s/sencryptXXXXXX", out_dir);
			if (len < 0) {
				warn(NULL);
				status = EXIT_FAILURE;
				goto out;
			}
			tmp_filename = malloc((size_t)len + 1);
			if (tmp_filename == NULL) {
				warn(NULL);
				status = EXIT_FAILURE;
				goto out;
			}
			if (snprintf(tmp_filename, (size_t)len + 1,
			    "%s/sencryptXXXXXX", out_dir) != len) {
				warn(NULL);
				status = EXIT_FAILURE;
				goto out;
			}
			old_mode = umask(077);
			fd_tmp = mkstemp(tmp_filename);
			umask(old_mode);
			if (fd_tmp == -1) {
				warn("could not create temporary file");
				status = EXIT_FAILURE;
				goto out;
			}
			fp_tmp = fdopen(fd_tmp, "w");
			if (fp_tmp == NULL) {
				warn("could not open temporary file");
				status = EXIT_FAILURE;
				goto out;
			}
			fd_tmp = -1;
			bio_out = BIO_new_fp(fp_tmp, BIO_CLOSE);
			if (bio_out == NULL) {
				warnx("could not open temporary file");
				openssl_warn();
				status = EXIT_FAILURE;
				goto out;
			}
			fp_tmp = NULL;
		} else {
			old_mode = umask(077);
			bio_out = BIO_new_file(out_filename, "w");
			umask(old_mode);
			if (bio_out == NULL) {
				warnx("could not open output file");
				openssl_warn();
				status = EXIT_FAILURE;
				goto out;
			}
		}
	} else {
		bio_out = BIO_new_fp(stdout, BIO_NOCLOSE);
		if (bio_out == NULL) {
			warnx("could not open output file");
			openssl_warn();
			status = EXIT_FAILURE;
			goto out;
		}
	}

	if (cmd == CMD_SENCRYPT) {
		if (sencrypt(cipher, bio_in, bio_out, key, key_len,
		    iv, salt) == -1) {
			status = EXIT_FAILURE;
		}
	} else {
		if (sdecrypt(cipher, bio_in, bio_out, key, key_len,
		    iv) == -1) {
			status = EXIT_FAILURE;
		}
	}

out:
	OPENSSL_cleanse(pwdata, pwdata_len);
	OPENSSL_cleanse(key, key_len);

	if (fd_tmp != -1) {
		close(fd_tmp);
	}

	if (fp_tmp != NULL) {
		fclose(fp_tmp);
	}

	if (bio_in != NULL) {
		BIO_free_all(bio_in);
	}

	if (bio_out != NULL) {
		BIO_free_all(bio_out);

		if (status == 0) {
			if (need_tmpfile) {
				if (rename(tmp_filename, out_filename) == -1) {
					warn("could not create output file");
					status = EXIT_FAILURE;
					unlink(tmp_filename);
				}
			}
		} else {
			if (need_tmpfile) {
				unlink(tmp_filename);
			} else if (oflag) {
				unlink(out_filename);
			}
		}
	}

	free(out_filename_tmp);
	free(tmp_filename);

	EVP_cleanup();
	ERR_free_strings();
	CONF_modules_free();

	exit(status);
}