fix(core/media_downloader): public stories

Signed-off-by: rhunk <101876869+rhunk@users.noreply.github.com>

common/src/main/kotlin/me/rhunk/snapenhance/common/data/download/DownloadRequest.kt

@@ -16,7 +16,7 @@ data class InputMedia(
     val isOverlay: Boolean = false,
 )
 
-class DownloadRequest(
+data class DownloadRequest(
     val inputMedias: Array<InputMedia>,
     val dashOptions: DashOptions? = null,
     val audioStreamFormat: AudioStreamFormat? = null,

common/src/main/kotlin/me/rhunk/snapenhance/common/data/download/MediaEncryptionKeyPair.kt

@@ -12,12 +12,17 @@ import kotlin.io.encoding.ExperimentalEncodingApi
 // key and iv are base64 encoded into url safe strings
 data class MediaEncryptionKeyPair(
     val key: String,
-    val iv: String
+    val iv: String,
+    val urlSafe: Boolean = true
 ) {
     @OptIn(ExperimentalEncodingApi::class)
     fun decryptInputStream(inputStream: InputStream): InputStream {
         val cipher = Cipher.getInstance("AES/CBC/PKCS5Padding")
-        cipher.init(Cipher.DECRYPT_MODE, SecretKeySpec(Base64.UrlSafe.decode(key), "AES"), IvParameterSpec(Base64.UrlSafe.decode(iv)))
+        cipher.init(
+            Cipher.DECRYPT_MODE,
+            SecretKeySpec(if (urlSafe) Base64.UrlSafe.decode(key) else Base64.Default.decode(key), "AES"),
+            IvParameterSpec(if (urlSafe) Base64.UrlSafe.decode(iv) else Base64.Default.decode(iv))
+        )
         return CipherInputStream(inputStream, cipher)
     }
 }

common/src/main/kotlin/me/rhunk/snapenhance/common/database/impl/StorySnapEntry.kt

@@ -0,0 +1,21 @@
+package me.rhunk.snapenhance.common.database.impl
+
+import android.database.Cursor
+import me.rhunk.snapenhance.common.database.DatabaseObject
+import me.rhunk.snapenhance.common.util.ktx.getStringOrNull
+
+data class StorySnapEntry(
+    var rawSnapId: String? = null,
+    var mediaUrl: String? = null,
+    var mediaKey: String? = null,
+    var mediaIv: String? = null,
+) : DatabaseObject {
+    override fun write(cursor: Cursor) {
+        with(cursor) {
+            rawSnapId = getStringOrNull("rawSnapId")!!
+            mediaUrl = getStringOrNull("mediaUrl")
+            mediaKey = getStringOrNull("mediaKey")?.takeIf { it.isNotEmpty() }
+            mediaIv = getStringOrNull("mediaIv")?.takeIf { it.isNotEmpty() }
+        }
+    }
+}

core/src/main/kotlin/me/rhunk/snapenhance/core/database/DatabaseAccess.kt

@@ -10,6 +10,7 @@ import me.rhunk.snapenhance.common.database.impl.ConversationMessage
 import me.rhunk.snapenhance.common.database.impl.FriendFeedEntry
 import me.rhunk.snapenhance.common.database.impl.FriendInfo
 import me.rhunk.snapenhance.common.database.impl.StoryEntry
+import me.rhunk.snapenhance.common.database.impl.StorySnapEntry
 import me.rhunk.snapenhance.common.database.impl.UserConversationLink
 import me.rhunk.snapenhance.common.util.ktx.getBlobOrNull
 import me.rhunk.snapenhance.common.util.ktx.getIntOrNull
@@ -24,7 +25,8 @@ enum class DatabaseType(
     val fileName: String
 ) {
     MAIN("main.db"),
-    ARROYO("arroyo.db")
+    ARROYO("arroyo.db"),
+    SIMPLE_DB_HELPER("simple_db_helper.db")
 }
 
 class DatabaseAccess(
@@ -495,4 +497,15 @@ class DatabaseAccess(
             }
         }?.close()
     }
+
+    fun getStorySnapEntry(rawSnapId: String): StorySnapEntry? {
+        return useDatabase(DatabaseType.SIMPLE_DB_HELPER)?.performOperation {
+            readDatabaseObject(
+                StorySnapEntry(),
+                "DiscoverStorySnap",
+                "rawSnapId = ?",
+                arrayOf(rawSnapId)
+            )
+        }
+    }
 }

core/src/main/kotlin/me/rhunk/snapenhance/core/features/impl/downloader/MediaDownloader.kt

@@ -207,9 +207,22 @@ class MediaDownloader : MessagingRuleFeature("MediaDownloader", MessagingRuleTyp
         }
     }
 
-    private fun downloadOperaMedia(downloadManagerClient: DownloadManagerClient, mediaInfoMap: Map<SplitMediaAssetType, MediaInfo>) {
+    private fun downloadOperaMedia(downloadManagerClient: DownloadManagerClient, mediaInfoMap: Map<SplitMediaAssetType, MediaInfo>, paramMap: ParamMap) {
         if (mediaInfoMap.isEmpty()) return
 
+        paramMap["SNAP_ID"]?.toString()?.let { snapId ->
+            context.database.getStorySnapEntry(snapId)?.let { storySnapEntry ->
+                downloadManagerClient.downloadSingleMedia(
+                    storySnapEntry.mediaUrl ?: throw Exception("Media URL not found"),
+                    DownloadMediaType.fromUri(Uri.parse(storySnapEntry.mediaUrl)),
+                    (storySnapEntry.mediaKey to storySnapEntry.mediaIv).takeIf { it.first != null && it.second != null }?.let { (key, iv) ->
+                        MediaEncryptionKeyPair(key!!, iv!!, urlSafe = false)
+                    }
+                )
+                return
+            }
+        }
+
         val originalMediaInfo = mediaInfoMap[SplitMediaAssetType.ORIGINAL]!!
         val originalMediaInfoReference = handleLocalReferences(originalMediaInfo.uri)
 
@@ -286,7 +299,7 @@ class MediaDownloader : MessagingRuleFeature("MediaDownloader", MessagingRuleTyp
                 downloadSource = MediaDownloadSource.CHAT_MEDIA,
                 friendInfo = author,
                 forceAllowDuplicate = forceAllowDuplicate
-            ), mediaInfoMap)
+            ), mediaInfoMap, paramMap)
 
             return
         }
@@ -333,40 +346,12 @@ class MediaDownloader : MessagingRuleFeature("MediaDownloader", MessagingRuleTyp
                 downloadSource = MediaDownloadSource.STORY,
                 friendInfo = author,
                 forceAllowDuplicate = forceAllowDuplicate,
-            ), mediaInfoMap)
+            ), mediaInfoMap, paramMap)
             return
         }
 
         val snapSource = paramMap["SNAP_SOURCE"].toString()
 
-        //public stories
-        if ((snapSource == "PUBLIC_USER" || snapSource == "SAVED_STORY") &&
-            (forceDownload || shouldAutoDownload("public_stories"))) {
-
-            val author = (
-                paramMap["USER_ID"]?.let { context.database.getFriendInfo(it.toString())?.mutableUsername } // only for following users
-                ?: paramMap["USERNAME"]?.toString()?.takeIf {
-                    it.contains("value=")
-                }?.substringAfter("value=")?.substringBefore(")")?.substringBefore(",")
-                ?: paramMap["CONTEXT_USER_IDENTITY"]?.toString()?.takeIf {
-                    it.contains("username=")
-                }?.substringAfter("username=")?.substringBefore(",")
-                // fallback display name
-                ?: paramMap["USER_DISPLAY_NAME"]?.toString()?.takeIf { it.isNotEmpty() }
-                ?: paramMap["TIME_STAMP"]?.toString()
-                ?: "unknown"
-            ).sanitizeForPath()
-
-            downloadOperaMedia(provideDownloadManagerClient(
-                mediaIdentifier = paramMap["SNAP_ID"].toString(),
-                mediaAuthor = author,
-                downloadSource = MediaDownloadSource.PUBLIC_STORY,
-                creationTimestamp = paramMap["SNAP_TIMESTAMP"]?.toString()?.toLongOrNull(),
-                forceAllowDuplicate = forceAllowDuplicate,
-            ), mediaInfoMap)
-            return
-        }
-
         //spotlight
         if (snapSource == "SINGLE_SNAP_STORY" && (forceDownload || shouldAutoDownload("spotlight"))) {
             downloadOperaMedia(provideDownloadManagerClient(
@@ -375,7 +360,7 @@ class MediaDownloader : MessagingRuleFeature("MediaDownloader", MessagingRuleTyp
                 mediaAuthor = paramMap["CREATOR_DISPLAY_NAME"].toString(),
                 creationTimestamp = paramMap["SNAP_TIMESTAMP"]?.toString()?.toLongOrNull(),
                 forceAllowDuplicate = forceAllowDuplicate,
-            ), mediaInfoMap)
+            ), mediaInfoMap, paramMap)
             return
         }
 
@@ -481,6 +466,29 @@ class MediaDownloader : MessagingRuleFeature("MediaDownloader", MessagingRuleTyp
                 }.show()
             }
         }
+
+        //public stories
+        val author = (
+                paramMap["USER_ID"]?.let { context.database.getFriendInfo(it.toString())?.mutableUsername } // only for following users
+                    ?: paramMap["USERNAME"]?.toString()?.takeIf {
+                        it.contains("value=")
+                    }?.substringAfter("value=")?.substringBefore(")")?.substringBefore(",")
+                    ?: paramMap["CONTEXT_USER_IDENTITY"]?.toString()?.takeIf {
+                        it.contains("username=")
+                    }?.substringAfter("username=")?.substringBefore(",")
+                    // fallback display name
+                    ?: paramMap["USER_DISPLAY_NAME"]?.toString()?.takeIf { it.isNotEmpty() }
+                    ?: paramMap["TIME_STAMP"]?.toString()
+                    ?: "unknown"
+                ).sanitizeForPath()
+
+        downloadOperaMedia(provideDownloadManagerClient(
+            mediaIdentifier = paramMap["SNAP_ID"].toString(),
+            mediaAuthor = author,
+            downloadSource = MediaDownloadSource.PUBLIC_STORY,
+            creationTimestamp = paramMap["SNAP_TIMESTAMP"]?.toString()?.toLongOrNull(),
+            forceAllowDuplicate = forceAllowDuplicate,
+        ), mediaInfoMap, paramMap)
     }
 
     private fun shouldAutoDownload(keyFilter: String? = null): Boolean {

core/src/main/kotlin/me/rhunk/snapenhance/core/wrapper/impl/media/MediaInfo.kt

@@ -16,7 +16,7 @@ class MediaInfo(obj: Any?) : AbstractWrapper(obj) {
     init {
         instance?.let {
             if (it is List<*>) {
-                if (it.size == 0) {
+                if (it.isEmpty()) {
                     throw RuntimeException("MediaInfo is empty")
                 }
                 instance = it[0]!!