From 98e04a81cb933204d436ec1a16a89a5a4448e46a Mon Sep 17 00:00:00 2001
From: Majjo <90888719+NoxRare@users.noreply.github.com>
Date: Mon, 20 Apr 2026 00:29:57 +0200
Subject: [PATCH] fix: capture story_lookup/batch_story_lookup endpoints for
 story media URLs

---
 .../features/impl/global/CompanionServer.kt   | 81 ++++++++++++-------
 1 file changed, 52 insertions(+), 29 deletions(-)

diff --git a/core/src/main/kotlin/me/eternal/purrfectsnap/core/features/impl/global/CompanionServer.kt b/core/src/main/kotlin/me/eternal/purrfectsnap/core/features/impl/global/CompanionServer.kt
index 7c0e76b5..fbc42409 100644
--- a/core/src/main/kotlin/me/eternal/purrfectsnap/core/features/impl/global/CompanionServer.kt
+++ b/core/src/main/kotlin/me/eternal/purrfectsnap/core/features/impl/global/CompanionServer.kt
@@ -185,49 +185,72 @@ class CompanionServer : Feature("CompanionServer") {
 
     // ---- Story capture ----
 
+    private fun storeStory(userId: String, url: String, postedAt: Long, createdAt: Long, keyB64: String, ivB64: String): Boolean {
+        val seen = capturedStoryUrls.getOrPut(userId) { java.util.concurrent.ConcurrentHashMap() }
+        if (seen.putIfAbsent(url, true) != null) return false
+        capturedStories.getOrPut(userId) { java.util.concurrent.CopyOnWriteArrayList() }.add(mapOf(
+            "userId"     to userId,
+            "url"        to url,
+            "postedAt"   to postedAt,
+            "createdAt"  to createdAt,
+            "key"        to keyB64,
+            "iv"         to ivB64,
+            "capturedAt" to System.currentTimeMillis()
+        ))
+        return true
+    }
+
     private fun subscribeToStoryEvents() {
         context.event.subscribe(NetworkApiRequestEvent::class) { event ->
-            if (!event.url.contains("df-mixer-prod") ||
-                (!event.url.endsWith("stories") && !event.url.endsWith("batch_stories"))) return@subscribe
+            val urlPath = event.url.substringBefore("?")
+            if (!urlPath.contains("df-mixer-prod")) return@subscribe
+            val endpoint = urlPath.substringAfterLast("/")
+            if (endpoint !in setOf("stories", "batch_stories", "story_lookup", "batch_story_lookup")) return@subscribe
             event.onSuccess { buffer ->
                 buffer ?: return@onSuccess
+                context.log.verbose("CompanionServer: [$endpoint] response ${buffer.size} bytes", "CompanionServer")
                 runCatching {
-                    ProtoReader(buffer).eachBuffer(3) {
+                    var parsed = 0
+                    // story_lookup / batch_story_lookup: field 1 → repeated story entries
+                    // Each entry: field 4 = userId, field 3 → media block → field 36 → field 1 → key/iv/url
+                    ProtoReader(buffer).eachBuffer(1) {
+                        val userId = getString(4, 1) ?: return@eachBuffer
                         eachBuffer(3) {
+                            followPath(36) {
+                                eachBuffer(1) {
+                                    val url       = getString(2, 2)?.substringBefore("?") ?: return@eachBuffer
+                                    val postedAt  = getVarInt(3) ?: return@eachBuffer
+                                    val createdAt = getVarInt(27) ?: -1L
+                                    val keyB64    = getString(2, 5) ?: return@eachBuffer
+                                    val ivB64     = getString(2, 4) ?: return@eachBuffer
+                                    if (storeStory(userId, url, postedAt, createdAt, keyB64, ivB64)) parsed++
+                                }
+                            }
+                        }
+                    }
+                    // batch_stories fallback: field 3 → field 3 → field 3 → field 36 → field 1
+                    if (parsed == 0) {
+                        ProtoReader(buffer).eachBuffer(3) {
                             eachBuffer(3) {
-                                followPath(36) {
-                                    eachBuffer(1) {
-                                        val userId    = getString(8, 1) ?: return@eachBuffer
-                                        val url       = getString(2, 2)?.substringBefore("?") ?: return@eachBuffer
-                                        val postedAt  = getVarInt(3) ?: return@eachBuffer
-                                        val createdAt = getVarInt(27) ?: -1L
-                                        val keyB64    = getString(2, 5) ?: return@eachBuffer
-                                        val ivB64     = getString(2, 4) ?: return@eachBuffer
-
-                                        // Deduplicate by URL
-                                        val seen = capturedStoryUrls.getOrPut(userId) {
-                                            java.util.concurrent.ConcurrentHashMap()
-                                        }
-                                        if (seen.putIfAbsent(url, true) == null) {
-                                            capturedStories.getOrPut(userId) {
-                                                java.util.concurrent.CopyOnWriteArrayList()
-                                            }.add(mapOf(
-                                                "userId"     to userId,
-                                                "url"        to url,
-                                                "postedAt"   to postedAt,
-                                                "createdAt"  to createdAt,
-                                                "key"        to keyB64,
-                                                "iv"         to ivB64,
-                                                "capturedAt" to System.currentTimeMillis()
-                                            ))
+                                eachBuffer(3) {
+                                    followPath(36) {
+                                        eachBuffer(1) {
+                                            val userId    = getString(8, 1) ?: return@eachBuffer
+                                            val url       = getString(2, 2)?.substringBefore("?") ?: return@eachBuffer
+                                            val postedAt  = getVarInt(3) ?: return@eachBuffer
+                                            val createdAt = getVarInt(27) ?: -1L
+                                            val keyB64    = getString(2, 5) ?: return@eachBuffer
+                                            val ivB64     = getString(2, 4) ?: return@eachBuffer
+                                            if (storeStory(userId, url, postedAt, createdAt, keyB64, ivB64)) parsed++
                                         }
                                     }
                                 }
                             }
                         }
                     }
+                    context.log.verbose("CompanionServer: [$endpoint] parsed $parsed new stories", "CompanionServer")
                 }.onFailure {
-                    context.log.warn("CompanionServer: story proto parse failed: ${it.message}", "CompanionServer")
+                    context.log.warn("CompanionServer: [$endpoint] proto parse failed: ${it.message}", "CompanionServer")
                 }
             }
         }