/*
 * QDigiDocCrypto
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
 *
 */

#include "CryptoDoc.h"

#include "Application.h"
#include "CDocSupport.h"
#include "TokenData.h"
#include "QCryptoBackend.h"
#include "QSigner.h"
#include "Settings.h"
#include "SslCertificate.h"
#include "Utils.h"
#include "dialogs/FileDialog.h"
#include "dialogs/WarningDialog.h"

#include <QtCore/QDir>
#include <QtCore/QFileInfo>
#include <QtCore/QRegularExpression>
#include <QtCore/QThread>
#include <QtCore/QUrl>
#include <QtCore/QUrlQuery>
#include <QtGui/QDesktopServices>
#include <QtNetwork/QSslKey>
#include <QtWidgets/QMessageBox>

#include <filesystem>

#include <cdoc/CDocReader.h>
#include <cdoc/CDocWriter.h>
#include <cdoc/Certificate.h>
#include <cdoc/Crypto.h>

using namespace ria::qdigidoc4;

Q_LOGGING_CATEGORY(CRYPTO, "CRYPTO")

CKey::CKey(QSslCertificate _rcpt_cert) : lock(libcdoc::Lock::PUBLIC_KEY), rcpt_cert(_rcpt_cert) {
	SslCertificate ssl(rcpt_cert);
	lock.pk_type = (rcpt_cert.publicKey().algorithm() == QSsl::Ec) ? libcdoc::PKType::ECC : libcdoc::PKType::RSA;
	QByteArray der = ssl.publicKeyDer();
	lock.setBytes(libcdoc::Lock::RCPT_KEY, std::vector<uint8_t>(der.cbegin(), der.cend()));
	der = rcpt_cert.toDer();
	lock.setBytes(libcdoc::Lock::CERT, std::vector<uint8_t>(der.cbegin(), der.cend()));
}

bool
CKey::operator==(const CKey& rhs) const
{
	if (!lock.isPKI() || !rhs.lock.isPKI()) return false;
	return lock.getBytes(libcdoc::Lock::RCPT_KEY) == rhs.lock.getBytes(libcdoc::Lock::RCPT_KEY);
}

struct CryptoDoc::Private
{
	std::unique_ptr<libcdoc::CDocReader> reader;

	QString			fileName;
	int version = -1;
	CDocumentModel	*documents = new CDocumentModel(this);
	QStringList		tempFiles;

	// libcdoc handlers
	DDConfiguration conf;
	DDCryptoBackend crypto;
	DDNetworkBackend network;

	std::vector<IOEntry> files;
	std::vector<CKey> keys;

	bool isEncryptedWarning(const QString &title) const;

	bool isEncrypted() const {
		return reader != nullptr;
	}

	void createCDocReader(const QString &filename) {
		reader = std::unique_ptr<libcdoc::CDocReader>(libcdoc::CDocReader::createReader(filename.toStdString(), &conf, &crypto, &network));
		if (!reader) 
			return;
		keys.clear();
		for (auto& key : reader->getLocks()) {
			keys.emplace_back(key);
		}
	}
};

bool CryptoDoc::Private::isEncryptedWarning(const QString &title) const
{
	if(fileName.isEmpty() )
		WarningDialog::create()->withTitle(title)->withText(CryptoDoc::tr("Container is not open"))->open();
	if(isEncrypted())
		WarningDialog::create()->withTitle(title)->withText(CryptoDoc::tr("Container is encrypted"))->open();
	return fileName.isEmpty() || isEncrypted();
}

CDocumentModel::CDocumentModel(CryptoDoc::Private *doc) : d(doc) {}

bool CDocumentModel::addFile(const QString &file, const QString &mime)
{
	if(d->isEncryptedWarning(DocumentModel::tr("Failed to add file")))
		return false;

	QFileInfo info(file);
	if(!addFileCheck(d->fileName, info))
		return false;

	if(d->fileName.last(3) == QLatin1String("cdoc") && info.size() > 120*1024*1024)
	{
		WarningDialog::create()
			->withTitle(DocumentModel::tr("Failed to add file"))
			->withText(tr("Added file(s) exceeds the maximum size limit of the container (∼120MB). "
				"<a href='https://www.id.ee/en/article/encrypting-large-120-mb-files/'>Read more about it</a>"))
			->open();
		return false;
	}

	auto data = std::make_unique<QFile>(file);
	if(!data->open(QFile::ReadOnly))
		return false;
	d->files.push_back({
		QFileInfo(file).fileName().toStdString(),
		mime.toStdString(),
		data->size(),
		std::move(data),
	});
	emit added(FileDialog::normalized(info.fileName()));
	return true;
}

void CDocumentModel::addTempReference(const QString &file)
{
	d->tempFiles.append(file);
}

QString CDocumentModel::copy(int row, const QString &dst) const
{
	auto &file = d->files.at(row);
	if( QFile::exists( dst ) )
		QFile::remove( dst );
	file.data->seek(0);
	if(QFile f(dst); f.open(QFile::WriteOnly) && copyIODevice(file.data.get(), &f) == file.size)
		return dst;
	WarningDialog::create()
		->withTitle(FileDialog::tr("Failed to save file"))
		->withText(dst)
		->open();
	return {};
}

QString CDocumentModel::data(int row) const
{
	return FileDialog::normalized(QString::fromStdString(d->files.at(row).name));
}

quint64 CDocumentModel::fileSize(int row) const
{
	return d->files.at(row).size;
}

QString CDocumentModel::mime(int row) const
{
	return FileDialog::normalized(QString::fromStdString(d->files.at(row).mime));
}

void CDocumentModel::open(int row)
{
	if(d->isEncrypted())
		return;
	QString path = FileDialog::tempPath(FileDialog::safeName(data(row)));
	if(!verifyFile(path))
		return;
	if(copy(row, path).isEmpty())
		return;
	d->tempFiles.append(path);
	FileDialog::setReadOnly(path);
	if(FileDialog::isSignedPDF(path))
		Application::showClient({ std::move(path) }, false, false, true);
	else
		QDesktopServices::openUrl(QUrl::fromLocalFile(path));
}

bool CDocumentModel::removeRow(int row)
{
	if(d->isEncryptedWarning(DocumentModel::tr("Failed remove document from container")))
		return false;

	if(d->files.empty() || row >= d->files.size())
	{
		WarningDialog::create()
			->withTitle(DocumentModel::tr("Failed remove document from container"))
			->withText(DocumentModel::tr("Internal error"))
			->open();
		return false;
	}

	d->files.erase(d->files.cbegin() + row);
	return true;
}

int CDocumentModel::rowCount() const
{
	return int(d->files.size());
}

QString CDocumentModel::save(int row, const QString &path) const
{
	if(d->isEncrypted())
		return {};

	QString fileName = copy(row, path);
	QFileInfo f(fileName);
	if(!f.exists())
		return {};
	FileDialog::setFileZone(fileName, d->fileName);
	return fileName;
}

CryptoDoc::CryptoDoc( QObject *parent )
	: QObject(parent)
	, d(new Private)
{
	const_cast<QLoggingCategory&>(CRYPTO()).setEnabled(QtDebugMsg,
		QFile::exists(QStringLiteral("%1/%2.log").arg(QDir::tempPath(), Application::applicationName())));
}

CryptoDoc::~CryptoDoc() {
	clear();
	delete d;
}

bool
CryptoDoc::supportsSymmetricKeys() const
{
	return !d->reader && (d->version == 2);
}

bool CryptoDoc::addEncryptionKey(const CKey& key) {
	if(d->isEncryptedWarning(tr("Failed to add key")))
		return false;
	for (auto &k : d->keys) {
		if (k == key) {
			WarningDialog::create()
				->withTitle(tr("Failed to add key"))
				->withText(tr("Key already exists"))
				->open();
			return false;
		}
	}
	d->keys.push_back(key);
	return true;
}

bool CryptoDoc::canDecrypt(const QSslCertificate &cert) {
	if (!d->reader)
		return false;
	if (cert.isNull())
		return false;
	QByteArray der = cert.toDer();
	return d->reader->getLockForCert(
		std::vector<uint8_t>(der.cbegin(), der.cend())) >= 0;
}

void CryptoDoc::clear(const QString &file, int version)
{
	for(const QString &f: qAsConst(d->tempFiles))
	{
		// reset read-only attribute to enable delete file
		FileDialog::setReadOnly(f, false);
		QFile::remove(f);
	}
	d->tempFiles.clear();
	d->reader.reset();
	d->files.clear();
	d->fileName = file;
	d->version = version;
}

ContainerState CryptoDoc::state() const
{
	return d->isEncrypted() ? EncryptedContainer : UnencryptedContainer;
}

bool CryptoDoc::decrypt(const libcdoc::Lock *lock, const QByteArray& secret)
{
	if(!d->reader)
	{
		WarningDialog::create()
			->withTitle(QSigner::tr("Failed to decrypt document"))
			->withText(tr("Container is not open"))
			->open();
		return false;
	}

	int lock_idx = -1;
	const std::vector<libcdoc::Lock> &locks = d->reader->getLocks();
	if (lock == nullptr) {
		QByteArray der = qApp->signer()->tokenauth().cert().toDer();
		lock_idx = d->reader->getLockForCert(
			std::vector<uint8_t>(der.cbegin(), der.cend()));
		if (lock_idx < 0) {
			WarningDialog::create()
				->withTitle(QSigner::tr("Failed to decrypt document"))
				->withText(tr("You do not have the key to decrypt this document"))
				->open();
			return false;
		}
		lock = &locks.at(lock_idx);
	} else {
		for (lock_idx = 0; lock_idx < locks.size(); lock_idx++) {
			if (lock->label == locks[lock_idx].label) {
				lock = &locks.at(lock_idx);
				break;
			}
		}
		if (lock_idx >= locks.size())
			lock_idx = -1;
	}
	if (!lock || (lock->isSymmetric() && secret.isEmpty())) {
		WarningDialog::create()
			->withTitle(QSigner::tr("Failed to decrypt document"))
			->withText(tr("You do not have the key to decrypt this document"))
			->open();
		return false;
	}

	if (d->reader->version == 2 &&
		(lock->type == libcdoc::Lock::Type::SERVER) &&
		!Settings::CDOC2_NOTIFICATION.isSet()) {
		auto *dlg = WarningDialog::create()
			->withTitle(tr("You must enter your PIN code twice in order to decrypt the CDOC2 container"))
			->withText(tr(
				"The first PIN entry is required for authentication to the key server referenced in the CDOC2 container. "
				"Second PIN entry is required to decrypt the CDOC2 container."))
			->setCancelText(WarningDialog::Cancel)
			->addButton(WarningDialog::OK, QMessageBox::Ok)
			->addButton(tr("Don't show again"), QMessageBox::Ignore);
		switch (dlg->exec())
		{
		case QMessageBox::Ok: break;
		case QMessageBox::Ignore:
			Settings::CDOC2_NOTIFICATION = true;
			break;
		default:
			return false;
		}
	}

	d->crypto.secret.assign(secret.cbegin(), secret.cend());

	TempListConsumer cons;
	libcdoc::result_t result = waitFor([&]{
		std::vector<uint8_t> fmk;
		auto scope = qScopeGuard([&] {
			std::fill(fmk.begin(), fmk.end(), 0);
		});
		libcdoc::result_t result = d->reader->getFMK(fmk, lock_idx);
		qCDebug(CRYPTO) << "getFMK result: " << result << " " << QString::fromStdString(d->reader->getLastErrorStr());
		if (result != libcdoc::OK)
			return result;
		result = d->reader->decrypt(fmk, &cons);
		qCDebug(CRYPTO) << "Decryption result: " << result << " " << QString::fromStdString(d->reader->getLastErrorStr());
		return result;
	});
	if (result != libcdoc::OK) {
		QString str;
		const std::string &msg = d->reader->getLastErrorStr();
		switch (result) {
		case libcdoc::WRONG_KEY:
			str = (lock->type == libcdoc::Lock::PASSWORD) ? tr("Wrong password.") : tr("Wrong key.");
			break;
		case libcdoc::HASH_MISMATCH:
		case libcdoc::DATA_FORMAT_ERROR:
			str = tr("Corrupted or tampered file.");
			break;
		case libcdoc::INPUT_ERROR:
		case libcdoc::INPUT_STREAM_ERROR:
			str = tr("Cannot read file.");
			break;
		case DDCryptoBackend::PIN_CANCELED:
			str = tr("PIN entry canceled");
			break;
		case DDCryptoBackend::PIN_INCORRECT:
			str = tr("PIN incorrect");
			break;
		case DDCryptoBackend::PIN_LOCKED:
			str = tr("PIN locked");
			break;
		default:
			str = tr("Please check your internet connection and network settings.");
			break;
		}
		WarningDialog::create()
			->withTitle(QSigner::tr("Failed to decrypt document"))
			->withText(str)
			->withDetails(QString::fromStdString(msg))
			->open();
		return false;
	}

	d->files = std::move(cons.files);
	// Success, immediately create writer from reader
	d->keys.clear();
	for(const libcdoc::Lock &l: d->reader->getLocks())
	{
		if(!l.isPKI())
			continue;
		d->keys.emplace_back(l);
	}
	d->reader.reset();
	return !d->isEncrypted();
}

DocumentModel *CryptoDoc::documentModel() const { return d->documents; }

bool CryptoDoc::encrypt(const QString &filename, const QString& label, const QByteArray& secret)
{
	// I think the correct semantics is to fail if container is already encrypted
	if(d->reader)
		return false;
	if(!filename.isEmpty() )
		d->fileName = filename;
	if(d->fileName.isEmpty() )
	{
		WarningDialog::create()
			->withTitle(tr("Failed to encrypt document"))
			->withText(tr("Container is not open"))
			->open();
		return false;
	}
	if(secret.isEmpty() && d->keys.empty())
	{
		WarningDialog::create()
			->withTitle(tr("Failed to encrypt document"))
			->withText(tr("No keys specified"))
			->open();
		return false;
	}
	QString writer_last_error;
	libcdoc::result_t result = waitFor([&] -> libcdoc::result_t {
		qCDebug(CRYPTO) << "Encrypt" << d->fileName;
		auto writer = std::unique_ptr<libcdoc::CDocWriter>(libcdoc::CDocWriter::createWriter(d->version, d->fileName.toStdString(), &d->conf, &d->crypto, &d->network));
		if (!writer)
			return libcdoc::OUTPUT_ERROR;
		std::vector<libcdoc::Recipient> enc_keys;
		std::string keyserver_id;
		if ((d->version == 2) && Settings::CDOC2_USE_KEYSERVER) {
			keyserver_id = Settings::CDOC2_DEFAULT_KEYSERVER;
		}
		// Encrypt for address list
		for (const auto &key : d->keys) {
			if (!key.rcpt_cert.isNull()) {
				QByteArray ba = key.rcpt_cert.toDer();
				enc_keys.push_back(keyserver_id.empty() ?
					libcdoc::Recipient::makeCertificate({}, {ba.cbegin(), ba.cend()}) :
					libcdoc::Recipient::makeServer({}, {ba.cbegin(), ba.cend()}, keyserver_id));
			} else {
				switch (key.lock.type) {
				case libcdoc::Lock::CDOC1:
					enc_keys.push_back(libcdoc::Recipient::makePublicKey(key.lock));
					break;
				case libcdoc::Lock::PUBLIC_KEY:
				case libcdoc::Lock::SERVER:
					if (keyserver_id.empty()) {
						enc_keys.push_back(libcdoc::Recipient::makePublicKey(key.lock));
					} else {
						enc_keys.push_back(libcdoc::Recipient::makeServer(key.lock, keyserver_id));
					}
					break;
				default:
					break;
				}
			}
		}
		// Encrypt with symmetric key
		if (!secret.isEmpty()) {
			// NIST recommends at least 600000 iterations for PBKDF2 with SHA-256, see https://csrc.nist.gov/publications/detail/sp/800-132/final
			d->crypto.secret.assign(secret.cbegin(), secret.cend());
			libcdoc::Recipient rcpt = libcdoc::Recipient::makeSymmetric({}, 600000);
			rcpt.setLabelValue(libcdoc::CDoc2::Label::LABEL, label.toStdString());
			enc_keys.push_back(std::move(rcpt));
		}
		StreamListSource slsrc(d->files);
		libcdoc::result_t result = writer->encrypt(slsrc, enc_keys);
		writer_last_error = QString::fromStdString(writer->getLastErrorStr());
		qCDebug(CRYPTO) << "Encryption result: " << result << ' ' << writer_last_error;
		if (result != libcdoc::OK)
			if(QFile::exists(d->fileName))
				QFile::remove(d->fileName);
		return result;
	});
	d->crypto.secret.clear();
	if (result != libcdoc::OK) {
		WarningDialog::create()
			->withTitle(tr("Failed to encrypt document"))
			->withText(tr("Please check your internet connection and network settings."))
			->withDetails(writer_last_error)
			->open();
	} else {
		this->clear(d->fileName);
		this->open(d->fileName);
	}
	return d->isEncrypted();
}

QString CryptoDoc::fileName() const { return d->fileName; }

const std::vector<CKey>&
CryptoDoc::keys() const
{
	return d->keys;
}

bool CryptoDoc::move(const QString &to)
{
	if(d->isEncrypted())
		return false;
	d->fileName = to;
	return true;
}

bool CryptoDoc::open(const QString &file)
{
	clear(file);
	d->createCDocReader(file);
	if(!d->reader) {
		WarningDialog::create()
			->withTitle(CryptoDoc::tr("Failed to open document"))
			->withText(CryptoDoc::tr("Unsupported file format"))
			->open();
		return false;
	}
	d->version = d->reader->version;
	std::vector<libcdoc::FileInfo> files = CDocSupport::getCDocFileList(file);
	for (auto& f : files) {
		d->files.push_back({f.name, {}, f.size, {}});
	}
	Application::addRecent( file );
	return true;
}

void CryptoDoc::removeKey(unsigned int id)
{
	if(!d->isEncryptedWarning(tr("Failed to remove key")))
		d->keys.erase(d->keys.begin() + id);
}

void CryptoDoc::clearKeys()
{
	if(!d->isEncryptedWarning(tr("Failed to remove key")))
		d->keys.clear();
}

bool CryptoDoc::saveCopy(const QString &filename)
{
	QFileInfo src(d->fileName);
	QFileInfo dst(filename);
	if(src == dst)
		return true;
	std::error_code ec;
	std::filesystem::copy_file(src.filesystemFilePath(), dst.filesystemFilePath(),
		std::filesystem::copy_options::overwrite_existing, ec);
	return !ec;
}
