// SPDX-FileCopyrightText: 2019 Alexey Andreyev <aa13q@ya.ru>
// SPDX-FileCopyrightText: 2019 Kitsune Ral <Kitsune-Ral@users.sf.net>
// SPDX-License-Identifier: LGPL-2.1-or-later

#ifdef Quotient_E2EE_ENABLED
#include "encryptionmanager.h"

#include "connection.h"
#include "events/encryptedfile.h"
#include "database.h"

#include "csapi/keys.h"

#include <QtCore/QHash>
#include <QtCore/QStringBuilder>
#include <QtCore/QCryptographicHash>

#include "e2ee/e2ee.h"
#include "e2ee/qolmaccount.h"
#include "e2ee/qolmsession.h"
#include "e2ee/qolmmessage.h"
#include "e2ee/qolmerrors.h"
#include "e2ee/qolmutils.h"
#include <functional>
#include <memory>

#include <openssl/evp.h>

using namespace Quotient;
using std::move;

class EncryptionManager::Private {
public:
    EncryptionManager* q;

    Connection* connection;

    // A map from SenderKey to vector of InboundSession
    UnorderedMap<QString, std::vector<QOlmSessionPtr>> sessions;

    void loadSessions() {
        sessions = connection->database()->loadOlmSessions(connection->picklingMode());
    }
    void saveSession(QOlmSessionPtr& session, const QString &senderKey) {
        auto pickleResult = session->pickle(connection->picklingMode());
        if (std::holds_alternative<QOlmError>(pickleResult)) {
            qCWarning(E2EE) << "Failed to pickle olm session. Error" << std::get<QOlmError>(pickleResult);
            return;
        }
        connection->database()->saveOlmSession(senderKey, session->sessionId(), std::get<QByteArray>(pickleResult));
    }
    QString sessionDecryptPrekey(const QOlmMessage& message, const QString &senderKey, std::unique_ptr<QOlmAccount>& olmAccount)
    {
        Q_ASSERT(message.type() == QOlmMessage::PreKey);
        for(auto& session : sessions[senderKey]) {
            const auto matches = session->matchesInboundSessionFrom(senderKey, message);
            if(std::holds_alternative<bool>(matches) && std::get<bool>(matches)) {
                qCDebug(E2EE) << "Found inbound session";
                const auto result = session->decrypt(message);
                if(std::holds_alternative<QString>(result)) {
                    return std::get<QString>(result);
                } else {
                    qCDebug(E2EE) << "Failed to decrypt prekey message";
                    return {};
                }
            }
        }
        qCDebug(E2EE) << "Creating new inbound session";
        auto newSessionResult = olmAccount->createInboundSessionFrom(senderKey.toUtf8(), message);
        if(std::holds_alternative<QOlmError>(newSessionResult)) {
            qCWarning(E2EE) << "Failed to create inbound session for" << senderKey << std::get<QOlmError>(newSessionResult);
            return {};
        }
        auto newSession = std::move(std::get<QOlmSessionPtr>(newSessionResult));
        auto error = olmAccount->removeOneTimeKeys(newSession);
        if (error) {
            qWarning(E2EE) << "Failed to remove one time key for session" << newSession->sessionId();
        }
        const auto result = newSession->decrypt(message);
        saveSession(newSession, senderKey);
        sessions[senderKey].push_back(std::move(newSession));
        if(std::holds_alternative<QString>(result)) {
            return std::get<QString>(result);
        } else {
            qCDebug(E2EE) << "Failed to decrypt prekey message with new session";
            return {};
        }
    }
    QString sessionDecryptGeneral(const QOlmMessage& message, const QString &senderKey)
    {
        Q_ASSERT(message.type() == QOlmMessage::General);
        for(auto& session : sessions[senderKey]) {
            const auto result = session->decrypt(message);
            if(std::holds_alternative<QString>(result)) {
                return std::get<QString>(result);
            }
        }
        qCWarning(E2EE) << "Failed to decrypt message";
        return {};
    }
};

EncryptionManager::EncryptionManager(QObject* parent)
    : QObject(parent)
    , d(std::make_unique<Private>())
{
    d->q = this;
    d->connection = static_cast<Connection *>(parent);
    d->loadSessions();
}

EncryptionManager::~EncryptionManager() = default;

QString EncryptionManager::sessionDecryptMessage(
    const QJsonObject& personalCipherObject, const QByteArray& senderKey, std::unique_ptr<QOlmAccount>& account)
{
    QString decrypted;
    int type = personalCipherObject.value(TypeKeyL).toInt(-1);
    QByteArray body = personalCipherObject.value(BodyKeyL).toString().toLatin1();
    if (type == 0) {
        QOlmMessage preKeyMessage(body, QOlmMessage::PreKey);
        decrypted = d->sessionDecryptPrekey(preKeyMessage, senderKey, account);
    } else if (type == 1) {
        QOlmMessage message(body, QOlmMessage::General);
        decrypted = d->sessionDecryptGeneral(message, senderKey);
    }
    return decrypted;
}

QByteArray EncryptionManager::decryptFile(const QByteArray &ciphertext, EncryptedFile* file)
{
    const auto key = QByteArray::fromBase64(file->key.k.replace(QLatin1Char('_'), QLatin1Char('/')).replace(QLatin1Char('-'), QLatin1Char('+')).toLatin1());
    const auto iv =  QByteArray::fromBase64(file->iv.toLatin1());
    const auto sha256 = QByteArray::fromBase64(file->hashes["sha256"].toLatin1());
    if(sha256 != QCryptographicHash::hash(ciphertext, QCryptographicHash::Sha256)) {
        qCWarning(E2EE) << "Hash verification failed for file";
        return QByteArray();
    }
    QByteArray plaintext(ciphertext.size(), 0);
    EVP_CIPHER_CTX *ctx;
    int length;
    ctx = EVP_CIPHER_CTX_new();
    EVP_DecryptInit_ex(ctx, EVP_aes_256_ctr(), NULL, (const unsigned char *)key.data(), (const unsigned char *)iv.data());
    EVP_DecryptUpdate(ctx, (unsigned char *)plaintext.data(), &length, (const unsigned char *)ciphertext.data(), ciphertext.size());
    EVP_DecryptFinal_ex(ctx, (unsigned char *)plaintext.data() + length, &length);
    EVP_CIPHER_CTX_free(ctx);
    return plaintext;
}
#endif // Quotient_E2EE_ENABLED