#include <openssl/conf.h>
#include <openssl/err.h>
#include <openssl/evp.h>
#include <stdio.h>
#include <string.h>

// based on OpenSSL sample
// refer to https://wiki.openssl.org/index.php/EVP_Authenticated_Encryption_and_Decryption

void handleErrors(void);
int gcm_encrypt(unsigned char *plaintext, int plaintext_len, unsigned char *key, unsigned char *iv, int iv_len, unsigned char *ciphertext, unsigned char *tag);
int gcm_decrypt(unsigned char *ciphertext, int ciphertext_len, unsigned char *tag, unsigned char *key, unsigned char *iv, int iv_len, unsigned char *plaintext);

int main(int argc, char **argv) {
	/*
	 * Set up the key and iv. Do I need to say to not hard code these in a
	 * real application? :-)
	 */

	FILE *f, *o;
	unsigned char signature[4] = {'C', 'C', 'A', 'T'};
	unsigned char key[32];
	unsigned char iv[12];
	unsigned char tag[16];
	unsigned char *plain;
	unsigned char *cipher;
	unsigned insize, mode;
	int decryptedtext_len, ciphertext_len;
	if (argc < 4) {
		printf("not enough args. use %s mode key in out\n", argv[0]);
		return 1;
	}

	mode = strtoul(argv[1], NULL, 10);

	if (!(f = fopen(argv[2], "rb"))) {
		printf("cannot open key\n");
		return 1;
	}
	fread(key, sizeof(key), 1, f);
	fclose(f);

	if (!(f = fopen("/dev/urandom", "rb"))) {
		printf("cannot open urandom\n");
		return 1;
	}
	fread(iv, sizeof(iv), 1, f);
	fclose(f);

	if (!(f = fopen(argv[3], "rb"))) {
		printf("cannot open in\n");
		return 1;
	}
	if (!(o = fopen(argv[4], "wb"))) {
		printf("cannot open out\n");
		return 1;
	}

	if (mode) {
		printf("using iv ");
		for (unsigned i = 0; i < sizeof(iv); i++)
			printf("%02x ", iv[i]);
		printf("\n");

		printf("using tag ");
		for (unsigned i = 0; i < sizeof(tag); i++)
			printf("%02x ", tag[i]);
		printf("\n");

		fseek(f, 0, SEEK_END);
		insize = ftell(f) + 4;
		fseek(f, 0, SEEK_SET);
		plain = malloc(insize);
		cipher = malloc(insize);
		fread(plain + 4, insize, 1, f);
		fclose(f);
		/* prepend signature */
		memcpy(plain, signature, 4);
		/* Encrypt the plaintext */
		ciphertext_len = gcm_encrypt(plain, insize, key, iv, sizeof(iv), cipher, tag);
		fwrite(iv, sizeof(iv), 1, o);
		fwrite(tag, sizeof(tag), 1, o);
		fwrite(cipher, ciphertext_len, 1, o);
		fclose(o);
	} else {
		fseek(f, 0, SEEK_END);
		insize = ftell(f) - sizeof(iv) - sizeof(tag);
		fseek(f, 0, SEEK_SET);
		cipher = malloc(insize);
		plain = malloc(insize);
		fread(iv, sizeof(iv), 1, f);
		fread(tag, sizeof(tag), 1, f);
		fread(cipher, insize, 1, f);
		fclose(f);

		printf("using iv ");
		for (unsigned i = 0; i < sizeof(iv); i++)
			printf("%02x ", iv[i]);
		printf("\n");

		printf("using tag ");
		for (unsigned i = 0; i < sizeof(tag); i++)
			printf("%02x ", tag[i]);
		printf("\n");

		decryptedtext_len = gcm_decrypt(cipher, insize, tag, key, iv, sizeof(iv), plain);
		if (decryptedtext_len < 0) {
			printf("decrypt failed\n");
			fclose(o);
			remove(argv[4]);
		} else if (memcmp(plain, signature, 4)) {
			printf("signature mismatch, expected ");
			for (int i = 0; i < sizeof(signature); i++)
				printf("%02x ", signature[i]);
			printf("but got ");
			for (int i = 0; i < sizeof(signature); i++)
				printf("%02x ", plain[i]);
			printf("\n");
			fclose(o);
			remove(argv[4]);
		} else {
			fwrite(plain + 4, decryptedtext_len - 4, 1, o);
			fclose(o);
		}
	}
	free(cipher);
	free(plain);
	printf("done\n");
	return 0;
}

void handleErrors(void) {
	ERR_print_errors_fp(stderr);
	abort();
}

int gcm_encrypt(unsigned char *plaintext, int plaintext_len, unsigned char *key, unsigned char *iv, int iv_len, unsigned char *ciphertext, unsigned char *tag) {
	EVP_CIPHER_CTX *ctx;

	int len;

	int ciphertext_len;

	/* Create and initialise the context */
	if (!(ctx = EVP_CIPHER_CTX_new()))
		handleErrors();

	/* Initialise the encryption operation. */
	if (1 != EVP_EncryptInit_ex(ctx, EVP_aes_256_gcm(), NULL, NULL, NULL))
		handleErrors();

	/*
	 * Set IV length if default 12 bytes (96 bits) is not appropriate
	 */
	if (1 != EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_IVLEN, iv_len, NULL))
		handleErrors();

	/* Initialise key and IV */
	if (1 != EVP_EncryptInit_ex(ctx, NULL, NULL, key, iv))
		handleErrors();

	/*
	 * Provide the message to be encrypted, and obtain the encrypted output.
	 * EVP_EncryptUpdate can be called multiple times if necessary
	 */
	if (1 != EVP_EncryptUpdate(ctx, ciphertext, &len, plaintext, plaintext_len))
		handleErrors();
	ciphertext_len = len;

	/*
	 * Finalise the encryption. Normally ciphertext bytes may be written at
	 * this stage, but this does not occur in GCM mode
	 */
	if (1 != EVP_EncryptFinal_ex(ctx, ciphertext + len, &len))
		handleErrors();
	ciphertext_len += len;

	/* Get the tag */
	if (1 != EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_GET_TAG, 16, tag))
		handleErrors();

	/* Clean up */
	EVP_CIPHER_CTX_free(ctx);

	return ciphertext_len;
}

int gcm_decrypt(unsigned char *ciphertext, int ciphertext_len, unsigned char *tag, unsigned char *key, unsigned char *iv, int iv_len,
                unsigned char *plaintext) {
	EVP_CIPHER_CTX *ctx;
	int len;
	int plaintext_len;
	int ret;

	/* Create and initialise the context */
	if (!(ctx = EVP_CIPHER_CTX_new()))
		handleErrors();

	/* Initialise the decryption operation. */
	if (!EVP_DecryptInit_ex(ctx, EVP_aes_256_gcm(), NULL, NULL, NULL))
		handleErrors();

	/* Set IV length. Not necessary if this is 12 bytes (96 bits) */
	if (!EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_IVLEN, iv_len, NULL))
		handleErrors();

	/* Initialise key and IV */
	if (!EVP_DecryptInit_ex(ctx, NULL, NULL, key, iv))
		handleErrors();

	/*
	 * Provide the message to be decrypted, and obtain the plaintext output.
	 * EVP_DecryptUpdate can be called multiple times if necessary
	 */
	if (!EVP_DecryptUpdate(ctx, plaintext, &len, ciphertext, ciphertext_len))
		handleErrors();
	plaintext_len = len;

	/* Set expected tag value. Works in OpenSSL 1.0.1d and later */
	if (!EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_TAG, 16, tag))
		handleErrors();

	/*
	 * Finalise the decryption. A positive return value indicates success,
	 * anything else is a failure - the plaintext is not trustworthy.
	 */
	ret = EVP_DecryptFinal_ex(ctx, plaintext + len, &len);

	/* Clean up */
	EVP_CIPHER_CTX_free(ctx);

	if (ret > 0) {
		/* Success */
		plaintext_len += len;
		return plaintext_len;
	} else {
		/* Verify failed */
		return -1;
	}
}