zclaw_openfang/crates/zclaw-saas/src/auth/totp.rs

//! TOTP 双因素认证

use axum::{
    extract::{Extension, State},
    Json,
};
use crate::state::AppState;
use crate::error::{SaasError, SaasResult};
use crate::auth::types::AuthContext;
use crate::auth::handlers::log_operation;
use crate::crypto;
use serde::{Deserialize, Serialize};

/// TOTP 设置响应
#[derive(Debug, Serialize)]
pub struct TotpSetupResponse {
    /// otpauth:// URI，用于扫码绑定
    pub otpauth_uri: String,
    /// Base32 编码的密钥（备用手动输入）
    pub secret: String,
    /// issuer 名称
    pub issuer: String,
}

/// TOTP 验证请求
#[derive(Debug, Deserialize)]
pub struct TotpVerifyRequest {
    pub code: String,
}

/// TOTP 禁用请求
#[derive(Debug, Deserialize)]
pub struct TotpDisableRequest {
    pub password: String,
}

/// 生成随机 Base32 密钥 (20 字节 = 32 字符 Base32)
fn generate_random_secret() -> String {
    use rand::Rng;
    let mut bytes = [0u8; 20];
    rand::thread_rng().fill(&mut bytes);
    data_encoding::BASE32.encode(&bytes)
}

/// Base32 解码
fn base32_decode(data: &str) -> Option<Vec<u8>> {
    data_encoding::BASE32.decode(data.as_bytes()).ok()
}

/// 加密 TOTP secret (AES-256-GCM，随机 nonce)
/// 存储格式: enc:<base64(nonce||ciphertext)>
/// 委托给 crypto::encrypt_value 统一加密
fn encrypt_totp_secret(plaintext: &str, key: &[u8; 32]) -> Result<String, SaasError> {
    crate::crypto::encrypt_value(plaintext, key)
        .map_err(|e| SaasError::Internal(e.to_string()))
}

/// 解密 TOTP secret (仅支持新格式: 随机 nonce)
/// 旧的固定 nonce 格式应通过启动时迁移转换。
fn decrypt_totp_secret(encrypted: &str, key: &[u8; 32]) -> Result<String, SaasError> {
    crate::crypto::decrypt_value(encrypted, key)
        .map_err(|e| SaasError::Internal(e.to_string()))
}

/// 生成 TOTP 密钥并返回 otpauth URI
pub fn generate_totp_secret(issuer: &str, account_name: &str) -> TotpSetupResponse {
    let secret = generate_random_secret();
    let otpauth_uri = format!(
        "otpauth://totp/{}:{}?secret={}&issuer={}&algorithm=SHA1&digits=6&period=30",
        urlencoding::encode(issuer),
        urlencoding::encode(account_name),
        secret,
        urlencoding::encode(issuer),
    );

    TotpSetupResponse {
        otpauth_uri,
        secret,
        issuer: issuer.to_string(),
    }
}

/// 验证 TOTP 6 位码
pub fn verify_totp_code(secret: &str, code: &str) -> bool {
    let secret_bytes = match base32_decode(secret) {
        Some(b) => b,
        None => return false,
    };

    let totp = match totp_rs::TOTP::new(
        totp_rs::Algorithm::SHA1,
        6,       // digits
        1,       // skew (允许 1 个周期偏差)
        30,      // step (秒)
        secret_bytes,
    ) {
        Ok(t) => t,
        Err(_) => return false,
    };

    totp.check_current(code).unwrap_or(false)
}

/// POST /api/v1/auth/totp/setup
/// 生成 TOTP 密钥并返回 otpauth URI
/// 用户扫码后需要调用 /verify 验证一个码才能激活
pub async fn setup_totp(
    State(state): State<AppState>,
    Extension(ctx): Extension<AuthContext>,
) -> SaasResult<Json<TotpSetupResponse>> {
    // 如果已启用 TOTP，先清除旧密钥
    let (username,): (String,) = sqlx::query_as(
        "SELECT username FROM accounts WHERE id = $1"
    )
    .bind(&ctx.account_id)
    .fetch_one(&state.db)
    .await?;

    let config = state.config.read().await;
    let setup = generate_totp_secret(&config.auth.totp_issuer, &username);

    // 加密后存储密钥 (但不启用，需要 /verify 确认)
    let enc_key = config.totp_encryption_key()
        .map_err(|e| SaasError::Internal(e.to_string()))?;
    let encrypted_secret = encrypt_totp_secret(&setup.secret, &enc_key)?;

    sqlx::query("UPDATE accounts SET totp_secret = $1 WHERE id = $2")
        .bind(&encrypted_secret)
        .bind(&ctx.account_id)
        .execute(&state.db)
        .await?;

    log_operation(&state.db, &ctx.account_id, "totp.setup", "account", &ctx.account_id,
        None, ctx.client_ip.as_deref()).await?;

    Ok(Json(setup))
}

/// POST /api/v1/auth/totp/verify
/// 验证 TOTP 码并启用 2FA
pub async fn verify_totp(
    State(state): State<AppState>,
    Extension(ctx): Extension<AuthContext>,
    Json(req): Json<TotpVerifyRequest>,
) -> SaasResult<Json<serde_json::Value>> {
    let code = req.code.trim();
    if code.len() != 6 || !code.chars().all(|c| c.is_ascii_digit()) {
        return Err(SaasError::InvalidInput("TOTP 码必须是 6 位数字".into()));
    }

    // TOTP 暴力破解保护: 10 分钟内最多 5 次失败
    const MAX_TOTP_FAILURES: u32 = 5;
    const TOTP_LOCKOUT_SECS: u64 = 600;
    let now = std::time::Instant::now();
    let lockout_duration = std::time::Duration::from_secs(TOTP_LOCKOUT_SECS);

    let is_locked = {
        if let Some(entry) = state.totp_fail_counts.get(&ctx.account_id) {
            let (count, first_fail) = entry.value();
            if *count >= MAX_TOTP_FAILURES && now.duration_since(*first_fail) < lockout_duration {
                true
            } else {
                // 窗口过期，重置
                drop(entry);
                state.totp_fail_counts.remove(&ctx.account_id);
                false
            }
        } else {
            false
        }
    };

    if is_locked {
        return Err(SaasError::RateLimited(
            format!("TOTP 验证失败次数过多，请 {} 秒后重试", TOTP_LOCKOUT_SECS)
        ));
    }

    // 获取存储的密钥
    let (totp_secret,): (Option<String>,) = sqlx::query_as(
        "SELECT totp_secret FROM accounts WHERE id = $1"
    )
    .bind(&ctx.account_id)
    .fetch_one(&state.db)
    .await?;

    let encrypted_secret = totp_secret.ok_or_else(|| {
        SaasError::InvalidInput("请先调用 /totp/setup 获取密钥".into())
    })?;

    // 解密 secret (兼容旧的明文格式)
    let config = state.config.read().await;
    let enc_key = config.totp_encryption_key()
        .map_err(|e| SaasError::Internal(e.to_string()))?;
    let secret = if encrypted_secret.starts_with(crypto::ENCRYPTED_PREFIX) {
        decrypt_totp_secret(&encrypted_secret, &enc_key)?
    } else {
        // 旧格式: 明文存储，需要迁移
        encrypted_secret.clone()
    };

    if !verify_totp_code(&secret, code) {
        // 记录失败次数
        let new_count = {
            let mut entry = state.totp_fail_counts
                .entry(ctx.account_id.clone())
                .or_insert((0, now));
            entry.value_mut().0 += 1;
            entry.value().0
        };
        tracing::warn!(
            "TOTP verify failed for account {} ({}/{} attempts)",
            ctx.account_id, new_count, MAX_TOTP_FAILURES
        );
        return Err(SaasError::Totp("TOTP 码验证失败".into()));
    }

    // 验证成功 → 清除失败计数
    state.totp_fail_counts.remove(&ctx.account_id);

    // 验证成功 → 启用 TOTP，同时确保密钥已加密
    let final_secret = if encrypted_secret.starts_with(crypto::ENCRYPTED_PREFIX) {
        encrypted_secret
    } else {
        // 迁移: 加密旧明文密钥
        encrypt_totp_secret(&secret, &enc_key)?
    };

    let now_ts = chrono::Utc::now();
    sqlx::query("UPDATE accounts SET totp_enabled = true, totp_secret = $1, updated_at = $2 WHERE id = $3")
        .bind(&final_secret)
        .bind(&now_ts)
        .bind(&ctx.account_id)
        .execute(&state.db)
        .await?;

    log_operation(&state.db, &ctx.account_id, "totp.verify", "account", &ctx.account_id,
        None, ctx.client_ip.as_deref()).await?;

    Ok(Json(serde_json::json!({"ok": true, "totp_enabled": true, "message": "TOTP 已启用"})))
}

/// POST /api/v1/auth/totp/disable
/// 禁用 TOTP (需要密码确认)
pub async fn disable_totp(
    State(state): State<AppState>,
    Extension(ctx): Extension<AuthContext>,
    Json(req): Json<TotpDisableRequest>,
) -> SaasResult<Json<serde_json::Value>> {
    // 验证密码
    let (password_hash,): (String,) = sqlx::query_as(
        "SELECT password_hash FROM accounts WHERE id = $1"
    )
    .bind(&ctx.account_id)
    .fetch_one(&state.db)
    .await?;

    if !crate::auth::password::verify_password_async(req.password.clone(), password_hash.clone()).await? {
        return Err(SaasError::AuthError("密码错误".into()));
    }

    // 清除 TOTP
    let now = chrono::Utc::now();
    sqlx::query("UPDATE accounts SET totp_enabled = false, totp_secret = NULL, updated_at = $1 WHERE id = $2")
        .bind(&now)
        .bind(&ctx.account_id)
        .execute(&state.db)
        .await?;

    log_operation(&state.db, &ctx.account_id, "totp.disable", "account", &ctx.account_id,
        None, ctx.client_ip.as_deref()).await?;

    Ok(Json(serde_json::json!({"ok": true, "totp_enabled": false, "message": "TOTP 已禁用"})))
}

/// 解密 TOTP secret (供 login handler 使用)
/// 返回解密后的明文 secret
pub fn decrypt_totp_for_login(encrypted_secret: &str, enc_key: &[u8; 32]) -> SaasResult<String> {
    if encrypted_secret.starts_with(crypto::ENCRYPTED_PREFIX) {
        decrypt_totp_secret(encrypted_secret, enc_key)
    } else {
        // 兼容旧的明文格式
        Ok(encrypted_secret.to_string())
    }
}