pgLab/pglab/ConnectionListModel.cpp

#include "ConnectionListModel.h"
#include "ScopeGuard.h"
#include "util.h"

#include <botan/cryptobox.h>
#include <QDir>
#include <QMimeData>
#include <QSettings>
#include <QSqlDatabase>
#include <QSqlError>
#include <QSqlQuery>
#include <QStandardPaths>
#include <QStringBuilder>


namespace {

	const char * const q_create_table_conngroup =
R"__(
CREATE TABLE IF NOT EXISTS conngroup (
	conngroup_id  INTEGER PRIMARY KEY,
	gname         TEXT    NOT NULL UNIQUE
);)__";

	const char * const q_create_table_connection =
R"__(
CREATE TABLE IF NOT EXISTS connection (
	uuid          TEXT    PRIMARY KEY,
	cname         TEXT,
	conngroup_id  INTEGER NOT NULL,
	host          TEXT    ,
	hostaddr      TEXT    ,
	port          INTEGER NOT NULL,
	user          TEXT    ,
	dbname        TEXT    ,
	sslmode       INTEGER NOT NULL,
	sslcert       TEXT    ,
	sslkey        TEXT    ,
	sslrootcert   TEXT    ,
	sslcrl        TEXT    ,
	password      TEXT
);)__";



	const char * const q_insert_or_replace_into_connection =
R"__(INSERT OR REPLACE INTO connection
	VALUES (:uuid, :name, :conngroup_id, :host, :hostaddr, :port, :user, :dbname,
		:sslmode, :sslcert, :sslkey, :sslrootcert, :sslcrl, :password);
)__"	;

	std::tuple<bool, QSqlError> InitConnectionTables(QSqlDatabase &db)
	{
		QSqlQuery q_create_table(db);
		q_create_table.prepare(q_create_table_conngroup);
		if (!q_create_table.exec()) {
			auto err = q_create_table.lastError();
			return { false, err };
		}
		q_create_table.prepare(q_create_table_connection);
		if (!q_create_table.exec()) {
			auto err = q_create_table.lastError();
			return { false, err };
		}
		return {true, {}};
	}

	std::optional<QSqlError> SaveConnectionConfig(QSqlDatabase &db, const ConnectionConfig &cc, int conngroup_id)
	{
		QSqlQuery q(db);
		q.prepare(q_insert_or_replace_into_connection);
		q.bindValue(":uuid", cc.uuid().toString());
		q.bindValue(":name", cc.name());
		q.bindValue(":conngroup_id", conngroup_id);
		q.bindValue(":host", cc.host());
		q.bindValue(":hostaddr", cc.hostAddr());
		q.bindValue(":port", (int)cc.port());
		q.bindValue(":user", cc.user());
		q.bindValue(":dbname", cc.dbname());
		q.bindValue(":sslmode", static_cast<int>(cc.sslMode()));
		q.bindValue(":sslcert", cc.sslCert());
		q.bindValue(":sslkey", cc.sslKey());
		q.bindValue(":sslrootcert", cc.sslRootCert());
		q.bindValue(":sslcrl", cc.sslCrl());
        auto& encodedPassword = cc.encodedPassword();
        if (encodedPassword.isEmpty())
            q.bindValue(":password", QVariant());
        else
            q.bindValue(":password", encodedPassword);

		if (!q.exec()) {
			auto sql_error = q.lastError();
			return { sql_error };
		}
		return {};
	}

} // end of unnamed namespace

ConnectionTreeModel::ConnectionTreeModel(QObject *parent, QSqlDatabase &db)
	: QAbstractItemModel(parent)
	, m_db(db)
{
}

void ConnectionTreeModel::load()
{
	InitConnectionTables(m_db);

	QSqlQuery q(m_db);
	q.prepare("SELECT conngroup_id, gname FROM conngroup;");
	if (!q.exec()) {
//		auto err = q_create_table.lastError();
//		return { false, err };
		throw std::runtime_error("Loading groups failed");
	}
	while (q.next()) {
		int id = q.value(0).toInt();
		QString name = q.value(1).toString();

		auto g = std::make_shared<ConnectionGroup>();
		g->conngroup_id = id;
		g->name = name;
		m_groups.push_back(g);
	}

	q.prepare("SELECT uuid, cname, conngroup_id, host, hostaddr, port, "
			  "    user, dbname, sslmode, sslcert, sslkey, sslrootcert, sslcrl, "
			  "    password "
			  "FROM connection ORDER BY conngroup_id, cname;");
	if (!q.exec()) {
//		auto err = q_create_table.lastError();
//		return { false, err };
		throw std::runtime_error("Loading groups failed");
	}
	while (q.next()) {
		auto cc = std::make_shared<ConnectionConfig>();
		cc->setUuid(q.value(0).toUuid());
		cc->setName(q.value(1).toString());
		cc->setHost(q.value(3).toString());
		cc->setHostAddr(q.value(4).toString());
		cc->setPort(static_cast<uint16_t>(q.value(5).toInt()));
		cc->setUser(q.value(6).toString());
		cc->setDbname(q.value(7).toString());
		cc->setSslMode(static_cast<SslMode>(q.value(8).toInt()));
		cc->setSslCert(q.value(9).toString());
		cc->setSslKey(q.value(10).toString());
		cc->setSslRootCert(q.value(11).toString());
		cc->setSslCrl(q.value(12).toString());
		cc->setEncodedPassword(q.value(13).toByteArray());

		int group_id = q.value(2).toInt();
		auto find_res = std::find_if(m_groups.begin(), m_groups.end(),
							[group_id] (auto item) { return item->conngroup_id == group_id; });
		if (find_res != m_groups.end()) {
			(*find_res)->add(cc);
		}
		else {
			throw std::runtime_error("conngroup missing");
		}
	}
}

QVariant ConnectionTreeModel::data(const QModelIndex &index, int role) const
{
	// Code below assumes two level tree groups/connections
	// it will fail for nested groups
	QVariant v;

	auto privdata =  static_cast<ConnectionNode*>(index.internalPointer());
	if (auto group = dynamic_cast<ConnectionGroup*>(privdata); group != nullptr) {
		// This is a group
		if (role == Qt::DisplayRole) {
			if (index.column() == Name) {
				v = group->name;
			}
		}
	}
	else if (auto conn = dynamic_cast<ConnectionConfig*>(privdata); conn != nullptr) {
		// This is a connection
		if (role == Qt::DisplayRole) {
			switch (index.column()) {
			case Name: v = conn->name();	break;
			case Host: v = conn->host(); break;
			case Port: v = conn->port(); break;
			case User: v = conn->user(); break;
			case DbName: v= conn->dbname(); break;
			}
		}
	}

	return v;
}

QVariant ConnectionTreeModel::headerData(int section, Qt::Orientation orientation, int role) const
{
	QVariant v;
	if (orientation == Qt::Horizontal) {
		if (role == Qt::DisplayRole) {
			switch (section) {
			case Name: v = tr("Name");	break;
			case Host: v = tr("Host"); break;
			case Port: v = tr("Port"); break;
			case User: v = tr("User"); break;
			case DbName: v= tr("Database"); break;
			}
		}
	}
	return v;
}

QModelIndex ConnectionTreeModel::index(int row, int column, const QModelIndex &parent) const
{
	if (!hasIndex(row, column, parent))
		return {};

	const ConnectionNode *node = nullptr;
	if (parent.isValid()) {
		auto privdata = static_cast<ConnectionNode*>(parent.internalPointer());
		if (auto group = dynamic_cast<ConnectionGroup*>(privdata); group != nullptr) {
			node = group->connections().at(row).get();
		}
		else {
			throw std::logic_error("Should never ask for a child index of a connectionconfig");
		}
	}
	else {
		node = m_groups[row].get();
	}
	return createIndex(row, column, const_cast<ConnectionNode*>(node));
}

QModelIndex ConnectionTreeModel::parent(const QModelIndex &index) const
{
	if (!index.isValid())
		return {};

	auto privdata =  static_cast<ConnectionNode*>(index.internalPointer());
	if (auto group = dynamic_cast<ConnectionGroup*>(privdata); group != nullptr) {
		return {};
	}
	else if (auto config = dynamic_cast<ConnectionConfig*>(privdata); config != nullptr) {
		auto p = config->parent();
		auto find_res = std::find_if(m_groups.begin(), m_groups.end(), [p] (auto item) -> bool { return *p == *item; });
		if (find_res != m_groups.end()) {
			return createIndex(find_res - m_groups.begin(), 0,
							   const_cast<ConnectionGroup*>(config->parent()));
		}
	}
	throw std::logic_error("Should never get here");
}

int ConnectionTreeModel::rowCount(const QModelIndex &parent) const
{
	int result = 0;
	if (parent.isValid()) {
		auto privdata =  static_cast<ConnectionNode*>(parent.internalPointer());
		if (auto group = dynamic_cast<ConnectionGroup*>(privdata); group != nullptr) {
			result = group->connections().size();
		}
		else if (auto config = dynamic_cast<ConnectionConfig*>(privdata); config != nullptr) {
			result = 0;
		}
	}
	else {
		result = m_groups.size();
	}
	return result;
}

int ConnectionTreeModel::columnCount(const QModelIndex &) const
{
	return ColCount;
}

bool ConnectionTreeModel::removeRows(int row, int count, const QModelIndex &parent)
{
	if (parent.isValid() && count == 1) {
		// should be a group
		auto grp = m_groups[parent.row()];
		for (int i = 0; i < count; ++i) {
			QUuid uuid = grp->connections().at(row + i)->uuid();
			QSqlQuery q(m_db);
			q.prepare(
				"DELETE FROM connection "
				"  WHERE uuid=:uuid");
			q.bindValue(":uuid", uuid);
			if (!q.exec()) {
				auto err = q.lastError();
				throw std::runtime_error("QqlError");
			}
		}
		beginRemoveRows(parent, row, row + count - 1);
		SCOPE_EXIT { endRemoveRows(); };
		grp->erase(row, count);
		return true;
	}
	return false;
}

void ConnectionTreeModel::save(const QString &group_name, const ConnectionConfig &cc)
{
	auto [grp_idx, conn_idx] = findConfig(cc.uuid());
	if (grp_idx >= 0) {
		auto grp = m_groups[grp_idx];
		if (grp->name == group_name) {
			// update config
			grp->update(conn_idx, cc);
			// send change event
			auto node = grp->connections().at(conn_idx);
			dataChanged(
				createIndex(conn_idx, 0, node.get()),
				createIndex(conn_idx, ColCount-1, node.get()));
			saveToDb(*node);
			return;
		}
		else {
			auto parent = createIndex(grp_idx, 0, grp.get());
			beginRemoveRows(parent, conn_idx, conn_idx);
			SCOPE_EXIT { endRemoveRows(); };
			grp->erase(conn_idx);
		}
	}
	// Here we can assume we have to find the new group or create a new group
	// because if the connection was in the right group the function has already returned.
	// We assume the model is in sync with the DB as the DB should not be shared!
	int new_grp_idx = findGroup(group_name);
	if (new_grp_idx < 0) {
		// Group not found we are g
		auto add_grp_res = addGroup(group_name);
		if (std::holds_alternative<int>(add_grp_res)) {
			new_grp_idx = std::get<int>(add_grp_res);
		}
		else {
			throw std::runtime_error("SqlError1");
		}
	}
	auto new_grp = m_groups[new_grp_idx];

	auto parent = createIndex(new_grp_idx, 0, new_grp.get());
	auto idx = new_grp->connections().size();
	beginInsertRows(parent, idx, idx);
	SCOPE_EXIT { endInsertRows(); };
	auto node = std::make_shared<ConnectionConfig>(cc);
	new_grp->add(node);
	auto save_res = saveToDb(*node);
	if (save_res) {
		QString msg = save_res->text()
			% "\n" % save_res->driverText()
			% "\n" % save_res->databaseText();

		throw std::runtime_error(msg.toUtf8().data());
	}
}

void ConnectionTreeModel::save(const ConnectionConfig &cc)
{
    saveToDb(cc);
}

void ConnectionTreeModel::clearAllPasswords()
{
    for (auto group : m_groups)
        for (auto cc : group->connections())
        {
            cc->setEncodedPassword({});
            saveToDb(*cc);
        }
}

std::tuple<int, int> ConnectionTreeModel::findConfig(const QUuid uuid) const
{
	int group_idx = -1, connection_idx = -1;

	for (int grp_idx = 0; grp_idx < m_groups.size(); ++grp_idx) {
		auto && grp = m_groups[grp_idx];
		auto && conns = grp->connections();
		auto find_res = std::find_if(conns.begin(), conns.end(),
			[&uuid] (auto item) -> bool { return item->uuid() == uuid; });
		if (find_res != conns.end()) {
			group_idx = grp_idx;
			connection_idx = find_res - conns.begin();
			break;
		}
	}

	return { group_idx, connection_idx };
}

int ConnectionTreeModel::findGroup(const QString &name) const
{
	for (int idx = 0; idx < m_groups.size(); ++idx) {
		if (m_groups[idx]->name == name) return idx;
	}
	return -1;
}

std::variant<int, QSqlError> ConnectionTreeModel::addGroup(const QString &group_name)
{
	QSqlQuery q(m_db);
	q.prepare("INSERT INTO conngroup (gname) VALUES (:name)");
	q.bindValue(":name", group_name);
	if (!q.exec()) {
		auto err = q.lastError();
		return { err };
	}
	auto cg = std::make_shared<ConnectionGroup>();
	cg->conngroup_id = q.lastInsertId().toInt();
	cg->name = group_name;

	int row = m_groups.size();
	beginInsertRows({}, row, row);
	SCOPE_EXIT { endInsertRows(); };
	m_groups.push_back(cg);
	return row;
}

std::optional<QSqlError> ConnectionTreeModel::removeGroup(int row)
{
	beginRemoveRows({}, row, row);
	SCOPE_EXIT { endRemoveRows(); };
	auto id = m_groups[row]->conngroup_id;
	QSqlQuery q(m_db);
	q.prepare("DELETE FROM connection WHERE conngroup_id=:id");
	q.bindValue(":id", id);
	if (!q.exec()) {
		auto err = q.lastError();
		return { err };
	}
	q.prepare("DELETE FROM conngroup WHERE conngroup_id=:id");
	q.bindValue(":id", id);
	if (!q.exec()) {
		auto err = q.lastError();
		return { err };
	}

	m_groups.remove(row);
	return {};
}

int ConnectionTreeModel::findGroup(int conngroup_id) const
{
	auto find_res = std::find_if(m_groups.begin(), m_groups.end(),
		[conngroup_id] (auto item) { return item->conngroup_id == conngroup_id; });
	if (find_res == m_groups.end())
		return -1;
	return find_res - m_groups.begin();
}

ConnectionConfig *ConnectionTreeModel::getConfigFromModelIndex(QModelIndex index)
{
	if (!index.isValid())
		return nullptr;
	auto node = static_cast<ConnectionNode*>(index.internalPointer());
	return dynamic_cast<ConnectionConfig*>(node);
}

ConnectionGroup *ConnectionTreeModel::getGroupFromModelIndex(QModelIndex index)
{
	if (!index.isValid())
		return nullptr;
	auto node = static_cast<ConnectionNode*>(index.internalPointer());
	return dynamic_cast<ConnectionGroup*>(node);
}

std::optional<QSqlError> ConnectionTreeModel::saveToDb(const ConnectionConfig &cc)
{
	return SaveConnectionConfig(m_db, cc, cc.parent()->conngroup_id);
}


Qt::DropActions ConnectionTreeModel::supportedDropActions() const
{
	return Qt::MoveAction;
}

Qt::DropActions ConnectionTreeModel::supportedDragActions() const
{
	return Qt::MoveAction;
}

Qt::ItemFlags ConnectionTreeModel::flags(const QModelIndex &index) const
{
	Qt::ItemFlags defaultFlags = QAbstractItemModel::flags(index);

	ConnectionConfig* cfg = getConfigFromModelIndex(index);

	if (cfg)
		return Qt::ItemIsDragEnabled | defaultFlags;
	else
		return Qt::ItemIsDropEnabled | defaultFlags;
}

//bool ConnectionTreeModel::insertRows(int row, int count, const QModelIndex &parent)
//{
//	return false;
//}

//bool ConnectionTreeModel::moveRows(const QModelIndex &sourceParent, int sourceRow, int count, const QModelIndex &destinationParent, int destinationChild)
//{
//	return false;
//}
namespace {
	const auto mimeType = "application/vnd.pgLab.connection";
}

QStringList ConnectionTreeModel::mimeTypes() const
{
	return { mimeType };
}

QMimeData *ConnectionTreeModel::mimeData(const QModelIndexList &indexes) const
{
	QMimeData *mimeData = new QMimeData;
	QByteArray encodedData;

	QDataStream stream(&encodedData, QIODevice::WriteOnly);

	for (const QModelIndex &index : indexes) {
		if (index.isValid()) {
			QString text = data(index, Qt::DisplayRole).toString();
			stream << text;
		}
	}

	mimeData->setData(mimeType, encodedData);
	return mimeData;
}

bool ConnectionTreeModel::dropMimeData(const QMimeData *data, Qt::DropAction action, int row, int column, const QModelIndex &parent)
{

	return false;
}