<?php
/**
 * Two-factor authentication: TOTP + single-use recovery codes.
 *
 * Login flow:
 *   1. User submits username/password → wp_authenticate_username_password
 *      validates → our `authenticate` filter (prio 30) checks if the user
 *      has 2FA enabled / required.
 *   2. If yes: stash a short-lived pending-token in a transient, set a
 *      cookie, redirect the browser to wp-login.php?action=sf_2fa, exit.
 *   3. The interstitial form posts the TOTP/recovery code back to the same
 *      action; on success we call wp_set_auth_cookie() ourselves and
 *      redirect to the original target.
 *
 * User-meta keys:
 *   shieldforge_2fa_secret           base32 TOTP secret (after confirm)
 *   shieldforge_2fa_pending_secret   base32 TOTP secret (before confirm)
 *   shieldforge_2fa_enabled          bool flag (only true after verify)
 *   shieldforge_2fa_recovery         array of sha256(normalized-code) hashes
 *   shieldforge_2fa_last_used        unix timestamp of last successful 2FA
 *
 * @package Brainwerk Security Suite
 */

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

class Shieldforge_TwoFactor {

    /** Login-action handled by wp-login.php?action=sf_2fa */
    const ACTION = 'sf_2fa';

    /** Cookie carrying the pending-token to the interstitial form. */
    const COOKIE = 'sf_2fa_pending';

    /** Transient prefix; full key = PREFIX . token. */
    const TRANSIENT_PREFIX = 'shieldforge_2fa_pending_';

    /** Pending-token TTL in seconds (5 minutes). */
    const PENDING_TTL = 300;

    /** Max failed second-factor attempts before a temporary lockout. */
    const MAX_ATTEMPTS = 5;

    /** Lockout window after exceeding MAX_ATTEMPTS, in seconds (15 minutes). */
    const LOCKOUT_TTL = 900;

    /** Recovery-code alphabet — no I/O/0/1 to avoid visual ambiguity. */
    const RECOVERY_ALPHABET = 'ABCDEFGHJKLMNPQRSTUVWXYZ23456789';

    /* ==================================================================
     *  Registration
     * ================================================================== */

    public static function register( Shieldforge_Loader $loader ) {
        // Step 1: catch successful credentials, demand 2FA.
        $loader->add_filter( 'authenticate', __CLASS__, 'maybe_demand_2fa', 30, 3 );

        // Step 2: handle the interstitial form.
        $loader->add_action( 'login_form_' . self::ACTION, __CLASS__, 'handle_interstitial' );
    }

    /* ==================================================================
     *  Public helpers
     * ================================================================== */

    public static function is_module_enabled() {
        return (bool) Shieldforge_Settings::get( 'twofactor_enabled', true );
    }

    public static function is_enabled_for_user( $user_id ) {
        return (bool) get_user_meta( (int) $user_id, 'shieldforge_2fa_enabled', true );
    }

    /**
     * True if the site policy says THIS user must have 2FA configured (their
     * role is in the force-list). Used for "you must set up 2FA" nag screens.
     */
    public static function is_required_for_user( $user_id_or_obj ) {
        $user = is_object( $user_id_or_obj ) ? $user_id_or_obj : get_user_by( 'id', (int) $user_id_or_obj );
        if ( ! $user || empty( $user->roles ) ) {
            return false;
        }
        $force_roles = (array) Shieldforge_Settings::get( 'twofactor_force_roles', array() );
        if ( empty( $force_roles ) ) {
            return false;
        }
        return (bool) array_intersect( (array) $user->roles, $force_roles );
    }

    /**
     * True if THIS request needs to go through the 2FA second step.
     */
    public static function needs_2fa( $user ) {
        if ( ! self::is_module_enabled() ) {
            return false;
        }
        if ( ! is_a( $user, 'WP_User' ) ) {
            return false;
        }
        // App passwords (REST/XML-RPC) bypass 2FA — they're a separate
        // credential the user opted into. Standard pattern in WP 2FA plugins.
        if ( did_action( 'application_password_did_authenticate' ) ) {
            return false;
        }
        return self::is_enabled_for_user( $user->ID );
    }

    /* ==================================================================
     *  Setup workflow (used from the user-profile UI)
     * ================================================================== */

    /**
     * Read or initialize the pending secret for setup. Call this from the
     * profile UI so the QR code is stable across page reloads (otherwise the
     * QR would change on every refresh and the user would see different
     * codes than they just entered).
     */
    public static function get_or_create_pending_secret( $user_id ) {
        $existing = get_user_meta( (int) $user_id, 'shieldforge_2fa_pending_secret', true );
        if ( $existing ) {
            return $existing;
        }
        $secret = Shieldforge_TOTP::generate_secret();
        update_user_meta( (int) $user_id, 'shieldforge_2fa_pending_secret', $secret );
        return $secret;
    }

    /**
     * Verify the TOTP code against the pending secret. On success: enable
     * 2FA, promote pending → real, generate recovery codes (returned as
     * plain strings so the UI can show them ONCE).
     *
     * @return array|false plain recovery codes on success, false otherwise
     */
    public static function confirm_setup( $user_id, $code ) {
        $user_id = (int) $user_id;
        $pending = (string) get_user_meta( $user_id, 'shieldforge_2fa_pending_secret', true );
        if ( '' === $pending ) {
            return false;
        }
        if ( ! Shieldforge_TOTP::verify( $pending, $code ) ) {
            return false;
        }
        update_user_meta( $user_id, 'shieldforge_2fa_secret', $pending );
        update_user_meta( $user_id, 'shieldforge_2fa_enabled', 1 );
        delete_user_meta( $user_id, 'shieldforge_2fa_pending_secret' );

        $plain = self::generate_recovery_codes();
        $hashes = array_map( array( __CLASS__, 'hash_recovery_code' ), $plain );
        update_user_meta( $user_id, 'shieldforge_2fa_recovery', $hashes );
        update_user_meta( $user_id, 'shieldforge_2fa_last_used', time() );

        Shieldforge_Logger::log( '2fa_enabled', array(
            'severity' => Shieldforge_Logger::SEV_INFO,
            'user_id'  => $user_id,
            'context'  => array(),
        ) );

        return $plain;
    }

    /**
     * Disable 2FA for a user — wipes every related meta-key.
     * Used both from the user's own profile (with their consent) and from
     * the admin "lockout-rescue" page (when an admin loses their 2FA device).
     */
    public static function disable_for_user( $user_id, $reason = 'user_disabled' ) {
        $user_id = (int) $user_id;
        delete_user_meta( $user_id, 'shieldforge_2fa_secret' );
        delete_user_meta( $user_id, 'shieldforge_2fa_pending_secret' );
        delete_user_meta( $user_id, 'shieldforge_2fa_enabled' );
        delete_user_meta( $user_id, 'shieldforge_2fa_recovery' );
        delete_user_meta( $user_id, 'shieldforge_2fa_last_used' );
        Shieldforge_Logger::log( '2fa_disabled', array(
            'severity' => Shieldforge_Logger::SEV_WARN,
            'user_id'  => $user_id,
            'context'  => array( 'reason' => $reason ),
        ) );
    }

    /**
     * Verify a code against a user's enabled secret OR consume a recovery
     * code. Returns a string indicating which path matched (for logging),
     * or false on failure.
     *
     * @return string|false 'totp' | 'recovery' | false
     */
    public static function verify_code( $user_id, $code ) {
        $user_id = (int) $user_id;
        $code    = (string) $code;

        $secret = (string) get_user_meta( $user_id, 'shieldforge_2fa_secret', true );
        if ( '' !== $secret ) {
            $totp_input = preg_replace( '/\D+/', '', $code );
            if ( strlen( $totp_input ) === Shieldforge_TOTP::DIGITS ) {
                $step = Shieldforge_TOTP::verify_step( $secret, $totp_input );
                if ( false !== $step ) {
                    // Replay protection (RFC 6238 §5.2): reject a code from a
                    // timestep we've already accepted. Stops a captured/shoulder-
                    // surfed code being reused inside its ~90s validity window.
                    $last_step = (int) get_user_meta( $user_id, 'shieldforge_2fa_last_step', true );
                    if ( $step <= $last_step ) {
                        return false;
                    }
                    update_user_meta( $user_id, 'shieldforge_2fa_last_step', $step );
                    update_user_meta( $user_id, 'shieldforge_2fa_last_used', time() );
                    return 'totp';
                }
            }
        }

        $recovery = (array) get_user_meta( $user_id, 'shieldforge_2fa_recovery', true );
        if ( ! empty( $recovery ) ) {
            $hash = self::hash_recovery_code( $code );
            foreach ( $recovery as $idx => $stored ) {
                if ( hash_equals( (string) $stored, $hash ) ) {
                    unset( $recovery[ $idx ] );
                    update_user_meta( $user_id, 'shieldforge_2fa_recovery', array_values( $recovery ) );
                    update_user_meta( $user_id, 'shieldforge_2fa_last_used', time() );
                    Shieldforge_Logger::log( '2fa_recovery_used', array(
                        'severity' => Shieldforge_Logger::SEV_WARN,
                        'user_id'  => $user_id,
                        'context'  => array( 'remaining' => count( $recovery ) ),
                    ) );
                    return 'recovery';
                }
            }
        }
        return false;
    }

    /**
     * 8 codes by default, formatted XXXX-XXXX (10 chars including hyphen).
     * Stored as sha256(normalized) hashes; original strings are never
     * persisted plain — they're shown to the user exactly once on
     * generation.
     */
    public static function generate_recovery_codes( $count = 8 ) {
        $codes    = array();
        $alphabet = self::RECOVERY_ALPHABET;
        $alpha_len = strlen( $alphabet );
        for ( $i = 0; $i < $count; $i++ ) {
            $code = '';
            for ( $j = 0; $j < 8; $j++ ) {
                $code .= $alphabet[ random_int( 0, $alpha_len - 1 ) ];
            }
            $codes[] = substr( $code, 0, 4 ) . '-' . substr( $code, 4 );
        }
        return $codes;
    }

    public static function recovery_codes_remaining( $user_id ) {
        $raw = get_user_meta( (int) $user_id, 'shieldforge_2fa_recovery', true );
        if ( ! is_array( $raw ) ) {
            return 0; // not set yet (get_user_meta returns '' for unset)
        }
        return count( $raw );
    }

    /**
     * Burn the existing recovery codes and issue a fresh batch.
     */
    public static function regenerate_recovery_codes( $user_id ) {
        $plain = self::generate_recovery_codes();
        $hashes = array_map( array( __CLASS__, 'hash_recovery_code' ), $plain );
        update_user_meta( (int) $user_id, 'shieldforge_2fa_recovery', $hashes );
        Shieldforge_Logger::log( '2fa_recovery_regen', array(
            'severity' => Shieldforge_Logger::SEV_INFO,
            'user_id'  => (int) $user_id,
            'context'  => array(),
        ) );
        return $plain;
    }

    private static function hash_recovery_code( $code ) {
        $normalized = strtoupper( preg_replace( '/[^A-Za-z0-9]/', '', (string) $code ) );
        return hash( 'sha256', $normalized );
    }

    /* ==================================================================
     *  Login flow — Step 1: detect 2FA need, redirect to interstitial
     * ================================================================== */

    public static function maybe_demand_2fa( $user, $username, $password ) {
        if ( ! $user || is_wp_error( $user ) ) {
            return $user;
        }
        if ( ! self::needs_2fa( $user ) ) {
            return $user;
        }

        // Hand off to the shared interstitial path. This exits (never returns)
        // so the rest of the authenticate chain never runs and wp-login.php
        // never sets the session cookie for a half-authenticated user.
        // phpcs:disable WordPress.Security.NonceVerification.Missing, WordPress.Security.NonceVerification.Recommended -- runs inside the core `authenticate` filter (wp-login.php), same pre-auth context as core itself, which carries no nonce here.
        self::stash_pending_and_redirect(
            $user,
            ! empty( $_POST['rememberme'] ),
            isset( $_REQUEST['redirect_to'] ) ? esc_url_raw( wp_unslash( $_REQUEST['redirect_to'] ) ) : admin_url()
        );
        // phpcs:enable WordPress.Security.NonceVerification.Missing, WordPress.Security.NonceVerification.Recommended
    }

    /**
     * Stash a short-lived pending-2FA token (transient + cookie) and redirect
     * the browser to the interstitial form, then exit. Shared entry point for
     * EVERY login path that must enforce the second factor — the password
     * filter (maybe_demand_2fa) as well as alternative flows such as the
     * passwordless Magic Link — so none of them can silently bypass 2FA.
     *
     * Callers MUST first confirm the user is fully credential-authenticated
     * and that needs_2fa() is true. This method never returns.
     */
    public static function stash_pending_and_redirect( $user, $remember = false, $redirect_to = '' ) {
        $token   = self::generate_token();
        $payload = array(
            'user_id'     => (int) $user->ID,
            'remember'    => (bool) $remember,
            'redirect_to' => $redirect_to ? esc_url_raw( $redirect_to ) : admin_url(),
            'created_at'  => time(),
        );
        set_transient( self::TRANSIENT_PREFIX . $token, $payload, self::PENDING_TTL );
        self::set_token_cookie( $token );

        // Audit log: credentials accepted, now waiting on the second factor.
        Shieldforge_Logger::log( '2fa_required', array(
            'severity' => Shieldforge_Logger::SEV_INFO,
            'user_id'  => (int) $user->ID,
            'username' => Shieldforge_Utils::store_username( $user->user_login ),
            'context'  => array(),
        ) );

        wp_safe_redirect( add_query_arg( 'action', self::ACTION, wp_login_url() ) );
        exit;
    }

    /* ==================================================================
     *  Login flow — Step 2: render + handle the interstitial form
     * ================================================================== */

    public static function handle_interstitial() {
        $token = self::read_token_cookie();
        if ( ! $token ) {
            wp_safe_redirect( wp_login_url() );
            exit;
        }
        $payload = get_transient( self::TRANSIENT_PREFIX . $token );
        if ( ! is_array( $payload ) || empty( $payload['user_id'] ) ) {
            self::clear_token_cookie();
            wp_safe_redirect( wp_login_url() );
            exit;
        }
        $user_id = (int) $payload['user_id'];

        $error = '';
        $request_method = isset( $_SERVER['REQUEST_METHOD'] ) ? sanitize_text_field( wp_unslash( $_SERVER['REQUEST_METHOD'] ) ) : '';
        if ( 'POST' === $request_method ) {
            check_admin_referer( 'shieldforge_2fa_verify' );

            // Brute-force throttle: cap failed code submissions per user. TOTP
            // is only 10^6 codes valid across a few timesteps, so an attacker
            // holding a valid pending cookie could otherwise grind guesses for
            // the whole 5-minute window. Keyed per-user so a fresh login
            // attempt can't reset the counter.
            $attempt_key = 'shieldforge_2fa_attempts_' . $user_id;
            $attempts    = (int) get_transient( $attempt_key );
            if ( $attempts >= self::MAX_ATTEMPTS ) {
                Shieldforge_Logger::log( '2fa_lockout', array(
                    'severity' => Shieldforge_Logger::SEV_WARN,
                    'user_id'  => $user_id,
                    'context'  => array( 'attempts' => $attempts ),
                ) );
                self::render_form( __( 'Too many invalid codes. Please wait a few minutes and try again.', 'brainwerk-security-suite' ) );
                exit;
            }

            $code   = isset( $_POST['shieldforge_2fa_code'] ) ? sanitize_text_field( wp_unslash( $_POST['shieldforge_2fa_code'] ) ) : '';
            $result = self::verify_code( $user_id, $code );
            if ( $result ) {
                // Success — finalize login.
                delete_transient( $attempt_key );
                delete_transient( self::TRANSIENT_PREFIX . $token );
                self::clear_token_cookie();
                wp_set_auth_cookie( $user_id, ! empty( $payload['remember'] ) );

                Shieldforge_Logger::log( '2fa_passed', array(
                    'severity' => Shieldforge_Logger::SEV_INFO,
                    'user_id'  => $user_id,
                    'context'  => array( 'method' => $result ),
                ) );

                $redirect = ! empty( $payload['redirect_to'] ) ? $payload['redirect_to'] : admin_url();
                wp_safe_redirect( $redirect );
                exit;
            }
            // Bad code — count the failure toward the lockout, log, stay on form.
            set_transient( $attempt_key, $attempts + 1, self::LOCKOUT_TTL );
            $error = __( 'Invalid 2FA code. Try again or use a recovery code.', 'brainwerk-security-suite' );
            Shieldforge_Logger::log( '2fa_failed', array(
                'severity' => Shieldforge_Logger::SEV_WARN,
                'user_id'  => $user_id,
                'context'  => array( 'attempt' => $attempts + 1 ),
            ) );
        }

        self::render_form( $error );
        exit;
    }

    /**
     * Render the interstitial form. Reuses login_header/login_footer so it
     * inherits the site's login screen styling (incl. custom logos via
     * any of the standard "branding" plugins).
     */
    private static function render_form( $error_message = '' ) {
        // login_header / login_footer are defined inside wp-login.php and
        // available when our login_form_<action> hook fires. If they're not,
        // the request reached us through some non-standard entry point —
        // bail rather than re-including wp-login.php (which would re-run
        // its top-level auth dispatch).
        if ( ! function_exists( 'login_header' ) ) {
            wp_die( esc_html__( 'Unable to render 2FA form (login subsystem unavailable).', 'brainwerk-security-suite' ) );
        }
        login_header( __( 'Two-factor authentication', 'brainwerk-security-suite' ), '', $error_message ? new WP_Error( 'sf_2fa', $error_message ) : null );
        ?>
        <form name="shieldforge_2fa_form" id="loginform" method="post">
            <p>
                <label for="shieldforge_2fa_code"><?php esc_html_e( 'Authentication code', 'brainwerk-security-suite' ); ?></label>
                <input type="text"
                       name="shieldforge_2fa_code"
                       id="shieldforge_2fa_code"
                       class="input"
                       autocomplete="one-time-code"
                       inputmode="numeric"
                       pattern="[0-9A-Za-z\- ]*"
                       autofocus
                       required>
            </p>
            <p class="description">
                <?php esc_html_e( 'Enter the 6-digit code from your authenticator app, or one of your recovery codes (e.g. ABCD-EFGH).', 'brainwerk-security-suite' ); ?>
            </p>
            <?php wp_nonce_field( 'shieldforge_2fa_verify' ); ?>
            <p class="submit">
                <input type="submit" class="button button-primary button-large" value="<?php esc_attr_e( 'Verify', 'brainwerk-security-suite' ); ?>">
            </p>
        </form>
        <p id="nav">
            <a href="<?php echo esc_url( wp_login_url() ); ?>"><?php esc_html_e( 'Back to login', 'brainwerk-security-suite' ); ?></a>
        </p>
        <?php
        login_footer( 'shieldforge_2fa_code' );
    }

    /* ==================================================================
     *  Token / cookie helpers
     * ================================================================== */

    private static function generate_token() {
        return wp_generate_password( 40, false, false );
    }

    private static function set_token_cookie( $token ) {
        if ( headers_sent() ) {
            return;
        }
        // PHP 7.3+ associative-array form is cleaner. WP requires PHP 7.4+
        // since v6.3 so this is safe.
        setcookie( self::COOKIE, $token, array(
            'expires'  => time() + self::PENDING_TTL,
            'path'     => '/',
            'secure'   => is_ssl(),
            'httponly' => true,
            'samesite' => 'Lax',
        ) );
    }

    private static function clear_token_cookie() {
        if ( headers_sent() ) {
            return;
        }
        setcookie( self::COOKIE, '', array(
            'expires'  => time() - 3600,
            'path'     => '/',
            'secure'   => is_ssl(),
            'httponly' => true,
            'samesite' => 'Lax',
        ) );
    }

    private static function read_token_cookie() {
        if ( empty( $_COOKIE[ self::COOKIE ] ) ) {
            return '';
        }
        return preg_replace( '/[^A-Za-z0-9]/', '', sanitize_text_field( wp_unslash( $_COOKIE[ self::COOKIE ] ) ) );
    }
}
