BIT-2018: Support org reset password enrollment in JIT provisioning (#1159)

2025-03-16 03:08:50 +03:00 · 2024-03-19 10:56:44 -06:00 · 2024-03-19 10:56:44 -06:00 · bd58dac0ff
commit bd58dac0ff
parent be127f5d49
16 changed files with 532 additions and 133 deletions
--- a/app/src/main/java/com/x8bit/bitwarden/data/auth/datasource/network/api/AuthenticatedOrganizationApi.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/auth/datasource/network/api/AuthenticatedOrganizationApi.kt
@ -0,0 +1,40 @@
+package com.x8bit.bitwarden.data.auth.datasource.network.api
+
+import com.x8bit.bitwarden.data.auth.datasource.network.model.OrganizationAutoEnrollStatusResponseJson
+import com.x8bit.bitwarden.data.auth.datasource.network.model.OrganizationKeysResponseJson
+import com.x8bit.bitwarden.data.auth.datasource.network.model.OrganizationResetPasswordEnrollRequestJson
+import retrofit2.http.Body
+import retrofit2.http.GET
+import retrofit2.http.PUT
+import retrofit2.http.Path
+
+/**
+ * Defines raw calls under the authenticated /organizations API.
+ */
+interface AuthenticatedOrganizationApi {
+    /**
+     * Enrolls this user in the organization's password reset.
+     */
+    @PUT("/organizations/{orgId}/users/{userId}/reset-password-enrollment")
+    suspend fun organizationResetPasswordEnroll(
+        @Path("orgId") organizationId: String,
+        @Path("userId") userId: String,
+        @Body body: OrganizationResetPasswordEnrollRequestJson,
+    ): Result<Unit>
+
+    /**
+     * Checks whether this organization auto enrolls users in password reset.
+     */
+    @GET("/organizations/{identifier}/auto-enroll-status")
+    suspend fun getOrganizationAutoEnrollResponse(
+        @Path("identifier") organizationIdentifier: String,
+    ): Result<OrganizationAutoEnrollStatusResponseJson>
+
+    /**
+     * Gets the public and private keys for this organization.
+     */
+    @GET("/organizations/{id}/keys")
+    suspend fun getOrganizationKeys(
+        @Path("id") organizationId: String,
+    ): Result<OrganizationKeysResponseJson>
+}
--- a/app/src/main/java/com/x8bit/bitwarden/data/auth/datasource/network/di/AuthNetworkModule.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/auth/datasource/network/di/AuthNetworkModule.kt
@ -92,6 +92,7 @@ object AuthNetworkModule {
    fun providesOrganizationService(
        retrofits: Retrofits,
    ): OrganizationService = OrganizationServiceImpl(
+        authenticatedOrganizationApi = retrofits.authenticatedApiRetrofit.create(),
        organizationApi = retrofits.unauthenticatedApiRetrofit.create(),
    )
 }
--- a/app/src/main/java/com/x8bit/bitwarden/data/auth/datasource/network/model/OrganizationAutoEnrollStatusResponseJson.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/auth/datasource/network/model/OrganizationAutoEnrollStatusResponseJson.kt
@ -0,0 +1,17 @@
+package com.x8bit.bitwarden.data.auth.datasource.network.model
+
+import kotlinx.serialization.SerialName
+import kotlinx.serialization.Serializable
+
+/**
+ * Response object returned when requesting organization domain SSO details.
+ *
+ * @property organizationId The ID of this organization.
+ * @property isResetPasswordEnabled Indicates whether the auto-enroll reset password functionality
+ * is enabled.
+ */
+@Serializable
+data class OrganizationAutoEnrollStatusResponseJson(
+    @SerialName("id") val organizationId: String,
+    @SerialName("resetPasswordEnabled") val isResetPasswordEnabled: Boolean,
+)
--- a/app/src/main/java/com/x8bit/bitwarden/data/auth/datasource/network/model/OrganizationKeysResponseJson.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/auth/datasource/network/model/OrganizationKeysResponseJson.kt
@ -0,0 +1,16 @@
+package com.x8bit.bitwarden.data.auth.datasource.network.model
+
+import kotlinx.serialization.SerialName
+import kotlinx.serialization.Serializable
+
+/**
+ * Response object containing keys for this organization.
+ *
+ * @property privateKey The private key for this organization.
+ * @property publicKey The public key for this organization.
+ */
+@Serializable
+data class OrganizationKeysResponseJson(
+    @SerialName("privateKey") val privateKey: String?,
+    @SerialName("publicKey") val publicKey: String,
+)
--- a/app/src/main/java/com/x8bit/bitwarden/data/auth/datasource/network/model/OrganizationResetPasswordEnrollRequestJson.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/auth/datasource/network/model/OrganizationResetPasswordEnrollRequestJson.kt
@ -0,0 +1,16 @@
+package com.x8bit.bitwarden.data.auth.datasource.network.model
+
+import kotlinx.serialization.SerialName
+import kotlinx.serialization.Serializable
+
+/**
+ * Request body object when enrolling a user in reset password functionality for this organization.
+ *
+ * @param passwordHash The hash of this user's password.
+ * @param resetPasswordKey The key used for password reset.
+ */
+@Serializable
+data class OrganizationResetPasswordEnrollRequestJson(
+    @SerialName("masterPasswordHash") val passwordHash: String,
+    @SerialName("resetPasswordKey") val resetPasswordKey: String,
+)
--- a/app/src/main/java/com/x8bit/bitwarden/data/auth/datasource/network/service/OrganizationService.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/auth/datasource/network/service/OrganizationService.kt
@ -1,15 +1,41 @@
 package com.x8bit.bitwarden.data.auth.datasource.network.service

+import com.x8bit.bitwarden.data.auth.datasource.network.model.OrganizationAutoEnrollStatusResponseJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.OrganizationDomainSsoDetailsResponseJson
+import com.x8bit.bitwarden.data.auth.datasource.network.model.OrganizationKeysResponseJson

 /**
 * Provides an API for querying organization endpoints.
 */
 interface OrganizationService {
+    /**
+     * Enrolls a user with the given [userId] in this organizations reset password functionality.
+     */
+    suspend fun organizationResetPasswordEnroll(
+        organizationId: String,
+        userId: String,
+        passwordHash: String,
+        resetPasswordKey: String,
+    ): Result<Unit>
+
    /**
     * Request claimed organization domain information for an [email] needed for SSO requests.
     */
    suspend fun getOrganizationDomainSsoDetails(
        email: String,
    ): Result<OrganizationDomainSsoDetailsResponseJson>
+
+    /**
+     * Gets info regarding whether this organization enforces reset password auto enrollment.
+     */
+    suspend fun getOrganizationAutoEnrollStatus(
+        organizationIdentifier: String,
+    ): Result<OrganizationAutoEnrollStatusResponseJson>
+
+    /**
+     * Gets the public and private keys for this organization.
+     */
+    suspend fun getOrganizationKeys(
+        organizationId: String,
+    ): Result<OrganizationKeysResponseJson>
 }
--- a/app/src/main/java/com/x8bit/bitwarden/data/auth/datasource/network/service/OrganizationServiceImpl.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/auth/datasource/network/service/OrganizationServiceImpl.kt
@ -1,15 +1,35 @@
 package com.x8bit.bitwarden.data.auth.datasource.network.service

+import com.x8bit.bitwarden.data.auth.datasource.network.api.AuthenticatedOrganizationApi
 import com.x8bit.bitwarden.data.auth.datasource.network.api.OrganizationApi
+import com.x8bit.bitwarden.data.auth.datasource.network.model.OrganizationAutoEnrollStatusResponseJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.OrganizationDomainSsoDetailsRequestJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.OrganizationDomainSsoDetailsResponseJson
+import com.x8bit.bitwarden.data.auth.datasource.network.model.OrganizationKeysResponseJson
+import com.x8bit.bitwarden.data.auth.datasource.network.model.OrganizationResetPasswordEnrollRequestJson

 /**
 * Default implementation of [OrganizationService].
 */
 class OrganizationServiceImpl(
+    private val authenticatedOrganizationApi: AuthenticatedOrganizationApi,
    private val organizationApi: OrganizationApi,
 ) : OrganizationService {
+    override suspend fun organizationResetPasswordEnroll(
+        organizationId: String,
+        userId: String,
+        passwordHash: String,
+        resetPasswordKey: String,
+    ): Result<Unit> = authenticatedOrganizationApi
+        .organizationResetPasswordEnroll(
+            organizationId = organizationId,
+            userId = userId,
+            body = OrganizationResetPasswordEnrollRequestJson(
+                passwordHash = passwordHash,
+                resetPasswordKey = resetPasswordKey,
+            ),
+        )
+
    override suspend fun getOrganizationDomainSsoDetails(
        email: String,
    ): Result<OrganizationDomainSsoDetailsResponseJson> = organizationApi
@ -18,4 +38,18 @@ class OrganizationServiceImpl(
                email = email,
            ),
        )
+
+    override suspend fun getOrganizationAutoEnrollStatus(
+        organizationIdentifier: String,
+    ): Result<OrganizationAutoEnrollStatusResponseJson> = authenticatedOrganizationApi
+        .getOrganizationAutoEnrollResponse(
+            organizationIdentifier = organizationIdentifier,
+        )
+
+    override suspend fun getOrganizationKeys(
+        organizationId: String,
+    ): Result<OrganizationKeysResponseJson> = authenticatedOrganizationApi
+        .getOrganizationKeys(
+            organizationId = organizationId,
+        )
 }
--- a/app/src/main/java/com/x8bit/bitwarden/data/auth/repository/AuthRepository.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/auth/repository/AuthRepository.kt
@ -69,9 +69,9 @@ interface AuthRepository : AuthenticatorProvider, AuthRequestManager {
    val yubiKeyResultFlow: Flow<YubiKeyResult>

    /**
-     * The organization identifier currently associated with this user.
+     * The organization identifier currently associated with this user's SSO flow.
     */
-    var organizationIdentifier: String?
+    val ssoOrganizationIdentifier: String?

    /**
     * The two-factor response data necessary for login and also to populate the
--- a/app/src/main/java/com/x8bit/bitwarden/data/auth/repository/AuthRepositoryImpl.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/auth/repository/AuthRepositoryImpl.kt
@ -73,11 +73,13 @@ import com.x8bit.bitwarden.data.platform.repository.EnvironmentRepository
 import com.x8bit.bitwarden.data.platform.repository.SettingsRepository
 import com.x8bit.bitwarden.data.platform.repository.util.bufferedMutableSharedFlow
 import com.x8bit.bitwarden.data.platform.util.asFailure
+import com.x8bit.bitwarden.data.platform.util.asSuccess
 import com.x8bit.bitwarden.data.platform.util.flatMap
 import com.x8bit.bitwarden.data.vault.datasource.network.model.PolicyTypeJson
 import com.x8bit.bitwarden.data.vault.datasource.network.model.SyncResponseJson
 import com.x8bit.bitwarden.data.vault.datasource.sdk.VaultSdkSource
 import com.x8bit.bitwarden.data.vault.repository.VaultRepository
+import com.x8bit.bitwarden.data.vault.repository.model.VaultUnlockResult
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
@ -143,6 +145,8 @@ class AuthRepositoryImpl(
     */
    private var resendEmailRequestJson: ResendEmailRequestJson? = null

+    private var organizationIdentifier: String? = null
+
    /**
     * The password that needs to be checked against any organization policies before
     * the user can complete the login flow.
@ -158,10 +162,9 @@ class AuthRepositoryImpl(

    private val ioScope = CoroutineScope(dispatcherManager.io)

-    override var organizationIdentifier: String? = null
-
    override var twoFactorResponse: TwoFactorRequired? = null

+    override val ssoOrganizationIdentifier: String? get() = organizationIdentifier
    override val activeUserId: String? get() = authDiskSource.userState?.activeUserId

    @OptIn(ExperimentalCoroutinesApi::class)
@ -877,12 +880,32 @@ class AuthRepositoryImpl(
                        )
                    }
            }
+            .flatMap {
+                when (vaultRepository.unlockVaultWithMasterPassword(password)) {
+                    is VaultUnlockResult.Success -> {
+                        enrollUserInPasswordReset(
+                            organizationIdentifier = organizationIdentifier,
+                            passwordHash = passwordHash,
+                        )
+                    }
+
+                    VaultUnlockResult.AuthenticationError,
+                    VaultUnlockResult.GenericError,
+                    VaultUnlockResult.InvalidStateError,
+                    -> {
+                        IllegalStateException("Failed to unlock vault").asFailure()
+                    }
+                }
+            }
            .onSuccess {
                authDiskSource.storeMasterPasswordHash(
                    userId = activeAccount.profile.userId,
                    passwordHash = passwordHash,
                )
+
                authDiskSource.userState = authDiskSource.userState?.toUserStateJsonWithPassword()
+
+                this.organizationIdentifier = null
            }
            .fold(
                onFailure = { SetPasswordResult.Error },
@ -1081,6 +1104,42 @@ class AuthRepositoryImpl(
        }
    }

+    /**
+     * Enrolls the active user in password reset if their organization requires it.
+     */
+    private suspend fun enrollUserInPasswordReset(
+        organizationIdentifier: String,
+        passwordHash: String,
+    ): Result<Unit> {
+        val userId = activeUserId ?: return IllegalStateException("No active user").asFailure()
+        return organizationService
+            .getOrganizationAutoEnrollStatus(
+                organizationIdentifier = organizationIdentifier,
+            )
+            .flatMap { statusResponse ->
+                if (statusResponse.isResetPasswordEnabled) {
+                    organizationService
+                        .getOrganizationKeys(statusResponse.organizationId)
+                        .flatMap { keys ->
+                            vaultSdkSource.getResetPasswordKey(
+                                orgPublicKey = keys.publicKey,
+                                userId = userId,
+                            )
+                        }
+                        .flatMap { key ->
+                            organizationService.organizationResetPasswordEnroll(
+                                organizationId = statusResponse.organizationId,
+                                passwordHash = passwordHash,
+                                resetPasswordKey = key,
+                                userId = userId,
+                            )
+                        }
+                } else {
+                    Unit.asSuccess()
+                }
+            }
+    }
+
    /**
     * Get the remembered two-factor token associated with the user's email, if applicable.
     */
--- a/app/src/main/java/com/x8bit/bitwarden/data/vault/datasource/sdk/VaultSdkSource.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/vault/datasource/sdk/VaultSdkSource.kt
@ -70,6 +70,17 @@ interface VaultSdkSource {
        userId: String,
    ): Result<String>

+    /**
+     * Get the reset password key for this [orgPublicKey] and [userId].
+     *
+     * This should only be called after a successful call to [initializeCrypto] for the associated
+     * user.
+     */
+    suspend fun getResetPasswordKey(
+        orgPublicKey: String,
+        userId: String,
+    ): Result<String>
+
    /**
     * Gets the user's encryption key, which can be used to later unlock their vault via a call to
     * [initializeCrypto] with [InitUserCryptoMethod.DecryptedKey].
--- a/app/src/main/java/com/x8bit/bitwarden/data/vault/datasource/sdk/VaultSdkSourceImpl.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/vault/datasource/sdk/VaultSdkSourceImpl.kt
@ -69,6 +69,15 @@ class VaultSdkSourceImpl(
                .approveAuthRequest(publicKey)
        }

+    override suspend fun getResetPasswordKey(
+        orgPublicKey: String,
+        userId: String,
+    ): Result<String> = runCatching {
+        getClient(userId = userId)
+            .crypto()
+            .enrollAdminPasswordReset(publicKey = orgPublicKey)
+    }
+
    override suspend fun getUserEncryptionKey(
        userId: String,
    ): Result<String> =
--- a/app/src/main/java/com/x8bit/bitwarden/ui/auth/feature/setpassword/SetPasswordViewModel.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/ui/auth/feature/setpassword/SetPasswordViewModel.kt
@ -6,8 +6,6 @@ import androidx.lifecycle.viewModelScope
 import com.x8bit.bitwarden.R
 import com.x8bit.bitwarden.data.auth.repository.AuthRepository
 import com.x8bit.bitwarden.data.auth.repository.model.SetPasswordResult
-import com.x8bit.bitwarden.data.vault.repository.VaultRepository
-import com.x8bit.bitwarden.data.vault.repository.model.VaultUnlockResult
 import com.x8bit.bitwarden.ui.auth.feature.resetpassword.util.toDisplayLabels
 import com.x8bit.bitwarden.ui.platform.base.BaseViewModel
 import com.x8bit.bitwarden.ui.platform.base.util.Text
@ -28,11 +26,10 @@ private const val MIN_PASSWORD_LENGTH = 12
@Suppress("TooManyFunctions")
 class SetPasswordViewModel @Inject constructor(
    private val authRepository: AuthRepository,
-    private val vaultRepository: VaultRepository,
    savedStateHandle: SavedStateHandle,
 ) : BaseViewModel<SetPasswordState, SetPasswordEvent, SetPasswordAction>(
    initialState = savedStateHandle[KEY_STATE] ?: run {
-        val organizationIdentifier = authRepository.organizationIdentifier
+        val organizationIdentifier = authRepository.ssoOrganizationIdentifier
        if (organizationIdentifier.isNullOrBlank()) authRepository.logout()
        SetPasswordState(
            dialogState = null,
@ -60,10 +57,6 @@ class SetPasswordViewModel @Inject constructor(
                handlePasswordHintInputChanged(action)
            }

-            is SetPasswordAction.Internal.ReceiveUnlockVaultResult -> {
-                handleReceiveUnlockVaultResult(action)
-            }
-
            is SetPasswordAction.Internal.ReceiveSetPasswordResult -> {
                handleReceiveSetPasswordResult(action)
            }
@ -180,30 +173,6 @@ class SetPasswordViewModel @Inject constructor(
        }
    }

-    private fun handleReceiveUnlockVaultResult(
-        action: SetPasswordAction.Internal.ReceiveUnlockVaultResult,
-    ) {
-        when (action.result) {
-            is VaultUnlockResult.Success -> {
-                mutableStateFlow.update { it.copy(dialogState = null) }
-            }
-
-            is VaultUnlockResult.AuthenticationError,
-            is VaultUnlockResult.InvalidStateError,
-            is VaultUnlockResult.GenericError,
-            -> {
-                mutableStateFlow.update {
-                    it.copy(
-                        dialogState = SetPasswordState.DialogState.Error(
-                            title = R.string.an_error_has_occurred.asText(),
-                            message = R.string.generic_error_message.asText(),
-                        ),
-                    )
-                }
-            }
-        }
-    }
-
    /**
     * Show an alert if the set password attempt failed, otherwise attempt to unlock the vault.
     */
@ -223,15 +192,7 @@ class SetPasswordViewModel @Inject constructor(
            }

            SetPasswordResult.Success -> {
-                viewModelScope.launch {
-                    sendAction(
-                        SetPasswordAction.Internal.ReceiveUnlockVaultResult(
-                            result = vaultRepository.unlockVaultWithMasterPassword(
-                                masterPassword = state.passwordInput,
-                            ),
-                        ),
-                    )
-                }
+                mutableStateFlow.update { it.copy(dialogState = null) }
            }
        }
    }
@ -361,13 +322,6 @@ sealed class SetPasswordAction {
     * Models actions that the [SetPasswordViewModel] might send itself.
     */
    sealed class Internal : SetPasswordAction() {
-        /**
-         * Indicates that a login result has been received.
-         */
-        data class ReceiveUnlockVaultResult(
-            val result: VaultUnlockResult,
-        ) : Internal()
-
        /**
         * Indicates that a set password result has been received.
         */
--- a/app/src/test/java/com/x8bit/bitwarden/data/auth/datasource/network/service/OrganizationServiceTest.kt
+++ b/app/src/test/java/com/x8bit/bitwarden/data/auth/datasource/network/service/OrganizationServiceTest.kt
@ -1,8 +1,12 @@
 package com.x8bit.bitwarden.data.auth.datasource.network.service

+import com.x8bit.bitwarden.data.auth.datasource.network.api.AuthenticatedOrganizationApi
 import com.x8bit.bitwarden.data.auth.datasource.network.api.OrganizationApi
+import com.x8bit.bitwarden.data.auth.datasource.network.model.OrganizationAutoEnrollStatusResponseJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.OrganizationDomainSsoDetailsResponseJson
+import com.x8bit.bitwarden.data.auth.datasource.network.model.OrganizationKeysResponseJson
 import com.x8bit.bitwarden.data.platform.base.BaseServiceTest
+import com.x8bit.bitwarden.data.platform.util.asSuccess
 import kotlinx.coroutines.test.runTest
 import okhttp3.mockwebserver.MockResponse
 import org.junit.jupiter.api.Assertions.assertEquals
@ -12,12 +16,40 @@ import retrofit2.create
 import java.time.ZonedDateTime

 class OrganizationServiceTest : BaseServiceTest() {
+    private val authenticatedOrganizationApi: AuthenticatedOrganizationApi = retrofit.create()
    private val organizationApi: OrganizationApi = retrofit.create()

    private val organizationService = OrganizationServiceImpl(
+        authenticatedOrganizationApi = authenticatedOrganizationApi,
        organizationApi = organizationApi,
    )

+    @Test
+    fun `organizationResetPasswordEnroll when response is success should return Unit as success`() =
+        runTest {
+            server.enqueue(MockResponse().setResponseCode(200))
+            val result = organizationService.organizationResetPasswordEnroll(
+                organizationId = "orgId",
+                userId = "userId",
+                passwordHash = "passwordHash",
+                resetPasswordKey = "resetPasswordKey",
+            )
+            assertEquals(Unit.asSuccess(), result)
+        }
+
+    @Test
+    fun `organizationResetPasswordEnroll when response is an error should return an error`() =
+        runTest {
+            server.enqueue(MockResponse().setResponseCode(400))
+            val result = organizationService.organizationResetPasswordEnroll(
+                organizationId = "orgId",
+                userId = "userId",
+                passwordHash = "passwordHash",
+                resetPasswordKey = "resetPasswordKey",
+            )
+            assertTrue(result.isFailure)
+        }
+
    @Suppress("MaxLineLength")
    @Test
    fun `getOrganizationDomainSsoDetails when response is success should return PrevalidateSsoResponseJson`() =
@ -38,8 +70,54 @@ class OrganizationServiceTest : BaseServiceTest() {
            val result = organizationService.getOrganizationDomainSsoDetails(email)
            assertTrue(result.isFailure)
        }
+
+    @Test
+    fun `getOrganizationAutoEnrollStatus when response is success should return valid response`() =
+        runTest {
+            server.enqueue(
+                MockResponse().setResponseCode(200).setBody(ORGANIZATION_AUTO_ENROLL_STATUS_JSON),
+            )
+            val result = organizationService.getOrganizationAutoEnrollStatus("orgId")
+            assertEquals(Result.success(ORGANIZATION_AUTO_ENROLL_STATUS_RESPONSE), result)
+        }
+
+    @Test
+    fun `getOrganizationAutoEnrollStatus when response is an error should return an error`() =
+        runTest {
+            server.enqueue(MockResponse().setResponseCode(400))
+            val result = organizationService.getOrganizationAutoEnrollStatus("orgId")
+            assertTrue(result.isFailure)
+        }
+
+    @Test
+    fun `getOrganizationKeys when response is success should return valid response`() = runTest {
+        server.enqueue(
+            MockResponse().setResponseCode(200).setBody(ORGANIZATION_KEYS_JSON),
+        )
+        val result = organizationService.getOrganizationKeys("orgId")
+        assertEquals(Result.success(ORGANIZATION_KEYS_RESPONSE), result)
+    }
+
+    @Test
+    fun `getOrganizationKeys when response is an error should return an error`() = runTest {
+        server.enqueue(MockResponse().setResponseCode(400))
+        val result = organizationService.getOrganizationKeys("orgId")
+        assertTrue(result.isFailure)
+    }
 }

+private const val ORGANIZATION_AUTO_ENROLL_STATUS_JSON = """
+{
+  "id": "orgId",
+  "resetPasswordEnabled": true
+}
+"""
+
+private val ORGANIZATION_AUTO_ENROLL_STATUS_RESPONSE = OrganizationAutoEnrollStatusResponseJson(
+    organizationId = "orgId",
+    isResetPasswordEnabled = true,
+)
+
 private const val ORGANIZATION_DOMAIN_SSO_DETAILS_JSON = """
 {
  "ssoAvailable": true,
@ -57,3 +135,15 @@ private val ORGANIZATION_DOMAIN_SSO_BODY = OrganizationDomainSsoDetailsResponseJ
    isSsoRequired = false,
    verifiedDate = ZonedDateTime.parse("2024-09-13T00:00Z"),
 )
+
+private const val ORGANIZATION_KEYS_JSON = """
+{
+  "privateKey": "privateKey",
+  "publicKey": "publicKey"
+}
+"""
+
+private val ORGANIZATION_KEYS_RESPONSE = OrganizationKeysResponseJson(
+    privateKey = "privateKey",
+    publicKey = "publicKey",
+)
--- a/app/src/test/java/com/x8bit/bitwarden/data/auth/repository/AuthRepositoryTest.kt
+++ b/app/src/test/java/com/x8bit/bitwarden/data/auth/repository/AuthRepositoryTest.kt
@ -18,7 +18,9 @@ import com.x8bit.bitwarden.data.auth.datasource.disk.util.FakeAuthDiskSource
 import com.x8bit.bitwarden.data.auth.datasource.network.model.GetTokenResponseJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.IdentityTokenAuthModel
 import com.x8bit.bitwarden.data.auth.datasource.network.model.KdfTypeJson
+import com.x8bit.bitwarden.data.auth.datasource.network.model.OrganizationAutoEnrollStatusResponseJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.OrganizationDomainSsoDetailsResponseJson
+import com.x8bit.bitwarden.data.auth.datasource.network.model.OrganizationKeysResponseJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.PasswordHintResponseJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.PreLoginResponseJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.PrevalidateSsoResponseJson
@ -2681,10 +2683,13 @@ class AuthRepositoryTest {
        val password = "password"
        val passwordHash = "passwordHash"
        val passwordHint = "passwordHint"
-        val organizationId = ORGANIZATION_IDENTIFIER
+        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(
@ -2695,7 +2700,7 @@ class AuthRepositoryTest {
        val setPasswordRequestJson = SetPasswordRequestJson(
            passwordHash = passwordHash,
            passwordHint = passwordHint,
-            organizationIdentifier = organizationId,
+            organizationIdentifier = organizationIdentifier,
            kdfIterations = profile.kdfIterations,
            kdfMemory = profile.kdfMemory,
            kdfParallelism = profile.kdfParallelism,
@ -2721,9 +2726,40 @@ class AuthRepositoryTest {
        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 = organizationId,
+            organizationIdentifier = organizationIdentifier,
            password = password,
            passwordHint = passwordHint,
        )
@ -2733,6 +2769,119 @@ class AuthRepositoryTest {
        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,
+                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,
+            )
+        }
+    }
+
+    @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,
+            )
+        } returns passwordHash.asSuccess()
+        coEvery {
+            authSdkSource.makeRegisterKeys(email = EMAIL, password = password, kdf = kdf)
+        } returns registerKeyResponse.asSuccess()
+        coEvery {
+            accountsService.setPassword(body = setPasswordRequestJson)
+        } returns Unit.asSuccess()
+        coEvery {
+            vaultRepository.unlockVaultWithMasterPassword(password)
+        } returns VaultUnlockResult.GenericError
+
+        val result = repository.setPassword(
+            organizationIdentifier = organizationIdentifier,
+            password = password,
+            passwordHint = passwordHint,
+        )
+
+        assertEquals(SetPasswordResult.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,
+                passwordHash = passwordHash,
+                resetPasswordKey = resetPasswordKey,
+            )
+            vaultSdkSource.getResetPasswordKey(
+                orgPublicKey = publicOrgKey,
+                userId = profile.userId,
+            )
+        }
    }

    @Test
--- a/app/src/test/java/com/x8bit/bitwarden/data/vault/datasource/sdk/VaultSdkSourceTest.kt
+++ b/app/src/test/java/com/x8bit/bitwarden/data/vault/datasource/sdk/VaultSdkSourceTest.kt
@ -128,6 +128,52 @@ class VaultSdkSourceTest {
            coVerify { sdkClientManager.getOrCreateClient(userId = userId) }
        }

+    @Test
+    fun `getAuthRequestKey should call SDK and return a Result with correct data`() =
+        runBlocking {
+            val publicKey = "key"
+            val userId = "userId"
+            val expectedResult = "authRequestKey"
+            coEvery {
+                clientAuth.approveAuthRequest(publicKey)
+            } returns expectedResult
+            val result = vaultSdkSource.getAuthRequestKey(
+                publicKey = publicKey,
+                userId = userId,
+            )
+            assertEquals(
+                expectedResult.asSuccess(),
+                result,
+            )
+            coVerify {
+                clientAuth.approveAuthRequest(publicKey)
+                sdkClientManager.getOrCreateClient(userId = userId)
+            }
+        }
+
+    @Test
+    fun `getResetPasswordKey should call SDK and return a Result with correct data`() =
+        runBlocking {
+            val orgPublicKey = "key"
+            val userId = "userId"
+            val expectedResult = "resetPasswordKey"
+            coEvery {
+                clientCrypto.enrollAdminPasswordReset(orgPublicKey)
+            } returns expectedResult
+            val result = vaultSdkSource.getResetPasswordKey(
+                orgPublicKey = orgPublicKey,
+                userId = userId,
+            )
+            assertEquals(
+                expectedResult.asSuccess(),
+                result,
+            )
+            coVerify {
+                clientCrypto.enrollAdminPasswordReset(orgPublicKey)
+                sdkClientManager.getOrCreateClient(userId = userId)
+            }
+        }
+
    @Test
    fun `getUserEncryptionKey should call SDK and return a Result with correct data`() =
        runBlocking {
@ -143,8 +189,8 @@ class VaultSdkSourceTest {
            )
            coVerify {
                clientCrypto.getUserEncryptionKey()
+                sdkClientManager.getOrCreateClient(userId = userId)
            }
-            coVerify { sdkClientManager.getOrCreateClient(userId = userId) }
        }

    @Test
--- a/app/src/test/java/com/x8bit/bitwarden/ui/auth/feature/setpassword/SetPasswordViewModelTest.kt
+++ b/app/src/test/java/com/x8bit/bitwarden/ui/auth/feature/setpassword/SetPasswordViewModelTest.kt
@ -5,8 +5,6 @@ import app.cash.turbine.test
 import com.x8bit.bitwarden.R
 import com.x8bit.bitwarden.data.auth.repository.AuthRepository
 import com.x8bit.bitwarden.data.auth.repository.model.SetPasswordResult
-import com.x8bit.bitwarden.data.vault.repository.VaultRepository
-import com.x8bit.bitwarden.data.vault.repository.model.VaultUnlockResult
 import com.x8bit.bitwarden.ui.platform.base.BaseViewModelTest
 import com.x8bit.bitwarden.ui.platform.base.util.asText
 import io.mockk.coEvery
@ -23,16 +21,18 @@ import org.junit.jupiter.api.Test
 class SetPasswordViewModelTest : BaseViewModelTest() {
    private val authRepository: AuthRepository = mockk {
        every { passwordPolicies } returns emptyList()
-        every { organizationIdentifier } returns ORGANIZATION_IDENTIFIER
+        every { ssoOrganizationIdentifier } returns ORGANIZATION_IDENTIFIER
    }
-    private val vaultRepository: VaultRepository = mockk()

    @Test
-    fun `null organizationIdentifier logs user out`() = runTest {
+    fun `null ssoOrganizationIdentifier logs user out`() = runTest {
        every { authRepository.logout() } just runs
-        every { authRepository.organizationIdentifier } returns null
+        every { authRepository.ssoOrganizationIdentifier } returns null
        createViewModel()
-        verify { authRepository.logout() }
+        verify {
+            authRepository.logout()
+            authRepository.ssoOrganizationIdentifier
+        }
    }

    @Test
@ -159,7 +159,7 @@ class SetPasswordViewModelTest : BaseViewModelTest() {
        }

    @Test
-    fun `SubmitClicked with all valid inputs and unlock vault success sets password`() = runTest {
+    fun `SubmitClicked with all valid inputs sets password`() = runTest {
        val password = "TestPassword123"
        coEvery {
            authRepository.setPassword(
@ -168,9 +168,6 @@ class SetPasswordViewModelTest : BaseViewModelTest() {
                passwordHint = "",
            )
        } returns SetPasswordResult.Success
-        coEvery {
-            vaultRepository.unlockVaultWithMasterPassword(password)
-        } returns VaultUnlockResult.Success

        val viewModel = createViewModel()
        viewModel.trySendAction(SetPasswordAction.PasswordInputChanged(password))
@ -215,71 +212,6 @@ class SetPasswordViewModelTest : BaseViewModelTest() {
                password = password,
                passwordHint = "",
            )
-            vaultRepository.unlockVaultWithMasterPassword(password)
-        }
-    }
-
-    @Test
-    fun `SubmitClicked with all valid inputs and unlock vault failure shows error`() = runTest {
-        val password = "TestPassword123"
-        coEvery {
-            authRepository.setPassword(
-                organizationIdentifier = ORGANIZATION_IDENTIFIER,
-                password = password,
-                passwordHint = "",
-            )
-        } returns SetPasswordResult.Success
-        coEvery {
-            vaultRepository.unlockVaultWithMasterPassword(password)
-        } returns VaultUnlockResult.InvalidStateError
-
-        val viewModel = createViewModel()
-        viewModel.trySendAction(SetPasswordAction.PasswordInputChanged(password))
-        viewModel.trySendAction(SetPasswordAction.RetypePasswordInputChanged(password))
-
-        viewModel.stateFlow.test {
-            assertEquals(
-                DEFAULT_STATE.copy(
-                    dialogState = null,
-                    passwordInput = password,
-                    retypePasswordInput = password,
-                ),
-                awaitItem(),
-            )
-
-            viewModel.trySendAction(SetPasswordAction.SubmitClick)
-
-            assertEquals(
-                DEFAULT_STATE.copy(
-                    dialogState = SetPasswordState.DialogState.Loading(
-                        message = R.string.updating_password.asText(),
-                    ),
-                    passwordInput = password,
-                    retypePasswordInput = password,
-                ),
-                awaitItem(),
-            )
-
-            assertEquals(
-                DEFAULT_STATE.copy(
-                    dialogState = SetPasswordState.DialogState.Error(
-                        title = R.string.an_error_has_occurred.asText(),
-                        message = R.string.generic_error_message.asText(),
-                    ),
-                    passwordInput = password,
-                    retypePasswordInput = password,
-                ),
-                awaitItem(),
-            )
-        }
-
-        coVerify {
-            authRepository.setPassword(
-                organizationIdentifier = ORGANIZATION_IDENTIFIER,
-                password = password,
-                passwordHint = "",
-            )
-            vaultRepository.unlockVaultWithMasterPassword(password)
        }
    }

@ -328,7 +260,6 @@ class SetPasswordViewModelTest : BaseViewModelTest() {
    ): SetPasswordViewModel =
        SetPasswordViewModel(
            authRepository = authRepository,
-            vaultRepository = vaultRepository,
            savedStateHandle = SavedStateHandle(mapOf("state" to state)),
        )
 }