PM-38625: Chore: Store the WrappedAccountCryptographicState (#7030)

app/src/main/kotlin/com/x8bit/bitwarden/data/auth/datasource/disk/AuthDiskSource.kt

@@ -1,6 +1,6 @@
 package com.x8bit.bitwarden.data.auth.datasource.disk
 
-import com.bitwarden.network.model.AccountKeysJson
+import com.bitwarden.core.WrappedAccountCryptographicState
 import com.bitwarden.network.model.SyncResponseJson
 import com.bitwarden.network.provider.AppIdProvider
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.AccountTokensJson
@@ -125,34 +125,16 @@ interface AuthDiskSource : AppIdProvider {
     fun storeLocalUserDataKey(userId: String, wrappedKey: String?)
 
     /**
-     * Retrieves a private key using a [userId].
+     * Returns the Wrapped Account Cryptographic State for the given [userId].
      */
-    @Deprecated(
-        message = "Use getAccountKeys instead.",
-        replaceWith = ReplaceWith("getAccountKeys"),
-    )
-    fun getPrivateKey(userId: String): String?
+    fun getAccountCryptographicState(userId: String): WrappedAccountCryptographicState?
 
     /**
-     * Stores a private key using a [userId].
+     * Stores the Wrapped Account Cryptographic State for a given [userId].
      */
-    @Deprecated(
-        message = "Use storeAccountKeys instead.",
-        replaceWith = ReplaceWith("storeAccountKeys"),
-    )
-    fun storePrivateKey(userId: String, privateKey: String?)
-
-    /**
-     * Returns the profile account keys for the given [userId].
-     */
-    fun getAccountKeys(userId: String): AccountKeysJson?
-
-    /**
-     * Stores the profile account keys for the given [userId].
-     */
-    fun storeAccountKeys(
+    fun storeAccountCryptographicState(
         userId: String,
-        accountKeys: AccountKeysJson?,
+        accountCryptographicState: WrappedAccountCryptographicState?,
     )
 
     /**

app/src/main/kotlin/com/x8bit/bitwarden/data/auth/datasource/disk/AuthDiskSourceImpl.kt

@@ -1,6 +1,7 @@
 package com.x8bit.bitwarden.data.auth.datasource.disk
 
 import android.content.SharedPreferences
+import com.bitwarden.core.WrappedAccountCryptographicState
 import com.bitwarden.core.data.repository.util.bufferedMutableSharedFlow
 import com.bitwarden.core.data.serializer.SafeMapSerializer
 import com.bitwarden.core.data.util.decodeFromStringOrNull
@@ -11,6 +12,8 @@ import com.x8bit.bitwarden.data.auth.datasource.disk.model.AccountTokensJson
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.OnboardingStatus
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.PendingAuthRequestJson
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.UserStateJson
+import com.x8bit.bitwarden.data.auth.datasource.disk.serializer.WrappedAccountCryptographicStateSerializer
+import com.x8bit.bitwarden.data.auth.repository.util.toAccountCryptographicState
 import com.x8bit.bitwarden.data.platform.datasource.disk.legacy.LegacySecureStorageMigrator
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableSharedFlow
@@ -54,6 +57,7 @@ private const val ONBOARDING_STATUS_KEY = "onboardingStatus"
 private const val SHOW_IMPORT_LOGINS_KEY = "showImportLogins"
 private const val LAST_LOCK_TIMESTAMP = "lastLockTimestamp"
 private const val PROFILE_ACCOUNT_KEYS_KEY = "profileAccountKeys"
+private const val ACCOUNT_CRYPTOGRAPHIC_STATE_KEY = "accountCryptographicState"
 
 /**
  * Primary implementation of [AuthDiskSource].
@@ -91,6 +95,10 @@ class AuthDiskSourceImpl(
         mutableMapOf<String, MutableSharedFlow<String?>>()
     private val mutableUserStateFlow = bufferedMutableSharedFlow<UserStateJson?>(replay = 1)
 
+    private val wrappedAccountCryptographicStateSerializer by lazy {
+        WrappedAccountCryptographicStateSerializer()
+    }
+
     override var userState: UserStateJson?
         get() = getString(key = STATE_KEY)?.let { json.decodeFromStringOrNull(it) }
         set(value) {
@@ -113,6 +121,10 @@ class AuthDiskSourceImpl(
         // We want to make sure that any left over encrypted user keys are scrubbed from storage
         // Since it is no longer supported.
         removeLegacyUserKeys()
+
+        // We must migrate the Private Key and Account Keys to use the Account Cryptographic state
+        // from now on.
+        migrateAccountKeys()
     }
 
     override var authenticatorSyncSymmetricKey: ByteArray?
@@ -152,8 +164,7 @@ class AuthDiskSourceImpl(
         storeInvalidUnlockAttempts(userId = userId, invalidUnlockAttempts = null)
         storeLocalUserDataKey(userId = userId, wrappedKey = null)
         storeUserAutoUnlockKey(userId = userId, userAutoUnlockKey = null)
-        storePrivateKey(userId = userId, privateKey = null)
-        storeAccountKeys(userId = userId, accountKeys = null)
+        storeAccountCryptographicState(userId = userId, accountCryptographicState = null)
         storeOrganizationKeys(userId = userId, organizationKeys = null)
         storeOrganizations(userId = userId, organizations = null)
         storeUserBiometricInitVector(userId = userId, iv = null)
@@ -241,29 +252,20 @@ class AuthDiskSourceImpl(
         putString(key = LOCAL_USER_DATA_KEY.appendIdentifier(userId), value = wrappedKey)
     }
 
-    @Deprecated("Use getAccountKeys instead.", replaceWith = ReplaceWith("getAccountKeys"))
-    override fun getPrivateKey(userId: String): String? =
-        getString(key = MASTER_KEY_ENCRYPTION_PRIVATE_KEY.appendIdentifier(userId))
+    override fun getAccountCryptographicState(userId: String): WrappedAccountCryptographicState? =
+        getEncryptedString(key = ACCOUNT_CRYPTOGRAPHIC_STATE_KEY.appendIdentifier(userId))?.let {
+            json.decodeFromStringOrNull(wrappedAccountCryptographicStateSerializer, it)
+        }
 
-    @Deprecated("Use storeAccountKeys instead.", replaceWith = ReplaceWith("storeAccountKeys"))
-    override fun storePrivateKey(userId: String, privateKey: String?) {
-        putString(
-            key = MASTER_KEY_ENCRYPTION_PRIVATE_KEY.appendIdentifier(userId),
-            value = privateKey,
-        )
-    }
-
-    override fun getAccountKeys(userId: String): AccountKeysJson? =
-        getEncryptedString(key = PROFILE_ACCOUNT_KEYS_KEY.appendIdentifier(userId))
-            ?.let { json.decodeFromStringOrNull(it) }
-
-    override fun storeAccountKeys(
+    override fun storeAccountCryptographicState(
         userId: String,
-        accountKeys: AccountKeysJson?,
+        accountCryptographicState: WrappedAccountCryptographicState?,
     ) {
         putEncryptedString(
-            key = PROFILE_ACCOUNT_KEYS_KEY.appendIdentifier(userId),
-            value = accountKeys?.let { json.encodeToString(it) },
+            key = ACCOUNT_CRYPTOGRAPHIC_STATE_KEY.appendIdentifier(userId),
+            value = accountCryptographicState?.let {
+                json.encodeToString(wrappedAccountCryptographicStateSerializer, it)
+            },
         )
     }
 
@@ -659,4 +661,31 @@ class AuthDiskSourceImpl(
     private fun removeLegacyUserKeys() {
         removeWithPrefix(prefix = MASTER_KEY_ENCRYPTION_USER_KEY)
     }
+
+    private fun migrateAccountKeys() {
+        userState
+            ?.accounts
+            .orEmpty()
+            .values
+            .forEach { account ->
+                val userId = account.profile.userId
+                val accountKeysKey = PROFILE_ACCOUNT_KEYS_KEY.appendIdentifier(userId)
+                val privateKeyKey = MASTER_KEY_ENCRYPTION_PRIVATE_KEY.appendIdentifier(userId)
+                val accountKeys = getEncryptedString(key = accountKeysKey)
+                    ?.let { json.decodeFromStringOrNull<AccountKeysJson>(it) }
+                val privateKey = accountKeys
+                    ?.publicKeyEncryptionKeyPair
+                    ?.wrappedPrivateKey
+                    ?: getString(key = privateKeyKey)
+                privateKey?.let {
+                    storeAccountCryptographicState(
+                        userId = userId,
+                        accountCryptographicState = accountKeys.toAccountCryptographicState(it),
+                    )
+                    // Remove the Account Keys and Private Key
+                    putEncryptedString(key = accountKeysKey, value = null)
+                    putString(key = privateKeyKey, value = null)
+                }
+            }
+    }
 }

app/src/main/kotlin/com/x8bit/bitwarden/data/auth/datasource/disk/serializer/WrappedAccountCryptographicStateSerializer.kt

@@ -0,0 +1,85 @@
+package com.x8bit.bitwarden.data.auth.datasource.disk.serializer
+
+import com.bitwarden.core.WrappedAccountCryptographicState
+import kotlinx.serialization.KSerializer
+import kotlinx.serialization.SerialName
+import kotlinx.serialization.Serializable
+import kotlinx.serialization.descriptors.SerialDescriptor
+import kotlinx.serialization.encoding.Decoder
+import kotlinx.serialization.encoding.Encoder
+
+/**
+ * Custom [KSerializer] for [WrappedAccountCryptographicState].
+ *
+ * Encodes the sealed class with a `"type"` discriminator field:
+ * - `"v1"`: [WrappedAccountCryptographicState.V1] — wrapped private key only.
+ * - `"v2"`: [WrappedAccountCryptographicState.V2] — wrapped private key, signing key, signed
+ *   public key, and signed security state.
+ */
+internal class WrappedAccountCryptographicStateSerializer :
+    KSerializer<WrappedAccountCryptographicState> {
+
+    private val surrogateSerializer = Surrogate.serializer()
+
+    override val descriptor: SerialDescriptor = surrogateSerializer.descriptor
+
+    override fun deserialize(decoder: Decoder): WrappedAccountCryptographicState =
+        when (val surrogate = decoder.decodeSerializableValue(surrogateSerializer)) {
+            is Surrogate.V1 -> {
+                WrappedAccountCryptographicState.V1(privateKey = surrogate.privateKey)
+            }
+
+            is Surrogate.V2 -> {
+                WrappedAccountCryptographicState.V2(
+                    privateKey = surrogate.privateKey,
+                    signingKey = surrogate.signingKey,
+                    signedPublicKey = surrogate.signedPublicKey,
+                    securityState = surrogate.securityState,
+                )
+            }
+        }
+
+    override fun serialize(encoder: Encoder, value: WrappedAccountCryptographicState) {
+        val surrogate = when (value) {
+            is WrappedAccountCryptographicState.V1 -> {
+                Surrogate.V1(privateKey = value.privateKey)
+            }
+
+            is WrappedAccountCryptographicState.V2 -> {
+                Surrogate.V2(
+                    privateKey = value.privateKey,
+                    signingKey = value.signingKey,
+                    signedPublicKey = value.signedPublicKey,
+                    securityState = value.securityState,
+                )
+            }
+        }
+        encoder.encodeSerializableValue(surrogateSerializer, surrogate)
+    }
+}
+
+@Serializable
+private sealed class Surrogate {
+    @Serializable
+    @SerialName("v1")
+    data class V1(
+        @SerialName("privateKey")
+        val privateKey: String,
+    ) : Surrogate()
+
+    @Serializable
+    @SerialName("v2")
+    data class V2(
+        @SerialName("privateKey")
+        val privateKey: String,
+
+        @SerialName("signingKey")
+        val signingKey: String,
+
+        @SerialName("signedPublicKey")
+        val signedPublicKey: String?,
+
+        @SerialName("securityState")
+        val securityState: String,
+    ) : Surrogate()
+}

app/src/main/kotlin/com/x8bit/bitwarden/data/auth/repository/AuthRepositoryImpl.kt

@@ -104,10 +104,8 @@ import com.x8bit.bitwarden.data.auth.repository.util.CookieCallbackResult
 import com.x8bit.bitwarden.data.auth.repository.util.DuoCallbackTokenResult
 import com.x8bit.bitwarden.data.auth.repository.util.SsoCallbackResult
 import com.x8bit.bitwarden.data.auth.repository.util.WebAuthResult
-import com.x8bit.bitwarden.data.auth.repository.util.accountKeysJson
 import com.x8bit.bitwarden.data.auth.repository.util.activeUserIdChangesFlow
 import com.x8bit.bitwarden.data.auth.repository.util.policyInformation
-import com.x8bit.bitwarden.data.auth.repository.util.privateKey
 import com.x8bit.bitwarden.data.auth.repository.util.toAccountCryptographicState
 import com.x8bit.bitwarden.data.auth.repository.util.toDeviceInfo
 import com.x8bit.bitwarden.data.auth.repository.util.toOrganizations
@@ -535,17 +533,13 @@ class AuthRepositoryImpl(
                     .map { registerTdeKeyResponse to createAccountKeysResponse }
             }
             .onSuccess { (registerTdeKeyResponse, createAccountKeysResponse) ->
-                authDiskSource.storeAccountKeys(
+                authDiskSource.storeAccountCryptographicState(
                     userId = userId,
-                    accountKeys = createAccountKeysResponse.accountKeys,
-                )
-                // TDE and SSO user creation still uses crypto-v1. These users are not expected to
-                // have the AEAD keys so we only store the private key for now.
-                // See https://github.com/bitwarden/android/pull/5682#discussion_r2273940332
-                // for more details.
-                authDiskSource.storePrivateKey(
-                    userId = userId,
-                    privateKey = registerTdeKeyResponse.privateKey,
+                    accountCryptographicState = createAccountKeysResponse
+                        .accountKeys
+                        .toAccountCryptographicState(
+                            privateKey = registerTdeKeyResponse.privateKey,
+                        ),
                 )
                 vaultRepository.syncVaultState(userId = userId)
                 registerTdeKeyResponse.deviceKey?.let { response ->
@@ -585,16 +579,9 @@ class AuthRepositoryImpl(
                     )
                     .also { result ->
                         if (result is VaultUnlockResult.Success) {
-                            authDiskSource.storeAccountKeys(
+                            authDiskSource.storeAccountCryptographicState(
                                 userId = userId,
-                                accountKeys = response.accountCryptographicState.accountKeysJson,
-                            )
-
-                            // Storing the private key here for legacy purposes, the
-                            // `accountKeysJson` stored above will be used for most purposes.
-                            authDiskSource.storePrivateKey(
-                                userId = userId,
-                                privateKey = response.accountCryptographicState.privateKey,
+                                accountCryptographicState = response.accountCryptographicState,
                             )
                             if (shouldTrustDevice) {
                                 authDiskSource.storeDeviceKey(
@@ -614,20 +601,16 @@ class AuthRepositoryImpl(
         val profile = authDiskSource.userState?.activeAccount?.profile
             ?: return LoginResult.Error(error = NoActiveUserException())
         val userId = profile.userId
-        val accountKeys = authDiskSource.getAccountKeys(userId = userId)
-        val privateKey = accountKeys?.publicKeyEncryptionKeyPair?.wrappedPrivateKey
-            ?: authDiskSource.getPrivateKey(userId = userId)
-            ?: return LoginResult.Error(error = MissingPropertyException("Private Key"))
-
+        val accountCryptographicState = authDiskSource
+            .getAccountCryptographicState(userId = userId)
+            ?: return LoginResult.Error(MissingPropertyException("Account Cryptographic State"))
         checkForVaultUnlockError(
             onVaultUnlockError = { error ->
                 return error.toLoginErrorResult()
             },
         ) {
             unlockVault(
-                accountCryptographicState = accountKeys.toAccountCryptographicState(
-                    privateKey = privateKey,
-                ),
+                accountCryptographicState = accountCryptographicState,
                 accountProfile = profile,
                 initUserCryptoMethod = InitUserCryptoMethod.AuthRequest(
                     requestPrivateKey = requestPrivateKey,
@@ -1279,17 +1262,9 @@ class AuthRepositoryImpl(
                 )
             }
             .onSuccess { response ->
-                authDiskSource.storeAccountKeys(
+                authDiskSource.storeAccountCryptographicState(
                     userId = userId,
-                    accountKeys = response.accountCryptographicState.accountKeysJson,
-                )
-                // TDE and SSO user creation still uses crypto-v1. These users are not
-                // expected to have the AEAD keys so we only store the private key for now.
-                // See https://github.com/bitwarden/android/pull/5682#discussion_r2273940332
-                // for more details.
-                authDiskSource.storePrivateKey(
-                    userId = userId,
-                    privateKey = response.accountCryptographicState.privateKey,
+                    accountCryptographicState = response.accountCryptographicState,
                 )
                 authDiskSource.userState = authDiskSource.userState?.toUserStateJsonWithPassword(
                     masterPasswordUnlock = response.masterPasswordUnlock,
@@ -1346,12 +1321,11 @@ class AuthRepositoryImpl(
                         ),
                     )
                     .onSuccess {
-                        // This process is used by TDE and Enterprise accounts during initial
-                        // login. We continue to store the locally generated keys
-                        // until TDE and Enterprise accounts support AEAD keys.
-                        authDiskSource.storePrivateKey(
+                        authDiskSource.storeAccountCryptographicState(
                             userId = userId,
-                            privateKey = response.keys.private,
+                            accountCryptographicState = WrappedAccountCryptographicState.V1(
+                                privateKey = response.keys.private,
+                            ),
                         )
                     }
                     .map { response.masterPasswordHash }
@@ -1963,17 +1937,16 @@ class AuthRepositoryImpl(
                     }
                 }
         }
-        // We continue to store the private key for backwards compatibility. Key connector
-        // conversion still relies on the private key.
-        loginResponse.privateKeyOrNull()?.let {
-            // Only set the value if it's present, since we may have set it already
-            // when we completed the key connector conversion.
-            authDiskSource.storePrivateKey(userId = userId, privateKey = it)
-        }
-        loginResponse.accountKeys?.let {
-            // Only set the value if it's present, since we may have set it already
-            // when we completed the key connector conversion.
-            authDiskSource.storeAccountKeys(userId = userId, accountKeys = it)
+
+        loginResponse.privateKeyOrNull()?.let { privateKey ->
+            // Only set the value if the private key is present, since we may have set
+            // the value already when we completed the key connector conversion.
+            authDiskSource.storeAccountCryptographicState(
+                userId = userId,
+                accountCryptographicState = loginResponse.accountKeys.toAccountCryptographicState(
+                    privateKey = privateKey,
+                ),
+            )
         }
         // If the user just authenticated with a two-factor code and selected the option to
         // remember it, then the API response will return a token that will be used in place
@@ -2088,9 +2061,10 @@ class AuthRepositoryImpl(
                     organizationIdentifier = orgIdentifier,
                 )
                 .map { keyConnector ->
+                    val accountCryptographicState = keyConnector.accountCryptographicState
                     this
                         .unlockVault(
-                            accountCryptographicState = keyConnector.accountCryptographicState,
+                            accountCryptographicState = accountCryptographicState,
                             accountProfile = profile,
                             initUserCryptoMethod = InitUserCryptoMethod.KeyConnector(
                                 masterKey = keyConnector.masterKey,
@@ -2099,15 +2073,9 @@ class AuthRepositoryImpl(
                         )
                         .also { result ->
                             if (result is VaultUnlockResult.Success) {
-                                // We continue to store the private key for backwards compatibility
-                                // since key connector conversion still relies on the private key.
-                                authDiskSource.storePrivateKey(
+                                authDiskSource.storeAccountCryptographicState(
                                     userId = userId,
-                                    privateKey = keyConnector.privateKey,
-                                )
-                                authDiskSource.storeAccountKeys(
-                                    userId = userId,
-                                    accountKeys = loginResponse.accountKeys,
+                                    accountCryptographicState = accountCryptographicState,
                                 )
                             }
                         }

app/src/main/kotlin/com/x8bit/bitwarden/data/auth/repository/util/AccountKeysJsonExtensions.kt

@@ -2,7 +2,6 @@ package com.x8bit.bitwarden.data.auth.repository.util
 
 import com.bitwarden.core.WrappedAccountCryptographicState
 import com.bitwarden.network.model.AccountKeysJson
-import com.x8bit.bitwarden.data.vault.repository.util.createWrappedAccountCryptographicState
 
 /**
  * Creates a [WrappedAccountCryptographicState] based on the available cryptographic parameters.
@@ -15,9 +14,20 @@ import com.x8bit.bitwarden.data.vault.repository.util.createWrappedAccountCrypto
  */
 fun AccountKeysJson?.toAccountCryptographicState(
     privateKey: String,
-): WrappedAccountCryptographicState = createWrappedAccountCryptographicState(
-    privateKey = privateKey,
-    securityState = this?.securityState?.securityState,
-    signingKey = this?.signatureKeyPair?.wrappedSigningKey,
-    signedPublicKey = this?.publicKeyEncryptionKeyPair?.signedPublicKey,
-)
+): WrappedAccountCryptographicState {
+    val securityState = this?.securityState?.securityState
+    val signingKey = this?.signatureKeyPair?.wrappedSigningKey
+    val signedPublicKey = this?.publicKeyEncryptionKeyPair?.signedPublicKey
+    return if (signingKey != null && securityState != null && signedPublicKey != null) {
+        WrappedAccountCryptographicState.V2(
+            privateKey = privateKey,
+            securityState = securityState,
+            signingKey = signingKey,
+            signedPublicKey = signedPublicKey,
+        )
+    } else {
+        WrappedAccountCryptographicState.V1(
+            privateKey = privateKey,
+        )
+    }
+}

app/src/main/kotlin/com/x8bit/bitwarden/data/platform/repository/AuthenticatorBridgeRepositoryImpl.kt

@@ -10,7 +10,6 @@ import com.bitwarden.core.data.util.flatMap
 import com.bitwarden.data.repository.util.toEnvironmentUrlsOrDefault
 import com.x8bit.bitwarden.data.auth.datasource.disk.AuthDiskSource
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.AccountJson
-import com.x8bit.bitwarden.data.auth.repository.util.toAccountCryptographicState
 import com.x8bit.bitwarden.data.auth.repository.util.toSdkParams
 import com.x8bit.bitwarden.data.vault.datasource.disk.VaultDiskSource
 import com.x8bit.bitwarden.data.vault.datasource.sdk.ScopedVaultSdkSource
@@ -135,20 +134,17 @@ class AuthenticatorBridgeRepositoryImpl(
         account: AccountJson,
         decryptedUserKey: String,
     ): VaultUnlockResult {
-        val accountKeys = authDiskSource.getAccountKeys(userId = userId)
-        val privateKey = accountKeys?.publicKeyEncryptionKeyPair?.wrappedPrivateKey
-            ?: authDiskSource.getPrivateKey(userId = userId)
+        val accountCryptographicState = authDiskSource
+            .getAccountCryptographicState(userId = userId)
             ?: return VaultUnlockResult.InvalidStateError(
-                MissingPropertyException("Private key"),
+                error = MissingPropertyException("Account Cryptographic State"),
             )
 
         return scopedVaultSdkSource
             .initializeCrypto(
                 userId = userId,
                 request = InitUserCryptoRequest(
-                    accountCryptographicState = accountKeys.toAccountCryptographicState(
-                        privateKey = privateKey,
-                    ),
+                    accountCryptographicState = accountCryptographicState,
                     userId = userId,
                     kdfParams = account.profile.toSdkParams(),
                     email = account.profile.email,

app/src/main/kotlin/com/x8bit/bitwarden/data/vault/manager/VaultLockManagerImpl.kt

@@ -28,7 +28,6 @@ import com.x8bit.bitwarden.data.auth.manager.UserLogoutManager
 import com.x8bit.bitwarden.data.auth.repository.model.LogoutReason
 import com.x8bit.bitwarden.data.auth.repository.model.UpdateKdfMinimumsResult
 import com.x8bit.bitwarden.data.auth.repository.util.activeUserIdChangesFlow
-import com.x8bit.bitwarden.data.auth.repository.util.toAccountCryptographicState
 import com.x8bit.bitwarden.data.auth.repository.util.toSdkParams
 import com.x8bit.bitwarden.data.auth.repository.util.userAccountTokens
 import com.x8bit.bitwarden.data.auth.repository.util.userSwitchingChangesFlow
@@ -675,17 +674,14 @@ class VaultLockManagerImpl(
     ): VaultUnlockResult {
         val account = authDiskSource.userState?.accounts?.get(userId)
             ?: return VaultUnlockResult.InvalidStateError(error = NoActiveUserException())
-        val accountKeys = authDiskSource.getAccountKeys(userId = userId)
-        val privateKey = accountKeys?.publicKeyEncryptionKeyPair?.wrappedPrivateKey
-            ?: authDiskSource.getPrivateKey(userId = userId)
+        val accountCryptographicState = authDiskSource
+            .getAccountCryptographicState(userId = userId)
             ?: return VaultUnlockResult.InvalidStateError(
-                error = MissingPropertyException("Private key"),
+                error = MissingPropertyException("Account Cryptographic State"),
             )
         val organizationKeys = authDiskSource.getOrganizationKeys(userId = userId)
         return unlockVault(
-            accountCryptographicState = accountKeys.toAccountCryptographicState(
-                privateKey = privateKey,
-            ),
+            accountCryptographicState = accountCryptographicState,
             userId = userId,
             email = account.profile.email,
             kdf = account.profile.toSdkParams(),

app/src/main/kotlin/com/x8bit/bitwarden/data/vault/manager/VaultSyncManagerImpl.kt

@@ -17,6 +17,7 @@ import com.x8bit.bitwarden.data.auth.datasource.disk.AuthDiskSource
 import com.x8bit.bitwarden.data.auth.manager.UserLogoutManager
 import com.x8bit.bitwarden.data.auth.manager.UserStateManager
 import com.x8bit.bitwarden.data.auth.repository.model.LogoutReason
+import com.x8bit.bitwarden.data.auth.repository.util.toAccountCryptographicState
 import com.x8bit.bitwarden.data.auth.repository.util.toUpdatedUserStateJson
 import com.x8bit.bitwarden.data.auth.repository.util.userSwitchingChangesFlow
 import com.x8bit.bitwarden.data.platform.datasource.disk.SettingsDiskSource
@@ -377,8 +378,14 @@ class VaultSyncManagerImpl(
         val profile = syncResponse.profile
         val userId = profile.id
         authDiskSource.apply {
-            storePrivateKey(userId = userId, privateKey = profile.privateKey)
-            storeAccountKeys(userId = userId, accountKeys = profile.accountKeys)
+            storeAccountCryptographicState(
+                userId = userId,
+                accountCryptographicState = profile.privateKeyOrNull()?.let {
+                    profile.accountKeys.toAccountCryptographicState(
+                        privateKey = it,
+                    )
+                },
+            )
             storeOrganizationKeys(
                 userId = userId,
                 organizationKeys = profile.organizations
@@ -539,6 +546,13 @@ class VaultSyncManagerImpl(
             ?: DataState.Loading
 }
 
+/**
+ * Convenience function to extract the private key from the [SyncResponseJson.Profile] response.
+ */
+private fun SyncResponseJson.Profile.privateKeyOrNull(): String? =
+    this.accountKeys?.publicKeyEncryptionKeyPair?.wrappedPrivateKey
+        ?: this.privateKey
+
 private fun <T> Throwable.toNetworkOrErrorState(
     data: T?,
 ): DataState<T> =

app/src/main/kotlin/com/x8bit/bitwarden/data/vault/repository/VaultRepositoryImpl.kt

@@ -19,7 +19,6 @@ import com.bitwarden.vault.CipherType
 import com.bitwarden.vault.CipherView
 import com.bitwarden.vault.FolderView
 import com.x8bit.bitwarden.data.auth.datasource.disk.AuthDiskSource
-import com.x8bit.bitwarden.data.auth.repository.util.toAccountCryptographicState
 import com.x8bit.bitwarden.data.auth.repository.util.toSdkParams
 import com.x8bit.bitwarden.data.autofill.util.login
 import com.x8bit.bitwarden.data.platform.error.NoActiveUserException
@@ -559,19 +558,14 @@ class VaultRepositoryImpl(
     ): VaultUnlockResult {
         val account = authDiskSource.userState?.accounts?.get(userId)
             ?: return VaultUnlockResult.InvalidStateError(error = NoActiveUserException())
-        val accountKeys = authDiskSource.getAccountKeys(userId = userId)
-        val privateKey = accountKeys
-            ?.publicKeyEncryptionKeyPair
-            ?.wrappedPrivateKey
-            ?: authDiskSource.getPrivateKey(userId = userId)
+        val accountCryptographicState = authDiskSource
+            .getAccountCryptographicState(userId = userId)
             ?: return VaultUnlockResult.InvalidStateError(
-                error = MissingPropertyException("Private key"),
+                error = MissingPropertyException("Account Cryptographic State"),
             )
         val organizationKeys = authDiskSource.getOrganizationKeys(userId = userId)
         return vaultLockManager.unlockVault(
-            accountCryptographicState = accountKeys.toAccountCryptographicState(
-                privateKey = privateKey,
-            ),
+            accountCryptographicState = accountCryptographicState,
             userId = userId,
             email = account.profile.email,
             kdf = account.profile.toSdkParams(),

app/src/main/kotlin/com/x8bit/bitwarden/data/vault/repository/util/WrappedAccountCryptographicStateExtensions.kt

@@ -1,34 +0,0 @@
-package com.x8bit.bitwarden.data.vault.repository.util
-
-import com.bitwarden.core.WrappedAccountCryptographicState
-
-/**
- * Creates a [WrappedAccountCryptographicState] based on the available cryptographic parameters.
- *
- * Returns [WrappedAccountCryptographicState.V2] if signing key, signed public key, and security
- * state are all present, otherwise returns [WrappedAccountCryptographicState.V1].
- *
- * @param privateKey The user's wrapped private key.
- * @param securityState The user's signed security state (V2 only).
- * @param signingKey The user's wrapped signing key (V2 only).
- * @param signedPublicKey The user's signed public key (V2 only).
- */
-fun createWrappedAccountCryptographicState(
-    privateKey: String,
-    securityState: String?,
-    signingKey: String?,
-    signedPublicKey: String?,
-): WrappedAccountCryptographicState {
-    return if (signingKey != null && securityState != null && signedPublicKey != null) {
-        WrappedAccountCryptographicState.V2(
-            privateKey = privateKey,
-            securityState = securityState,
-            signingKey = signingKey,
-            signedPublicKey = signedPublicKey,
-        )
-    } else {
-        WrappedAccountCryptographicState.V1(
-            privateKey = privateKey,
-        )
-    }
-}

app/src/test/kotlin/com/x8bit/bitwarden/data/auth/datasource/disk/AuthDiskSourceTest.kt

@@ -3,6 +3,7 @@ package com.x8bit.bitwarden.data.auth.datasource.disk
 import androidx.core.content.edit
 import app.cash.turbine.test
 import com.bitwarden.authenticatorbridge.util.generateSecretKey
+import com.bitwarden.core.WrappedAccountCryptographicState
 import com.bitwarden.core.di.CoreModule
 import com.bitwarden.data.datasource.disk.base.FakeSharedPreferences
 import com.bitwarden.data.datasource.disk.model.EnvironmentUrlDataJson
@@ -10,7 +11,6 @@ import com.bitwarden.network.model.KdfTypeJson
 import com.bitwarden.network.model.KeyConnectorUserDecryptionOptionsJson
 import com.bitwarden.network.model.TrustedDeviceUserDecryptionOptionsJson
 import com.bitwarden.network.model.UserDecryptionOptionsJson
-import com.bitwarden.network.model.createMockAccountKeysJson
 import com.bitwarden.network.model.createMockOrganizationNetwork
 import com.bitwarden.network.model.createMockPolicy
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.AccountJson
@@ -19,6 +19,7 @@ import com.x8bit.bitwarden.data.auth.datasource.disk.model.ForcePasswordResetRea
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.OnboardingStatus
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.PendingAuthRequestJson
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.UserStateJson
+import com.x8bit.bitwarden.data.auth.datasource.disk.serializer.WrappedAccountCryptographicStateSerializer
 import com.x8bit.bitwarden.data.platform.datasource.disk.legacy.LegacySecureStorageMigrator
 import io.mockk.every
 import io.mockk.just
@@ -294,10 +295,9 @@ class AuthDiskSourceTest {
             userId = userId,
             userAutoUnlockKey = "userAutoUnlockKey",
         )
-        authDiskSource.storePrivateKey(userId = userId, privateKey = "privateKey")
-        authDiskSource.storeAccountKeys(
+        authDiskSource.storeAccountCryptographicState(
             userId = userId,
-            accountKeys = createMockAccountKeysJson(number = 1),
+            accountCryptographicState = WrappedAccountCryptographicState.V1("privateKey"),
         )
         authDiskSource.storeOrganizationKeys(
             userId = userId,
@@ -345,8 +345,7 @@ class AuthDiskSourceTest {
         assertNull(authDiskSource.getUserBiometricUnlockKey(userId = userId))
         assertNull(authDiskSource.getInvalidUnlockAttempts(userId = userId))
         assertNull(authDiskSource.getUserAutoUnlockKey(userId = userId))
-        assertNull(authDiskSource.getPrivateKey(userId = userId))
-        assertNull(authDiskSource.getAccountKeys(userId = userId))
+        assertNull(authDiskSource.getAccountCryptographicState(userId = userId))
         assertNull(authDiskSource.getOrganizationKeys(userId = userId))
         assertNull(authDiskSource.getOrganizations(userId = userId))
         assertNull(authDiskSource.getPolicies(userId = userId))
@@ -441,76 +440,48 @@ class AuthDiskSourceTest {
     }
 
     @Test
-    fun `getPrivateKey should pull from SharedPreferences`() {
-        val privateKeyBaseKey = "bwPreferencesStorage:encPrivateKey"
+    fun `getAccountCryptographicState should pull from SharedPreferences`() {
+        val accountKeysBaseKey = "bwSecureStorage:accountCryptographicState"
         val mockUserId = "mockUserId"
-        val mockPrivateKey = "mockPrivateKey"
-        fakeSharedPreferences
-            .edit {
-                putString(
-                    "${privateKeyBaseKey}_$mockUserId",
-                    mockPrivateKey,
-                )
-            }
-        val actual = authDiskSource.getPrivateKey(userId = mockUserId)
-        assertEquals(
-            mockPrivateKey,
-            actual,
+        val mockAccountCryptographicState = WrappedAccountCryptographicState.V1(
+            privateKey = "privateKey",
         )
-    }
-
-    @Test
-    fun `storePrivateKey should update SharedPreferences`() {
-        val privateKeyBaseKey = "bwPreferencesStorage:encPrivateKey"
-        val mockUserId = "mockUserId"
-        val mockPrivateKey = "mockPrivateKey"
-        authDiskSource.storePrivateKey(
-            userId = mockUserId,
-            privateKey = mockPrivateKey,
-        )
-        val actual = fakeSharedPreferences.getString(
-            "${privateKeyBaseKey}_$mockUserId",
-            null,
-        )
-        assertEquals(
-            mockPrivateKey,
-            actual,
-        )
-    }
-
-    @Test
-    fun `getAccountKeys should pull from SharedPreferences`() {
-        val accountKeysBaseKey = "bwSecureStorage:profileAccountKeys"
-        val mockUserId = "mockUserId"
-        val mockAccountKeys = createMockAccountKeysJson(number = 1)
         fakeEncryptedSharedPreferences.edit {
             putString(
                 "${accountKeysBaseKey}_$mockUserId",
-                json.encodeToString(mockAccountKeys),
+                json.encodeToString(
+                    WrappedAccountCryptographicStateSerializer(),
+                    mockAccountCryptographicState,
+                ),
             )
         }
-        val actual = authDiskSource.getAccountKeys(userId = mockUserId)
+        val actual = authDiskSource.getAccountCryptographicState(userId = mockUserId)
         assertEquals(
-            mockAccountKeys,
+            mockAccountCryptographicState,
             actual,
         )
     }
 
     @Test
-    fun `storeAccountKeys should update sharedPreferences`() {
-        val accountKeysBaseKey = "bwSecureStorage:profileAccountKeys"
+    fun `storeAccountCryptographicState should update sharedPreferences`() {
+        val accountKeysBaseKey = "bwSecureStorage:accountCryptographicState"
         val mockUserId = "mockUserId"
-        val mockAccountKeys = createMockAccountKeysJson(number = 1)
-        authDiskSource.storeAccountKeys(
+        val mockAccountCryptographicState = WrappedAccountCryptographicState.V1(
+            privateKey = "privateKey",
+        )
+        authDiskSource.storeAccountCryptographicState(
             userId = mockUserId,
-            accountKeys = mockAccountKeys,
+            accountCryptographicState = mockAccountCryptographicState,
         )
         val actual = fakeEncryptedSharedPreferences.getString(
             "${accountKeysBaseKey}_$mockUserId",
             null,
         )
         assertEquals(
-            json.encodeToJsonElement(mockAccountKeys),
+            json.encodeToJsonElement(
+                WrappedAccountCryptographicStateSerializer(),
+                mockAccountCryptographicState,
+            ),
             json.parseToJsonElement(requireNotNull(actual)),
         )
     }

app/src/test/kotlin/com/x8bit/bitwarden/data/auth/datasource/disk/serializer/WrappedAccountCryptographicStateSerializerTest.kt

@@ -0,0 +1,163 @@
+package com.x8bit.bitwarden.data.auth.datasource.disk.serializer
+
+import com.bitwarden.core.WrappedAccountCryptographicState
+import com.bitwarden.core.di.CoreModule
+import io.mockk.mockk
+import kotlinx.serialization.json.buildJsonObject
+import kotlinx.serialization.json.put
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Test
+
+class WrappedAccountCryptographicStateSerializerTest {
+    private val json = CoreModule.providesJson(buildInfoManager = mockk(relaxed = true))
+    private val serializer = WrappedAccountCryptographicStateSerializer()
+
+    @Test
+    fun `serialize V1 produces JSON with v1 type discriminator`() {
+        assertEquals(
+            buildJsonObject {
+                put(key = "type", value = "v1")
+                put(key = "privateKey", value = "mockPrivateKey")
+            },
+            json.encodeToJsonElement(
+                serializer = serializer,
+                value = WrappedAccountCryptographicState.V1(privateKey = "mockPrivateKey"),
+            ),
+        )
+    }
+
+    @Test
+    fun `serialize V2 produces JSON with v2 type discriminator and all properties`() {
+        assertEquals(
+            buildJsonObject {
+                put(key = "type", value = "v2")
+                put(key = "privateKey", value = "mockPrivateKey")
+                put(key = "signingKey", value = "mockSigningKey")
+                put(key = "signedPublicKey", value = "mockSignedPublicKey")
+                put(key = "securityState", value = "mockSecurityState")
+            },
+            json.encodeToJsonElement(
+                serializer = serializer,
+                value = WrappedAccountCryptographicState.V2(
+                    privateKey = "mockPrivateKey",
+                    signingKey = "mockSigningKey",
+                    signedPublicKey = "mockSignedPublicKey",
+                    securityState = "mockSecurityState",
+                ),
+            ),
+        )
+    }
+
+    @Test
+    fun `serialize V2 with null signedPublicKey omits the signedPublicKey key`() {
+        assertEquals(
+            buildJsonObject {
+                put(key = "type", value = "v2")
+                put(key = "privateKey", value = "mockPrivateKey")
+                put(key = "signingKey", value = "mockSigningKey")
+                put(key = "securityState", value = "mockSecurityState")
+            },
+            json.encodeToJsonElement(
+                serializer = serializer,
+                value = WrappedAccountCryptographicState.V2(
+                    privateKey = "mockPrivateKey",
+                    signingKey = "mockSigningKey",
+                    signedPublicKey = null,
+                    securityState = "mockSecurityState",
+                ),
+            ),
+        )
+    }
+
+    @Test
+    fun `deserialize v1 JSON produces a V1 state`() {
+        assertEquals(
+            WrappedAccountCryptographicState.V1(privateKey = "mockPrivateKey"),
+            json.decodeFromString(
+                deserializer = serializer,
+                string = """
+                    {
+                      "type": "v1",
+                      "privateKey": "mockPrivateKey"
+                    }
+                """,
+            ),
+        )
+    }
+
+    @Test
+    fun `deserialize v2 JSON produces a V2 state`() {
+        assertEquals(
+            WrappedAccountCryptographicState.V2(
+                privateKey = "mockPrivateKey",
+                signingKey = "mockSigningKey",
+                signedPublicKey = "mockSignedPublicKey",
+                securityState = "mockSecurityState",
+            ),
+            json.decodeFromString(
+                deserializer = serializer,
+                string = """
+                    {
+                      "type": "v2",
+                      "privateKey": "mockPrivateKey",
+                      "signingKey": "mockSigningKey",
+                      "signedPublicKey": "mockSignedPublicKey",
+                      "securityState": "mockSecurityState"
+                    }
+                """,
+            ),
+        )
+    }
+
+    @Test
+    fun `deserialize v2 JSON with missing signedPublicKey produces a V2 state with null`() {
+        assertEquals(
+            WrappedAccountCryptographicState.V2(
+                privateKey = "mockPrivateKey",
+                signingKey = "mockSigningKey",
+                signedPublicKey = null,
+                securityState = "mockSecurityState",
+            ),
+            json.decodeFromString(
+                deserializer = serializer,
+                string = """
+                    {
+                      "type": "v2",
+                      "privateKey": "mockPrivateKey",
+                      "signingKey": "mockSigningKey",
+                      "securityState": "mockSecurityState"
+                    }
+                """,
+            ),
+        )
+    }
+
+    @Test
+    fun `serialize then deserialize V1 returns an equivalent state`() {
+        val original = WrappedAccountCryptographicState.V1(privateKey = "mockPrivateKey")
+        assertEquals(
+            original,
+            json.decodeFromString(
+                deserializer = serializer,
+                string = json.encodeToString(serializer = serializer, value = original),
+            ),
+        )
+    }
+
+    @Test
+    fun `serialize then deserialize V2 returns an equivalent state`() {
+        val original = WrappedAccountCryptographicState.V2(
+            privateKey = "mockPrivateKey",
+            signingKey = "mockSigningKey",
+            signedPublicKey = "mockSignedPublicKey",
+            securityState = "mockSecurityState",
+        )
+        assertEquals(
+            original,
+            json.decodeFromString(
+                deserializer = serializer,
+                string = json.encodeToString(serializer = serializer, value = original),
+            ),
+        )
+    }
+}

app/src/test/kotlin/com/x8bit/bitwarden/data/auth/datasource/disk/util/FakeAuthDiskSource.kt

@@ -1,7 +1,7 @@
 package com.x8bit.bitwarden.data.auth.datasource.disk.util
 
+import com.bitwarden.core.WrappedAccountCryptographicState
 import com.bitwarden.core.data.repository.util.bufferedMutableSharedFlow
-import com.bitwarden.network.model.AccountKeysJson
 import com.bitwarden.network.model.SyncResponseJson
 import com.x8bit.bitwarden.data.auth.datasource.disk.AuthDiskSource
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.AccountTokensJson
@@ -45,7 +45,6 @@ class FakeAuthDiskSource : AuthDiskSource {
     private val storedShouldTrustDevice = mutableMapOf<String, Boolean?>()
     private val storedInvalidUnlockAttempts = mutableMapOf<String, Int?>()
     private val storedLocalUserDataKeys = mutableMapOf<String, String?>()
-    private val storedPrivateKeys = mutableMapOf<String, String?>()
     private val storedTwoFactorTokens = mutableMapOf<String, String?>()
     private val storedUserAutoUnlockKeys = mutableMapOf<String, String?>()
     private val storedPinProtectedUserKeys = mutableMapOf<String, Pair<String?, Boolean>>()
@@ -64,7 +63,8 @@ class FakeAuthDiskSource : AuthDiskSource {
     private val storedOnboardingStatus = mutableMapOf<String, OnboardingStatus?>()
     private val storedShowImportLogins = mutableMapOf<String, Boolean?>()
     private val storedLastLockTimestampState = mutableMapOf<String, Instant?>()
-    private val storedAccountKeys = mutableMapOf<String, AccountKeysJson?>()
+    private val storedAccountCryptographicState =
+        mutableMapOf<String, WrappedAccountCryptographicState?>()
     private val storedPinProtectedUserKeyEnvelopes = mutableMapOf<String, Pair<String?, Boolean>>()
     private val mutablePinProtectedUserKeyEnvelopesFlowMap =
         mutableMapOf<String, MutableSharedFlow<String?>>()
@@ -81,7 +81,7 @@ class FakeAuthDiskSource : AuthDiskSource {
     override fun clearData(userId: String) {
         storedInvalidUnlockAttempts.remove(userId)
         storedLocalUserDataKeys.remove(userId)
-        storedPrivateKeys.remove(userId)
+        storedAccountCryptographicState.remove(userId)
         storedTwoFactorTokens.clear()
         storedUserAutoUnlockKeys.remove(userId)
         storedOrganizations.remove(userId)
@@ -150,22 +150,15 @@ class FakeAuthDiskSource : AuthDiskSource {
         storedLocalUserDataKeys[userId] = wrappedKey
     }
 
-    @Deprecated("Use getAccountKeys instead.", replaceWith = ReplaceWith("getAccountKeys"))
-    override fun getPrivateKey(userId: String): String? = storedPrivateKeys[userId]
-
-    @Deprecated("Use storeAccountKeys instead.", replaceWith = ReplaceWith("storeAccountKeys"))
-    override fun storePrivateKey(userId: String, privateKey: String?) {
-        storedPrivateKeys[userId] = privateKey
-    }
-
-    override fun getAccountKeys(userId: String): AccountKeysJson? =
-        storedAccountKeys[userId]
-
-    override fun storeAccountKeys(
+    override fun getAccountCryptographicState(
         userId: String,
-        accountKeys: AccountKeysJson?,
+    ): WrappedAccountCryptographicState? = storedAccountCryptographicState[userId]
+
+    override fun storeAccountCryptographicState(
+        userId: String,
+        accountCryptographicState: WrappedAccountCryptographicState?,
     ) {
-        storedAccountKeys[userId] = accountKeys
+        storedAccountCryptographicState[userId] = accountCryptographicState
     }
 
     override fun getTwoFactorToken(email: String): String? = storedTwoFactorTokens[email]
@@ -415,18 +408,13 @@ class FakeAuthDiskSource : AuthDiskSource {
     }
 
     /**
-     * Assert that the [privateKey] was stored successfully using the [userId].
+     * Assert that the [accountCryptographicState] was stored successfully using the [userId].
      */
-    @Deprecated("Use assertAccountKeys instead.", replaceWith = ReplaceWith("assertAccountKeys"))
-    fun assertPrivateKey(userId: String, privateKey: String?) {
-        assertEquals(privateKey, storedPrivateKeys[userId])
-    }
-
-    /**
-     * Assert that the [accountKeys] was stored successfully using the [userId].
-     */
-    fun assertAccountKeys(userId: String, accountKeys: AccountKeysJson?) {
-        assertEquals(accountKeys, storedAccountKeys[userId])
+    fun assertAccountCryptographicState(
+        userId: String,
+        accountCryptographicState: WrappedAccountCryptographicState?,
+    ) {
+        assertEquals(accountCryptographicState, storedAccountCryptographicState[userId])
     }
 
     /**

app/src/test/kotlin/com/x8bit/bitwarden/data/auth/repository/model/AccountCryptographicStateUtil.kt

@@ -0,0 +1,20 @@
+package com.x8bit.bitwarden.data.auth.repository.model
+
+import com.bitwarden.core.WrappedAccountCryptographicState
+
+/**
+ * Create a mock [WrappedAccountCryptographicState] with a given number.
+ */
+fun createMockWrappedAccountCryptographicState(
+    number: Int,
+    signedPublicKey: String? = "mockSignedPublicKey-$number",
+    privateKey: String = "mockWrappedPrivateKey-$number",
+    signingKey: String = "mockWrappedSigningKey-$number",
+    securityState: String = "mockSecurityState-$number",
+): WrappedAccountCryptographicState =
+    WrappedAccountCryptographicState.V2(
+        signedPublicKey = signedPublicKey,
+        privateKey = privateKey,
+        signingKey = signingKey,
+        securityState = securityState,
+    )

app/src/test/kotlin/com/x8bit/bitwarden/data/auth/repository/util/AccountKeysExtensionsTest.kt

@@ -0,0 +1,98 @@
+package com.x8bit.bitwarden.data.auth.repository.util
+
+import com.bitwarden.core.WrappedAccountCryptographicState
+import com.bitwarden.network.model.AccountKeysJson
+import com.bitwarden.network.model.AccountKeysJson.PublicKeyEncryptionKeyPair
+import com.bitwarden.network.model.AccountKeysJson.SignatureKeyPair
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Test
+
+class AccountKeysExtensionsTest {
+
+    @Suppress("MaxLineLength")
+    @Test
+    fun `toAccountCryptographicState returns V2 when securityState, signedPublicKey and signingKey are non-null`() {
+        val result = ACCOUNT_KEYS.toAccountCryptographicState(privateKey = PRIVATE_KEY)
+
+        assertEquals(
+            WrappedAccountCryptographicState.V2(
+                privateKey = PRIVATE_KEY,
+                signedPublicKey = SIGNED_PUBLIC_KEY,
+                signingKey = SIGNING_KEY,
+                securityState = SECURITY_STATE,
+            ),
+            result,
+        )
+    }
+
+    @Test
+    fun `toAccountCryptographicState returns V1 when securityState is null`() {
+        val result = ACCOUNT_KEYS
+            .copy(securityState = null)
+            .toAccountCryptographicState(privateKey = PRIVATE_KEY)
+
+        assertEquals(
+            WrappedAccountCryptographicState.V1(privateKey = PRIVATE_KEY),
+            result,
+        )
+    }
+
+    @Test
+    fun `toAccountCryptographicState returns V1 when signedPublicKey is null`() {
+        val result = ACCOUNT_KEYS
+            .copy(
+                publicKeyEncryptionKeyPair = ACCOUNT_KEYS
+                    .publicKeyEncryptionKeyPair
+                    .copy(signedPublicKey = null),
+            )
+            .toAccountCryptographicState(privateKey = PRIVATE_KEY)
+
+        assertEquals(
+            WrappedAccountCryptographicState.V1(privateKey = PRIVATE_KEY),
+            result,
+        )
+    }
+
+    @Test
+    fun `toAccountCryptographicState returns V1 when signingKey is null`() {
+        val result = ACCOUNT_KEYS
+            .copy(signatureKeyPair = null)
+            .toAccountCryptographicState(privateKey = PRIVATE_KEY)
+
+        assertEquals(
+            WrappedAccountCryptographicState.V1(privateKey = PRIVATE_KEY),
+            result,
+        )
+    }
+
+    @Test
+    fun `toAccountCryptographicState returns V1 when AccountKeysJson is null`() {
+        val result = null.toAccountCryptographicState(privateKey = PRIVATE_KEY)
+
+        assertEquals(
+            WrappedAccountCryptographicState.V1(privateKey = PRIVATE_KEY),
+            result,
+        )
+    }
+}
+
+private const val PRIVATE_KEY = "test-private-key"
+private const val SECURITY_STATE = "test-security-state"
+private const val SIGNING_KEY = "test-signing-key"
+private const val SIGNED_PUBLIC_KEY = "test-signed-public-key"
+
+private val ACCOUNT_KEYS = AccountKeysJson(
+    signatureKeyPair = SignatureKeyPair(
+        wrappedSigningKey = SIGNING_KEY,
+        verifyingKey = "verifying-key",
+    ),
+    publicKeyEncryptionKeyPair = PublicKeyEncryptionKeyPair(
+        wrappedPrivateKey = PRIVATE_KEY,
+        signedPublicKey = SIGNED_PUBLIC_KEY,
+        publicKey = "public-key",
+    ),
+    securityState = AccountKeysJson.SecurityState(
+        securityState = SECURITY_STATE,
+        securityVersion = 2,
+    ),
+)

app/src/test/kotlin/com/x8bit/bitwarden/data/platform/repository/AuthenticatorBridgeRepositoryTest.kt

@@ -6,6 +6,7 @@ import com.bitwarden.authenticatorbridge.util.toSymmetricEncryptionKeyData
 import com.bitwarden.core.InitOrgCryptoRequest
 import com.bitwarden.core.InitUserCryptoMethod
 import com.bitwarden.core.InitUserCryptoRequest
+import com.bitwarden.core.WrappedAccountCryptographicState
 import com.bitwarden.core.data.util.asSuccess
 import com.bitwarden.crypto.Kdf
 import com.bitwarden.data.datasource.disk.model.EnvironmentUrlDataJson
@@ -22,7 +23,6 @@ import com.x8bit.bitwarden.data.auth.datasource.disk.util.FakeAuthDiskSource
 import com.x8bit.bitwarden.data.vault.datasource.disk.VaultDiskSource
 import com.x8bit.bitwarden.data.vault.datasource.sdk.ScopedVaultSdkSource
 import com.x8bit.bitwarden.data.vault.datasource.sdk.model.InitializeCryptoResult
-import com.x8bit.bitwarden.data.vault.repository.util.createWrappedAccountCryptographicState
 import com.x8bit.bitwarden.data.vault.repository.util.toEncryptedSdkCipher
 import io.mockk.coEvery
 import io.mockk.coVerify
@@ -80,18 +80,19 @@ class AuthenticatorBridgeRepositoryTest {
             userId = USER_2_ID,
             authenticatorSyncUnlockKey = USER_2_UNLOCK_KEY,
         )
-        fakeAuthDiskSource.storePrivateKey(userId = USER_1_ID, privateKey = USER_1_PRIVATE_KEY)
-        fakeAuthDiskSource.storePrivateKey(userId = USER_2_ID, privateKey = USER_2_PRIVATE_KEY)
+        fakeAuthDiskSource.storeAccountCryptographicState(
+            userId = USER_1_ID,
+            accountCryptographicState = USER_1_ACCOUNT_CRYPTOGRAPHIC_STATE,
+        )
+        fakeAuthDiskSource.storeAccountCryptographicState(
+            userId = USER_2_ID,
+            accountCryptographicState = USER_2_ACCOUNT_CRYPTOGRAPHIC_STATE,
+        )
         coEvery {
             scopedVaultSdkSource.initializeCrypto(
                 userId = USER_1_ID,
                 request = InitUserCryptoRequest(
-                    accountCryptographicState = createWrappedAccountCryptographicState(
-                        privateKey = USER_1_PRIVATE_KEY,
-                        securityState = null,
-                        signedPublicKey = null,
-                        signingKey = null,
-                    ),
+                    accountCryptographicState = USER_1_ACCOUNT_CRYPTOGRAPHIC_STATE,
                     userId = USER_1_ID,
                     kdfParams = Kdf.Argon2id(iterations = 0U, memory = 0U, parallelism = 0U),
                     email = USER_1_EMAIL,
@@ -106,12 +107,7 @@ class AuthenticatorBridgeRepositoryTest {
             scopedVaultSdkSource.initializeCrypto(
                 userId = USER_2_ID,
                 request = InitUserCryptoRequest(
-                    accountCryptographicState = createWrappedAccountCryptographicState(
-                        privateKey = USER_2_PRIVATE_KEY,
-                        securityState = null,
-                        signedPublicKey = null,
-                        signingKey = null,
-                    ),
+                    accountCryptographicState = USER_2_ACCOUNT_CRYPTOGRAPHIC_STATE,
                     userId = USER_2_ID,
                     kdfParams = Kdf.Argon2id(iterations = 0U, memory = 0U, parallelism = 0U),
                     email = USER_2_EMAIL,
@@ -203,12 +199,7 @@ class AuthenticatorBridgeRepositoryTest {
                 scopedVaultSdkSource.initializeCrypto(
                     userId = USER_2_ID,
                     request = InitUserCryptoRequest(
-                        accountCryptographicState = createWrappedAccountCryptographicState(
-                            privateKey = USER_2_PRIVATE_KEY,
-                            securityState = null,
-                            signedPublicKey = null,
-                            signingKey = null,
-                        ),
+                        accountCryptographicState = USER_2_ACCOUNT_CRYPTOGRAPHIC_STATE,
                         userId = USER_2_ID,
                         kdfParams = Kdf.Argon2id(iterations = 0U, memory = 0U, parallelism = 0U),
                         email = USER_2_EMAIL,
@@ -246,12 +237,7 @@ class AuthenticatorBridgeRepositoryTest {
                 scopedVaultSdkSource.initializeCrypto(
                     userId = USER_1_ID,
                     request = InitUserCryptoRequest(
-                        accountCryptographicState = createWrappedAccountCryptographicState(
-                            privateKey = USER_1_PRIVATE_KEY,
-                            securityState = null,
-                            signedPublicKey = null,
-                            signingKey = null,
-                        ),
+                        accountCryptographicState = USER_1_ACCOUNT_CRYPTOGRAPHIC_STATE,
                         userId = USER_1_ID,
                         kdfParams = Kdf.Argon2id(iterations = 0U, memory = 0U, parallelism = 0U),
                         email = USER_1_EMAIL,
@@ -273,12 +259,7 @@ class AuthenticatorBridgeRepositoryTest {
                 scopedVaultSdkSource.initializeCrypto(
                     userId = USER_2_ID,
                     request = InitUserCryptoRequest(
-                        accountCryptographicState = createWrappedAccountCryptographicState(
-                            privateKey = USER_2_PRIVATE_KEY,
-                            securityState = null,
-                            signedPublicKey = null,
-                            signingKey = null,
-                        ),
+                        accountCryptographicState = USER_2_ACCOUNT_CRYPTOGRAPHIC_STATE,
                         userId = USER_2_ID,
                         kdfParams = Kdf.Argon2id(iterations = 0U, memory = 0U, parallelism = 0U),
                         email = USER_2_EMAIL,
@@ -312,12 +293,7 @@ class AuthenticatorBridgeRepositoryTest {
                 scopedVaultSdkSource.initializeCrypto(
                     userId = USER_1_ID,
                     request = InitUserCryptoRequest(
-                        accountCryptographicState = createWrappedAccountCryptographicState(
-                            privateKey = USER_1_PRIVATE_KEY,
-                            securityState = null,
-                            signedPublicKey = null,
-                            signingKey = null,
-                        ),
+                        accountCryptographicState = USER_1_ACCOUNT_CRYPTOGRAPHIC_STATE,
                         userId = USER_1_ID,
                         kdfParams = Kdf.Argon2id(iterations = 0U, memory = 0U, parallelism = 0U),
                         email = USER_1_EMAIL,
@@ -338,12 +314,7 @@ class AuthenticatorBridgeRepositoryTest {
                 scopedVaultSdkSource.initializeCrypto(
                     userId = USER_1_ID,
                     request = InitUserCryptoRequest(
-                        accountCryptographicState = createWrappedAccountCryptographicState(
-                            privateKey = USER_1_PRIVATE_KEY,
-                            securityState = null,
-                            signedPublicKey = null,
-                            signingKey = null,
-                        ),
+                        accountCryptographicState = USER_1_ACCOUNT_CRYPTOGRAPHIC_STATE,
                         userId = USER_1_ID,
                         kdfParams = Kdf.Argon2id(iterations = 0U, memory = 0U, parallelism = 0U),
                         email = USER_1_EMAIL,
@@ -356,12 +327,7 @@ class AuthenticatorBridgeRepositoryTest {
                 scopedVaultSdkSource.initializeCrypto(
                     userId = USER_2_ID,
                     request = InitUserCryptoRequest(
-                        accountCryptographicState = createWrappedAccountCryptographicState(
-                            privateKey = USER_2_PRIVATE_KEY,
-                            securityState = null,
-                            signedPublicKey = null,
-                            signingKey = null,
-                        ),
+                        accountCryptographicState = USER_2_ACCOUNT_CRYPTOGRAPHIC_STATE,
                         userId = USER_2_ID,
                         kdfParams = Kdf.Argon2id(iterations = 0U, memory = 0U, parallelism = 0U),
                         email = USER_2_EMAIL,
@@ -440,7 +406,13 @@ private const val USER_1_EMAIL = "john@doe.com"
 private const val USER_2_EMAIL = "jane@doe.com"
 
 private const val USER_1_PRIVATE_KEY = "user1PrivateKey"
+private val USER_1_ACCOUNT_CRYPTOGRAPHIC_STATE = WrappedAccountCryptographicState.V1(
+    privateKey = USER_1_PRIVATE_KEY,
+)
 private const val USER_2_PRIVATE_KEY = "user2PrivateKey"
+private val USER_2_ACCOUNT_CRYPTOGRAPHIC_STATE = WrappedAccountCryptographicState.V1(
+    privateKey = USER_2_PRIVATE_KEY,
+)
 
 private const val USER_1_UNLOCK_KEY = "user1UnlockKey"
 private const val USER_2_UNLOCK_KEY = "user2UnlockKey"

app/src/test/kotlin/com/x8bit/bitwarden/data/vault/manager/VaultLockManagerTest.kt

@@ -9,6 +9,7 @@ import com.bitwarden.core.InitOrgCryptoRequest
 import com.bitwarden.core.InitUserCryptoMethod
 import com.bitwarden.core.InitUserCryptoRequest
 import com.bitwarden.core.MasterPasswordUnlockData
+import com.bitwarden.core.WrappedAccountCryptographicState
 import com.bitwarden.core.data.manager.dispatcher.FakeDispatcherManager
 import com.bitwarden.core.data.manager.realtime.RealtimeManager
 import com.bitwarden.core.data.util.asFailure
@@ -37,7 +38,6 @@ import com.x8bit.bitwarden.data.vault.datasource.sdk.model.InitializeCryptoResul
 import com.x8bit.bitwarden.data.vault.manager.model.VaultStateEvent
 import com.x8bit.bitwarden.data.vault.repository.model.VaultUnlockData
 import com.x8bit.bitwarden.data.vault.repository.model.VaultUnlockResult
-import com.x8bit.bitwarden.data.vault.repository.util.createWrappedAccountCryptographicState
 import io.mockk.clearMocks
 import io.mockk.coEvery
 import io.mockk.coVerify
@@ -725,12 +725,12 @@ class VaultLockManagerTest {
     @Suppress("MaxLineLength")
     @Test
     fun `vaultTimeout update to Never for a locked account when there is no stored auto-unlock key should do nothing`() {
-        val privateKey = "privateKey"
+        val accountCryptographicState = WrappedAccountCryptographicState.V1("privateKey")
         fakeAuthDiskSource.apply {
             userState = MOCK_USER_STATE
-            storePrivateKey(
+            storeAccountCryptographicState(
                 userId = USER_ID,
-                privateKey = privateKey,
+                accountCryptographicState = accountCryptographicState,
             )
         }
         assertFalse(vaultLockManager.isVaultUnlocked(userId = USER_ID))
@@ -743,13 +743,13 @@ class VaultLockManagerTest {
     @Suppress("MaxLineLength")
     @Test
     fun `vaultTimeout update to Never for a locked account when there is a stored auto-unlock key should unlock the vault`() {
-        val privateKey = "privateKey"
+        val accountCryptographicState = WrappedAccountCryptographicState.V1("privateKey")
         val userAutoUnlockKey = "userAutoUnlockKey"
         fakeAuthDiskSource.apply {
             userState = MOCK_USER_STATE
-            storePrivateKey(
+            storeAccountCryptographicState(
                 userId = USER_ID,
-                privateKey = privateKey,
+                accountCryptographicState = accountCryptographicState,
             )
             storeUserAutoUnlockKey(
                 userId = USER_ID,
@@ -760,12 +760,7 @@ class VaultLockManagerTest {
             vaultSdkSource.initializeCrypto(
                 userId = USER_ID,
                 request = InitUserCryptoRequest(
-                    accountCryptographicState = createWrappedAccountCryptographicState(
-                        privateKey = privateKey,
-                        securityState = null,
-                        signedPublicKey = null,
-                        signingKey = null,
-                    ),
+                    accountCryptographicState = accountCryptographicState,
                     userId = USER_ID,
                     kdfParams = MOCK_PROFILE.toSdkParams(),
                     email = MOCK_PROFILE.email,
@@ -790,12 +785,7 @@ class VaultLockManagerTest {
             vaultSdkSource.initializeCrypto(
                 userId = USER_ID,
                 request = InitUserCryptoRequest(
-                    accountCryptographicState = createWrappedAccountCryptographicState(
-                        privateKey = privateKey,
-                        securityState = null,
-                        signedPublicKey = null,
-                        signingKey = null,
-                    ),
+                    accountCryptographicState = accountCryptographicState,
                     userId = USER_ID,
                     kdfParams = MOCK_PROFILE.toSdkParams(),
                     email = MOCK_PROFILE.email,
@@ -920,18 +910,12 @@ class VaultLockManagerTest {
             val kdf = MOCK_PROFILE.toSdkParams()
             val email = MOCK_PROFILE.email
             val masterPassword = "mockValue"
-            val privateKey = "54321"
             val organizationKeys = mapOf("orgId1" to "orgKey1")
             coEvery {
                 vaultSdkSource.initializeCrypto(
                     userId = USER_ID,
                     request = InitUserCryptoRequest(
-                        accountCryptographicState = createWrappedAccountCryptographicState(
-                            privateKey = privateKey,
-                            securityState = null,
-                            signedPublicKey = null,
-                            signingKey = null,
-                        ),
+                        accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
                         userId = USER_ID,
                         kdfParams = kdf,
                         email = email,
@@ -963,12 +947,7 @@ class VaultLockManagerTest {
             )
 
             val result = vaultLockManager.unlockVault(
-                accountCryptographicState = createWrappedAccountCryptographicState(
-                    privateKey = privateKey,
-                    securityState = null,
-                    signedPublicKey = null,
-                    signingKey = null,
-                ),
+                accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
                 userId = USER_ID,
                 email = email,
                 kdf = kdf,
@@ -1002,12 +981,7 @@ class VaultLockManagerTest {
                 vaultSdkSource.initializeCrypto(
                     userId = USER_ID,
                     request = InitUserCryptoRequest(
-                        accountCryptographicState = createWrappedAccountCryptographicState(
-                            privateKey = privateKey,
-                            securityState = null,
-                            signedPublicKey = null,
-                            signingKey = null,
-                        ),
+                        accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
                         userId = USER_ID,
                         kdfParams = kdf,
                         email = email,
@@ -1034,19 +1008,13 @@ class VaultLockManagerTest {
             val kdf = MOCK_PROFILE.toSdkParams()
             val email = MOCK_PROFILE.email
             val masterPassword = "mockValue"
-            val privateKey = "54321"
             val organizationKeys = mapOf("orgId1" to "orgKey1")
             val userAutoUnlockKey = "userAutoUnlockKey"
             coEvery {
                 vaultSdkSource.initializeCrypto(
                     userId = USER_ID,
                     request = InitUserCryptoRequest(
-                        accountCryptographicState = createWrappedAccountCryptographicState(
-                            privateKey = privateKey,
-                            securityState = null,
-                            signedPublicKey = null,
-                            signingKey = null,
-                        ),
+                        accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
                         userId = USER_ID,
                         kdfParams = kdf,
                         email = email,
@@ -1087,12 +1055,7 @@ class VaultLockManagerTest {
             }
 
             val result = vaultLockManager.unlockVault(
-                accountCryptographicState = createWrappedAccountCryptographicState(
-                    privateKey = privateKey,
-                    securityState = null,
-                    signedPublicKey = null,
-                    signingKey = null,
-                ),
+                accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
                 userId = USER_ID,
                 email = email,
                 kdf = kdf,
@@ -1128,12 +1091,7 @@ class VaultLockManagerTest {
                 vaultSdkSource.initializeCrypto(
                     userId = USER_ID,
                     request = InitUserCryptoRequest(
-                        accountCryptographicState = createWrappedAccountCryptographicState(
-                            privateKey = privateKey,
-                            securityState = null,
-                            signedPublicKey = null,
-                            signingKey = null,
-                        ),
+                        accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
                         userId = USER_ID,
                         kdfParams = kdf,
                         email = email,
@@ -1160,19 +1118,13 @@ class VaultLockManagerTest {
             val kdf = MOCK_PROFILE.toSdkParams()
             val email = MOCK_PROFILE.email
             val masterPassword = "mockValue"
-            val privateKey = "54321"
             val organizationKeys = mapOf("orgId1" to "orgKey1")
             val error = Throwable("Fail")
             coEvery {
                 vaultSdkSource.initializeCrypto(
                     userId = USER_ID,
                     request = InitUserCryptoRequest(
-                        accountCryptographicState = createWrappedAccountCryptographicState(
-                            privateKey = privateKey,
-                            securityState = null,
-                            signedPublicKey = null,
-                            signingKey = null,
-                        ),
+                        accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
                         userId = USER_ID,
                         kdfParams = kdf,
                         email = email,
@@ -1195,12 +1147,7 @@ class VaultLockManagerTest {
             )
 
             val result = vaultLockManager.unlockVault(
-                accountCryptographicState = createWrappedAccountCryptographicState(
-                    privateKey = privateKey,
-                    securityState = null,
-                    signedPublicKey = null,
-                    signingKey = null,
-                ),
+                accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
                 userId = USER_ID,
                 email = email,
                 kdf = kdf,
@@ -1224,12 +1171,7 @@ class VaultLockManagerTest {
                 vaultSdkSource.initializeCrypto(
                     userId = USER_ID,
                     request = InitUserCryptoRequest(
-                        accountCryptographicState = createWrappedAccountCryptographicState(
-                            privateKey = privateKey,
-                            securityState = null,
-                            signedPublicKey = null,
-                            signingKey = null,
-                        ),
+                        accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
                         userId = USER_ID,
                         kdfParams = kdf,
                         email = email,
@@ -1250,18 +1192,12 @@ class VaultLockManagerTest {
             val kdf = MOCK_PROFILE.toSdkParams()
             val email = MOCK_PROFILE.email
             val masterPassword = "mockValue"
-            val privateKey = "54321"
             val organizationKeys = mapOf("orgId1" to "orgKey1")
             coEvery {
                 vaultSdkSource.initializeCrypto(
                     userId = USER_ID,
                     request = InitUserCryptoRequest(
-                        accountCryptographicState = createWrappedAccountCryptographicState(
-                            privateKey = privateKey,
-                            securityState = null,
-                            signedPublicKey = null,
-                            signingKey = null,
-                        ),
+                        accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
                         userId = USER_ID,
                         kdfParams = kdf,
                         email = email,
@@ -1291,12 +1227,7 @@ class VaultLockManagerTest {
             )
 
             val result = vaultLockManager.unlockVault(
-                accountCryptographicState = createWrappedAccountCryptographicState(
-                    privateKey = privateKey,
-                    securityState = null,
-                    signedPublicKey = null,
-                    signingKey = null,
-                ),
+                accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
                 userId = USER_ID,
                 email = email,
                 kdf = kdf,
@@ -1320,12 +1251,7 @@ class VaultLockManagerTest {
                 vaultSdkSource.initializeCrypto(
                     userId = USER_ID,
                     request = InitUserCryptoRequest(
-                        accountCryptographicState = createWrappedAccountCryptographicState(
-                            privateKey = privateKey,
-                            securityState = null,
-                            signedPublicKey = null,
-                            signingKey = null,
-                        ),
+                        accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
                         userId = USER_ID,
                         kdfParams = kdf,
                         email = email,
@@ -1352,19 +1278,13 @@ class VaultLockManagerTest {
             val kdf = MOCK_PROFILE.toSdkParams()
             val email = MOCK_PROFILE.email
             val masterPassword = "mockValue"
-            val privateKey = "54321"
             val organizationKeys = mapOf("orgId1" to "orgKey1")
             val error = Throwable("Fail")
             coEvery {
                 vaultSdkSource.initializeCrypto(
                     userId = USER_ID,
                     request = InitUserCryptoRequest(
-                        accountCryptographicState = createWrappedAccountCryptographicState(
-                            privateKey = privateKey,
-                            securityState = null,
-                            signedPublicKey = null,
-                            signingKey = null,
-                        ),
+                        accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
                         userId = USER_ID,
                         kdfParams = kdf,
                         email = email,
@@ -1386,12 +1306,7 @@ class VaultLockManagerTest {
             )
 
             val result = vaultLockManager.unlockVault(
-                accountCryptographicState = createWrappedAccountCryptographicState(
-                    privateKey = privateKey,
-                    securityState = null,
-                    signedPublicKey = null,
-                    signingKey = null,
-                ),
+                accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
                 userId = USER_ID,
                 email = email,
                 kdf = kdf,
@@ -1415,12 +1330,7 @@ class VaultLockManagerTest {
                 vaultSdkSource.initializeCrypto(
                     userId = USER_ID,
                     request = InitUserCryptoRequest(
-                        accountCryptographicState = createWrappedAccountCryptographicState(
-                            privateKey = privateKey,
-                            securityState = null,
-                            signedPublicKey = null,
-                            signingKey = null,
-                        ),
+                        accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
                         userId = USER_ID,
                         kdfParams = kdf,
                         email = email,
@@ -1441,18 +1351,12 @@ class VaultLockManagerTest {
             val kdf = MOCK_PROFILE.toSdkParams()
             val email = MOCK_PROFILE.email
             val masterPassword = "mockValue"
-            val privateKey = "54321"
             val organizationKeys = mapOf("orgId1" to "orgKey1")
             coEvery {
                 vaultSdkSource.initializeCrypto(
                     userId = USER_ID,
                     request = InitUserCryptoRequest(
-                        accountCryptographicState = createWrappedAccountCryptographicState(
-                            privateKey = privateKey,
-                            securityState = null,
-                            signedPublicKey = null,
-                            signingKey = null,
-                        ),
+                        accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
                         userId = USER_ID,
                         kdfParams = kdf,
                         email = email,
@@ -1481,12 +1385,7 @@ class VaultLockManagerTest {
             )
 
             val result = vaultLockManager.unlockVault(
-                accountCryptographicState = createWrappedAccountCryptographicState(
-                    privateKey = privateKey,
-                    securityState = null,
-                    signedPublicKey = null,
-                    signingKey = null,
-                ),
+                accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
                 userId = USER_ID,
                 email = email,
                 kdf = kdf,
@@ -1510,12 +1409,7 @@ class VaultLockManagerTest {
                 vaultSdkSource.initializeCrypto(
                     userId = USER_ID,
                     request = InitUserCryptoRequest(
-                        accountCryptographicState = createWrappedAccountCryptographicState(
-                            privateKey = privateKey,
-                            securityState = null,
-                            signedPublicKey = null,
-                            signingKey = null,
-                        ),
+                        accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
                         userId = USER_ID,
                         kdfParams = kdf,
                         email = email,
@@ -1542,18 +1436,12 @@ class VaultLockManagerTest {
             val kdf = MOCK_PROFILE.toSdkParams()
             val email = MOCK_PROFILE.email
             val masterPassword = "mockValue"
-            val privateKey = "54321"
             val organizationKeys = mapOf("orgId1" to "orgKey1")
             coEvery {
                 vaultSdkSource.initializeCrypto(
                     userId = USER_ID,
                     request = InitUserCryptoRequest(
-                        accountCryptographicState = createWrappedAccountCryptographicState(
-                            privateKey = privateKey,
-                            securityState = null,
-                            signedPublicKey = null,
-                            signingKey = null,
-                        ),
+                        accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
                         userId = USER_ID,
                         kdfParams = kdf,
                         email = email,
@@ -1582,12 +1470,7 @@ class VaultLockManagerTest {
             )
 
             val result = vaultLockManager.unlockVault(
-                accountCryptographicState = createWrappedAccountCryptographicState(
-                    privateKey = privateKey,
-                    securityState = null,
-                    signedPublicKey = null,
-                    signingKey = null,
-                ),
+                accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
                 userId = USER_ID,
                 email = email,
                 kdf = kdf,
@@ -1618,12 +1501,7 @@ class VaultLockManagerTest {
                 vaultSdkSource.initializeCrypto(
                     userId = USER_ID,
                     request = InitUserCryptoRequest(
-                        accountCryptographicState = createWrappedAccountCryptographicState(
-                            privateKey = privateKey,
-                            securityState = null,
-                            signedPublicKey = null,
-                            signingKey = null,
-                        ),
+                        accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
                         userId = USER_ID,
                         kdfParams = kdf,
                         email = email,
@@ -1731,12 +1609,7 @@ class VaultLockManagerTest {
                 vaultSdkSource.initializeCrypto(
                     userId = USER_ID,
                     request = InitUserCryptoRequest(
-                        accountCryptographicState = createWrappedAccountCryptographicState(
-                            privateKey = privateKey,
-                            securityState = null,
-                            signedPublicKey = null,
-                            signingKey = null,
-                        ),
+                        accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
                         userId = USER_ID,
                         kdfParams = kdf,
                         email = email,
@@ -1764,12 +1637,7 @@ class VaultLockManagerTest {
             mutableVaultTimeoutStateFlow.value = VaultTimeout.ThirtyMinutes
 
             val result = vaultLockManager.unlockVault(
-                accountCryptographicState = createWrappedAccountCryptographicState(
-                    privateKey = privateKey,
-                    securityState = null,
-                    signedPublicKey = null,
-                    signingKey = null,
-                ),
+                accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
                 userId = USER_ID,
                 email = email,
                 kdf = kdf,
@@ -1795,12 +1663,7 @@ class VaultLockManagerTest {
                 vaultSdkSource.initializeCrypto(
                     userId = USER_ID,
                     request = InitUserCryptoRequest(
-                        accountCryptographicState = createWrappedAccountCryptographicState(
-                            privateKey = privateKey,
-                            securityState = null,
-                            signedPublicKey = null,
-                            signingKey = null,
-                        ),
+                        accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
                         userId = USER_ID,
                         kdfParams = kdf,
                         email = email,
@@ -1827,7 +1690,6 @@ class VaultLockManagerTest {
             val kdf = MOCK_PROFILE.toSdkParams()
             val email = MOCK_PROFILE.email
             val masterPassword = "mockValue"
-            val privateKey = "54321"
             val organizationKeys = mapOf("orgId1" to "orgKey1")
             val initUserCryptoMethod = InitUserCryptoMethod.MasterPasswordUnlock(
                 password = masterPassword,
@@ -1837,12 +1699,7 @@ class VaultLockManagerTest {
                 vaultSdkSource.initializeCrypto(
                     userId = USER_ID,
                     request = InitUserCryptoRequest(
-                        accountCryptographicState = createWrappedAccountCryptographicState(
-                            privateKey = privateKey,
-                            securityState = null,
-                            signedPublicKey = null,
-                            signingKey = null,
-                        ),
+                        accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
                         userId = USER_ID,
                         kdfParams = kdf,
                         email = email,
@@ -1871,12 +1728,7 @@ class VaultLockManagerTest {
             )
 
             val result = vaultLockManager.unlockVault(
-                accountCryptographicState = createWrappedAccountCryptographicState(
-                    privateKey = privateKey,
-                    securityState = null,
-                    signedPublicKey = null,
-                    signingKey = null,
-                ),
+                accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
                 userId = USER_ID,
                 email = email,
                 kdf = kdf,
@@ -1907,12 +1759,7 @@ class VaultLockManagerTest {
                 vaultSdkSource.initializeCrypto(
                     userId = USER_ID,
                     request = InitUserCryptoRequest(
-                        accountCryptographicState = createWrappedAccountCryptographicState(
-                            privateKey = privateKey,
-                            securityState = null,
-                            signedPublicKey = null,
-                            signingKey = null,
-                        ),
+                        accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
                         userId = USER_ID,
                         kdfParams = kdf,
                         email = email,
@@ -1953,18 +1800,12 @@ class VaultLockManagerTest {
         val email = MOCK_PROFILE.email
         val userKey = "12345"
         val masterPassword = "mockValue"
-        val privateKey = "54321"
         val organizationKeys = null
         coEvery {
             vaultSdkSource.initializeCrypto(
                 userId = userId,
                 request = InitUserCryptoRequest(
-                    accountCryptographicState = createWrappedAccountCryptographicState(
-                        privateKey = privateKey,
-                        securityState = null,
-                        signedPublicKey = null,
-                        signingKey = null,
-                    ),
+                    accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
                     userId = userId,
                     kdfParams = kdf,
                     email = email,
@@ -1981,12 +1822,7 @@ class VaultLockManagerTest {
         }
 
         vaultLockManager.unlockVault(
-            accountCryptographicState = createWrappedAccountCryptographicState(
-                privateKey = privateKey,
-                securityState = null,
-                signedPublicKey = null,
-                signingKey = null,
-            ),
+            accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
             userId = userId,
             email = email,
             kdf = kdf,
@@ -2009,7 +1845,6 @@ class VaultLockManagerTest {
         val kdf = MOCK_PROFILE.toSdkParams()
         val email = MOCK_PROFILE.email
         val masterPassword = "mockValue"
-        val privateKey = "54321"
         val organizationKeys = null
         val userAutoUnlockKey = "userAutoUnlockKey"
         // Clear recorded calls so this helper can be called multiple times and assert a unique
@@ -2019,12 +1854,7 @@ class VaultLockManagerTest {
             vaultSdkSource.initializeCrypto(
                 userId = userId,
                 request = InitUserCryptoRequest(
-                    accountCryptographicState = createWrappedAccountCryptographicState(
-                        privateKey = privateKey,
-                        securityState = null,
-                        signedPublicKey = null,
-                        signingKey = null,
-                    ),
+                    accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
                     userId = userId,
                     kdfParams = kdf,
                     email = email,
@@ -2044,12 +1874,7 @@ class VaultLockManagerTest {
         } returns true.asSuccess()
 
         val result = vaultLockManager.unlockVault(
-            accountCryptographicState = createWrappedAccountCryptographicState(
-                privateKey = privateKey,
-                securityState = null,
-                signedPublicKey = null,
-                signingKey = null,
-            ),
+            accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
             userId = userId,
             email = email,
             kdf = kdf,
@@ -2065,12 +1890,7 @@ class VaultLockManagerTest {
             vaultSdkSource.initializeCrypto(
                 userId = userId,
                 request = InitUserCryptoRequest(
-                    accountCryptographicState = createWrappedAccountCryptographicState(
-                        privateKey = privateKey,
-                        securityState = null,
-                        signedPublicKey = null,
-                        signingKey = null,
-                    ),
+                    accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
                     userId = userId,
                     kdfParams = kdf,
                     email = email,
@@ -2167,3 +1987,7 @@ private val MOCK_MASTER_PASSWORD_UNLOCK_DATA = MasterPasswordUnlockData(
     masterKeyWrappedUserKey = "12345",
     salt = "mockSalt",
 )
+
+private val ACCOUNT_CRYPTOGRAPHIC_STATE = WrappedAccountCryptographicState.V1(
+    privateKey = "privateKey",
+)

app/src/test/kotlin/com/x8bit/bitwarden/data/vault/manager/VaultSyncManagerTest.kt

@@ -35,6 +35,7 @@ import com.x8bit.bitwarden.data.auth.datasource.disk.util.FakeAuthDiskSource
 import com.x8bit.bitwarden.data.auth.manager.UserLogoutManager
 import com.x8bit.bitwarden.data.auth.manager.UserStateManager
 import com.x8bit.bitwarden.data.auth.repository.model.LogoutReason
+import com.x8bit.bitwarden.data.auth.repository.model.createMockWrappedAccountCryptographicState
 import com.x8bit.bitwarden.data.platform.datasource.disk.SettingsDiskSource
 import com.x8bit.bitwarden.data.platform.error.NoActiveUserException
 import com.x8bit.bitwarden.data.platform.manager.DatabaseSchemeManager
@@ -747,7 +748,10 @@ class VaultSyncManagerTest {
                 ),
             )
             fakeAuthDiskSource.assertUserState(userState = updatedUserState)
-            fakeAuthDiskSource.assertPrivateKey(userId = userId, privateKey = "mockPrivateKey-1")
+            fakeAuthDiskSource.assertAccountCryptographicState(
+                userId = userId,
+                accountCryptographicState = createMockWrappedAccountCryptographicState(number = 1),
+            )
             fakeAuthDiskSource.assertOrganizationKeys(
                 userId = userId,
                 organizationKeys = mapOf(userId to "mockKey-1"),

app/src/test/kotlin/com/x8bit/bitwarden/data/vault/repository/VaultRepositoryTest.kt

@@ -32,6 +32,7 @@ import com.x8bit.bitwarden.data.auth.datasource.disk.model.AccountTokensJson
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.UserStateJson
 import com.x8bit.bitwarden.data.auth.datasource.disk.util.FakeAuthDiskSource
 import com.x8bit.bitwarden.data.auth.datasource.sdk.util.toKdfRequestModel
+import com.x8bit.bitwarden.data.auth.repository.model.createMockWrappedAccountCryptographicState
 import com.x8bit.bitwarden.data.auth.repository.util.toSdkParams
 import com.x8bit.bitwarden.data.platform.error.NoActiveUserException
 import com.x8bit.bitwarden.data.vault.datasource.disk.VaultDiskSource
@@ -58,7 +59,6 @@ import com.x8bit.bitwarden.data.vault.repository.model.ImportCredentialsResult
 import com.x8bit.bitwarden.data.vault.repository.model.SendData
 import com.x8bit.bitwarden.data.vault.repository.model.VaultData
 import com.x8bit.bitwarden.data.vault.repository.model.VaultUnlockResult
-import com.x8bit.bitwarden.data.vault.repository.util.createWrappedAccountCryptographicState
 import com.x8bit.bitwarden.data.vault.repository.util.toEncryptedSdkCipher
 import com.x8bit.bitwarden.data.vault.repository.util.toSdkMasterPasswordUnlock
 import com.x8bit.bitwarden.ui.vault.feature.verificationcode.util.createVerificationCodeItem
@@ -199,7 +199,6 @@ class VaultRepositoryTest {
     fun `unlockVaultWithBiometrics with failure to decode biometrics key should return BiometricDecodingError`() =
         runTest {
             val userId = MOCK_USER_STATE.activeUserId
-            val privateKey = "mockPrivateKey-1"
             val biometricsKey = "asdf1234"
             fakeAuthDiskSource.userState = MOCK_USER_STATE
             val initVector = byteArrayOf(2, 2)
@@ -210,7 +209,10 @@ class VaultRepositoryTest {
             fakeAuthDiskSource.apply {
                 storeUserBiometricInitVector(userId = userId, iv = initVector)
                 storeUserBiometricUnlockKey(userId = userId, biometricsKey = biometricsKey)
-                storePrivateKey(userId = userId, privateKey = privateKey)
+                storeAccountCryptographicState(
+                    userId = userId,
+                    accountCryptographicState = MOCK_ACCOUNT_CRYPTOGRAPHIC_STATE,
+                )
             }
 
             val result = vaultRepository.unlockVaultWithBiometrics(cipher = cipher)
@@ -250,7 +252,6 @@ class VaultRepositoryTest {
     fun `unlockVaultWithBiometrics with an IV and VaultLockManager Success should store the updated key and IV and unlock for the current user and return Success`() =
         runTest {
             val userId = MOCK_USER_STATE.activeUserId
-            val privateKey = "mockPrivateKey-1"
             val biometricsKey = "asdf1234"
             fakeAuthDiskSource.userState = MOCK_USER_STATE
             val encryptedBytes = byteArrayOf(1, 1)
@@ -260,12 +261,7 @@ class VaultRepositoryTest {
             }
             coEvery {
                 vaultLockManager.unlockVault(
-                    accountCryptographicState = createWrappedAccountCryptographicState(
-                        privateKey = privateKey,
-                        securityState = null,
-                        signedPublicKey = null,
-                        signingKey = null,
-                    ),
+                    accountCryptographicState = MOCK_ACCOUNT_CRYPTOGRAPHIC_STATE,
                     userId = userId,
                     email = "email",
                     kdf = MOCK_PROFILE.toSdkParams(),
@@ -278,7 +274,10 @@ class VaultRepositoryTest {
             fakeAuthDiskSource.apply {
                 storeUserBiometricInitVector(userId = userId, iv = initVector)
                 storeUserBiometricUnlockKey(userId = userId, biometricsKey = biometricsKey)
-                storePrivateKey(userId = userId, privateKey = privateKey)
+                storeAccountCryptographicState(
+                    userId = userId,
+                    accountCryptographicState = MOCK_ACCOUNT_CRYPTOGRAPHIC_STATE,
+                )
             }
 
             val result = vaultRepository.unlockVaultWithBiometrics(cipher = cipher)
@@ -286,12 +285,7 @@ class VaultRepositoryTest {
             assertEquals(VaultUnlockResult.Success, result)
             coVerify(exactly = 1) {
                 vaultLockManager.unlockVault(
-                    accountCryptographicState = createWrappedAccountCryptographicState(
-                        privateKey = privateKey,
-                        securityState = null,
-                        signedPublicKey = null,
-                        signingKey = null,
-                    ),
+                    accountCryptographicState = MOCK_ACCOUNT_CRYPTOGRAPHIC_STATE,
                     userId = userId,
                     email = "email",
                     kdf = MOCK_PROFILE.toSdkParams(),
@@ -308,7 +302,6 @@ class VaultRepositoryTest {
     fun `unlockVaultWithBiometrics with VaultLockManager Success and a stored encrypted pin should unlock for the current user, derive a new pin-protected key, and return Success`() =
         runTest {
             val userId = MOCK_USER_STATE.activeUserId
-            val privateKey = "mockPrivateKey-1"
             val biometricsKey = "asdf1234"
             fakeAuthDiskSource.userState = MOCK_USER_STATE
             val encryptedBytes = byteArrayOf(1, 1)
@@ -319,12 +312,7 @@ class VaultRepositoryTest {
             }
             coEvery {
                 vaultLockManager.unlockVault(
-                    accountCryptographicState = createWrappedAccountCryptographicState(
-                        privateKey = privateKey,
-                        securityState = null,
-                        signedPublicKey = null,
-                        signingKey = null,
-                    ),
+                    accountCryptographicState = MOCK_ACCOUNT_CRYPTOGRAPHIC_STATE,
                     userId = userId,
                     email = "email",
                     kdf = MOCK_PROFILE.toSdkParams(),
@@ -337,7 +325,10 @@ class VaultRepositoryTest {
             fakeAuthDiskSource.apply {
                 storeUserBiometricInitVector(userId = userId, iv = null)
                 storeUserBiometricUnlockKey(userId = userId, biometricsKey = biometricsKey)
-                storePrivateKey(userId = userId, privateKey = privateKey)
+                storeAccountCryptographicState(
+                    userId = userId,
+                    accountCryptographicState = MOCK_ACCOUNT_CRYPTOGRAPHIC_STATE,
+                )
             }
 
             val result = vaultRepository.unlockVaultWithBiometrics(cipher = cipher)
@@ -345,12 +336,7 @@ class VaultRepositoryTest {
             assertEquals(VaultUnlockResult.Success, result)
             coVerify {
                 vaultLockManager.unlockVault(
-                    accountCryptographicState = createWrappedAccountCryptographicState(
-                        privateKey = "mockPrivateKey-1",
-                        securityState = null,
-                        signedPublicKey = null,
-                        signingKey = null,
-                    ),
+                    accountCryptographicState = MOCK_ACCOUNT_CRYPTOGRAPHIC_STATE,
                     userId = userId,
                     email = "email",
                     kdf = MOCK_PROFILE.toSdkParams(),
@@ -378,16 +364,10 @@ class VaultRepositoryTest {
         runTest {
             val userId = MOCK_USER_STATE.activeUserId
             val authenticatorSyncUnlockKey = "asdf1234"
-            val privateKey = "mockPrivateKey-1"
             fakeAuthDiskSource.userState = MOCK_USER_STATE
             coEvery {
                 vaultLockManager.unlockVault(
-                    accountCryptographicState = createWrappedAccountCryptographicState(
-                        privateKey = privateKey,
-                        securityState = null,
-                        signedPublicKey = null,
-                        signingKey = null,
-                    ),
+                    accountCryptographicState = MOCK_ACCOUNT_CRYPTOGRAPHIC_STATE,
                     userId = userId,
                     email = "email",
                     kdf = MOCK_PROFILE.toSdkParams(),
@@ -402,7 +382,10 @@ class VaultRepositoryTest {
                     userId = userId,
                     authenticatorSyncUnlockKey = authenticatorSyncUnlockKey,
                 )
-                storePrivateKey(userId = userId, privateKey = privateKey)
+                storeAccountCryptographicState(
+                    userId = userId,
+                    accountCryptographicState = MOCK_ACCOUNT_CRYPTOGRAPHIC_STATE,
+                )
             }
 
             val result = vaultRepository.unlockVaultWithDecryptedUserKey(
@@ -412,12 +395,7 @@ class VaultRepositoryTest {
             assertEquals(VaultUnlockResult.Success, result)
             coVerify {
                 vaultLockManager.unlockVault(
-                    accountCryptographicState = createWrappedAccountCryptographicState(
-                        privateKey = "mockPrivateKey-1",
-                        securityState = null,
-                        signedPublicKey = null,
-                        signingKey = null,
-                    ),
+                    accountCryptographicState = MOCK_ACCOUNT_CRYPTOGRAPHIC_STATE,
                     userId = userId,
                     email = "email",
                     kdf = MOCK_PROFILE.toSdkParams(),
@@ -435,17 +413,11 @@ class VaultRepositoryTest {
         runTest {
             val userId = MOCK_USER_STATE.activeUserId
             val authenticatorSyncUnlockKey = "asdf1234"
-            val privateKey = "mockPrivateKey-1"
             fakeAuthDiskSource.userState = MOCK_USER_STATE
             val error = Throwable("Fail")
             coEvery {
                 vaultLockManager.unlockVault(
-                    accountCryptographicState = createWrappedAccountCryptographicState(
-                        privateKey = privateKey,
-                        securityState = null,
-                        signedPublicKey = null,
-                        signingKey = null,
-                    ),
+                    accountCryptographicState = MOCK_ACCOUNT_CRYPTOGRAPHIC_STATE,
                     userId = userId,
                     email = "email",
                     kdf = MOCK_PROFILE.toSdkParams(),
@@ -460,7 +432,10 @@ class VaultRepositoryTest {
                     userId = userId,
                     authenticatorSyncUnlockKey = authenticatorSyncUnlockKey,
                 )
-                storePrivateKey(userId = userId, privateKey = privateKey)
+                storeAccountCryptographicState(
+                    userId = userId,
+                    accountCryptographicState = MOCK_ACCOUNT_CRYPTOGRAPHIC_STATE,
+                )
             }
 
             val result = vaultRepository.unlockVaultWithDecryptedUserKey(
@@ -470,12 +445,7 @@ class VaultRepositoryTest {
             assertEquals(VaultUnlockResult.InvalidStateError(error = error), result)
             coVerify {
                 vaultLockManager.unlockVault(
-                    accountCryptographicState = createWrappedAccountCryptographicState(
-                        privateKey = "mockPrivateKey-1",
-                        securityState = null,
-                        signedPublicKey = null,
-                        signingKey = null,
-                    ),
+                    accountCryptographicState = MOCK_ACCOUNT_CRYPTOGRAPHIC_STATE,
                     userId = userId,
                     email = "email",
                     kdf = MOCK_PROFILE.toSdkParams(),
@@ -504,9 +474,9 @@ class VaultRepositoryTest {
     @Test
     fun `unlockVaultWithMasterPassword with missing private key should return InvalidStateError`() =
         runTest {
-            fakeAuthDiskSource.storePrivateKey(
+            fakeAuthDiskSource.storeAccountCryptographicState(
                 userId = "mockId-1",
-                privateKey = null,
+                accountCryptographicState = null,
             )
             fakeAuthDiskSource.userState = MOCK_USER_STATE
 
@@ -536,12 +506,7 @@ class VaultRepositoryTest {
             )
             coVerify {
                 vaultLockManager.unlockVault(
-                    accountCryptographicState = createWrappedAccountCryptographicState(
-                        privateKey = "mockPrivateKey-1",
-                        securityState = null,
-                        signedPublicKey = null,
-                        signingKey = null,
-                    ),
+                    accountCryptographicState = MOCK_ACCOUNT_CRYPTOGRAPHIC_STATE,
                     userId = userId,
                     email = "email",
                     kdf = MOCK_PROFILE.toSdkParams(),
@@ -585,16 +550,14 @@ class VaultRepositoryTest {
             )
 
             fakeAuthDiskSource.userState = userState
-            fakeAuthDiskSource.storePrivateKey(userId = userId, privateKey = "mockPrivateKey-1")
+            fakeAuthDiskSource.storeAccountCryptographicState(
+                userId = userId,
+                accountCryptographicState = MOCK_ACCOUNT_CRYPTOGRAPHIC_STATE,
+            )
 
             coEvery {
                 vaultLockManager.unlockVault(
-                    accountCryptographicState = createWrappedAccountCryptographicState(
-                        privateKey = "mockPrivateKey-1",
-                        securityState = null,
-                        signedPublicKey = null,
-                        signingKey = null,
-                    ),
+                    accountCryptographicState = MOCK_ACCOUNT_CRYPTOGRAPHIC_STATE,
                     userId = userId,
                     email = "email",
                     kdf = MOCK_PROFILE.toSdkParams(),
@@ -613,12 +576,7 @@ class VaultRepositoryTest {
             assertEquals(VaultUnlockResult.Success, result)
             coVerify {
                 vaultLockManager.unlockVault(
-                    accountCryptographicState = createWrappedAccountCryptographicState(
-                        privateKey = "mockPrivateKey-1",
-                        securityState = null,
-                        signedPublicKey = null,
-                        signingKey = null,
-                    ),
+                    accountCryptographicState = MOCK_ACCOUNT_CRYPTOGRAPHIC_STATE,
                     userId = userId,
                     email = "email",
                     kdf = MOCK_PROFILE.toSdkParams(),
@@ -637,7 +595,6 @@ class VaultRepositoryTest {
         runTest {
             val userId = "mockId-1"
             val masterPassword = "mockPassword-1"
-            val userKey = "mockUserKey-1"
             val userState = MOCK_USER_STATE.copy(
                 accounts = mapOf(
                     "mockId-1" to MOCK_ACCOUNT.copy(
@@ -648,7 +605,10 @@ class VaultRepositoryTest {
                 ),
             )
             fakeAuthDiskSource.userState = userState
-            fakeAuthDiskSource.storePrivateKey(userId = userId, privateKey = "mockPrivateKey-1")
+            fakeAuthDiskSource.storeAccountCryptographicState(
+                userId = userId,
+                accountCryptographicState = MOCK_ACCOUNT_CRYPTOGRAPHIC_STATE,
+            )
 
             val result = vaultRepository.unlockVaultWithMasterPassword(
                 masterPassword = masterPassword,
@@ -678,12 +638,7 @@ class VaultRepositoryTest {
             )
             coVerify {
                 vaultLockManager.unlockVault(
-                    accountCryptographicState = createWrappedAccountCryptographicState(
-                        privateKey = "mockPrivateKey-1",
-                        securityState = null,
-                        signedPublicKey = null,
-                        signingKey = null,
-                    ),
+                    accountCryptographicState = MOCK_ACCOUNT_CRYPTOGRAPHIC_STATE,
                     userId = userId,
                     email = "email",
                     kdf = MOCK_PROFILE.toSdkParams(),
@@ -724,9 +679,9 @@ class VaultRepositoryTest {
                 userId = "mockId-1",
                 pinProtectedUserKeyEnvelope = null,
             )
-            fakeAuthDiskSource.storePrivateKey(
+            fakeAuthDiskSource.storeAccountCryptographicState(
                 userId = "mockId-1",
-                privateKey = "mockPrivateKey-1",
+                accountCryptographicState = MOCK_ACCOUNT_CRYPTOGRAPHIC_STATE,
             )
             fakeAuthDiskSource.userState = MOCK_USER_STATE
 
@@ -748,9 +703,9 @@ class VaultRepositoryTest {
             userId = "mockId-1",
             pinProtectedUserKeyEnvelope = null,
         )
-        fakeAuthDiskSource.storePrivateKey(
+        fakeAuthDiskSource.storeAccountCryptographicState(
             userId = "mockId-1",
-            privateKey = null,
+            accountCryptographicState = null,
         )
         fakeAuthDiskSource.userState = MOCK_USER_STATE
 
@@ -778,12 +733,7 @@ class VaultRepositoryTest {
             )
             coVerify {
                 vaultLockManager.unlockVault(
-                    accountCryptographicState = createWrappedAccountCryptographicState(
-                        privateKey = "mockPrivateKey-1",
-                        securityState = null,
-                        signedPublicKey = null,
-                        signingKey = null,
-                    ),
+                    accountCryptographicState = MOCK_ACCOUNT_CRYPTOGRAPHIC_STATE,
                     userId = userId,
                     email = "email",
                     kdf = MOCK_PROFILE.toSdkParams(),
@@ -816,12 +766,7 @@ class VaultRepositoryTest {
             )
             coVerify {
                 vaultLockManager.unlockVault(
-                    accountCryptographicState = createWrappedAccountCryptographicState(
-                        privateKey = "mockPrivateKey-1",
-                        securityState = null,
-                        signedPublicKey = null,
-                        signingKey = null,
-                    ),
+                    accountCryptographicState = MOCK_ACCOUNT_CRYPTOGRAPHIC_STATE,
                     userId = userId,
                     email = "email",
                     kdf = MOCK_PROFILE.toSdkParams(),
@@ -850,12 +795,7 @@ class VaultRepositoryTest {
             )
             coVerify {
                 vaultLockManager.unlockVault(
-                    accountCryptographicState = createWrappedAccountCryptographicState(
-                        privateKey = "mockPrivateKey-1",
-                        securityState = null,
-                        signedPublicKey = null,
-                        signingKey = null,
-                    ),
+                    accountCryptographicState = MOCK_ACCOUNT_CRYPTOGRAPHIC_STATE,
                     userId = userId,
                     email = "email",
                     kdf = MOCK_PROFILE.toSdkParams(),
@@ -1637,9 +1577,9 @@ class VaultRepositoryTest {
                 sendList = listOf(createMockSdkSend(number = 1)),
             )
         } returns listOf(createMockSendView(number = 1)).asSuccess()
-        fakeAuthDiskSource.storePrivateKey(
+        fakeAuthDiskSource.storeAccountCryptographicState(
             userId = userId,
-            privateKey = "mockPrivateKey-1",
+            accountCryptographicState = MOCK_ACCOUNT_CRYPTOGRAPHIC_STATE,
         )
         fakeAuthDiskSource.storePinProtectedUserKey(
             userId = userId,
@@ -1658,12 +1598,7 @@ class VaultRepositoryTest {
         // Master password unlock
         coEvery {
             vaultLockManager.unlockVault(
-                accountCryptographicState = createWrappedAccountCryptographicState(
-                    privateKey = "mockPrivateKey-1",
-                    securityState = null,
-                    signedPublicKey = null,
-                    signingKey = null,
-                ),
+                accountCryptographicState = MOCK_ACCOUNT_CRYPTOGRAPHIC_STATE,
                 userId = userId,
                 email = "email",
                 kdf = MOCK_PROFILE.toSdkParams(),
@@ -1682,12 +1617,7 @@ class VaultRepositoryTest {
         // PIN unlock
         coEvery {
             vaultLockManager.unlockVault(
-                accountCryptographicState = createWrappedAccountCryptographicState(
-                    privateKey = "mockPrivateKey-1",
-                    securityState = null,
-                    signedPublicKey = null,
-                    signingKey = null,
-                ),
+                accountCryptographicState = MOCK_ACCOUNT_CRYPTOGRAPHIC_STATE,
                 userId = userId,
                 email = "email",
                 kdf = MOCK_PROFILE.toSdkParams(),
@@ -1702,12 +1632,7 @@ class VaultRepositoryTest {
         // PIN ENVELOPE unlock
         coEvery {
             vaultLockManager.unlockVault(
-                accountCryptographicState = createWrappedAccountCryptographicState(
-                    privateKey = "mockPrivateKey-1",
-                    securityState = null,
-                    signedPublicKey = null,
-                    signingKey = null,
-                ),
+                accountCryptographicState = MOCK_ACCOUNT_CRYPTOGRAPHIC_STATE,
                 userId = userId,
                 email = "email",
                 kdf = MOCK_PROFILE.toSdkParams(),
@@ -1778,3 +1703,7 @@ private val MOCK_MASTER_PASSWORD_UNLOCK_DATA = MasterPasswordUnlockDataJson(
     kdf = MOCK_ACCOUNT.profile.toSdkParams().toKdfRequestModel(),
     masterKeyWrappedUserKey = "masterKeyWrappedUserKeyMock",
 )
+
+private val MOCK_ACCOUNT_CRYPTOGRAPHIC_STATE = createMockWrappedAccountCryptographicState(
+    number = 1,
+)

app/src/test/kotlin/com/x8bit/bitwarden/data/vault/repository/util/WrappedAccountCryptographicStateExtensionsTest.kt

@@ -1,111 +0,0 @@
-package com.x8bit.bitwarden.data.vault.repository.util
-
-import com.bitwarden.core.WrappedAccountCryptographicState
-import org.junit.jupiter.api.Assertions.assertEquals
-import org.junit.jupiter.api.Test
-
-private const val PRIVATE_KEY = "test-private-key"
-private const val SECURITY_STATE = "test-security-state"
-private const val SIGNING_KEY = "test-signing-key"
-private const val SIGNED_PUBLIC_KEY = "test-signed-public-key"
-
-class WrappedAccountCryptographicStateExtensionsTest {
-
-    @Suppress("MaxLineLength")
-    @Test
-    fun `createWrappedAccountCryptographicState returns V2 when securityState, signedPublicKey and signingKey are non-null`() {
-        val result = createWrappedAccountCryptographicState(
-            privateKey = PRIVATE_KEY,
-            securityState = SECURITY_STATE,
-            signedPublicKey = SIGNED_PUBLIC_KEY,
-            signingKey = SIGNING_KEY,
-        )
-
-        val v2State = result as WrappedAccountCryptographicState.V2
-        assertEquals(PRIVATE_KEY, v2State.privateKey)
-        assertEquals(SECURITY_STATE, v2State.securityState)
-        assertEquals(SIGNED_PUBLIC_KEY, v2State.signedPublicKey)
-        assertEquals(SIGNING_KEY, v2State.signingKey)
-    }
-
-    @Test
-    fun `createWrappedAccountCryptographicState returns V1 when securityState is null`() {
-        val result = createWrappedAccountCryptographicState(
-            privateKey = PRIVATE_KEY,
-            securityState = null,
-            signedPublicKey = SIGNED_PUBLIC_KEY,
-            signingKey = SIGNING_KEY,
-        )
-
-        val v1State = result as WrappedAccountCryptographicState.V1
-        assertEquals(PRIVATE_KEY, v1State.privateKey)
-    }
-
-    @Test
-    fun `createWrappedAccountCryptographicState returns V1 when signedPublicKey is null`() {
-        val result = createWrappedAccountCryptographicState(
-            privateKey = PRIVATE_KEY,
-            securityState = SECURITY_STATE,
-            signedPublicKey = null,
-            signingKey = SIGNING_KEY,
-        )
-
-        val v1State = result as WrappedAccountCryptographicState.V1
-        assertEquals(PRIVATE_KEY, v1State.privateKey)
-    }
-
-    @Test
-    fun `createWrappedAccountCryptographicState returns V1 when signingKey is null`() {
-        val result = createWrappedAccountCryptographicState(
-            privateKey = PRIVATE_KEY,
-            securityState = SECURITY_STATE,
-            signedPublicKey = SIGNED_PUBLIC_KEY,
-            signingKey = null,
-        )
-
-        val v1State = result as WrappedAccountCryptographicState.V1
-        assertEquals(PRIVATE_KEY, v1State.privateKey)
-    }
-
-    @Suppress("MaxLineLength")
-    @Test
-    fun `createWrappedAccountCryptographicState returns V1 when both securityState and signedPublicKey are null`() {
-        val result = createWrappedAccountCryptographicState(
-            privateKey = PRIVATE_KEY,
-            securityState = null,
-            signedPublicKey = null,
-            signingKey = SIGNING_KEY,
-        )
-
-        val v1State = result as WrappedAccountCryptographicState.V1
-        assertEquals(PRIVATE_KEY, v1State.privateKey)
-    }
-
-    @Suppress("MaxLineLength")
-    @Test
-    fun `createWrappedAccountCryptographicState returns V1 when both securityState and signingKey are null`() {
-        val result = createWrappedAccountCryptographicState(
-            privateKey = PRIVATE_KEY,
-            securityState = null,
-            signedPublicKey = SIGNED_PUBLIC_KEY,
-            signingKey = null,
-        )
-
-        val v1State = result as WrappedAccountCryptographicState.V1
-        assertEquals(PRIVATE_KEY, v1State.privateKey)
-    }
-
-    @Suppress("MaxLineLength")
-    @Test
-    fun `createWrappedAccountCryptographicState returns V1 when both signedPublicKey and signingKey are null`() {
-        val result = createWrappedAccountCryptographicState(
-            privateKey = PRIVATE_KEY,
-            securityState = SECURITY_STATE,
-            signedPublicKey = null,
-            signingKey = null,
-        )
-
-        val v1State = result as WrappedAccountCryptographicState.V1
-        assertEquals(PRIVATE_KEY, v1State.privateKey)
-    }
-}