From b21163a421a7588a12c7f7eda2e3285e8764e54b Mon Sep 17 00:00:00 2001
From: Niels Andriesse <andriesseniels@gmail.com>
Date: Tue, 2 Mar 2021 15:14:00 +1100
Subject: [PATCH] Handle media saved notifications

---
 .../Conversations/ConversationVC+Interaction.swift   | 12 +++++++++++-
 .../Message Cells/InfoMessageCell.swift              |  1 +
 .../MessageReceiver+Handling.swift                   |  8 ++++++--
 3 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/Session/Conversations/ConversationVC+Interaction.swift b/Session/Conversations/ConversationVC+Interaction.swift
index 680c600e4..17bc8b56b 100644
--- a/Session/Conversations/ConversationVC+Interaction.swift
+++ b/Session/Conversations/ConversationVC+Interaction.swift
@@ -487,6 +487,7 @@ extension ConversationVC : InputViewDelegate, MessageCellDelegate, ContextMenuAc
     func save(_ viewItem: ConversationViewItem) {
         guard viewItem.canSaveMedia() else { return }
         viewItem.saveMediaAction()
+        sendMediaSavedNotificationIfNeeded(with: viewItem.interaction.timestamp)
     }
     
     func ban(_ viewItem: ConversationViewItem) {
@@ -654,7 +655,7 @@ extension ConversationVC : InputViewDelegate, MessageCellDelegate, ContextMenuAc
         audioSession.endAudioActivity(recordVoiceMessageActivity)
     }
     
-    // MARK: Screenshot Notifications
+    // MARK: Data Extraction Notifications
     @objc func sendScreenshotNotificationIfNeeded() {
         guard thread is TSContactThread else { return }
         let message = DataExtractionNotification()
@@ -663,6 +664,15 @@ extension ConversationVC : InputViewDelegate, MessageCellDelegate, ContextMenuAc
             MessageSender.send(message, in: self.thread, using: transaction)
         }
     }
+    
+    func sendMediaSavedNotificationIfNeeded(with timestamp: UInt64) {
+        guard thread is TSContactThread else { return }
+        let message = DataExtractionNotification()
+        message.kind = .mediaSaved(timestamp: timestamp)
+        Storage.write { transaction in
+            MessageSender.send(message, in: self.thread, using: transaction)
+        }
+    }
 
     // MARK: Requesting Permission
     func requestCameraPermissionIfNeeded() -> Bool {
diff --git a/Session/Conversations/Message Cells/InfoMessageCell.swift b/Session/Conversations/Message Cells/InfoMessageCell.swift
index 1391dc952..6520435aa 100644
--- a/Session/Conversations/Message Cells/InfoMessageCell.swift	
+++ b/Session/Conversations/Message Cells/InfoMessageCell.swift	
@@ -57,6 +57,7 @@ final class InfoMessageCell : MessageCell {
             } else {
                 icon = nil
             }
+        case .mediaSavedNotification: icon = UIImage(named: "ic_download")
         default: icon = nil
         }
         if let icon = icon {
diff --git a/SessionMessagingKit/Sending & Receiving/MessageReceiver+Handling.swift b/SessionMessagingKit/Sending & Receiving/MessageReceiver+Handling.swift
index 4ae6838a3..a389b9d5d 100644
--- a/SessionMessagingKit/Sending & Receiving/MessageReceiver+Handling.swift	
+++ b/SessionMessagingKit/Sending & Receiving/MessageReceiver+Handling.swift	
@@ -111,8 +111,12 @@ extension MessageReceiver {
         let transaction = transaction as! YapDatabaseReadWriteTransaction
         guard message.groupPublicKey == nil,
             let thread = TSContactThread.getWithContactId(message.sender!, transaction: transaction) else { return }
-        // TODO: Handle media saved type notifications
-        let message = DataExtractionNotificationInfoMessage(type: .screenshotNotification, sentTimestamp: message.sentTimestamp!, thread: thread, referencedAttachmentTimestamp: nil)
+        let type: TSInfoMessageType
+        switch message.kind! {
+        case .screenshot: type = .screenshotNotification
+        case .mediaSaved: type = .mediaSavedNotification
+        }
+        let message = DataExtractionNotificationInfoMessage(type: type, sentTimestamp: message.sentTimestamp!, thread: thread, referencedAttachmentTimestamp: nil)
         message.save(with: transaction)
     }