gitdataai/libs/service/auth/totp.rs

use sha2::{Sha256, Digest};

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<String>,
}

#[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<String>,
    pub has_backup_codes: bool,
}

impl AppService {
    pub async fn auth_2fa_enable(&self, context: &Session) -> Result<Enable2FAResponse, AppError> {
        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!(Self::hash_backup_codes(&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?;
        }

        tracing::info!(user_uid = %user_uid, "2FA setup initiated");
        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, &params.code)? {
            tracing::warn!(user_uid = %user_uid, ip = ?context.ip_address(), "2FA verification failed during setup");
            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?;

        tracing::info!(user_uid = %user_uid, ip = ?context.ip_address(), "2FA enabled");
        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<String> =
            serde_json::from_value(two_fa.backup_codes.clone()).unwrap_or_default();
        let hashed_code = Self::hash_backup_code(&params.code);
        let is_valid =
            self.verify_totp_code(secret, &params.code)? || backup_codes.contains(&hashed_code);

        if !is_valid {
            return Err(AppError::InvalidTwoFactorCode);
        }

        user_2fa::Entity::delete_by_id(user_uid)
            .exec(&self.db)
            .await?;

        tracing::info!(user_uid = %user_uid, ip = ?context.ip_address(), "2FA disabled");
        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<bool, AppError> {
        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<String> =
            serde_json::from_value(two_fa.backup_codes.clone()).unwrap_or_default();
        let hashed_code = Self::hash_backup_code(code);
        if backup_codes.contains(&hashed_code) {
            backup_codes.retain(|c| c != &hashed_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!(Self::hash_backup_codes(&backup_codes)));
            active_model.updated_at = Set(chrono::Utc::now());
            active_model.update(&self.db).await?;

            return Ok(true);
        }

        Ok(false)
    }

    /// Look up 2FA status by explicit user_uid. Used in login flow where session.user is not set yet.
    pub async fn auth_2fa_status_by_uid(
        &self,
        user_uid: Uuid,
    ) -> Result<Get2FAStatusResponse, AppError> {
        let two_fa = user_2fa::Entity::find_by_id(user_uid).one(&self.db).await?;

        match two_fa {
            Some(fa) => {
                let backup_codes: Vec<String> =
                    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,
            }),
        }
    }

    /// Look up 2FA status from session context (requires authenticated user).
    pub async fn auth_2fa_status(
        &self,
        context: &Session,
    ) -> Result<Get2FAStatusResponse, AppError> {
        let user_uid = context.user().ok_or(AppError::Unauthorized)?;
        self.auth_2fa_status_by_uid(user_uid).await
    }

    pub async fn auth_2fa_verify_login(
        &self,
        context: &Session,
        code: &str,
    ) -> Result<bool, AppError> {
        let totp_key: String = context
            .get::<String>(Self::TOTP_KEY)
            .ok()
            .flatten()
            .ok_or(AppError::TwoFactorNotSetup)?;

        if let Ok(mut conn) = self.cache.conn().await {
            let stored_user_uid: Option<String> = 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();
                            tracing::info!(user_uid = %user_uid, ip = ?context.ip_address(), "2FA verification succeeded during login");
                            return Ok(true);
                        } else {
                            tracing::warn!(user_uid = %user_uid, ip = ?context.ip_address(), "2FA verification failed during login");
                        }
                    }
                }
            }
        }

        Ok(false)
    }

    pub async fn auth_2fa_regenerate_backup_codes(
        &self,
        context: &Session,
        password: String,
    ) -> Result<Vec<String>, 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;

        tracing::info!(user_uid = %user_uid, ip = ?context.ip_address(), "2FA backup codes regenerated");
        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<String> {
        #[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 hash_backup_code(code: &str) -> String {
        let mut hasher = Sha256::new();
        hasher.update(code.as_bytes());
        hasher.finalize().iter().map(|b| format!("{:02x}", b)).collect::<String>()
    }

    fn hash_backup_codes(codes: &[String]) -> Vec<String> {
        codes.iter().map(|c| Self::hash_backup_code(c)).collect()
    }

    fn verify_totp_code(&self, secret: &str, code: &str) -> Result<bool, AppError> {
        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<String, AppError> {
        use hmac::{Hmac, Mac, KeyInit};
        use sha1::Sha1;

        let secret_bytes = self.decode_base32(secret)?;

        let counter_bytes = counter.to_be_bytes();

        let mut mac = Hmac::<Sha1>::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<Vec<u8>, 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(())
    }
}