Mercury-IM/data/src/main/java/org/mercury_im/messenger/data/repository/RxOpenPgpRepository.java

package org.mercury_im.messenger.data.repository;

import org.bouncycastle.openpgp.PGPPublicKeyRing;
import org.bouncycastle.openpgp.PGPPublicKeyRingCollection;
import org.bouncycastle.openpgp.PGPSecretKeyRingCollection;
import org.jivesoftware.smackx.ox.store.definition.OpenPgpTrustStore;
import org.jxmpp.jid.EntityBareJid;
import org.jxmpp.jid.impl.JidCreate;
import org.mercury_im.messenger.core.data.repository.OpenPgpRepository;
import org.mercury_im.messenger.data.model.AnnouncedOpenPgpContactKey;
import org.mercury_im.messenger.data.model.OpenPgpKeyFetchDate;
import org.mercury_im.messenger.data.model.OpenPgpKeyTrust;
import org.mercury_im.messenger.data.model.OpenPgpPublicKeyRing;
import org.mercury_im.messenger.data.model.OpenPgpSecretKeyRing;
import org.pgpainless.PGPainless;
import org.pgpainless.key.OpenPgpV4Fingerprint;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Date;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.logging.Level;
import java.util.logging.Logger;

import javax.inject.Inject;

import io.reactivex.Completable;
import io.reactivex.Observable;
import io.reactivex.Single;
import io.requery.Persistable;
import io.requery.query.ResultDelegate;
import io.requery.reactivex.ReactiveEntityStore;

public class RxOpenPgpRepository implements OpenPgpRepository {

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

    private final ReactiveEntityStore<Persistable> data;

    @Inject
    public RxOpenPgpRepository(ReactiveEntityStore<Persistable> data) {
        this.data = data;
    }

    @Override
    public Completable storePublicKeysOf(UUID accountId, EntityBareJid owner, PGPPublicKeyRingCollection keys) {
        return Single.fromCallable(() -> {
            OpenPgpPublicKeyRing keyRing = new OpenPgpPublicKeyRing();
            keyRing.setAccountId(accountId);
            keyRing.setOwner(owner);
            keyRing.setBytes(keys.getEncoded());
            return keyRing;
        }).flatMap(data::upsert).ignoreElement()
                .doOnComplete(() -> LOGGER.log(Level.INFO, "Successfully stored public keys of " + owner + " for account " + accountId));
    }

    @Override
    public Single<PGPPublicKeyRingCollection> loadPublicKeysOf(UUID accountId, EntityBareJid owner) {
        return data.select(OpenPgpPublicKeyRing.class)
                .where(OpenPgpPublicKeyRing.ACCOUNT_ID.eq(accountId))
                .and(OpenPgpPublicKeyRing.OWNER.eq(owner))
                .limit(1).get()
                .maybe().toSingle()
                .map(keyring -> PGPainless.readKeyRing().publicKeyRingCollection(keyring.getBytes()));
    }

    @Override
    public Single<Integer> deletePublicKeysOf(UUID accountId, EntityBareJid owner) {
        return data.delete(OpenPgpPublicKeyRing.class)
                .where(OpenPgpPublicKeyRing.ACCOUNT_ID.eq(accountId))
                .and(OpenPgpPublicKeyRing.OWNER.eq(owner))
                .get().single()
                .doOnSuccess(count -> LOGGER.log(Level.INFO, "Successfully deleted " + count + " public keys of " + owner + " for account " + accountId));
    }

    @Override
    public Completable storeSecretKeysOf(UUID accountId, EntityBareJid owner, PGPSecretKeyRingCollection keys) {
        return Single.fromCallable(() -> {
            OpenPgpSecretKeyRing keyRing = new OpenPgpSecretKeyRing();
            keyRing.setAccountId(accountId);
            keyRing.setOwner(owner);
            keyRing.setBytes(keys.getEncoded());
            return keyRing;
        }).flatMap(data::upsert).ignoreElement()
                .doOnComplete(() -> LOGGER.log(Level.INFO, "Successfully stored secret keys of " + owner + " for account " + accountId));
    }

    @Override
    public Single<PGPSecretKeyRingCollection> loadSecretKeysOf(UUID accountId, EntityBareJid owner) {
        return data.select(OpenPgpSecretKeyRing.class)
                .where(OpenPgpSecretKeyRing.ACCOUNT_ID.eq(accountId))
                .and(OpenPgpSecretKeyRing.OWNER.eq(owner))
                .limit(1).get()
                .maybe().toSingle()
                .map(keyring -> PGPainless.readKeyRing().secretKeyRingCollection(keyring.getBytes()));
    }

    @Override
    public Single<Integer> deleteSecretKeysOf(UUID accountId, EntityBareJid owner) {
        return data.delete(OpenPgpSecretKeyRing.class)
                .where(OpenPgpSecretKeyRing.ACCOUNT_ID.eq(accountId))
                .and(OpenPgpSecretKeyRing.OWNER.eq(owner))
                .get().single()
                .doOnSuccess(count -> LOGGER.log(Level.INFO, "Successfully deleted " + count + " secret keys of " + owner + " for account " + accountId));
    }

    @Override
    public Completable storePublicKeyFetchDates(UUID accountId, EntityBareJid owner, Map<OpenPgpV4Fingerprint, Date> dates) {
        List<OpenPgpKeyFetchDate> entities = new LinkedList<>();
        for (Map.Entry<OpenPgpV4Fingerprint, Date> entry : dates.entrySet()) {
            OpenPgpKeyFetchDate entity = new OpenPgpKeyFetchDate();
            entity.setAccountId(accountId);
            entity.setOwner(owner);
            entity.setFingerprint(entry.getKey());
            entity.setFetchDate(entry.getValue());
            entities.add(entity);
        }
        return data.upsert(entities).ignoreElement();
    }

    @Override
    public Completable storeAnnouncedFingerprints(UUID accountId, EntityBareJid owner, Map<OpenPgpV4Fingerprint, Date> metadata) {
        List<AnnouncedOpenPgpContactKey> entities = new LinkedList<>();
        for (Map.Entry<OpenPgpV4Fingerprint, Date> entry : metadata.entrySet()) {
            AnnouncedOpenPgpContactKey entity = new AnnouncedOpenPgpContactKey();
            entity.setAccountId(accountId);
            entity.setOwner(owner);
            entity.setFingerprint(entry.getKey());
            entity.setModificationDate(entry.getValue());
            entities.add(entity);
        }
        return data.upsert(entities).ignoreElement()
                .doOnComplete(() -> LOGGER.log(Level.INFO, "Successfully stored announced fingerprints of " +
                        owner + " for account " + accountId + ": " + Arrays.toString(metadata.keySet().toArray())));
    }

    @Override
    public Single<Map<OpenPgpV4Fingerprint, Date>> loadAnnouncedFingerprints(UUID accountId, EntityBareJid contact) {
        return data.select(AnnouncedOpenPgpContactKey.class)
                .where(AnnouncedOpenPgpContactKey.ACCOUNT_ID.eq(accountId))
                .and(AnnouncedOpenPgpContactKey.OWNER.eq(contact))
                .get()
                .observableResult()
                .map(ResultDelegate::toList)
                .single(Collections.emptyList())
                .map(list -> {
                    Map<OpenPgpV4Fingerprint, Date> map = new ConcurrentHashMap<>();
                    for (AnnouncedOpenPgpContactKey key : list) {
                        map.put(key.getFingerprint(), key.getModificationDate());
                    }
                    return map;
                });
    }

    @Override
    public Single<OpenPgpTrustStore.Trust> loadTrust(UUID accountId, EntityBareJid owner, OpenPgpV4Fingerprint fingerprint) {
        return data.select(OpenPgpKeyTrust.class)
                .where(OpenPgpKeyTrust.ACCOUNT_ID.eq(accountId))
                .and(OpenPgpKeyTrust.OWNER.eq(owner))
                .and(OpenPgpKeyTrust.FINGERPRINT.eq(fingerprint))
                .limit(1)
                .get()
                .maybe()
                .toSingle(new OpenPgpKeyTrust())
                .map(entity -> entity.getTrust() != null ? entity.getTrust() : OpenPgpTrustStore.Trust.undecided);
    }

    @Override
    public Completable storeTrust(UUID accountId, EntityBareJid owner, OpenPgpV4Fingerprint fingerprint, OpenPgpTrustStore.Trust trust) {
        OpenPgpKeyTrust entity = new OpenPgpKeyTrust();
        entity.setAccountId(accountId);
        entity.setOwner(owner);
        entity.setFingerprint(fingerprint);
        entity.setTrust(trust);
        return data.upsert(entity).ignoreElement();
    }

    @Override
    public Observable<List<OpenPgpV4Fingerprint>> observeFingerprintsOf(UUID accountId, String peerAddress) {
        return data.select(OpenPgpPublicKeyRing.class)
                .where(OpenPgpPublicKeyRing.ACCOUNT_ID.eq(accountId))
                .and(OpenPgpPublicKeyRing.OWNER.eq(JidCreate.entityBareFromOrThrowUnchecked(peerAddress)))
                .get().observableResult()
                .map(result -> {
                    OpenPgpPublicKeyRing model = new ResultDelegate<>(result).firstOrNull();
                    if (model == null) {
                        return Collections.emptyList();
                    }
                    List<OpenPgpV4Fingerprint> fingerprints = new ArrayList<>();
                    PGPPublicKeyRingCollection keys = PGPainless.readKeyRing().publicKeyRingCollection(model.getBytes());
                    for (PGPPublicKeyRing key : keys) {
                        fingerprints.add(new OpenPgpV4Fingerprint(key));
                    }
                    return fingerprints;
                });
    }

    @Override
    public Single<Map<OpenPgpV4Fingerprint, Date>> loadPublicKeyFetchDates(UUID accountId, EntityBareJid owner) {
        List<OpenPgpKeyFetchDate> fetchDates =
                data.select(OpenPgpKeyFetchDate.class)
                        .where(OpenPgpKeyFetchDate.ACCOUNT_ID.eq(accountId))
                        .and(OpenPgpKeyFetchDate.OWNER.eq(owner))
                        .get()
                        .toList();

        Map<OpenPgpV4Fingerprint, Date> map = new ConcurrentHashMap<>();
        for (OpenPgpKeyFetchDate date : fetchDates) {
            map.put(date.getFingerprint(), date.getFetchDate());
        }
        return Single.just(map);
    }
}