001// SPDX-FileCopyrightText: 2021 Paul Schaub <vanitasvitae@fsfe.org>, 2021 Flowcrypt a.s.
002//
003// SPDX-License-Identifier: Apache-2.0
004
005package org.pgpainless.key.collection;
006
007import java.io.ByteArrayInputStream;
008import java.io.IOException;
009import java.io.InputStream;
010import java.util.ArrayList;
011import java.util.Collection;
012import java.util.List;
013import javax.annotation.Nonnull;
014
015import org.bouncycastle.openpgp.PGPException;
016import org.bouncycastle.openpgp.PGPKeyRing;
017import org.bouncycastle.openpgp.PGPMarker;
018import org.bouncycastle.openpgp.PGPObjectFactory;
019import org.bouncycastle.openpgp.PGPPublicKeyRing;
020import org.bouncycastle.openpgp.PGPPublicKeyRingCollection;
021import org.bouncycastle.openpgp.PGPSecretKeyRing;
022import org.bouncycastle.openpgp.PGPSecretKeyRingCollection;
023import org.pgpainless.implementation.ImplementationFactory;
024import org.pgpainless.util.ArmorUtils;
025
026/**
027 * This class describes a logic of handling a collection of different {@link PGPKeyRing}. The logic was inspired by
028 * {@link PGPSecretKeyRingCollection} and {@link PGPPublicKeyRingCollection}.
029 */
030public class PGPKeyRingCollection {
031
032    private final PGPSecretKeyRingCollection pgpSecretKeyRingCollection;
033    private final PGPPublicKeyRingCollection pgpPublicKeyRingCollection;
034
035    public PGPKeyRingCollection(@Nonnull byte[] encoding, boolean isSilent) throws IOException, PGPException {
036        this(new ByteArrayInputStream(encoding), isSilent);
037    }
038
039    /**
040     * Build a {@link PGPKeyRingCollection} from the passed in input stream.
041     *
042     * @param in       input stream containing data
043     * @param isSilent flag indicating that unsupported objects will be ignored
044     * @throws IOException  if a problem parsing the base stream occurs
045     * @throws PGPException if an object is encountered which isn't a {@link PGPSecretKeyRing} or {@link PGPPublicKeyRing}
046     */
047    public PGPKeyRingCollection(@Nonnull InputStream in, boolean isSilent) throws IOException, PGPException {
048        // Double getDecoderStream because of #96
049        InputStream decoderStream = ArmorUtils.getDecoderStream(in);
050        PGPObjectFactory pgpFact = ImplementationFactory.getInstance().getPGPObjectFactory(decoderStream);
051        Object obj;
052
053        List<PGPSecretKeyRing> secretKeyRings = new ArrayList<>();
054        List<PGPPublicKeyRing> publicKeyRings = new ArrayList<>();
055
056        while ((obj = pgpFact.nextObject()) != null) {
057            if (obj instanceof PGPMarker) {
058                // Skip marker packets
059                continue;
060            }
061            if (obj instanceof PGPSecretKeyRing) {
062                secretKeyRings.add((PGPSecretKeyRing) obj);
063            } else if (obj instanceof PGPPublicKeyRing) {
064                publicKeyRings.add((PGPPublicKeyRing) obj);
065            } else if (!isSilent) {
066                throw new PGPException(obj.getClass().getName() + " found where " +
067                        PGPSecretKeyRing.class.getSimpleName() + " or " +
068                        PGPPublicKeyRing.class.getSimpleName() + " expected");
069            }
070        }
071
072        pgpSecretKeyRingCollection = new PGPSecretKeyRingCollection(secretKeyRings);
073        pgpPublicKeyRingCollection = new PGPPublicKeyRingCollection(publicKeyRings);
074    }
075
076    public PGPKeyRingCollection(@Nonnull Collection<PGPKeyRing> collection, boolean isSilent)
077            throws IOException, PGPException {
078        List<PGPSecretKeyRing> secretKeyRings = new ArrayList<>();
079        List<PGPPublicKeyRing> publicKeyRings = new ArrayList<>();
080
081        for (PGPKeyRing pgpKeyRing : collection) {
082            if (pgpKeyRing instanceof PGPSecretKeyRing) {
083                secretKeyRings.add((PGPSecretKeyRing) pgpKeyRing);
084            } else if (pgpKeyRing instanceof PGPPublicKeyRing) {
085                publicKeyRings.add((PGPPublicKeyRing) pgpKeyRing);
086            } else if (!isSilent) {
087                throw new PGPException(pgpKeyRing.getClass().getName() + " found where " +
088                        PGPSecretKeyRing.class.getSimpleName() + " or " +
089                        PGPPublicKeyRing.class.getSimpleName() + " expected");
090            }
091        }
092
093        pgpSecretKeyRingCollection = new PGPSecretKeyRingCollection(secretKeyRings);
094        pgpPublicKeyRingCollection = new PGPPublicKeyRingCollection(publicKeyRings);
095    }
096
097    public @Nonnull PGPSecretKeyRingCollection getPGPSecretKeyRingCollection() {
098        return pgpSecretKeyRingCollection;
099    }
100
101    public @Nonnull PGPPublicKeyRingCollection getPgpPublicKeyRingCollection() {
102        return pgpPublicKeyRingCollection;
103    }
104
105    /**
106     * Return the number of rings in this collection.
107     *
108     * @return total size of {@link PGPSecretKeyRingCollection} and {@link PGPPublicKeyRingCollection}
109     * in this collection
110     */
111    public int size() {
112        return pgpSecretKeyRingCollection.size() + pgpPublicKeyRingCollection.size();
113    }
114}