Merge pull request #3 from n3d1117/2fa-test

Handle 2 factor authentication
n3d1117 · Mar 23, 2021 · 1df44d5 · 1df44d5
2 parents 2059ef9 + 1d43d26
commit 1df44d5
Show file tree

Hide file tree

Showing 5 changed files with 26 additions and 7 deletions.
diff --git a/README.md b/README.md
@@ -22,6 +22,8 @@ To authenticate you can pass the following arguments to any recipe:
 
 If you prefer, you can set these environment variables instead: `COOK_APPLEID_EMAIL` and `COOK_APPLEID_PASSWORD`.
 
+Pass `--2fa-code` to specify the six-digit two factor authentication code.
+
 You can also specify `--base64-anisette-data` to use custom base64 encoded anisette data in your requests (for example, generated with `anisette_server` recipe).
 
 **NOTE:** Your Apple ID is **never** sent to anyone but Apple. Feel free to create a new Apple ID account to test it.

diff --git a/cook/Stuff/Authenticator.swift b/cook/Stuff/Authenticator.swift
@@ -34,9 +34,8 @@ class Authenticator {
 
         if let anisetteData = customAnisetteData {
             logger.log(.info, "Logging in using custom anisette data...")
-            ALTAppleAPI.shared.authenticate(appleID: self.appleId, password: self.password, anisetteData: anisetteData, verificationHandler: nil, completionHandler: { [weak self] (account, session, error) in
-                self?.completionHandler?(account, session, error)
-            })
+            logger.log(.verbose, "customAnisetteData is \(anisetteData)")
+            authenticateAndComplete(anisetteData: anisetteData)
         } else {
             // Open Mail.app
             Utils.shell("open", "-j", "-g", "-a", "Mail")
@@ -64,15 +63,29 @@ class Authenticator {
                 adjustedDescription += "(com.apple.dt.Xcode/3594.4.19)>"
                 anisetteData.deviceDescription = String(adjustedDescription)
             }
-
             logger.log(.info, "Logging in...")
-            ALTAppleAPI.shared.authenticate(appleID: self.appleId, password: self.password, anisetteData: anisetteData, verificationHandler: nil, completionHandler: { [weak self] (account, session, error) in
-                self?.completionHandler?(account, session, error)
-            })
+            authenticateAndComplete(anisetteData: anisetteData)
 
         } catch {
             return abort(.unableToUnarchive)
         }
     }
 
+    fileprivate func handle2FAVerificationCode(_ completionHandler: @escaping (String?) -> Void) {
+        DispatchQueue.main.async {
+            let twoFactorCode = CLI.parseArgument(.twoFactorCode)
+
+            guard twoFactorCode != nil else { return _abort(UsageError.missing2FACode) }
+            guard twoFactorCode!.count == 6 else { return _abort(AuthError.malformed2FACode) }
+
+            logger.log(.verbose, "2fa code is \(twoFactorCode!)")
+            completionHandler(twoFactorCode)
+        }
+    }
+
+    fileprivate func authenticateAndComplete(anisetteData: ALTAnisetteData) {
+        ALTAppleAPI.shared.authenticate(appleID: self.appleId, password: self.password, anisetteData: anisetteData, verificationHandler: handle2FAVerificationCode, completionHandler: { [weak self] (account, session, error) in
+            self?.completionHandler?(account, session, error)
+        })
+    }
 }
diff --git a/cook/Stuff/CLI.swift b/cook/Stuff/CLI.swift
@@ -14,6 +14,7 @@ enum CLI {
         // Authentication
         case appleId
         case password
+        case twoFactorCode = "2fa-code"
 
         // Certificates
         case machinePrefix = "machine-name"

diff --git a/cook/Stuff/Errors.swift b/cook/Stuff/Errors.swift
@@ -9,6 +9,7 @@
 enum UsageError: Error {
     case missingAppleId
     case missingPassword
+    case missing2FACode
 
     case missingOutput
 
@@ -36,6 +37,7 @@ enum AuthError: Error {
     case malformedAnisetteData
     case unableToUnarchive
     case unknownAuthFailure
+    case malformed2FACode
 }
 
 enum TeamError: Error {

diff --git a/cook/Stuff/Utils.swift b/cook/Stuff/Utils.swift
@@ -83,6 +83,7 @@ enum Utils {
         print("  To authenticate pass the following arguments to any recipe:")
         print("    --appleId             Apple ID Email")
         print("    --password            Apple ID Password\n")
+        print("    --2fa-code            Two Factor Authentication code\n")
         print("  or, if you prefer, you can set these environment variables:")
         print("    COOK_APPLEID_EMAIL")
         print("    COOK_APPLEID_PASSWORD\n")