PM-27234: feat: jit password v2 encryption (#6835)

app/src/main/kotlin/com/x8bit/bitwarden/data/auth/datasource/sdk/AuthSdkSource.kt

@@ -1,5 +1,6 @@
 package com.x8bit.bitwarden.data.auth.datasource.sdk
 
+import com.bitwarden.auth.JitMasterPasswordRegistrationResponse
 import com.bitwarden.auth.KeyConnectorRegistrationResult
 import com.bitwarden.core.AuthRequestResponse
 import com.bitwarden.core.KeyConnectorResponse
@@ -14,6 +15,21 @@ import com.x8bit.bitwarden.data.auth.datasource.sdk.model.PasswordStrength
  * Source of authentication information and functionality from the Bitwarden SDK.
  */
 interface AuthSdkSource {
+    /**
+     * Enrolls the user to master password unlock.
+     */
+    @Suppress("LongParameterList")
+    suspend fun postKeysForJitPasswordRegistration(
+        userId: String,
+        organizationId: String,
+        organizationPublicKey: String,
+        organizationSsoIdentifier: String,
+        salt: String,
+        masterPassword: String,
+        masterPasswordHint: String?,
+        shouldResetPasswordEnroll: Boolean,
+    ): Result<JitMasterPasswordRegistrationResponse>
+
     /**
      * Enrolls the user to key connector unlock.
      */

app/src/main/kotlin/com/x8bit/bitwarden/data/auth/datasource/sdk/AuthSdkSourceImpl.kt

@@ -1,5 +1,7 @@
 package com.x8bit.bitwarden.data.auth.datasource.sdk
 
+import com.bitwarden.auth.JitMasterPasswordRegistrationRequest
+import com.bitwarden.auth.JitMasterPasswordRegistrationResponse
 import com.bitwarden.auth.KeyConnectorRegistrationResult
 import com.bitwarden.core.AuthRequestResponse
 import com.bitwarden.core.FingerprintRequest
@@ -25,6 +27,33 @@ class AuthSdkSourceImpl(
 ) : BaseSdkSource(sdkClientManager = sdkClientManager),
     AuthSdkSource {
 
+    override suspend fun postKeysForJitPasswordRegistration(
+        userId: String,
+        organizationId: String,
+        organizationPublicKey: String,
+        organizationSsoIdentifier: String,
+        salt: String,
+        masterPassword: String,
+        masterPasswordHint: String?,
+        shouldResetPasswordEnroll: Boolean,
+    ): Result<JitMasterPasswordRegistrationResponse> = runCatchingWithLogs {
+        getClient(userId = userId)
+            .auth()
+            .registration()
+            .postKeysForJitPasswordRegistration(
+                request = JitMasterPasswordRegistrationRequest(
+                    orgId = organizationId,
+                    orgPublicKey = organizationPublicKey,
+                    userId = userId,
+                    organizationSsoIdentifier = organizationSsoIdentifier,
+                    salt = salt,
+                    masterPassword = masterPassword,
+                    masterPasswordHint = masterPasswordHint,
+                    resetPasswordEnroll = shouldResetPasswordEnroll,
+                ),
+            )
+    }
+
     override suspend fun postKeysForKeyConnectorRegistration(
         userId: String,
         accessToken: String,

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

@@ -5,6 +5,7 @@ import com.bitwarden.core.InitUserCryptoMethod
 import com.bitwarden.core.RegisterTdeKeyResponse
 import com.bitwarden.core.WrappedAccountCryptographicState
 import com.bitwarden.core.data.manager.dispatcher.DispatcherManager
+import com.bitwarden.core.data.manager.model.FlagKey
 import com.bitwarden.core.data.manager.toast.ToastManager
 import com.bitwarden.core.data.repository.error.MissingPropertyException
 import com.bitwarden.core.data.repository.util.bufferedMutableSharedFlow
@@ -100,8 +101,10 @@ 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.toOrganizations
 import com.x8bit.bitwarden.data.auth.repository.util.toRemovedPasswordUserStateJson
@@ -115,6 +118,7 @@ import com.x8bit.bitwarden.data.auth.util.toSdkParams
 import com.x8bit.bitwarden.data.platform.datasource.disk.SettingsDiskSource
 import com.x8bit.bitwarden.data.platform.error.NoActiveUserException
 import com.x8bit.bitwarden.data.platform.manager.BiometricsEncryptionManager
+import com.x8bit.bitwarden.data.platform.manager.FeatureFlagManager
 import com.x8bit.bitwarden.data.platform.manager.LogsManager
 import com.x8bit.bitwarden.data.platform.manager.PolicyManager
 import com.x8bit.bitwarden.data.platform.manager.PushManager
@@ -145,6 +149,7 @@ import kotlinx.coroutines.flow.merge
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.receiveAsFlow
 import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.withContext
 import timber.log.Timber
 import java.time.Clock
 import javax.inject.Singleton
@@ -178,9 +183,10 @@ class AuthRepositoryImpl(
     private val userStateManager: UserStateManager,
     private val kdfManager: KdfManager,
     private val toastManager: ToastManager,
+    private val featureFlagManager: FeatureFlagManager,
     logsManager: LogsManager,
     pushManager: PushManager,
-    dispatcherManager: DispatcherManager,
+    private val dispatcherManager: DispatcherManager,
 ) : AuthRepository,
     AuthRequestManager by authRequestManager,
     BiometricsEncryptionManager by biometricsEncryptionManager,
@@ -1105,85 +1111,73 @@ class AuthRepositoryImpl(
             )
     }
 
-    @Suppress("LongMethod")
     override suspend fun setPassword(
         organizationIdentifier: String,
         password: String,
         passwordHint: String?,
     ): SetPasswordResult {
-        val activeAccount = authDiskSource
-            .userState
-            ?.activeAccount
+        val profile = authDiskSource.userState?.activeAccount?.profile
             ?: return SetPasswordResult.Error(error = NoActiveUserException())
-        val userId = activeAccount.profile.userId
-
-        // Update the saved master password hash.
-        val passwordHash = authSdkSource
-            .hashPassword(
-                email = activeAccount.profile.email,
-                password = password,
-                kdf = activeAccount.profile.toSdkParams(),
-                purpose = HashPurpose.SERVER_AUTHORIZATION,
-            )
-            .getOrElse { return@setPassword SetPasswordResult.Error(error = it) }
-
-        return when (activeAccount.profile.forcePasswordResetReason) {
+        return when (profile.forcePasswordResetReason) {
             ForcePasswordResetReason.TDE_USER_WITHOUT_PASSWORD_HAS_PASSWORD_RESET_PERMISSION -> {
-                vaultSdkSource
-                    .updatePassword(userId = userId, newPassword = password)
-                    .map { it.newKey to null }
+                setUpdatedPassword(
+                    profile = profile,
+                    organizationIdentifier = organizationIdentifier,
+                    password = password,
+                    passwordHint = passwordHint,
+                )
             }
 
             ForcePasswordResetReason.ADMIN_FORCE_PASSWORD_RESET,
             ForcePasswordResetReason.WEAK_MASTER_PASSWORD_ON_LOGIN,
             null,
                 -> {
-                authSdkSource
-                    .makeRegisterKeys(
-                        email = activeAccount.profile.email,
-                        password = password,
-                        kdf = activeAccount.profile.toSdkParams(),
-                    )
-                    .map { it.encryptedUserKey to it.keys }
+                setPasswordForJit(
+                    profile = profile,
+                    organizationIdentifier = organizationIdentifier,
+                    password = password,
+                    passwordHint = passwordHint,
+                )
             }
         }
-            .flatMap { (encryptedUserKey, rsaKeys) ->
+    }
+
+    private suspend fun setUpdatedPassword(
+        profile: AccountJson.Profile,
+        organizationIdentifier: String,
+        password: String,
+        passwordHint: String?,
+    ): SetPasswordResult {
+        val userId = profile.userId
+        return vaultSdkSource
+            .updatePassword(userId = userId, newPassword = password)
+            .flatMap { response ->
                 accountsService
                     .setPassword(
                         body = SetPasswordRequestJson(
-                            passwordHash = passwordHash,
+                            passwordHash = response.passwordHash,
                             passwordHint = passwordHint,
                             organizationIdentifier = organizationIdentifier,
-                            kdfIterations = activeAccount.profile.kdfIterations,
-                            kdfMemory = activeAccount.profile.kdfMemory,
-                            kdfParallelism = activeAccount.profile.kdfParallelism,
-                            kdfType = activeAccount.profile.kdfType,
-                            key = encryptedUserKey,
-                            keys = rsaKeys?.let {
-                                RegisterRequestJson.Keys(
-                                    publicKey = it.public,
-                                    encryptedPrivateKey = it.private,
-                                )
-                            },
+                            kdfIterations = profile.kdfIterations,
+                            kdfMemory = profile.kdfMemory,
+                            kdfParallelism = profile.kdfParallelism,
+                            kdfType = profile.kdfType,
+                            key = response.newKey,
+                            keys = null,
                         ),
                     )
                     .onSuccess {
-                        rsaKeys?.private?.let {
-                            // 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(userId = userId, privateKey = it)
-                        }
-                        authDiskSource.storeUserKey(userId = userId, userKey = encryptedUserKey)
+                        authDiskSource.storeUserKey(userId = userId, userKey = response.newKey)
                     }
+                    .map { response.passwordHash }
             }
-            .flatMap {
+            .flatMap { masterPasswordHash ->
                 when (val result = vaultRepository.unlockVaultWithMasterPassword(password)) {
                     is VaultUnlockResult.Success -> {
                         enrollUserInPasswordReset(
                             userId = userId,
                             organizationIdentifier = organizationIdentifier,
-                            passwordHash = passwordHash,
+                            passwordHash = masterPasswordHash,
                         )
                     }
 
@@ -1194,8 +1188,155 @@ class AuthRepositoryImpl(
                 }
             }
             .onSuccess {
-                authDiskSource.storeMasterPasswordHash(userId = userId, passwordHash = passwordHash)
-                authDiskSource.userState = authDiskSource.userState?.toUserStateJsonWithPassword()
+                authDiskSource.userState = authDiskSource.userState?.toUserStateJsonWithPassword(
+                    masterPasswordUnlock = null,
+                )
+                this.organizationIdentifier = null
+            }
+            .fold(
+                onFailure = { SetPasswordResult.Error(error = it) },
+                onSuccess = { SetPasswordResult.Success },
+            )
+    }
+
+    private suspend fun setPasswordForJit(
+        profile: AccountJson.Profile,
+        organizationIdentifier: String,
+        password: String,
+        passwordHint: String?,
+    ): SetPasswordResult {
+        if (!featureFlagManager.getFeatureFlag(FlagKey.V2EncryptionJitPassword)) {
+            return setPasswordForJitV1(
+                profile = profile,
+                organizationIdentifier = organizationIdentifier,
+                password = password,
+                passwordHint = passwordHint,
+            )
+        }
+        val userId = profile.userId
+        return organizationService
+            .getOrganizationAutoEnrollStatus(organizationIdentifier = organizationIdentifier)
+            .flatMap { enrollStatus ->
+                organizationService
+                    .getOrganizationKeys(organizationId = enrollStatus.organizationId)
+                    .map { orgKeys -> enrollStatus to orgKeys }
+            }
+            .flatMap { (enrollStatus, orgKeys) ->
+                withContext(dispatcherManager.io) {
+                    authSdkSource.postKeysForJitPasswordRegistration(
+                        userId = userId,
+                        organizationId = enrollStatus.organizationId,
+                        organizationPublicKey = orgKeys.publicKey,
+                        organizationSsoIdentifier = organizationIdentifier,
+                        salt = profile.email,
+                        masterPassword = password,
+                        masterPasswordHint = passwordHint,
+                        shouldResetPasswordEnroll = enrollStatus.isResetPasswordEnabled,
+                    )
+                }
+            }
+            .onSuccess { response ->
+                authDiskSource.storeAccountKeys(
+                    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,
+                )
+                authDiskSource.userState = authDiskSource.userState?.toUserStateJsonWithPassword(
+                    masterPasswordUnlock = response.masterPasswordUnlock,
+                )
+                this.organizationIdentifier = null
+            }
+            .flatMap { response ->
+                // Logging in with the password instead of the decrypted userKey will store
+                // the master password hash automatically.
+                when (val result = vaultRepository.unlockVaultWithMasterPassword(password)) {
+                    VaultUnlockResult.Success -> response.asSuccess()
+                    is VaultUnlockError -> {
+                        (result.error ?: IllegalStateException("Failed to unlock vault"))
+                            .asFailure()
+                    }
+                }
+            }
+            .fold(
+                onFailure = { SetPasswordResult.Error(error = it) },
+                onSuccess = { SetPasswordResult.Success },
+            )
+    }
+
+    @Suppress("LongMethod")
+    private suspend fun setPasswordForJitV1(
+        profile: AccountJson.Profile,
+        organizationIdentifier: String,
+        password: String,
+        passwordHint: String?,
+    ): SetPasswordResult {
+        val userId = profile.userId
+        return authSdkSource
+            .makeRegisterKeys(
+                email = profile.email,
+                password = password,
+                kdf = profile.toSdkParams(),
+            )
+            .flatMap { response ->
+                accountsService
+                    .setPassword(
+                        body = SetPasswordRequestJson(
+                            passwordHash = response.masterPasswordHash,
+                            passwordHint = passwordHint,
+                            organizationIdentifier = organizationIdentifier,
+                            kdfIterations = profile.kdfIterations,
+                            kdfMemory = profile.kdfMemory,
+                            kdfParallelism = profile.kdfParallelism,
+                            kdfType = profile.kdfType,
+                            key = response.encryptedUserKey,
+                            keys = RegisterRequestJson.Keys(
+                                publicKey = response.keys.public,
+                                encryptedPrivateKey = response.keys.private,
+                            ),
+                        ),
+                    )
+                    .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(
+                            userId = userId,
+                            privateKey = response.keys.private,
+                        )
+                        authDiskSource.storeUserKey(
+                            userId = userId,
+                            userKey = response.encryptedUserKey,
+                        )
+                    }
+                    .map { response.masterPasswordHash }
+            }
+            .flatMap { masterPasswordHash ->
+                when (val result = vaultRepository.unlockVaultWithMasterPassword(password)) {
+                    is VaultUnlockResult.Success -> {
+                        enrollUserInPasswordReset(
+                            userId = userId,
+                            organizationIdentifier = organizationIdentifier,
+                            passwordHash = masterPasswordHash,
+                        )
+                    }
+
+                    is VaultUnlockError -> {
+                        (result.error ?: IllegalStateException("Failed to unlock vault"))
+                            .asFailure()
+                    }
+                }
+            }
+            .onSuccess {
+                authDiskSource.userState = authDiskSource.userState?.toUserStateJsonWithPassword(
+                    masterPasswordUnlock = null,
+                )
                 this.organizationIdentifier = null
             }
             .fold(

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

@@ -21,6 +21,7 @@ import com.x8bit.bitwarden.data.auth.repository.AuthRepository
 import com.x8bit.bitwarden.data.auth.repository.AuthRepositoryImpl
 import com.x8bit.bitwarden.data.platform.datasource.disk.SettingsDiskSource
 import com.x8bit.bitwarden.data.platform.manager.BiometricsEncryptionManager
+import com.x8bit.bitwarden.data.platform.manager.FeatureFlagManager
 import com.x8bit.bitwarden.data.platform.manager.FirstTimeActionManager
 import com.x8bit.bitwarden.data.platform.manager.LogsManager
 import com.x8bit.bitwarden.data.platform.manager.PolicyManager
@@ -73,6 +74,7 @@ object AuthRepositoryModule {
         userStateManager: UserStateManager,
         kdfManager: KdfManager,
         toastManager: ToastManager,
+        featureFlagManager: FeatureFlagManager,
     ): AuthRepository = AuthRepositoryImpl(
         clock = clock,
         accountsService = accountsService,
@@ -100,6 +102,7 @@ object AuthRepositoryModule {
         userStateManager = userStateManager,
         kdfManager = kdfManager,
         toastManager = toastManager,
+        featureFlagManager = featureFlagManager,
     )
 
     @Provides

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

@@ -1,7 +1,9 @@
 package com.x8bit.bitwarden.data.auth.repository.util
 
+import com.bitwarden.core.MasterPasswordUnlockData
 import com.bitwarden.data.repository.util.toEnvironmentUrlsOrDefault
 import com.bitwarden.network.model.KdfTypeJson
+import com.bitwarden.network.model.MasterPasswordUnlockDataJson
 import com.bitwarden.network.model.OrganizationType
 import com.bitwarden.network.model.PolicyTypeJson
 import com.bitwarden.network.model.SyncResponseJson
@@ -9,6 +11,7 @@ import com.bitwarden.network.model.UserDecryptionOptionsJson
 import com.bitwarden.ui.platform.base.util.toHexColorRepresentation
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.OnboardingStatus
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.UserStateJson
+import com.x8bit.bitwarden.data.auth.datasource.sdk.util.toKdfRequestModel
 import com.x8bit.bitwarden.data.auth.repository.model.UserAccountTokens
 import com.x8bit.bitwarden.data.auth.repository.model.UserKeyConnectorState
 import com.x8bit.bitwarden.data.auth.repository.model.UserOrganizations
@@ -108,20 +111,34 @@ fun UserStateJson.toUpdatedUserStateJson(
  * Updates the [UserStateJson] to set the `hasMasterPassword` value to `true` after a user sets
  * their password.
  */
-fun UserStateJson.toUserStateJsonWithPassword(): UserStateJson {
+fun UserStateJson.toUserStateJsonWithPassword(
+    masterPasswordUnlock: MasterPasswordUnlockData?,
+): UserStateJson {
     val account = this.activeAccount
     val profile = account.profile
+    val userDecryptionOptions = profile.userDecryptionOptions
+    val masterPasswordUnlockJson = masterPasswordUnlock
+        ?.let {
+            MasterPasswordUnlockDataJson(
+                salt = it.salt,
+                kdf = it.kdf.toKdfRequestModel(),
+                masterKeyWrappedUserKey = it.masterKeyWrappedUserKey,
+            )
+        }
+        ?: userDecryptionOptions?.masterPasswordUnlock
     val updatedProfile = profile
         .copy(
             forcePasswordResetReason = null,
-            userDecryptionOptions = profile
-                .userDecryptionOptions
-                ?.copy(hasMasterPassword = true)
+            userDecryptionOptions = userDecryptionOptions
+                ?.copy(
+                    hasMasterPassword = true,
+                    masterPasswordUnlock = masterPasswordUnlockJson,
+                )
                 ?: UserDecryptionOptionsJson(
                     hasMasterPassword = true,
                     keyConnectorUserDecryptionOptions = null,
                     trustedDeviceUserDecryptionOptions = null,
-                    masterPasswordUnlock = null,
+                    masterPasswordUnlock = masterPasswordUnlockJson,
                 ),
         )
     val updatedAccount = account.copy(profile = updatedProfile)

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

@@ -0,0 +1,41 @@
+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.SecurityState
+import com.bitwarden.network.model.AccountKeysJson.SignatureKeyPair
+
+/**
+ * The user's encryption private key, wrapped by the user key.
+ */
+val WrappedAccountCryptographicState.privateKey: String
+    get() = when (this) {
+        is WrappedAccountCryptographicState.V1 -> this.privateKey
+        is WrappedAccountCryptographicState.V2 -> this.privateKey
+    }
+
+/**
+ * Converts the [WrappedAccountCryptographicState] into a [AccountKeysJson].
+ *
+ * @receiver `WrappedAccountCryptographicState` to convert to `AccountEncryptionKeysJson`.
+ */
+val WrappedAccountCryptographicState.accountKeysJson: AccountKeysJson?
+    get() = when (this) {
+        is WrappedAccountCryptographicState.V1 -> null
+        is WrappedAccountCryptographicState.V2 -> AccountKeysJson(
+            publicKeyEncryptionKeyPair = PublicKeyEncryptionKeyPair(
+                publicKey = "",
+                signedPublicKey = this.signedPublicKey,
+                wrappedPrivateKey = this.privateKey,
+            ),
+            signatureKeyPair = SignatureKeyPair(
+                wrappedSigningKey = this.signingKey,
+                verifyingKey = "",
+            ),
+            securityState = SecurityState(
+                securityState = this.securityState,
+                securityVersion = 2,
+            ),
+        )
+    }

app/src/test/kotlin/com/x8bit/bitwarden/data/auth/datasource/sdk/AuthSdkSourceTest.kt

@@ -1,5 +1,7 @@
 package com.x8bit.bitwarden.data.auth.datasource.sdk
 
+import com.bitwarden.auth.JitMasterPasswordRegistrationRequest
+import com.bitwarden.auth.JitMasterPasswordRegistrationResponse
 import com.bitwarden.auth.KeyConnectorRegistrationResult
 import com.bitwarden.core.AuthRequestResponse
 import com.bitwarden.core.FingerprintRequest
@@ -47,6 +49,63 @@ class AuthSdkSourceTest {
         sdkClientManager = sdkClientManager,
     )
 
+    @Suppress("MaxLineLength")
+    @Test
+    fun `postKeysForJitPasswordRegistration should call SDK and return a Result with correct data`() =
+        runBlocking {
+            val userId = "userId"
+            val organizationId = "organizationId"
+            val organizationPublicKey = "organizationPublicKey"
+            val organizationSsoIdentifier = "organizationSsoIdentifier"
+            val salt = "salt"
+            val masterPassword = "masterPassword"
+            val masterPasswordHint = "masterPasswordHint"
+            val shouldResetPasswordEnroll = false
+            val expectedResult = mockk<JitMasterPasswordRegistrationResponse>()
+            coEvery { sdkClientManager.getOrCreateClient(userId = userId) } returns client
+            coEvery {
+                clientRegistration.postKeysForJitPasswordRegistration(
+                    request = JitMasterPasswordRegistrationRequest(
+                        orgId = organizationId,
+                        orgPublicKey = organizationPublicKey,
+                        organizationSsoIdentifier = organizationSsoIdentifier,
+                        userId = userId,
+                        salt = salt,
+                        masterPassword = masterPassword,
+                        masterPasswordHint = masterPasswordHint,
+                        resetPasswordEnroll = shouldResetPasswordEnroll,
+                    ),
+                )
+            } returns expectedResult
+
+            val result = authSkdSource.postKeysForJitPasswordRegistration(
+                organizationId = organizationId,
+                organizationPublicKey = organizationPublicKey,
+                organizationSsoIdentifier = organizationSsoIdentifier,
+                userId = userId,
+                salt = salt,
+                masterPassword = masterPassword,
+                masterPasswordHint = masterPasswordHint,
+                shouldResetPasswordEnroll = shouldResetPasswordEnroll,
+            )
+
+            assertEquals(expectedResult, result.getOrThrow())
+            coVerify(exactly = 1) {
+                clientRegistration.postKeysForJitPasswordRegistration(
+                    request = JitMasterPasswordRegistrationRequest(
+                        orgId = organizationId,
+                        orgPublicKey = organizationPublicKey,
+                        organizationSsoIdentifier = organizationSsoIdentifier,
+                        userId = userId,
+                        salt = salt,
+                        masterPassword = masterPassword,
+                        masterPasswordHint = masterPasswordHint,
+                        resetPasswordEnroll = shouldResetPasswordEnroll,
+                    ),
+                )
+            }
+        }
+
     @Suppress("MaxLineLength")
     @Test
     fun `postKeysForKeyConnectorRegistration should call SDK and return a Result with correct data`() =

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

@@ -1,6 +1,7 @@
 package com.x8bit.bitwarden.data.auth.repository
 
 import app.cash.turbine.test
+import com.bitwarden.auth.JitMasterPasswordRegistrationResponse
 import com.bitwarden.core.AuthRequestMethod
 import com.bitwarden.core.AuthRequestResponse
 import com.bitwarden.core.InitUserCryptoMethod
@@ -13,6 +14,7 @@ import com.bitwarden.core.UpdatePasswordResponse
 import com.bitwarden.core.WrappedAccountCryptographicState
 import com.bitwarden.core.data.manager.dispatcher.DispatcherManager
 import com.bitwarden.core.data.manager.dispatcher.FakeDispatcherManager
+import com.bitwarden.core.data.manager.model.FlagKey
 import com.bitwarden.core.data.manager.toast.ToastManager
 import com.bitwarden.core.data.repository.error.MissingPropertyException
 import com.bitwarden.core.data.repository.util.bufferedMutableSharedFlow
@@ -124,6 +126,7 @@ 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.toRemovedPasswordUserStateJson
 import com.x8bit.bitwarden.data.auth.repository.util.toSdkParams
 import com.x8bit.bitwarden.data.auth.repository.util.toUserState
@@ -131,6 +134,7 @@ import com.x8bit.bitwarden.data.auth.util.YubiKeyResult
 import com.x8bit.bitwarden.data.auth.util.toSdkParams
 import com.x8bit.bitwarden.data.platform.datasource.disk.util.FakeSettingsDiskSource
 import com.x8bit.bitwarden.data.platform.error.NoActiveUserException
+import com.x8bit.bitwarden.data.platform.manager.FeatureFlagManager
 import com.x8bit.bitwarden.data.platform.manager.LogsManager
 import com.x8bit.bitwarden.data.platform.manager.PolicyManager
 import com.x8bit.bitwarden.data.platform.manager.PushManager
@@ -289,6 +293,9 @@ class AuthRepositoryTest {
     private val toastManager: ToastManager = mockk {
         every { show(messageId = any(), duration = any()) } just runs
     }
+    private val featureFlagManager: FeatureFlagManager = mockk {
+        every { getFeatureFlag(FlagKey.V2EncryptionJitPassword) } returns true
+    }
 
     private val repository: AuthRepository = AuthRepositoryImpl(
         clock = FIXED_CLOCK,
@@ -317,6 +324,7 @@ class AuthRepositoryTest {
         userStateManager = userStateManager,
         kdfManager = kdfManager,
         toastManager = toastManager,
+        featureFlagManager = featureFlagManager,
     )
 
     @BeforeEach
@@ -5682,73 +5690,331 @@ class AuthRepositoryTest {
     }
 
     @Test
-    fun `setPassword with authSdkSource hashPassword failure should return Error`() = runTest {
-        val password = "password"
-        val error = Throwable("Fail")
-        fakeAuthDiskSource.userState = SINGLE_USER_STATE_1
-        coEvery {
-            authSdkSource.hashPassword(
-                email = EMAIL,
+    fun `setPassword with authSdkSource makeRegisterKeys failure should return Error for v1`() =
+        runTest {
+            every {
+                featureFlagManager.getFeatureFlag(FlagKey.V2EncryptionJitPassword)
+            } returns false
+            val password = "password"
+            val error = Throwable("Fail")
+            val kdf = SINGLE_USER_STATE_1.activeAccount.profile.toSdkParams()
+            fakeAuthDiskSource.userState = SINGLE_USER_STATE_1
+            coEvery {
+                authSdkSource.makeRegisterKeys(
+                    email = EMAIL,
+                    password = password,
+                    kdf = kdf,
+                )
+            } returns error.asFailure()
+
+            val result = repository.setPassword(
+                organizationIdentifier = "organizationId",
                 password = password,
-                kdf = SINGLE_USER_STATE_1.activeAccount.profile.toSdkParams(),
-                purpose = HashPurpose.SERVER_AUTHORIZATION,
+                passwordHint = "passwordHint",
             )
-        } returns error.asFailure()
 
-        val result = repository.setPassword(
-            organizationIdentifier = "organizationId",
-            password = password,
-            passwordHint = "passwordHint",
-        )
-
-        assertEquals(SetPasswordResult.Error(error = error), result)
-        fakeAuthDiskSource.assertMasterPasswordHash(userId = USER_ID_1, passwordHash = null)
-        fakeAuthDiskSource.assertPrivateKey(userId = USER_ID_1, privateKey = null)
-        fakeAuthDiskSource.assertUserKey(userId = USER_ID_1, userKey = null)
-    }
+            assertEquals(SetPasswordResult.Error(error = error), result)
+            fakeAuthDiskSource.assertPrivateKey(userId = USER_ID_1, privateKey = null)
+            fakeAuthDiskSource.assertAccountKeys(userId = USER_ID_1, accountKeys = null)
+            fakeAuthDiskSource.assertUserKey(userId = USER_ID_1, userKey = null)
+        }
 
     @Test
-    fun `setPassword with authSdkSource makeRegisterKeys failure should return Error`() = runTest {
-        val password = "password"
-        val passwordHash = "passwordHash"
-        val error = Throwable("Fail")
-        val kdf = SINGLE_USER_STATE_1.activeAccount.profile.toSdkParams()
-        fakeAuthDiskSource.userState = SINGLE_USER_STATE_1
-        coEvery {
-            authSdkSource.hashPassword(
-                email = EMAIL,
-                password = password,
-                kdf = kdf,
-                purpose = HashPurpose.SERVER_AUTHORIZATION,
-            )
-        } returns passwordHash.asSuccess()
-        coEvery {
-            authSdkSource.makeRegisterKeys(
-                email = EMAIL,
-                password = password,
-                kdf = kdf,
-            )
-        } returns error.asFailure()
+    fun `setPassword with getOrganizationAutoEnrollStatus failure should return Error`() =
+        runTest {
+            val error = Throwable("Fail")
+            val organizationIdentifier = "organizationIdentifier"
+            fakeAuthDiskSource.userState = SINGLE_USER_STATE_1
+            coEvery {
+                organizationService.getOrganizationAutoEnrollStatus(organizationIdentifier)
+            } returns error.asFailure()
 
-        val result = repository.setPassword(
-            organizationIdentifier = "organizationId",
-            password = password,
-            passwordHint = "passwordHint",
-        )
+            val result = repository.setPassword(
+                organizationIdentifier = organizationIdentifier,
+                password = "password",
+                passwordHint = "passwordHint",
+            )
 
-        assertEquals(SetPasswordResult.Error(error = error), result)
-        fakeAuthDiskSource.assertMasterPasswordHash(userId = USER_ID_1, passwordHash = null)
-        fakeAuthDiskSource.assertPrivateKey(userId = USER_ID_1, privateKey = null)
-        fakeAuthDiskSource.assertAccountKeys(userId = USER_ID_1, accountKeys = null)
-        fakeAuthDiskSource.assertUserKey(userId = USER_ID_1, userKey = null)
-    }
+            assertEquals(SetPasswordResult.Error(error = error), result)
+            coVerify(exactly = 1) {
+                organizationService.getOrganizationAutoEnrollStatus(organizationIdentifier)
+            }
+            fakeAuthDiskSource.assertPrivateKey(userId = USER_ID_1, privateKey = null)
+            fakeAuthDiskSource.assertAccountKeys(userId = USER_ID_1, accountKeys = null)
+        }
+
+    @Test
+    fun `setPassword with getOrganizationKeys failure should return Error`() =
+        runTest {
+            val error = Throwable("Fail")
+            val organizationIdentifier = "organizationIdentifier"
+            val organizationId = "organizationId"
+            val enrollResponse = OrganizationAutoEnrollStatusResponseJson(
+                organizationId = organizationId,
+                isResetPasswordEnabled = true,
+            )
+            fakeAuthDiskSource.userState = SINGLE_USER_STATE_1
+            coEvery {
+                organizationService.getOrganizationAutoEnrollStatus(organizationIdentifier)
+            } returns enrollResponse.asSuccess()
+            coEvery {
+                organizationService.getOrganizationKeys(organizationId)
+            } returns error.asFailure()
+
+            val result = repository.setPassword(
+                organizationIdentifier = organizationIdentifier,
+                password = "password",
+                passwordHint = "passwordHint",
+            )
+
+            assertEquals(SetPasswordResult.Error(error = error), result)
+            coVerify(exactly = 1) {
+                organizationService.getOrganizationAutoEnrollStatus(organizationIdentifier)
+                organizationService.getOrganizationKeys(organizationId)
+            }
+            fakeAuthDiskSource.assertPrivateKey(userId = USER_ID_1, privateKey = null)
+            fakeAuthDiskSource.assertAccountKeys(userId = USER_ID_1, accountKeys = null)
+        }
+
+    @Test
+    fun `setPassword with postKeysForJitPasswordRegistration failure should return Error`() =
+        runTest {
+            val error = Throwable("Fail")
+            val password = "password"
+            val passwordHint = "passwordHint"
+            val organizationIdentifier = "organizationIdentifier"
+            val organizationId = "organizationId"
+            val isResetPasswordEnabled = true
+            val enrollResponse = OrganizationAutoEnrollStatusResponseJson(
+                organizationId = organizationId,
+                isResetPasswordEnabled = isResetPasswordEnabled,
+            )
+            val orgPublicKey = "orgPublicKey"
+            val orgKeysResponse = OrganizationKeysResponseJson(
+                privateKey = "orgPrivateKey",
+                publicKey = orgPublicKey,
+            )
+            fakeAuthDiskSource.userState = SINGLE_USER_STATE_1
+            coEvery {
+                organizationService.getOrganizationAutoEnrollStatus(organizationIdentifier)
+            } returns enrollResponse.asSuccess()
+            coEvery {
+                organizationService.getOrganizationKeys(organizationId)
+            } returns orgKeysResponse.asSuccess()
+            coEvery {
+                authSdkSource.postKeysForJitPasswordRegistration(
+                    userId = USER_ID_1,
+                    organizationId = organizationId,
+                    organizationPublicKey = orgPublicKey,
+                    organizationSsoIdentifier = organizationIdentifier,
+                    salt = EMAIL,
+                    masterPassword = password,
+                    masterPasswordHint = passwordHint,
+                    shouldResetPasswordEnroll = isResetPasswordEnabled,
+                )
+            } returns error.asFailure()
+
+            val result = repository.setPassword(
+                organizationIdentifier = organizationIdentifier,
+                password = password,
+                passwordHint = passwordHint,
+            )
+
+            assertEquals(SetPasswordResult.Error(error = error), result)
+            coVerify(exactly = 1) {
+                organizationService.getOrganizationAutoEnrollStatus(organizationIdentifier)
+                organizationService.getOrganizationKeys(organizationId)
+                authSdkSource.postKeysForJitPasswordRegistration(
+                    userId = USER_ID_1,
+                    organizationId = organizationId,
+                    organizationPublicKey = orgPublicKey,
+                    organizationSsoIdentifier = organizationIdentifier,
+                    salt = EMAIL,
+                    masterPassword = password,
+                    masterPasswordHint = passwordHint,
+                    shouldResetPasswordEnroll = isResetPasswordEnabled,
+                )
+            }
+            fakeAuthDiskSource.assertPrivateKey(userId = USER_ID_1, privateKey = null)
+            fakeAuthDiskSource.assertAccountKeys(userId = USER_ID_1, accountKeys = null)
+        }
+
+    @Test
+    fun `setPassword with unlockVaultWithMasterPassword failure should return Error`() =
+        runTest {
+            val error = Throwable("Fail")
+            val unlockError = VaultUnlockResult.GenericError(error = error)
+            val password = "password"
+            val passwordHint = "passwordHint"
+            val organizationIdentifier = "organizationIdentifier"
+            val organizationId = "organizationId"
+            val isResetPasswordEnabled = true
+            val enrollResponse = OrganizationAutoEnrollStatusResponseJson(
+                organizationId = organizationId,
+                isResetPasswordEnabled = isResetPasswordEnabled,
+            )
+            val orgPublicKey = "orgPublicKey"
+            val orgKeysResponse = OrganizationKeysResponseJson(
+                privateKey = "orgPrivateKey",
+                publicKey = orgPublicKey,
+            )
+            val privateKey = "privateKey"
+            val accountCryptographicState = WrappedAccountCryptographicState.V2(
+                privateKey = privateKey,
+                securityState = "securityState",
+                signedPublicKey = "signedPublicKey",
+                signingKey = "signingKey",
+            )
+            val jitMasterPasswordResponse = JitMasterPasswordRegistrationResponse(
+                accountCryptographicState = accountCryptographicState,
+                masterPasswordUnlock = MasterPasswordUnlockData(
+                    kdf = Kdf.Pbkdf2(iterations = 1u),
+                    masterKeyWrappedUserKey = "masterKeyWrappedUserKey",
+                    salt = EMAIL,
+                ),
+                userKey = "userKey",
+            )
+            fakeAuthDiskSource.userState = SINGLE_USER_STATE_1
+            coEvery {
+                organizationService.getOrganizationAutoEnrollStatus(organizationIdentifier)
+            } returns enrollResponse.asSuccess()
+            coEvery {
+                organizationService.getOrganizationKeys(organizationId)
+            } returns orgKeysResponse.asSuccess()
+            coEvery {
+                authSdkSource.postKeysForJitPasswordRegistration(
+                    userId = USER_ID_1,
+                    organizationId = organizationId,
+                    organizationPublicKey = orgPublicKey,
+                    organizationSsoIdentifier = organizationIdentifier,
+                    salt = EMAIL,
+                    masterPassword = password,
+                    masterPasswordHint = passwordHint,
+                    shouldResetPasswordEnroll = isResetPasswordEnabled,
+                )
+            } returns jitMasterPasswordResponse.asSuccess()
+            coEvery {
+                vaultRepository.unlockVaultWithMasterPassword(password)
+            } returns unlockError
+
+            val result = repository.setPassword(
+                organizationIdentifier = organizationIdentifier,
+                password = password,
+                passwordHint = passwordHint,
+            )
+
+            assertEquals(SetPasswordResult.Error(error = error), result)
+            coVerify(exactly = 1) {
+                organizationService.getOrganizationAutoEnrollStatus(organizationIdentifier)
+                organizationService.getOrganizationKeys(organizationId)
+                authSdkSource.postKeysForJitPasswordRegistration(
+                    userId = USER_ID_1,
+                    organizationId = organizationId,
+                    organizationPublicKey = orgPublicKey,
+                    organizationSsoIdentifier = organizationIdentifier,
+                    salt = EMAIL,
+                    masterPassword = password,
+                    masterPasswordHint = passwordHint,
+                    shouldResetPasswordEnroll = isResetPasswordEnabled,
+                )
+                vaultRepository.unlockVaultWithMasterPassword(password)
+            }
+            fakeAuthDiskSource.assertPrivateKey(userId = USER_ID_1, privateKey = privateKey)
+            fakeAuthDiskSource.assertAccountKeys(
+                userId = USER_ID_1,
+                accountKeys = accountCryptographicState.accountKeysJson,
+            )
+        }
+
+    @Test
+    fun `setPassword with no failures should return Success`() =
+        runTest {
+            val password = "password"
+            val passwordHint = "passwordHint"
+            val organizationIdentifier = "organizationIdentifier"
+            val organizationId = "organizationId"
+            val isResetPasswordEnabled = true
+            val enrollResponse = OrganizationAutoEnrollStatusResponseJson(
+                organizationId = organizationId,
+                isResetPasswordEnabled = isResetPasswordEnabled,
+            )
+            val orgPublicKey = "orgPublicKey"
+            val orgKeysResponse = OrganizationKeysResponseJson(
+                privateKey = "orgPrivateKey",
+                publicKey = orgPublicKey,
+            )
+            val privateKey = "privateKey"
+            val accountCryptographicState = WrappedAccountCryptographicState.V2(
+                privateKey = privateKey,
+                securityState = "securityState",
+                signedPublicKey = "signedPublicKey",
+                signingKey = "signingKey",
+            )
+            val jitMasterPasswordResponse = JitMasterPasswordRegistrationResponse(
+                accountCryptographicState = accountCryptographicState,
+                masterPasswordUnlock = MasterPasswordUnlockData(
+                    kdf = Kdf.Pbkdf2(iterations = 1u),
+                    masterKeyWrappedUserKey = "masterKeyWrappedUserKey",
+                    salt = EMAIL,
+                ),
+                userKey = "userKey",
+            )
+            fakeAuthDiskSource.userState = SINGLE_USER_STATE_1
+            coEvery {
+                organizationService.getOrganizationAutoEnrollStatus(organizationIdentifier)
+            } returns enrollResponse.asSuccess()
+            coEvery {
+                organizationService.getOrganizationKeys(organizationId)
+            } returns orgKeysResponse.asSuccess()
+            coEvery {
+                authSdkSource.postKeysForJitPasswordRegistration(
+                    userId = USER_ID_1,
+                    organizationId = organizationId,
+                    organizationPublicKey = orgPublicKey,
+                    organizationSsoIdentifier = organizationIdentifier,
+                    salt = EMAIL,
+                    masterPassword = password,
+                    masterPasswordHint = passwordHint,
+                    shouldResetPasswordEnroll = isResetPasswordEnabled,
+                )
+            } returns jitMasterPasswordResponse.asSuccess()
+            coEvery {
+                vaultRepository.unlockVaultWithMasterPassword(password)
+            } returns VaultUnlockResult.Success
+
+            val result = repository.setPassword(
+                organizationIdentifier = organizationIdentifier,
+                password = password,
+                passwordHint = passwordHint,
+            )
+
+            assertEquals(SetPasswordResult.Success, result)
+            coVerify(exactly = 1) {
+                organizationService.getOrganizationAutoEnrollStatus(organizationIdentifier)
+                organizationService.getOrganizationKeys(organizationId)
+                authSdkSource.postKeysForJitPasswordRegistration(
+                    userId = USER_ID_1,
+                    organizationId = organizationId,
+                    organizationPublicKey = orgPublicKey,
+                    organizationSsoIdentifier = organizationIdentifier,
+                    salt = EMAIL,
+                    masterPassword = password,
+                    masterPasswordHint = passwordHint,
+                    shouldResetPasswordEnroll = isResetPasswordEnabled,
+                )
+                vaultRepository.unlockVaultWithMasterPassword(password)
+            }
+            fakeAuthDiskSource.assertPrivateKey(userId = USER_ID_1, privateKey = privateKey)
+            fakeAuthDiskSource.assertAccountKeys(
+                userId = USER_ID_1,
+                accountKeys = accountCryptographicState.accountKeysJson,
+            )
+        }
 
     @Test
     fun `setPassword with vaultSdkSource updatePassword failure should return Error`() = runTest {
         val password = "password"
-        val passwordHash = "passwordHash"
         val error = Throwable("Fail")
-        val kdf = SINGLE_USER_STATE_1.activeAccount.profile.toSdkParams()
         fakeAuthDiskSource.userState = SINGLE_USER_STATE_1.copy(
             accounts = mapOf(
                 USER_ID_1 to ACCOUNT_1.copy(
@@ -5759,14 +6025,6 @@ class AuthRepositoryTest {
                 ),
             ),
         )
-        coEvery {
-            authSdkSource.hashPassword(
-                email = EMAIL,
-                password = password,
-                kdf = kdf,
-                purpose = HashPurpose.SERVER_AUTHORIZATION,
-            )
-        } returns passwordHash.asSuccess()
         coEvery {
             vaultSdkSource.updatePassword(userId = USER_ID_1, newPassword = password)
         } returns error.asFailure()
@@ -5778,187 +6036,170 @@ class AuthRepositoryTest {
         )
 
         assertEquals(SetPasswordResult.Error(error = error), result)
-        fakeAuthDiskSource.assertMasterPasswordHash(userId = USER_ID_1, passwordHash = null)
         fakeAuthDiskSource.assertPrivateKey(userId = USER_ID_1, privateKey = null)
         fakeAuthDiskSource.assertAccountKeys(userId = USER_ID_1, accountKeys = null)
         fakeAuthDiskSource.assertUserKey(userId = USER_ID_1, userKey = null)
     }
 
     @Test
-    fun `setPassword with accountsService setPassword failure should return Error`() = runTest {
-        val password = "password"
-        val passwordHash = "passwordHash"
-        val passwordHint = "passwordHint"
-        val organizationId = ORGANIZATION_IDENTIFIER
-        val encryptedUserKey = "encryptedUserKey"
-        val privateRsaKey = "privateRsaKey"
-        val publicRsaKey = "publicRsaKey"
-        val profile = SINGLE_USER_STATE_1.activeAccount.profile
-        val kdf = profile.toSdkParams()
-        val error = Throwable("Fail")
-        val registerKeyResponse = RegisterKeyResponse(
-            masterPasswordHash = passwordHash,
-            encryptedUserKey = encryptedUserKey,
-            keys = RsaKeyPair(public = publicRsaKey, private = privateRsaKey),
-        )
-        val setPasswordRequestJson = SetPasswordRequestJson(
-            passwordHash = passwordHash,
-            passwordHint = passwordHint,
-            organizationIdentifier = organizationId,
-            kdfIterations = profile.kdfIterations,
-            kdfMemory = profile.kdfMemory,
-            kdfParallelism = profile.kdfParallelism,
-            kdfType = profile.kdfType,
-            key = encryptedUserKey,
-            keys = RegisterRequestJson.Keys(
-                publicKey = publicRsaKey,
-                encryptedPrivateKey = privateRsaKey,
-            ),
-        )
-        fakeAuthDiskSource.userState = SINGLE_USER_STATE_1
-        coEvery {
-            authSdkSource.hashPassword(
-                email = EMAIL,
-                password = password,
-                kdf = kdf,
-                purpose = HashPurpose.SERVER_AUTHORIZATION,
+    fun `setPassword with accountsService setPassword failure should return Error for v1`() =
+        runTest {
+            every {
+                featureFlagManager.getFeatureFlag(FlagKey.V2EncryptionJitPassword)
+            } returns false
+            val password = "password"
+            val passwordHash = "passwordHash"
+            val passwordHint = "passwordHint"
+            val organizationId = ORGANIZATION_IDENTIFIER
+            val encryptedUserKey = "encryptedUserKey"
+            val privateRsaKey = "privateRsaKey"
+            val publicRsaKey = "publicRsaKey"
+            val profile = SINGLE_USER_STATE_1.activeAccount.profile
+            val kdf = profile.toSdkParams()
+            val error = Throwable("Fail")
+            val registerKeyResponse = RegisterKeyResponse(
+                masterPasswordHash = passwordHash,
+                encryptedUserKey = encryptedUserKey,
+                keys = RsaKeyPair(public = publicRsaKey, private = privateRsaKey),
             )
-        } returns passwordHash.asSuccess()
-        coEvery {
-            authSdkSource.makeRegisterKeys(email = EMAIL, password = password, kdf = kdf)
-        } returns registerKeyResponse.asSuccess()
-        coEvery {
-            accountsService.setPassword(body = setPasswordRequestJson)
-        } returns error.asFailure()
-
-        val result = repository.setPassword(
-            organizationIdentifier = organizationId,
-            password = password,
-            passwordHint = passwordHint,
-        )
-
-        assertEquals(SetPasswordResult.Error(error = error), result)
-        fakeAuthDiskSource.assertMasterPasswordHash(userId = USER_ID_1, passwordHash = null)
-        fakeAuthDiskSource.assertPrivateKey(userId = USER_ID_1, privateKey = null)
-        fakeAuthDiskSource.assertAccountKeys(userId = USER_ID_1, accountKeys = null)
-        fakeAuthDiskSource.assertUserKey(userId = USER_ID_1, userKey = null)
-    }
-
-    @Test
-    fun `setPassword with accountsService setPassword success should return Success`() = runTest {
-        val password = "password"
-        val passwordHash = "passwordHash"
-        val passwordHint = "passwordHint"
-        val organizationIdentifier = ORGANIZATION_IDENTIFIER
-        val organizationId = "orgId"
-        val encryptedUserKey = "encryptedUserKey"
-        val privateRsaKey = "privateRsaKey"
-        val publicRsaKey = "publicRsaKey"
-        val publicOrgKey = "publicOrgKey"
-        val resetPasswordKey = "resetPasswordKey"
-        val profile = SINGLE_USER_STATE_1.activeAccount.profile
-        val kdf = profile.toSdkParams()
-        val registerKeyResponse = RegisterKeyResponse(
-            masterPasswordHash = passwordHash,
-            encryptedUserKey = encryptedUserKey,
-            keys = RsaKeyPair(public = publicRsaKey, private = privateRsaKey),
-        )
-        val setPasswordRequestJson = SetPasswordRequestJson(
-            passwordHash = passwordHash,
-            passwordHint = passwordHint,
-            organizationIdentifier = organizationIdentifier,
-            kdfIterations = profile.kdfIterations,
-            kdfMemory = profile.kdfMemory,
-            kdfParallelism = profile.kdfParallelism,
-            kdfType = profile.kdfType,
-            key = encryptedUserKey,
-            keys = RegisterRequestJson.Keys(
-                publicKey = publicRsaKey,
-                encryptedPrivateKey = privateRsaKey,
-            ),
-        )
-        fakeAuthDiskSource.userState = SINGLE_USER_STATE_1
-        coEvery {
-            authSdkSource.hashPassword(
-                email = EMAIL,
-                password = password,
-                kdf = kdf,
-                purpose = HashPurpose.SERVER_AUTHORIZATION,
-            )
-        } returns passwordHash.asSuccess()
-        coEvery {
-            authSdkSource.makeRegisterKeys(email = EMAIL, password = password, kdf = kdf)
-        } returns registerKeyResponse.asSuccess()
-        coEvery {
-            accountsService.setPassword(body = setPasswordRequestJson)
-        } returns Unit.asSuccess()
-        coEvery {
-            organizationService.getOrganizationAutoEnrollStatus(organizationIdentifier)
-        } returns OrganizationAutoEnrollStatusResponseJson(
-            organizationId = organizationId,
-            isResetPasswordEnabled = true,
-        )
-            .asSuccess()
-        coEvery {
-            organizationService.getOrganizationKeys(organizationId)
-        } returns OrganizationKeysResponseJson(
-            privateKey = "",
-            publicKey = publicOrgKey,
-        )
-            .asSuccess()
-        coEvery {
-            organizationService.organizationResetPasswordEnroll(
-                organizationId = organizationId,
-                userId = profile.userId,
+            val setPasswordRequestJson = SetPasswordRequestJson(
                 passwordHash = passwordHash,
-                resetPasswordKey = resetPasswordKey,
+                passwordHint = passwordHint,
+                organizationIdentifier = organizationId,
+                kdfIterations = profile.kdfIterations,
+                kdfMemory = profile.kdfMemory,
+                kdfParallelism = profile.kdfParallelism,
+                kdfType = profile.kdfType,
+                key = encryptedUserKey,
+                keys = RegisterRequestJson.Keys(
+                    publicKey = publicRsaKey,
+                    encryptedPrivateKey = privateRsaKey,
+                ),
             )
-        } returns Unit.asSuccess()
-        coEvery {
-            vaultSdkSource.getResetPasswordKey(
-                orgPublicKey = publicOrgKey,
-                userId = profile.userId,
-            )
-        } returns resetPasswordKey.asSuccess()
-        coEvery {
-            vaultRepository.unlockVaultWithMasterPassword(password)
-        } returns VaultUnlockResult.Success
+            fakeAuthDiskSource.userState = SINGLE_USER_STATE_1
+            coEvery {
+                authSdkSource.makeRegisterKeys(email = EMAIL, password = password, kdf = kdf)
+            } returns registerKeyResponse.asSuccess()
+            coEvery {
+                accountsService.setPassword(body = setPasswordRequestJson)
+            } returns error.asFailure()
 
-        val result = repository.setPassword(
-            organizationIdentifier = organizationIdentifier,
-            password = password,
-            passwordHint = passwordHint,
-        )
-
-        assertEquals(SetPasswordResult.Success, result)
-        fakeAuthDiskSource.assertMasterPasswordHash(userId = USER_ID_1, passwordHash = passwordHash)
-        fakeAuthDiskSource.assertPrivateKey(userId = USER_ID_1, privateKey = privateRsaKey)
-        fakeAuthDiskSource.assertUserKey(userId = USER_ID_1, userKey = encryptedUserKey)
-        fakeAuthDiskSource.assertUserState(SINGLE_USER_STATE_1_WITH_PASS)
-        coVerify {
-            authSdkSource.hashPassword(
-                email = EMAIL,
+            val result = repository.setPassword(
+                organizationIdentifier = organizationId,
                 password = password,
-                kdf = kdf,
-                purpose = HashPurpose.SERVER_AUTHORIZATION,
-            )
-            authSdkSource.makeRegisterKeys(email = EMAIL, password = password, kdf = kdf)
-            accountsService.setPassword(body = setPasswordRequestJson)
-            organizationService.getOrganizationAutoEnrollStatus(organizationIdentifier)
-            organizationService.getOrganizationKeys(organizationId)
-            organizationService.organizationResetPasswordEnroll(
-                organizationId = organizationId,
-                userId = profile.userId,
-                passwordHash = passwordHash,
-                resetPasswordKey = resetPasswordKey,
-            )
-            vaultRepository.unlockVaultWithMasterPassword(password)
-            vaultSdkSource.getResetPasswordKey(
-                orgPublicKey = publicOrgKey,
-                userId = profile.userId,
+                passwordHint = passwordHint,
             )
+
+            assertEquals(SetPasswordResult.Error(error = error), result)
+            fakeAuthDiskSource.assertPrivateKey(userId = USER_ID_1, privateKey = null)
+            fakeAuthDiskSource.assertAccountKeys(userId = USER_ID_1, accountKeys = null)
+            fakeAuthDiskSource.assertUserKey(userId = USER_ID_1, userKey = null)
+        }
+
+    @Test
+    fun `setPassword with accountsService setPassword success should return Success for v1`() =
+        runTest {
+            every {
+                featureFlagManager.getFeatureFlag(FlagKey.V2EncryptionJitPassword)
+            } returns false
+            val password = "password"
+            val passwordHash = "passwordHash"
+            val passwordHint = "passwordHint"
+            val organizationIdentifier = ORGANIZATION_IDENTIFIER
+            val organizationId = "orgId"
+            val encryptedUserKey = "encryptedUserKey"
+            val privateRsaKey = "privateRsaKey"
+            val publicRsaKey = "publicRsaKey"
+            val publicOrgKey = "publicOrgKey"
+            val resetPasswordKey = "resetPasswordKey"
+            val profile = SINGLE_USER_STATE_1.activeAccount.profile
+            val kdf = profile.toSdkParams()
+            val registerKeyResponse = RegisterKeyResponse(
+                masterPasswordHash = passwordHash,
+                encryptedUserKey = encryptedUserKey,
+                keys = RsaKeyPair(public = publicRsaKey, private = privateRsaKey),
+            )
+            val setPasswordRequestJson = SetPasswordRequestJson(
+                passwordHash = passwordHash,
+                passwordHint = passwordHint,
+                organizationIdentifier = organizationIdentifier,
+                kdfIterations = profile.kdfIterations,
+                kdfMemory = profile.kdfMemory,
+                kdfParallelism = profile.kdfParallelism,
+                kdfType = profile.kdfType,
+                key = encryptedUserKey,
+                keys = RegisterRequestJson.Keys(
+                    publicKey = publicRsaKey,
+                    encryptedPrivateKey = privateRsaKey,
+                ),
+            )
+            fakeAuthDiskSource.userState = SINGLE_USER_STATE_1
+            coEvery {
+                authSdkSource.makeRegisterKeys(email = EMAIL, password = password, kdf = kdf)
+            } returns registerKeyResponse.asSuccess()
+            coEvery {
+                accountsService.setPassword(body = setPasswordRequestJson)
+            } returns Unit.asSuccess()
+            coEvery {
+                organizationService.getOrganizationAutoEnrollStatus(organizationIdentifier)
+            } returns OrganizationAutoEnrollStatusResponseJson(
+                organizationId = organizationId,
+                isResetPasswordEnabled = true,
+            )
+                .asSuccess()
+            coEvery {
+                organizationService.getOrganizationKeys(organizationId)
+            } returns OrganizationKeysResponseJson(
+                privateKey = "",
+                publicKey = publicOrgKey,
+            )
+                .asSuccess()
+            coEvery {
+                organizationService.organizationResetPasswordEnroll(
+                    organizationId = organizationId,
+                    userId = profile.userId,
+                    passwordHash = passwordHash,
+                    resetPasswordKey = resetPasswordKey,
+                )
+            } returns Unit.asSuccess()
+            coEvery {
+                vaultSdkSource.getResetPasswordKey(
+                    orgPublicKey = publicOrgKey,
+                    userId = profile.userId,
+                )
+            } returns resetPasswordKey.asSuccess()
+            coEvery {
+                vaultRepository.unlockVaultWithMasterPassword(password)
+            } returns VaultUnlockResult.Success
+
+            val result = repository.setPassword(
+                organizationIdentifier = organizationIdentifier,
+                password = password,
+                passwordHint = passwordHint,
+            )
+
+            assertEquals(SetPasswordResult.Success, result)
+            fakeAuthDiskSource.assertPrivateKey(userId = USER_ID_1, privateKey = privateRsaKey)
+            fakeAuthDiskSource.assertUserKey(userId = USER_ID_1, userKey = encryptedUserKey)
+            fakeAuthDiskSource.assertUserState(SINGLE_USER_STATE_1_WITH_PASS)
+            coVerify(exactly = 1) {
+                authSdkSource.makeRegisterKeys(email = EMAIL, password = password, kdf = kdf)
+                accountsService.setPassword(body = setPasswordRequestJson)
+                organizationService.getOrganizationAutoEnrollStatus(organizationIdentifier)
+                organizationService.getOrganizationKeys(organizationId)
+                organizationService.organizationResetPasswordEnroll(
+                    organizationId = organizationId,
+                    userId = profile.userId,
+                    passwordHash = passwordHash,
+                    resetPasswordKey = resetPasswordKey,
+                )
+                vaultRepository.unlockVaultWithMasterPassword(password)
+                vaultSdkSource.getResetPasswordKey(
+                    orgPublicKey = publicOrgKey,
+                    userId = profile.userId,
+                )
+            }
         }
-    }
 
     @Test
     fun `setPassword with updatePassword success should return Success`() = runTest {
@@ -5981,7 +6222,6 @@ class AuthRepositoryTest {
             ),
         )
         val profile = userState.activeAccount.profile
-        val kdf = profile.toSdkParams()
         val updatePasswordResponse = UpdatePasswordResponse(
             passwordHash = passwordHash,
             newKey = encryptedUserKey,
@@ -5998,14 +6238,6 @@ class AuthRepositoryTest {
             keys = null,
         )
         fakeAuthDiskSource.userState = userState
-        coEvery {
-            authSdkSource.hashPassword(
-                email = EMAIL,
-                password = password,
-                kdf = kdf,
-                purpose = HashPurpose.SERVER_AUTHORIZATION,
-            )
-        } returns passwordHash.asSuccess()
         coEvery {
             vaultSdkSource.updatePassword(userId = USER_ID_1, newPassword = password)
         } returns updatePasswordResponse.asSuccess()
@@ -6051,21 +6283,11 @@ class AuthRepositoryTest {
         )
 
         assertEquals(SetPasswordResult.Success, result)
-        fakeAuthDiskSource.assertMasterPasswordHash(
-            userId = USER_ID_1,
-            passwordHash = passwordHash,
-        )
         fakeAuthDiskSource.assertPrivateKey(userId = USER_ID_1, privateKey = null)
         fakeAuthDiskSource.assertAccountKeys(userId = USER_ID_1, accountKeys = null)
         fakeAuthDiskSource.assertUserKey(userId = USER_ID_1, userKey = encryptedUserKey)
         fakeAuthDiskSource.assertUserState(SINGLE_USER_STATE_1_WITH_PASS)
         coVerify(exactly = 1) {
-            authSdkSource.hashPassword(
-                email = EMAIL,
-                password = password,
-                kdf = kdf,
-                purpose = HashPurpose.SERVER_AUTHORIZATION,
-            )
             vaultSdkSource.updatePassword(userId = USER_ID_1, newPassword = password)
             accountsService.setPassword(body = setPasswordRequestJson)
             organizationService.getOrganizationAutoEnrollStatus(organizationIdentifier)
@@ -6085,95 +6307,84 @@ class AuthRepositoryTest {
     }
 
     @Test
-    fun `setPassword with unlockVaultWithMasterPassword error should return Failure`() = runTest {
-        val password = "password"
-        val passwordHash = "passwordHash"
-        val passwordHint = "passwordHint"
-        val organizationIdentifier = ORGANIZATION_IDENTIFIER
-        val organizationId = "orgId"
-        val encryptedUserKey = "encryptedUserKey"
-        val privateRsaKey = "privateRsaKey"
-        val publicRsaKey = "publicRsaKey"
-        val publicOrgKey = "publicOrgKey"
-        val resetPasswordKey = "resetPasswordKey"
-        val profile = SINGLE_USER_STATE_1.activeAccount.profile
-        val kdf = profile.toSdkParams()
-        val registerKeyResponse = RegisterKeyResponse(
-            masterPasswordHash = passwordHash,
-            encryptedUserKey = encryptedUserKey,
-            keys = RsaKeyPair(public = publicRsaKey, private = privateRsaKey),
-        )
-        val setPasswordRequestJson = SetPasswordRequestJson(
-            passwordHash = passwordHash,
-            passwordHint = passwordHint,
-            organizationIdentifier = organizationIdentifier,
-            kdfIterations = profile.kdfIterations,
-            kdfMemory = profile.kdfMemory,
-            kdfParallelism = profile.kdfParallelism,
-            kdfType = profile.kdfType,
-            key = encryptedUserKey,
-            keys = RegisterRequestJson.Keys(
-                publicKey = publicRsaKey,
-                encryptedPrivateKey = privateRsaKey,
-            ),
-        )
-        fakeAuthDiskSource.userState = SINGLE_USER_STATE_1
-        coEvery {
-            authSdkSource.hashPassword(
-                email = EMAIL,
-                password = password,
-                kdf = kdf,
-                purpose = HashPurpose.SERVER_AUTHORIZATION,
+    fun `setPassword with unlockVaultWithMasterPassword error should return Failure for v1`() =
+        runTest {
+            every {
+                featureFlagManager.getFeatureFlag(FlagKey.V2EncryptionJitPassword)
+            } returns false
+            val password = "password"
+            val passwordHash = "passwordHash"
+            val passwordHint = "passwordHint"
+            val organizationIdentifier = ORGANIZATION_IDENTIFIER
+            val organizationId = "orgId"
+            val encryptedUserKey = "encryptedUserKey"
+            val privateRsaKey = "privateRsaKey"
+            val publicRsaKey = "publicRsaKey"
+            val publicOrgKey = "publicOrgKey"
+            val resetPasswordKey = "resetPasswordKey"
+            val profile = SINGLE_USER_STATE_1.activeAccount.profile
+            val kdf = profile.toSdkParams()
+            val registerKeyResponse = RegisterKeyResponse(
+                masterPasswordHash = passwordHash,
+                encryptedUserKey = encryptedUserKey,
+                keys = RsaKeyPair(public = publicRsaKey, private = privateRsaKey),
             )
-        } returns passwordHash.asSuccess()
-        coEvery {
-            authSdkSource.makeRegisterKeys(email = EMAIL, password = password, kdf = kdf)
-        } returns registerKeyResponse.asSuccess()
-        coEvery {
-            accountsService.setPassword(body = setPasswordRequestJson)
-        } returns Unit.asSuccess()
-        val error = Throwable("Fail")
-        coEvery {
-            vaultRepository.unlockVaultWithMasterPassword(password)
-        } returns VaultUnlockResult.GenericError(error = error)
-
-        val result = repository.setPassword(
-            organizationIdentifier = organizationIdentifier,
-            password = password,
-            passwordHint = passwordHint,
-        )
-
-        assertEquals(SetPasswordResult.Error(error = error), result)
-        fakeAuthDiskSource.assertMasterPasswordHash(userId = USER_ID_1, passwordHash = null)
-        fakeAuthDiskSource.assertPrivateKey(userId = USER_ID_1, privateKey = privateRsaKey)
-        fakeAuthDiskSource.assertUserKey(userId = USER_ID_1, userKey = encryptedUserKey)
-        fakeAuthDiskSource.assertUserState(SINGLE_USER_STATE_1)
-        coVerify {
-            authSdkSource.hashPassword(
-                email = EMAIL,
-                password = password,
-                kdf = kdf,
-                purpose = HashPurpose.SERVER_AUTHORIZATION,
-            )
-            authSdkSource.makeRegisterKeys(email = EMAIL, password = password, kdf = kdf)
-            accountsService.setPassword(body = setPasswordRequestJson)
-            vaultRepository.unlockVaultWithMasterPassword(password)
-        }
-        coVerify(exactly = 0) {
-            organizationService.getOrganizationAutoEnrollStatus(organizationIdentifier)
-            organizationService.getOrganizationKeys(organizationId)
-            organizationService.organizationResetPasswordEnroll(
-                organizationId = organizationId,
-                userId = profile.userId,
+            val setPasswordRequestJson = SetPasswordRequestJson(
                 passwordHash = passwordHash,
-                resetPasswordKey = resetPasswordKey,
+                passwordHint = passwordHint,
+                organizationIdentifier = organizationIdentifier,
+                kdfIterations = profile.kdfIterations,
+                kdfMemory = profile.kdfMemory,
+                kdfParallelism = profile.kdfParallelism,
+                kdfType = profile.kdfType,
+                key = encryptedUserKey,
+                keys = RegisterRequestJson.Keys(
+                    publicKey = publicRsaKey,
+                    encryptedPrivateKey = privateRsaKey,
+                ),
             )
-            vaultSdkSource.getResetPasswordKey(
-                orgPublicKey = publicOrgKey,
-                userId = profile.userId,
+            fakeAuthDiskSource.userState = SINGLE_USER_STATE_1
+            coEvery {
+                authSdkSource.makeRegisterKeys(email = EMAIL, password = password, kdf = kdf)
+            } returns registerKeyResponse.asSuccess()
+            coEvery {
+                accountsService.setPassword(body = setPasswordRequestJson)
+            } returns Unit.asSuccess()
+            val error = Throwable("Fail")
+            coEvery {
+                vaultRepository.unlockVaultWithMasterPassword(password)
+            } returns VaultUnlockResult.GenericError(error = error)
+
+            val result = repository.setPassword(
+                organizationIdentifier = organizationIdentifier,
+                password = password,
+                passwordHint = passwordHint,
             )
+
+            assertEquals(SetPasswordResult.Error(error = error), result)
+            fakeAuthDiskSource.assertPrivateKey(userId = USER_ID_1, privateKey = privateRsaKey)
+            fakeAuthDiskSource.assertUserKey(userId = USER_ID_1, userKey = encryptedUserKey)
+            fakeAuthDiskSource.assertUserState(SINGLE_USER_STATE_1)
+            coVerify(exactly = 1) {
+                authSdkSource.makeRegisterKeys(email = EMAIL, password = password, kdf = kdf)
+                accountsService.setPassword(body = setPasswordRequestJson)
+                vaultRepository.unlockVaultWithMasterPassword(password)
+            }
+            coVerify(exactly = 0) {
+                organizationService.getOrganizationAutoEnrollStatus(organizationIdentifier)
+                organizationService.getOrganizationKeys(organizationId)
+                organizationService.organizationResetPasswordEnroll(
+                    organizationId = organizationId,
+                    userId = profile.userId,
+                    passwordHash = passwordHash,
+                    resetPasswordKey = resetPasswordKey,
+                )
+                vaultSdkSource.getResetPasswordKey(
+                    orgPublicKey = publicOrgKey,
+                    userId = profile.userId,
+                )
+            }
         }
-    }
 
     @Test
     fun `passwordHintRequest with valid email should return Success`() = runTest {

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

@@ -1,5 +1,7 @@
 package com.x8bit.bitwarden.data.auth.repository.util
 
+import com.bitwarden.core.MasterPasswordUnlockData
+import com.bitwarden.crypto.Kdf
 import com.bitwarden.data.datasource.disk.model.EnvironmentUrlDataJson
 import com.bitwarden.data.repository.model.Environment
 import com.bitwarden.network.model.KdfJson
@@ -18,6 +20,7 @@ import com.x8bit.bitwarden.data.auth.datasource.disk.model.AccountTokensJson
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.ForcePasswordResetReason
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.OnboardingStatus
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.UserStateJson
+import com.x8bit.bitwarden.data.auth.datasource.sdk.util.toKdfRequestModel
 import com.x8bit.bitwarden.data.auth.repository.model.UserAccountTokens
 import com.x8bit.bitwarden.data.auth.repository.model.UserKeyConnectorState
 import com.x8bit.bitwarden.data.auth.repository.model.UserOrganizations
@@ -29,13 +32,27 @@ import com.x8bit.bitwarden.data.platform.manager.model.FirstTimeState
 import com.x8bit.bitwarden.data.vault.repository.model.VaultUnlockData
 import io.mockk.every
 import io.mockk.mockk
+import io.mockk.mockkStatic
+import io.mockk.unmockkStatic
+import org.junit.jupiter.api.AfterEach
 import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.BeforeEach
 import org.junit.jupiter.api.Test
 import java.time.Instant
 
 @Suppress("LargeClass")
 class UserStateJsonExtensionsTest {
 
+    @BeforeEach
+    fun setup() {
+        mockkStatic(Kdf::toKdfRequestModel)
+    }
+
+    @AfterEach
+    fun tearDown() {
+        unmockkStatic(Kdf::toKdfRequestModel)
+    }
+
     @Suppress("MaxLineLength")
     @Test
     fun `toUpdatedUserStateJson should do nothing for a non-matching account using toRemovedPasswordUserStateJson`() {
@@ -288,7 +305,73 @@ class UserStateJsonExtensionsTest {
                     "activeUserId" to originalAccount,
                 ),
             )
-                .toUserStateJsonWithPassword(),
+                .toUserStateJsonWithPassword(masterPasswordUnlock = null),
+        )
+    }
+
+    @Suppress("MaxLineLength")
+    @Test
+    fun `toUserStateJsonWithPassword with masterPasswordUnlock should update active account to set hasMasterPassword and masterPasswordUnlock`() {
+        val originalProfile = AccountJson.Profile(
+            userId = "activeUserId",
+            email = "email",
+            isEmailVerified = true,
+            name = "name",
+            stamp = null,
+            organizationId = null,
+            avatarColorHex = null,
+            hasPremium = true,
+            forcePasswordResetReason = ForcePasswordResetReason
+                .TDE_USER_WITHOUT_PASSWORD_HAS_PASSWORD_RESET_PERMISSION,
+            kdfType = KdfTypeJson.ARGON2_ID,
+            kdfIterations = 600000,
+            kdfMemory = 16,
+            kdfParallelism = 4,
+            userDecryptionOptions = null,
+            isTwoFactorEnabled = false,
+            creationDate = Instant.parse("2024-09-13T01:00:00.00Z"),
+        )
+        val originalAccount = AccountJson(
+            profile = originalProfile,
+            tokens = mockk(),
+            settings = mockk(),
+        )
+        val kdf = mockk<KdfJson>()
+        val masterKeyWrappedUserKey = "masterKeyWrappedUserKey"
+        val salt = "salt"
+        val masterPasswordUnlock = MasterPasswordUnlockData(
+            kdf = mockk<Kdf> { every { toKdfRequestModel() } returns kdf },
+            masterKeyWrappedUserKey = masterKeyWrappedUserKey,
+            salt = salt,
+        )
+        assertEquals(
+            UserStateJson(
+                activeUserId = "activeUserId",
+                accounts = mapOf(
+                    "activeUserId" to originalAccount.copy(
+                        profile = originalProfile.copy(
+                            forcePasswordResetReason = null,
+                            userDecryptionOptions = UserDecryptionOptionsJson(
+                                hasMasterPassword = true,
+                                keyConnectorUserDecryptionOptions = null,
+                                trustedDeviceUserDecryptionOptions = null,
+                                masterPasswordUnlock = MasterPasswordUnlockDataJson(
+                                    kdf = kdf,
+                                    masterKeyWrappedUserKey = masterKeyWrappedUserKey,
+                                    salt = salt,
+                                ),
+                            ),
+                        ),
+                    ),
+                ),
+            ),
+            UserStateJson(
+                activeUserId = "activeUserId",
+                accounts = mapOf(
+                    "activeUserId" to originalAccount,
+                ),
+            )
+                .toUserStateJsonWithPassword(masterPasswordUnlock = masterPasswordUnlock),
         )
     }
 
@@ -352,7 +435,7 @@ class UserStateJsonExtensionsTest {
                     "activeUserId" to originalAccount,
                 ),
             )
-                .toUserStateJsonWithPassword(),
+                .toUserStateJsonWithPassword(masterPasswordUnlock = null),
         )
     }
 

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

@@ -0,0 +1,54 @@
+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.SecurityState
+import com.bitwarden.network.model.AccountKeysJson.SignatureKeyPair
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.assertNull
+
+class WrappedAccountCryptographicStateExtensionsTest {
+    @Test
+    fun `privateKey should return correct value`() {
+        assertEquals("v1PrivateKey", V1_WRAPPED_ACCOUNT_CRYPTOGRAPHIC_STATE.privateKey)
+        assertEquals("v2PrivateKey", V2_WRAPPED_ACCOUNT_CRYPTOGRAPHIC_STATE.privateKey)
+    }
+
+    @Test
+    fun `accountKeysJson should return correct value`() {
+        assertNull(V1_WRAPPED_ACCOUNT_CRYPTOGRAPHIC_STATE.accountKeysJson)
+        assertEquals(
+            AccountKeysJson(
+                publicKeyEncryptionKeyPair = PublicKeyEncryptionKeyPair(
+                    publicKey = "",
+                    signedPublicKey = "signedPublicKey",
+                    wrappedPrivateKey = "v2PrivateKey",
+                ),
+                signatureKeyPair = SignatureKeyPair(
+                    wrappedSigningKey = "signingKey",
+                    verifyingKey = "",
+                ),
+                securityState = SecurityState(
+                    securityState = "securityState",
+                    securityVersion = 2,
+                ),
+            ),
+            V2_WRAPPED_ACCOUNT_CRYPTOGRAPHIC_STATE.accountKeysJson,
+        )
+    }
+}
+
+private val V1_WRAPPED_ACCOUNT_CRYPTOGRAPHIC_STATE: WrappedAccountCryptographicState =
+    WrappedAccountCryptographicState.V1(
+        privateKey = "v1PrivateKey",
+    )
+
+private val V2_WRAPPED_ACCOUNT_CRYPTOGRAPHIC_STATE: WrappedAccountCryptographicState =
+    WrappedAccountCryptographicState.V2(
+        privateKey = "v2PrivateKey",
+        securityState = "securityState",
+        signingKey = "signingKey",
+        signedPublicKey = "signedPublicKey",
+    )