appks/service/auth/totp.rs

use argon2::{Argon2, PasswordHash, password_hash::PasswordVerifier};
use hmac::{Hmac, Mac};
use rand::Rng;
use serde::{Deserialize, Serialize};
use sha1::Sha1;
use sha2::Sha256;
use sqlx::Row;
use uuid::Uuid;

use crate::error::AppError;
use crate::models::users::User2Fa;
use crate::service::AuthService;
use crate::session::Session;

#[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 AuthService {
    pub async fn auth_2fa_enable(&self, context: &Session) -> Result<Enable2FAResponse, AppError> {
        let user_uid = context.user().ok_or(AppError::Unauthorized)?;
        let user = self.auth_find_user_by_uid(user_uid).await?;

        let existing = self.find_2fa(user_uid).await?;
        if existing.as_ref().is_some_and(|f| f.enabled) {
            return Err(AppError::TwoFactorAlreadyEnabled);
        }

        let secret = Self::generate_totp_secret();
        let backup_codes = Self::generate_backup_codes(10);
        let qr_code = format!(
            "otpauth://totp/AppKS:{}?secret={}&issuer=AppKS",
            user.username, secret
        );
        let now = chrono::Utc::now();
        let hashed_backup_codes = self.hash_backup_codes(&backup_codes)?.join(".");

        if existing.is_some() {
            sqlx::query(
                "UPDATE user_2fa SET secret = $1, backup_codes = $2, enabled = false, updated_at = $3 \
                 WHERE user_id = $4",
            )
            .bind(&secret)
            .bind(&hashed_backup_codes)
            .bind(now)
            .bind(user_uid)
            .execute(self.ctx.db.writer())
            .await
            .map_err(AppError::Database)?;
        } else {
            sqlx::query(
                "INSERT INTO user_2fa (user_id, secret, backup_codes, enabled, created_at, updated_at) \
                 VALUES ($1, $2, $3, false, $4, $4)",
            )
            .bind(user_uid)
            .bind(&secret)
            .bind(&hashed_backup_codes)
            .bind(now)
            .execute(self.ctx.db.writer())
            .await
            .map_err(AppError::Database)?;
        }

        Ok(Enable2FAResponse {
            secret,
            qr_code,
            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 = self
            .find_2fa(user_uid)
            .await?
            .ok_or(AppError::TwoFactorNotSetup)?;
        if two_fa.enabled {
            return Err(AppError::TwoFactorAlreadyEnabled);
        }
        let secret = two_fa.secret.as_ref().ok_or(AppError::TwoFactorNotSetup)?;
        if !self.verify_totp_code(secret, &params.code)? {
            return Err(AppError::InvalidTwoFactorCode);
        }

        sqlx::query("UPDATE user_2fa SET enabled = true, updated_at = $1 WHERE user_id = $2")
            .bind(chrono::Utc::now())
            .bind(user_uid)
            .execute(self.ctx.db.writer())
            .await
            .map_err(AppError::Database)?;
        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 = self
            .find_2fa(user_uid)
            .await?
            .ok_or(AppError::TwoFactorNotSetup)?;
        if !two_fa.enabled {
            return Err(AppError::TwoFactorNotEnabled);
        }
        if !self
            .verify_2fa_or_backup_code(&two_fa, &params.code)
            .await?
        {
            return Err(AppError::InvalidTwoFactorCode);
        }

        sqlx::query("DELETE FROM user_2fa WHERE user_id = $1")
            .bind(user_uid)
            .execute(self.ctx.db.writer())
            .await
            .map_err(AppError::Database)?;
        Ok(())
    }

    pub async fn auth_2fa_verify(&self, user_uid: Uuid, code: &str) -> Result<bool, AppError> {
        let Some(two_fa) = self.find_2fa(user_uid).await? else {
            return Ok(true);
        };
        if !two_fa.enabled {
            return Ok(true);
        }
        self.verify_2fa_or_backup_code(&two_fa, code).await
    }

    pub async fn auth_2fa_status_by_uid(
        &self,
        user_uid: Uuid,
    ) -> Result<Get2FAStatusResponse, AppError> {
        let Some(two_fa) = self.find_2fa(user_uid).await? else {
            return Ok(Get2FAStatusResponse {
                is_enabled: false,
                method: None,
                has_backup_codes: false,
            });
        };
        Ok(Get2FAStatusResponse {
            is_enabled: two_fa.enabled,
            method: Some("totp".into()),
            has_backup_codes: !two_fa.backup_codes.is_empty(),
        })
    }

    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,
        expected_user_uid: Uuid,
        code: &str,
    ) -> Result<bool, AppError> {
        let Some(totp_key) = context.get::<String>(Self::TOTP_KEY).ok().flatten() else {
            return Ok(false);
        };
        let Some(user_uid) = self.ctx.cache.get::<Uuid>(&totp_key) else {
            context.remove(Self::TOTP_KEY);
            return Ok(false);
        };
        if user_uid != expected_user_uid {
            context.remove(Self::TOTP_KEY);
            let _ = self.ctx.cache.delete(&totp_key);
            tracing::warn!(expected_user_uid = %expected_user_uid, pending_user_uid = %user_uid, "2FA pending user mismatch");
            return Ok(false);
        }

        let verified = self.auth_2fa_verify(user_uid, code).await?;
        if verified {
            context.remove(Self::TOTP_KEY);
            let _ = self.ctx.cache.delete(&totp_key);
        }
        Ok(verified)
    }

    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 = self
            .find_2fa(user_uid)
            .await?
            .ok_or(AppError::TwoFactorNotSetup)?;
        if !two_fa.enabled {
            return Err(AppError::TwoFactorNotEnabled);
        }

        let backup_codes = Self::generate_backup_codes(10);
        sqlx::query("UPDATE user_2fa SET backup_codes = $1, updated_at = $2 WHERE user_id = $3")
            .bind(self.hash_backup_codes(&backup_codes)?.join("."))
            .bind(chrono::Utc::now())
            .bind(user_uid)
            .execute(self.ctx.db.writer())
            .await
            .map_err(AppError::Database)?;
        Ok(backup_codes)
    }

    fn generate_totp_secret() -> String {
        const CHARSET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZ234567";
        let mut rng = rand::thread_rng();
        (0..32)
            .map(|_| {
                let idx = rng.gen_range(0..CHARSET.len());
                CHARSET[idx] as char
            })
            .collect()
    }

    fn generate_backup_codes(count: usize) -> Vec<String> {
        let mut rng = rand::thread_rng();
        (0..count)
            .map(|_| {
                format!(
                    "{:04}-{:04}-{:04}",
                    rng.gen_range(0..10000),
                    rng.gen_range(0..10000),
                    rng.gen_range(0..10000)
                )
            })
            .collect()
    }

    fn backup_code_pepper(&self) -> Result<String, AppError> {
        self.ctx
            .config
            .env
            .get("APP_SESSION_SECRET")
            .map(|s| s.trim().to_string())
            .filter(|s| !s.is_empty())
            .ok_or_else(|| AppError::Config("APP_SESSION_SECRET is required".into()))
    }

    fn hash_backup_code(&self, code: &str) -> Result<String, AppError> {
        let pepper = self.backup_code_pepper()?;
        let mut mac = Hmac::<Sha256>::new_from_slice(pepper.as_bytes())
            .map_err(|_| AppError::InternalServerError("invalid backup code pepper".into()))?;
        mac.update(code.trim().as_bytes());
        Ok(mac
            .finalize()
            .into_bytes()
            .iter()
            .map(|b| format!("{:02x}", b))
            .collect::<String>())
    }

    fn hash_backup_codes(&self, codes: &[String]) -> Result<Vec<String>, AppError> {
        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 constant_time_eq(&expected_code, code) {
                return Ok(true);
            }
        }
        Ok(false)
    }

    fn generate_totp_code(&self, secret: &str, counter: u64) -> Result<String, AppError> {
        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(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> {
        let row = sqlx::query("SELECT password_hash FROM user_password WHERE user_id = $1")
            .bind(user_uid)
            .fetch_optional(self.ctx.db.reader())
            .await
            .map_err(AppError::Database)?
            .ok_or(AppError::UserNotFound)?;
        let hash: String = row.try_get("password_hash").map_err(AppError::Database)?;

        let password_hash = PasswordHash::new(&hash).map_err(|_| AppError::InvalidPassword)?;
        Argon2::default()
            .verify_password(password.as_bytes(), &password_hash)
            .map_err(|_| AppError::InvalidPassword)?;
        Ok(())
    }

    async fn find_2fa(&self, user_uid: Uuid) -> Result<Option<User2Fa>, AppError> {
        sqlx::query_as::<_, User2Fa>(
            "SELECT user_id, secret, backup_codes, enabled, created_at, updated_at \
             FROM user_2fa WHERE user_id = $1",
        )
        .bind(user_uid)
        .fetch_optional(self.ctx.db.reader())
        .await
        .map_err(AppError::Database)
    }

    async fn verify_2fa_or_backup_code(
        &self,
        two_fa: &User2Fa,
        code: &str,
    ) -> Result<bool, AppError> {
        let secret = two_fa.secret.as_ref().ok_or(AppError::TwoFactorNotSetup)?;
        if self.verify_totp_code(secret, code)? {
            return Ok(true);
        }

        let hashed_code = self.hash_backup_code(code)?;
        let mut backup_codes: Vec<String> = two_fa
            .backup_codes
            .split('.')
            .filter(|c| !c.is_empty())
            .map(ToOwned::to_owned)
            .collect();
        if backup_codes
            .iter()
            .any(|stored| constant_time_eq(stored, &hashed_code))
        {
            backup_codes.retain(|stored| stored != &hashed_code);
            sqlx::query(
                "UPDATE user_2fa SET backup_codes = $1, updated_at = $2 WHERE user_id = $3",
            )
            .bind(backup_codes.join("."))
            .bind(chrono::Utc::now())
            .bind(two_fa.user_id)
            .execute(self.ctx.db.writer())
            .await
            .map_err(AppError::Database)?;
            return Ok(true);
        }
        Ok(false)
    }
}

use crate::service::util::constant_time_eq;