prvng_platform/crates/control-center/tests/mfa_integration_test.rs

//! MFA Integration Tests
//!
//! Comprehensive tests for the MFA system including TOTP and WebAuthn flows
//!
//! NOTE: These tests are disabled due to accessing private MfaService.storage
//! field. Need to add pub(crate) or test helper methods to MfaService to
//! re-enable.

// Disabled: accesses private fields
#![allow(unexpected_cfgs)]
#![cfg(feature = "mfa_integration_tests")]
#![allow(unused_imports, unused_variables)]

use control_center::mfa::{MfaService, TotpService};
use control_center::storage::{Database, DatabaseConfig};

/// Helper function to create test database
async fn create_test_db() -> Database {
    let config = DatabaseConfig {
        url: ":memory:".to_string(),
        max_connections: 5,
    };

    let db = Database::new(config).await.unwrap();

    // Create users table for foreign key constraints
    sqlx::query(
        r#"
        CREATE TABLE users (
            id TEXT PRIMARY KEY,
            email TEXT NOT NULL UNIQUE,
            name TEXT NOT NULL,
            password_hash TEXT NOT NULL,
            created_at TEXT NOT NULL,
            updated_at TEXT NOT NULL
        )
        "#,
    )
    .execute(db.pool())
    .await
    .unwrap();

    db
}

/// Helper function to create test MFA service
async fn create_test_mfa_service() -> MfaService {
    let db = create_test_db().await;

    MfaService::new(
        "TestApp".to_string(),
        "example.com".to_string(),
        "Test Service".to_string(),
        "https://example.com".to_string(),
        db,
    )
    .await
    .unwrap()
}

#[tokio::test]
async fn test_totp_full_enrollment_flow() {
    let service = create_test_mfa_service().await;
    let user_id = "test_user_1";
    let user_email = "test1@example.com";

    // Enroll TOTP
    let enrollment = service.enroll_totp(user_id, user_email).await.unwrap();

    // Verify enrollment response
    assert!(!enrollment.device_id.is_empty());
    assert!(!enrollment.secret.is_empty());
    assert!(enrollment.qr_code.starts_with("data:image/png;base64,"));
    assert_eq!(enrollment.backup_codes.len(), 10);
    assert!(enrollment.otpauth_url.contains("TestApp"));
    assert!(enrollment.otpauth_url.contains(user_email));

    // Verify device was created
    let status = service.get_mfa_status(user_id).await.unwrap();
    assert_eq!(status.totp_devices.len(), 1);
    assert!(!status.totp_devices[0].enabled); // Not enabled until verified
    assert!(status.has_backup_codes);
}

#[tokio::test]
async fn test_totp_verification_flow() {
    let service = create_test_mfa_service().await;
    let user_id = "test_user_2";
    let user_email = "test2@example.com";

    // Enroll TOTP
    let enrollment = service.enroll_totp(user_id, user_email).await.unwrap();

    // Generate valid TOTP code
    let totp = TotpService::new("TestApp".to_string());
    let devices = service
        .storage
        .get_totp_devices_by_user(user_id)
        .await
        .unwrap();
    let device = &devices[0];

    // Create TOTP and generate current code
    let totp_instance = totp_rs::TOTP::new(
        totp_rs::Algorithm::SHA1,
        6,
        1,
        30,
        totp_rs::Secret::Encoded(device.secret.clone())
            .to_bytes()
            .unwrap(),
        Some("TestApp".to_string()),
        user_email.to_string(),
    )
    .unwrap();

    let code = totp_instance.generate_current().unwrap();

    // Verify code
    let verification = service
        .verify_totp(user_id, user_email, &code, None)
        .await
        .unwrap();

    assert!(verification.verified);
    assert!(!verification.backup_code_used);
    assert_eq!(verification.device_id, Some(device.id.clone()));

    // Verify device is now enabled
    let status = service.get_mfa_status(user_id).await.unwrap();
    assert!(status.totp_devices[0].enabled);
}

#[tokio::test]
async fn test_backup_code_verification() {
    let service = create_test_mfa_service().await;
    let user_id = "test_user_3";
    let user_email = "test3@example.com";

    // Enroll TOTP
    let enrollment = service.enroll_totp(user_id, user_email).await.unwrap();
    let backup_code = enrollment.backup_codes[0].clone();

    // Verify backup code
    let verification = service
        .verify_totp(user_id, user_email, &backup_code, None)
        .await
        .unwrap();

    assert!(verification.verified);
    assert!(verification.backup_code_used);

    // Try to use same backup code again (should fail)
    let verification2 = service
        .verify_totp(user_id, user_email, &backup_code, None)
        .await
        .unwrap();

    assert!(!verification2.verified);
}

#[tokio::test]
async fn test_backup_code_regeneration() {
    let service = create_test_mfa_service().await;
    let user_id = "test_user_4";
    let user_email = "test4@example.com";

    // Enroll TOTP
    let enrollment = service.enroll_totp(user_id, user_email).await.unwrap();
    let old_codes = enrollment.backup_codes;

    // Regenerate backup codes
    let new_codes = service.regenerate_backup_codes(user_id).await.unwrap();

    // Verify new codes are different
    assert_eq!(new_codes.len(), 10);
    assert_ne!(old_codes, new_codes);

    // Old codes should not work
    let old_verification = service
        .verify_totp(user_id, user_email, &old_codes[0], None)
        .await
        .unwrap();
    assert!(!old_verification.verified);

    // New codes should work
    let new_verification = service
        .verify_totp(user_id, user_email, &new_codes[0], None)
        .await
        .unwrap();
    assert!(new_verification.verified);
}

#[tokio::test]
async fn test_mfa_status_tracking() {
    let service = create_test_mfa_service().await;
    let user_id = "test_user_5";
    let user_email = "test5@example.com";

    // Initially no MFA
    let status = service.get_mfa_status(user_id).await.unwrap();
    assert!(!status.enabled);
    assert_eq!(status.totp_devices.len(), 0);
    assert_eq!(status.webauthn_devices.len(), 0);

    // Enroll TOTP
    service.enroll_totp(user_id, user_email).await.unwrap();

    // Status should show TOTP device (not enabled yet)
    let status = service.get_mfa_status(user_id).await.unwrap();
    assert!(!status.enabled); // Not enabled until verified
    assert_eq!(status.totp_devices.len(), 1);
    assert!(status.has_backup_codes);
}

#[tokio::test]
async fn test_disable_totp() {
    let service = create_test_mfa_service().await;
    let user_id = "test_user_6";
    let user_email = "test6@example.com";

    // Enroll TOTP
    service.enroll_totp(user_id, user_email).await.unwrap();

    // Verify device exists
    let status = service.get_mfa_status(user_id).await.unwrap();
    assert_eq!(status.totp_devices.len(), 1);

    // Disable TOTP
    service.disable_totp(user_id).await.unwrap();

    // Verify device removed
    let status = service.get_mfa_status(user_id).await.unwrap();
    assert_eq!(status.totp_devices.len(), 0);
}

#[tokio::test]
async fn test_disable_all_mfa() {
    let service = create_test_mfa_service().await;
    let user_id = "test_user_7";
    let user_email = "test7@example.com";

    // Enroll TOTP
    service.enroll_totp(user_id, user_email).await.unwrap();

    // Disable all MFA
    service.disable_all_mfa(user_id).await.unwrap();

    // Verify all removed
    let status = service.get_mfa_status(user_id).await.unwrap();
    assert!(!status.enabled);
    assert_eq!(status.totp_devices.len(), 0);
    assert_eq!(status.webauthn_devices.len(), 0);
}

#[tokio::test]
async fn test_invalid_totp_code() {
    let service = create_test_mfa_service().await;
    let user_id = "test_user_8";
    let user_email = "test8@example.com";

    // Enroll TOTP
    service.enroll_totp(user_id, user_email).await.unwrap();

    // Try invalid code
    let verification = service
        .verify_totp(user_id, user_email, "000000", None)
        .await
        .unwrap();

    assert!(!verification.verified);
}

#[tokio::test]
async fn test_multiple_totp_devices_not_allowed() {
    let service = create_test_mfa_service().await;
    let user_id = "test_user_9";
    let user_email = "test9@example.com";

    // Enroll first TOTP
    service.enroll_totp(user_id, user_email).await.unwrap();

    // Try to enroll second TOTP (should succeed, but only one is recommended)
    let result = service.enroll_totp(user_id, user_email).await;
    assert!(result.is_ok()); // Multiple devices are allowed in this implementation

    let status = service.get_mfa_status(user_id).await.unwrap();
    assert_eq!(status.totp_devices.len(), 2);
}

// Note: WebAuthn tests would require mocking the browser credential API
// or using a test harness that simulates authenticator responses.
// For now, we test the service creation and basic flows.

#[tokio::test]
async fn test_webauthn_service_creation() {
    let service = create_test_mfa_service().await;
    // If service was created, WebAuthn is initialized
    assert!(true);
}

#[tokio::test]
async fn test_user_has_mfa_check() {
    let service = create_test_mfa_service().await;
    let user_id = "test_user_10";
    let user_email = "test10@example.com";

    // Initially no MFA
    let has_mfa = service.user_has_mfa(user_id).await.unwrap();
    assert!(!has_mfa);

    // Enroll and enable TOTP
    let enrollment = service.enroll_totp(user_id, user_email).await.unwrap();

    // Generate and verify code to enable
    let totp = TotpService::new("TestApp".to_string());
    let devices = service
        .storage
        .get_totp_devices_by_user(user_id)
        .await
        .unwrap();
    let device = &devices[0];

    let totp_instance = totp_rs::TOTP::new(
        totp_rs::Algorithm::SHA1,
        6,
        1,
        30,
        totp_rs::Secret::Encoded(device.secret.clone())
            .to_bytes()
            .unwrap(),
        Some("TestApp".to_string()),
        user_email.to_string(),
    )
    .unwrap();

    let code = totp_instance.generate_current().unwrap();
    service
        .verify_totp(user_id, user_email, &code, None)
        .await
        .unwrap();

    // Now should have MFA
    let has_mfa = service.user_has_mfa(user_id).await.unwrap();
    assert!(has_mfa);
}