115 lines · 4.9 KB · Swift
RawDownload
import Foundation public enum FileManagedCodexAccountStoreError: Error, Equatable, Sendable {    case unsupportedVersion(Int)} public protocol ManagedCodexAccountStoring: Sendable {    func loadAccounts() throws -> ManagedCodexAccountSet    func storeAccounts(_ accounts: ManagedCodexAccountSet) throws    func ensureFileExists() throws -> URL} public struct FileManagedCodexAccountStore: ManagedCodexAccountStoring, @unchecked Sendable {    public static let currentVersion = 3     private let fileURL: URL    private let fileManager: FileManager     public init(fileURL: URL = Self.defaultURL(), fileManager: FileManager = .default) {        self.fileURL = fileURL        self.fileManager = fileManager    }     public func loadAccounts() throws -> ManagedCodexAccountSet {        guard self.fileManager.fileExists(atPath: self.fileURL.path) else {            return Self.emptyAccountSet()        }         let data = try Data(contentsOf: self.fileURL)        let decoder = JSONDecoder()        let accounts = try decoder.decode(ManagedCodexAccountSet.self, from: data)        guard (1...Self.currentVersion).contains(accounts.version) else {            throw FileManagedCodexAccountStoreError.unsupportedVersion(accounts.version)        }        if accounts.version == Self.currentVersion {            return ManagedCodexAccountSet(version: Self.currentVersion, accounts: accounts.accounts)        }        return self.migrateLegacyAccounts(accounts)    }     public func storeAccounts(_ accounts: ManagedCodexAccountSet) throws {        let normalizedAccounts = ManagedCodexAccountSet(            version: Self.currentVersion,            accounts: accounts.accounts)        let encoder = JSONEncoder()        encoder.outputFormatting = [.prettyPrinted, .sortedKeys]        let data = try encoder.encode(normalizedAccounts)        let directory = self.fileURL.deletingLastPathComponent()        if !self.fileManager.fileExists(atPath: directory.path) {            try self.fileManager.createDirectory(at: directory, withIntermediateDirectories: true)        }        try data.write(to: self.fileURL, options: [.atomic])        try self.applySecurePermissionsIfNeeded()    }     public func ensureFileExists() throws -> URL {        if self.fileManager.fileExists(atPath: self.fileURL.path) { return self.fileURL }        try self.storeAccounts(Self.emptyAccountSet())        return self.fileURL    }     private func applySecurePermissionsIfNeeded() throws {        #if os(macOS)        try self.fileManager.setAttributes([            .posixPermissions: NSNumber(value: Int16(0o600)),        ], ofItemAtPath: self.fileURL.path)        #endif    }     private static func emptyAccountSet() -> ManagedCodexAccountSet {        ManagedCodexAccountSet(version: self.currentVersion, accounts: [])    }     private func migrateLegacyAccounts(_ accounts: ManagedCodexAccountSet) -> ManagedCodexAccountSet {        let migratedAccounts = accounts.accounts.map { account in            let hydratedProviderAccountID = account.providerAccountID ?? self.hydrateProviderAccountID(for: account)            return ManagedCodexAccount(                id: account.id,                email: account.email,                providerAccountID: hydratedProviderAccountID,                workspaceLabel: account.workspaceLabel,                workspaceAccountID: account.workspaceAccountID,                authFingerprint: account.authFingerprint ?? CodexAuthFingerprint.fingerprint(                    homePath: account.managedHomePath,                    fileManager: self.fileManager),                managedHomePath: account.managedHomePath,                createdAt: account.createdAt,                updatedAt: account.updatedAt,                lastAuthenticatedAt: account.lastAuthenticatedAt)        }        return ManagedCodexAccountSet(version: Self.currentVersion, accounts: migratedAccounts)    }     private func hydrateProviderAccountID(for account: ManagedCodexAccount) -> String? {        guard let credentials = try? CodexOAuthCredentialsStore.load(            env: ["CODEX_HOME": account.managedHomePath])        else {            return nil        }        let payload = credentials.idToken.flatMap(UsageFetcher.parseJWT)        let authDict = payload?["https://api.openai.com/auth"] as? [String: Any]        let providerAccountID = credentials.accountId            ?? (authDict?["chatgpt_account_id"] as? String)            ?? (payload?["chatgpt_account_id"] as? String)        return ManagedCodexAccount.normalizeProviderAccountID(providerAccountID)    }     public static func defaultURL() -> URL {        let base = FileManager.default.urls(for: .applicationSupportDirectory, in: .userDomainMask).first            ?? FileManager.default.homeDirectoryForCurrentUser        return base            .appendingPathComponent("CodexBar", isDirectory: true)            .appendingPathComponent("managed-codex-accounts.json")    }}