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

package org.mercury_im.messenger.core.crypto;

import org.bouncycastle.openpgp.PGPPublicKeyRing;
import org.bouncycastle.openpgp.PGPPublicKeyRingCollection;
import org.jivesoftware.smack.AbstractConnectionListener;
import org.jivesoftware.smack.XMPPConnection;
import org.jivesoftware.smackx.ox.OpenPgpManager;
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.util.SecretKeyBackupHelper;
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.data.repository.OpenPgpRepository;
import org.mercury_im.messenger.core.data.repository.Repositories;
import org.mercury_im.messenger.core.store.crypto.MercuryOpenPgpStore;
import org.mercury_im.messenger.core.store.message.MercuryMessageStore;
import org.mercury_im.messenger.core.xmpp.MercuryConnection;
import org.pgpainless.key.OpenPgpV4Fingerprint;

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 Repositories repositories;
    private final OpenPgpRepository openPgpRepository;
    private final SchedulersFacade schedulers;

    @Inject
    public MercuryOpenPgpManager(OpenPgpRepository openPgpRepository,
                                 SchedulersFacade schedulers,
                                 Repositories repositories) {
        this.openPgpRepository = openPgpRepository;
        this.schedulers = schedulers;
        this.repositories = repositories;
    }

    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);
        SecretKeyBackupHelper.setBackupCodeGenerator(
                () -> "71ZA-Y416-UA7A-7NCE-3SNM-88EF"
        );
        try {
            if (!oxManager.hasSecretKeysAvailable()) {
                try {
                    oxManager.restoreSecretKeyServerBackup(
                            //() -> "RW8X-367S-A2C3-QYAL-VG6E-Z2IM");
                            //() -> "KISJ-5Z1T-FGDW-WMDK-SC2U-SQUA");
                            () -> "71ZA-Y416-UA7A-7NCE-3SNM-88EF");
                    PGPPublicKeyRingCollection keys = oxManager.getOpenPgpSelf().getAnyPublicKeys();
                    for (PGPPublicKeyRing key : keys) {
                        oxManager.getOpenPgpSelf().trust(new OpenPgpV4Fingerprint(key));
                    }
                    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);
                    oxManager.generateAndImportKeyPair(connection.getAccount().getJid());
                    oxManager.backupSecretKeyToServer(
                            backupCode -> LOGGER.log(Level.INFO, "OpenPGP Backup Code: " + backupCode),
                            availableSecretKeys -> availableSecretKeys);
                }
            }
            oxManager.announceSupportAndPublish();
            OXInstantMessagingManager oximManager = OXInstantMessagingManager.getInstanceFor(connection.getConnection());
            oximManager.addOxMessageListener(new MercuryMessageStore(connection.getAccount(),
                    repositories.getPeerRepository(), repositories.getDirectChatRepository(),
                    repositories.getMessageRepository(), schedulers));
            oximManager.announceSupportForOxInstantMessaging();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}