Statistics
| Branch: | Tag: | Revision:

chatsecureios / ChatSecure / Classes / Controllers / FileTransferManager.swift @ 9a71f7ed

History | View | Annotate | Download (37.3 KB)

1
//
2
//  FileTransferManager.swift
3
//  ChatSecure
4
//
5
//  Created by Chris Ballinger on 3/28/17.
6
//  Copyright © 2017 Chris Ballinger. All rights reserved.
7
//
8

    
9
import Foundation
10
import XMPPFramework
11
import CocoaLumberjack
12
import OTRKit
13
import Alamofire
14
import OTRAssets
15

    
16
extension UIImage {
17
    enum DataSize {
18
        case unlimited
19
        case maxBytes(UInt)
20
        var numBytes: UInt {
21
            switch self {
22
            case .unlimited:
23
                return UInt.max
24
            case .maxBytes(let limit):
25
                return limit
26
            }
27
        }
28
    }
29
     struct Quality {
30
        static let low = Quality(initial: 0.4, decrementFactor: 0.65)
31
        static let medium = Quality(initial: 0.65, decrementFactor: 0.65)
32
        static let high = Quality(initial: 0.75, decrementFactor: 0.75)
33
        
34
        /// This value cannot be > 1 or bad things will happen
35
        let initial: CGFloat
36
        /// Multiplied to reduce the initial value. This value cannot be > 1 or bad things will happen
37
        let decrementFactor: CGFloat
38
    }
39
    func jpegData(dataSize: DataSize,
40
                  resize: Quality = Quality.medium,
41
                  jpeg: Quality = Quality.medium,
42
                  maxTries: UInt = 10) -> Data? {
43
        let image = self
44
        var sizeInBytes: UInt = 0
45
        var scaleFactor: CGFloat = resize.initial
46
        var jpegQuality: CGFloat = jpeg.initial
47
        let qualityDecrement: CGFloat = jpeg.decrementFactor
48
        let scaleDecrement: CGFloat = resize.decrementFactor
49
        var scaledImageData: Data? = nil
50
        var newSize = CGSize.zero
51
        let maxTries = 10
52
        var numTries = 0
53
        let maxSize = dataSize.numBytes
54
        while (sizeInBytes == 0 || sizeInBytes > maxSize) &&
55
            numTries < maxTries {
56
                numTries = numTries + 1
57
                newSize = CGSize(width: image.size.width * scaleFactor, height: image.size.height * scaleFactor)
58
                let scaledImage = UIImage.otr_image(with: image, scaledTo: newSize)
59
                scaledImageData = UIImageJPEGRepresentation(scaledImage, jpegQuality)
60
                if let imageData = scaledImageData {
61
                    sizeInBytes = UInt(imageData.count)
62
                    scaleFactor = scaleFactor * scaleDecrement
63
                    jpegQuality = jpegQuality * qualityDecrement
64
                } else {
65
                    DDLogError("Could not make JPEG out of image!")
66
                    return nil
67
                }
68
        }
69
        return scaledImageData
70
    }
71
}
72

    
73
public enum FileTransferError: CustomNSError {
74
    case unknown
75
    case noServers
76
    case serverError
77
    case exceedsMaxSize
78
    case urlFormatting
79
    case fileNotFound
80
    case keyGenerationError
81
    case cryptoError
82
    case automaticDownloadsDisabled
83
    
84
    public var localizedDescription: String {
85
        switch self {
86
        case .unknown:
87
            return UNKNOWN_ERROR_STRING()
88
        case .noServers:
89
            return NO_HTTP_UPLOAD_SERVERS_STRING() + " " + PLEASE_CONTACT_SERVER_OP_STRING()
90
        case .serverError:
91
            return UNKNOWN_ERROR_STRING() + " " + PLEASE_CONTACT_SERVER_OP_STRING()
92
        case .exceedsMaxSize:
93
            return FILE_EXCEEDS_MAX_SIZE_STRING()
94
        case .urlFormatting:
95
            return COULD_NOT_PARSE_URL_STRING()
96
        case .fileNotFound:
97
            return FILE_NOT_FOUND_STRING()
98
        case .cryptoError, .keyGenerationError:
99
            return errSSLCryptoString()
100
        case .automaticDownloadsDisabled:
101
            return AUTOMATIC_DOWNLOADS_DISABLED_STRING()
102
        }
103
    }
104
}
105

    
106
public class FileTransferManager: NSObject, OTRServerCapabilitiesDelegate {
107

    
108
    let httpFileUpload: XMPPHTTPFileUpload
109
    let serverCapabilities: OTRServerCapabilities
110
    let connection: YapDatabaseConnection
111
    let internalQueue = DispatchQueue(label: "FileTransferManager Queue")
112
    let callbackQueue = DispatchQueue.main
113
    let sessionManager: SessionManager
114
    private var servers: [HTTPServer] = []
115
    
116
    @objc public var canUploadFiles: Bool {
117
        return self.servers.first != nil
118
    }
119
    
120
    deinit {
121
        httpFileUpload.removeDelegate(self)
122
        serverCapabilities.removeDelegate(self)
123
    }
124
    
125
    @objc public init(connection: YapDatabaseConnection,
126
                serverCapabilities: OTRServerCapabilities,
127
                sessionConfiguration: URLSessionConfiguration) {
128
        self.serverCapabilities = serverCapabilities
129
        self.httpFileUpload = XMPPHTTPFileUpload()
130
        self.connection = connection
131
        self.sessionManager = Alamofire.SessionManager(configuration: sessionConfiguration)
132
        super.init()
133
        if let stream = serverCapabilities.xmppStream {
134
            httpFileUpload.activate(stream)
135
        }
136
        httpFileUpload.addDelegate(self, delegateQueue: DispatchQueue.main)
137
        serverCapabilities.addDelegate(self, delegateQueue: DispatchQueue.main)
138
        self.refreshCapabilities()
139
    }
140
    
141
    // MARK: - Public Methods
142
    
143
    /// This will fetch capabilities and setup XMPP transfer module if needed
144
    @objc public func refreshCapabilities() {
145
        guard let allCapabilities = serverCapabilities.allCapabilities else {
146
            serverCapabilities.fetchAllCapabilities()
147
            return
148
        }
149
        servers = serversFromCapabilities(capabilities: allCapabilities)
150
        serverCapabilities.fetchAllCapabilities()
151
    }
152

    
153
    private func upload(mediaItem: OTRMediaItem,
154
                        shouldEncrypt: Bool,
155
                       prefetchedData: Data?,
156
                       completion: @escaping (_ url: URL?, _ error: Error?) -> ()) {
157
        internalQueue.async {
158
            if let data = prefetchedData {
159
                self.upload(media: mediaItem, data: data, shouldEncrypt: shouldEncrypt, filename: mediaItem.filename, contentType: mediaItem.mimeType, completion: completion)
160
            } else {
161
                var url: URL? = nil
162
                self.connection.read({ (transaction) in
163
                    url = mediaItem.mediaServerURL(with: transaction)
164
                })
165
                if let url = url {
166
                    self.upload(media: mediaItem, file: url, shouldEncrypt: shouldEncrypt, completion: completion)
167
                } else {
168
                    let error = FileTransferError.fileNotFound
169
                    DDLogError("Upload filed: File not found \(error)")
170
                    self.callbackQueue.async {
171
                        completion(nil, error)
172
                    }
173
                }
174
            }
175
        }
176
    }
177
    
178
    /// Currently just a wrapper around sendData
179
    private func upload(media: OTRMediaItem,
180
                        file: URL,
181
                        shouldEncrypt: Bool,
182
                     completion: @escaping (_ url: URL?, _ error: Error?) -> ()) {
183
        internalQueue.async {
184
            do {
185
                let data = try Data(contentsOf: file)
186
                let mimeType = OTRKitGetMimeTypeForExtension(file.pathExtension)
187
                self.upload(media: media, data: data, shouldEncrypt: shouldEncrypt, filename: file.lastPathComponent, contentType: mimeType, completion: completion)
188
            } catch let error {
189
                DDLogError("Error sending file URL \(file): \(error)")
190
            }
191
        }
192
        
193
    }
194
    
195
    private func upload(media: OTRMediaItem,
196
                        data inData: Data,
197
                        shouldEncrypt: Bool,
198
                 filename: String,
199
                 contentType: String,
200
                 completion: @escaping (_ url: URL?, _ error: Error?) -> ()) {
201
        internalQueue.async {
202
            guard let service = self.servers.first else {
203
                DDLogWarn("No HTTP upload servers available")
204
                self.callbackQueue.async {
205
                    completion(nil, FileTransferError.noServers)
206
                }
207
                return
208
            }
209
            var data = inData
210
            
211
            // When resending images, sometimes we need to recompress them
212
            // to fit the max upload limit
213
            if UInt(data.count) > service.maxSize,
214
                let _ = media as? OTRImageItem,
215
                let image = UIImage(data: inData),
216
                let imageData = image.jpegData(dataSize: .maxBytes(service.maxSize), resize: UIImage.Quality.medium, jpeg: UIImage.Quality.medium, maxTries: 10)
217
                {
218
                    data = imageData
219
            }
220
            
221
            if UInt(data.count) > service.maxSize {
222
                DDLogError("HTTP Upload exceeds max size \(data.count) > \(service.maxSize)")
223
                self.callbackQueue.async {
224
                    completion(nil, FileTransferError.exceedsMaxSize)
225
                }
226
                return
227
            }
228
            
229
            // TODO: Refactor to use streaming encryption
230
            var outData = data
231
            var outKeyIv: Data? = nil
232
            if shouldEncrypt {
233
                guard let key = OTRPasswordGenerator.randomData(withLength: 32), let iv = OTRPasswordGenerator.randomData(withLength: 16) else {
234
                    DDLogError("Could not generate key/iv")
235
                    self.callbackQueue.async {
236
                        completion(nil, FileTransferError.keyGenerationError)
237
                    }
238
                    return
239
                }
240
                outKeyIv = iv + key
241
                do {
242
                    let crypted = try OTRCryptoUtility.encryptAESGCMData(data, key: key, iv: iv)
243
                    outData = crypted.data + crypted.authTag
244
                } catch let error {
245
                    outData = Data()
246
                    DDLogError("Could not encrypt data for file transfer \(error)")
247
                    self.callbackQueue.async {
248
                        completion(nil, error)
249
                    }
250
                    return
251
                }
252
            }
253
            
254
            
255
            self.httpFileUpload.requestSlot(fromService: service.jid, filename: filename, size: UInt(outData.count), contentType: contentType, completion: { (slot: XMPPSlot?, iq: XMPPIQ?, error: Error?) in
256
                guard let slot = slot else {
257
                    let outError = error ?? FileTransferError.serverError
258
                    DDLogError("\(service) failed to assign upload slot: \(outError)")
259
                    self.callbackQueue.async {
260
                        completion(nil, outError)
261
                    }
262
                    return
263
                }
264
                self.sessionManager.upload(outData, to: slot.putURL, method: .put)
265
                    .validate()
266
                    .responseData(queue: self.callbackQueue) { response in
267
                        switch response.result {
268
                        case .success:
269
                            if let outKeyIv = outKeyIv {
270
                                // If there's a AES-GCM key, we gotta put it in the url
271
                                // and change the scheme to `aesgcm`
272
                                if var components = URLComponents(url: slot.getURL, resolvingAgainstBaseURL: true) {
273
                                    components.scheme = URLScheme.aesgcm.rawValue
274
                                    components.fragment = outKeyIv.hexString()
275
                                    if let outURL = components.url {
276
                                        completion(outURL, nil)
277
                                    } else {
278
                                        completion(nil, FileTransferError.urlFormatting)
279
                                    }
280
                                } else {
281
                                    completion(nil, FileTransferError.urlFormatting)
282
                                }
283
                            } else {
284
                                // The plaintext case
285
                                completion(slot.getURL, nil)
286
                            }
287
                        case .failure(let error):
288
                            completion(nil, error)
289
                            DDLogError("Upload error: \(error)")
290
                        }
291
                    }.uploadProgress(queue: self.internalQueue) { progress in
292
                        DDLogVerbose("Download progress \(progress.fractionCompleted)")
293
                        self.connection.asyncReadWrite { transaction in
294
                            if let media = media.refetch(with: transaction) {
295
                                media.transferProgress = Float(progress.fractionCompleted)
296
                                media.save(with: transaction)
297
                                media.touchParentMessage(with: transaction)
298
                            }
299
                        }
300
                }
301
            })
302
        }
303
    }
304
    
305
    @objc public func send(videoURL url: URL, thread: OTRThreadOwner) {
306
        internalQueue.async {
307
            self.send(url: url, thread: thread, type: .video)
308
        }
309
    }
310
    
311
    private enum MediaURLType {
312
        case audio
313
        case video
314
        //case image
315
    }
316
    
317
    private func send(url: URL, thread: OTRThreadOwner, type: MediaURLType) {
318
        internalQueue.async {
319
            var item: OTRMediaItem? = nil
320
            switch type {
321
            case .audio:
322
                item = OTRAudioItem(audioURL: url, isIncoming: false)
323
            case .video:
324
                item = OTRVideoItem(videoURL: url, isIncoming: false)
325
            }
326
            guard let mediaItem = item else {
327
                DDLogError("No media item to share for URL: \(url)")
328
                return
329
            }
330
            
331
            guard let message = self.newOutgoingMessage(to: thread, mediaItem: mediaItem) else {
332
                DDLogError("No message could be created for \(thread) \(mediaItem)")
333
                return
334
            }
335
            mediaItem.parentObjectKey = message.messageKey
336
            mediaItem.parentObjectCollection = message.messageCollection
337
            let newPath = OTRMediaFileManager.path(for: mediaItem, buddyUniqueId: thread.threadIdentifier)
338
            self.connection.readWrite { transaction in
339
                message.save(with: transaction)
340
                mediaItem.save(with: transaction)
341
                
342
            }
343
            OTRMediaFileManager.shared.copyData(fromFilePath: url.path, toEncryptedPath: newPath, completion: { (result, copyError: Error?) in
344
                var prefetchedData: Data? = nil
345
                if FileManager.default.fileExists(atPath: url.path) {
346
                    do {
347
                        let attributes = try FileManager.default.attributesOfItem(atPath: url.path)
348
                        if let size = attributes[FileAttributeKey.size] as? NSNumber, size.uint64Value < 1024 * 1024 * 1 {
349
                            prefetchedData = try Data(contentsOf: url)
350
                        }
351
                    } catch let error {
352
                        DDLogError("Error prefetching data: \(error)")
353
                    }
354
                    do {
355
                        try FileManager.default.removeItem(atPath: url.path)
356
                    } catch let error {
357
                        DDLogError("Error removing video: \(error)")
358
                    }
359
                }
360
                message.messageError = copyError
361
                self.connection.readWrite({ (transaction) in
362
                    mediaItem.save(with: transaction)
363
                    message.save(with: transaction)
364
                })
365
                self.send(mediaItem: mediaItem, prefetchedData: prefetchedData, message: message)
366
            }, completionQueue: self.internalQueue)
367
        }
368
    }
369
    
370
    @objc public func send(audioURL url: URL, thread: OTRThreadOwner) {
371
        internalQueue.async {
372
            self.send(url: url, thread: thread, type: .audio)
373
        }
374
    }
375
    
376
    @objc public func send(image: UIImage, thread: OTRThreadOwner) {
377
        internalQueue.async {
378
            guard let service = self.servers.first, service.maxSize > 0 else {
379
                DDLogError("No HTTP upload service available!")
380
                return
381
            }
382
            let filename = "\(UUID().uuidString).jpg"
383
            let imageItem = OTRImageItem(filename: filename, size: image.size, mimeType: "image/jpeg", isIncoming: false)
384
            guard let message = self.newOutgoingMessage(to: thread, mediaItem: imageItem) else {
385
                DDLogError("No message could be created")
386
                return
387
            }
388
            imageItem.parentObjectKey = message.messageKey
389
            imageItem.parentObjectCollection = message.messageCollection
390
            self.connection.readWrite { transaction in
391
                message.save(with: transaction)
392
                imageItem.save(with: transaction)
393
            }
394
            guard let ourImageData = image.jpegData(dataSize: .unlimited, resize: UIImage.Quality.high, jpeg: UIImage.Quality.high, maxTries: 10) else {
395
                DDLogError("Could not make JPEG out of image!")
396
                return
397
            }
398
            OTRMediaFileManager.shared.setData(ourImageData, for: imageItem, buddyUniqueId: thread.threadIdentifier, completion: { (bytesWritten: Int, error: Error?) in
399
                self.connection.readWrite({ (transaction) in
400
                    imageItem.touchParentMessage(with: transaction)
401
                    if let error = error {
402
                        message.messageError = error
403
                        message.save(with: transaction)
404
                    }
405
                })
406
                if let imageData = image.jpegData(dataSize: .maxBytes(service.maxSize), resize: UIImage.Quality.medium, jpeg: UIImage.Quality.medium, maxTries: 10) {
407
                    self.send(mediaItem: imageItem, prefetchedData: imageData, message: message)
408
                } else {
409
                    DDLogError("Could not make JPEG out of image! Bad size")
410
                    message.messageError = FileTransferError.exceedsMaxSize
411
                    self.connection.readWrite { transaction in
412
                        message.save(with: transaction)
413
                    }
414
                }
415
            }, completionQueue: self.internalQueue)
416
        }
417
    }
418
    
419
    private func newOutgoingMessage(to thread: OTRThreadOwner, mediaItem: OTRMediaItem) -> OTRMessageProtocol? {
420
        if let buddy = thread as? OTRBuddy {
421
            let message = OTROutgoingMessage()!
422
            var security: OTRMessageTransportSecurity = .invalid
423
            self.connection.read({ (transaction) in
424
                security = buddy.preferredTransportSecurity(with: transaction)
425
            })
426
            message.buddyUniqueId = buddy.uniqueId
427
            message.mediaItemUniqueId = mediaItem.uniqueId
428
            message.messageSecurityInfo = OTRMessageEncryptionInfo(messageSecurity: security)
429
            return message
430
        } else if let room = thread as? OTRXMPPRoom {
431
            let message = OTRXMPPRoomMessage()!
432
            message.messageDate = Date()
433
            message.roomUniqueId = room.threadIdentifier
434
            message.roomJID = room.jid
435
            message.senderJID = room.ownJID
436
            message.state = .needsSending
437
            message.mediaItemId = mediaItem.uniqueId
438
            return message
439
        }
440
        return nil
441
    }
442
    
443
    @objc public func send(mediaItem: OTRMediaItem, prefetchedData: Data?, message: OTRMessageProtocol) {
444
        var shouldEncrypt = false
445
        switch message.messageSecurity {
446
        case .OMEMO, .OTR:
447
            shouldEncrypt = true
448
        case .invalid, .plaintext, .plaintextWithOTR:
449
            shouldEncrypt = false
450
        }
451
        
452
        self.upload(mediaItem: mediaItem, shouldEncrypt: shouldEncrypt, prefetchedData: prefetchedData, completion: { (_url: URL?, error: Error?) in
453
            guard let url = _url else {
454
                if let error = error {
455
                    DDLogError("Error uploading: \(error)")
456
                }
457
                self.connection.readWrite({ (transaction) in
458
                    message.messageError = error
459
                    message.save(with: transaction)
460
                })
461
                return
462
            }
463
            self.connection.readWrite({ (transaction) in
464
                mediaItem.transferProgress = 1.0
465
                message.messageText = url.absoluteString
466
                mediaItem.save(with: transaction)
467
                message.save(with: transaction)
468
            })
469
            self.queueOutgoingMessage(message: message)
470
        })
471
    }
472
    
473
    private func queueOutgoingMessage(message: OTRMessageProtocol) {
474
        let sendAction = OTRYapMessageSendAction(forMessage: message, date: nil)
475
        self.connection.readWrite { transaction in
476
            message.save(with: transaction)
477
            sendAction.save(with: transaction)
478
            if let thread = message.threadOwner(with: transaction)  {
479
                if let buddy = thread as? OTRBuddy {
480
                    buddy.lastMessageId = message.uniqueId
481
                    buddy.save(with: transaction)
482
                }
483
            }
484
        }
485
    }
486
    
487
    // MARK: - Private Methods
488
    
489
    private func serversFromCapabilities(capabilities: [XMPPJID : XMLElement]) -> [HTTPServer] {
490
        var servers: [HTTPServer] = []
491
        for (jid, element) in capabilities {
492
            let supported = element.supportsHTTPUpload()
493
            let maxSize = element.maxHTTPUploadSize()
494
            if supported && maxSize > 0 {
495
                let server = HTTPServer(jid: jid, maxSize: maxSize)
496
                servers.append(server)
497
            }
498
        }
499
        return servers
500
    }
501

    
502
    // MARK: - OTRServerCapabilitiesDelegate
503
    
504
    public func serverCapabilities(_ sender: OTRServerCapabilities, didDiscoverCapabilities capabilities: [XMPPJID : XMLElement]) {
505
        servers = serversFromCapabilities(capabilities: capabilities)
506
    }
507
}
508

    
509
// MARK: - Scanning and downloading incoming media
510
extension FileTransferManager {
511
    
512
    /** creates downloadmessages and then downloads if needed. parent message should already be saved! @warn Do not call from within an existing db transaction! */
513
    @objc public func createAndDownloadItemsIfNeeded(message: OTRMessageProtocol, readConnection: YapDatabaseConnection, force: Bool) {
514
        DispatchQueue.global(qos: .default).async {
515
            if message.messageMediaItemKey != nil || message.messageText?.count == 0 || message.downloadableURLs.count == 0 {
516
                //DDLogVerbose("Download of message not needed \(message.messageKey)")
517
                return
518
            }
519
            var downloads: [OTRDownloadMessage] = []
520
            var disableAutomaticURLFetching = false
521
            if !force {
522
                readConnection.read { (transaction) in
523
                    downloads = message.existingDownloads(with: transaction)
524
                    if let thread = message.threadOwner(with: transaction), let account = OTRAccount.fetchObject(withUniqueID: thread.threadAccountIdentifier, transaction: transaction) {
525
                        disableAutomaticURLFetching = account.disableAutomaticURLFetching
526
                    }
527
                }
528
            }
529
            if downloads.count == 0 {
530
                downloads = message.downloads()
531
                if downloads.count == 0 {
532
                    return
533
                }
534
                self.connection.readWrite({ (transaction) in
535
                    for download in downloads {
536
                        if disableAutomaticURLFetching,
537
                            let filename = download.downloadableURL?.absoluteString {
538
                            let media = OTRMediaItem.incomingItem(withFilename: filename, mimeType: nil)
539
                            media.parentObjectKey = download.uniqueId
540
                            media.parentObjectCollection = download.messageCollection
541
                            media.save(with: transaction)
542
                            download.messageMediaItemKey = media.uniqueId
543
                            download.messageError = FileTransferError.automaticDownloadsDisabled
544
                        }
545
                        download.save(with: transaction)
546
                    }
547
                    message.touch(with: transaction)
548
                })
549
            }
550
            if disableAutomaticURLFetching {
551
                DDLogVerbose("Automatic URL fetching disabled \(message.messageKey)")
552
                return
553
            }
554
            for download in downloads {
555
                self.downloadMediaIfNeeded(download)
556
            }
557
        }
558
    }
559
    
560
    /** Downloads media for a single downloadmessage */
561
    public func downloadMediaIfNeeded(_ downloadMessage: OTRDownloadMessage) {
562
        // Bail out if we've already downloaded the media
563
        if downloadMessage.messageMediaItemKey != nil &&
564
            downloadMessage.messageError == nil {
565
            // DDLogWarn("Already downloaded media for this item")
566
            return
567
        }
568
        guard let url = downloadMessage.downloadableURL else {
569
            DDLogWarn("Attempted to download message but couldn't parse a URL \(downloadMessage)")
570
            return
571
        }
572
        self.sessionManager.session.getTasksWithCompletionHandler { (tasks, _, _) in
573
            // Bail out if we've already got a task for this
574
            for task in tasks where task.originalRequest?.url == url {
575
                DDLogWarn("Already have outstanding task: \(task)")
576
                return
577
            }
578
            self.sessionManager.request(url, method: .head)
579
                .validate()
580
                .responseData(queue: self.internalQueue) { response in
581
                switch response.result {
582
                case .success:
583
                    DDLogInfo("HEAD response: \(String(describing: response.response?.allHeaderFields))")
584
                    if let headers = response.response?.allHeaderFields {
585
                        let contentType = headers["Content-Type"] as? String
586
                        let contentLength = headers["Content-Length"] as? UInt ?? 0
587
                        self.continueDownloading(downloadMessage: downloadMessage, url: url, contentType: contentType, contentLength: contentLength)
588
                    }
589
                case .failure(let error):
590
                    self.setError(error, onMessage: downloadMessage)
591
                    DDLogError("HEAD error: \(error)")
592
                }
593
            }
594
        }
595
    }
596
    
597
    private func setError(_ error: Error, onMessage downloadMessage: OTRDownloadMessage) {
598
        self.connection.readWrite { transaction in
599
            if let message = downloadMessage.refetch(with: transaction) {
600
                message.messageError = error
601
                message.save(with: transaction)
602
            }
603
        }
604
    }
605
    
606
    private func continueDownloading(downloadMessage: OTRDownloadMessage, url: URL, contentType: String?, contentLength: UInt) {
607
        var mediaItem: OTRMediaItem? = nil
608
        self.connection.readWrite { transaction in
609
            // Remove placeholder media item
610
            mediaItem = OTRMediaItem(forMessage: downloadMessage, transaction: transaction)
611
            mediaItem?.remove(with: transaction)
612
            mediaItem = OTRMediaItem.incomingItem(withFilename: url.lastPathComponent, mimeType: contentType)
613
            mediaItem?.parentObjectKey = downloadMessage.uniqueId
614
            mediaItem?.parentObjectCollection = downloadMessage.messageCollection
615
            mediaItem?.save(with: transaction)
616
            downloadMessage.messageMediaItemKey = mediaItem?.uniqueId
617
            downloadMessage.save(with: transaction)
618
        }
619
        guard let media = mediaItem else {
620
            DDLogError("Could not unwrap media item")
621
            self.setError(FileTransferError.fileNotFound, onMessage: downloadMessage)
622
            return
623
        }
624
        DDLogVerbose("Downloading media item at URL: \(url)")
625
        self.sessionManager.request(url)
626
            .validate()
627
            .responseData(queue: self.internalQueue) { response in
628
                self.finishDownload(downloadMessage: downloadMessage, mediaItem: media, inData: response.data, urlResponse: response.response, error: response.error)
629
                switch response.result {
630
                case .success:
631
                    DDLogVerbose("Download Successful")
632
                case .failure(let error):
633
                    self.setError(error, onMessage: downloadMessage)
634
                    DDLogError("Download Error \(error)")
635
                }
636
            }.downloadProgress(queue: self.internalQueue) { progress in
637
                DDLogVerbose("Download progress \(progress.fractionCompleted)")
638
                self.connection.asyncReadWrite { transaction in
639
                    if let media = media.refetch(with: transaction) {
640
                        media.transferProgress = Float(progress.fractionCompleted)
641
                        media.save(with: transaction)
642
                        media.touchParentMessage(with: transaction)
643
                    }
644
                }
645
        }
646
    }
647
    
648
    private func finishDownload(downloadMessage: OTRDownloadMessage, mediaItem: OTRMediaItem, inData: Data?, urlResponse: URLResponse?, error: Error?) {
649
        if let error = error {
650
            self.setError(error, onMessage: downloadMessage)
651
            DDLogError("Error downloading file \(error)")
652
            return
653
        }
654
        guard var data = inData, let response = urlResponse, let url = response.url else {
655
            self.setError(FileTransferError.fileNotFound, onMessage: downloadMessage)
656
            DDLogError("No data or response")
657
            return
658
        }
659
        DDLogVerbose("Received response \(response)")
660
        let authTagSize = 16 // i'm not sure if this can be assumed, but how else would we know the size?
661
        if let (key, iv) = url.aesGcmKey, data.count > authTagSize {
662
            DDLogVerbose("Received encrypted response, attempting decryption...")
663
            
664
            let cryptedData = data.subdata(in: 0..<data.count - authTagSize)
665
            let authTag = data.subdata(in: data.count - authTagSize..<data.count)
666
            let cryptoData = OTRCryptoData(data: cryptedData, authTag: authTag)
667
            do {
668
                data = try OTRCryptoUtility.decryptAESGCMData(cryptoData, key: key, iv: iv)
669
            } catch let error {
670
                self.setError(error, onMessage: downloadMessage)
671
                DDLogError("Error decrypting data: \(error)")
672
                return
673
            }
674
            DDLogVerbose("Decrpytion successful")
675
        }
676
        OTRMediaFileManager.sharedInstance().setData(data, for: mediaItem, buddyUniqueId: downloadMessage.threadId, completion: { (bytesWritten, error) in
677
            if let error = error {
678
                self.setError(error, onMessage: downloadMessage)
679
                DDLogError("Error copying data: \(error)")
680
                return
681
            }
682
            self.connection.asyncReadWrite({ (transaction) in
683
                mediaItem.transferProgress = 1.0
684
                
685
                if let audioItem = mediaItem as? OTRAudioItem, let url = OTRMediaServer.sharedInstance().url(for: mediaItem, buddyUniqueId: downloadMessage.threadId) {
686
                    audioItem.populateFromData(at: url)
687
                }
688
                
689
                mediaItem.save(with: transaction)
690
                // downloadMessage.save(with: transaction)
691
                if let message = downloadMessage.refetch(with: transaction) {
692
                    message.messageError = nil
693
                    message.save(with: transaction)
694
                } else {
695
                    DDLogError("Failed to refetch download message WTF \(downloadMessage)")
696
                }
697
            }, completionQueue: DispatchQueue.main,
698
               completionBlock: {
699
                UIApplication.shared.showLocalNotification(downloadMessage)
700
            })
701
        }, completionQueue: nil)
702
    }
703
}
704

    
705
extension OTRDownloadMessage {
706
    /// Turn aesgcm links into https links
707
    var downloadableURL: URL? {
708
        guard var downloadableURL = url else { return nil }
709
        if downloadableURL.isAesGcm, var components = URLComponents(url: downloadableURL, resolvingAgainstBaseURL: true) {
710
            components.scheme = URLScheme.https.rawValue
711
            if let rawURL = components.url {
712
                downloadableURL = rawURL
713
            }
714
        }
715
        return downloadableURL
716
    }
717
}
718

    
719
public extension OTRMessageProtocol {
720
    public var downloadableURLs: [URL] {
721
        return self.messageText?.downloadableURLs ?? []
722
    }
723
}
724

    
725
public extension OTRBaseMessage {
726
    @objc public var downloadableNSURLs: [NSURL] {
727
        return self.downloadableURLs as [NSURL]
728
    }
729
}
730

    
731
public extension OTRXMPPRoomMessage {
732
    @objc public var downloadableNSURLs: [NSURL] {
733
        return self.downloadableURLs as [NSURL]
734
    }
735
}
736

    
737
// MARK: - Extensions
738

    
739
fileprivate struct HTTPServer {
740
    /// service jid for upload service
741
    let jid: XMPPJID
742
    /// max upload size in bytes
743
    let maxSize: UInt
744
}
745

    
746
public extension XMLElement {
747
    
748
    // For use on a <query> element
749
    func supportsHTTPUpload() -> Bool {
750
        let features = self.elements(forName: "feature")
751
        var supported = false
752
        for feature in features {
753
            if let value = feature.attributeStringValue(forName: "var"),
754
                value == XMPPHTTPFileUploadNamespace  {
755
                supported = true
756
                break
757
            }
758
        }
759
        return supported
760
    }
761
    
762
    /// Returns 0 on failure, or max file size in bytes
763
    func maxHTTPUploadSize() -> UInt {
764
        var maxSize: UInt = 0
765
        let xes = self.elements(forXmlns: "jabber:x:data")
766
        
767
        for x in xes {
768
            let fields = x.elements(forName: "field")
769
            var correctXEP = false
770
            for field in fields {
771
                if let value = field.forName("value") {
772
                    if value.stringValue == XMPPHTTPFileUploadNamespace {
773
                        correctXEP = true
774
                    }
775
                    if let varMaxFileSize = field.attributeStringValue(forName: "var"), varMaxFileSize == "max-file-size" {
776
                        maxSize = value.stringValueAsNSUInteger()
777
                    }
778
                }
779
            }
780
            if correctXEP && maxSize > 0 {
781
                break
782
            }
783
        }
784
        
785
        return maxSize
786
    }
787
}
788

    
789
enum URLScheme: String {
790
    case https = "https"
791
    case aesgcm = "aesgcm"
792
    static let downloadableSchemes: [URLScheme] = [.https, .aesgcm]
793
}
794

    
795
extension URL {
796
    
797
    /** URL scheme matches aesgcm:// */
798
    var isAesGcm: Bool {
799
        return scheme == URLScheme.aesgcm.rawValue
800
    }
801
    
802
    /** Has hex anchor with key and IV. 48 bytes w/ 16 iv + 32 key */
803
    var anchorData: Data? {
804
        guard let anchor = self.fragment else { return nil }
805
        let data = anchor.dataFromHex()
806
        return data
807
    }
808
    
809
    var aesGcmKey: (key: Data, iv: Data)? {
810
        guard let data = self.anchorData, data.count == 48 else { return nil }
811
        let iv = data.subdata(in: 0..<16)
812
        let key = data.subdata(in: 16..<48)
813
        return (key, iv)
814
    }
815
}
816

    
817
public extension NSString {
818
    public var isSingleURLOnly: Bool {
819
        return (self as String).isSingleURLOnly
820
    }
821
}
822

    
823
public extension String {
824
    
825
    private var urlRanges: ([URL], [NSRange]) {
826
        guard let detector = try? NSDataDetector(types: NSTextCheckingResult.CheckingType.link.rawValue) else {
827
            return ([], [])
828
        }
829
        var urls: [URL] = []
830
        var ranges: [NSRange] = []
831
        let matches = detector.matches(in: self, options: NSRegularExpression.MatchingOptions(rawValue: 0), range: NSMakeRange(0, self.count))
832
        for match in matches where match.resultType == .link {
833
            if let url = match.url {
834
                urls.append(url)
835
                ranges.append(match.range)
836
            }
837
        }
838
        return (urls, ranges)
839
    }
840
    
841
    /** Grab any URLs from a string */
842
    public var urls: [URL] {
843
        let (urls, _) = urlRanges
844
        return urls
845
    }
846
    
847
    /** Returns true if the message is ONLY a single URL */
848
    public var isSingleURLOnly: Bool {
849
        let (_, ranges) = urlRanges
850
        guard ranges.count == 1,
851
            let range = ranges.first,
852
            range.length == self.count else {
853
            return false
854
        }
855
        return true
856
    }
857
    
858
    /** Use this for extracting potentially downloadable URLs from a message. Currently checks for https:// and aesgcm:// */
859
    public var downloadableURLs: [URL] {
860
        return urlsMatchingSchemes(URLScheme.downloadableSchemes)
861
    }
862
    
863
    fileprivate func urlsMatchingSchemes(_ schemes: [URLScheme]) -> [URL] {
864
        let urls = self.urls.filter {
865
            guard let scheme = $0.scheme else { return false }
866
            for inScheme in schemes where inScheme.rawValue == scheme {
867
                return true
868
            }
869
            return false
870
        }
871
        return urls
872
    }
873
}
874

    
875
public extension FileTransferManager {
876
    /// Returns whether or not message should be displayed or hidden from collection. Single incoming URLs should be hidden, for example.
877
    @objc public static func shouldDisplayMessage(_ message: OTRMessageProtocol, transaction: YapDatabaseReadTransaction) -> Bool {
878
        // Always show media messages
879
        if message.messageMediaItemKey != nil {
880
            return true
881
        }
882
        // Always show downloads
883
        if message is OTRDownloadMessage {
884
            return true
885
        }
886
        // Hide non-media messages that have no text
887
        guard let messageText = message.messageText else {
888
            return false
889
        }
890
        
891
        // Filter out messages that are aesgcm scheme file transfers
892
        if messageText.contains("aesgcm://"),
893
            message.messageError == nil {
894
            return false
895
        }
896
        
897
        // Filter out messages that are just URLs and have downloads
898
        if messageText.isSingleURLOnly,
899
            message.hasExistingDownloads(with: transaction) {
900
            return false
901
        }
902

    
903
        return true
904
    }
905
}