433 lines · 16.5 KB · Swift
RawDownload
import Foundation#if canImport(FoundationNetworking)import FoundationNetworking#endif public struct GrokWebBillingSnapshot: Sendable, Equatable {    public let usedPercent: Double?    public let resetsAt: Date?     public init(usedPercent: Double?, resetsAt: Date?) {        self.usedPercent = usedPercent        self.resetsAt = resetsAt    }} public enum GrokWebBillingError: LocalizedError, Sendable {    case missingCredentials    case emptyResponse    case invalidResponse    case requestFailed(Int, String)    case rpcFailed(Int, String)    case parseFailed     public var errorDescription: String? {        switch self {        case .missingCredentials:            "Grok web billing requires a signed-in grok.com browser session or `grok login`."        case .emptyResponse:            "Grok web billing returned no protobuf payload."        case .invalidResponse:            "Grok web billing returned an invalid response."        case let .requestFailed(status, body):            if status == 401 || status == 403 {                Self.reauthMessage            } else {                "Grok web billing request failed with HTTP \(status): \(body)"            }        case let .rpcFailed(status, message):            if Self.isAuthenticationFailure(status: status, message: message) {                Self.reauthMessage            } else {                "Grok web billing RPC failed with status \(status): \(message)"            }        case .parseFailed:            "Could not parse Grok web billing usage."        }    }     private static let reauthMessage =        "Grok web billing rejected credentials. Sign in to grok.com in Chrome or run `grok login` to refresh xAI auth."     static func isAuthenticationFailure(status: Int, message: String) -> Bool {        if status == 16 { return true }        guard status == 7 else { return false }        let lower = message.lowercased()        return lower.contains("bad-credentials") ||            lower.contains("unauthenticated") ||            (lower.contains("oauth2") && lower.contains("could not be validated")) ||            (lower.contains("access token") &&                (lower.contains("invalid") ||                    lower.contains("expired") ||                    lower.contains("could not be validated")))    }} public enum GrokWebBillingFetcher {    public static let defaultEndpoint =        URL(string: "https://grok.com/grok_api_v2.GrokBuildBilling/GetGrokCreditsConfig")!    private static let requestTimeoutSeconds: TimeInterval = 15     public static func fetch(        credentials: GrokCredentials,        session transport: any ProviderHTTPTransport = ProviderHTTPClient.shared,        endpoint: URL = Self.defaultEndpoint) async throws -> GrokWebBillingSnapshot    {        try await self.fetch(            authorizationHeader: "Bearer \(credentials.accessToken)",            cookieHeader: nil,            transport: transport,            endpoint: endpoint)    }     public static func fetch(        cookieHeader: String,        session transport: any ProviderHTTPTransport = ProviderHTTPClient.shared,        endpoint: URL = Self.defaultEndpoint) async throws -> GrokWebBillingSnapshot    {        try await self.fetch(            cookieHeader: cookieHeader,            credentials: nil,            session: transport,            endpoint: endpoint)    }     public static func fetch(        cookieHeader: String,        credentials: GrokCredentials?,        session transport: any ProviderHTTPTransport = ProviderHTTPClient.shared,        endpoint: URL = Self.defaultEndpoint) async throws -> GrokWebBillingSnapshot    {        let authorizationHeader = credentials.flatMap { credential in            credential.isExpired ? nil : "Bearer \(credential.accessToken)"        }        return try await self.fetch(            authorizationHeader: authorizationHeader,            cookieHeader: cookieHeader,            transport: transport,            endpoint: endpoint)    }     private static func fetch(        authorizationHeader: String?,        cookieHeader: String?,        transport: any ProviderHTTPTransport,        endpoint: URL) async throws -> GrokWebBillingSnapshot    {        do {            return try await self.fetchOnce(                authorizationHeader: authorizationHeader,                cookieHeader: cookieHeader,                transport: transport,                endpoint: endpoint)        } catch where self.shouldRetry(error) {            return try await self.fetchOnce(                authorizationHeader: authorizationHeader,                cookieHeader: cookieHeader,                transport: transport,                endpoint: endpoint)        }    }     private static func fetchOnce(        authorizationHeader: String?,        cookieHeader: String?,        transport: any ProviderHTTPTransport,        endpoint: URL) async throws -> GrokWebBillingSnapshot    {        var request = URLRequest(url: endpoint)        request.httpMethod = "POST"        request.timeoutInterval = Self.requestTimeoutSeconds        request.httpBody = Data([0x00, 0x00, 0x00, 0x00, 0x00])        if let authorizationHeader {            request.setValue(authorizationHeader, forHTTPHeaderField: "Authorization")        }        if let cookieHeader {            request.setValue(cookieHeader, forHTTPHeaderField: "Cookie")        }        request.setValue("https://grok.com", forHTTPHeaderField: "Origin")        request.setValue("https://grok.com/?_s=usage", forHTTPHeaderField: "Referer")        request.setValue("*/*", forHTTPHeaderField: "Accept")        request.setValue("application/grpc-web+proto", forHTTPHeaderField: "Content-Type")        request.setValue("1", forHTTPHeaderField: "x-grpc-web")        request.setValue("connect-es/2.1.1", forHTTPHeaderField: "x-user-agent")        request.setValue("CodexBar", forHTTPHeaderField: "User-Agent")         let response: ProviderHTTPResponse        do {            response = try await transport.response(for: request)        } catch let error as URLError where error.code == .badServerResponse {            throw GrokWebBillingError.invalidResponse        } catch {            throw error        }        guard response.statusCode == 200 else {            let body = String(data: response.data.prefix(400), encoding: .utf8) ?? ""            throw GrokWebBillingError.requestFailed(response.statusCode, body)        }        try Self.validateGRPCStatusFields(Self.grpcHeaderFields(from: response.response.allHeaderFields))        try Self.validateGRPCWebTrailers(response.data)         return try Self.parseGRPCWebResponse(response.data)    }     private static func shouldRetry(_ error: Error) -> Bool {        if let urlError = error as? URLError {            return urlError.code == .timedOut || urlError.code == .networkConnectionLost        }        if case let GrokWebBillingError.requestFailed(status, body) = error {            if [408, 502, 503, 504].contains(status) { return true }            return body.localizedCaseInsensitiveContains("timeout")                || body.localizedCaseInsensitiveContains("deadline")        }        guard case let GrokWebBillingError.rpcFailed(status, message) = error else { return false }        if status == 4 { return true }        guard status == 1 else { return false }        return message.localizedCaseInsensitiveContains("timeout")            || message.localizedCaseInsensitiveContains("deadline")            || message.localizedCaseInsensitiveContains("expired")    }     static func parseGRPCWebResponse(_ data: Data, now: Date = Date()) throws -> GrokWebBillingSnapshot {        var payloads = Self.grpcWebDataFrames(from: data)        if payloads.isEmpty, Self.looksLikeProtobufPayload(data) {            payloads = [data]        }        guard !payloads.isEmpty else { throw GrokWebBillingError.emptyResponse }         var scan = ProtobufScan()        for payload in payloads {            scan.merge(Self.scanProtobuf(payload, depth: 0))        }         let parsedPercent = scan.fixed32Fields            .filter { field in                field.path.last == 1 && field.value.isFinite && field.value >= 0 && field.value <= 100            }            .min { lhs, rhs in                lhs.path.count == rhs.path.count ? lhs.order < rhs.order : lhs.path.count < rhs.path.count            }            .map { Double($0.value) }         let resetFields = scan.varintFields.compactMap { field -> (path: [UInt64], date: Date)? in            let raw = field.value            guard raw >= 1_700_000_000, raw <= 2_100_000_000 else { return nil }            return (field.path, Date(timeIntervalSince1970: TimeInterval(raw)))        }        let futureResetFields = resetFields.filter { $0.date > now }        let reset = futureResetFields            .filter { $0.path == [1, 5, 1] }            .map(\.date)            .min() ?? futureResetFields            .map(\.date)            .min()         let hasUsagePeriod = scan.varintFields.contains { field in            field.path.starts(with: [1, 6]) ||                (field.path == [1, 8, 1] && (field.value == 1 || field.value == 2))        }        let noUsageYet = parsedPercent == nil &&            scan.fixed32Fields.isEmpty &&            reset != nil &&            hasUsagePeriod        guard let percent = parsedPercent ?? (noUsageYet ? 0 : nil) else {            throw GrokWebBillingError.parseFailed        }        return GrokWebBillingSnapshot(usedPercent: percent, resetsAt: reset)    }     static func looksLikeProtobufPayload(_ data: Data) -> Bool {        guard let first = data.first else { return false }        let fieldNumber = first >> 3        let wireType = first & 0x07        return fieldNumber > 0 && (wireType == 0 || wireType == 1 || wireType == 2 || wireType == 5)    }     static func grpcWebDataFrames(from data: Data) -> [Data] {        let bytes = [UInt8](data)        var frames: [Data] = []        var index = 0        while index < bytes.count {            guard index + 5 <= bytes.count else { return [] }            let flags = bytes[index]            let length = (Int(bytes[index + 1]) << 24)                | (Int(bytes[index + 2]) << 16)                | (Int(bytes[index + 3]) << 8)                | Int(bytes[index + 4])            let start = index + 5            let end = start + length            guard length >= 0, end <= bytes.count else { return [] }            if flags & 0x80 == 0 {                frames.append(Data(bytes[start..<end]))            }            index = end        }        return frames    }     static func validateGRPCWebTrailers(_ data: Data) throws {        try self.validateGRPCStatusFields(self.grpcWebTrailerFields(from: data))    }     static func validateGRPCStatusFields(_ fields: [String: String]) throws {        guard let rawStatus = fields["grpc-status"],              let status = Int(rawStatus),              status != 0        else {            return        }        throw GrokWebBillingError.rpcFailed(status, fields["grpc-message"] ?? "")    }     static func grpcHeaderFields(from headers: [AnyHashable: Any]) -> [String: String] {        var fields: [String: String] = [:]        for (key, value) in headers {            let normalizedKey = String(describing: key)                .trimmingCharacters(in: .whitespacesAndNewlines)                .lowercased()            guard normalizedKey.hasPrefix("grpc-") else { continue }            fields[normalizedKey] = String(describing: value)                .trimmingCharacters(in: .whitespacesAndNewlines)                .removingPercentEncoding ?? ""        }        return fields    }     static func grpcWebTrailerFields(from data: Data) -> [String: String] {        let bytes = [UInt8](data)        var fields: [String: String] = [:]        var index = 0        while index + 5 <= bytes.count {            let flags = bytes[index]            let length = (Int(bytes[index + 1]) << 24)                | (Int(bytes[index + 2]) << 16)                | (Int(bytes[index + 3]) << 8)                | Int(bytes[index + 4])            let start = index + 5            let end = start + length            guard length >= 0, end <= bytes.count else { break }            if flags & 0x80 != 0, let text = String(data: Data(bytes[start..<end]), encoding: .utf8) {                for line in text.components(separatedBy: .newlines) where !line.isEmpty {                    guard let separator = line.firstIndex(of: ":") else { continue }                    let key = line[..<separator]                        .trimmingCharacters(in: .whitespacesAndNewlines)                        .lowercased()                    let value = line[line.index(after: separator)...]                        .trimmingCharacters(in: .whitespacesAndNewlines)                        .removingPercentEncoding ?? ""                    fields[key] = value                }            }            index = end        }        return fields    }     private struct ProtobufScan {        struct Fixed32Field {            var path: [UInt64]            var value: Float            var order: Int        }         struct VarintField {            var path: [UInt64]            var value: UInt64        }         var fixed32Fields: [Fixed32Field] = []        var varintFields: [VarintField] = []         mutating func merge(_ other: ProtobufScan) {            self.fixed32Fields.append(contentsOf: other.fixed32Fields)            self.varintFields.append(contentsOf: other.varintFields)        }    }     private static func scanProtobuf(_ data: Data, depth: Int) -> ProtobufScan {        self.scanProtobuf(data, depth: depth, path: [], order: 0).scan    }     private static func scanProtobuf(        _ data: Data,        depth: Int,        path: [UInt64],        order: Int) -> (scan: ProtobufScan, order: Int)    {        let bytes = [UInt8](data)        var scan = ProtobufScan()        var index = 0        var nextOrder = order         while index < bytes.count {            let fieldStart = index            guard let key = Self.readVarint(bytes, index: &index), key != 0 else {                index = fieldStart + 1                continue            }            let fieldNumber = key >> 3            let wireType = key & 0x07            let fieldPath = path + [fieldNumber]             switch wireType {            case 0:                if let value = Self.readVarint(bytes, index: &index) {                    scan.varintFields.append(ProtobufScan.VarintField(path: fieldPath, value: value))                } else {                    index = fieldStart + 1                }            case 1:                guard index + 8 <= bytes.count else { return (scan, nextOrder) }                index += 8            case 2:                guard let length = Self.readVarint(bytes, index: &index),                      length <= UInt64(bytes.count - index)                else {                    index = fieldStart + 1                    continue                }                let start = index                let end = index + Int(length)                if depth < 4 {                    let nested = Self.scanProtobuf(                        Data(bytes[start..<end]),                        depth: depth + 1,                        path: fieldPath,                        order: nextOrder)                    scan.merge(nested.scan)                    nextOrder = nested.order                }                index = end            case 5:                guard index + 4 <= bytes.count else { return (scan, nextOrder) }                let bitPattern = UInt32(bytes[index])                    | (UInt32(bytes[index + 1]) << 8)                    | (UInt32(bytes[index + 2]) << 16)                    | (UInt32(bytes[index + 3]) << 24)                scan.fixed32Fields.append(ProtobufScan.Fixed32Field(                    path: fieldPath,                    value: Float(bitPattern: bitPattern),                    order: nextOrder))                nextOrder += 1                index += 4            default:                index = fieldStart + 1            }        }         return (scan, nextOrder)    }     private static func readVarint(_ bytes: [UInt8], index: inout Int) -> UInt64? {        var value: UInt64 = 0        var shift: UInt64 = 0        while index < bytes.count, shift < 64 {            let byte = bytes[index]            index += 1            value |= UInt64(byte & 0x7F) << shift            if byte & 0x80 == 0 { return value }            shift += 7        }        return nil    }}