robotstxt-2fa/includes/class-geo-restrictions.php

<?php
/**
 * GeoIP-based country restrictions for the 2FA challenge.
 *
 * @package Robotstxt_2FA
 */

namespace Robotstxt\TwoFA;

use Robotstxt\TwoFA\User\Two_Factor_Config;

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

/**
 * Applies country-based allow/deny/force-challenge rules using MaxMind GeoIP.
 *
 * @since 1.5.0
 */
class Geo_Restrictions {
	/**
	 * Global two-factor configuration.
	 *
	 * @var Two_Factor_Config
	 */
	private Two_Factor_Config $config;

	/**
	 * Constructor.
	 *
	 * @param Two_Factor_Config $config Global two-factor configuration.
	 */
	public function __construct( Two_Factor_Config $config ) {
		$this->config = $config;
	}

	/**
	 * Register WordPress hooks.
	 *
	 * @since 1.5.0
	 *
	 * @return void
	 */
	public function register_hooks(): void {
		add_filter( 'robotstxt_2fa_skip_challenge', array( $this, 'maybe_skip_for_allowed_country' ), 6, 2 );
		add_filter( 'authenticate', array( $this, 'maybe_block_denied_country' ), 2, 3 );
		add_filter( 'robotstxt_2fa_force_challenge', array( $this, 'maybe_force_challenge_for_country' ), 10, 2 );
	}

	/**
	 * Skip the 2FA challenge for users logging in from an allowed country.
	 *
	 * @since 1.5.0
	 *
	 * @param bool     $skip Whether the challenge is already being skipped.
	 * @param \WP_User $user Authenticated user about to be challenged.
	 *
	 * @return bool
	 */
	public function maybe_skip_for_allowed_country( bool $skip, \WP_User $user ): bool {
		unset( $user );

		if ( $skip ) {
			return true;
		}

		$allow_list = $this->config->get_geoip_country_allow();

		if ( '' === $allow_list ) {
			return false;
		}

		$country = $this->get_country_for_ip();

		if ( '' === $country ) {
			return false;
		}

		return $this->country_in_list( $country, $allow_list );
	}

	/**
	 * Block logins from denied countries before credentials are checked.
	 *
	 * @since 1.5.0
	 *
	 * @param mixed       $user     Previously authenticated user or null.
	 * @param string|null $username Submitted username.
	 * @param string|null $password Submitted password.
	 *
	 * @return mixed
	 */
	public function maybe_block_denied_country( mixed $user, ?string $username, ?string $password ): mixed {
		unset( $username, $password );

		$deny_list = $this->config->get_geoip_country_deny();

		if ( '' === $deny_list ) {
			return $user;
		}

		$country = $this->get_country_for_ip();

		if ( '' === $country ) {
			return $user;
		}

		if ( $this->country_in_list( $country, $deny_list ) ) {
			return new \WP_Error(
				'robotstxt_2fa_country_blocked',
				__( 'Login is not allowed from your country.', 'robotstxt-2fa' )
			);
		}

		return $user;
	}

	/**
	 * Force a fresh 2FA challenge for users logging in from a high-risk country.
	 *
	 * @since 1.5.0
	 *
	 * @param bool     $force Whether the challenge is already being forced.
	 * @param \WP_User $user  User about to be challenged.
	 *
	 * @return bool
	 */
	public function maybe_force_challenge_for_country( bool $force, \WP_User $user ): bool {
		unset( $user );

		if ( $force ) {
			return true;
		}

		$always_challenge = $this->config->get_geoip_always_challenge();

		if ( '' === $always_challenge ) {
			return false;
		}

		$country = $this->get_country_for_ip();

		if ( '' === $country ) {
			return false;
		}

		return $this->country_in_list( $country, $always_challenge );
	}

	/**
	 * Resolve the ISO 3166-1 alpha-2 country code for the current request IP.
	 *
	 * Returns an empty string when GeoIP is unavailable or the lookup fails.
	 *
	 * @since 1.5.0
	 *
	 * @return string
	 */
	private function get_country_for_ip(): string {
		$db_path = $this->config->get_geoip_database_path();

		if ( '' === $db_path || ! file_exists( $db_path ) || ! is_file( $db_path ) || is_link( $db_path ) ) {
			return '';
		}

		if ( ! class_exists( '\MaxMind\Db\Reader' ) ) {
			return '';
		}

		$ip = '';

		if ( isset( $_SERVER['REMOTE_ADDR'] ) && is_string( $_SERVER['REMOTE_ADDR'] ) ) {
			$ip = sanitize_text_field( wp_unslash( $_SERVER['REMOTE_ADDR'] ) );
		}

		if ( '' === $ip ) {
			return '';
		}

		try {
			$reader = new \MaxMind\Db\Reader( $db_path );
			$record = $reader->get( $ip );
			$reader->close();

			if ( ! is_array( $record ) ) {
				return '';
			}

			$country = $record['country'] ?? null;

			if ( ! is_array( $country ) || ! isset( $country['iso_code'] ) || ! is_string( $country['iso_code'] ) ) {
				return '';
			}

			return strtoupper( $country['iso_code'] );
		} catch ( \Throwable $e ) {
			return '';
		}
	}

	/**
	 * Check whether a country code appears in a newline/comma-separated list.
	 *
	 * @since 1.5.0
	 *
	 * @param string $country  ISO 3166-1 alpha-2 country code (uppercase).
	 * @param string $codelist Newline- or comma-separated country codes.
	 *
	 * @return bool
	 */
	private function country_in_list( string $country, string $codelist ): bool {
		if ( '' === $country || '' === $codelist ) {
			return false;
		}

		$split_result = preg_split( '/[\s,]+/', strtoupper( $codelist ) );
		$codes        = array_filter( array_map( 'trim', is_array( $split_result ) ? $split_result : array() ) );

		return in_array( $country, $codes, true );
	}
}