From 33cd3d10a467be147a4a26bff7d3a1da1a262812 Mon Sep 17 00:00:00 2001
From: nab138 <nab@nabdev.me>
Date: Sun, 25 Jan 2026 23:26:30 -0500
Subject: [PATCH] Add extra auth step to login

---
 examples/minimal/src/main.rs        |  2 +-
 isideload/src/auth/apple_account.rs | 46 +++++++++++++++++++++++++++--
 2 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/examples/minimal/src/main.rs b/examples/minimal/src/main.rs
index d5fb149..d36dc2f 100644
--- a/examples/minimal/src/main.rs
+++ b/examples/minimal/src/main.rs
@@ -37,7 +37,7 @@ async fn main() {
         .await;
 
     match account {
-        Ok(_account) => println!("Successfully logged in to Apple ID"),
+        Ok(a) => println!("Logged in. {}", a),
         Err(e) => eprintln!("Failed to log in to Apple ID: {:?}", e),
     }
 }
diff --git a/isideload/src/auth/apple_account.rs b/isideload/src/auth/apple_account.rs
index ddaff81..6e8782d 100644
--- a/isideload/src/auth/apple_account.rs
+++ b/isideload/src/auth/apple_account.rs
@@ -178,7 +178,7 @@ impl AppleAccount {
                     self.login_state
                 );
             }
-            match self.login_state {
+            match &self.login_state {
                 LoginState::LoggedIn => {
                     info!("Successfully logged in to Apple ID");
                     return Ok(());
@@ -198,7 +198,13 @@ impl AppleAccount {
                     debug!("SMS 2FA completed, need to login again");
                     self.login_state = LoginState::NeedsLogin;
                 }
-                LoginState::NeedsExtraStep(_) => todo!(),
+                LoginState::NeedsExtraStep(s) => {
+                    info!("Additional authentication step required: {}", s);
+                    if self.get_pet().is_err() {
+                        bail!("Additional authentication required: {}", s);
+                    }
+                    self.login_state = LoginState::LoggedIn;
+                }
                 LoginState::NeedsLogin => {
                     debug!("Logging in again...");
                     self.login_state = self
@@ -210,6 +216,29 @@ impl AppleAccount {
         }
     }
 
+    pub fn get_name(&self) -> Result<(String, String), Report> {
+        let spd = self
+            .spd
+            .as_ref()
+            .ok_or_else(|| report!("SPD not available, cannot get name"))?;
+
+        Ok((spd.get_string("fn")?, spd.get_string("ln")?))
+    }
+
+    fn get_pet(&self) -> Result<String, Report> {
+        let spd = self
+            .spd
+            .as_ref()
+            .ok_or_else(|| report!("SPD not available, cannot get pet"))?;
+
+        let pet = spd
+            .get_dict("t")?
+            .get_dict("com.apple.gs.idms.pet")?
+            .get_string("token")?;
+
+        Ok(pet)
+    }
+
     async fn trusted_device_2fa(
         &mut self,
         two_factor_callback: impl Fn() -> Option<String>,
@@ -471,7 +500,7 @@ impl AppleAccount {
 
         if let Some(plist::Value::String(s)) = status.get("au") {
             return Ok(match s.as_str() {
-                "trustedDeviceSecondaryAuth" => LoginState::NeedsSMS2FA,
+                "trustedDeviceSecondaryAuth" => LoginState::NeedsDevice2FA,
                 "secondaryAuth" => LoginState::NeedsSMS2FA,
                 "repair" => LoginState::LoggedIn, // Just means that you don't have 2FA set up
                 unknown => LoginState::NeedsExtraStep(unknown.to_string()),
@@ -500,3 +529,14 @@ impl AppleAccount {
         )
     }
 }
+
+impl std::fmt::Display for AppleAccount {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "Apple Account: ")?;
+        match self.get_name() {
+            Ok((first, last)) => write!(f, "{} {} ", first, last),
+            Err(_) => Ok(()),
+        }?;
+        write!(f, "{} ({:?})", self.email, self.login_state)
+    }
+}