#include "Cryptor.h"

bool Cryptor::haveKey = false;
unsigned char Cryptor::key[CryptoPP::AES::DEFAULT_KEYLENGTH];

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

	// IV
	unsigned char iv[CryptoPP::AES::BLOCKSIZE];
	CryptoPP::AutoSeededRandomPool randl;
	randl.GenerateBlock(iv, sizeof(iv));

	// Encrypt
	std::string cipher;
	CryptoPP::CBC_Mode<CryptoPP::AES>::Encryption enc(key, sizeof(key), iv);
	CryptoPP::StringSource(payload, true, new CryptoPP::StreamTransformationFilter(enc, new CryptoPP::StringSink(cipher)));

	// Save Data
	std::string encIV, encCipher;
	CryptoPP::StringSource(iv, sizeof(iv), true, new CryptoPP::HexEncoder(new CryptoPP::StringSink(encIV)));
	CryptoPP::StringSource(cipher, true, new CryptoPP::HexEncoder(new CryptoPP::StringSink(encCipher)));
	if (remoteHost == "") {
		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();
	}
	else {
		Socket s;
		std::string err = "";
		s.conn(remoteHost, port);
		s.sendline("store");
		s.sendline(directory);
		s.sendline(encIV);
		s.sendline(encCipher);
		err = s.readline();
		s.clo();
		if (err != "OK") throw 1;
	}
}

std::string Cryptor::loadAndDecrypt(std::string remoteHost, std::string port, std::string directory) {
	// Load Data
	std::string encIV, encCipher, ivstr, cipher;
	if (remoteHost == "") {
		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();
	}
	else {
		Socket s;
		std::string err = "";
		s.conn(remoteHost, port);
		s.sendline("fetch");
		s.sendline(directory);
		err = s.readline();
		if (err != "OK") {
			s.clo();
			throw 1;
		}
		encIV = s.readline();
		encCipher = s.readline();
		s.clo();
	}
	CryptoPP::StringSource(encIV, true, new CryptoPP::HexDecoder(new CryptoPP::StringSink(ivstr)));
	CryptoPP::StringSource(encCipher, true, new CryptoPP::HexDecoder(new CryptoPP::StringSink(cipher)));

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

	// IV
	unsigned char iv[CryptoPP::AES::BLOCKSIZE];
	memcpy(iv, ivstr.c_str(), ivstr.size());

	// Decrypt
	std::string payload;
	CryptoPP::CBC_Mode<CryptoPP::AES>::Decryption dec(key, sizeof(key), iv);
	CryptoPP::StringSource(cipher, true, new CryptoPP::StreamTransformationFilter(dec, new CryptoPP::StringSink(payload)));

	return payload;
}

std::string Cryptor::createRandomPassword(PasswordSpec spec) {
	CryptoPP::AutoSeededRandomPool randl;
	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 r[1];
		char c;
		randl.GenerateBlock(r, sizeof(r));
		c = validChars[r[0] % validChars.size()];
		password += c;
	}

	return password;
}

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

void Cryptor::sha256(std::string str) {
	CryptoPP::SHA256 hash;
	std::string rtr;
	CryptoPP::StringSource(str, true, new CryptoPP::HashFilter(hash, new CryptoPP::StringSink(rtr), false, CryptoPP::AES::DEFAULT_KEYLENGTH));
	memcpy(key, rtr.c_str(), rtr.size());
	haveKey = true;
}

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;
}