Refactor and simplify KeyManager API (#1650)

5 files changed, 125 insertions, 99 deletions

diff --git a/crypto-pgpainless/src/main/kotlin/dev/msfjarvis/aps/crypto/PGPKeyManager.kt b/crypto-pgpainless/src/main/kotlin/dev/msfjarvis/aps/crypto/PGPKeyManager.kt
index fd886843..f1c53721 100644
--- a/crypto-pgpainless/src/main/kotlin/dev/msfjarvis/aps/crypto/PGPKeyManager.kt
+++ b/crypto-pgpainless/src/main/kotlin/dev/msfjarvis/aps/crypto/PGPKeyManager.kt
@@ -2,46 +2,59 @@
  * Copyright © 2014-2021 The Android Password Store Authors. All Rights Reserved.
  * SPDX-License-Identifier: GPL-3.0-only
  */
+@file:Suppress("BlockingMethodInNonBlockingContext")
 
 package dev.msfjarvis.aps.crypto
 
 import androidx.annotation.VisibleForTesting
 import com.github.michaelbull.result.Result
+import com.github.michaelbull.result.get
 import com.github.michaelbull.result.runCatching
 import java.io.File
+import java.util.Locale
+import javax.inject.Inject
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.withContext
+import org.bouncycastle.openpgp.PGPKeyRing
 import org.pgpainless.PGPainless
+import org.pgpainless.util.selection.userid.SelectUserId
 
-public class PGPKeyManager(
+public class PGPKeyManager
+@Inject
+constructor(
   filesDir: String,
   private val dispatcher: CoroutineDispatcher,
-) : KeyManager<PGPKeyPair> {
+) : KeyManager {
 
   private val keyDir = File(filesDir, KEY_DIR_NAME)
 
-  override suspend fun addKey(key: PGPKeyPair, replace: Boolean): Result<PGPKeyPair, Throwable> =
+  override suspend fun addKey(key: Key, replace: Boolean): Result<Key, Throwable> =
     withContext(dispatcher) {
       runCatching {
         if (!keyDirExists()) throw KeyManagerException.KeyDirectoryUnavailableException
-        val keyFile = File(keyDir, "${key.getKeyId()}.$KEY_EXTENSION")
+        if (tryParseKeyring(key) == null) throw KeyManagerException.InvalidKeyException
+        val keyFile = File(keyDir, "${tryGetId(key)}.$KEY_EXTENSION")
         if (keyFile.exists()) {
           // Check for replace flag first and if it is false, throw an error
-          if (!replace) throw KeyManagerException.KeyAlreadyExistsException(key.getKeyId())
+          if (!replace)
+            throw KeyManagerException.KeyAlreadyExistsException(
+              tryGetId(key) ?: "Failed to retrieve key ID"
+            )
           if (!keyFile.delete()) throw KeyManagerException.KeyDeletionFailedException
         }
 
-        keyFile.writeBytes(key.getPrivateKey())
+        keyFile.writeBytes(key.contents)
 
         key
       }
     }
 
-  override suspend fun removeKey(key: PGPKeyPair): Result<PGPKeyPair, Throwable> =
+  override suspend fun removeKey(key: Key): Result<Key, Throwable> =
     withContext(dispatcher) {
       runCatching {
         if (!keyDirExists()) throw KeyManagerException.KeyDirectoryUnavailableException
-        val keyFile = File(keyDir, "${key.getKeyId()}.$KEY_EXTENSION")
+        if (tryParseKeyring(key) == null) throw KeyManagerException.InvalidKeyException
+        val keyFile = File(keyDir, "${tryGetId(key)}.$KEY_EXTENSION")
         if (keyFile.exists()) {
           if (!keyFile.delete()) throw KeyManagerException.KeyDeletionFailedException
         }
@@ -50,63 +63,105 @@ public class PGPKeyManager(
       }
     }
 
-  override suspend fun getKeyById(id: String): Result<PGPKeyPair, Throwable> =
+  override suspend fun getKeyById(id: String): Result<Key, Throwable> =
     withContext(dispatcher) {
       runCatching {
         if (!keyDirExists()) throw KeyManagerException.KeyDirectoryUnavailableException
-        val keys = keyDir.listFiles()
-        if (keys.isNullOrEmpty()) throw KeyManagerException.NoKeysAvailableException
-
-        for (keyFile in keys) {
-          val secretKeyRing = PGPainless.readKeyRing().secretKeyRing(keyFile.inputStream())
-          val secretKey = secretKeyRing.secretKey
-          val keyPair = PGPKeyPair(secretKey)
-          if (keyPair.getKeyId() == id) return@runCatching keyPair
+        val keyFiles = keyDir.listFiles()
+        if (keyFiles.isNullOrEmpty()) throw KeyManagerException.NoKeysAvailableException
+
+        val keys = keyFiles.map { file -> Key(file.readBytes()) }
+        // Try to parse the key ID as an email
+        val selector = SelectUserId.byEmail(id)
+        val userIdMatch =
+          keys.map { key -> key to tryParseKeyring(key) }.firstOrNull { (_, keyRing) ->
+            selector.firstMatch(keyRing) != null
+          }
+
+        if (userIdMatch != null) {
+          return@runCatching userIdMatch.first
+        }
+
+        val keyIdMatch =
+          keys.map { key -> key to tryGetId(key) }.firstOrNull { (_, keyId) ->
+            keyId == id || keyId == "0x$id"
+          }
+
+        if (keyIdMatch != null) {
+          return@runCatching keyIdMatch.first
         }
 
         throw KeyManagerException.KeyNotFoundException(id)
       }
     }
 
-  override suspend fun getAllKeys(): Result<List<PGPKeyPair>, Throwable> =
+  override suspend fun getAllKeys(): Result<List<Key>, Throwable> =
     withContext(dispatcher) {
       runCatching {
         if (!keyDirExists()) throw KeyManagerException.KeyDirectoryUnavailableException
-        val keys = keyDir.listFiles()
-        if (keys.isNullOrEmpty()) return@runCatching listOf()
-
-        keys
-          .map { keyFile ->
-            val secretKeyRing = PGPainless.readKeyRing().secretKeyRing(keyFile.inputStream())
-            val secretKey = secretKeyRing.secretKey
-
-            PGPKeyPair(secretKey)
-          }
-          .toList()
+        val keyFiles = keyDir.listFiles()
+        if (keyFiles.isNullOrEmpty()) return@runCatching emptyList()
+        keyFiles.map { keyFile -> Key(keyFile.readBytes()) }.toList()
       }
     }
 
+  override suspend fun getKeyId(key: Key): String? = tryGetId(key)
+
+  // TODO: This is a temp hack for now and in future it should check that the GPGKeyManager can
+  // decrypt the file.
   override fun canHandle(fileName: String): Boolean {
-    // TODO: This is a temp hack for now and in future it should check that the GPGKeyManager can
-    // decrypt the file
     return fileName.endsWith(KEY_EXTENSION)
   }
 
+  /** Checks if [keyDir] exists and attempts to create it if not. */
   private fun keyDirExists(): Boolean {
     return keyDir.exists() || keyDir.mkdirs()
   }
 
+  /**
+   * Attempts to parse a [PGPKeyRing] from a given [key]. The key is first tried as a secret key and
+   * then as a public one before the method gives up and returns null.
+   */
+  private fun tryParseKeyring(key: Key): PGPKeyRing? {
+    val secKeyRing = runCatching { PGPainless.readKeyRing().secretKeyRing(key.contents) }.get()
+    if (secKeyRing != null) {
+      return secKeyRing
+    }
+    val pubKeyRing = runCatching { PGPainless.readKeyRing().publicKeyRing(key.contents) }.get()
+    if (pubKeyRing != null) {
+      return pubKeyRing
+    }
+    return null
+  }
+
+  /** Parses a [PGPKeyRing] from the given [key] and returns its hex-formatted key ID. */
+  private fun tryGetId(key: Key): String? {
+    val keyRing = tryParseKeyring(key) ?: return null
+    return convertKeyIdToHex(keyRing.publicKey.keyID)
+  }
+
+  /** Convert a [Long] key ID to a formatted string. */
+  private fun convertKeyIdToHex(keyId: Long): String {
+    return "0x" + convertKeyIdToHex32bit(keyId shr 32) + convertKeyIdToHex32bit(keyId)
+  }
+
+  /**
+   * Converts [keyId] to an unsigned [Long] then uses [java.lang.Long.toHexString] to convert it to
+   * a lowercase hex ID.
+   */
+  private fun convertKeyIdToHex32bit(keyId: Long): String {
+    var hexString = java.lang.Long.toHexString(keyId and 0xffffffffL).lowercase(Locale.ENGLISH)
+    while (hexString.length < 8) {
+      hexString = "0$hexString"
+    }
+    return hexString
+  }
+
   public companion object {
 
     @VisibleForTesting(otherwise = VisibleForTesting.PRIVATE)
     internal const val KEY_DIR_NAME: String = "keys"
     @VisibleForTesting(otherwise = VisibleForTesting.PRIVATE)
     internal const val KEY_EXTENSION: String = "key"
-
-    @JvmStatic
-    public fun makeKey(armoredKey: String): PGPKeyPair {
-      val secretKey = PGPainless.readKeyRing().secretKeyRing(armoredKey).secretKey
-      return PGPKeyPair(secretKey)
-    }
   }
 }
diff --git a/crypto-pgpainless/src/main/kotlin/dev/msfjarvis/aps/crypto/PGPKeyPair.kt b/crypto-pgpainless/src/main/kotlin/dev/msfjarvis/aps/crypto/PGPKeyPair.kt
deleted file mode 100644
index 03bcf515..00000000
--- a/crypto-pgpainless/src/main/kotlin/dev/msfjarvis/aps/crypto/PGPKeyPair.kt
+++ /dev/null
@@ -1,31 +0,0 @@
-/*
- * Copyright © 2014-2021 The Android Password Store Authors. All Rights Reserved.
- * SPDX-License-Identifier: GPL-3.0-only
- */
-
-package dev.msfjarvis.aps.crypto
-
-import org.bouncycastle.openpgp.PGPSecretKey
-
-public class PGPKeyPair(private val secretKey: PGPSecretKey) : KeyPair {
-
-  init {
-    if (secretKey.isPrivateKeyEmpty) throw KeyPairException.PrivateKeyUnavailableException
-  }
-
-  override fun getPrivateKey(): ByteArray {
-    return secretKey.encoded
-  }
-  override fun getPublicKey(): ByteArray {
-    return secretKey.publicKey.encoded
-  }
-  override fun getKeyId(): String {
-    var keyId = secretKey.keyID.toString(radix = 16)
-    if (keyId.length < KEY_ID_LENGTH) keyId = "0$keyId"
-    return keyId
-  }
-
-  private companion object {
-    private const val KEY_ID_LENGTH = 16
-  }
-}
diff --git a/crypto-pgpainless/src/test/kotlin/dev/msfjarvis/aps/crypto/CryptoConstants.kt b/crypto-pgpainless/src/test/kotlin/dev/msfjarvis/aps/crypto/CryptoConstants.kt
index 7aef4675..6d5eaf01 100644
--- a/crypto-pgpainless/src/test/kotlin/dev/msfjarvis/aps/crypto/CryptoConstants.kt
+++ b/crypto-pgpainless/src/test/kotlin/dev/msfjarvis/aps/crypto/CryptoConstants.kt
@@ -10,5 +10,5 @@ object CryptoConstants {
   const val PLAIN_TEXT = "encryption worthy content"
   const val KEY_NAME = "John Doe"
   const val KEY_EMAIL = "john.doe@example.com"
-  const val KEY_ID = "08edf7567183ce27"
+  const val KEY_ID = "0x08edf7567183ce27"
 }
diff --git a/crypto-pgpainless/src/test/kotlin/dev/msfjarvis/aps/crypto/PGPKeyManagerTest.kt b/crypto-pgpainless/src/test/kotlin/dev/msfjarvis/aps/crypto/PGPKeyManagerTest.kt
index 4b022e5e..c547bdd4 100644
--- a/crypto-pgpainless/src/test/kotlin/dev/msfjarvis/aps/crypto/PGPKeyManagerTest.kt
+++ b/crypto-pgpainless/src/test/kotlin/dev/msfjarvis/aps/crypto/PGPKeyManagerTest.kt
@@ -8,6 +8,7 @@ import kotlin.test.BeforeTest
 import kotlin.test.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertIs
+import kotlin.test.assertNotNull
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.StandardTestDispatcher
@@ -27,7 +28,7 @@ class PGPKeyManagerTest {
   private val dispatcher = StandardTestDispatcher()
   private val scope = TestScope(dispatcher)
   private val keyManager by unsafeLazy { PGPKeyManager(filesDir.absolutePath, dispatcher) }
-  private val key = PGPKeyManager.makeKey(TestUtils.getArmoredPrivateKey())
+  private val key = Key(TestUtils.getArmoredPrivateKey().encodeToByteArray())
 
   private fun <T> unsafeLazy(initializer: () -> T) =
     lazy(LazyThreadSafetyMode.NONE) { initializer.invoke() }
@@ -46,7 +47,7 @@ class PGPKeyManagerTest {
   fun testAddingKey() =
     scope.runTest {
       // Check if the key id returned is correct
-      val keyId = keyManager.addKey(key).unwrap().getKeyId()
+      val keyId = keyManager.getKeyId(keyManager.addKey(key).unwrap())
       assertEquals(CryptoConstants.KEY_ID, keyId)
 
       // Check if the keys directory have one file
@@ -72,7 +73,7 @@ class PGPKeyManagerTest {
     scope.runTest {
       // Check adding the keys twice
       keyManager.addKey(key, true).unwrap()
-      val keyId = keyManager.addKey(key, true).unwrap().getKeyId()
+      val keyId = keyManager.getKeyId(keyManager.addKey(key, true).unwrap())
 
       assertEquals(CryptoConstants.KEY_ID, keyId)
     }
@@ -84,7 +85,7 @@ class PGPKeyManagerTest {
       keyManager.addKey(key).unwrap()
 
       // Check if the key id returned is correct
-      val keyId = keyManager.removeKey(key).unwrap().getKeyId()
+      val keyId = keyManager.getKeyId(keyManager.removeKey(key).unwrap())
       assertEquals(CryptoConstants.KEY_ID, keyId)
 
       // Check if the keys directory have 0 files
@@ -93,15 +94,38 @@ class PGPKeyManagerTest {
     }
 
   @Test
-  fun testGetExistingKey() =
+  fun testGetExistingKeyById() =
     scope.runTest {
       // Add key using KeyManager
       keyManager.addKey(key).unwrap()
 
+      val keyId = keyManager.getKeyId(key)
+      assertNotNull(keyId)
+      assertEquals(CryptoConstants.KEY_ID, keyManager.getKeyId(key))
+
       // Check returned key id matches the expected id and the created key id
-      val returnedKeyPair = keyManager.getKeyById(key.getKeyId()).unwrap()
-      assertEquals(CryptoConstants.KEY_ID, key.getKeyId())
-      assertEquals(key.getKeyId(), returnedKeyPair.getKeyId())
+      val returnedKey = keyManager.getKeyById(keyId).unwrap()
+      assertEquals(keyManager.getKeyId(key), keyManager.getKeyId(returnedKey))
+    }
+
+  @Test
+  fun testGetExistingKeyByFullUserId() =
+    scope.runTest {
+      keyManager.addKey(key).unwrap()
+
+      val keyId = "${CryptoConstants.KEY_NAME} <${CryptoConstants.KEY_EMAIL}>"
+      val returnedKey = keyManager.getKeyById(keyId).unwrap()
+      assertEquals(keyManager.getKeyId(key), keyManager.getKeyId(returnedKey))
+    }
+
+  @Test
+  fun testGetExistingKeyByEmailUserId() =
+    scope.runTest {
+      keyManager.addKey(key).unwrap()
+
+      val keyId = CryptoConstants.KEY_EMAIL
+      val returnedKey = keyManager.getKeyById(keyId).unwrap()
+      assertEquals(keyManager.getKeyId(key), keyManager.getKeyId(returnedKey))
     }
 
   @Test
diff --git a/crypto-pgpainless/src/test/kotlin/dev/msfjarvis/aps/crypto/PGPKeyPairTest.kt b/crypto-pgpainless/src/test/kotlin/dev/msfjarvis/aps/crypto/PGPKeyPairTest.kt
deleted file mode 100644
index 45b25f30..00000000
--- a/crypto-pgpainless/src/test/kotlin/dev/msfjarvis/aps/crypto/PGPKeyPairTest.kt
+++ /dev/null
@@ -1,22 +0,0 @@
-/*
- * Copyright © 2014-2021 The Android Password Store Authors. All Rights Reserved.
- * SPDX-License-Identifier: GPL-3.0-only
- */
-
-package dev.msfjarvis.aps.crypto
-
-import kotlin.test.Test
-import kotlin.test.assertEquals
-import org.pgpainless.PGPainless
-
-class PGPKeyPairTest {
-
-  @Test
-  fun testIfKeyIdIsCorrect() {
-    val secretKey =
-      PGPainless.readKeyRing().secretKeyRing(TestUtils.getArmoredPrivateKey()).secretKey
-    val keyPair = PGPKeyPair(secretKey)
-
-    assertEquals(CryptoConstants.KEY_ID, keyPair.getKeyId())
-  }
-}