Mercury-IM/domain/src/main/java/org/jivesoftware/smackx/ikey/IkeyManager.java

package org.jivesoftware.smackx.ikey;

import org.apache.xml.security.c14n.CanonicalizationException;
import org.apache.xml.security.c14n.Canonicalizer;
import org.apache.xml.security.c14n.InvalidCanonicalizerException;
import org.apache.xml.security.parser.XMLParserException;
import org.bouncycastle.openpgp.PGPPublicKeyRing;
import org.jivesoftware.smack.Manager;
import org.jivesoftware.smack.SmackException;
import org.jivesoftware.smack.XMPPConnection;
import org.jivesoftware.smack.XMPPException;
import org.jivesoftware.smack.packet.Message;
import org.jivesoftware.smack.util.Async;
import org.jivesoftware.smackx.ikey.element.IkeyElement;
import org.jivesoftware.smackx.ikey.mechanism.IkeySignatureVerificationMechanism;
import org.jivesoftware.smackx.ikey.util.IkeyConstants;
import org.jivesoftware.smackx.ikey.util.UnsupportedSignatureAlgorithmException;
import org.jivesoftware.smackx.ikey.util.canonicalization.ElementCanonicalizer;
import org.jivesoftware.smackx.ikey.util.canonicalization.XmlSecElementCanonicalizer;
import org.jivesoftware.smackx.ikey_ox.OxIkeySignatureVerificationMechanism;
import org.jivesoftware.smackx.pep.PepEventListener;
import org.jivesoftware.smackx.pep.PepManager;
import org.jivesoftware.smackx.pubsub.LeafNode;
import org.jivesoftware.smackx.pubsub.PayloadItem;
import org.jivesoftware.smackx.pubsub.PubSubException;
import org.jivesoftware.smackx.pubsub.PubSubManager;
import org.jxmpp.jid.EntityBareJid;
import org.pgpainless.PGPainless;

import java.io.IOException;
import java.util.Date;
import java.util.List;
import java.util.Map;
import java.util.WeakHashMap;
import java.util.logging.Level;
import java.util.logging.Logger;

public final class IkeyManager extends Manager {

    private static final Logger LOGGER = Logger.getLogger(IkeyManager.class.getName());
    private static final Map<XMPPConnection, IkeyManager> INSTANCES = new WeakHashMap<>();

    private IkeyStore store;
    private final ElementCanonicalizer canonicalizer;

    private IkeyManager(XMPPConnection connection) {
        super(connection);
        try {
            this.canonicalizer = new XmlSecElementCanonicalizer(Canonicalizer.getInstance(Canonicalizer.ALGO_ID_C14N11_OMIT_COMMENTS));
        } catch (InvalidCanonicalizerException e) {
            throw new AssertionError(e);
        }
    }

    public static synchronized IkeyManager getInstanceFor(XMPPConnection connection) {
        IkeyManager manager = INSTANCES.get(connection);
        if (manager == null) {
            manager = new IkeyManager(connection);
            INSTANCES.put(connection, manager);
        }
        return manager;
    }

    public void startListeners() {
        PepManager.getInstanceFor(connection())
                .addPepEventListener(IkeyConstants.IKEY_NODE, IkeyElement.class, pepEventListener);
    }

    public void stopListeners() {
        PepManager.getInstanceFor(connection())
                .removePepEventListener(pepEventListener);
    }

    public void publishIkeyElement(IkeyElement ikeyElement)
            throws InterruptedException, PubSubException.NotALeafNodeException, XMPPException.XMPPErrorException,
            SmackException.NotConnectedException, SmackException.NoResponseException {
        PepManager.getInstanceFor(connection())
                .publish(IkeyConstants.IKEY_NODE, new PayloadItem<>(ikeyElement));
    }

    public IkeyElement fetchIkeyElementOf(EntityBareJid jid)
            throws InterruptedException, PubSubException.NotALeafNodeException, SmackException.NoResponseException,
            SmackException.NotConnectedException, XMPPException.XMPPErrorException, PubSubException.NotAPubSubNodeException {
        PubSubManager pubSubManager = PubSubManager.getInstanceFor(connection(), jid);
        return fetchIkeyElementFrom(pubSubManager);
    }

    public IkeyElement fetchOwnIkeyElement()
            throws InterruptedException, PubSubException.NotALeafNodeException, SmackException.NoResponseException,
            SmackException.NotConnectedException, XMPPException.XMPPErrorException, PubSubException.NotAPubSubNodeException {
        PubSubManager pubSubManager = PubSubManager.getInstanceFor(connection());
        return fetchIkeyElementFrom(pubSubManager);
    }

    private static IkeyElement fetchIkeyElementFrom(PubSubManager pubSubManager)
            throws PubSubException.NotALeafNodeException, SmackException.NoResponseException,
            SmackException.NotConnectedException, InterruptedException, XMPPException.XMPPErrorException,
            PubSubException.NotAPubSubNodeException {
        LeafNode node = pubSubManager.getLeafNode(IkeyConstants.IKEY_NODE);
        List<PayloadItem<IkeyElement>> items = node.getItems(1);
        if (items.isEmpty()) {
            return null;
        } else {
            return items.get(0).getPayload();
        }
    }

    private void processIkeyElement(EntityBareJid from, IkeyElement element)
            throws XMLParserException, IOException, CanonicalizationException, UnsupportedSignatureAlgorithmException {
        if (isFromTheFuture(element)) {
            LOGGER.log(Level.WARNING, "Received ikey element appears to be from the future: " + element.getSubordinates().getTimestamp());
            return;
        }

        if (existsSameOrNewerRecord(element)) {
            LOGGER.log(Level.WARNING, "There exists this exact, or a newer ikey record in the database for " + from);
            return;
        }

        if (!verifyIkeyElement(from, element)) {
            LOGGER.log(Level.WARNING, "Invalid signature on ikey element of " + from);
            return;
        }

        store.storeIkeyRecord(from, element);
    }

    private boolean verifyIkeyElement(EntityBareJid from, IkeyElement element)
            throws XMLParserException, IOException, CanonicalizationException, UnsupportedSignatureAlgorithmException {
        IkeySignatureVerificationMechanism verificationMechanism = getSignatureVerificationMechanismFor(element);
        IkeySignatureVerifier verifier = new IkeySignatureVerifier(verificationMechanism, canonicalizer);
        return verifier.verify(element, from);
    }

    private static IkeySignatureVerificationMechanism getSignatureVerificationMechanismFor(IkeyElement ikeyElement)
            throws IOException, UnsupportedSignatureAlgorithmException {
        switch (ikeyElement.getType()) {
            case OX:
                PGPPublicKeyRing ikey = PGPainless.readKeyRing().publicKeyRing(ikeyElement.getSuperordinate().getPubKeyBytes());
                return new OxIkeySignatureVerificationMechanism(ikey);

            case X509:
                throw new UnsupportedSignatureAlgorithmException("X.509");
        }
        throw new AssertionError("Unknown verification algorithm encountered: " + ikeyElement.getType());
    }

    private static boolean isFromTheFuture(IkeyElement element) {
        Date elementTimestamp = element.getSubordinates().getTimestamp();
        Date now = new Date();
        return elementTimestamp.after(now);
    }

    private boolean existsSameOrNewerRecord(IkeyElement ikeyElement) throws IOException {
        IkeyElement existingRecord = store.loadIkeyRecord(ikeyElement.getSubordinates().getJid());
        if (existingRecord == null) {
            return false;
        }
        Date latestTimestamp = existingRecord.getSubordinates().getTimestamp();
        Date eventTimestamp = ikeyElement.getSubordinates().getTimestamp();
        return latestTimestamp.equals(eventTimestamp) // same
                || latestTimestamp.after(eventTimestamp); // newer
    }

    @SuppressWarnings("UnnecessaryAnonymousClass")
    private final PepEventListener<IkeyElement> pepEventListener = new PepEventListener<IkeyElement>() {
        @Override
        public void onPepEvent(EntityBareJid from, IkeyElement event, String id, Message carrierMessage) {
            try {
                processIkeyElement(from, event);
            } catch (XMLParserException | CanonicalizationException | IOException | UnsupportedSignatureAlgorithmException e) {
                LOGGER.log(Level.WARNING, "Error:", e);
            }
        }
    };
}