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

import Foundation
import Network

/// Lokale HTTP-proxy: haalt DASH MPD/segmenten op, strip DRM-metadata, decrypt CENC met ClearKey.
/// VLC speelt daarna plain DASH vanaf 127.0.0.1.
final class ClearKeyDashProxy {
    static let shared = ClearKeyDashProxy()

    private let queue = DispatchQueue(label: "nl.vonas.vmc.clearkey-proxy")
    private var listener: NWListener?
    private var port: UInt16 = 0
    private var remoteManifest: URL?
    private var keyMap: ClearKeyKeyMap?
    private var cryptoState = CencMp4Decryptor.SessionState()
    private var generation = 0

    private init() {}

    var isRunning: Bool { listener != nil && port > 0 }

    /// Start (of herstart) proxy voor dit manifest + keys. Geeft lokale MPD-URL terug.
    func start(manifestURL: URL, keys: [ClearKeyEntry]) throws -> URL {
        let map = ClearKeyKeyMap(entries: keys)
        guard !map.isEmpty else {
            throw ProxyError.noKeys
        }
        stop()
        remoteManifest = manifestURL
        keyMap = map
        cryptoState = CencMp4Decryptor.SessionState()
        generation += 1

        let params = NWParameters.tcp
        params.allowLocalEndpointReuse = true
        let listener = try NWListener(using: params, on: .any)
        self.listener = listener

        let ready = DispatchSemaphore(value: 0)
        var readyError: Error?

        listener.stateUpdateHandler = { [weak self] state in
            guard let self else { return }
            switch state {
            case .ready:
                if let p = listener.port?.rawValue {
                    self.port = p
                }
                ready.signal()
            case .failed(let err):
                readyError = err
                ready.signal()
            default:
                break
            }
        }

        listener.newConnectionHandler = { [weak self] conn in
            self?.handle(connection: conn)
        }

        listener.start(queue: queue)
        _ = ready.wait(timeout: .now() + 5)
        if let readyError { throw readyError }
        guard port > 0 else { throw ProxyError.bindFailed }

        return URL(string: "http://127.0.0.1:\(port)/manifest.mpd")!
    }

    func stop() {
        listener?.cancel()
        listener = nil
        port = 0
        remoteManifest = nil
        keyMap = nil
        cryptoState = CencMp4Decryptor.SessionState()
    }

    // MARK: - HTTP

    private func handle(connection: NWConnection) {
        connection.start(queue: queue)
        receiveRequest(on: connection, buffer: Data())
    }

    private func receiveRequest(on connection: NWConnection, buffer: Data) {
        connection.receive(minimumIncompleteLength: 1, maximumLength: 64 * 1024) { [weak self] data, _, isComplete, error in
            guard let self else {
                connection.cancel()
                return
            }
            if error != nil {
                connection.cancel()
                return
            }
            var buf = buffer
            if let data { buf.append(data) }
            if let range = buf.range(of: Data("\r\n\r\n".utf8)) {
                let headerData = buf.subdata(in: 0..<range.lowerBound)
                self.serve(connection: connection, headerData: headerData)
                return
            }
            if isComplete {
                connection.cancel()
                return
            }
            if buf.count > 256 * 1024 {
                connection.cancel()
                return
            }
            self.receiveRequest(on: connection, buffer: buf)
        }
    }

    private func serve(connection: NWConnection, headerData: Data) {
        guard let header = String(data: headerData, encoding: .utf8),
              let firstLine = header.split(separator: "\r\n", omittingEmptySubsequences: false).first
        else {
            respond(connection, status: 400, contentType: "text/plain", body: Data("bad request".utf8))
            return
        }
        let parts = firstLine.split(separator: " ")
        guard parts.count >= 2 else {
            respond(connection, status: 400, contentType: "text/plain", body: Data("bad request".utf8))
            return
        }
        let method = String(parts[0])
        let path = String(parts[1])
        guard method == "GET" || method == "HEAD" else {
            respond(connection, status: 405, contentType: "text/plain", body: Data("method".utf8))
            return
        }

        Task {
            do {
                let (body, contentType) = try await self.payload(for: path)
                if method == "HEAD" {
                    self.respond(connection, status: 200, contentType: contentType, body: Data(), headOnly: true, contentLength: body.count)
                } else {
                    self.respond(connection, status: 200, contentType: contentType, body: body)
                }
            } catch {
                let msg = error.localizedDescription
                self.respond(connection, status: 502, contentType: "text/plain; charset=utf-8", body: Data(msg.utf8))
            }
        }
    }

    private func payload(for pathAndQuery: String) async throws -> (Data, String) {
        guard let keyMap, let remoteManifest else {
            throw ProxyError.notStarted
        }
        let pathOnly = pathAndQuery.split(separator: "?", maxSplits: 1).first.map(String.init) ?? pathAndQuery

        if pathOnly == "/" || pathOnly == "/manifest.mpd" {
            let raw = try await fetch(remoteManifest)
            let rewritten = rewriteMpd(String(data: raw, encoding: .utf8) ?? "", mpdURL: remoteManifest)
            return (Data(rewritten.utf8), "application/dash+xml")
        }

        if pathOnly.hasPrefix("/a/") {
            let b64 = String(pathOnly.dropFirst(3))
            guard let url = decodeURL(b64) else { throw ProxyError.badPath }
            return try await fetchAndDecrypt(url, keys: keyMap)
        }

        if pathOnly.hasPrefix("/p/") {
            let rest = String(pathOnly.dropFirst(3))
            guard let slash = rest.firstIndex(of: "/") else { throw ProxyError.badPath }
            let b64 = String(rest[..<slash])
            let rel = String(rest[rest.index(after: slash)...])
            guard let base = decodeURL(b64) else { throw ProxyError.badPath }
            let joined = join(base: base, relative: rel.removingPercentEncoding ?? rel)
            return try await fetchAndDecrypt(joined, keys: keyMap)
        }

        throw ProxyError.badPath
    }

    private func fetchAndDecrypt(_ url: URL, keys: ClearKeyKeyMap) async throws -> (Data, String) {
        var data = try await fetch(url)
        queue.sync {
            _ = CencMp4Decryptor.process(data: &data, keys: keys, state: &cryptoState)
        }
        let ct: String
        let lower = url.absoluteString.lowercased()
        if lower.contains(".mpd") {
            ct = "application/dash+xml"
        } else if lower.contains(".m4s") || lower.contains(".mp4") || lower.contains("init") {
            ct = "video/mp4"
        } else {
            ct = "application/octet-stream"
        }
        return (data, ct)
    }

    private func fetch(_ url: URL) async throws -> Data {
        var req = URLRequest(url: url)
        req.timeoutInterval = 20
        req.setValue("VMC-iOS-ClearKey/0.5.1", forHTTPHeaderField: "User-Agent")
        let (data, resp) = try await URLSession.shared.data(for: req)
        if let http = resp as? HTTPURLResponse, !(200...299).contains(http.statusCode) {
            throw ProxyError.upstream(http.statusCode)
        }
        return data
    }

    private func rewriteMpd(_ xml: String, mpdURL: URL) -> String {
        var s = xml
        // Strip ContentProtection (ClearKey/Widevine/PlayReady/mp4protection).
        s = s.replacingOccurrences(
            of: #"<ContentProtection\b[^>]*/>"#,
            with: "",
            options: .regularExpression
        )
        s = s.replacingOccurrences(
            of: #"<ContentProtection\b[^>]*>[\s\S]*?</ContentProtection>"#,
            with: "",
            options: .regularExpression
        )

        let originDir = mpdURL.deletingLastPathComponent()
        let proxyBase = "http://127.0.0.1:\(port)/p/\(encodeURL(originDir))/"

        // Herschrijf bestaande BaseURL-tekst naar proxy.
        s = s.replacingOccurrences(
            of: #"<BaseURL>([\s\S]*?)</BaseURL>"#,
            with: "<BaseURL>\(proxyBase)</BaseURL>",
            options: .regularExpression
        )

        if !s.localizedCaseInsensitiveContains("<BaseURL") {
            if let regex = try? NSRegularExpression(pattern: #"<MPD\b[^>]*>"#, options: []) {
                let range = NSRange(s.startIndex..<s.endIndex, in: s)
                if let match = regex.firstMatch(in: s, options: [], range: range),
                   let r = Range(match.range, in: s) {
                    s.insert(contentsOf: "\n  <BaseURL>\(proxyBase)</BaseURL>", at: r.upperBound)
                }
            }
        }

        // Absolute http(s) URL's in attributen → /a/<b64>
        if let urlRegex = try? NSRegularExpression(pattern: #"https?://[^\s"'<>]+"#, options: []) {
            let ns = s as NSString
            let matches = urlRegex.matches(in: s, options: [], range: NSRange(location: 0, length: ns.length))
            for match in matches.reversed() {
                let abs = ns.substring(with: match.range)
                // Laat lokale proxy-URL's met rust.
                if abs.contains("127.0.0.1:\(port)") { continue }
                let replacement = "http://127.0.0.1:\(port)/a/\(encodeURLString(abs))"
                s = (s as NSString).replacingCharacters(in: match.range, with: replacement)
            }
        }

        return s
    }

    private func join(base: URL, relative: String) -> URL {
        if relative.hasPrefix("http://") || relative.hasPrefix("https://") {
            return URL(string: relative) ?? base
        }
        return URL(string: relative, relativeTo: base)?.absoluteURL ?? base.appendingPathComponent(relative)
    }

    private func encodeURL(_ url: URL) -> String {
        encodeURLString(url.absoluteString)
    }

    private func encodeURLString(_ s: String) -> String {
        Data(s.utf8).base64EncodedString()
            .replacingOccurrences(of: "+", with: "-")
            .replacingOccurrences(of: "/", with: "_")
            .replacingOccurrences(of: "=", with: "")
    }

    private func decodeURL(_ b64url: String) -> URL? {
        var b64 = b64url
            .replacingOccurrences(of: "-", with: "+")
            .replacingOccurrences(of: "_", with: "/")
        let pad = (4 - b64.count % 4) % 4
        if pad > 0 { b64 += String(repeating: "=", count: pad) }
        guard let data = Data(base64Encoded: b64),
              let s = String(data: data, encoding: .utf8)
        else { return nil }
        return URL(string: s)
    }

    private func respond(
        _ connection: NWConnection,
        status: Int,
        contentType: String,
        body: Data,
        headOnly: Bool = false,
        contentLength: Int? = nil
    ) {
        let reason: String
        switch status {
        case 200: reason = "OK"
        case 400: reason = "Bad Request"
        case 405: reason = "Method Not Allowed"
        case 502: reason = "Bad Gateway"
        default: reason = "Error"
        }
        let len = contentLength ?? body.count
        var header = "HTTP/1.1 \(status) \(reason)\r\n"
        header += "Content-Type: \(contentType)\r\n"
        header += "Content-Length: \(len)\r\n"
        header += "Connection: close\r\n"
        header += "Access-Control-Allow-Origin: *\r\n"
        header += "Cache-Control: no-store\r\n"
        header += "\r\n"
        var payload = Data(header.utf8)
        if !headOnly { payload.append(body) }
        connection.send(content: payload, completion: .contentProcessed { _ in
            connection.cancel()
        })
    }

    enum ProxyError: LocalizedError {
        case noKeys
        case bindFailed
        case notStarted
        case badPath
        case upstream(Int)

        var errorDescription: String? {
            switch self {
            case .noKeys: return "Geen geldige ClearKey-sleutels"
            case .bindFailed: return "Kon lokale ClearKey-proxy niet starten"
            case .notStarted: return "ClearKey-proxy is niet actief"
            case .badPath: return "Ongeldig proxy-pad"
            case .upstream(let code): return "Upstream HTTP \(code)"
            }
        }
    }
}