From 8baec519d747ff667642954e377e510a24d05795 Mon Sep 17 00:00:00 2001
From: Ian Dunn <ian@iandunn.name>
Date: Mon, 6 Feb 2023 13:13:40 -0800
Subject: [PATCH] Convert Backup Codes AJAX to REST API

---
 class-two-factor-core.php                     |  42 ++++
 providers/class-two-factor-backup-codes.php   | 115 +++++++----
 .../class-two-factor-backup-codes-ajax.php    | 110 ----------
 ...class-two-factor-backup-codes-rest-api.php | 193 ++++++++++++++++++
 .../class-two-factor-backup-codes.php         |   6 +-
 5 files changed, 311 insertions(+), 155 deletions(-)
 delete mode 100644 tests/providers/class-two-factor-backup-codes-ajax.php
 create mode 100644 tests/providers/class-two-factor-backup-codes-rest-api.php

diff --git a/class-two-factor-core.php b/class-two-factor-core.php
index b70d619d..ee15420a 100644
--- a/class-two-factor-core.php
+++ b/class-two-factor-core.php
@@ -49,6 +49,13 @@ class Two_Factor_Core {
 	 */
 	const USER_SETTINGS_ACTION_NONCE_QUERY_ARG = '_two_factor_action_nonce';
 
+	/**
+	 * Namespace for plugin rest api endpoints.
+	 *
+	 * @var string
+	 */
+	const REST_NAMESPACE = 'two-factor/1.0';
+
 	/**
 	 * Keep track of all the password-based authentication sessions that
 	 * need to invalidated before the second factor authentication.
@@ -1076,6 +1083,41 @@ 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 );
+
+		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;
+	}
+
 	/**
 	 * Update the user meta value.
 	 *
diff --git a/providers/class-two-factor-backup-codes.php b/providers/class-two-factor-backup-codes.php
index 62707f2f..7a8da399 100644
--- a/providers/class-two-factor-backup-codes.php
+++ b/providers/class-two-factor-backup-codes.php
@@ -48,13 +48,49 @@ public static function get_instance() {
 	 * @since 0.1-dev
 	 */
 	protected function __construct() {
+		add_action( 'rest_api_init', array( $this, 'register_rest_routes' ) );
 		add_action( 'two_factor_user_options_' . __CLASS__, array( $this, 'user_options' ) );
 		add_action( 'admin_notices', array( $this, 'admin_notices' ) );
-		add_action( 'wp_ajax_two_factor_backup_codes_generate', array( $this, 'ajax_generate_json' ) );
 
 		return parent::__construct();
 	}
 
+	/**
+	 * Register the rest-api endpoints required for this provider.
+	 */
+	public function register_rest_routes() {
+		register_rest_route(
+			Two_Factor_Core::REST_NAMESPACE,
+			'/generate-backup-codes',
+			array(
+				'methods'             => WP_REST_Server::CREATABLE,
+				'callback'            => array( $this, 'rest_generate_codes' ),
+				'permission_callback' => function( $request ) {
+					return current_user_can( 'edit_user', $request['user_id'] );
+				},
+				'args'                => array(
+					'user_id' => array(
+						'required' => true,
+						'type'     => 'number',
+					),
+					'number' => array(
+						'type'    => 'number',
+						'default' => self::NUMBER_OF_CODES,
+					),
+					'append' => array(
+						'type'    => 'boolean',
+						'default' => false,
+					),
+					'enable_provider' => array(
+						'required' => false,
+						'type'     => 'boolean',
+						'default'  => false,
+					),
+				),
+			)
+		);
+	}
+
 	/**
 	 * Displays an admin notice when backup codes have run out.
 	 *
@@ -125,8 +161,7 @@ public function is_available_for_user( $user ) {
 	 * @param WP_User $user WP_User object of the logged-in user.
 	 */
 	public function user_options( $user ) {
-		$ajax_nonce = wp_create_nonce( 'two-factor-backup-codes-generate-json-' . $user->ID );
-		$count      = self::codes_remaining_for_user( $user );
+		$count = self::codes_remaining_for_user( $user );
 		?>
 		<p id="two-factor-backup-codes">
 			<button type="button" class="button button-two-factor-backup-codes-generate button-secondary hide-if-no-js">
@@ -154,30 +189,26 @@ public function user_options( $user ) {
 		<script type="text/javascript">
 			( function( $ ) {
 				$( '.button-two-factor-backup-codes-generate' ).click( function() {
-					$.ajax( {
+					wp.apiRequest( {
 						method: 'POST',
-						url: ajaxurl,
+						path: <?php echo wp_json_encode( Two_Factor_Core::REST_NAMESPACE . '/generate-backup-codes' ); ?>,
 						data: {
-							action: 'two_factor_backup_codes_generate',
-							user_id: '<?php echo esc_js( $user->ID ); ?>',
-							nonce: '<?php echo esc_js( $ajax_nonce ); ?>'
-						},
-						dataType: 'JSON',
-						success: function( response ) {
-							var $codesList = $( '.two-factor-backup-codes-unused-codes' );
-
-							$( '.two-factor-backup-codes-wrapper' ).show();
-							$codesList.html( '' );
-
-							// Append the codes.
-							for ( i = 0; i < response.data.codes.length; i++ ) {
-								$codesList.append( '<li>' + response.data.codes[ i ] + '</li>' );
-							}
-
-							// Update counter.
-							$( '.two-factor-backup-codes-count' ).html( response.data.i18n.count );
-							$( '#two-factor-backup-codes-download-link' ).attr( 'href', response.data.download_link );
+							user_id: <?php echo wp_json_encode( $user->ID ); ?>
 						}
+					} ).then( function( response ) {
+						var $codesList = $( '.two-factor-backup-codes-unused-codes' );
+
+						$( '.two-factor-backup-codes-wrapper' ).show();
+						$codesList.html( '' );
+
+						// Append the codes.
+						for ( i = 0; i < response.codes.length; i++ ) {
+							$codesList.append( '<li>' + response.codes[ i ] + '</li>' );
+						}
+
+						// Update counter.
+						$( '.two-factor-backup-codes-count' ).html( response.i18n.count );
+						$( '#two-factor-backup-codes-download-link' ).attr( 'href', response.download_link );
 					} );
 				} );
 			} )( jQuery );
@@ -224,21 +255,21 @@ public function generate_codes( $user, $args = '' ) {
 	}
 
 	/**
-	 * Generates a JSON object of backup codes.
+	 * Generates Backup Codes for returning through the WordPress Rest API.
 	 *
-	 * @since 0.1-dev
+	 * @since 0.8.0
 	 */
-	public function ajax_generate_json() {
-		$user_id = 0;
-		if ( ! empty( $_POST['user_id'] ) ) {
-			$user_id = absint( $_POST['user_id'] );
-		}
+	public function rest_generate_codes( $request ) {
+		$user_id = $request['user_id'];
+		$user    = get_user_by( 'id', $user_id );
 
-		$user = get_user_by( 'id', $user_id );
-		check_ajax_referer( 'two-factor-backup-codes-generate-json-' . $user->ID, 'nonce' );
+		$args =  array(
+			'number' => $request['number'],
+			'method' => wp_validate_boolean( $request['append'] ) ? 'append' : 'replace',
+		);
 
 		// Setup the return data.
-		$codes = $this->generate_codes( $user );
+		$codes = $this->generate_codes( $user, $args );
 		$count = self::codes_remaining_for_user( $user );
 		$title = sprintf(
 			/* translators: %s: the site's domain */
@@ -261,13 +292,15 @@ public function ajax_generate_json() {
 			'count' => esc_html( sprintf( _n( '%s unused code remaining.', '%s unused codes remaining.', $count, 'two-factor' ), $count ) ),
 		);
 
-		// Send the response.
-		wp_send_json_success(
-			array(
-				'codes'         => $codes,
-				'download_link' => $download_link,
-				'i18n'          => $i18n,
-			)
+		if ( $request->get_param( 'enable_provider' ) && ! Two_Factor_Core::enable_provider_for_user( $user_id, 'Two_Factor_Backup_Codes' ) ) {
+			return new WP_Error( 'db_error', __( 'Unable to enable Backup Codes provider for this user.', 'two-factor' ), array( 'status' => 500 ) );
+		}
+
+		return array(
+			'codes'         => $codes,
+			'download_link' => $download_link,
+			'remaining'     => $count,
+			'i18n'          => $i18n,
 		);
 	}
 
diff --git a/tests/providers/class-two-factor-backup-codes-ajax.php b/tests/providers/class-two-factor-backup-codes-ajax.php
deleted file mode 100644
index 1c9cbd55..00000000
--- a/tests/providers/class-two-factor-backup-codes-ajax.php
+++ /dev/null
@@ -1,110 +0,0 @@
-<?php
-/**
- * Test Two Factor Backup Codes.
- *
- * @package Two_Factor
- */
-
-/**
- * Class Tests_Two_Factor_Backup_Codes_AJAX
- *
- * @package Two_Factor
- * @group providers
- */
-class Tests_Two_Factor_Backup_Codes_AJAX extends WP_Ajax_UnitTestCase {
-
-	/**
-	 * Instance of our provider class.
-	 *
-	 * @var Two_Factor_Backup_Codes
-	 */
-	protected $provider;
-
-	/**
-	 * Set up a test case.
-	 *
-	 * @see WP_UnitTestCase_Base::set_up()
-	 */
-	public function set_up() {
-		parent::set_up();
-		$this->provider = Two_Factor_Backup_Codes::get_instance();
-	}
-
-	/**
-	 * Verify that the downloaded file contains the codes.
-	 *
-	 * @covers Two_Factor_Backup_Codes::ajax_generate_json
-	 */
-	public function test_generate_code_and_validate_in_download_file() {
-		$this->_setRole( 'administrator' );
-
-		$user             = wp_get_current_user();
-		$_POST['user_id'] = $user->ID;
-		$_POST['nonce']   = wp_create_nonce( 'two-factor-backup-codes-generate-json-' . $user->ID );
-
-		try {
-			$this->_handleAjax( 'two_factor_backup_codes_generate' );
-		} catch ( WPAjaxDieContinueException $e ) {
-			unset( $e );
-		}
-
-		$this->assertStringContainsString( 'download_link', $this->_last_response );
-
-		$response = json_decode( $this->_last_response );
-
-		$this->assertTrue( $response->success );
-		$this->assertNotEmpty( $response->data->codes );
-		$this->assertTrue( $this->provider->validate_code( $user, $response->data->codes[0] ) );
-		$this->assertStringContainsString( $response->data->codes[0], $response->data->download_link );
-	}
-
-	/**
-	 * Verify that a different user cannot generate codes for another.
-	 *
-	 * @covers Two_Factor_Backup_Codes::ajax_generate_json
-	 */
-	public function test_cannot_generate_code_for_different_user() {
-		$this->_setRole( 'administrator' );
-
-		$user           = wp_get_current_user();
-		$_POST['nonce'] = wp_create_nonce( 'two-factor-backup-codes-generate-json-' . $user->ID );
-
-		// Create a new user
-		$user             = new WP_User( self::factory()->user->create() );
-		$_POST['user_id'] = $user->ID;
-
-		$this->expectException( 'WPAjaxDieStopException' );
-		$this->expectExceptionMessage( '-1' );
-		$this->_handleAjax( 'two_factor_backup_codes_generate' );
-	}
-
-	/**
-	 * Verify that an admin can create Backup codes for another user.
-	 *
-	 * @covers Two_Factor_Backup_Codes::ajax_generate_json
-	 */
-	public function test_generate_codes_for_other_users() {
-		$this->_setRole( 'administrator' );
-
-		$current_user     = wp_get_current_user();
-		$user             =  new WP_User( self::factory()->user->create() );
-		$_POST['user_id'] = $user->ID;
-		$_POST['nonce']   = wp_create_nonce( 'two-factor-backup-codes-generate-json-' . $user->ID );
-
-		try {
-			$this->_handleAjax( 'two_factor_backup_codes_generate' );
-		} catch ( WPAjaxDieContinueException $e ) {
-			unset( $e );
-		}
-
-		$this->assertStringContainsString( 'codes', $this->_last_response );
-
-		$response = json_decode( $this->_last_response );
-
-		$this->assertTrue( $response->success );
-		$this->assertNotEmpty( $response->data->codes );
-
-		$this->assertFalse( $this->provider->validate_code( $current_user, $response->data->codes[0] ) );
-		$this->assertTrue( $this->provider->validate_code( $user, $response->data->codes[0] ) );
-	}
-}
diff --git a/tests/providers/class-two-factor-backup-codes-rest-api.php b/tests/providers/class-two-factor-backup-codes-rest-api.php
new file mode 100644
index 00000000..8adf074e
--- /dev/null
+++ b/tests/providers/class-two-factor-backup-codes-rest-api.php
@@ -0,0 +1,193 @@
+<?php
+/**
+ * Test Two Factor Backup Codes.
+ *
+ * @package Two_Factor
+ */
+
+/**
+ * Class Tests_Two_Factor_Backup_Codes_REST_API.
+ *
+ * @package Two_Factor
+ * @group providers
+ */
+class Tests_Two_Factor_Backup_Codes_REST_API extends WP_Test_REST_TestCase {
+
+	/**
+	 * Instance of our provider class.
+	 *
+	 * @var Two_Factor_Backup_Codes
+	 */
+	protected static $provider;
+
+	/**
+	 * Administrator user ID.
+	 *
+	 * @var int
+	 */
+	protected static $admin_id;
+
+	/**
+	 * Editor user ID.
+	 *
+	 * @var int
+	 */
+	protected static $editor_id;
+
+	public static function wpSetUpBeforeClass( WP_UnitTest_Factory $factory ) {
+		self::$admin_id = $factory->user->create(
+				array(
+						'role' => 'administrator',
+				)
+		);
+
+		self::$editor_id = $factory->user->create(
+				array(
+						'role' => 'editor',
+				)
+		);
+
+		self::$provider = Two_Factor_Backup_Codes::get_instance();
+	}
+
+	public static function wpTearDownAfterClass() {
+			self::delete_user( self::$admin_id );
+			self::delete_user( self::$editor_id );
+	}
+
+	/**
+	 * Verify that the downloaded file contains the requested number of codes.
+	 *
+	 * @covers Two_Factor_Backup_Codes::rest_generate_codes
+	 */
+	public function test_generate_code_and_validate_in_download_file() {
+		wp_set_current_user( self::$admin_id );
+
+		$request = new WP_REST_Request( 'POST', '/' . Two_Factor_Core::REST_NAMESPACE . '/generate-backup-codes' );
+		$request->set_body_params(
+			array(
+				'user_id' => self::$admin_id,
+				'number'  => 5,
+			)
+		);
+
+		$response = rest_do_request( $request );
+		$data     = $response->get_data();
+
+		$this->assertEquals( 200, $response->get_status() );
+		$this->assertNotEmpty( $data['download_link'] );
+		$this->assertNotEmpty( $data['codes'] );
+		$this->assertCount( 5, $data['codes'] );
+		$this->assertTrue( self::$provider->validate_code( wp_get_current_user(), $data['codes'][0] ) );
+		$this->assertStringContainsString( $data['codes'][0], $data['download_link'] );
+	}
+
+	/**
+	 * Verify that overwriting, and appending works.
+	 *
+	 * @covers Two_Factor_Backup_Codes::rest_generate_codes
+	 */
+	public function test_generate_code_append() {
+		wp_set_current_user( self::$admin_id );
+
+		$request  = new WP_REST_Request( 'POST', '/' . Two_Factor_Core::REST_NAMESPACE . '/generate-backup-codes' );
+		$request->set_body_params(
+			array(
+				'user_id' => self::$admin_id,
+				'number'  => 5,
+			)
+		);
+
+		$response  = rest_do_request( $request );
+		$discarded = $response->get_data();
+
+		$this->assertEquals( 200, $response->get_status() );
+		$this->assertEquals( 5, $discarded['remaining'] );
+
+		$request = new WP_REST_Request( 'POST', '/' . Two_Factor_Core::REST_NAMESPACE . '/generate-backup-codes' );
+		$request->set_body_params(
+			array(
+				'user_id' => self::$admin_id,
+				'number'  => 5,
+			)
+		);
+
+		$response = rest_do_request( $request );
+		$first    = $response->get_data();
+
+		$this->assertEquals( 200, $response->get_status() );
+		$this->assertNotEmpty( $first['codes'] );
+		$this->assertEquals( 5, $first['remaining'] );
+
+		$request = new WP_REST_Request( 'POST', '/' . Two_Factor_Core::REST_NAMESPACE . '/generate-backup-codes' );
+		$request->set_body_params(
+			array(
+				'user_id' => self::$admin_id,
+				'number'  => 1,
+				'append'  => true,
+			)
+		);
+
+		$response = rest_do_request( $request );
+		$second   = $response->get_data();
+
+		$this->assertEquals( 200, $response->get_status() );
+		$this->assertNotEmpty( $second['codes'] );
+		$this->assertEquals( 6, $second['remaining'] );
+
+		$this->assertEquals( $second['remaining'], self::$provider->codes_remaining_for_user( wp_get_current_user() ) );
+
+		$this->assertFalse( self::$provider->validate_code( wp_get_current_user(), $discarded['codes'][0] ) );
+		$this->assertTrue( self::$provider->validate_code( wp_get_current_user(), $first['codes'][0] ) );
+		$this->assertTrue( self::$provider->validate_code( wp_get_current_user(), $second['codes'][0] ) );
+	}
+
+	/**
+	 * Verify that a user without edit_user capabilities cannot generate codes for another.
+	 *
+	 * @covers Two_Factor_Backup_Codes::rest_generate_codes
+	 */
+	public function test_cannot_generate_code_for_different_user() {
+		wp_set_current_user( self::$editor_id );
+
+		$request = new WP_REST_Request( 'POST', '/' . Two_Factor_Core::REST_NAMESPACE . '/generate-backup-codes' );
+		$request->set_body_params(
+			array(
+				'user_id' => self::$admin_id,
+			)
+		);
+
+		$response = rest_do_request( $request );
+		$data     = $response->get_data();
+
+		$this->assertErrorResponse( 'rest_forbidden', $response, 403 );
+
+		$this->assertArrayNotHasKey( 'download_link', $data );
+		$this->assertArrayNotHasKey( 'codes', $data );
+	}
+
+	/**
+	 * Verify that an admin can create Backup codes for another user.
+	 *
+	 * @covers Two_Factor_Backup_Codes::ajax_generate_json
+	 */
+	public function test_generate_codes_for_other_users() {
+		wp_set_current_user( self::$admin_id );
+
+		$request = new WP_REST_Request( 'POST', '/' . Two_Factor_Core::REST_NAMESPACE . '/generate-backup-codes' );
+		$request->set_body_params(
+			array(
+				'user_id' => self::$editor_id,
+			)
+		);
+
+		$response = rest_do_request( $request );
+		$data     = $response->get_data();
+
+		$this->assertEquals( 200, $response->get_status() );
+		$this->assertNotEmpty( $data['codes'] );
+
+		$this->assertFalse( self::$provider->validate_code( wp_get_current_user(), $data['codes'][0] ) );
+		$this->assertTrue( self::$provider->validate_code( get_user_by( 'id', self::$editor_id ), $data['codes'][0] ) );
+	}
+}
diff --git a/tests/providers/class-two-factor-backup-codes.php b/tests/providers/class-two-factor-backup-codes.php
index e43e5055..a198f6f0 100644
--- a/tests/providers/class-two-factor-backup-codes.php
+++ b/tests/providers/class-two-factor-backup-codes.php
@@ -27,6 +27,7 @@ class Tests_Two_Factor_Backup_Codes extends WP_UnitTestCase {
 	 */
 	public function set_up() {
 		parent::set_up();
+
 		$this->provider = Two_Factor_Backup_Codes::get_instance();
 	}
 
@@ -152,7 +153,6 @@ public function test_generate_code_and_validate_code_false_different_users() {
 	 */
 	public function test_user_options() {
 		$user  = new WP_User( self::factory()->user->create() );
-		$nonce = wp_create_nonce( 'two-factor-backup-codes-generate-json-' . $user->ID );
 
 		ob_start();
 		$this->provider->user_options( $user );
@@ -160,8 +160,7 @@ public function test_user_options() {
 
 		$this->assertStringContainsString( '<p id="two-factor-backup-codes">', $buffer );
 		$this->assertStringContainsString( '<div class="two-factor-backup-codes-wrapper" style="display:none;">', $buffer );
-		$this->assertStringContainsString( "user_id: '{$user->ID}'", $buffer );
-		$this->assertStringContainsString( "nonce: '{$nonce}'", $buffer );
+		$this->assertStringContainsString( "user_id: {$user->ID}", $buffer );
 	}
 
 	/**
@@ -190,5 +189,4 @@ public function test_delete_code() {
 		$this->provider->delete_code( $user, $backup_codes[0] );
 		$this->assertEquals( 1, $this->provider->codes_remaining_for_user( $user ) );
 	}
-
 }