mockforge-registry-server 0.3.124

//! Two-Factor Authentication (2FA) handlers
//!
//! Handles 2FA setup, verification, and management

use axum::{extract::State, Json};
use serde::{Deserialize, Serialize};

use crate::{
    error::{ApiError, ApiResult},
    middleware::AuthUser,
    models::AuditEventType,
    redis::{two_factor_backup_codes_key, two_factor_setup_key, TWO_FACTOR_SETUP_TTL_SECONDS},
    two_factor::{
        generate_backup_codes, generate_qr_code_data_url, generate_secret, hash_backup_code,
        verify_totp_code,
    },
    AppState,
};

#[derive(Debug, Serialize)]
pub struct Setup2FAResponse {
    pub secret: String,            // Base32-encoded secret (for manual entry)
    pub qr_code_url: String,       // Data URL for QR code
    pub backup_codes: Vec<String>, // Plain text backup codes (shown once)
}

#[derive(Debug, Deserialize)]
pub struct Verify2FASetupRequest {
    pub code: String, // 6-digit TOTP code
}

#[derive(Debug, Serialize)]
pub struct Verify2FASetupResponse {
    pub success: bool,
    pub message: String,
}

#[derive(Debug, Deserialize)]
pub struct Disable2FARequest {
    pub password: String, // Require password confirmation
}

#[derive(Debug, Serialize)]
pub struct Disable2FAResponse {
    pub success: bool,
    pub message: String,
}

/// Start 2FA setup process
/// Generates a secret and QR code for the user to scan
pub async fn setup_2fa(
    State(state): State<AppState>,
    AuthUser(user_id): AuthUser,
) -> ApiResult<Json<Setup2FAResponse>> {
    // Get user
    let user = state
        .store
        .find_user_by_id(user_id)
        .await?
        .ok_or_else(|| ApiError::InvalidRequest("User not found".to_string()))?;

    // Check if 2FA is already enabled
    if user.two_factor_enabled {
        return Err(ApiError::InvalidRequest(
            "2FA is already enabled. Disable it first to set up a new device.".to_string(),
        ));
    }

    // Generate TOTP secret
    let secret = generate_secret()
        .map_err(|e| ApiError::Internal(anyhow::anyhow!("Failed to generate 2FA secret: {}", e)))?;

    // Generate backup codes (10 codes)
    let backup_codes = generate_backup_codes(10).map_err(|e| {
        ApiError::Internal(anyhow::anyhow!("Failed to generate backup codes: {}", e))
    })?;

    // Generate QR code
    let issuer = "MockForge";
    let account_name = format!("{}:{}", user.email, user.username);
    let qr_code_url = generate_qr_code_data_url(&secret, &account_name, issuer)
        .map_err(|e| ApiError::Internal(anyhow::anyhow!("Failed to generate QR code: {}", e)))?;

    // Store secret and backup codes temporarily in Redis (5 minute TTL)
    if let Some(ref redis) = state.redis {
        // Store the secret
        let secret_key = two_factor_setup_key(&user_id);
        redis
            .set_with_expiry(&secret_key, &secret, TWO_FACTOR_SETUP_TTL_SECONDS)
            .await
            .map_err(|e| {
                ApiError::Internal(anyhow::anyhow!("Failed to store 2FA secret: {}", e))
            })?;

        // Store backup codes as JSON
        let backup_codes_key = two_factor_backup_codes_key(&user_id);
        let backup_codes_json = serde_json::to_string(&backup_codes).map_err(|e| {
            ApiError::Internal(anyhow::anyhow!("Failed to serialize backup codes: {}", e))
        })?;
        redis
            .set_with_expiry(&backup_codes_key, &backup_codes_json, TWO_FACTOR_SETUP_TTL_SECONDS)
            .await
            .map_err(|e| {
                ApiError::Internal(anyhow::anyhow!("Failed to store backup codes: {}", e))
            })?;

        tracing::debug!("Stored 2FA setup data in Redis for user {}", user_id);
    } else {
        tracing::warn!("Redis not configured - 2FA setup will require secret to be passed in verification request");
    }

    Ok(Json(Setup2FAResponse {
        secret: secret.clone(),
        qr_code_url,
        backup_codes: backup_codes.clone(),
    }))
}

/// Verify and enable 2FA
/// User must provide a valid TOTP code to confirm setup
pub async fn verify_2fa_setup(
    State(state): State<AppState>,
    AuthUser(user_id): AuthUser,
    Json(request): Json<Verify2FASetupRequest>,
) -> ApiResult<Json<Verify2FASetupResponse>> {
    // Get user
    let user = state
        .store
        .find_user_by_id(user_id)
        .await?
        .ok_or_else(|| ApiError::InvalidRequest("User not found".to_string()))?;

    // Check if 2FA is already enabled
    if user.two_factor_enabled {
        return Err(ApiError::InvalidRequest("2FA is already enabled".to_string()));
    }

    // Retrieve secret and backup codes from Redis
    let redis = state.redis.as_ref().ok_or_else(|| {
        ApiError::Internal(anyhow::anyhow!(
            "Redis not configured. Use verify_2fa_setup_with_secret endpoint instead."
        ))
    })?;

    let secret_key = two_factor_setup_key(&user_id);
    let secret = redis
        .get(&secret_key)
        .await
        .map_err(|e| ApiError::Internal(anyhow::anyhow!("Failed to retrieve 2FA secret: {}", e)))?
        .ok_or_else(|| {
            ApiError::InvalidRequest(
                "2FA setup expired or not started. Please call setup_2fa first.".to_string(),
            )
        })?;

    let backup_codes_key = two_factor_backup_codes_key(&user_id);
    let backup_codes_json = redis
        .get(&backup_codes_key)
        .await
        .map_err(|e| ApiError::Internal(anyhow::anyhow!("Failed to retrieve backup codes: {}", e)))?
        .ok_or_else(|| {
            ApiError::InvalidRequest(
                "2FA setup expired or not started. Please call setup_2fa first.".to_string(),
            )
        })?;

    let backup_codes: Vec<String> = serde_json::from_str(&backup_codes_json)
        .map_err(|e| ApiError::Internal(anyhow::anyhow!("Failed to parse backup codes: {}", e)))?;

    // Verify TOTP code
    let valid = verify_totp_code(&secret, &request.code, Some(1))
        .map_err(|e| ApiError::Internal(anyhow::anyhow!("TOTP verification error: {}", e)))?;

    if !valid {
        return Err(ApiError::InvalidRequest(
            "Invalid verification code. Please try again.".to_string(),
        ));
    }

    // Hash backup codes for storage
    let hashed_backup_codes: Vec<String> = backup_codes
        .iter()
        .map(|code| hash_backup_code(code))
        .collect::<Result<Vec<_>, _>>()
        .map_err(|e| ApiError::Internal(anyhow::anyhow!("Failed to hash backup codes: {}", e)))?;

    // Enable 2FA in database
    state.store.enable_user_2fa(user_id, &secret, &hashed_backup_codes).await?;

    // Clean up Redis keys
    let _ = redis.delete(&secret_key).await;
    let _ = redis.delete(&backup_codes_key).await;

    // Record audit log
    let user_org_id = uuid::Uuid::nil();
    state
        .store
        .record_audit_event(
            user_org_id,
            Some(user_id),
            AuditEventType::TwoFactorEnabled,
            "Two-factor authentication enabled".to_string(),
            None,
            None,
            None,
        )
        .await;

    tracing::info!("2FA enabled for user {}", user_id);

    send_2fa_security_alert(&user, true).await;

    Ok(Json(Verify2FASetupResponse {
        success: true,
        message:
            "2FA has been enabled successfully. Please save your backup codes in a safe place."
                .to_string(),
    }))
}

/// Send a best-effort security-alert email when 2FA is enabled or disabled.
/// Gated on the user's `security_alerts` preference. Never fails the request.
async fn send_2fa_security_alert(user: &mockforge_registry_core::models::User, enabled: bool) {
    if !user.security_alerts {
        return;
    }
    let Ok(email_service) = crate::email::EmailService::from_env() else {
        return;
    };
    let (headline, detail) = if enabled {
        (
            "Two-factor authentication was enabled",
            "If you did not enable 2FA, contact support immediately — your account may be compromised.",
        )
    } else {
        (
            "Two-factor authentication was disabled",
            "If you did not disable 2FA, reset your password immediately and contact support.",
        )
    };
    let msg = crate::email::EmailService::generate_security_alert_email(
        &user.username,
        &user.email,
        headline,
        detail,
    );
    if let Err(e) = email_service.send(msg).await {
        tracing::warn!("Failed to send 2FA security alert: {}", e);
    }
}

/// Simplified verify_2fa_setup that accepts secret
/// In production, this would retrieve the secret from a temporary store
pub async fn verify_2fa_setup_with_secret(
    State(state): State<AppState>,
    AuthUser(user_id): AuthUser,
    Json(request): Json<Verify2FASetupRequestWithSecret>,
) -> ApiResult<Json<Verify2FASetupResponse>> {
    // Get user
    let user = state
        .store
        .find_user_by_id(user_id)
        .await?
        .ok_or_else(|| ApiError::InvalidRequest("User not found".to_string()))?;

    // Check if 2FA is already enabled
    if user.two_factor_enabled {
        return Err(ApiError::InvalidRequest("2FA is already enabled".to_string()));
    }

    // Verify TOTP code
    let valid = verify_totp_code(&request.secret, &request.code, Some(1))
        .map_err(|e| ApiError::Internal(anyhow::anyhow!("TOTP verification error: {}", e)))?;

    if !valid {
        return Err(ApiError::InvalidRequest(
            "Invalid verification code. Please try again.".to_string(),
        ));
    }

    // Generate and hash backup codes
    let backup_codes = generate_backup_codes(10).map_err(|e| {
        ApiError::Internal(anyhow::anyhow!("Failed to generate backup codes: {}", e))
    })?;
    let hashed_backup_codes: Vec<String> = backup_codes
        .iter()
        .map(|code| hash_backup_code(code))
        .collect::<Result<Vec<_>, _>>()
        .map_err(|e| ApiError::Internal(anyhow::anyhow!("Failed to hash backup codes: {}", e)))?;

    // Enable 2FA
    state
        .store
        .enable_user_2fa(user_id, &request.secret, &hashed_backup_codes)
        .await?;

    // Record audit log
    // Use a sentinel UUID for user-level actions (no org)
    let user_org_id = uuid::Uuid::nil();
    state
        .store
        .record_audit_event(
            user_org_id,
            Some(user_id),
            AuditEventType::TwoFactorEnabled,
            "Two-factor authentication enabled".to_string(),
            None,
            None,
            None,
        )
        .await;

    send_2fa_security_alert(&user, true).await;

    Ok(Json(Verify2FASetupResponse {
        success: true,
        message:
            "2FA has been enabled successfully. Please save your backup codes in a safe place."
                .to_string(),
    }))
}

#[derive(Debug, Deserialize)]
pub struct Verify2FASetupRequestWithSecret {
    pub secret: String,
    pub code: String,
    pub backup_codes: Vec<String>, // Plain text backup codes from setup
}

/// Disable 2FA
/// Requires password confirmation for security
pub async fn disable_2fa(
    State(state): State<AppState>,
    AuthUser(user_id): AuthUser,
    Json(request): Json<Disable2FARequest>,
) -> ApiResult<Json<Disable2FAResponse>> {
    // Get user
    let user = state
        .store
        .find_user_by_id(user_id)
        .await?
        .ok_or_else(|| ApiError::InvalidRequest("User not found".to_string()))?;

    // Verify password
    use crate::auth::verify_password;
    let valid =
        verify_password(&request.password, &user.password_hash).map_err(ApiError::Internal)?;

    if !valid {
        return Err(ApiError::InvalidRequest("Invalid password".to_string()));
    }

    // Check if 2FA is enabled
    if !user.two_factor_enabled {
        return Err(ApiError::InvalidRequest("2FA is not enabled".to_string()));
    }

    // Disable 2FA
    state.store.disable_user_2fa(user_id).await?;

    // Record audit log
    // Use a sentinel UUID for user-level actions (no org)
    let user_org_id = uuid::Uuid::nil();
    state
        .store
        .record_audit_event(
            user_org_id,
            Some(user_id),
            AuditEventType::TwoFactorDisabled,
            "Two-factor authentication disabled".to_string(),
            None,
            None,
            None,
        )
        .await;

    send_2fa_security_alert(&user, false).await;

    Ok(Json(Disable2FAResponse {
        success: true,
        message: "2FA has been disabled successfully.".to_string(),
    }))
}

/// Get 2FA status
pub async fn get_2fa_status(
    State(state): State<AppState>,
    AuthUser(user_id): AuthUser,
) -> ApiResult<Json<serde_json::Value>> {
    // Get user
    let user = state
        .store
        .find_user_by_id(user_id)
        .await?
        .ok_or_else(|| ApiError::InvalidRequest("User not found".to_string()))?;

    Ok(Json(serde_json::json!({
        "enabled": user.two_factor_enabled,
        "verified_at": user.two_factor_verified_at,
        "backup_codes_count": user.two_factor_backup_codes.as_ref().map(|c| c.len()).unwrap_or(0),
    })))
}