Add support for federated authentication with Apple Business Manager

Sources/AppleAPI/Client.swift

@@ -26,7 +26,8 @@ public class Client {
         case accountUsesHardwareKey
         case srpInvalidPublicKey
         case srpError(String)
-
+        case federatedAuthenticationRequired
+
         public var errorDescription: String? {
             switch self {
             case .invalidUsernameOrPassword(let username):
@@ -47,6 +48,8 @@ public class Client {
                 return "Expected security code info but didn't receive any."
             case .accountUsesHardwareKey:
                 return "Account uses a hardware key for authentication but this is not supported yet."
+            case .federatedAuthenticationRequired:
+                return "This account uses federated authentication (e.g. Apple Business Manager). Browser-based login is required."
             default:
                 return String(describing: self)
             }
@@ -81,10 +84,10 @@ public class Client {
             struct ServiceKeyResponse: Decodable {
                 let authServiceKey: String?
             }
-            
+
             let response = try JSONDecoder().decode(ServiceKeyResponse.self, from: data)
             serviceKey = response.authServiceKey
-            
+
             /// Load a hashcash of the account name
             return self.loadHashcash(accountName: accountName, serviceKey: serviceKey).map { (serviceKey, $0) }
         }
@@ -348,6 +351,50 @@ public class Client {
         }
     }
 
+    public func checkFederation(accountName: String, serviceKey: String) -> Promise<FederationResponse> {
+        return Current.network.dataTask(with: URLRequest.checkFederation(serviceKey: serviceKey, accountName: accountName))
+            .map { data, _ in
+                try JSONDecoder().decode(FederationResponse.self, from: data)
+            }
+    }
+
+    /// Fetches the service key and checks whether the account uses federated authentication.
+    /// Call this before prompting for a password to avoid unnecessary password prompts for federated accounts.
+    public func checkIsFederated(accountName: String) -> Promise<FederationResponse> {
+        return firstly { () -> Promise<(data: Data, response: URLResponse)> in
+            Current.network.dataTask(with: URLRequest.itcServiceKey)
+        }
+        .then { (data, _) -> Promise<FederationResponse> in
+            struct ServiceKeyResponse: Decodable {
+                let authServiceKey: String?
+            }
+            let response = try JSONDecoder().decode(ServiceKeyResponse.self, from: data)
+            guard let serviceKey = response.authServiceKey else {
+                return .value(FederationResponse(federated: false))
+            }
+            return self.checkFederation(accountName: accountName, serviceKey: serviceKey)
+        }
+    }
+
+    /// Validates a federated authentication token by calling Apple's /federate/validate endpoint.
+    /// This exchanges the token and relay state from the IdP callback URL for session cookies.
+    public func validateFederatedToken(widgetKey: String, token: String, relayState: String) -> Promise<Void> {
+        return Current.network.dataTask(with: URLRequest.federateValidate(widgetKey: widgetKey, token: token, relayState: relayState))
+            .then { (data, response) -> Promise<Void> in
+                let httpResponse = response as! HTTPURLResponse
+                switch httpResponse.statusCode {
+                case 200...299:
+                    return Current.network.dataTask(with: URLRequest.olympusSession).asVoid()
+                case 409:
+                    // May need 2FA even after federated login
+                    let serviceKey = widgetKey
+                    return self.handleTwoStepOrFactor(data: data, response: response, serviceKey: serviceKey)
+                default:
+                    throw Error.unexpectedSignInResponse(statusCode: httpResponse.statusCode, message: nil)
+                }
+            }
+    }
+
     // Fixes issue https://github.com/RobotsAndPencils/XcodesApp/issues/360
     // On 2023-02-23, Apple added a custom implementation of hashcash to their auth flow
     // Without this addition, Apple ID's would get set to locked
@@ -516,3 +563,55 @@ public struct ServerSRPInitResponse: Decodable {
     let b: String
     let c: String
 }
+
+public struct FederationResponse: Decodable, Equatable {
+    public let federated: Bool
+    public let showFederatedIdpConfirmation: Bool?
+    public let federatedIdpRequest: FederatedIdpRequest?
+    public let federatedAuthIntro: FederatedAuthIntro?
+
+    public init(federated: Bool, showFederatedIdpConfirmation: Bool? = nil, federatedIdpRequest: FederatedIdpRequest? = nil, federatedAuthIntro: FederatedAuthIntro? = nil) {
+        self.federated = federated
+        self.showFederatedIdpConfirmation = showFederatedIdpConfirmation
+        self.federatedIdpRequest = federatedIdpRequest
+        self.federatedAuthIntro = federatedAuthIntro
+    }
+
+    /// Builds the full IdP URL with request parameters as query items.
+    public var idpURL: URL? {
+        guard let idpRequest = federatedIdpRequest else { return nil }
+        var components = URLComponents(string: idpRequest.idPUrl)
+        components?.queryItems = idpRequest.requestParams.map { key, value in
+            URLQueryItem(name: key, value: value)
+        }
+        return components?.url
+    }
+}
+
+public struct FederatedIdpRequest: Decodable, Equatable {
+    public let idPUrl: String
+    public let requestParams: [String: String]
+    public let httpMethod: String?
+
+    public init(idPUrl: String, requestParams: [String: String], httpMethod: String?) {
+        self.idPUrl = idPUrl
+        self.requestParams = requestParams
+        self.httpMethod = httpMethod
+    }
+}
+
+public struct FederatedAuthIntro: Decodable, Equatable {
+    public let orgName: String?
+    public let idpName: String?
+    public let idpUrl: String?
+    public let orgType: String?
+    public let accountManagementUrl: String?
+
+    public init(orgName: String?, idpName: String?, idpUrl: String?, orgType: String?, accountManagementUrl: String?) {
+        self.orgName = orgName
+        self.idpName = idpName
+        self.idpUrl = idpUrl
+        self.orgType = orgType
+        self.accountManagementUrl = accountManagementUrl
+    }
+}

Sources/AppleAPI/URLRequest+Apple.swift

@@ -12,6 +12,8 @@ extension URL {
 
     static let srpInit = URL(string: "https://idmsa.apple.com/appleauth/auth/signin/init")!
     static let srpComplete = URL(string: "https://idmsa.apple.com/appleauth/auth/signin/complete?isRememberMeEnabled=false")!
+    static let federateCheck = URL(string: "https://idmsa.apple.com/appleauth/auth/federate")!
+    static let federateValidate = URL(string: "https://idmsa.apple.com/appleauth/auth/federate/validate")!
 }
 
 extension URLRequest {
@@ -133,6 +135,36 @@ extension URLRequest {
         return request
     }
 
+    static func checkFederation(serviceKey: String, accountName: String) -> URLRequest {
+        struct Body: Encodable {
+            let accountName: String
+            let rememberMe = true
+        }
+
+        var request = URLRequest(url: .federateCheck)
+        request.allHTTPHeaderFields = request.allHTTPHeaderFields ?? [:]
+        request.allHTTPHeaderFields?["Content-Type"] = "application/json"
+        request.allHTTPHeaderFields?["X-Requested-With"] = "XMLHttpRequest"
+        request.allHTTPHeaderFields?["X-Apple-Widget-Key"] = serviceKey
+        request.allHTTPHeaderFields?["Accept"] = "application/json"
+        request.httpMethod = "POST"
+        request.httpBody = try! JSONEncoder().encode(Body(accountName: accountName))
+        return request
+    }
+
+    static func federateValidate(widgetKey: String, token: String, relayState: String) -> URLRequest {
+        var components = URLComponents(url: .federateValidate, resolvingAgainstBaseURL: false)!
+        components.queryItems = [
+            URLQueryItem(name: "widgetKey", value: widgetKey),
+            URLQueryItem(name: "token", value: token),
+            URLQueryItem(name: "relayState", value: relayState),
+        ]
+        var request = URLRequest(url: components.url!)
+        request.allHTTPHeaderFields = request.allHTTPHeaderFields ?? [:]
+        request.allHTTPHeaderFields?["Accept"] = "application/json"
+        return request
+    }
+
     static func SRPInit(serviceKey: String, a: String, accountName: String) -> URLRequest {
         struct ServerSRPInitRequest: Encodable {
             public let a: String

Sources/XcodesKit/AppleSessionService.swift

@@ -51,34 +51,41 @@ public class AppleSessionService {
             }
             guard let username = possibleUsername else { throw Error.missingUsernameOrPassword }
 
-            let passwordPrompt: String
-            if hasPromptedForUsername {
-                passwordPrompt = "Apple ID Password: "
-            } else {
-                // If the user wasn't prompted for their username, also explain which Apple ID password they need to enter
-                passwordPrompt = "Apple ID Password (\(username)): "
-            }
-            var possiblePassword = self.findPassword(withUsername: username)
-            if possiblePassword == nil || shouldPromptForPassword {
-                possiblePassword = Current.shell.readSecureLine(prompt: passwordPrompt)
-            }
-            guard let password = possiblePassword else { throw Error.missingUsernameOrPassword }
-
-            return firstly { () -> Promise<Void> in
-                self.login(username, password: password)
-            }
-            .recover { error -> Promise<Void> in
-                Current.logging.log(error.legibleLocalizedDescription.red)
-
-                if case Client.Error.invalidUsernameOrPassword = error {
-                    Current.logging.log("Try entering your password again")
-                    // Prompt for the password next time to avoid being stuck in a loop of using an incorrect XCODES_PASSWORD environment variable
-                    return self.loginIfNeeded(withUsername: username, shouldPromptForPassword: true)
-                }
-                else {
-                    return Promise(error: error)
+            // Check if this account uses federated authentication before prompting for a password
+            return Current.network.checkIsFederated(accountName: username)
+                .then { federationResponse -> Promise<Void> in
+                    if federationResponse.federated {
+                        return self.handleFederatedLogin(username: username, federationResponse: federationResponse)
+                    }
+
+                    // Not federated — proceed with normal password-based login
+                    let passwordPrompt: String
+                    if hasPromptedForUsername {
+                        passwordPrompt = "Apple ID Password: "
+                    } else {
+                        passwordPrompt = "Apple ID Password (\(username)): "
+                    }
+                    var possiblePassword = self.findPassword(withUsername: username)
+                    if possiblePassword == nil || shouldPromptForPassword {
+                        possiblePassword = Current.shell.readSecureLine(prompt: passwordPrompt)
+                    }
+                    guard let password = possiblePassword else { throw Error.missingUsernameOrPassword }
+
+                    return firstly { () -> Promise<Void> in
+                        self.login(username, password: password)
+                    }
+                    .recover { error -> Promise<Void> in
+                        Current.logging.log(error.legibleLocalizedDescription.red)
+
+                        if case Client.Error.invalidUsernameOrPassword = error {
+                            Current.logging.log("Try entering your password again")
+                            return self.loginIfNeeded(withUsername: username, shouldPromptForPassword: true)
+                        }
+                        else {
+                            return Promise(error: error)
+                        }
+                    }
                 }
-            }
         }
     }
 
@@ -89,7 +96,7 @@ public class AppleSessionService {
         .recover { error -> Promise<Void> in
 
             if let error = error as? Client.Error {
-                switch error  {
+                switch error {
                     case .invalidUsernameOrPassword(_):
                         // remove any keychain password if we fail to log with an invalid username or password so it doesn't try again.
                         try? Current.keychain.remove(username)
@@ -110,6 +117,85 @@ public class AppleSessionService {
         }
     }
 
+    private func handleFederatedLogin(username: String, federationResponse: FederationResponse) -> Promise<Void> {
+        guard let idpURL = federationResponse.idpURL else {
+            return Promise(error: Client.Error.federatedAuthenticationRequired)
+        }
+
+        let orgName = federationResponse.federatedAuthIntro?.orgName ?? "your organization"
+        let idpName = federationResponse.federatedAuthIntro?.idpName
+        let orgNameWithIdp = idpName.map { "\(orgName) (\($0))" } ?? orgName
+
+        Current.logging.log("\n- This account uses federated authentication via \(orgNameWithIdp)")
+        Current.logging.log("- Your browser will open to complete sign-in")
+        Current.logging.log("- After signing in, you will be redirected to a blank page")
+        Current.logging.log("- Copy the URL from your browser's address bar, then return here and paste it")
+
+        Current.shell.waitForKeypress(prompt: "\nPress any key to open your browser...")
+        Current.shell.openURL(idpURL)
+
+        let callbackURLString = Current.shell.readLongLine(prompt: "\nPaste the URL here: ")
+        guard let callbackURLString = callbackURLString,
+              let callbackURL = URL(string: callbackURLString),
+              let components = URLComponents(url: callbackURL, resolvingAgainstBaseURL: false),
+              let queryItems = components.queryItems else {
+            return Promise(error: Error.missingUsernameOrPassword)
+        }
+
+        let widgetKey = queryItems.first(where: { $0.name == "widgetKey" })?.value
+        let token = queryItems.first(where: { $0.name == "token" })?.value
+        let relayState = queryItems.first(where: { $0.name == "relayState" })?.value
+
+        guard let widgetKey = widgetKey, let token = token, let relayState = relayState else {
+            Current.logging.log("The URL is missing required parameters (widgetKey, token, relayState).")
+            return Promise(error: Client.Error.invalidSession)
+        }
+
+        return Current.network.validateFederatedToken(widgetKey: widgetKey, token: token, relayState: relayState)
+            .done {
+                // Apple's auth cookies (e.g. myacinfo) are session-only, so HTTPCookieStorage
+                // discards them when the process exits. Since federated accounts have no password
+                // in the keychain to re-authenticate silently, we replace session-only cookies
+                // with persistent copies so the session survives across runs.
+                self.persistSessionOnlyCookies()
+
+                if self.configuration.defaultUsername != username {
+                    self.configuration.defaultUsername = username
+                    try? self.configuration.save()
+                }
+            }
+    }
+
+    /// Replaces session-only Apple cookies with persistent copies by adding an expiry date.
+    /// This ensures HTTPCookieStorage keeps them across process launches.
+    private func persistSessionOnlyCookies() {
+        let appleDomains = [".apple.com", ".idmsa.apple.com", "appstoreconnect.apple.com"]
+        guard let cookieStorage = AppleAPI.Current.network.session.configuration.httpCookieStorage else { return }
+        let cookies = cookieStorage.cookies ?? []
+        let expiry = Date(timeIntervalSinceNow: 24 * 60 * 60) // 1 day
+
+        for cookie in cookies where cookie.isSessionOnly {
+            guard appleDomains.contains(where: { cookie.domain.hasSuffix($0) }) else { continue }
+
+            var properties: [HTTPCookiePropertyKey: Any] = [
+                .name: cookie.name,
+                .value: cookie.value,
+                .domain: cookie.domain,
+                .path: cookie.path,
+                .secure: cookie.isSecure,
+                .expires: expiry,
+            ]
+            if let version = cookie.properties?[.version] {
+                properties[.version] = version
+            }
+
+            cookieStorage.deleteCookie(cookie)
+            if let persistent = HTTPCookie(properties: properties) {
+                cookieStorage.setCookie(persistent)
+            }
+        }
+    }
+
     public func logout() -> Promise<Void> {
         guard let username = findUsername() else { return Promise<Void>(error: Client.Error.notAuthenticated) }