From 8c97b6ead1601d4765690177ee81736e526508ef Mon Sep 17 00:00:00 2001
From: Paul Schaub <vanitasvitae@fsfe.org>
Date: Thu, 18 Mar 2021 21:28:08 +0100
Subject: [PATCH] In PasswordBasedSecretKeyRingProtector.forKey(ring,
 passphrase): Return passphrase also for subkeys

Fixes #97, thanks @DenBond7
---
 .../PasswordBasedSecretKeyRingProtector.java    |  9 +++++++--
 .../protection/PassphraseProtectedKeyTest.java  | 17 +++++++++++++++++
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/pgpainless-core/src/main/java/org/pgpainless/key/protection/PasswordBasedSecretKeyRingProtector.java b/pgpainless-core/src/main/java/org/pgpainless/key/protection/PasswordBasedSecretKeyRingProtector.java
index 0bec8b00..e5678fb5 100644
--- a/pgpainless-core/src/main/java/org/pgpainless/key/protection/PasswordBasedSecretKeyRingProtector.java
+++ b/pgpainless-core/src/main/java/org/pgpainless/key/protection/PasswordBasedSecretKeyRingProtector.java
@@ -15,11 +15,13 @@
  */
 package org.pgpainless.key.protection;
 
+import java.util.Iterator;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 
 import org.bouncycastle.openpgp.PGPException;
 import org.bouncycastle.openpgp.PGPKeyRing;
+import org.bouncycastle.openpgp.PGPPublicKey;
 import org.bouncycastle.openpgp.PGPSecretKey;
 import org.bouncycastle.openpgp.operator.PBESecretKeyDecryptor;
 import org.bouncycastle.openpgp.operator.PBESecretKeyEncryptor;
@@ -55,8 +57,11 @@ public class PasswordBasedSecretKeyRingProtector implements SecretKeyRingProtect
             @Override
             @Nullable
             public Passphrase getPassphraseFor(Long keyId) {
-                if (keyRing.getPublicKey().getKeyID() == keyId) {
-                    return passphrase;
+                for (Iterator<PGPPublicKey> it = keyRing.getPublicKeys(); it.hasNext(); ) {
+                    PGPPublicKey key = it.next();
+                    if (key.getKeyID() == keyId) {
+                        return passphrase;
+                    }
                 }
                 return null;
             }
diff --git a/pgpainless-core/src/test/java/org/pgpainless/key/protection/PassphraseProtectedKeyTest.java b/pgpainless-core/src/test/java/org/pgpainless/key/protection/PassphraseProtectedKeyTest.java
index 437ef12f..8dd23bf9 100644
--- a/pgpainless-core/src/test/java/org/pgpainless/key/protection/PassphraseProtectedKeyTest.java
+++ b/pgpainless-core/src/test/java/org/pgpainless/key/protection/PassphraseProtectedKeyTest.java
@@ -18,10 +18,16 @@ package org.pgpainless.key.protection;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
 
+import java.security.InvalidAlgorithmParameterException;
+import java.security.NoSuchAlgorithmException;
+import java.util.Iterator;
 import javax.annotation.Nullable;
 
 import org.bouncycastle.openpgp.PGPException;
+import org.bouncycastle.openpgp.PGPPublicKey;
+import org.bouncycastle.openpgp.PGPSecretKeyRing;
 import org.junit.jupiter.api.Test;
+import org.pgpainless.PGPainless;
 import org.pgpainless.key.TestKeys;
 import org.pgpainless.key.protection.passphrase_provider.SecretKeyPassphraseProvider;
 import org.pgpainless.util.Passphrase;
@@ -56,4 +62,15 @@ public class PassphraseProtectedKeyTest {
         assertNull(protector.getEncryptor(TestKeys.JULIET_KEY_ID));
         assertNull(protector.getDecryptor(TestKeys.JULIET_KEY_ID));
     }
+
+    @Test
+    public void testReturnsNonNullDecryptorForSubkeys() throws InvalidAlgorithmParameterException, NoSuchAlgorithmException, PGPException {
+        PGPSecretKeyRing secretKeys = PGPainless.generateKeyRing().modernKeyRing("alice", "passphrase");
+        SecretKeyRingProtector protector = PasswordBasedSecretKeyRingProtector.forKey(secretKeys, Passphrase.fromPassword("passphrase"));
+        for (Iterator<PGPPublicKey> it = secretKeys.getPublicKeys(); it.hasNext(); ) {
+            PGPPublicKey subkey = it.next();
+            assertNotNull(protector.getEncryptor(subkey.getKeyID()));
+            assertNotNull(protector.getDecryptor(subkey.getKeyID()));
+        }
+    }
 }