From 5aa8369ac5c9d9d6cf5e67c143e84d522a09a23c Mon Sep 17 00:00:00 2001
From: Patrick Honkonen <1883101+SaintPatrck@users.noreply.github.com>
Date: Fri, 20 Dec 2024 12:48:01 -0500
Subject: [PATCH] [PM-15863] Request master password before revealing private
 SSH key (#4481)

---
 .../vault/feature/item/VaultItemViewModel.kt  | 24 +++++++
 .../feature/item/VaultItemViewModelTest.kt    | 66 +++++++++++++++++--
 2 files changed, 84 insertions(+), 6 deletions(-)

diff --git a/app/src/main/java/com/x8bit/bitwarden/ui/vault/feature/item/VaultItemViewModel.kt b/app/src/main/java/com/x8bit/bitwarden/ui/vault/feature/item/VaultItemViewModel.kt
index c6ba8363f38..4b04861e406 100644
--- a/app/src/main/java/com/x8bit/bitwarden/ui/vault/feature/item/VaultItemViewModel.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/ui/vault/feature/item/VaultItemViewModel.kt
@@ -804,6 +804,16 @@ class VaultItemViewModel @Inject constructor(
         action: VaultItemAction.ItemType.SshKey.PrivateKeyVisibilityClicked,
     ) {
         onSshKeyContent { content, sshKey ->
+            if (content.common.requiresReprompt) {
+                updateDialogState(
+                    VaultItemState.DialogState.MasterPasswordDialog(
+                        action = PasswordRepromptAction.ViewPrivateKeyClicked(
+                            isVisible = action.isVisible,
+                        ),
+                    ),
+                )
+                return@onSshKeyContent
+            }
             mutableStateFlow.update { currentState ->
                 currentState.copy(
                     viewState = content.copy(
@@ -2231,4 +2241,18 @@ sealed class PasswordRepromptAction : Parcelable {
         override val vaultItemAction: VaultItemAction
             get() = VaultItemAction.Common.RestoreVaultItemClick
     }
+
+    /**
+     * Indicates that we should launch the
+     * [VaultItemAction.ItemType.SshKey.PrivateKeyVisibilityClicked] upon password validation.
+     */
+    @Parcelize
+    data class ViewPrivateKeyClicked(
+        val isVisible: Boolean,
+    ) : PasswordRepromptAction() {
+        override val vaultItemAction: VaultItemAction
+            get() = VaultItemAction.ItemType.SshKey.PrivateKeyVisibilityClicked(
+                isVisible = isVisible,
+            )
+    }
 }
diff --git a/app/src/test/java/com/x8bit/bitwarden/ui/vault/feature/item/VaultItemViewModelTest.kt b/app/src/test/java/com/x8bit/bitwarden/ui/vault/feature/item/VaultItemViewModelTest.kt
index 81a3ac8c989..9eb5496eeba 100644
--- a/app/src/test/java/com/x8bit/bitwarden/ui/vault/feature/item/VaultItemViewModelTest.kt
+++ b/app/src/test/java/com/x8bit/bitwarden/ui/vault/feature/item/VaultItemViewModelTest.kt
@@ -2490,11 +2490,58 @@ class VaultItemViewModelTest : BaseViewModelTest() {
 
         @Suppress("MaxLineLength")
         @Test
-        fun `on PrivateKeyVisibilityClick should show password dialog when re-prompt is required`() =
+        fun `on PrivateKeyVisibilityClick should show private key when re-prompt is not required`() =
             runTest {
                 val sshKeyViewState = createViewState(
                     common = DEFAULT_COMMON.copy(requiresReprompt = false),
+                    type = DEFAULT_SSH_KEY_TYPE,
                 )
+                val sshKeyState = DEFAULT_STATE.copy(viewState = sshKeyViewState)
+                every {
+                    mockCipherView.toViewState(
+                        previousState = null,
+                        isPremiumUser = true,
+                        hasMasterPassword = true,
+                        totpCodeItemData = null,
+                        canDelete = true,
+                        canAssignToCollections = true,
+                    )
+                } returns sshKeyViewState
+                mutableVaultItemFlow.value = DataState.Loaded(data = mockCipherView)
+                mutableAuthCodeItemFlow.value = DataState.Loaded(data = null)
+                mutableCollectionsStateFlow.value = DataState.Loaded(emptyList())
+
+                assertEquals(sshKeyState, viewModel.stateFlow.value)
+                viewModel.trySendAction(
+                    VaultItemAction.ItemType.SshKey.PrivateKeyVisibilityClicked(
+                        isVisible = true,
+                    ),
+                )
+                assertEquals(
+                    sshKeyState.copy(
+                        viewState = sshKeyViewState.copy(
+                            common = DEFAULT_COMMON.copy(requiresReprompt = false),
+                            type = DEFAULT_SSH_KEY_TYPE.copy(showPrivateKey = true),
+                        ),
+                    ),
+                    viewModel.stateFlow.value,
+                )
+                verify(exactly = 1) {
+                    mockCipherView.toViewState(
+                        previousState = null,
+                        isPremiumUser = true,
+                        hasMasterPassword = true,
+                        totpCodeItemData = null,
+                        canDelete = true,
+                        canAssignToCollections = true,
+                    )
+                }
+            }
+
+        @Suppress("MaxLineLength")
+        @Test
+        fun `on PrivateKeyVisibilityClick should show password dialog when re-prompt is required`() =
+            runTest {
                 val sshKeyState = DEFAULT_STATE.copy(viewState = SSH_KEY_VIEW_STATE)
                 every {
                     mockCipherView.toViewState(
@@ -2518,15 +2565,22 @@ class VaultItemViewModelTest : BaseViewModelTest() {
                 )
                 assertEquals(
                     sshKeyState.copy(
-                        viewState = sshKeyViewState.copy(
-                            common = DEFAULT_COMMON,
-                            type = DEFAULT_SSH_KEY_TYPE.copy(
-                                showPrivateKey = true,
-                            ),
+                        dialog = VaultItemState.DialogState.MasterPasswordDialog(
+                            PasswordRepromptAction.ViewPrivateKeyClicked(isVisible = true),
                         ),
                     ),
                     viewModel.stateFlow.value,
                 )
+                verify(exactly = 1) {
+                    mockCipherView.toViewState(
+                        previousState = null,
+                        isPremiumUser = true,
+                        hasMasterPassword = true,
+                        totpCodeItemData = null,
+                        canDelete = true,
+                        canAssignToCollections = true,
+                    )
+                }
             }
 
         @Test