#include "PasswordManager.h"

#include <QSqlQuery>
#include <QSqlError>
#include <QDebug>
#include <QVariant>
#include <botan/hash.h>
#include <botan/auto_rng.h>
#include <botan/base64.h>
#include <botan/scrypt.h>
#include <botan/nist_keywrap.h>
#include <botan/base64.h>
#include <botan/mac.h>
#include <botan/block_cipher.h>
#include <boost/lexical_cast.hpp>

namespace {

	class SqlException : public std::runtime_error {
	public:
		QSqlError error;
		SqlException(const QSqlError &err)
			: std::runtime_error(err.text().toUtf8().data())
			, error(err)
		{}

	};

}

using namespace Botan;

class PasswordCryptoEngine {
public:
	PasswordCryptoEngine(const secure_vector<uint8_t>& master_key)
	{
		m_cipher = BlockCipher::create_or_throw("AES-256");
		m_hmac = MessageAuthenticationCode::create_or_throw("HMAC(SHA-256)");
		m_hmac->set_key(master_key);

		m_cipher->set_key(m_hmac->process("wrap"));
		m_hmac->set_key(m_hmac->process("hmac"));
	}

	std::string  set(const std::string& name, const uint8_t val[], size_t len) const
	{
		/*
		* Both as a basic precaution wrt key seperation, and specifically to prevent
		* cut-and-paste attacks against the database, each PSK is encrypted with a
		* distinct key which is derived by hashing the wrapped key name with HMAC.
		*/
		const std::vector<uint8_t> wrapped_name =
		   nist_key_wrap_padded(cast_char_ptr_to_uint8(name.data()),
								name.size(),
								*m_cipher);

		std::unique_ptr<BlockCipher> wrap_cipher(m_cipher->clone());
		wrap_cipher->set_key(m_hmac->process(wrapped_name));
		const std::vector<uint8_t> wrapped_key = nist_key_wrap_padded(val, len, *wrap_cipher);

		return base64_encode(wrapped_key);
	}

	secure_vector<uint8_t> get(const std::string& name, const std::string_view &wrapped_key) const
	{
	   const std::vector<uint8_t> wrapped_name =
		  nist_key_wrap_padded(cast_char_ptr_to_uint8(name.data()),
							   name.size(),
							   *m_cipher);

	   const secure_vector<uint8_t> val = base64_decode(wrapped_key.data(), wrapped_key.size());

	   std::unique_ptr<BlockCipher> wrap_cipher(m_cipher->clone());
	   wrap_cipher->set_key(m_hmac->process(wrapped_name));

	   return nist_key_unwrap_padded(val.data(), val.size(), *wrap_cipher);
	}
private:
	std::unique_ptr<BlockCipher> m_cipher;
	std::unique_ptr<MessageAuthenticationCode> m_hmac;
};

// -------------------------

PasswordManager::PasswordManager() = default;
PasswordManager::~PasswordManager() = default;


bool PasswordManager::initialized(QSqlDatabase& db)
{
	return isPskStoreInitialized(db);
}

std::tuple<Botan::secure_vector<uint8_t>, Botan::secure_vector<uint8_t>>
PasswordManager::deriveKey(KeyStrengthener &ks, QString passphrase)
{
	auto master_key = ks.derive(passphrase.toUtf8().data());

	std::unique_ptr<Botan::HashFunction> hash3(Botan::HashFunction::create("SHA-3"));
	hash3->update(master_key);
	auto mkh = hash3->final();
	return { master_key, mkh };
}

bool PasswordManager::createDatabase(QSqlDatabase &db, QString passphrase)
{
	m_cryptoEngine.reset();
	if (!isPskStoreInitialized(db)) {
		initializeNewPskStore(db);
		auto ks = createKeyStrengthener();
		ks.saveParams(db, m_secretAlgoTableName);

		auto [master_key, mkh_bin] = deriveKey(ks, passphrase);
		auto mkh = QString::fromUtf8(Botan::base64_encode(mkh_bin).c_str());

		QSqlQuery q_ins_hash(db);
		q_ins_hash.prepare("INSERT INTO " + m_secretHashTableName + "(id, hash) VALUES(:id, :hash)");
		q_ins_hash.bindValue(":id", 1);
		q_ins_hash.bindValue(":hash", mkh);
		if (!q_ins_hash.exec()) {
			auto err = q_ins_hash.lastError();
			qDebug() << err.text();
			throw SqlException(err);
		}

		m_cryptoEngine = std::make_unique<PasswordCryptoEngine>(master_key);
		return true;
	}
	return false;
}

bool PasswordManager::openDatabase(QSqlDatabase &db, QString passphrase)
{
	m_cryptoEngine.reset();
	if (isPskStoreInitialized(db)) {
		auto ks = getKeyStrengthener(db);
		auto [master_key, mkh_bin] = deriveKey(ks, passphrase);

		QSqlQuery q("SELECT hash FROM " + m_secretHashTableName + " WHERE id=1", db);
		if (q.next()) {
			auto hash_b64 = q.value(0).toString().toUtf8();
			auto hash_bin = Botan::base64_decode(hash_b64.data(), static_cast<size_t>(hash_b64.size()));
			if (hash_bin == mkh_bin) {
				m_cryptoEngine = std::make_unique<PasswordCryptoEngine>(master_key);
				return true;
			}
		}
	}
	return false;
}

void PasswordManager::closeDatabase()
{
	m_cryptoEngine.reset();
}

bool PasswordManager::locked() const
{
    return m_cryptoEngine == nullptr;
}

void PasswordManager::resetMasterPassword(QSqlDatabase &db)
{
    if (!isPskStoreInitialized(db))
        return;

    closeDatabase();
    QSqlQuery del_algo("DELETE FROM " + m_secretAlgoTableName + " WHERE id=1", db);
    del_algo.exec();
    QSqlQuery del_hash("DELETE FROM " + m_secretHashTableName + " WHERE id=1", db);
    del_hash.exec();
}

std::string PasswordManager::encrypt(const std::string &name, const std::string &passwd)
{
	if (m_cryptoEngine) {
		return m_cryptoEngine->set(name, reinterpret_cast<const uint8_t*>(passwd.data()), passwd.length());
	}
	else {
		throw PasswordManagerLockedException();
	}
}

std::string PasswordManager::decrypt(const std::string &id, const std::string_view &encpwd)
{
    if (m_cryptoEngine)
    {
        try
        {
            secure_vector<uint8_t> decoded = m_cryptoEngine->get(id, encpwd);
            return std::string(reinterpret_cast<const char*>(decoded.data()), decoded.size());
        }
        catch (const Botan::Exception &ex)
        {
            throw PasswordManagerException(ex.what());
        }
	}
	else {
		throw PasswordManagerLockedException();
	}
}

void PasswordManager::initializeNewPskStore(QSqlDatabase &db)
{
//	// Create tables
//	// - psk_masterkey_algo
//	// - psk_passwd
	{
		QSqlQuery create_tbl(db);
		create_tbl.prepare(
			"CREATE TABLE IF NOT EXISTS " + m_secretAlgoTableName + "( \n"
			" id INTEGER PRIMARY KEY, \n"
			" algo TEXT, \n"
			" i1 INTEGER, \n"
			" i2 INTEGER, \n"
			" i3 INTEGER, \n"
			" ks INTEGER, \n"
			" salt TEXT \n"
			");");
		if (!create_tbl.exec()) {
			auto err = create_tbl.lastError();
			throw SqlException(err);
		}
	}

	QSqlQuery create_tbl(db);
	create_tbl.prepare(
		"CREATE TABLE IF NOT EXISTS " + m_secretHashTableName + "( \n"
		" id INTEGER PRIMARY KEY, \n"
		" hash TEXT \n"
		");");
	if (!create_tbl.exec()) {
		auto err = create_tbl.lastError();
		throw SqlException(err);
	}
}

bool PasswordManager::isPskStoreInitialized(QSqlDatabase& db)
{
	// Is the table with the secret data present and filled?
	QSqlQuery query(db);
	query.prepare("SELECT name FROM sqlite_master WHERE type='table' AND name=:name");
	query.bindValue(":name", m_secretAlgoTableName);
	if (!query.exec()) {
		auto err = query.lastError();
		throw SqlException(err);
	}
	if (!query.next()) {
		return false;
	}

	query.bindValue(":name", m_secretHashTableName);
	if (!query.exec()) {
		auto err = query.lastError();
		throw SqlException(err);
	}
	if (!query.next()) {
		return false;
	}

    QSqlQuery sel_algo("SELECT algo FROM " + m_secretAlgoTableName + " WHERE id=1", db);
	if (!sel_algo.next()) {
		return false;
	}

	QSqlQuery sel_hash("SELECT hash FROM " + m_secretHashTableName + " WHERE id=1", db);
	if (!sel_hash.next()) {
		return false;
	}

	return true;
}

KeyStrengthener PasswordManager::getKeyStrengthener(QSqlDatabase &db)
{
	QSqlQuery query("SELECT algo, i1, i2, i3, ks, salt FROM " + m_secretAlgoTableName + " WHERE id=1", db);
	if (query.next()) {
		std::string algo = query.value(0).toString().toUtf8().data();
		size_t i1 = query.value(1).toUInt();
		size_t i2 = query.value(2).toUInt();
		size_t i3 = query.value(3).toUInt();
		size_t ks = query.value(4).toUInt();
		auto salt = query.value(5).toString().toUtf8();

		auto pwh_fam = Botan::PasswordHashFamily::create(algo);
		return KeyStrengthener(
					pwh_fam->from_params(i1, i2, i3),
					Botan::base64_decode(salt.data(), static_cast<size_t>(salt.size())),
					ks
					);
	}
	else {
		throw std::runtime_error("fail");
	}
}

KeyStrengthener PasswordManager::createKeyStrengthener()
{
	size_t key_size = 64;
	Botan::secure_vector<uint8_t> salt(key_size);
	Botan::AutoSeeded_RNG rng;
	rng.randomize(salt.data(), salt.size());

	const std::string algo = "Scrypt";
	auto pwh_fam = Botan::PasswordHashFamily::create(algo);
	return KeyStrengthener(
				pwh_fam->tune(key_size, std::chrono::seconds(2), 130),
				salt,
				key_size
				);
}