Skip to content

Commit

Permalink
Enable TOTP method when method is configured
Browse files Browse the repository at this point in the history
  • Loading branch information
@kasparsd
kasparsd authored Jan 9, 2025
2 parents abd8289 + ff7cfb5 commit 4c733e4
Show file tree
Hide file tree
Showing 3 changed files with 47 additions and 44 deletions.
47 changes: 21 additions & 26 deletions class-two-factor-core.php
View file
Original file line number Diff line number Diff line change
Expand Up @@ -249,7 +249,8 @@ private static function get_providers_classes( $providers ) {
}

/**
* Get all enabled two-factor providers.
* Get all enabled two-factor providers with keys as the original
* provider class names and the values as the provider class instances.
*
* @since 0.1-dev
*
Expand Down Expand Up @@ -1789,10 +1790,9 @@ public static function user_two_factor_options( $user ) {
$enabled_providers = array_keys( self::get_available_providers_for_user( $user ) );
$primary_provider = self::get_primary_provider_for_user( $user->ID );

$primary_provider_key = null;
if ( ! empty( $primary_provider ) && is_object( $primary_provider ) ) {
$primary_provider_key = $primary_provider->get_key();
} else {
$primary_provider_key = null;
}

// This is specific to the current session, not the displayed user.
Expand Down Expand Up @@ -1896,30 +1896,21 @@ public static function user_two_factor_options( $user ) {
* @return bool True if the provider was enabled, false otherwise.
*/
public static function enable_provider_for_user( $user_id, $new_provider ) {
$available_providers = self::get_providers();

if ( ! array_key_exists( $new_provider, $available_providers ) ) {
// Ensure the provider is even available.
if ( ! array_key_exists( $new_provider, self::get_providers() ) ) {
return false;
}

$user = get_userdata( $user_id );
$enabled_providers = self::get_enabled_providers_for_user( $user );
$enabled_providers = self::get_enabled_providers_for_user( $user_id );

// Check if this is enabled already.
if ( in_array( $new_provider, $enabled_providers ) ) {
return true;
}

$enabled_providers[] = $new_provider;
$enabled = update_user_meta( $user_id, self::ENABLED_PROVIDERS_USER_META_KEY, $enabled_providers );

// Primary provider must be enabled.
$has_primary = is_object( self::get_primary_provider_for_user( $user_id ) );

if ( ! $has_primary ) {
$has_primary = update_user_meta( $user_id, self::PROVIDER_USER_META_KEY, $new_provider );
}

return $enabled && $has_primary;
return (bool) update_user_meta( $user_id, self::ENABLED_PROVIDERS_USER_META_KEY, $enabled_providers );
}

/**
Expand All @@ -1936,23 +1927,27 @@ public static function enable_provider_for_user( $user_id, $new_provider ) {
* @return bool True if the provider was disabled, false otherwise.
*/
public static function disable_provider_for_user( $user_id, $provider_to_delete ) {
$is_registered = array_key_exists( $provider_to_delete, self::get_providers() );

if ( ! $is_registered ) {
// Check if the provider is even enabled.
if ( ! array_key_exists( $provider_to_delete, self::get_providers() ) ) {
return false;
}

$old_enabled_providers = self::get_enabled_providers_for_user( $user_id );
$is_enabled = in_array( $provider_to_delete, $old_enabled_providers );
$enabled_providers = self::get_enabled_providers_for_user( $user_id );

if ( ! $is_enabled ) {
// Check if this is disabled already.
if ( ! in_array( $provider_to_delete, $enabled_providers ) ) {
return true;
}

$new_enabled_providers = array_diff( $old_enabled_providers, array( $provider_to_delete ) );
$was_disabled = update_user_meta( $user_id, self::ENABLED_PROVIDERS_USER_META_KEY, $new_enabled_providers );
$enabled_providers = array_diff( $enabled_providers, array( $provider_to_delete ) );

// Remove this from being a primary provider, if set.
$primary_provider = self::get_primary_provider_for_user( $user_id );
if ( $primary_provider && $primary_provider->get_key() === $provider_to_delete ) {
delete_user_meta( $user_id, self::PROVIDER_USER_META_KEY );
}

return (bool) $was_disabled;
return (bool) update_user_meta( $user_id, self::ENABLED_PROVIDERS_USER_META_KEY, $enabled_providers );
}

/**
Expand Down
8 changes: 8 additions & 0 deletions providers/class-two-factor-totp.php
View file
Original file line number Diff line number Diff line change
Expand Up @@ -151,6 +151,10 @@ public function rest_delete_totp( $request ) {

$this->delete_user_totp_key( $user_id );

if ( ! Two_Factor_Core::disable_provider_for_user( $user_id, 'Two_Factor_Totp' ) ) {
return new WP_Error( 'db_error', __( 'Unable to disable TOTP provider for this user.', 'two-factor' ), array( 'status' => 500 ) );
}

ob_start();
$this->user_two_factor_options( $user );
$html = ob_get_clean();
Expand Down Expand Up @@ -375,6 +379,7 @@ public function user_two_factor_options( $user ) {
user_id: <?php echo wp_json_encode( $user->ID ); ?>,
key: key,
code: code,
enable_provider: true,
}
} ).fail( function( response, status ) {
var errorMessage = response.responseJSON.message || status,
Expand All @@ -386,8 +391,10 @@ public function user_two_factor_options( $user ) {

$error.find('p').text( errorMessage );

$( '#enabled-Two_Factor_Totp' ).prop( 'checked', false );
$('#two-factor-totp-authcode').val('');
} ).then( function( response ) {
$( '#enabled-Two_Factor_Totp' ).prop( 'checked', true );
$( '#two-factor-totp-options' ).html( response.html );
} );
} );
Expand All @@ -414,6 +421,7 @@ public function user_two_factor_options( $user ) {
user_id: <?php echo wp_json_encode( $user->ID ); ?>,
}
} ).then( function( response ) {
$( '#enabled-Two_Factor_Totp' ).prop( 'checked', false );
$( '#two-factor-totp-options' ).html( response.html );
} );
} );
Expand Down
36 changes: 18 additions & 18 deletions tests/class-two-factor-core.php
View file
Original file line number Diff line number Diff line change
Expand Up @@ -818,52 +818,52 @@ public function test_show_password_reset_error() {
public function test_enable_disable_provider_for_user() {
$user = self::factory()->user->create_and_get();
$enabled_providers = Two_Factor_Core::get_enabled_providers_for_user( $user->ID );
$this->assertEmpty( $enabled_providers );
$this->assertEmpty( $enabled_providers, 'No providers are enabled by default' );

// Disabling one that's already disabled should succeed.
$totp_disabled = Two_Factor_Core::disable_provider_for_user( $user->ID, 'Two_Factor_Totp' );
$this->assertTrue( $totp_disabled );
$this->assertTrue( $totp_disabled, 'Disabling something that wasn\'t enabled should succeed' );

// Disabling one that doesn't exist should fail.
$nonexistent_enabled = Two_Factor_Core::enable_provider_for_user( $user->ID, 'Nonexistent_Provider' );
$enabled_providers = Two_Factor_Core::get_enabled_providers_for_user( $user->ID );
$this->assertFalse( $nonexistent_enabled );
$this->assertEmpty( $enabled_providers );
$this->assertNull( Two_Factor_Core::get_primary_provider_for_user( $user->ID ) );
$this->assertFalse( $nonexistent_enabled, 'Nonexistent shouldn\'t be allowed to be enabled' );
$this->assertEmpty( $enabled_providers, 'Nonexistent wasn\'t enabled' );
$this->assertNull( Two_Factor_Core::get_primary_provider_for_user( $user->ID ), 'Nonexistent wasn\'t set as primary' );

// Enabling a valid one should succeed. The first one that's enabled and configured should be the default primary.
$totp = Two_Factor_Totp::get_instance();
$totp->set_user_totp_key( $user->ID, 'foo' );
$totp_enabled = Two_Factor_Core::enable_provider_for_user( $user->ID, 'Two_Factor_Totp' );
$enabled_providers = Two_Factor_Core::get_enabled_providers_for_user( $user->ID );
$this->assertTrue( $totp_enabled );
$this->assertSame( array( 'Two_Factor_Totp' ), $enabled_providers );
$this->assertSame( 'Two_Factor_Totp', Two_Factor_Core::get_primary_provider_for_user( $user->ID )->get_key() );
$this->assertTrue( $totp_enabled, 'Can enable a valid provider' );
$this->assertSame( array( 'Two_Factor_Totp' ), $enabled_providers, 'Enabled provider is now listed as enabled' );
$this->assertSame( 'Two_Factor_Totp', Two_Factor_Core::get_primary_provider_for_user( $user->ID )->get_key(), 'Primary is now the only enabled provider' );

// Enabling one that's already enabled should succeed.
$totp_enabled = Two_Factor_Core::enable_provider_for_user( $user->ID, 'Two_Factor_Totp' );
$this->assertTrue( $totp_enabled );
$this->assertTrue( $totp_enabled, 'Can enable a provider that is already enabled' );

// Enabling another should succeed, and not change the primary.
$dummy_enabled = Two_Factor_Core::enable_provider_for_user( $user->ID, 'Two_Factor_Dummy' );
$enabled_providers = Two_Factor_Core::get_enabled_providers_for_user( $user->ID );
$this->assertTrue( $dummy_enabled );
$this->assertSame( array( 'Two_Factor_Totp', 'Two_Factor_Dummy' ), $enabled_providers );
$this->assertSame( 'Two_Factor_Totp', Two_Factor_Core::get_primary_provider_for_user( $user->ID )->get_key() );
$this->assertTrue( $dummy_enabled, 'Can enable valid provider' );
$this->assertSame( array( 'Two_Factor_Totp', 'Two_Factor_Dummy' ), $enabled_providers, 'Multiple can be enabled at the same time' );
$this->assertSame( 'Two_Factor_Totp', Two_Factor_Core::get_primary_provider_for_user( $user->ID )->get_key(), 'The primary not changed upon additional providers being enabled' );

// Disabling one that doesn't exist should fail.
$nonexistent_disabled = Two_Factor_Core::disable_provider_for_user( $user->ID, 'Nonexistent_Provider' );
$enabled_providers = Two_Factor_Core::get_enabled_providers_for_user( $user->ID );
$this->assertFalse( $nonexistent_disabled );
$this->assertSame( array( 'Two_Factor_Totp', 'Two_Factor_Dummy' ), $enabled_providers );
$this->assertSame( 'Two_Factor_Totp', Two_Factor_Core::get_primary_provider_for_user( $user->ID )->get_key() );
$this->assertFalse( $nonexistent_disabled, 'Unavailable provider can\'t be disabled' );
$this->assertSame( array( 'Two_Factor_Totp', 'Two_Factor_Dummy' ), $enabled_providers, 'Unavailable wasn\'t added to the list of enabled proviers' );
$this->assertSame( 'Two_Factor_Totp', Two_Factor_Core::get_primary_provider_for_user( $user->ID )->get_key(), 'The primary is still the same after unavailable disable attempt' );

// Disabling one that's enabled should succeed, and change the primary to the next available one.
$totp_disabled = Two_Factor_Core::disable_provider_for_user( $user->ID, 'Two_Factor_Totp' );
$enabled_providers = Two_Factor_Core::get_enabled_providers_for_user( $user->ID );
$this->assertTrue( $totp_disabled ); //todo enable and fix
$this->assertSame( array( 1 => 'Two_Factor_Dummy' ), $enabled_providers );
$this->assertSame( 'Two_Factor_Dummy', Two_Factor_Core::get_primary_provider_for_user( $user->ID )->get_key() );
$this->assertTrue( $totp_disabled, 'Can disable a provider that is enabled' );
$this->assertSame( array( 1 => 'Two_Factor_Dummy' ), $enabled_providers, 'The other providers are kept enabled' );
$this->assertSame( 'Two_Factor_Dummy', Two_Factor_Core::get_primary_provider_for_user( $user->ID )->get_key(), 'Primary is updated to the first available' );
}

/**
Expand Down

0 comments on commit 4c733e4

Please sign in to comment.