#include "PasswordManager.h"

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

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(SQLiteConnection& 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(SQLiteConnection &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());

        auto  q_ins_hash = db.Prepare(
            "INSERT INTO " + m_secretHashTableName + "(id, hash) VALUES(?1, ?2)");
        q_ins_hash.Bind(1, 1);
        q_ins_hash.Bind(2, mkh);
        q_ins_hash.Step();

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

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

        auto q = db.Prepare("SELECT hash FROM " + m_secretHashTableName + " WHERE id=1");
        if (q.Step()) {
            QByteArray hash_b64 = q.ColumnCharPtr(0);
			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(SQLiteConnection &db)
{
    if (!isPskStoreInitialized(db))
        return;

    closeDatabase();
    auto del_algo = db.Prepare("DELETE FROM " + m_secretAlgoTableName + " WHERE id=1");
    del_algo.Step();
    auto del_hash = db.Prepare("DELETE FROM " + m_secretHashTableName + " WHERE id=1");
    del_hash.Step();
}

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(SQLiteConnection &db)
{
//	// Create tables
//	// - psk_masterkey_algo
//	// - psk_passwd
	{
        auto create_tbl = db.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"
			");");
        create_tbl.Step();
	}

    auto create_tbl = db.Prepare(
		"CREATE TABLE IF NOT EXISTS " + m_secretHashTableName + "( \n"
		" id INTEGER PRIMARY KEY, \n"
		" hash TEXT \n"
		");");
    create_tbl.Step();
}

bool PasswordManager::isPskStoreInitialized(SQLiteConnection& db)
{
	// Is the table with the secret data present and filled?
    auto query = db.Prepare("SELECT name FROM sqlite_master WHERE type='table' AND name=?1");
    query.Bind(1, m_secretAlgoTableName);
    if (!query.Step()) {
		return false;
	}

    query.Reset();
    query.Bind(1, m_secretHashTableName);
    if (!query.Step()) {
		return false;
	}

    auto sel_algo = db.Prepare("SELECT algo FROM " + m_secretAlgoTableName + " WHERE id=1");
    if (!sel_algo.Step()) {
		return false;
	}

    auto sel_hash = db.Prepare("SELECT hash FROM " + m_secretHashTableName + " WHERE id=1");
    if (!sel_hash.Step()) {
		return false;
	}

	return true;
}

KeyStrengthener PasswordManager::getKeyStrengthener(SQLiteConnection &db)
{
    auto query = db.Prepare("SELECT algo, i1, i2, i3, ks, salt FROM " + m_secretAlgoTableName + " WHERE id=1");
    if (query.Step()) {
        std::string algo = query.ColumnCharPtr(0);
        size_t i1 = query.ColumnInteger(1);
        size_t i2 = query.ColumnInteger(2);
        size_t i3 = query.ColumnInteger(3);
        size_t ks = query.ColumnInteger(4);
        QByteArray salt = query.ColumnCharPtr(5);

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