commit 904d9175b69b20156d6f083a4d389149655ab908 parent 5767f9333132cf9b91007546a8ff88334bf481aa Author: rhunk <101876869+rhunk@users.noreply.github.com> Date: Thu, 14 Sep 2023 23:55:06 +0200 feat(download_manager): download multiple attachment - experimental message decoder - fix async calls for downloads database - add debug view for ProtoReader Diffstat:
15 files changed, 483 insertions(+), 161 deletions(-)
diff --git 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( diff --git 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() + } } } \ No newline at end of file diff --git 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}", diff --git 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 ) ) ) diff --git 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, ) diff --git 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 diff --git 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) diff --git 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) } \ No newline at end of file diff --git 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) diff --git 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); diff --git 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) + } + } } diff --git 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) + val previewCoroutine = async { + val downloadedMediaList = MediaDownloaderHelper.downloadMediaFromReference(Base64.decode(firstAttachment.mediaKey!!)) { + EncryptionHelper.decryptInputStream( + it, + decodedAttachments.first().attachmentInfo?.encryption + ) + } - if (bitmap == null) { - context.shortToast(translations["failed_to_create_preview_toast"]) - return@async null - } + val originalMedia = downloadedMediaList[SplitMediaAssetType.ORIGINAL] ?: return@async null + val overlay = downloadedMediaList[SplitMediaAssetType.OVERLAY] - overlay?.also { - bitmap = PreviewUtils.mergeBitmapOverlay(bitmap!!, BitmapFactory.decodeByteArray(it, 0, it.size)) - } + var bitmap: Bitmap? = PreviewUtils.createPreview(originalMedia, isVideo = FileType.fromByteArray(originalMedia).isVideo) - bitmap + if (bitmap == null) { + context.shortToast(translations["failed_to_create_preview_toast"]) + return@async null } - 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 - }) - } + overlay?.also { + bitmap = PreviewUtils.mergeBitmapOverlay(bitmap!!, BitmapFactory.decodeByteArray(it, 0, it.size)) + } - setOnDismissListener { - previewCoroutine.cancel() - } + bitmap + } - 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 - } + 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() + } - viewGroup.addView(ImageView(context).apply { - setImageBitmap(previewCoroutine.getCompleted()) - layoutParams = LinearLayout.LayoutParams( - LinearLayout.LayoutParams.MATCH_PARENT, - LinearLayout.LayoutParams.MATCH_PARENT - ) - adjustViewBounds = true + 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 } - } - runOnUiThread { - show().apply { - setContentView(viewGroup) - requestWindowFeature(Window.FEATURE_NO_TITLE) - window?.setLayout( - context.resources.displayMetrics.widthPixels, - context.resources.displayMetrics.heightPixels + viewGroup.addView(ImageView(context).apply { + setImageBitmap(previewCoroutine.getCompleted()) + layoutParams = LinearLayout.LayoutParams( + LinearLayout.LayoutParams.MATCH_PARENT, + LinearLayout.LayoutParams.MATCH_PARENT ) - } - previewCoroutine.start() + adjustViewBounds = true + }) + } + } + + 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) } } diff --git 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)" +}+ \ No newline at end of file diff --git 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"), +}+ \ No newline at end of file diff --git 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 + } +}+ \ No newline at end of file