#include "WedgeSecurity.h"
#include <openssl/evp.h>
#include <openssl/rand.h>
#include <QByteArray>

static const int ITERATIONS = 35690;
static const int KEY_LENGTH = 32;
static const int SALT_LENGTH = 16;
static const int IV_LENGTH = 12;
static const QString PREFIX = "wgd:";

bool WedgeSecurity::isEncrypted(const QString& text) {
	return text.startsWith(PREFIX);
}

QString WedgeSecurity::encrypt(const QString& plainText, const QString& passPhrase) {
	QByteArray salt(SALT_LENGTH, 0);
	RAND_bytes((unsigned char*)salt.data(), SALT_LENGTH);

	QByteArray iv(IV_LENGTH, 0);
	RAND_bytes((unsigned char*)iv.data(), IV_LENGTH);

	QByteArray key(KEY_LENGTH, 0);
	QByteArray phrase = passPhrase.toUtf8();
	PKCS5_PBKDF2_HMAC(phrase.constData(), phrase.length(),
		(const unsigned char*)salt.constData(), salt.length(),
		ITERATIONS, EVP_sha256(),
		KEY_LENGTH, (unsigned char*)key.data());

	EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();
	EVP_EncryptInit_ex(ctx, EVP_aes_256_gcm(), nullptr, nullptr, nullptr);
	EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_IVLEN, IV_LENGTH, nullptr);
	EVP_EncryptInit_ex(ctx, nullptr, nullptr, (unsigned char*)key.data(), (unsigned char*)iv.data());

	QByteArray pText = plainText.toUtf8();
	QByteArray cText(pText.length() + EVP_MAX_BLOCK_LENGTH, 0);
	int len = 0, ciphertext_len = 0;

	EVP_EncryptUpdate(ctx, (unsigned char*)cText.data(), &len, (const unsigned char*)pText.constData(), pText.length());
	ciphertext_len = len;

	EVP_EncryptFinal_ex(ctx, (unsigned char*)cText.data() + len, &len);
	ciphertext_len += len;
	cText.resize(ciphertext_len);

	QByteArray tag(16, 0);
	EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_GET_TAG, 16, tag.data());
	EVP_CIPHER_CTX_free(ctx);

	QByteArray combined = salt + iv + cText + tag;
	return PREFIX + QString::fromLatin1(combined.toBase64());
}

QString WedgeSecurity::decrypt(const QString& cipherTextBase64, const QString& passPhrase) {
	if (!isEncrypted(cipherTextBase64)) return QString();

	QString cleanInput = cipherTextBase64.mid(PREFIX.length());
	QByteArray combined = QByteArray::fromBase64(cleanInput.toLatin1());
	if (combined.length() < SALT_LENGTH + IV_LENGTH + 16) return QString();

	QByteArray salt = combined.left(SALT_LENGTH);
	QByteArray iv = combined.mid(SALT_LENGTH, IV_LENGTH);
	QByteArray encrypted = combined.mid(SALT_LENGTH + IV_LENGTH, combined.length() - SALT_LENGTH - IV_LENGTH - 16);
	QByteArray tag = combined.right(16);

	QByteArray key(KEY_LENGTH, 0);
	QByteArray phrase = passPhrase.toUtf8();
	PKCS5_PBKDF2_HMAC(phrase.constData(), phrase.length(),
		(const unsigned char*)salt.constData(), salt.length(),
		ITERATIONS, EVP_sha256(),
		KEY_LENGTH, (unsigned char*)key.data());

	EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();
	EVP_DecryptInit_ex(ctx, EVP_aes_256_gcm(), nullptr, nullptr, nullptr);
	EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_IVLEN, IV_LENGTH, nullptr);
	EVP_DecryptInit_ex(ctx, nullptr, nullptr, (unsigned char*)key.data(), (unsigned char*)iv.data());

	QByteArray pText(encrypted.length() + EVP_MAX_BLOCK_LENGTH, 0);
	int len = 0, plaintext_len = 0;

	EVP_DecryptUpdate(ctx, (unsigned char*)pText.data(), &len, (const unsigned char*)encrypted.constData(), encrypted.length());
	plaintext_len = len;

	EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_TAG, 16, tag.data());
	int ret = EVP_DecryptFinal_ex(ctx, (unsigned char*)pText.data() + len, &len);
	EVP_CIPHER_CTX_free(ctx);

	if (ret > 0) {
		plaintext_len += len;
		pText.resize(plaintext_len);
		return QString::fromUtf8(pText);
	}
	return QString();
}