wehub-resource-sync
347 lines
13 KiB
Swift
347 lines
13 KiB
Swift
import CodexBarCore
|
|
import Foundation
|
|
|
|
struct PreparedIdentity: Equatable {
|
|
let email: String?
|
|
let identity: CodexIdentity
|
|
let providerAccountID: String?
|
|
let workspaceLabel: String?
|
|
let workspaceAccountID: String?
|
|
}
|
|
|
|
struct PreparedAuthMaterial {
|
|
let homeURL: URL
|
|
let rawData: Data
|
|
let credentials: CodexOAuthCredentials
|
|
let runtimeAccount: CodexAuthBackedAccount
|
|
let authIdentity: PreparedIdentity
|
|
}
|
|
|
|
enum PreparedManagedHomeState {
|
|
case readable(PreparedAuthMaterial)
|
|
case missing(homeURL: URL)
|
|
case unreadable(homeURL: URL)
|
|
}
|
|
|
|
struct PreparedStoredManagedAccount {
|
|
let persisted: ManagedCodexAccount
|
|
let persistedIdentity: PreparedIdentity
|
|
let homeState: PreparedManagedHomeState
|
|
|
|
var authIdentity: PreparedIdentity? {
|
|
switch self.homeState {
|
|
case let .readable(authMaterial):
|
|
authMaterial.authIdentity
|
|
case .missing, .unreadable:
|
|
nil
|
|
}
|
|
}
|
|
}
|
|
|
|
enum PreparedLiveHomeState {
|
|
case missing(homeURL: URL)
|
|
case unreadable(homeURL: URL)
|
|
case apiKeyOnly(PreparedAuthMaterial)
|
|
case readable(PreparedAuthMaterial)
|
|
}
|
|
|
|
struct PreparedLiveAccount {
|
|
let homeState: PreparedLiveHomeState
|
|
|
|
var homeURL: URL {
|
|
switch self.homeState {
|
|
case let .missing(homeURL), let .unreadable(homeURL):
|
|
homeURL
|
|
case let .apiKeyOnly(authMaterial), let .readable(authMaterial):
|
|
authMaterial.homeURL
|
|
}
|
|
}
|
|
|
|
var authIdentity: PreparedIdentity? {
|
|
switch self.homeState {
|
|
case let .apiKeyOnly(authMaterial), let .readable(authMaterial):
|
|
authMaterial.authIdentity
|
|
case .missing, .unreadable:
|
|
nil
|
|
}
|
|
}
|
|
}
|
|
|
|
struct PreparedPromotionContext {
|
|
let snapshot: CodexAccountReconciliationSnapshot
|
|
let managedAccounts: ManagedCodexAccountSet
|
|
let storedManagedAccounts: [PreparedStoredManagedAccount]
|
|
let target: PreparedStoredManagedAccount
|
|
let live: PreparedLiveAccount
|
|
}
|
|
|
|
@MainActor
|
|
struct PreparedPromotionContextBuilder {
|
|
private let store: any ManagedCodexAccountStoring
|
|
private let workspaceResolver: any ManagedCodexWorkspaceResolving
|
|
private let snapshotLoader: any CodexAccountReconciliationSnapshotLoading
|
|
private let authMaterialReader: any CodexAuthMaterialReading
|
|
private let baseEnvironment: [String: String]
|
|
private let fileManager: FileManager
|
|
|
|
init(
|
|
store: any ManagedCodexAccountStoring,
|
|
workspaceResolver: any ManagedCodexWorkspaceResolving,
|
|
snapshotLoader: any CodexAccountReconciliationSnapshotLoading,
|
|
authMaterialReader: any CodexAuthMaterialReading,
|
|
baseEnvironment: [String: String] = ProcessInfo.processInfo.environment,
|
|
fileManager: FileManager = .default)
|
|
{
|
|
self.store = store
|
|
self.workspaceResolver = workspaceResolver
|
|
self.snapshotLoader = snapshotLoader
|
|
self.authMaterialReader = authMaterialReader
|
|
self.baseEnvironment = baseEnvironment
|
|
self.fileManager = fileManager
|
|
}
|
|
|
|
func build(targetID: UUID) async throws -> PreparedPromotionContext {
|
|
let snapshot = self.snapshotLoader.loadSnapshot()
|
|
let managedAccounts = try self.store.loadAccounts()
|
|
var preparedAccounts: [PreparedStoredManagedAccount] = []
|
|
preparedAccounts.reserveCapacity(managedAccounts.accounts.count)
|
|
for account in managedAccounts.accounts {
|
|
let preparedAccount = try await self.prepareStoredManagedAccount(account)
|
|
preparedAccounts.append(preparedAccount)
|
|
}
|
|
|
|
guard let target = preparedAccounts.first(where: { $0.persisted.id == targetID }) else {
|
|
throw CodexAccountPromotionError.targetManagedAccountNotFound
|
|
}
|
|
|
|
let live = await self.prepareLiveAccount()
|
|
return PreparedPromotionContext(
|
|
snapshot: snapshot,
|
|
managedAccounts: managedAccounts,
|
|
storedManagedAccounts: preparedAccounts,
|
|
target: target,
|
|
live: live)
|
|
}
|
|
|
|
private func prepareStoredManagedAccount(
|
|
_ account: ManagedCodexAccount) async throws
|
|
-> PreparedStoredManagedAccount
|
|
{
|
|
let homeURL = URL(fileURLWithPath: account.managedHomePath, isDirectory: true)
|
|
let persistedIdentity = Self.persistedIdentity(from: account)
|
|
let homeState = await self.prepareManagedHomeState(homeURL: homeURL)
|
|
|
|
return PreparedStoredManagedAccount(
|
|
persisted: account,
|
|
persistedIdentity: persistedIdentity,
|
|
homeState: homeState)
|
|
}
|
|
|
|
private func prepareManagedHomeState(homeURL: URL) async -> PreparedManagedHomeState {
|
|
let readResult = self.readAuthData(homeURL: homeURL)
|
|
switch readResult {
|
|
case .missing:
|
|
return .missing(homeURL: homeURL)
|
|
case .unreadable:
|
|
return .unreadable(homeURL: homeURL)
|
|
case let .readable(rawData):
|
|
guard let authMaterial = await self.inspectAuthMaterial(homeURL: homeURL, rawData: rawData) else {
|
|
return .unreadable(homeURL: homeURL)
|
|
}
|
|
return .readable(authMaterial)
|
|
}
|
|
}
|
|
|
|
private func prepareLiveAccount() async -> PreparedLiveAccount {
|
|
let liveHomeURL = self.liveHomeURL()
|
|
let readResult = self.readAuthData(homeURL: liveHomeURL)
|
|
switch readResult {
|
|
case .missing:
|
|
return PreparedLiveAccount(homeState: .missing(homeURL: liveHomeURL))
|
|
case .unreadable:
|
|
return PreparedLiveAccount(homeState: .unreadable(homeURL: liveHomeURL))
|
|
case let .readable(rawData):
|
|
guard let authMaterial = await self.inspectAuthMaterial(homeURL: liveHomeURL, rawData: rawData) else {
|
|
return PreparedLiveAccount(homeState: .unreadable(homeURL: liveHomeURL))
|
|
}
|
|
if Self.isAPIKeyOnly(credentials: authMaterial.credentials, rawData: authMaterial.rawData) {
|
|
return PreparedLiveAccount(homeState: .apiKeyOnly(authMaterial))
|
|
}
|
|
return PreparedLiveAccount(homeState: .readable(authMaterial))
|
|
}
|
|
}
|
|
|
|
private func inspectAuthMaterial(homeURL: URL, rawData: Data) async -> PreparedAuthMaterial? {
|
|
guard let credentials = try? CodexOAuthCredentialsStore.parse(data: rawData),
|
|
let runtimeAccount = try? Self.runtimeAccount(from: rawData)
|
|
else {
|
|
return nil
|
|
}
|
|
|
|
let authIdentity = await self.derivedIdentity(
|
|
homePath: homeURL.path,
|
|
runtimeAccount: runtimeAccount)
|
|
|
|
return PreparedAuthMaterial(
|
|
homeURL: homeURL,
|
|
rawData: rawData,
|
|
credentials: credentials,
|
|
runtimeAccount: runtimeAccount,
|
|
authIdentity: authIdentity)
|
|
}
|
|
|
|
private func derivedIdentity(homePath: String, runtimeAccount: CodexAuthBackedAccount) async -> PreparedIdentity {
|
|
let normalizedEmail = Self.normalizeEmail(runtimeAccount.email)
|
|
let normalizedIdentity = Self.normalizedIdentity(runtimeAccount.identity, email: normalizedEmail)
|
|
let providerAccountID: String? = switch normalizedIdentity {
|
|
case let .providerAccount(id):
|
|
ManagedCodexAccount.normalizeProviderAccountID(id)
|
|
case .emailOnly, .unresolved:
|
|
nil
|
|
}
|
|
let workspaceIdentity: CodexOpenAIWorkspaceIdentity? = if let providerAccountID {
|
|
await self.workspaceResolver.resolveWorkspaceIdentity(
|
|
homePath: homePath,
|
|
providerAccountID: providerAccountID)
|
|
} else {
|
|
nil
|
|
}
|
|
|
|
return PreparedIdentity(
|
|
email: normalizedEmail,
|
|
identity: normalizedIdentity,
|
|
providerAccountID: providerAccountID,
|
|
workspaceLabel: workspaceIdentity?.workspaceLabel,
|
|
workspaceAccountID: workspaceIdentity?.workspaceAccountID ?? providerAccountID)
|
|
}
|
|
|
|
private static func persistedIdentity(from account: ManagedCodexAccount) -> PreparedIdentity {
|
|
let normalizedEmail = Self.normalizeEmail(account.email)
|
|
let providerAccountID = ManagedCodexAccount.normalizeProviderAccountID(account.providerAccountID)
|
|
let identity = Self.normalizedIdentity(
|
|
CodexIdentityResolver.resolve(accountId: providerAccountID, email: normalizedEmail),
|
|
email: normalizedEmail)
|
|
|
|
return PreparedIdentity(
|
|
email: normalizedEmail,
|
|
identity: identity,
|
|
providerAccountID: providerAccountID,
|
|
workspaceLabel: account.workspaceLabel,
|
|
workspaceAccountID: account.workspaceAccountID)
|
|
}
|
|
|
|
private func liveHomeURL() -> URL {
|
|
CodexHomeScope.ambientHomeURL(env: self.baseEnvironment, fileManager: self.fileManager)
|
|
}
|
|
|
|
private func readAuthData(homeURL: URL) -> PreparedAuthReadState {
|
|
do {
|
|
let rawData = try self.authMaterialReader.readAuthData(homeURL: homeURL)
|
|
guard let rawData else {
|
|
return .missing
|
|
}
|
|
return .readable(rawData)
|
|
} catch {
|
|
return .unreadable
|
|
}
|
|
}
|
|
|
|
private static func runtimeAccount(from rawData: Data) throws -> CodexAuthBackedAccount {
|
|
guard let json = try JSONSerialization.jsonObject(with: rawData) as? [String: Any] else {
|
|
throw CodexOAuthCredentialsError.decodeFailed("Invalid JSON")
|
|
}
|
|
|
|
let tokens = json["tokens"] as? [String: Any]
|
|
let idToken = tokens.flatMap {
|
|
Self.nonEmptyString(in: $0, snakeCaseKey: "id_token", camelCaseKey: "idToken")
|
|
}
|
|
let payload = idToken.flatMap(UsageFetcher.parseJWT)
|
|
let authDict = payload?["https://api.openai.com/auth"] as? [String: Any]
|
|
let profileDict = payload?["https://api.openai.com/profile"] as? [String: Any]
|
|
|
|
let email = Self.normalizeEmail(
|
|
(payload?["email"] as? String) ?? (profileDict?["email"] as? String))
|
|
let plan = Self.normalizedField(
|
|
(authDict?["chatgpt_plan_type"] as? String) ?? (payload?["chatgpt_plan_type"] as? String))
|
|
let accountID = ManagedCodexAccount.normalizeProviderAccountID(
|
|
tokens.flatMap {
|
|
Self.nonEmptyString(in: $0, snakeCaseKey: "account_id", camelCaseKey: "accountId")
|
|
}
|
|
?? (authDict?["chatgpt_account_id"] as? String)
|
|
?? (payload?["chatgpt_account_id"] as? String))
|
|
let identity = Self.normalizedIdentity(
|
|
CodexIdentityResolver.resolve(accountId: accountID, email: email),
|
|
email: email)
|
|
|
|
return CodexAuthBackedAccount(identity: identity, email: email, plan: plan)
|
|
}
|
|
|
|
private static func normalizedField(_ value: String?) -> String? {
|
|
guard let value = value?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else {
|
|
return nil
|
|
}
|
|
return value
|
|
}
|
|
|
|
private static func normalizeEmail(_ email: String?) -> String? {
|
|
CodexIdentityResolver.normalizeEmail(email)
|
|
}
|
|
|
|
private static func normalizedIdentity(_ identity: CodexIdentity, email: String?) -> CodexIdentity {
|
|
guard let email else { return identity }
|
|
return CodexIdentityMatcher.normalized(identity, fallbackEmail: email)
|
|
}
|
|
|
|
private static func isAPIKeyOnly(credentials: CodexOAuthCredentials, rawData: Data) -> Bool {
|
|
guard self.hasUsableOAuthTokens(in: rawData) == false else {
|
|
return false
|
|
}
|
|
return credentials.refreshToken.isEmpty
|
|
&& credentials.idToken == nil
|
|
&& credentials.accountId == nil
|
|
&& credentials.lastRefresh == nil
|
|
}
|
|
|
|
private static func hasUsableOAuthTokens(in rawData: Data) -> Bool {
|
|
guard let json = try? JSONSerialization.jsonObject(with: rawData) as? [String: Any],
|
|
let tokens = json["tokens"] as? [String: Any]
|
|
else {
|
|
return false
|
|
}
|
|
let accessToken = self.nonEmptyString(
|
|
in: tokens,
|
|
snakeCaseKey: "access_token",
|
|
camelCaseKey: "accessToken")
|
|
let refreshToken = self.nonEmptyString(
|
|
in: tokens,
|
|
snakeCaseKey: "refresh_token",
|
|
camelCaseKey: "refreshToken")
|
|
return accessToken != nil && refreshToken != nil
|
|
}
|
|
|
|
private static func nonEmptyString(
|
|
in dictionary: [String: Any],
|
|
snakeCaseKey: String,
|
|
camelCaseKey: String)
|
|
-> String?
|
|
{
|
|
if let value = dictionary[snakeCaseKey] as? String,
|
|
value.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty == false
|
|
{
|
|
return value
|
|
}
|
|
if let value = dictionary[camelCaseKey] as? String,
|
|
value.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty == false
|
|
{
|
|
return value
|
|
}
|
|
return nil
|
|
}
|
|
}
|
|
|
|
private enum PreparedAuthReadState {
|
|
case missing
|
|
case unreadable
|
|
case readable(Data)
|
|
}
|