// SPDX-FileCopyrightText: 2021 Melvin Keskin <melvo@olomono.de>
//
// SPDX-License-Identifier: LGPL-2.1-or-later

#include "QXmppAtmManager.h"

#include "QXmppCarbonManager.h"
#include "QXmppCarbonManagerV2.h"
#include "QXmppClient.h"
#include "QXmppConstants_p.h"
#include "QXmppE2eeMetadata.h"
#include "QXmppMessage.h"
#include "QXmppTrustMessageElement.h"
#include "QXmppTrustMessageKeyOwner.h"
#include "QXmppUtils.h"
#include "QXmppUtils_p.h"

#include "Async.h"

using namespace QXmpp;
using namespace QXmpp::Private;

/*!
    \class QXmppAtmManager
    \inmodule QXmpp

    \brief The QXmppAtmManager class represents a manager for
    \xep{0450}{Automatic Trust Management (ATM)}.

    For interacting with the storage, a corresponding implementation of the
    storage interface must be added.
    That implementation has to be adapted to your storage such as a database.
    In case you only need memory and no peristent storage, you can use the
    existing implementation and add the storage with it:

    \code
    QXmppAtmTrustStorage *trustStorage = new QXmppAtmTrustMemoryStorage;
    QXmppAtmManager *manager = new QXmppAtmManager(trustStorage);
    client->addExtension(manager);
    \endcode

    It is strongly recommended to enable \xep{0280}{Message Carbons} with
    \code
    client->addNewExtension<QXmppCarbonManagerV2>();
    \endcode
    and \xep{0313}{Message Archive Management} with
    \code
    QXmppMamManager *mamManager = new QXmppMamManager;
    client->addExtension(mamManager);
    \endcode
    for delivering trust messages to all online and offline endpoints.

    In addition, archiving via MAM must be enabled on the server.

    \ingroup Managers

    \since QXmpp 1.5
*/

/*!
    Constructs an ATM manager with the given \a trustStorage trust storage implementation.
*/
QXmppAtmManager::QXmppAtmManager(QXmppAtmTrustStorage *trustStorage)
    : QXmppTrustManager(trustStorage)
{
}

/*!
    Authenticates or distrusts keys manually (e.g., by the Trust Message URI of
    a scanned QR code or after entering key IDs by hand) and sends corresponding
    trust messages.

    \a encryption is the encryption protocol namespace. \a keyOwnerJid is the JID of the key
    owner. \a keyIdsForAuthentication contains the IDs of the keys being authenticated and
    \a keyIdsForDistrusting contains the IDs of the keys being distrusted.
*/
QXmppTask<void> QXmppAtmManager::makeTrustDecisions(QString encryption, QString keyOwnerJid, QList<QByteArray> keyIdsForAuthentication, QList<QByteArray> keyIdsForDistrusting)
{
    auto retrievedKeys = co_await keys(encryption, TrustLevel::Authenticated | TrustLevel::ManuallyDistrusted).withContext(this);

    const auto authenticatedKeys = retrievedKeys.value(TrustLevel::Authenticated);
    const auto manuallyDistrustedKeys = retrievedKeys.value(TrustLevel::ManuallyDistrusted);
    const auto ownJid = client()->configuration().jidBare();
    const auto ownAuthenticatedKeys = authenticatedKeys.values(ownJid);

    // Create a key owner for the keys being authenticated or
    // distrusted.
    QXmppTrustMessageKeyOwner keyOwner;
    keyOwner.setJid(keyOwnerJid);

    QList<QByteArray> modifiedAuthenticatedKeys;
    QList<QByteArray> modifiedManuallyDistrustedKeys;

    for (const auto &keyId : keyIdsForAuthentication) {
        if (!authenticatedKeys.contains(keyOwnerJid, keyId)) {
            modifiedAuthenticatedKeys.append(keyId);
        }
    }

    for (const auto &keyId : keyIdsForDistrusting) {
        if (!manuallyDistrustedKeys.contains(keyOwnerJid, keyId)) {
            modifiedManuallyDistrustedKeys.append(keyId);
        }
    }

    if (modifiedAuthenticatedKeys.isEmpty() && modifiedManuallyDistrustedKeys.isEmpty()) {
        // Skip further processing if there are no changes.
        co_return;
    }
    keyOwner.setTrustedKeys(modifiedAuthenticatedKeys);
    keyOwner.setDistrustedKeys(modifiedManuallyDistrustedKeys);

    QMultiHash<QString, QByteArray> keysBeingAuthenticated;
    QMultiHash<QString, QByteArray> keysBeingDistrusted;

    for (const auto &key : std::as_const(modifiedAuthenticatedKeys)) {
        keysBeingAuthenticated.insert(keyOwnerJid, key);
    }

    for (const auto &key : std::as_const(modifiedManuallyDistrustedKeys)) {
        keysBeingDistrusted.insert(keyOwnerJid, key);
    }

    // Create a key owner for authenticated and distrusted keys of own
    // endpoints.
    QXmppTrustMessageKeyOwner ownKeyOwner;
    ownKeyOwner.setJid(ownJid);

    if (!ownAuthenticatedKeys.isEmpty()) {
        ownKeyOwner.setTrustedKeys(ownAuthenticatedKeys);
    }

    const auto ownManuallyDistrustedKeys = manuallyDistrustedKeys.values(ownJid);
    if (!ownManuallyDistrustedKeys.isEmpty()) {
        ownKeyOwner.setDistrustedKeys(ownManuallyDistrustedKeys);
    }

    const auto areOwnKeysProcessed = keyOwnerJid == ownJid;
    if (areOwnKeysProcessed) {
        auto contactsAuthenticatedKeys = authenticatedKeys;
        contactsAuthenticatedKeys.remove(ownJid);

        const auto contactsWithAuthenticatedKeys = contactsAuthenticatedKeys.uniqueKeys();

        // Send trust messages for the keys of the own endpoints being
        // authenticated or distrusted to endpoints of contacts with
        // authenticated keys.
        // Own endpoints with authenticated keys can receive the trust
        // messages via Message Carbons.
        for (const auto &contactJid : contactsWithAuthenticatedKeys) {
            sendTrustMessage(encryption, { keyOwner }, contactJid);
        }

        // Send a trust message for the keys of the own endpoints being
        // authenticated or distrusted to other own endpoints with
        // authenticated keys.
        // It is skipped if a trust message is already delivered via
        // Message Carbons or there are no other own endpoints with
        // authenticated keys.
        auto isMessageCarbonsEnabled = [this]() {
            if (const auto *v2 = client()->findExtension<QXmppCarbonManagerV2>()) {
                return v2->enabled().value();
            }
            if (const auto *v1 = client()->findExtension<QXmppCarbonManager>()) {
                return v1->carbonsEnabled();
            }
            return false;
        };
        const auto isMessageCarbonsDisabled = !isMessageCarbonsEnabled();
        if (isMessageCarbonsDisabled || (contactsAuthenticatedKeys.isEmpty() && !ownAuthenticatedKeys.isEmpty())) {
            sendTrustMessage(encryption, { keyOwner }, ownJid);
        }

        co_await makeTrustDecisions(encryption, keysBeingAuthenticated, keysBeingDistrusted)
            .withContext(this);

        // Send a trust message for all authenticated or distrusted
        // keys to the own endpoints whose keys have been
        // authenticated.
        // It is skipped if no keys of own endpoints have been
        // authenticated.
        if (!keyOwner.trustedKeys().isEmpty()) {
            auto contactsManuallyDistrustedKeys = manuallyDistrustedKeys;
            contactsManuallyDistrustedKeys.remove(ownJid);

            auto contactJids = contactsManuallyDistrustedKeys.uniqueKeys() + contactsWithAuthenticatedKeys;

            // Remove duplicates from contactJids.
            std::sort(contactJids.begin(), contactJids.end());
            contactJids.erase(std::unique(contactJids.begin(), contactJids.end()), contactJids.end());

            QList<QXmppTrustMessageKeyOwner> allKeyOwners;

            for (const auto &contactJid : std::as_const(contactJids)) {
                QXmppTrustMessageKeyOwner contactKeyOwner;
                contactKeyOwner.setJid(contactJid);
                contactKeyOwner.setTrustedKeys(contactsAuthenticatedKeys.values(contactJid));

                if (const auto contactManuallyDistrustedKeys = contactsManuallyDistrustedKeys.values(contactJid); !contactsManuallyDistrustedKeys.isEmpty()) {
                    contactKeyOwner.setDistrustedKeys(contactManuallyDistrustedKeys);
                }

                allKeyOwners.append(contactKeyOwner);
            }

            if (!(ownKeyOwner.trustedKeys().isEmpty() && ownKeyOwner.distrustedKeys().isEmpty())) {
                allKeyOwners.append(ownKeyOwner);
            }

            if (!allKeyOwners.isEmpty()) {
                sendTrustMessage(encryption, allKeyOwners, ownJid);
            }
        }
    } else {
        // Send a trust message for the keys of the contact's endpoints
        // being authenticated or distrusted to own endpoints
        // with authenticated keys.
        if (!ownAuthenticatedKeys.isEmpty()) {
            sendTrustMessage(encryption, { keyOwner }, ownJid);
        }

        co_await makeTrustDecisions(encryption, keysBeingAuthenticated, keysBeingDistrusted).withContext(this);

        // Send a trust message for own authenticated or distrusted
        // keys to the contact's endpoints whose keys have been
        // authenticated.
        // It is skipped if no keys of contacts have been
        // authenticated or there are no keys for the trust message.
        if (!keyOwner.trustedKeys().isEmpty() && !(ownKeyOwner.trustedKeys().isEmpty() && ownKeyOwner.distrustedKeys().isEmpty())) {
            sendTrustMessage(encryption, { ownKeyOwner }, keyOwnerJid);
        }
    }
}

void QXmppAtmManager::onRegistered(QXmppClient *client)
{
    connect(client, &QXmppClient::messageReceived, this, &QXmppAtmManager::handleMessage);
}

void QXmppAtmManager::onUnregistered(QXmppClient *client)
{
    disconnect(client, &QXmppClient::messageReceived, this, &QXmppAtmManager::handleMessage);
}

/*!
    Authenticates or distrusts keys.

    \a encryption is the encryption protocol namespace. \a keyIdsForAuthentication maps key
    owners' bare JIDs to the IDs of their keys being authenticated.
    \a keyIdsForDistrusting maps key owners' bare JIDs to the IDs of their keys being distrusted.
*/
QXmppTask<void> QXmppAtmManager::makeTrustDecisions(QString encryption, QMultiHash<QString, QByteArray> keyIdsForAuthentication, QMultiHash<QString, QByteArray> keyIdsForDistrusting)
{
    // authenticate
    if (!keyIdsForAuthentication.isEmpty()) {
        co_await setTrustLevel(encryption, keyIdsForAuthentication, TrustLevel::Authenticated).withContext(this);
        if (co_await securityPolicy(encryption).withContext(this) == Toakafa) {
            co_await distrustAutomaticallyTrustedKeys(encryption, keyIdsForAuthentication.uniqueKeys()).withContext(this);
        }
        co_await makePostponedTrustDecisions(encryption, keyIdsForAuthentication.values()).withContext(this);
    }
    // distrust
    if (!keyIdsForDistrusting.isEmpty()) {
        co_await setTrustLevel(encryption, keyIdsForDistrusting, TrustLevel::ManuallyDistrusted).withContext(this);
        co_await trustStorage()->removeKeysForPostponedTrustDecisions(encryption, keyIdsForDistrusting.values());
    }
}

/*!
    Handles incoming messages and uses included trust message elements for
    making automatic trust decisions. \a message is a message that can contain a trust message
    element.
*/
QXmppTask<void> QXmppAtmManager::handleMessage(const QXmppMessage &message)
{
    const auto trustMessageElement = message.trustMessageElement();

    // Skip further processing in the following cases:
    // 1. The message does not contain an ATM trust message element.
    // 2. The trust message is sent by this endpoint and reflected via
    //    Message Carbons.
    if (!trustMessageElement ||
        trustMessageElement->usage() != ns_atm ||
        message.from() == client()->configuration().jid()) {
        co_return;
    }

    const auto senderJid = QXmppUtils::jidToBareJid(message.from());
    const auto e2eeMetadata = message.e2eeMetadata();
    const auto senderKey = e2eeMetadata ? e2eeMetadata->senderKey() : QByteArray();
    const auto encryption = trustMessageElement->encryption();

    auto senderKeyTrustLevel = co_await trustLevel(encryption, senderJid, senderKey).withContext(this);

    // key owner JIDs mapped to key IDs
    QMultiHash<QString, QByteArray> keysBeingAuthenticated;
    QMultiHash<QString, QByteArray> keysBeingDistrusted;

    QList<QXmppTrustMessageKeyOwner> keyOwnersForPostponedTrustDecisions;

    const auto ownJid = client()->configuration().jidBare();
    const auto isOwnTrustMessage = senderJid == ownJid;
    const auto keyOwners = trustMessageElement->keyOwners();

    for (const auto &keyOwner : keyOwners) {
        const auto keyOwnerJid = keyOwner.jid();

        // A trust message from an own endpoint is allowed to
        // authenticate or distrust the keys of own endpoints and
        // endpoints of contacts.
        // Whereas a trust message from an endpoint of a contact is
        // only allowed to authenticate or distrust the keys of that
        // contact's own endpoints.
        const auto isSenderQualifiedForTrustDecisions = isOwnTrustMessage || senderJid == keyOwnerJid;
        if (isSenderQualifiedForTrustDecisions) {
            // Make trust decisions if the key of the sender is
            // authenticated.
            // Othwerwise, store the keys of the trust message for
            // making the trust decisions as soon as the key of the
            // sender is authenticated.
            if (senderKeyTrustLevel == TrustLevel::Authenticated) {
                const auto trustedKeys = keyOwner.trustedKeys();
                for (const auto &key : trustedKeys) {
                    keysBeingAuthenticated.insert(keyOwnerJid, key);
                }

                const auto distrustedKeys = keyOwner.distrustedKeys();
                for (const auto &key : distrustedKeys) {
                    keysBeingDistrusted.insert(keyOwnerJid, key);
                }
            } else {
                keyOwnersForPostponedTrustDecisions.append(keyOwner);
            }
        }
    }

    auto addPostponedTask = trustStorage()->addKeysForPostponedTrustDecisions(encryption, senderKey, keyOwnersForPostponedTrustDecisions);
    auto trustDecisionTask = makeTrustDecisions(encryption, keysBeingAuthenticated, keysBeingDistrusted);
    co_await addPostponedTask;
    co_await trustDecisionTask;
}

/*!
    Distrusts all formerly automatically trusted keys (as specifed by the
    security policy TOAKAFA).

    \a encryption is the encryption protocol namespace. \a keyOwnerJids contains the bare JIDs
    of the key owners.
*/
QXmppTask<void> QXmppAtmManager::distrustAutomaticallyTrustedKeys(const QString &encryption, const QList<QString> &keyOwnerJids)
{
    return setTrustLevel(encryption, keyOwnerJids, TrustLevel::AutomaticallyTrusted, TrustLevel::AutomaticallyDistrusted);
}

/*!
    Authenticates or distrusts keys for whom earlier trust messages were
    received but not used for authenticating or distrusting at that time.

    As soon as the senders' keys have been authenticated, all postponed trust
    decisions can be performed by this method.

    \a encryption is the encryption protocol namespace. \a senderKeyIds contains the IDs of the
    keys that were used by the senders.
*/
QXmppTask<void> QXmppAtmManager::makePostponedTrustDecisions(QString encryption, const QList<QByteArray> &senderKeyIds)
{
    auto postponedTrustDecisionKeys =
        co_await trustStorage()
            ->keysForPostponedTrustDecisions(encryption, senderKeyIds)
            .withContext(this);

    // JIDs of key owners mapped to the IDs of their keys
    auto keysBeingAuthenticated = postponedTrustDecisionKeys.value(true);
    auto keysBeingDistrusted = postponedTrustDecisionKeys.value(false);

    auto cleanupTask = trustStorage()->removeKeysForPostponedTrustDecisions(encryption, keysBeingAuthenticated.values(), keysBeingDistrusted.values());
    auto trustDecisionTask = makeTrustDecisions(encryption, keysBeingAuthenticated, keysBeingDistrusted);
    co_await cleanupTask;
    co_await trustDecisionTask;
}

/*!
    Sends a trust message.

    \a encryption is the namespace of the encryption. \a keyOwners contains the key owners with
    the data for authentication or distrusting. \a recipientJid is the JID of the recipient.
*/
QXmppTask<QXmpp::SendResult> QXmppAtmManager::sendTrustMessage(const QString &encryption, const QList<QXmppTrustMessageKeyOwner> &keyOwners, const QString &recipientJid)
{
    QXmppTrustMessageElement trustMessageElement;
    trustMessageElement.setUsage(ns_atm.toString());
    trustMessageElement.setEncryption(encryption);
    trustMessageElement.setKeyOwners(keyOwners);

    QXmppMessage message;
    message.setTo(recipientJid);
    message.setTrustMessageElement(trustMessageElement);

    QXmppSendStanzaParams params;
    params.setAcceptedTrustLevels(TrustLevel::Authenticated);

    return client()->sendSensitive(std::move(message), params);
}
