BIT-2408: Migrate attachments without keys when moving a cipher to an organization (#674)

Author: matt-livefront

BitwardenShared/Core/Vault/Repositories/VaultRepository.swift

@@ -392,6 +392,37 @@ class DefaultVaultRepository { // swiftlint:disable:this type_body_length
         return cipher
     }
 
+    /// Downloads, re-encrypts, and re-uploads an attachment without an attachment key so that it
+    /// can be shared to an organization.
+    ///
+    /// - Parameters:
+    ///   - attachment: The attachment that will be shared with the organization.
+    ///   - cipher: The cipher containing the attachment.
+    /// - Returns: The updated attachment with an attachment key that can be moved into the organization.
+    ///
+    private func fixCipherAttachment(
+        _ attachment: AttachmentView,
+        cipher: CipherView
+    ) async throws -> CipherView {
+        guard let downloadUrl = try await downloadAttachment(attachment, cipher: cipher) else {
+            throw BitwardenError.dataError("Unable to download attachment")
+        }
+
+        guard let cipherId = cipher.id else { throw CipherAPIServiceError.updateMissingId }
+        guard let fileName = attachment.fileName else { throw BitwardenError.dataError("Missing filename") }
+
+        let attachmentData = try Data(contentsOf: downloadUrl)
+        var updatedCipher = try await saveAttachment(cipherView: cipher, fileData: attachmentData, fileName: fileName)
+        try FileManager.default.removeItem(at: downloadUrl)
+
+        if let attachmentId = attachment.id,
+           let cipher = try await deleteAttachment(withId: attachmentId, cipherId: cipherId) {
+            updatedCipher = cipher
+        }
+
+        return updatedCipher
+    }
+
     /// Returns a list of `VaultListItem`s for the folders within a nested tree. By default, this
     /// will return the list items for the folders at the root of the tree. Specifying a
     /// `nestedFolderId` will return the list items for the children of the folder with the
@@ -1080,15 +1111,32 @@ extension DefaultVaultRepository: VaultRepository {
         return try await clientService.vault().ciphers().decrypt(cipher: updatedCipher)
     }
 
-    func shareCipher(_ cipher: CipherView, newOrganizationId: String, newCollectionIds: [String]) async throws {
+    func shareCipher(_ cipherView: CipherView, newOrganizationId: String, newCollectionIds: [String]) async throws {
+        // Ensure the cipher has a cipher key.
+        let encryptedCipher = try await encryptAndUpdateCipher(cipherView)
+        var cipherView = try await clientService.vault().ciphers().decrypt(cipher: encryptedCipher)
+
+        if let attachments = cipherView.attachments {
+            for attachment in attachments where attachment.key == nil {
+                // When moving a cipher to an organization, any attachments without an encryption
+                // key need to be re-encrypted with an attachment key.
+                cipherView = try await fixCipherAttachment(
+                    attachment,
+                    cipher: cipherView
+                )
+            }
+        }
+
         let organizationCipher = try await clientService.vault().ciphers()
             .moveToOrganization(
-                cipher: cipher,
+                cipher: cipherView,
                 organizationId: newOrganizationId
             )
             .update(collectionIds: newCollectionIds) // The SDK updates the cipher's organization ID.
-        let encryptedCipher = try await clientService.vault().ciphers().encrypt(cipherView: organizationCipher)
-        try await cipherService.shareCipherWithServer(encryptedCipher)
+
+        let encryptedOrganizationCipher = try await clientService.vault().ciphers()
+            .encrypt(cipherView: organizationCipher)
+        try await cipherService.shareCipherWithServer(encryptedOrganizationCipher)
     }
 
     func shouldShowUnassignedCiphersAlert() async -> Bool {

BitwardenShared/Core/Vault/Repositories/VaultRepositoryTests.swift

@@ -1037,13 +1037,80 @@ class VaultRepositoryTests: BitwardenTestCase { // swiftlint:disable:this type_b
         let updatedCipher = cipher.update(collectionIds: ["6", "7"])
 
         XCTAssertEqual(cipherService.shareCipherWithServerCiphers, [Cipher(cipherView: updatedCipher)])
-        XCTAssertEqual(clientCiphers.encryptedCiphers, [updatedCipher])
+        XCTAssertEqual(clientCiphers.encryptedCiphers.last, updatedCipher)
         XCTAssertEqual(clientCiphers.moveToOrganizationCipher, cipher)
         XCTAssertEqual(clientCiphers.moveToOrganizationOrganizationId, "5")
 
         XCTAssertEqual(cipherService.shareCipherWithServerCiphers.last, Cipher(cipherView: updatedCipher))
     }
 
+    /// `shareCipher()` migrates any attachments without an attachment key.
+    func test_shareCipher_attachmentMigration() async throws {
+        let account = Account.fixtureAccountLogin()
+        stateService.activeAccount = account
+
+        // The original cipher.
+        let cipherViewOriginal = CipherView.fixture(
+            attachments: [
+                .fixture(fileName: "file.txt", id: "1", key: nil),
+                .fixture(fileName: "existing-attachment-key.txt", id: "2", key: "abc"),
+            ],
+            id: "1"
+        )
+
+        // The cipher after saving the new attachment, encrypted with an attachment key.
+        let cipherAfterAttachmentSave = Cipher.fixture(
+            attachments: [
+                .fixture(id: "1", fileName: "file.txt", key: nil),
+                .fixture(id: "2", fileName: "existing-attachment-key.txt", key: "abc"),
+                .fixture(id: "3", fileName: "file.txt", key: "def"),
+            ],
+            id: "1"
+        )
+        cipherService.saveAttachmentWithServerResult = .success(cipherAfterAttachmentSave)
+
+        // The cipher after deleting the old attachment without an attachment key.
+        let cipherAfterAttachmentDelete = Cipher.fixture(
+            attachments: [
+                .fixture(id: "2", fileName: "existing-attachment-key.txt", key: "abc"),
+                .fixture(id: "3", fileName: "file.txt", key: "def"),
+            ],
+            id: "1"
+        )
+        cipherService.deleteAttachmentWithServerResult = .success(cipherAfterAttachmentDelete)
+        clientService.mockVault.clientCiphers.moveToOrganizationResult = .success(
+            CipherView(cipher: cipherAfterAttachmentDelete)
+        )
+
+        // Temporary download file (would normally be created by the network layer).
+        let downloadUrl = FileManager.default.temporaryDirectory.appendingPathComponent("file.txt")
+        try Data("📁".utf8).write(to: downloadUrl)
+        cipherService.downloadAttachmentResult = .success(downloadUrl)
+
+        // Decrypted download file (would normally be created by the SDK when decrypting the attachment).
+        let attachmentsUrl = try FileManager.default.attachmentsUrl(for: account.profile.userId)
+        try FileManager.default.createDirectory(at: attachmentsUrl, withIntermediateDirectories: true)
+        let decryptUrl = attachmentsUrl.appendingPathComponent("file.txt")
+        try Data("🗂️".utf8).write(to: decryptUrl)
+
+        try await subject.shareCipher(cipherViewOriginal, newOrganizationId: "5", newCollectionIds: ["6", "7"])
+
+        let updatedCipherView = CipherView(cipher: cipherAfterAttachmentDelete).update(collectionIds: ["6", "7"])
+
+        // Attachment migration: download attachment, save updated and delete old.
+        XCTAssertEqual(cipherService.downloadAttachmentId, "1")
+        XCTAssertEqual(cipherService.saveAttachmentWithServerCipher, Cipher(cipherView: cipherViewOriginal))
+        XCTAssertEqual(cipherService.deleteAttachmentWithServerAttachmentId, "1")
+        XCTAssertThrowsError(try Data(contentsOf: downloadUrl))
+        XCTAssertThrowsError(try Data(contentsOf: decryptUrl))
+
+        // Share cipher with updated attachments.
+        XCTAssertEqual(cipherService.shareCipherWithServerCiphers, [Cipher(cipherView: updatedCipherView)])
+        XCTAssertEqual(clientCiphers.encryptedCiphers.last, updatedCipherView)
+        XCTAssertEqual(clientCiphers.moveToOrganizationCipher, CipherView(cipher: cipherAfterAttachmentDelete))
+        XCTAssertEqual(clientCiphers.moveToOrganizationOrganizationId, "5")
+    }
+
     /// `shouldShowUnassignedCiphersAlert` is true if the feature flag is on,
     /// we should check for this user, the user has organizations, and the user has unassigned ciphers.
     func test_shouldShowUnassignedCiphersAlert() async {