Fix tests

pgpainless-core/src/main/java/org/pgpainless/decryption_verification/OpenPgpMessageInputStream.java

@@ -1,6 +1,5 @@
 package org.pgpainless.decryption_verification;
 
-import com.sun.tools.javac.code.Attribute;
 import org.bouncycastle.bcpg.BCPGInputStream;
 import org.bouncycastle.bcpg.BCPGOutputStream;
 import org.bouncycastle.bcpg.OnePassSignaturePacket;
@@ -25,13 +24,14 @@ import org.bouncycastle.openpgp.PGPSignature;
 import org.bouncycastle.openpgp.PGPSignatureList;
 import org.bouncycastle.openpgp.operator.PBEDataDecryptorFactory;
 import org.bouncycastle.openpgp.operator.PublicKeyDataDecryptorFactory;
+import org.bouncycastle.openpgp.operator.SessionKeyDataDecryptorFactory;
 import org.pgpainless.PGPainless;
 import org.pgpainless.algorithm.EncryptionPurpose;
 import org.pgpainless.algorithm.OpenPgpPacket;
 import org.pgpainless.decryption_verification.automaton.InputAlphabet;
 import org.pgpainless.decryption_verification.automaton.PDA;
-import org.pgpainless.exception.MalformedOpenPgpMessageException;
 import org.pgpainless.exception.MessageNotIntegrityProtectedException;
+import org.pgpainless.exception.MissingDecryptionMethodException;
 import org.pgpainless.implementation.ImplementationFactory;
 import org.pgpainless.key.info.KeyRingInfo;
 import org.pgpainless.key.protection.SecretKeyRingProtector;
@@ -114,7 +114,27 @@ public class OpenPgpMessageInputStream extends InputStream {
 
                     SortedESKs esks = new SortedESKs(encDataList);
 
-                    // TODO: try session keys
+                    if (options.getSessionKey() != null) {
+                        SessionKeyDataDecryptorFactory decryptorFactory = ImplementationFactory.getInstance()
+                                .getSessionKeyDataDecryptorFactory(options.getSessionKey());
+                        // TODO: Replace with encDataList.addSessionKeyDecryptionMethod(sessionKey)
+                        PGPEncryptedData esk = esks.all().get(0);
+                        try {
+                            if (esk instanceof PGPPBEEncryptedData) {
+                                PGPPBEEncryptedData skesk = (PGPPBEEncryptedData) esk;
+                                in = skesk.getDataStream(decryptorFactory);
+                                break loop;
+                            } else if (esk instanceof PGPPublicKeyEncryptedData) {
+                                PGPPublicKeyEncryptedData pkesk = (PGPPublicKeyEncryptedData) esk;
+                                in = pkesk.getDataStream(decryptorFactory);
+                                break loop;
+                            } else {
+                                throw new RuntimeException("Unknown ESK class type: " + esk.getClass().getName());
+                            }
+                        } catch (PGPException e) {
+                            // Session key mismatch?
+                        }
+                    }
 
                     // Try passwords
                     for (PGPPBEEncryptedData skesk : esks.skesks) {
@@ -174,7 +194,7 @@ public class OpenPgpMessageInputStream extends InputStream {
 
                     // TODO: try interactive password callbacks
 
-                    break loop;
+                    throw new MissingDecryptionMethodException("No working decryption method found.");
 
                 case MARKER:
                     bcpgIn.readPacket(); // skip marker packet
@@ -256,6 +276,7 @@ public class OpenPgpMessageInputStream extends InputStream {
             if (tag == PacketTags.SIGNATURE) {
                 SignaturePacket sigPacket = (SignaturePacket) packet;
                 sigPacket.encode(bcpgOut);
+                tag = bcpgIn.nextPacketTag();
             }
         }
         bcpgOut.close();
@@ -270,16 +291,21 @@ public class OpenPgpMessageInputStream extends InputStream {
     @Override
     public int read() throws IOException {
         int r = -1;
+        if (in != null) {
             try {
                 r = in.read();
             } catch (IOException e) {
                 //
             }
+        }
         if (r == -1) {
             if (in instanceof OpenPgpMessageInputStream) {
+                System.out.println("Read -1: close " + automaton);
                 in.close();
+                in = null;
             } else {
                 try {
+                    System.out.println("Walk " + automaton);
                     walk();
                 } catch (PGPException e) {
                     throw new RuntimeException(e);
@@ -291,8 +317,15 @@ public class OpenPgpMessageInputStream extends InputStream {
 
     @Override
     public void close() throws IOException {
+        if (in == null) {
+            System.out.println("Close " + automaton);
+            automaton.next(InputAlphabet.EndOfSequence);
+            automaton.assertValid();
+            return;
+        }
         try {
             in.close();
+            in = null;
             // Nested streams (except LiteralData) need to be closed.
             if (automaton.getState() != PDA.State.LiteralMessage) {
                 automaton.next(InputAlphabet.EndOfSequence);
@@ -327,5 +360,13 @@ public class OpenPgpMessageInputStream extends InputStream {
                 }
             }
         }
+
+        public List<PGPEncryptedData> all() {
+            List<PGPEncryptedData> esks = new ArrayList<>();
+            esks.addAll(skesks);
+            esks.addAll(pkesks);
+            esks.addAll(anonPkesks);
+            return esks;
+        }
     }
 }

pgpainless-core/src/main/java/org/pgpainless/decryption_verification/automaton/PDA.java

@@ -9,6 +9,9 @@ import static org.pgpainless.decryption_verification.automaton.StackAlphabet.ops
 import static org.pgpainless.decryption_verification.automaton.StackAlphabet.terminus;
 
 public class PDA {
+
+    private static int ID = 0;
+
     /**
      * Set of states of the automaton.
      * Each state defines its valid transitions in their {@link NestingPDA.State#transition(InputAlphabet, NestingPDA)}
@@ -174,18 +177,20 @@ public class PDA {
 
     private final Stack<StackAlphabet> stack = new Stack<>();
     private State state;
+    private int id;
 
     public PDA() {
         state = State.OpenPgpMessage;
         stack.push(terminus);
         stack.push(msg);
+        this.id = ID++;
     }
 
     public void next(InputAlphabet input) throws MalformedOpenPgpMessageException {
         State old = state;
         StackAlphabet stackItem = stack.isEmpty() ? null : stack.peek();
         state = state.transition(input, this);
-        System.out.println("Transition from " + old + " to " + state + " via " + input + " with stack " + stackItem);
+        System.out.println(id + ": Transition from " + old + " to " + state + " via " + input + " with stack " + stackItem);
     }
 
     /**
@@ -232,6 +237,6 @@ public class PDA {
 
     @Override
     public String toString() {
-        return "State: " + state + " Stack: " + stack;
+        return "PDA " + id + ": State: " + state + " Stack: " + stack;
     }
 }

pgpainless-core/src/test/java/org/pgpainless/decryption_verification/OpenPgpMessageInputStreamTest.java

@@ -2,8 +2,10 @@ package org.pgpainless.decryption_verification;
 
 import org.bouncycastle.bcpg.ArmoredInputStream;
 import org.bouncycastle.openpgp.PGPException;
+import org.bouncycastle.openpgp.PGPSecretKeyRing;
 import org.bouncycastle.util.io.Streams;
 import org.junit.jupiter.api.Test;
+import org.pgpainless.PGPainless;
 import org.pgpainless.exception.MalformedOpenPgpMessageException;
 import org.pgpainless.util.ArmoredInputStreamFactory;
 import org.pgpainless.util.Passphrase;
@@ -18,9 +20,11 @@ import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.pgpainless.decryption_verification.PGPDecryptionStreamTest.COMP;
 import static org.pgpainless.decryption_verification.PGPDecryptionStreamTest.COMP_COMP_LIT;
 import static org.pgpainless.decryption_verification.PGPDecryptionStreamTest.COMP_LIT;
+import static org.pgpainless.decryption_verification.PGPDecryptionStreamTest.KEY;
 import static org.pgpainless.decryption_verification.PGPDecryptionStreamTest.LIT;
 import static org.pgpainless.decryption_verification.PGPDecryptionStreamTest.LIT_LIT;
 import static org.pgpainless.decryption_verification.PGPDecryptionStreamTest.PASSPHRASE;
+import static org.pgpainless.decryption_verification.PGPDecryptionStreamTest.PENC_COMP_LIT;
 import static org.pgpainless.decryption_verification.PGPDecryptionStreamTest.PLAINTEXT;
 import static org.pgpainless.decryption_verification.PGPDecryptionStreamTest.SENC_LIT;
 import static org.pgpainless.decryption_verification.PGPDecryptionStreamTest.SIG_LIT;
@@ -69,6 +73,14 @@ public class OpenPgpMessageInputStreamTest {
         assertEquals(PLAINTEXT, plain);
     }
 
+    @Test
+    public void testProcessPENC_COMP_LIT() throws IOException, PGPException {
+        PGPSecretKeyRing secretKeys = PGPainless.readKeyRing().secretKeyRing(KEY);
+        String plain = process(PENC_COMP_LIT, ConsumerOptions.get()
+                .addDecryptionKey(secretKeys));
+        assertEquals(PLAINTEXT, plain);
+    }
+
     private String process(String armoredMessage, ConsumerOptions options) throws PGPException, IOException {
         OpenPgpMessageInputStream in = get(armoredMessage, options);
         ByteArrayOutputStream out = new ByteArrayOutputStream();

pgpainless-core/src/test/java/org/pgpainless/decryption_verification/PGPDecryptionStreamTest.java

@@ -7,6 +7,7 @@ import org.bouncycastle.openpgp.PGPCompressedDataGenerator;
 import org.bouncycastle.openpgp.PGPException;
 import org.bouncycastle.openpgp.PGPLiteralData;
 import org.bouncycastle.openpgp.PGPLiteralDataGenerator;
+import org.bouncycastle.openpgp.PGPPublicKeyRing;
 import org.bouncycastle.openpgp.PGPSecretKeyRing;
 import org.bouncycastle.openpgp.PGPSignature;
 import org.bouncycastle.util.io.Streams;
@@ -124,6 +125,17 @@ public class PGPDecryptionStreamTest {
             "=i4Y0\n" +
             "-----END PGP MESSAGE-----";
 
+    public static final String PENC_COMP_LIT = "" +
+            "-----BEGIN PGP MESSAGE-----\n" +
+            "Version: PGPainless\n" +
+            "\n" +
+            "hF4Dyqa/GWUy6WsSAQdAQ62BwmUt8Iby0+jvrLhMgST79KR/as+dyl0nf1uki2sw\n" +
+            "Thg1Ojtf0hOyJgcpQ4nP2Q0wYFR0F1sCydaIlTGreYZHlGtybP7/Ml6KNZILTRWP\n" +
+            "0kYBkGBgK7oQWRIVyoF2POvEP6EX1X8nvQk7O3NysVdRVbnia7gE3AzRYuha4kxs\n" +
+            "pI6xJkntLMS3K6him1Y9FHINIASFSB+C\n" +
+            "=5p00\n" +
+            "-----END PGP MESSAGE-----";
+
     @Test
     public void genLIT() throws IOException {
         ArmoredOutputStream armorOut = new ArmoredOutputStream(System.out);
@@ -328,4 +340,22 @@ public class PGPDecryptionStreamTest {
 
         System.out.println(out);
     }
+
+    @Test
+    public void genPENC_COMP_LIT() throws IOException, PGPException {
+        PGPSecretKeyRing secretKeys = PGPainless.readKeyRing().secretKeyRing(KEY);
+        PGPPublicKeyRing cert = PGPainless.extractCertificate(secretKeys);
+
+        ByteArrayOutputStream out = new ByteArrayOutputStream();
+        EncryptionStream enc = PGPainless.encryptAndOrSign()
+                .onOutputStream(out)
+                .withOptions(ProducerOptions.encrypt(EncryptionOptions.get()
+                                .addRecipient(cert))
+                        .overrideCompressionAlgorithm(CompressionAlgorithm.ZLIB));
+
+        Streams.pipeAll(new ByteArrayInputStream(PLAINTEXT.getBytes(StandardCharsets.UTF_8)), enc);
+        enc.close();
+
+        System.out.println(out);
+    }
 }