001// SPDX-FileCopyrightText: 2021 Paul Schaub <vanitasvitae@fsfe.org>
002//
003// SPDX-License-Identifier: Apache-2.0
004
005package org.pgpainless.util;
006
007import java.io.IOException;
008import java.io.OutputStream;
009import java.util.Date;
010
011import javax.annotation.Nonnull;
012
013import org.bouncycastle.openpgp.PGPCanonicalizedDataGenerator;
014import org.bouncycastle.openpgp.PGPLiteralDataGenerator;
015import org.pgpainless.algorithm.StreamEncoding;
016
017/**
018 * Literal Data can be encoded in different ways.
019 * BINARY encoding leaves the data as is and is generated through the {@link PGPLiteralDataGenerator}.
020 * However, if the data is encoded in TEXT or UTF8 encoding, we need to use the {@link PGPCanonicalizedDataGenerator}
021 * instead.
022 *
023 * This wrapper class acts as a handle for both options and provides a unified interface for them.
024 */
025public final class StreamGeneratorWrapper {
026
027    private final StreamEncoding encoding;
028    private final PGPLiteralDataGenerator literalDataGenerator;
029    private final PGPCanonicalizedDataGenerator canonicalizedDataGenerator;
030
031    /**
032     * Create a new instance for the given encoding.
033     *
034     * @param encoding stream encoding
035     * @return wrapper
036     */
037    public static StreamGeneratorWrapper forStreamEncoding(@Nonnull StreamEncoding encoding) {
038        if (encoding == StreamEncoding.BINARY) {
039            return new StreamGeneratorWrapper(encoding, new PGPLiteralDataGenerator());
040        } else {
041            return new StreamGeneratorWrapper(encoding, new PGPCanonicalizedDataGenerator());
042        }
043    }
044
045    private StreamGeneratorWrapper(@Nonnull StreamEncoding encoding, @Nonnull PGPLiteralDataGenerator literalDataGenerator) {
046        if (encoding != StreamEncoding.BINARY) {
047            throw new IllegalArgumentException("PGPLiteralDataGenerator can only be used with BINARY encoding.");
048        }
049        this.encoding = encoding;
050        this.literalDataGenerator = literalDataGenerator;
051        this.canonicalizedDataGenerator = null;
052    }
053
054    private StreamGeneratorWrapper(@Nonnull StreamEncoding encoding, @Nonnull PGPCanonicalizedDataGenerator canonicalizedDataGenerator) {
055        if (encoding != StreamEncoding.TEXT && encoding != StreamEncoding.UTF8) {
056            throw new IllegalArgumentException("PGPCanonicalizedDataGenerator can only be used with TEXT or UTF8 encoding.");
057        }
058        this.encoding = encoding;
059        this.canonicalizedDataGenerator = canonicalizedDataGenerator;
060        this.literalDataGenerator = null;
061    }
062
063    /**
064     * Open a new encoding stream.
065     *
066     * @param outputStream wrapped output stream
067     * @param filename file name
068     * @param modificationDate modification date
069     * @param buffer buffer
070     * @return encoding stream
071     */
072    public OutputStream open(OutputStream outputStream, String filename, Date modificationDate, byte[] buffer) throws IOException {
073        if (literalDataGenerator != null) {
074            return literalDataGenerator.open(outputStream, encoding.getCode(), filename, modificationDate, buffer);
075        } else {
076            return canonicalizedDataGenerator.open(outputStream, encoding.getCode(), filename, modificationDate, buffer);
077        }
078    }
079
080    /**
081     * Close all encoding streams opened by this generator wrapper.
082     */
083    public void close() throws IOException {
084        if (literalDataGenerator != null) {
085            literalDataGenerator.close();
086        }
087        if (canonicalizedDataGenerator != null) {
088            canonicalizedDataGenerator.close();
089        }
090    }
091}