337 lines
16 KiB
Swift
337 lines
16 KiB
Swift
// Copyright (C) 2026 ysicing
|
|
// SPDX-License-Identifier: AGPL-3.0-or-later
|
|
|
|
import Foundation
|
|
import LocalAuthentication
|
|
import Security
|
|
|
|
public enum Sub2APIConfigurationError: Error, Equatable {
|
|
case invalidAddress, invalidAccountID, missingCredentials, providerMismatch, keychain(OSStatus)
|
|
}
|
|
|
|
public struct Sub2APIConfiguration: Equatable, Sendable {
|
|
public let provider: AIProviderID
|
|
public let baseURL: URL
|
|
public let email: String
|
|
public let password: String
|
|
public let accountID: Int
|
|
|
|
public init(address: String, email: String, password: String, accountID: String,
|
|
provider: AIProviderID = .codex) throws {
|
|
let trimmedAddress = address.trimmingCharacters(in: .whitespacesAndNewlines)
|
|
guard var parts = URLComponents(string: trimmedAddress),
|
|
parts.scheme?.lowercased() == "https", parts.host?.isEmpty == false,
|
|
parts.user == nil, parts.password == nil, parts.query == nil, parts.fragment == nil else {
|
|
throw Sub2APIConfigurationError.invalidAddress
|
|
}
|
|
let prefix = parts.path.trimmingCharacters(in: CharacterSet(charactersIn: "/"))
|
|
parts.path = prefix.isEmpty ? "" : "/" + prefix
|
|
guard let baseURL = parts.url else { throw Sub2APIConfigurationError.invalidAddress }
|
|
let trimmedAccount = accountID.trimmingCharacters(in: .whitespacesAndNewlines)
|
|
guard let id = Int(trimmedAccount), id > 0 else { throw Sub2APIConfigurationError.invalidAccountID }
|
|
let trimmedEmail = email.trimmingCharacters(in: .whitespacesAndNewlines)
|
|
guard !trimmedEmail.isEmpty, !password.isEmpty else { throw Sub2APIConfigurationError.missingCredentials }
|
|
self.provider = provider
|
|
self.baseURL = baseURL
|
|
self.email = trimmedEmail
|
|
self.password = password
|
|
self.accountID = id
|
|
}
|
|
|
|
func endpoint(_ path: String) -> URL { baseURL.appending(path: path) }
|
|
}
|
|
|
|
@MainActor
|
|
public protocol Sub2APISecretStore {
|
|
func read() throws -> String?
|
|
func write(_ password: String) throws
|
|
func delete() throws
|
|
}
|
|
|
|
/// 管理员密码只保存在本机钥匙串；写入失败时不提交其余配置。
|
|
@MainActor
|
|
public struct KeychainSub2APISecretStore: Sub2APISecretStore {
|
|
private let provider: AIProviderID
|
|
|
|
public init(provider: AIProviderID = .codex) { self.provider = provider }
|
|
|
|
private var query: [String: Any] {
|
|
[kSecClass as String: kSecClassGenericPassword,
|
|
kSecAttrService as String: "work.12306.xstats.sub2api",
|
|
kSecAttrAccount as String: provider == .codex ? "admin" : "admin.claude"]
|
|
}
|
|
|
|
public func read() throws -> String? {
|
|
var item = query
|
|
item[kSecReturnData as String] = true
|
|
item[kSecMatchLimit as String] = kSecMatchLimitOne
|
|
let context = LAContext()
|
|
context.interactionNotAllowed = true
|
|
item[kSecUseAuthenticationContext as String] = context
|
|
var result: AnyObject?
|
|
let status = SecItemCopyMatching(item as CFDictionary, &result)
|
|
if status == errSecItemNotFound { return nil }
|
|
guard status == errSecSuccess else { throw Sub2APIConfigurationError.keychain(status) }
|
|
guard let data = result as? Data, let value = String(data: data, encoding: .utf8) else {
|
|
throw Sub2APIConfigurationError.keychain(errSecDecode)
|
|
}
|
|
return value
|
|
}
|
|
|
|
public func write(_ password: String) throws {
|
|
let data = Data(password.utf8)
|
|
var status = SecItemUpdate(query as CFDictionary, [kSecValueData as String: data] as CFDictionary)
|
|
if status == errSecItemNotFound {
|
|
var item = query
|
|
item[kSecValueData as String] = data
|
|
item[kSecAttrLabel as String] = "XStats Sub2API"
|
|
status = SecItemAdd(item as CFDictionary, nil)
|
|
}
|
|
guard status == errSecSuccess else { throw Sub2APIConfigurationError.keychain(status) }
|
|
}
|
|
|
|
public func delete() throws {
|
|
let status = SecItemDelete(query as CFDictionary)
|
|
guard status == errSecSuccess || status == errSecItemNotFound else {
|
|
throw Sub2APIConfigurationError.keychain(status)
|
|
}
|
|
}
|
|
}
|
|
|
|
@MainActor
|
|
public struct Sub2APISettingsStore {
|
|
private let provider: AIProviderID
|
|
private let defaults: UserDefaults
|
|
private let secrets: any Sub2APISecretStore
|
|
private var key: String { provider == .codex ? "aiUsage.sub2api" : "aiUsage.sub2api.claude" }
|
|
|
|
public init(provider: AIProviderID = .codex, defaults: UserDefaults = .standard,
|
|
secrets: (any Sub2APISecretStore)? = nil) {
|
|
self.provider = provider
|
|
self.defaults = defaults
|
|
self.secrets = secrets ?? KeychainSub2APISecretStore(provider: provider)
|
|
}
|
|
|
|
public func hasConfiguration() -> Bool { defaults.dictionary(forKey: key) != nil }
|
|
|
|
public func load() throws -> Sub2APIConfiguration? {
|
|
guard let saved = defaults.dictionary(forKey: key) else { return nil }
|
|
guard let address = saved["address"] as? String,
|
|
let email = saved["email"] as? String,
|
|
let accountID = saved["accountID"] as? Int else {
|
|
throw Sub2APIConfigurationError.missingCredentials
|
|
}
|
|
guard let password = try secrets.read() else { throw Sub2APIConfigurationError.missingCredentials }
|
|
return try Sub2APIConfiguration(address: address, email: email,
|
|
password: password, accountID: String(accountID), provider: provider)
|
|
}
|
|
|
|
public func save(_ configuration: Sub2APIConfiguration) throws {
|
|
guard configuration.provider == provider else { throw Sub2APIConfigurationError.providerMismatch }
|
|
try secrets.write(configuration.password)
|
|
guard try secrets.read() == configuration.password else {
|
|
throw Sub2APIConfigurationError.missingCredentials
|
|
}
|
|
defaults.set(["address": configuration.baseURL.absoluteString,
|
|
"email": configuration.email, "accountID": configuration.accountID], forKey: key)
|
|
}
|
|
|
|
public func clear() throws {
|
|
try secrets.delete()
|
|
defaults.removeObject(forKey: key)
|
|
}
|
|
}
|
|
|
|
protocol Sub2APIQuotaFetching: Sendable {
|
|
func fetch(configuration: Sub2APIConfiguration) async throws -> AIQuotaSnapshot
|
|
}
|
|
|
|
public actor Sub2APIQuotaClient: Sub2APIQuotaFetching {
|
|
private struct AuthIdentity: Hashable {
|
|
let baseURL: URL
|
|
let email: String
|
|
let password: String
|
|
|
|
init(_ configuration: Sub2APIConfiguration) {
|
|
baseURL = configuration.baseURL
|
|
email = configuration.email
|
|
password = configuration.password
|
|
}
|
|
}
|
|
|
|
public static let shared = Sub2APIQuotaClient()
|
|
private let http: any QuotaHTTPClient
|
|
private var sessions: [AuthIdentity: (token: String, expiry: Date)] = [:]
|
|
private var loginTasks: [AuthIdentity: Task<(String, Date), Error>] = [:]
|
|
private var loginMarkers: [AuthIdentity: UUID] = [:]
|
|
|
|
public init() { http = URLSessionQuotaHTTPClient() }
|
|
|
|
init(http: any QuotaHTTPClient) { self.http = http }
|
|
|
|
public func clearSession() {
|
|
for task in loginTasks.values { task.cancel() }
|
|
loginTasks.removeAll()
|
|
loginMarkers.removeAll()
|
|
sessions.removeAll()
|
|
}
|
|
|
|
public func fetch(configuration: Sub2APIConfiguration) async throws -> AIQuotaSnapshot {
|
|
let token = try await validToken(configuration)
|
|
let usageURL = configuration.endpoint("api/v1/admin/accounts/\(configuration.accountID)/usage")
|
|
let accountURL = configuration.endpoint("api/v1/admin/accounts/\(configuration.accountID)")
|
|
let accountData = try await get(accountURL, token: token, configuration: configuration)
|
|
guard let account = Self.envelope(accountData),
|
|
let platform = account["platform"] as? String else {
|
|
throw AIQuotaFailure.sub2apiInvalidResponse
|
|
}
|
|
let expectedPlatform = configuration.provider == .codex ? "openai" : "anthropic"
|
|
guard platform == expectedPlatform else { throw AIQuotaFailure.sub2apiAccountMismatch }
|
|
var parts = URLComponents(url: usageURL, resolvingAgainstBaseURL: false)!
|
|
parts.queryItems = [URLQueryItem(name: "source", value: "active"),
|
|
URLQueryItem(name: "force", value: "false"),
|
|
URLQueryItem(name: "timezone", value: TimeZone.current.identifier)]
|
|
do {
|
|
let data = try await get(parts.url!, token: token, configuration: configuration)
|
|
let snapshot = try Self.parseUsage(data, now: Date(), provider: configuration.provider)
|
|
// Claude 的 Fable 周窗口可能只存在于账号 extra，详细接口缺失时回填。
|
|
if configuration.provider == .claude, snapshot.window(.fableWeekly) == nil,
|
|
let fallback = try? Self.parseAccount(accountData, now: Date(), provider: .claude),
|
|
let extra = fallback.window(.fableWeekly) {
|
|
return AIQuotaSnapshot(provider: .claude, windows: snapshot.windows + [extra],
|
|
fetchedAt: snapshot.fetchedAt, source: .sub2api)
|
|
}
|
|
return snapshot
|
|
} catch AIQuotaFailure.sub2apiInvalidResponse {
|
|
// 较旧的 sub2api 没有 /usage 时，账号 extra 仍可提供较粗的窗口数据。
|
|
return try Self.parseAccount(accountData, now: Date(), provider: configuration.provider)
|
|
}
|
|
}
|
|
|
|
private func validToken(_ configuration: Sub2APIConfiguration) async throws -> String {
|
|
let identity = AuthIdentity(configuration)
|
|
if let session = sessions[identity], session.expiry > Date() { return session.token }
|
|
if let task = loginTasks[identity] { return try await task.value.0 }
|
|
let marker = UUID()
|
|
let task = Task { try await login(configuration) }
|
|
loginTasks[identity] = task
|
|
loginMarkers[identity] = marker
|
|
defer {
|
|
if loginMarkers[identity] == marker {
|
|
loginTasks[identity] = nil
|
|
loginMarkers[identity] = nil
|
|
}
|
|
}
|
|
let (token, expiry) = try await task.value
|
|
if loginMarkers[identity] == marker { sessions[identity] = (token, expiry) }
|
|
return token
|
|
}
|
|
|
|
private func login(_ configuration: Sub2APIConfiguration) async throws -> (String, Date) {
|
|
var request = URLRequest(url: configuration.endpoint("api/v1/auth/login"))
|
|
request.httpMethod = "POST"
|
|
request.setValue("application/json", forHTTPHeaderField: "Content-Type")
|
|
request.httpBody = try JSONSerialization.data(withJSONObject: ["email": configuration.email,
|
|
"password": configuration.password])
|
|
let (data, status) = try await send(request)
|
|
// 只有明确的拒绝才算凭据错误；服务故障不能清空上次额度，也不能误导用户去改密码。
|
|
switch status {
|
|
case 200..<300: break
|
|
case 400, 401, 403: throw AIQuotaFailure.sub2apiUnauthorized
|
|
case 404: throw AIQuotaFailure.sub2apiInvalidResponse
|
|
default: throw AIQuotaFailure.sub2apiNetwork
|
|
}
|
|
guard let payload = Self.envelope(data),
|
|
let token = payload["access_token"] as? String, !token.isEmpty else {
|
|
throw AIQuotaFailure.sub2apiUnauthorized
|
|
}
|
|
let seconds = (payload["expires_in"] as? NSNumber)?.doubleValue ?? 600
|
|
return (token, Date().addingTimeInterval(max(0, seconds - 60)))
|
|
}
|
|
|
|
private func get(_ url: URL, token: String, configuration: Sub2APIConfiguration) async throws -> Data {
|
|
var token = token
|
|
for attempt in 0...1 {
|
|
var request = URLRequest(url: url)
|
|
request.timeoutInterval = url.path.hasSuffix("/usage") ? 20 : 8
|
|
request.setValue("Bearer \(token)", forHTTPHeaderField: "Authorization")
|
|
request.setValue("application/json", forHTTPHeaderField: "Accept")
|
|
let (data, status) = try await send(request)
|
|
if status == 401, attempt == 0 {
|
|
sessions[AuthIdentity(configuration)] = nil
|
|
token = try await validToken(configuration)
|
|
continue
|
|
}
|
|
if status == 401 || status == 403 { throw AIQuotaFailure.sub2apiUnauthorized }
|
|
if status == 404 { throw AIQuotaFailure.sub2apiInvalidResponse }
|
|
guard (200..<300).contains(status) else { throw AIQuotaFailure.sub2apiNetwork }
|
|
return data
|
|
}
|
|
throw AIQuotaFailure.sub2apiUnauthorized
|
|
}
|
|
|
|
private func send(_ request: URLRequest) async throws -> (Data, Int) {
|
|
do { return try await http.send(request) }
|
|
catch is CancellationError { throw CancellationError() }
|
|
catch { throw AIQuotaFailure.sub2apiNetwork }
|
|
}
|
|
|
|
private static func envelope(_ data: Data) -> [String: Any]? {
|
|
guard let object = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
|
|
(object["code"] as? NSNumber)?.intValue == 0 else { return nil }
|
|
return object["data"] as? [String: Any]
|
|
}
|
|
|
|
static func parseUsage(_ data: Data, now: Date, provider: AIProviderID = .codex) throws -> AIQuotaSnapshot {
|
|
guard let payload = envelope(data) else { throw AIQuotaFailure.sub2apiInvalidResponse }
|
|
var windows: [AIQuotaWindow] = []
|
|
var fields: [(String, AIQuotaKind)] = [("five_hour", .session), ("seven_day", .weekly)]
|
|
if provider == .claude {
|
|
fields += [("seven_day_sonnet", .sonnetWeekly), ("seven_day_fable", .fableWeekly)]
|
|
}
|
|
for (key, kind) in fields {
|
|
guard let item = payload[key] as? [String: Any],
|
|
let used = (item["utilization"] as? NSNumber)?.doubleValue,
|
|
used.isFinite, (0...100).contains(used) else { continue }
|
|
let reset = (item["resets_at"] as? String).flatMap(parseDate)
|
|
?? (item["remaining_seconds"] as? NSNumber).map { now.addingTimeInterval($0.doubleValue) }
|
|
windows.append(AIQuotaWindow(kind: kind, usedPercent: used, resetsAt: reset))
|
|
}
|
|
guard !windows.isEmpty else { throw AIQuotaFailure.sub2apiInvalidResponse }
|
|
return AIQuotaSnapshot(provider: provider, windows: windows, fetchedAt: now, source: .sub2api)
|
|
}
|
|
|
|
static func parseAccount(_ data: Data, now: Date, provider: AIProviderID = .codex) throws -> AIQuotaSnapshot {
|
|
guard let payload = envelope(data), let extra = payload["extra"] as? [String: Any] else {
|
|
throw AIQuotaFailure.sub2apiInvalidResponse
|
|
}
|
|
var windows: [AIQuotaWindow] = []
|
|
if let fraction = (extra["session_window_utilization"] as? NSNumber)?.doubleValue,
|
|
fraction.isFinite, (0...1).contains(fraction) {
|
|
windows.append(AIQuotaWindow(kind: .session, usedPercent: fraction * 100,
|
|
resetsAt: (payload["session_window_end"] as? String).flatMap(parseDate)))
|
|
}
|
|
if let fraction = (extra["passive_usage_7d_utilization"] as? NSNumber)?.doubleValue,
|
|
fraction.isFinite, (0...1).contains(fraction) {
|
|
let reset = (extra["passive_usage_7d_reset"] as? NSNumber)
|
|
.map { Date(timeIntervalSince1970: $0.doubleValue) }
|
|
windows.append(AIQuotaWindow(kind: .weekly, usedPercent: fraction * 100, resetsAt: reset))
|
|
}
|
|
if provider == .claude,
|
|
let fraction = (extra["passive_usage_7d_oi_utilization"] as? NSNumber)?.doubleValue,
|
|
fraction.isFinite, (0...1).contains(fraction) {
|
|
let reset = (extra["passive_usage_7d_oi_reset"] as? NSNumber)
|
|
.map { Date(timeIntervalSince1970: $0.doubleValue) }
|
|
windows.append(AIQuotaWindow(kind: .fableWeekly, usedPercent: fraction * 100, resetsAt: reset))
|
|
}
|
|
guard !windows.isEmpty else { throw AIQuotaFailure.sub2apiInvalidResponse }
|
|
return AIQuotaSnapshot(provider: provider, windows: windows, fetchedAt: now, source: .sub2api)
|
|
}
|
|
|
|
private static func parseDate(_ value: String) -> Date? {
|
|
let formatter = ISO8601DateFormatter()
|
|
formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds]
|
|
return formatter.date(from: value) ?? ISO8601DateFormatter().date(from: value)
|
|
}
|
|
}
|