stremio/apps/ios/VMC/Playback/CencMp4Decryptor.swift

import Foundation

/// In-place CENC/CBCS decryptie van gefragmenteerde MP4 (init + media segments).
enum CencMp4Decryptor {
    struct TrackCrypto {
        var kid: Data
        var perSampleIvSize: Int
        var constantIv: Data?
        var scheme: String // cenc | cbcs | cbc1 | cens
        var cryptByteBlock: Int
        var skipByteBlock: Int
    }

    struct SessionState {
        var byTrackId: [UInt32: TrackCrypto] = [:]
        var fallback: TrackCrypto?
    }

    @discardableResult
    static func process(data: inout Data, keys: ClearKeyKeyMap, state: inout SessionState) -> Bool {
        guard !data.isEmpty, !keys.isEmpty else { return false }
        var touched = false
        if containsBox(data, type: "moov") {
            harvestAndClearInit(&data, keys: keys, state: &state)
            touched = true
        }
        if containsBox(data, type: "moof") {
            decryptFragments(&data, keys: keys, state: state)
            touched = true
        }
        return touched
    }

    // MARK: - Box helpers

    private struct Box {
        let type: String
        let start: Int
        let headerSize: Int
        let size: Int
        var payloadStart: Int { start + headerSize }
        var payloadEnd: Int { start + size }
        var payloadLength: Int { size - headerSize }
    }

    private static func containsBox(_ data: Data, type: String) -> Bool {
        findBoxes(data, type: type, recursive: true).isEmpty == false
    }

    private static func findBoxes(_ data: Data, type: String, recursive: Bool) -> [Box] {
        var out: [Box] = []
        walk(data, from: 0, to: data.count, recursive: recursive) { box in
            if box.type == type { out.append(box) }
            return true
        }
        return out
    }

    private static func walk(
        _ data: Data,
        from: Int,
        to: Int,
        recursive: Bool,
        visit: (Box) -> Bool
    ) {
        var offset = from
        let containers: Set<String> = [
            "moov", "trak", "mdia", "minf", "stbl", "moof", "traf",
            "mvex", "edts", "sinf", "schi", "dinf", "udta", "meta",
        ]
        while offset + 8 <= to {
            guard let box = readBox(data, at: offset, limit: to) else { break }
            let cont = visit(box)
            if !cont { return }
            if recursive, containers.contains(box.type), box.payloadLength > 0 {
                walk(data, from: box.payloadStart, to: box.payloadEnd, recursive: true, visit: visit)
            }
            offset = box.payloadEnd
            if box.size <= 0 { break }
        }
    }

    private static func readBox(_ data: Data, at offset: Int, limit: Int) -> Box? {
        guard offset + 8 <= limit else { return nil }
        var size = Int(readU32(data, offset))
        let type = readFourCC(data, offset + 4)
        var header = 8
        if size == 1 {
            guard offset + 16 <= limit else { return nil }
            let large = readU64(data, offset + 8)
            size = Int(large)
            header = 16
        } else if size == 0 {
            size = limit - offset
        }
        guard size >= header, offset + size <= limit else { return nil }
        return Box(type: type, start: offset, headerSize: header, size: size)
    }

    private static func readU16(_ data: Data, _ o: Int) -> UInt16 {
        UInt16(data[o]) << 8 | UInt16(data[o + 1])
    }

    private static func readU32(_ data: Data, _ o: Int) -> UInt32 {
        UInt32(data[o]) << 24 | UInt32(data[o + 1]) << 16 | UInt32(data[o + 2]) << 8 | UInt32(data[o + 3])
    }

    private static func readU64(_ data: Data, _ o: Int) -> UInt64 {
        (UInt64(readU32(data, o)) << 32) | UInt64(readU32(data, o + 4))
    }

    private static func readFourCC(_ data: Data, _ o: Int) -> String {
        String(bytes: data[o..<(o + 4)], encoding: .ascii) ?? "????"
    }

    private static func writeFourCC(_ data: inout Data, _ o: Int, _ s: String) {
        let chars = Array(s.utf8)
        guard chars.count == 4, o + 4 <= data.count else { return }
        for i in 0..<4 { data[o + i] = chars[i] }
    }

    // MARK: - Init segment

    private static func harvestAndClearInit(_ data: inout Data, keys: ClearKeyKeyMap, state: inout SessionState) {
        // Track id from tkhd, crypto from sinf/tenc/frma/schm inside same trak.
        for trak in findBoxes(data, type: "trak", recursive: true) {
            var trackId: UInt32?
            walk(data, from: trak.payloadStart, to: trak.payloadEnd, recursive: true) { box in
                if box.type == "tkhd", box.payloadLength >= 8 {
                    let version = Int(data[box.payloadStart])
                    let idOff = box.payloadStart + (version == 1 ? 20 : 12)
                    if idOff + 4 <= box.payloadEnd {
                        trackId = readU32(data, idOff)
                    }
                }
                return true
            }
            var crypto: TrackCrypto?
            var originalFormat: String?
            walk(data, from: trak.payloadStart, to: trak.payloadEnd, recursive: true) { box in
                if box.type == "frma", box.payloadLength >= 4 {
                    originalFormat = readFourCC(data, box.payloadStart)
                }
                if box.type == "schm", box.payloadLength >= 8 {
                    // version/flags (4) + scheme_type (4)
                    let scheme = readFourCC(data, box.payloadStart + 4)
                    if crypto == nil {
                        crypto = TrackCrypto(
                            kid: Data(count: 16),
                            perSampleIvSize: 8,
                            constantIv: nil,
                            scheme: scheme,
                            cryptByteBlock: 0,
                            skipByteBlock: 0
                        )
                    } else {
                        crypto?.scheme = scheme
                    }
                }
                if box.type == "tenc", box.payloadLength >= 20 {
                    let version = Int(data[box.payloadStart])
                    let base = box.payloadStart + 4 // na version+flags
                    var o = base
                    o += 1 // reserved
                    var cryptBlocks = 0
                    var skipBlocks = 0
                    if version >= 1 {
                        let pattern = data[o]
                        cryptBlocks = Int((pattern >> 4) & 0x0F)
                        skipBlocks = Int(pattern & 0x0F)
                    }
                    o += 1 // reserved (v0) of pattern (v1)
                    let isProtectedIndex = o
                    let isProtected = Int(data[o]); o += 1
                    let ivSize = Int(data[o]); o += 1
                    let kid = data.subdata(in: o..<(o + 16)); o += 16
                    _ = isProtected
                    var constantIv: Data?
                    if ivSize == 0, o < box.payloadEnd {
                        let constSize = Int(data[o]); o += 1
                        if constSize > 0, o + constSize <= box.payloadEnd {
                            constantIv = data.subdata(in: o..<(o + constSize))
                        }
                    }
                    crypto = TrackCrypto(
                        kid: kid,
                        perSampleIvSize: ivSize,
                        constantIv: constantIv,
                        scheme: crypto?.scheme ?? "cenc",
                        cryptByteBlock: cryptBlocks,
                        skipByteBlock: skipBlocks
                    )
                    data[isProtectedIndex] = 0
                }
                return true
            }

            // Replace encv/enca with original format (same 4 bytes → no size rewrite).
            if let originalFormat {
                walk(data, from: trak.payloadStart, to: trak.payloadEnd, recursive: true) { box in
                    // Sample entries live in stsd; type is at box.type for encv/enca boxes themselves.
                    return true
                }
                // Find encv/enca boxes by scanning stsd children manually.
                for stsd in findBoxes(data, type: "stsd", recursive: true) {
                    guard stsd.start >= trak.start, stsd.payloadEnd <= trak.payloadEnd else { continue }
                    var o = stsd.payloadStart + 4 // version/flags + entry_count follows version/flags (4) then count (4)
                    // FullBox (4) + entry_count (4)
                    o = stsd.payloadStart + 8
                    while o + 8 <= stsd.payloadEnd {
                        guard let entry = readBox(data, at: o, limit: stsd.payloadEnd) else { break }
                        if entry.type == "encv" || entry.type == "enca" {
                            writeFourCC(&data, entry.start + 4, originalFormat)
                        }
                        o = entry.payloadEnd
                    }
                }
            }

            if let crypto {
                if let trackId {
                    state.byTrackId[trackId] = crypto
                }
                state.fallback = crypto
                _ = keys.key(forKid: crypto.kid)
            }
        }
    }

    // MARK: - Media fragments

    private static func decryptFragments(_ data: inout Data, keys: ClearKeyKeyMap, state: SessionState) {
        // Pair each moof with the following mdat.
        var offset = 0
        while offset + 8 <= data.count {
            guard let box = readBox(data, at: offset, limit: data.count) else { break }
            if box.type == "moof" {
                let moof = box
                var mdat: Box?
                var scan = moof.payloadEnd
                while scan + 8 <= data.count {
                    guard let next = readBox(data, at: scan, limit: data.count) else { break }
                    if next.type == "mdat" {
                        mdat = next
                        break
                    }
                    if next.type == "moof" { break }
                    scan = next.payloadEnd
                }
                if let mdat {
                    decryptMoofMdat(moof: moof, mdat: mdat, data: &data, keys: keys, state: state)
                }
            }
            offset = box.payloadEnd
        }
    }

    private static func decryptMoofMdat(
        moof: Box,
        mdat: Box,
        data: inout Data,
        keys: ClearKeyKeyMap,
        state: SessionState
    ) {
        let trafs = findBoxesInRange(data, type: "traf", from: moof.payloadStart, to: moof.payloadEnd)
        var mdatCursor = mdat.payloadStart

        for traf in trafs {
            var trackId: UInt32 = 0
            var defaultSampleSize: UInt32 = 0
            var baseDataOffset: Int64 = Int64(mdat.payloadStart)
            var dataOffsetPresent = false
            var sampleSizes: [Int] = []
            var sencSamples: [SencSample] = []
            var perSampleIvSize = state.fallback?.perSampleIvSize ?? 8

            // tfhd
            for tfhd in findBoxesInRange(data, type: "tfhd", from: traf.payloadStart, to: traf.payloadEnd) {
                let versionFlags = readU32(data, tfhd.payloadStart)
                let flags = versionFlags & 0x00FF_FFFF
                var o = tfhd.payloadStart + 4
                trackId = readU32(data, o); o += 4
                if let crypto = state.byTrackId[trackId] {
                    perSampleIvSize = crypto.perSampleIvSize
                }
                if flags & 0x000001 != 0 { // base-data-offset
                    baseDataOffset = Int64(readU64(data, o)); o += 8
                    dataOffsetPresent = true
                }
                if flags & 0x000002 != 0 { o += 4 } // sample-description-index
                if flags & 0x000008 != 0 {
                    defaultSampleSize = readU32(data, o); o += 4
                }
                _ = versionFlags
            }

            // trun
            for trun in findBoxesInRange(data, type: "trun", from: traf.payloadStart, to: traf.payloadEnd) {
                let versionFlags = readU32(data, trun.payloadStart)
                let version = Int(versionFlags >> 24)
                let flags = versionFlags & 0x00FF_FFFF
                var o = trun.payloadStart + 4
                let sampleCount = Int(readU32(data, o)); o += 4
                var dataOffset: Int32 = 0
                if flags & 0x000001 != 0 {
                    dataOffset = Int32(bitPattern: readU32(data, o)); o += 4
                    dataOffsetPresent = true
                }
                if flags & 0x000004 != 0 { o += 4 } // first-sample-flags
                sampleSizes.removeAll(keepingCapacity: true)
                for _ in 0..<sampleCount {
                    if flags & 0x000100 != 0 { // sample-duration
                        o += 4
                    }
                    var size = Int(defaultSampleSize)
                    if flags & 0x000200 != 0 {
                        size = Int(readU32(data, o)); o += 4
                    }
                    sampleSizes.append(size)
                    if flags & 0x000400 != 0 { o += 4 } // sample-flags
                    if flags & 0x000800 != 0 {
                        o += version == 0 ? 4 : 4 // sample-composition-time-offset
                    }
                }
                if dataOffsetPresent {
                    // data_offset is relative to moof start when base-data-offset-present is NOT set;
                    // when base-data-offset is set, absolute = base + dataOffset.
                    // Common CMAF: no base-data-offset, data_offset relative to moof.
                    let tfhd = findBoxesInRange(data, type: "tfhd", from: traf.payloadStart, to: traf.payloadEnd).first
                    let tfhdFlags: UInt32 = {
                        guard let tfhd else { return 0 }
                        return readU32(data, tfhd.payloadStart) & 0x00FF_FFFF
                    }()
                    if tfhdFlags & 0x000001 != 0 {
                        mdatCursor = Int(baseDataOffset + Int64(dataOffset))
                    } else {
                        mdatCursor = moof.start + Int(dataOffset)
                    }
                }
                _ = version
            }

            // senc
            let crypto = state.byTrackId[trackId] ?? state.fallback
            if let crypto {
                perSampleIvSize = crypto.perSampleIvSize
            }
            for senc in findBoxesInRange(data, type: "senc", from: traf.payloadStart, to: traf.payloadEnd) {
                sencSamples = parseSenc(data, box: senc, perSampleIvSize: perSampleIvSize, constantIv: crypto?.constantIv)
            }

            // saiz/saio fallback when no senc
            if sencSamples.isEmpty {
                sencSamples = parseSaizSaio(
                    data,
                    traf: traf,
                    moof: moof,
                    perSampleIvSize: perSampleIvSize,
                    constantIv: crypto?.constantIv,
                    sampleCount: sampleSizes.count
                )
            }

            guard !sampleSizes.isEmpty, sencSamples.count == sampleSizes.count || sencSamples.isEmpty == false else {
                // Advance cursor even if we cannot decrypt.
                for sz in sampleSizes { mdatCursor += sz }
                continue
            }

            let key: Data? = {
                if let crypto, let k = keys.key(forKid: crypto.kid) { return k }
                return keys.primaryKey
            }()
            guard let key else {
                for sz in sampleSizes { mdatCursor += sz }
                continue
            }

            let scheme = (crypto?.scheme ?? "cenc").lowercased()
            let cryptBlocks = crypto?.cryptByteBlock ?? 0
            let skipBlocks = crypto?.skipByteBlock ?? 0

            for i in 0..<sampleSizes.count {
                let size = sampleSizes[i]
                guard size > 0, mdatCursor + size <= data.count else { break }
                let sampleRange = mdatCursor..<(mdatCursor + size)
                var sample = data.subdata(in: sampleRange)
                let senc = i < sencSamples.count ? sencSamples[i] : nil
                if let senc {
                    if let decrypted = decryptSample(
                        sample,
                        key: key,
                        iv: senc.iv,
                        subsamples: senc.subsamples,
                        scheme: scheme,
                        cryptByteBlock: cryptBlocks,
                        skipByteBlock: skipBlocks
                    ) {
                        sample = decrypted
                        data.replaceSubrange(sampleRange, with: sample)
                    }
                }
                mdatCursor += size
            }
        }
    }

    private struct SencSample {
        var iv: Data
        var subsamples: [(clear: Int, protected: Int)]?
    }

    private static func parseSenc(
        _ data: Data,
        box: Box,
        perSampleIvSize: Int,
        constantIv: Data?
    ) -> [SencSample] {
        let versionFlags = readU32(data, box.payloadStart)
        let flags = versionFlags & 0x00FF_FFFF
        var o = box.payloadStart + 4
        let sampleCount = Int(readU32(data, o)); o += 4
        var out: [SencSample] = []
        out.reserveCapacity(sampleCount)
        let useSubsamples = (flags & 0x2) != 0
        for _ in 0..<sampleCount {
            var iv = Data(count: 16)
            if perSampleIvSize > 0 {
                let n = min(perSampleIvSize, 16)
                guard o + n <= box.payloadEnd else { break }
                iv.replaceSubrange(0..<n, with: data.subdata(in: o..<(o + n)))
                o += n
            } else if let constantIv {
                let n = min(constantIv.count, 16)
                iv.replaceSubrange(0..<n, with: constantIv.prefix(n))
            }
            var subs: [(clear: Int, protected: Int)]?
            if useSubsamples {
                guard o + 2 <= box.payloadEnd else { break }
                let count = Int(readU16(data, o)); o += 2
                var list: [(clear: Int, protected: Int)] = []
                for _ in 0..<count {
                    guard o + 6 <= box.payloadEnd else { break }
                    let clear = Int(readU16(data, o)); o += 2
                    let protected = Int(readU32(data, o)); o += 4
                    list.append((clear, protected))
                }
                subs = list
            }
            out.append(SencSample(iv: iv, subsamples: subs))
        }
        return out
    }

    private static func parseSaizSaio(
        _ data: Data,
        traf: Box,
        moof: Box,
        perSampleIvSize: Int,
        constantIv: Data?,
        sampleCount: Int
    ) -> [SencSample] {
        let saizBoxes = findBoxesInRange(data, type: "saiz", from: traf.payloadStart, to: traf.payloadEnd)
        let saioBoxes = findBoxesInRange(data, type: "saio", from: traf.payloadStart, to: traf.payloadEnd)
        guard let saiz = saizBoxes.first, let saio = saioBoxes.first else { return [] }

        let saizFlags = readU32(data, saiz.payloadStart) & 0x00FF_FFFF
        var o = saiz.payloadStart + 4
        if saizFlags & 0x1 != 0 { o += 8 } // aux_info_type + parameter
        let defaultSize = Int(data[o]); o += 1
        let count = Int(readU32(data, o)); o += 4
        var sizes: [Int] = []
        if defaultSize == 0 {
            for _ in 0..<count {
                sizes.append(Int(data[o])); o += 1
            }
        } else {
            sizes = Array(repeating: defaultSize, count: count)
        }

        let saioVersionFlags = readU32(data, saio.payloadStart)
        let saioVersion = Int(saioVersionFlags >> 24)
        let saioFlags = saioVersionFlags & 0x00FF_FFFF
        var so = saio.payloadStart + 4
        if saioFlags & 0x1 != 0 { so += 8 }
        let entryCount = Int(readU32(data, so)); so += 4
        guard entryCount >= 1 else { return [] }
        let auxOffset: Int
        if saioVersion == 0 {
            auxOffset = Int(readU32(data, so))
        } else {
            auxOffset = Int(readU64(data, so))
        }
        // Offsets are usually relative to moof start (ISO) or absolute file offset.
        var auxPos = auxOffset
        if auxPos < moof.start || auxPos > data.count {
            auxPos = moof.start + auxOffset
        }
        guard auxPos >= 0, auxPos < data.count else { return [] }

        var out: [SencSample] = []
        let n = min(sizes.count, sampleCount > 0 ? sampleCount : sizes.count)
        for i in 0..<n {
            let entrySize = sizes[i]
            guard auxPos + entrySize <= data.count else { break }
            var p = auxPos
            var iv = Data(count: 16)
            if perSampleIvSize > 0 {
                let ivn = min(perSampleIvSize, 16)
                iv.replaceSubrange(0..<ivn, with: data.subdata(in: p..<(p + ivn)))
                p += ivn
            } else if let constantIv {
                let ivn = min(constantIv.count, 16)
                iv.replaceSubrange(0..<ivn, with: constantIv.prefix(ivn))
            }
            var subs: [(clear: Int, protected: Int)]?
            if entrySize > perSampleIvSize {
                let sc = Int(readU16(data, p)); p += 2
                var list: [(clear: Int, protected: Int)] = []
                for _ in 0..<sc {
                    let clear = Int(readU16(data, p)); p += 2
                    let protected = Int(readU32(data, p)); p += 4
                    list.append((clear, protected))
                }
                subs = list
            }
            out.append(SencSample(iv: iv, subsamples: subs))
            auxPos += entrySize
        }
        return out
    }

    private static func decryptSample(
        _ sample: Data,
        key: Data,
        iv: Data,
        subsamples: [(clear: Int, protected: Int)]?,
        scheme: String,
        cryptByteBlock: Int,
        skipByteBlock: Int
    ) -> Data? {
        if scheme == "cbcs" || scheme == "cbc1" {
            return decryptCbcsSample(
                sample,
                key: key,
                iv: iv,
                subsamples: subsamples,
                cryptByteBlock: max(cryptByteBlock, 1),
                skipByteBlock: skipByteBlock
            )
        }
        // cenc / cens: AES-CTR. Concatenate protected ranges → one CTR stream → merge.
        guard let ranges = protectedRanges(sampleCount: sample.count, subsamples: subsamples), !ranges.isEmpty else {
            return AesCtr.crypt(key: key, iv: iv, data: sample)
        }
        var encrypted = Data()
        for r in ranges {
            encrypted.append(sample.subdata(in: r))
        }
        guard let plain = AesCtr.crypt(key: key, iv: iv, data: encrypted) else { return nil }
        var out = sample
        var read = 0
        for r in ranges {
            let len = r.count
            out.replaceSubrange(r, with: plain.subdata(in: read..<(read + len)))
            read += len
        }
        return out
    }

    private static func decryptCbcsSample(
        _ sample: Data,
        key: Data,
        iv: Data,
        subsamples: [(clear: Int, protected: Int)]?,
        cryptByteBlock: Int,
        skipByteBlock: Int
    ) -> Data? {
        var out = sample
        if let subsamples {
            var cursor = 0
            for sub in subsamples {
                cursor += sub.clear
                let end = cursor + sub.protected
                guard end <= out.count else { break }
                let slice = out.subdata(in: cursor..<end)
                guard let dec = AesCtr.cryptCbcs(
                    key: key,
                    iv: iv,
                    data: slice,
                    cryptByteBlock: cryptByteBlock,
                    skipByteBlock: skipByteBlock
                ) else { return nil }
                out.replaceSubrange(cursor..<end, with: dec)
                cursor = end
            }
            return out
        }
        return AesCtr.cryptCbcs(
            key: key,
            iv: iv,
            data: sample,
            cryptByteBlock: cryptByteBlock,
            skipByteBlock: skipByteBlock
        )
    }

    private static func protectedRanges(
        sampleCount: Int,
        subsamples: [(clear: Int, protected: Int)]?
    ) -> [Range<Int>]? {
        guard let subsamples else {
            return [0..<sampleCount]
        }
        var ranges: [Range<Int>] = []
        var cursor = 0
        for sub in subsamples {
            cursor += sub.clear
            if sub.protected > 0 {
                ranges.append(cursor..<(cursor + sub.protected))
            }
            cursor += sub.protected
        }
        return ranges
    }

    private static func findBoxesInRange(_ data: Data, type: String, from: Int, to: Int) -> [Box] {
        var out: [Box] = []
        walk(data, from: from, to: to, recursive: true) { box in
            if box.type == type { out.append(box) }
            return true
        }
        return out
    }
}