#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/psk_db_sql.h>
#include <botan/psk_db.h>

//#include <botan/sqlite3.h>
#include <botan/scrypt.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)
		{}

	};

	class QPSK_Database : public Botan::Encrypted_PSK_Database
	{
	public:
		/**
		* @param master_key specifies the master key used to encrypt all
		* keys and value. It can be of any length, but should be at least 256 bits.
		*
		* Subkeys for the cryptographic algorithms used are derived from this
		* master key. No key stretching is performed; if encrypting a PSK database
		* using a password, it is recommended to use PBKDF2 to derive the database
		* master key.
		*/
		QPSK_Database(const Botan::secure_vector<uint8_t>& master_key, QSqlDatabase &db, const QString &table_name)
			: Encrypted_PSK_Database(master_key)
			, m_db(db)
			, m_tableName(table_name)
		{
			QSqlQuery q_create_table(m_db);
			q_create_table.prepare("CREATE TABLE IF NOT EXISTS " + table_name +
									"(psk_name TEXT PRIMARY KEY, psk_value TEXT)");
			if (!q_create_table.exec()) {
				auto err = q_create_table.lastError();
				throw SqlException(err);
			}
		}

	protected:
		/// Save a encrypted (name.value) pair to the database. Both will be base64 encoded strings.
		virtual void kv_set(const std::string& index, const std::string& value) override
		{
			QSqlQuery q(m_db);
			q.prepare("insert or replace into " + m_tableName + " values(:name, :value)");
			q.bindValue(":name", QString::fromUtf8(index.c_str()));
			q.bindValue(":value", QString::fromUtf8(value.c_str()));
			if (!q.exec()) {
				auto err = q.lastError();
				throw SqlException(err);
			}
		}

		/// Get a value previously saved with set_raw_value. Should return an empty
		/// string if index is not found.
		virtual std::string kv_get(const std::string& index) const override
		{
			QSqlQuery q(m_db);
			q.prepare("SELECT psk_value FROM " + m_tableName +
					  "  WHERE psk_name = :name");
			q.bindValue(":name", QString::fromUtf8(index.c_str()));
			if (q.exec()) {
				if (q.next()) {
					return q.value(0).toString().toUtf8().data();
				}
			}
			else {
				auto err = q.lastError();
				throw SqlException(err);
			}
			return std::string();
		}

		/// Remove an index
		virtual void kv_del(const std::string& index) override
		{
			QSqlQuery q(m_db);
			q.prepare("DELETE FROM " + m_tableName + " WHERE psk_name=:name");
			q.bindValue(":name", QString::fromUtf8(index.c_str()));
			if (!q.exec()) {
				auto err = q.lastError();
				throw SqlException(err);
			}
		}

		/// Return all indexes in the table.
		virtual std::set<std::string> kv_get_all() const override
		{
			QSqlQuery q(m_db);
			q.prepare("SELECT psk_name FROM " + m_tableName);

			std::set<std::string> result;
			if (q.exec()) {
				while (q.next()) {
					result.insert(q.value(0).toString().toUtf8().data());
				}
			}
			else {
				auto err = q.lastError();
				throw SqlException(err);
			}
			return result;
		}
	private:
		QSqlDatabase &m_db;
		QString m_tableName;
	};


}


Botan::secure_vector<uint8_t> PasswordManager::KeyStrengthener::derive(const std::string &passphrase)
{
	Botan::secure_vector<uint8_t> master_key(m_keySize);
	m_hasher->derive_key(master_key.data(), master_key.size(), passphrase.c_str(), passphrase.length(), m_salt.data(), m_salt.size());

	return master_key;
}

void PasswordManager::KeyStrengthener::saveParams(QSqlDatabase &db, const QString &table_name)
{
	auto sc = dynamic_cast<Botan::Scrypt*>(m_hasher.get());
	size_t i1 = sc->N();
	size_t i2 = sc->r();
	size_t i3 = sc->p();


	auto salt_str = QString::fromUtf8(Botan::base64_encode(m_salt).c_str());
	// SAVE parameters in database
	QSqlQuery insert_statement(db);
	insert_statement.prepare("INSERT OR REPLACE INTO " + table_name + "(id, algo, i1, i2, i3, ks, salt) "
								+ "VALUES(:id, :algo, :i1, :i2, :i3, :ks, :salt)");
	insert_statement.bindValue(":id", 1);
	insert_statement.bindValue(":algo", "Scrypt");
	insert_statement.bindValue(":i1", i1);
	insert_statement.bindValue(":i2", i2);
	insert_statement.bindValue(":i3", i3);
	insert_statement.bindValue(":ks", m_keySize);
	insert_statement.bindValue(":salt", salt_str);
	if (!insert_statement.exec()) {
		//throw std::runtime_error("PasswordManager::KeyStrengthener::saveParams failed");
		auto err = insert_statement.lastError();
		throw SqlException(err);

	}

}

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

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


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

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

		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 = QString::fromUtf8(Botan::base64_encode(hash3->final()).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_pskDatabase = std::make_unique<QPSK_Database>(master_key, db, m_passwordTableName);
		return true;
	}
	return false;
}

bool PasswordManager::openDatabase(QSqlDatabase &db, QString passphrase)
{
	if (isPskStoreInitialized(db)) {
		auto ks = getKeyStrengthener(db);

		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_bin = hash3->final();

		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_pskDatabase = std::make_unique<QPSK_Database>(master_key, db, m_passwordTableName);
				return true;
			}
		}
	}
	return false;
}

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

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

void PasswordManager::set(const std::string &id, const std::string &passwd)
{
	if (m_pskDatabase) {
		m_pskDatabase->set_str(id, passwd);
	}
	else {
		throw PasswordManagerLockedException();
	}
}

bool PasswordManager::get(const std::string &id, std::string &password)
{
	if (m_pskDatabase) {
		try {
			password =  m_pskDatabase->get_str(id);
			return true;
		}
		catch (const Botan::Invalid_Argument &) {
			// not present
			return false;
		}
	}
	else {
		throw PasswordManagerLockedException();
	}
}

void PasswordManager::remove(const std::string &id)
{
	if (m_pskDatabase) {
		m_pskDatabase->remove(id);
	}
	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 sql_error = create_tbl.lastError();
//			throw std::runtime_error("create table failed");
			auto err = create_tbl.lastError();
			throw SqlException(err);
		}
	}

//	db->create_table(
	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;
}

PasswordManager::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");
	}
}

PasswordManager::KeyStrengthener PasswordManager::createKeyStrengthener()
{
//	std::unique_ptr<Botan::PasswordHash> pwh;

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