Parameterize key and key identifier types for KeyManager (#1669)

16 files changed, 202 insertions, 87 deletions

diff --git a/app/src/main/java/dev/msfjarvis/aps/injection/crypto/CryptoHandlerModule.kt b/app/src/main/java/dev/msfjarvis/aps/injection/crypto/CryptoHandlerModule.kt
index 63a860d1..ef6a11ce 100644
--- a/app/src/main/java/dev/msfjarvis/aps/injection/crypto/CryptoHandlerModule.kt
+++ b/app/src/main/java/dev/msfjarvis/aps/injection/crypto/CryptoHandlerModule.kt
@@ -9,23 +9,10 @@ import dagger.Module
 import dagger.Provides
 import dagger.hilt.InstallIn
 import dagger.hilt.components.SingletonComponent
-import dagger.multibindings.IntoSet
-import dev.msfjarvis.aps.crypto.CryptoHandler
 import dev.msfjarvis.aps.crypto.PGPainlessCryptoHandler
 
-/**
- * This module adds all [CryptoHandler] implementations into a Set which makes it easier to build
- * generic UIs which are not tied to a specific implementation because of injection.
- */
 @Module
 @InstallIn(SingletonComponent::class)
 object CryptoHandlerModule {
-  @Provides
-  @IntoSet
-  fun providePgpCryptoHandler(): CryptoHandler {
-    return PGPainlessCryptoHandler()
-  }
+  @Provides fun providePgpCryptoHandler() = PGPainlessCryptoHandler()
 }
-
-/** Typealias for a [Set] of [CryptoHandler] instances injected by Dagger. */
-typealias CryptoSet = Set<@JvmSuppressWildcards CryptoHandler>
diff --git a/app/src/main/java/dev/msfjarvis/aps/injection/crypto/KeyManagerModule.kt b/app/src/main/java/dev/msfjarvis/aps/injection/crypto/KeyManagerModule.kt
index 0bc1c43d..a1119a1c 100644
--- a/app/src/main/java/dev/msfjarvis/aps/injection/crypto/KeyManagerModule.kt
+++ b/app/src/main/java/dev/msfjarvis/aps/injection/crypto/KeyManagerModule.kt
@@ -12,8 +12,8 @@ import dagger.hilt.InstallIn
 import dagger.hilt.android.qualifiers.ApplicationContext
 import dagger.hilt.components.SingletonComponent
 import dev.msfjarvis.aps.crypto.PGPKeyManager
+import dev.msfjarvis.aps.util.coroutines.DispatcherProvider
 import javax.inject.Qualifier
-import kotlinx.coroutines.Dispatchers
 
 @Module
 @InstallIn(SingletonComponent::class)
@@ -21,10 +21,11 @@ object KeyManagerModule {
   @Provides
   fun providePGPKeyManager(
     @PGPKeyDir keyDir: String,
+    dispatcherProvider: DispatcherProvider,
   ): PGPKeyManager {
     return PGPKeyManager(
       keyDir,
-      Dispatchers.IO,
+      dispatcherProvider.io(),
     )
   }
 
diff --git a/app/src/main/java/dev/msfjarvis/aps/ui/pgp/PGPKeyImportActivity.kt b/app/src/main/java/dev/msfjarvis/aps/ui/pgp/PGPKeyImportActivity.kt
index 3cd6930c..b97de5e0 100644
--- a/app/src/main/java/dev/msfjarvis/aps/ui/pgp/PGPKeyImportActivity.kt
+++ b/app/src/main/java/dev/msfjarvis/aps/ui/pgp/PGPKeyImportActivity.kt
@@ -14,8 +14,8 @@ import com.github.michaelbull.result.runCatching
 import com.google.android.material.dialog.MaterialAlertDialogBuilder
 import dagger.hilt.android.AndroidEntryPoint
 import dev.msfjarvis.aps.R
-import dev.msfjarvis.aps.crypto.Key
 import dev.msfjarvis.aps.crypto.KeyUtils.tryGetId
+import dev.msfjarvis.aps.crypto.PGPKey
 import dev.msfjarvis.aps.crypto.PGPKeyManager
 import javax.inject.Inject
 import kotlinx.coroutines.runBlocking
@@ -35,7 +35,7 @@ class PGPKeyImportActivity : AppCompatActivity() {
             contentResolver.openInputStream(uri)
               ?: throw IllegalStateException("Failed to open selected file")
           val bytes = keyInputStream.readBytes()
-          val (key, error) = runBlocking { keyManager.addKey(Key(bytes)) }
+          val (key, error) = runBlocking { keyManager.addKey(PGPKey(bytes)) }
           if (error != null) throw error
           key
         }
diff --git a/crypto-common/src/main/kotlin/dev/msfjarvis/aps/crypto/CryptoException.kt b/crypto-common/src/main/kotlin/dev/msfjarvis/aps/crypto/CryptoException.kt
index 431cd90c..42429a14 100644
--- a/crypto-common/src/main/kotlin/dev/msfjarvis/aps/crypto/CryptoException.kt
+++ b/crypto-common/src/main/kotlin/dev/msfjarvis/aps/crypto/CryptoException.kt
@@ -15,7 +15,7 @@ public sealed class KeyManagerException(message: String? = null) : CryptoExcepti
   /** Failed to delete given key. */
   public object KeyDeletionFailedException : KeyManagerException("Couldn't delete the key file")
 
-  /** Failed to parse a [Key] as a known type. */
+  /** Failed to parse the key as a known type. */
   public object InvalidKeyException :
     KeyManagerException("Given key cannot be parsed as a known key type")
 
diff --git a/crypto-common/src/main/kotlin/dev/msfjarvis/aps/crypto/CryptoHandler.kt b/crypto-common/src/main/kotlin/dev/msfjarvis/aps/crypto/CryptoHandler.kt
index eb04ee7a..fc4c500d 100644
--- a/crypto-common/src/main/kotlin/dev/msfjarvis/aps/crypto/CryptoHandler.kt
+++ b/crypto-common/src/main/kotlin/dev/msfjarvis/aps/crypto/CryptoHandler.kt
@@ -9,7 +9,7 @@ import java.io.InputStream
 import java.io.OutputStream
 
 /** Generic interface to implement cryptographic operations on top of. */
-public interface CryptoHandler {
+public interface CryptoHandler<Key> {
 
   /**
    * Decrypt the given [ciphertextStream] using a [privateKey] and [passphrase], and writes the
diff --git a/crypto-common/src/main/kotlin/dev/msfjarvis/aps/crypto/Key.kt b/crypto-common/src/main/kotlin/dev/msfjarvis/aps/crypto/Key.kt
deleted file mode 100644
index 73dee199..00000000
--- a/crypto-common/src/main/kotlin/dev/msfjarvis/aps/crypto/Key.kt
+++ /dev/null
@@ -1,14 +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
-
-/**
- * A simple value class wrapping over a [ByteArray] that can be used as a key type for cryptographic
- * purposes. The public/private distinction is elided specifically to defer that decision to
- * implementations of [KeyManager]. Similarly, identification of the key's identities is also
- * deferred to [KeyManager] to ensure maximum flexibility.
- */
-@JvmInline public value class Key(public val contents: ByteArray)
diff --git a/crypto-common/src/main/kotlin/dev/msfjarvis/aps/crypto/KeyManager.kt b/crypto-common/src/main/kotlin/dev/msfjarvis/aps/crypto/KeyManager.kt
index dacdfc6a..aa71abf6 100644
--- a/crypto-common/src/main/kotlin/dev/msfjarvis/aps/crypto/KeyManager.kt
+++ b/crypto-common/src/main/kotlin/dev/msfjarvis/aps/crypto/KeyManager.kt
@@ -12,7 +12,7 @@ import com.github.michaelbull.result.Result
  * used by an implementation of [CryptoHandler] to obtain eligible public or private keys as
  * required.
  */
-public interface KeyManager {
+public interface KeyManager<Key, KeyIdentifier> {
 
   /**
    * Inserts a [key] into the store. If the key already exists, this method will return
@@ -28,7 +28,7 @@ public interface KeyManager {
    * implementations to figure out for themselves. For example, in GPG this can be a full
    * hexadecimal key ID, an email, a short hex key ID, and probably a few more things.
    */
-  public suspend fun getKeyById(id: String): Result<Key, Throwable>
+  public suspend fun getKeyById(id: KeyIdentifier): Result<Key, Throwable>
 
   /** Returns all keys currently in the store as a [List]. */
   public suspend fun getAllKeys(): Result<List<Key>, Throwable>
@@ -37,7 +37,7 @@ public interface KeyManager {
    * Get a stable identifier for the given [key]. The returned key ID should be suitable to be used
    * as an identifier for the cryptographic identity tied to this key.
    */
-  public suspend fun getKeyId(key: Key): String?
+  public suspend fun getKeyId(key: Key): KeyIdentifier?
 
   /** Given a [fileName], return whether this instance can handle it. */
   public fun canHandle(fileName: String): Boolean
diff --git a/crypto-pgpainless/src/main/kotlin/dev/msfjarvis/aps/crypto/GpgIdentifier.kt b/crypto-pgpainless/src/main/kotlin/dev/msfjarvis/aps/crypto/GpgIdentifier.kt
new file mode 100644
index 00000000..20337d3e
--- /dev/null
+++ b/crypto-pgpainless/src/main/kotlin/dev/msfjarvis/aps/crypto/GpgIdentifier.kt
@@ -0,0 +1,82 @@
+/*
+ * Copyright © 2014-2022 The Android Password Store Authors. All Rights Reserved.
+ * SPDX-License-Identifier: GPL-3.0-only
+ */
+
+package dev.msfjarvis.aps.crypto
+
+import java.util.regex.Pattern
+
+public sealed class GpgIdentifier {
+  public data class KeyId(val id: Long) : GpgIdentifier() {
+    override fun toString(): String {
+      return java.lang.Long.toHexString(id)
+    }
+  }
+  public data class UserId(val email: String) : GpgIdentifier() {
+    override fun toString(): String {
+      return email
+    }
+  }
+
+  public companion object {
+    @OptIn(ExperimentalUnsignedTypes::class)
+    public fun fromString(identifier: String): GpgIdentifier? {
+      if (identifier.isEmpty()) return null
+      // Match long key IDs:
+      // FF22334455667788 or 0xFF22334455667788
+      val maybeLongKeyId =
+        identifier.removePrefix("0x").takeIf { it.matches("[a-fA-F0-9]{16}".toRegex()) }
+      if (maybeLongKeyId != null) {
+        val keyId = maybeLongKeyId.toULong(16)
+        return KeyId(keyId.toLong())
+      }
+
+      // Match fingerprints:
+      // FF223344556677889900112233445566778899 or 0xFF223344556677889900112233445566778899
+      val maybeFingerprint =
+        identifier.removePrefix("0x").takeIf { it.matches("[a-fA-F0-9]{40}".toRegex()) }
+      if (maybeFingerprint != null) {
+        // Truncating to the long key ID is not a security issue since OpenKeychain only
+        // accepts
+        // non-ambiguous key IDs.
+        val keyId = maybeFingerprint.takeLast(16).toULong(16)
+        return KeyId(keyId.toLong())
+      }
+
+      return splitUserId(identifier)?.let { UserId(it) }
+    }
+
+    private val USER_ID_PATTERN = Pattern.compile("^(.*?)(?: \\((.*)\\))?(?: <(.*)>)?$")
+    private val EMAIL_PATTERN = Pattern.compile("^<?\"?([^<>\"]*@[^<>\"]*[.]?[^<>\"]*)\"?>?$")
+
+    /**
+     * Takes a 'Name (Comment) <Email>' user ID in any of its permutations and attempts to extract
+     * an email from it.
+     */
+    private fun splitUserId(userId: String): String? {
+      if (userId.isNotEmpty()) {
+        val matcher = USER_ID_PATTERN.matcher(userId)
+        if (matcher.matches()) {
+          var name = if (matcher.group(1)?.isEmpty() == true) null else matcher.group(1)
+          var email = matcher.group(3)
+          if (email != null && name != null) {
+            val emailMatcher = EMAIL_PATTERN.matcher(name)
+            if (emailMatcher.matches() && email == emailMatcher.group(1)) {
+              email = emailMatcher.group(1)
+              name = null
+            }
+          }
+          if (email == null && name != null) {
+            val emailMatcher = EMAIL_PATTERN.matcher(name)
+            if (emailMatcher.matches()) {
+              email = emailMatcher.group(1)
+            }
+          }
+          return email
+        }
+      }
+      return null
+    }
+  }
+}
diff --git a/crypto-pgpainless/src/main/kotlin/dev/msfjarvis/aps/crypto/KeyUtils.kt b/crypto-pgpainless/src/main/kotlin/dev/msfjarvis/aps/crypto/KeyUtils.kt
index e7a3c387..852ed628 100644
--- a/crypto-pgpainless/src/main/kotlin/dev/msfjarvis/aps/crypto/KeyUtils.kt
+++ b/crypto-pgpainless/src/main/kotlin/dev/msfjarvis/aps/crypto/KeyUtils.kt
@@ -7,17 +7,18 @@ package dev.msfjarvis.aps.crypto
 
 import com.github.michaelbull.result.get
 import com.github.michaelbull.result.runCatching
+import dev.msfjarvis.aps.crypto.GpgIdentifier.KeyId
 import java.util.Locale
 import org.bouncycastle.openpgp.PGPKeyRing
 import org.pgpainless.PGPainless
 
-/** Utility methods to deal with PGP [Key]s. */
+/** Utility methods to deal with [PGPKey]s. */
 public object KeyUtils {
   /**
    * 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.
    */
-  public fun tryParseKeyring(key: Key): PGPKeyRing? {
+  public fun tryParseKeyring(key: PGPKey): PGPKeyRing? {
     val secKeyRing = runCatching { PGPainless.readKeyRing().secretKeyRing(key.contents) }.get()
     if (secKeyRing != null) {
       return secKeyRing
@@ -29,15 +30,15 @@ public object KeyUtils {
     return null
   }
 
-  /** Parses a [PGPKeyRing] from the given [key] and returns its hex-formatted key ID. */
-  public fun tryGetId(key: Key): String? {
+  /** Parses a [PGPKeyRing] from the given [key] and calculates its long key ID */
+  public fun tryGetId(key: PGPKey): KeyId? {
     val keyRing = tryParseKeyring(key) ?: return null
-    return convertKeyIdToHex(keyRing.publicKey.keyID)
+    return KeyId(convertKeyIdToHex(keyRing.publicKey.keyID).toLong(radix = 16))
   }
 
   /** Convert a [Long] key ID to a formatted string. */
   private fun convertKeyIdToHex(keyId: Long): String {
-    return "0x" + convertKeyIdToHex32bit(keyId shr 32) + convertKeyIdToHex32bit(keyId)
+    return convertKeyIdToHex32bit(keyId shr 32) + convertKeyIdToHex32bit(keyId)
   }
 
   /**
diff --git a/crypto-pgpainless/src/main/kotlin/dev/msfjarvis/aps/crypto/PGPKey.kt b/crypto-pgpainless/src/main/kotlin/dev/msfjarvis/aps/crypto/PGPKey.kt
new file mode 100644
index 00000000..2e1311a1
--- /dev/null
+++ b/crypto-pgpainless/src/main/kotlin/dev/msfjarvis/aps/crypto/PGPKey.kt
@@ -0,0 +1,12 @@
+/*
+ * Copyright © 2014-2021 The Android Password Store Authors. All Rights Reserved.
+ * SPDX-License-Identifier: GPL-3.0-only
+ */
+
+package dev.msfjarvis.aps.crypto
+
+/**
+ * A simple value class wrapping over a [ByteArray] that can represent a PGP key. The implementation
+ * details of public/ private parts as well as identities are deferred to [PGPKeyManager].
+ */
+@JvmInline public value class PGPKey(public val contents: ByteArray)
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 2053ecd7..adcce22b 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
@@ -22,11 +22,11 @@ public class PGPKeyManager
 constructor(
   filesDir: String,
   private val dispatcher: CoroutineDispatcher,
-) : KeyManager {
+) : KeyManager<PGPKey, GpgIdentifier> {
 
   private val keyDir = File(filesDir, KEY_DIR_NAME)
 
-  override suspend fun addKey(key: Key, replace: Boolean): Result<Key, Throwable> =
+  override suspend fun addKey(key: PGPKey, replace: Boolean): Result<PGPKey, Throwable> =
     withContext(dispatcher) {
       runCatching {
         if (!keyDirExists()) throw KeyManagerException.KeyDirectoryUnavailableException
@@ -36,7 +36,7 @@ constructor(
           // Check for replace flag first and if it is false, throw an error
           if (!replace)
             throw KeyManagerException.KeyAlreadyExistsException(
-              tryGetId(key) ?: "Failed to retrieve key ID"
+              tryGetId(key)?.toString() ?: "Failed to retrieve key ID"
             )
           if (!keyFile.delete()) throw KeyManagerException.KeyDeletionFailedException
         }
@@ -47,7 +47,7 @@ constructor(
       }
     }
 
-  override suspend fun removeKey(key: Key): Result<Key, Throwable> =
+  override suspend fun removeKey(key: PGPKey): Result<PGPKey, Throwable> =
     withContext(dispatcher) {
       runCatching {
         if (!keyDirExists()) throw KeyManagerException.KeyDirectoryUnavailableException
@@ -61,49 +61,52 @@ constructor(
       }
     }
 
-  override suspend fun getKeyById(id: String): Result<Key, Throwable> =
+  override suspend fun getKeyById(id: GpgIdentifier): Result<PGPKey, Throwable> =
     withContext(dispatcher) {
       runCatching {
         if (!keyDirExists()) throw KeyManagerException.KeyDirectoryUnavailableException
         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"
+        val keys = keyFiles.map { file -> PGPKey(file.readBytes()) }
+
+        val matchResult =
+          when (id) {
+            is GpgIdentifier.KeyId -> {
+              val keyIdMatch =
+                keys.map { key -> key to tryGetId(key) }.firstOrNull { (_, keyId) ->
+                  keyId?.id == id.id
+                }
+              keyIdMatch?.first
+            }
+            is GpgIdentifier.UserId -> {
+              val selector = SelectUserId.byEmail(id.email)
+              val userIdMatch =
+                keys.map { key -> key to tryParseKeyring(key) }.firstOrNull { (_, keyRing) ->
+                  selector.firstMatch(keyRing) != null
+                }
+              userIdMatch?.first
+            }
           }
 
-        if (keyIdMatch != null) {
-          return@runCatching keyIdMatch.first
+        if (matchResult != null) {
+          return@runCatching matchResult
         }
 
-        throw KeyManagerException.KeyNotFoundException(id)
+        throw KeyManagerException.KeyNotFoundException("$id")
       }
     }
 
-  override suspend fun getAllKeys(): Result<List<Key>, Throwable> =
+  override suspend fun getAllKeys(): Result<List<PGPKey>, Throwable> =
     withContext(dispatcher) {
       runCatching {
         if (!keyDirExists()) throw KeyManagerException.KeyDirectoryUnavailableException
         val keyFiles = keyDir.listFiles()
         if (keyFiles.isNullOrEmpty()) return@runCatching emptyList()
-        keyFiles.map { keyFile -> Key(keyFile.readBytes()) }.toList()
+        keyFiles.map { keyFile -> PGPKey(keyFile.readBytes()) }.toList()
       }
     }
 
-  override suspend fun getKeyId(key: Key): String? = tryGetId(key)
+  override suspend fun getKeyId(key: PGPKey): GpgIdentifier? = tryGetId(key)
 
   // TODO: This is a temp hack for now and in future it should check that the GPGKeyManager can
   // decrypt the file.
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 427cd555..637c8586 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
@@ -18,10 +18,10 @@ import org.pgpainless.encryption_signing.ProducerOptions
 import org.pgpainless.key.protection.PasswordBasedSecretKeyRingProtector
 import org.pgpainless.util.Passphrase
 
-public class PGPainlessCryptoHandler @Inject constructor() : CryptoHandler {
+public class PGPainlessCryptoHandler @Inject constructor() : CryptoHandler<PGPKey> {
 
   public override fun decrypt(
-    privateKey: Key,
+    privateKey: PGPKey,
     passphrase: String,
     ciphertextStream: InputStream,
     outputStream: OutputStream,
@@ -44,7 +44,7 @@ public class PGPainlessCryptoHandler @Inject constructor() : CryptoHandler {
   }
 
   public override fun encrypt(
-    keys: List<Key>,
+    keys: List<PGPKey>,
     plaintextStream: InputStream,
     outputStream: OutputStream,
   ) {
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 6d5eaf01..ede8f08d 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 = "0x08edf7567183ce27"
+  const val KEY_ID = 0x08edf7567183ce27
 }
diff --git a/crypto-pgpainless/src/test/kotlin/dev/msfjarvis/aps/crypto/GpgIdentifierTest.kt b/crypto-pgpainless/src/test/kotlin/dev/msfjarvis/aps/crypto/GpgIdentifierTest.kt
new file mode 100644
index 00000000..3860873c
--- /dev/null
+++ b/crypto-pgpainless/src/test/kotlin/dev/msfjarvis/aps/crypto/GpgIdentifierTest.kt
@@ -0,0 +1,41 @@
+/*
+ * Copyright © 2014-2022 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.assertNotNull
+import kotlin.test.assertTrue
+
+class GpgIdentifierTest {
+
+  @Test
+  fun `parses hexadecimal key id without leading 0x`() {
+    val identifier = GpgIdentifier.fromString("79E8208280490C77")
+    assertNotNull(identifier)
+    assertTrue { identifier is GpgIdentifier.KeyId }
+  }
+
+  @Test
+  fun `parses hexadecimal key id`() {
+    val identifier = GpgIdentifier.fromString("0x79E8208280490C77")
+    assertNotNull(identifier)
+    assertTrue { identifier is GpgIdentifier.KeyId }
+  }
+
+  @Test
+  fun `parses email as user id`() {
+    val identifier = GpgIdentifier.fromString("john.doe@example.org")
+    assertNotNull(identifier)
+    assertTrue { identifier is GpgIdentifier.UserId }
+  }
+
+  @Test
+  fun `parses user@host without TLD`() {
+    val identifier = GpgIdentifier.fromString("john.doe@example")
+    assertNotNull(identifier)
+    assertTrue { identifier is GpgIdentifier.UserId }
+  }
+}
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 ed2b00c3..24b27ef8 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
@@ -2,6 +2,8 @@ package dev.msfjarvis.aps.crypto
 
 import com.github.michaelbull.result.unwrap
 import com.github.michaelbull.result.unwrapError
+import dev.msfjarvis.aps.crypto.GpgIdentifier.KeyId
+import dev.msfjarvis.aps.crypto.GpgIdentifier.UserId
 import java.io.File
 import kotlin.test.AfterTest
 import kotlin.test.BeforeTest
@@ -28,7 +30,7 @@ class PGPKeyManagerTest {
   private val dispatcher = StandardTestDispatcher()
   private val scope = TestScope(dispatcher)
   private val keyManager by unsafeLazy { PGPKeyManager(filesDir.absolutePath, dispatcher) }
-  private val key = Key(TestUtils.getArmoredPrivateKey())
+  private val key = PGPKey(TestUtils.getArmoredPrivateKey())
 
   private fun <T> unsafeLazy(initializer: () -> T) =
     lazy(LazyThreadSafetyMode.NONE) { initializer.invoke() }
@@ -48,7 +50,7 @@ class PGPKeyManagerTest {
     scope.runTest {
       // Check if the key id returned is correct
       val keyId = keyManager.getKeyId(keyManager.addKey(key).unwrap())
-      assertEquals(CryptoConstants.KEY_ID, keyId)
+      assertEquals(KeyId(CryptoConstants.KEY_ID), keyId)
 
       // Check if the keys directory have one file
       assertEquals(1, filesDir.list()?.size)
@@ -75,7 +77,7 @@ class PGPKeyManagerTest {
       keyManager.addKey(key, true).unwrap()
       val keyId = keyManager.getKeyId(keyManager.addKey(key, true).unwrap())
 
-      assertEquals(CryptoConstants.KEY_ID, keyId)
+      assertEquals(KeyId(CryptoConstants.KEY_ID), keyId)
     }
 
   @Test
@@ -86,7 +88,7 @@ class PGPKeyManagerTest {
 
       // Check if the key id returned is correct
       val keyId = keyManager.getKeyId(keyManager.removeKey(key).unwrap())
-      assertEquals(CryptoConstants.KEY_ID, keyId)
+      assertEquals(KeyId(CryptoConstants.KEY_ID), keyId)
 
       // Check if the keys directory have 0 files
       val keysDir = File(filesDir, PGPKeyManager.KEY_DIR_NAME)
@@ -101,7 +103,7 @@ class PGPKeyManagerTest {
 
       val keyId = keyManager.getKeyId(key)
       assertNotNull(keyId)
-      assertEquals(CryptoConstants.KEY_ID, keyManager.getKeyId(key))
+      assertEquals(KeyId(CryptoConstants.KEY_ID), keyManager.getKeyId(key))
 
       // Check returned key id matches the expected id and the created key id
       val returnedKey = keyManager.getKeyById(keyId).unwrap()
@@ -114,7 +116,7 @@ class PGPKeyManagerTest {
       keyManager.addKey(key).unwrap()
 
       val keyId = "${CryptoConstants.KEY_NAME} <${CryptoConstants.KEY_EMAIL}>"
-      val returnedKey = keyManager.getKeyById(keyId).unwrap()
+      val returnedKey = keyManager.getKeyById(UserId(keyId)).unwrap()
       assertEquals(keyManager.getKeyId(key), keyManager.getKeyId(returnedKey))
     }
 
@@ -124,7 +126,7 @@ class PGPKeyManagerTest {
       keyManager.addKey(key).unwrap()
 
       val keyId = CryptoConstants.KEY_EMAIL
-      val returnedKey = keyManager.getKeyById(keyId).unwrap()
+      val returnedKey = keyManager.getKeyById(UserId(keyId)).unwrap()
       assertEquals(keyManager.getKeyId(key), keyManager.getKeyId(returnedKey))
     }
 
@@ -134,19 +136,19 @@ class PGPKeyManagerTest {
       // Add key using KeyManager
       keyManager.addKey(key).unwrap()
 
-      val randomKeyId = "0x123456789"
+      val keyId = KeyId(0x08edf7567183ce44)
 
       // Check returned key
-      val error = keyManager.getKeyById(randomKeyId).unwrapError()
+      val error = keyManager.getKeyById(keyId).unwrapError()
       assertIs<KeyManagerException.KeyNotFoundException>(error)
-      assertEquals("No key found with id: $randomKeyId", error.message)
+      assertEquals("No key found with id: $keyId", error.message)
     }
 
   @Test
   fun testFindKeysWithoutAdding() =
     scope.runTest {
       // Check returned key
-      val error = keyManager.getKeyById("0x123456789").unwrapError()
+      val error = keyManager.getKeyById(KeyId(0x08edf7567183ce44)).unwrapError()
       assertIs<KeyManagerException.NoKeysAvailableException>(error)
       assertEquals("No keys were found", error.message)
     }
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 6ec6ccee..9b4cb664 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
@@ -14,8 +14,8 @@ import kotlin.test.assertTrue
 class PGPainlessCryptoHandlerTest {
 
   private val cryptoHandler = PGPainlessCryptoHandler()
-  private val privateKey = Key(TestUtils.getArmoredPrivateKey())
-  private val publicKey = Key(TestUtils.getArmoredPublicKey())
+  private val privateKey = PGPKey(TestUtils.getArmoredPrivateKey())
+  private val publicKey = PGPKey(TestUtils.getArmoredPublicKey())
 
   @Test
   fun encryptAndDecrypt() {