#include "Cryptor.h"

bool Cryptor::haveKey = false;
unsigned char Cryptor::key[AES_BLOCK_LENGTH];

static void toHex(char *output, const void *input, size_t size) {
	const unsigned char *_input = (const unsigned char *)input;
	for (size_t i = 0; i < size; i++) {
		sprintf(output+(i*2), "%02hhX", _input[i]);
	}
}

static void fromHex(void *output, const char *input) {
	unsigned char *_output = (unsigned char *)output;
	size_t size = strlen(input);
	if (size & 1) {
		throw 1;
	} else {
		size /= 2;
	}
	for (size_t i = 0; i < size; i++) {
		if (sscanf(input+(i*2), "%02hhx", &_output[i]) != 1) {
			throw 1;
		}
	}
}

void Cryptor::encryptAndSave(std::string directory, std::string payload) {
	// Key
	if (!haveKey) {
		assembleKey(true);
	}

	// IV
	unsigned char iv[AES_BLOCK_LENGTH];
	generateRandom(iv, sizeof(iv));

	// Encrypt
	size_t ciphertextLen = 0;
	unsigned char *ciphertext = new unsigned char[payload.size() + 64];
	const mbedtls_cipher_info_t *cipherInfo = mbedtls_cipher_info_from_type(MBEDTLS_CIPHER_AES_128_CBC);
	mbedtls_cipher_context_t cipher;

	mbedtls_cipher_init(&cipher);
	mbedtls_cipher_setup(&cipher, cipherInfo);
	mbedtls_cipher_setkey(&cipher, key, sizeof(key)*8, MBEDTLS_ENCRYPT);
	mbedtls_cipher_set_padding_mode(&cipher, MBEDTLS_PADDING_PKCS7);
	mbedtls_cipher_crypt(&cipher, iv, sizeof(iv),
			(const unsigned char *)payload.c_str(), payload.size(),
			ciphertext, &ciphertextLen);
	mbedtls_cipher_free(&cipher);

	// Save Data
	char encIV[(AES_BLOCK_LENGTH * 2) + 1];
	char *encCipher = new char[(ciphertextLen * 2) + 1];
	toHex(encIV, iv, sizeof(iv));
	toHex(encCipher, ciphertext, ciphertextLen);
	std::string _encIV(encIV);
	std::string _encCipher(encCipher);

	delete[] encCipher;
	delete[] ciphertext;

	if (directory[directory.size() - 1] != '/') {
		directory += "/";
	}
	directory += KEYCHAIN_FILE;
	std::ofstream f(directory.c_str());
	f << _encIV << std::endl;
	f << _encCipher << std::endl;
	f.close();
}

std::string Cryptor::loadAndDecrypt(std::string directory) {
	// Load Data
	std::string encIV, encCipher;

	if (directory[directory.size() - 1] != '/') {
		directory += "/";
	}
	directory += KEYCHAIN_FILE;
	std::ifstream f(directory.c_str());
	if (!f.good()) {
		f.close();
		throw 1;
	}
	f >> encIV;
	f >> encCipher;
	f.close();

	// Decode data
	unsigned char *ciphertext = new unsigned char[encCipher.size() / 2];
	unsigned char iv[AES_BLOCK_LENGTH];
	fromHex(ciphertext, encCipher.c_str());
	fromHex(iv, encIV.c_str());

	// Key
	if (!haveKey) {
		assembleKey();
	}

	// Decrypt
	size_t plaintextLen = 0;
	char *plaintext = new char[(encCipher.size() / 2) + 64];
	const mbedtls_cipher_info_t *cipherInfo = mbedtls_cipher_info_from_type(MBEDTLS_CIPHER_AES_128_CBC);
	mbedtls_cipher_context_t cipher;

	mbedtls_cipher_init(&cipher);
	mbedtls_cipher_setup(&cipher, cipherInfo);
	mbedtls_cipher_setkey(&cipher, key, sizeof(key)*8, MBEDTLS_DECRYPT);
	mbedtls_cipher_set_padding_mode(&cipher, MBEDTLS_PADDING_PKCS7);
	mbedtls_cipher_crypt(&cipher, iv, sizeof(iv),
			ciphertext, (encCipher.size() / 2),
			(unsigned char *)plaintext, &plaintextLen);
	mbedtls_cipher_free(&cipher);

	std::string payload(plaintext, plaintextLen);

	delete[] plaintext;
	delete[] ciphertext;

	return payload;
}

std::string Cryptor::createRandomPassword(PasswordSpec spec) {
	std::string password;
	std::vector<char> validChars;

	// Always allow lower-case alphabetic characters
	for (int i = 0; i < 26; i++)
		validChars.push_back('a' + i);

	// Allow special chars?
	if (!spec.ns) {
		validChars.push_back('`');
		validChars.push_back('~');
		validChars.push_back('!');
		validChars.push_back('@');
		validChars.push_back('#');
		validChars.push_back('$');
		validChars.push_back('%');
		validChars.push_back('^');
		validChars.push_back('&');
		validChars.push_back('*');
		validChars.push_back('(');
		validChars.push_back(')');
		validChars.push_back('-');
		validChars.push_back('_');
		validChars.push_back('=');
		validChars.push_back('+');
		validChars.push_back('[');
		validChars.push_back('{');
		validChars.push_back(']');
		validChars.push_back('}');
		validChars.push_back('\\');
		validChars.push_back('|');
		validChars.push_back(';');
		validChars.push_back(':');
		validChars.push_back(',');
		validChars.push_back('<');
		validChars.push_back('.');
		validChars.push_back('>');
		validChars.push_back('/');
		validChars.push_back('?');
	}

	// Allow capital letters?
	if (!spec.nc) {
		for (int i = 0; i < 26; i++)
			validChars.push_back('A' + i);
	}

	// Allow numbers?
	if (!spec.nn) {
		for (int i = 0; i < 10; i++)
			validChars.push_back('0' + i);
	}

	// Build string
	for (int i = 0; i < spec.ml; i++) {
		unsigned char c;
		generateRandom(&c, sizeof(c));
		password += validChars[c % validChars.size()];
	}

	return password;
}

void Cryptor::rekey() {
	haveKey = false;
}

void Cryptor::sha256(std::string str) {
	unsigned char hashbuf[32];
	mbedtls_sha256((const unsigned char *)str.c_str(), str.size(), hashbuf, 0);
	memcpy(key, hashbuf, sizeof(key));
	haveKey = true;
}

void Cryptor::generateRandom(void *output, size_t size) {
	mbedtls_entropy_context entropy;
	mbedtls_ctr_drbg_context random;

	mbedtls_entropy_init(&entropy);
	mbedtls_ctr_drbg_init(&random);
	mbedtls_ctr_drbg_seed(&random, mbedtls_entropy_func, &entropy, NULL, 0);
	mbedtls_ctr_drbg_random(&random, (unsigned char *)output, size);
	mbedtls_ctr_drbg_free(&random);
	mbedtls_entropy_free(&entropy);
}

std::string Cryptor::readPassword(bool confirm) {
	std::string password;

	#ifdef WIN32

		HANDLE hstdin = GetStdHandle(STD_INPUT_HANDLE);
		DWORD mode = 0;
		GetConsoleMode(hstdin, &mode);
		SetConsoleMode(hstdin, mode & ~ENABLE_ECHO_INPUT);

		password = promptPassword(confirm);

		SetConsoleMode(hstdin, mode);

	#else

		termios oldt;
		tcgetattr(STDIN_FILENO, &oldt);
		termios newt = oldt;
		newt.c_lflag &= ~ECHO;
		tcsetattr(STDIN_FILENO, TCSANOW, &newt);

		password = promptPassword(confirm);

		tcsetattr(STDIN_FILENO, TCSANOW, &oldt);

	#endif // WIN32

	return password;
}

std::string Cryptor::readPassword() {
	return readPassword(false);
}

void Cryptor::assembleKey(bool confirm) {
	sha256(readPassword(confirm));
}

void Cryptor::assembleKey() {
	assembleKey(false);
}

std::string Cryptor::promptPassword(bool confirm) {
	std::string password;

	while (true) {
		std::cout << PASSWORD_PROMPT;
		std::cin >> password;
		std::cout << std::endl;

		if (confirm) {
			std::string conf;
			std::cout << PASSWORD_CONF;
			std::cin >> conf;
			std::cout << std::endl;

			if (password == conf)
				break;
			else
				std::cout << PASSWORD_ERROR << std::endl;
		}
		else
			break;
	}

	return password;
}