Refactor EncryptionStream

pgpainless-core/src/main/java/org/pgpainless/encryption_signing/EncryptionStream.java

@@ -21,7 +21,6 @@ import java.io.OutputStream;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Date;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
 import java.util.logging.Level;
@@ -57,22 +56,26 @@ public final class EncryptionStream extends OutputStream {
 
     private static final int BUFFER_SIZE = 1 << 8;
 
+    private final SymmetricKeyAlgorithm symmetricKeyAlgorithm;
+    private final HashAlgorithm hashAlgorithm;
+    private final CompressionAlgorithm compressionAlgorithm;
+    private final Set<PGPPublicKey> encryptionKeys;
+    private final Set<PGPPrivateKey> signingKeys;
+    private final boolean asciiArmor;
+
     private final OpenPgpMetadata.Builder resultBuilder = OpenPgpMetadata.getBuilder();
 
     private List<PGPSignatureGenerator> signatureGenerators = new ArrayList<>();
     private boolean closed = false;
 
-    // ASCII Armor
+    OutputStream outermostStream = null;
-    private ArmoredOutputStream armorOutputStream = null;
 
-    // Public Key Encryption of Symmetric Session Key
+    private ArmoredOutputStream armorOutputStream = null;
     private OutputStream publicKeyEncryptedStream = null;
 
-    // Data Compression
     private PGPCompressedDataGenerator compressedDataGenerator;
     private BCPGOutputStream basicCompressionStream;
 
-    // Literal Data
     private PGPLiteralDataGenerator literalDataGenerator;
     private OutputStream literalDataStream;
 
@@ -85,26 +88,44 @@ public final class EncryptionStream extends OutputStream {
                      boolean asciiArmor)
             throws IOException, PGPException {
 
-        // Currently outermost Stream
+        this.symmetricKeyAlgorithm = symmetricKeyAlgorithm;
-        OutputStream outerMostStream;
+        this.hashAlgorithm = hashAlgorithm;
-        if (asciiArmor) {
+        this.compressionAlgorithm = compressionAlgorithm;
-            LOGGER.log(LEVEL, "Wrap encryption output in ASCII armor");
+        this.encryptionKeys = Collections.unmodifiableSet(encryptionKeys);
-            armorOutputStream = new ArmoredOutputStream(targetOutputStream);
+        this.signingKeys = Collections.unmodifiableSet(signingKeys);
-            outerMostStream = armorOutputStream;
+        this.asciiArmor = asciiArmor;
-        } else {
+
-            LOGGER.log(LEVEL, "Encryption output will be binary");
+        outermostStream = targetOutputStream;
-            outerMostStream = targetOutputStream;
+        prepareArmor();
+        prepareEncryption();
+        prepareSigning();
+        prepareCompression();
+        prepareOnePassSignatures();
+        prepareLiteralDataProcessing();
+        prepareResultBuilder();
+    }
+
+    private void prepareArmor() {
+        if (!asciiArmor) {
+            LOGGER.log(LEVEL, "Encryption output will be binary");
+            return;
+        }
+
+        LOGGER.log(LEVEL, "Wrap encryption output in ASCII armor");
+        armorOutputStream = new ArmoredOutputStream(outermostStream);
+        outermostStream = armorOutputStream;
+    }
+
+    private void prepareEncryption() throws IOException, PGPException {
+        if (encryptionKeys.isEmpty()) {
+            return;
         }
 
-        // If we want to encrypt
-        if (!encryptionKeys.isEmpty()) {
         LOGGER.log(LEVEL, "At least one encryption key is available -> encrypt using " + symmetricKeyAlgorithm);
         BcPGPDataEncryptorBuilder dataEncryptorBuilder =
                 new BcPGPDataEncryptorBuilder(symmetricKeyAlgorithm.getAlgorithmId());
 
-            LOGGER.log(LEVEL, "Integrity protection enabled");
         dataEncryptorBuilder.setWithIntegrityPacket(true);
-
         PGPEncryptedDataGenerator encryptedDataGenerator =
                 new PGPEncryptedDataGenerator(dataEncryptorBuilder);
 
@@ -113,12 +134,15 @@ public final class EncryptionStream extends OutputStream {
             encryptedDataGenerator.addMethod(new BcPublicKeyKeyEncryptionMethodGenerator(key));
         }
 
-            publicKeyEncryptedStream = encryptedDataGenerator.open(outerMostStream, new byte[BUFFER_SIZE]);
+        publicKeyEncryptedStream = encryptedDataGenerator.open(outermostStream, new byte[BUFFER_SIZE]);
-            outerMostStream = publicKeyEncryptedStream;
+        outermostStream = publicKeyEncryptedStream;
+    }
+
+    private void prepareSigning() throws PGPException {
+        if (signingKeys.isEmpty()) {
+            return;
         }
 
-        // If we want to sign, prepare for signing
-        if (!signingKeys.isEmpty()) {
         LOGGER.log(LEVEL, "At least one signing key is available -> sign " + hashAlgorithm + " hash of message");
         for (PGPPrivateKey privateKey : signingKeys) {
             LOGGER.log(LEVEL, "Sign using key " + Long.toHexString(privateKey.getKeyID()));
@@ -131,22 +155,26 @@ public final class EncryptionStream extends OutputStream {
         }
     }
 
+    private void prepareCompression() throws IOException {
         LOGGER.log(LEVEL, "Compress using " + compressionAlgorithm);
-        // Compression
         compressedDataGenerator = new PGPCompressedDataGenerator(
                 compressionAlgorithm.getAlgorithmId());
-        basicCompressionStream = new BCPGOutputStream(compressedDataGenerator.open(outerMostStream));
+        basicCompressionStream = new BCPGOutputStream(compressedDataGenerator.open(outermostStream));
+    }
 
-        // If we want to sign, sign!
+    private void prepareOnePassSignatures() throws IOException, PGPException {
         for (PGPSignatureGenerator signatureGenerator : signatureGenerators) {
             signatureGenerator.generateOnePassVersion(false).encode(basicCompressionStream);
         }
+    }
 
+    private void prepareLiteralDataProcessing() throws IOException {
         literalDataGenerator = new PGPLiteralDataGenerator();
         literalDataStream = literalDataGenerator.open(basicCompressionStream,
                 PGPLiteralData.BINARY, PGPLiteralData.CONSOLE, new Date(), new byte[BUFFER_SIZE]);
+    }
 
-        // Prepare result
+    private void prepareResultBuilder() {
         for (PGPPublicKey recipient : encryptionKeys) {
             resultBuilder.addRecipientKeyId(recipient.getKeyID());
         }
@@ -185,24 +213,16 @@ public final class EncryptionStream extends OutputStream {
 
     @Override
     public void close() throws IOException {
-        if (!closed) {
+        if (closed) {
+            return;
+        }
 
         // Literal Data
         literalDataStream.flush();
         literalDataStream.close();
         literalDataGenerator.close();
 
-            // Signing
+        writeSignatures();
-            for (PGPSignatureGenerator signatureGenerator : signatureGenerators) {
-                try {
-                    PGPSignature signature = signatureGenerator.generate();
-                    signature.encode(basicCompressionStream);
-                    resultBuilder.addSignature(signature);
-                    resultBuilder.addUnverifiedSignatureKeyId(signature.getKeyID());
-                } catch (PGPException e) {
-                    throw new IOException(e);
-                }
-            }
 
         // Compressed Data
         compressedDataGenerator.close();
@@ -220,6 +240,18 @@ public final class EncryptionStream extends OutputStream {
         }
         closed = true;
     }
+
+    private void writeSignatures() throws IOException {
+        for (PGPSignatureGenerator signatureGenerator : signatureGenerators) {
+            try {
+                PGPSignature signature = signatureGenerator.generate();
+                signature.encode(basicCompressionStream);
+                resultBuilder.addSignature(signature);
+                resultBuilder.addUnverifiedSignatureKeyId(signature.getKeyID());
+            } catch (PGPException e) {
+                throw new IOException(e);
+            }
+        }
     }
 
     public OpenPgpMetadata getResult() {