Code Coverage
 
Lines
Functions and Methods
Classes and Traits
Total
0.00% covered (danger)
0.00%
0 / 78
0.00% covered (danger)
0.00%
0 / 4
CRAP
0.00% covered (danger)
0.00%
0 / 1
MoveRecoveryCodesFromTOTP
0.00% covered (danger)
0.00%
0 / 78
0.00% covered (danger)
0.00%
0 / 4
90
0.00% covered (danger)
0.00%
0 / 1
 __construct
0.00% covered (danger)
0.00%
0 / 3
0.00% covered (danger)
0.00%
0 / 1
2
 doDBUpdates
0.00% covered (danger)
0.00%
0 / 68
0.00% covered (danger)
0.00%
0 / 1
42
 updateRow
0.00% covered (danger)
0.00%
0 / 6
0.00% covered (danger)
0.00%
0 / 1
2
 getUpdateKey
0.00% covered (danger)
0.00%
0 / 1
0.00% covered (danger)
0.00%
0 / 1
2
1<?php
2/**
3 * @license GPL-2.0-or-later
4 *
5 * @file
6 * @ingroup Maintenance
7 */
8
9declare( strict_types = 1 );
10
11namespace MediaWiki\Extension\OATHAuth\Maintenance;
12
13use JsonSerializable;
14use MediaWiki\Extension\OATHAuth\Key\RecoveryCodeKeys;
15use MediaWiki\Extension\OATHAuth\Key\TOTPKey;
16use MediaWiki\Extension\OATHAuth\Module\RecoveryCodes;
17use MediaWiki\Extension\OATHAuth\Module\TOTP;
18use MediaWiki\Extension\OATHAuth\OATHAuthServices;
19use MediaWiki\Json\FormatJson;
20use MediaWiki\Maintenance\LoggedUpdateMaintenance;
21use Wikimedia\Rdbms\IDatabase;
22
23// @codeCoverageIgnoreStart
24$IP = getenv( 'MW_INSTALL_PATH' );
25if ( $IP === false ) {
26    $IP = __DIR__ . '/../../..';
27}
28
29require_once "$IP/maintenance/Maintenance.php";
30// @codeCoverageIgnoreEnd
31
32/**
33 * Moves TOTP scratch_tokens to their own recoverykeys rows
34 *
35 * Usage: php MoveRecoveryCodesFromTOTP.php
36 */
37class MoveRecoveryCodesFromTOTP extends LoggedUpdateMaintenance {
38
39    public function __construct() {
40        parent::__construct();
41        $this->requireExtension( 'OATHAuth' );
42        $this->addDescription( 'Moves TOTP scratch_tokens to their own recoverykeys rows' );
43    }
44
45    /** @inheritDoc */
46    protected function doDBUpdates() {
47        $startTime = time();
48        $updatedCount = 0;
49        $totalRows = 0;
50
51        $services = $this->getServiceContainer();
52
53        $moduleRegistry = OATHAuthServices::getInstance()->getModuleRegistry();
54        $recoveryModuleId = $moduleRegistry->getModuleId( RecoveryCodes::MODULE_NAME );
55        $totpModuleId = $moduleRegistry->getModuleId( TOTP::MODULE_NAME );
56
57        $dbw = $services
58            ->getDBLoadBalancerFactory()
59            ->getPrimaryDatabase( 'virtual-oathauth' );
60        $res = $dbw->newSelectQueryBuilder()
61            ->select( [ 'oad_id', 'oad_data', 'oad_user', 'oad_created' ] )
62            ->from( 'oathauth_devices' )
63            ->where( [ 'oad_type' => $totpModuleId ] )
64            ->caller( __METHOD__ )
65            ->fetchResultSet();
66
67        foreach ( $res as $row ) {
68            $totalRows++;
69            $data = FormatJson::decode( $row->oad_data, true );
70
71            if ( !isset( $data['scratch_tokens'] ) ) {
72                // No scratch_tokens in oad_data to move, skip the row
73                continue;
74            }
75
76            $recoveryData = [
77                'recoverycodekeys' => $data['scratch_tokens'],
78            ];
79
80            // Remove scratch_tokens from oad_data because they're being transferred
81            unset( $data['scratch_tokens'] );
82
83            $recoveryCodesRow = $dbw->newSelectQueryBuilder()
84                ->select( [ 'oad_id', 'oad_data' ] )
85                ->from( 'oathauth_devices' )
86                ->where( [ 'oad_user' => $row->oad_user, 'oad_type' => $recoveryModuleId ] )
87                ->caller( __METHOD__ )
88                ->fetchRow();
89
90            $this->beginTransactionRound( __METHOD__ );
91
92            // Update totp row without scratch_tokens
93            // T406953 - Also explicitly remove the empty scratch_tokens array from oad_data; these were incorrectly
94            // added for a while. This means we don't need to keep WMF back compat code around for a long time.
95            $this->updateRow(
96                $dbw,
97                (int)$row->oad_id,
98                TOTPKey::newFromArray( $data )
99            );
100
101            if ( $recoveryData['recoverycodekeys'] === [] ) {
102                // No rows to actually migrate, but we've cleaned up the TOTP row, so we can skip to the next
103                $this->commitTransactionRound( __METHOD__ );
104                continue;
105            }
106
107            if ( $recoveryCodesRow ) {
108                $keys = FormatJson::decode( $recoveryCodesRow->oad_data, true );
109                // Prepend new style recovery codes to existing from TOTP row
110                $recoveryData['recoverycodekeys'] = array_merge(
111                    $keys['recoverycodekeys'],
112                    $recoveryData['recoverycodekeys']
113                );
114
115                $this->updateRow(
116                    $dbw,
117                    (int)$recoveryCodesRow->oad_id,
118                    RecoveryCodeKeys::newFromArray( $recoveryData )
119                );
120            } else {
121                // Insert a new recoverykeys row for this user
122                $dbw->newInsertQueryBuilder()
123                    ->insertInto( 'oathauth_devices' )
124                    ->row( [
125                        'oad_user' => $row->oad_user,
126                        'oad_type' => $recoveryModuleId,
127                        'oad_data' => FormatJson::encode( RecoveryCodeKeys::newFromArray( $recoveryData ) ),
128                        // Use the existing timestamp if available, otherwise use the current timestamp
129                        'oad_created' => $row->oad_created ?? $dbw->timestamp(),
130                    ] )
131                    ->caller( __METHOD__ )
132                    ->execute();
133            }
134
135            $this->commitTransactionRound( __METHOD__ );
136            $updatedCount++;
137            if ( $updatedCount % 50 === 0 ) {
138                $this->output( "{$updatedCount}\n" );
139            }
140        }
141
142        $totalTimeInSeconds = time() - $startTime;
143        $this->output( "Done. Updated {$updatedCount} of {$totalRows} rows in {$totalTimeInSeconds} seconds.\n" );
144        return true;
145    }
146
147    private function updateRow( IDatabase $dbw, int $id, JsonSerializable $data ): void {
148        $dbw->newUpdateQueryBuilder()
149            ->update( 'oathauth_devices' )
150            ->set( [ 'oad_data' => FormatJson::encode( $data->jsonSerialize() ) ] )
151            ->where( [ 'oad_id' => $id ] )
152            ->caller( __METHOD__ )
153            ->execute();
154    }
155
156    /** @return string */
157    protected function getUpdateKey() {
158        return __CLASS__;
159    }
160}
161
162// @codeCoverageIgnoreStart
163$maintClass = MoveRecoveryCodesFromTOTP::class;
164require_once RUN_MAINTENANCE_IF_MAIN;
165// @codeCoverageIgnoreEnd