#include <openssl/evp.h>
#include <openssl/err.h>
#include <openssl/aes.h>
#include <string>
#include <iostream>
#include <fstream>

#include "cryptography.h"


Cryptography::Cryptography(const char* password, size_t size)
{
	OpenSSL_add_all_algorithms();
	ERR_load_crypto_strings();

	if (!generate_key_and_iv_from_password(password, size))
	{
		printf_s("Error generating key and IV from password\n");
		return;
	}
}

Cryptography::~Cryptography()
{
	ERR_free_strings();
	EVP_cleanup();
}

bool Cryptography::encrypt(Buffer* plain, Buffer* encrypted)
{
	encrypted->resize(plain->taken + AES_BLOCK_SIZE);

	ctx = EVP_CIPHER_CTX_new();
	if (!ctx) return handleErrors();
	if (1 != EVP_EncryptInit_ex(ctx, EVP_aes_256_cbc(), NULL, key, iv)) return handleErrors();

	if (1 != EVP_EncryptUpdate(ctx, encrypted->buffer, (int*)& encrypted->taken, plain->buffer, plain->taken)) return handleErrors();
	int final_len;
	if (1 != EVP_EncryptFinal_ex(ctx, encrypted->buffer + encrypted->taken, &final_len)) return handleErrors();
	encrypted->taken += final_len;

	EVP_CIPHER_CTX_free(ctx);


	return true;
}

bool Cryptography::decrypt(Buffer* encrypted, Buffer* decrypted)
{
	decrypted->resize(encrypted->taken + AES_BLOCK_SIZE);


	ctx = EVP_CIPHER_CTX_new();
	if (!ctx) return handleErrors();
	if (1 != EVP_DecryptInit_ex(ctx, EVP_aes_256_cbc(), NULL, key, iv)) return handleErrors();

	if (1 != EVP_DecryptUpdate(ctx, decrypted->buffer, (int*) & decrypted->taken, encrypted->buffer, encrypted->taken)) return handleErrors();
	int final_len;
	if (1 != EVP_DecryptFinal_ex(ctx, decrypted->buffer + decrypted->taken, &final_len)) return handleErrors();
	decrypted->taken += final_len;

	EVP_CIPHER_CTX_free(ctx);
	return true;
}

bool Cryptography::generate_key_and_iv_from_password(const char* password, size_t size)
{
	int iterations = 10000;

	// Derive key and IV using PBKDF2
	if (1 != PKCS5_PBKDF2_HMAC(password, size, nullptr, 0, iterations, EVP_sha256(), 32, key)) return false;
	if (1 != PKCS5_PBKDF2_HMAC(password, size, nullptr, 0, iterations, EVP_sha256(), 16, iv)) return false;
	return true;
}

bool Cryptography::handleErrors()
{
	EVP_CIPHER_CTX_free(ctx);

	while (auto error = ERR_get_error())
	{
		char* error_string = ERR_error_string(error, nullptr);
		printf_s("%s\n", error_string);
	}
	return false;
}