DownloadProcessor.kt (22812B) - raw


      1 package me.rhunk.snapenhance.download
      2 
      3 import android.content.Intent
      4 import android.graphics.Bitmap
      5 import android.graphics.BitmapFactory
      6 import android.net.Uri
      7 import android.widget.Toast
      8 import androidx.documentfile.provider.DocumentFile
      9 import com.google.gson.GsonBuilder
     10 import kotlinx.coroutines.Job
     11 import kotlinx.coroutines.job
     12 import kotlinx.coroutines.joinAll
     13 import kotlinx.coroutines.launch
     14 import kotlinx.coroutines.runBlocking
     15 import me.rhunk.snapenhance.RemoteSideContext
     16 import me.rhunk.snapenhance.bridge.DownloadCallback
     17 import me.rhunk.snapenhance.common.Constants
     18 import me.rhunk.snapenhance.common.ReceiversConfig
     19 import me.rhunk.snapenhance.common.data.FileType
     20 import me.rhunk.snapenhance.common.data.download.DownloadMediaType
     21 import me.rhunk.snapenhance.common.data.download.DownloadMetadata
     22 import me.rhunk.snapenhance.common.data.download.DownloadRequest
     23 import me.rhunk.snapenhance.common.data.download.InputMedia
     24 import me.rhunk.snapenhance.common.data.download.SplitMediaAssetType
     25 import me.rhunk.snapenhance.common.util.snap.MediaDownloaderHelper
     26 import me.rhunk.snapenhance.common.util.snap.RemoteMediaResolver
     27 import me.rhunk.snapenhance.core.features.impl.downloader.decoder.AttachmentType
     28 import me.rhunk.snapenhance.task.PendingTask
     29 import me.rhunk.snapenhance.task.PendingTaskListener
     30 import me.rhunk.snapenhance.task.Task
     31 import me.rhunk.snapenhance.task.TaskStatus
     32 import me.rhunk.snapenhance.task.TaskType
     33 import java.io.File
     34 import java.io.InputStream
     35 import java.net.HttpURLConnection
     36 import java.net.URL
     37 import java.util.concurrent.ConcurrentHashMap
     38 import javax.xml.parsers.DocumentBuilderFactory
     39 import javax.xml.transform.TransformerFactory
     40 import javax.xml.transform.dom.DOMSource
     41 import javax.xml.transform.stream.StreamResult
     42 import kotlin.coroutines.coroutineContext
     43 import kotlin.io.encoding.Base64
     44 import kotlin.io.encoding.ExperimentalEncodingApi
     45 
     46 data class DownloadedFile(
     47     val file: File,
     48     val fileType: FileType
     49 )
     50 
     51 /**
     52  * DownloadProcessor handles the download requests of the user
     53  */
     54 @OptIn(ExperimentalEncodingApi::class)
     55 class DownloadProcessor (
     56     private val remoteSideContext: RemoteSideContext,
     57     private val callback: DownloadCallback
     58 ) {
     59 
     60     private val translation by lazy {
     61         remoteSideContext.translation.getCategory("download_processor")
     62     }
     63 
     64     private val gson by lazy { GsonBuilder().setPrettyPrinting().create() }
     65 
     66     private fun fallbackToast(message: Any) {
     67         android.os.Handler(remoteSideContext.androidContext.mainLooper).post {
     68             Toast.makeText(remoteSideContext.androidContext, message.toString(), Toast.LENGTH_SHORT).show()
     69         }
     70     }
     71 
     72     private fun callbackOnSuccess(path: String) = runCatching {
     73         callback.onSuccess(path)
     74     }.onFailure {
     75         fallbackToast(it)
     76     }
     77 
     78     private fun callbackOnFailure(message: String, throwable: String? = null) = runCatching {
     79         callback.onFailure(message, throwable)
     80     }.onFailure {
     81         fallbackToast("$message\n$throwable")
     82     }
     83 
     84     private fun callbackOnProgress(message: String) = runCatching {
     85         callback.onProgress(message)
     86     }.onFailure {
     87         fallbackToast(it)
     88     }
     89 
     90     private fun newFFMpegProcessor(pendingTask: PendingTask) = FFMpegProcessor.newFFMpegProcessor(remoteSideContext, pendingTask)
     91 
     92     suspend fun saveMediaToGallery(pendingTask: PendingTask, inputFile: File, metadata: DownloadMetadata) {
     93         if (coroutineContext.job.isCancelled) return
     94 
     95         runCatching {
     96             var fileType = FileType.fromFile(inputFile)
     97 
     98             if (fileType.isImage) {
     99                 remoteSideContext.config.root.downloader.forceImageFormat.getNullable()?.let { format ->
    100                     val bitmap = BitmapFactory.decodeFile(inputFile.absolutePath) ?: throw Exception("Failed to decode bitmap")
    101                     @Suppress("DEPRECATION") val compressFormat = when (format) {
    102                         "png" -> Bitmap.CompressFormat.PNG
    103                         "jpg" -> Bitmap.CompressFormat.JPEG
    104                         "webp" -> Bitmap.CompressFormat.WEBP
    105                         else -> throw Exception("Invalid image format")
    106                     }
    107 
    108                     pendingTask.updateProgress("Converting image to $format")
    109                     inputFile.outputStream().use {
    110                         bitmap.compress(compressFormat, 100, it)
    111                     }
    112                     fileType = FileType.fromFile(inputFile)
    113                 }
    114             }
    115 
    116             val fileName = metadata.outputPath.substringAfterLast("/") + "." + fileType.fileExtension
    117 
    118             val outputFolder = DocumentFile.fromTreeUri(remoteSideContext.androidContext, Uri.parse(remoteSideContext.config.root.downloader.saveFolder.get()))
    119                 ?: throw Exception("Failed to open output folder")
    120 
    121             val outputFileFolder = metadata.outputPath.let {
    122                 if (it.contains("/")) {
    123                     it.substringBeforeLast("/").split("/").fold(outputFolder) { folder, name ->
    124                         folder.findFile(name) ?: folder.createDirectory(name)!!
    125                     }
    126                 } else {
    127                     outputFolder
    128                 }
    129             }
    130 
    131             // checks if the file already exists and if it does, compares its contents with the input file, if contents differ, deletes existing file.
    132             outputFileFolder.findFile(fileName)?.let { existingFile ->
    133                 pendingTask.updateProgress("Comparing existing media")
    134                 if (existingFile.length() != inputFile.length()) {
    135                     existingFile.delete()
    136                     return@let
    137                 }
    138 
    139                 remoteSideContext.androidContext.contentResolver.openInputStream(existingFile.uri)?.use { existingInputStream ->
    140                     val buffer1 = ByteArray(1024 * 1024)
    141                     val buffer2 = ByteArray(1024 * 1024)
    142                     var read1: Int
    143                     var read2: Int
    144 
    145                     inputFile.inputStream().use { inputStream ->
    146                         while (true) {
    147                             read1 = inputStream.read(buffer1)
    148                             read2 = existingInputStream.read(buffer2)
    149                             if (read1 != read2 || !buffer1.contentEquals(buffer2)) {
    150                                 existingFile.delete()
    151                                 return@let
    152                             }
    153                             if (read1 == -1) break
    154                         }
    155                     }
    156                 }
    157 
    158                 pendingTask.task.extra = existingFile.uri.toString()
    159                 pendingTask.success()
    160                 callbackOnFailure(translation["already_downloaded_toast"])
    161                 return
    162             }
    163 
    164             val outputFile = outputFileFolder.createFile(fileType.mimeType, fileName)!!
    165 
    166             pendingTask.updateProgress("Saving media to gallery")
    167             remoteSideContext.androidContext.contentResolver.openOutputStream(outputFile.uri)!!.use { outputStream ->
    168                 inputFile.inputStream().use { inputStream ->
    169                     inputStream.copyTo(outputStream)
    170                 }
    171             }
    172 
    173             pendingTask.task.extra = outputFile.uri.toString()
    174             pendingTask.success()
    175 
    176             runCatching {
    177                 remoteSideContext.androidContext.sendBroadcast(Intent("android.intent.action.MEDIA_SCANNER_SCAN_FILE").apply {
    178                     data = outputFile.uri
    179                 })
    180             }.onFailure {
    181                 remoteSideContext.log.error("Failed to scan media file", it)
    182                 callbackOnFailure(translation.format("failed_gallery_toast", "error" to it.toString()), it.message)
    183             }
    184 
    185             remoteSideContext.log.verbose("download complete")
    186             callbackOnSuccess(fileName)
    187         }.onFailure { exception ->
    188             remoteSideContext.log.error("Failed to save media to gallery", exception)
    189             callbackOnFailure(translation.format("failed_gallery_toast", "error" to exception.toString()), exception.message)
    190             pendingTask.fail("Failed to save media to gallery")
    191         }
    192     }
    193 
    194     private fun createMediaTempFile(): File {
    195         return File.createTempFile("media", ".tmp")
    196     }
    197 
    198     private fun downloadInputMedias(pendingTask: PendingTask, downloadRequest: DownloadRequest) = runBlocking {
    199         val jobs = mutableListOf<Job>()
    200         val downloadedMedias = mutableMapOf<InputMedia, File>()
    201         var totalSize = 1L
    202         val inputMediaDownloadedBytes = mutableMapOf<InputMedia, Long>()
    203         val inputMediaProgress = ConcurrentHashMap<InputMedia, String>()
    204 
    205         fun updateDownloadProgress() {
    206             pendingTask.updateProgress(
    207                 inputMediaProgress.values.joinToString("\n"),
    208                 progress = (inputMediaDownloadedBytes.values.sum() * 100 / totalSize.coerceAtLeast(1)).toInt().coerceIn(0, 100)
    209             )
    210         }
    211 
    212         downloadRequest.inputMedias.forEach { inputMedia ->
    213             fun setProgress(progress: String) {
    214                 inputMediaProgress[inputMedia] = progress
    215                 updateDownloadProgress()
    216             }
    217 
    218             fun handleInputStream(inputStream: InputStream, estimatedSize: Long = 0L) {
    219                 createMediaTempFile().apply {
    220                     val decryptedInputStream = (inputMedia.encryption?.decryptInputStream(inputStream) ?: inputStream).buffered()
    221                     val buffer = ByteArray(1024 * 1024 * 2) // 2MB
    222                     var read: Int
    223                     var totalRead = 0L
    224 
    225                     outputStream().use { outputStream ->
    226                         while (decryptedInputStream.read(buffer).also { read = it } != -1) {
    227                             outputStream.write(buffer, 0, read)
    228                             totalRead += read
    229                             inputMediaDownloadedBytes[inputMedia] = totalRead
    230                             setProgress("${totalRead / 1024}KB/${estimatedSize / 1024}KB")
    231                         }
    232                     }
    233                 }.also { downloadedMedias[inputMedia] = it }
    234             }
    235 
    236             launch {
    237                 when (inputMedia.type) {
    238                     DownloadMediaType.PROTO_MEDIA -> {
    239                         RemoteMediaResolver.downloadBoltMedia(Base64.UrlSafe.decode(inputMedia.content), decryptionCallback = { it }, resultCallback = { inputStream, length ->
    240                             totalSize += length
    241                             inputStream.use {
    242                                 handleInputStream(it, estimatedSize = length)
    243                             }
    244                         })
    245                     }
    246                     DownloadMediaType.REMOTE_MEDIA -> {
    247                         with(URL(inputMedia.content).openConnection() as HttpURLConnection) {
    248                             requestMethod = "GET"
    249                             setRequestProperty("User-Agent", Constants.USER_AGENT)
    250                             connect()
    251                             totalSize += contentLength.toLong()
    252                             inputStream.use {
    253                                 handleInputStream(it, estimatedSize = contentLength.toLong())
    254                             }
    255                         }
    256                     }
    257                     DownloadMediaType.DIRECT_MEDIA -> {
    258                         val decoded = Base64.UrlSafe.decode(inputMedia.content)
    259                         totalSize += decoded.size.toLong()
    260                         handleInputStream(decoded.inputStream(), estimatedSize = decoded.size.toLong())
    261                     }
    262                     else -> {
    263                         File(inputMedia.content).inputStream().use {
    264                             totalSize += it.available().toLong()
    265                             handleInputStream(it, estimatedSize = it.available().toLong())
    266                         }
    267                     }
    268                 }
    269             }.also { jobs.add(it) }
    270         }
    271 
    272         jobs.joinAll()
    273         downloadedMedias
    274     }
    275 
    276     private suspend fun downloadRemoteMedia(pendingTask: PendingTask, metadata: DownloadMetadata, downloadedMedias: Map<InputMedia, DownloadedFile>, downloadRequest: DownloadRequest) {
    277         downloadRequest.inputMedias.first().let { inputMedia ->
    278             val mediaType = inputMedia.type
    279             val media = downloadedMedias[inputMedia]!!
    280 
    281             if (!downloadRequest.isDashPlaylist) {
    282                 if (inputMedia.attachmentType == AttachmentType.NOTE.key) {
    283                     remoteSideContext.config.root.downloader.forceVoiceNoteFormat.getNullable()?.let { format ->
    284                         val outputFile = File.createTempFile("voice_note", ".$format")
    285                         newFFMpegProcessor(pendingTask).execute(FFMpegProcessor.Request(
    286                             action = FFMpegProcessor.Action.CONVERSION,
    287                             inputs = listOf(media.file.absolutePath),
    288                             output = outputFile
    289                         ))
    290                         media.file.delete()
    291                         saveMediaToGallery(pendingTask, outputFile, metadata)
    292                         outputFile.delete()
    293                         return
    294                     }
    295                 }
    296 
    297                 saveMediaToGallery(pendingTask, media.file, metadata)
    298                 media.file.delete()
    299                 return
    300             }
    301 
    302             assert(mediaType == DownloadMediaType.REMOTE_MEDIA)
    303 
    304             val playlistXml = DocumentBuilderFactory.newInstance().newDocumentBuilder().parse(media.file)
    305             val baseUrlNodeList = playlistXml.getElementsByTagName("BaseURL")
    306             for (i in 0 until baseUrlNodeList.length) {
    307                 val baseUrlNode = baseUrlNodeList.item(i)
    308                 val baseUrl = baseUrlNode.textContent
    309                 baseUrlNode.textContent = "${RemoteMediaResolver.CF_ST_CDN_D}$baseUrl"
    310             }
    311 
    312             val dashOptions = downloadRequest.dashOptions!!
    313 
    314             val dashPlaylistFile = renameFromFileType(media.file, FileType.MPD)
    315             dashPlaylistFile.outputStream().use {
    316                 TransformerFactory.newInstance().newTransformer().transform(DOMSource(playlistXml), StreamResult(it))
    317             }
    318 
    319             callbackOnProgress(translation.format("download_toast", "path" to dashPlaylistFile.nameWithoutExtension))
    320             val outputFile = File.createTempFile("dash", ".mp4")
    321             runCatching {
    322                 newFFMpegProcessor(pendingTask).execute(FFMpegProcessor.Request(
    323                     action = FFMpegProcessor.Action.DOWNLOAD_DASH,
    324                     inputs = listOf(dashPlaylistFile.absolutePath),
    325                     output = outputFile,
    326                     startTime = dashOptions.offsetTime,
    327                     duration = dashOptions.duration
    328                 ))
    329                 saveMediaToGallery(pendingTask, outputFile, metadata)
    330             }.onFailure { exception ->
    331                 if (coroutineContext.job.isCancelled) return@onFailure
    332                 remoteSideContext.log.error("Failed to download dash media", exception)
    333                 callbackOnFailure(translation.format("failed_processing_toast", "error" to exception.toString()), exception.message)
    334                 pendingTask.fail("Failed to download dash media")
    335             }
    336 
    337             dashPlaylistFile.delete()
    338             outputFile.delete()
    339             media.file.delete()
    340         }
    341     }
    342 
    343     private fun renameFromFileType(file: File, fileType: FileType): File {
    344         val newFile = File(file.parentFile, file.nameWithoutExtension + "." + fileType.fileExtension)
    345         file.renameTo(newFile)
    346         return newFile
    347     }
    348 
    349     fun enqueue(downloadRequest: DownloadRequest, downloadMetadata: DownloadMetadata) {
    350         remoteSideContext.coroutineScope.launch {
    351             remoteSideContext.taskManager.getTaskByHash(downloadMetadata.mediaIdentifier)?.let { task ->
    352                 remoteSideContext.log.debug("already queued or downloaded")
    353 
    354                 if (task.status.isFinalStage()) {
    355                     if (task.status != TaskStatus.SUCCESS) return@let
    356                     // check if the media file has been deleted
    357                     if (task.type == TaskType.DOWNLOAD) {
    358                         val outputFile = runCatching {
    359                             DocumentFile.fromTreeUri(remoteSideContext.androidContext, Uri.parse(task.extra))
    360                         }.getOrNull()
    361 
    362                         if (outputFile != null && !outputFile.exists()) {
    363                             return@let
    364                         }
    365                     }
    366                     callbackOnFailure(translation["already_downloaded_toast"], null)
    367                 } else {
    368                     callbackOnFailure(translation["already_queued_toast"], null)
    369                 }
    370                 return@launch
    371             }
    372 
    373             remoteSideContext.log.debug("downloading media")
    374             val pendingTask = remoteSideContext.taskManager.createPendingTask(
    375                 Task(
    376                     type = TaskType.DOWNLOAD,
    377                     title = downloadMetadata.downloadSource,
    378                     author = downloadMetadata.mediaAuthor,
    379                     hash = downloadMetadata.mediaIdentifier
    380                 )
    381             ).apply {
    382                 status = TaskStatus.RUNNING
    383                 addListener(PendingTaskListener(onCancel = {
    384                     coroutineContext.job.cancel()
    385                 }))
    386                 updateProgress("Downloading...")
    387             }
    388 
    389             runCatching {
    390                 if (downloadRequest.isAudioStream) {
    391                     val streamUrl = downloadRequest.inputMedias.first().content
    392                     val outputFile = File.createTempFile("audio_stream", ".mp3")
    393 
    394                     callbackOnProgress("Downloading audio stream")
    395                     pendingTask.updateProgress("Downloading audio stream")
    396                     newFFMpegProcessor(pendingTask).execute(FFMpegProcessor.Request(
    397                         action = FFMpegProcessor.Action.DOWNLOAD_AUDIO_STREAM,
    398                         inputs = listOf(streamUrl),
    399                         output = outputFile,
    400                         audioStreamFormat = downloadRequest.audioStreamFormat
    401                     ))
    402                     saveMediaToGallery(pendingTask, outputFile, downloadMetadata)
    403                     return@launch
    404                 }
    405 
    406                 //first download all input medias into cache
    407                 val downloadedMedias = downloadInputMedias(pendingTask, downloadRequest).map {
    408                     it.key to DownloadedFile(it.value, FileType.fromFile(it.value))
    409                 }.toMap().toMutableMap()
    410                 remoteSideContext.log.verbose("downloaded ${downloadedMedias.size} medias")
    411 
    412                 var shouldMergeOverlay = downloadRequest.shouldMergeOverlay
    413 
    414                 //if there is a zip file, extract it and replace the downloaded media with the extracted ones
    415                 downloadedMedias.values.find { it.fileType == FileType.ZIP }?.let { zipFile ->
    416                     val oldDownloadedMedias = downloadedMedias.toMap()
    417                     downloadedMedias.clear()
    418 
    419                     zipFile.file.inputStream().use { zipFileInputStream ->
    420                         MediaDownloaderHelper.getSplitElements(zipFileInputStream) { type, inputStream ->
    421                             createMediaTempFile().apply {
    422                                 outputStream().use {
    423                                     inputStream.copyTo(it)
    424                                 }
    425                             }.also {
    426                                 downloadedMedias[InputMedia(
    427                                     type = DownloadMediaType.LOCAL_MEDIA,
    428                                     content = it.absolutePath,
    429                                     isOverlay = type == SplitMediaAssetType.OVERLAY
    430                                 )] = DownloadedFile(it, FileType.fromFile(it))
    431                             }
    432                         }
    433                     }
    434 
    435                     oldDownloadedMedias.forEach { (_, value) ->
    436                         value.file.delete()
    437                     }
    438 
    439                     shouldMergeOverlay = true
    440                 }
    441 
    442                 if (shouldMergeOverlay) {
    443                     assert(downloadedMedias.size == 2)
    444                     val media = downloadedMedias.entries.first { !it.key.isOverlay }.value
    445                     val overlayMedia = downloadedMedias.entries.first { it.key.isOverlay }.value
    446 
    447                     val renamedMedia = renameFromFileType(media.file, media.fileType)
    448                     val renamedOverlayMedia = renameFromFileType(overlayMedia.file, overlayMedia.fileType)
    449                     val mergedOverlay: File = File.createTempFile("merged", ".mp4")
    450                     runCatching {
    451                         callbackOnProgress(translation.format("processing_toast", "path" to media.file.nameWithoutExtension))
    452 
    453                         newFFMpegProcessor(pendingTask).execute(FFMpegProcessor.Request(
    454                             action = FFMpegProcessor.Action.MERGE_OVERLAY,
    455                             inputs = listOf(renamedMedia.absolutePath),
    456                             output = mergedOverlay,
    457                             overlay = renamedOverlayMedia
    458                         ))
    459 
    460                         saveMediaToGallery(pendingTask, mergedOverlay, downloadMetadata)
    461                     }.onFailure { exception ->
    462                         if (coroutineContext.job.isCancelled) return@onFailure
    463                         remoteSideContext.log.error("Failed to merge overlay", exception)
    464                         callbackOnFailure(translation.format("failed_processing_toast", "error" to exception.toString()), exception.message)
    465                         pendingTask.fail("Failed to merge overlay")
    466                     }
    467 
    468                     mergedOverlay.delete()
    469                     renamedOverlayMedia.delete()
    470                     renamedMedia.delete()
    471                     return@launch
    472                 }
    473 
    474                 downloadRemoteMedia(pendingTask, downloadMetadata, downloadedMedias, downloadRequest)
    475             }.onFailure { exception ->
    476                 pendingTask.fail("Failed to download media")
    477                 remoteSideContext.log.error("Failed to download media", exception)
    478                 callbackOnFailure(translation["failed_generic_toast"], exception.message)
    479             }
    480         }
    481     }
    482 
    483     fun onReceive(intent: Intent) {
    484         val downloadMetadata = gson.fromJson(intent.getStringExtra(ReceiversConfig.DOWNLOAD_METADATA_EXTRA)!!, DownloadMetadata::class.java)
    485         val downloadRequest = gson.fromJson(intent.getStringExtra(ReceiversConfig.DOWNLOAD_REQUEST_EXTRA)!!, DownloadRequest::class.java)
    486 
    487         enqueue(downloadRequest, downloadMetadata)
    488     }
    489 }