clinch/test/controllers/totp_security_test.rb

require "test_helper"

class TotpSecurityTest < ActionDispatch::IntegrationTest
  # ====================
  # TOTP CODE REPLAY PREVENTION TESTS
  # ====================

  test "TOTP code cannot be reused" do
    user = User.create!(email_address: "totp_replay_test@example.com", password: "password123")
    user.enable_totp!

    # Generate a valid TOTP code
    totp = ROTP::TOTP.new(user.totp_secret)
    valid_code = totp.now

    # Set up pending TOTP session
    post signin_path, params: {email_address: "totp_replay_test@example.com", password: "password123"}
    assert_redirected_to totp_verification_path

    # First use of the code should succeed
    post totp_verification_path, params: {code: valid_code}
    assert_response :redirect
    assert_redirected_to root_path

    # Sign out
    delete session_path
    assert_response :redirect

    # Note: In the current implementation, TOTP codes CAN be reused within the 60-second time window
    # This is standard TOTP behavior. For enhanced security, you could implement used code tracking.
    # This test documents the current behavior - codes work within their time window

    user.sessions.delete_all
    user.destroy
  end

  # ====================
  # BACKUP CODE SINGLE-USE ENFORCEMENT TESTS
  # ====================

  test "backup code can only be used once" do
    user = User.create!(email_address: "backup_code_test@example.com", password: "password123")

    # Enable TOTP and generate backup codes
    user.totp_secret = ROTP::Base32.random
    backup_codes = user.send(:generate_backup_codes) # Call private method
    user.save!

    # Store the original backup codes for comparison
    original_codes = user.reload.backup_codes

    # Set up pending TOTP session
    post signin_path, params: {email_address: "backup_code_test@example.com", password: "password123"}
    assert_redirected_to totp_verification_path

    # Use a backup code
    backup_code = backup_codes.first
    post totp_verification_path, params: {code: backup_code}

    # Should successfully sign in
    assert_response :redirect
    assert_redirected_to root_path

    # Verify the backup code was marked as used
    user.reload
    assert_not_equal original_codes, user.backup_codes

    # Try to use the same backup code again
    delete session_path
    assert_response :redirect

    # Sign in again
    post signin_path, params: {email_address: "backup_code_test@example.com", password: "password123"}
    assert_redirected_to totp_verification_path

    # Try the same backup code
    post totp_verification_path, params: {code: backup_code}

    # Should fail - backup code already used
    assert_response :redirect
    assert_redirected_to totp_verification_path
    follow_redirect!
    assert_match(/invalid/i, flash[:alert].to_s)

    user.sessions.delete_all
    user.destroy
  end

  test "backup codes are hashed and not stored in plaintext" do
    user = User.create!(email_address: "backup_hash_test@example.com", password: "password123")

    # Generate backup codes
    user.totp_secret = ROTP::Base32.random
    user.send(:generate_backup_codes) # Call private method
    user.save!

    # Check that stored codes are BCrypt hashes (start with $2a$)
    # backup_codes is already an Array (JSON column), no need to parse
    user.backup_codes.each do |code|
      assert_match(/^\$2[aby]\$/, code, "Backup codes should be BCrypt hashed")
    end

    user.destroy
  end

  # ====================
  # TIME WINDOW VALIDATION TESTS
  # ====================

  test "TOTP code outside valid time window is rejected" do
    user = User.create!(email_address: "totp_time_test@example.com", password: "password123")

    # Enable TOTP with backup codes
    user.totp_secret = ROTP::Base32.random
    user.send(:generate_backup_codes)
    user.save!

    # Set up pending TOTP session
    post signin_path, params: {email_address: "totp_time_test@example.com", password: "password123"}
    assert_redirected_to totp_verification_path

    # Generate a TOTP code for a time far in the future (outside valid window)
    totp = ROTP::TOTP.new(user.totp_secret)
    future_code = totp.at(Time.now.to_i + 300) # 5 minutes in the future

    # Try to use the future code
    post totp_verification_path, params: {code: future_code}

    # Should fail - code is outside valid time window
    assert_response :redirect
    assert_redirected_to totp_verification_path
    follow_redirect!
    assert_match(/invalid/i, flash[:alert].to_s)

    user.destroy
  end

  # ====================
  # TOTP SECRET SECURITY TESTS
  # ====================

  test "TOTP secret is not exposed in API responses" do
    user = User.create!(email_address: "totp_secret_test@example.com", password: "password123")
    user.enable_totp!

    # Verify the TOTP secret exists (sanity check)
    assert user.totp_secret.present?
    user.totp_secret

    # Sign in with TOTP
    post signin_path, params: {email_address: "totp_secret_test@example.com", password: "password123"}
    assert_redirected_to totp_verification_path

    # Complete TOTP verification
    totp = ROTP::TOTP.new(user.totp_secret)
    valid_code = totp.now
    post totp_verification_path, params: {code: valid_code}
    assert_response :redirect

    # The TOTP secret should never be exposed in the response body or headers
    # This is enforced at the model level - the secret is a private attribute

    user.sessions.delete_all
    user.destroy
  end

  test "TOTP secret is rotated when re-enabling" do
    user = User.create!(email_address: "totp_rotate_test@example.com", password: "password123")

    # Enable TOTP first time
    user.enable_totp!
    first_secret = user.totp_secret

    # Disable and re-enable TOTP
    user.update!(totp_secret: nil, backup_codes: nil)
    user.enable_totp!
    second_secret = user.totp_secret

    # Secrets should be different
    assert_not_equal first_secret, second_secret, "TOTP secret should be rotated when re-enabled"

    user.destroy
  end

  # ====================
  # TOTP REQUIRED BY ADMIN TESTS
  # ====================

  test "user with TOTP required cannot disable it" do
    user = User.create!(email_address: "totp_required_test@example.com", password: "password123")
    user.update!(totp_required: true)
    user.enable_totp!

    # Verify TOTP is enabled and required
    assert user.totp_enabled?
    assert user.totp_required?

    # The disable_totp! method will clear the secret, but totp_required flag remains
    # This is enforced in the controller - users can't disable TOTP if it's required
    # The controller check is at app/controllers/totp_controller.rb:121-124

    # Verify that totp_required flag prevents disabling
    # (This is a controller-level check, not model-level)

    user.destroy
  end

  test "user with TOTP required is prompted to set it up on first login" do
    user = User.create!(email_address: "totp_setup_test@example.com", password: "password123")
    user.update!(totp_required: true, totp_secret: nil)

    # Sign in
    post signin_path, params: {email_address: "totp_setup_test@example.com", password: "password123"}

    # Should redirect to TOTP setup, not verification
    assert_response :redirect
    assert_redirected_to new_totp_path

    user.destroy
  end

  # ====================
  # TOTP CODE FORMAT VALIDATION TESTS
  # ====================

  test "invalid TOTP code formats are rejected" do
    user = User.create!(email_address: "totp_format_test@example.com", password: "password123")

    # Enable TOTP with backup codes
    user.totp_secret = ROTP::Base32.random
    user.send(:generate_backup_codes)
    user.save!

    # Set up pending TOTP session
    post signin_path, params: {email_address: "totp_format_test@example.com", password: "password123"}
    assert_redirected_to totp_verification_path

    # Try invalid formats
    invalid_codes = [
      "12345",    # Too short
      "1234567",  # Too long
      "abcdef",   # Non-numeric (6 chars, won't match backup code format)
      "12 3456",  # Contains space
      ""          # Empty
    ]

    invalid_codes.each do |invalid_code|
      post totp_verification_path, params: {code: invalid_code}
      assert_response :redirect
      assert_redirected_to totp_verification_path
    end

    user.destroy
  end

  # ====================
  # TOTP RECOVERY FLOW TESTS
  # ====================

  # ====================
  # TOTP ENROLLMENT — SECRET BINDING TESTS (H-1)
  # ====================

  test "enrollment uses the server-issued secret, not any client-submitted one" do
    user = User.create!(email_address: "totp_enroll_binding@example.com", password: "password123")
    post signin_path, params: {email_address: "totp_enroll_binding@example.com", password: "password123"}
    assert_response :redirect

    # Start enrollment: server generates a secret and stores it in the session
    get new_totp_path
    assert_response :success

    # Attacker-supplied secret + a code that's valid for THAT secret must not
    # succeed. The server should only ever consider its own session-stored secret.
    attacker_secret = ROTP::Base32.random
    attacker_code = ROTP::TOTP.new(attacker_secret).now

    assert_no_enqueued_emails do
      post totp_path, params: {totp_secret: attacker_secret, code: attacker_code}
    end

    user.reload
    assert_nil user.totp_secret, "attacker-chosen secret must not be saved"
    assert_not user.totp_enabled?

    user.sessions.delete_all
    user.destroy
  end

  test "enrollment succeeds, clears pending secret, and notifies the user by email" do
    user = User.create!(email_address: "totp_enroll_success@example.com", password: "password123")
    post signin_path, params: {email_address: "totp_enroll_success@example.com", password: "password123"}
    assert_response :redirect

    get new_totp_path
    assert_response :success

    # Pull the session-stored secret and produce a valid code for it
    pending_secret = session[:pending_totp_secret]
    assert pending_secret.present?, "new action should stash the secret in session"
    valid_code = ROTP::TOTP.new(pending_secret).now

    assert_enqueued_email_with TotpMailer, :enabled, args: [user] do
      post totp_path, params: {code: valid_code}
    end

    user.reload
    assert_equal pending_secret, user.totp_secret
    assert user.totp_enabled?
    assert_nil session[:pending_totp_secret], "pending secret must be cleared after enrollment"

    user.sessions.delete_all
    user.destroy
  end

  test "enrollment without a prior GET new is rejected" do
    user = User.create!(email_address: "totp_enroll_no_session@example.com", password: "password123")
    post signin_path, params: {email_address: "totp_enroll_no_session@example.com", password: "password123"}
    assert_response :redirect

    # Skip GET /totp/new — no session[:pending_totp_secret] is set
    post totp_path, params: {code: "123456"}
    assert_redirected_to new_totp_path

    user.reload
    assert_nil user.totp_secret
    assert_not user.totp_enabled?

    user.sessions.delete_all
    user.destroy
  end

  test "user can sign in with backup code when TOTP device is lost" do
    user = User.create!(email_address: "totp_recovery_test@example.com", password: "password123")

    # Enable TOTP and generate backup codes
    user.totp_secret = ROTP::Base32.random
    backup_codes = user.send(:generate_backup_codes) # Call private method
    user.save!

    # Sign in
    post signin_path, params: {email_address: "totp_recovery_test@example.com", password: "password123"}
    assert_redirected_to totp_verification_path

    # Use backup code instead of TOTP
    post totp_verification_path, params: {code: backup_codes.first}

    # Should successfully sign in
    assert_response :redirect
    assert_redirected_to root_path

    user.sessions.delete_all
    user.destroy
  end
end