From d23b0c5d6fe1b862c28084576bbef4369196a4bf Mon Sep 17 00:00:00 2001
From: Harsh Shandilya <me@msfjarvis.dev>
Date: Thu, 14 Jul 2022 00:42:23 +0530
Subject: Fix PGPainless backend key handling (#2000)

---
 crypto-pgpainless/build.gradle.kts                 |  1 +
 .../aps/crypto/PGPainlessCryptoHandler.kt          | 30 +++++++++++++++-------
 .../aps/crypto/PGPainlessCryptoHandlerTest.kt      | 16 +++++++++---
 3 files changed, 35 insertions(+), 12 deletions(-)

(limited to 'crypto-pgpainless')

diff --git a/crypto-pgpainless/build.gradle.kts b/crypto-pgpainless/build.gradle.kts
index b3ace3d9..93d56eff 100644
--- a/crypto-pgpainless/build.gradle.kts
+++ b/crypto-pgpainless/build.gradle.kts
@@ -18,4 +18,5 @@ dependencies {
   implementation(libs.thirdparty.pgpainless)
   testImplementation(libs.bundles.testDependencies)
   testImplementation(libs.kotlin.coroutines.test)
+  testImplementation(libs.testing.testparameterinjector)
 }
diff --git a/crypto-pgpainless/src/main/kotlin/dev/msfjarvis/aps/crypto/PGPainlessCryptoHandler.kt b/crypto-pgpainless/src/main/kotlin/dev/msfjarvis/aps/crypto/PGPainlessCryptoHandler.kt
index 3ebdb44c..24b3e665 100644
--- a/crypto-pgpainless/src/main/kotlin/dev/msfjarvis/aps/crypto/PGPainlessCryptoHandler.kt
+++ b/crypto-pgpainless/src/main/kotlin/dev/msfjarvis/aps/crypto/PGPainlessCryptoHandler.kt
@@ -16,6 +16,8 @@ import java.io.ByteArrayInputStream
 import java.io.InputStream
 import java.io.OutputStream
 import javax.inject.Inject
+import org.bouncycastle.openpgp.PGPPublicKeyRing
+import org.bouncycastle.openpgp.PGPPublicKeyRingCollection
 import org.bouncycastle.openpgp.PGPSecretKeyRingCollection
 import org.pgpainless.PGPainless
 import org.pgpainless.decryption_verification.ConsumerOptions
@@ -65,14 +67,25 @@ public class PGPainlessCryptoHandler @Inject constructor() : CryptoHandler<PGPKe
   ): Result<Unit, CryptoHandlerException> =
     runCatching {
         if (keys.isEmpty()) throw NoKeysProvided("No keys provided for encryption")
-        val armoredKeys = keys.map { key -> key.contents.decodeToString() }
-        val pubKeysStream = ByteArrayInputStream(armoredKeys.joinToString("\n").toByteArray())
-        val publicKeyRingCollection =
-          pubKeysStream.use { PGPainless.readKeyRing().publicKeyRingCollection(pubKeysStream) }
-        val encryptionOptions =
-          EncryptionOptions.encryptCommunications()
-            .addRecipients(publicKeyRingCollection.asIterable())
-        val producerOptions = ProducerOptions.encrypt(encryptionOptions).setAsciiArmor(true)
+        val publicKeyRings = arrayListOf<PGPPublicKeyRing>()
+        val armoredKeys =
+          keys.joinToString("\n") { key -> key.contents.decodeToString() }.toByteArray()
+        val secKeysStream = ByteArrayInputStream(armoredKeys)
+        val secretKeyRingCollection =
+          PGPainless.readKeyRing().secretKeyRingCollection(secKeysStream)
+        secretKeyRingCollection.forEach { secretKeyRing ->
+          publicKeyRings.add(PGPainless.extractCertificate(secretKeyRing))
+        }
+        if (publicKeyRings.isEmpty()) {
+          val pubKeysStream = ByteArrayInputStream(armoredKeys)
+          val publicKeyRingCollection =
+            PGPainless.readKeyRing().publicKeyRingCollection(pubKeysStream)
+          publicKeyRings.addAll(publicKeyRingCollection)
+        }
+        require(publicKeyRings.isNotEmpty()) { "No public keys to encrypt message to" }
+        val publicKeyRingCollection = PGPPublicKeyRingCollection(publicKeyRings)
+        val encryptionOptions = EncryptionOptions().addRecipients(publicKeyRingCollection)
+        val producerOptions = ProducerOptions.encrypt(encryptionOptions).setAsciiArmor(false)
         val encryptor =
           PGPainless.encryptAndOrSign().onOutputStream(outputStream).withOptions(producerOptions)
         plaintextStream.copyTo(encryptor)
@@ -83,7 +96,6 @@ public class PGPainlessCryptoHandler @Inject constructor() : CryptoHandler<PGPKe
             "Stream should be encrypted for ${keyRing.publicKey.keyID} but wasn't"
           }
         }
-        return@runCatching
       }
       .mapError { error ->
         when (error) {
diff --git a/crypto-pgpainless/src/test/kotlin/dev/msfjarvis/aps/crypto/PGPainlessCryptoHandlerTest.kt b/crypto-pgpainless/src/test/kotlin/dev/msfjarvis/aps/crypto/PGPainlessCryptoHandlerTest.kt
index a9484317..e39bc06e 100644
--- a/crypto-pgpainless/src/test/kotlin/dev/msfjarvis/aps/crypto/PGPainlessCryptoHandlerTest.kt
+++ b/crypto-pgpainless/src/test/kotlin/dev/msfjarvis/aps/crypto/PGPainlessCryptoHandlerTest.kt
@@ -7,6 +7,8 @@ package dev.msfjarvis.aps.crypto
 
 import com.github.michaelbull.result.Err
 import com.github.michaelbull.result.getError
+import com.google.testing.junit.testparameterinjector.TestParameter
+import com.google.testing.junit.testparameterinjector.TestParameterInjector
 import dev.msfjarvis.aps.crypto.errors.IncorrectPassphraseException
 import java.io.ByteArrayOutputStream
 import kotlin.test.Test
@@ -14,18 +16,26 @@ import kotlin.test.assertEquals
 import kotlin.test.assertFalse
 import kotlin.test.assertIs
 import kotlin.test.assertTrue
+import org.junit.runner.RunWith
 
+@Suppress("Unused") // Test runner handles it internally
+enum class EncryptionKey(val key: PGPKey) {
+  PUBLIC(PGPKey(TestUtils.getArmoredPublicKey())),
+  SECRET(PGPKey(TestUtils.getArmoredPrivateKey())),
+}
+
+@RunWith(TestParameterInjector::class)
 class PGPainlessCryptoHandlerTest {
 
+  @TestParameter private lateinit var encryptionKey: EncryptionKey
   private val cryptoHandler = PGPainlessCryptoHandler()
   private val privateKey = PGPKey(TestUtils.getArmoredPrivateKey())
-  private val publicKey = PGPKey(TestUtils.getArmoredPublicKey())
 
   @Test
   fun encryptAndDecrypt() {
     val ciphertextStream = ByteArrayOutputStream()
     cryptoHandler.encrypt(
-      listOf(publicKey),
+      listOf(encryptionKey.key),
       CryptoConstants.PLAIN_TEXT.byteInputStream(Charsets.UTF_8),
       ciphertextStream,
     )
@@ -43,7 +53,7 @@ class PGPainlessCryptoHandlerTest {
   fun decryptWithWrongPassphrase() {
     val ciphertextStream = ByteArrayOutputStream()
     cryptoHandler.encrypt(
-      listOf(publicKey),
+      listOf(encryptionKey.key),
       CryptoConstants.PLAIN_TEXT.byteInputStream(Charsets.UTF_8),
       ciphertextStream,
     )
-- 
cgit v1.2.3