Implement KeyRingInfo.getKeysWithFlag() and KeyRingInfo.getExpirationDateForUse()

pgpainless-core/src/main/java/org/pgpainless/key/info/KeyRingInfo.java

@@ -25,6 +25,7 @@ import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.NoSuchElementException;
 import java.util.Set;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
@@ -587,24 +588,16 @@ public class KeyRingInfo {
      * @return expiration date
      */
     public @Nullable Date getPrimaryKeyExpirationDate() {
-        Date lastExpiration = null;
+        PGPSignature primaryUserIdCertification = getLatestUserIdCertification(getPrimaryUserId());
-        if (getLatestDirectKeySelfSignature() != null) {
+        if (primaryUserIdCertification != null) {
-            lastExpiration = SignatureUtils.getKeyExpirationDate(getCreationDate(), getLatestDirectKeySelfSignature());
+            return SignatureSubpacketsUtil.getKeyExpirationTimeAsDate(primaryUserIdCertification, getPublicKey());
         }
 
-        for (String userId : getValidUserIds()) {
+        PGPSignature directKeySig = getLatestDirectKeySelfSignature();
-
+        if (directKeySig != null) {
-            PGPSignature signature = getLatestUserIdCertification(userId);
+            return SignatureSubpacketsUtil.getKeyExpirationTimeAsDate(directKeySig, getPublicKey());
-            if (signature == null) {
-                continue;
-            }
-
-            Date expiration = SignatureUtils.getKeyExpirationDate(getCreationDate(), signature);
-            if (expiration != null && (lastExpiration == null || expiration.after(lastExpiration))) {
-                lastExpiration = expiration;
-            }
         }
-        return lastExpiration;
+        throw new NoSuchElementException("No suitable signatures found on the key.");
     }
 
     /**
@@ -620,15 +613,47 @@ public class KeyRingInfo {
 
         PGPPublicKey subkey = getPublicKey(fingerprint.getKeyId());
         if (subkey == null) {
-            throw new IllegalArgumentException("No subkey with fingerprint " + fingerprint + " found.");
+            throw new NoSuchElementException("No subkey with fingerprint " + fingerprint + " found.");
         }
+
         PGPSignature bindingSig = getCurrentSubkeyBindingSignature(fingerprint.getKeyId());
         if (bindingSig == null) {
-            return null;
+            throw new AssertionError("Subkey has no valid binding signature.");
         }
+
         return SignatureUtils.getKeyExpirationDate(subkey.getCreationTime(), bindingSig);
     }
 
+    public Date getExpirationDateForUse(KeyFlag use) {
+        if (use == KeyFlag.SPLIT || use == KeyFlag.SHARED) {
+            throw new IllegalArgumentException("SPLIT and SHARED are not uses, but properties.");
+        }
+
+        Date primaryExpiration = getPrimaryKeyExpirationDate();
+        List<PGPPublicKey> nonExpiringSubkeys = new ArrayList<>();
+        Date latestSubkeyExpirationDate = null;
+
+        List<PGPPublicKey> keysWithFlag = getKeysWithKeyFlag(use);
+        for (PGPPublicKey key : keysWithFlag) {
+            Date subkeyExpirationDate = getSubkeyExpirationDate(new OpenPgpV4Fingerprint(key));
+            if (subkeyExpirationDate == null) {
+                nonExpiringSubkeys.add(key);
+            } else {
+                if (latestSubkeyExpirationDate == null || subkeyExpirationDate.after(latestSubkeyExpirationDate)) {
+                    latestSubkeyExpirationDate = subkeyExpirationDate;
+                }
+            }
+        }
+
+        if (nonExpiringSubkeys.isEmpty()) {
+            if (latestSubkeyExpirationDate.before(primaryExpiration)) {
+                return latestSubkeyExpirationDate;
+            }
+            return primaryExpiration;
+        }
+        return null;
+    }
+
     /**
      * Return true if the key ring is a {@link PGPSecretKeyRing}.
      * If it is a {@link PGPPublicKeyRing} return false and if it is neither, throw an {@link AssertionError}.
@@ -727,6 +752,18 @@ public class KeyRingInfo {
         return encryptionKeys;
     }
 
+    public List<PGPPublicKey> getKeysWithKeyFlag(KeyFlag flag) {
+        List<PGPPublicKey> keysWithFlag = new ArrayList<>();
+        for (PGPPublicKey key : getPublicKeys()) {
+            List<KeyFlag> keyFlags = getKeyFlagsOf(key.getKeyID());
+            if (keyFlags.contains(flag)) {
+                keysWithFlag.add(key);
+            }
+        }
+
+        return keysWithFlag;
+    }
+
     /**
      * Return a list of all subkeys that can be used for encryption with the given user-id.
      * This list does not include expired or revoked keys.

pgpainless-core/src/test/java/org/pgpainless/key/info/KeyRingInfoTest.java

@@ -22,18 +22,33 @@ import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import java.io.IOException;
+import java.security.InvalidAlgorithmParameterException;
+import java.security.NoSuchAlgorithmException;
+import java.util.Calendar;
 import java.util.Collections;
 import java.util.Date;
+import java.util.Iterator;
+import java.util.List;
 
 import org.bouncycastle.openpgp.PGPException;
+import org.bouncycastle.openpgp.PGPPublicKey;
 import org.bouncycastle.openpgp.PGPPublicKeyRing;
+import org.bouncycastle.openpgp.PGPSecretKey;
 import org.bouncycastle.openpgp.PGPSecretKeyRing;
 import org.junit.jupiter.api.Test;
 import org.pgpainless.PGPainless;
+import org.pgpainless.algorithm.KeyFlag;
 import org.pgpainless.algorithm.PublicKeyAlgorithm;
+import org.pgpainless.key.OpenPgpV4Fingerprint;
 import org.pgpainless.key.TestKeys;
+import org.pgpainless.key.generation.KeySpec;
+import org.pgpainless.key.generation.type.KeyType;
+import org.pgpainless.key.generation.type.ecc.EllipticCurve;
+import org.pgpainless.key.generation.type.eddsa.EdDSACurve;
+import org.pgpainless.key.protection.SecretKeyRingProtector;
 import org.pgpainless.key.protection.UnprotectedKeysProtector;
 import org.pgpainless.key.util.KeyRingUtils;
+import org.pgpainless.key.util.UserId;
 import org.pgpainless.util.ArmorUtils;
 import org.pgpainless.util.Passphrase;
 
@@ -179,4 +194,62 @@ public class KeyRingInfoTest {
         PGPSecretKeyRing secretKeys = PGPainless.readKeyRing().secretKeyRing(withDummyS2K);
         assertTrue(new KeyInfo(secretKeys.getSecretKey()).hasDummyS2K());
     }
+
+    @Test
+    public void testGetKeysWithFlagsAndExpiry() throws PGPException, InvalidAlgorithmParameterException, NoSuchAlgorithmException {
+        PGPSecretKeyRing secretKeys = PGPainless.generateKeyRing()
+                .withSubKey(KeySpec.getBuilder(KeyType.ECDH(EllipticCurve._BRAINPOOLP384R1)).withKeyFlags(KeyFlag.ENCRYPT_STORAGE).withDefaultAlgorithms())
+                .withSubKey(KeySpec.getBuilder(KeyType.ECDSA(EllipticCurve._BRAINPOOLP384R1)).withKeyFlags(KeyFlag.SIGN_DATA).withDefaultAlgorithms())
+                .withPrimaryKey(KeySpec.getBuilder(KeyType.EDDSA(EdDSACurve._Ed25519)).withKeyFlags(KeyFlag.CERTIFY_OTHER).withDefaultAlgorithms())
+                .withPrimaryUserId(UserId.newBuilder().withName("Alice").withEmail("alice@pgpainless.org").build())
+                .withoutPassphrase()
+                .build();
+
+        Iterator<PGPSecretKey> keys = secretKeys.iterator();
+        Date now = new Date();
+
+        Calendar calendar = Calendar.getInstance();
+        calendar.setTime(now);
+        calendar.add(Calendar.DATE, 5);
+        Date primaryKeyExpiration = calendar.getTime(); // in 5 days
+        PGPSecretKey primaryKey = keys.next();
+
+        calendar.setTime(now);
+        calendar.add(Calendar.DATE, 10);
+        Date encryptionKeyExpiration = calendar.getTime(); // in 10 days
+        PGPSecretKey encryptionKey = keys.next();
+
+        calendar.setTime(now);
+        calendar.add(Calendar.DATE, 3);
+        Date signingKeyExpiration = calendar.getTime(); // in 3 days
+        PGPSecretKey signingKey = keys.next();
+
+        SecretKeyRingProtector protector = SecretKeyRingProtector.unprotectedKeys();
+        secretKeys = PGPainless.modifyKeyRing(secretKeys)
+                .setExpirationDate(new OpenPgpV4Fingerprint(primaryKey), primaryKeyExpiration, protector)
+                .setExpirationDate(new OpenPgpV4Fingerprint(encryptionKey), encryptionKeyExpiration, protector)
+                .setExpirationDate(new OpenPgpV4Fingerprint(signingKey), signingKeyExpiration, protector)
+                .done();
+
+        KeyRingInfo info = new KeyRingInfo(secretKeys);
+
+        List<PGPPublicKey> encryptionKeys = info.getKeysWithKeyFlag(KeyFlag.ENCRYPT_STORAGE);
+        assertEquals(1, encryptionKeys.size());
+        assertEquals(encryptionKey.getKeyID(), encryptionKeys.get(0).getKeyID());
+
+        List<PGPPublicKey> signingKeys = info.getKeysWithKeyFlag(KeyFlag.SIGN_DATA);
+        assertEquals(1, signingKeys.size());
+        assertEquals(signingKey.getKeyID(), signingKeys.get(0).getKeyID());
+
+        List<PGPPublicKey> certKeys = info.getKeysWithKeyFlag(KeyFlag.CERTIFY_OTHER);
+        assertEquals(1, certKeys.size());
+        assertEquals(primaryKey.getKeyID(), certKeys.get(0).getKeyID());
+
+        assertEquals(primaryKeyExpiration.getTime(), info.getPrimaryKeyExpirationDate().getTime(), 5);
+        assertEquals(signingKeyExpiration.getTime(), info.getExpirationDateForUse(KeyFlag.SIGN_DATA).getTime(), 5);
+
+        // Encryption key expires after primary key, so we return primary key expiration instead.
+        assertEquals(primaryKeyExpiration.getTime(), info.getExpirationDateForUse(KeyFlag.ENCRYPT_STORAGE).getTime(), 5);
+
+    }
 }