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

2026-06-03 03:06:21 -05:00 · 2024-03-19 10:56:44 -06:00
parent be127f5d49
commit bd58dac0ff
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.
         */