#include "PasswordManager.h"

#include <botan/filters.h>
#include <botan/pipe.h>
#include <botan/sha2_64.h>
#include <botan/hash.h>
#include <botan/hmac.h>
#include <botan/pbkdf2.h>
#include <botan/rng.h>
#include <botan/base64.h>
#include <botan/loadstor.h>
#include <botan/mem_ops.h>

#include <boost/assert.hpp>


using namespace Botan;

namespace {

	/*
	First 24 bits of SHA-256("Botan Cryptobox"), followed by 8 0 bits
	for later use as flags, etc if needed
	*/
	const uint8_t c_PasswordVersionCode = 0x10;

	const size_t c_VersionCodeLen = 1;

	const size_t CIPHER_KEY_LEN = 32;
	const size_t CIPHER_IV_LEN = 16;
	const size_t MAC_KEY_LEN = 32;
	const size_t MAC_OUTPUT_LEN = 20;
	const size_t PBKDF_SALT_LEN = 10;
	const size_t PBKDF_ITERATIONS = 8 * 1024;

	const size_t PBKDF_OUTPUT_LEN = CIPHER_KEY_LEN + CIPHER_IV_LEN + MAC_KEY_LEN;

	const char * const c_Cipher = "Serpent/CTR-BE";

	const char * const c_IniGroupSecurity = "Security";



	StrengthenedKey generateKey(const std::string &passphrase, const uint8_t *salt, int saltlength)
	{
		PKCS5_PBKDF2 pbkdf(new HMAC(new SHA_512));
		OctetString master_key = pbkdf.derive_key(
		   PBKDF_OUTPUT_LEN,
		   passphrase,
		   salt, saltlength,
		   PBKDF_ITERATIONS);

		const uint8_t* mk = master_key.begin();

		return StrengthenedKey(
			SymmetricKey(mk, CIPHER_KEY_LEN),
			SymmetricKey(mk + CIPHER_KEY_LEN, MAC_KEY_LEN),
			InitializationVector(mk + CIPHER_KEY_LEN + MAC_KEY_LEN, CIPHER_IV_LEN));
	}

//	secure_vector<uint8_t> pbkdf_salt(PBKDF_SALT_LEN);
//	rng.randomize( pbkdf_salt.data(), pbkdf_salt.size());
//	StrengthenedKey strengthened_key = generateKey(passphrase, pbkdf_salt.data(), pbkdf_salt.size());


	std::string encrypt(const std::string &input,
						const StrengthenedKey &strengthened_key)
	{

		Pipe pipe(get_cipher(c_Cipher, strengthened_key.cipher_key,
			strengthened_key.iv, ENCRYPTION),
			new Fork(
				nullptr,
				new MAC_Filter(new HMAC(new SHA_512),
						strengthened_key.mac_key, MAC_OUTPUT_LEN)));

		pipe.process_msg((const uint8_t*)input.data(), input.length());

	   /*
	   Output format is:
		  mac (20 bytes)
		  ciphertext
	   */
		const size_t ciphertext_len = pipe.remaining(0);
		std::vector<uint8_t> out_buf(MAC_OUTPUT_LEN + ciphertext_len);

		BOTAN_ASSERT_EQUAL(
			pipe.read(&out_buf[0], MAC_OUTPUT_LEN, 1),
				MAC_OUTPUT_LEN, "MAC output");
		BOTAN_ASSERT_EQUAL(
			pipe.read(&out_buf[MAC_OUTPUT_LEN],	ciphertext_len, 0),
				ciphertext_len, "Ciphertext size");

		return base64_encode(out_buf.data(), out_buf.size());
	}

	std::string decrypt(const std::string &input, const StrengthenedKey &strengthened_key)
	{
		secure_vector<uint8_t> ciphertext = base64_decode(input);

		if(ciphertext.size() < (MAC_OUTPUT_LEN)) {
			throw Decoding_Error("Invalid encrypted password input");
		}

		Pipe pipe(new Fork(
					get_cipher(c_Cipher, strengthened_key.cipher_key, strengthened_key.iv, DECRYPTION),
					new MAC_Filter(new HMAC(new SHA_512), strengthened_key.mac_key, MAC_OUTPUT_LEN)
					  ));

		const size_t ciphertext_offset = MAC_OUTPUT_LEN;
		pipe.process_msg(&ciphertext[ciphertext_offset], ciphertext.size() - ciphertext_offset);

		uint8_t computed_mac[MAC_OUTPUT_LEN];
		BOTAN_ASSERT_EQUAL(MAC_OUTPUT_LEN, pipe.read(computed_mac, MAC_OUTPUT_LEN, 1), "MAC size");

		if(!same_mem(computed_mac, &ciphertext[0], MAC_OUTPUT_LEN)) {
			throw Decoding_Error("Encrypted password integrity failure");
		}

		return pipe.read_all_as_string(0);
	}

	struct constants {
		const int pbkdf_salt_len;
	};

	constants v1_consts = {
		10
	};

} // end of unnamed namespace

/*
 * File layout:
 *
 * Header
 *   version
 *   key_salt
 *   hash_salt
 *   master_hash
 *
 *
 * Passwords
 *   key = pw
 */



PasswordManager::PasswordManager()
{

}

Expected<bool> PasswordManager::unlock(const std::string &master_password)
{
	try {
		bool result = false;
		if (m_masterHash.length() == 0 && master_password.empty()) {
			result = true;
		} else {
			StrengthenedKey key = generateKey(master_password, m_keySalt.begin(), m_keySalt.length());
			OctetString hash = hashStrengthenedKey(key, m_hashSalt);

			BOOST_ASSERT_MSG(hash.length() == m_masterHash.length(), "Both hashes should have the same length! Versioning error?");

			if (same_mem(m_masterHash.begin(), hash.begin(), hash.length())) {
				result = true;
				m_masterKey = key;
			}
		}
		return result;
	} catch (...) {
		return Expected<bool>::fromException();
	}
}

Expected<bool> PasswordManager::changeMasterPassword(const std::string &old_master_password,
									const std::string &new_master_password)
{
	try {
		bool result = false;
		if (m_masterHash.length() == 0 && old_master_password.empty()) {
			// Nothing set yet so we initialize for first use
			m_keySalt  = OctetString(m_rng, v1_consts.pbkdf_salt_len);
			m_masterKey = generateKey(new_master_password, m_keySalt.begin(), m_keySalt.length());

			m_hashSalt = OctetString(m_rng, v1_consts.pbkdf_salt_len);
			m_masterHash = hashStrengthenedKey(m_masterKey, m_hashSalt);
			result = true;
		}
		return result;
	} catch (...) {
		return Expected<bool>::fromException();
	}
}

void PasswordManager::lock()
{
	m_masterKey = StrengthenedKey();
}

bool PasswordManager::locked() const
{
	return m_masterKey.cipher_key.size() == 0;
}

Expected<void> PasswordManager::savePassword(const std::string &key, const std::string &password)
{
	if (locked()) {
		return Expected<void>::fromException(std::logic_error("Need to unlock the password manager first"));
	}
	std::string epw = encrypt(password, m_masterKey);
	m_store.emplace(key, epw);

	return Expected<void>();
}

Expected<bool> PasswordManager::getPassword(const std::string &key, std::string &out)
{
	if (locked()) {
		return Expected<bool>::fromException(std::logic_error("Need to unlock the password manager first"));
	}
	auto fi = m_store.find(key);

	bool result = false;
	if (fi != m_store.end()) {
		out = decrypt(fi->second, m_masterKey);
		result = true;
	}

	return result;
}

Botan::OctetString PasswordManager::hashStrengthenedKey(const StrengthenedKey &key, const OctetString &salt)
{
	std::unique_ptr<Botan::HashFunction> hash3(Botan::HashFunction::create("SHA-3"));
	BOOST_ASSERT_MSG(hash3 != nullptr, "SHA-3 algorithm not available");
	hash3->update(salt.begin(), salt.length());
	hash3->update(key.cipher_key.begin(), key.cipher_key.length());
	hash3->update(key.mac_key.begin(),    key.mac_key.length());
	hash3->update(key.iv.begin(),         key.iv.length());
	return hash3->final();
}