PM-14995 check for org ownership and premium status on items with TOT…

…P function

app/src/main/java/com/x8bit/bitwarden/data/auth/repository/model/Organization.kt

@@ -18,4 +18,5 @@ data class Organization(
     val shouldManageResetPassword: Boolean,
     val shouldUseKeyConnector: Boolean,
     val role: OrganizationType,
+    val shouldUsersGetPremium: Boolean,
 )

app/src/main/java/com/x8bit/bitwarden/data/auth/repository/util/SyncResponseJsonExtensions.kt

@@ -22,6 +22,7 @@ fun SyncResponseJson.Profile.Organization.toOrganization(): Organization =
         shouldUseKeyConnector = this.shouldUseKeyConnector,
         role = this.type,
         shouldManageResetPassword = this.permissions.shouldManageResetPassword,
+        shouldUsersGetPremium = this.shouldUsersGetPremium,
     )
 
 /**

app/src/main/java/com/x8bit/bitwarden/data/autofill/manager/AutofillTotpManagerImpl.kt

@@ -9,6 +9,7 @@ import com.x8bit.bitwarden.data.platform.manager.clipboard.BitwardenClipboardMan
 import com.x8bit.bitwarden.data.platform.repository.SettingsRepository
 import com.x8bit.bitwarden.data.vault.repository.VaultRepository
 import com.x8bit.bitwarden.data.vault.repository.model.GenerateTotpResult
+import com.x8bit.bitwarden.ui.vault.feature.vault.util.getOrganizationPremiumStatusMap
 import java.time.Clock
 
 /**
@@ -24,8 +25,15 @@ class AutofillTotpManagerImpl(
 ) : AutofillTotpManager {
     override suspend fun tryCopyTotpToClipboard(cipherView: CipherView) {
         if (settingsRepository.isAutoCopyTotpDisabled) return
+        val organizationPremiumStatusMap = authRepository
+            .userStateFlow
+            .value
+            ?.activeAccount
+            ?.getOrganizationPremiumStatusMap()
+            ?: emptyMap()
         val isPremium = authRepository.userStateFlow.value?.activeAccount?.isPremium == true
-        if (!isPremium && !cipherView.organizationUseTotp) return
+        val premiumStatus = organizationPremiumStatusMap[cipherView.organizationId] ?: isPremium
+        if (!premiumStatus && !cipherView.organizationUseTotp) return
         val totpCode = cipherView.login?.totp ?: return
 
         val totpResult = vaultRepository.generateTotp(

app/src/main/java/com/x8bit/bitwarden/data/vault/manager/TotpCodeManagerImpl.kt

@@ -122,6 +122,7 @@ class TotpCodeManagerImpl(
                                         CipherRepromptType.NONE -> false
                                     },
                                     orgUsesTotp = cipher.organizationUseTotp,
+                                    orgId = cipher.organizationId,
                                 )
                             }
                             .onFailure {

app/src/main/java/com/x8bit/bitwarden/data/vault/manager/model/VerificationCodeItem.kt

@@ -29,4 +29,5 @@ data class VerificationCodeItem(
     val username: String?,
     val hasPasswordReprompt: Boolean,
     val orgUsesTotp: Boolean,
+    val orgId: String?,
 )

app/src/main/java/com/x8bit/bitwarden/ui/platform/feature/search/SearchViewModel.kt

@@ -46,6 +46,7 @@ import com.x8bit.bitwarden.ui.platform.feature.search.util.updateWithAdditionalD
 import com.x8bit.bitwarden.ui.vault.feature.itemlisting.model.ListingItemOverflowAction
 import com.x8bit.bitwarden.ui.vault.feature.vault.model.VaultFilterData
 import com.x8bit.bitwarden.ui.vault.feature.vault.model.VaultFilterType
+import com.x8bit.bitwarden.ui.vault.feature.vault.util.getOrganizationPremiumStatusMap
 import com.x8bit.bitwarden.ui.vault.feature.vault.util.toFilteredList
 import com.x8bit.bitwarden.ui.vault.feature.vault.util.toVaultFilterData
 import com.x8bit.bitwarden.ui.vault.model.TotpData
@@ -107,6 +108,9 @@ class SearchViewModel @Inject constructor(
                 totpData = specialCircumstance?.toTotpDataOrNull(),
                 hasMasterPassword = userState.activeAccount.hasMasterPassword,
                 isPremium = userState.activeAccount.isPremium,
+                organizationPremiumStatusMap = userState
+                    .activeAccount
+                    .getOrganizationPremiumStatusMap(),
             )
         },
 ) {
@@ -687,6 +691,7 @@ class SearchViewModel @Inject constructor(
                                 isAutofill = state.isAutofill,
                                 isTotp = state.isTotp,
                                 isPremiumUser = state.isPremium,
+                                organizationPremiumStatusMap = state.organizationPremiumStatusMap,
                             )
                     }
 
@@ -733,6 +738,7 @@ data class SearchState(
     val totpData: TotpData?,
     val hasMasterPassword: Boolean,
     val isPremium: Boolean,
+    val organizationPremiumStatusMap: Map<String, Boolean>,
 ) : Parcelable {
 
     /**

app/src/main/java/com/x8bit/bitwarden/ui/platform/feature/search/util/SearchTypeDataExtensions.kt

@@ -152,6 +152,7 @@ fun List<CipherView>.toViewState(
     isAutofill: Boolean,
     isTotp: Boolean,
     isPremiumUser: Boolean,
+    organizationPremiumStatusMap: Map<String, Boolean>,
 ): SearchState.ViewState =
     when {
         searchTerm.isEmpty() -> SearchState.ViewState.Empty(message = null)
@@ -164,6 +165,7 @@ fun List<CipherView>.toViewState(
                     isAutofill = isAutofill,
                     isTotp = isTotp,
                     isPremiumUser = isPremiumUser,
+                    organizationPremiumStatusMap = organizationPremiumStatusMap,
                 )
                     .sortAlphabetically(),
             )
@@ -184,15 +186,17 @@ private fun List<CipherView>.toDisplayItemList(
     isAutofill: Boolean,
     isTotp: Boolean,
     isPremiumUser: Boolean,
+    organizationPremiumStatusMap: Map<String, Boolean>,
 ): List<SearchState.DisplayItem> =
     this.map {
+        val premiumStatus = organizationPremiumStatusMap[it.organizationId] ?: isPremiumUser
         it.toDisplayItem(
             baseIconUrl = baseIconUrl,
             hasMasterPassword = hasMasterPassword,
             isIconLoadingDisabled = isIconLoadingDisabled,
             isAutofill = isAutofill,
             isTotp = isTotp,
-            isPremiumUser = isPremiumUser,
+            isPremiumUser = premiumStatus,
         )
     }
 

app/src/main/java/com/x8bit/bitwarden/ui/vault/feature/item/VaultItemViewModel.kt

@@ -8,6 +8,7 @@ import com.bitwarden.vault.CipherView
 import com.x8bit.bitwarden.R
 import com.x8bit.bitwarden.data.auth.repository.AuthRepository
 import com.x8bit.bitwarden.data.auth.repository.model.BreachCountResult
+import com.x8bit.bitwarden.data.auth.repository.model.Organization
 import com.x8bit.bitwarden.data.auth.repository.model.UserState
 import com.x8bit.bitwarden.data.auth.repository.model.ValidatePasswordResult
 import com.x8bit.bitwarden.data.platform.manager.clipboard.BitwardenClipboardManager
@@ -1028,14 +1029,19 @@ class VaultItemViewModel @Inject constructor(
     ): VaultItemState.ViewState = this
         .data
         ?.cipher
-        ?.toViewState(
-            previousState = state.viewState.asContentOrNull(),
-            isPremiumUser = account.isPremium,
-            hasMasterPassword = account.hasMasterPassword,
-            totpCodeItemData = this.data?.totpCodeItemData,
-            canDelete = this.data?.canDelete == true,
-            canAssignToCollections = this.data?.canAssociateToCollections == true,
-        )
+        ?.let { cipher ->
+            val ownerOrg: Organization? = account.organizations.find {
+                cipher.organizationId == it.id
+            }
+            cipher.toViewState(
+                previousState = state.viewState.asContentOrNull(),
+                isPremiumUser = ownerOrg?.shouldUsersGetPremium ?: account.isPremium,
+                hasMasterPassword = account.hasMasterPassword,
+                totpCodeItemData = this.data?.totpCodeItemData,
+                canDelete = this.data?.canDelete == true,
+                canAssignToCollections = this.data?.canAssociateToCollections == true,
+            )
+        }
         ?: VaultItemState.ViewState.Error(message = errorText)
 
     private fun handleValidatePasswordReceive(

app/src/main/java/com/x8bit/bitwarden/ui/vault/feature/itemlisting/VaultItemListingViewModel.kt

@@ -13,9 +13,9 @@ import com.x8bit.bitwarden.data.auth.repository.model.ValidatePasswordResult
 import com.x8bit.bitwarden.data.auth.repository.model.ValidatePinResult
 import com.x8bit.bitwarden.data.autofill.accessibility.manager.AccessibilitySelectionManager
 import com.x8bit.bitwarden.data.autofill.fido2.manager.Fido2CredentialManager
+import com.x8bit.bitwarden.data.autofill.fido2.model.Fido2CreateCredentialRequest
 import com.x8bit.bitwarden.data.autofill.fido2.model.Fido2CredentialAssertionRequest
 import com.x8bit.bitwarden.data.autofill.fido2.model.Fido2CredentialAssertionResult
-import com.x8bit.bitwarden.data.autofill.fido2.model.Fido2CreateCredentialRequest
 import com.x8bit.bitwarden.data.autofill.fido2.model.Fido2GetCredentialsRequest
 import com.x8bit.bitwarden.data.autofill.fido2.model.Fido2GetCredentialsResult
 import com.x8bit.bitwarden.data.autofill.fido2.model.Fido2RegisterCredentialResult
@@ -70,6 +70,7 @@ import com.x8bit.bitwarden.ui.vault.feature.itemlisting.util.toVaultItemCipherTy
 import com.x8bit.bitwarden.ui.vault.feature.itemlisting.util.toViewState
 import com.x8bit.bitwarden.ui.vault.feature.itemlisting.util.updateWithAdditionalDataIfNecessary
 import com.x8bit.bitwarden.ui.vault.feature.vault.model.VaultFilterType
+import com.x8bit.bitwarden.ui.vault.feature.vault.util.getOrganizationPremiumStatusMap
 import com.x8bit.bitwarden.ui.vault.feature.vault.util.toAccountSummaries
 import com.x8bit.bitwarden.ui.vault.feature.vault.util.toActiveAccountSummary
 import com.x8bit.bitwarden.ui.vault.feature.vault.util.toFilteredList
@@ -138,6 +139,9 @@ class VaultItemListingViewModel @Inject constructor(
             fido2GetCredentialsRequest = specialCircumstance?.toFido2GetCredentialsRequestOrNull(),
             isPremium = userState.activeAccount.isPremium,
             isRefreshing = false,
+            organizationPremiumStatusMap = userState
+                .activeAccount
+                .getOrganizationPremiumStatusMap(),
         )
     },
 ) {
@@ -1580,6 +1584,7 @@ class VaultItemListingViewModel @Inject constructor(
                                 .fido2CredentialAutofillViewList,
                             totpData = state.totpData,
                             isPremiumUser = state.isPremium,
+                            organizationPremiumStatusMap = state.organizationPremiumStatusMap,
                         )
                     }
 
@@ -1745,6 +1750,7 @@ data class VaultItemListingState(
     val hasMasterPassword: Boolean,
     val isPremium: Boolean,
     val isRefreshing: Boolean,
+    val organizationPremiumStatusMap: Map<String, Boolean>,
 ) {
     /**
      * Whether or not the add FAB should be shown.

app/src/main/java/com/x8bit/bitwarden/ui/vault/feature/itemlisting/util/VaultItemListingDataExtensions.kt

@@ -111,6 +111,7 @@ fun VaultData.toViewState(
     fido2CredentialAutofillViews: List<Fido2CredentialAutofillView>?,
     totpData: TotpData?,
     isPremiumUser: Boolean,
+    organizationPremiumStatusMap: Map<String, Boolean>,
 ): VaultItemListingState.ViewState {
     val filteredCipherViewList = cipherViewList
         .filter { cipherView ->
@@ -142,6 +143,7 @@ fun VaultData.toViewState(
                 fido2CredentialAutofillViews = fido2CredentialAutofillViews,
                 isPremiumUser = isPremiumUser,
                 isTotp = totpData != null,
+                organizationPremiumStatusMap = organizationPremiumStatusMap,
             ),
             displayFolderList = folderList.map { folderView ->
                 VaultItemListingState.FolderDisplayItem(
@@ -290,8 +292,10 @@ private fun List<CipherView>.toDisplayItemList(
     fido2CredentialAutofillViews: List<Fido2CredentialAutofillView>?,
     isPremiumUser: Boolean,
     isTotp: Boolean,
+    organizationPremiumStatusMap: Map<String, Boolean>,
 ): List<VaultItemListingState.DisplayItem> =
     this.map {
+        val premiumStatus = organizationPremiumStatusMap[it.organizationId] ?: isPremiumUser
         it.toDisplayItem(
             baseIconUrl = baseIconUrl,
             hasMasterPassword = hasMasterPassword,
@@ -302,7 +306,7 @@ private fun List<CipherView>.toDisplayItemList(
                 ?.firstOrNull { fido2CredentialAutofillView ->
                     fido2CredentialAutofillView.cipherId == it.id
                 },
-            isPremiumUser = isPremiumUser,
+            isPremiumUser = premiumStatus,
             isTotp = isTotp,
         )
     }

app/src/main/java/com/x8bit/bitwarden/ui/vault/feature/vault/VaultViewModel.kt

@@ -37,6 +37,7 @@ import com.x8bit.bitwarden.ui.platform.manager.snackbar.SnackbarRelayManager
 import com.x8bit.bitwarden.ui.vault.feature.itemlisting.model.ListingItemOverflowAction
 import com.x8bit.bitwarden.ui.vault.feature.vault.model.VaultFilterData
 import com.x8bit.bitwarden.ui.vault.feature.vault.model.VaultFilterType
+import com.x8bit.bitwarden.ui.vault.feature.vault.util.getOrganizationPremiumStatusMap
 import com.x8bit.bitwarden.ui.vault.feature.vault.util.initials
 import com.x8bit.bitwarden.ui.vault.feature.vault.util.toAccountSummaries
 import com.x8bit.bitwarden.ui.vault.feature.vault.util.toActiveAccountSummary
@@ -102,6 +103,9 @@ class VaultViewModel @Inject constructor(
             isRefreshing = false,
             showImportActionCard = false,
             showSshKeys = showSshKeys,
+            organizationPremiumStatusMap = userState
+                .activeAccount
+                .getOrganizationPremiumStatusMap(),
         )
     },
 ) {
@@ -608,6 +612,7 @@ class VaultViewModel @Inject constructor(
                     hasMasterPassword = state.hasMasterPassword,
                     vaultFilterType = vaultFilterTypeOrDefault,
                     showSshKeys = showSshKeys,
+                    organizationPremiumStatusMap = state.organizationPremiumStatusMap,
                 ),
                 dialog = null,
                 isRefreshing = false,
@@ -644,6 +649,7 @@ class VaultViewModel @Inject constructor(
                     hasMasterPassword = state.hasMasterPassword,
                     vaultFilterType = vaultFilterTypeOrDefault,
                     showSshKeys = state.showSshKeys,
+                    organizationPremiumStatusMap = state.organizationPremiumStatusMap,
                 ),
             )
         }
@@ -715,6 +721,7 @@ data class VaultState(
     val isRefreshing: Boolean,
     val showImportActionCard: Boolean,
     val showSshKeys: Boolean,
+    val organizationPremiumStatusMap: Map<String, Boolean>,
 ) : Parcelable {
 
     /**
@@ -1334,6 +1341,7 @@ private fun MutableStateFlow<VaultState>.updateToErrorStateOrDialog(
                     vaultFilterType = vaultFilterType,
                     isIconLoadingDisabled = isIconLoadingDisabled,
                     showSshKeys = it.showSshKeys,
+                    organizationPremiumStatusMap = it.organizationPremiumStatusMap,
                 ),
                 dialog = VaultState.DialogState.Error(
                     title = errorTitle,

app/src/main/java/com/x8bit/bitwarden/ui/vault/feature/vault/util/UserStateExtensions.kt

@@ -71,3 +71,11 @@ fun UserState.Account.toVaultFilterData(
                 ),
             )
         }
+
+/**
+ * Returns a map of organization IDs and if they provide a premium status to the user for
+ * items owned by that organization.
+ */
+fun UserState.Account.getOrganizationPremiumStatusMap(): Map<String, Boolean> {
+    return organizations.associate { it.id to it.shouldUsersGetPremium }
+}

app/src/main/java/com/x8bit/bitwarden/ui/vault/feature/vault/util/VaultDataExtensions.kt

@@ -40,6 +40,7 @@ fun VaultData.toViewState(
     baseIconUrl: String,
     vaultFilterType: VaultFilterType,
     showSshKeys: Boolean,
+    organizationPremiumStatusMap: Map<String, Boolean>,
 ): VaultState.ViewState {
 
     val filteredCipherViewListWithDeletedItems =
@@ -73,14 +74,24 @@ fun VaultData.toViewState(
     return if (filteredCipherViewListWithDeletedItems.isEmpty()) {
         VaultState.ViewState.NoItems
     } else {
-        val totpItems = filteredCipherViewList.filter { it.login?.totp != null }
+        val totpItemsGroupedByOwnership = filteredCipherViewList.groupBy {
+            !it.organizationId.isNullOrBlank()
+        }
+        val userOwnedTotpItems = totpItemsGroupedByOwnership[false]
+            ?.filter {
+                it.login?.totp != null && isPremium
+            }
+            ?: emptyList()
+        val organizationOwnedTotpItems = totpItemsGroupedByOwnership[true]
+            ?.filter {
+            it.login?.totp != null &&
+                (organizationPremiumStatusMap[it.id] == true || it.organizationUseTotp)
+            }
+            ?: emptyList()
         VaultState.ViewState.Content(
             itemTypesCount = itemTypesCount,
-            totpItemsCount = if (isPremium) {
-                totpItems.count()
-            } else {
-                totpItems.count { it.organizationUseTotp }
-            },
+            totpItemsCount = userOwnedTotpItems.count() +
+                organizationOwnedTotpItems.count(),
             loginItemsCount = filteredCipherViewList.count { it.type == CipherType.LOGIN },
             cardItemsCount = filteredCipherViewList.count { it.type == CipherType.CARD },
             identityItemsCount = filteredCipherViewList.count { it.type == CipherType.IDENTITY },
@@ -94,7 +105,8 @@ fun VaultData.toViewState(
                         hasMasterPassword = hasMasterPassword,
                         isIconLoadingDisabled = isIconLoadingDisabled,
                         baseIconUrl = baseIconUrl,
-                        isPremiumUser = isPremium,
+                        isPremiumUser = organizationPremiumStatusMap[it.organizationId]
+                            ?: isPremium,
                     )
                 },
             folderItems = filteredFolderViewList
@@ -128,7 +140,8 @@ fun VaultData.toViewState(
                         hasMasterPassword = hasMasterPassword,
                         isIconLoadingDisabled = isIconLoadingDisabled,
                         baseIconUrl = baseIconUrl,
-                        isPremiumUser = isPremium,
+                        isPremiumUser = organizationPremiumStatusMap[it.organizationId]
+                            ?: isPremium,
                     )
                 }
                 .takeIf { it.size < NO_FOLDER_ITEM_THRESHOLD }

app/src/main/java/com/x8bit/bitwarden/ui/vault/feature/verificationcode/VerificationCodeViewModel.kt

@@ -18,6 +18,7 @@ import com.x8bit.bitwarden.ui.platform.base.util.asText
 import com.x8bit.bitwarden.ui.platform.base.util.concat
 import com.x8bit.bitwarden.ui.platform.components.model.IconData
 import com.x8bit.bitwarden.ui.vault.feature.vault.model.VaultFilterType
+import com.x8bit.bitwarden.ui.vault.feature.vault.util.getOrganizationPremiumStatusMap
 import com.x8bit.bitwarden.ui.vault.feature.vault.util.toLoginIconData
 import dagger.hilt.android.lifecycle.HiltViewModel
 import kotlinx.coroutines.flow.combine
@@ -319,8 +320,11 @@ class VerificationCodeViewModel @Inject constructor(
         authCodes: List<VerificationCodeItem>,
         userAccount: UserState.Account?,
     ): DataState<List<VerificationCodeItem>> {
+        val orgPremiumStatusMap = userAccount?.getOrganizationPremiumStatusMap() ?: emptyMap()
         val filteredAuthCodes = authCodes.mapNotNull { authCode ->
-            if (userAccount?.isPremium == true) {
+            val premiumStatus =
+                (authCode.orgId?.let { orgPremiumStatusMap[it] } ?: userAccount?.isPremium) == true
+            if (premiumStatus) {
                 authCode
             } else {
                 authCode.takeIf { it.orgUsesTotp }