/* $OpenBSD: crypto.c,v 1.4 2013/12/26 17:25:32 eric Exp $	 */

/*
 * Copyright (c) 2013 Gilles Chehade <gilles@openbsd.org>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#include <sys/types.h>
#include <sys/stat.h>

#include <stdlib.h>
#include <string.h>

#include <openssl/evp.h>


#define	CRYPTO_BUFFER_SIZE	16384

#define	GCM_TAG_SIZE		16
#define	IV_SIZE			12
#define	KEY_SIZE		32

/* bump if we ever switch from aes-256-gcm to anything else */
#define	API_VERSION    		1


int	crypto_setup(const char *, size_t);
int	crypto_encrypt_file(FILE *, FILE *);
int	crypto_decrypt_file(FILE *, FILE *);
size_t	crypto_encrypt_buffer(const char *, size_t, char *, size_t);
size_t	crypto_decrypt_buffer(const char *, size_t, char *, size_t);

static struct crypto_ctx {
	const EVP_CIPHER       *cipher;
	unsigned char  		key[KEY_SIZE];
} cp;

int
crypto_setup(const char *key, size_t len)
{
	if (len != KEY_SIZE)
		return 0;

	memset(&cp, 0, sizeof cp);
	cp.cipher = EVP_aes_256_gcm();

	/* openssl rand -hex 16 */
	memcpy(cp.key, key, sizeof cp.key);

	return 1;
}

int
crypto_encrypt_file(FILE * in, FILE * out)
{
	EVP_CIPHER_CTX	ctx;
	uint8_t		ibuf[CRYPTO_BUFFER_SIZE];
	uint8_t		obuf[CRYPTO_BUFFER_SIZE];
	uint8_t		iv[IV_SIZE];
	uint8_t		tag[GCM_TAG_SIZE];
	uint8_t		version = API_VERSION;
	size_t		r, w;
	int		len;
	int		ret = 0;
	struct stat	sb;

	/* XXX - Do NOT encrypt files bigger than 64GB */
	if (fstat(fileno(in), &sb) < 0)
		return 0;
	if (sb.st_size >= 0x1000000000LL)
		return 0;

	/* prepend version byte*/
	if ((w = fwrite(&version, 1, sizeof version, out)) != sizeof version)
		return 0;

	/* generate and prepend IV */
	memset(iv, 0, sizeof iv);
	arc4random_buf(iv, sizeof iv);
	if ((w = fwrite(iv, 1, sizeof iv, out)) != sizeof iv)
		return 0;

	EVP_CIPHER_CTX_init(&ctx);
	EVP_EncryptInit(&ctx, cp.cipher, cp.key, iv);

	/* encrypt until end of file */
	while ((r = fread(ibuf, 1, CRYPTO_BUFFER_SIZE, in)) != 0) {
		if (!EVP_EncryptUpdate(&ctx, obuf, &len, ibuf, r))
			goto end;
		if (len && (w = fwrite(obuf, len, 1, out)) != 1)
			goto end;
	}
	if (!feof(in))
		goto end;

	/* finalize and write last chunk if any */
	if (!EVP_EncryptFinal(&ctx, obuf, &len))
		goto end;
	if (len && (w = fwrite(obuf, len, 1, out)) != 1)
		goto end;

	/* get and append tag */
	EVP_CIPHER_CTX_ctrl(&ctx, EVP_CTRL_GCM_GET_TAG, sizeof tag, tag);
	if ((w = fwrite(tag, sizeof tag, 1, out)) != 1)
		goto end;

	fflush(out);
	ret = 1;

end:
	EVP_CIPHER_CTX_cleanup(&ctx);
	return ret;
}

int
crypto_decrypt_file(FILE * in, FILE * out)
{
	EVP_CIPHER_CTX	ctx;
	uint8_t		ibuf[CRYPTO_BUFFER_SIZE];
	uint8_t		obuf[CRYPTO_BUFFER_SIZE];
	uint8_t		iv[IV_SIZE];
	uint8_t		tag[GCM_TAG_SIZE];
	uint8_t		version;
	size_t		r, w;
	off_t		sz;
	int		len;
	int		ret = 0;
	struct stat	sb;

	/* input file too small to be an encrypted file */
	if (fstat(fileno(in), &sb) < 0)
		return 0;
	if (sb.st_size <= (off_t) (sizeof version + sizeof tag + sizeof iv))
		return 0;
	sz = sb.st_size;

	/* extract tag */
	if (fseek(in, -sizeof(tag), SEEK_END) == -1)
		return 0;
	if ((r = fread(tag, 1, sizeof tag, in)) != sizeof tag)
		return 0;

	if (fseek(in, 0, SEEK_SET) == -1)
		return 0;

	/* extract version */
	if ((r = fread(&version, 1, sizeof version, in)) != sizeof version)
		return 0;
	if (version != API_VERSION)
		return 0;

	/* extract IV */
	memset(iv, 0, sizeof iv);
	if ((r = fread(iv, 1, sizeof iv, in)) != sizeof iv)
		return 0;

	/* real ciphertext length */
	sz -= sizeof version;
	sz -= sizeof iv;
	sz -= sizeof tag;


	EVP_CIPHER_CTX_init(&ctx);
	EVP_DecryptInit(&ctx, cp.cipher, cp.key, iv);

	/* set expected tag */
	EVP_CIPHER_CTX_ctrl(&ctx, EVP_CTRL_GCM_SET_TAG, sizeof tag, tag);

	/* decrypt until end of ciphertext */
	while (sz) {
		if (sz > CRYPTO_BUFFER_SIZE)
			r = fread(ibuf, 1, CRYPTO_BUFFER_SIZE, in);
		else
			r = fread(ibuf, 1, sz, in);
		if (!r)
			break;
		if (!EVP_DecryptUpdate(&ctx, obuf, &len, ibuf, r))
			goto end;
		if (len && (w = fwrite(obuf, len, 1, out)) != 1)
			goto end;
		sz -= r;
	}
	if (ferror(in))
		goto end;

	/* finalize, write last chunk if any and perform authentication check */
	if (!EVP_DecryptFinal(&ctx, obuf, &len))
		goto end;
	if (len && (w = fwrite(obuf, len, 1, out)) != 1)
		goto end;

	fflush(out);
	ret = 1;

end:
	EVP_CIPHER_CTX_cleanup(&ctx);
	return ret;
}

size_t
crypto_encrypt_buffer(const char *in, size_t inlen, char *out, size_t outlen)
{
	EVP_CIPHER_CTX	ctx;
	uint8_t		iv[IV_SIZE];
	uint8_t		tag[GCM_TAG_SIZE];
	uint8_t		version = API_VERSION;
	off_t		sz;
	int		olen;
	int		len = 0;
	int		ret = 0;

	/* output buffer does not have enough room */
	if (outlen < inlen + sizeof version + sizeof tag + sizeof iv)
		return 0;

	/* input should not exceed 64GB */
	sz = inlen;
	if (sz >= 0x1000000000LL)
		return 0;

	/* prepend version */
	*out = version;
	len++;

	/* generate IV */
	memset(iv, 0, sizeof iv);
	arc4random_buf(iv, sizeof iv);
	memcpy(out + len, iv, sizeof iv);
	len += sizeof iv;

	EVP_CIPHER_CTX_init(&ctx);
	EVP_EncryptInit(&ctx, cp.cipher, cp.key, iv);

	/* encrypt buffer */
	if (!EVP_EncryptUpdate(&ctx, out + len, &olen, in, inlen))
		goto end;
	len += olen;

	/* finalize and write last chunk if any */
	if (!EVP_EncryptFinal(&ctx, out + len, &olen))
		goto end;
	len += olen;

	/* get and append tag */
	EVP_CIPHER_CTX_ctrl(&ctx, EVP_CTRL_GCM_GET_TAG, sizeof tag, tag);
	memcpy(out + len, tag, sizeof tag);
	ret = len + sizeof tag;

end:
	EVP_CIPHER_CTX_cleanup(&ctx);
	return ret;
}

size_t
crypto_decrypt_buffer(const char *in, size_t inlen, char *out, size_t outlen)
{
	EVP_CIPHER_CTX	ctx;
	uint8_t		iv[IV_SIZE];
	uint8_t		tag[GCM_TAG_SIZE];
	int		olen;
	int		len = 0;
	int		ret = 0;

	/* out does not have enough room */
	if (outlen < inlen - sizeof tag + sizeof iv)
		return 0;

	/* extract tag */
	memcpy(tag, in + inlen - sizeof tag, sizeof tag);
	inlen -= sizeof tag;

	/* check version */
	if (*in != API_VERSION)
		return 0;
	in++;
	inlen--;

	/* extract IV */
	memset(iv, 0, sizeof iv);
	memcpy(iv, in, sizeof iv);
	inlen -= sizeof iv;
	in += sizeof iv;

	EVP_CIPHER_CTX_init(&ctx);
	EVP_DecryptInit(&ctx, cp.cipher, cp.key, iv);

	/* set expected tag */
	EVP_CIPHER_CTX_ctrl(&ctx, EVP_CTRL_GCM_SET_TAG, sizeof tag, tag);

	/* decrypt buffer */
	if (!EVP_DecryptUpdate(&ctx, out, &olen, in, inlen))
		goto end;
	len += olen;

	/* finalize, write last chunk if any and perform authentication check */
	if (!EVP_DecryptFinal(&ctx, out + len, &olen))
		goto end;
	ret = len + olen;

end:
	EVP_CIPHER_CTX_cleanup(&ctx);
	return ret;
}

#if 0
int
main(int argc, char *argv[])
{
	if (argc != 3) {
		printf("usage: crypto <key> <buffer>\n");
		return 1;
	}

	if (! crypto_setup(argv[1], strlen(argv[1]))) {
		printf("crypto_setup failed\n");
		return 1;
	}

	{
		char            encbuffer[4096];
		size_t          enclen;
		char            decbuffer[4096];
		size_t          declen;

		printf("encrypt/decrypt buffer: ");
		enclen = crypto_encrypt_buffer(argv[2], strlen(argv[2]),
					       encbuffer, sizeof encbuffer);

		/* uncomment below to provoke integrity check failure */
		/*
		 * encbuffer[13] = 0x42;
		 * encbuffer[14] = 0x42;
		 * encbuffer[15] = 0x42;
		 * encbuffer[16] = 0x42;
		 */

		declen = crypto_decrypt_buffer(encbuffer, enclen,
					       decbuffer, sizeof decbuffer);
		if (declen != 0 && !strncmp(argv[2], decbuffer, declen))
			printf("ok\n");
		else
			printf("nope\n");
	}

	{
		FILE           *fpin;
		FILE           *fpout;
		printf("encrypt/decrypt file: ");

		fpin = fopen("/etc/passwd", "r");
		fpout = fopen("/tmp/passwd.enc", "w");
		if (!crypto_encrypt_file(fpin, fpout)) {
			printf("encryption failed\n");
			return 1;
		}
		fclose(fpin);
		fclose(fpout);

		/* uncomment below to provoke integrity check failure */
		/*
		 * fpin = fopen("/tmp/passwd.enc", "a");
		 * fprintf(fpin, "borken");
		 * fclose(fpin);
		 */
		fpin = fopen("/tmp/passwd.enc", "r");
		fpout = fopen("/tmp/passwd.dec", "w");
		if (!crypto_decrypt_file(fpin, fpout))
			printf("nope\n");
		else
			printf("ok\n");
		fclose(fpin);
		fclose(fpout);
	}


	return 0;
}
#endif