use crate::AppService; use crate::error::AppError; use models::users::{user_2fa, user_activity_log, user_password}; use rand::RngExt; use redis::AsyncCommands; use sea_orm::*; use serde::{Deserialize, Serialize}; use serde_json::json; use session::Session; use uuid::Uuid; #[derive(Deserialize, Serialize, Clone, Debug, utoipa::ToSchema)] pub struct Enable2FAResponse { pub secret: String, pub qr_code: String, pub backup_codes: Vec, } #[derive(Deserialize, Serialize, Clone, Debug, utoipa::ToSchema)] pub struct Verify2FAParams { pub code: String, } #[derive(Deserialize, Serialize, Clone, Debug, utoipa::ToSchema)] pub struct Disable2FAParams { pub code: String, pub password: String, } #[derive(Deserialize, Serialize, Clone, Debug, utoipa::ToSchema)] pub struct Get2FAStatusResponse { pub is_enabled: bool, pub method: Option, pub has_backup_codes: bool, } impl AppService { pub async fn auth_2fa_enable(&self, context: &Session) -> Result { let user_uid = context.user().ok_or(AppError::Unauthorized)?; let user = self.utils_find_user_by_uid(user_uid).await?; let existing_2fa = user_2fa::Entity::find_by_id(user_uid).one(&self.db).await?; if let Some(ref existing) = existing_2fa { if existing.is_enabled { return Err(AppError::TwoFactorAlreadyEnabled); } } let secret = self.generate_totp_secret(); let backup_codes = self.generate_backup_codes(10); let issuer = "GitDataAI"; let account_name = format!("{}:{}", issuer, user.username); let qr_data = format!( "otpauth://totp/{}?secret={}&issuer={}", account_name, secret, issuer ); let now = chrono::Utc::now(); let model = user_2fa::ActiveModel { user: Set(user_uid), method: Set("totp".to_string()), secret: Set(Some(secret.clone())), backup_codes: Set(serde_json::json!(backup_codes)), is_enabled: Set(false), created_at: Set(now), updated_at: Set(now), }; if existing_2fa.is_some() { model.update(&self.db).await?; } else { model.insert(&self.db).await?; } slog::info!(self.logs, "2FA setup initiated"; "user_uid" => %user_uid); Ok(Enable2FAResponse { secret, qr_code: qr_data, backup_codes, }) } pub async fn auth_2fa_verify_and_enable( &self, context: &Session, params: Verify2FAParams, ) -> Result<(), AppError> { let user_uid = context.user().ok_or(AppError::Unauthorized)?; let two_fa = user_2fa::Entity::find_by_id(user_uid) .one(&self.db) .await? .ok_or(AppError::TwoFactorNotSetup)?; if two_fa.is_enabled { return Err(AppError::TwoFactorAlreadyEnabled); } let secret = two_fa.secret.as_ref().ok_or(AppError::TwoFactorNotSetup)?; if !self.verify_totp_code(secret, ¶ms.code)? { slog::warn!(self.logs, "2FA verification failed during setup"; "user_uid" => %user_uid, "ip" => context.ip_address()); return Err(AppError::InvalidTwoFactorCode); } let mut active_model: user_2fa::ActiveModel = two_fa.into(); active_model.is_enabled = Set(true); active_model.updated_at = Set(chrono::Utc::now()); active_model.update(&self.db).await?; slog::info!(self.logs, "2FA enabled"; "user_uid" => %user_uid, "ip" => context.ip_address()); let _ = user_activity_log::ActiveModel { user_uid: Set(Option::from(user_uid)), action: Set("2fa_enabled".to_string()), ip_address: Set(context.ip_address()), user_agent: Set(context.user_agent()), details: Set(serde_json::json!({ "method": "totp" })), created_at: Set(chrono::Utc::now()), ..Default::default() } .insert(&self.db) .await; Ok(()) } pub async fn auth_2fa_disable( &self, context: &Session, params: Disable2FAParams, ) -> Result<(), AppError> { let user_uid = context.user().ok_or(AppError::Unauthorized)?; let password = self.auth_rsa_decode(context, params.password).await?; self.verify_user_password(user_uid, &password).await?; let two_fa = user_2fa::Entity::find_by_id(user_uid) .one(&self.db) .await? .ok_or(AppError::TwoFactorNotSetup)?; if !two_fa.is_enabled { return Err(AppError::TwoFactorNotEnabled); } let secret = two_fa.secret.as_ref().ok_or(AppError::TwoFactorNotSetup)?; let backup_codes: Vec = serde_json::from_value(two_fa.backup_codes.clone()).unwrap_or_default(); let is_valid = self.verify_totp_code(secret, ¶ms.code)? || backup_codes.contains(¶ms.code); if !is_valid { return Err(AppError::InvalidTwoFactorCode); } user_2fa::Entity::delete_by_id(user_uid) .exec(&self.db) .await?; slog::info!(self.logs, "2FA disabled"; "user_uid" => %user_uid, "ip" => context.ip_address()); let _ = user_activity_log::ActiveModel { user_uid: Set(Some(user_uid)), action: Set("2fa_disabled".to_string()), ip_address: Set(context.ip_address()), user_agent: Set(context.user_agent()), details: Set(json!({})), created_at: Set(chrono::Utc::now()), ..Default::default() } .insert(&self.db) .await; Ok(()) } pub async fn auth_2fa_verify(&self, user_uid: Uuid, code: &str) -> Result { let two_fa = user_2fa::Entity::find_by_id(user_uid).one(&self.db).await?; let Some(two_fa) = two_fa else { return Ok(true); }; if !two_fa.is_enabled { return Ok(true); } let secret = two_fa.secret.as_ref().ok_or(AppError::TwoFactorNotSetup)?; if self.verify_totp_code(secret, code)? { return Ok(true); } let mut backup_codes: Vec = serde_json::from_value(two_fa.backup_codes.clone()).unwrap_or_default(); if backup_codes.contains(&code.to_string()) { backup_codes.retain(|c| c != code); let mut active_model: user_2fa::ActiveModel = user_2fa::Entity::find_by_id(user_uid) .one(&self.db) .await? .ok_or(AppError::TwoFactorNotSetup)? .into(); active_model.backup_codes = Set(serde_json::json!(backup_codes)); active_model.updated_at = Set(chrono::Utc::now()); active_model.update(&self.db).await?; return Ok(true); } Ok(false) } pub async fn auth_2fa_status( &self, context: &Session, ) -> Result { let user_uid = context.user().ok_or(AppError::Unauthorized)?; let two_fa = user_2fa::Entity::find_by_id(user_uid).one(&self.db).await?; match two_fa { Some(fa) => { let backup_codes: Vec = serde_json::from_value(fa.backup_codes).unwrap_or_default(); Ok(Get2FAStatusResponse { is_enabled: fa.is_enabled, method: Some(fa.method), has_backup_codes: !backup_codes.is_empty(), }) } None => Ok(Get2FAStatusResponse { is_enabled: false, method: None, has_backup_codes: false, }), } } pub async fn auth_2fa_verify_login( &self, context: &Session, code: &str, ) -> Result { let totp_key: String = context .get::(Self::TOTP_KEY) .ok() .flatten() .ok_or(AppError::TwoFactorNotSetup)?; if let Ok(mut conn) = self.cache.conn().await { let stored_user_uid: Option = conn.get(totp_key.as_str()).await.ok(); if let Some(user_uid_str) = stored_user_uid { let user_uid = Uuid::parse_str(&user_uid_str).map_err(|_| AppError::UserNotFound)?; let two_fa = user_2fa::Entity::find_by_id(user_uid).one(&self.db).await?; if let Some(two_fa) = two_fa { if two_fa.is_enabled { let secret = two_fa.secret.as_ref().ok_or(AppError::TwoFactorNotSetup)?; if self.verify_totp_code(secret, code)? { let _: Option<()> = conn.del(totp_key.as_str()).await.ok(); slog::info!(self.logs, "2FA verification succeeded during login"; "user_uid" => %user_uid, "ip" => context.ip_address()); return Ok(true); } else { slog::warn!(self.logs, "2FA verification failed during login"; "user_uid" => %user_uid, "ip" => context.ip_address()); } } } } } Ok(false) } pub async fn auth_2fa_regenerate_backup_codes( &self, context: &Session, password: String, ) -> Result, AppError> { let user_uid = context.user().ok_or(AppError::Unauthorized)?; let password = self.auth_rsa_decode(context, password).await?; self.verify_user_password(user_uid, &password).await?; let two_fa = user_2fa::Entity::find_by_id(user_uid) .one(&self.db) .await? .ok_or(AppError::TwoFactorNotSetup)?; if !two_fa.is_enabled { return Err(AppError::TwoFactorNotEnabled); } let backup_codes = self.generate_backup_codes(10); let mut active_model: user_2fa::ActiveModel = two_fa.into(); active_model.backup_codes = Set(serde_json::json!(backup_codes)); active_model.updated_at = Set(chrono::Utc::now()); active_model.update(&self.db).await?; let _ = user_activity_log::ActiveModel { user_uid: Set(Some(user_uid)), action: Set("2fa_backup_codes_regenerated".to_string()), ip_address: Set(context.ip_address()), user_agent: Set(context.user_agent()), details: Set(json!({})), created_at: Set(chrono::Utc::now()), ..Default::default() } .insert(&self.db) .await; slog::info!(self.logs, "2FA backup codes regenerated"; "user_uid" => %user_uid, "ip" => context.ip_address()); Ok(backup_codes) } fn generate_totp_secret(&self) -> String { const CHARSET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZ234567"; #[allow(deprecated)] let mut rng = rand::rng(); (0..32) .map(|_| { #[allow(deprecated)] let idx = rng.random_range(0..CHARSET.len()); CHARSET[idx] as char }) .collect() } fn generate_backup_codes(&self, count: usize) -> Vec { #[allow(deprecated)] let mut rng = rand::rng(); (0..count) .map(|_| { format!( "{:04}-{:04}-{:04}", rng.random_range(0..10000), rng.random_range(0..10000), rng.random_range(0..10000) ) }) .collect() } fn verify_totp_code(&self, secret: &str, code: &str) -> Result { let now = chrono::Utc::now().timestamp() as u64; let time_step = 30; let counter = now / time_step; for offset in [-1i64, 0, 1] { let test_counter = (counter as i64 + offset) as u64; let expected_code = self.generate_totp_code(secret, test_counter)?; if expected_code == code { return Ok(true); } } Ok(false) } fn generate_totp_code(&self, secret: &str, counter: u64) -> Result { use hmac::{Hmac, Mac}; use sha1::Sha1; let secret_bytes = self.decode_base32(secret)?; let counter_bytes = counter.to_be_bytes(); let mut mac = Hmac::::new_from_slice(&secret_bytes) .map_err(|_| AppError::InvalidTwoFactorCode)?; mac.update(&counter_bytes); let result = mac.finalize().into_bytes(); let offset = (result[19] & 0x0f) as usize; let code = u32::from_be_bytes([ result[offset] & 0x7f, result[offset + 1], result[offset + 2], result[offset + 3], ]); Ok(format!("{:06}", code % 1_000_000)) } fn decode_base32(&self, input: &str) -> Result, AppError> { const CHARSET: &str = "ABCDEFGHIJKLMNOPQRSTUVWXYZ234567"; let input = input.to_uppercase().replace("=", ""); let mut bits = 0u64; let mut bit_count = 0; let mut output = Vec::new(); for c in input.chars() { let val = CHARSET.find(c).ok_or(AppError::InvalidTwoFactorCode)? as u64; bits = (bits << 5) | val; bit_count += 5; if bit_count >= 8 { bit_count -= 8; output.push((bits >> bit_count) as u8); bits &= (1 << bit_count) - 1; } } Ok(output) } async fn verify_user_password(&self, user_uid: Uuid, password: &str) -> Result<(), AppError> { use argon2::{Argon2, PasswordHash, PasswordVerifier}; let user_password = user_password::Entity::find() .filter(user_password::Column::User.eq(user_uid)) .one(&self.db) .await? .ok_or(AppError::UserNotFound)?; let password_hash = PasswordHash::new(&user_password.password_hash) .map_err(|_| AppError::InvalidPassword)?; Argon2::default() .verify_password(password.as_bytes(), &password_hash) .map_err(|_| AppError::InvalidPassword)?; Ok(()) } }