extern "C" {
#include <openssl/evp.h>
#include <openssl/hmac.h>
#include <openssl/aes.h>
#include <openssl/rand.h>
}

#include "totpgen.h"

#define UNUSED(x) (void(x))

#define SALT_BYTES 8
#define KEY_BYTES  32
#define IV_BYTES   16


#define AES_BUFFER(LEN) \
	reinterpret_cast<unsigned char *>( \
		calloc(AES_BLOCK_SIZE * (1 + ((LEN) / AES_BLOCK_SIZE)), \
			sizeof(unsigned char)));

bool
TOTPGen::GenTOTP(std::string& secret_str)
{
	uint32_t bits, code, count, hmacsize;
	uint64_t clock;
	size_t length;
	uint8_t msg[8], hmac[EVP_MAX_MD_SIZE];
	char totp_buf[7];

	size_t keysize = 0;
	const EVP_MD *digest = EVP_sha1();
	uint8_t *key = NULL;

	char *secret = strdup(secret_str.c_str());

	for (count = bits = length = 0; *secret; secret++) {
		if (!strchr(base32.c_str(), *secret)) {
			return false;
		}
		bits = (bits << 5) | (strchr(base32.c_str(), *secret) - base32.c_str());
		count += 5;

		while (count >= 8) {
			while (length >= keysize) {
				keysize = keysize ? keysize << 1 : 64;
				if (!(key = (uint8_t *)realloc(key, keysize))) {
					return false;
				}
			}
			count -= 8, key[length++] = bits >> count;
		}
	}

	if (length > 0) {
		clock = time(NULL) / 30;

		for (count = 0; count < 8; count++) {
			msg[7 - count] = clock >> 8 * count;
		}
		if (!HMAC(digest, key, length, msg, sizeof(msg), hmac, &hmacsize)) {
			return false;
		}

		for (code = count = 0; count < 4; count++) {
			code += hmac[(hmac[hmacsize - 1] & 0x0f) + 3 - count] << 8 * count;
		}
		code &= 0x7fffffff;

		snprintf((char *)&totp_buf, 7, "%06u", code % 1000000);
	}
	else {
		return false;
	}

	TOTP = totp_buf;
	TOTPField->SetLabel(TOTP);

	free(key);

	return true;
}

void
TOTPGen::InitAES(const char *passphrase, unsigned char *salt,
                 unsigned char *key, unsigned char *iv)
{
	char *key_data = strdup(passphrase);
	int key_text_len = strlen(passphrase);
	int nrounds = 0xFF;

	EVP_BytesToKey(EVP_aes_256_cbc(), EVP_sha1(), salt,
	               reinterpret_cast<unsigned char *>(key_data),
	               key_text_len, nrounds, key, iv);

	free(key_data);
}

bool
TOTPGen::EncryptSecret(std::string& secret_str)
{
	unsigned char salt[SALT_BYTES], key[KEY_BYTES], iv[IV_BYTES];

	RAND_bytes(salt, SALT_BYTES);

	memset(&key, 0, KEY_BYTES);
	memset(&iv, 0, IV_BYTES);

	InitAES(Passphrase.c_str(), reinterpret_cast<unsigned char *>(&salt),
	        reinterpret_cast<unsigned char *>(&key),
	        reinterpret_cast<unsigned char *>(&iv));

	EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();
	EVP_EncryptInit_ex(ctx, EVP_aes_256_cbc(), NULL,
	                   reinterpret_cast<unsigned char *>(&key),
	                   reinterpret_cast<unsigned char *>(&iv));

	size_t plaintext_len = strlen(secret_str.c_str()) ;

	unsigned char *plaintext = AES_BUFFER(plaintext_len);
	strncpy(reinterpret_cast<char *>(plaintext),
	        secret_str.c_str(), plaintext_len);

	unsigned char *ciphertext = AES_BUFFER(plaintext_len);

	int len, ciphertext_len;
	EVP_EncryptUpdate(ctx, ciphertext, &len, plaintext, plaintext_len);
	ciphertext_len = len;
	EVP_EncryptFinal_ex(ctx, ciphertext + len, &len);
	ciphertext_len += len;

	FILE *secret_file = fopen(SecretFile.c_str(), "wb");
	if (!secret_file) {
		return false;
	}

	fwrite(salt, SALT_BYTES, 1, secret_file);
	fwrite(ciphertext, ciphertext_len, 1, secret_file);

	fclose(secret_file);

	EVP_CIPHER_CTX_free(ctx);
	free(plaintext);
	free(ciphertext);

	return true;
}


bool
TOTPGen::DecryptSecret(void)
{
	unsigned char salt[SALT_BYTES], key[KEY_BYTES], iv[IV_BYTES];

	memset(&key, 0, KEY_BYTES);
	memset(&iv, 0, IV_BYTES);

	FILE *secret_file = fopen(SecretFile.c_str(), "rb");
	if (!secret_file) {
		return false;
	}

	fseek(secret_file, 0L, SEEK_END);
	size_t ciphertext_len = ftell(secret_file) - SALT_BYTES;
	rewind(secret_file);

	fread(&salt, SALT_BYTES, 1, secret_file);

	InitAES(Passphrase.c_str(), reinterpret_cast<unsigned char *>(&salt),
	        reinterpret_cast<unsigned char *>(&key),
	        reinterpret_cast<unsigned char *>(&iv));

	EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();
	EVP_DecryptInit_ex(ctx, EVP_aes_256_cbc(), NULL,
	                   reinterpret_cast<unsigned char *>(&key),
	                   reinterpret_cast<unsigned char *>(&iv));

	int len;
	unsigned char *ciphertext = AES_BUFFER(ciphertext_len);
	unsigned char *plaintext = AES_BUFFER(ciphertext_len + 8);

	fread(ciphertext, 1, ciphertext_len, secret_file);
	fclose(secret_file);

	if (!EVP_DecryptUpdate(ctx, plaintext, &len, ciphertext, ciphertext_len)) {
		return false;
	}

	if (!EVP_DecryptFinal_ex(ctx, plaintext + len, &len)) {
		return false;
	}

	Secret = reinterpret_cast<const char *>(plaintext);

	EVP_CIPHER_CTX_free(ctx);
	free(ciphertext);
	free(plaintext);

	return true;
}