From 1d43d26dd439f60337359840f295e292ade5d85d Mon Sep 17 00:00:00 2001
From: ned <ned@me.com>
Date: Tue, 23 Mar 2021 14:04:24 +0100
Subject: [PATCH] handle 2fa

by zond80
---
 README.md                      |  2 ++
 cook/Stuff/Authenticator.swift | 27 ++++++++++++++++++++-------
 cook/Stuff/CLI.swift           |  1 +
 cook/Stuff/Errors.swift        |  2 ++
 cook/Stuff/Utils.swift         |  1 +
 5 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/README.md b/README.md
index 1b8f2b2..a2e6c35 100644
--- 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
index c71fc4a..5c29639 100644
--- 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
index e58ae09..18c31b3 100644
--- 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
index 67432a2..c8bd8d0 100644
--- 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
index 40480e4..1b825b7 100644
--- 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")