PM-11226: Wrap Key Connector APIs (#3794)

2024-11-24 02:15:53 +03:00 · 2024-08-21 12:26:20 -05:00 · 2024-08-21 12:26:20 -05:00 · 17fd3ec0f0
commit 17fd3ec0f0
parent 43a6495b98
12 changed files with 606 additions and 22 deletions
--- a/app/src/main/java/com/x8bit/bitwarden/data/auth/datasource/network/api/AuthenticatedAccountsApi.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/auth/datasource/network/api/AuthenticatedAccountsApi.kt
@ -2,7 +2,6 @@ package com.x8bit.bitwarden.data.auth.datasource.network.api

 import com.x8bit.bitwarden.data.auth.datasource.network.model.CreateAccountKeysRequest
 import com.x8bit.bitwarden.data.auth.datasource.network.model.DeleteAccountRequestJson
-import com.x8bit.bitwarden.data.auth.datasource.network.model.KeyConnectorKeyRequestJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.ResetPasswordRequestJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.SetPasswordRequestJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.VerifyOtpRequestJson
@ -53,12 +52,6 @@ interface AuthenticatedAccountsApi {
    @HTTP(method = "POST", path = "/accounts/password", hasBody = true)
    suspend fun resetPassword(@Body body: ResetPasswordRequestJson): Result<Unit>

-    /**
-     * Sets the key connector key.
-     */
-    @POST("/accounts/set-key-connector-key")
-    suspend fun setKeyConnectorKey(@Body body: KeyConnectorKeyRequestJson): Result<Unit>
-
    /**
     * Sets the password.
     */
--- a/app/src/main/java/com/x8bit/bitwarden/data/auth/datasource/network/api/AuthenticatedKeyConnectorApi.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/auth/datasource/network/api/AuthenticatedKeyConnectorApi.kt
@ -2,9 +2,7 @@ package com.x8bit.bitwarden.data.auth.datasource.network.api

 import androidx.annotation.Keep
 import com.x8bit.bitwarden.data.auth.datasource.network.model.KeyConnectorMasterKeyRequestJson
-import com.x8bit.bitwarden.data.auth.datasource.network.model.KeyConnectorMasterKeyResponseJson
 import retrofit2.http.Body
-import retrofit2.http.GET
 import retrofit2.http.POST
 import retrofit2.http.Url

@ -18,9 +16,4 @@ interface AuthenticatedKeyConnectorApi {
        @Url url: String,
        @Body body: KeyConnectorMasterKeyRequestJson,
    ): Result<Unit>
-
-    @GET
-    suspend fun getMasterKeyFromKeyConnector(
-        @Url url: String,
-    ): Result<KeyConnectorMasterKeyResponseJson>
 }
--- a/app/src/main/java/com/x8bit/bitwarden/data/auth/datasource/network/api/UnauthenticatedAccountsApi.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/auth/datasource/network/api/UnauthenticatedAccountsApi.kt
@ -1,8 +1,11 @@
 package com.x8bit.bitwarden.data.auth.datasource.network.api

+import com.x8bit.bitwarden.data.auth.datasource.network.model.KeyConnectorKeyRequestJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.PasswordHintRequestJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.ResendEmailRequestJson
+import com.x8bit.bitwarden.data.platform.datasource.network.util.HEADER_KEY_AUTHORIZATION
 import retrofit2.http.Body
+import retrofit2.http.Header
 import retrofit2.http.POST

 /**
@ -18,4 +21,10 @@ interface UnauthenticatedAccountsApi {
    suspend fun resendVerificationCodeEmail(
        @Body body: ResendEmailRequestJson,
    ): Result<Unit>
+
+    @POST("/accounts/set-key-connector-key")
+    suspend fun setKeyConnectorKey(
+        @Body body: KeyConnectorKeyRequestJson,
+        @Header(HEADER_KEY_AUTHORIZATION) bearerToken: String,
+    ): Result<Unit>
 }
--- a/app/src/main/java/com/x8bit/bitwarden/data/auth/datasource/network/api/UnauthenticatedKeyConnectorApi.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/auth/datasource/network/api/UnauthenticatedKeyConnectorApi.kt
@ -0,0 +1,30 @@
+package com.x8bit.bitwarden.data.auth.datasource.network.api
+
+import androidx.annotation.Keep
+import com.x8bit.bitwarden.data.auth.datasource.network.model.KeyConnectorMasterKeyRequestJson
+import com.x8bit.bitwarden.data.auth.datasource.network.model.KeyConnectorMasterKeyResponseJson
+import com.x8bit.bitwarden.data.platform.datasource.network.util.HEADER_KEY_AUTHORIZATION
+import retrofit2.http.Body
+import retrofit2.http.GET
+import retrofit2.http.Header
+import retrofit2.http.POST
+import retrofit2.http.Url
+
+/**
+ * Defines raw calls specific for key connectors that use custom urls.
+ */
+@Keep
+interface UnauthenticatedKeyConnectorApi {
+    @POST
+    suspend fun storeMasterKeyToKeyConnector(
+        @Url url: String,
+        @Header(HEADER_KEY_AUTHORIZATION) bearerToken: String,
+        @Body body: KeyConnectorMasterKeyRequestJson,
+    ): Result<Unit>
+
+    @GET
+    suspend fun getMasterKeyFromKeyConnector(
+        @Url url: String,
+        @Header(HEADER_KEY_AUTHORIZATION) bearerToken: String,
+    ): Result<KeyConnectorMasterKeyResponseJson>
+}
--- 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
@ -38,6 +38,7 @@ object AuthNetworkModule {
    ): AccountsService = AccountsServiceImpl(
        unauthenticatedAccountsApi = retrofits.unauthenticatedApiRetrofit.create(),
        authenticatedAccountsApi = retrofits.authenticatedApiRetrofit.create(),
+        unauthenticatedKeyConnectorApi = retrofits.createStaticRetrofit().create(),
        authenticatedKeyConnectorApi = retrofits
            .createStaticRetrofit(isAuthenticated = true)
            .create(),
--- a/app/src/main/java/com/x8bit/bitwarden/data/auth/datasource/network/service/AccountsService.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/auth/datasource/network/service/AccountsService.kt
@ -59,8 +59,14 @@ interface AccountsService {

    /**
     * Set the key connector key.
+     *
+     * This API requires the [accessToken] to be passed in manually because it occurs during the
+     * login process.
     */
-    suspend fun setKeyConnectorKey(body: KeyConnectorKeyRequestJson): Result<Unit>
+    suspend fun setKeyConnectorKey(
+        accessToken: String,
+        body: KeyConnectorKeyRequestJson,
+    ): Result<Unit>

    /**
     * Set the password.
@ -69,13 +75,32 @@ interface AccountsService {

    /**
     * Retrieves the master key from the key connector.
+     *
+     * This API requires the [accessToken] to be passed in manually because it occurs during the
+     * login process.
     */
    suspend fun getMasterKeyFromKeyConnector(
        url: String,
+        accessToken: String,
    ): Result<KeyConnectorMasterKeyResponseJson>

    /**
     * Stores the master key to the key connector.
     */
-    suspend fun storeMasterKeyToKeyConnector(url: String, masterKey: String): Result<Unit>
+    suspend fun storeMasterKeyToKeyConnector(
+        url: String,
+        masterKey: String,
+    ): Result<Unit>
+
+    /**
+     * Stores the master key to the key connector.
+     *
+     * This API requires the [accessToken] to be passed in manually because it occurs during the
+     * login process.
+     */
+    suspend fun storeMasterKeyToKeyConnector(
+        url: String,
+        accessToken: String,
+        masterKey: String,
+    ): Result<Unit>
 }
--- a/app/src/main/java/com/x8bit/bitwarden/data/auth/datasource/network/service/AccountsServiceImpl.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/auth/datasource/network/service/AccountsServiceImpl.kt
@ -3,6 +3,7 @@ package com.x8bit.bitwarden.data.auth.datasource.network.service
 import com.x8bit.bitwarden.data.auth.datasource.network.api.AuthenticatedAccountsApi
 import com.x8bit.bitwarden.data.auth.datasource.network.api.AuthenticatedKeyConnectorApi
 import com.x8bit.bitwarden.data.auth.datasource.network.api.UnauthenticatedAccountsApi
+import com.x8bit.bitwarden.data.auth.datasource.network.api.UnauthenticatedKeyConnectorApi
 import com.x8bit.bitwarden.data.auth.datasource.network.model.CreateAccountKeysRequest
 import com.x8bit.bitwarden.data.auth.datasource.network.model.DeleteAccountRequestJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.DeleteAccountResponseJson
@ -16,6 +17,7 @@ import com.x8bit.bitwarden.data.auth.datasource.network.model.ResetPasswordReque
 import com.x8bit.bitwarden.data.auth.datasource.network.model.SetPasswordRequestJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.VerifyOtpRequestJson
 import com.x8bit.bitwarden.data.platform.datasource.network.model.toBitwardenError
+import com.x8bit.bitwarden.data.platform.datasource.network.util.HEADER_VALUE_BEARER_PREFIX
 import com.x8bit.bitwarden.data.platform.datasource.network.util.parseErrorBodyOrNull
 import kotlinx.serialization.json.Json

@ -26,6 +28,7 @@ import kotlinx.serialization.json.Json
 class AccountsServiceImpl(
    private val unauthenticatedAccountsApi: UnauthenticatedAccountsApi,
    private val authenticatedAccountsApi: AuthenticatedAccountsApi,
+    private val unauthenticatedKeyConnectorApi: UnauthenticatedKeyConnectorApi,
    private val authenticatedKeyConnectorApi: AuthenticatedKeyConnectorApi,
    private val json: Json,
 ) : AccountsService {
@ -109,8 +112,12 @@ class AccountsServiceImpl(
    }

    override suspend fun setKeyConnectorKey(
+        accessToken: String,
        body: KeyConnectorKeyRequestJson,
-    ): Result<Unit> = authenticatedAccountsApi.setKeyConnectorKey(body)
+    ): Result<Unit> = unauthenticatedAccountsApi.setKeyConnectorKey(
+        body = body,
+        bearerToken = "$HEADER_VALUE_BEARER_PREFIX$accessToken",
+    )

    override suspend fun setPassword(
        body: SetPasswordRequestJson,
@ -118,8 +125,12 @@ class AccountsServiceImpl(

    override suspend fun getMasterKeyFromKeyConnector(
        url: String,
+        accessToken: String,
    ): Result<KeyConnectorMasterKeyResponseJson> =
-        authenticatedKeyConnectorApi.getMasterKeyFromKeyConnector(url = "$url/user-keys")
+        unauthenticatedKeyConnectorApi.getMasterKeyFromKeyConnector(
+            url = "$url/user-keys",
+            bearerToken = "$HEADER_VALUE_BEARER_PREFIX$accessToken",
+        )

    override suspend fun storeMasterKeyToKeyConnector(
        url: String,
@ -129,4 +140,15 @@ class AccountsServiceImpl(
            url = "$url/user-keys",
            body = KeyConnectorMasterKeyRequestJson(masterKey = masterKey),
        )
+
+    override suspend fun storeMasterKeyToKeyConnector(
+        url: String,
+        accessToken: String,
+        masterKey: String,
+    ): Result<Unit> =
+        unauthenticatedKeyConnectorApi.storeMasterKeyToKeyConnector(
+            url = "$url/user-keys",
+            bearerToken = "$HEADER_VALUE_BEARER_PREFIX$accessToken",
+            body = KeyConnectorMasterKeyRequestJson(masterKey = masterKey),
+        )
 }
--- a/app/src/main/java/com/x8bit/bitwarden/data/auth/manager/KeyConnectorManager.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/auth/manager/KeyConnectorManager.kt
@ -0,0 +1,46 @@
+package com.x8bit.bitwarden.data.auth.manager
+
+import com.bitwarden.core.KeyConnectorResponse
+import com.bitwarden.crypto.Kdf
+import com.x8bit.bitwarden.data.auth.datasource.network.model.KdfTypeJson
+import com.x8bit.bitwarden.data.auth.datasource.network.model.KeyConnectorMasterKeyResponseJson
+
+/**
+ * Manager used to interface with a key connector.
+ */
+interface KeyConnectorManager {
+    /**
+     * Retrieves the master key from the key connector.
+     */
+    suspend fun getMasterKeyFromKeyConnector(
+        url: String,
+        accessToken: String,
+    ): Result<KeyConnectorMasterKeyResponseJson>
+
+    /**
+     * Migrates an existing user to use the key connector.
+     */
+    @Suppress("LongParameterList")
+    suspend fun migrateExistingUserToKeyConnector(
+        userId: String,
+        url: String,
+        userKeyEncrypted: String,
+        email: String,
+        masterPassword: String,
+        kdf: Kdf,
+    ): Result<Unit>
+
+    /**
+     * Migrates a new user to use the key connector.
+     */
+    @Suppress("LongParameterList")
+    suspend fun migrateNewUserToKeyConnector(
+        url: String,
+        accessToken: String,
+        kdfType: KdfTypeJson,
+        kdfIterations: Int?,
+        kdfMemory: Int?,
+        kdfParallelism: Int?,
+        organizationIdentifier: String,
+    ): Result<KeyConnectorResponse>
+}
--- a/app/src/main/java/com/x8bit/bitwarden/data/auth/manager/KeyConnectorManagerImpl.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/auth/manager/KeyConnectorManagerImpl.kt
@ -0,0 +1,88 @@
+package com.x8bit.bitwarden.data.auth.manager
+
+import com.bitwarden.core.KeyConnectorResponse
+import com.bitwarden.crypto.Kdf
+import com.x8bit.bitwarden.data.auth.datasource.network.model.KdfTypeJson
+import com.x8bit.bitwarden.data.auth.datasource.network.model.KeyConnectorKeyRequestJson
+import com.x8bit.bitwarden.data.auth.datasource.network.model.KeyConnectorMasterKeyResponseJson
+import com.x8bit.bitwarden.data.auth.datasource.network.service.AccountsService
+import com.x8bit.bitwarden.data.auth.datasource.sdk.AuthSdkSource
+import com.x8bit.bitwarden.data.platform.util.flatMap
+import com.x8bit.bitwarden.data.vault.datasource.sdk.VaultSdkSource
+
+/**
+ * The default implementation of the [KeyConnectorManager].
+ */
+class KeyConnectorManagerImpl(
+    private val accountsService: AccountsService,
+    private val authSdkSource: AuthSdkSource,
+    private val vaultSdkSource: VaultSdkSource,
+) : KeyConnectorManager {
+    override suspend fun getMasterKeyFromKeyConnector(
+        url: String,
+        accessToken: String,
+    ): Result<KeyConnectorMasterKeyResponseJson> =
+        accountsService.getMasterKeyFromKeyConnector(
+            url = url,
+            accessToken = accessToken,
+        )
+
+    override suspend fun migrateExistingUserToKeyConnector(
+        userId: String,
+        url: String,
+        userKeyEncrypted: String,
+        email: String,
+        masterPassword: String,
+        kdf: Kdf,
+    ): Result<Unit> =
+        vaultSdkSource
+            .deriveKeyConnector(
+                userId = userId,
+                userKeyEncrypted = userKeyEncrypted,
+                email = email,
+                password = masterPassword,
+                kdf = kdf,
+            )
+            .flatMap { masterKey ->
+                accountsService.storeMasterKeyToKeyConnector(url = url, masterKey = masterKey)
+            }
+            .flatMap { accountsService.convertToKeyConnector() }
+
+    override suspend fun migrateNewUserToKeyConnector(
+        url: String,
+        accessToken: String,
+        kdfType: KdfTypeJson,
+        kdfIterations: Int?,
+        kdfMemory: Int?,
+        kdfParallelism: Int?,
+        organizationIdentifier: String,
+    ): Result<KeyConnectorResponse> =
+        authSdkSource
+            .makeKeyConnectorKeys()
+            .flatMap { keyConnectorResponse ->
+                accountsService
+                    .storeMasterKeyToKeyConnector(
+                        url = url,
+                        accessToken = accessToken,
+                        masterKey = keyConnectorResponse.masterKey,
+                    )
+                    .flatMap {
+                        accountsService.setKeyConnectorKey(
+                            accessToken = accessToken,
+                            body = KeyConnectorKeyRequestJson(
+                                userKey = keyConnectorResponse.encryptedUserKey,
+                                keys = KeyConnectorKeyRequestJson.Keys(
+                                    publicKey = keyConnectorResponse.keys.public,
+                                    encryptedPrivateKey = keyConnectorResponse.keys.private,
+                                ),
+                                kdfType = kdfType,
+                                kdfIterations = kdfIterations,
+                                kdfMemory = kdfMemory,
+                                kdfParallelism = kdfParallelism,
+                                organizationIdentifier = organizationIdentifier,
+                            ),
+                        )
+                    }
+                    .map { keyConnectorResponse }
+            }
+}
--- a/app/src/main/java/com/x8bit/bitwarden/data/auth/manager/di/AuthManagerModule.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/auth/manager/di/AuthManagerModule.kt
@ -2,6 +2,7 @@ package com.x8bit.bitwarden.data.auth.manager.di

 import android.content.Context
 import com.x8bit.bitwarden.data.auth.datasource.disk.AuthDiskSource
+import com.x8bit.bitwarden.data.auth.datasource.network.service.AccountsService
 import com.x8bit.bitwarden.data.auth.datasource.network.service.AuthRequestsService
 import com.x8bit.bitwarden.data.auth.datasource.network.service.DevicesService
 import com.x8bit.bitwarden.data.auth.datasource.network.service.NewAuthRequestService
@ -10,6 +11,8 @@ import com.x8bit.bitwarden.data.auth.manager.AuthRequestManager
 import com.x8bit.bitwarden.data.auth.manager.AuthRequestManagerImpl
 import com.x8bit.bitwarden.data.auth.manager.AuthRequestNotificationManager
 import com.x8bit.bitwarden.data.auth.manager.AuthRequestNotificationManagerImpl
+import com.x8bit.bitwarden.data.auth.manager.KeyConnectorManager
+import com.x8bit.bitwarden.data.auth.manager.KeyConnectorManagerImpl
 import com.x8bit.bitwarden.data.auth.manager.TrustedDeviceManager
 import com.x8bit.bitwarden.data.auth.manager.TrustedDeviceManagerImpl
 import com.x8bit.bitwarden.data.auth.manager.UserLogoutManager
@ -71,6 +74,19 @@ object AuthManagerModule {
            authDiskSource = authDiskSource,
        )

+    @Provides
+    @Singleton
+    fun provideKeyConnectorManager(
+        accountsService: AccountsService,
+        authSdkSource: AuthSdkSource,
+        vaultSdkSource: VaultSdkSource,
+    ): KeyConnectorManager =
+        KeyConnectorManagerImpl(
+            accountsService = accountsService,
+            authSdkSource = authSdkSource,
+            vaultSdkSource = vaultSdkSource,
+        )
+
    @Provides
    @Singleton
    fun provideTrustedDeviceManager(
--- a/app/src/test/java/com/x8bit/bitwarden/data/auth/datasource/network/service/AccountsServiceTest.kt
+++ b/app/src/test/java/com/x8bit/bitwarden/data/auth/datasource/network/service/AccountsServiceTest.kt
@ -3,6 +3,7 @@ package com.x8bit.bitwarden.data.auth.datasource.network.service
 import com.x8bit.bitwarden.data.auth.datasource.network.api.AuthenticatedAccountsApi
 import com.x8bit.bitwarden.data.auth.datasource.network.api.AuthenticatedKeyConnectorApi
 import com.x8bit.bitwarden.data.auth.datasource.network.api.UnauthenticatedAccountsApi
+import com.x8bit.bitwarden.data.auth.datasource.network.api.UnauthenticatedKeyConnectorApi
 import com.x8bit.bitwarden.data.auth.datasource.network.model.KdfTypeJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.KeyConnectorKeyRequestJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.KeyConnectorMasterKeyResponseJson
@ -24,10 +25,12 @@ class AccountsServiceTest : BaseServiceTest() {

    private val unauthenticatedAccountsApi: UnauthenticatedAccountsApi = retrofit.create()
    private val authenticatedAccountsApi: AuthenticatedAccountsApi = retrofit.create()
+    private val unauthenticatedKeyConnectorApi: UnauthenticatedKeyConnectorApi = retrofit.create()
    private val authenticatedKeyConnectorApi: AuthenticatedKeyConnectorApi = retrofit.create()
    private val service = AccountsServiceImpl(
        unauthenticatedAccountsApi = unauthenticatedAccountsApi,
        authenticatedAccountsApi = authenticatedAccountsApi,
+        unauthenticatedKeyConnectorApi = unauthenticatedKeyConnectorApi,
        authenticatedKeyConnectorApi = authenticatedKeyConnectorApi,
        json = json,
    )
@ -189,10 +192,11 @@ class AccountsServiceTest : BaseServiceTest() {
    }

    @Test
-    fun `setKeyConnectorKey with empty response is success`() = runTest {
+    fun `setKeyConnectorKey with token and empty response is success`() = runTest {
        val response = MockResponse().setBody("")
        server.enqueue(response)
        val result = service.setKeyConnectorKey(
+            accessToken = "token",
            body = KeyConnectorKeyRequestJson(
                organizationIdentifier = "organizationId",
                kdfIterations = 7,
@ -210,11 +214,14 @@ class AccountsServiceTest : BaseServiceTest() {
    }

    @Test
-    fun `getMasterKeyFromKeyConnector with empty response is success`() = runTest {
+    fun `getMasterKeyFromKeyConnector with token and empty response is success`() = runTest {
        val masterKey = "masterKey"
        val response = MockResponse().setBody("""{ "key": "$masterKey" }""")
        server.enqueue(response)
-        val result = service.getMasterKeyFromKeyConnector(url = "$url/test")
+        val result = service.getMasterKeyFromKeyConnector(
+            url = "$url/test",
+            accessToken = "token",
+        )
        assertEquals(
            KeyConnectorMasterKeyResponseJson(masterKey = masterKey).asSuccess(),
            result,
@ -222,7 +229,7 @@ class AccountsServiceTest : BaseServiceTest() {
    }

    @Test
-    fun `storeMasterKeyToKeyConnector success should return Success`() = runTest {
+    fun `storeMasterKeyToKeyConnector without token success should return Success`() = runTest {
        val response = MockResponse()
        server.enqueue(response)
        val result = service.storeMasterKeyToKeyConnector(
@ -231,4 +238,16 @@ class AccountsServiceTest : BaseServiceTest() {
        )
        assertEquals(Unit.asSuccess(), result)
    }
+
+    @Test
+    fun `storeMasterKeyToKeyConnector with token success should return Success`() = runTest {
+        val response = MockResponse()
+        server.enqueue(response)
+        val result = service.storeMasterKeyToKeyConnector(
+            url = "$url/test",
+            masterKey = "masterKey",
+            accessToken = "token",
+        )
+        assertEquals(Unit.asSuccess(), result)
+    }
 }
--- a/app/src/test/java/com/x8bit/bitwarden/data/auth/manager/KeyConnectorManagerTest.kt
+++ b/app/src/test/java/com/x8bit/bitwarden/data/auth/manager/KeyConnectorManagerTest.kt
@ -0,0 +1,342 @@
+package com.x8bit.bitwarden.data.auth.manager
+
+import com.bitwarden.core.KeyConnectorResponse
+import com.bitwarden.crypto.Kdf
+import com.bitwarden.crypto.RsaKeyPair
+import com.x8bit.bitwarden.data.auth.datasource.network.model.KdfTypeJson
+import com.x8bit.bitwarden.data.auth.datasource.network.model.KeyConnectorKeyRequestJson
+import com.x8bit.bitwarden.data.auth.datasource.network.model.KeyConnectorMasterKeyResponseJson
+import com.x8bit.bitwarden.data.auth.datasource.network.service.AccountsService
+import com.x8bit.bitwarden.data.auth.datasource.sdk.AuthSdkSource
+import com.x8bit.bitwarden.data.platform.util.asFailure
+import com.x8bit.bitwarden.data.platform.util.asSuccess
+import com.x8bit.bitwarden.data.vault.datasource.sdk.VaultSdkSource
+import io.mockk.coEvery
+import io.mockk.every
+import io.mockk.mockk
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+import org.junit.jupiter.api.Assertions.assertEquals
+
+class KeyConnectorManagerTest {
+    private val accountsService: AccountsService = mockk()
+    private val authSdkSource: AuthSdkSource = mockk()
+    private val vaultSdkSource: VaultSdkSource = mockk()
+
+    private val keyConnectorManager: KeyConnectorManager = KeyConnectorManagerImpl(
+        accountsService = accountsService,
+        authSdkSource = authSdkSource,
+        vaultSdkSource = vaultSdkSource,
+    )
+
+    @Test
+    fun `getMasterKeyFromKeyConnector with service failure should return failure`() = runTest {
+        val expectedResult = Throwable("Fail").asFailure()
+        coEvery {
+            accountsService.getMasterKeyFromKeyConnector(url = URL, accessToken = ACCESS_TOKEN)
+        } returns expectedResult
+
+        val result = keyConnectorManager.getMasterKeyFromKeyConnector(
+            url = URL,
+            accessToken = ACCESS_TOKEN,
+        )
+
+        assertEquals(expectedResult, result)
+    }
+
+    @Test
+    fun `getMasterKeyFromKeyConnector with service success should return success`() = runTest {
+        val expectedResult = mockk<KeyConnectorMasterKeyResponseJson>().asSuccess()
+        coEvery {
+            accountsService.getMasterKeyFromKeyConnector(url = URL, accessToken = ACCESS_TOKEN)
+        } returns expectedResult
+
+        val result = keyConnectorManager.getMasterKeyFromKeyConnector(
+            url = URL,
+            accessToken = ACCESS_TOKEN,
+        )
+
+        assertEquals(expectedResult, result)
+    }
+
+    @Suppress("MaxLineLength")
+    @Test
+    fun `migrateExistingUserToKeyConnector with deriveKeyConnector failure should return failure`() =
+        runTest {
+            val expectedResult = Throwable("Fail").asFailure()
+            coEvery {
+                vaultSdkSource.deriveKeyConnector(
+                    userId = USER_ID,
+                    userKeyEncrypted = ENCRYPTED_USER_KEY,
+                    email = EMAIL,
+                    password = MASTER_PASSWORD,
+                    kdf = KDF,
+                )
+            } returns expectedResult
+
+            val result = keyConnectorManager.migrateExistingUserToKeyConnector(
+                userId = USER_ID,
+                url = URL,
+                userKeyEncrypted = ENCRYPTED_USER_KEY,
+                email = EMAIL,
+                masterPassword = MASTER_PASSWORD,
+                kdf = KDF,
+            )
+
+            assertEquals(expectedResult, result)
+        }
+
+    @Suppress("MaxLineLength")
+    @Test
+    fun `migrateExistingUserToKeyConnector with storeMasterKeyToKeyConnector failure should return failure`() =
+        runTest {
+            val expectedResult = Throwable("Fail").asFailure()
+            coEvery {
+                vaultSdkSource.deriveKeyConnector(
+                    userId = USER_ID,
+                    userKeyEncrypted = ENCRYPTED_USER_KEY,
+                    email = EMAIL,
+                    password = MASTER_PASSWORD,
+                    kdf = KDF,
+                )
+            } returns MASTER_KEY.asSuccess()
+            coEvery {
+                accountsService.storeMasterKeyToKeyConnector(url = URL, masterKey = MASTER_KEY)
+            } returns expectedResult
+
+            val result = keyConnectorManager.migrateExistingUserToKeyConnector(
+                userId = USER_ID,
+                url = URL,
+                userKeyEncrypted = ENCRYPTED_USER_KEY,
+                email = EMAIL,
+                masterPassword = MASTER_PASSWORD,
+                kdf = KDF,
+            )
+
+            assertEquals(expectedResult, result)
+        }
+
+    @Suppress("MaxLineLength")
+    @Test
+    fun `migrateExistingUserToKeyConnector with convertToKeyConnector failure should return failure`() =
+        runTest {
+            val expectedResult = Throwable("Fail").asFailure()
+            coEvery {
+                vaultSdkSource.deriveKeyConnector(
+                    userId = USER_ID,
+                    userKeyEncrypted = ENCRYPTED_USER_KEY,
+                    email = EMAIL,
+                    password = MASTER_PASSWORD,
+                    kdf = KDF,
+                )
+            } returns MASTER_KEY.asSuccess()
+            coEvery {
+                accountsService.storeMasterKeyToKeyConnector(url = URL, masterKey = MASTER_KEY)
+            } returns Unit.asSuccess()
+            coEvery { accountsService.convertToKeyConnector() } returns expectedResult
+
+            val result = keyConnectorManager.migrateExistingUserToKeyConnector(
+                userId = USER_ID,
+                url = URL,
+                userKeyEncrypted = ENCRYPTED_USER_KEY,
+                email = EMAIL,
+                masterPassword = MASTER_PASSWORD,
+                kdf = KDF,
+            )
+
+            assertEquals(expectedResult, result)
+        }
+
+    @Test
+    fun `migrateExistingUserToKeyConnector should return success`() = runTest {
+        coEvery {
+            vaultSdkSource.deriveKeyConnector(
+                userId = USER_ID,
+                userKeyEncrypted = ENCRYPTED_USER_KEY,
+                email = EMAIL,
+                password = MASTER_PASSWORD,
+                kdf = KDF,
+            )
+        } returns MASTER_KEY.asSuccess()
+        coEvery {
+            accountsService.storeMasterKeyToKeyConnector(url = URL, masterKey = MASTER_KEY)
+        } returns Unit.asSuccess()
+        coEvery { accountsService.convertToKeyConnector() } returns Unit.asSuccess()
+
+        val result = keyConnectorManager.migrateExistingUserToKeyConnector(
+            userId = USER_ID,
+            url = URL,
+            userKeyEncrypted = ENCRYPTED_USER_KEY,
+            email = EMAIL,
+            masterPassword = MASTER_PASSWORD,
+            kdf = KDF,
+        )
+
+        assertEquals(Unit.asSuccess(), result)
+    }
+
+    @Test
+    fun `migrateNewUserToKeyConnector with makeKeyConnectorKeys failure should return failure`() =
+        runTest {
+            val expectedResult = Throwable("Fail").asFailure()
+            coEvery { authSdkSource.makeKeyConnectorKeys() } returns expectedResult
+
+            val result = keyConnectorManager.migrateNewUserToKeyConnector(
+                url = URL,
+                accessToken = ACCESS_TOKEN,
+                kdfType = KDF_TYPE,
+                kdfIterations = KDF_ITERATIONS,
+                kdfMemory = KDF_MEMORY,
+                kdfParallelism = KDF_PARALLELISM,
+                organizationIdentifier = ORGANIZATION_IDENTIFIER,
+            )
+
+            assertEquals(expectedResult, result)
+        }
+
+    @Suppress("MaxLineLength")
+    @Test
+    fun `migrateNewUserToKeyConnector with storeMasterKeyToKeyConnector failure should return failure`() =
+        runTest {
+            val keyConnectorResponse: KeyConnectorResponse = mockk {
+                every { masterKey } returns MASTER_KEY
+            }
+            val expectedResult = Throwable("Fail").asFailure()
+            coEvery {
+                authSdkSource.makeKeyConnectorKeys()
+            } returns keyConnectorResponse.asSuccess()
+            coEvery {
+                accountsService.storeMasterKeyToKeyConnector(
+                    url = URL,
+                    accessToken = ACCESS_TOKEN,
+                    masterKey = MASTER_KEY,
+                )
+            } returns expectedResult
+
+            val result = keyConnectorManager.migrateNewUserToKeyConnector(
+                url = URL,
+                accessToken = ACCESS_TOKEN,
+                kdfType = KDF_TYPE,
+                kdfIterations = KDF_ITERATIONS,
+                kdfMemory = KDF_MEMORY,
+                kdfParallelism = KDF_PARALLELISM,
+                organizationIdentifier = ORGANIZATION_IDENTIFIER,
+            )
+
+            assertEquals(expectedResult, result)
+        }
+
+    @Test
+    fun `migrateNewUserToKeyConnector with setKeyConnectorKey failure should return failure`() =
+        runTest {
+            val keyConnectorResponse: KeyConnectorResponse = mockk {
+                every { masterKey } returns MASTER_KEY
+                every { encryptedUserKey } returns ENCRYPTED_USER_KEY
+                every { keys } returns RsaKeyPair(public = PUBLIC_KEY, private = PRIVATE_KEY)
+            }
+            val expectedResult = Throwable("Fail").asFailure()
+            coEvery {
+                authSdkSource.makeKeyConnectorKeys()
+            } returns keyConnectorResponse.asSuccess()
+            coEvery {
+                accountsService.storeMasterKeyToKeyConnector(
+                    url = URL,
+                    accessToken = ACCESS_TOKEN,
+                    masterKey = MASTER_KEY,
+                )
+            } returns Unit.asSuccess()
+            coEvery {
+                accountsService.setKeyConnectorKey(
+                    accessToken = ACCESS_TOKEN,
+                    body = KeyConnectorKeyRequestJson(
+                        userKey = ENCRYPTED_USER_KEY,
+                        keys = KeyConnectorKeyRequestJson.Keys(
+                            publicKey = PUBLIC_KEY,
+                            encryptedPrivateKey = PRIVATE_KEY,
+                        ),
+                        kdfType = KDF_TYPE,
+                        kdfIterations = KDF_ITERATIONS,
+                        kdfMemory = KDF_MEMORY,
+                        kdfParallelism = KDF_PARALLELISM,
+                        organizationIdentifier = ORGANIZATION_IDENTIFIER,
+                    ),
+                )
+            } returns expectedResult
+
+            val result = keyConnectorManager.migrateNewUserToKeyConnector(
+                url = URL,
+                accessToken = ACCESS_TOKEN,
+                kdfType = KDF_TYPE,
+                kdfIterations = KDF_ITERATIONS,
+                kdfMemory = KDF_MEMORY,
+                kdfParallelism = KDF_PARALLELISM,
+                organizationIdentifier = ORGANIZATION_IDENTIFIER,
+            )
+
+            assertEquals(expectedResult, result)
+        }
+
+    @Test
+    fun `migrateNewUserToKeyConnector should return succeed`() = runTest {
+        val keyConnectorResponse: KeyConnectorResponse = mockk {
+            every { masterKey } returns MASTER_KEY
+            every { encryptedUserKey } returns ENCRYPTED_USER_KEY
+            every { keys } returns RsaKeyPair(public = PUBLIC_KEY, private = PRIVATE_KEY)
+        }
+        coEvery {
+            authSdkSource.makeKeyConnectorKeys()
+        } returns keyConnectorResponse.asSuccess()
+        coEvery {
+            accountsService.storeMasterKeyToKeyConnector(
+                url = URL,
+                accessToken = ACCESS_TOKEN,
+                masterKey = MASTER_KEY,
+            )
+        } returns Unit.asSuccess()
+        coEvery {
+            accountsService.setKeyConnectorKey(
+                accessToken = ACCESS_TOKEN,
+                body = KeyConnectorKeyRequestJson(
+                    userKey = ENCRYPTED_USER_KEY,
+                    keys = KeyConnectorKeyRequestJson.Keys(
+                        publicKey = PUBLIC_KEY,
+                        encryptedPrivateKey = PRIVATE_KEY,
+                    ),
+                    kdfType = KDF_TYPE,
+                    kdfIterations = KDF_ITERATIONS,
+                    kdfMemory = KDF_MEMORY,
+                    kdfParallelism = KDF_PARALLELISM,
+                    organizationIdentifier = ORGANIZATION_IDENTIFIER,
+                ),
+            )
+        } returns Unit.asSuccess()
+
+        val result = keyConnectorManager.migrateNewUserToKeyConnector(
+            url = URL,
+            accessToken = ACCESS_TOKEN,
+            kdfType = KDF_TYPE,
+            kdfIterations = KDF_ITERATIONS,
+            kdfMemory = KDF_MEMORY,
+            kdfParallelism = KDF_PARALLELISM,
+            organizationIdentifier = ORGANIZATION_IDENTIFIER,
+        )
+
+        assertEquals(keyConnectorResponse.asSuccess(), result)
+    }
+}
+
+private const val ACCESS_TOKEN: String = "token"
+private const val USER_ID: String = "userId"
+private const val URL: String = "www.example.com"
+private const val ENCRYPTED_USER_KEY: String = "userKeyEncrypted"
+private const val EMAIL: String = "email@email.com"
+private const val MASTER_PASSWORD: String = "masterPassword"
+private const val MASTER_KEY: String = "masterKey"
+private const val PUBLIC_KEY: String = "publicKey"
+private const val PRIVATE_KEY: String = "privateKey"
+private const val ORGANIZATION_IDENTIFIER: String = "org_identifier"
+private val KDF: Kdf = mockk()
+private val KDF_TYPE: KdfTypeJson = KdfTypeJson.ARGON2_ID
+private const val KDF_ITERATIONS: Int = 1
+private const val KDF_MEMORY: Int = 2
+private const val KDF_PARALLELISM: Int = 3