Add post processing support for FileDownloadManager.

MessageDownloadManager now extracts source from json and proceeds saving
devwaseem · Oct 10, 2021 · 42c9517 · 42c9517
1 parent 4d3264a
commit 42c9517
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 9 deletions.
diff --git 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,15 +68,19 @@ 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 {
             fileURL = saveLocation
         } 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
@@ -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
+    }
 }