From 42c95177730e59f08731f9e75bdcefc0989c2bd6 Mon Sep 17 00:00:00 2001
From: Waseem akram <waseem07799@gmail.com>
Date: Sun, 10 Oct 2021 20:02:37 +0530
Subject: [PATCH] Add post processing support for `FileDownloadManager`.

MessageDownloadManager now extracts source from json and proceeds saving
---
 macOS/Services/FileDownloadManager.swift    | 24 +++++++++----
 macOS/Services/MessageDownloadManager.swift | 37 +++++++++++++++++++--
 2 files changed, 52 insertions(+), 9 deletions(-)

diff --git a/macOS/Services/FileDownloadManager.swift b/macOS/Services/FileDownloadManager.swift
index 704c222..bb7152e 100644
--- a/macOS/Services/FileDownloadManager.swift
+++ b/macOS/Services/FileDownloadManager.swift
@@ -26,6 +26,7 @@ class FileDownloadTask: ObservableObject {
     let savedFileLocation: URL
     let fileName: String
     var error: Error?
+    var beforeSave: ((URL) -> URL?)?
     var afterDownload: ((FileDownloadTask) -> Void)?
     
     @Published var state: State = .idle
@@ -34,11 +35,13 @@ class FileDownloadTask: ObservableObject {
         task.progress
     }
     
-    init (task: URLSessionDownloadTask, fileName: String, savedFileLocation: URL, afterDownload: ((FileDownloadTask) -> Void)? = nil) {
+    init (task: URLSessionDownloadTask, fileName: String, savedFileLocation: URL,
+          beforeSave: ((URL) -> URL?)? = nil, afterDownload: ((FileDownloadTask) -> Void)? = nil) {
         self.task = task
         self.savedFileLocation = savedFileLocation
         self.fileName = fileName
         self.error = nil
+        self.beforeSave = beforeSave
         self.afterDownload = afterDownload
     }
     
@@ -65,7 +68,7 @@ final class FileDownloadManager: NSObject {
     }
   
     func schedule(with request: URLRequest, fileName: String, saveLocation: URL? = nil,
-                  afterDownload: ((FileDownloadTask) -> Void)? = nil) -> FileDownloadTask {
+                  beforeSave: ((URL) -> URL?)? = nil, afterDownload: ((FileDownloadTask) -> Void)? = nil) -> FileDownloadTask {
         let downloadTask = session.downloadTask(with: request)
         let fileURL: URL
         if let saveLocation = saveLocation {
@@ -73,7 +76,11 @@ final class FileDownloadManager: NSObject {
         } else {
             fileURL = Self.downloadDirectoryURL.appendingPathComponent(fileName)
         }
-        let fileDownloadTask = FileDownloadTask(task: downloadTask, fileName: fileName, savedFileLocation: fileURL, afterDownload: afterDownload)
+        let fileDownloadTask = FileDownloadTask(task: downloadTask,
+                                                fileName: fileName,
+                                                savedFileLocation: fileURL,
+                                                beforeSave: beforeSave,
+                                                afterDownload: afterDownload)
         tasks[fileDownloadTask.id] = fileDownloadTask
         return fileDownloadTask
     }
@@ -87,13 +94,16 @@ extension FileDownloadManager: URLSessionDownloadDelegate {
             return
         }
         
-        let desiredUrl = task.savedFileLocation
+        let sourceFile = task.beforeSave?(location) ?? location
+        
+        let savingLocation = task.savedFileLocation
+        
         do {
-            let isFileExists = FileManager.default.fileExists(atPath: desiredUrl.path)
+            let isFileExists = FileManager.default.fileExists(atPath: savingLocation.path)
             if isFileExists {
-                _ = try FileManager.default.replaceItemAt(desiredUrl, withItemAt: location)
+                _ = try FileManager.default.replaceItemAt(savingLocation, withItemAt: sourceFile)
             } else {
-                try FileManager.default.moveItem(at: location, to: desiredUrl)
+                try FileManager.default.moveItem(at: sourceFile, to: savingLocation)
             }
             task.state = .saved
             task.afterDownload?(task)
diff --git a/macOS/Services/MessageDownloadManager.swift b/macOS/Services/MessageDownloadManager.swift
index 3db902f..5bf30a0 100644
--- a/macOS/Services/MessageDownloadManager.swift
+++ b/macOS/Services/MessageDownloadManager.swift
@@ -10,11 +10,13 @@ import Resolver
 import MailTMSwift
 
 class MessageDownloadManager {
-                
+    
     private var messageDownloadTasks: [Message.ID: FileDownloadTask] = [:]
     
     private let fileDownloadManager: FileDownloadManager
     
+    private let decoder = JSONDecoder()
+    
     init(fileDownloadManger: FileDownloadManager = Resolver.resolve()) {
         self.fileDownloadManager = fileDownloadManger
     }
@@ -33,9 +35,40 @@ class MessageDownloadManager {
         } else {
             fileName = "\(message.data.subject).eml"
         }
-        let task = fileDownloadManager.schedule(with: request, fileName: fileName, saveLocation: saveLocation, afterDownload: afterDownload)
+        let task = fileDownloadManager.schedule(with: request, fileName: fileName, saveLocation: saveLocation,
+                                                beforeSave: extractSource(location:), afterDownload: afterDownload)
         messageDownloadTasks[message.id] = task
         task.download()
         return task
     }
+    
+    func extractSource(location: URL) -> URL? {
+        guard FileManager.default.fileExists(atPath: location.path) else {
+            return nil
+        }
+        
+        do {
+            guard let data = try String(contentsOf: location).data(using: .utf8) else {
+                return nil
+            }
+            let sourceObj = try decoder.decode(MTMessageSource.self, from: data)
+            
+            let temporaryDirectoryURL =
+            try FileManager.default.url(for: .itemReplacementDirectory,
+                                            in: .userDomainMask,
+                                            appropriateFor: location,
+                                            create: true)
+
+            let temporaryFilename = UUID().uuidString
+
+            let temporaryFileURL =
+                temporaryDirectoryURL.appendingPathComponent(temporaryFilename)
+
+            try sourceObj.data.write(to: temporaryFileURL, atomically: true, encoding: .utf8)
+            return temporaryFileURL
+        } catch {
+            print(error)
+        }
+        return nil
+    }
 }