<?php
/**
 * Two-factor authentication (2FA) with an authenticator app.
 *
 * @package Aslansheed\OTP
 */

namespace Aslansheed\OTP\Auth;

use Aslansheed\OTP\Core\Crypto;
use Aslansheed\OTP\Core\Options;
use Aslansheed\OTP\Core\Request;
use Aslansheed\OTP\Security\Audit_Log;
use WP_Error;

if ( ! defined( 'ABSPATH' ) ) {
	exit;
}

/**
 * Decides who needs a second step, stores the secrets and checks the codes.
 *
 * The site owner picks the mode (off, optional, required) and who it covers
 * (everyone, chosen roles or chosen users). The second step runs after the
 * first proof of identity (code, password or Google) and before the sign-in
 * cookie is set, so nobody is half signed in.
 */
final class Two_Factor {

	const META_SECRET    = 'aslansheed_totp_secret';
	const META_ENABLED   = 'aslansheed_totp_enabled';
	const META_LAST_STEP = 'aslansheed_totp_last_step';
	const META_RECOVERY  = 'aslansheed_totp_recovery';
	const META_SKIPPED   = 'aslansheed_totp_skipped';

	const PENDING_PREFIX = '_aslansheed_otp_2fa_';
	const SETUP_PREFIX   = '_aslansheed_otp_2fa_setup_';
	const PENDING_TTL    = 600;
	const RECOVERY_COUNT = 8;

	/**
	 * Tells whether the feature is switched off for the whole site.
	 *
	 * The constant is an escape hatch for a site owner who lost both the
	 * authenticator and the recovery codes: add it to wp-config.php.
	 *
	 * @return bool
	 */
	public static function is_off() {
		return 'off' === self::mode() || ( defined( 'ASLANSHEED_OTP_DISABLE_2FA' ) && ASLANSHEED_OTP_DISABLE_2FA );
	}

	/**
	 * Returns the configured mode.
	 *
	 * @return string off, optional or required.
	 */
	public static function mode() {
		return (string) Options::get( 'twofa.mode' );
	}

	/**
	 * Returns the role slugs chosen by the site owner.
	 *
	 * @return string[]
	 */
	public static function roles() {
		return array_values( array_filter( array_map( 'sanitize_key', explode( ',', (string) Options::get( 'twofa.roles' ) ) ) ) );
	}

	/**
	 * Tells whether the settings cover a user.
	 *
	 * @param \WP_User $user User to check.
	 * @return bool
	 */
	public static function applies_to( $user ) {
		if ( self::is_off() ) {
			return false;
		}

		$scope = (string) Options::get( 'twofa.scope' );

		if ( 'roles' === $scope ) {
			return array() !== array_intersect( (array) $user->roles, self::roles() );
		}

		if ( 'users' === $scope ) {
			$wanted = array_filter( array_map( 'trim', preg_split( '/[\s,]+/', strtolower( (string) Options::get( 'twofa.users' ) ) ) ) );

			return in_array( (string) $user->ID, $wanted, true )
				|| in_array( strtolower( (string) $user->user_login ), $wanted, true )
				|| in_array( strtolower( (string) $user->user_email ), $wanted, true );
		}

		return true;
	}

	/**
	 * Tells whether the user finished setting up an authenticator.
	 *
	 * @param int $user_id User id.
	 * @return bool
	 */
	public static function is_enabled( $user_id ) {
		return '' !== (string) get_user_meta( $user_id, self::META_ENABLED, true ) && '' !== (string) get_user_meta( $user_id, self::META_SECRET, true );
	}

	/**
	 * Tells what a user must do at sign-in.
	 *
	 * @param \WP_User $user User signing in.
	 * @return string none, challenge, setup_required or setup_offered.
	 */
	public static function state( $user ) {
		if ( ! self::applies_to( $user ) ) {
			return 'none';
		}

		if ( self::is_enabled( (int) $user->ID ) ) {
			return 'challenge';
		}

		if ( 'required' === self::mode() ) {
			return 'setup_required';
		}

		return '' === (string) get_user_meta( $user->ID, self::META_SKIPPED, true ) ? 'setup_offered' : 'none';
	}

	/**
	 * Tells whether a user may switch 2FA off.
	 *
	 * @param \WP_User $user User.
	 * @return bool
	 */
	public static function can_disable( $user ) {
		return ! ( 'required' === self::mode() && self::applies_to( $user ) );
	}

	/**
	 * Called before a sign-in is completed. Returns a response when a step is needed.
	 *
	 * @param \WP_User $user        User who proved their identity.
	 * @param string   $method      First-step method slug.
	 * @param string   $redirect_to Requested redirect.
	 * @return array|null Null when the sign-in may continue.
	 */
	public static function gate( $user, $method, $redirect_to ) {
		$state = self::state( $user );

		if ( 'none' === $state ) {
			return null;
		}

		$pending = array(
			'user_id'  => (int) $user->ID,
			'method'   => (string) $method,
			'redirect' => (string) $redirect_to,
			'state'    => $state,
		);

		if ( 'challenge' !== $state ) {
			$secret            = Totp::generate_secret();
			$pending['secret'] = Crypto::encrypt( $secret );
		}

		$token = wp_generate_password( 32, false );

		set_transient( self::PENDING_PREFIX . hash( 'sha256', $token ), $pending, self::PENDING_TTL );

		$response = array(
			'status' => 'challenge' === $state ? 'two_factor' : 'two_factor_setup',
			'token'  => $token,
		);

		if ( 'challenge' !== $state ) {
			$response['secret']    = $secret;
			$response['uri']       = self::uri( $user, $secret );
			$response['skippable'] = 'setup_offered' === $state;
		}

		return $response;
	}

	/**
	 * Checks the code of a pending sign-in and finishes it.
	 *
	 * @param string $token Pending token.
	 * @param string $code  Authenticator code or recovery code.
	 * @return array|WP_Error
	 */
	public static function verify_login( $token, $code ) {
		$pending = self::pending( $token, 'challenge' );

		if ( is_wp_error( $pending ) ) {
			return $pending;
		}

		$user = get_userdata( $pending['user_id'] );

		if ( ! $user ) {
			return self::expired();
		}

		$checked = self::check_code( $user, $code );

		if ( is_wp_error( $checked ) ) {
			return $checked;
		}

		self::drop( $token );

		return Flow::finish_login( $user, $pending['method'] . '+2fa', $pending['redirect'] );
	}

	/**
	 * Confirms the first code of a setup that happens during sign-in, then finishes it.
	 *
	 * @param string $token Pending token.
	 * @param string $code  Code from the new authenticator entry.
	 * @return array|WP_Error
	 */
	public static function confirm_login_setup( $token, $code ) {
		$pending = self::pending( $token, 'setup' );

		if ( is_wp_error( $pending ) ) {
			return $pending;
		}

		$user = get_userdata( $pending['user_id'] );

		if ( ! $user ) {
			return self::expired();
		}

		$secret = Crypto::decrypt( (string) $pending['secret'] );
		$codes  = self::activate( $user, (string) $secret, $code );

		if ( is_wp_error( $codes ) ) {
			return $codes;
		}

		self::drop( $token );

		$result = Flow::finish_login( $user, $pending['method'] . '+2fa', $pending['redirect'] );

		if ( is_array( $result ) ) {
			$result['status'] = 'two_factor_codes';
			$result['codes']  = $codes;
		}

		return $result;
	}

	/**
	 * Lets the user leave the optional setup for later, then finishes the sign-in.
	 *
	 * @param string $token Pending token.
	 * @return array|WP_Error
	 */
	public static function skip_login_setup( $token ) {
		$pending = self::pending( $token, 'setup' );

		if ( is_wp_error( $pending ) ) {
			return $pending;
		}

		if ( 'setup_offered' !== $pending['state'] ) {
			return new WP_Error( 'aslansheed_2fa_required', __( 'Two-factor authentication (2FA) is required on this site.', 'aslansheed-login' ), array( 'status' => 403 ) );
		}

		$user = get_userdata( $pending['user_id'] );

		if ( ! $user ) {
			return self::expired();
		}

		update_user_meta( $user->ID, self::META_SKIPPED, (string) time() );
		self::drop( $token );

		return Flow::finish_login( $user, $pending['method'], $pending['redirect'] );
	}

	/**
	 * Starts a setup for a signed-in user, from their profile.
	 *
	 * @param int $user_id User id.
	 * @return array|WP_Error
	 */
	public static function account_start( $user_id ) {
		$user = get_userdata( $user_id );

		if ( ! $user || ! self::applies_to( $user ) ) {
			return new WP_Error( 'aslansheed_2fa_unavailable', __( 'Two-factor authentication (2FA) is not available for your account.', 'aslansheed-login' ), array( 'status' => 403 ) );
		}

		if ( self::is_enabled( $user_id ) ) {
			return new WP_Error( 'aslansheed_2fa_already', __( 'Two-factor authentication (2FA) is already on.', 'aslansheed-login' ) );
		}

		$secret = Totp::generate_secret();

		set_transient( self::SETUP_PREFIX . $user_id, Crypto::encrypt( $secret ), self::PENDING_TTL );

		return array(
			'secret' => $secret,
			'uri'    => self::uri( $user, $secret ),
		);
	}

	/**
	 * Finishes a setup started from the profile.
	 *
	 * @param int    $user_id User id.
	 * @param string $code    Code from the new authenticator entry.
	 * @return array|WP_Error
	 */
	public static function account_confirm( $user_id, $code ) {
		$user   = get_userdata( $user_id );
		$stored = get_transient( self::SETUP_PREFIX . $user_id );

		if ( ! $user || ! is_string( $stored ) || '' === $stored ) {
			return new WP_Error( 'aslansheed_2fa_expired', __( 'The setup expired. Please start again.', 'aslansheed-login' ) );
		}

		$codes = self::activate( $user, (string) Crypto::decrypt( $stored ), $code );

		if ( is_wp_error( $codes ) ) {
			return $codes;
		}

		delete_transient( self::SETUP_PREFIX . $user_id );

		return array( 'codes' => $codes );
	}

	/**
	 * Switches 2FA off for a signed-in user after a valid code.
	 *
	 * @param int    $user_id User id.
	 * @param string $code    Authenticator code or recovery code.
	 * @return array|WP_Error
	 */
	public static function account_disable( $user_id, $code ) {
		$user = get_userdata( $user_id );

		if ( ! $user || ! self::is_enabled( $user_id ) ) {
			return new WP_Error( 'aslansheed_2fa_off', __( 'Two-factor authentication (2FA) is not on.', 'aslansheed-login' ) );
		}

		if ( ! self::can_disable( $user ) ) {
			return new WP_Error( 'aslansheed_2fa_required', __( 'Two-factor authentication (2FA) is required on this site.', 'aslansheed-login' ), array( 'status' => 403 ) );
		}

		$checked = self::check_code( $user, $code );

		if ( is_wp_error( $checked ) ) {
			return $checked;
		}

		self::reset( $user_id, $user_id );

		return array();
	}

	/**
	 * Creates a new set of recovery codes after a valid code.
	 *
	 * @param int    $user_id User id.
	 * @param string $code    Authenticator code.
	 * @return array|WP_Error
	 */
	public static function account_regenerate( $user_id, $code ) {
		$user = get_userdata( $user_id );

		if ( ! $user || ! self::is_enabled( $user_id ) ) {
			return new WP_Error( 'aslansheed_2fa_off', __( 'Two-factor authentication (2FA) is not on.', 'aslansheed-login' ) );
		}

		$checked = self::check_code( $user, $code );

		if ( is_wp_error( $checked ) ) {
			return $checked;
		}

		return array( 'codes' => self::new_recovery_codes( $user_id ) );
	}

	/**
	 * Removes the authenticator of a user. Used by the user and by administrators.
	 *
	 * @param int $user_id  User whose settings are cleared.
	 * @param int $actor_id User who did it.
	 * @return void
	 */
	public static function reset( $user_id, $actor_id ) {
		foreach ( array( self::META_SECRET, self::META_ENABLED, self::META_LAST_STEP, self::META_RECOVERY, self::META_SKIPPED ) as $key ) {
			delete_user_meta( $user_id, $key );
		}

		delete_transient( self::SETUP_PREFIX . $user_id );

		Audit_Log::log(
			'twofa_disabled',
			array(
				'user_id' => (int) $user_id,
				'context' => array( 'by' => (int) $actor_id ),
			)
		);
	}

	/**
	 * Counts the recovery codes a user still has.
	 *
	 * @param int $user_id User id.
	 * @return int
	 */
	public static function recovery_left( $user_id ) {
		$stored = get_user_meta( $user_id, self::META_RECOVERY, true );

		return is_array( $stored ) ? count( $stored ) : 0;
	}

	/**
	 * Checks the code typed on the classic login form.
	 *
	 * @param \WP_User $user User who gave the right password.
	 * @param string   $code Authenticator code or recovery code.
	 * @return true|WP_Error
	 */
	public static function check_for_login( $user, $code ) {
		return self::check_code( $user, $code );
	}

	/**
	 * Checks an authenticator code or a recovery code for a user.
	 *
	 * @param \WP_User $user User.
	 * @param string   $code Code typed by the user.
	 * @return true|WP_Error
	 */
	private static function check_code( $user, $code ) {
		$identifier = 'twofa:' . (int) $user->ID;
		$ip         = Request::ip();
		$allowed    = Rate_Limiter::check( Rate_Limiter::ACTION_VERIFY, $identifier, $ip );

		if ( is_wp_error( $allowed ) ) {
			return $allowed;
		}

		$code   = trim( (string) $code );
		$secret = (string) Crypto::decrypt( (string) get_user_meta( $user->ID, self::META_SECRET, true ) );
		$step   = 0;
		$used   = false;

		if ( preg_match( '/^[\d\s]{6,7}$/', Phone_Normalizer::ascii_digits( $code ) ) ) {
			$step = Totp::verify( $secret, $code, (int) get_user_meta( $user->ID, self::META_LAST_STEP, true ) );
		} else {
			$used = self::use_recovery_code( (int) $user->ID, $code );
		}

		if ( 0 === $step && ! $used ) {
			Rate_Limiter::record( Rate_Limiter::ACTION_VERIFY, $identifier, $ip, false );
			Audit_Log::log( 'twofa_failed', array( 'user_id' => (int) $user->ID ) );

			return new WP_Error( 'aslansheed_2fa_invalid', __( 'That code is not correct. Check the app and try again.', 'aslansheed-login' ) );
		}

		if ( $step > 0 ) {
			update_user_meta( $user->ID, self::META_LAST_STEP, $step );
		}

		Rate_Limiter::record( Rate_Limiter::ACTION_VERIFY, $identifier, $ip, true );
		Rate_Limiter::reset_failures( $identifier );

		if ( $used ) {
			Audit_Log::log( 'twofa_recovery_used', array( 'user_id' => (int) $user->ID ) );
		}

		return true;
	}

	/**
	 * Saves a secret after its first code was typed correctly and returns new recovery codes.
	 *
	 * @param \WP_User $user   User.
	 * @param string   $secret Base32 secret.
	 * @param string   $code   Code typed by the user.
	 * @return string[]|WP_Error
	 */
	private static function activate( $user, $secret, $code ) {
		$identifier = 'twofa:' . (int) $user->ID;
		$ip         = Request::ip();
		$allowed    = Rate_Limiter::check( Rate_Limiter::ACTION_VERIFY, $identifier, $ip );

		if ( is_wp_error( $allowed ) ) {
			return $allowed;
		}

		$step = '' === $secret ? 0 : Totp::verify( $secret, $code );

		if ( 0 === $step ) {
			Rate_Limiter::record( Rate_Limiter::ACTION_VERIFY, $identifier, $ip, false );

			return new WP_Error( 'aslansheed_2fa_invalid', __( 'That code is not correct. Check the app and try again.', 'aslansheed-login' ) );
		}

		update_user_meta( $user->ID, self::META_SECRET, Crypto::encrypt( $secret ) );
		update_user_meta( $user->ID, self::META_LAST_STEP, $step );
		update_user_meta( $user->ID, self::META_ENABLED, (string) time() );
		delete_user_meta( $user->ID, self::META_SKIPPED );

		Rate_Limiter::record( Rate_Limiter::ACTION_VERIFY, $identifier, $ip, true );
		Audit_Log::log( 'twofa_enabled', array( 'user_id' => (int) $user->ID ) );

		return self::new_recovery_codes( (int) $user->ID );
	}

	/**
	 * Creates recovery codes, stores only their hashes and returns the readable codes.
	 *
	 * @param int $user_id User id.
	 * @return string[]
	 */
	private static function new_recovery_codes( $user_id ) {
		$codes  = array();
		$hashes = array();

		for ( $i = 0; $i < self::RECOVERY_COUNT; $i++ ) {
			$raw      = substr( Totp::base32_encode( random_bytes( 7 ) ), 0, 10 );
			$codes[]  = substr( $raw, 0, 5 ) . '-' . substr( $raw, 5 );
			$hashes[] = Crypto::hmac( 'recovery|' . $raw );
		}

		update_user_meta( $user_id, self::META_RECOVERY, $hashes );

		return $codes;
	}

	/**
	 * Uses a recovery code once.
	 *
	 * @param int    $user_id User id.
	 * @param string $input   Code typed by the user.
	 * @return bool
	 */
	private static function use_recovery_code( $user_id, $input ) {
		$raw    = strtoupper( preg_replace( '/[^A-Za-z0-9]/', '', $input ) );
		$stored = get_user_meta( $user_id, self::META_RECOVERY, true );

		if ( 10 !== strlen( $raw ) || ! is_array( $stored ) ) {
			return false;
		}

		$hash = Crypto::hmac( 'recovery|' . $raw );

		foreach ( $stored as $index => $candidate ) {
			if ( hash_equals( (string) $candidate, $hash ) ) {
				unset( $stored[ $index ] );
				update_user_meta( $user_id, self::META_RECOVERY, array_values( $stored ) );

				return true;
			}
		}

		return false;
	}

	/**
	 * Builds the authenticator address for a user.
	 *
	 * @param \WP_User $user   User.
	 * @param string   $secret Base32 secret.
	 * @return string
	 */
	private static function uri( $user, $secret ) {
		$issuer = trim( (string) Options::get( 'twofa.issuer' ) );

		if ( '' === $issuer ) {
			$issuer = wp_specialchars_decode( get_bloginfo( 'name' ), ENT_QUOTES );
		}

		return Totp::uri( '' !== (string) $user->user_email ? (string) $user->user_email : (string) $user->user_login, $secret, '' !== $issuer ? $issuer : 'WordPress' );
	}

	/**
	 * Loads a pending sign-in.
	 *
	 * @param string $token Token.
	 * @param string $kind  "challenge" or "setup".
	 * @return array|WP_Error
	 */
	private static function pending( $token, $kind ) {
		$stored = get_transient( self::PENDING_PREFIX . hash( 'sha256', (string) $token ) );

		if ( ! is_array( $stored ) || empty( $stored['user_id'] ) ) {
			return self::expired();
		}

		$is_setup = 'challenge' !== $stored['state'];

		if ( ( 'setup' === $kind ) !== $is_setup ) {
			return self::expired();
		}

		return $stored;
	}

	/**
	 * Removes a pending sign-in.
	 *
	 * @param string $token Token.
	 * @return void
	 */
	private static function drop( $token ) {
		delete_transient( self::PENDING_PREFIX . hash( 'sha256', (string) $token ) );
	}

	/**
	 * The error for a missing or old pending sign-in.
	 *
	 * @return WP_Error
	 */
	private static function expired() {
		return new WP_Error( 'aslansheed_2fa_expired', __( 'Your sign-in expired. Please start again.', 'aslansheed-login' ) );
	}
}