foxhunt/services/api/tests/mfa_enrollment_integration_test.rs

#![cfg(feature = "mfa")]
//! MFA Enrollment Integration Test
//!
//! Agent H3: Tests complete MFA enrollment flow including:
//! - QR code generation
//! - TOTP verification
//! - Backup code generation
//! - Account lockout after failed attempts
//! - Admin enforcement

use anyhow::Result;
use secrecy::ExposeSecret;
use sqlx::PgPool;
use uuid::Uuid;

use api::auth::mfa::MfaManager;

/// Helper to create test database connection
async fn setup_test_db() -> Result<PgPool> {
    let database_url = std::env::var("DATABASE_URL").unwrap_or_else(|_| {
        "postgresql://foxhunt:foxhunt_dev_password@localhost:5432/foxhunt".to_string()
    });

    let pool = PgPool::connect(&database_url).await?;
    Ok(pool)
}

/// Helper to create test user with admin role
async fn create_test_admin_user(pool: &PgPool) -> Result<Uuid> {
    let user_id = Uuid::new_v4();
    let username = format!(
        "test_admin_{}",
        Uuid::new_v4().to_string().split('-').next().expect("INVARIANT: Iterator should have next element")
    );
    let email = format!("{}@test.local", username);

    // Create user
    sqlx::query(
        r#"
        INSERT INTO users (
            id, username, email, password_hash, salt, must_change_password, active
        ) VALUES ($1, $2, $3, '$2b$12$test_hash', 'test_salt', FALSE, TRUE)
        "#,
    )
    .bind(user_id)
    .bind(&username)
    .bind(&email)
    .execute(pool)
    .await?;

    // Assign system_admin role
    let admin_role_id = Uuid::parse_str("00000000-0000-0000-0000-000000000001")?;
    sqlx::query("INSERT INTO user_roles (user_id, role_id, granted_by) VALUES ($1, $2, $1)")
        .bind(user_id)
        .bind(admin_role_id)
        .execute(pool)
        .await?;

    Ok(user_id)
}

/// Helper to cleanup test user
async fn cleanup_test_user(pool: &PgPool, user_id: Uuid) -> Result<()> {
    // Foreign key constraints will cascade delete related records
    sqlx::query("DELETE FROM users WHERE id = $1")
        .bind(user_id)
        .execute(pool)
        .await?;

    Ok(())
}

#[tokio::test]
async fn test_mfa_enrollment_complete_flow() -> Result<()> {
    let pool = setup_test_db().await?;
    let user_id = create_test_admin_user(&pool).await?;

    // Create MFA manager
    let encryption_key = "test_encryption_key_32_bytes_long!!".to_string();
    let mfa_manager = MfaManager::new(pool.clone(), encryption_key)?;

    // Step 1: Start enrollment
    let enrollment = mfa_manager
        .start_enrollment(user_id, "Foxhunt Test", "test_user@test.local")
        .await?;

    println!("✓ MFA enrollment session created");
    println!("  Session ID: {}", enrollment.session_id);
    println!("  QR Code URI: {}", enrollment.qr_code_uri);
    println!("  QR Code PNG: {} bytes", enrollment.qr_code_png.len());
    println!("  Manual Entry Key: {}", enrollment.manual_entry_key);

    assert!(
        !enrollment.qr_code_uri.is_empty(),
        "QR code URI should not be empty"
    );
    assert!(
        !enrollment.qr_code_png.is_empty(),
        "QR code PNG should not be empty"
    );
    assert!(
        enrollment.qr_code_png.len() > 100,
        "QR code PNG should be at least 100 bytes"
    );
    assert!(
        enrollment.manual_entry_key.len() >= 16,
        "TOTP secret should be at least 16 characters"
    );

    // Step 2: Generate valid TOTP code from secret
    use totp_rs::{Algorithm, TOTP};
    let totp = TOTP::new(
        Algorithm::SHA1,
        6,
        1,
        30,
        enrollment.manual_entry_key.as_bytes().to_vec(),
    )?;

    let valid_code = totp.generate_current()?;
    println!("✓ Generated TOTP code: {}", valid_code);

    // Step 3: Complete enrollment with valid TOTP code
    let backup_codes = mfa_manager
        .complete_enrollment(enrollment.session_id, user_id, &valid_code)
        .await?;

    println!("✓ MFA enrollment completed");
    println!("  Backup codes generated: {}", backup_codes.len());

    assert_eq!(backup_codes.len(), 10, "Should generate 10 backup codes");

    for (i, code) in backup_codes.iter().enumerate() {
        println!(
            "  Backup Code {}: {} (hint: {})",
            i + 1,
            code.code.expose_secret(),
            code.hint
        );
        assert!(
            code.code.expose_secret().len() >= 8,
            "Backup code should be at least 8 characters"
        );
        assert_eq!(code.hint.len(), 4, "Hint should be 4 characters");
    }

    // Step 4: Verify MFA is enabled
    let config = mfa_manager.get_mfa_config(user_id).await?;
    assert!(config.is_some(), "MFA config should exist");

    let config = config.unwrap();
    assert!(config.is_enabled, "MFA should be enabled");
    assert!(config.is_verified, "MFA should be verified");
    assert_eq!(
        config.backup_codes_remaining, 10,
        "Should have 10 backup codes"
    );

    println!("✓ MFA config verified");
    println!("  Enabled: {}", config.is_enabled);
    println!("  Verified: {}", config.is_verified);
    println!(
        "  Backup codes remaining: {}",
        config.backup_codes_remaining
    );

    // Cleanup
    cleanup_test_user(&pool, user_id).await?;

    Ok(())
}

#[tokio::test]
async fn test_mfa_totp_verification() -> Result<()> {
    let pool = setup_test_db().await?;
    let user_id = create_test_admin_user(&pool).await?;

    let encryption_key = "test_encryption_key_32_bytes_long!!".to_string();
    let mfa_manager = MfaManager::new(pool.clone(), encryption_key)?;

    // Enroll user in MFA
    let enrollment = mfa_manager
        .start_enrollment(user_id, "Foxhunt Test", "test@test.local")
        .await?;

    use totp_rs::{Algorithm, TOTP};
    let totp = TOTP::new(
        Algorithm::SHA1,
        6,
        1,
        30,
        enrollment.manual_entry_key.as_bytes().to_vec(),
    )?;

    let valid_code = totp.generate_current()?;
    let _backup_codes = mfa_manager
        .complete_enrollment(enrollment.session_id, user_id, &valid_code)
        .await?;

    // Test valid TOTP verification
    let new_code = totp.generate_current()?;
    let is_valid = mfa_manager
        .verify_totp(user_id, &new_code, Some("127.0.0.1".to_string()))
        .await?;

    assert!(is_valid, "Valid TOTP code should verify successfully");
    println!("✓ Valid TOTP code verified");

    // Test invalid TOTP code
    let invalid_result = mfa_manager
        .verify_totp(user_id, "000000", Some("127.0.0.1".to_string()))
        .await;

    assert!(
        invalid_result.is_ok(),
        "Invalid code should return Ok(false)"
    );
    assert!(
        !invalid_result.unwrap(),
        "Invalid TOTP code should not verify"
    );
    println!("✓ Invalid TOTP code rejected");

    // Cleanup
    cleanup_test_user(&pool, user_id).await?;

    Ok(())
}

#[tokio::test]
async fn test_mfa_backup_code_recovery() -> Result<()> {
    let pool = setup_test_db().await?;
    let user_id = create_test_admin_user(&pool).await?;

    let encryption_key = "test_encryption_key_32_bytes_long!!".to_string();
    let mfa_manager = MfaManager::new(pool.clone(), encryption_key)?;

    // Enroll user in MFA
    let enrollment = mfa_manager
        .start_enrollment(user_id, "Foxhunt Test", "test@test.local")
        .await?;

    use totp_rs::{Algorithm, TOTP};
    let totp = TOTP::new(
        Algorithm::SHA1,
        6,
        1,
        30,
        enrollment.manual_entry_key.as_bytes().to_vec(),
    )?;

    let valid_code = totp.generate_current()?;
    let backup_codes = mfa_manager
        .complete_enrollment(enrollment.session_id, user_id, &valid_code)
        .await?;

    // Test backup code usage
    let first_backup_code = backup_codes[0].code.expose_secret();
    let is_valid = mfa_manager
        .verify_backup_code(user_id, first_backup_code, Some("127.0.0.1".to_string()))
        .await?;

    assert!(is_valid, "Valid backup code should verify successfully");
    println!("✓ Backup code verified and consumed");

    // Verify backup code count decreased
    let status = mfa_manager.get_backup_codes_status(user_id).await?;
    assert_eq!(status.remaining, 9, "Should have 9 remaining backup codes");
    assert_eq!(status.used, 1, "Should have 1 used backup code");
    println!("✓ Backup code status updated (9 remaining, 1 used)");

    // Test reusing same backup code (should fail)
    let is_valid = mfa_manager
        .verify_backup_code(user_id, first_backup_code, Some("127.0.0.1".to_string()))
        .await?;

    assert!(!is_valid, "Used backup code should not verify again");
    println!("✓ Reused backup code rejected");

    // Cleanup
    cleanup_test_user(&pool, user_id).await?;

    Ok(())
}

#[tokio::test]
async fn test_mfa_account_lockout() -> Result<()> {
    let pool = setup_test_db().await?;
    let user_id = create_test_admin_user(&pool).await?;

    let encryption_key = "test_encryption_key_32_bytes_long!!".to_string();
    let mfa_manager = MfaManager::new(pool.clone(), encryption_key)?;

    // Enroll user in MFA
    let enrollment = mfa_manager
        .start_enrollment(user_id, "Foxhunt Test", "test@test.local")
        .await?;

    use totp_rs::{Algorithm, TOTP};
    let totp = TOTP::new(
        Algorithm::SHA1,
        6,
        1,
        30,
        enrollment.manual_entry_key.as_bytes().to_vec(),
    )?;

    let valid_code = totp.generate_current()?;
    let _backup_codes = mfa_manager
        .complete_enrollment(enrollment.session_id, user_id, &valid_code)
        .await?;

    // Make 5 failed attempts to trigger lockout
    for i in 1..=5 {
        let result = mfa_manager
            .verify_totp(user_id, "000000", Some("127.0.0.1".to_string()))
            .await;

        if i < 5 {
            assert!(
                result.is_ok(),
                "Failed attempt {} should not lock account yet",
                i
            );
            println!("✓ Failed attempt {} recorded", i);
        }
    }

    // Check if account is locked
    let is_locked = mfa_manager.is_mfa_locked(user_id).await?;
    assert!(
        is_locked,
        "Account should be locked after 5 failed attempts"
    );
    println!("✓ Account locked after 5 failed attempts");

    // Verify locked account cannot authenticate even with valid code
    let new_code = totp.generate_current()?;
    let locked_result = mfa_manager
        .verify_totp(user_id, &new_code, Some("127.0.0.1".to_string()))
        .await;

    assert!(
        locked_result.is_err(),
        "Locked account should not allow authentication"
    );
    assert!(
        locked_result.unwrap_err().to_string().contains("locked"),
        "Error should mention account lock"
    );
    println!("✓ Locked account rejected valid TOTP code");

    // Cleanup
    cleanup_test_user(&pool, user_id).await?;

    Ok(())
}

#[tokio::test]
async fn test_mfa_admin_enforcement() -> Result<()> {
    let pool = setup_test_db().await?;
    let user_id = create_test_admin_user(&pool).await?;

    // Test 1: Verify MFA is required for admin user
    let is_required: bool = sqlx::query_scalar("SELECT is_mfa_required($1)")
        .bind(user_id)
        .fetch_one(&pool)
        .await?;

    assert!(is_required, "MFA should be required for admin user");
    println!("✓ MFA requirement detected for admin user");

    // Test 2: Verify session creation is blocked without MFA
    let session_id = Uuid::new_v4();
    let token_hash = format!("test_token_{}", Uuid::new_v4());

    let result = sqlx::query(
        r#"
        INSERT INTO sessions (
            id, token_hash, user_id, expires_at, client_ip, session_type
        ) VALUES ($1, $2, $3, NOW() + INTERVAL '1 hour', '127.0.0.1'::inet, 'test')
        "#,
    )
    .bind(session_id)
    .bind(token_hash)
    .bind(user_id)
    .execute(&pool)
    .await;

    assert!(result.is_err(), "Session creation should fail without MFA");
    let error_msg = result.unwrap_err().to_string();
    assert!(
        error_msg.contains("MFA_REQUIRED"),
        "Error should indicate MFA requirement"
    );
    println!("✓ Session creation blocked without MFA enrollment");

    // Test 3: Enroll in MFA and verify session creation succeeds
    let encryption_key = "test_encryption_key_32_bytes_long!!".to_string();
    let mfa_manager = MfaManager::new(pool.clone(), encryption_key)?;

    let enrollment = mfa_manager
        .start_enrollment(user_id, "Foxhunt Test", "test@test.local")
        .await?;

    use totp_rs::{Algorithm, TOTP};
    let totp = TOTP::new(
        Algorithm::SHA1,
        6,
        1,
        30,
        enrollment.manual_entry_key.as_bytes().to_vec(),
    )?;

    let valid_code = totp.generate_current()?;
    let _backup_codes = mfa_manager
        .complete_enrollment(enrollment.session_id, user_id, &valid_code)
        .await?;

    println!("✓ MFA enrollment completed");

    // Now session creation should succeed
    let session_id = Uuid::new_v4();
    let token_hash = format!("test_token_{}", Uuid::new_v4());

    let result = sqlx::query(
        r#"
        INSERT INTO sessions (
            id, token_hash, user_id, expires_at, client_ip, session_type
        ) VALUES ($1, $2, $3, NOW() + INTERVAL '1 hour', '127.0.0.1'::inet, 'test')
        "#,
    )
    .bind(session_id)
    .bind(token_hash)
    .bind(user_id)
    .execute(&pool)
    .await;

    assert!(
        result.is_ok(),
        "Session creation should succeed with MFA enrolled"
    );
    println!("✓ Session creation allowed after MFA enrollment");

    // Cleanup
    cleanup_test_user(&pool, user_id).await?;

    Ok(())
}