Mercury-IM/domain/src/main/java/org/mercury_im/messenger/core/crypto/MercuryOpenPgpManager.java

package org.mercury_im.messenger.core.crypto;

import org.jivesoftware.smack.AbstractConnectionListener;
import org.jivesoftware.smack.XMPPConnection;
import org.jivesoftware.smackx.ox.OpenPgpManager;
import org.jivesoftware.smackx.ox.OpenPgpSecretKeyBackupPassphrase;
import org.jivesoftware.smackx.ox.callback.SecretKeyPassphraseCallback;
import org.jivesoftware.smackx.ox.crypto.OpenPgpProvider;
import org.jivesoftware.smackx.ox.crypto.PainlessOpenPgpProvider;
import org.jivesoftware.smackx.ox.exception.InvalidBackupCodeException;
import org.jivesoftware.smackx.ox.exception.NoBackupFoundException;
import org.jivesoftware.smackx.ox.store.definition.OpenPgpStore;
import org.jivesoftware.smackx.ox_im.OXInstantMessagingManager;
import org.jivesoftware.smackx.pubsub.PubSubException;
import org.mercury_im.messenger.core.SchedulersFacade;
import org.mercury_im.messenger.core.connection.MercuryConnection;
import org.mercury_im.messenger.core.data.repository.DirectChatRepository;
import org.mercury_im.messenger.core.data.repository.MessageRepository;
import org.mercury_im.messenger.core.data.repository.OpenPgpRepository;
import org.mercury_im.messenger.core.data.repository.PeerRepository;
import org.mercury_im.messenger.core.store.crypto.MercuryOpenPgpStore;
import org.mercury_im.messenger.core.store.message.MercuryMessageStore;

import java.util.logging.Level;
import java.util.logging.Logger;

import javax.inject.Inject;

public class MercuryOpenPgpManager {

    private static final Logger LOGGER = Logger.getLogger(MercuryOpenPgpManager.class.getName());

    private final PeerRepository peerRepository;
    private final DirectChatRepository directChatRepository;
    private final MessageRepository messageRepository;
    private final OpenPgpRepository openPgpRepository;
    private final SchedulersFacade schedulers;
    private final OpenPgpSecretKeyBackupPassphraseGenerator passphraseGenerator;
    private final LocalOxKeyGenerationStrategy keyGenerationStrategy;

    @Inject
    public MercuryOpenPgpManager(PeerRepository peerRepository,
                                 DirectChatRepository directChatRepository,
                                 MessageRepository messageRepository,
                                 OpenPgpRepository openPgpRepository,
                                 OpenPgpSecretKeyBackupPassphraseGenerator passphraseGenerator,
                                 LocalOxKeyGenerationStrategy keyGenerationStrategy,
                                 SchedulersFacade schedulers) {
        this.peerRepository = peerRepository;
        this.directChatRepository = directChatRepository;
        this.messageRepository = messageRepository;
        this.openPgpRepository = openPgpRepository;
        this.schedulers = schedulers;
        this.keyGenerationStrategy = keyGenerationStrategy;
        this.passphraseGenerator = passphraseGenerator;
    }

    public void initialize(MercuryConnection connection) {
        if (connection.getConnection().isAuthenticated()) {
            setup(connection);
        } else {
            connection.getConnection().addConnectionListener(new AbstractConnectionListener() {
                @Override
                public void authenticated(XMPPConnection con, boolean resumed) {
                    if (!resumed) {
                        setup(connection);
                    }
                }
            });
        }
    }

    private void setup(MercuryConnection connection) {
        OpenPgpStore store = new MercuryOpenPgpStore(connection.getAccountId(), openPgpRepository, schedulers);
        OpenPgpProvider provider = new PainlessOpenPgpProvider(store);
        OpenPgpManager oxManager = OpenPgpManager.getInstanceFor(connection.getConnection());
        oxManager.setOpenPgpProvider(provider);
        OpenPgpSecretKeyBackupPassphrase passphrase = passphraseGenerator.generateBackupPassphrase();


        generateAndPublish(connection, oxManager, passphrase);
    }

    private void generateAndPublish(MercuryConnection connection, OpenPgpManager oxManager, OpenPgpSecretKeyBackupPassphrase passphrase) {
        boolean mustGenerate = false;
        try {
            if (!oxManager.hasSecretKeysAvailable()) {
                oxManager.generateAndImportKeyPair(connection.getAccount().getJid());
                //if (OpenPgpManager.serverSupportsSecretKeyBackups(connection.getConnection())) {
                //    oxManager.backupSecretKeyToServer(
                //            availableSecretKeys -> availableSecretKeys,
                //            passphrase);
                //}
            }
            oxManager.announceSupportAndPublish();
            OXInstantMessagingManager oximManager = OXInstantMessagingManager.getInstanceFor(connection.getConnection());
            oximManager.addOxMessageListener(new MercuryMessageStore(connection.getAccount(),
                    peerRepository, directChatRepository, messageRepository, schedulers));
            oximManager.announceSupportForOxInstantMessaging();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    private void tryRestoringAndPublish(MercuryConnection connection, OpenPgpManager oxManager, OpenPgpSecretKeyBackupPassphrase passphrase) {
        boolean mustGenerate = false;
        try {
            if (!oxManager.hasSecretKeysAvailable()) {
                mustGenerate = true;
                if (OpenPgpManager.serverSupportsSecretKeyBackups(connection.getConnection())) {
                    try {
                        oxManager.restoreSecretKeyServerBackup(
                                () -> passphrase);
                        mustGenerate = false;
                        LOGGER.log(Level.INFO, "Successfully restored secret key backup!");
                    } catch (NoBackupFoundException | PubSubException.NotALeafNodeException | InvalidBackupCodeException e) {
                        LOGGER.log(Level.INFO, "Error restoring secret key backup.", e);
                    }
                }
            }

            if (mustGenerate) {
                oxManager.generateAndImportKeyPair(connection.getAccount().getJid());
                if (OpenPgpManager.serverSupportsSecretKeyBackups(connection.getConnection())) {
                    oxManager.backupSecretKeyToServer(
                            availableSecretKeys -> availableSecretKeys,
                            passphrase);
                }
            }
            oxManager.announceSupportAndPublish();
            OXInstantMessagingManager oximManager = OXInstantMessagingManager.getInstanceFor(connection.getConnection());
            oximManager.addOxMessageListener(new MercuryMessageStore(connection.getAccount(),
                    peerRepository, directChatRepository, messageRepository, schedulers));
            oximManager.announceSupportForOxInstantMessaging();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public boolean mustPromptForRestore() {
        return keyGenerationStrategy.promptForBackupRestoreIfNoLocalKeyPresent();
    }
}