Enable the provider when setting up TOTP

class-two-factor-core.php

@@ -1051,6 +1051,36 @@ public static function user_two_factor_options( $user ) {
 		do_action( 'show_user_security_settings', $user );
 	}
 
+	/**
+	 * Enable a provider for a user.
+	 *
+	 * @param int    $user_id      The ID of the user.
+	 * @param string $new_provider The name of the provider class.
+	 *
+	 * @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 ) ) {
+			return false;
+		}
+
+		$user                = get_userdata( $user_id );
+		$enabled_providers   = self::get_enabled_providers_for_user( $user );
+		$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 = 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;
+	}
+
 	/**
 	 * Update the user meta value.
 	 *

providers/class-two-factor-totp.php

@@ -174,6 +174,10 @@ public function rest_setup_totp( $request ) {
 			return new WP_Error( 'db_error', __( 'Unable to save Two Factor Authentication code. Please re-scan the QR code and enter the code provided by your application.', 'two-factor' ), array( 'status' => 500 ) );
 		}
 
+		if ( ! Two_Factor_Core::enable_provider_for_user( $user_id, 'Two_Factor_Totp' ) ) {
+			return new WP_Error( 'db_error', __( 'Unable to enable TOTP provider for this user.', 'two-factor' ), array( 'status' => 500 ) );
+		}
+
 		ob_start();
 		$this->user_two_factor_options( $user );
 		$html = ob_get_clean();