From b17a319859dbd83a2fb351b6a3645ae9ca27d100 Mon Sep 17 00:00:00 2001
From: Artemis Tosini <me@artem.ist>
Date: Thu, 4 Apr 2024 07:23:35 +0000
Subject: [PATCH] extract_ktsl2stbin_opus.py: ...so that's what granule pos is

---
 scripts/extract_ktsl2stbin_opus.py | 39 +++++++++++++++++++++++++-----
 1 file changed, 33 insertions(+), 6 deletions(-)

diff --git a/scripts/extract_ktsl2stbin_opus.py b/scripts/extract_ktsl2stbin_opus.py
index 11fc965..a29746e 100755
--- a/scripts/extract_ktsl2stbin_opus.py
+++ b/scripts/extract_ktsl2stbin_opus.py
@@ -5,6 +5,18 @@ import sys
 import pathlib
 
 
+FRAME_SIZES = [
+    480, 960, 1920, 3840, # SILK NB
+    480, 960, 1920, 3840, # SILK MB
+    480, 960, 1920, 3840, # SILK WB
+    480, 960, # Hybrid SWB
+    480, 960, # Hybrid FB
+    240, 480, 960, 1920, # CELT NB
+    240, 480, 960, 1920, # CELT WB
+    240, 480, 960, 1920, # CELT SWB
+    240, 480, 960, 1920, # CELT FB
+]
+
 def crc32ogg(seq):
     crc = 0
     for b in seq:
@@ -13,8 +25,7 @@ def crc32ogg(seq):
             crc = (crc << 1) ^ 0x104C11DB7 if crc & 0x80000000 else crc << 1
     return crc
 
-
-def paginate(sequence: int, is_last: bool, content: bytes):
+def paginate(sequence: int, is_last: bool, granule_pos, content: bytes):
     # Version, flags, position, serial number, sequence number, checksum, segments, segment table
     flags = 2 if sequence == 0 else 0
     if is_last:
@@ -22,7 +33,7 @@ def paginate(sequence: int, is_last: bool, content: bytes):
     page = bytearray(
         b"OggS"
         + struct.pack(
-            "<BBQIIIB", 0, flags, 0, 0xACAB1234, sequence, 0, (len(content) // 255) + 1
+            "<BBQIIIB", 0, flags, granule_pos, 0xACAB1234, sequence, 0, (len(content) // 255) + 1
         )
         + (b"\xff" * (len(content) // 255))
         + bytes([len(content) % 255])
@@ -64,19 +75,35 @@ def write_opus(ktss: bytes, filename: pathlib.Path):
         + channel_mapping
     )
     # Channel mapping is apparently incorrect for 6 channels but I don't care
-    out.write(paginate(0, False, opus_header))
+    out.write(paginate(0, False, 0, opus_header))
 
     comment_header = b"OpusTags\4\0\0\0ktss\0\0\0\0\0\0\0\0"
-    out.write(paginate(1, False, comment_header))
+    out.write(paginate(1, False, 0, comment_header))
 
     # Weird length encoding here, not opus standard
     sequence = 2
     offset = start_offset
+    granule_pos = 0
     while offset < len(ktss):
         (packet_len,) = struct.unpack(">I", ktss[offset : offset + 4])
         is_last = offset + 8 + packet_len >= len(ktss)
+        toc = ktss[offset + 8]
+        if toc & 3 == 0:
+            num_frames = 1
+        elif toc & 3 == 3:
+            num_frames = ktss[offset + 9] & 0x3f
+        else:
+            num_frames = 2
+        
+        granule_len = FRAME_SIZES[toc >> 3] * num_frames
+        if toc & 4 != 0:
+            # We have to divide by 2 if packet is stereo
+            granule_len //= 2
+
+        granule_pos += granule_len
+
         out.write(
-            paginate(sequence, is_last, ktss[offset + 8 : offset + 8 + packet_len])
+            paginate(sequence, is_last, granule_pos, ktss[offset + 8 : offset + 8 + packet_len])
         )
 
         offset += packet_len + 8