diff --git a/grammers-client/Cargo.toml b/grammers-client/Cargo.toml index 4ccdda12..afdf90fb 100644 --- a/grammers-client/Cargo.toml +++ b/grammers-client/Cargo.toml @@ -22,6 +22,7 @@ fs = ["tokio/fs"] default = ["fs"] [dependencies] +base64 = { version = "0.22", features = ["std"] } chrono = "0.4.42" futures-util = { version = "0.3.31", default-features = false, features = [ "alloc" diff --git a/grammers-client/src/client/auth.rs b/grammers-client/src/client/auth.rs index cac65076..8b3b695f 100644 --- a/grammers-client/src/client/auth.rs +++ b/grammers-client/src/client/auth.rs @@ -7,12 +7,15 @@ // except according to those terms. use std::fmt; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; use grammers_crypto::two_factor_auth::{calculate_2fa, check_p_and_g}; use grammers_mtsender::InvocationError; use grammers_session::types::{PeerInfo, UpdateState, UpdatesState}; use grammers_tl_types as tl; +use base64::Engine; + use super::Client; use crate::peer::User; use crate::utils; @@ -495,3 +498,484 @@ impl Client { self.0.handle.quit(); } } + +// QR Login functionality + +/// Status of the QR login process +#[derive(Debug, Clone, PartialEq)] +pub enum QrLoginStatus { + Idle, + Waiting, + Expired, + Success, + PasswordRequired(Option), // 2FA required, with optional hint + Error(String), +} + +/// Information about the current QR login state +pub struct QrLoginInfo { + pub qr_url: String, + pub expires_unix: u64, + pub expires_in_seconds: i64, + pub status: QrLoginStatus, +} + +impl Client { + /// Export login token for QR code login + pub async fn export_login_token( + &self, + api_id: i32, + api_hash: &str, + ) -> Result { + let request = tl::functions::auth::ExportLoginToken { + api_id, + api_hash: api_hash.to_string(), + except_ids: vec![], + }; + + self.invoke(&request).await + } + + /// Import login token for DC migration + /// Import login token for DC migration - switches home DC before importing + pub async fn import_login_token( + &self, + token: Vec, + dc_id: i32, + ) -> Result { + // Handle DC migration by switching home DC + self.handle_qr_login_migration(dc_id).await?; + + let request = tl::functions::auth::ImportLoginToken { token }; + + // Import on the new home DC (after migration) + match self.invoke(&request).await { + Ok(result) => Ok(result), + Err(InvocationError::Rpc(ref err)) if err.name == "SESSION_PASSWORD_NEEDED" => { + // If password is needed, we return an error that can be handled by the caller + Err(InvocationError::Rpc(err.clone())) + } + Err(e) => Err(e), + } + } + + /// Handle DC migration during QR login by switching home DC + pub(crate) async fn handle_qr_login_migration( + &self, + new_dc_id: i32, + ) -> Result<(), InvocationError> { + let old_dc = self.0.session.home_dc_id(); + self.0.handle.disconnect_from_dc(old_dc); + self.0.session.set_home_dc_id(new_dc_id).await; + Ok(()) + } + + /// Finalize QR login by completing the authorization process + pub async fn finalize_qr_login( + &self, + auth: tl::types::auth::Authorization, + ) -> Result { + self.complete_login(auth).await + } + + /// Get password information for 2FA authentication + pub async fn qr_get_password_token(&self) -> Result { + self.get_password_information().await + } + + /// Fetches 2FA password token (hint, SRP params) so the app can prompt for password. + pub async fn get_password_token(&self) -> Result { + self.get_password_information().await + } + + /// Convert raw token bytes to base64url encoded string + fn encode_token_to_base64url(&self, token_bytes: &[u8]) -> String { + // Use URL-safe base64 encoding without padding + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(token_bytes) + } + + /// Generate QR code URL from token bytes + fn generate_qr_url(&self, token_bytes: &[u8]) -> String { + let encoded_token = self.encode_token_to_base64url(token_bytes); + format!("tg://login?token={}", encoded_token) + } + + /// Start QR login process and return initial QR information + pub async fn start_qr_login( + &self, + api_id: i32, + api_hash: &str, + ) -> Result { + let login_token = match self.export_login_token(api_id, api_hash).await { + Ok(token) => token, + Err(InvocationError::Rpc(err)) if err.name == "SESSION_PASSWORD_NEEDED" => { + // Return PasswordRequired status instead of failing + return Ok(QrLoginInfo { + qr_url: "".to_string(), + expires_unix: 0, + expires_in_seconds: 0, + status: QrLoginStatus::PasswordRequired(None), // We'll get the hint separately + }); + } + Err(e) => return Err(e), + }; + + match login_token { + tl::enums::auth::LoginToken::Token(token) => { + let qr_url = self.generate_qr_url(&token.token); + let expires_unix = token.expires as u64; + let current_time = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs(); + let expires_in_seconds = token.expires as i64 - current_time as i64; + + Ok(QrLoginInfo { + qr_url, + expires_unix, + expires_in_seconds, + status: QrLoginStatus::Waiting, + }) + } + tl::enums::auth::LoginToken::MigrateTo(migrate_to) => { + // Handle migration by switching home DC and importing the token on the new DC + self.handle_qr_login_migration(migrate_to.dc_id).await?; + + let import_request = tl::functions::auth::ImportLoginToken { + token: migrate_to.token, + }; + + let import_result = match self.invoke(&import_request).await { + Ok(result) => result, + Err(InvocationError::Rpc(err)) if err.name == "SESSION_PASSWORD_NEEDED" => { + // Return PasswordRequired status instead of failing + return Ok(QrLoginInfo { + qr_url: "".to_string(), + expires_unix: 0, + expires_in_seconds: 0, + status: QrLoginStatus::PasswordRequired(None), // We'll get the hint separately + }); + } + Err(e) => return Err(e), + }; + match import_result { + tl::enums::auth::LoginToken::Success(success) => { + // Successfully logged in after migration + match success.authorization { + tl::enums::auth::Authorization::Authorization(auth) => { + // Complete the login and establish session on the new DC + let _user = self.complete_login(auth).await?; + + // Ensure session is properly established on the new DC + // This ensures is_authorized() will work correctly + Ok(QrLoginInfo { + qr_url: "".to_string(), + expires_unix: 0, + expires_in_seconds: 0, + status: QrLoginStatus::Success, + }) + } + tl::enums::auth::Authorization::SignUpRequired(_) => { + Err(InvocationError::Rpc(grammers_mtsender::RpcError { + code: 400, + name: "SIGN_UP_REQUIRED".to_string(), + value: None, + caused_by: None, + })) + } + } + } + tl::enums::auth::LoginToken::Token(token) => { + // Got a new token after migration + let qr_url = self.generate_qr_url(&token.token); + let expires_unix = token.expires as u64; + let current_time = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs(); + let expires_in_seconds = token.expires as i64 - current_time as i64; + + Ok(QrLoginInfo { + qr_url, + expires_unix, + expires_in_seconds, + status: QrLoginStatus::Waiting, + }) + } + tl::enums::auth::LoginToken::MigrateTo(_) => { + // Unexpected double migration + Err(InvocationError::Rpc(grammers_mtsender::RpcError { + code: 400, + name: "UNEXPECTED_MIGRATION".to_string(), + value: None, + caused_by: None, + })) + } + } + } + tl::enums::auth::LoginToken::Success(success) => { + // Already logged in + match success.authorization { + tl::enums::auth::Authorization::Authorization(auth) => self + .complete_login(auth) + .await + .map(|_| QrLoginInfo { + qr_url: "".to_string(), + expires_unix: 0, + expires_in_seconds: 0, + status: QrLoginStatus::Success, + }) + .map_err(|e| e), + tl::enums::auth::Authorization::SignUpRequired(_) => { + Err(InvocationError::Rpc(grammers_mtsender::RpcError { + code: 400, + name: "SIGN_UP_REQUIRED".to_string(), + value: None, + caused_by: None, + })) + } + } + } + } + } + + /// Perform continuous QR login with automatic token refresh + pub async fn start_continuous_qr_login( + &self, + api_id: i32, + api_hash: String, + ) -> Result< + ( + tokio::sync::watch::Sender, + tokio::sync::broadcast::Sender<()>, + tokio::task::JoinHandle>, + ), + InvocationError, + > { + // Create a QR info channel and a cancellation channel + let (qr_info_tx, _) = tokio::sync::watch::channel(QrLoginInfo { + qr_url: "".to_string(), + expires_unix: 0, + expires_in_seconds: 0, + status: QrLoginStatus::Idle, + }); + let (cancellation_tx, _) = tokio::sync::broadcast::channel(1); + + // Create a clone of the client to move into the task + let client_clone = self.clone(); + let qr_info_tx_clone = qr_info_tx.clone(); + let cancellation_tx_clone = cancellation_tx.clone(); + + let join_handle = tokio::spawn(async move { + let mut _current_expiration = 0; + let mut _current_token = Vec::::new(); + // Used to track current token and expiration during QR refresh + + loop { + // Check for cancellation + let mut cancel_subscriber = cancellation_tx_clone.subscribe(); + let cancellation_result = tokio::select! { + _ = cancel_subscriber.recv() => { + return Err(InvocationError::Rpc(grammers_mtsender::RpcError { + code: 406, + name: "LOGIN_CANCELLED".to_string(), + value: None, + caused_by: None, + })); + }, + result = client_clone.export_login_token(api_id, &api_hash) => { + match result { + Ok(token) => Ok(token), + Err(InvocationError::Rpc(err)) if err.name == "SESSION_PASSWORD_NEEDED" => { + // Return error indicating password is required + Err(InvocationError::Rpc(err)) + }, + Err(e) => Err(e), + } + }, + }; + + let login_token = match cancellation_result { + Ok(token) => token, + Err(InvocationError::Rpc(err)) if err.name == "SESSION_PASSWORD_NEEDED" => { + // Propagate the password needed error + return Err(InvocationError::Rpc(err)); + } + Err(e) => return Err(e), + }; + + match login_token { + tl::enums::auth::LoginToken::Token(token) => { + _current_token = token.token.clone(); + _current_expiration = token.expires as u64; + + // Send QR info update + let qr_url = client_clone.generate_qr_url(&token.token); + let expires_unix = token.expires as u64; + let current_time = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs(); + let expires_in_seconds = token.expires as i64 - current_time as i64; + + let qr_info = QrLoginInfo { + qr_url, + expires_unix, + expires_in_seconds, + status: QrLoginStatus::Waiting, + }; + + // Send the QR info update (ignore if no receivers) + let _ = qr_info_tx_clone.send(qr_info); + + // Sleep until 5 seconds before expiration to refresh + let time_until_expiry = if _current_expiration > current_time { + _current_expiration - current_time + } else { + 0 + }; + + // Sleep for 1 second or until close to expiry + let sleep_duration = std::cmp::min(time_until_expiry, 5) as u64; + if sleep_duration > 0 { + let sleep_future = + tokio::time::sleep(Duration::from_secs(sleep_duration)); + let mut cancel_subscriber = cancellation_tx_clone.subscribe(); + tokio::select! { + _ = sleep_future => {}, + _ = cancel_subscriber.recv() => { + return Err(InvocationError::Rpc(grammers_mtsender::RpcError { + code: 406, + name: "LOGIN_CANCELLED".to_string(), + value: None, + caused_by: None, + })); + } + } + } + } + tl::enums::auth::LoginToken::MigrateTo(migrate_to) => { + // Handle migration by switching home DC and importing the token on the new DC + client_clone + .handle_qr_login_migration(migrate_to.dc_id) + .await?; + + let import_request = tl::functions::auth::ImportLoginToken { + token: migrate_to.token, + }; + + match client_clone.invoke(&import_request).await { + Ok(import_result) => { + match import_result { + tl::enums::auth::LoginToken::Success(success) => { + match success.authorization { + tl::enums::auth::Authorization::Authorization(auth) => { + return client_clone.finalize_qr_login(auth).await; + } + tl::enums::auth::Authorization::SignUpRequired(_) => { + return Err(InvocationError::Rpc( + grammers_mtsender::RpcError { + code: 400, + name: "SIGN_UP_REQUIRED".to_string(), + value: None, + caused_by: None, + }, + )); + } + } + } + tl::enums::auth::LoginToken::Token(token) => { + // Got a new token, continue the loop + _current_token = token.token.clone(); + _current_expiration = token.expires as u64; + + // Send QR info update + let qr_url = client_clone.generate_qr_url(&token.token); + let expires_unix = token.expires as u64; + let current_time = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs(); + let expires_in_seconds = + token.expires as i64 - current_time as i64; + + let qr_info = QrLoginInfo { + qr_url, + expires_unix, + expires_in_seconds, + status: QrLoginStatus::Waiting, + }; + + // Send the QR info update (ignore if no receivers) + let _ = qr_info_tx_clone.send(qr_info); + + let current_time = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs(); + let time_until_expiry = + if _current_expiration > current_time { + _current_expiration - current_time + } else { + 0 + }; + + let sleep_duration = + std::cmp::min(time_until_expiry, 5) as u64; + if sleep_duration > 0 { + let sleep_future = tokio::time::sleep( + Duration::from_secs(sleep_duration), + ); + let mut cancel_subscriber = + cancellation_tx_clone.subscribe(); + tokio::select! { + _ = sleep_future => {}, + _ = cancel_subscriber.recv() => { + return Err(InvocationError::Rpc(grammers_mtsender::RpcError { + code: 406, + name: "LOGIN_CANCELLED".to_string(), + value: None, + caused_by: None, + })); + } + } + } + } + tl::enums::auth::LoginToken::MigrateTo(_) => { + // This shouldn't happen, but continue anyway + continue; + } + } + } + Err(InvocationError::Rpc(err)) + if err.name == "SESSION_PASSWORD_NEEDED" => + { + // Return error indicating password is required + return Err(InvocationError::Rpc(err)); + } + Err(e) => return Err(e), + } + } + tl::enums::auth::LoginToken::Success(success) => { + // Login successful + match success.authorization { + tl::enums::auth::Authorization::Authorization(auth) => { + return client_clone.complete_login(auth).await; + } + tl::enums::auth::Authorization::SignUpRequired(_) => { + return Err(InvocationError::Rpc(grammers_mtsender::RpcError { + code: 400, + name: "SIGN_UP_REQUIRED".to_string(), + value: None, + caused_by: None, + })); + } + } + } + } + } + }); + + Ok((qr_info_tx, cancellation_tx, join_handle)) + } +} diff --git a/grammers-client/src/client/mod.rs b/grammers-client/src/client/mod.rs index 0df6942c..2568abf4 100644 --- a/grammers-client/src/client/mod.rs +++ b/grammers-client/src/client/mod.rs @@ -21,7 +21,7 @@ mod net; mod retry_policy; mod updates; -pub use auth::{LoginToken, PasswordToken, SignInError}; +pub use auth::{LoginToken, PasswordToken, QrLoginInfo, QrLoginStatus, SignInError}; pub use bots::{InlineResult, InlineResultIter}; pub use chats::{ParticipantIter, ParticipantPermissions, ProfilePhotoIter}; pub(crate) use client::ClientInner; diff --git a/grammers-client/tests/qr_login_test.rs b/grammers-client/tests/qr_login_test.rs new file mode 100644 index 00000000..a7f04a39 --- /dev/null +++ b/grammers-client/tests/qr_login_test.rs @@ -0,0 +1,668 @@ +use base64::Engine; +use grammers_client::Client; +use grammers_client::peer::User; +use grammers_mtsender::SenderPool; +use grammers_session::storages::MemorySession; +use std::sync::Arc; + +#[tokio::test] +#[ignore] +async fn test_qr_login_functionality() { + let api_id = std::env::var("API_ID") + .expect("API_ID must be set") + .parse::() + .expect("API_ID must be a valid integer"); + let api_hash = std::env::var("API_HASH").expect("API_HASH must be set"); + + // Create a session + let session = Arc::new(MemorySession::default()); + + // Create a sender pool + let SenderPool { runner, handle, .. } = SenderPool::new(Arc::clone(&session), api_id); + let client = Client::new(handle); + + // Spawn the runner + let _runner_handle = tokio::spawn(runner.run()); + + // Wait for QR login with a timeout of 2 minutes + let user = wait_for_qr_login( + &client, + api_id, + &api_hash, + std::time::Duration::from_secs(120), + ) + .await + .expect("Failed to complete QR login within timeout"); + + // Verify successful login + println!( + "Successfully logged in as: {:?}", + user.first_name().unwrap_or("Unknown") + ); + + // Make a test API call to ensure the connection is established on the new DC + // This helps ensure the client is properly authenticated on the migrated DC + match client + .invoke(&grammers_client::tl::functions::updates::GetState {}) + .await + { + Ok(_) => println!("GetState call succeeded, connection established"), + Err(e) => println!("GetState call failed: {}", e), + } + + // Check authorization status immediately + let is_auth_immediate = client.is_authorized().await.unwrap_or(false); + println!("Immediate authorization status: {}", is_auth_immediate); + + // Allow time for session synchronization after migration + for i in 0..10 { + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + let is_authorized = client.is_authorized().await.unwrap_or(false); + println!("Attempt {}: is_authorized = {}", i + 1, is_authorized); + + if is_authorized { + println!("Authorization confirmed after {} attempts", i + 1); + break; + } + + if i == 9 { + println!("Authorization still not detected after 10 attempts"); + } + } + + let final_auth_status = client.is_authorized().await.unwrap(); + println!("Final authorization status: {}", final_auth_status); + assert!(final_auth_status); + + println!("QR login completed successfully!"); +} + +use std::time::Duration; +use tokio::time::Instant; + +async fn wait_for_qr_login( + client: &Client, + api_id: i32, + api_hash: &str, + max_wait: Duration, +) -> Result> { + println!("Starting QR login process..."); + + // Get initial QR token + let mut qr_info = client.start_qr_login(api_id, api_hash).await?; + println!("QR URL: {}", qr_info.qr_url); + println!("Expires in: {} seconds", qr_info.expires_in_seconds); + + let start_time = Instant::now(); + let mut current_token = Vec::::new(); + let mut loop_count = 0; + + loop { + // Check for timeout + if start_time.elapsed() >= max_wait { + return Err(Box::new(grammers_mtsender::InvocationError::Rpc( + grammers_mtsender::RpcError { + code: 408, + name: "TIMEOUT".to_string(), + value: None, + caused_by: None, + }, + ))); + } + + // Check if we're close to expiration, refresh if needed + if qr_info.expires_in_seconds <= 5 { + println!("Refreshing QR token..."); + match client.export_login_token(api_id, api_hash).await { + Ok(login_token) => { + match login_token { + grammers_tl_types::enums::auth::LoginToken::Token(token) => { + let new_qr_url = format!( + "tg://login?token={}", + base64::engine::general_purpose::URL_SAFE_NO_PAD + .encode(&token.token) + ); + println!("New QR URL: {}", new_qr_url); + + let expires_unix = token.expires as u64; + let current_time = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs(); + let expires_in_seconds = token.expires as i64 - current_time as i64; + + qr_info.qr_url = new_qr_url; + qr_info.expires_unix = expires_unix; + qr_info.expires_in_seconds = expires_in_seconds; + + current_token = token.token.clone(); + } + grammers_tl_types::enums::auth::LoginToken::Success(success) => { + println!("Login successful via refresh!"); + match success.authorization { + grammers_tl_types::enums::auth::Authorization::Authorization( + auth, + ) => { + return client + .finalize_qr_login(auth) + .await + .map_err(|e| Box::new(e) as Box); + } + grammers_tl_types::enums::auth::Authorization::SignUpRequired( + _, + ) => { + return Err(Box::new(grammers_mtsender::InvocationError::Rpc( + grammers_mtsender::RpcError { + code: 400, + name: "SIGN_UP_REQUIRED".to_string(), + value: None, + caused_by: None, + }, + ))); + } + } + } + grammers_tl_types::enums::auth::LoginToken::MigrateTo(migrate_to) => { + println!("Handling DC migration to DC {}", migrate_to.dc_id); + // Handle migration by importing the token on the new DC + match client + .import_login_token(migrate_to.token, migrate_to.dc_id) + .await + { + Ok(import_result) => { + match import_result { + grammers_tl_types::enums::auth::LoginToken::Success( + success, + ) => { + println!("Login successful after migration!"); + match success.authorization { + grammers_tl_types::enums::auth::Authorization::Authorization(auth) => { + return client.finalize_qr_login(auth).await.map_err(|e| Box::new(e) as Box); + } + grammers_tl_types::enums::auth::Authorization::SignUpRequired(_) => { + return Err(Box::new(grammers_mtsender::InvocationError::Rpc(grammers_mtsender::RpcError { + code: 400, + name: "SIGN_UP_REQUIRED".to_string(), + value: None, + caused_by: None, + }))); + } + } + } + grammers_tl_types::enums::auth::LoginToken::Token( + token, + ) => { + let new_qr_url = format!( + "tg://login?token={}", + base64::engine::general_purpose::URL_SAFE_NO_PAD + .encode(&token.token) + ); + println!("New QR URL after migration: {}", new_qr_url); + + let expires_unix = token.expires as u64; + let current_time = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs(); + let expires_in_seconds = + token.expires as i64 - current_time as i64; + + qr_info.qr_url = new_qr_url; + qr_info.expires_unix = expires_unix; + qr_info.expires_in_seconds = expires_in_seconds; + + current_token = token.token.clone(); + } + grammers_tl_types::enums::auth::LoginToken::MigrateTo( + _, + ) => { + println!("Unexpected double migration"); + // Continue loop + } + } + } + Err(grammers_mtsender::InvocationError::Rpc(ref err)) + if err.name == "SESSION_PASSWORD_NEEDED" => + { + // Handle 2FA required case after migration + println!( + "2FA password required after migration. Getting password information..." + ); + match client.qr_get_password_token().await { + Ok(password_token) => { + if let Some(hint) = password_token.hint() { + println!("Password hint: {}", hint); + } + + // Try to get password from environment variable + match std::env::var("TG_2FA_PASSWORD") { + Ok(password) => { + println!( + "Using 2FA password from environment variable" + ); + // Complete login with password + match client + .check_password( + password_token, + password.as_bytes(), + ) + .await + { + Ok(user) => { + println!( + "Successfully logged in with 2FA after migration" + ); + return Ok(user); + } + Err(sign_in_error) => match sign_in_error { + grammers_client::SignInError::Other( + invocation_error, + ) => { + println!( + "Failed to complete login with 2FA after migration: {}", + invocation_error + ); + return Err(Box::new( + invocation_error, + )); + } + _ => { + println!( + "Failed to complete login with 2FA after migration: {}", + sign_in_error + ); + return Err(Box::new( + sign_in_error, + )); + } + }, + } + } + Err(_) => { + println!( + "TG_2FA_PASSWORD environment variable not set. Please set it to complete 2FA login." + ); + return Err(Box::new( + grammers_mtsender::InvocationError::Rpc( + grammers_mtsender::RpcError { + code: 401, + name: "SESSION_PASSWORD_NEEDED" + .to_string(), + value: None, + caused_by: None, + }, + ), + )); + } + } + } + Err(get_pw_error) => { + println!( + "Failed to get password information: {}", + get_pw_error + ); + return Err(Box::new(get_pw_error)); + } + } + } + Err(e) => { + println!("Failed to import login token after migration: {}", e); + return Err(Box::new(e)); + } + } + } + } + } + Err(grammers_mtsender::InvocationError::Rpc(ref err)) + if err.name == "SESSION_PASSWORD_NEEDED" => + { + // Handle 2FA required case + println!("2FA password required. Getting password information..."); + match client.get_password_token().await { + Ok(password_token) => { + if let Some(hint) = password_token.hint() { + println!("Password hint: {}", hint); + } + + // Try to get password from environment variable + match std::env::var("TG_2FA_PASSWORD") { + Ok(password) => { + println!("Using 2FA password from environment variable"); + // Complete login with password + match client + .check_password(password_token, password.as_bytes()) + .await + { + Ok(user) => { + println!("Successfully logged in with 2FA"); + return Ok(user); + } + Err(sign_in_error) => match sign_in_error { + grammers_client::SignInError::Other( + invocation_error, + ) => { + println!( + "Failed to complete login with 2FA: {}", + invocation_error + ); + return Err(Box::new(invocation_error)); + } + _ => { + println!( + "Failed to complete login with 2FA: {}", + sign_in_error + ); + return Err(Box::new(sign_in_error)); + } + }, + } + } + Err(_) => { + println!( + "TG_2FA_PASSWORD environment variable not set. Please set it to complete 2FA login." + ); + return Err(Box::new(grammers_mtsender::InvocationError::Rpc( + grammers_mtsender::RpcError { + code: 401, + name: "SESSION_PASSWORD_NEEDED".to_string(), + value: None, + caused_by: None, + }, + ))); + } + } + } + Err(get_pw_error) => { + println!("Failed to get password information: {}", get_pw_error); + return Err(Box::new(get_pw_error)); + } + } + } + Err(e) => { + println!("Failed to refresh token: {}", e); + return Err(Box::new(e)); + } + } + } + + // Poll for login status + match client.export_login_token(api_id, api_hash).await { + Ok(login_token) => { + match login_token { + grammers_tl_types::enums::auth::LoginToken::Token(token) => { + // Update QR info if token has changed + let token_bytes = &token.token; + if token_bytes != ¤t_token { + let new_qr_url = format!( + "tg://login?token={}", + base64::engine::general_purpose::URL_SAFE_NO_PAD + .encode(token_bytes) + ); + println!("QR token updated: {}", new_qr_url); + + let expires_unix = token.expires as u64; + let current_time = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs(); + let expires_in_seconds = token.expires as i64 - current_time as i64; + + qr_info.qr_url = new_qr_url; + qr_info.expires_unix = expires_unix; + qr_info.expires_in_seconds = expires_in_seconds; + + current_token = token.token.clone(); + } + } + grammers_tl_types::enums::auth::LoginToken::Success(success) => { + println!("Login successful!"); + match success.authorization { + grammers_tl_types::enums::auth::Authorization::Authorization(auth) => { + return client + .finalize_qr_login(auth) + .await + .map_err(|e| Box::new(e) as Box); + } + grammers_tl_types::enums::auth::Authorization::SignUpRequired(_) => { + return Err(Box::new(grammers_mtsender::InvocationError::Rpc( + grammers_mtsender::RpcError { + code: 400, + name: "SIGN_UP_REQUIRED".to_string(), + value: None, + caused_by: None, + }, + ))); + } + } + } + grammers_tl_types::enums::auth::LoginToken::MigrateTo(migrate_to) => { + println!("Detected DC migration to DC {}", migrate_to.dc_id); + // Handle migration by importing the token on the new DC + match client + .import_login_token(migrate_to.token, migrate_to.dc_id) + .await + { + Ok(import_result) => { + match import_result { + grammers_tl_types::enums::auth::LoginToken::Success( + success, + ) => { + println!("Login successful after migration!"); + match success.authorization { + grammers_tl_types::enums::auth::Authorization::Authorization(auth) => { + return client.finalize_qr_login(auth).await.map_err(|e| Box::new(e) as Box); + }, + grammers_tl_types::enums::auth::Authorization::SignUpRequired(_) => { + return Err(Box::new(grammers_mtsender::InvocationError::Rpc(grammers_mtsender::RpcError { + code: 400, + name: "SIGN_UP_REQUIRED".to_string(), + value: None, + caused_by: None, + }))); + } + } + } + grammers_tl_types::enums::auth::LoginToken::Token(token) => { + let new_qr_url = format!( + "tg://login?token={}", + base64::engine::general_purpose::URL_SAFE_NO_PAD + .encode(&token.token) + ); + println!("New QR URL after migration: {}", new_qr_url); + + let expires_unix = token.expires as u64; + let current_time = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs(); + let expires_in_seconds = + token.expires as i64 - current_time as i64; + + qr_info.qr_url = new_qr_url; + qr_info.expires_unix = expires_unix; + qr_info.expires_in_seconds = expires_in_seconds; + + current_token = token.token.clone(); + } + grammers_tl_types::enums::auth::LoginToken::MigrateTo(_) => { + println!("Unexpected double migration"); + // Continue loop + } + } + } + Err(grammers_mtsender::InvocationError::Rpc(ref err)) + if err.name == "SESSION_PASSWORD_NEEDED" => + { + // Handle 2FA required case after migration + println!( + "2FA password required after migration. Getting password information..." + ); + match client.get_password_token().await { + Ok(password_token) => { + if let Some(hint) = password_token.hint() { + println!("Password hint: {}", hint); + } + + // Try to get password from environment variable + match std::env::var("TG_2FA_PASSWORD") { + Ok(password) => { + println!( + "Using 2FA password from environment variable" + ); + // Complete login with password + match client + .check_password( + password_token, + password.as_bytes(), + ) + .await + { + Ok(user) => { + println!( + "Successfully logged in with 2FA after migration" + ); + return Ok(user); + } + Err(sign_in_error) => match sign_in_error { + grammers_client::SignInError::Other( + invocation_error, + ) => { + println!( + "Failed to complete login with 2FA after migration: {}", + invocation_error + ); + return Err(Box::new(invocation_error)); + } + _ => { + println!( + "Failed to complete login with 2FA after migration: {}", + sign_in_error + ); + return Err(Box::new(sign_in_error)); + } + }, + } + } + Err(_) => { + println!( + "TG_2FA_PASSWORD environment variable not set. Please set it to complete 2FA login." + ); + return Err(Box::new( + grammers_mtsender::InvocationError::Rpc( + grammers_mtsender::RpcError { + code: 401, + name: "SESSION_PASSWORD_NEEDED" + .to_string(), + value: None, + caused_by: None, + }, + ), + )); + } + } + } + Err(get_pw_error) => { + println!( + "Failed to get password information: {}", + get_pw_error + ); + return Err(Box::new(get_pw_error)); + } + } + } + Err(e) => { + println!("Failed to import login token after migration: {}", e); + return Err(Box::new(e)); + } + } + } + } + } + Err(grammers_mtsender::InvocationError::Rpc(ref err)) + if err.name == "SESSION_PASSWORD_NEEDED" => + { + // Handle 2FA required case + println!("2FA password required. Getting password information..."); + match client.get_password_token().await { + Ok(password_token) => { + if let Some(hint) = password_token.hint() { + println!("Password hint: {}", hint); + } + + // Try to get password from environment variable + match std::env::var("TG_2FA_PASSWORD") { + Ok(password) => { + println!("Using 2FA password from environment variable"); + // Complete login with password + match client + .check_password(password_token, password.as_bytes()) + .await + { + Ok(user) => { + println!("Successfully logged in with 2FA"); + return Ok(user); + } + Err(sign_in_error) => match sign_in_error { + grammers_client::SignInError::Other(invocation_error) => { + println!( + "Failed to complete login with 2FA: {}", + invocation_error + ); + return Err(Box::new(invocation_error)); + } + _ => { + println!( + "Failed to complete login with 2FA: {}", + sign_in_error + ); + return Err(Box::new(sign_in_error)); + } + }, + } + } + Err(_) => { + println!( + "TG_2FA_PASSWORD environment variable not set. Please set it to complete 2FA login." + ); + return Err(Box::new(grammers_mtsender::InvocationError::Rpc( + grammers_mtsender::RpcError { + code: 401, + name: "SESSION_PASSWORD_NEEDED".to_string(), + value: None, + caused_by: None, + }, + ))); + } + } + } + Err(get_pw_error) => { + println!("Failed to get password information: {}", get_pw_error); + return Err(Box::new(get_pw_error)); + } + } + } + Err(e) => { + println!("Error polling for login status: {}", e); + return Err(Box::new(e)); + } + } + + // Sleep briefly before next check + tokio::time::sleep(Duration::from_secs(1)).await; + loop_count += 1; + + // Update expiration countdown + let elapsed = start_time.elapsed().as_secs(); + let remaining = max_wait.as_secs().saturating_sub(elapsed); + if loop_count % 5 == 0 { + // Print every 5 seconds + println!("Still waiting... {} seconds remaining", remaining); + } + } +} + +// Add a comment about how to run this test +/* +To run this test: +cargo test test_qr_login_functionality -- --ignored --nocapture +*/