gitdataai/libs/service/auth/login.rs

use crate::AppService;
use crate::error::AppError;
use argon2::{Argon2, PasswordHash, PasswordVerifier};
use models::users::{user_activity_log, user_password};
use rand::RngExt;
use redis::AsyncCommands;
use sea_orm::*;
use serde::{Deserialize, Serialize};
use session::Session;
use sha1::Digest;

#[derive(Deserialize, Serialize, Clone, Debug, utoipa::ToSchema)]
pub struct LoginParams {
    pub username: String,
    pub password: String,
    pub captcha: String,
    pub totp_code: Option<String>,
}

impl AppService {
    pub const TOTP_KEY: &'static str = "totp_key";
    pub async fn auth_login(&self, params: LoginParams, context: Session) -> Result<(), AppError> {
        self.auth_check_captcha(&context, params.captcha).await?;
        let password = self.auth_rsa_decode(&context, params.password).await?;
        let user = match self
            .utils_find_user_by_username(params.username.clone())
            .await
        {
            Ok(user) => user,
            Err(_) => {
                self.utils_find_user_by_email(params.username.clone())
                    .await?
            }
        };
        let user_password = user_password::Entity::find()
            .filter(user_password::Column::User.eq(user.uid))
            .one(&self.db)
            .await
            .ok()
            .flatten()
            .ok_or(AppError::UserNotFound)?;
        let password_hash =
            PasswordHash::new(&user_password.password_hash).map_err(|_| AppError::UserNotFound)?;

        if let Err(_e) = Argon2::default().verify_password(password.as_bytes(), &password_hash) {
            slog::warn!(self.logs, "Login failed: invalid password"; "username" => &params.username, "ip" => context.ip_address());
            return Err(AppError::UserNotFound);
        }

        let needs_totp_verification = context
            .get::<String>(Self::TOTP_KEY)
            .ok()
            .flatten()
            .is_some();

        if needs_totp_verification {
            if let Some(ref totp_code) = params.totp_code {
                if !self.auth_2fa_verify_login(&context, totp_code).await? {
                    slog::warn!(self.logs, "Login failed: invalid 2FA code"; "username" => &params.username, "ip" => context.ip_address());
                    return Err(AppError::InvalidTwoFactorCode);
                }
            }
        } else if !self.auth_2fa_status_by_uid(user.uid).await?.is_enabled {
            let user_uid = user.uid;
            let mut rng = rand::rng();
            let mut sha = sha1::Sha1::default();
            for _ in 0..5 {
                sha.update(
                    (0..1024)
                        .map(|_| {
                            format!(
                                "{:04}-{:04}-{:04}",
                                rng.random_range(0..10000),
                                rng.random_range(0..10000),
                                rng.random_range(0..10000)
                            )
                        })
                        .collect::<String>()
                        .as_bytes(),
                )
            }
            let key = format!("{:?}", sha.finalize());
            context.insert(Self::TOTP_KEY, key.clone()).ok();
            if let Ok(mut conn) = self.cache.conn().await {
                conn.set_ex::<String, String, ()>(key, user_uid.to_string(), 60 * 5)
                    .await
                    .ok();
            }
            slog::info!(self.logs, "Login 2FA triggered for new 2FA user"; "username" => &params.username, "ip" => context.ip_address());
            return Err(AppError::TwoFactorRequired);
        }

        let mut arch = user.clone().into_active_model();
        arch.last_sign_in_at = Set(Some(chrono::Utc::now()));
        arch.update(&self.db)
            .await
            .map_err(|_| AppError::UserNotFound)?;

        let _ = user_activity_log::ActiveModel {
            user_uid: Set(Some(user.uid)),
            action: Set("login".to_string()),
            ip_address: Set(context.ip_address()),
            user_agent: Set(context.user_agent()),
            details: Set(Some(serde_json::json!({
                "method": "password",
                "username": user.username,
                "2fa_used": params.totp_code.is_some()
            }))
            .into()),
            created_at: Set(chrono::Utc::now()),
            ..Default::default()
        }
        .insert(&self.db)
        .await;

        context.set_user(user.uid);
        context.remove(Self::RSA_PRIVATE_KEY);
        context.remove(Self::RSA_PUBLIC_KEY);
        slog::info!(self.logs, "User logged in successfully"; "user_uid" => %user.uid, "username" => &user.username, "ip" => context.ip_address(), "2fa_used" => params.totp_code.is_some());
        Ok(())
    }
}