diff --git a/Sources/CodexBarCore/Providers/Kiro/KiroAPIFetchStrategy.swift b/Sources/CodexBarCore/Providers/Kiro/KiroAPIFetchStrategy.swift new file mode 100644 index 0000000000..621a798e46 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Kiro/KiroAPIFetchStrategy.swift @@ -0,0 +1,35 @@ +import Foundation + +/// Direct API fetch strategy for Kiro usage that bypasses `kiro-cli` and calls +/// the AWS Q `getUsageLimits` REST endpoint directly. +/// +/// This fixes enterprise/IdC account usage fetching where `kiro-cli` fails because +/// it incorrectly passes `profileArn` to the legacy REST API, causing a 400 error. +/// The Kiro 0.9.2 version of these APIs does NOT accept `profileArn`. +/// +/// Reference: https://github.com/ZyphrZero/kiro.rs (v0.6.11 fix) +struct KiroAPIFetchStrategy: ProviderFetchStrategy { + let id: String = "kiro.api" + let kind: ProviderFetchKind = .apiToken + + private let fetcher: KiroAPIUsageFetcher + + init(fetcher: KiroAPIUsageFetcher = KiroAPIUsageFetcher()) { + self.fetcher = fetcher + } + + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + self.fetcher.hasCredentials(allowSocial: context.sourceMode == .api) + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let snapshot = try await self.fetcher.fetchUsage(allowSocial: context.sourceMode == .api) + return self.makeResult( + usage: snapshot.toUsageSnapshot(), + sourceLabel: "api") + } + + func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool { + context.sourceMode == .auto + } +} diff --git a/Sources/CodexBarCore/Providers/Kiro/KiroAPIUsageFetcher.swift b/Sources/CodexBarCore/Providers/Kiro/KiroAPIUsageFetcher.swift new file mode 100644 index 0000000000..6e706d6cc7 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Kiro/KiroAPIUsageFetcher.swift @@ -0,0 +1,616 @@ +import Foundation +#if canImport(CryptoKit) +import CryptoKit +#else +import Crypto +#endif + +struct KiroAPIUsageFetcher: Sendable { + private static let usageAPIKiroVersion = "0.9.2" + private static let apiRegions = ["us-east-1", "eu-central-1"] + private static let idcAmzUserAgent = + "aws-sdk-js/3.738.0 ua/2.1 os/other lang/js md/browser#unknown_unknown api/sso-oidc#3.738.0 m/E KiroIDE" + private static let maxErrorBodyLength = 240 + + let credentialStore: KiroCLICredentialStore + let transport: any ProviderHTTPTransport + + init( + credentialStore: KiroCLICredentialStore = KiroCLICredentialStore(), + transport: any ProviderHTTPTransport = ProviderHTTPClient.shared) + { + self.credentialStore = credentialStore + self.transport = transport + } + + func hasCredentials(allowSocial: Bool) -> Bool { + self.credentialStore.loadCredentials(allowSocial: allowSocial) != nil + } + + func fetchUsage(allowSocial: Bool) async throws -> KiroUsageSnapshot { + guard var credentials = self.credentialStore.loadCredentials(allowSocial: allowSocial) else { + throw KiroAPIError.credentialsNotFound + } + + if credentials.needsRefresh() { + credentials = try await self.refreshCredentials(credentials) + } + + return try await self.fetchUsage(credentials: credentials) + } + + // MARK: - Token refresh + + private func refreshCredentials(_ credentials: KiroCLICredentials) async throws -> KiroCLICredentials { + guard let refreshToken = credentials.refreshToken, !refreshToken.isEmpty else { + throw KiroAPIError.authenticationFailed + } + guard refreshToken.count >= 100, !refreshToken.contains("...") else { + throw KiroAPIError.refreshTokenUnavailable + } + + if credentials.isExternalIDP { + return try await self.refreshExternalIDPToken(credentials, refreshToken: refreshToken) + } + + let authMethod = credentials.canonicalAuthMethod + if authMethod == "idc" { + return try await self.refreshIDCToken(credentials, refreshToken: refreshToken) + } + return try await self.refreshSocialToken(credentials, refreshToken: refreshToken) + } + + private func refreshSocialToken( + _ credentials: KiroCLICredentials, + refreshToken: String) async throws -> KiroCLICredentials + { + let region = credentials.effectiveAuthRegion + let host = "prod.\(region).auth.desktop.kiro.dev" + guard let url = URL(string: "https://\(host)/refreshToken") else { + throw KiroAPIError.invalidURL + } + + let machineID = Self.machineID(for: credentials) + var request = URLRequest(url: url) + request.httpMethod = "POST" + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + request.setValue("application/json, text/plain, */*", forHTTPHeaderField: "Accept") + request.setValue( + "KiroIDE-\(Self.usageAPIKiroVersion)-\(machineID)", + forHTTPHeaderField: "User-Agent") + request.setValue(host, forHTTPHeaderField: "host") + request.httpBody = try JSONSerialization.data(withJSONObject: ["refreshToken": refreshToken]) + + let response = try await self.transport.response(for: request) + try Self.validateRefreshResponse(response) + + let json = try Self.decodeJSON(response.data) + guard let accessToken = Self.string(json, keys: "access_token", "accessToken"), !accessToken.isEmpty else { + throw KiroAPIError.authenticationFailed + } + + return credentials.refreshed( + accessToken: accessToken, + refreshToken: Self.string(json, keys: "refresh_token", "refreshToken") ?? refreshToken, + expiresAt: Self.expiresAt(from: json)) + } + + private func refreshIDCToken( + _ credentials: KiroCLICredentials, + refreshToken: String) async throws -> KiroCLICredentials + { + guard let clientID = credentials.clientID, !clientID.isEmpty else { + throw KiroAPIError.authenticationFailed + } + guard let clientSecret = credentials.clientSecret, !clientSecret.isEmpty else { + throw KiroAPIError.authenticationFailed + } + + let region = credentials.effectiveAuthRegion + let host = "oidc.\(region).amazonaws.com" + guard let url = URL(string: "https://\(host)/token") else { + throw KiroAPIError.invalidURL + } + + var request = URLRequest(url: url) + request.httpMethod = "POST" + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + request.setValue(host, forHTTPHeaderField: "Host") + request.setValue(Self.idcAmzUserAgent, forHTTPHeaderField: "x-amz-user-agent") + request.setValue("node", forHTTPHeaderField: "User-Agent") + request.httpBody = try JSONSerialization.data(withJSONObject: [ + "clientId": clientID, + "clientSecret": clientSecret, + "refreshToken": refreshToken, + "grantType": "refresh_token", + ]) + + let response = try await self.transport.response(for: request) + try Self.validateRefreshResponse(response) + + let json = try Self.decodeJSON(response.data) + guard let accessToken = Self.string(json, keys: "access_token", "accessToken"), !accessToken.isEmpty else { + throw KiroAPIError.authenticationFailed + } + + return credentials.refreshed( + accessToken: accessToken, + refreshToken: Self.string(json, keys: "refresh_token", "refreshToken") ?? refreshToken, + expiresAt: Self.expiresAt(from: json)) + } + + private func refreshExternalIDPToken( + _ credentials: KiroCLICredentials, + refreshToken: String) async throws -> KiroCLICredentials + { + guard let clientID = credentials.clientID, !clientID.isEmpty else { + throw KiroAPIError.authenticationFailed + } + guard let tokenEndpoint = credentials.tokenEndpoint, !tokenEndpoint.isEmpty else { + throw KiroAPIError.authenticationFailed + } + try KiroCLICredentials.validateExternalIDPEndpoint(tokenEndpoint) + + guard let url = URL(string: tokenEndpoint) else { + throw KiroAPIError.invalidURL + } + + var bodyComponents = [ + URLQueryItem(name: "client_id", value: clientID), + URLQueryItem(name: "grant_type", value: "refresh_token"), + URLQueryItem(name: "refresh_token", value: refreshToken), + ] + if let scopes = credentials.scopes, !scopes.isEmpty { + bodyComponents.append(URLQueryItem(name: "scope", value: scopes)) + } + + var request = URLRequest(url: url) + request.httpMethod = "POST" + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.setValue("application/x-www-form-urlencoded", forHTTPHeaderField: "Content-Type") + request.httpBody = bodyComponents + .percentEncoded() + .data(using: .utf8) + + let response = try await self.transport.response(for: request) + try Self.validateRefreshResponse(response) + + let json = try Self.decodeJSON(response.data) + guard let accessToken = Self.string(json, keys: "access_token", "accessToken"), !accessToken.isEmpty else { + throw KiroAPIError.authenticationFailed + } + + return credentials.refreshed( + accessToken: accessToken, + refreshToken: Self.string(json, keys: "refresh_token", "refreshToken") ?? refreshToken, + expiresAt: Self.expiresAt(from: json)) + } + + // MARK: - Usage fetch + + private func fetchUsage(credentials: KiroCLICredentials) async throws -> KiroUsageSnapshot { + let regions = Self.regionCandidates(for: credentials.region) + var lastError: Error? + + for region in regions { + do { + return try await self.fetchUsageFromRegion(region: region, credentials: credentials) + } catch let error as URLError where error.code == .badServerResponse { + lastError = error + continue + } catch { + throw error + } + } + + throw lastError ?? KiroAPIError.allRegionsFailed + } + + private func fetchUsageFromRegion( + region: String, + credentials: KiroCLICredentials) async throws -> KiroUsageSnapshot + { + let host = "q.\(region).amazonaws.com" + let urlString = + "https://\(host)/getUsageLimits?origin=AI_EDITOR&resourceType=AGENTIC_REQUEST&isEmailRequired=true" + guard let url = URL(string: urlString) else { + throw KiroAPIError.invalidURL + } + + let machineID = Self.machineID(for: credentials) + var request = URLRequest(url: url) + request.httpMethod = "GET" + request.timeoutInterval = 30 + + let userAgent = [ + "aws-sdk-js/1.0.0 ua/2.1 os/macos lang/js md/nodejs#20.0.0", + "api/codewhispererruntime#1.0.0 m/N,E KiroIDE-\(Self.usageAPIKiroVersion)-\(machineID)", + ].joined(separator: " ") + let amzUserAgent = "aws-sdk-js/1.0.0 KiroIDE-\(Self.usageAPIKiroVersion)-\(machineID)" + + request.setValue(amzUserAgent, forHTTPHeaderField: "x-amz-user-agent") + request.setValue(userAgent, forHTTPHeaderField: "user-agent") + request.setValue(host, forHTTPHeaderField: "host") + request.setValue(UUID().uuidString, forHTTPHeaderField: "amz-sdk-invocation-id") + request.setValue("attempt=1; max=1", forHTTPHeaderField: "amz-sdk-request") + request.setValue("Bearer \(credentials.accessToken)", forHTTPHeaderField: "Authorization") + request.setValue("close", forHTTPHeaderField: "Connection") + if let tokenType = credentials.tokenTypeHeader { + request.setValue(tokenType, forHTTPHeaderField: "tokentype") + } + + let response = try await self.transport.response( + for: request, + retryPolicy: .transientIdempotent) + + switch response.statusCode { + case 200: + let usageResponse = try JSONDecoder().decode(KiroUsageLimitsResponse.self, from: response.data) + return try usageResponse.toSnapshot() + case 403: + throw URLError(.badServerResponse) + case 401: + throw KiroAPIError.authenticationFailed + case 429: + throw KiroAPIError.rateLimited + default: + throw KiroAPIError.httpError( + statusCode: response.statusCode, + summary: Self.sanitizedResponseBodySummary(response.data)) + } + } + + // MARK: - Helpers + + private static func regionCandidates(for ssoRegion: String?) -> [String] { + guard let ssoRegion else { return self.apiRegions } + if ssoRegion == "eu-central-1" || ssoRegion.hasPrefix("eu-") { + return ["eu-central-1", "us-east-1"] + } + return self.apiRegions + } + + static func machineID(for credentials: KiroCLICredentials) -> String { + if let configured = credentials.machineID, + let normalized = normalizeMachineID(configured) + { + return normalized + } + if let refreshToken = credentials.refreshToken, !refreshToken.isEmpty { + return Self.sha256Hex("KotlinNativeAPI/\(refreshToken)") + } + return Self.sha256Hex("KiroFallback/\(credentials.storageKey)") + } + + static func normalizeMachineID(_ raw: String) -> String? { + let trimmed = raw.trimmingCharacters(in: .whitespacesAndNewlines) + if trimmed.count == 64, trimmed.allSatisfy(\.isHexDigit) { + return trimmed.lowercased() + } + let withoutDashes = trimmed.filter { $0 != "-" } + if withoutDashes.count == 32, withoutDashes.allSatisfy(\.isHexDigit) { + return String(repeating: withoutDashes.lowercased(), count: 2) + } + return nil + } + + private static func sha256Hex(_ input: String) -> String { + let digest = SHA256.hash(data: Data(input.utf8)) + return digest.map { String(format: "%02x", $0) }.joined() + } + + private static func validateRefreshResponse(_ response: ProviderHTTPResponse) throws { + switch response.statusCode { + case 200...299: + return + case 400 where String(data: response.data, encoding: .utf8)?.contains("invalid_grant") == true: + throw KiroAPIError.refreshTokenExpired + case 401: + throw KiroAPIError.authenticationFailed + case 429: + throw KiroAPIError.rateLimited + default: + throw KiroAPIError.httpError( + statusCode: response.statusCode, + summary: self.sanitizedResponseBodySummary(response.data)) + } + } + + private static func decodeJSON(_ data: Data) throws -> [String: Any] { + guard let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] else { + throw KiroAPIError.invalidResponse + } + return json + } + + private static func string(_ json: [String: Any], keys: String...) -> String? { + for key in keys { + guard let value = json[key] as? String else { continue } + let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines) + if !trimmed.isEmpty { return trimmed } + } + return nil + } + + private static func expiresAt(from json: [String: Any]) -> Date? { + if let expiresIn = json["expires_in"] as? Int ?? json["expiresIn"] as? Int { + return Date().addingTimeInterval(TimeInterval(expiresIn)) + } + if let expiresIn = json["expires_in"] as? Double ?? json["expiresIn"] as? Double { + return Date().addingTimeInterval(expiresIn) + } + if let expiresAt = self.string(json, keys: "expires_at", "expiresAt") { + return Self.parseISO8601Date(expiresAt) + } + return nil + } + + static func sanitizedResponseBodySummary(_ data: Data) -> String { + guard !data.isEmpty else { return "empty body" } + guard let rawBody = String(data: data, encoding: .utf8) else { + return "non-text body (\(data.count) bytes)" + } + + let body = Self.redactSensitiveBodyContent(rawBody) + .replacingOccurrences(of: #"\s+"#, with: " ", options: .regularExpression) + .trimmingCharacters(in: .whitespacesAndNewlines) + guard !body.isEmpty else { return "non-text body (\(data.count) bytes)" } + guard body.count > Self.maxErrorBodyLength else { return body } + let index = body.index(body.startIndex, offsetBy: Self.maxErrorBodyLength) + return "\(body[.. String { + let replacements: [(String, String)] = [ + (#"(?i)(bearer\s+)[A-Za-z0-9._\-]+"#, "$1[REDACTED]"), + ( + #"(?i)(\"(?:api_?key|authorization|token|access_token|refresh_token)\"\s*:\s*\")([^\"]+)(\")"#, + "$1[REDACTED]$3"), + ( + #"(?i)((?:api_?key|authorization|token|access_token|refresh_token)\s*[=:]\s*)([^,\s]+)"#, + "$1[REDACTED]"), + ] + return replacements.reduce(text) { partial, replacement in + partial.replacingOccurrences( + of: replacement.0, + with: replacement.1, + options: .regularExpression) + } + } + + #if DEBUG + static func _machineIDForTesting(_ credentials: KiroCLICredentials) -> String { + self.machineID(for: credentials) + } + + static func _sanitizedResponseBodySummaryForTesting(_ body: String) -> String { + self.sanitizedResponseBodySummary(Data(body.utf8)) + } + #endif +} + +extension KiroCLICredentials { + fileprivate func refreshed(accessToken: String, refreshToken: String, expiresAt: Date?) -> KiroCLICredentials { + KiroCLICredentials( + storageKey: self.storageKey, + accessToken: accessToken, + refreshToken: refreshToken, + expiresAt: expiresAt, + region: self.region, + authRegion: self.authRegion, + startURL: self.startURL, + tokenEndpoint: self.tokenEndpoint, + scopes: self.scopes, + clientID: self.clientID, + clientSecret: self.clientSecret, + authMethod: self.authMethod, + provider: self.provider, + machineID: self.machineID) + } +} + +extension [URLQueryItem] { + fileprivate func percentEncoded() -> String { + var components = URLComponents() + components.queryItems = self + return components.percentEncodedQuery ?? "" + } +} + +extension Character { + fileprivate var isHexDigit: Bool { + ("0"..."9").contains(self) || ("a"..."f").contains(self) || ("A"..."F").contains(self) + } +} + +enum KiroAPIError: Error, LocalizedError, Equatable { + case credentialsNotFound + case invalidURL + case invalidResponse + case authenticationFailed + case refreshTokenUnavailable + case refreshTokenExpired + case rateLimited + case allRegionsFailed + case invalidExternalIDPEndpoint + case httpError(statusCode: Int, summary: String) + + var errorDescription: String? { + switch self { + case .credentialsNotFound: + "Kiro credentials not found. Please run 'kiro-cli login' first." + case .invalidURL: + "Invalid Kiro API URL" + case .invalidResponse: + "Invalid response from Kiro API" + case .authenticationFailed: + "Kiro authentication failed. Please run 'kiro-cli login' again." + case .refreshTokenUnavailable: + "Kiro refresh token is unavailable. Please run 'kiro-cli login' again." + case .refreshTokenExpired: + "Kiro session expired. Please run 'kiro-cli login' again." + case .rateLimited: + "Kiro API rate limited. Please try again later." + case .allRegionsFailed: + "All Kiro API regions failed" + case .invalidExternalIDPEndpoint: + "Kiro external IdP endpoint is not allowed" + case let .httpError(statusCode, summary): + "Kiro API error (\(statusCode)): \(summary)" + } + } +} + +// MARK: - API Response Models + +struct KiroUsageLimitsResponse: Decodable { + let nextDateReset: KiroFlexibleResetDate? + let subscriptionInfo: KiroSubscriptionInfo? + let usageBreakdownList: [KiroUsageBreakdown]? + let overageConfiguration: KiroOverageConfiguration? + let userInfo: KiroUserInfo? + + func toSnapshot() throws -> KiroUsageSnapshot { + let planName = self.subscriptionInfo?.subscriptionTitle ?? "KIRO FREE" + let email = self.userInfo?.email + + guard let breakdown = self.usageBreakdownList?.first else { + return KiroUsageSnapshot( + planName: planName, + displayPlanName: KiroStatusProbe.displayPlanName(planName), + accountEmail: email, + authMethod: nil, + creditsUsed: 0, + creditsTotal: 0, + creditsPercent: 0, + bonusCreditsUsed: nil, + bonusCreditsTotal: nil, + bonusExpiryDays: nil, + overagesStatus: self.overageConfiguration?.overageStatus, + overageCreditsUsed: nil, + estimatedOverageCostUSD: nil, + manageURL: nil, + contextUsage: nil, + resetsAt: self.nextDateReset?.date, + updatedAt: Date()) + } + + let creditsUsed = breakdown.currentUsageWithPrecision ?? Double(breakdown.currentUsage ?? 0) + let creditsTotal = breakdown.usageLimitWithPrecision ?? Double(breakdown.usageLimit ?? 0) + let creditsPercent = creditsTotal > 0 ? (creditsUsed / creditsTotal) * 100.0 : 0 + + var bonusUsed: Double? + var bonusTotal: Double? + let activeBonuses = (breakdown.bonuses ?? []).filter { $0.status == "ACTIVE" } + if !activeBonuses.isEmpty { + bonusUsed = activeBonuses.reduce(0) { $0 + ($1.currentUsage ?? 0) } + bonusTotal = activeBonuses.reduce(0) { $0 + ($1.usageLimit ?? 0) } + } + + return KiroUsageSnapshot( + planName: planName, + displayPlanName: KiroStatusProbe.displayPlanName(planName), + accountEmail: email, + authMethod: nil, + creditsUsed: creditsUsed, + creditsTotal: creditsTotal, + creditsPercent: creditsPercent, + bonusCreditsUsed: bonusUsed, + bonusCreditsTotal: bonusTotal, + bonusExpiryDays: nil, + overagesStatus: self.overageConfiguration?.overageStatus, + overageCreditsUsed: nil, + estimatedOverageCostUSD: nil, + manageURL: nil, + contextUsage: nil, + resetsAt: self.nextDateReset?.date, + updatedAt: Date()) + } +} + +struct KiroFlexibleResetDate: Decodable { + let date: Date? + + init(from decoder: Decoder) throws { + let container = try decoder.singleValueContainer() + if container.decodeNil() { + self.date = nil + return + } + if let timestamp = try? container.decode(Double.self) { + self.date = Date(timeIntervalSince1970: timestamp) + return + } + if let stringValue = try? container.decode(String.self) { + self.date = KiroAPIUsageFetcher.parseISO8601Date(stringValue) + return + } + self.date = nil + } +} + +struct KiroSubscriptionInfo: Decodable { + let subscriptionTitle: String? + let overageCapability: String? +} + +struct KiroUsageBreakdown: Decodable { + let currentUsage: Int? + let currentUsageWithPrecision: Double? + let usageLimit: Int? + let usageLimitWithPrecision: Double? + let bonuses: [KiroBonus]? + let freeTrialInfo: KiroFreeTrialInfo? + + enum CodingKeys: String, CodingKey { + case currentUsage + case currentUsageWithPrecision + case usageLimit + case usageLimitWithPrecision + case bonuses + case freeTrialInfo + } + + init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + self.currentUsage = try container.decodeIfPresent(Int.self, forKey: .currentUsage) + self.currentUsageWithPrecision = try container.decodeIfPresent(Double.self, forKey: .currentUsageWithPrecision) + self.usageLimit = try container.decodeIfPresent(Int.self, forKey: .usageLimit) + self.usageLimitWithPrecision = try container.decodeIfPresent(Double.self, forKey: .usageLimitWithPrecision) + self.bonuses = try container.decodeIfPresent([KiroBonus].self, forKey: .bonuses) + self.freeTrialInfo = try container.decodeIfPresent(KiroFreeTrialInfo.self, forKey: .freeTrialInfo) + } +} + +struct KiroBonus: Decodable { + let currentUsage: Double? + let usageLimit: Double? + let status: String? +} + +struct KiroFreeTrialInfo: Decodable { + let currentUsage: Int? + let currentUsageWithPrecision: Double? + let usageLimit: Int? + let usageLimitWithPrecision: Double? + let freeTrialExpiry: Double? + let freeTrialStatus: String? +} + +struct KiroOverageConfiguration: Decodable { + let overageEnabled: Bool? + let overageStatus: String? +} + +struct KiroUserInfo: Decodable { + let email: String? +} + +extension KiroAPIUsageFetcher { + static func parseISO8601Date(_ value: String) -> Date? { + let formatter = ISO8601DateFormatter() + if let date = formatter.date(from: value) { return date } + formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + return formatter.date(from: value) + } +} diff --git a/Sources/CodexBarCore/Providers/Kiro/KiroCLICredentialStore.swift b/Sources/CodexBarCore/Providers/Kiro/KiroCLICredentialStore.swift new file mode 100644 index 0000000000..cc8e132d0f --- /dev/null +++ b/Sources/CodexBarCore/Providers/Kiro/KiroCLICredentialStore.swift @@ -0,0 +1,224 @@ +import Foundation +import SQLite3 + +/// Reads Kiro CLI OAuth material from the local SQLite auth store. +/// +/// ponytail: injectable `databaseURL` keeps filesystem access out of fetch strategies and +/// mirrors `OpenCodeGoLocalUsageReader`. +struct KiroCLICredentialStore: Sendable { + static let enterpriseStorageKeySuffixes = ["oidc:token", "odic:token"] + static let socialStorageKey = "kirocli:social:token" + + private static let storageKeys = [ + "kirocli:odic:token", + "kirocli:oidc:token", + socialStorageKey, + ] + + let databaseURL: URL + + init(homeDirectory: URL = FileManager.default.homeDirectoryForCurrentUser) { + self.databaseURL = homeDirectory + .appendingPathComponent("Library/Application Support/kiro-cli", isDirectory: true) + .appendingPathComponent("data.sqlite3", isDirectory: false) + } + + init(databaseURL: URL) { + self.databaseURL = databaseURL + } + + func loadCredentials(allowSocial: Bool) -> KiroCLICredentials? { + guard FileManager.default.fileExists(atPath: self.databaseURL.path) else { + return nil + } + + var db: OpaquePointer? + guard sqlite3_open_v2(self.databaseURL.path, &db, SQLITE_OPEN_READONLY, nil) == SQLITE_OK else { + sqlite3_close(db) + return nil + } + defer { sqlite3_close(db) } + + for key in Self.storageKeys { + let isSocial = key == Self.socialStorageKey + if isSocial, !allowSocial { continue } + if let credentials = Self.queryCredentials(db: db, storageKey: key) { + if !allowSocial, !credentials.isEnterpriseAffected { continue } + return credentials + } + } + return nil + } + + private static func queryCredentials(db: OpaquePointer?, storageKey: String) -> KiroCLICredentials? { + guard let db else { return nil } + let query = "SELECT value FROM auth_kv WHERE key = ?" + var stmt: OpaquePointer? + guard sqlite3_prepare_v2(db, query, -1, &stmt, nil) == SQLITE_OK else { + sqlite3_finalize(stmt) + return nil + } + defer { sqlite3_finalize(stmt) } + + sqlite3_bind_text(stmt, 1, storageKey, -1, unsafeBitCast(-1, to: sqlite3_destructor_type.self)) + + guard sqlite3_step(stmt) == SQLITE_ROW, + let blob = sqlite3_column_text(stmt, 0) + else { + return nil + } + + let jsonString = String(cString: blob) + guard let data = jsonString.data(using: .utf8), + let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] + else { + return nil + } + + guard let accessToken = Self.string(json, keys: "access_token", "accessToken"), + !accessToken.isEmpty + else { + return nil + } + + return KiroCLICredentials( + storageKey: storageKey, + accessToken: accessToken, + refreshToken: Self.string(json, keys: "refresh_token", "refreshToken"), + expiresAt: Self.parseExpires(json), + region: Self.string(json, keys: "region"), + authRegion: Self.string(json, keys: "auth_region", "authRegion"), + startURL: Self.string(json, keys: "start_url", "startUrl"), + tokenEndpoint: Self.string(json, keys: "token_endpoint", "tokenEndpoint"), + scopes: Self.string(json, keys: "scopes", "scope"), + clientID: Self.string(json, keys: "client_id", "clientId"), + clientSecret: Self.string(json, keys: "client_secret", "clientSecret"), + authMethod: Self.string(json, keys: "auth_method", "authMethod"), + provider: Self.string(json, keys: "provider"), + machineID: Self.string(json, keys: "machine_id", "machineId")) + } + + private static func string(_ json: [String: Any], keys: String...) -> String? { + for key in keys { + guard let value = json[key] as? String else { continue } + let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines) + if !trimmed.isEmpty { return trimmed } + } + return nil + } + + private static func parseExpires(_ json: [String: Any]) -> Date? { + if let expiresAt = self.string(json, keys: "expires_at", "expiresAt") { + return parseISO8601Date(expiresAt) + } + if let expiresIn = json["expires_in"] as? Int ?? json["expiresIn"] as? Int { + return Date().addingTimeInterval(TimeInterval(expiresIn)) + } + if let expiresIn = json["expires_in"] as? Double ?? json["expiresIn"] as? Double { + return Date().addingTimeInterval(expiresIn) + } + return nil + } +} + +struct KiroCLICredentials: Sendable, Equatable { + let storageKey: String + let accessToken: String + let refreshToken: String? + let expiresAt: Date? + let region: String? + let authRegion: String? + let startURL: String? + let tokenEndpoint: String? + let scopes: String? + let clientID: String? + let clientSecret: String? + let authMethod: String? + let provider: String? + let machineID: String? + + var canonicalAuthMethod: String { + Self.canonicalizeAuthMethod(self.authMethod, tokenEndpoint: self.tokenEndpoint) + } + + var isExternalIDP: Bool { + self.canonicalAuthMethod == "external_idp" + } + + var isEnterpriseAffected: Bool { + if KiroCLICredentialStore.enterpriseStorageKeySuffixes.contains(where: { self.storageKey.hasSuffix($0) }) { + return true + } + let method = self.canonicalAuthMethod + if method == "idc" || method == "external_idp" { return true } + if let startURL = self.startURL, !startURL.isEmpty { return true } + if let tokenEndpoint = self.tokenEndpoint, !tokenEndpoint.isEmpty { return true } + return false + } + + var tokenTypeHeader: String? { + if self.canonicalAuthMethod == "api_key" { return "API_KEY" } + if self.isExternalIDP { return "EXTERNAL_IDP" } + return nil + } + + var effectiveAuthRegion: String { + self.authRegion ?? self.region ?? "us-east-1" + } + + func needsRefresh(now: Date = Date(), buffer: TimeInterval = 5 * 60) -> Bool { + guard let expiresAt else { return true } + return expiresAt.timeIntervalSince(now) <= buffer + } + + static func canonicalizeAuthMethod(_ raw: String?, tokenEndpoint: String?) -> String { + let trimmed = raw?.trimmingCharacters(in: .whitespacesAndNewlines) ?? "" + let normalized: String = switch trimmed.lowercased() { + case "builder-id", "iam": "idc" + case "api_key", "apikey": "api_key" + case "external_idp", "azuread", "azure", "entra", "entra-id", "microsoft", "m365", "office365", "external": + "external_idp" + case "": "" + default: trimmed.lowercased() + } + + if normalized == "external_idp" { return "external_idp" } + if let tokenEndpoint, !tokenEndpoint.isEmpty { return "external_idp" } + return normalized.isEmpty ? "social" : normalized + } + + static func validateExternalIDPEndpoint(_ rawURL: String) throws { + guard let url = URL(string: rawURL.trimmingCharacters(in: .whitespacesAndNewlines)), + let host = url.host?.lowercased(), + !host.isEmpty + else { + throw KiroAPIError.invalidExternalIDPEndpoint + } + + guard url.scheme?.lowercased() == "https" else { + throw KiroAPIError.invalidExternalIDPEndpoint + } + + if host.contains(":") || host.split(separator: ".").allSatisfy({ $0.allSatisfy(\.isNumber) }) { + throw KiroAPIError.invalidExternalIDPEndpoint + } + + let allowedSuffixes = [ + ".microsoftonline.com", + ".microsoftonline.us", + ".microsoftonline.cn", + ] + guard allowedSuffixes.contains(where: { host.hasSuffix($0) }) else { + throw KiroAPIError.invalidExternalIDPEndpoint + } + } +} + +extension KiroCLICredentialStore { + fileprivate static func parseISO8601Date(_ value: String) -> Date? { + let formatter = ISO8601DateFormatter() + if let date = formatter.date(from: value) { return date } + formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + return formatter.date(from: value) + } +} diff --git a/Sources/CodexBarCore/Providers/Kiro/KiroProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Kiro/KiroProviderDescriptor.swift index 804b346c08..a9651cc9f8 100644 --- a/Sources/CodexBarCore/Providers/Kiro/KiroProviderDescriptor.swift +++ b/Sources/CodexBarCore/Providers/Kiro/KiroProviderDescriptor.swift @@ -31,8 +31,19 @@ public enum KiroProviderDescriptor { supportsTokenCost: false, noDataMessage: { "Kiro cost summary is not supported." }), fetchPlan: ProviderFetchPlan( - sourceModes: [.auto, .cli], - pipeline: ProviderFetchPipeline(resolveStrategies: { _ in [KiroCLIFetchStrategy()] })), + sourceModes: [.auto, .api, .cli], + pipeline: ProviderFetchPipeline(resolveStrategies: { context in + switch context.sourceMode { + case .auto: + [KiroAPIFetchStrategy(), KiroCLIFetchStrategy()] + case .api: + [KiroAPIFetchStrategy()] + case .cli: + [KiroCLIFetchStrategy()] + case .web, .oauth: + [] + } + })), cli: ProviderCLIConfig( name: "kiro", aliases: ["kiro-cli"], @@ -56,7 +67,7 @@ struct KiroCLIFetchStrategy: ProviderFetchStrategy { sourceLabel: "cli") } - func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { - false + func shouldFallback(on _: Error, context: ProviderFetchContext) -> Bool { + context.sourceMode == .auto } } diff --git a/Tests/CodexBarTests/KiroAPIUsageFetcherTests.swift b/Tests/CodexBarTests/KiroAPIUsageFetcherTests.swift new file mode 100644 index 0000000000..9f82e5d656 --- /dev/null +++ b/Tests/CodexBarTests/KiroAPIUsageFetcherTests.swift @@ -0,0 +1,384 @@ +import Foundation +import SQLite3 +import Testing +@testable import CodexBarCore + +@Suite(.serialized) +struct KiroAPIUsageFetcherTests { + private struct StubClaudeFetcher: ClaudeUsageFetching { + func loadLatestUsage(model _: String) async throws -> ClaudeUsageSnapshot { + throw ClaudeUsageError.parseFailed("stub") + } + + func debugRawProbe(model _: String) async -> String { + "stub" + } + + func detectVersion() -> String? { + nil + } + } + + private func makeFetchContext(sourceMode: ProviderSourceMode) -> ProviderFetchContext { + let env: [String: String] = [:] + return ProviderFetchContext( + runtime: .cli, + sourceMode: sourceMode, + includeCredits: false, + webTimeout: 30, + webDebugDumpHTML: false, + verbose: false, + env: env, + settings: ProviderSettingsSnapshot.make(), + fetcher: UsageFetcher(environment: env), + claudeFetcher: StubClaudeFetcher(), + browserDetection: BrowserDetection(cacheTTL: 0)) + } + + @Test + func `auto mode ignores social credentials`() throws { + let env = try self.makeCredentialEnvironment( + key: KiroCLICredentialStore.socialStorageKey, + json: Self.socialCredentialJSON) + defer { try? FileManager.default.removeItem(at: env.root) } + + let store = KiroCLICredentialStore(databaseURL: env.databaseURL) + #expect(store.loadCredentials(allowSocial: false) == nil) + #expect(store.loadCredentials(allowSocial: true) != nil) + } + + @Test + func `credential keys prefer enterprise oidc before social`() throws { + let env = try self.makeCredentialEnvironment( + entries: [ + (KiroCLICredentialStore.socialStorageKey, Self.socialCredentialJSON), + ("kirocli:oidc:token", Self.idcCredentialJSON), + ]) + defer { try? FileManager.default.removeItem(at: env.root) } + + let store = KiroCLICredentialStore(databaseURL: env.databaseURL) + let credentials = store.loadCredentials(allowSocial: false) + #expect(credentials?.storageKey == "kirocli:oidc:token") + #expect(credentials?.isEnterpriseAffected == true) + } + + @Test + func `odic typo key remains supported`() throws { + let env = try self.makeCredentialEnvironment( + key: "kirocli:odic:token", + json: Self.idcCredentialJSON) + defer { try? FileManager.default.removeItem(at: env.root) } + + let credentials = KiroCLICredentialStore(databaseURL: env.databaseURL) + .loadCredentials(allowSocial: false) + #expect(credentials?.storageKey == "kirocli:odic:token") + } + + @Test + func `external idp auth method sends tokentype header`() async throws { + let env = try self.makeCredentialEnvironment( + key: "kirocli:oidc:token", + json: Self.externalIDPCredentialJSON) + defer { try? FileManager.default.removeItem(at: env.root) } + + let transport = ProviderHTTPTransportStub { request in + #expect(request.value(forHTTPHeaderField: "tokentype") == "EXTERNAL_IDP") + #expect(request.url?.absoluteString.contains("profileArn") == false) + return try Self.httpResponse(data: Self.usageFixture(reset: 1_700_000_000), statusCode: 200) + } + + let fetcher = KiroAPIUsageFetcher( + credentialStore: KiroCLICredentialStore(databaseURL: env.databaseURL), + transport: transport) + let snapshot = try await fetcher.fetchUsage(allowSocial: false) + #expect(snapshot.planName == "KIRO POWER") + #expect(snapshot.creditsUsed == 12.5) + } + + @Test + func `idc credentials do not send tokentype header`() async throws { + let env = try self.makeCredentialEnvironment( + key: "kirocli:oidc:token", + json: Self.idcCredentialJSON) + defer { try? FileManager.default.removeItem(at: env.root) } + + let transport = ProviderHTTPTransportStub { request in + #expect(request.value(forHTTPHeaderField: "tokentype") == nil) + return try Self.httpResponse(data: Self.usageFixture(reset: "2026-07-01T00:00:00.000Z"), statusCode: 200) + } + + let fetcher = KiroAPIUsageFetcher( + credentialStore: KiroCLICredentialStore(databaseURL: env.databaseURL), + transport: transport) + let snapshot = try await fetcher.fetchUsage(allowSocial: false) + #expect(snapshot.resetsAt != nil) + } + + @Test + func `region fallback tries second endpoint after 403`() async throws { + let env = try self.makeCredentialEnvironment( + key: "kirocli:oidc:token", + json: Self.idcCredentialJSON) + defer { try? FileManager.default.removeItem(at: env.root) } + + let transport = ProviderHTTPTransportStub { request in + guard let host = request.url?.host else { + return try Self.httpResponse(data: Data(), statusCode: 500) + } + if host == "q.us-east-1.amazonaws.com" { + return try Self.httpResponse(data: Data("forbidden".utf8), statusCode: 403) + } + return try Self.httpResponse(data: Self.usageFixture(reset: 1_700_000_000), statusCode: 200) + } + + let fetcher = KiroAPIUsageFetcher( + credentialStore: KiroCLICredentialStore(databaseURL: env.databaseURL), + transport: transport) + let snapshot = try await fetcher.fetchUsage(allowSocial: false) + #expect(snapshot.creditsTotal == 50) + } + + @Test + func `expired idc credentials refresh before usage fetch`() async throws { + let env = try self.makeCredentialEnvironment( + key: "kirocli:oidc:token", + json: Self.expiredIDCCredentialJSON) + defer { try? FileManager.default.removeItem(at: env.root) } + + let transport = ProviderHTTPTransportStub { request in + if request.url?.host?.hasPrefix("oidc.") == true { + let body = """ + {"accessToken":"fresh-token","refreshToken":"refresh-abc","expiresIn":3600} + """ + return try Self.httpResponse(data: Data(body.utf8), statusCode: 200) + } + #expect(request.value(forHTTPHeaderField: "Authorization") == "Bearer fresh-token") + return try Self.httpResponse(data: Self.usageFixture(reset: 1_700_000_000), statusCode: 200) + } + + let fetcher = KiroAPIUsageFetcher( + credentialStore: KiroCLICredentialStore(databaseURL: env.databaseURL), + transport: transport) + _ = try await fetcher.fetchUsage(allowSocial: false) + } + + @Test + func `refresh failure surfaces authentication error`() async throws { + let env = try self.makeCredentialEnvironment( + key: "kirocli:oidc:token", + json: Self.expiredIDCCredentialJSON) + defer { try? FileManager.default.removeItem(at: env.root) } + + let transport = ProviderHTTPTransportStub { request in + if request.url?.host?.hasPrefix("oidc.") == true { + return try Self.httpResponse(data: Data("{\"error\":\"invalid_grant\"}".utf8), statusCode: 400) + } + return try Self.httpResponse(data: Data(), statusCode: 500) + } + + let fetcher = KiroAPIUsageFetcher( + credentialStore: KiroCLICredentialStore(databaseURL: env.databaseURL), + transport: transport) + + await #expect(throws: KiroAPIError.refreshTokenExpired) { + _ = try await fetcher.fetchUsage(allowSocial: false) + } + } + + @Test + func `machine id derives from refresh token instead of hardware uuid`() { + let credentials = KiroCLICredentials( + storageKey: "kirocli:oidc:token", + accessToken: "token", + refreshToken: "refresh-token-value", + expiresAt: nil, + region: "us-east-1", + authRegion: nil, + startURL: nil, + tokenEndpoint: nil, + scopes: nil, + clientID: nil, + clientSecret: nil, + authMethod: "idc", + provider: nil, + machineID: nil) + + let machineID = KiroAPIUsageFetcher._machineIDForTesting(credentials) + #expect(machineID.count == 64) + #expect(machineID == KiroAPIUsageFetcher._machineIDForTesting(credentials)) + } + + @Test + func `http errors redact sensitive response bodies`() { + let summary = KiroAPIUsageFetcher._sanitizedResponseBodySummaryForTesting( + #"{"access_token":"secret-token","message":"bad request"}"#) + #expect(summary.contains("secret-token") == false) + #expect(summary.contains("[REDACTED]")) + } + + @Test + func `provider descriptor respects explicit source modes`() async { + let descriptor = KiroProviderDescriptor.descriptor + let autoContext = self.makeFetchContext(sourceMode: .auto) + + let apiStrategies = await descriptor.fetchPlan.pipeline.resolveStrategies( + self.makeFetchContext(sourceMode: .api)) + #expect(apiStrategies.count == 1) + #expect(apiStrategies[0].id == "kiro.api") + + let cliStrategies = await descriptor.fetchPlan.pipeline.resolveStrategies( + self.makeFetchContext(sourceMode: .cli)) + #expect(cliStrategies.count == 1) + #expect(cliStrategies[0].id == "kiro.cli") + + let autoStrategies = await descriptor.fetchPlan.pipeline.resolveStrategies(autoContext) + #expect(autoStrategies.count == 2) + #expect(autoStrategies[0].id == "kiro.api") + #expect(autoStrategies[1].id == "kiro.cli") + } + + @Test + func `api strategy only falls back in auto mode`() { + let strategy = KiroAPIFetchStrategy() + let autoContext = self.makeFetchContext(sourceMode: .auto) + let apiContext = self.makeFetchContext(sourceMode: .api) + + #expect(strategy.shouldFallback(on: KiroAPIError.authenticationFailed, context: autoContext)) + #expect(!strategy.shouldFallback(on: KiroAPIError.authenticationFailed, context: apiContext)) + } + + // MARK: - Fixtures + + private static let socialCredentialJSON = """ + {"access_token":"social-token","refresh_token":"\(String(repeating: "s", count: 120))","auth_method":"social"} + """ + + private static let idcCredentialJSON = """ + { + "access_token": "idc-token", + "refresh_token": "\(String(repeating: "r", count: 120))", + "auth_method": "idc", + "region": "us-east-1", + "client_id": "client-id", + "client_secret": "client-secret", + "expires_at": "2099-01-01T00:00:00Z" + } + """ + + private static let expiredIDCCredentialJSON = """ + { + "access_token": "stale-token", + "refresh_token": "\(String(repeating: "r", count: 120))", + "auth_method": "idc", + "region": "us-east-1", + "client_id": "client-id", + "client_secret": "client-secret", + "expires_at": "2020-01-01T00:00:00Z" + } + """ + + private static let externalIDPCredentialJSON = """ + { + "access_token": "external-token", + "refresh_token": "\(String(repeating: "e", count: 120))", + "auth_method": "external_idp", + "client_id": "client-id", + "token_endpoint": "https://login.microsoftonline.com/tenant/oauth2/v2.0/token", + "scopes": "openid offline_access", + "expires_at": "2099-01-01T00:00:00Z" + } + """ + + private static func usageFixture(reset: Double) -> Data { + self.usageFixture(resetJSON: String(reset)) + } + + private static func usageFixture(reset: String) -> Data { + self.usageFixture(resetJSON: "\"\(reset)\"") + } + + private static func usageFixture(resetJSON: String) -> Data { + let body = """ + { + "nextDateReset": \(resetJSON), + "subscriptionInfo": { "subscriptionTitle": "KIRO POWER" }, + "usageBreakdownList": [{ + "currentUsage": 12, + "currentUsageWithPrecision": 12.5, + "usageLimit": 50, + "usageLimitWithPrecision": 50.0 + }] + } + """ + return Data(body.utf8) + } + + private static func httpResponse(data: Data, statusCode: Int) throws -> (Data, URLResponse) { + let response = HTTPURLResponse( + url: URL(string: "https://q.us-east-1.amazonaws.com/getUsageLimits")!, + statusCode: statusCode, + httpVersion: "HTTP/1.1", + headerFields: nil)! + return (data, response) + } + + private struct CredentialEnvironment { + let root: URL + let databaseURL: URL + } + + private func makeCredentialEnvironment( + key: String, + json: String) throws -> CredentialEnvironment + { + try self.makeCredentialEnvironment(entries: [(key, json)]) + } + + private func makeCredentialEnvironment(entries: [(String, String)]) throws -> CredentialEnvironment { + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("KiroAPIUsageFetcherTests-\(UUID().uuidString)", isDirectory: true) + try FileManager.default.createDirectory(at: root, withIntermediateDirectories: true) + let databaseURL = root.appendingPathComponent("data.sqlite3") + + var db: OpaquePointer? + guard sqlite3_open(databaseURL.path, &db) == SQLITE_OK else { + throw CredentialTestError.sqliteOpen + } + defer { sqlite3_close(db) } + + let createSQL = "CREATE TABLE auth_kv (key TEXT PRIMARY KEY, value TEXT NOT NULL);" + guard sqlite3_exec(db, createSQL, nil, nil, nil) == SQLITE_OK else { + throw CredentialTestError.sqliteExec + } + + for (key, json) in entries { + var stmt: OpaquePointer? + guard sqlite3_prepare_v2( + db, + "INSERT INTO auth_kv(key, value) VALUES(?, ?);", + -1, + &stmt, + nil) == SQLITE_OK + else { + throw CredentialTestError.sqlitePrepare + } + defer { sqlite3_finalize(stmt) } + + let transient = unsafeBitCast(-1, to: sqlite3_destructor_type.self) + sqlite3_bind_text(stmt, 1, key, -1, transient) + sqlite3_bind_text(stmt, 2, json, -1, transient) + guard sqlite3_step(stmt) == SQLITE_DONE else { + throw CredentialTestError.sqliteStep + } + } + + return CredentialEnvironment(root: root, databaseURL: databaseURL) + } + + private enum CredentialTestError: Error { + case sqliteOpen + case sqliteExec + case sqlitePrepare + case sqliteStep + } +}