Update TOTP Scratch Tokens to array

Bug: T237303
Change-Id: Ie8de059888363bf1cea4f0b268a46faaa5671904

maintenance/updateTOTPScratchTokensToArray.php

@@ -0,0 +1,57 @@
+<?php
+/**
+ * Updates TOTP Scratch Tokens to an array
+ *
+ * Usage: php updateTOTPScratchTokensToArray.php
+ *
+ * This program is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation; either version 2 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License along
+ * with this program; if not, write to the Free Software Foundation, Inc.,
+ * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
+ * http://www.gnu.org/copyleft/gpl.html
+ *
+ * @file
+ * @ingroup Maintenance
+ */
+
+use MediaWiki\Extension\OATHAuth\Hook\LoadExtensionSchemaUpdates\UpdateTables;
+use MediaWiki\MediaWikiServices;
+
+if ( getenv( 'MW_INSTALL_PATH' ) ) {
+	$IP = getenv( 'MW_INSTALL_PATH' );
+} else {
+	$IP = __DIR__ . '/../../..';
+}
+require_once "$IP/maintenance/Maintenance.php";
+
+class UpdateTOTPScratchTokensToArray extends Maintenance {
+	public function __construct() {
+		parent::__construct();
+		$this->addDescription( 'Script to update TOTP Scratch Tokens to an array' );
+		$this->requireExtension( 'OATHAuth' );
+	}
+
+	public function execute() {
+		global $wgOATHAuthDatabase;
+		$lb = MediaWikiServices::getInstance()->getDBLoadBalancerFactory()
+			->getMainLB( $wgOATHAuthDatabase );
+		$dbw = $lb->getConnectionRef( DB_MASTER, [], $wgOATHAuthDatabase );
+
+		if ( !UpdateTables::switchTOTPScratchTokensToArray( $dbw ) ) {
+			$this->error( "Failed to update TOTP Scratch Tokens.\n", 1 );
+		}
+		$this->output( "Done.\n" );
+	}
+}
+
+$maintClass = UpdateTOTPScratchTokensToArray::class;
+require_once RUN_MAINTENANCE_IF_MAIN;

src/Hook/LoadExtensionSchemaUpdates/UpdateTables.php

@@ -71,6 +71,10 @@ class UpdateTables {
 					[ [ __CLASS__, 'schemaUpdateTOTPToMultipleKeys' ] ]
 				);
 
+				$this->updater->addExtensionUpdate(
+					[ [ __CLASS__, 'schemaUpdateTOTPScratchTokensToArray' ] ]
+				);
+
 				break;
 
 			case 'postgres':
@@ -124,6 +128,16 @@ class UpdateTables {
 		return self::switchTOTPToMultipleKeys( self::getDatabase() );
 	}
 
+	/**
+	 * Helper function for converting single TOTP keys to multi-key system
+	 * @param DatabaseUpdater $updater
+	 * @return bool
+	 * @throws ConfigException
+	 */
+	public static function schemaUpdateTOTPScratchTokensToArray( DatabaseUpdater $updater ) {
+		return self::switchTOTPScratchTokensToArray( self::getDatabase() );
+	}
+
 	/**
 	 * Converts old, TOTP specific, column values to new structure
 	 * @param IDatabase $db
@@ -219,6 +233,49 @@ class UpdateTables {
 		return true;
 	}
 
+	/**
+	 * Switch scratch tokens from string to an array
+	 *
+	 * @param IDatabase $db
+	 * @return bool
+	 * @throws ConfigException
+	 */
+	public static function switchTOTPScratchTokensToArray( IDatabase $db ) {
+		if ( !$db->fieldExists( 'oathauth_users', 'data' ) ) {
+			return true;
+		}
+
+		$res = $db->select(
+			'oathauth_users',
+			[ 'id', 'data' ],
+			[
+				'module' => 'totp'
+			],
+			__METHOD__
+		);
+
+		foreach ( $res as $row ) {
+			$data = FormatJson::decode( $row->data, true );
+
+			foreach ( $data['keys'] as &$k ) {
+				if ( is_string( $k['scratch_tokens'] ) ) {
+					$k['scratch_tokens'] = explode( ',', $k['scratch_tokens'] );
+				}
+			}
+
+			$db->update(
+				'oathauth_users',
+				[
+					'data' => FormatJson::encode( $data )
+				],
+				[ 'id' => $row->id ],
+				__METHOD__
+			);
+		}
+
+		return true;
+	}
+
 	/**
 	 * Helper function for converting old users to the new schema
 	 *