feat(download_manager): download multiple attachment

- experimental message decoder - fix async calls for downloads database - add debug view for ProtoReader
2023-09-14 23:55:06 +02:00
parent 5767f93331
commit 904d9175b6
15 changed files with 487 additions and 168 deletions
--- a/app/src/main/kotlin/me/rhunk/snapenhance/download/DownloadProcessor.kt
+++ b/app/src/main/kotlin/me/rhunk/snapenhance/download/DownloadProcessor.kt
@@ -8,8 +8,6 @@ import android.net.Uri
 import android.widget.Toast
 import androidx.documentfile.provider.DocumentFile
 import com.google.gson.GsonBuilder
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.job
 import kotlinx.coroutines.joinAll
@@ -254,7 +252,7 @@ class DownloadProcessor (
            val media = downloadedMedias[inputMedia]!!

            if (!downloadRequest.isDashPlaylist) {
-                if (inputMedia.messageContentType == "NOTE") {
+                if (inputMedia.attachmentType == "NOTE") {
                    remoteSideContext.config.root.downloader.forceVoiceNoteFormat.getNullable()?.let { format ->
                        val outputFile = File.createTempFile("voice_note", ".$format")
                        ffmpegProcessor.execute(FFMpegProcessor.Request(
--- a/app/src/main/kotlin/me/rhunk/snapenhance/download/DownloadTaskManager.kt
+++ b/app/src/main/kotlin/me/rhunk/snapenhance/download/DownloadTaskManager.kt
@@ -9,11 +9,13 @@ import me.rhunk.snapenhance.core.download.data.MediaDownloadSource
 import me.rhunk.snapenhance.core.util.SQLiteDatabaseHelper
 import me.rhunk.snapenhance.core.util.ktx.getIntOrNull
 import me.rhunk.snapenhance.core.util.ktx.getStringOrNull
+import java.util.concurrent.Executors

 class DownloadTaskManager {
    private lateinit var taskDatabase: SQLiteDatabase
    private val pendingTasks = mutableMapOf<Int, DownloadObject>()
    private val cachedTasks = mutableMapOf<Int, DownloadObject>()
+    private val executor = Executors.newSingleThreadExecutor()

    @SuppressLint("Range")
    fun init(context: Context) {
@@ -34,39 +36,43 @@ class DownloadTaskManager {
        }
    }

-    fun addTask(task: DownloadObject): Int {
-        taskDatabase.execSQL("INSERT INTO tasks (hash, outputPath, outputFile, downloadSource, mediaAuthor, iconUrl, downloadStage) VALUES (?, ?, ?, ?, ?, ?, ?)",
-            arrayOf(
-                task.metadata.mediaIdentifier,
-                task.metadata.outputPath,
-                task.outputFile,
-                task.metadata.downloadSource,
-                task.metadata.mediaAuthor,
-                task.metadata.iconUrl,
-                task.downloadStage.name
+    fun addTask(task: DownloadObject) {
+        executor.execute {
+            taskDatabase.execSQL("INSERT INTO tasks (hash, outputPath, outputFile, downloadSource, mediaAuthor, iconUrl, downloadStage) VALUES (?, ?, ?, ?, ?, ?, ?)",
+                arrayOf(
+                    task.metadata.mediaIdentifier,
+                    task.metadata.outputPath,
+                    task.outputFile,
+                    task.metadata.downloadSource,
+                    task.metadata.mediaAuthor,
+                    task.metadata.iconUrl,
+                    task.downloadStage.name
+                )
            )
-        )
-        task.downloadId = taskDatabase.rawQuery("SELECT last_insert_rowid()", null).use {
-            it.moveToFirst()
-            it.getInt(0)
+            task.downloadId = taskDatabase.rawQuery("SELECT last_insert_rowid()", null).use {
+                it.moveToFirst()
+                it.getInt(0)
+            }
+            pendingTasks[task.downloadId] = task
        }
-        pendingTasks[task.downloadId] = task
-        return task.downloadId
    }

    fun updateTask(task: DownloadObject) {
-        taskDatabase.execSQL("UPDATE tasks SET hash = ?, outputPath = ?, outputFile = ?, downloadSource = ?, mediaAuthor = ?, iconUrl = ?, downloadStage = ? WHERE id = ?",
-            arrayOf(
-                task.metadata.mediaIdentifier,
-                task.metadata.outputPath,
-                task.outputFile,
-                task.metadata.downloadSource,
-                task.metadata.mediaAuthor,
-                task.metadata.iconUrl,
-                task.downloadStage.name,
-                task.downloadId
+        executor.execute {
+            taskDatabase.execSQL(
+                "UPDATE tasks SET hash = ?, outputPath = ?, outputFile = ?, downloadSource = ?, mediaAuthor = ?, iconUrl = ?, downloadStage = ? WHERE id = ?",
+                arrayOf(
+                    task.metadata.mediaIdentifier,
+                    task.metadata.outputPath,
+                    task.outputFile,
+                    task.metadata.downloadSource,
+                    task.metadata.mediaAuthor,
+                    task.metadata.iconUrl,
+                    task.downloadStage.name,
+                    task.downloadId
+                )
            )
-        )
+        }
        //if the task is no longer active, move it to the cached tasks
        if (task.isJobActive()) {
            pendingTasks[task.downloadId] = task
@@ -103,9 +109,11 @@ class DownloadTaskManager {
    }

    private fun removeTask(id: Int) {
-        taskDatabase.execSQL("DELETE FROM tasks WHERE id = ?", arrayOf(id))
-        cachedTasks.remove(id)
-        pendingTasks.remove(id)
+        executor.execute {
+            taskDatabase.execSQL("DELETE FROM tasks WHERE id = ?", arrayOf(id))
+            cachedTasks.remove(id)
+            pendingTasks.remove(id)
+        }
    }

    fun removeTask(task: DownloadObject) {
@@ -174,8 +182,10 @@ class DownloadTaskManager {
    }

    fun removeAllTasks() {
-        taskDatabase.execSQL("DELETE FROM tasks")
-        cachedTasks.clear()
-        pendingTasks.clear()
+        executor.execute {
+            taskDatabase.execSQL("DELETE FROM tasks")
+            cachedTasks.clear()
+            pendingTasks.clear()
+        }
    }
 }
--- a/core/src/main/assets/lang/en_US.json
+++ b/core/src/main/assets/lang/en_US.json
@@ -97,6 +97,9 @@
            },
            "hide_chat_feed": {
                "name": "Hide from Chat feed"
+            },
+            "pin_conversation": {
+                "name": "Pin Conversation"
            }
        }
    },
@@ -696,9 +699,18 @@
    },

    "download_processor": {
+        "attachment_type": {
+            "snap": "Snap",
+            "sticker": "Sticker",
+            "external_media": "External Media",
+            "note": "Note",
+            "original_story": "Original Story"
+        },
+        "select_attachments_title": "Select attachments to download",
        "download_started_toast": "Download started",
        "unsupported_content_type_toast": "Unsupported content type!",
        "failed_no_longer_available_toast": "Media no longer available",
+        "no_attachments_toast": "No attachments found!",
        "already_queued_toast": "Media already in queue!",
        "already_downloaded_toast": "Media already downloaded!",
        "saved_toast": "Saved to {path}",
--- a/core/src/main/kotlin/me/rhunk/snapenhance/core/download/DownloadManagerClient.kt
+++ b/core/src/main/kotlin/me/rhunk/snapenhance/core/download/DownloadManagerClient.kt
@@ -10,7 +10,7 @@ import me.rhunk.snapenhance.core.download.data.DownloadMetadata
 import me.rhunk.snapenhance.core.download.data.DownloadRequest
 import me.rhunk.snapenhance.core.download.data.InputMedia
 import me.rhunk.snapenhance.core.download.data.MediaEncryptionKeyPair
-import me.rhunk.snapenhance.data.ContentType
+import me.rhunk.snapenhance.features.impl.downloader.decoder.AttachmentType

 class DownloadManagerClient (
    private val context: ModContext,
@@ -50,7 +50,7 @@ class DownloadManagerClient (
        mediaData: String,
        mediaType: DownloadMediaType,
        encryption: MediaEncryptionKeyPair? = null,
-        messageContentType: ContentType? = null
+        attachmentType: AttachmentType? = null
    ) {
        enqueueDownloadRequest(
            DownloadRequest(
@@ -59,7 +59,7 @@ class DownloadManagerClient (
                        content = mediaData,
                        type = mediaType,
                        encryption = encryption,
-                        messageContentType = messageContentType?.name
+                        attachmentType = attachmentType?.name
                    )
                )
            )
--- a/core/src/main/kotlin/me/rhunk/snapenhance/core/download/data/DownloadRequest.kt
+++ b/core/src/main/kotlin/me/rhunk/snapenhance/core/download/data/DownloadRequest.kt
@@ -6,7 +6,7 @@ data class InputMedia(
    val content: String,
    val type: DownloadMediaType,
    val encryption: MediaEncryptionKeyPair? = null,
-    val messageContentType: String? = null,
+    val attachmentType: String? = null,
    val isOverlay: Boolean = false,
 )

--- a/core/src/main/kotlin/me/rhunk/snapenhance/core/download/data/MediaEncryptionKeyPair.kt
+++ b/core/src/main/kotlin/me/rhunk/snapenhance/core/download/data/MediaEncryptionKeyPair.kt
@@ -6,7 +6,7 @@ import me.rhunk.snapenhance.data.wrapper.impl.media.EncryptionWrapper
 import kotlin.io.encoding.Base64
 import kotlin.io.encoding.ExperimentalEncodingApi

-// key and iv are base64 encoded
+// key and iv are base64 encoded into url safe strings
 data class MediaEncryptionKeyPair(
    val key: String,
    val iv: String
--- a/core/src/main/kotlin/me/rhunk/snapenhance/core/util/protobuf/ProtoEditor.kt
+++ b/core/src/main/kotlin/me/rhunk/snapenhance/core/util/protobuf/ProtoEditor.kt
@@ -14,7 +14,7 @@ class EditorContext(
    }
    fun addVarInt(id: Int, value: Int) = addVarInt(id, value.toLong())
    fun addVarInt(id: Int, value: Long) = addWire(Wire(id, WireType.VARINT, value))
-    fun addBuffer(id: Int, value: ByteArray) = addWire(Wire(id, WireType.LENGTH_DELIMITED, value))
+    fun addBuffer(id: Int, value: ByteArray) = addWire(Wire(id, WireType.CHUNK, value))
    fun add(id: Int, content: ProtoWriter.() -> Unit) = addBuffer(id, ProtoWriter().apply(content).toByteArray())
    fun addString(id: Int, value: String) = addBuffer(id, value.toByteArray())
    fun addFixed64(id: Int, value: Long) = addWire(Wire(id, WireType.FIXED64, value))
@@ -48,7 +48,7 @@ class ProtoEditor(
                    wires.getOrPut(wireId) { mutableListOf() }.add(value)
                    return@forEach
                }
-                wires[wireId]!!.add(Wire(wireId, WireType.LENGTH_DELIMITED, writeAtPath(path, currentIndex + 1, childReader, wireToWriteCallback)))
+                wires[wireId]!!.add(Wire(wireId, WireType.CHUNK, writeAtPath(path, currentIndex + 1, childReader, wireToWriteCallback)))
                return@forEach
            }
            wires[wireId]!!.add(value)
--- a/core/src/main/kotlin/me/rhunk/snapenhance/core/util/protobuf/ProtoReader.kt
+++ b/core/src/main/kotlin/me/rhunk/snapenhance/core/util/protobuf/ProtoReader.kt
@@ -1,5 +1,7 @@
 package me.rhunk.snapenhance.core.util.protobuf

+import java.util.UUID
+
 data class Wire(val id: Int, val type: WireType, val value: Any)

 class ProtoReader(private val buffer: ByteArray) {
@@ -43,7 +45,7 @@ class ProtoReader(private val buffer: ByteArray) {
                        }
                        bytes
                    }
-                    WireType.LENGTH_DELIMITED -> {
+                    WireType.CHUNK -> {
                        val length = readVarInt().toInt()
                        val bytes = ByteArray(length)
                        for (i in 0 until length) {
@@ -135,7 +137,7 @@ class ProtoReader(private val buffer: ByteArray) {
    fun eachBuffer(reader: (Int, ByteArray) -> Unit) {
        values.forEach { (id, wires) ->
            wires.forEach { wire ->
-                if (wire.type == WireType.LENGTH_DELIMITED) {
+                if (wire.type == WireType.CHUNK) {
                    reader(id, wire.value as ByteArray)
                }
            }
@@ -172,4 +174,82 @@ class ProtoReader(private val buffer: ByteArray) {
        }
        return value
    }
+
+    private fun prettyPrint(tabSize: Int): String {
+        val tabLine = "    ".repeat(tabSize)
+        val stringBuilder = StringBuilder()
+        values.forEach { (id, wires) ->
+            wires.forEach { wire ->
+                stringBuilder.append(tabLine)
+                stringBuilder.append("$id <${wire.type.name.lowercase()}> = ")
+                when (wire.type) {
+                    WireType.VARINT -> stringBuilder.append("${wire.value}\n")
+                    WireType.FIXED64, WireType.FIXED32 -> {
+                        //print as double, int, floating point
+                        val doubleValue = run {
+                            val bytes = wire.value as ByteArray
+                            var value = 0L
+                            for (i in bytes.indices) {
+                                value = value or ((bytes[i].toLong() and 0xFF) shl (i * 8))
+                            }
+                            value
+                        }.let {
+                            if (wire.type == WireType.FIXED32) {
+                                it.toInt()
+                            } else {
+                                it
+                            }
+                        }
+
+                        stringBuilder.append("$doubleValue/${doubleValue.toDouble().toBits().toString(16)}\n")
+                    }
+                    WireType.CHUNK -> {
+                        fun printArray() {
+                            stringBuilder.append("\n")
+                            stringBuilder.append("$tabLine    ")
+                            stringBuilder.append((wire.value as ByteArray).joinToString(" ") { byte -> "%02x".format(byte) })
+                            stringBuilder.append("\n")
+                        }
+                        runCatching {
+                            val array = (wire.value as ByteArray)
+                            if (array.isEmpty()) {
+                                stringBuilder.append("empty\n")
+                                return@runCatching
+                            }
+                            //auto detect ascii strings
+                            if (array.all { it in 0x20..0x7E }) {
+                                stringBuilder.append("string: ${array.toString(Charsets.UTF_8)}\n")
+                                return@runCatching
+                            }
+
+                            // auto detect uuids
+                            if (array.size == 16) {
+                                val longs = LongArray(2)
+                                for (i in 0..7) {
+                                    longs[0] = longs[0] or ((array[i].toLong() and 0xFF) shl (i * 8))
+                                }
+                                for (i in 8..15) {
+                                    longs[1] = longs[1] or ((array[i].toLong() and 0xFF) shl ((i - 8) * 8))
+                                }
+                                stringBuilder.append("uuid: ${UUID(longs[0], longs[1])}\n")
+                                return@runCatching
+                            }
+
+                            ProtoReader(array).prettyPrint(tabSize + 1).takeIf { it.isNotEmpty() }?.let {
+                                stringBuilder.append("message:\n")
+                                stringBuilder.append(it)
+                            } ?: printArray()
+                        }.onFailure {
+                            printArray()
+                        }
+                    }
+                    else -> stringBuilder.append("unknown\n")
+                }
+            }
+        }
+
+        return stringBuilder.toString()
+    }
+
+    override fun toString() = prettyPrint(0)
 }
--- a/core/src/main/kotlin/me/rhunk/snapenhance/core/util/protobuf/ProtoWriter.kt
+++ b/core/src/main/kotlin/me/rhunk/snapenhance/core/util/protobuf/ProtoWriter.kt
@@ -24,7 +24,7 @@ class ProtoWriter {
    }

    fun addBuffer(id: Int, value: ByteArray) {
-        writeVarInt(id shl 3 or WireType.LENGTH_DELIMITED.value)
+        writeVarInt(id shl 3 or WireType.CHUNK.value)
        writeVarInt(value.size)
        stream.write(value)
    }
@@ -98,7 +98,7 @@ class ProtoWriter {
                    is ByteArray -> stream.write(wire.value)
                }
            }
-            WireType.LENGTH_DELIMITED -> {
+            WireType.CHUNK -> {
                val value = wire.value as ByteArray
                writeVarInt(value.size)
                stream.write(value)
--- a/core/src/main/kotlin/me/rhunk/snapenhance/core/util/protobuf/WireType.kt
+++ b/core/src/main/kotlin/me/rhunk/snapenhance/core/util/protobuf/WireType.kt
@@ -3,7 +3,7 @@ package me.rhunk.snapenhance.core.util.protobuf;
 enum class WireType(val value: Int) {
    VARINT(0),
    FIXED64(1),
-    LENGTH_DELIMITED(2),
+    CHUNK(2),
    START_GROUP(3),
    END_GROUP(4),
    FIXED32(5);
--- a/core/src/main/kotlin/me/rhunk/snapenhance/core/util/snap/EncryptionHelper.kt
+++ b/core/src/main/kotlin/me/rhunk/snapenhance/core/util/snap/EncryptionHelper.kt
@@ -1,15 +1,18 @@
 package me.rhunk.snapenhance.core.util.snap

 import me.rhunk.snapenhance.Constants
+import me.rhunk.snapenhance.core.download.data.MediaEncryptionKeyPair
 import me.rhunk.snapenhance.core.util.protobuf.ProtoReader
 import me.rhunk.snapenhance.data.ContentType
 import java.io.InputStream
-import java.util.Base64
 import javax.crypto.Cipher
 import javax.crypto.CipherInputStream
 import javax.crypto.spec.IvParameterSpec
 import javax.crypto.spec.SecretKeySpec
+import kotlin.io.encoding.Base64
+import kotlin.io.encoding.ExperimentalEncodingApi

+@OptIn(ExperimentalEncodingApi::class)
 object EncryptionHelper {
    fun getEncryptionKeys(contentType: ContentType, messageProto: ProtoReader, isArroyo: Boolean): Pair<ByteArray, ByteArray>? {
        val mediaEncryptionInfo = MediaDownloaderHelper.getMessageMediaEncryptionInfo(
@@ -28,9 +31,8 @@ object EncryptionHelper {
        var iv: ByteArray = encryptionProto.getByteArray(2)!!

        if (encryptionProtoIndex == Constants.ENCRYPTION_PROTO_INDEX_V2) {
-            val decoder = Base64.getMimeDecoder()
-            key = decoder.decode(key)
-            iv = decoder.decode(iv)
+            key = Base64.UrlSafe.decode(key)
+            iv = Base64.UrlSafe.decode(iv)
        }

        return Pair(key, iv)
@@ -50,4 +52,22 @@ object EncryptionHelper {
            return CipherInputStream(inputStream, cipher)
        }
    }
+
+    fun decryptInputStream(
+        inputStream: InputStream,
+        mediaEncryptionKeyPair: MediaEncryptionKeyPair?
+    ): InputStream {
+        if (mediaEncryptionKeyPair == null) {
+            return inputStream
+        }
+
+        Cipher.getInstance("AES/CBC/PKCS5Padding").apply {
+            init(Cipher.DECRYPT_MODE,
+                SecretKeySpec(Base64.UrlSafe.decode(mediaEncryptionKeyPair.key), "AES"),
+                IvParameterSpec(Base64.UrlSafe.decode(mediaEncryptionKeyPair.iv))
+            )
+        }.let { cipher ->
+            return CipherInputStream(inputStream, cipher)
+        }
+    }
 }
--- a/core/src/main/kotlin/me/rhunk/snapenhance/features/impl/downloader/MediaDownloader.kt
+++ b/core/src/main/kotlin/me/rhunk/snapenhance/features/impl/downloader/MediaDownloader.kt
@@ -6,7 +6,6 @@ import android.graphics.BitmapFactory
 import android.net.Uri
 import android.view.Gravity
 import android.view.ViewGroup.MarginLayoutParams
-import android.view.Window
 import android.widget.ImageView
 import android.widget.LinearLayout
 import android.widget.ProgressBar
@@ -15,6 +14,7 @@ import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.async
 import kotlinx.coroutines.runBlocking
 import me.rhunk.snapenhance.bridge.DownloadCallback
+import me.rhunk.snapenhance.core.database.objects.ConversationMessage
 import me.rhunk.snapenhance.core.database.objects.FriendInfo
 import me.rhunk.snapenhance.core.download.DownloadManagerClient
 import me.rhunk.snapenhance.core.download.data.DownloadMediaType
@@ -31,7 +31,6 @@ import me.rhunk.snapenhance.core.util.snap.BitmojiSelfie
 import me.rhunk.snapenhance.core.util.snap.EncryptionHelper
 import me.rhunk.snapenhance.core.util.snap.MediaDownloaderHelper
 import me.rhunk.snapenhance.core.util.snap.PreviewUtils
-import me.rhunk.snapenhance.data.ContentType
 import me.rhunk.snapenhance.data.FileType
 import me.rhunk.snapenhance.data.wrapper.impl.media.MediaInfo
 import me.rhunk.snapenhance.data.wrapper.impl.media.dash.LongformVideoPlaylistItem
@@ -41,6 +40,8 @@ import me.rhunk.snapenhance.data.wrapper.impl.media.opera.ParamMap
 import me.rhunk.snapenhance.features.FeatureLoadParams
 import me.rhunk.snapenhance.features.MessagingRuleFeature
 import me.rhunk.snapenhance.features.impl.Messaging
+import me.rhunk.snapenhance.features.impl.downloader.decoder.DecodedAttachment
+import me.rhunk.snapenhance.features.impl.downloader.decoder.MessageDecoder
 import me.rhunk.snapenhance.features.impl.spying.MessageLogger
 import me.rhunk.snapenhance.hook.HookAdapter
 import me.rhunk.snapenhance.hook.HookStage
@@ -69,6 +70,10 @@ class MediaDownloader : MessagingRuleFeature("MediaDownloader", MessagingRuleTyp
    private var lastSeenMediaInfoMap: MutableMap<SplitMediaAssetType, MediaInfo>? = null
    private var lastSeenMapParams: ParamMap? = null

+    private val translations by lazy {
+        context.translation.getCategory("download_processor")
+    }
+
    private fun provideDownloadManagerClient(
        mediaIdentifier: String,
        mediaAuthor: String,
@@ -81,7 +86,7 @@ class MediaDownloader : MessagingRuleFeature("MediaDownloader", MessagingRuleTyp

        val downloadLogging by context.config.downloader.logging
        if (downloadLogging.contains("started")) {
-            context.shortToast(context.translation["download_processor.download_started_toast"])
+            context.shortToast(translations["download_started_toast"])
        }

        val outputPath = createNewFilePath(generatedHash, downloadSource, mediaAuthor)
@@ -101,7 +106,7 @@ class MediaDownloader : MessagingRuleFeature("MediaDownloader", MessagingRuleTyp
                override fun onSuccess(outputFile: String) {
                    if (!downloadLogging.contains("success")) return
                    context.log.verbose("onSuccess: outputFile=$outputFile")
-                    context.shortToast(context.translation.format("download_processor.saved_toast", "path" to outputFile.split("/").takeLast(2).joinToString("/")))
+                    context.shortToast(translations.format("saved_toast", "path" to outputFile.split("/").takeLast(2).joinToString("/")))
                }

                override fun onProgress(message: String) {
@@ -483,6 +488,34 @@ class MediaDownloader : MessagingRuleFeature("MediaDownloader", MessagingRuleTyp
        }
    }

+    private fun downloadMessageAttachments(
+        friendInfo: FriendInfo,
+        message: ConversationMessage,
+        authorName: String,
+        attachments: List<DecodedAttachment>
+    ) {
+        //TODO: stickers
+        attachments.forEach { attachment ->
+            runCatching {
+                provideDownloadManagerClient(
+                    mediaIdentifier = "${message.clientConversationId}${message.senderId}${message.serverMessageId}${attachment.attachmentInfo?.encryption?.iv}",
+                    downloadSource = MediaDownloadSource.CHAT_MEDIA,
+                    mediaAuthor = authorName,
+                    friendInfo = friendInfo
+                ).downloadSingleMedia(
+                    mediaData = attachment.mediaKey!!,
+                    mediaType = DownloadMediaType.PROTO_MEDIA,
+                    encryption = attachment.attachmentInfo?.encryption,
+                    attachmentType = attachment.type
+                )
+            }.onFailure {
+                context.longToast(translations["failed_generic_toast"])
+                context.log.error("Failed to download", it)
+            }
+        }
+    }
+
+
    @SuppressLint("SetTextI18n")
    @OptIn(ExperimentalCoroutinesApi::class)
    fun downloadMessageId(messageId: Long, isPreview: Boolean = false) {
@@ -494,15 +527,10 @@ class MediaDownloader : MessagingRuleFeature("MediaDownloader", MessagingRuleTyp
        val authorName = friendInfo.usernameForSorting!!

        var messageContent = message.messageContent!!
-        var isArroyoMessage = true
-        var deletedMediaReference: ByteArray? = null
-
-        //check if the messageId
-        var contentType: ContentType = ContentType.fromId(message.contentType)
+        var customMediaReferences = mutableListOf<String>()

        if (messageLogger.isMessageRemoved(message.clientConversationId!!, message.serverMessageId.toLong())) {
            val messageObject = messageLogger.getMessageObject(message.clientConversationId!!, message.serverMessageId.toLong()) ?: throw Exception("Message not found in database")
-            isArroyoMessage = false
            val messageContentObject = messageObject.getAsJsonObject("mMessageContent")

            messageContent = messageContentObject
@@ -510,137 +538,138 @@ class MediaDownloader : MessagingRuleFeature("MediaDownloader", MessagingRuleTyp
                .map { it.asByte }
                .toByteArray()

-            contentType = ContentType.valueOf(messageContentObject
-                .getAsJsonPrimitive("mContentType").asString
-            )
-
-            deletedMediaReference = messageContentObject.getAsJsonArray("mRemoteMediaReferences")
+            customMediaReferences = messageContentObject
+                .getAsJsonArray("mRemoteMediaReferences")
                .map { it.asJsonObject.getAsJsonArray("mMediaReferences") }
-                .flatten().let {  reference ->
-                    if (reference.isEmpty()) return@let null
-                    reference[0].asJsonObject.getAsJsonArray("mContentObject").map { it.asByte }.toByteArray()
+                .flatten().map { reference ->
+                    Base64.UrlSafe.encode(
+                        reference.asJsonObject.getAsJsonArray("mContentObject").map { it.asByte }.toByteArray()
+                    )
                }
-        }
-
-        val translations = context.translation.getCategory("download_processor")
-
-        if (contentType != ContentType.NOTE &&
-            contentType != ContentType.SNAP &&
-            contentType != ContentType.EXTERNAL_MEDIA) {
-            context.shortToast(translations["unsupported_content_type_toast"])
-            return
+                .toMutableList()
        }

        val messageReader = ProtoReader(messageContent)
-        val urlProto: ByteArray? = if (isArroyoMessage) {
-            var finalProto: ByteArray? = null
-            messageReader.eachBuffer(4, 5) {
-                finalProto = getByteArray(1, 3)
-            }
-            finalProto
-        } else deletedMediaReference
+        val decodedAttachments = MessageDecoder.decode(
+            protoReader = messageReader,
+            customMediaReferences = customMediaReferences.takeIf { it.isNotEmpty() }
+        )

-        if (urlProto == null) {
-            context.shortToast(translations["unsupported_content_type_toast"])
+        if (decodedAttachments.isEmpty()) {
+            context.shortToast(translations["no_attachments_toast"])
            return
        }

-        runCatching {
-            if (!isPreview) {
-                val encryptionKeys = EncryptionHelper.getEncryptionKeys(contentType, messageReader, isArroyo = isArroyoMessage)
-                provideDownloadManagerClient(
-                    mediaIdentifier = "${message.clientConversationId}${message.senderId}${message.serverMessageId}",
-                    downloadSource = MediaDownloadSource.CHAT_MEDIA,
-                    mediaAuthor = authorName,
-                    friendInfo = friendInfo
-                ).downloadSingleMedia(
-                    mediaData = Base64.UrlSafe.encode(urlProto),
-                    mediaType = DownloadMediaType.PROTO_MEDIA,
-                    encryption = encryptionKeys?.toKeyPair(),
-                    messageContentType = contentType
+        if (!isPreview) {
+            if (decodedAttachments.size == 1) {
+                downloadMessageAttachments(friendInfo, message, authorName,
+                    listOf(decodedAttachments.first())
                )
                return
            }

-            if (EncryptionHelper.getEncryptionKeys(contentType, messageReader, isArroyo = isArroyoMessage) == null) {
-                context.shortToast(translations["failed_no_longer_available_toast"])
-                return
+            runOnUiThread {
+                ViewAppearanceHelper.newAlertDialogBuilder(context.mainActivity).apply {
+                    val selectedAttachments = mutableListOf<Int>().apply {
+                        addAll(decodedAttachments.indices)
+                    }
+                    setMultiChoiceItems(
+                        decodedAttachments.mapIndexed { index, decodedAttachment ->
+                            "${index + 1}: ${translations["attachment_type.${decodedAttachment.type.key}"]} ${decodedAttachment.attachmentInfo?.resolution?.let { "(${it.first}x${it.second})" } ?: ""}"
+                        }.toTypedArray(),
+                        decodedAttachments.map { true }.toBooleanArray()
+                    ) { _, which, isChecked ->
+                        if (isChecked) {
+                            selectedAttachments.add(which)
+                        } else if (selectedAttachments.contains(which)) {
+                            selectedAttachments.remove(which)
+                        }
+                    }
+                    setTitle(translations["select_attachments_title"])
+                    setNegativeButton(this@MediaDownloader.context.translation["button.cancel"]) { dialog, _ -> dialog.dismiss() }
+                    setPositiveButton(this@MediaDownloader.context.translation["button.download"]) { _, _ ->
+                        downloadMessageAttachments(friendInfo, message, authorName, selectedAttachments.map { decodedAttachments[it] })
+                    }
+                }.show()
            }

-            runBlocking {
-                val previewCoroutine = async {
-                    val downloadedMediaList = MediaDownloaderHelper.downloadMediaFromReference(urlProto) {
-                        EncryptionHelper.decryptInputStream(it, contentType, messageReader, isArroyo = isArroyoMessage)
-                    }
+            return
+        }

-                    val originalMedia = downloadedMediaList[SplitMediaAssetType.ORIGINAL] ?: return@async null
-                    val overlay = downloadedMediaList[SplitMediaAssetType.OVERLAY]
+        runBlocking {
+            val firstAttachment = decodedAttachments.first()

-                    var bitmap: Bitmap? = PreviewUtils.createPreview(originalMedia, isVideo = FileType.fromByteArray(originalMedia).isVideo)
-
-                    if (bitmap == null) {
-                        context.shortToast(translations["failed_to_create_preview_toast"])
-                        return@async null
-                    }
-
-                    overlay?.also {
-                        bitmap = PreviewUtils.mergeBitmapOverlay(bitmap!!, BitmapFactory.decodeByteArray(it, 0, it.size))
-                    }
-
-                    bitmap
+            val previewCoroutine = async {
+                val downloadedMediaList = MediaDownloaderHelper.downloadMediaFromReference(Base64.decode(firstAttachment.mediaKey!!)) {
+                    EncryptionHelper.decryptInputStream(
+                        it,
+                        decodedAttachments.first().attachmentInfo?.encryption
+                    )
                }

-                with(ViewAppearanceHelper.newAlertDialogBuilder(context.mainActivity)) {
-                    val viewGroup = LinearLayout(context).apply {
-                        layoutParams = MarginLayoutParams(MarginLayoutParams.MATCH_PARENT, MarginLayoutParams.MATCH_PARENT)
-                        gravity = Gravity.CENTER_HORIZONTAL or Gravity.CENTER_VERTICAL
-                        addView(ProgressBar(context).apply {
-                            isIndeterminate = true
+                val originalMedia = downloadedMediaList[SplitMediaAssetType.ORIGINAL] ?: return@async null
+                val overlay = downloadedMediaList[SplitMediaAssetType.OVERLAY]
+
+                var bitmap: Bitmap? = PreviewUtils.createPreview(originalMedia, isVideo = FileType.fromByteArray(originalMedia).isVideo)
+
+                if (bitmap == null) {
+                    context.shortToast(translations["failed_to_create_preview_toast"])
+                    return@async null
+                }
+
+                overlay?.also {
+                    bitmap = PreviewUtils.mergeBitmapOverlay(bitmap!!, BitmapFactory.decodeByteArray(it, 0, it.size))
+                }
+
+                bitmap
+            }
+
+            with(ViewAppearanceHelper.newAlertDialogBuilder(context.mainActivity)) {
+                val viewGroup = LinearLayout(context).apply {
+                    layoutParams = MarginLayoutParams(MarginLayoutParams.MATCH_PARENT, MarginLayoutParams.MATCH_PARENT)
+                    gravity = Gravity.CENTER_HORIZONTAL or Gravity.CENTER_VERTICAL
+                    addView(ProgressBar(context).apply {
+                        isIndeterminate = true
+                    })
+                }
+
+                setOnDismissListener {
+                    previewCoroutine.cancel()
+                }
+
+                previewCoroutine.invokeOnCompletion { cause ->
+                    runOnUiThread {
+                        viewGroup.removeAllViews()
+                        if (cause != null) {
+                            viewGroup.addView(TextView(context).apply {
+                                text = translations["failed_to_create_preview_toast"] + "\n" + cause.message
+                                setPadding(30, 30, 30, 30)
+                            })
+                            return@runOnUiThread
+                        }
+
+                        viewGroup.addView(ImageView(context).apply {
+                            setImageBitmap(previewCoroutine.getCompleted())
+                            layoutParams = LinearLayout.LayoutParams(
+                                LinearLayout.LayoutParams.MATCH_PARENT,
+                                LinearLayout.LayoutParams.MATCH_PARENT
+                            )
+                            adjustViewBounds = true
                        })
                    }
+                }

-                    setOnDismissListener {
-                        previewCoroutine.cancel()
-                    }
-
-                    previewCoroutine.invokeOnCompletion { cause ->
-                        runOnUiThread {
-                            viewGroup.removeAllViews()
-                            if (cause != null) {
-                                viewGroup.addView(TextView(context).apply {
-                                    text = translations["failed_to_create_preview_toast"] + "\n" + cause.message
-                                    setPadding(30, 30, 30, 30)
-                                })
-                                return@runOnUiThread
-                            }
-
-                            viewGroup.addView(ImageView(context).apply {
-                                setImageBitmap(previewCoroutine.getCompleted())
-                                layoutParams = LinearLayout.LayoutParams(
-                                    LinearLayout.LayoutParams.MATCH_PARENT,
-                                    LinearLayout.LayoutParams.MATCH_PARENT
-                                )
-                                adjustViewBounds = true
-                            })
-                        }
-                    }
-
-                    runOnUiThread {
-                        show().apply {
-                            setContentView(viewGroup)
-                            requestWindowFeature(Window.FEATURE_NO_TITLE)
-                            window?.setLayout(
-                                context.resources.displayMetrics.widthPixels,
-                                context.resources.displayMetrics.heightPixels
-                            )
-                        }
-                        previewCoroutine.start()
+                runOnUiThread {
+                    show().apply {
+                        setContentView(viewGroup)
+                        window?.setLayout(
+                            context.resources.displayMetrics.widthPixels,
+                            context.resources.displayMetrics.heightPixels
+                        )
                    }
+                    previewCoroutine.start()
                }
            }
-        }.onFailure {
-            context.longToast(translations["failed_generic_toast"])
-            context.log.error("Failed to download message", it)
        }
    }

--- a/core/src/main/kotlin/me/rhunk/snapenhance/features/impl/downloader/decoder/AttachmentInfo.kt
+++ b/core/src/main/kotlin/me/rhunk/snapenhance/features/impl/downloader/decoder/AttachmentInfo.kt
@@ -0,0 +1,15 @@
+package me.rhunk.snapenhance.features.impl.downloader.decoder
+
+import me.rhunk.snapenhance.core.download.data.MediaEncryptionKeyPair
+
+data class BitmojiSticker(
+    val reference: String,
+) : AttachmentInfo()
+
+open class AttachmentInfo(
+    val encryption: MediaEncryptionKeyPair? = null,
+    val resolution: Pair<Int, Int>? = null,
+    val duration: Long? = null
+) {
+    override fun toString() = "AttachmentInfo(encryption=$encryption, resolution=$resolution, duration=$duration)"
+}
--- a/core/src/main/kotlin/me/rhunk/snapenhance/features/impl/downloader/decoder/AttachmentType.kt
+++ b/core/src/main/kotlin/me/rhunk/snapenhance/features/impl/downloader/decoder/AttachmentType.kt
@@ -0,0 +1,11 @@
+package me.rhunk.snapenhance.features.impl.downloader.decoder
+
+enum class AttachmentType(
+    val key: String,
+) {
+    SNAP("snap"),
+    STICKER("sticker"),
+    EXTERNAL_MEDIA("external_media"),
+    NOTE("note"),
+    ORIGINAL_STORY("original_story"),
+}
--- a/core/src/main/kotlin/me/rhunk/snapenhance/features/impl/downloader/decoder/MessageDecoder.kt
+++ b/core/src/main/kotlin/me/rhunk/snapenhance/features/impl/downloader/decoder/MessageDecoder.kt
@@ -0,0 +1,144 @@
+package me.rhunk.snapenhance.features.impl.downloader.decoder
+
+import me.rhunk.snapenhance.core.download.data.toKeyPair
+import me.rhunk.snapenhance.core.util.protobuf.ProtoReader
+import kotlin.io.encoding.Base64
+import kotlin.io.encoding.ExperimentalEncodingApi
+
+data class DecodedAttachment(
+    val mediaKey: String?,
+    val type: AttachmentType,
+    val attachmentInfo: AttachmentInfo?
+)
+
+@OptIn(ExperimentalEncodingApi::class)
+object MessageDecoder {
+    private fun decodeAttachment(protoReader: ProtoReader): AttachmentInfo? {
+        val mediaInfo =  protoReader.followPath(1, 1) ?: return null
+
+        return AttachmentInfo(
+            encryption = run {
+                val encryptionProtoIndex = if (mediaInfo.contains(19)) 19 else 4
+                val encryptionProto = mediaInfo.followPath(encryptionProtoIndex) ?: return@run null
+
+                var key = encryptionProto.getByteArray(1) ?: return@run null
+                var iv = encryptionProto.getByteArray(2) ?: return@run null
+
+                if (encryptionProtoIndex == 4) {
+                    key = Base64.decode(encryptionProto.getString(1)?.replace("\n","") ?: return@run null)
+                    iv = Base64.decode(encryptionProto.getString(2)?.replace("\n","") ?: return@run null)
+                }
+
+                Pair(key, iv).toKeyPair()
+            },
+            resolution = mediaInfo.followPath(5)?.let {
+                (it.getVarInt(1)?.toInt() ?: 0) to (it.getVarInt(2)?.toInt() ?: 0)
+            },
+            duration = mediaInfo.getVarInt(15)   // external medias
+                ?: mediaInfo.getVarInt(13) // audio notes
+        )
+    }
+
+    fun decode(
+        protoReader: ProtoReader,
+        customMediaReferences: List<String>? = null // when customReferences is null it means that the message is from arroyo database
+    ): List<DecodedAttachment> {
+        val decodedAttachment = mutableListOf<DecodedAttachment>()
+        val mediaReferences = mutableListOf<String>()
+        customMediaReferences?.let { mediaReferences.addAll(it) }
+        var mediaKeyIndex = 0
+
+        fun decodeMedia(type: AttachmentType, protoReader: ProtoReader) {
+            decodedAttachment.add(
+                DecodedAttachment(
+                    mediaKey = mediaReferences.getOrNull(mediaKeyIndex++),
+                    type = type,
+                    attachmentInfo = decodeAttachment(protoReader) ?: return
+                )
+            )
+        }
+
+        // for snaps, external media, and original story replies
+        fun decodeDirectMedia(type: AttachmentType, protoReader: ProtoReader) {
+            protoReader.followPath(5) { decodeMedia(type,this) }
+        }
+
+        fun decodeSticker(protoReader: ProtoReader) {
+            protoReader.followPath(1) {
+                decodedAttachment.add(
+                    DecodedAttachment(
+                        mediaKey = null,
+                        type = AttachmentType.STICKER,
+                        attachmentInfo = BitmojiSticker(
+                            reference = getString(2) ?: return@followPath
+                        )
+                    )
+                )
+            }
+        }
+
+        // media keys
+        protoReader.eachBuffer(4, 5) {
+            getByteArray(1, 3)?.also { mediaKey ->
+                mediaReferences.add(Base64.UrlSafe.encode(mediaKey))
+            }
+        }
+
+        val mediaReader = customMediaReferences?.let { protoReader } ?: protoReader.followPath(4, 4) ?: return emptyList()
+
+        mediaReader.apply {
+            // external media
+            eachBuffer(3, 3) {
+                decodeDirectMedia(AttachmentType.EXTERNAL_MEDIA, this)
+            }
+
+            // stickers
+            followPath(4) { decodeSticker(this) }
+
+            // shares
+            followPath(5, 24, 2) {
+                decodeDirectMedia(AttachmentType.EXTERNAL_MEDIA, this)
+            }
+
+            // audio notes
+            followPath(6) note@{
+                val audioNote = decodeAttachment(this) ?: return@note
+
+                decodedAttachment.add(
+                    DecodedAttachment(
+                        mediaKey = mediaReferences.getOrNull(mediaKeyIndex++),
+                        type = AttachmentType.NOTE,
+                        attachmentInfo = audioNote
+                    )
+                )
+            }
+
+            // story replies
+            followPath(7) {
+                // original story reply
+                followPath(3) {
+                    decodeDirectMedia(AttachmentType.ORIGINAL_STORY, this)
+                }
+
+                // external medias
+                followPath(12) {
+                    eachBuffer(3) { decodeDirectMedia(AttachmentType.EXTERNAL_MEDIA, this) }
+                }
+
+                // attached sticker
+                followPath(13) { decodeSticker(this) }
+
+                // attached audio note
+                followPath(15) { decodeMedia(AttachmentType.NOTE, this) }
+            }
+
+            // snaps
+            followPath(11) {
+                decodeDirectMedia(AttachmentType.SNAP, this)
+            }
+        }
+
+
+        return decodedAttachment
+    }
+}