Adding Simple Math Captcha to WordPress Login, Register or Forms (UPDATED)

// ══════════════════════════════════════════════════════════════════════════════
//  MATH CAPTCHA — SETTINGS
// ══════════════════════════════════════════════════════════════════════════════

define( 'SNN_CAPTCHA_MIN', 1 );
define( 'SNN_CAPTCHA_MAX', 9 );
define( 'SNN_CAPTCHA_ERROR', '<strong>ERROR</strong>: Incorrect or empty math captcha.' );

// Enable built-in WordPress forms
define( 'SNN_CAPTCHA_LOGIN',    true );
define( 'SNN_CAPTCHA_REGISTER', true );
define( 'SNN_CAPTCHA_LOST_PW',  true );
define( 'SNN_CAPTCHA_COMMENTS', true );

// Extra WordPress action hooks (third-party plugins)
// e.g. 'woocommerce_login_form', 'um_after_register_fields'
$snn_captcha_hooks = [
    // 'woocommerce_login_form',
    // 'woocommerce_register_form',
];

// Frontend form selectors — captcha injected via JS before the submit button
// Use any CSS selector: '#my-form', '.contact-form', '[data-form="signup"]'
$snn_captcha_selectors = [
    // '#my-custom-form',
    // '.wpcf7-form',
];

// ══════════════════════════════════════════════════════════════════════════════

function snn_add_math_captcha() {
    $n1  = rand( SNN_CAPTCHA_MIN, SNN_CAPTCHA_MAX );
    $n2  = rand( SNN_CAPTCHA_MIN, SNN_CAPTCHA_MAX );
    $uid = uniqid( 'captcha_' );
    ?>
    <p id="<?php echo esc_attr( $uid ); ?>" style="display:none">
        <label for="<?php echo esc_attr( $uid ); ?>_input"></label>
        <input type="text" name="math_captcha" id="<?php echo esc_attr( $uid ); ?>_input" class="input" value="" size="20" autocomplete="off" required>
        <input type="hidden" name="captcha_solution" value="">
        <input type="hidden" name="js_enabled" value="no">
    </p>
    <script>
    document.addEventListener('DOMContentLoaded', function () {
        var wrap  = document.getElementById('<?php echo esc_js( $uid ); ?>');
        var input = document.getElementById('<?php echo esc_js( $uid ); ?>_input');
        var sol   = wrap.querySelector('[name="captcha_solution"]');
        var jsf   = wrap.querySelector('[name="js_enabled"]');
        var form  = wrap.closest('form');
        var btn   = form && form.querySelector('input[type="submit"],button[type="submit"]');
        var n1 = <?php echo (int) $n1; ?>, n2 = <?php echo (int) $n2; ?>, ans = n1 + n2;

        wrap.style.display = 'block';
        jsf.value = 'yes';
        if ( btn ) btn.disabled = true;

        var canvas = document.createElement('canvas');
        canvas.width = 150; canvas.height = 24;
        var ctx = canvas.getContext('2d');
        ctx.font = '20px Arial'; ctx.fillStyle = '#333';
        ctx.fillText( n1 + ' + ' + n2 + ' = ?', 2, 20 );
        wrap.querySelector('label').appendChild( canvas );

        sol.value = ans;
        input.addEventListener('input', function () {
            if ( btn ) btn.disabled = ( parseInt( this.value.trim(), 10 ) !== ans );
        });
    });
    </script>
    <?php
}

function snn_check_captcha() {
    if ( empty( $_POST['js_enabled'] ) || $_POST['js_enabled'] !== 'yes' ) return false;
    if ( ! isset( $_POST['math_captcha'], $_POST['captcha_solution'] ) )   return false;
    return (int) trim( $_POST['math_captcha'] ) === (int) trim( $_POST['captcha_solution'] );
}

// Built-in hooks
if ( SNN_CAPTCHA_LOGIN ) {
    add_action( 'login_form', 'snn_add_math_captcha' );
    add_filter( 'authenticate', function ( $r, $u, $p ) {
        if ( $_SERVER['REQUEST_METHOD'] === 'POST' && in_array( $_REQUEST['action'] ?? '', [ 'login', '' ] ) && ! snn_check_captcha() )
            return new WP_Error( 'captcha_error', SNN_CAPTCHA_ERROR );
        return $r;
    }, 30, 3 );
}

if ( SNN_CAPTCHA_REGISTER ) {
    add_action( 'register_form', 'snn_add_math_captcha' );
    add_filter( 'registration_errors', function ( $e, $l, $m ) {
        if ( ! snn_check_captcha() ) $e->add( 'captcha_error', SNN_CAPTCHA_ERROR );
        return $e;
    }, 10, 3 );
}

if ( SNN_CAPTCHA_LOST_PW ) {
    add_action( 'lostpassword_form', 'snn_add_math_captcha' );
    add_filter( 'lostpassword_post_errors', function ( $e, $l ) {
        if ( ! snn_check_captcha() ) $e->add( 'captcha_error', SNN_CAPTCHA_ERROR );
        return $e;
    }, 10, 2 );
}

if ( SNN_CAPTCHA_COMMENTS ) {
    add_filter( 'comment_form_field_comment', function ( $f ) {
        if ( is_user_logged_in() ) return $f;
        ob_start(); snn_add_math_captcha(); return $f . ob_get_clean();
    } );
    add_filter( 'preprocess_comment', function ( $d ) {
        if ( ! is_user_logged_in() && ! snn_check_captcha() ) wp_die( SNN_CAPTCHA_ERROR );
        return $d;
    } );
}

// Extra plugin hooks
foreach ( $snn_captcha_hooks as $hook ) {
    add_action( $hook, 'snn_add_math_captcha' );
}

// Frontend form selectors — JS injection
if ( ! empty( $snn_captcha_selectors ) ) {
    add_action( 'wp_footer', function () use ( $snn_captcha_selectors ) {
        $selectors_json = json_encode( $snn_captcha_selectors );
        ?>
        <script>
        document.addEventListener('DOMContentLoaded', function () {
            <?php echo $selectors_json; ?>.forEach( function ( sel ) {
                var form = document.querySelector( sel );
                if ( ! form ) return;
                var btn  = form.querySelector('input[type="submit"],button[type="submit"]');
                var n1   = Math.ceil( Math.random() * <?php echo SNN_CAPTCHA_MAX; ?> );
                var n2   = Math.ceil( Math.random() * <?php echo SNN_CAPTCHA_MAX; ?> );
                var ans  = n1 + n2;

                var canvas = document.createElement('canvas');
                canvas.width = 150; canvas.height = 24;
                var ctx = canvas.getContext('2d');
                ctx.font = '20px Arial'; ctx.fillStyle = '#333';
                ctx.fillText( n1 + ' + ' + n2 + ' = ?', 2, 20 );

                var input  = document.createElement('input');
                input.type = 'text'; input.name = 'math_captcha';
                input.size = 20; input.required = true; input.autocomplete = 'off';

                var sol  = document.createElement('input');
                sol.type = 'hidden'; sol.name = 'captcha_solution'; sol.value = ans;

                var jsf  = document.createElement('input');
                jsf.type = 'hidden'; jsf.name = 'js_enabled'; jsf.value = 'yes';

                var wrap = document.createElement('p');
                wrap.appendChild( canvas );
                wrap.appendChild( input );
                wrap.appendChild( sol );
                wrap.appendChild( jsf );

                if ( btn ) btn.before( wrap );
                else form.appendChild( wrap );

                if ( btn ) btn.disabled = true;
                input.addEventListener('input', function () {
                    if ( btn ) btn.disabled = ( parseInt( this.value.trim(), 10 ) !== ans );
                });
            });
        });
        </script>
        <?php
    } );
}